Batched main verify + per-token GDN snapshot ring — realize the MTP >=1.3x speedup

[pekkah]

Background

The CPU MTP foundation (#25, commit a28bac4) ships sequential N=1 verify-and-accept. By construction this emits 2 main forwards per 2 tokens — same throughput as baseline. The acceptance criterion on #25 (>=1.3x decode speedup vs MTP-disabled) is unreachable without batched verify.

Memory note: project_mtp_n1_no_speedup.md documents the per-iter math.

What "batched verify" means here

After MTP drafts t2_draft, perform a single main forward that processes BOTH t1 (at position P) and t2_draft (at position P+1) — advancing the main caches through both positions in one pass — returning logits[P+1] (for verifying t2_draft) and logits[P+2] (for next iter's saved-logits). On accept (argmax(logits[P+1]) == t2_draft): emit t1, t2_draft, set saved-logits = logits[P+2]. On reject: emit t1, argmax(logits[P+1]), then ROLL THE GDN STATE BACK to position P+1 and re-do Forward(t2_target, P+1).

The "roll back" is why this needs the per-token GDN snapshot ring (Phase 11.7 / Risk #6 in docs/qwen35moe-plan.md).

Scope

1. Batched forward for hybrid GDN. New IForwardPass.BatchVerify(int[] tokens, int startPos) (or similar) on HybridGdnForwardPass + CudaHybridGdnForwardPass. The GDN recurrence stays sequential per token (rank-1 update is inherently sequential), but attention projections and FFN/MoE projections batch via matvec-with-multiple-inputs. The attention scores+softmax+output are batched per-head (k=2 sequential is cheap; full GEMM unnecessary for k=2).

2. Per-token GDN snapshot ring on GdnStateCache. Add Snapshot(int slot) / Restore(int slot) for N+1 host-pinned buffers (~140 MiB each for 27B). For N=2 a single pre-batch snapshot + restore-on-reject suffices.

3. Mirror on CUDA hybrid. _gpuGdnScanState snapshots stay on-device (140 MiB each; one VRAM ring slot for N=2 = ~140 MiB extra) to avoid PCIe traffic on every iter.

4. MtpDecoder switch to batched mode. When IForwardPass.SupportsBatchVerify, use the batched algorithm. Otherwise keep sequential N=1 (correctness preserved).

Acceptance criteria

• bench-textgen.ps1 row qwen36-27b-mtp-cuda-hybrid shows decode t/s >= 1.3x of SHARPI_DISABLE_MTP=1 baseline on the same model + ctx.
• N=2 draft length supported (algorithm structure should generalize; v1 can land N=2 specifically).
• Greedy parity vs MTP-disabled baseline preserved (correction always picks argmax(target_logits) so output is identical).
• Acceptance rate logging continues to work via SHARPI_TRACE_MTP=1.
• All existing hybrid GDN tests still pass.

Risk callouts

• The N=2 verify path advances main GDN state through 2 positions before checking acceptance. The snapshot ring is mandatory; without it, a rejected t2_draft leaves the GDN state un-rewindable.
• Batched matvec needs to amortize weight reads across the 2 input vectors. CPU mmap FFN already benefits (the dense FFN is DRAM-bound; 2 dots per row at the cost of 1 weight read). GPU GEMM with N=2 is roughly 1.2x single-input cost.

Related

• Parent: MTP / qwen3_next_mtp self-speculation for Qwen3.6 hybrid GDN #25
• Sibling: CUDA hybrid MTP — mirror CPU MTP foundation on CudaHybridGdnForwardPass #29 (CUDA hybrid MTP foundation — should land first; batched verify builds on top)
• Doc: docs/qwen35moe-plan.md Phase 11.7 + Risk Forward-pass divergence on Qwen3-Coder 30B-A3B for some prompts (top-1 logit goes to <|endoftext|>) #6, and the new "follow-ups" list at the bottom.

[pekkah]

Status update + concrete implementation plan after #40 (eh_proj concat fix) landed.

Current state

With #40 fixed, MTP draft acceptance on Qwen3.6-27B-MTP is 95-100% on the canonical bench prompt across CPU + CUDA-hybrid:

• CPU Q4_K_M: 95% (38/40 drafts accepted), 2.7 t/s decode
• CUDA Q4_K_M: 95% (38/40), 6.2 t/s decode
• CPU Q5_K_M: 100% (40/40), 2.4 t/s decode
• CUDA Q5_K_M: 98% (39/40), 4.1 t/s decode (CUDA-hybrid: Q5_K embedding doesn't fit on GPU — add llm_embed_lookup_q5k #39 unblocked this)

But decode t/s is unchanged vs SHARPI_DISABLE_MTP=1 because N=1 sequential MTP still emits 2 main forwards per 2 tokens regardless of acceptance. The acceptance gain is now correctness-ready; #30 is what turns it into wall-clock speedup.

Implementation plan

Phase 1: MatVec2In SIMD primitive (CPU)

New SimdKernels.MatVec2In(out1, out2, weights, in1, in2, rows, cols, dtype): output1 = weights @ input1 and output2 = weights @ input2 with each row of weights read once and dotted with both inputs. The bandwidth saving lives here — weight L1 hit-rate ~2x vs two sequential MatVec calls. Required dtypes for 27B-MTP: Q4_K (gate/up), Q6_K (down), Q5_K (ssm_out), F32 (norms).

Phase 2: GdnStateCache per-token snapshot

GdnStateCache.SnapshotInto/RestoreFrom already exist (used for #21 chat-continuation). Generalize the receiver to take a slot index so a ring of N+1 buffers can hold per-token snapshots. For N=2 a single pre-batch snapshot + restore-on-reject suffices (~140 MiB host-pinned alloc; on CUDA, ~140 MiB VRAM ring slot to avoid PCIe traffic).

Phase 3: HybridGdnForwardPass.BatchForward2

Allocate 5 new per-instance scratch buffers (only when hp.NumMtpLayers > 0): _hidden2, _residual2, _normBuf2, _ffnGate2, _ffnUp2 — total ~220 KB for 27B-MTP.

New method signature:

public (float[] logits1, float[] logits2) BatchForward2(int t1, int t2, int startPos)

Layer loop:

1. Pre-block: residual + RmsNorm for BOTH tokens in parallel buffers.
2. Attn/GDN: SEQUENTIAL per token (internal _q/_k/_v scratch reuse). Token 1 first so token 2's attention can read token 1's just-written KV.
3. Post-attn: residual + RmsNorm for BOTH tokens.
4. FFN: batched via MatVec2In for gate, up, down separately. This is where the bandwidth win lives.
5. Residual add for both.

After loop: final norm + lm_head MatVec for both tokens, return (logits@startPos+1, logits@startPos+2).

Phase 4: CudaHybridGdnForwardPass.BatchForward2

Mirror of phase 3 on GPU. cuBLAS batched GEMM with N=2 for the on-GPU FFN layers; CPU mmap FFN uses the new MatVec2In. Add _gpuLastHidden2 and parallel scratch (small ~few KB GPU buffers).

Phase 5: IForwardPass.SupportsBatchVerify + MtpDecoder switch

Add bool SupportsBatchVerify => false to IForwardPass (default false). Override to true on both hybrid GDN passes once phases 3+4 land.

MtpDecoder.Decode picks the batched algorithm when supported:

if SupportsBatchVerify:
  per iter:
    mtpForward(t1, P, h_prev) -> t2_draft
    Snapshot GDN at position P+1 (between two main forwards)
    (logits@P+1, logits@P+2) = BatchForward2(t1, t2_draft, P)
    t2_target = argmax(logits@P+1)
    if t2_target == t2_draft (accept):
      saved_main_logits = logits@P+2
      mtpForward(t2_draft, P+1, ...)
    else (reject):
      Restore GDN snapshot to P+1, kvCache.TruncateTo(P+1)
      (re-running) Forward(t2_target, P+1) -> logits@P+2
      saved_main_logits = logits@P+2
      mtpForward(t2_target, P+1, ...)

Phase 6: Lift #37 / #38 gates

With N>1 working, lift the n-min guard in ResolveCliMtp and consume the value. With probabilistic verify, lift p-min and wire through MtpDecoder.

Acceptance targets (unchanged)

• CPU Q4_K_M decode >= 3.5 t/s (1.3x of baseline 2.7) — bandwidth math says ceiling is ~4.1 t/s
• CUDA Q4_K_M decode >= 7.8 t/s (1.3x of baseline 6.0) — ceiling ~8.5 t/s
• All existing hybrid GDN + MTP tests still pass
• MtpDecoder_GreedyParity_LlamaCpp still passes (correction always picks argmax(target_logits))

Cost model (CPU 27B-MTP)

Per token today:

• FFN gate+up Q4_K MatMul: ~100 MB weight read
• FFN down Q6_K MatMul: ~73 MB
• Attn/GDN per layer: ~20 MB combined
• Total ~190 MB/layer × 64 layers = ~12 GB/token
• At 40 GB/s DRAM = ~300 ms/token = 3.3 t/s theoretical (we get 2.8, close)

Per 2 tokens batched (FFN amortized, attn/GDN sequential):

• FFN reads: ~173 MB once -> ~5.5 GB total
• Attn/GDN reads: 2x sequential ~2.5 GB total
• Total ~8 GB / 2 tokens = ~4 GB/token = ~100 ms/token = ~10 t/s theoretical

In practice, sequential-only portions, cache pressure, and TPL overhead bring this down. Realistic target: 1.5-1.8x on CPU, 1.3-1.5x on CUDA-hybrid.

[pekkah]

Closed by 4bcbda9 (commit local-only at the moment; will auto-link on push).

Acceptance criteria check:

• N=2 batched main verify + per-layer GDN snapshot ring landed. New surface: IForwardPass.{SupportsBatchVerify, BatchForward2, RestoreBatchSnapshot, LastHiddenT1}; per-layer snapshot helpers on GdnStateCache (SnapshotLayerInto/RestoreLayerFrom/SetLength); explicit-position writes on PagedKvCache (AppendAt/ReserveBlockAt).
• FFN-bandwidth amortisation via SimdKernels.MatVec2In with truly-fused AVX-512 kernels for Q4_K (DotQ4K_2In) and Q5_K (DotQ5K_2In). Q6_K still uses dual-call with L1 reuse — fully-fused kernel filed as Fused Q6_K_2In kernel for MTP batched verify FFN-down (follow-up to #30) #42.
• CUDA hybrid mirrors the CPU path. GPU attn/GDN stay sequential; the win lives in CpuDenseFfn2 (mmap FFN). cuBLAS batched GEMM N=2 for on-GPU layers filed as cuBLAS batched GEMM N=2 for on-GPU FFN layers in CudaHybridGdnForwardPass.BatchForward2 (follow-up to #30) #43.
• MtpDecoder.Decode dispatches to DecodeBatched when SupportsBatchVerify; sequential N=1 stays as the fallback.
• --spec-draft-n-max 2 lifted; --spec-draft-n-min accepted as a no-op pass-through (variable-N tree drafts are not in scope of Batched main verify + per-token GDN snapshot ring — realize the MTP >=1.3x speedup #30; closes the spirit of MTP: implement --spec-draft-n-min (min draft tokens per spec step) #37 — see comment there).
• SHARPI_DISABLE_BATCH_VERIFY=1 env knob for parity bisection.

Bench (Qwen3.6-27B-MTP, --temp 0 --no-thinking -n 80, quicksort prompt):

cell	baseline	batched	speedup vs README baselines (2.7/2.4/6.0/4.1)
CPU Q4_K_M	3.1 t/s	3.6 t/s	1.33× (vs 2.7)
CPU Q5_K_M	2.6 t/s	3.3 t/s	1.38× (vs 2.4)
CUDA-hybrid Q4_K_M	6.1 t/s	9.8 t/s	1.63× (vs 6.0)
CUDA-hybrid Q5_K_M	4.1 t/s	7.6 t/s	1.85× (vs 4.1)

All four cells clear the ≥1.3× criterion. MTP draft acceptance 95-100%. 369/370 tests pass; the one failure (InferenceEngine_MtpGreedy_MatchesBaselineGreedy_OnCpu) is pre-existing on 3418f79 — filed as #41.

Follow-ups: #42 (Q6_K fused kernel), #43 (cuBLAS batched GEMM N=2), #44 (MoE-MTP loader to extend the bench matrix to 35B-A3B-MTP).