diff --git a/docs/cli.md b/docs/cli.md index feda4f7a34..2df0308322 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 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 | 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/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index e49346dd5d..28effa5ea2 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/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, + ) 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/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..e51f3637ec 100644 --- a/python/freetoken/models/blocks.py +++ b/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: 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/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/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py new file mode 100644 index 0000000000..d3b257b651 --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/mtp.py @@ -0,0 +1,108 @@ +"""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 +import os +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.config import FullAttentionGroupConfig +from freetoken.utils import torch_dtype + +from .model import Qwen3_5DecoderLayer +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 + + +def draft_config(config: ModelConfig) -> ModelConfig: + """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,)),), + ) + + +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 + + @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 + 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 + + # 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_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 + finalize_quant(draft.fc) + finalize_quant(draft.layers) + return draft + + @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.""" + 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..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"}) @@ -271,6 +271,51 @@ 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." +_MTP_EXPERT_RE = re.compile(re.escape(_MTP_EXPERTS_PREFIX) + r"(?:(?P\d+)\.(?Pgate|up|down)_proj\.(?Pweight|weight_scale_inv)|(?Pgate_up|down)_proj)$") + + +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_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): @@ -398,6 +443,8 @@ def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: __all__ = [ + "iter_mtp_expert_pieces", + "iter_mtp_weights", "iter_weights", "iter_weights_parallel", "iter_expert_pieces", 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..623389825a 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,10 +124,16 @@ 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)) + 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] @@ -185,11 +192,14 @@ 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), - max_extend_tokens=engine_config.max_extend_tokens, + # 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), + # 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: @@ -210,7 +220,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..e86fc29d26 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/mtp.py @@ -0,0 +1,204 @@ +"""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 os +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.config import FullAttentionGroupConfig +from freetoken.utils import torch_dtype + +from .hc import GatedResidual +from .model import Qwen4ExpDecoderLayer +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 + + +_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.""" + 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, + ) + + +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 + + @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 + 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.models.gguf.reader import is_gguf_path + from freetoken.moe.expert_banks import attach_resident_banks, build_expert_banks + + 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() + 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_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 + for module in (draft.fc_embedding, draft.fc_hidden, draft.layers, draft.hyper_connection_mixer): + finalize_quant(module) + return draft + + @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.""" + 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/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 5a3b321f08..bc2afe1f60 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: @@ -318,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. @@ -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,13 +377,14 @@ def build_ple_metadata( state_slots=slots, fresh_slots=fresh, is_decode=batch.is_decode, + is_verify=getattr(batch, "phase", None) == "verify", ) 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. """ @@ -647,9 +654,11 @@ 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: + 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..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) @@ -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/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 734eb80a83..15b7b050dd 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 @@ -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"}) @@ -296,6 +296,51 @@ 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." +_MTP_EXPERT_RE = re.compile(re.escape(_MTP_EXPERTS_PREFIX) + r"(?:(?P\d+)\.(?Pgate|up|down)_proj\.(?Pweight|weight_scale_inv)|(?Pgate_up|down)_proj)$") + + +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_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): @@ -468,6 +513,8 @@ 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/models/register.py b/python/freetoken/models/register.py index 3ed67c99bf..7625b5abcf 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,8 @@ 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"),) +_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")), @@ -288,6 +292,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 +300,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", @@ -336,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( @@ -347,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 @@ -359,6 +369,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..e25fa760e3 100644 --- a/python/freetoken/moe/expert_pieces.py +++ b/python/freetoken/moe/expert_pieces.py @@ -169,10 +169,52 @@ 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, 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 + + 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 (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 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)}") + piece = {} + for role, (file, name) in stacked.items(): + with safetensors.safe_open(file, framework="pt", device="cpu") as f: + piece[role] = f.get_tensor(name) + return iter([(0, 0, piece["down"].shape[0], piece)]) + + def tensors(): + 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 found.items() if where[0] == file) + + roles = {role for _, _, role in found.values()} + return per_expert_pieces(tensors(), lambda name: (0, *found[name][1:]), tensors_per_expert=len(roles)) + + __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..2b73e61f9b --- /dev/null +++ b/python/freetoken/spec/mtp/chain_graph.py @@ -0,0 +1,165 @@ +"""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 + +from typing import TYPE_CHECKING, Dict, List + +import torch + +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 + +if TYPE_CHECKING: + from freetoken.core import Req + + 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) + 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.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 = self.backend_steps.metadata(bs, step, batch.positions) + 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.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): + 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 + + 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..9c2057235a --- /dev/null +++ b/python/freetoken/spec/mtp/proposer.py @@ -0,0 +1,385 @@ +"""MTPProposer: drafts with the checkpoint's own MTP layer, one draft token per chained forward.""" + +from __future__ import annotations + +import dataclasses +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.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"]``.""" + + # 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.""" + ... + + @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: ... + + @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: ... + + def head_input(self, hidden: torch.Tensor) -> torch.Tensor: + """What the target's LM head reads from a draft hidden.""" + ... + + +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).""" + + 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: + import os + + from freetoken.checkpoint.ftw import is_ftw_checkpoint + + 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 {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 + 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): + raise ValueError("MTP cannot load its draft from an FTW checkpoint yet") + cls._draft_class(config).check_source(path, config.model_path) + + @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(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: + 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 + 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) + 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.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 + 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, num_speculative_tokens=self.k, + ) + 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 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: + 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_tokens.forward(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_tokens.forward(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 + 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 + 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.""" + sink, self._prefill_sink = self._prefill_sink, None + hidden = sink.hidden + # image rows take the embeddings the target used + 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)): + 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 = e == req.prompt_len + 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_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]] + 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/mtp/relay.py b/python/freetoken/spec/mtp/relay.py new file mode 100644 index 0000000000..ddbb101388 --- /dev/null +++ b/python/freetoken/spec/mtp/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) 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: 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/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py index 958d6daf66..381e5a4f49 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,39 @@ 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, 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)) + 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(): + assert tensor.shape == state[key].shape, key + if key.endswith(".weight"): + assert tensor.dtype is state[key].dtype, key + 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"]) + 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"]) 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"]))