Axis-2 of the flagship stack ([[flagship-full-tuning-stack]]): the ONLY lever that beats the decode bandwidth wall — verify K tokens per forward. Multiplies on top of the Axis-1 byte cut (Axis-1 cuts per-forward cost; spec-decode cuts forward count). Industry: n-gram 2–4× (training-free), EAGLE-3 3–5×.
export_{lfm2_moe,qwen3_6_moe}_*) do NOT set
last_token_only, so the model default (last_token_only = False) holds → the head runs on ALL
positions → the graph outputs [1, S, vocab]. Feeding [1, K] returns K per-position logit vectors
in ONE forward. The bundle is causal, so position i’s logits are conditioned on the prefix ≤ i — exactly
what verification needs. No re-export required (the LFM #2 bundle on the A19 already qualifies).PipelinedBench.chunkStep already feeds [1, q] input_ids +
[1, processed+q] position_ids via InferenceFunction.run with the 4 states threaded as MutableViews.
It currently resolves the output to [1, 1, vocab]; for verify, resolve to [1, q, vocab] and read all q.State per step: accepted context ctx (token ids), processed = KV length, the 4 device states.
loop:
# 1. DRAFT (n-gram / prompt-lookup, no model): find the longest suffix of ctx that occurred earlier;
# propose the K tokens that followed it. (∅ → fall back to 1 normal decode step.)
draft = ngram_lookup(ctx, K) # e.g. K=4
# 2. SNAPSHOT the recurrent SSM state (conv_state, rec_state) — small, cheap to copy.
snap = (conv.copy(), rec.copy()); base = processed
# 3. VERIFY-FORWARD: one forward over [last_accepted, *draft] = [1, 1+len(draft)].
logits = chunkStep(tokens=[ctx[-1], *draft], processed=base) # -> [1, 1+len(draft), vocab]
argmax = [argmaxF(logits[:, i]) for i in range(1+len(draft))]
# 4. ACCEPT the longest prefix where draft[i] == argmax[i]; argmax at first mismatch = the free "bonus".
j = longest_prefix_match(draft, argmax[1:])
accepted = draft[:j] + [argmax[j]] # j drafts + 1 bonus token (always ≥1 token/forward)
ctx += accepted
# 5. ROLLBACK to the accepted length:
# - KV cache: append-only per position → just set processed = base + len(accepted) (stale slots
# beyond it are overwritten next forward). Trivial.
# - SSM (conv/rec): the verify-forward advanced them by 1+len(draft) tokens in-place. Restore `snap`
# then re-apply the `len(accepted)` accepted tokens as one small [1,len(accepted)] forward
# (cheap: len(accepted) ≤ K+1 ≪ the K weight-reads spec-decode saved). OR keep per-token SSM
# snapshots and truncate. Restore-and-replay is simplest and correct.
processed = base + len(accepted)
Lossless: every emitted token equals greedy argmax of the target model (verify), so the output distribution is identical to plain greedy — spec-decode only changes SPEED, never quality.
len(accepted) = j+1 (1 ≤ · ≤ K+1). Avg acceptance ā → speedup ≈ ā (minus the
small SSM restore-replay overhead). n-gram: ā ~2–4 on code/RAG/structured, ~1 on free chat.PB_SPEC mode: extend chunkStep to resolve the output as [1, q, vocab] and return all q rows;
implement ngram_lookup + the accept/rollback loop above; report tok/s + a token-match vs plain greedy
(must be EXACT — lossless gate). Reuse the LFM #2 bundle already on the A19 (no re-export).ā is high; then wire the vanilla-draft variant
(second engine feeding a small model’s argmax as the draft); then EAGLE-3.[1,len(accepted)] forward per step. If it erodes
the win, keep per-token conv/rec snapshots (K small states) and truncate instead. Measure first.ā: n-gram wins only on input-grounded tasks; EAGLE-3 is the general 3–5×.[1,K] forward = the same weight read as one decode step (weights dominate; the extra K−1
positions add ~0 weight bytes), so a rejected round costs ≈ 1 normal step — spec-decode is ~never a loss.Axis-1 kernels cap at ~2× (byte floor). Spec-decode is the only lever that cuts the FORWARD COUNT, so it stacks multiplicatively — the ~3× that turns “~2× flagship” into “~5–6× flagship”. Feasibility is high (verify-forward native; machinery exists); the build is host-loop + rollback, not a kernel or re-export. ```