Metal Compute Kernels

The Metal kernel is S-MoE's execution fist—the relentless component that transforms quantised expert blobs into brilliant floating-point output vectors, weaponizing Apple Silicon's GPU and Unified Memory Architecture against the Monolith.

---

Zero-Copy Metal Integration

The absolute constraint that defines our Metal integration: no data copy between CPU and GPU memory. We reject the inefficient bus transfers of traditional discrete GPUs.

The kernels live as MSL source in src/compute/kernels.metal and are JIT-compiled at engine startup via newLibraryWithSource: — no pre-built .metallib is shipped or required. One binary serves both Q2 and Q4 vaults; the dispatch adapts to the vault header's bits field.

On Apple Silicon, CPU and GPU natively share the same physical Unified Memory. A pointer to a posix_memalign-allocated buffer is flawlessly wrapped as a MTLBuffer using newBufferWithBytesNoCopy:

id<MTLBuffer> buf = [device newBufferWithBytesNoCopy:ptr
                                               length:size
                                              options:MTLResourceStorageModeShared
                                          deallocator:nil];

The GPU reads directly from the pristine address where the Streamer wrote the expert blob. The sovereign data path:

NVMe → DMA → UMA (ring slot) → MTLBuffer (same exact address) → GPU registers

No memcpy. No blit. The bytes move exactly once.

Alignment is Non-Negotiable

newBufferWithBytesNoCopy demands that the pointer be flawlessly aligned to the 16 KB hardware page boundary. Any misaligned pointer will silently fail or trigger undefined behavior. posix_memalign(16384) is the only acceptable allocation method for ring slot data.

---

Fused FFN Kernels — One Simdgroup Per Output Row

The decode hot path executes each routed expert as a two-pass fused dispatch inside one command buffer: smoe_gate_up computes both the gate and up projections in a single pass over the input and applies SiLU in register space; behind one barrier, smoe_down projects the hidden vector back to d_model:

Input: heavy_normed[d_model]  (pre-normalised activations)

    hidden[i] = SiLU(gate_proj·x)[i] × (up_proj·x)[i]   // one pass, shared x
    out[j]    = (down_proj · hidden)[j]                  // shape [d_model]

The kernels are shaped for the memory system, not the arithmetic — a quantised GEMV is pure bandwidth:

• One simdgroup owns one output row. Its 32 lanes read consecutive uint4 chunks of the packed row — 512 coalesced bytes per simdgroup per step. (A thread-per-row layout would leave the lanes reading addresses ~2 KB apart, one byte at a time — the memory system punishes that at roughly 7× the cost.)
• A 256-thread threadgroup covers 8 rows, so the grid is ceil(rows/8) — 192 threadgroups per gate_up expert, enough to actually occupy an M-class GPU (a 1536-row projection dispatched row-per-threadgroup-of-rows yields single-digit threadgroups and starves it).
• The input vector is staged once per threadgroup in threadgroup memory (cols × 4 B; 16 KB at d_model 4096) and shared by all 8 rows.
• Dequantisation is folded into an affine epilogue. With w = (code − Z)·(1/Z)·s (Z = 7.5 at Q4, 1.5 at Q2), the row dot product becomes Σw·x = s·((Σcode·x)/Z − Σx) — the inner loop is a pure fma(code4, x4, acc) per 4 codes, and the scale is touched once per 16-byte chunk, not once per 2 codes.
• Gate and up are fused (qrow_dot2): both matrices share x, so the threadgroup loads and the Σx term are paid once for the two dots — worth ~17% over two single passes.

Measured with a standalone microbench replaying the production dispatch against real vault blobs (Apple M4 Pro, Q4, group 64): 158 GB/s effective bandwidth — the same memory system's practical ceiling for this kernel family — 0.51 ms per 8-expert group, 0.08 ms for a single late-miss expert, ~48 ms/token of GPU-side FFN across 94 layers. In-engine, the decode gpu-wait bucket sits at ~46 ms/token.

No intermediate weight buffer is ever created. The float32 weight never even exists individually — whole 16-byte chunks of codes flow through the fma stream and are gone.

Layout assumptions the kernels rely on

cols % 32 == 0 (Q4) / % 64 == 0 (Q2) — no chunk tail handling; group_size a multiple of 32/64 — a 16-byte chunk never straddles a scale group; packed rows 16-byte aligned. Every .smoe vault guarantees all three (dims are multiples of 256, group size 64). The batch prefill kernels below keep the older thread-per-row contract for now.

---

Coalesced Dense Matvec — the Same Contract for bf16

The dense bf16 projections ride the FFN kernels' dispatch contract through scout_matvec_bf16_sg: one simdgroup per output row, 32 lanes reading consecutive uint4 chunks (8 bf16 each — two register shifts unpack a pair, no dequant epilogue needed), the input vector staged whole in threadgroup memory. Measured on the standalone microbench (M4 Pro, synthetic bf16 weights, cross-verified against the thread-per-row kernel):

Shape	Role	GB/s effective
8192 × 4096	Q projection	238
4096 × 8192	O projection	207
512 × 4096	K/V projection	285 — a thread-per-row grid strands this shape at 2 threadgroups
151936 × 4096	LM head	234

The bridge chooses per dispatch: the coalesced kernel whenever its contract holds — weights 16-byte aligned, cols % 8 == 0, staged input within the 32 KB threadgroup-memory limit (cols ≤ 8192) — and the thread-per-row twin for everything else (e.g. the wider dense/shared FFN matrices of DeepSeek-family backbones). Capability selection, never failure.

Not bit-exact — by construction

simd_sum reassociates the row reduction, so outputs are not bitwise identical to the previous kernels. Verified instead at max |diff| ~1e-6 against a CPU reference dequant FFN, and token-identical on the canonical greedy prompt and across a 3-turn serve session.

---

Token-Batch Fused FFN

Layer-major prefill demanded a new weapon: apply one expert to many token activations in a single command buffer. The smoe_gate_up_batch / smoe_down_batch kernel pair takes a row-major input matrix [batch × d_model] (up to 64 tokens) and dispatches a 2D threadgroup grid — (⌈rows/256⌉, batch) — so all routed tokens of an expert are computed in one two-pass submission instead of batch separate dispatch/wait round-trips:

per claimed expert, per chunk:
  input  [B × gate_cols] ──► gate+up+SiLU ──► hidden [B × gate_rows]
                                   │ (barrier)
                                   └────────► down ──► output [B × down_rows]

At today's NVMe-bound prefill the wall-clock effect is neutral — GPU dispatch hides entirely under I/O — but every future I/O gain (Q2 vaults, faster storage) lands on a GPU path that is already batched. If the batch or dimensions exceed the pre-allocated staging buffers, the bridge returns NULL and the caller falls back to the per-token path — capability degradation, never failure.

One source of truth

The MSL source is JIT-compiled from kernels.metal itself — the build embeds the file verbatim as a generated raw-string header (build/kernels_msl.h) that metal_bridge.mm includes, so the .metal file is the single source of truth and a kernel change lands in exactly one place. One MSL landmine worth engraving: grid-input attributes must agree in dimensionality — a uint2 [[threadgroup_position_in_grid]] requires a uint2 [[threads_per_threadgroup]], or the JIT compile fails at engine boot.

---

GPU-Resident Decode Attention

Each layer's entire decode attention block executes as one command buffer, one CPU sync — 94 sync points per generated token, where separate QKV and o_proj submissions with CPU work between them would cost 188:

① scout_matvec_bf16_sg × 3  Q / K / V projections (coalesced)
        │ memoryBarrierWithScope
② attn_prep                 per-head QK-RMSNorm (simd-reduced mean-of-squares,
        │                   +1e-6, matching the CPU formula) → RoPE →
        │                   K/V append into the KV ring slot
③ attn_decode               causal GQA attention: one threadgroup per query
        │                   head, scores staged in threadgroup memory
        │                   (context ≤ 4096 fits the 32 KB limit), softmax
        │                   via simd_max/simd_sum + threadgroup reductions
④ scout_matvec_bf16_sg      o_proj → one encodeSignalEvent, one spin-wait

The KV ring — ~1.5 GB — is registered once at boot as a single zero-copy MTLBuffer, so the GPU appends K/V into the very memory the CPU prefill path writes. The two writers coexist in the same ring coherently (UMA shared mode + event ordering), verified token-for-token across a multi-turn serve session. Attention is O(context) work that lives entirely on the GPU — it does not grow on the CPU as the conversation lengthens.

Not bit-exact — by construction

Simd-reduced norms and softmax reassociate floating-point sums, and the GPU's exp is not the CPU's expf. The kernels are verified at the token level: deterministic across runs, token-identical to the CPU reference on the canonical greedy prompt. Any change here must re-clear that bar — bitwise diffing against the CPU path will always fail, and that failure means nothing.

---

Grouped Expert Dispatch

All experts ready in the same claim sweep are encoded into one command buffer, with params passed via setBytes — no allocation, no wrapper object, honouring the no-allocation law in the hottest loop of the program. The honest sizing: command-buffer launch overhead measures only ~5 ms/token, so the grouping is hot-loop hygiene rather than a bandwidth win — the gpu-wait bucket is real kernel execution on late-arriving misses, and only faster kernels or fewer misses move it.

---

The LM Head Rides the Same Path

The heavy model's LM head — a 151,936 × 4,096 bf16 matvec — executes through the same coalesced GPU pipeline as every other dense projection, at ~234 GB/s (the scalar CPU matvec survives only as the no-Metal fallback). It runs only for generating positions: prompt positions before the last already know their next token, so the engine never spends the vocab-sized matvec on them.

---

Batched Scout Matvec

The workhorse smoe_metal_scout_matvec_batch_bf16 encodes N independent matvec operations into a single Metal command buffer:

smoe_metal_scout_matvec_batch_bf16(
    metal_ctx,
    batch_weights[N],   // N bf16 weight matrices
    batch_inputs[N],    // N input vectors
    batch_outputs[N],   // N output vectors
    batch_rows[N],
    batch_cols[N],
    N
);

(The scout_ prefix in the kernel names refers to the Scout weights these matvecs read — the resident dense backbone.) The dense layer-0 MLP and shared-expert gate/up pairs ride it, amortizing command-buffer overhead across every matvec that can share a submission. A barrier-free sibling, smoe_metal_scout_matvec_group_bf16, carries layer-major prefill's chunk-wide dispatches — all B tokens' QKV in one command buffer, all B o_proj rows in another — where every output plane is distinct, so the GPU is free to overlap the dispatches. Two dense matvecs deliberately live elsewhere: decode's QKV is encoded inside the per-layer attention command buffer above, and the exact-routing gate matvec runs on the CPU's NEON path — a 128 × d_model matvec is too small to be worth a GPU round-trip.

---

Buffer Pre-Registration

At dawn, all long-lived buffers are aggressively registered with the Metal context as permanent MTLBuffer handles:

• gate_weights_storage — all routing gate matrices
• hidden, normed — the shared hidden state and scratch buffers
• qbuf, kbuf, vbuf, attn_out — attention projection scratch
• gate_scores, gate_scores_batch — routing logit buffers
• The full KV ring — one ~1.5 GB zero-copy wrap covering every layer's K/V cache, written by GPU (attn_prep) during decode and by the CPU during prefill
• Streamer pool data — the entire ring buffer pool

Pre-registration obliterates newBufferWithBytesNoCopy overhead in the token generation loop. By the time the very first token fires, every pointer the GPU will ever need is already wrapped, primed, and accessible.