feat(mtp): chain-of-K on hybrid targets + drafter-hidden cascade

[raullenchai]

Draft — spec_decode/ is the core-team lane. Opening for coordination/review before landing.

Why

Native-MTP speculative decoding on Qwen3.5/3.6 (hybrid GatedDeltaNet MoE) was clamped to K=1, and even K=1 left throughput on the table. This lifts the clamp and closes the gap toward llama.cpp.

What (3 pieces)

• Multi-slot recurrent rollback (cache_patch.py): the GatedDeltaNet chunk-split now snapshots (conv, ssm) at every draft boundary (rollback_states[n_from_end]) instead of one offset, so _rollback_draft(n) restores to any accepted depth. This lets generator.py drop the SSM K=1 clamp — chain-of-K works on hybrid targets. K=1 reduces to the exact prior single-snapshot behavior.
• Drafter-hidden cascade (qwen3_5_inject.py): mtp_forward gains return_hidden, returning the head's pre-lm_head hidden. The generator already feeds it back as the next iteration's hidden, so each successive draft conditions on the drafter's own refined hidden rather than the target's frozen hidden. K=3 acceptance ~45% → ~71%.
• Lean-Python cascade (generator.py): the draft cascade no longer host-syncs (mx.eval) per step — it stays lazy and materializes once at the verify's single eval, removing 2*K GPU round-trips/round.

Tests

• Full MTP unit suite green (200 passed, 2 skipped).
• Correctness: isolated test shows the chunk-split snapshot/restore is numerically exact (diffs 0.0). Greedy served output coherent.
• Measured (Qwen3.6-35B-A3B Q4, M3 Ultra): 1.24× → 1.36×, ~73–77% acceptance. Also validated on Qwen3.5-9B (dense hybrid).

Notes

Three files are interdependent (multi-slot API is shared by cache_patch + generator; cascade spans inject + generator), so kept as one PR. Depends functionally on the sidecar fix (#1214) to have a head to draft with.

🤖 Prepared with AI assistance (Claude Code); spec_decode/ review requested.

[github-actions]

PR #1216 validation scorecard

Title: feat(mtp): chain-of-K on hybrid targets + drafter-hidden cascade
Author: raullenchai
Diff: 4 file(s), +429/-198 LOC, blast radius: medium

Verdict: MERGE-SAFE

step	status	summary	time
fetch	PASS	4 files, +429/-198 LOC, blast=medium	1.3s
test_plan_check	PASS	no checkbox-style test plan found	0.0s
cl_description_quality	PASS	title OK + body has rationale (1908 chars)	0.0s
supply_chain	PASS	15 warning(s) — human review wanted	0.0s
test_env_check	PASS	installed trusted-pins (2 missing plugin(s) recovered)	4.9s
codex_review	skip	codex CLI not found on PATH (install: npm i -g @openai/codex)	0.0s
lint	PASS	clean (4 file(s))	0.0s
targeted_tests	PASS	exit 2 (in 2 target file(s))	0.8s
diff_coverage	skip	diff-cover not installed — advisory coverage unavailable	0.0s

Details

supply_chain — PASS

Findings:

• tests/test_mtp_multislot_rollback.py near l45: eval() — usually wrong; never on untrusted data

  # eval() exercises the SAME Metal-kernel segmented path production uses

• tests/test_mtp_multislot_rollback.py near l47: eval() — usually wrong; never on untrusted data

  layer.eval()

• tests/test_mtp_multislot_rollback.py near l48: eval() — usually wrong; never on untrusted data

  mx.eval(layer.parameters())

• tests/test_mtp_multislot_rollback.py near l72: eval() — usually wrong; never on untrusted data

  mx.eval(prefix, verify)

• tests/test_mtp_multislot_rollback.py near l77: eval() — usually wrong; never on untrusted data

  mx.eval(c[0], c[1])

• tests/test_mtp_multislot_rollback.py near l84: eval() — usually wrong; never on untrusted data

  mx.eval(c_pat[0], c_pat[1])

• tests/test_mtp_multislot_rollback.py near l99: eval() — usually wrong; never on untrusted data

  mx.eval(c_ref[0], c_ref[1])

• tests/test_mtp_multislot_rollback.py near l101: eval() — usually wrong; never on untrusted data

  mx.eval(conv_keep, ssm_keep)

• tests/test_mtp_multislot_rollback.py near l125: eval() — usually wrong; never on untrusted data

  mx.eval(prefix, verify)

• tests/test_mtp_multislot_rollback.py near l130: eval() — usually wrong; never on untrusted data

  mx.eval(c[0], c[1])

• tests/test_mtp_multislot_rollback.py near l135: eval() — usually wrong; never on untrusted data

  mx.eval(out_orig, c_orig[0], c_orig[1])

• tests/test_mtp_multislot_rollback.py near l140: eval() — usually wrong; never on untrusted data

  mx.eval(out_pat, c_pat[0], c_pat[1])

• tests/test_mtp_multislot_rollback.py near l176: eval() — usually wrong; never on untrusted data

  mx.eval(prefix, verify)

• tests/test_mtp_multislot_rollback.py near l180: eval() — usually wrong; never on untrusted data

  mx.eval(c[0], c[1])

• tests/test_mtp_multislot_rollback.py near l187: eval() — usually wrong; never on untrusted data

  mx.eval(c[0], c[1])

Artifacts:

• /tmp/pr_validate/pr-1216/run-20260726T171337Z-2143-1785086017733874851/supply-chain-scan.log

[raullenchai]

Folded in the single-pass GatedDeltaNet rollback perf fix (commit 8dbe4fe): replaces the per-boundary chunk-split with one fused gated_delta_update + lazy recompute-on-reject. Validated on M3 Ultra (Qwen3.6-35B-A3B Q4, warm, same OpenAI-server harness): baseline 95 → MTP 149 tok/s = 1.57x (was 114/1.18x), ~90% accept, output byte-identical to baseline (lossless). This is what makes chain-of-K actually pay off — without it the chunk-split ate the K≥2 gain. 479 MTP/spec tests pass; full unit suite green.