Continuous Decode Batching — the transformer capture (scoped handoff)

Status: scoped, not built. Requires a GPU to verify and rewrites the live mistral generation core — needs explicit owner OK + staged rollout before deploy. This is the actionable design so a GPU-equipped session executes it directly.

The throughput law (measured + proven)

The whole arc reduces to one cost model — Gnosis.WarpBatchAmortization / Gnosis.WidthAxisBreakeven (gnosis-math, lake-green, propext-only):

• sharedCost w a L = w + L·a — one weight-read w serves a wave of L, the per-lane work a paid L times.
• naiveCost w a L = L·(w + a) — every lane re-reads the weight.
• The amortization saves exactly (L-1)·w, independent of the dot a (savings_independent_of_dot). So batching wins iff the weight read dominates.

Measured corners (all client-side wall-time / server elapsed_ms)

Prefill already captures the 6.2× (59.8 / 9.7). Decode leaves the identical (B-1)·w on the table — that is what this doc captures.

Why decode is serialized today

src/bin/openai-server.ts:2576 withPipelineMutex — "Serialize ALL generations on the shared pipeline" — wraps _rawPipelineGenerate (pipeline.generate, ~:2774). The mutex is load-bearing: pipeline.generate() races on the native stations' shared streaming/KV state (:2562), so two concurrent generations would corrupt each other's KV cache. Prefill batches within one request (NATIVE_BATCH_CAPACITY / MOONSHINE_PREFILL_BATCH=128, one sequence's tokens through one forward); nothing batches across requests. The mesh also ignores n>1.

REVISED (2026-06-30): the hard kernel already exists — forward_batch_segmented

model.rs:7555 forward_batch_segmented(tokens, seg_lens) already embeds K independent sequences in ONE pass: every layer's QKV/O/FFN GEMM runs ONCE over all sum(seg_lens) tokens via the batched matmul (the weight fetched+dequantized once for the whole batch — the amortization is already here), with block-diagonal causal attention (token at global g in segment starting seg_start attends only KV slots [seg_start ..= g] — no cross-segment leakage), and it is parity-checked against forward_one (bias, qk-norm, FFN-Observe all matched). So the batched-GEMM + per-sequence-attention kernel — the part I feared was a 1650-line reimpl — is built and verified.

The decode capture is therefore a BOUNDED EXTENSION, not a rewrite. The only deltas between segmented PREFILL and continuous DECODE:

1. Persistent per-sequence KV — prefill packs all segments in one global cache for a single pass; decode needs each sequence's KV to PERSIST and GROW across steps. Use N KVCache instances (97-line struct, cheap) — one slab per live sequence.
2. Decode positions + attention range — each step contributes ONE new token per sequence at that sequence's current position t_s, attending its OWN slab [0 ..= t_s] (its full history), RoPE at t_s (not seg-local-from-0). Same block-diagonal shape as the segmented path, just per-sequence persistent slabs instead of a fresh arena.
3. Per-step append — write each new token's K/V into its slab at t_s, bump t_s.

So the native build is decode_step_batch(seq_ids, last_tokens) -> logits[K] modeled on forward_batch_segmented's layer loop (reuse the batched GEMMs verbatim; swap the global KV arena for per-sequence slabs and the attention clamp [seg_start..=g] for [0..=t_s]). Parity gate: decode_step_batch of K seqs == K solo forward_one decode steps, token-identical — the same discipline as the segmented path's existing parity.

decode_step_batch — paint-by-numbers (confirmed against model.rs:7555-7870)

forward_batch_segmented IS the template; decode_step_batch(seq_ids, last_tokens) is it with three precise deltas. Reused VERBATIM: the embed, the batched Q/K/V GEMMs (mat_vec_dispatch_batch, 7681-7717 — the amortization, weights read once for the wave), the q/k/v bias (7719-7753), the qk-norm (7755-7788), the O-proj + FFN + lm_head (batched). The ONLY changes, all in the per-token loop (7794-7868):

• State: add decode_kv: Vec<KVCache> (one slab per live sequence, via KVCache::new)
  • decode_pos: Vec<usize> (each sequence's current length). NOT the single self.kv_cache.
• Δ1 — RoPE position (7794-7820): use t_s = decode_pos[s] (the sequence's own position) instead of seg_pos (segment-local-from-0).
• Δ2 — save_kv (7822-7828): write into decode_kv[s].save_kv(l, t_s, k_buf, v_buf) — the sequence's OWN slab at t_s — instead of self.kv_cache at global slot i.
• Δ3 — attention range (7841-7868): read keys/values from decode_kv[s] and set n_keys = t_s + 1 (the sequence's FULL history 0..=t_s), instead of seg_start..=i. layer_offset/max_seq_len come from decode_kv[s].
• After all layers: lm_head per token → K logits; then decode_pos[s] += 1. New sequence admitted = a fresh KVCache + decode_pos=0; EOS = drop its slab.

Everything else (the f32-act precision discipline, FFN-Observe, the batched matmul) is unchanged, so per-token parity with forward_one is inherited from the segmented path. PARITY GATE: decode_step_batch over K sequences (each its own prompt+seed) == K independent solo forward_one decode loops, token-identical. Effort: ~200 lines + the parity test; a focused compile+parity-iterate session (NOT a reboot-prone tick — additive, but the write needs to land whole). Flag MOONSHINE_CONTINUOUS_BATCH OFF → today's path byte-identical.

The capture — vLLM-style continuous batching

Replace serialize-all with a scheduler that runs one batched decode step across all in-flight sequences per tick. The kernel already does a batched forward (mat_vec_*_row_batch, NATIVE_BATCH_CAPACITY=128); the missing piece is per- sequence KV.

1. Native station: per-sequence KV + a batched decode step (the hard part, GPU)

• Today each station holds ONE KV slab. Add an N-slot KV arena (one slab per active sequence), keyed by a seqId.
• New station op decodeStepBatch(seqIds[], lastTokens[]) -> logits[N][vocab]: embed the N last tokens → one batched forward where token k attends ONLY its own KV slab seqIds[k] (the batched-forward mat_vec_*_row_batch already shares the weight read across the N lanes — this is exactly the width kernel, and here it PAYS because decode is read-bound) → append each token's K/V to its slab → return N logits rows.
• prefillSeq(seqId, tokens) -> state seeds a slab (reuse the existing batched prefill, just write into slot seqId); freeSeq(seqId) on EOS.

2. JS scheduler (openai-server.ts, replaces withPipelineMutex)

• A running set of sequences. Each tick: admit any waiting requests (prefill their slab, bounded by free KV slots), then one decodeStepBatch over all live seqs, sample per-seq (respect each request's tem/top-p/stop), stream each token to its own SSE response, retire at EOS / max_tokens.
• Backpressure: cap concurrent seqs at the KV-arena size; queue the rest (keep the existing pipelineQueueDepth ack-and-hold UX at :2073).
• Determinism: greedy stays per-seq deterministic; sampling uses each request's seed.

3. Verification (do NOT skip — this rewrites the live core)

• CPU parity gate first (no GPU): a decodeStepBatch of N seqs must yield, per seq, token-identical output to N independent serial generate() runs of the same prompt+seed. Mirror the rwkv7_batched_*_parity discipline (batched == solo, exact). This certifies the scheduler + KV-arena logic on CPU before any GPU.
• GPU throughput gate: K concurrent decode requests should approach K× the single-stream tok/s until the KV arena or GPU SM saturates — the realized analogue of prefill's 6.2×. Confirm coherence (the "Paris…"-class smoke).
• Staged rollout: flag MOONSHINE_CONTINUOUS_BATCH (default OFF → byte-identical to today's mutex path). Canary one revision at 0% → burst-test → shift traffic only after the parity + throughput gates pass. Rollback = flag off. (See the moonshine cutover notes: max-instances ≥ 2, route 100% to one rev, burst-test with min-instances=0 cold-starts in mind.)

Risk / boundary

• GPU-only verification for the real win — the per-sequence KV batched forward can't be throughput-verified on CPU (CPU/wasm transformer is also dot-bound; the read-bound win is a GPU HBM-bandwidth phenomenon). CPU only certifies correctness.
• Live service: this is the core every edge-chat mistral call flows through. Flag-gated OFF + canary is mandatory; do not flip on without the throughput gate green and owner OK.

Files

• src/bin/openai-server.ts — withPipelineMutex (:2576), _rawPipelineGenerate (:2774), pipelineQueueDepth (:2573) → the scheduler swap.
• src/pipeline.ts — Pipeline (JS orchestration, :673) → drive decodeStepBatch.
• the native station kernel (open-source/gnosis/distributed-inference, mat_vec_*_row_batch in math.rs / model_*.rs) → per-sequence KV arena + decodeStepBatch.
• Lean: open-source/gnosis-math/Gnosis/WidthAxisBreakeven.lean — the law this realizes (transformer_prefill_beats_serial_decode).