perf(qwen): full-context MTP draft prefill and replay-free rewind

[alan13367]

Summary

Qwen native MTP was a net loss on SSD-offloaded decode. This PR fixes why and makes it faster:

1. The draft layer saw almost none of the prompt. After layer-major prefill, the MTP layer's attention cache was filled with only the last mtp_slots // top_k positions (3–4 tokens at default budgets). It is now filled from the whole prompt. On Code, draft acceptance rose from 65% to 81%.
2. Draft-layer prefill runs through one whole-layer expert buffer (batched_layer) in prefill_step_size chunks, instead of many slot-sized chunks.
3. Grouped expert prefill is no longer disabled with MTP. Grouping only reorders gather_qmm rows, and the existing bit-identity test covers it. The App toggle and its hint text (en / zh-Hans / zh-Hant) no longer say "requires MTP off".
4. Replay-free rollback of rejected drafts. Previously each accepted token was replayed one at a time through the full model, including SSD expert reads. Now attention layers trim the verified KV cache. Linear-attention layers rerun only PLE and the Gated DeltaNet mixer on their captured verification inputs, starting from the pre-verification state, with no routed experts. Rollback time per request fell from 1.2–5.7 s to 0.02–0.06 s.
5. Pooled QSA index keys are truncated on trim instead of rebuilt; complete blocks before the new offset are unchanged.

The installed model format, sampling, and the MTP-off path are unchanged.

Measurements

Apple M4 Max, 36 GB, macOS 27.0.1, MLX 0.32.2. I ran the runtime directly, not through the App or HTTP server. Each run is one request in a fresh process: bundled Code or Novel corpus, 4,096 input tokens, 256 output tokens, greedy seed 42, Prompt Cache off.

Settings: 7.5 GiB expert cache, LFU, 16 read workers, 1,024-token batched layer-major prefill. MTP runs use a 1.3 GiB MTP cache, 2 draft tokens, and a zero-acceptance limit of 32. Base ba9e126 and this branch were interleaved A/B/B/A, two runs per arm. The OS page cache was not purged.

Workload	Config	Decode tok/s, base	Decode tok/s, PR	TTFT s, base	TTFT s, PR	Acceptance
Code	MTP on	7.38 / 7.42	9.87 / 9.90	25.3 / 25.2	21.5 / 21.5	64.9% → 80.9%
Novel	MTP on	5.77 / 5.79	8.37 / 8.27	32.2 / 24.7	21.4 / 29.5	34.8% → 42.5%
Code	MTP off	8.04 / 8.59	8.14 / 8.54	26.4 / 26.0	26.1 / 26.4	–

• MTP-off output tokens are identical to the base, and speed is unchanged within noise.
• There was no swap in any run. Sampled peak footprint was 19.6 GiB with MTP and 18.4 GiB without.

Caveats

• MTP depends on the workload. With this PR it is +18.5% decode over MTP off on Code, but −14% on Novel prose (9.70 tok/s with MTP off vs 8.32 with it), where acceptance is about 43%.
• The zero-acceptance-streak fallback is fragile. With the App's 2/2 custom strategy, one Code run fell back after two early rounds. A rate-aware fallback is left as follow-up.
• MTP output differs from MTP-off output starting at the second token, on base as well. I checked the target logits at that position: the top three are within one bf16 step (17.375 / 17.25 / 17.25), so the verification batch shape flips the argmax. This is numerical, not a correctness bug.

BENCHMARK.md has the full section, including a rejected decode-time next-layer expert prefetch. It kept outputs identical but was slower, because this SSD is bandwidth-bound: one 2.6 MB expert read takes about 0.65 ms, and four concurrent reads take about 1.95 ms.

Tests

• New runtime/tests/test_qwen_mtp_rewind.py: a tiny 4-layer Qwen with PLE, with and without pooled index keys. It checks rewound caches and next-token logits against a direct forward over the accepted prefix, and asserts no routed experts are used. I confirmed it fails if the PLE rerun or the KV trim is removed.
• Updated the tests that encoded the old behaviour: the MTP prefill window, the grouped-experts MTP guard, and pooled-key reset on trim. Added a test for the whole-layer-buffer MTP prefill.
• make test-python: 618 tests, 0 failures, 10 optional skips.
• make test: 147 App + 26 repack tests, 0 failures, 4 optional skips.

🤖 Generated with Claude Code

[yanun0323]

@alan13367
Thanks for the improvements! I found one timing issue: the timer stopped before some of the rollback calculations had finished, leaving that work for the next step.
I pushed a fix in eba1f5b so the timer now waits for that work to finish. I also added three tests that catch the issue and marked the benchmark numbers as measurements taken before this fix.

[alan13367]

@alan13367 Thanks for the improvements! I found one timing issue: the timer stopped before some of the rollback calculations had finished, leaving that work for the next step. I pushed a fix in eba1f5b so the timer now waits for that work to finish. I also added three tests that catch the issue and marked the benchmark numbers as measurements taken before this fix.

You're welcome, thanks for the awesome work in this project! I'll keep using it and contributing if I see good opportunities to do so :D