From b062a991097695cd4c602cbe3b862f0675856045 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 26 Oct 2026 07:12:07 +0000
Subject: [PATCH 01/20] feat(spec): add the MTP proposer for the qwen3_5 family

---
 python/freetoken/kernel/triton/attention.py   |  42 +++
 python/freetoken/models/qwen3_5_moe/mtp.py    | 103 +++++
 python/freetoken/models/qwen3_5_moe/weight.py |  30 ++
 python/freetoken/models/register.py           |   8 +
 python/freetoken/moe/expert_pieces.py         |  37 ++
 python/freetoken/server/args.py               |   2 +-
 python/freetoken/spec/__init__.py             |   1 +
 python/freetoken/spec/mtp/__init__.py         |   0
 python/freetoken/spec/mtp/chain_graph.py      | 106 ++++++
 python/freetoken/spec/mtp/proposer.py         | 355 ++++++++++++++++++
 python/freetoken/spec/relay.py                |  14 +
 11 files changed, 697 insertions(+), 1 deletion(-)
 create mode 100644 python/freetoken/models/qwen3_5_moe/mtp.py
 create mode 100644 python/freetoken/spec/mtp/__init__.py
 create mode 100644 python/freetoken/spec/mtp/chain_graph.py
 create mode 100644 python/freetoken/spec/mtp/proposer.py
 create mode 100644 python/freetoken/spec/relay.py

diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index e49346dd5d..e56a9184a5 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -1058,3 +1058,45 @@ def paged_attention(
         num_stages=2,
     )
     return o
+
+
+@triton.jit
+def _gather_kv_indices_kernel(
+    page_table,
+    rows,
+    indptr,
+    out,
+    pt_stride,
+    indptr_stride,
+    out_stride,
+    NUM_CHUNKS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    batch_id = tl.program_id(1)
+    indptr_row = indptr + tl.program_id(0) * indptr_stride
+    start = tl.load(indptr_row + batch_id)
+    n = tl.load(indptr_row + batch_id + 1) - start
+    src = page_table + tl.load(rows + batch_id).to(tl.int64) * pt_stride
+    dst = out + tl.program_id(0).to(tl.int64) * out_stride + start
+    for off in range(tl.program_id(2) * BLOCK, n, NUM_CHUNKS * BLOCK):
+        cols = off + tl.arange(0, BLOCK)
+        mask = cols < n
+        tl.store(dst + cols, tl.load(src + cols, mask=mask), mask=mask)
+
+
+def gather_kv_indices(page_table: torch.Tensor, rows: torch.Tensor, indptr: torch.Tensor, out: torch.Tensor) -> None:
+    """Flat KV slot lists of several decode batches, built on the device: ``out[s, indptr[s, i] + j] = page_table[rows[i], j]`` for ``j < indptr[s, i + 1] - indptr[s, i]``."""
+    steps, batch = indptr.shape[0], indptr.shape[1] - 1
+    # a fixed grid keeps the launch graph-capturable; long rows loop over it
+    chunks = max(1, triton.cdiv(2 * _num_sms(page_table.device.index), steps * batch))
+    _gather_kv_indices_kernel[(steps, batch, chunks)](
+        page_table,
+        rows,
+        indptr,
+        out,
+        page_table.stride(0),
+        indptr.stride(0),
+        out.stride(0),
+        NUM_CHUNKS=chunks,
+        BLOCK=1024,
+    )
diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
new file mode 100644
index 0000000000..4799751f1f
--- /dev/null
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -0,0 +1,103 @@
+"""Qwen3.5-family MTP draft: the checkpoint's own ``mtp.*`` next-token layer, built from the target's layer classes."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+
+import torch
+
+from freetoken.layers import GemmaPlusOneRMSNorm, GemmaPlusOneRMSNormFused, LinearReplicated, OPList
+from freetoken.layers.quantization import finalize_quant
+from freetoken.layers.base import BaseOP
+from freetoken.models.blocks import embed_input_ids
+from freetoken.models.config import FullAttentionGroupConfig
+from freetoken.utils import torch_dtype
+
+from .model import Qwen3_5DecoderLayer
+from .weight import iter_mtp_weights
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+
+
+def draft_config(config: ModelConfig) -> ModelConfig:
+    """One full-attention layer with the target's geometry; its MoE is resident on the GPU whatever the target's strategy."""
+    groups = [g for g in config.attention_groups if isinstance(g, FullAttentionGroupConfig)]
+    group = groups[0] if groups else config.default_full_attention_group()
+    return dataclasses.replace(
+        config,
+        num_layers=1,
+        attention_groups=(dataclasses.replace(group, layer_ids=(0,)),),
+        moe_strategy="fused",
+        decode_target="gpu",
+    )
+
+
+class Qwen3_5MTPDraft(BaseOP):
+    """``h' = layer(fc(cat(norm_e(embed(token)), norm_h(h))))``: the next token's hidden from a token and the hidden before it."""
+
+    # the draft MTPProposer picks from the target's backends
+    attention_backend = None
+    draft_config = staticmethod(draft_config)
+
+    @staticmethod
+    def relay_width(config: ModelConfig) -> int:
+        return config.hidden_size
+
+    def __init__(self, config: ModelConfig, target_model) -> None:
+        self.config = draft_config(config)
+        hidden = config.hidden_size
+        self.fc = LinearReplicated(2 * hidden, hidden, has_bias=False, quant_config=config.quant, prefix="mtp.fc")
+        self.pre_fc_norm_embedding = GemmaPlusOneRMSNorm(hidden, eps=config.rms_norm_eps)
+        self.pre_fc_norm_hidden = GemmaPlusOneRMSNorm(hidden, eps=config.rms_norm_eps)
+        self.layers = OPList([Qwen3_5DecoderLayer(self.config, 0, prefix="mtp.layers.0")])
+        self.norm = GemmaPlusOneRMSNormFused(hidden, eps=config.rms_norm_eps)
+        # the target's own module; an underscore name keeps it out of this module's state dict
+        self._embed_tokens = target_model.model.embed_tokens
+
+    @classmethod
+    def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> Qwen3_5MTPDraft:
+        with torch.device("meta"), torch_dtype(dtype):
+            draft = cls(config, target_model)
+        expected = draft.state_dict()
+        loaded = {}
+        for name, tensor in iter_mtp_weights(model_path, device):
+            if name not in expected:
+                raise ValueError(f"{model_path}: unexpected MTP tensor {name!r}")
+            loaded[name] = tensor.to(device=device, dtype=expected[name].dtype)
+        if missing := sorted(set(expected) - set(loaded)):
+            raise ValueError(f"{model_path}: MTP tensors missing: {missing}")
+        draft.load_state_dict(loaded)
+        if draft.config.moe_enabled:
+            from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks
+            from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+
+            # packed by the layer's own expert method, as the target's resident layers are
+            moe = draft.layers.op_list[0].mlp.experts
+            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts", device)
+            banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
+            attach_resident_banks([moe], banks)
+        # not the root: it reaches the target's embedding, which the target finalized already
+        finalize_quant(draft.fc)
+        finalize_quant(draft.layers)
+        return draft
+
+    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self._embed_tokens.forward(input_ids)
+
+    def embed_batch(self, batch) -> torch.Tensor:
+        """The target batch's token embeddings with its multimodal rows merged in, as the target embeds them."""
+        return embed_input_ids(self._embed_tokens, batch.input_ids, batch)
+
+    def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor:
+        """Draft hidden of each row from the embedding of its token and the hidden of the position before it."""
+        x = self.fc.forward(torch.cat([self.pre_fc_norm_embedding.forward(embeds), self.pre_fc_norm_hidden.forward(prev_hidden)], dim=-1))
+        residual: torch.Tensor | None = None
+        for layer in self.layers.op_list:
+            x, residual = layer.forward(x, residual)
+        x, _ = self.norm.forward(x, residual)
+        return x
+
+    def head_input(self, hidden: torch.Tensor) -> torch.Tensor:
+        return hidden
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index 947e1b4833..77883c57dc 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -271,6 +271,35 @@ def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | N
         raise ValueError(f"checkpoint is missing tensors the quant config declares for {len(lines)} modules:\n  {shown}")
 
 
+# routed experts: expert_pieces.iter_extra_layer_expert_pieces
+_MTP_EXPERTS_PREFIX = "mtp.layers.0.mlp.experts."
+
+
+def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
+    """The ``mtp.*`` head but its routed experts, named as Qwen3_5MTPDraft names it, projections fused like the target's."""
+    hf_config = cached_load_hf_config(model_path)
+    reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0]))
+    seen = False
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+            for raw_name in f.keys():
+                if not raw_name.startswith("mtp.") or raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
+                    continue
+                seen = True
+                if raw_name.startswith(_MTP_EXPERTS_PREFIX):
+                    continue
+                tensor = f.get_tensor(raw_name)
+                emitted = reader.add(raw_name, tensor)
+                if emitted is None:
+                    yield raw_name[len("mtp."):], tensor
+                else:
+                    yield from ((n.removeprefix("mtp."), t) for n, t in emitted)
+    if not seen:
+        raise ValueError(f"{model_path} has no mtp.* weights")
+    if reader.pending:
+        raise ValueError(f"{model_path}: incomplete MTP projections: {reader.missing()}")
+
+
 def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
     """The vision tower alone, named as iter_weights names it."""
     for file in iter_weight_files(model_path):
@@ -398,6 +427,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec:
 
 
 __all__ = [
+    "iter_mtp_weights",
     "iter_weights",
     "iter_weights_parallel",
     "iter_expert_pieces",
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 3ed67c99bf..1a04a43934 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -32,6 +32,8 @@ class ModelSpec:
     # "module:Class" turning the checkpoint's media into items; None: the family takes no multimodal input
     mm_processor: str | None = None
     encoders: tuple[EncoderSpec, ...] = ()
+    # speculative method -> "module:Class" of the family's draft for it
+    spec_drafts: tuple[tuple[str, str], ...] = ()
 
 
 # Multimodal wrappers store the text tower under model.language_model.
@@ -40,6 +42,7 @@ class ModelSpec:
 _QWEN_VL_PROCESSOR = "freetoken.mm.processors.qwen_vl:QwenVLMMProcessor"
 # TODO: video once the processor samples frames; the tower already takes grid_thw with t > 1
 _QWEN_VL_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
+_QWEN3_5_DRAFTS = (("mtp", "freetoken.models.qwen3_5_moe.mtp:Qwen3_5MTPDraft"),)
 # Fused projections split back into their HF leaves so the quantization_config lookup sees the stored names.
 _DENSE_PACKED = (
     ("qkv_proj", ("q_proj", "k_proj", "v_proj")),
@@ -288,6 +291,7 @@ class ModelSpec:
         encoders=_QWEN_VL_ENCODERS,
         packed_modules_mapping=_QWEN3_5_PACKED,
         unquantized_modules=_QWEN3_5_UNQUANTIZED,
+        spec_drafts=_QWEN3_5_DRAFTS,
     ),
     # text-only Qwen3.5 releases: text_config is the top-level config and there is no tower
     "Qwen3_5MoeForCausalLM": ModelSpec(
@@ -295,14 +299,17 @@ class ModelSpec:
         "Qwen3_5MoeForCausalLM",
         packed_modules_mapping=_QWEN3_5_PACKED,
         unquantized_modules=_QWEN3_5_UNQUANTIZED,
+        spec_drafts=_QWEN3_5_DRAFTS,
     ),
     "Qwen3_5ForCausalLM": ModelSpec(
         "freetoken.models.qwen3_5_moe",
         "Qwen3_5ForCausalLM",
         packed_modules_mapping=_QWEN3_5_PACKED,
         unquantized_modules=_QWEN3_5_UNQUANTIZED,
+        spec_drafts=_QWEN3_5_DRAFTS,
     ),
     # llama.cpp qwen35 / qwen35moe GGUF: the same model classes with every tensor in its packed ggml type; the vision tower comes from an mmproj file and runs in bf16
+    # no MTP draft yet: the nextn blocks and mtp-*.gguf files need a loader of their own
     "Qwen3_5GGUFForCausalLM": ModelSpec(
         "freetoken.models.qwen3_5_moe",
         "Qwen3_5ForConditionalGeneration",
@@ -359,6 +366,7 @@ class ModelSpec:
         encoders=_QWEN_VL_ENCODERS,
         packed_modules_mapping=_QWEN3_5_PACKED,
         unquantized_modules=_QWEN3_5_UNQUANTIZED,
+        spec_drafts=_QWEN3_5_DRAFTS,
     ),
     # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in
     # text_config, weights under model.language_model.); the windowed ViT under
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index 74c4a41b0f..e55a9d0ca0 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import re
 from typing import Callable, Collection, Iterable, Iterator
 
 import torch
@@ -169,10 +170,46 @@ def per_expert_pieces(
         raise ValueError(f"Incomplete expert tensors in checkpoint (layer, expert): {missing}")
 
 
+def iter_extra_layer_expert_pieces(model_path: str, prefix: str, device: torch.device) -> Iterator[Piece]:
+    """The routed experts stored under ``prefix``, one MoE layer outside the model's stack (an MTP head), as pieces of bank layer 0: stacked ``gate_up``/``down``, or ``gate``/``up``/``down`` per expert with ``_scale`` companions in a block-fp8 checkpoint."""
+    import safetensors
+
+    from freetoken.models.loader import iter_weight_files
+
+    per_expert_re = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.(gate|up|down)_proj\.(weight|weight_scale_inv)$")
+    stacked_re = re.compile(rf"^{re.escape(prefix)}\.(gate_up|down)_proj$")
+    names: dict[str, str] = {}
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device="cpu") as f:
+            names.update((name, file) for name in f.keys() if per_expert_re.match(name) or stacked_re.match(name))
+    if not names:
+        raise ValueError(f"{model_path} has no routed experts under {prefix}")
+    stacked = {stacked_re.match(n)[1]: f for n, f in names.items() if stacked_re.match(n)}
+    if stacked:
+        piece = {}
+        for role, file in stacked.items():
+            with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+                piece[role] = f.get_tensor(f"{prefix}.{role}_proj")
+        return iter([(0, 0, piece["down"].shape[0], piece)])
+
+    def locate(name: str) -> tuple[int, int, str]:
+        m = per_expert_re.match(name)
+        return 0, int(m[1]), m[2] + ("_scale" if m[3] == "weight_scale_inv" else "")
+
+    def tensors():
+        for file in dict.fromkeys(names.values()):
+            with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+                yield from ((name, f.get_tensor(name)) for name, where in names.items() if where == file)
+
+    scaled = any(name.endswith("_scale_inv") for name in names)
+    return per_expert_pieces(tensors(), locate, tensors_per_expert=6 if scaled else 3)
+
+
 __all__ = [
     "Piece",
     "bank_layer_of",
     "iter_expert_pieces",
+    "iter_extra_layer_expert_pieces",
     "num_moe_layers",
     "packed_expert_source_info",
     "per_expert_pieces",
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index 8fbdf748c2..a3293d14a4 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -924,7 +924,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "--spec-algo",
         type=str,
         default=ServerArgs.speculative_algorithm,
-        help="Speculative decoding with this draft method. Off by default.",
+        help="Speculative decoding with this draft method (mtp: the checkpoint's own MTP layers). Off by default.",
     )
 
     parser.add_argument(
diff --git a/python/freetoken/spec/__init__.py b/python/freetoken/spec/__init__.py
index 97e1673ad6..9b0077989f 100644
--- a/python/freetoken/spec/__init__.py
+++ b/python/freetoken/spec/__init__.py
@@ -4,5 +4,6 @@
 
 # method name -> "module:Class" of its Proposer, loaded on use so this package imports no method's code
 PROPOSERS = Registry[str]("Speculative proposer")
+PROPOSERS.register("mtp")("freetoken.spec.mtp.proposer:MTPProposer")
 
 __all__ = ["PROPOSERS"]
diff --git a/python/freetoken/spec/mtp/__init__.py b/python/freetoken/spec/mtp/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/python/freetoken/spec/mtp/chain_graph.py b/python/freetoken/spec/mtp/chain_graph.py
new file mode 100644
index 0000000000..0634c1b07d
--- /dev/null
+++ b/python/freetoken/spec/mtp/chain_graph.py
@@ -0,0 +1,106 @@
+"""MTP chain steps as one CUDA graph per batch size: the k-1 draft decode steps on the triton backend, with each step's metadata built on the device."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, List
+
+import torch
+
+from freetoken.attention.triton import TritonMetadata
+from freetoken.core import Batch
+from freetoken.kernel.triton.attention import gather_kv_indices
+
+if TYPE_CHECKING:
+    from freetoken.core import Req
+
+    from .proposer import MTPProposer
+
+
+class ChainGraphs:
+    """Replays ``MTPProposer.propose``: the inputs are each request's table row, entry position and mrope delta, everything else derives from them in the graph."""
+
+    def __init__(self, proposer: MTPProposer, bs_list: List[int], pool) -> None:
+        self.p = proposer
+        ctx = proposer.ctx
+        device, steps, max_bs = ctx.device, proposer.k - 1, max(bs_list)
+        self.dummy_row = ctx.dummy_req.table_idx
+        # rows: table row, entry position, mrope delta
+        self.inputs = torch.zeros(3, max_bs, dtype=torch.int64, device=device)
+        self.drafts = torch.zeros(max_bs, proposer.k, dtype=torch.int64, device=device)
+        self.positions = torch.zeros(steps, max_bs, dtype=torch.int32, device=device)
+        self.mrope = torch.zeros(steps, 3, max_bs, dtype=torch.int32, device=device) if proposer.mrope else None
+        self.out_loc = torch.zeros(steps, max_bs, dtype=torch.int32, device=device)
+        self.indptr = torch.zeros(steps, max_bs + 1, dtype=torch.int32, device=device)
+        self.indices = torch.zeros(steps, max_bs * ctx.page_table.shape[1], dtype=torch.int32, device=device)
+        self.arange = torch.arange(max_bs + 1, dtype=torch.int32, device=device)
+        self.graphs: Dict[int, torch.cuda.CUDAGraph] = {}
+        # the metadata of these batches holds the decode scratch the graphs read
+        self._batches: Dict[int, List[Batch]] = {}
+        for bs in sorted(bs_list, reverse=True):
+            self.inputs[0, :bs] = self.dummy_row
+            batches = [self._batch(bs, s) for s in range(steps)]
+            # outside the graph: triton compiles and the decode scratch is allocated here
+            self._run(bs, batches)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool, stream=ctx.stream):
+                self._run(bs, batches)
+            if pool is None:
+                pool = graph.pool()
+            self.graphs[bs] = graph
+            self._batches[bs] = batches
+
+    def _batch(self, bs: int, step: int) -> Batch:
+        batch = Batch(reqs=[None] * bs, phase="decode")
+        batch.padded_reqs = batch.reqs
+        batch.positions = self.positions[step, :bs]
+        if self.mrope is not None:
+            batch.mrope_positions = self.mrope[step, :, :bs]
+        batch.out_loc = self.out_loc[step, :bs]
+        batch.attn_metadata = TritonMetadata(
+            cu_seqlens_q_gpu=self.arange[: bs + 1],
+            indptr=self.indptr[step, : bs + 1],
+            indices=self.indices[step],
+            q_to_req=self.arange[:bs],
+            q_positions=self.positions[step, :bs],
+            is_decode=True,
+            prefix_lens=self.positions[step, :bs],
+            max_q_len=1,
+        )
+        return batch
+
+    def _run(self, bs: int, batches: List[Batch]) -> None:
+        p = self.p
+        page_table = p.ctx.page_table
+        rows, entry, delta = self.inputs[:, :bs]
+        for step in range(len(batches)):
+            # step s drafts position c+s over keys [0, c+s]
+            pos = entry + step
+            self.positions[step, :bs] = pos
+            if self.mrope is not None:
+                self.mrope[step, :, :bs] = (pos + delta).unsqueeze(0)
+            self.out_loc[step, :bs] = page_table[rows, pos]
+            self.indptr[step, 1 : bs + 1] = torch.cumsum(pos + 1, 0)
+        gather_kv_indices(page_table, rows, self.indptr[:, : bs + 1], self.indices)
+        token, hidden = p.relay.d1[rows], p.relay.hidden[rows]
+        self.drafts[:bs, 0] = token
+        for step, batch in enumerate(batches):
+            hidden = p._forward(batch, p.chain_backend, p.draft.embed(token), hidden)
+            token = p._draft_token(hidden)
+            self.drafts[:bs, step + 1] = token
+
+    def replay(self, reqs: List[Req], cs: List[int]) -> torch.Tensor | None:
+        """The drafts [bs, k], or None when no captured batch size fits."""
+        bs = len(reqs)
+        padded = next((b for b in sorted(self.graphs) if b >= bs), None)
+        if padded is None:
+            return None
+        pad = padded - bs
+        # padded rows must point at the dummy page: a stale real row would get draft KV written into it
+        host = torch.tensor(
+            [[r.table_idx for r in reqs] + [self.dummy_row] * pad, cs + [0] * pad,
+             [getattr(r, "mrope_delta", 0) for r in reqs] + [0] * pad],
+            dtype=torch.int64, pin_memory=True,
+        )
+        self.inputs[:, :padded].copy_(host, non_blocking=True)
+        self.graphs[padded].replay()
+        return self.drafts[:bs].clone()
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
new file mode 100644
index 0000000000..5eee3a2269
--- /dev/null
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -0,0 +1,355 @@
+"""MTPProposer: drafts with the checkpoint's own MTP layer, one draft token per chained forward."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar, List, Protocol
+
+import torch
+
+from freetoken.core import Batch, get_global_ctx
+from freetoken.engine.sample import make_device_tensor
+from freetoken.spec.proposer import DraftContext, Proposer
+from freetoken.spec.relay import HiddenRelay
+
+if TYPE_CHECKING:
+    from freetoken.core import Req
+    from freetoken.models.config import ModelConfig
+
+
+class MTPDraft(Protocol):
+    """The class a family registers as ``ModelSpec.spec_drafts["mtp"]``."""
+
+    # the draft's own config, from draft_config
+    config: ModelConfig
+    # None: the proposer picks from the target's backends
+    attention_backend: ClassVar[str | None]
+
+    @staticmethod
+    def draft_config(config: ModelConfig) -> ModelConfig: ...
+
+    @staticmethod
+    def relay_width(config: ModelConfig) -> int:
+        """Width of the hidden a draft row reads and relays to the next step."""
+        ...
+
+    @classmethod
+    def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> MTPDraft: ...
+
+    def embed(self, input_ids: torch.Tensor) -> torch.Tensor: ...
+
+    def embed_batch(self, batch: Batch) -> torch.Tensor: ...
+
+    def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor: ...
+
+    def head_input(self, hidden: torch.Tensor) -> torch.Tensor:
+        """What the target's LM head reads from a draft hidden."""
+        ...
+
+
+@dataclass
+class _DraftReq:
+    """The request fields the attention backends and rope read, for a draft batch over positions [cached_len, device_len)."""
+
+    cached_len: int
+    device_len: int
+    table_idx: int
+    mrope_delta: int = 0
+    mrope_positions_full: torch.Tensor | None = None
+
+    @property
+    def extend_len(self) -> int:
+        return self.device_len - self.cached_len
+
+
+class _PrefillSink:
+    """Every prefill row's final hidden, and each request's [cached_len, device_len) when the batch was formed."""
+
+    wanted_layers = frozenset()
+
+    def __init__(self, spans: List[tuple[int, int]]) -> None:
+        self.spans = spans
+        self.hidden: torch.Tensor | None = None
+
+    def final_hidden(self, x: torch.Tensor) -> None:
+        self.hidden = x
+
+    def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None:
+        pass
+
+
+class _VerifySink:
+    """Copies the verify rows' final hidden into a persistent buffer, so a captured verify graph writes it too."""
+
+    wanted_layers = frozenset()
+
+    def __init__(self, buffer: torch.Tensor) -> None:
+        self.buffer = buffer
+
+    def final_hidden(self, x: torch.Tensor) -> None:
+        self.buffer[: x.shape[0]].copy_(x)
+
+    def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None:
+        pass
+
+
+class MTPProposer(Proposer):
+    default_k = 3
+    max_k = 8
+    radix_lookahead = True
+
+    @classmethod
+    def adjust_config(cls, config, override) -> None:
+        from freetoken.checkpoint.ftw import is_ftw_checkpoint
+
+        super().adjust_config(config, override)
+        if "mtp" not in dict(config.model_spec.spec_drafts):
+            raise ValueError(f"MTP is not supported for {config.hf_config.architectures[0]} yet")
+        # FTW keeps only what the target's iter_weights yields, which drops mtp.*
+        if is_ftw_checkpoint(config.speculative_draft_model_path or config.model_path):
+            raise ValueError("MTP cannot load its draft from an FTW checkpoint yet")
+
+    @classmethod
+    def _draft_class(cls, config) -> type[MTPDraft]:
+        from freetoken.models.register import _load_attr
+
+        return _load_attr(*dict(config.model_spec.spec_drafts)["mtp"].split(":"))
+
+    @classmethod
+    def budget(cls, config) -> tuple[int, int]:
+        from freetoken.kvcache.base import spec_kv_bytes_per_token
+
+        draft_cls = cls._draft_class(config)
+        relay_bytes = draft_cls.relay_width(config.model_config) * config.dtype.itemsize
+        rows = config.max_running_req + 1
+        # relay hidden and carry, and the verify hiddens
+        fixed = rows * (config.speculative_num_steps + 3) * relay_bytes
+        groups = draft_cls.draft_config(config.model_config).kv_cache_group_specs()
+        return fixed, sum(spec_kv_bytes_per_token(group, config) for group in groups)
+
+    def __init__(self, config, target_model) -> None:
+        super().__init__(config, target_model)
+        self.device = torch.device("cuda", torch.cuda.current_device())
+        self.dtype = config.dtype
+        path = config.speculative_draft_model_path or config.model_path
+        self.draft = self._draft_class(config).load(config.model_config, target_model, path, self.device, config.dtype)
+        self.lm_head = target_model.lm_head
+        self.mrope = config.model_config.model_is_mrope
+        self._prefill_state: dict[int, tuple[int, int]] = {}  # table row -> (uid, first position not drafted yet)
+        self._prefill_sink: _PrefillSink | None = None
+        self.kv = None
+
+    # ------------------------------------------------------------------ lifecycle
+
+    def attach(self, ctx: DraftContext) -> None:
+        from freetoken.attention import create_attention_backend
+        from freetoken.kvcache import create_kvcache_pool
+
+        super().attach(ctx)
+        cfg = self.draft.config
+        rows = self.config.max_running_req + 1
+        # the draft's own layers on the target's slot ids: same page size, same pages plus the dummy page padded rows write to
+        # a rebuild that keeps the target KV pool keeps this one, which the prefix cache still indexes
+        if self.kv is None:
+            self.kv = create_kvcache_pool(
+                model_config=cfg, num_pages=ctx.num_pages + 1, page_size=self.config.page_size, dtype=self.dtype,
+                device=ctx.device, num_req_slots=rows,
+            )
+        width = self.draft.relay_width(self.config.model_config)
+        self.relay = HiddenRelay(rows, width, self.dtype, ctx.device)
+        self.verify_hidden = torch.zeros(rows * (self.k + 1), width, dtype=self.dtype, device=ctx.device)
+        # verify row of each request's last accepted position; padded rows keep stale but in-range values
+        self.picked_rows = torch.arange(rows, dtype=torch.int64, device=ctx.device) * (self.k + 1)
+        self.extend_graphs = self.chain_graphs = None
+        parts = [p.strip() for p in self.config.attention_backend.split(",")]
+        fixed = self.draft.attention_backend
+        # backends bind the context's KV pool when they are built
+        with get_global_ctx().override(kv_cache=self.kv):
+            self.extend_backend = create_attention_backend(fixed or ("fa" if "fa" in parts else "triton"), cfg)
+            self.chain_backend = create_attention_backend(fixed or "triton", cfg)
+
+    def detach(self, *, keep_kv: bool) -> None:
+        if self.extend_graphs is not None:
+            self.extend_graphs.destroy_cuda_graphs()
+        self.extend_graphs = self.chain_graphs = None
+        if not keep_kv:
+            self.kv = None
+        self.relay = self.verify_hidden = self.picked_rows = None
+        self.extend_backend = self.chain_backend = None
+        self._prefill_state.clear()
+        super().detach(keep_kv=keep_kv)
+
+    def release(self, req: Req) -> None:
+        self._prefill_state.pop(req.table_idx, None)
+
+    # ------------------------------------------------------------------ draft forwards
+
+    def _draft_batch(self, reqs: List[_DraftReq], phase: str, rows_per_req: int | None, backend) -> Batch:
+        batch = Batch(reqs=reqs, phase=phase, rows_per_req=rows_per_req)
+        batch.padded_reqs = reqs
+        host = torch.cat([torch.arange(r.cached_len, r.device_len, dtype=torch.int32) for r in reqs])
+        batch.positions = host.pin_memory().to(self.device, non_blocking=True)
+        if self.mrope:
+            batch.mrope_positions = self._mrope_positions(reqs).pin_memory().to(self.device, non_blocking=True)
+        rows = torch.cat([torch.full((r.extend_len,), r.table_idx, dtype=torch.int64) for r in reqs])
+        rows = rows.pin_memory().to(self.device, non_blocking=True)
+        batch.out_loc = self.ctx.page_table[rows, batch.positions.to(torch.int64)]
+        backend.prepare_metadata(batch)
+        return batch
+
+    @staticmethod
+    def _mrope_positions(reqs: List[_DraftReq]) -> torch.Tensor:
+        """[3, n] rope rows of the draft positions, by the target's rule: a prompt position's precomputed column, else position + delta."""
+        cols = []
+        for r in reqs:
+            pos = torch.arange(r.cached_len, r.device_len, dtype=torch.int32)
+            col = (pos + r.mrope_delta).unsqueeze(0).expand(3, -1).clone()
+            full = r.mrope_positions_full
+            if full is not None:
+                inside = pos < full.shape[1]
+                col[:, inside] = full[:, pos[inside].long()]
+            cols.append(col)
+        return torch.cat(cols, dim=1)
+
+    def _forward(self, batch: Batch, backend, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor:
+        with get_global_ctx().forward_batch(batch, kv_cache=self.kv, attn_backend=backend, hide_target_state=True):
+            return self.draft.forward_hidden(embeds, prev_hidden)
+
+    def _draft_token(self, hidden: torch.Tensor) -> torch.Tensor:
+        return self.lm_head.logits(self.draft.head_input(hidden)).argmax(dim=-1)
+
+    def _view(self, req, cached_len: int, device_len: int) -> _DraftReq:
+        return _DraftReq(cached_len, device_len, req.table_idx, getattr(req, "mrope_delta", 0), getattr(req, "mrope_positions_full", None))
+
+    # ------------------------------------------------------------------ one step
+
+    def capture_graphs(self, bs_list: List[int], pool) -> None:
+        from freetoken.engine.graph import GraphRunner
+
+        from .chain_graph import ChainGraphs
+
+        if not bs_list:
+            return
+        ctx = self.ctx
+        self.extend_graphs = GraphRunner(
+            stream=ctx.stream,
+            device=ctx.device,
+            forward=self._extend_forward,
+            attn_backend=self.extend_backend,
+            cuda_graph_bs=bs_list,
+            cuda_graph_max_bs=max(bs_list),
+            free_memory=0,
+            # the table width, not the request limit: verify widens a request k rows past the limit
+            max_seq_len=ctx.page_table.shape[1],
+            dummy_req=ctx.dummy_req,
+            mrope=self.mrope,
+            rows_per_req=self.k + 1,
+            phase="verify",
+            ctx_override=dict(kv_cache=self.kv, attn_backend=self.extend_backend, hide_target_state=True),
+            pool=pool,
+        )
+        if self.k > 1:
+            self.chain_graphs = ChainGraphs(self, bs_list, self.extend_graphs.graph_pool)
+
+    def propose(self, reqs: List[Req], rows: torch.Tensor) -> torch.Tensor:
+        cs = [r.cached_len for r in reqs]
+        if self.chain_graphs is not None and (drafts := self.chain_graphs.replay(reqs, cs)) is not None:
+            return drafts
+        token = self.relay.d1[rows]
+        hidden = self.relay.hidden[rows]
+        drafts = [token]
+        for j in range(1, self.k):
+            # chain step j drafts at position C+j-1 with the previous step's draft hidden standing in for the target's
+            views = [self._view(r, c + j - 1, c + j) for r, c in zip(reqs, cs)]
+            batch = self._draft_batch(views, "decode", 1, self.chain_backend)
+            hidden = self._forward(batch, self.chain_backend, self.draft.embed(token), hidden)
+            token = self._draft_token(hidden)
+            drafts.append(token)
+        return torch.stack(drafts, dim=1)
+
+    def feature_sink(self, batch: Batch) -> _PrefillSink | _VerifySink:
+        if batch.phase == "verify":
+            return _VerifySink(self.verify_hidden)
+        self._prefill_sink = _PrefillSink([(r.cached_len, r.device_len) for r in batch.reqs])
+        return self._prefill_sink
+
+    def _extend_forward(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Draft hiddens of a verify-shaped batch; returns each request's row at its last accepted position and that row's argmax."""
+        input_ids = get_global_ctx().batch.input_ids
+        n = input_ids.shape[0]
+        hidden = self.draft.forward_hidden(self.draft.embed(input_ids), self.verify_hidden[:n])
+        picked = hidden[self.picked_rows[: n // (self.k + 1)]]
+        return picked, self._draft_token(picked)
+
+    def post_verify(self, batch: Batch, out: torch.Tensor, emitted: torch.Tensor) -> None:
+        """Redraft positions C..C+k from the target's verify hiddens, then relay each request's row at its last accepted position."""
+        k1 = self.k + 1
+        bs, padded = batch.size, batch.padded_size
+        tokens = torch.zeros(padded, k1, dtype=torch.int64, device=self.device)
+        tokens[:bs] = out
+        views = [self._view(r, r.cached_len, r.device_len) for r in batch.padded_reqs]
+        draft = Batch(reqs=views[:bs], phase="verify", rows_per_req=k1)
+        draft.padded_reqs = views
+        draft.input_ids = tokens.view(-1)
+        draft.positions = batch.positions
+        draft.mrope_positions = batch.mrope_positions
+        draft.out_loc = batch.out_loc
+        self.picked_rows[:bs] = torch.arange(bs, device=self.device) * k1 + emitted
+        self.extend_backend.prepare_metadata(draft)
+        graphs = self.extend_graphs
+        if graphs is not None and graphs.can_use_cuda_graph(draft) and padded in graphs.graph_map:
+            picked, d1 = graphs.replay(draft)
+        else:
+            with get_global_ctx().forward_batch(draft, kv_cache=self.kv, attn_backend=self.extend_backend, hide_target_state=True):
+                picked, d1 = self._extend_forward()
+            picked, d1 = picked[:bs], d1[:bs]
+        table = make_device_tensor([r.table_idx for r in batch.reqs], torch.int64, self.device)
+        self.relay.hidden[table] = picked
+        self.relay.d1[table] = d1
+
+    def on_prefill(self, batch: Batch, next_tokens_gpu: torch.Tensor) -> None:
+        """Draft the chunk's positions: row p pairs token p+1 with the target hidden at p; a final chunk adds the sampled token's row and relays it."""
+        from freetoken.scheduler.prefill import ChunkedReq
+
+        sink, self._prefill_sink = self._prefill_sink, None
+        hidden = sink.hidden
+        # image rows take the embeddings the target used
+        embeds = self.draft.embed_batch(batch)
+        views, emb, prev, finals = [], [], [], []
+        offset = drafted = 0
+        for i, (req, (s, e)) in enumerate(zip(batch.reqs, sink.spans)):
+            rows = hidden[offset : offset + e - s]
+            chunk_embeds = embeds[offset : offset + e - s]
+            offset += e - s
+            row = req.table_idx
+            state = self._prefill_state.get(row)
+            # a continuation chunk resumes where the last one stopped; a first chunk starts at its prefix hit, whose draft KV the tree already holds
+            frontier = state[1] if state is not None and state[0] == req.uid else s
+            final = not isinstance(req, ChunkedReq)
+            end = e if final else e - 1
+            if frontier < end:
+                # draft row p embeds the token at p+1
+                emb.append(chunk_embeds[frontier + 1 - s : min(end + 1, e) - s])
+                if final:
+                    emb.append(self.draft.embed(next_tokens_gpu[i : i + 1]))
+                if frontier == s - 1:
+                    # cloned: the carry of this row is overwritten below, before the draft forward reads it
+                    prev += [self.relay.carry[row : row + 1].clone(), rows[: end - s]]
+                else:
+                    prev.append(rows[frontier - s : end - s])
+                views.append(self._view(req, frontier, end))
+                drafted += end - frontier
+                if final:
+                    finals.append((drafted - 1, row))
+            self.relay.carry[row] = rows[-1]
+            self._prefill_state[row] = (req.uid, max(frontier, end))
+        if not views:
+            return
+        draft = self._draft_batch(views, "prefill", None, self.extend_backend)
+        out = self._forward(draft, self.extend_backend, torch.cat(emb), torch.cat(prev))
+        if finals:
+            idx = make_device_tensor([i for i, _ in finals], torch.int64, self.device)
+            table = make_device_tensor([r for _, r in finals], torch.int64, self.device)
+            picked = out[idx]
+            self.relay.hidden[table] = picked
+            self.relay.d1[table] = self._draft_token(picked)
diff --git a/python/freetoken/spec/relay.py b/python/freetoken/spec/relay.py
new file mode 100644
index 0000000000..ddbb101388
--- /dev/null
+++ b/python/freetoken/spec/relay.py
@@ -0,0 +1,14 @@
+"""Per-request draft state carried from one step to the next, by page-table row."""
+
+from __future__ import annotations
+
+import torch
+
+
+class HiddenRelay:
+    """For each table row: the draft hidden the next proposal starts from, the first draft token it predicts, and the target hidden that closes a prefill chunk."""
+
+    def __init__(self, rows: int, hidden_size: int, dtype: torch.dtype, device: torch.device) -> None:
+        self.hidden = torch.zeros(rows, hidden_size, dtype=dtype, device=device)
+        self.d1 = torch.zeros(rows, dtype=torch.int64, device=device)
+        self.carry = torch.zeros(rows, hidden_size, dtype=dtype, device=device)

From 407fd5e4182901d7976efaba876ec0dda880df2b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 26 Oct 2026 08:46:35 +0000
Subject: [PATCH 02/20] test(models): read the qwen3_5 MTP head of bf16 and
 block-fp8 checkpoints

---
 tests/models/test_qwen3_5_moe_weight.py | 70 +++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 4 deletions(-)

diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index 958d6daf66..2ed0ad9229 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -113,6 +113,34 @@ def _quantize(raw: dict[str, torch.Tensor], modules: list[str], quantize) -> Non
         raw.update({f"{module}.{suffix}": t for suffix, t in quantize(weight).items()})
 
 
+MTP = "mtp.layers.0"
+MTP_LINEARS = [f"{MTP}.self_attn.{p}_proj" for p in "qkvo"] + [f"{MTP}.mlp.shared_expert.{p}_proj" for p in ("gate", "up", "down")]
+# the released FP8 checkpoints keep these MTP modules bf16
+MTP_UNQUANTIZED = ["mtp.fc", "mtp.norm", "mtp.pre_fc_norm_embedding", "mtp.pre_fc_norm_hidden", f"{MTP}.mlp.gate", f"{MTP}.mlp.shared_expert_gate",
+                   f"{MTP}.input_layernorm", f"{MTP}.post_attention_layernorm", f"{MTP}.self_attn.q_norm", f"{MTP}.self_attn.k_norm"]
+
+
+def _mtp_head(raw: dict[str, torch.Tensor], quantize=None) -> None:
+    """The ``mtp.*`` head of a MoE checkpoint: stacked bf16 experts, or per-expert tensors with the linears under ``quantize``."""
+    raw.update({"mtp.fc.weight": _bf16(H, 2 * H), "mtp.norm.weight": _bf16(H), "mtp.pre_fc_norm_embedding.weight": _bf16(H),
+                "mtp.pre_fc_norm_hidden.weight": _bf16(H), f"{MTP}.input_layernorm.weight": _bf16(H),
+                f"{MTP}.post_attention_layernorm.weight": _bf16(H), f"{MTP}.self_attn.q_norm.weight": _bf16(AHD),
+                f"{MTP}.self_attn.k_norm.weight": _bf16(AHD), f"{MTP}.mlp.gate.weight": _bf16(E, H),
+                f"{MTP}.mlp.shared_expert_gate.weight": _bf16(1, H)})
+    for proj, shape in (("q", (2 * QH * AHD, H)), ("k", (KVH * AHD, H)), ("v", (KVH * AHD, H)), ("o", (H, QH * AHD))):
+        raw[f"{MTP}.self_attn.{proj}_proj.weight"] = _bf16(*shape)
+    for proj, shape in (("gate", (I, H)), ("up", (I, H)), ("down", (H, I))):
+        raw[f"{MTP}.mlp.shared_expert.{proj}_proj.weight"] = _bf16(*shape)
+    if quantize is None:
+        raw[f"{MTP}.mlp.experts.gate_up_proj"] = _bf16(E, 2 * MI, H)
+        raw[f"{MTP}.mlp.experts.down_proj"] = _bf16(E, H, MI)
+        return
+    _quantize(raw, MTP_LINEARS, quantize)
+    for expert in range(E):
+        for proj, shape in (("gate_proj", (MI, H)), ("up_proj", (MI, H)), ("down_proj", (H, MI))):
+            raw.update({f"{MTP}.mlp.experts.{expert}.{proj}.{suffix}": t for suffix, t in quantize(_bf16(*shape)).items()})
+
+
 def _experts(raw: dict[str, torch.Tensor], quantize=None) -> None:
     """Routed experts: stacked bf16 per layer, or per-expert tensors under ``quantize``."""
     for layer in (0, 1):
@@ -163,7 +191,7 @@ def _ct(config_groups: dict, ignore: list[str], fmt: str) -> dict:
 # Qwen/Qwen3.6-35B-A3B-FP8: 128x128 block-fp8 everywhere but the listed modules
 QWEN_FP8 = {
     "quant_method": "fp8", "activation_scheme": "dynamic", "fmt": "e4m3", "weight_block_size": [128, 128],
-    "modules_to_not_convert": ["lm_head", "model.embed_tokens", *GDN_BA, *ROUTERS,
+    "modules_to_not_convert": ["lm_head", "model.embed_tokens", *GDN_BA, *ROUTERS, *MTP_UNQUANTIZED,
                                *[f"{LM}.layers.{l}.{n}" for l in (0, 1) for n in ("input_layernorm", "post_attention_layernorm")]],
 }
 # nvidia/Qwen3.6-35B-A3B-NVFP4: per-tensor FP8 attention / GDN, NVFP4 shared expert, experts and lm_head
@@ -208,10 +236,12 @@ def _layout(name: str) -> tuple[bool, dict | None, dict[str, torch.Tensor]]:
     raw = _dense_bf16(moe)
     if name == "bf16":
         _experts(raw)
+        _mtp_head(raw)
         return moe, None, raw
     if name == "fp8_block":
         _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, lambda w: _fp8_block(w, ct=False))
         _experts(raw, lambda w: _fp8_block(w, ct=False))
+        _mtp_head(raw, lambda w: _fp8_block(w, ct=False))
         return moe, QWEN_FP8, raw
     if name.startswith("modelopt_mixed"):
         _quantize(raw, GDN_QKVZ_OUT + ATTN, _fp8_tensor)
@@ -294,8 +324,8 @@ def _load(folder: str, *, experts: bool = False, vision: bool = True) -> dict[st
     return {n: t.clone() for n, t in iter_weights(folder, torch.device("cpu"), include_moe_experts=experts, include_non_moe=True, include_vision=vision)}
 
 
-def _meta_state_dict(folder: str) -> dict[str, torch.Tensor]:
-    """State dict of the model the engine builds for ``folder`` (routed experts offloaded), on the meta device."""
+def _meta_state_dict(folder: str, build=None) -> dict[str, torch.Tensor]:
+    """State dict of the model the engine builds for ``folder`` (routed experts offloaded), or of ``build(model_config)``, on the meta device."""
     from freetoken.engine.config import EngineConfig
     from freetoken.engine.engine import _decode_target
     from freetoken.layers import rotary
@@ -313,7 +343,7 @@ def _meta_state_dict(folder: str) -> dict[str, torch.Tensor]:
     rotary.get_rope.cache_clear()
     try:
         with torch.device("meta"), torch_dtype(torch.bfloat16):
-            return create_model(config.model_config).state_dict()
+            return (build or create_model)(config.model_config).state_dict()
     finally:
         rotary.set_rope_device(saved)
         rotary.get_rope.cache_clear()
@@ -598,3 +628,35 @@ def test_a_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, qu
     (tmp_path / "config.json").write_text(json.dumps(_config_json(True, quantization_config)))
     with pytest.raises(ValueError, match=match):
         _load(str(tmp_path))
+
+
+def test_mtp_head_fills_the_draft_and_its_experts(checkpoint):
+    """The MTP reader fills exactly the draft the engine builds, and the routed experts keep the form the checkpoint stores."""
+    from types import SimpleNamespace
+
+    from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft
+    from freetoken.models.qwen3_5_moe.weight import iter_mtp_weights
+    from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+
+    name, folder, raw = checkpoint
+    if name not in ("bf16", "fp8_block"):
+        pytest.skip("layouts with an MTP head only")
+    _install(folder)
+    target = SimpleNamespace(model=SimpleNamespace(embed_tokens=None))
+    state = _meta_state_dict(folder, lambda config: Qwen3_5MTPDraft(config, target))
+    loaded = dict(iter_mtp_weights(folder, torch.device("cpu")))
+    assert set(loaded) == set(state)
+    for key, tensor in loaded.items():
+        assert tensor.shape == state[key].shape, key
+        if key.endswith(".weight"):
+            assert tensor.dtype is state[key].dtype, key
+    pieces = list(iter_extra_layer_expert_pieces(folder, f"{MTP}.mlp.experts", torch.device("cpu")))
+    if name == "bf16":
+        assert [p[:3] for p in pieces] == [(0, 0, E)]
+        assert torch.equal(pieces[0][3]["gate_up"], raw[f"{MTP}.mlp.experts.gate_up_proj"])
+        return
+    assert sorted(p[1] for p in pieces) == list(range(E))
+    for _layer, expert, _end, piece in pieces:
+        assert set(piece) == {"gate", "up", "down", "gate_scale", "up_scale", "down_scale"}
+        assert _same(piece["down"][0], raw[f"{MTP}.mlp.experts.{expert}.down_proj.weight"])
+        assert _same(piece["gate_scale"][0], raw[f"{MTP}.mlp.experts.{expert}.gate_proj.weight_scale_inv"])

From 0939740e557f00e4bc2a452303d7cdf7500bff22 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 28 Oct 2026 02:21:03 +0000
Subject: [PATCH 03/20] feat(spec): verify qwen4_exp targets on qsa_sparse
 graphs with PLE verify handlers

---
 python/freetoken/attention/__init__.py        |  1 +
 python/freetoken/attention/qsa_sparse.py      | 27 +++---
 python/freetoken/engine/engine.py             |  3 +
 python/freetoken/kvcache/__init__.py          |  3 +
 python/freetoken/kvcache/linear_state_pool.py | 31 ++++++-
 python/freetoken/kvcache/qsa_pool.py          |  3 +-
 python/freetoken/models/config.py             | 27 +++++-
 python/freetoken/models/qwen4_exp/config.py   |  5 +
 python/freetoken/models/qwen4_exp/model.py    | 10 +-
 python/freetoken/models/qwen4_exp/ple.py      | 77 +++++++++++++++-
 python/freetoken/models/qwen4_exp/ple_disk.py | 28 +++---
 tests/models/qwen4_exp/common.py              |  8 +-
 tests/models/qwen4_exp/test_ple.py            | 62 ++++++++++++-
 tests/models/qwen4_exp/test_ple_disk.py       | 27 +++---
 tests/models/qwen4_exp/test_qsa_backend.py    | 66 ++++++++++----
 tests/spec/test_verify_contract.py            | 91 +++++++++++++++----
 16 files changed, 386 insertions(+), 83 deletions(-)

diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py
index 8e52845003..8bc52f92a4 100644
--- a/python/freetoken/attention/__init__.py
+++ b/python/freetoken/attention/__init__.py
@@ -155,6 +155,7 @@ def create_m3_sparse_backend(config: ModelConfig):
         # 64-token pages: a 4-token compress group never straddles a page, so the
         # compressed row of a group is page_base // 4 + block-in-page.
         page_sizes=(64,),
+        uniform_extend_graph="device",
     ),
 )
 def create_qsa_sparse_backend(config: ModelConfig):
diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index 3e37d5f888..18860104cb 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -145,6 +145,7 @@ def __init__(self, config: ModelConfig) -> None:
         # decode staging (static buffers under CUDA graphs; eager decode snapshots per step)
         self._graph: dict[str, torch.Tensor] = {}
         self.capture_bs: List[int] = []
+        self._capture_rows_per_req = 1
 
     @staticmethod
     def _qsa_group(config: ModelConfig):
@@ -242,7 +243,7 @@ def _stage_decode(self, md: QSASparseMetadata, bs: int, table_idx: torch.Tensor)
         md.block_table = self._graph["block_table"][:bs]
         md.seq_lens = self._graph["kvlen"][:bs]
         md.ring_slots = self._graph["table_idx"][:bs]
-        md.token_to_req = self._graph["token_to_req"][:bs]
+        md.token_to_req = self._graph["token_to_req"][: bs * self._capture_rows_per_req]
         md.cu_seqlens = self._graph["cu_seqlens"][: bs + 1]
 
     def _snapshot_decode(self, md: QSASparseMetadata, batch: Batch) -> None:
@@ -490,14 +491,14 @@ def _scratch(self, name: str, rows: int, *shape: int, dtype: torch.dtype) -> tor
 
     # ----- CUDA graph (decode) --------------------------------------------------------------
     def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
-        if rows_per_req != 1:
-            raise NotImplementedError("qsa_sparse captures decode graphs only")
         self.capture_bs = sorted(bs_list)
+        self._capture_rows_per_req = rows_per_req
         max_bs = max(bs_list)
+        rows = max_bs * rows_per_req
         width = get_global_ctx().page_table.shape[1]
         pages = -(-width // self.page_size)
         columns = pages * self.cmp_page_size
-        chunk = max(1, min(max_bs, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)))
+        chunk = max(1, min(rows, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)))
         topk_scratch = self._topk_scratch_width(columns)
 
         def empty(*shape: int, dtype: torch.dtype) -> torch.Tensor:
@@ -507,15 +508,15 @@ def empty(*shape: int, dtype: torch.dtype) -> torch.Tensor:
             "block_table": torch.zeros((max_bs, pages), dtype=torch.int32, device=self.device),
             "kvlen": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
             "table_idx": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
-            "token_to_req": torch.arange(max_bs, dtype=torch.int32, device=self.device),
-            "cu_seqlens": torch.arange(max_bs + 1, dtype=torch.int32, device=self.device),
+            "token_to_req": torch.arange(max_bs, dtype=torch.int32, device=self.device).repeat_interleave(rows_per_req),
+            "cu_seqlens": torch.arange(max_bs + 1, dtype=torch.int32, device=self.device) * rows_per_req,
             "logits": empty(chunk, columns, dtype=torch.float32),
-            "visible": empty(max_bs, dtype=torch.int32),
-            "blocks": empty(max_bs, self.block_topk, dtype=torch.int32),
-            "indices": empty(max_bs, self.select_width, dtype=torch.int32),
-            "pooled": empty(max_bs, self.index_head_dim, dtype=self.dtype),
-            "first_pos": empty(max_bs, dtype=torch.int32),
-            "q_index": empty(max_bs, self.index_heads, self.index_head_dim, dtype=self.dtype),
+            "visible": empty(rows, dtype=torch.int32),
+            "blocks": empty(rows, self.block_topk, dtype=torch.int32),
+            "indices": empty(rows, self.select_width, dtype=torch.int32),
+            "pooled": empty(rows, self.index_head_dim, dtype=self.dtype),
+            "first_pos": empty(rows, dtype=torch.int32),
+            "q_index": empty(rows, self.index_heads, self.index_head_dim, dtype=self.dtype),
         }
         if topk_scratch:
             self._graph["topk_scratch"] = empty(chunk, topk_scratch, dtype=torch.int32)
@@ -533,7 +534,7 @@ def prepare_for_capture(self, batch: Batch) -> None:
     def prepare_for_replay(self, batch: Batch) -> None:
         md = batch.attn_metadata
         assert isinstance(md, QSASparseMetadata)
-        assert batch.active_table_idx is not None, "decode batch is missing its page-table rows"
+        assert batch.active_table_idx is not None, "graph batch is missing its page-table rows"
         self._stage_decode(md, batch.padded_size, batch.active_table_idx.to(torch.int64))
 
     def reset_capture(self) -> None:
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 590cf80cca..6e3422a164 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -210,6 +210,9 @@ def _adjust_spec_config(config, override, windowed_kv: bool) -> None:
             "speculative decoding does not support targets with sliding-window or DSV4 window-tier KV yet: "
             "their window frees would release slots a verify step still reads"
         )
+    # at config time: the state pool that holds the handlers is built only after the weights load
+    if unverified := [s.name for s in config.model_config.slot_states if s.verify is None]:
+        raise ValueError(f"speculative decoding does not support per-request state {unverified} yet: it declares no verify handler")
     parts = [p.strip() for p in config.attention_backend.split(",")]
     graph = attention_backend_info(parts[-1]).uniform_extend_graph
     if graph == "host_plan" and len(parts) == 2 and attention_backend_info(parts[0]).uniform_extend_graph == "device":
diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py
index e0cf606db1..8d038e2fb6 100644
--- a/python/freetoken/kvcache/__init__.py
+++ b/python/freetoken/kvcache/__init__.py
@@ -136,6 +136,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt
         device=device,
         dtype=dtype,
         num_req_slots=config.max_running_req + 1,  # + 1 for the dummy request row
+        num_speculative_tokens=getattr(config, "decode_rows_per_req", 1) - 1,
     )
 
 
@@ -147,6 +148,7 @@ def create_kvcache_pool(
     device: torch.device,
     num_swa_tokens: int | None = None,
     num_req_slots: int | None = None,
+    num_speculative_tokens: int = 0,
 ) -> BaseKVCachePool:
     if model_config.has_swa_attention:
         from .hybrid_swa_pool import HybridSWAKVCache
@@ -222,6 +224,7 @@ def create_kvcache_pool(
             num_index_layers=spec.num_index_layers,
             index_ratio=spec.index_ratio,
             num_req_slots=num_req_slots,
+            ring_capacity=QSAKVCache.ring_capacity_for(spec.index_ratio, num_speculative_tokens),
             layer_ids=spec.layer_ids,
             mrope=model_config.model_is_mrope,
         )
diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py
index fa9675f9d5..374a9dab54 100644
--- a/python/freetoken/kvcache/linear_state_pool.py
+++ b/python/freetoken/kvcache/linear_state_pool.py
@@ -6,7 +6,7 @@
 import torch
 from freetoken.distributed import get_tp_info
 from freetoken.env import ENV
-from freetoken.models.config import LinearGatedDeltaGroupConfig, SlotStateSpec
+from freetoken.models.config import LinearGatedDeltaGroupConfig, SlotStateSpec, SlotVerify
 from freetoken.utils import div_even
 
 _SSM_DTYPES = {
@@ -96,6 +96,7 @@ def __init__(
         self.padding_slot = 0
         # verify batches only read the live state; Engine attaches this when a decode step has more than one row
         self.verify_scratch: LinearVerifyScratch | None = None
+        self._verify_handlers: dict[str, SlotVerify] = {}
         self._free_slots: list[int] = list(range(1, num_slots))
 
     def _alloc_slot_states(self, num_slots: int) -> dict[str, torch.Tensor]:
@@ -103,12 +104,15 @@ def _alloc_slot_states(self, num_slots: int) -> dict[str, torch.Tensor]:
             spec.name: torch.full(
                 (max(1, len(spec.layer_ids)), num_slots, *spec.shape),
                 spec.fill_value,
-                dtype=spec.dtype if spec.dtype is not None else self._conv_dtype,
+                dtype=self._state_dtype(spec),
                 device=self._device,
             )
             for spec in self._slot_specs
         }
 
+    def _state_dtype(self, spec: SlotStateSpec) -> torch.dtype:
+        return spec.dtype if spec.dtype is not None else self._conv_dtype
+
     def has_slot_state(self, name: str) -> bool:
         return name in self.slot_states
 
@@ -228,12 +232,26 @@ def device(self) -> torch.device:
         return self._device
 
     def attach_verify_scratch(self, max_bs: int, rows_per_req: int) -> None:
-        """Build the scratch for verify batches of up to ``max_bs`` requests with ``rows_per_req`` rows each; build_fla_metadata hands it to the GDN layers."""
+        """Build the scratch for verify batches of up to ``max_bs`` requests with ``rows_per_req`` rows each, and the verify handler of every slot state; build_fla_metadata hands the GDN scratch to the GDN layers."""
         self.verify_scratch = LinearVerifyScratch(self, max_bs, rows_per_req)
+        self._verify_handlers = {
+            spec.name: spec.verify(spec, max_bs, rows_per_req, self._device, self._state_dtype(spec))
+            for spec in self._slot_specs
+            if spec.verify is not None
+        }
+
+    def verify_handler(self, name: str) -> SlotVerify:
+        """The verify handler of slot state ``name``, for its owner's verify forward."""
+        return self._verify_handlers[name]
 
     def commit_verify(self, slots: List[int], counts: List[int]) -> None:
         """Advance live slot slots[i] over verify row i's first counts[i] tokens."""
         self.verify_scratch.commit(self, slots, counts)
+        if slots and self._verify_handlers:
+            index = torch.tensor([slots, counts], dtype=torch.int64, pin_memory=self._device.type == "cuda")
+            index = index.to(self._device, non_blocking=True)
+            for handler in self._verify_handlers.values():
+                handler.commit(self, *index)
 
     def bytes_per_slot(self) -> int:
         """Total state bytes for one request (all linear layers)."""
@@ -337,7 +355,12 @@ def verify_scratch_bytes(config, rows_per_req: int) -> int:
     rows = n_layers * config.max_running_req * rows_per_req
     stash = rows * (hk * group.key_head_dim + hv * group.value_head_dim + 2 * hv) * 4
     window = rows * conv_dim * (group.conv_kernel_dim - 1) * config.dtype.itemsize
-    return stash + window
+    handlers = sum(
+        spec.verify.nbytes(spec, config.max_running_req, rows_per_req, spec.dtype if spec.dtype is not None else config.dtype)
+        for spec in getattr(config.model_config, "slot_states", ())
+        if spec.verify is not None
+    )
+    return stash + window + handlers
 
 
 def _linear_pool_num_slots(config) -> int:
diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py
index 7de70977d5..84cab5d97f 100644
--- a/python/freetoken/kvcache/qsa_pool.py
+++ b/python/freetoken/kvcache/qsa_pool.py
@@ -173,7 +173,8 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]:
             if spec.attn_type is AttnType.QSA:
                 # One index-key row = all index layers at one position.
                 row = spec.index_head_dim * spec.num_index_layers * _INDEX_DTYPE_BYTES
-                fixed += num_req_slots * row * (cls.ring_capacity_for(spec.index_ratio) + 1)
+                ring = cls.ring_capacity_for(spec.index_ratio, getattr(config, "decode_rows_per_req", 1) - 1)
+                fixed += num_req_slots * row * (ring + 1)
                 if config.model_config.model_is_mrope:
                     per_token += _ROPE_POS_BYTES
         return per_token * config.page_size, fixed, config.page_size, 0
diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py
index 25bc73f5e4..7e954c06df 100644
--- a/python/freetoken/models/config.py
+++ b/python/freetoken/models/config.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 from dataclasses import dataclass
-from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias
+from typing import Any, ClassVar, Dict, List, Literal, Protocol, Tuple, TypeAlias
 
 from freetoken.attention.base import AttnType
 
@@ -242,6 +242,29 @@ def _full_group_attn_type(group: FullAttentionGroupConfig) -> AttnType:
     return AttnType.MLA
 
 
+class SlotVerify(Protocol):
+    """How a speculative verify step stages and commits one slot state; the state's owner supplies it and the pool only calls it.
+
+    The pool builds one per engine as ``cls(spec, max_bs, rows_per_req, device, dtype)``, ``dtype`` being the state's own, and keeps it across ``rebuild``.
+    A verify forward reads the live state only and writes what the commit needs into this handler's scratch; the owner gets it back with ``pool.verify_handler(spec.name)``.
+    Scratch row = verify-batch row, never a pool slot.
+    The scratch is single-buffered, which is safe because speculative decoding forces the normal loop.
+    A handler reads only its own scratch and writes only its own state, so the handlers commit in any order.
+    """
+
+    @staticmethod
+    def nbytes(spec: SlotStateSpec, max_bs: int, rows_per_req: int, dtype) -> int:
+        """Scratch bytes, from the spec alone, before any pool exists."""
+        ...
+
+    def commit(self, pool, slots, counts) -> None:
+        """Roll live slot ``slots[i]`` to where it stands after the first ``counts[i]`` rows of verify request i; int64 device tensors, one entry per request.
+
+        Reads the scratch without consuming it and may run more than once per step; reads the state through ``pool.slot_state`` on every call, since ``rebuild`` replaces it.
+        """
+        ...
+
+
 @dataclass(frozen=True)
 class SlotStateSpec:
     """One extra per-request tensor riding the LinearStatePool slots.
@@ -256,6 +279,8 @@ class SlotStateSpec:
     layer_ids: Tuple[int, ...] = ()
     dtype: Any | None = None  # a torch dtype; None -> the pool's compute dtype
     fill_value: float = 0.0
+    # None: speculative decoding cannot serve a model carrying this state
+    verify: type[SlotVerify] | None = None
 
 
 @dataclass(frozen=True)
diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py
index 2bc96f67a1..7a2743b5ad 100644
--- a/python/freetoken/models/qwen4_exp/config.py
+++ b/python/freetoken/models/qwen4_exp/config.py
@@ -81,12 +81,16 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]:
     """Per-request PLE state riding the linear-state slots (see LinearStatePool.slot_states)."""
     if not args.ple_layer_ids:
         return ()
+    # ple.py imports this module
+    from .ple import PLEConvVerify, PLENgramVerify
+
     return (
         # dilated-conv left context; replicated (not TP-sharded), model dtype
         SlotStateSpec(
             name=PLE_CONV_STATE,
             shape=(args.ple_state_width, args.ple_conv_state_len),
             layer_ids=args.ple_layer_ids,
+            verify=PLEConvVerify,
         ),
         # last ngram_size-1 token ids, shared by every PLE layer; eos = hash boundary
         SlotStateSpec(
@@ -94,6 +98,7 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]:
             shape=(args.ngram_size - 1,),
             dtype=torch.int32,
             fill_value=float(args.ngram_boundary_token_id),
+            verify=PLENgramVerify,
         ),
     )
 
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index cca354d369..bc1d06b488 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -24,6 +24,7 @@
 from freetoken.utils import nvtx_annotate
 
 from .attention import Qwen4ExpAttention
+from .config import PLE_NGRAM_STATE
 from .hc import GatedResidual
 from .moe import Qwen4ExpMoE
 from .ple import PLELayer
@@ -123,7 +124,10 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor:
                 ple.start_prefetch(batch, meta)
         for layer in self.layers.op_list:
             hidden = layer.forward(hidden, batch)
-        if meta is not None:
+        if meta is not None and meta.is_verify:
+            # the commit rolls the context by the accepted rows only
+            get_global_ctx().linear_state_pool.verify_handler(PLE_NGRAM_STATE).stash(meta)
+        elif meta is not None:
             # single writer: the layers only read the context, so a second PLE layer's
             # prefetch sees the un-rolled window
             commit_ngram_context(meta, getattr(batch, "fla_metadata", None))
@@ -185,10 +189,12 @@ def load_host_tables(self, engine_config) -> int:
                 "eos_token_id": args.ngram_boundary_token_id,
                 "image_token_id": args.image_token_id,
             }
+            rows = engine_config.decode_rows_per_req
             disk_table = DiskRowTable(
                 resolve_row_source(engine_config.model_path),
                 constants,
-                max_graph_rows=max(256, engine_config.cuda_graph_max_bs or 0),
+                # verify graphs stop at max_running_req requests, decode graphs at cuda_graph_max_bs
+                max_graph_rows=engine_config.max_running_req * rows if rows > 1 else max(256, engine_config.cuda_graph_max_bs or 0),
                 max_extend_tokens=engine_config.max_extend_tokens,
             )
             self._ple_table = disk_table
diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py
index 5a3b321f08..9114eef602 100644
--- a/python/freetoken/models/qwen4_exp/ple.py
+++ b/python/freetoken/models/qwen4_exp/ple.py
@@ -285,6 +285,7 @@ class PLEMetadata:
                      ``Req.table_idx``); keys every PLE slot state
       fresh_slots    [B] bool device or None -- request starts a new sequence, so read a zero state
       is_decode      one token per request (the batched 4-tap path)
+      is_verify      the same number of rows per request over read-only state; the verify handlers stage what the commit rolls
     """
 
     input_ids: torch.Tensor
@@ -294,6 +295,7 @@ class PLEMetadata:
     state_slots: torch.Tensor
     fresh_slots: torch.Tensor | None
     is_decode: bool
+    is_verify: bool = False
 
 
 def _state_slot(req) -> int:
@@ -356,13 +358,17 @@ def build_ple_metadata(
         cu = fla.cu_seqlens
         slots = fla.cache_indices.long()
         fresh = ~fla.has_initial_state
+    elif fla is not None:
+        # a verify: every request continues its live state
+        cu, slots, fresh = fla.cu_seqlens, fla.cache_indices.long(), None
     else:  # direct-op callers (tests) with no scheduler metadata
         pin = {"device": "cpu", "pin_memory": torch.cuda.is_available()}
         cu = torch.tensor([0, *lens], dtype=torch.int64, **pin).cumsum_(0).to(device, non_blocking=True)
         slots = torch.tensor([_state_slot(r) for r in reqs], dtype=torch.int64, **pin).to(device, non_blocking=True)
         fresh = torch.tensor([r.cached_len == 0 for r in reqs], dtype=torch.bool, **pin).to(device, non_blocking=True)
     context = context_pool.index_select(0, slots).long()
-    context = torch.where(fresh.unsqueeze(1), context.new_full((), eos), context)
+    if fresh is not None:
+        context = torch.where(fresh.unsqueeze(1), context.new_full((), eos), context)
     return PLEMetadata(
         input_ids=input_ids,
         cu_seqlens=cu,
@@ -371,6 +377,7 @@ def build_ple_metadata(
         state_slots=slots,
         fresh_slots=fresh,
         is_decode=batch.is_decode,
+        is_verify=getattr(batch, "phase", None) == "verify",
     )
 
 
@@ -650,6 +657,8 @@ def _short_conv(
         """silu of the dilated depthwise conv over [state | x], and roll the per-request state."""
         if meta.is_decode:
             return self._decode_conv(x, meta, states)
+        if meta.is_verify:
+            return self._verify_conv(x, meta, states)
         return self._prefill_conv(x, meta, states)
 
     def _decode_conv(
@@ -666,6 +675,21 @@ def _decode_conv(
         )
         return F.silu(out.to(x.dtype))
 
+    def _verify_conv(
+        self, x: torch.Tensor, meta: PLEMetadata, states: torch.Tensor
+    ) -> torch.Tensor:
+        """The decode taps over every verify row, the live state read-only; device ops only, so a verify graph captures it."""
+        bs = meta.state_slots.numel()
+        rows = x.shape[0] // bs
+        handler = get_global_ctx().linear_state_pool.verify_handler(PLE_CONV_STATE)
+        history = handler.history[self.layer_id][:bs]
+        history[..., : self.state_len].copy_(self._read_state(meta, states, x.dtype))
+        history[..., self.state_len :].copy_(x.view(bs, rows, -1).transpose(1, 2))
+        # contiguous taps reduce in the order the decode taps do
+        window = history.unfold(2, self.state_len + 1, 1)[..., :: self.dilation].contiguous().float()
+        out = (window * self.conv1d.weight.squeeze(1).float().unsqueeze(1)).sum(-1)
+        return F.silu(out.to(x.dtype)).transpose(1, 2).reshape(x.shape)
+
     def _prefill_conv(
         self, x: torch.Tensor, meta: PLEMetadata, states: torch.Tensor
     ) -> torch.Tensor:
@@ -715,12 +739,63 @@ def _prefill_indices(self, lens: List[int], device: torch.device):
         return packed[:n_out], packed[n_out : n_out + n_state], packed[n_out + n_state :]
 
 
+class PLEConvVerify:
+    """``SlotVerify`` of the conv history: per PLE layer and verify row, the live history then every verify conv input."""
+
+    def __init__(self, spec, max_bs: int, rows_per_req: int, device: torch.device, dtype: torch.dtype) -> None:
+        self.name = spec.name
+        width, self.state_len = spec.shape
+        self.history = {
+            layer_id: torch.zeros(max_bs, width, self.state_len + rows_per_req, dtype=dtype, device=device)
+            for layer_id in spec.layer_ids
+        }
+
+    @staticmethod
+    def nbytes(spec, max_bs: int, rows_per_req: int, dtype: torch.dtype) -> int:
+        width, state_len = spec.shape
+        return len(spec.layer_ids) * max_bs * width * (state_len + rows_per_req) * dtype.itemsize
+
+    def commit(self, pool, slots: torch.Tensor, counts: torch.Tensor) -> None:
+        # after n accepted rows the history is columns n .. n + state_len - 1
+        cols = (counts.unsqueeze(1) + torch.arange(self.state_len, device=counts.device)).unsqueeze(1)
+        for layer_id, history in self.history.items():
+            states = pool.slot_state(self.name, layer_id)
+            window = history[: slots.numel()].gather(2, cols.expand(-1, history.shape[1], -1))
+            states.index_copy_(0, slots, window.to(states.dtype))
+
+
+class PLENgramVerify:
+    """``SlotVerify`` of the n-gram context: per verify row, the live context then every verify input token."""
+
+    def __init__(self, spec, max_bs: int, rows_per_req: int, device: torch.device, dtype: torch.dtype) -> None:
+        self.name = spec.name
+        (self.ctx_len,) = spec.shape
+        self.ids = torch.zeros(max_bs, self.ctx_len + rows_per_req, dtype=dtype, device=device)
+
+    @staticmethod
+    def nbytes(spec, max_bs: int, rows_per_req: int, dtype: torch.dtype) -> int:
+        return max_bs * (spec.shape[0] + rows_per_req) * dtype.itemsize
+
+    def stash(self, meta: PLEMetadata) -> None:
+        """Stage a verify forward's context and input tokens; the forward itself leaves the live context as it was."""
+        bs = meta.state_slots.numel()
+        self.ids[:bs, : self.ctx_len].copy_(meta.ngram_context)
+        self.ids[:bs, self.ctx_len :].copy_(meta.input_ids.view(bs, -1))
+
+    def commit(self, pool, slots: torch.Tensor, counts: torch.Tensor) -> None:
+        context = pool.slot_state(self.name)
+        cols = counts.unsqueeze(1) + torch.arange(self.ctx_len, device=counts.device)
+        context.index_copy_(0, slots, self.ids[: slots.numel()].gather(1, cols).to(context.dtype))
+
+
 __all__ = [
     "GpuResidentTable",
     "NGramEmbedding",
     "ZeroTable",
+    "PLEConvVerify",
     "PLELayer",
     "PLEMetadata",
+    "PLENgramVerify",
     "PLETableBackend",
     "PinnedUVATable",
     "build_ple_metadata",
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index 7ec067bd81..65217608f8 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -217,12 +217,12 @@ def __init__(
         self._pending: Future | None = None
         self._failed = False
 
-    def _fill_decode(self, histories: Sequence[Sequence[int]], bs: int) -> None:
+    def _fill_decode(self, histories: Sequence[Sequence[int]], bs: int, rows: int) -> None:
         try:
             self._readback_event.synchronize()
-            tokens = self._token_readback[:bs].to(torch.int64).tolist()
-            runs = [torch.tensor([*history, t], dtype=torch.int64)
-                    for history, t in zip(histories, tokens)]
+            tokens = self._token_readback[: bs * rows].to(torch.int64).tolist()
+            runs = [torch.tensor([*history, *tokens[i * rows : (i + 1) * rows]], dtype=torch.int64)
+                    for i, history in enumerate(histories)]
             self.fill(runs, graph=True)
         except BaseException:
             from freetoken.kernel.row_store import signal
@@ -231,10 +231,10 @@ def _fill_decode(self, histories: Sequence[Sequence[int]], bs: int) -> None:
             signal(self._flag)
             raise
 
-    def _start_fill(self, histories: Sequence[Sequence[int]], bs: int) -> None:
+    def _start_fill(self, histories: Sequence[Sequence[int]], bs: int, rows: int) -> None:
         @torch.inference_mode()
         def _complete() -> None:
-            self._fill_decode(histories, bs)
+            self._fill_decode(histories, bs, rows)
 
         # Like Engram, release HIP's WAIT independently of a queue-blocked graph launch.
         self._pending = self._worker.submit(_complete)
@@ -290,25 +290,27 @@ def host_fill_batch(self, batch: Batch, use_graph: bool):
             if self._failed:
                 raise RuntimeError("PLE disk table failed; recreate the engine before dispatch")
         eos = self.eos_token_id
-        if batch.is_decode:
+        # decode and verify: rows_per_req tokens per request on the device, the first at device_len - rows_per_req
+        rows = getattr(batch, "rows_per_req", None) or 1
+        if batch.is_decode or rows > 1:
             reqs = list(batch.reqs)
+            histories = [self._ple_context(r.input_ids, r.device_len - rows) for r in reqs]
             if use_graph and self._wait_sync:
                 bs = batch.padded_size
-                self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
+                self._token_readback[: bs * rows].copy_(batch.input_ids, non_blocking=True)
                 self._readback_event.record(torch.cuda.current_stream(self._device))
-                histories = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs]
                 if self._worker is not None:
-                    self._start_fill(histories, bs)
+                    self._start_fill(histories, bs, rows)
                     return None
 
                 def _complete() -> None:
-                    self._fill_decode(histories, bs)
+                    self._fill_decode(histories, bs, rows)
 
                 return _complete
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
             tokens = batch.input_ids.to("cpu").to(torch.int64).tolist()
-            runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64)
-                    for r, t in zip(reqs, tokens)]
+            runs = [torch.tensor([*history, *tokens[i * rows : (i + 1) * rows]], dtype=torch.int64)
+                    for i, history in enumerate(histories)]
             self.fill(runs, graph=use_graph)
             return None
         runs = [
diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py
index 0a52fb0e01..0997f5d490 100644
--- a/tests/models/qwen4_exp/common.py
+++ b/tests/models/qwen4_exp/common.py
@@ -154,6 +154,7 @@ def __init__(
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         page_size: int = 64,
+        num_speculative_tokens: int = 0,
     ) -> None:
         from freetoken.attention.qsa_sparse import QSASparseAttnBackend
         from freetoken.kvcache import create_kvcache_pool
@@ -170,6 +171,7 @@ def __init__(
             dtype=dtype,
             device=self.device,
             num_req_slots=self.num_req_slots,
+            num_speculative_tokens=num_speculative_tokens,
         )
         self.page_table = torch.zeros(
             (self.num_req_slots, num_pages * page_size), dtype=torch.int32, device=self.device
@@ -208,9 +210,9 @@ def req(self, table_idx: int, cached_len: int, device_len: int) -> SimpleNamespa
             extend_len=device_len - cached_len,
         )
 
-    def step(self, req: SimpleNamespace) -> None:
-        self.allocate(req.table_idx, req.device_len, req.device_len + 1)
-        req.cached_len, req.device_len, req.extend_len = req.device_len, req.device_len + 1, 1
+    def step(self, req: SimpleNamespace, rows: int = 1) -> None:
+        self.allocate(req.table_idx, req.device_len, req.device_len + rows)
+        req.cached_len, req.device_len, req.extend_len = req.device_len, req.device_len + rows, rows
 
     def batch(self, reqs, phase: str) -> SimpleNamespace:
         positions = torch.cat(
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index 350a798cb4..4eeeca061c 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -30,7 +30,7 @@
     short_conv_reference,
 )
 
-from .common import EOS, VOCAB, hash_constants, requires_cuda, toy_hf_config
+from .common import EOS, VOCAB, fresh_ctx, hash_constants, requires_cuda, toy_hf_config
 
 _HF_REF_PYTHON = os.environ.get("FREETOKEN_QWEN4_HF_PYTHON", "")
 _HF_REF_SCRIPT = Path(__file__).with_name("ple_hf_ref.py")
@@ -623,6 +623,66 @@ def test_context_matches_the_token_history_across_chunks_and_decode():
         cached += 1
 
 
+@requires_cuda
+@pytest.mark.parametrize("accepted", [(1, 4), (2, 1), (4, 3)])
+def test_verify_commit_matches_stepwise_decode(accepted):
+    """A verify forward over k+1 rows per request leaves the PLE state as it was; committing n rows leaves the conv history and n-gram context of n decode steps, and every row's conv output is its decode step's, bitwise."""
+    from freetoken.attention.linear import FLAMetadata
+    from freetoken.core import Batch, Req, SamplingParams
+    from freetoken.kvcache.linear_state_pool import LinearStatePool
+    from freetoken.models.qwen4_exp.config import PLE_CONV_STATE, PLE_NGRAM_STATE
+
+    torch.manual_seed(19)
+    config = _config()
+    args = config.qwen4_args
+    device = torch.device("cuda")
+    layer = _make_layer(config, device="cuda")
+    pool = LinearStatePool(config.linear_attention_group(), 6, torch.float32, device, tp_size=1, slot_states=config.slot_states)
+    fresh_ctx(linear_state_pool=pool)
+    rows, slots = 4, [3, 1]
+    bs = len(slots)
+    conv = pool.slot_state(PLE_CONV_STATE, layer.layer_id)
+    context = pool.slot_state(PLE_NGRAM_STATE)
+    conv.copy_(torch.randn_like(conv))
+    context.copy_(torch.randint(8, VOCAB, context.shape, device=device, dtype=context.dtype))
+    conv0, context0 = conv.clone(), context.clone()
+    tokens = torch.tensor([[11, EOS, 12, 13], [14, 15, 16, EOS]], dtype=torch.int32, device=device)
+    x = torch.randn(bs, rows, args.ple_state_width, device=device)
+    slot_ids = torch.tensor(slots, dtype=torch.int32, device=device)
+
+    def batch(phase, input_ids, extend):
+        reqs = [Req(input_ids=torch.zeros(8 + extend, dtype=torch.int32), table_idx=slot, cached_len=8, output_len=8,
+                    uid=i, sampling_params=SamplingParams(), cache_handle=None) for i, slot in enumerate(slots)]
+        b = Batch(reqs=reqs, phase=phase, rows_per_req=extend)
+        b.padded_reqs, b.input_ids = reqs, input_ids
+        b.fla_metadata = FLAMetadata(
+            cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device) * extend, cache_indices=slot_ids)
+        return b
+
+    want, after = [], []
+    for t in range(rows):
+        meta = build_ple_metadata(batch("decode", tokens[:, t].contiguous(), 1), args, device)
+        want.append(layer._short_conv(x[:, t].contiguous(), meta, conv))
+        commit_ngram_context(meta, None)
+        after.append((conv.clone(), context.clone()))
+    conv.copy_(conv0)
+    context.copy_(context0)
+
+    pool.attach_verify_scratch(bs, rows)
+    meta = build_ple_metadata(batch("verify", tokens.view(-1), rows), args, device)
+    got = layer._short_conv(x.view(bs * rows, -1), meta, conv)
+    pool.verify_handler(PLE_NGRAM_STATE).stash(meta)
+    assert torch.equal(conv, conv0) and torch.equal(context, context0)
+    assert torch.equal(got.view(bs, rows, -1), torch.stack(want, dim=1))
+
+    pool.commit_verify(slots, list(accepted))
+    for slot, n in zip(slots, accepted):
+        assert torch.equal(conv[slot], after[n - 1][0][slot]), f"conv history of slot {slot} after {n} rows"
+        assert torch.equal(context[slot], after[n - 1][1][slot]), f"n-gram context of slot {slot} after {n} rows"
+    others = [s for s in range(conv.shape[0]) if s not in slots]
+    assert torch.equal(conv[others], conv0[others]) and torch.equal(context[others], context0[others])
+
+
 # --------------------------------------------------------------------------------------
 # CUDA graph + prefetch overlap
 # --------------------------------------------------------------------------------------
diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py
index 85a98dfdae..13a30438f5 100644
--- a/tests/models/qwen4_exp/test_ple_disk.py
+++ b/tests/models/qwen4_exp/test_ple_disk.py
@@ -364,7 +364,7 @@ def test_graph_sync_protocol(tmp_path, monkeypatch):
     assert gated.host_fill_batch(_decode_batch([3, 4], 5), use_graph=True) is None
 
 
-def _wait_sync_large_graph(tmp_path, kernels):
+def _wait_sync_large_graph(tmp_path, kernels, rows_per_req=1):
     from freetoken.core import Batch, Context, Req, set_global_ctx
     from freetoken.distributed import set_tp_info
     from freetoken.engine.graph import GraphRunner
@@ -373,7 +373,7 @@ def _wait_sync_large_graph(tmp_path, kernels):
     set_global_ctx(Context(page_size=1))
     disk, oracle, args = _make_table(tmp_path)
     emb = _embedding()
-    row_ids = torch.zeros((8, args.num_ngram_heads), dtype=torch.int64, device="cuda")
+    row_ids = torch.zeros((8 * rows_per_req, args.num_ngram_heads), dtype=torch.int64, device="cuda")
     pad = torch.zeros(256, device="cuda")
     release = threading.Event()
     fill = disk.fill
@@ -390,8 +390,10 @@ def forward():
         return out
 
     disk.fill = delayed_fill
-    req = Req(input_ids=torch.tensor([3, 4, 9], dtype=torch.int32), table_idx=0,
+    # a verify request widened by rows_per_req - 1 drafts; the store reads its rows from the device
+    req = Req(input_ids=torch.tensor([3, 4] + [9] * rows_per_req, dtype=torch.int32), table_idx=0,
               cached_len=2, output_len=3, uid=0, sampling_params=None, cache_handle=None)
+    phase = "decode" if rows_per_req == 1 else "verify"
     attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None,
                            prepare_for_capture=lambda batch: None, prepare_for_replay=lambda batch: None)
     stream = torch.cuda.Stream()
@@ -400,14 +402,16 @@ def forward():
         # Exercise the capture-time upload from #697 as well as the PLE dispatch hook.
         runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward,
                              attn_backend=attn, cuda_graph_bs=[8], cuda_graph_max_bs=None,
-                             free_memory=1024**3, max_seq_len=16, dummy_req=req, out_dtype=torch.float32)
-        batch = Batch(reqs=[req] * 8, phase="decode")
+                             free_memory=1024**3, max_seq_len=16, dummy_req=req, out_dtype=torch.float32,
+                             rows_per_req=rows_per_req, phase=phase)
+        batch = Batch(reqs=[req] * 8, phase=phase, rows_per_req=rows_per_req)
         runner.pad_batch(batch)
-        batch.positions = torch.full((8,), 2, dtype=torch.int32, device="cuda")
-        batch.out_loc = torch.arange(8, dtype=torch.int32, device="cuda")
+        batch.positions = torch.full((8 * rows_per_req,), 2, dtype=torch.int32, device="cuda")
+        batch.out_loc = torch.arange(8 * rows_per_req, dtype=torch.int32, device="cuda")
         for step, token in enumerate((9, 13, 5)):
             release.clear()
-            batch.input_ids = torch.full((8,), token, dtype=torch.int32, device="cuda")
+            tokens = [token + row for row in range(rows_per_req)]
+            batch.input_ids = torch.tensor(tokens * 8, dtype=torch.int32, device="cuda")
             with disk.forward_host_ctx(batch, use_graph=True):
                 # Keep the WAIT unsatisfied when the first replay starts.
                 assert int(disk._flag[0]) == 0
@@ -418,7 +422,7 @@ def forward():
                 finally:
                     timer.join()
             stream.synchronize()
-            ids = emb.row_ids(_meta([[token]] * 8, [[3, 4]] * 8, decode=True)).cuda()
+            ids = emb.row_ids(_meta([tokens] * 8, [[3, 4]] * 8, decode=rows_per_req == 1)).cuda()
             assert _bitwise_equal(out.to(disk.dtype), oracle.lookup(ids))
             assert torch.equal(pad, torch.full_like(pad, (step + 1) * kernels))
             assert int(disk._flag[0]) == 0
@@ -426,7 +430,8 @@ def forward():
 
 
 @requires_cuda
-def test_wait_sync_large_graph(tmp_path):
+@pytest.mark.parametrize("rows_per_req", [1, 4], ids=["decode", "verify"])
+def test_wait_sync_large_graph(tmp_path, rows_per_req):
     from freetoken.kernel.row_store import probe_wait_sync
 
     if not probe_wait_sync("auto", torch.device("cuda")):
@@ -437,7 +442,7 @@ def test_wait_sync_large_graph(tmp_path):
     script = (
         "from pathlib import Path; "
         "from tests.models.qwen4_exp.test_ple_disk import _wait_sync_large_graph; "
-        f"_wait_sync_large_graph(Path({str(tmp_path)!r}), 16384)"
+        f"_wait_sync_large_graph(Path({str(tmp_path)!r}), 16384, {rows_per_req})"
     )
     result = subprocess.run(
         [sys.executable, "-c", script], cwd=root, env=env, capture_output=True, text=True, timeout=45,
diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py
index 236610791d..642649316e 100644
--- a/tests/models/qwen4_exp/test_qsa_backend.py
+++ b/tests/models/qwen4_exp/test_qsa_backend.py
@@ -5,7 +5,7 @@
     exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32)
     and a flashinfer dense run over the same pool;
 (b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress);
-(c) a captured decode replay equals the eager decode step.
+(c) a captured decode or verify replay equals the eager step.
 """
 
 from __future__ import annotations
@@ -149,31 +149,34 @@ def test_chunked_prefill_matches_one_shot(cut: int):
 
 
 @requires_cuda
-def test_decode_graph_replay_matches_eager():
+@pytest.mark.parametrize("rows_per_req", [1, 4], ids=["decode", "verify"])
+def test_decode_graph_replay_matches_eager(rows_per_req):
     config = parsed_config()
-    fixture = Fixture(config, num_pages=256)
+    # the pending ring holds the rows a verify may reject
+    fixture = Fixture(config, num_pages=256, num_speculative_tokens=rows_per_req - 1)
     attn = fixture.layer(QSA_LAYER)
     lengths, steps = [300, 411], 4
-    bs = len(lengths)
-    inputs = _inputs(fixture, lengths, extra=steps)
+    bs, phase = len(lengths), "decode" if rows_per_req == 1 else "verify"
+    inputs = _inputs(fixture, lengths, extra=steps * rows_per_req)
     reqs = [fixture.req(i, 0, n) for i, n in enumerate(lengths)]
     attn.forward(
         torch.cat([row[:n] for row, n in zip(inputs, lengths)]),
         fixture.batch(reqs, "prefill"),
     )
 
-    fixture.backend.init_capture_graph(max_seq_len=fixture.page_table.shape[1], bs_list=[bs])
+    fixture.backend.init_capture_graph(max_seq_len=fixture.page_table.shape[1], bs_list=[bs], rows_per_req=rows_per_req)
     dummy = SimpleNamespace(
-        table_idx=fixture.num_req_slots - 1, cached_len=1, device_len=2, extend_len=1
+        table_idx=fixture.num_req_slots - 1, cached_len=1, device_len=1 + rows_per_req, extend_len=rows_per_req
     )
+    rows = bs * rows_per_req
     static = {
-        "x": torch.zeros(bs, config.hidden_size, device=fixture.device, dtype=fixture.dtype),
-        "positions": torch.zeros(bs, dtype=torch.int32, device=fixture.device),
-        "out_loc": torch.zeros(bs, dtype=torch.int32, device=fixture.device),
+        "x": torch.zeros(rows, config.hidden_size, device=fixture.device, dtype=fixture.dtype),
+        "positions": torch.zeros(rows, dtype=torch.int32, device=fixture.device),
+        "out_loc": torch.zeros(rows, dtype=torch.int32, device=fixture.device),
     }
     capture_batch = SimpleNamespace(
-        padded_reqs=[dummy] * bs, reqs=[dummy] * bs, phase="decode", size=bs, padded_size=bs,
-        is_prefill=False, is_decode=True, positions=static["positions"],
+        padded_reqs=[dummy] * bs, reqs=[dummy] * bs, phase=phase, size=bs, padded_size=bs,
+        is_prefill=False, is_decode=rows_per_req == 1, positions=static["positions"],
         get_attn_positions=lambda: static["positions"],
         out_loc=static["out_loc"], attn_metadata=None, active_table_idx=None,
     )
@@ -187,9 +190,10 @@ def test_decode_graph_replay_matches_eager():
 
     for step in range(steps):
         for req in reqs:
-            fixture.step(req)
-        x = torch.stack([row[n + step] for row, n in zip(inputs, lengths)])
-        batch = fixture.batch(reqs, "decode")
+            fixture.step(req, rows_per_req)
+        start = step * rows_per_req
+        x = torch.cat([row[n + start : n + start + rows_per_req] for row, n in zip(inputs, lengths)])
+        batch = fixture.batch(reqs, phase)
         static["x"].copy_(x)
         static["positions"].copy_(batch.positions)
         static["out_loc"].copy_(batch.out_loc)
@@ -199,8 +203,36 @@ def test_decode_graph_replay_matches_eager():
         assert md.block_table.data_ptr() == fixture.backend._graph["block_table"].data_ptr()
         graph.replay()
         replayed = captured_out.clone()
-        eager = attn.forward(x, fixture.batch(reqs, "decode"))
-        assert torch.equal(replayed, eager), f"graph replay diverged at decode step {step}"
+        eager = attn.forward(x, fixture.batch(reqs, phase))
+        assert torch.equal(replayed, eager), f"graph replay diverged at {phase} step {step}"
+
+
+@requires_cuda
+def test_verify_after_rejected_rows_compresses_like_decode():
+    """A verify whose last rows get rejected must leave the pending ring able to close the open group: the next verify compresses it as decode steps do."""
+    config = parsed_config()
+    rows, ratio = 4, config.qwen4_args.index_ratio
+    fixture = Fixture(config, num_pages=64, num_speculative_tokens=rows - 1)
+    attn = fixture.layer(QSA_LAYER)
+    # group 300..303 is open after the prompt, its first member in the ring
+    length = 301
+    x = _inputs(fixture, [length], extra=rows + 1)[0]
+    spec, decode = fixture.req(0, 0, length), fixture.req(1, 0, length)
+    attn.forward(torch.cat([x[:length]] * 2), fixture.batch([spec, decode], "prefill"))
+    # verify 301..304 and keep only 301, then verify 302..305, which closes the group at 303
+    fixture.step(spec, rows)
+    attn.forward(x[length : length + rows], fixture.batch([spec], "verify"))
+    spec.device_len = length + 1
+    fixture.step(spec, rows)
+    attn.forward(x[length + 1 : length + 1 + rows], fixture.batch([spec], "verify"))
+    for position in range(length, length + 3):
+        fixture.step(decode)
+        attn.forward(x[position : position + 1], fixture.batch([decode], "decode"))
+
+    group = length - length % ratio
+    slab = fixture.pool.cmp_k_cache(0)
+    got, want = (slab[int(fixture.page_table[r.table_idx, group]) // ratio].float() for r in (spec, decode))
+    torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2)
 
 
 @requires_cuda
diff --git a/tests/spec/test_verify_contract.py b/tests/spec/test_verify_contract.py
index f0976e093c..5e1b84758d 100644
--- a/tests/spec/test_verify_contract.py
+++ b/tests/spec/test_verify_contract.py
@@ -1,10 +1,11 @@
 """A k+1-row verify step commits exactly what k+1 decode steps would: with drafts taken from a plain greedy run and wrong at a different place each step (so every accept count occurs), speculative decoding must reproduce that run token for token.
 
-A tiny hybrid (GatedDeltaNet + attention) model with dummy weights, served end to end through the scheduler; the drafts come from a stub proposer, so this pins the target side alone.
-A verify row and a decode step run different GEMM and attention shapes, so their logits can differ by an ulp: the runs may part at a near tie, and the test checks they part only there. fp16 keeps such ties rare (bf16 has 3 fewer mantissa bits).
+Tiny hybrid models with dummy weights, served end to end through the scheduler: qwen3_5 (GatedDeltaNet + attention) and qwen4_exp (GatedDeltaNet + QSA, PLE, hyper-connections). The drafts come from a stub proposer, so this pins the target side alone.
+A verify row and a decode step run different GEMM and attention shapes, so their logits can differ by an ulp: the runs may part at a near tie, and the test checks they part only there. fp16 keeps such ties rare (bf16 has 3 fewer mantissa bits), and fan-in scaled dummy weights keep the activations O(1), so ulp noise does not grow through the layers.
 """
 
 import json
+import math
 import os
 import shutil
 import subprocess
@@ -99,6 +100,44 @@ def propose(self, reqs, rows):
         return torch.tensor(drafts, dtype=torch.int64, device=rows.device)
 
 
+def fan_in_scaled(tensor):
+    # unit-variance dummy matrices grow the activations by sqrt(fan_in) per layer, until fp16 noise decides the argmax
+    if tensor.is_floating_point() and tensor.dim() >= 2:
+        tensor.mul_(tensor.shape[-1] ** -0.5)
+    return tensor
+
+
+dummy_state = E._make_dummy_weight_state_dict
+E._make_dummy_weight_state_dict = lambda *args, **kwargs: {
+    name: fan_in_scaled(tensor) for name, tensor in dummy_state(*args, **kwargs).items()}
+create_model = E.create_model
+
+
+def create_model_with_ple_rows(*args, **kwargs):
+    # imported here, once CUDA is up
+    from freetoken.models.qwen4_exp import ple
+    from freetoken.moe import expert_banks
+
+    dummy_fill = expert_banks._dummy_fill
+
+    def scaled_fill(role, tensor):
+        dummy_fill(role, tensor)
+        if not role.endswith(("_scale", "_global")):
+            fan_in_scaled(tensor)
+
+    expert_banks._dummy_fill = scaled_fill
+
+    # seeded PLE rows in place of the dummy zeros, so the PLE path carries signal
+    class RandomTable(ple.GpuResidentTable):
+        def __init__(self, num_rows, head_dim, dtype=torch.bfloat16):
+            generator = torch.Generator(device="cuda").manual_seed(0)
+            super().__init__(torch.randn(num_rows, head_dim, generator=generator, device="cuda", dtype=dtype))
+
+    ple.ZeroTable = RandomTable
+    return create_model(*args, **kwargs)
+
+
+E.create_model = create_model_with_ple_rows
 if reference != "-":
     PROPOSERS.register("oracle")("__main__:OracleProposer")
 llm = LLM(model_path=model, dtype=torch.float16, use_dummy_weight=True, max_running_req=4, **kwargs)
@@ -110,22 +149,40 @@ def propose(self, reqs, rows):
 """
 
 
-def _tiny_model(path: str) -> str:
+def _tiny_model(path: str, family: str) -> str:
     os.makedirs(path, exist_ok=True)
     for name in ("tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"):
         if os.path.exists(os.path.join(TOKENIZER_SOURCE, name)):
             shutil.copy(os.path.join(TOKENIZER_SOURCE, name), path)
-    config = {
-        "architectures": ["Qwen3_5ForCausalLM"], "model_type": "qwen3_5_text", "vocab_size": 248320,
-        "hidden_size": 256, "intermediate_size": 512, "num_hidden_layers": 4,
+    text = {
+        "vocab_size": 248320, "hidden_size": 256, "num_hidden_layers": 4,
         "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"],
-        "full_attention_interval": 4, "num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 64,
+        "full_attention_interval": 4, "num_attention_heads": 4, "num_key_value_heads": 2,
         "linear_conv_kernel_dim": 4, "linear_key_head_dim": 32, "linear_num_key_heads": 2,
-        "linear_num_value_heads": 4, "linear_value_head_dim": 32, "attn_output_gate": True, "hidden_act": "silu",
+        "linear_num_value_heads": 4, "linear_value_head_dim": 32, "hidden_act": "silu",
         "rms_norm_eps": 1e-6, "max_position_embeddings": 4096, "tie_word_embeddings": False, "eos_token_id": 248044,
-        "rope_parameters": {"rope_type": "default", "rope_theta": 10000000, "partial_rotary_factor": 0.25,
-                            "mrope_section": [3, 3, 2], "mrope_interleaved": True},
     }
+    if family == "qwen3_5":
+        config = {
+            "architectures": ["Qwen3_5ForCausalLM"], "model_type": "qwen3_5_text", **text,
+            "intermediate_size": 512, "head_dim": 64, "attn_output_gate": True,
+            "rope_parameters": {"rope_type": "default", "rope_theta": 10000000, "partial_rotary_factor": 0.25,
+                                "mrope_section": [3, 3, 2], "mrope_interleaved": True},
+        }
+    else:
+        # the shipping QSA geometry (head_dim 256, index head_dim 128, ratio 4), scaled-down widths
+        text.update({
+            "model_type": "qwen4_exp_text", "head_dim": 256,
+            "rope_parameters": {"rope_type": "default", "rope_theta": 10000000, "partial_rotary_factor": 0.25,
+                                "mrope_section": [11, 11, 10], "mrope_interleaved": True},
+            "num_experts": 8, "num_experts_per_tok": 2, "moe_intermediate_size": 64, "shared_expert_intermediate_size": 64,
+            "output_gate_type": "sigmoid", "indexer_n_heads": 4, "indexer_kv_heads": 1, "indexer_head_dim": 128,
+            "indexer_budget": 2048, "indexer_compress_ratio": 4, "hc_count": 4, "hc_lowrank": 16,
+            "ple_layer_ids": [2], "ple_embed_dim": 64, "ple_conv_kernel_size": 4, "ngram_size": 3, "heads_per_ngram": 2,
+            "ngram_vocab_size_base": 1000, "make_ngram_vocab_size_divisible_by": 8, "split_ngram_parts": 4,
+        })
+        config = {"architectures": ["Qwen4ExpForConditionalGeneration"], "model_type": "qwen4_exp",
+                  "tie_word_embeddings": False, "text_config": text}
     with open(os.path.join(path, "config.json"), "w") as f:
         json.dump(config, f)
     return path
@@ -140,12 +197,13 @@ def _run(model: str, out: str, reference: str, **kwargs) -> dict:
 
 
 @pytest.mark.parametrize("graphs", [True, False])
-@pytest.mark.parametrize("backend", ["fa", "triton"])
-def test_verify_steps_reproduce_greedy_decode(tmp_path, backend, graphs):
+@pytest.mark.parametrize("family, backend", [("qwen3_5", "fa"), ("qwen3_5", "triton"), ("qwen4_exp", "qsa_sparse")])
+def test_verify_steps_reproduce_greedy_decode(tmp_path, family, backend, graphs):
     if backend == "fa":
         pytest.importorskip("sgl_kernel")
-    model = _tiny_model(str(tmp_path / "model"))
-    kwargs = {"attention_backend": backend, "cuda_graph_max_bs": 4 if graphs else 0}
+    model = _tiny_model(str(tmp_path / "model"), family)
+    # resident experts take the model dtype
+    kwargs = {"attention_backend": backend, "cuda_graph_max_bs": 4 if graphs else 0, "moe_strategy": "fused"}
     reference = _run(model, str(tmp_path / "decode.json"), "-", **kwargs)
     spec = _run(model, str(tmp_path / "spec.json"), str(tmp_path / "decode.json"),
                 speculative_algorithm="oracle", speculative_num_steps=K, **kwargs)
@@ -155,10 +213,11 @@ def test_verify_steps_reproduce_greedy_decode(tmp_path, backend, graphs):
         split = next((i for i, (a, b) in enumerate(zip(want, got)) if a != b), len(want))
         agreed += split
         if split < len(want):
-            # same prefix, so both runs must rank the two tokens top-2 and within two fp16 ulps of logits in [64, 128)
+            # same prefix, so both runs must rank the two tokens top-2 and within two fp16 ulps of each other
             pair = {want[split], got[split]}
             for run in (reference, spec):
                 a, b, va, vb = next(t[2:] for t in run["top2"] if t[:2] == [uid, split])
-                assert {a, b} == pair and va - vb <= 0.125, f"request {uid} parts at {split}: top-2 {(a, b)} gap {va - vb}"
+                ulp = 2.0 ** (math.floor(math.log2(abs(va) or 1.0)) - 10)
+                assert {a, b} == pair and va - vb <= 2 * ulp, f"request {uid} parts at {split}: top-2 {(a, b)} gap {va - vb}"
     # most of the run must still be compared, or a tie early in every request would hide a bug
     assert agreed >= 0.5 * sum(map(len, reference["outputs"]))

From 395c2cdf719af67632fa38fab203734dc531be7e Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Wed, 28 Oct 2026 02:23:13 +0000
Subject: [PATCH 04/20] feat(spec): add the qwen4_exp MTP draft from the
 checkpoint or a Strata pack

---
 python/freetoken/layers/embedding.py        |   5 +-
 python/freetoken/models/blocks.py           |   2 +-
 python/freetoken/models/gguf/experts.py     |  10 +-
 python/freetoken/models/qwen4_exp/model.py  |   6 +-
 python/freetoken/models/qwen4_exp/mtp.py    | 194 ++++++++++++++++++++
 python/freetoken/models/qwen4_exp/weight.py |  32 +++-
 python/freetoken/models/register.py         |   3 +
 python/freetoken/spec/mtp/chain_graph.py    |  93 ++++++++--
 python/freetoken/spec/mtp/proposer.py       |  11 +-
 9 files changed, 330 insertions(+), 26 deletions(-)
 create mode 100644 python/freetoken/models/qwen4_exp/mtp.py

diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index 469b752da0..5f786e51f9 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -125,11 +125,12 @@ def state_dict(
         return {} if result is None else result
 
     @nvtx_annotate("LMHead")
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, *, export_hidden: bool = True) -> torch.Tensor:
+        """``export_hidden=False``: the model hands its own final hidden to ``batch.features``."""
         ctx = get_global_ctx()
         batch = ctx.batch
         bs = batch.size
-        if (features := getattr(batch, "features", None)) is not None:
+        if export_hidden and (features := getattr(batch, "features", None)) is not None:
             features.final_hidden(x)
         if batch.is_prefill:
             indices = batch.attn_metadata.get_last_indices(bs)
diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py
index a7a095a291..755a93ae10 100644
--- a/python/freetoken/models/blocks.py
+++ b/python/freetoken/models/blocks.py
@@ -55,7 +55,7 @@ class FeatureSink(Protocol):
     wanted_layers: frozenset[int]
 
     def final_hidden(self, x: torch.Tensor) -> None:
-        """Every row's hidden state before the LM head; the LM head calls it once per forward."""
+        """Every row's final hidden state, once per forward: the LM head's input, or what the model hands over itself before the head."""
         ...
 
     def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None:
diff --git a/python/freetoken/models/gguf/experts.py b/python/freetoken/models/gguf/experts.py
index 0efb5982fb..4cb3b86126 100644
--- a/python/freetoken/models/gguf/experts.py
+++ b/python/freetoken/models/gguf/experts.py
@@ -48,6 +48,14 @@ def _device_bank(shape: tuple[int, ...], ggml_type: int, k: int, dtype: torch.dt
     return flat[: experts * rows].view(shape)
 
 
+def resident_expert_bank(info: GgufTensorInfo, device: torch.device) -> torch.Tensor:
+    """One stacked ``[E, rows, K]`` expert tensor as a resident device bank, in its packed ggml rows."""
+    rows = _expert_rows(info)
+    bank = _device_bank(tuple(rows.shape), info.ggml_type, info.ne[0], rows.dtype, device)
+    bank.copy_(rows)
+    return bank
+
+
 def _expert_scales(gate_up: list[GgufTensorInfo], down: GgufTensorInfo, experts: int) -> torch.Tensor:
     """[E, 4] fp32 (gate, up, down, 1) weight_scale_2 of one layer's experts."""
     gate, up = gate_up if len(gate_up) == 2 else (gate_up[0], gate_up[0])
@@ -128,4 +136,4 @@ def load_gguf_expert_banks(model_path: str, num_layers: int, *, first_layer: int
     return ExpertBanks("gguf_scaled" if scaled else "gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None, layer_tags=tags, row_tail_bytes=tails)
 
 
-__all__ = ["load_gguf_expert_banks"]
+__all__ = ["load_gguf_expert_banks", "resident_expert_bank"]
diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index bc1d06b488..59be84fde8 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -131,6 +131,9 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor:
             # single writer: the layers only read the context, so a second PLE layer's
             # prefetch sees the un-rolled window
             commit_ngram_context(meta, getattr(batch, "fla_metadata", None))
+        if (features := getattr(batch, "features", None)) is not None:
+            # a draft reads the residual streams before the final mix
+            features.final_hidden(hidden)
         return self.hyper_connection_mixer.mix(hidden)[0]
 
 
@@ -216,7 +219,8 @@ def load_host_tables(self, engine_config) -> int:
 
     def forward(self) -> torch.Tensor:
         batch = get_global_ctx().batch
-        return self.lm_head.forward(self.model.forward(batch.input_ids, batch))
+        # the model hands its pre-mix streams to batch.features itself
+        return self.lm_head.forward(self.model.forward(batch.input_ids, batch), export_hidden=False)
 
 
 class Qwen4ExpForConditionalGeneration(QwenVLVisionMixin, Qwen4ExpForCausalLM):
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
new file mode 100644
index 0000000000..b3b7f6611b
--- /dev/null
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -0,0 +1,194 @@
+"""Qwen3.8-Flash-Next MTP draft: the checkpoint's own ``mtp.*`` layer, one QSA decoder layer over the hyper-connection streams.
+
+Draft row c reads the target's residual streams before the final mix, ``R_c [hc_count*hidden]``, and the token after it::
+
+    e     = fc_embedding(norm_e(embed(x_{c+1})))
+    R_in  = fc_hidden(norm_h(R_c)) per stream + e      # norm_h: one RMS statistic over all hc_count*hidden
+    R_out = layer(R_in)                                # QSA attention + routed MoE, no PLE
+    logits = lm_head(mixer(R_out))
+
+``R_out`` relays to the next draft step. The weights come from the checkpoint's ``mtp.*`` tensors, or from a Strata MTP pack (a GGUF with the same names, stacked ggml experts).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import re
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.core import get_global_ctx
+from freetoken.layers import GemmaPlusOneRMSNorm, LinearReplicated, OPList
+from freetoken.layers.base import BaseOP
+from freetoken.layers.quantization import finalize_quant
+from freetoken.models.blocks import embed_input_ids
+from freetoken.models.config import FullAttentionGroupConfig
+from freetoken.utils import torch_dtype
+
+from .hc import GatedResidual
+from .model import Qwen4ExpDecoderLayer
+from .weight import iter_mtp_weights
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+
+
+_PACK_EXPERT_RE = re.compile(r"^(mtp\.layers\.0\.mlp\.experts)\.0\.(gate|up|down)_proj$")
+
+
+def _pack_tensors(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Module name -> the pack's tensors: the HF name itself, but stacked experts with gate and up in one tensor."""
+    m = _PACK_EXPERT_RE.match(name)
+    if m is None:
+        return (name,)
+    return (f"{m[1]}.{'down_proj' if m[2] == 'down' else 'gate_up_proj'}",)
+
+
+def _iter_pack_weights(model_path: str, config: ModelConfig) -> Iterator[tuple[str, torch.Tensor]]:
+    """The dense tensors of a Strata MTP pack, named and fused as Qwen4ExpMTPDraft names them; its norms are stored zero-centered, as the layers read them."""
+    from freetoken.models.gguf.reader import gguf_tensor_index
+    from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+    index = gguf_tensor_index(model_path)
+    hidden, width = config.hidden_size, config.qwen4_args.hc_count * config.hidden_size
+
+    def t(name: str):
+        return index[f"mtp.{name}.weight"]
+
+    def linear(name: str, parts: tuple[str, ...], k: int, pad_to: int = 0):
+        infos = [t(p) for p in parts]
+        return fused(name, infos, [matrix(i) for i in infos], k, pad_to)
+
+    sa, mlp = "layers.0.self_attn", "layers.0.mlp"
+    for name in ("pre_fc_norm_embedding", "pre_fc_norm_hidden", f"{sa}.q_norm", f"{sa}.k_norm", f"{sa}.indexer.q_layernorm",
+                 f"{sa}.indexer.k_layernorm", f"{mlp}.gate", f"{mlp}.shared_expert_gate", "hyper_connection_mixer.hc_norm"):
+        yield f"{name}.weight", float_tensor(t(name))
+    yield from linear("fc_embedding", ("fc_embedding",), hidden)
+    yield from linear("fc_hidden", ("fc_hidden",), hidden)
+    yield from linear(f"{sa}.qkv_proj", tuple(f"{sa}.{p}_proj" for p in "qkv"), hidden)
+    yield from linear(f"{sa}.o_proj", (f"{sa}.o_proj",), int(t(f"{sa}.o_proj").ne[0]))
+    yield from linear(f"{sa}.indexer.index_qk_proj", (f"{sa}.indexer.index_qk_proj",), hidden)
+    yield from linear(f"{mlp}.shared_expert.gate_up_proj", (f"{mlp}.shared_expert.gate_proj", f"{mlp}.shared_expert.up_proj"), hidden)
+    yield from linear(f"{mlp}.shared_expert.down_proj", (f"{mlp}.shared_expert.down_proj",), config.shared_expert_intermediate_size)
+    for hc in ("layers.0.attn_hyper_connection", "layers.0.mlp_hyper_connection"):
+        down = (f"{hc}.input_mix_weight_down", f"{hc}.block_inject_weight")
+        yield from linear(f"{hc}.input_mix_weight_down_block_inject", down, width, pad_to=16)
+        yield from linear(f"{hc}.input_mix_weight_up", (f"{hc}.input_mix_weight_up",), config.qwen4_args.hc_lowrank)
+        yield f"{hc}.hc_norm.weight", float_tensor(t(f"{hc}.hc_norm"))
+    yield from linear("hyper_connection_mixer.input_mix_weight_down", ("hyper_connection_mixer.input_mix_weight_down",), width)
+    yield from linear("hyper_connection_mixer.input_mix_weight_up", ("hyper_connection_mixer.input_mix_weight_up",), config.qwen4_args.hc_lowrank)
+
+
+def _pack_expert_banks(model_path: str, device: torch.device):
+    """The pack's stacked experts as the resident banks of bank layer 0, in their packed ggml rows."""
+    from freetoken.layers.quantization import QuantKind
+    from freetoken.models.gguf.experts import resident_expert_bank
+    from freetoken.models.gguf.reader import gguf_tensor_index
+    from freetoken.moe.expert_banks import ExpertBanks
+
+    index = gguf_tensor_index(model_path)
+    sources = {role: [resident_expert_bank(index[f"mtp.layers.0.mlp.experts.{role}_proj"], device)] for role in ("gate_up", "down")}
+    return ExpertBanks("gguf", sources, kind=QuantKind.GGUF, kernel="ggml", layout=None)
+
+
+def _pack_quant(model_path: str):
+    from freetoken.layers.quantization import GGUFConfig, NameMap
+    from freetoken.models.gguf.reader import gguf_tensor_index
+    from freetoken.models.register import get_model_spec
+
+    packed = get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping
+    types = {name: info.ggml_type for name, info in gguf_tensor_index(model_path).items()}
+    return GGUFConfig(types, _pack_tensors, name_map=NameMap(packed=packed))
+
+
+def draft_config(config: ModelConfig) -> ModelConfig:
+    """The target config cut to one QSA layer (layer id 0) without PLE; its MoE is resident on the GPU whatever the target's strategy."""
+    group = next(g for g in config.attention_groups if isinstance(g, FullAttentionGroupConfig) and g.index_ratio > 1)
+    return dataclasses.replace(
+        config,
+        num_layers=1,
+        attention_groups=(dataclasses.replace(group, layer_ids=(0,), num_index_layers=1),),
+        qwen4_args=dataclasses.replace(config.qwen4_args, ple_layer_ids=()),
+        slot_states=(),
+        vision_config=None,
+        moe_strategy="fused",
+        decode_target="gpu",
+    )
+
+
+class Qwen4ExpMTPDraft(BaseOP):
+    attention_backend = "qsa_sparse"
+    draft_config = staticmethod(draft_config)
+
+    @staticmethod
+    def relay_width(config: ModelConfig) -> int:
+        return config.qwen4_args.hc_count * config.hidden_size
+
+    def __init__(self, config: ModelConfig, target_model, quant=None) -> None:
+        draft = draft_config(config)
+        # a Strata pack carries its own ggml types
+        self.config = draft = draft if quant is None else dataclasses.replace(draft, quant=quant)
+        hidden = draft.hidden_size
+        self.hc_count = draft.qwen4_args.hc_count
+        self.fc_embedding = LinearReplicated(hidden, hidden, has_bias=False, quant_config=draft.quant, prefix="mtp.fc_embedding")
+        self.fc_hidden = LinearReplicated(hidden, hidden, has_bias=False, quant_config=draft.quant, prefix="mtp.fc_hidden")
+        self.pre_fc_norm_embedding = GemmaPlusOneRMSNorm(hidden, eps=draft.rms_norm_eps)
+        self.pre_fc_norm_hidden = GemmaPlusOneRMSNorm(self.hc_count * hidden, eps=draft.rms_norm_eps)
+        self.layers = OPList([Qwen4ExpDecoderLayer(draft, 0, prefix="mtp.layers.0")])
+        self.hyper_connection_mixer = GatedResidual(draft, use_combine=False, prefix="mtp.hyper_connection_mixer")
+        # the target's own module; an underscore name keeps it out of this module's state dict
+        self._embed_tokens = target_model.model.embed_tokens
+
+    @classmethod
+    def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> Qwen4ExpMTPDraft:
+        from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks
+        from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+
+        pack = model_path.endswith(".gguf")
+        with torch.device("meta"), torch_dtype(dtype):
+            draft = cls(config, target_model, _pack_quant(model_path) if pack else None)
+        expected = draft.state_dict()
+        loaded = {}
+        weights = _iter_pack_weights(model_path, draft.config) if pack else iter_mtp_weights(model_path, device)
+        for name, tensor in weights:
+            if name not in expected:
+                raise ValueError(f"{model_path}: unexpected MTP tensor {name!r}")
+            loaded[name] = tensor.to(device=device, dtype=expected[name].dtype)
+        if missing := sorted(set(expected) - set(loaded)):
+            raise ValueError(f"{model_path}: MTP tensors missing: {missing}")
+        draft.load_state_dict(loaded)
+        # packed by the layer's own expert method, as the target's resident layers are
+        moe = draft.layers.op_list[0].mlp.experts
+        if pack:
+            banks = _pack_expert_banks(model_path, device)
+        else:
+            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts", device)
+            banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
+        attach_resident_banks([moe], banks)
+        # not the root: it reaches the target's embedding, which the target finalized already
+        for module in (draft.fc_embedding, draft.fc_hidden, draft.layers, draft.hyper_connection_mixer):
+            finalize_quant(module)
+        return draft
+
+    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self._embed_tokens.forward(input_ids)
+
+    def embed_batch(self, batch) -> torch.Tensor:
+        """The target batch's token embeddings with its multimodal rows merged in, as the target embeds them."""
+        return embed_input_ids(self._embed_tokens, batch.input_ids, batch)
+
+    def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor:
+        """``R_out`` of each row from the embedding of its token and the residual streams of the position before it."""
+        rows, hidden = embeds.shape
+        e = self.fc_embedding.forward(self.pre_fc_norm_embedding.forward(embeds))
+        h = self.pre_fc_norm_hidden.forward(prev_hidden)
+        r = self.fc_hidden.forward(h.view(rows * self.hc_count, hidden)).view(rows, self.hc_count, hidden)
+        r += e.unsqueeze(1)
+        return self.layers.op_list[0].forward(r.view(rows, -1), get_global_ctx().batch)
+
+    def head_input(self, hidden: torch.Tensor) -> torch.Tensor:
+        return self.hyper_connection_mixer.mix(hidden)[0]
+
+
+__all__ = ["Qwen4ExpMTPDraft", "draft_config"]
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 734eb80a83..7824ab0613 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -6,7 +6,7 @@
 * :func:`load_ple_table` -- the n-gram table (47.7 GiB fp8 or 95.4 GiB bf16), 128 checkpoint shards concatenated into one pinned :class:`HostBank`.
 * :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named under the checkpoint's dialect, for the offload cache's expert reader.
 
-Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``); ``model.visual.*`` is kept only when the model built the tower.
+Dropped: ``mtp.*`` (speculative head, read by :func:`iter_mtp_weights`); ``model.visual.*`` is kept only when the model built the tower.
 """
 
 from __future__ import annotations
@@ -296,6 +296,35 @@ def iter_weights(
         raise ValueError(f"checkpoint is missing tensors the quant config declares for {len(lines)} modules:\n  {shown}")
 
 
+# routed experts: expert_pieces.iter_extra_layer_expert_pieces
+_MTP_EXPERTS_PREFIX = "mtp.layers.0.mlp.experts."
+
+
+def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
+    """The ``mtp.*`` head but its routed experts, named as Qwen4ExpMTPDraft names it and fused like the target's dense weights."""
+    hf_config = cached_load_hf_config(model_path)
+    reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0]).packed_modules_mapping)
+    seen = False
+    for file in iter_weight_files(model_path):
+        with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+            for raw_name in f.keys():
+                if not raw_name.startswith("mtp.") or raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
+                    continue
+                seen = True
+                if raw_name.startswith(_MTP_EXPERTS_PREFIX):
+                    continue
+                tensor = f.get_tensor(raw_name)
+                emitted = reader.add(raw_name, tensor)
+                if emitted is None:
+                    yield raw_name[len("mtp."):], tensor
+                else:
+                    yield from ((n.removeprefix("mtp."), t) for n, t in emitted)
+    if not seen:
+        raise ValueError(f"{model_path} has no mtp.* weights")
+    if reader.pending:
+        raise ValueError(f"{model_path}: incomplete MTP projections: {reader.missing()}")
+
+
 def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
     """The vision tower alone, named as iter_weights names it."""
     for file in iter_weight_files(model_path):
@@ -468,6 +497,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec:
 __all__ = [
     "nvfp4_expert_spec",
     "PleTable",
+    "iter_mtp_weights",
     "iter_weights",
     "load_ple_table",
 ]
diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py
index 1a04a43934..7625b5abcf 100644
--- a/python/freetoken/models/register.py
+++ b/python/freetoken/models/register.py
@@ -43,6 +43,7 @@ class ModelSpec:
 # TODO: video once the processor samples frames; the tower already takes grid_thw with t > 1
 _QWEN_VL_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _QWEN3_5_DRAFTS = (("mtp", "freetoken.models.qwen3_5_moe.mtp:Qwen3_5MTPDraft"),)
+_QWEN4_EXP_DRAFTS = (("mtp", "freetoken.models.qwen4_exp.mtp:Qwen4ExpMTPDraft"),)
 # Fused projections split back into their HF leaves so the quantization_config lookup sees the stored names.
 _DENSE_PACKED = (
     ("qkv_proj", ("q_proj", "k_proj", "v_proj")),
@@ -343,6 +344,7 @@ class ModelSpec:
         mm_processor=_QWEN_VL_PROCESSOR,
         encoders=_QWEN_VL_ENCODERS,
         packed_modules_mapping=_QWEN4_EXP_PACKED,
+        spec_drafts=_QWEN4_EXP_DRAFTS,
     ),
     # llama.cpp qwen4exp GGUF: the same model classes with every tensor in its packed ggml type; the vision tower comes from an mmproj file and runs in bf16
     "Qwen4ExpGGUFForCausalLM": ModelSpec(
@@ -354,6 +356,7 @@ class ModelSpec:
         packed_modules_mapping=_QWEN4_EXP_GGUF_PACKED,
         mm_processor=_QWEN_VL_PROCESSOR,
         encoders=_QWEN_VL_ENCODERS,
+        spec_drafts=_QWEN4_EXP_DRAFTS,
     ),
     # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the
     # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the
diff --git a/python/freetoken/spec/mtp/chain_graph.py b/python/freetoken/spec/mtp/chain_graph.py
index 0634c1b07d..5fa61ba8ce 100644
--- a/python/freetoken/spec/mtp/chain_graph.py
+++ b/python/freetoken/spec/mtp/chain_graph.py
@@ -1,4 +1,4 @@
-"""MTP chain steps as one CUDA graph per batch size: the k-1 draft decode steps on the triton backend, with each step's metadata built on the device."""
+"""MTP chain steps as one CUDA graph per batch size: the k-1 draft decode steps, with each step's metadata built on the device."""
 
 from __future__ import annotations
 
@@ -6,7 +6,8 @@
 
 import torch
 
-from freetoken.attention.triton import TritonMetadata
+from freetoken.attention.qsa_sparse import QSASparseAttnBackend, QSASparseMetadata
+from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata
 from freetoken.core import Batch
 from freetoken.kernel.triton.attention import gather_kv_indices
 
@@ -16,13 +17,84 @@
     from .proposer import MTPProposer
 
 
+class _TritonSteps:
+    """Each step's flat KV slot list, gathered from the page table in the graph."""
+
+    def __init__(self, ctx, steps: int, max_bs: int) -> None:
+        self.page_table = ctx.page_table
+        self.indptr = torch.zeros(steps, max_bs + 1, dtype=torch.int32, device=ctx.device)
+        self.indices = torch.zeros(steps, max_bs * ctx.page_table.shape[1], dtype=torch.int32, device=ctx.device)
+        self.arange = torch.arange(max_bs + 1, dtype=torch.int32, device=ctx.device)
+
+    def metadata(self, bs: int, step: int, positions: torch.Tensor) -> TritonMetadata:
+        return TritonMetadata(
+            cu_seqlens_q_gpu=self.arange[: bs + 1],
+            indptr=self.indptr[step, : bs + 1],
+            indices=self.indices[step],
+            q_to_req=self.arange[:bs],
+            q_positions=positions,
+            is_decode=True,
+            prefix_lens=positions,
+            max_q_len=1,
+        )
+
+    def stage(self, bs: int, rows: torch.Tensor, entry: torch.Tensor) -> None:
+        for step in range(self.indptr.shape[0]):
+            self.indptr[step, 1 : bs + 1] = torch.cumsum(entry + step + 1, 0)
+        gather_kv_indices(self.page_table, rows, self.indptr[:, : bs + 1], self.indices)
+
+
+class _QSASteps:
+    """The block table and ring slots of the rows, shared by the steps, and one KV length row per step."""
+
+    def __init__(self, backend: QSASparseAttnBackend, ctx, steps: int, bs_list: List[int]) -> None:
+        # static scratch for the sparse kernels, so the captured graphs keep their addresses
+        backend.init_capture_graph(ctx.page_table.shape[1], bs_list)
+        max_bs = max(bs_list)
+        self.page_table = ctx.page_table
+        self.page_size = backend.page_size
+        pages = -(-ctx.page_table.shape[1] // self.page_size)
+        self.block_table = torch.zeros(max_bs, pages, dtype=torch.int32, device=ctx.device)
+        self.slots = torch.zeros(max_bs, dtype=torch.int32, device=ctx.device)
+        self.kv_len = torch.zeros(steps, max_bs, dtype=torch.int32, device=ctx.device)
+        self.arange = torch.arange(max_bs + 1, dtype=torch.int32, device=ctx.device)
+        self.cpu = torch.zeros(max_bs + 1, dtype=torch.int32)
+
+    def metadata(self, bs: int, step: int, positions: torch.Tensor) -> QSASparseMetadata:
+        return QSASparseMetadata(
+            is_decode=True,
+            last_indices=self.arange[:bs],
+            qo_indptr_cpu=self.cpu[: bs + 1],
+            kv_len_cpu=self.cpu[:bs],
+            token_to_req=self.arange[:bs],
+            cu_seqlens=self.arange[: bs + 1],
+            seq_lens=self.kv_len[step, :bs],
+            ring_slots=self.slots[:bs],
+            block_table=self.block_table[:bs],
+        )
+
+    def stage(self, bs: int, rows: torch.Tensor, entry: torch.Tensor) -> None:
+        self.block_table[:bs] = self.page_table[rows, :: self.page_size] // self.page_size
+        self.slots[:bs] = rows
+        for step in range(self.kv_len.shape[0]):
+            self.kv_len[step, :bs] = entry + step + 1
+
+
 class ChainGraphs:
     """Replays ``MTPProposer.propose``: the inputs are each request's table row, entry position and mrope delta, everything else derives from them in the graph."""
 
+    @staticmethod
+    def supports(backend) -> bool:
+        return isinstance(backend, (TritonAttentionBackend, QSASparseAttnBackend))
+
     def __init__(self, proposer: MTPProposer, bs_list: List[int], pool) -> None:
         self.p = proposer
         ctx = proposer.ctx
         device, steps, max_bs = ctx.device, proposer.k - 1, max(bs_list)
+        backend = proposer.chain_backend
+        self.backend_steps = (
+            _QSASteps(backend, ctx, steps, bs_list) if isinstance(backend, QSASparseAttnBackend) else _TritonSteps(ctx, steps, max_bs)
+        )
         self.dummy_row = ctx.dummy_req.table_idx
         # rows: table row, entry position, mrope delta
         self.inputs = torch.zeros(3, max_bs, dtype=torch.int64, device=device)
@@ -30,9 +102,6 @@ def __init__(self, proposer: MTPProposer, bs_list: List[int], pool) -> None:
         self.positions = torch.zeros(steps, max_bs, dtype=torch.int32, device=device)
         self.mrope = torch.zeros(steps, 3, max_bs, dtype=torch.int32, device=device) if proposer.mrope else None
         self.out_loc = torch.zeros(steps, max_bs, dtype=torch.int32, device=device)
-        self.indptr = torch.zeros(steps, max_bs + 1, dtype=torch.int32, device=device)
-        self.indices = torch.zeros(steps, max_bs * ctx.page_table.shape[1], dtype=torch.int32, device=device)
-        self.arange = torch.arange(max_bs + 1, dtype=torch.int32, device=device)
         self.graphs: Dict[int, torch.cuda.CUDAGraph] = {}
         # the metadata of these batches holds the decode scratch the graphs read
         self._batches: Dict[int, List[Batch]] = {}
@@ -56,16 +125,7 @@ def _batch(self, bs: int, step: int) -> Batch:
         if self.mrope is not None:
             batch.mrope_positions = self.mrope[step, :, :bs]
         batch.out_loc = self.out_loc[step, :bs]
-        batch.attn_metadata = TritonMetadata(
-            cu_seqlens_q_gpu=self.arange[: bs + 1],
-            indptr=self.indptr[step, : bs + 1],
-            indices=self.indices[step],
-            q_to_req=self.arange[:bs],
-            q_positions=self.positions[step, :bs],
-            is_decode=True,
-            prefix_lens=self.positions[step, :bs],
-            max_q_len=1,
-        )
+        batch.attn_metadata = self.backend_steps.metadata(bs, step, batch.positions)
         return batch
 
     def _run(self, bs: int, batches: List[Batch]) -> None:
@@ -79,8 +139,7 @@ def _run(self, bs: int, batches: List[Batch]) -> None:
             if self.mrope is not None:
                 self.mrope[step, :, :bs] = (pos + delta).unsqueeze(0)
             self.out_loc[step, :bs] = page_table[rows, pos]
-            self.indptr[step, 1 : bs + 1] = torch.cumsum(pos + 1, 0)
-        gather_kv_indices(page_table, rows, self.indptr[:, : bs + 1], self.indices)
+        self.backend_steps.stage(bs, rows, entry)
         token, hidden = p.relay.d1[rows], p.relay.hidden[rows]
         self.drafts[:bs, 0] = token
         for step, batch in enumerate(batches):
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 5eee3a2269..ca3bebd6c2 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -101,13 +101,17 @@ class MTPProposer(Proposer):
     @classmethod
     def adjust_config(cls, config, override) -> None:
         from freetoken.checkpoint.ftw import is_ftw_checkpoint
+        from freetoken.models.gguf.reader import gguf_tensor_index, is_gguf_path
 
         super().adjust_config(config, override)
         if "mtp" not in dict(config.model_spec.spec_drafts):
             raise ValueError(f"MTP is not supported for {config.hf_config.architectures[0]} yet")
+        path = config.speculative_draft_model_path or config.model_path
         # FTW keeps only what the target's iter_weights yields, which drops mtp.*
-        if is_ftw_checkpoint(config.speculative_draft_model_path or config.model_path):
+        if is_ftw_checkpoint(path):
             raise ValueError("MTP cannot load its draft from an FTW checkpoint yet")
+        if is_gguf_path(path) and not any(name.startswith("mtp.") for name in gguf_tensor_index(path)):
+            raise ValueError(f"{path} has no mtp.* tensors; pass a GGUF that carries them (a Strata MTP pack) with --speculative-draft-model-path")
 
     @classmethod
     def _draft_class(cls, config) -> type[MTPDraft]:
@@ -153,7 +157,7 @@ def attach(self, ctx: DraftContext) -> None:
         if self.kv is None:
             self.kv = create_kvcache_pool(
                 model_config=cfg, num_pages=ctx.num_pages + 1, page_size=self.config.page_size, dtype=self.dtype,
-                device=ctx.device, num_req_slots=rows,
+                device=ctx.device, num_req_slots=rows, num_speculative_tokens=self.k,
             )
         width = self.draft.relay_width(self.config.model_config)
         self.relay = HiddenRelay(rows, width, self.dtype, ctx.device)
@@ -248,7 +252,7 @@ def capture_graphs(self, bs_list: List[int], pool) -> None:
             ctx_override=dict(kv_cache=self.kv, attn_backend=self.extend_backend, hide_target_state=True),
             pool=pool,
         )
-        if self.k > 1:
+        if self.k > 1 and ChainGraphs.supports(self.chain_backend):
             self.chain_graphs = ChainGraphs(self, bs_list, self.extend_graphs.graph_pool)
 
     def propose(self, reqs: List[Req], rows: torch.Tensor) -> torch.Tensor:
@@ -294,6 +298,7 @@ def post_verify(self, batch: Batch, out: torch.Tensor, emitted: torch.Tensor) ->
         draft.positions = batch.positions
         draft.mrope_positions = batch.mrope_positions
         draft.out_loc = batch.out_loc
+        draft.active_table_idx = batch.active_table_idx
         self.picked_rows[:bs] = torch.arange(bs, device=self.device) * k1 + emitted
         self.extend_backend.prepare_metadata(draft)
         graphs = self.extend_graphs

From 5de3d77f59f987cfaaee8c5dc62c7c6f12ea7c8f Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 03:28:45 +0000
Subject: [PATCH 05/20] docs(cli): document the speculative decoding flags

---
 docs/cli.md | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/docs/cli.md b/docs/cli.md
index feda4f7a34..5b3cb057a3 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -98,6 +98,16 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf
 | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) |
 | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap |
 
+### Speculative decoding
+
+| Flag | Default | Meaning |
+|---|---|---|
+| `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF checkpoints, and Qwen3.8-Flash-Next HF or GGUF |
+| `--speculative-num-steps`, `--spec-num-steps` | the method's default (3 for `mtp`) | Draft tokens per decode step, 1 to 8 for `mtp` |
+| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from; a GGUF target without `mtp.*` tensors needs a GGUF that carries them (a Strata MTP pack) |
+
+Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` also refuses FTW checkpoints and Qwen3.5 / Qwen3.6 GGUF.
+
 ### API behaviour
 
 | Flag | Default | Meaning |

From 495643f9250383c7d14b0173dc7829e9d664f345 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 03:59:52 +0000
Subject: [PATCH 06/20] fix(spec): refuse MTP draft sources the draft cannot
 read before the target loads

---
 docs/cli.md                                |  4 ++--
 python/freetoken/models/qwen3_5_moe/mtp.py |  1 +
 python/freetoken/models/qwen4_exp/mtp.py   |  1 +
 python/freetoken/moe/expert_pieces.py      |  9 ++++++++-
 python/freetoken/spec/mtp/proposer.py      | 23 +++++++++++++++++++---
 5 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 5b3cb057a3..2df0308322 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -104,9 +104,9 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf
 |---|---|---|
 | `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF checkpoints, and Qwen3.8-Flash-Next HF or GGUF |
 | `--speculative-num-steps`, `--spec-num-steps` | the method's default (3 for `mtp`) | Draft tokens per decode step, 1 to 8 for `mtp` |
-| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from; a GGUF target without `mtp.*` tensors needs a GGUF that carries them (a Strata MTP pack) |
+| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from. A Qwen3.8-Flash-Next target also takes a Strata MTP pack (a GGUF that keeps the checkpoint's `mtp.*` tensors under their HF names), and a GGUF target needs one; the `nextn` MTP blocks llama.cpp writes are not read; a separate checkpoint folder is read with the target's quantization config |
 
-Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` also refuses FTW checkpoints and Qwen3.5 / Qwen3.6 GGUF.
+Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` also refuses Qwen3.5 / Qwen3.6 GGUF targets and an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors.
 
 ### API behaviour
 
diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index 4799751f1f..b51e400b83 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -39,6 +39,7 @@ class Qwen3_5MTPDraft(BaseOP):
 
     # the draft MTPProposer picks from the target's backends
     attention_backend = None
+    reads_gguf_pack = False
     draft_config = staticmethod(draft_config)
 
     @staticmethod
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index b3b7f6611b..ece6315b20 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -119,6 +119,7 @@ def draft_config(config: ModelConfig) -> ModelConfig:
 
 class Qwen4ExpMTPDraft(BaseOP):
     attention_backend = "qsa_sparse"
+    reads_gguf_pack = True
     draft_config = staticmethod(draft_config)
 
     @staticmethod
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index e55a9d0ca0..e3049f6a6d 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -179,9 +179,16 @@ def iter_extra_layer_expert_pieces(model_path: str, prefix: str, device: torch.d
     per_expert_re = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.(gate|up|down)_proj\.(weight|weight_scale_inv)$")
     stacked_re = re.compile(rf"^{re.escape(prefix)}\.(gate_up|down)_proj$")
     names: dict[str, str] = {}
+    unknown: list[str] = []
     for file in iter_weight_files(model_path):
         with safetensors.safe_open(file, framework="pt", device="cpu") as f:
-            names.update((name, file) for name in f.keys() if per_expert_re.match(name) or stacked_re.match(name))
+            for name in f.keys():
+                if per_expert_re.match(name) or stacked_re.match(name):
+                    names[name] = file
+                elif name.startswith(prefix + "."):
+                    unknown.append(name)
+    if unknown:
+        raise ValueError(f"{model_path}: the experts under {prefix} use a format this reader does not support, e.g. {sorted(unknown)[:3]}")
     if not names:
         raise ValueError(f"{model_path} has no routed experts under {prefix}")
     stacked = {stacked_re.match(n)[1]: f for n, f in names.items() if stacked_re.match(n)}
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index ca3bebd6c2..2a31cb6cf5 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -24,6 +24,8 @@ class MTPDraft(Protocol):
     config: ModelConfig
     # None: the proposer picks from the target's backends
     attention_backend: ClassVar[str | None]
+    # True: the draft can also load from a GGUF that carries mtp.* (a Strata MTP pack)
+    reads_gguf_pack: ClassVar[bool]
 
     @staticmethod
     def draft_config(config: ModelConfig) -> ModelConfig: ...
@@ -100,18 +102,33 @@ class MTPProposer(Proposer):
 
     @classmethod
     def adjust_config(cls, config, override) -> None:
+        import os
+
         from freetoken.checkpoint.ftw import is_ftw_checkpoint
         from freetoken.models.gguf.reader import gguf_tensor_index, is_gguf_path
+        from freetoken.models.loader import safetensors_weight_map
 
         super().adjust_config(config, override)
+        arch = config.hf_config.architectures[0]
         if "mtp" not in dict(config.model_spec.spec_drafts):
-            raise ValueError(f"MTP is not supported for {config.hf_config.architectures[0]} yet")
+            raise ValueError(f"MTP is not supported for {arch} yet")
         path = config.speculative_draft_model_path or config.model_path
+        # checked here because a draft that cannot load otherwise fails only after the whole target has loaded
+        if config.speculative_draft_model_path and (os.path.isabs(path) or path.endswith(".gguf")) and not os.path.exists(path):
+            raise ValueError(f"--speculative-draft-model-path {path} does not exist")
         # FTW keeps only what the target's iter_weights yields, which drops mtp.*
         if is_ftw_checkpoint(path):
             raise ValueError("MTP cannot load its draft from an FTW checkpoint yet")
-        if is_gguf_path(path) and not any(name.startswith("mtp.") for name in gguf_tensor_index(path)):
-            raise ValueError(f"{path} has no mtp.* tensors; pass a GGUF that carries them (a Strata MTP pack) with --speculative-draft-model-path")
+        if is_gguf_path(path):
+            if not cls._draft_class(config).reads_gguf_pack:
+                raise ValueError(f"MTP for {arch} loads its draft from safetensors, not from the GGUF {path}")
+            if not any(name.startswith("mtp.") for name in gguf_tensor_index(path)):
+                raise ValueError(f"{path} has no mtp.* tensors; pass a GGUF that carries them (a Strata MTP pack) with --speculative-draft-model-path")
+        elif is_gguf_path(config.model_path):
+            # a draft read from a folder takes the target's quant config, which for a GGUF target describes GGUF tensors
+            raise ValueError(f"a GGUF target needs a GGUF draft that carries mtp.* tensors (a Strata MTP pack), not {path}")
+        elif os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)):
+            raise ValueError(f"{path} has no mtp.* tensors")
 
     @classmethod
     def _draft_class(cls, config) -> type[MTPDraft]:

From 8815c0bb2f869aea87f80cceaa8e3b7a7ce61eaa Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 03:59:52 +0000
Subject: [PATCH 07/20] docs(qwen4_exp): describe the verify path in the PLE
 docstrings

---
 python/freetoken/models/qwen4_exp/ple.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py
index 9114eef602..bc2afe1f60 100644
--- a/python/freetoken/models/qwen4_exp/ple.py
+++ b/python/freetoken/models/qwen4_exp/ple.py
@@ -320,8 +320,8 @@ def build_ple_metadata(
     """Build ``PLEMetadata`` from a scheduler batch.
 
     The n-gram context is per-request device state (``ple_ngram_ctx`` [num_slots, ngram_size-1],
-    rolled forward once per forward by ``commit_ngram_context``), so it never lags the sampled
-    token under overlap scheduling and follows the slot on COW/snapshot. A decode batch takes its
+    rolled forward once per forward by ``commit_ngram_context``, or by ``PLENgramVerify.commit`` after a verify), so it never lags the sampled
+    token under overlap scheduling and follows the slot on COW/snapshot. A decode or verify batch takes its
     slots straight off ``batch.fla_metadata.cache_indices`` (persistent under a graph), so the
     build is capture-safe and sync-free. Reuses ``batch.fla_metadata`` (slots / indptr / fresh
     mask) when the scheduler built it.
@@ -384,7 +384,7 @@ def build_ple_metadata(
 def commit_ngram_context(meta: PLEMetadata, fla, context_pool: torch.Tensor | None = None) -> None:
     """Roll each request's ``ple_ngram_ctx`` forward past this forward's tokens.
 
-    Called ONCE per forward after every PLE layer ran (the layers only read the context);
+    Called ONCE per prefill or decode forward after every PLE layer ran (the layers only read the context; a verify commits through ``PLENgramVerify.commit``);
     also writes the boundary-aligned window to the track slot so a donated snapshot restores
     the context together with the conv state. Pure device arithmetic, capture-safe.
     """
@@ -654,7 +654,7 @@ def _read_state(
     def _short_conv(
         self, x: torch.Tensor, meta: PLEMetadata, states: torch.Tensor
     ) -> torch.Tensor:
-        """silu of the dilated depthwise conv over [state | x], and roll the per-request state."""
+        """silu of the dilated depthwise conv over [state | x], and roll the per-request state (a verify leaves it to ``PLEConvVerify.commit``)."""
         if meta.is_decode:
             return self._decode_conv(x, meta, states)
         if meta.is_verify:

From 144ea92608c1d217b07e82ac95514c74c00193d8 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 05:49:15 +0000
Subject: [PATCH 08/20] fix(spec): refuse a missing relative MTP draft path
 before the target loads

---
 python/freetoken/spec/mtp/proposer.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 2a31cb6cf5..3957e3cda4 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -114,7 +114,9 @@ def adjust_config(cls, config, override) -> None:
             raise ValueError(f"MTP is not supported for {arch} yet")
         path = config.speculative_draft_model_path or config.model_path
         # checked here because a draft that cannot load otherwise fails only after the whole target has loaded
-        if config.speculative_draft_model_path and (os.path.isabs(path) or path.endswith(".gguf")) and not os.path.exists(path):
+        # a HF repo id is "name" or "org/name"; anything else names a local path
+        local = os.path.isabs(path) or path.startswith(".") or path.count("/") > 1 or path.endswith(".gguf")
+        if config.speculative_draft_model_path and local and not os.path.exists(path):
             raise ValueError(f"--speculative-draft-model-path {path} does not exist")
         # FTW keeps only what the target's iter_weights yields, which drops mtp.*
         if is_ftw_checkpoint(path):

From 674d671480ce4ce3c183f1d9a04a65cc28873064 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 05:49:15 +0000
Subject: [PATCH 09/20] fix(moe): read MTP head experts on the host and require
 both stacked tensors

---
 python/freetoken/models/qwen3_5_moe/mtp.py | 2 +-
 python/freetoken/models/qwen4_exp/mtp.py   | 2 +-
 python/freetoken/moe/expert_pieces.py      | 8 +++++---
 tests/models/test_qwen3_5_moe_weight.py    | 2 +-
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index b51e400b83..e9b954132f 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -76,7 +76,7 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
 
             # packed by the layer's own expert method, as the target's resident layers are
             moe = draft.layers.op_list[0].mlp.experts
-            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts", device)
+            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts")
             banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
             attach_resident_banks([moe], banks)
         # not the root: it reaches the target's embedding, which the target finalized already
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index ece6315b20..5aec06b67f 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -164,7 +164,7 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
         if pack:
             banks = _pack_expert_banks(model_path, device)
         else:
-            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts", device)
+            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts")
             banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
         attach_resident_banks([moe], banks)
         # not the root: it reaches the target's embedding, which the target finalized already
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index e3049f6a6d..62f829c885 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -170,7 +170,7 @@ def per_expert_pieces(
         raise ValueError(f"Incomplete expert tensors in checkpoint (layer, expert): {missing}")
 
 
-def iter_extra_layer_expert_pieces(model_path: str, prefix: str, device: torch.device) -> Iterator[Piece]:
+def iter_extra_layer_expert_pieces(model_path: str, prefix: str) -> Iterator[Piece]:
     """The routed experts stored under ``prefix``, one MoE layer outside the model's stack (an MTP head), as pieces of bank layer 0: stacked ``gate_up``/``down``, or ``gate``/``up``/``down`` per expert with ``_scale`` companions in a block-fp8 checkpoint."""
     import safetensors
 
@@ -193,9 +193,11 @@ def iter_extra_layer_expert_pieces(model_path: str, prefix: str, device: torch.d
         raise ValueError(f"{model_path} has no routed experts under {prefix}")
     stacked = {stacked_re.match(n)[1]: f for n, f in names.items() if stacked_re.match(n)}
     if stacked:
+        if missing := {"gate_up", "down"} - set(stacked):
+            raise ValueError(f"{model_path}: stacked experts under {prefix} lack {sorted(missing)}")
         piece = {}
         for role, file in stacked.items():
-            with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+            with safetensors.safe_open(file, framework="pt", device="cpu") as f:
                 piece[role] = f.get_tensor(f"{prefix}.{role}_proj")
         return iter([(0, 0, piece["down"].shape[0], piece)])
 
@@ -205,7 +207,7 @@ def locate(name: str) -> tuple[int, int, str]:
 
     def tensors():
         for file in dict.fromkeys(names.values()):
-            with safetensors.safe_open(file, framework="pt", device=str(device)) as f:
+            with safetensors.safe_open(file, framework="pt", device="cpu") as f:
                 yield from ((name, f.get_tensor(name)) for name, where in names.items() if where == file)
 
     scaled = any(name.endswith("_scale_inv") for name in names)
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index 2ed0ad9229..9ad3781d4a 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -650,7 +650,7 @@ def test_mtp_head_fills_the_draft_and_its_experts(checkpoint):
         assert tensor.shape == state[key].shape, key
         if key.endswith(".weight"):
             assert tensor.dtype is state[key].dtype, key
-    pieces = list(iter_extra_layer_expert_pieces(folder, f"{MTP}.mlp.experts", torch.device("cpu")))
+    pieces = list(iter_extra_layer_expert_pieces(folder, f"{MTP}.mlp.experts"))
     if name == "bf16":
         assert [p[:3] for p in pieces] == [(0, 0, E)]
         assert torch.equal(pieces[0][3]["gate_up"], raw[f"{MTP}.mlp.experts.gate_up_proj"])

From 85bace806aeef58057af3c020561231499a8cb19 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 21:41:24 +0000
Subject: [PATCH 10/20] refactor(spec): tell the final prefill chunk from
 prompt_len instead of the scheduler's ChunkedReq

---
 python/freetoken/spec/mtp/proposer.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 3957e3cda4..f5223bc082 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -333,8 +333,6 @@ def post_verify(self, batch: Batch, out: torch.Tensor, emitted: torch.Tensor) ->
 
     def on_prefill(self, batch: Batch, next_tokens_gpu: torch.Tensor) -> None:
         """Draft the chunk's positions: row p pairs token p+1 with the target hidden at p; a final chunk adds the sampled token's row and relays it."""
-        from freetoken.scheduler.prefill import ChunkedReq
-
         sink, self._prefill_sink = self._prefill_sink, None
         hidden = sink.hidden
         # image rows take the embeddings the target used
@@ -349,7 +347,7 @@ def on_prefill(self, batch: Batch, next_tokens_gpu: torch.Tensor) -> None:
             state = self._prefill_state.get(row)
             # a continuation chunk resumes where the last one stopped; a first chunk starts at its prefix hit, whose draft KV the tree already holds
             frontier = state[1] if state is not None and state[0] == req.uid else s
-            final = not isinstance(req, ChunkedReq)
+            final = e == req.prompt_len
             end = e if final else e - 1
             if frontier < end:
                 # draft row p embeds the token at p+1

From 41e4f679508c9ac063cd4addfd152584fc27b985 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Thu, 29 Oct 2026 23:52:24 +0000
Subject: [PATCH 11/20] refactor(spec): let each MTP draft family check its own
 weight source

---
 python/freetoken/models/qwen3_5_moe/mtp.py | 12 +++++++++++-
 python/freetoken/models/qwen4_exp/mtp.py   | 19 +++++++++++++++++--
 python/freetoken/spec/mtp/proposer.py      | 20 ++++++--------------
 3 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index e9b954132f..a181c5fc04 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import dataclasses
+import os
 from typing import TYPE_CHECKING
 
 import torch
@@ -39,13 +40,22 @@ class Qwen3_5MTPDraft(BaseOP):
 
     # the draft MTPProposer picks from the target's backends
     attention_backend = None
-    reads_gguf_pack = False
     draft_config = staticmethod(draft_config)
 
     @staticmethod
     def relay_width(config: ModelConfig) -> int:
         return config.hidden_size
 
+    @staticmethod
+    def check_source(path: str, target_path: str) -> None:
+        from freetoken.models.gguf.reader import is_gguf_path
+        from freetoken.models.loader import safetensors_weight_map
+
+        if is_gguf_path(path):
+            raise ValueError(f"the Qwen3.5 MTP draft loads from safetensors, not from the GGUF {path}")
+        if os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)):
+            raise ValueError(f"{path} has no mtp.* tensors")
+
     def __init__(self, config: ModelConfig, target_model) -> None:
         self.config = draft_config(config)
         hidden = config.hidden_size
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index 5aec06b67f..f0ff34566f 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import dataclasses
+import os
 import re
 from typing import TYPE_CHECKING, Iterator
 
@@ -119,13 +120,26 @@ def draft_config(config: ModelConfig) -> ModelConfig:
 
 class Qwen4ExpMTPDraft(BaseOP):
     attention_backend = "qsa_sparse"
-    reads_gguf_pack = True
     draft_config = staticmethod(draft_config)
 
     @staticmethod
     def relay_width(config: ModelConfig) -> int:
         return config.qwen4_args.hc_count * config.hidden_size
 
+    @staticmethod
+    def check_source(path: str, target_path: str) -> None:
+        from freetoken.models.gguf.reader import gguf_tensor_index, is_gguf_path
+        from freetoken.models.loader import safetensors_weight_map
+
+        if is_gguf_path(path):
+            if not any(name.startswith("mtp.") for name in gguf_tensor_index(path)):
+                raise ValueError(f"{path} has no mtp.* tensors; pass a GGUF that carries them (a Strata MTP pack) with --speculative-draft-model-path")
+        elif is_gguf_path(target_path):
+            # a draft read from a folder takes the target's quant config, which for a GGUF target describes GGUF tensors
+            raise ValueError(f"a GGUF target needs a GGUF draft that carries mtp.* tensors (a Strata MTP pack), not {path}")
+        elif os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)):
+            raise ValueError(f"{path} has no mtp.* tensors")
+
     def __init__(self, config: ModelConfig, target_model, quant=None) -> None:
         draft = draft_config(config)
         # a Strata pack carries its own ggml types
@@ -143,10 +157,11 @@ def __init__(self, config: ModelConfig, target_model, quant=None) -> None:
 
     @classmethod
     def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> Qwen4ExpMTPDraft:
+        from freetoken.models.gguf.reader import is_gguf_path
         from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks
         from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
 
-        pack = model_path.endswith(".gguf")
+        pack = is_gguf_path(model_path)
         with torch.device("meta"), torch_dtype(dtype):
             draft = cls(config, target_model, _pack_quant(model_path) if pack else None)
         expected = draft.state_dict()
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index f5223bc082..6a3003d308 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -24,8 +24,6 @@ class MTPDraft(Protocol):
     config: ModelConfig
     # None: the proposer picks from the target's backends
     attention_backend: ClassVar[str | None]
-    # True: the draft can also load from a GGUF that carries mtp.* (a Strata MTP pack)
-    reads_gguf_pack: ClassVar[bool]
 
     @staticmethod
     def draft_config(config: ModelConfig) -> ModelConfig: ...
@@ -35,6 +33,11 @@ def relay_width(config: ModelConfig) -> int:
         """Width of the hidden a draft row reads and relays to the next step."""
         ...
 
+    @staticmethod
+    def check_source(path: str, target_path: str) -> None:
+        """Raise ValueError if ``path`` cannot give this draft its weights; called before the target loads."""
+        ...
+
     @classmethod
     def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> MTPDraft: ...
 
@@ -105,8 +108,6 @@ def adjust_config(cls, config, override) -> None:
         import os
 
         from freetoken.checkpoint.ftw import is_ftw_checkpoint
-        from freetoken.models.gguf.reader import gguf_tensor_index, is_gguf_path
-        from freetoken.models.loader import safetensors_weight_map
 
         super().adjust_config(config, override)
         arch = config.hf_config.architectures[0]
@@ -121,16 +122,7 @@ def adjust_config(cls, config, override) -> None:
         # FTW keeps only what the target's iter_weights yields, which drops mtp.*
         if is_ftw_checkpoint(path):
             raise ValueError("MTP cannot load its draft from an FTW checkpoint yet")
-        if is_gguf_path(path):
-            if not cls._draft_class(config).reads_gguf_pack:
-                raise ValueError(f"MTP for {arch} loads its draft from safetensors, not from the GGUF {path}")
-            if not any(name.startswith("mtp.") for name in gguf_tensor_index(path)):
-                raise ValueError(f"{path} has no mtp.* tensors; pass a GGUF that carries them (a Strata MTP pack) with --speculative-draft-model-path")
-        elif is_gguf_path(config.model_path):
-            # a draft read from a folder takes the target's quant config, which for a GGUF target describes GGUF tensors
-            raise ValueError(f"a GGUF target needs a GGUF draft that carries mtp.* tensors (a Strata MTP pack), not {path}")
-        elif os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)):
-            raise ValueError(f"{path} has no mtp.* tensors")
+        cls._draft_class(config).check_source(path, config.model_path)
 
     @classmethod
     def _draft_class(cls, config) -> type[MTPDraft]:

From d04a86be607c56e8f646010725e5bc414c32eb0b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 01:48:38 +0000
Subject: [PATCH 12/20] fix(qwen4_exp): size the disk PLE eager staging for
 decode and verify rows

---
 python/freetoken/models/qwen4_exp/model.py    | 3 ++-
 python/freetoken/models/qwen4_exp/ple_disk.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py
index 59be84fde8..623389825a 100644
--- a/python/freetoken/models/qwen4_exp/model.py
+++ b/python/freetoken/models/qwen4_exp/model.py
@@ -198,7 +198,8 @@ def load_host_tables(self, engine_config) -> int:
                 constants,
                 # verify graphs stop at max_running_req requests, decode graphs at cuda_graph_max_bs
                 max_graph_rows=engine_config.max_running_req * rows if rows > 1 else max(256, engine_config.cuda_graph_max_bs or 0),
-                max_extend_tokens=engine_config.max_extend_tokens,
+                # an eager decode or verify step stages max_running_req * rows tokens, which can exceed a prefill chunk
+                max_eager_tokens=max(engine_config.max_extend_tokens, engine_config.max_running_req * rows),
             )
             self._ple_table = disk_table
             for ple in ple_layers:
diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py
index 65217608f8..bc0b4927fa 100644
--- a/python/freetoken/models/qwen4_exp/ple_disk.py
+++ b/python/freetoken/models/qwen4_exp/ple_disk.py
@@ -139,7 +139,7 @@ def __init__(
         hash_constants: dict,
         *,
         max_graph_rows: int = 256,
-        max_extend_tokens: int = 8192,
+        max_eager_tokens: int = 8192,
         dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         from freetoken.kernel.row_store import PleStore
@@ -188,7 +188,7 @@ def __init__(
         # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the
         # dispatch context after the step's fill: a record inside the capture would not be observable)
         self._graph_consumed = torch.cuda.Event()
-        eager_bytes = max_extend_tokens * self._token_bytes
+        eager_bytes = max_eager_tokens * self._token_bytes
         # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the
         # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy
         # that last read it has run (``_eager_read[i]``, recorded right after that copy)

From 801ada1def88b7b442558f3f7d96f54d75ab39b5 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 01:48:38 +0000
Subject: [PATCH 13/20] perf(kernel): read the KV-index gather's chunk count
 from the grid, not a constexpr

---
 python/freetoken/kernel/triton/attention.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py
index e56a9184a5..28effa5ea2 100644
--- a/python/freetoken/kernel/triton/attention.py
+++ b/python/freetoken/kernel/triton/attention.py
@@ -1069,7 +1069,6 @@ def _gather_kv_indices_kernel(
     pt_stride,
     indptr_stride,
     out_stride,
-    NUM_CHUNKS: tl.constexpr,
     BLOCK: tl.constexpr,
 ):
     batch_id = tl.program_id(1)
@@ -1078,7 +1077,7 @@ def _gather_kv_indices_kernel(
     n = tl.load(indptr_row + batch_id + 1) - start
     src = page_table + tl.load(rows + batch_id).to(tl.int64) * pt_stride
     dst = out + tl.program_id(0).to(tl.int64) * out_stride + start
-    for off in range(tl.program_id(2) * BLOCK, n, NUM_CHUNKS * BLOCK):
+    for off in range(tl.program_id(2) * BLOCK, n, tl.num_programs(2) * BLOCK):
         cols = off + tl.arange(0, BLOCK)
         mask = cols < n
         tl.store(dst + cols, tl.load(src + cols, mask=mask), mask=mask)
@@ -1097,6 +1096,5 @@ def gather_kv_indices(page_table: torch.Tensor, rows: torch.Tensor, indptr: torc
         page_table.stride(0),
         indptr.stride(0),
         out.stride(0),
-        NUM_CHUNKS=chunks,
         BLOCK=1024,
     )

From c099cbb24e13979400163feb8075fb0fbedce7e5 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 01:48:38 +0000
Subject: [PATCH 14/20] refactor(moe): let each family name its MTP expert
 tensors

---
 python/freetoken/models/qwen3_5_moe/mtp.py    |  5 +--
 python/freetoken/models/qwen3_5_moe/weight.py | 17 +++++++
 python/freetoken/models/qwen4_exp/mtp.py      |  5 +--
 python/freetoken/models/qwen4_exp/weight.py   | 17 +++++++
 python/freetoken/moe/expert_pieces.py         | 44 +++++++++----------
 tests/models/test_qwen3_5_moe_weight.py       |  5 +--
 6 files changed, 60 insertions(+), 33 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index a181c5fc04..a246c71778 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -16,7 +16,7 @@
 from freetoken.utils import torch_dtype
 
 from .model import Qwen3_5DecoderLayer
-from .weight import iter_mtp_weights
+from .weight import iter_mtp_expert_pieces, iter_mtp_weights
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -82,11 +82,10 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
         draft.load_state_dict(loaded)
         if draft.config.moe_enabled:
             from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks
-            from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
 
             # packed by the layer's own expert method, as the target's resident layers are
             moe = draft.layers.op_list[0].mlp.experts
-            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts")
+            pieces = iter_mtp_expert_pieces(model_path)
             banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
             attach_resident_banks([moe], banks)
         # not the root: it reaches the target's embedding, which the target finalized already
diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index 77883c57dc..c5acc775bf 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -273,6 +273,7 @@ def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | N
 
 # routed experts: expert_pieces.iter_extra_layer_expert_pieces
 _MTP_EXPERTS_PREFIX = "mtp.layers.0.mlp.experts."
+_MTP_EXPERT_RE = re.compile(re.escape(_MTP_EXPERTS_PREFIX) + r"(?:(?P<expert>\d+)\.(?P<proj>gate|up|down)_proj\.(?P<kind>weight|weight_scale_inv)|(?P<stacked>gate_up|down)_proj)$")
 
 
 def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
@@ -300,6 +301,21 @@ def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[st
         raise ValueError(f"{model_path}: incomplete MTP projections: {reader.missing()}")
 
 
+def iter_mtp_expert_pieces(model_path: str):
+    """The MTP head's routed experts: the stacked bf16 pair, or per-expert projections with their block-fp8 scales."""
+    from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+
+    def locate(name: str):
+        m = _MTP_EXPERT_RE.match(name)
+        if m is None:
+            return None
+        if m["stacked"]:
+            return None, m["stacked"]
+        return int(m["expert"]), m["proj"] + ("_scale" if m["kind"] == "weight_scale_inv" else "")
+
+    return iter_extra_layer_expert_pieces(model_path, _MTP_EXPERTS_PREFIX, locate)
+
+
 def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
     """The vision tower alone, named as iter_weights names it."""
     for file in iter_weight_files(model_path):
@@ -427,6 +443,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec:
 
 
 __all__ = [
+    "iter_mtp_expert_pieces",
     "iter_mtp_weights",
     "iter_weights",
     "iter_weights_parallel",
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index f0ff34566f..c630fa116d 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -29,7 +29,7 @@
 
 from .hc import GatedResidual
 from .model import Qwen4ExpDecoderLayer
-from .weight import iter_mtp_weights
+from .weight import iter_mtp_expert_pieces, iter_mtp_weights
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -159,7 +159,6 @@ def __init__(self, config: ModelConfig, target_model, quant=None) -> None:
     def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> Qwen4ExpMTPDraft:
         from freetoken.models.gguf.reader import is_gguf_path
         from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks
-        from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
 
         pack = is_gguf_path(model_path)
         with torch.device("meta"), torch_dtype(dtype):
@@ -179,7 +178,7 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
         if pack:
             banks = _pack_expert_banks(model_path, device)
         else:
-            pieces = iter_extra_layer_expert_pieces(model_path, "mtp.layers.0.mlp.experts")
+            pieces = iter_mtp_expert_pieces(model_path)
             banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True)
         attach_resident_banks([moe], banks)
         # not the root: it reaches the target's embedding, which the target finalized already
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index 7824ab0613..d0534cb23c 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -298,6 +298,7 @@ def iter_weights(
 
 # routed experts: expert_pieces.iter_extra_layer_expert_pieces
 _MTP_EXPERTS_PREFIX = "mtp.layers.0.mlp.experts."
+_MTP_EXPERT_RE = re.compile(re.escape(_MTP_EXPERTS_PREFIX) + r"(?:(?P<expert>\d+)\.(?P<proj>gate|up|down)_proj\.(?P<kind>weight|weight_scale_inv)|(?P<stacked>gate_up|down)_proj)$")
 
 
 def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
@@ -325,6 +326,21 @@ def iter_mtp_weights(model_path: str, device: torch.device) -> Iterator[tuple[st
         raise ValueError(f"{model_path}: incomplete MTP projections: {reader.missing()}")
 
 
+def iter_mtp_expert_pieces(model_path: str):
+    """The MTP head's routed experts: the stacked bf16 pair, or per-expert projections with their block-fp8 scales."""
+    from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+
+    def locate(name: str):
+        m = _MTP_EXPERT_RE.match(name)
+        if m is None:
+            return None
+        if m["stacked"]:
+            return None, m["stacked"]
+        return int(m["expert"]), m["proj"] + ("_scale" if m["kind"] == "weight_scale_inv" else "")
+
+    return iter_extra_layer_expert_pieces(model_path, _MTP_EXPERTS_PREFIX, locate)
+
+
 def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]:
     """The vision tower alone, named as iter_weights names it."""
     for file in iter_weight_files(model_path):
@@ -497,6 +513,7 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec:
 __all__ = [
     "nvfp4_expert_spec",
     "PleTable",
+    "iter_mtp_expert_pieces",
     "iter_mtp_weights",
     "iter_weights",
     "load_ple_table",
diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py
index 62f829c885..e25fa760e3 100644
--- a/python/freetoken/moe/expert_pieces.py
+++ b/python/freetoken/moe/expert_pieces.py
@@ -8,7 +8,6 @@
 
 from __future__ import annotations
 
-import re
 from typing import Callable, Collection, Iterable, Iterator
 
 import torch
@@ -170,48 +169,45 @@ def per_expert_pieces(
         raise ValueError(f"Incomplete expert tensors in checkpoint (layer, expert): {missing}")
 
 
-def iter_extra_layer_expert_pieces(model_path: str, prefix: str) -> Iterator[Piece]:
-    """The routed experts stored under ``prefix``, one MoE layer outside the model's stack (an MTP head), as pieces of bank layer 0: stacked ``gate_up``/``down``, or ``gate``/``up``/``down`` per expert with ``_scale`` companions in a block-fp8 checkpoint."""
+def iter_extra_layer_expert_pieces(model_path: str, prefix: str, locate: Callable[[str], tuple[int | None, str] | None]) -> Iterator[Piece]:
+    """The routed experts of one MoE layer outside the model's stack (an MTP head), as pieces of bank layer 0.
+
+    The family names the tensors: ``locate(name)`` -> ``(expert, role)`` of a per-expert tensor, ``(None, "gate_up" | "down")`` of a stacked one, or None. A name under ``prefix`` that ``locate`` does not read is refused.
+    """
     import safetensors
 
     from freetoken.models.loader import iter_weight_files
 
-    per_expert_re = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.(gate|up|down)_proj\.(weight|weight_scale_inv)$")
-    stacked_re = re.compile(rf"^{re.escape(prefix)}\.(gate_up|down)_proj$")
-    names: dict[str, str] = {}
+    found: dict[str, tuple[str, int | None, str]] = {}
     unknown: list[str] = []
     for file in iter_weight_files(model_path):
         with safetensors.safe_open(file, framework="pt", device="cpu") as f:
             for name in f.keys():
-                if per_expert_re.match(name) or stacked_re.match(name):
-                    names[name] = file
-                elif name.startswith(prefix + "."):
+                if (where := locate(name)) is not None:
+                    found[name] = (file, *where)
+                elif name.startswith(prefix):
                     unknown.append(name)
     if unknown:
-        raise ValueError(f"{model_path}: the experts under {prefix} use a format this reader does not support, e.g. {sorted(unknown)[:3]}")
-    if not names:
-        raise ValueError(f"{model_path} has no routed experts under {prefix}")
-    stacked = {stacked_re.match(n)[1]: f for n, f in names.items() if stacked_re.match(n)}
+        raise ValueError(f"{model_path}: the experts under {prefix}* use a format this reader does not support, e.g. {sorted(unknown)[:3]}")
+    if not found:
+        raise ValueError(f"{model_path} has no routed experts under {prefix}*")
+    stacked = {role: (file, name) for name, (file, expert, role) in found.items() if expert is None}
     if stacked:
         if missing := {"gate_up", "down"} - set(stacked):
-            raise ValueError(f"{model_path}: stacked experts under {prefix} lack {sorted(missing)}")
+            raise ValueError(f"{model_path}: stacked experts under {prefix}* lack {sorted(missing)}")
         piece = {}
-        for role, file in stacked.items():
+        for role, (file, name) in stacked.items():
             with safetensors.safe_open(file, framework="pt", device="cpu") as f:
-                piece[role] = f.get_tensor(f"{prefix}.{role}_proj")
+                piece[role] = f.get_tensor(name)
         return iter([(0, 0, piece["down"].shape[0], piece)])
 
-    def locate(name: str) -> tuple[int, int, str]:
-        m = per_expert_re.match(name)
-        return 0, int(m[1]), m[2] + ("_scale" if m[3] == "weight_scale_inv" else "")
-
     def tensors():
-        for file in dict.fromkeys(names.values()):
+        for file in dict.fromkeys(file for file, _, _ in found.values()):
             with safetensors.safe_open(file, framework="pt", device="cpu") as f:
-                yield from ((name, f.get_tensor(name)) for name, where in names.items() if where == file)
+                yield from ((name, f.get_tensor(name)) for name, where in found.items() if where[0] == file)
 
-    scaled = any(name.endswith("_scale_inv") for name in names)
-    return per_expert_pieces(tensors(), locate, tensors_per_expert=6 if scaled else 3)
+    roles = {role for _, _, role in found.values()}
+    return per_expert_pieces(tensors(), lambda name: (0, *found[name][1:]), tensors_per_expert=len(roles))
 
 
 __all__ = [
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index 9ad3781d4a..a82c581b9e 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -635,8 +635,7 @@ def test_mtp_head_fills_the_draft_and_its_experts(checkpoint):
     from types import SimpleNamespace
 
     from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft
-    from freetoken.models.qwen3_5_moe.weight import iter_mtp_weights
-    from freetoken.moe.expert_pieces import iter_extra_layer_expert_pieces
+    from freetoken.models.qwen3_5_moe.weight import iter_mtp_expert_pieces, iter_mtp_weights
 
     name, folder, raw = checkpoint
     if name not in ("bf16", "fp8_block"):
@@ -650,7 +649,7 @@ def test_mtp_head_fills_the_draft_and_its_experts(checkpoint):
         assert tensor.shape == state[key].shape, key
         if key.endswith(".weight"):
             assert tensor.dtype is state[key].dtype, key
-    pieces = list(iter_extra_layer_expert_pieces(folder, f"{MTP}.mlp.experts"))
+    pieces = list(iter_mtp_expert_pieces(folder))
     if name == "bf16":
         assert [p[:3] for p in pieces] == [(0, 0, E)]
         assert torch.equal(pieces[0][3]["gate_up"], raw[f"{MTP}.mlp.experts.gate_up_proj"])

From 5b46b07ed60f3a1ee7dc0da70e3469ba9a891083 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 01:48:38 +0000
Subject: [PATCH 15/20] refactor(spec): force the MTP draft's MoE resident in
 the proposer and check it after load

---
 python/freetoken/models/qwen3_5_moe/mtp.py |  4 +---
 python/freetoken/models/qwen4_exp/mtp.py   |  4 +---
 python/freetoken/spec/mtp/proposer.py      | 17 +++++++++++++++--
 tests/models/test_qwen3_5_moe_weight.py    |  9 +++++++--
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index a246c71778..01c0a1865f 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -23,15 +23,13 @@
 
 
 def draft_config(config: ModelConfig) -> ModelConfig:
-    """One full-attention layer with the target's geometry; its MoE is resident on the GPU whatever the target's strategy."""
+    """One full-attention layer with the target's geometry."""
     groups = [g for g in config.attention_groups if isinstance(g, FullAttentionGroupConfig)]
     group = groups[0] if groups else config.default_full_attention_group()
     return dataclasses.replace(
         config,
         num_layers=1,
         attention_groups=(dataclasses.replace(group, layer_ids=(0,)),),
-        moe_strategy="fused",
-        decode_target="gpu",
     )
 
 
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index c630fa116d..728c04ac18 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -104,7 +104,7 @@ def _pack_quant(model_path: str):
 
 
 def draft_config(config: ModelConfig) -> ModelConfig:
-    """The target config cut to one QSA layer (layer id 0) without PLE; its MoE is resident on the GPU whatever the target's strategy."""
+    """The target config cut to one QSA layer (layer id 0) without PLE."""
     group = next(g for g in config.attention_groups if isinstance(g, FullAttentionGroupConfig) and g.index_ratio > 1)
     return dataclasses.replace(
         config,
@@ -113,8 +113,6 @@ def draft_config(config: ModelConfig) -> ModelConfig:
         qwen4_args=dataclasses.replace(config.qwen4_args, ple_layer_ids=()),
         slot_states=(),
         vision_config=None,
-        moe_strategy="fused",
-        decode_target="gpu",
     )
 
 
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 6a3003d308..73b30dce0b 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar, List, Protocol
 
@@ -52,6 +53,11 @@ def head_input(self, hidden: torch.Tensor) -> torch.Tensor:
         ...
 
 
+def with_resident_moe(config: ModelConfig) -> ModelConfig:
+    """``config`` with its MoE resident on the GPU, the strategy every draft builds with whatever the target's."""
+    return dataclasses.replace(config, moe_strategy="fused", decode_target="gpu")
+
+
 @dataclass
 class _DraftReq:
     """The request fields the attention backends and rope read, for a draft batch over positions [cached_len, device_len)."""
@@ -139,7 +145,7 @@ def budget(cls, config) -> tuple[int, int]:
         rows = config.max_running_req + 1
         # relay hidden and carry, and the verify hiddens
         fixed = rows * (config.speculative_num_steps + 3) * relay_bytes
-        groups = draft_cls.draft_config(config.model_config).kv_cache_group_specs()
+        groups = draft_cls.draft_config(with_resident_moe(config.model_config)).kv_cache_group_specs()
         return fixed, sum(spec_kv_bytes_per_token(group, config) for group in groups)
 
     def __init__(self, config, target_model) -> None:
@@ -147,7 +153,14 @@ def __init__(self, config, target_model) -> None:
         self.device = torch.device("cuda", torch.cuda.current_device())
         self.dtype = config.dtype
         path = config.speculative_draft_model_path or config.model_path
-        self.draft = self._draft_class(config).load(config.model_config, target_model, path, self.device, config.dtype)
+        from freetoken.layers.moe import OffloadMoELayer, iter_moe_layers
+
+        draft_cls = self._draft_class(config)
+        self.draft = draft_cls.load(with_resident_moe(config.model_config), target_model, path, self.device, config.dtype)
+        # an offloaded or shared MoE layer here would read the target's experts as its own, silently
+        target_moe = {id(m) for m in iter_moe_layers(target_model)}
+        if any(isinstance(m, OffloadMoELayer) or id(m) in target_moe for m in iter_moe_layers(self.draft)):
+            raise RuntimeError(f"{draft_cls.__name__} must build its own MoE layers, resident on the GPU")
         self.lm_head = target_model.lm_head
         self.mrope = config.model_config.model_is_mrope
         self._prefill_state: dict[int, tuple[int, int]] = {}  # table row -> (uid, first position not drafted yet)
diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py
index a82c581b9e..381e5a4f49 100644
--- a/tests/models/test_qwen3_5_moe_weight.py
+++ b/tests/models/test_qwen3_5_moe_weight.py
@@ -631,18 +631,23 @@ def test_a_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, qu
 
 
 def test_mtp_head_fills_the_draft_and_its_experts(checkpoint):
-    """The MTP reader fills exactly the draft the engine builds, and the routed experts keep the form the checkpoint stores."""
+    """The MTP reader fills exactly the draft the engine builds, whose MoE stays resident though the target offloads, and the routed experts keep the form the checkpoint stores."""
     from types import SimpleNamespace
 
+    from freetoken.layers.moe import OffloadMoELayer, iter_moe_layers
     from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft
     from freetoken.models.qwen3_5_moe.weight import iter_mtp_expert_pieces, iter_mtp_weights
+    from freetoken.spec.mtp.proposer import with_resident_moe
 
     name, folder, raw = checkpoint
     if name not in ("bf16", "fp8_block"):
         pytest.skip("layouts with an MTP head only")
     _install(folder)
     target = SimpleNamespace(model=SimpleNamespace(embed_tokens=None))
-    state = _meta_state_dict(folder, lambda config: Qwen3_5MTPDraft(config, target))
+    built = {}
+    state = _meta_state_dict(folder, lambda config: built.setdefault("draft", Qwen3_5MTPDraft(with_resident_moe(config), target)))
+    moe = list(iter_moe_layers(built["draft"]))
+    assert moe and not any(isinstance(m, OffloadMoELayer) for m in moe)
     loaded = dict(iter_mtp_weights(folder, torch.device("cpu")))
     assert set(loaded) == set(state)
     for key, tensor in loaded.items():

From 8d34d2275216c9deb925c91a6e55885887495b50 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 01:48:38 +0000
Subject: [PATCH 16/20] fix(spec): refuse MTP with dummy weights

---
 python/freetoken/spec/mtp/proposer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 73b30dce0b..df5c741b85 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -119,6 +119,8 @@ def adjust_config(cls, config, override) -> None:
         arch = config.hf_config.architectures[0]
         if "mtp" not in dict(config.model_spec.spec_drafts):
             raise ValueError(f"MTP is not supported for {arch} yet")
+        if config.use_dummy_weight:
+            raise ValueError("MTP reads its draft from the checkpoint, so it cannot run with --dummy-weight")
         path = config.speculative_draft_model_path or config.model_path
         # checked here because a draft that cannot load otherwise fails only after the whole target has loaded
         # a HF repo id is "name" or "org/name"; anything else names a local path

From 0fb975da9cc1e338240b2a3d56c2f6dc9da1f872 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Fri, 30 Oct 2026 03:50:24 +0000
Subject: [PATCH 17/20] refactor(spec): move HiddenRelay under spec/mtp, its
 only user

---
 python/freetoken/spec/mtp/proposer.py    | 2 +-
 python/freetoken/spec/{ => mtp}/relay.py | 0
 python/freetoken/spec/proposer.py        | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename python/freetoken/spec/{ => mtp}/relay.py (100%)

diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index df5c741b85..48d982b1b6 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -11,7 +11,7 @@
 from freetoken.core import Batch, get_global_ctx
 from freetoken.engine.sample import make_device_tensor
 from freetoken.spec.proposer import DraftContext, Proposer
-from freetoken.spec.relay import HiddenRelay
+from freetoken.spec.mtp.relay import HiddenRelay
 
 if TYPE_CHECKING:
     from freetoken.core import Req
diff --git a/python/freetoken/spec/relay.py b/python/freetoken/spec/mtp/relay.py
similarity index 100%
rename from python/freetoken/spec/relay.py
rename to python/freetoken/spec/mtp/relay.py
diff --git a/python/freetoken/spec/proposer.py b/python/freetoken/spec/proposer.py
index 67ecb41386..c91282e0e2 100644
--- a/python/freetoken/spec/proposer.py
+++ b/python/freetoken/spec/proposer.py
@@ -49,7 +49,7 @@ def __init__(self, config, target_model) -> None:
         self.k = config.speculative_num_steps
 
     def attach(self, ctx: DraftContext) -> None:
-        """Build the draft KV pool, backends and relays; called again with a new context after a rebuild."""
+        """Build the draft KV pool, backends and buffers; called again with a new context after a rebuild."""
         self.ctx = ctx
 
     def detach(self, *, keep_kv: bool) -> None:

From 242eec95354e53e166a6dcbf8865bf09d6a89551 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 2 Nov 2026 00:19:43 +0000
Subject: [PATCH 18/20] refactor(spec): drop MTPDraft.config, fold its embed
 hooks into embed_tokens

---
 python/freetoken/models/qwen3_5_moe/mtp.py | 11 ++++-------
 python/freetoken/models/qwen4_exp/mtp.py   | 11 ++++-------
 python/freetoken/spec/mtp/chain_graph.py   |  2 +-
 python/freetoken/spec/mtp/proposer.py      | 21 +++++++++++----------
 4 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py
index 01c0a1865f..d3b257b651 100644
--- a/python/freetoken/models/qwen3_5_moe/mtp.py
+++ b/python/freetoken/models/qwen3_5_moe/mtp.py
@@ -11,7 +11,6 @@
 from freetoken.layers import GemmaPlusOneRMSNorm, GemmaPlusOneRMSNormFused, LinearReplicated, OPList
 from freetoken.layers.quantization import finalize_quant
 from freetoken.layers.base import BaseOP
-from freetoken.models.blocks import embed_input_ids
 from freetoken.models.config import FullAttentionGroupConfig
 from freetoken.utils import torch_dtype
 
@@ -19,6 +18,7 @@
 from .weight import iter_mtp_expert_pieces, iter_mtp_weights
 
 if TYPE_CHECKING:
+    from freetoken.layers.embedding import VocabParallelEmbedding
     from freetoken.models.config import ModelConfig
 
 
@@ -91,12 +91,9 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
         finalize_quant(draft.layers)
         return draft
 
-    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self._embed_tokens.forward(input_ids)
-
-    def embed_batch(self, batch) -> torch.Tensor:
-        """The target batch's token embeddings with its multimodal rows merged in, as the target embeds them."""
-        return embed_input_ids(self._embed_tokens, batch.input_ids, batch)
+    @property
+    def embed_tokens(self) -> VocabParallelEmbedding:
+        return self._embed_tokens
 
     def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor:
         """Draft hidden of each row from the embedding of its token and the hidden of the position before it."""
diff --git a/python/freetoken/models/qwen4_exp/mtp.py b/python/freetoken/models/qwen4_exp/mtp.py
index 728c04ac18..e86fc29d26 100644
--- a/python/freetoken/models/qwen4_exp/mtp.py
+++ b/python/freetoken/models/qwen4_exp/mtp.py
@@ -23,7 +23,6 @@
 from freetoken.layers import GemmaPlusOneRMSNorm, LinearReplicated, OPList
 from freetoken.layers.base import BaseOP
 from freetoken.layers.quantization import finalize_quant
-from freetoken.models.blocks import embed_input_ids
 from freetoken.models.config import FullAttentionGroupConfig
 from freetoken.utils import torch_dtype
 
@@ -32,6 +31,7 @@
 from .weight import iter_mtp_expert_pieces, iter_mtp_weights
 
 if TYPE_CHECKING:
+    from freetoken.layers.embedding import VocabParallelEmbedding
     from freetoken.models.config import ModelConfig
 
 
@@ -184,12 +184,9 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.
             finalize_quant(module)
         return draft
 
-    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self._embed_tokens.forward(input_ids)
-
-    def embed_batch(self, batch) -> torch.Tensor:
-        """The target batch's token embeddings with its multimodal rows merged in, as the target embeds them."""
-        return embed_input_ids(self._embed_tokens, batch.input_ids, batch)
+    @property
+    def embed_tokens(self) -> VocabParallelEmbedding:
+        return self._embed_tokens
 
     def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor:
         """``R_out`` of each row from the embedding of its token and the residual streams of the position before it."""
diff --git a/python/freetoken/spec/mtp/chain_graph.py b/python/freetoken/spec/mtp/chain_graph.py
index 5fa61ba8ce..2b73e61f9b 100644
--- a/python/freetoken/spec/mtp/chain_graph.py
+++ b/python/freetoken/spec/mtp/chain_graph.py
@@ -143,7 +143,7 @@ def _run(self, bs: int, batches: List[Batch]) -> None:
         token, hidden = p.relay.d1[rows], p.relay.hidden[rows]
         self.drafts[:bs, 0] = token
         for step, batch in enumerate(batches):
-            hidden = p._forward(batch, p.chain_backend, p.draft.embed(token), hidden)
+            hidden = p._forward(batch, p.chain_backend, p.draft.embed_tokens.forward(token), hidden)
             token = p._draft_token(hidden)
             self.drafts[:bs, step + 1] = token
 
diff --git a/python/freetoken/spec/mtp/proposer.py b/python/freetoken/spec/mtp/proposer.py
index 48d982b1b6..9c2057235a 100644
--- a/python/freetoken/spec/mtp/proposer.py
+++ b/python/freetoken/spec/mtp/proposer.py
@@ -10,19 +10,19 @@
 
 from freetoken.core import Batch, get_global_ctx
 from freetoken.engine.sample import make_device_tensor
+from freetoken.models.blocks import embed_input_ids
 from freetoken.spec.proposer import DraftContext, Proposer
 from freetoken.spec.mtp.relay import HiddenRelay
 
 if TYPE_CHECKING:
     from freetoken.core import Req
+    from freetoken.layers.embedding import VocabParallelEmbedding
     from freetoken.models.config import ModelConfig
 
 
 class MTPDraft(Protocol):
     """The class a family registers as ``ModelSpec.spec_drafts["mtp"]``."""
 
-    # the draft's own config, from draft_config
-    config: ModelConfig
     # None: the proposer picks from the target's backends
     attention_backend: ClassVar[str | None]
 
@@ -42,9 +42,10 @@ def check_source(path: str, target_path: str) -> None:
     @classmethod
     def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> MTPDraft: ...
 
-    def embed(self, input_ids: torch.Tensor) -> torch.Tensor: ...
-
-    def embed_batch(self, batch: Batch) -> torch.Tensor: ...
+    @property
+    def embed_tokens(self) -> VocabParallelEmbedding:
+        """The target's token embedding, which the draft shares."""
+        ...
 
     def forward_hidden(self, embeds: torch.Tensor, prev_hidden: torch.Tensor) -> torch.Tensor: ...
 
@@ -176,7 +177,7 @@ def attach(self, ctx: DraftContext) -> None:
         from freetoken.kvcache import create_kvcache_pool
 
         super().attach(ctx)
-        cfg = self.draft.config
+        cfg = self.draft.draft_config(with_resident_moe(self.config.model_config))
         rows = self.config.max_running_req + 1
         # the draft's own layers on the target's slot ids: same page size, same pages plus the dummy page padded rows write to
         # a rebuild that keeps the target KV pool keeps this one, which the prefix cache still indexes
@@ -292,7 +293,7 @@ def propose(self, reqs: List[Req], rows: torch.Tensor) -> torch.Tensor:
             # chain step j drafts at position C+j-1 with the previous step's draft hidden standing in for the target's
             views = [self._view(r, c + j - 1, c + j) for r, c in zip(reqs, cs)]
             batch = self._draft_batch(views, "decode", 1, self.chain_backend)
-            hidden = self._forward(batch, self.chain_backend, self.draft.embed(token), hidden)
+            hidden = self._forward(batch, self.chain_backend, self.draft.embed_tokens.forward(token), hidden)
             token = self._draft_token(hidden)
             drafts.append(token)
         return torch.stack(drafts, dim=1)
@@ -307,7 +308,7 @@ def _extend_forward(self) -> tuple[torch.Tensor, torch.Tensor]:
         """Draft hiddens of a verify-shaped batch; returns each request's row at its last accepted position and that row's argmax."""
         input_ids = get_global_ctx().batch.input_ids
         n = input_ids.shape[0]
-        hidden = self.draft.forward_hidden(self.draft.embed(input_ids), self.verify_hidden[:n])
+        hidden = self.draft.forward_hidden(self.draft.embed_tokens.forward(input_ids), self.verify_hidden[:n])
         picked = hidden[self.picked_rows[: n // (self.k + 1)]]
         return picked, self._draft_token(picked)
 
@@ -343,7 +344,7 @@ def on_prefill(self, batch: Batch, next_tokens_gpu: torch.Tensor) -> None:
         sink, self._prefill_sink = self._prefill_sink, None
         hidden = sink.hidden
         # image rows take the embeddings the target used
-        embeds = self.draft.embed_batch(batch)
+        embeds = embed_input_ids(self.draft.embed_tokens, batch.input_ids, batch)
         views, emb, prev, finals = [], [], [], []
         offset = drafted = 0
         for i, (req, (s, e)) in enumerate(zip(batch.reqs, sink.spans)):
@@ -360,7 +361,7 @@ def on_prefill(self, batch: Batch, next_tokens_gpu: torch.Tensor) -> None:
                 # draft row p embeds the token at p+1
                 emb.append(chunk_embeds[frontier + 1 - s : min(end + 1, e) - s])
                 if final:
-                    emb.append(self.draft.embed(next_tokens_gpu[i : i + 1]))
+                    emb.append(self.draft.embed_tokens.forward(next_tokens_gpu[i : i + 1]))
                 if frontier == s - 1:
                     # cloned: the carry of this row is overwritten below, before the draft forward reads it
                     prev += [self.relay.carry[row : row + 1].clone(), rows[: end - s]]

From e5ac81f32720599ffb62be94e72310dede964704 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 2 Nov 2026 00:19:43 +0000
Subject: [PATCH 19/20] refactor(models): note FeatureSink's graph binding,
 drop its runtime_checkable

---
 python/freetoken/models/blocks.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py
index 755a93ae10..e51f3637ec 100644
--- a/python/freetoken/models/blocks.py
+++ b/python/freetoken/models/blocks.py
@@ -48,9 +48,11 @@ def place_encoder_weights(self, mode: str) -> None:
         ...
 
 
-@runtime_checkable
 class FeatureSink(Protocol):
-    """Where a forward hands per-row hidden states to a draft model; set on ``batch.features``."""
+    """Where a forward hands per-row hidden states to a draft model; set on ``batch.features``.
+
+    A CUDA graph replays the sink it was captured with, never the one set at replay: write only to buffers the draft owns, keep no per-step state on the sink.
+    """
 
     wanted_layers: frozenset[int]
 

From 1e3f41887fd15b0f960935edb45a3066292617bd Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 2 Nov 2026 03:14:50 +0000
Subject: [PATCH 20/20] fix(models): read quantized MTP fc projections under
 the quant config

---
 python/freetoken/models/qwen3_5_moe/weight.py | 2 +-
 python/freetoken/models/qwen4_exp/weight.py   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py
index c5acc775bf..9231f1ad64 100644
--- a/python/freetoken/models/qwen3_5_moe/weight.py
+++ b/python/freetoken/models/qwen3_5_moe/weight.py
@@ -38,7 +38,7 @@
 # leaves the model builds as Linear layers: only their tensors are read under the QuantConfig, the rest passes through as stored
 _LINEAR_LEAVES = frozenset({
     "q_proj", "k_proj", "v_proj", "o_proj", "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj",
-    "gate_proj", "up_proj", "down_proj", "gate", "shared_expert_gate", "lm_head",
+    "gate_proj", "up_proj", "down_proj", "gate", "shared_expert_gate", "lm_head", "fc",
 })
 # activation scales of modules whose scheme carries no input_scale role
 _DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"})
diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py
index d0534cb23c..15b7b050dd 100644
--- a/python/freetoken/models/qwen4_exp/weight.py
+++ b/python/freetoken/models/qwen4_exp/weight.py
@@ -51,7 +51,7 @@
     "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj",
     "gate_proj", "up_proj", "down_proj",
     "input_mix_weight_down", "block_inject_weight", "input_mix_weight_up",
-    "key_proj", "value_proj", "lm_head",
+    "key_proj", "value_proj", "lm_head", "fc_embedding", "fc_hidden",
 })
 # activation scales of modules whose scheme carries no input_scale role
 _DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"})