[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 draft path that does not exist, unless it has the form of a Hugging Face repo id (name or org/name).
• --dummy-weight ("MTP reads its draft from the checkpoint, so it cannot run with --dummy-weight").
• 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), or that stores them stacked without both gate_up_proj and down_proj, stops at draft load with an error naming the 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
• fix(spec): refuse a missing relative MTP draft path before the target loads
• fix(moe): read MTP head experts on the host and require both stacked tensors
• refactor(spec): tell the final prefill chunk from prompt_len instead of the scheduler's ChunkedReq
• refactor(spec): let each MTP draft family check its own weight source
• fix(qwen4_exp): size the disk PLE eager staging for decode and verify rows
• perf(kernel): read the KV-index gather's chunk count from the grid, not a constexpr
• refactor(moe): let each family name its MTP expert tensors
• refactor(spec): force the MTP draft's MoE resident in the proposer and check it after load
• fix(spec): refuse MTP with dummy weights
• refactor(spec): move HiddenRelay under spec/mtp, its only user
• refactor(spec): drop MTPDraft.config, fold its embed hooks into embed_tokens
• refactor(models): note FeatureSink's graph binding, drop its runtime_checkable
• fix(models): read quantized MTP fc projections under the quant config

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: attention_backend, draft_config, relay_width, check_source, load, embed_tokens, forward_hidden and head_input. Each family checks its own draft source in check_source; MTPProposer.adjust_config keeps the checks that name no family (the registry entry, --dummy-weight, the path, FTW).
• This PR adds no import between model families. Each family names its MTP expert tensors and passes a locate function to the shared reader moe/expert_pieces.iter_extra_layer_expert_pieces, which reads them on the host.
• The MTP input projections (mtp.fc; mtp.fc_embedding and mtp.fc_hidden on Flash-Next) go through the family's dense reader like the target's linears, so a checkpoint that quantizes them is read under its quantization config.
• The draft layer is built from the target's own layer classes and shares the target's embedding and LM head. The proposer builds it with with_resident_moe (moe_strategy="fused", decode_target="gpu"), so its experts are resident on the GPU and never enter the target's expert cache or offload, whatever the target's strategy; after load it refuses a draft whose MoE layers are offloaded or shared with the target.
• 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 and spec/mtp/relay.py; nothing method-specific goes on the batch. A prefill chunk is the final one when it ends at the prompt length; spec/ imports nothing from the scheduler. Host-to-device index uploads are pinned and non-blocking.
• A captured verify graph replays the FeatureSink it was captured with, so the MTP verify sink writes into a buffer the proposer owns; the FeatureSink docstring now states this for every draft method.
• The Triton chain graphs gather each step's KV indices with gather_kv_indices, which reads its chunk count from the launch grid: over 7 grid shapes a reference check compiled 1 kernel variant instead of 6, with the same output.

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. The eager staging holds max(max_extend_tokens, max_running_req * (k + 1)) rows, since an eager verify can stage more rows than a prefill chunk.

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; the draft's MoE stays resident though the target offloads its experts; 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_mtp_head_fills_the_draft_and_its_experts	with_resident_moe returns the target's config unchanged	2 failed: the draft builds an OffloadMoELayer
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) and the with_resident_moe row on e5ac81f; 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 1e3f418, on main c2c3e82.
• The full runs below used e02a5890: the first seven 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 8d34d22: 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).
• Rerun on e5ac81f: tests/spec, tests/models/qwen4_exp and tests/models/test_qwen3_5_moe_weight.py (187 passed, 154 skipped). The MTP-on greedy runs of Qwen3.6-35B-A3B and the Flash-Next NVFP4 checkpoint are token for token identical to the ones on e02a5890 in the tables below, and the Flash-Next GGUF with the Strata pack to its run on 8815c0b (846 tokens, answer 70, accept length 3.44). The same greedy run with MTP also completed on the block-fp8 Qwen/Qwen3.6-35B-A3B-FP8 (16384 tokens, the limit; accept length 3.51), on Qwen3.6-35B-A3B with its experts offloaded (moe_strategy="offload", automatic MoE cache; identical to the fused run), and on the Flash-Next NVFP4 checkpoint without CUDA graphs (identical to the run with graphs).
• 1e3f418 changes only the MTP reader. On 12 local checkpoints (Qwen3.5-35B-A3B-FP8; Qwen3.6-35B-A3B and Qwen3.6-27B in bf16, block-fp8 and ModelOpt NVFP4; RedHatAI's compressed-tensors NVFP4 Qwen3.6-35B-A3B; Qwen3.8-27B in bf16 and block-fp8; the Flash-Next NVFP4 and FP4-Local checkpoints) it emits the same tensors, byte for byte, as e5ac81f; all of them store the fc projections in bf16. On synthetic checkpoints that quantize them (HF block-fp8, compressed-tensors NVFP4 and ModelOpt FP8 for qwen3_5; compressed-tensors block-fp8 and ModelOpt NVFP4 for Flash-Next), its output fills the draft's quantized fc parameters; on e5ac81f every one of these but HF block-fp8 emits names or shapes the draft does not have. The unit tests above rerun on 1e3f418 with the same result (187 passed, 154 skipped).

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. Rerun through the config resolution on e5ac81f, the FTW, Flash-Next GGUF and CPU cases above give the same results, --dummy-weight is refused, and of three draft paths that do not exist, ./nonexistent-draft and models/x/nonexistent-draft are refused while nonexistent-org/nonexistent-draft passes as a repo id.

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 block-fp8 Qwen3.6-35B-A3B-FP8 (greedy only), the Flash-Next NVFP4 checkpoint, and one Flash-Next GGUF with one Strata pack. NVFP4 Qwen3.5 and Qwen3.6 checkpoints were not run.
• 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 checkpoint at hand quantizes the MTP fc projections, so that path is checked only on synthetic checkpoints and has not run end to end.
• 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.