Layer-Major Batched Prefill
Prefill — the silence between pressing Enter and the first token — is where a streaming engine can bleed the most. A naive token-major walk gives every token its own private audience with the vault: 94 layers × 8 experts × ~9.4 MB ≈ 7 GB of Direct I/O per prompt token — a 45-token prompt would re-read a third of a terabyte from NVMe.
The layer-major prefill (src/prefill.cpp) inverts the loop — and with it, the economics.
The Inversion
Token-major prefill asks: "for this token, which experts?" — 94 times per token, once per token.
Layer-major prefill asks: "for this layer, which experts does the whole prompt need?" — once per layer, for up to 64 tokens at a time.
TOKEN-MAJOR (naive) LAYER-MAJOR (S-MoE)
for each token t: for each chunk of ≤64 tokens:
for each layer l: for each layer l:
fetch t's 8 experts ──► attention for ALL tokens
compute exact routing for ALL tokens
dedupe → expert UNION
I/O: tokens × layers × 8 fetch each union expert ONCE
blob reads apply it to every routed token
I/O: layers × |union| blob readsFor a 64-token chunk drawing 8 experts per token from a 128-expert layer, the union saturates near the layer's full population — meaning the I/O bound becomes "read each needed expert once per layer", regardless of prompt length. The longer the prompt, the greater the theft from the Monolith.
Choreography Per Chunk
For every layer L, all B tokens of the chunk move together — and the dense math moves as batches, never as per-token GPU round-trips:
- RMS norm (CPU), then ALL
Btokens' Q/K/V projections in one command buffer — one GPU sync for the whole chunk instead of one per token, and the coalesced matvec kernel does the reading (Metal Compute Kernels). - QK-norm + RoPE + K/V append per token (CPU) — RoPE at each token's absolute stream position (the tables are precomputed once; the angle depends only on position and dimension, never on head or layer). The KV ring is the engine's single cache, shared zero-copy with the GPU decode path.
- Causal attention per token (CPU/NEON): token
i's window is capped at its own history, so causality survives even though all chunk K/V rows are already resident. - ALL
Bo_proj rows in one command buffer, then residual + FFN norm per token (CPU). - Exact routing, firing as it goes: the real router gate (a 128 × d_model matvec, ~0.5M MACs) is evaluated on each token's true hidden state —
norm_topk_probre-normalisation per the Qwen3 specification. Every expert fires at the Streamer the moment it first enters the union, so the NVMe starts on the layer's blobs while later tokens are still routing. No prediction. No approximation. Prefill routing miss rate: zero, by construction. - Token-batch fused FFN: each claimed expert is applied to every token routed to it in a single two-pass GPU dispatch (
smoe_gate_up_batch/smoe_down_batch) — one dequantisation stream,Bactivations.
The last prompt token deliberately stays out of the batch: it runs through the ordinary serial step, so sampling and token emission live in exactly one place.
Why There Is No Prediction in Prefill
Prediction here would be a tax with no purpose: the next token is already known — it's the prompt. The heavy hidden state produces the routing exactly, for the price of one tiny matvec per layer, so nothing is guessed and nothing can miss. Decode uses the same exact routing at every layer of every generated token (see Surface Scout).
Retention, never replay
Replaying the previous token's routing as prefill prefetch hints measures slower (114 s vs 84 s on the 45-token benchmark): the hints stream extra experts ahead of demand fetches that must happen anyway. Adjacent-token expert overlap is real — 46.4% — but the way to harvest it is retention (keeping what was already read), never replay (reading on a guess). The distinction between skipping reads and adding them is the entire game; the negative result is documented in main.cpp so nobody walks that path twice. (The incremental fire in step 5 is the opposite kind of move: it starts reads that must happen anyway a few milliseconds earlier — demand pulled forward, never speculation added.)
Invariants, Untouched
The layer-major path defends every sacred invariant:
- Zero heap allocation: all activation planes (
hidden,normed,qbuf,kbuf,vbuf,attn_out, batch staging) are pre-allocated at startup —CHUNK × dimrows,posix_memalign(16384), Metal-registered once. - Atomics only: expert claims ride the Streamer's existing lock-free state machine.
- Zero-copy: expert blobs are consumed exactly where
pread()landed them.
Measured Result
45-token prompt, greedy, seed 1337, Qwen3-235B Q4 vault: TTFT ~29.4 s on the 48 GB reference machine — several times faster than a token-serial walk of the same prompt, and independent of prompt length beyond the chunk size.
Output token IDs are bit-identical to the token-serial exact-routing path (which still exists — the last prompt token runs through it) — same KV slots, same RoPE positions, same routing, same accumulation. The speedup is pure I/O arithmetic, not numerical drift.