Skip to content

Data Flow

This is the precise choreography of tokens through the S-MoE engine. The engine lives in two distinct regimes: layer-major prefill (the prompt, where the future is known) and exact-routing decode (generation — where the future is more knowable than it looks). Each regime gets the choreography it deserves.


Regime 1 — Prefill (the prompt)

The prompt's tokens are already decided, so prediction is waste and token-serial I/O is robbery. All prompt tokens except the last move through the model layer by layer, in chunks of up to 64, with each layer's routed experts deduplicated and read once (full treatment: Layer-Major Batched Prefill):

chunk of B ≤ 64 tokens, for each layer L:

  ① attention block for ALL B tokens        (norm → QKV → QK-norm →
      RoPE at absolute positions → causal GQA → o_proj → residual)
  ② exact routing for ALL B tokens          (real gate matvec on the
      true hidden state — miss rate 0 by construction)
  ③ UNION of the chunk's experts ──► Streamer prefetches everything,
      workers stream ahead while…
  ④ …each expert is claimed ONCE and applied to every routed token
      in a single token-batch GPU dispatch

The last prompt token runs through the serial step below, so sampling and emission live in one place. During prefill no LM head runs and no sampler runs — positions before the last already know their next token.

In server mode, prefill applies only to the suffix beyond the longest common prefix with the engine's stored stream — turn 2 of a conversation never re-pays for turn 1.


Regime 2 — Decode (generation)

Decode uses the same exact routing as prefill — the real router gate evaluated on the true hidden state at every MoE layer. What keeps the Streamer ahead of the GPU is not a predictor but the routing archive:

Time ──────────────────────────────────────────────────────────►

Step N, layer l:
  Heavy    ─── [ attention → gate matvec → exact top-8 for (N, l) ]───►
  Ring     ─── [ ~46% of those experts are already retained from    ]
               [ steps N−1, N−2, … — claimed with zero bytes read   ]
  Streamer ─── [ demand pread() for the misses; optional --spec N   ]───►
               [ ranks 9..8+N queued LOW-PRIORITY as bets for N+1   ]
  Metal    ─── [ fused FFN per claimed expert → weighted residual   ]───►

The GPU and the SSD remain sovereign entities. The measured fact this design rests on: adjacent tokens reuse 46.4% of each other's experts, and the current token's top-16 gate ranking covers 63.7% of the next token's top-8 — a free oracle that no predictive pass beats (Surface Scout).

Serial Step Trace

1. Token Embedding

heavy_cur_token ──► embed_tokens[token_id]
                            │ memcpy d_model floats

                      heavy_hidden[d_model]    (pre-allocated, 16KB-aligned)

RoPE cos/sin for the step's absolute stream position are computed once here and reused by every head of every layer — the angle depends only on (position, dimension).

2. Layer Loop (l = 0 … num_layers-1)

heavy_hidden

    ├── RMS Norm (input_layernorm) ──► heavy_normed

    ├── ATTENTION — ONE Metal command buffer, ONE CPU sync:
    │     ① QKV projections (batched bf16 matvec)
    │     ② attn_prep: per-head QK-RMSNorm → RoPE → K/V append
    │          into the layer's KV ring slot (GPU writes UMA directly)
    │     ③ attn_decode: causal GQA attention, one threadgroup
    │          per query head, simd-reduced softmax
    │     ④ O Proj ──► residual
    │     (the CPU path survives verbatim as the no-Metal fallback)

    ├── Residual add ──► heavy_hidden

    ├── RMS Norm (post_attention_layernorm) ──► heavy_normed

    └── FFN block:
          if l == 0 && has_dense_layer_0:
              Dense MLP (gate × up SiLU × down) ──► residual
          else (MoE layer):
              ┌── EXACT ROUTING: gate matvec on heavy_normed ──► ranked top-(8+spec)
              │       ranks 0–7  → this token's experts + norm_topk_prob weights
              │       ranks 8+   → speculative prefetch for the NEXT token
              │                    (low-priority queue; --spec 0 disables)
              ├── Shared Expert FFN (CPU, concurrent with GPU) ──► shared_out
              │      (when the architecture has one — Qwen3-235B does not)
              └── Routed Experts (the exact top-8):
                      sweep over pred.expert_ids[]:
                          slot = streamer.claim_specific(layer, expert_id)
                          ── retained-ring hit → zero bytes read
                          all experts READY in the same sweep are encoded
                          into ONE command buffer (setBytes params — zero
                          allocation) ──► fused FFN dispatches, GPU async
                          [spin-wait + re-prefetch if not yet loaded]
                          gather GPU outputs × routing weights ──► routed_out

              └── residual += shared_out + routed_out

3. Final Norm + LM Head — generating positions only

heavy_hidden

    ├── RMS Norm (model.norm) ──► heavy_hidden

    └── LM Head matvec on METAL ──► logit_scores[vocab_size]
                                │     (151,936 × 4,096 — the same GPU
                                │      matvec path as QKV/o_proj; the
                                │      CPU matvec is the fallback)
                                └── sampler ──► next_token_id
                                    (repetition penalty → temperature softmax
                                     → top-k / top-p truncation; temperature
                                     below 1e-4 collapses to greedy argmax)

Prompt positions before the last skip this block entirely: their next token is already known, so the vocab-sized matvec and the sampling pass would be ~622M discarded MACs per token.


Ring Buffer State Machine

Each ring slot moves through four atomic states — but the cycle is not a naive round trip. The ring is a cache, not a conveyor belt:

StateMeaning
EMPTYUnclaimed. Waiting for the Streamer's command.
LOADINGA worker is executing pread() into this slot.
READYData has materialized — claimable by the main thread.
CONSUMEDClaimed at least once. The data remains valid — the slot is retained with its expert identity, and a later request for the same (layer, expert) re-claims it as a cache hit.

Slots leave CONSUMED/READY only through LRU eviction: when prefetch() finds no EMPTY slot, it evicts the least-recently-used idle slot (ref_count == 0 — a slot the GPU is actively reading is untouchable). Freshly loaded slots are stamped with a fresh LRU tick so a large prefetch burst can never evict its own head.

There is no wholesale pruning anywhere. Retention + lazy LRU is the single eviction discipline in both regimes, and it survives across serve-mode turns — expert weights are immutable, so a retained slot can never go stale. Any per-token wipe would silently destroy the measured 46% cross-token reuse and re-read gigabytes per generated token.

Transitions occur solely via std::atomic compare-exchange. No mutexes. No condition variables.


Handling Misses

There is no divergence to handle — exact routing cannot be wrong. The only miss left in the system is an I/O miss: an expert the ring does not hold when its layer arrives. The claim loop spin-waits (with periodic re-prefetch) until the demand fetch lands; correctness never depends on any cache. Speculative bets that never get claimed simply age out of the ring through LRU eviction — nothing is ever explicitly cleaned up, by design.

MIT License.