[jason-fxz]

Refs #828

Third PR of the speculative decoding plan in #828: the first draft method, MTP, which drafts with the checkpoint's own mtp.* layer. It covers the Qwen3.5 family (qwen3_5 and qwen3_5_moe architectures, dense and MoE) from HF checkpoints, and Qwen3.8-Flash-Next (qwen4_exp) from its HF checkpoint or from a GGUF target with a Strata MTP pack. It builds on the verify step from #836, the seams from #833, the radix_lookahead keys from #834 and the shared GatedDeltaNet layer from #841. The full stack is in the draft PoC #832.

Summary

--speculative-algorithm mtp (--spec-algo mtp) turns it on; k defaults to 3 (--speculative-num-steps, up to 8). --speculative-draft-model-path (--spec-draft-model) points at the draft weights when they are not in the target checkpoint, which is the case for Flash-Next GGUF targets: pass a Strata MTP pack, Strata's GGUF that keeps the checkpoint's mtp.* tensors under their HF names. The nextn MTP blocks llama.cpp writes are not read, for either family. A separate draft checkpoint folder is read with the target's quantization config. docs/cli.md documents the three flags.

Rejected at config time, before any weights load, so a draft that cannot load does not fail only after the whole target has loaded:

• Qwen3.5-family GGUF targets ("MTP is not supported for ... yet"): their MTP weights need a loader of their own.
• An FTW checkpoint as the draft source ("MTP cannot load its draft from an FTW checkpoint yet"): the FTW conversion drops mtp.*.
• A GGUF draft source without mtp.* tensors, e.g. a Flash-Next GGUF target with no --speculative-draft-model-path; a GGUF draft for a Qwen3.5-family target, whose draft reads safetensors only; a GGUF target with a draft source that is not a GGUF.
• A local draft folder without mtp.* tensors, and an explicit local draft path that does not exist.
• Any model whose per-request slot state declares no verify handler (see below).

A checkpoint whose MTP routed experts use a layout the reader does not handle (per-expert NVFP4 such as nvidia/Qwen3.8-Flash-Next-FP4-Local, or per-tensor FP8) stops at draft load with an error naming the unsupported tensors.

Commits:

• feat(spec): add the MTP proposer for the qwen3_5 family
• test(models): read the qwen3_5 MTP head of bf16 and block-fp8 checkpoints
• feat(spec): verify qwen4_exp targets on qsa_sparse graphs with PLE verify handlers
• feat(spec): add the qwen4_exp MTP draft from the checkpoint or a Strata pack
• docs(cli): document the speculative decoding flags
• fix(spec): refuse MTP draft sources the draft cannot read before the target loads
• docs(qwen4_exp): describe the verify path in the PLE docstrings

Design

prefill    target forward -> final hidden of every prompt row
           draft over the same rows (row p pairs token p+1 with the target hidden at p) -> draft KV, d1
verify     target scores [pending, d1..dk]; n tokens accepted
post       draft extend over the k+1 verify rows (one graph), keep the hidden of the last accepted row
           k-1 chained draft decode steps (one graph) -> d1..dk for the next verify

MTP proposer (spec/mtp/):

• A family registers its draft class as ModelSpec.spec_drafts["mtp"]; the interface is the MTPDraft protocol in spec/mtp/proposer.py. This PR adds no import between model families; the routed-expert reader both use is moe/expert_pieces.iter_extra_layer_expert_pieces.
• The draft layer is built from the target's own layer classes and shares the target's embedding and LM head. Its MoE is built with moe_strategy="fused", so its experts are resident on the GPU and never enter the target's expert cache or offload, whatever the target's strategy.
• The draft has its own KV pool keyed by the target's slot ids, with the target's page size and page count, so the target's page table, prefix cache and rollback cover it. It is counted in the KV budget next to the target pool. A cache rebuild keeps the draft KV exactly when it keeps the target KV pool, and recaptures the draft graphs.
• The draft's KV at position i depends on token i+1, so MTP sets radix_lookahead = True: a prefix hit never reuses a draft KV slot written for a different next token.
• Draft forwards run under Context.override with the target's GDN pool hidden. The draft's per-step state (relayed hidden, first draft token, picked rows) stays inside MTPProposer; nothing method-specific goes on the batch. Host-to-device index uploads are pinned and non-blocking.

Flash-Next target verify:

• Both families use the shared GatedDeltaNet layer from refactor(layers): share one GatedDeltaNet op in layers/linear_attn #841, whose verify path comes with feat(spec): verify draft tokens on the target side #836; this PR adds nothing GDN-specific.
• qsa_sparse captures multi-row verify graphs (uniform_extend_graph="device"); before this, speculative decoding rejected Flash-Next at startup. The per-request pending ring of index keys is sized with the draft depth (QSAKVCache.ring_capacity_for(index_ratio, k)), so the rows a verify writes and then rejects never overwrite a ring entry an open group still needs. The same applies to the QSA draft pool.
• PLE state rides the GDN state slots as two slot states (ple_conv, ple_ngram_ctx). The pool does not know how to restore them: each SlotStateSpec now names a SlotVerify handler class (models/config.py). The pool builds one handler per state next to the GDN verify scratch, the owner gets it back with pool.verify_handler(name) at forward time, and commit_verify runs the GDN commit and then every handler's commit; the handlers share one pinned, non-blocking upload of the slot and count indices; verify request i's scratch is row i. Handler scratch bytes are counted with the GDN verify scratch. Speculative decoding is rejected at config time, before the weights load, if any slot state has no handler.
• The PLE verify forward reads the live state only: the conv runs the decode taps over every verify row on the device (PLELayer._verify_conv) and stages [live history | verify conv inputs] per row; the n-gram context is not rolled, and the model stages [live context | verify input ids] per row instead. The commit restores both as they stand after the accepted rows. PLEMetadata.is_verify carries the flag, and a verify takes its PLE slots and indptr from batch.fla_metadata as decode does.
• --ple-backend disk (the default): a verify graph replay reads its rows from the graph staging buffer like decode does, rows_per_req tokens per request read back from the device. The graph buffer holds max_running_req * (k + 1) rows, since verify graphs stop at max_running_req requests.

Flash-Next MTP draft (models/qwen4_exp/mtp.py):

• One QSA decoder layer over the hyper-connection streams. The target hands its residual streams before the final mix to the draft (FeatureSink.final_hidden from the model, with the LM head told not to export again), and the draft maps its output through its own mixer before the target's LM head.
• Weights come from the checkpoint's mtp.* tensors, or from a Strata MTP pack: a GGUF with the same tensor names, ggml-typed linears and stacked ggml experts, passed with --speculative-draft-model-path.
• The chained draft steps capture one graph per batch size on the QSA backend as on triton.

Tests

New or extended:

• tests/models/test_qwen3_5_moe_weight.py: the MTP reader fills exactly the draft the engine builds, for bf16 (stacked experts) and block-fp8 (per-expert, quantized linears) checkpoints, and the routed experts keep the form the checkpoint stores.
• tests/models/qwen4_exp/test_ple.py::test_verify_commit_matches_stepwise_decode: a PLE verify over k+1 rows leaves the conv history and n-gram context untouched; committing n rows (different n per request, verify rows mapped to other slots) leaves exactly what n decode steps leave, and each row's conv output equals its decode step bitwise.
• tests/models/qwen4_exp/test_qsa_backend.py::test_decode_graph_replay_matches_eager[verify]: a captured 4-row verify replay equals the eager verify bitwise.
• tests/models/qwen4_exp/test_qsa_backend.py::test_verify_after_rejected_rows_compresses_like_decode: after a verify whose last rows are rejected, the next verify compresses the open index group as decode steps do.
• tests/models/qwen4_exp/test_ple_disk.py::test_wait_sync_large_graph[verify]: a verify graph replay on the disk PLE backend reads every row it needs, bitwise against the in-memory table.
• tests/spec/test_verify_contract.py: a tiny qwen4_exp model (GDN + QSA, PLE, hyper-connections) joins the contract test: speculative decoding with stub drafts must reproduce plain greedy decoding token for token, with and without graphs. The dummy weights are scaled by fan-in and the PLE table gets seeded rows, so the PLE path carries signal and ulp noise stays small; the tie check is now relative to the logit magnitude (the same 0.125 bound as before for logits in [64, 128)).

Each new check, run with only the logic it guards reverted:

Test	Logic reverted	Result
test_mtp_head_fills_the_draft_and_its_experts	the MTP reader stops skipping the routed experts	2 failed: emitted names differ from the draft's state dict
test_verify_commit_matches_stepwise_decode	a verify runs the prefill conv (the path before this PR)	3 failed: the live PLE state changed during the verify
test_verify_commit_matches_stepwise_decode	the conv commit restores the window one column early	3 failed: wrong conv history after n rows
test_verify_after_rejected_rows_compresses_like_decode	the KV pool ignores the draft depth when sizing the ring	failed: 125 of 128 elements differ, max abs diff 1.91
test_wait_sync_large_graph[verify]	host_fill_batch takes the graph path for decode only	failed with "fill was not released": the verify batch took the eager fill path
test_verify_steps_reproduce_greedy_decode[qwen4_exp]	the model does not stage the n-gram context on a verify	2 failed: the runs part at token 2 with a top-2 gap of 0.133

Reverting the staged token_to_req slice in the QSA graph test makes the capture crash on a shape mismatch rather than fail an assertion. The two test_verify_commit_matches_stepwise_decode rows ran on 4f8f0cda (this PR on an earlier head of #836); the other rows ran on this branch before its first rebase onto #836.

Test environment

• GPU: 1x NVIDIA H100 80GB HBM3, driver 580.95.05
• CPU: Intel Xeon Platinum 8480+ (26 vCPUs), 221 GB RAM, Ubuntu 22.04.5
• CUDA 13.0, torch 2.14.1+cu130, Python 3.12.12
• This branch at 8815c0b, on main c2c3e82.
• The full runs below used e02a5890: the same commits on the last head of feat(spec): verify draft tokens on the target side #836 (fe28362, on main d043fb6). The rebase onto main changed no line of this PR's diff (one hunk in kernel/triton/attention.py moved by 6 lines) besides one sentence in docs/cli.md. Rerun on 8815c0b: tests/spec, tests/layers, tests/models/qwen4_exp, tests/models/test_qwen3_5_moe_weight.py, tests/moe and tests/kernels (758 passed, 244 skipped), and the MTP-on greedy runs of Qwen3.6-35B-A3B and the Flash-Next NVFP4 checkpoint, both token for token identical to the ones on e02a5890; the Flash-Next GGUF with the Strata pack also ran greedy with MTP (846 tokens, answer 70, accept length 3.44).

Unit tests

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B PYTHONPATH=python pytest tests/spec tests/layers tests/scheduler tests/engine tests/kvcache tests/attention tests/moe/test_offload.py tests/server tests/tokenizer tests/models tests/kernels/test_kda.py

1705 passed, 211 skipped, 1 failed. The failure is tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]], the same one #836 reports (a local checkpoint with quant_method: modelopt_mixed); this PR does not change the quantization code it checks.

Smoke tests, both passed:

FREETOKEN_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl PYTHONPATH=python pytest tests/e2e/test_aime.py
FREETOKEN_REBUILD_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_REBUILD_EXTRA_ARGS="--speculative-algorithm mtp --max-running-requests 4" PYTHONPATH=python pytest tests/e2e/test_cache_rebuild.py

The second one boots ft serve with MTP on and grows the KV pool (the draft KV is rebuilt), grows the GDN state pool (the draft KV is kept), flushes the prefix cache and is refused an unfittable size, generating after each step.

End-to-end

All runs use AIME25 problem 1 (problems 1-8 for the concurrent run) through freetoken.llm.LLM(model_path, dtype=torch.bfloat16, max_running_req=4) with the chat template, MTP with speculative_algorithm="mtp" and the default k = 3. Checkpoints: Qwen/Qwen3.6-35B-A3B and Qwen/Qwen3.6-27B (bf16), nvidia/Qwen3.8-Flash-Next-NVFP4 (default --ple-backend disk, experts offloaded by the automatic MoE cache), and the Flash-Next GGUF ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF (IQ3_XXS) with the Strata pack mtp-q2_0.gguf as speculative_draft_model_path.

Speculative decoding off must match main. Greedy (temperature=0, max_tokens=16384, natural end) on Qwen3.6-35B-A3B, compared token by token with the same run on main a2ce636 (which #841 and #836 reproduced token for token): both runs hit the 16384-token limit and are identical.

MTP off vs on, greedy, same settings:

Checkpoint	off	MTP on	accept length	first different token
Qwen3.6-35B-A3B	16384 tokens (limit), 85.4 s	16384 tokens (limit), 40.8 s	3.74	27
Qwen3.6-27B	15676 tokens, answer 70, 312.5 s	16384 tokens (limit), 109.9 s	3.67	559
Qwen3.8-Flash-Next-NVFP4	727 tokens, answer 70, 9.8 s	739 tokens, answer 70, 6.5 s	3.36	102

The accept length is committed tokens per request per verify step (at most k + 1 = 4). The Qwen3.6 divergence points are the ones the PoC run showed.

8 requests at once with MTP (problems 1-8 submitted together, recommended sampling temperature=1.0, top_p=0.95, top_k=20, max_tokens=4096): every request finished without error and the end of every output reads as coherent reasoning. Where a request finished before the limit, its boxed answer is the reference answer.

Checkpoint	wall time	accept length	boxed answers before the limit
Qwen3.6-35B-A3B	41.1 s	3.33	none (all 8 hit 4096 tokens)
Qwen3.6-27B	63.7 s	3.34	none (all 8 hit 4096 tokens)
Qwen3.8-Flash-Next-NVFP4	135.9 s	2.60	5 of 8: 70, 16, 117, 504, 77
Flash-Next GGUF + Strata pack	100.2 s	2.46	3 of 8: 70, 16, 117

Decode speed, one request: problem 1, recommended sampling, 1024 decode tokens (ignore_eos), one warmup then 3 runs, median. The KV column is the pool the engine sized at startup with default memory settings; the draft's weights and KV come out of the same budget.

Checkpoint	KV tokens off / on	off tok/s	MTP tok/s	accept length	PoC #832 off -> on (accept)
Qwen3.6-27B	257595 / 230251	50.3	142.2	3.50	50.3 -> 137.7 (3.50)
Qwen3.6-35B-A3B	224198 / 126644	159.0	261.2	3.47	159.0 -> 243.5 (3.47)
Qwen3.8-Flash-Next-NVFP4	8256 / 8256	96.7	161.1	3.08	96.4 -> 153.7 (3.08)
Flash-Next GGUF + Strata pack	1019072 / 906560	105.6	168.2	2.96	105.7 -> 144.4 (2.87)

Startup rejections checked on this branch: an FTW checkpoint (Qwen3.6-27B-NVFP4 converted to FTW) and the Flash-Next GGUF without a draft path both stop at config time with the errors above. No Qwen3.5-family GGUF was at hand, so that rejection was checked by calling MTPProposer.adjust_config with the Qwen3_5GGUFForCausalLM registry entry ("MTP is not supported for Qwen3_5GGUFForCausalLM yet"). The other config-time checks were run through the engine's config resolution on local checkpoints, on CPU: a Strata pack as the draft of Qwen3.6-27B, the Flash-Next GGUF with the Flash-Next HF folder as its draft, Qwen3-30B-A3B (no mtp.*) as the draft of Qwen3.6-27B, and a draft path that does not exist are refused; Qwen3.6-35B-A3B, the Flash-Next NVFP4 checkpoint (alone and with the Strata pack) and the Flash-Next GGUF with the Strata pack pass. The expert-format error was checked against the mtp.* expert names of local checkpoints: the bf16 and block-fp8 Qwen3.5-35B-A3B and Qwen3.6-35B-A3B checkpoints, three NVFP4 releases of Qwen3.6-35B-A3B and the Flash-Next NVFP4 checkpoint pass, nvidia/Qwen3.8-Flash-Next-FP4-Local (per-expert NVFP4) is refused.

Known gaps

• MTP on and off greedy outputs part early (tokens 27, 559 and 102 above). The contract test shows the target side reproduces plain decoding on tiny models except at near ties; whether these real-model divergences are only near ties is not checked.
• The draft KV budget prices the draft's paged KV per token only. For the Flash-Next QSA draft pool it leaves out the 12 B/token mrope position table (2112 of 2124 B/token budgeted) and the per-request ring and scratch rows (11520 B at max_running_req=4). The chain graphs' slot-index buffers ((k-1) x max batch x page-table width x 4 B) and the draft extend graph's capture buffers are not budgeted either; all of these come out of the memory_ratio headroom.
• Qwen3.5-family GGUF targets and FTW draft sources are rejected; their MTP weights need loaders of their own. MTP routed experts stored as per-expert NVFP4 or per-tensor FP8 are not read.
• Run end to end: bf16 Qwen3.6 checkpoints, the Flash-Next NVFP4 checkpoint, and one Flash-Next GGUF with one Strata pack. The block-fp8 Qwen3.5 MTP head is covered only by the reader test; NVFP4 Qwen3.5 checkpoints were not tried.
• There is no reader test for the Flash-Next MTP head (checkpoint or Strata pack), and the qwen3_5 reader test does not cover the NVFP4 releases (whose MTP head is bf16); only the end-to-end runs cover them.
• No test exercises MTPProposer itself (the contract test drives a stub proposer): chunked against one-shot prefill, chain-graph against eager drafting, the post-verify relay over padded rows, and detach/attach. A bug there would lower the accept rate, not the output, since the verify is exact.
• TP > 1 and the tool-call anchor snapshot under speculative decoding stay as in feat(spec): verify draft tokens on the target side #836.