feat(spec): add MTP drafts for the Qwen3.5 family and Qwen3.8-Flash-Next

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 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 Qwen3.5 / Qwen3.6 GGUF targets and an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors.
+
 ### API behaviour
 
 | Flag | Default | Meaning |

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):

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:

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":

python/freetoken/kernel/triton/attention.py

@@ -1058,3 +1058,43 @@ 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,
+    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, 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)
+
+
+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),
+        BLOCK=1024,
+    )

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,
         )

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,19 +96,23 @@ 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]:
         return {
             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:

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

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)

python/freetoken/models/blocks.py

@@ -48,14 +48,16 @@ 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]
 
     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:

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)

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"]