diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index 8915af4c6475..f2f278dea788 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -6,11 +6,13 @@ import side effect. The zoo is imported lazily, so these tables record, without importing anything, which ``modeling_*`` module provides which architecture (``MODEL_ARCH_TO_MODULE``), which public class (``MODEL_CLASS_TO_MODULE``), -and which multimodal ``model_type`` (``MULTIMODAL_MODEL_TYPE_TO_MODULE``). +which multimodal ``model_type`` (``MULTIMODAL_MODEL_TYPE_TO_MODULE``), and +which speculative-decoding mode (``SPEC_MODE_TO_MODULE``). Regenerate after adding/moving a model: add the new entry by hand next to its neighbors, mirroring the ``@register_auto_model("")`` / -``@register_input_processor(..., model_type="")`` decorators and the +``@register_input_processor(..., model_type="")`` / +``@register_draft_model(SpeculativeDecodingMode.)`` decorators and the public class name. ``test_lazy_model_zoo.py`` fails on any drift between these tables and the decorators. """ @@ -223,3 +225,58 @@ def is_builtin_zoo_module(module_name: str) -> bool: "step3p7": "modeling_step3p7vl", "whisper": "modeling_whisper", } + +# ``SpeculativeDecodingMode`` member name -> module providing that mode's draft +# model builder (registered via ``@register_draft_model``). Keyed by the enum +# member *name* rather than the enum itself so this module keeps importing +# nothing. Modes absent from this table have no one-engine draft model to build +# (two-model / drafter-loop modes such as NGRAM, SA and USER_PROVIDED). +# +# Adding a speculative decoding mode +# ---------------------------------- +# 1. Write the builder in *your own* ``modeling_*.py``, next to the draft model +# it constructs -- never in ``modeling_speculative.py``. Keeping builders out +# of the factory file is the entire point of this table: ``get_draft_model`` +# imports no concrete draft implementation, which is what let DSpark drop the +# lazy import it needed while ``modeling_dspark`` imports back into +# ``modeling_speculative`` through ``modeling_deepseekv4``. +# 2. Decorate it with ``@register_draft_model(SpeculativeDecodingMode.)``. +# Stack the decorator to serve several modes with one builder. +# 3. Add the ``"": "modeling_"`` row below. +# 4. ``test_lazy_model_zoo.py`` and +# ``tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py`` +# fail in both directions on any drift between decorators and this table. +# +# The builder signature is fixed at +# ``(model_config, draft_config, lm_head, model) -> nn.Module`` -- byte-for-byte +# the arguments of ``get_draft_model``, so the factory is pure forwarding with +# no per-mode glue. Everything a builder needs is reachable from those four +# (``model_config.pretrained_config.num_hidden_layers``, ``model.aux_stream_dict``, +# ``model_config.spec_config.*``). Do not widen it: an extra parameter has to be +# populated by the factory, which puts mode-specific knowledge straight back +# into the shared file this table exists to keep generic. +# +# Three rules the registry inherits from ``register_auto_model`` (see +# ``modeling_utils.py``), each written down because it was learned the hard way: +# - Look up only through ``get_registered_draft_model_builder``. It triggers +# the on-demand import before reading the mapping; a raw ``.get()`` silently +# misses every provider that has not been imported yet, and the zoo is +# imported lazily. +# - Built-in builders only fill empty slots, never overwrite. Lazy loading +# means a built-in's decorator can run *after* an external registration +# (``--custom_module_dirs``), so last-wins would let the built-in clobber a +# user's drafter. +# - Map a builder back to its modes via its ``_registered_spec_modes`` +# attribute, not by scanning the mapping for it. A built-in that lost its +# slot to an external registration is absent from the mapping but still has +# the attribute, so an identity scan reports it as unregistered. +SPEC_MODE_TO_MODULE = { + "DFLASH": "modeling_dflash", + "DRAFT_TARGET_ONE_MODEL": "modeling_speculative", + "DSPARK": "modeling_dspark", + "EAGLE3_ONE_MODEL": "modeling_speculative", + "MTP": "modeling_speculative", + "MTP_EAGLE": "modeling_speculative", + "MTP_EAGLE_ONE_MODEL": "modeling_speculative", + "PARD": "modeling_speculative", +} diff --git a/tensorrt_llm/_torch/models/dspark/__init__.py b/tensorrt_llm/_torch/models/dspark/__init__.py deleted file mode 100644 index b33d7553d877..000000000000 --- a/tensorrt_llm/_torch/models/dspark/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""DSpark draft-model components.""" diff --git a/tensorrt_llm/_torch/models/dspark/attention.py b/tensorrt_llm/_torch/models/dspark/attention.py deleted file mode 100644 index f3cfcedfbef4..000000000000 --- a/tensorrt_llm/_torch/models/dspark/attention.py +++ /dev/null @@ -1,567 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# The DSpark captured-context attention primitives are ported from DeepSeek's -# DeepSpec reference ``inference/kernel.py`` (``sparse_attn``) and -# ``inference/model.py`` (``get_dspark_topk_idxs``). The reference computes these -# with a TileLang kernel; this is a functional-first pure-PyTorch port with the -# same math (index-gather + online softmax + a learnable attention sink that -# contributes only to the softmax denominator). -"""DSpark draft captured-context attention primitives (hardware-agnostic). - -The DSpark draft uses *dense* sliding-window MLA (``compress_ratio == 0``): the -query comes from the block's draft tokens, while the keys/values are gathered -from a small per-request set of positions (a sliding window of the projected -captured context plus the current block's own positions). Two primitives capture -the parts that differ from the standard MLA path: - -* :func:`get_dspark_topk_idxs` — the (window-context + block) position list. -* :func:`dspark_sparse_attn` — index-gathered attention with an attention sink. -""" - -from functools import lru_cache - -import torch -import torch.nn.functional as F - -from ...._utils import is_sm_100f -from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE - -if IS_CUTLASS_DSL_AVAILABLE: - from ...custom_ops.dspark_attention_custom_op import ( - cute_dsl_dspark_attention, - is_fused_dspark_attention_supported, - ) - from ...custom_ops.dspark_rmsnorm_rope_custom_op import ( - cute_dsl_dspark_rmsnorm_rope, - is_fused_dspark_rmsnorm_rope_supported, - ) - -__all__ = [ - "get_dspark_topk_idxs", - "get_dspark_topk_idxs_batched", - "dspark_sparse_attn", - "precompute_dspark_freqs_cis", - "apply_dspark_rotary", - "apply_dspark_rotary_batched", - "dspark_attention_forward", - "dspark_attention_forward_batched", -] - - -def precompute_dspark_freqs_cis( - rope_head_dim: int, - seqlen: int, - rope_theta: float = 10000.0, - device: torch.device | str = "cpu", -) -> torch.Tensor: - """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. - - The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the - base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with - ``original_seq_len == 0``). - - Returns: - complex64 tensor ``[seqlen, rope_head_dim // 2]``. - """ - freqs = 1.0 / ( - rope_theta - ** (torch.arange(0, rope_head_dim, 2, dtype=torch.float32, device=device) / rope_head_dim) - ) - t = torch.arange(seqlen, dtype=torch.float32, device=device) - freqs = torch.outer(t, freqs) - return torch.polar(torch.ones_like(freqs), freqs) - - -def apply_dspark_rotary( - x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. - - Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the - last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the - sequence axis, and conjugates for the inverse (de-rotation applied to the - attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or - ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. - """ - orig_dtype = x.dtype - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - if xc.ndim == 3: - fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) - else: - fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) - out = torch.view_as_real(xc * fc).flatten(-2) - return out.to(orig_dtype) - - -def apply_dspark_rotary_batched( - x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Per-row (batched) variant of :func:`apply_dspark_rotary`. - - Identical math, but ``freqs_cis`` carries a leading batch axis so each row of - ``x`` is rotated by its own per-request phases (the generation draft runs each - request at a different absolute ``start_pos``). ``x`` is the rope-dim slice - only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape - ``[G, s, rd // 2]``. - """ - orig_dtype = x.dtype - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - g, s, half = freqs_cis.shape - if xc.ndim == 3: - fc = freqs_cis.view(g, s, half) - else: - fc = freqs_cis.view(g, s, 1, half) - out = torch.view_as_real(xc * fc).flatten(-2) - return out.to(orig_dtype) - - -@lru_cache(maxsize=64) -def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: - # [min(window, start_pos+1)] context positions in the rolling KV window, - # followed by [block_size] positions for the current block's own K/V (which - # the caller appends to the window at offset ``window_size``). - ctx = torch.arange(min(window_size, start_pos + 1)) - blk = window_size + torch.arange(block_size) - return torch.cat([ctx, blk]).int() - - -def get_dspark_topk_idxs( - window_size: int, - bsz: int, - block_size: int, - start_pos: int, - device: torch.device | str = "cpu", -) -> torch.Tensor: - """Per-query attended-position indices for the DSpark draft block. - - Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` - query positions attends to the same set — the ``min(window_size, start_pos+1)`` - most-recent context positions in the rolling KV window, then the - ``block_size`` positions of the current block (stored at offset - ``window_size`` in the concatenated KV). Note this is *non-causal* within the - block (every position sees every block position), matching the reference. - - Args: - window_size: sliding-window length of the captured-context KV cache. - bsz: batch size. - block_size: number of draft positions per request. - start_pos: absolute decode position (must be > 0); bounds the context. - device: device for the returned index tensor. - - Returns: - int32 tensor ``[bsz, block_size, topk]`` with - ``topk = min(window_size, start_pos+1) + block_size``. - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) - return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() - - -def get_dspark_topk_idxs_batched( - window_size: int, - block_size: int, - start_pos: torch.Tensor, - valid_len: torch.Tensor | None = None, -) -> torch.Tensor: - """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - - Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width - ``min(window_size, start_pos+1) + block_size`` depends on the host int - ``start_pos``), this always returns the **fixed** width ``window_size + - block_size`` and masks the unfilled context slots with ``-1``. The masked - slots are excluded by :func:`dspark_sparse_attn` exactly as if they were - absent while the shape remains CUDA-graph safe. - - Every query attends to the actually written circular-window suffix, followed - by the current-block positions. Without ``valid_len`` this preserves the - legacy ``start_pos``-only behavior. - - Args: - window_size: sliding-window length of the captured-context KV cache. - block_size: number of draft positions per request. - start_pos: ``[G]`` int tensor of per-request absolute decode positions. - valid_len: optional ``[G]`` count of actually written rolling-window - entries. When omitted, preserve the legacy ``start_pos`` mask. - - Returns: - int32 tensor ``[G, block_size, window_size + block_size]``. - """ - device = start_pos.device - g = start_pos.shape[0] - ctx_cols = torch.arange(window_size, device=device) # [win] - if valid_len is None: - valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] - else: - # The valid entries are the contiguous logical suffix ending at - # start_pos, but their physical slots wrap modulo window_size. - valid_len = valid_len.clamp(min=0, max=window_size) - age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) - valid = age < valid_len.unsqueeze(1) - ctx_idx = torch.where( - valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) - ) - blk_idx = window_size + torch.arange(block_size, device=device) # [block] - blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] - row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] - return row.unsqueeze(1).expand(g, block_size, -1).contiguous() - - -def dspark_sparse_attn( - q: torch.Tensor, - kv: torch.Tensor, - attn_sink: torch.Tensor, - topk_idxs: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - """Index-gathered multi-query attention with an attention sink. - - Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For - each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by - ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled - dot-product softmax over them, and adds a per-head learnable *sink* logit that - participates only in the softmax denominator (i.e. an "attend-to-nothing" - option with a zero value vector). KV is shared across query heads (MQA). - - Args: - q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). - kv: ``[b, n, d]`` keys/values (shared across heads). - attn_sink: ``[h]`` per-head sink logits (fp32). - topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). - softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). - - Returns: - ``[b, m, h, d]`` attention output, in ``q.dtype``. - """ - b, m, h, d = q.shape - idx = topk_idxs.long() # [b, m, topk] - valid = idx >= 0 - safe = idx.clamp(min=0) - - # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below - # zeros their softmax probs, so the einsum nullifies them. - kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) - gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() - - # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. - scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale - scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) - - # Online-softmax max is taken over gathered positions only (the sink is added - # to the denominator afterwards), matching the kernel's reduce order. - smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] - smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) - probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 - sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) - denom = probs.sum(dim=-1) + sink # [b, m, h] - out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) - return out.to(q.dtype) - - -def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" - dtype = x.dtype - xf = x.float() - xf = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) - return (weight.float() * xf).to(dtype) - - -def _rope_last_dims( - t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" - nope = t[..., :-rope_head_dim] - rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) - return torch.cat([nope, rope], dim=-1) - - -def _rope_last_dims_batched( - t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Per-row variant of :func:`_rope_last_dims` (``freqs_cis`` has a batch axis).""" - nope = t[..., :-rope_head_dim] - rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) - return torch.cat([nope, rope], dim=-1) - - -def _rmsnorm_rope_batched( - t: torch.Tensor, - weight: torch.Tensor, - eps: float, - rope_head_dim: int, - freqs_cis: torch.Tensor, - *, - num_heads: int = 1, - apply_weight: bool = True, - apply_rmsnorm: bool = True, - inverse_rope: bool = False, -) -> torch.Tensor: - """Fuse DSpark RMSNorm and last-dimension RoPE when supported.""" - if IS_CUTLASS_DSL_AVAILABLE and is_sm_100f(): - freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) - if is_fused_dspark_rmsnorm_rope_supported(t, weight, freqs_real, num_heads, rope_head_dim): - return cute_dsl_dspark_rmsnorm_rope( - t, - weight, - freqs_real, - num_heads, - rope_head_dim, - eps, - apply_weight, - apply_rmsnorm, - inverse_rope, - ) - - if apply_rmsnorm: - if apply_weight: - t = _rmsnorm(t, weight, eps) - else: - t = t * torch.rsqrt(t.square().mean(-1, keepdim=True) + eps) - elif apply_weight: - t = (t.float() * weight.float()).to(t.dtype) - if rope_head_dim > 0: - t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) - return t - - -def dspark_attention_forward( - x: torch.Tensor, - main_x: torch.Tensor, - start_pos: int, - kv_cache: torch.Tensor, - *, - wq_a: torch.Tensor, - q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, - kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, - attn_sink: torch.Tensor, - n_heads: int, - head_dim: int, - rope_head_dim: int, - n_groups: int, - o_lora_rank: int, - window_size: int, - eps: float, - softmax_scale: float, - freqs_cis: torch.Tensor, - persist: bool = False, -) -> torch.Tensor: - """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). - - Functional port of DeepSpec ``DSparkAttention.forward`` for the dense - (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) - with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys - gathered from a rolling captured-context window (``kv_cache``, into which the - projected ``main_x`` context is written at ``start_pos % window_size``) plus the - block's own positions, attention-sink softmax, inverse-RoPE on the output, and a - grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). - - Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - b, block, _ = x.shape - rd = rope_head_dim - main_freqs = freqs_cis[start_pos : start_pos + 1] - blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] - - # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] - main_kv = _rope_last_dims(main_kv, rd, main_freqs) - - # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] - # Per-head RMS in the query dtype (matches the reference inline normalization, - # which is NOT the fp32 RMSNorm path). - q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) - q = _rope_last_dims(q, rd, blk_freqs) - - # Block K/V. - kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] - kv = _rope_last_dims(kv, rd, blk_freqs) - - # Write the context K/V into the rolling window, then attend over - # [window context | block] with the sink. ``persist=True`` writes through - # to the caller's buffer (cross-step decode, worker-owned window); the - # default clones so single-shot callers (golden / unit tests) stay pure. - cache = kv_cache if persist else kv_cache.clone() - cache[:, start_pos % window_size] = main_kv.squeeze(1) - kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] - topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) - o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] - o = _rope_last_dims(o, rd, blk_freqs, inverse=True) - - # Grouped low-rank O projection. - o = o.reshape(b, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) - - -def dspark_attention_forward_batched( - x: torch.Tensor, - main_x: torch.Tensor, - start_pos: torch.Tensor, - kv_cache: torch.Tensor, - slots: torch.Tensor, - valid_len: torch.Tensor | None = None, - *, - wq_a: torch.Tensor, - q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, - kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, - attn_sink: torch.Tensor, - n_heads: int, - head_dim: int, - rope_head_dim: int, - n_groups: int, - o_lora_rank: int, - window_size: int, - eps: float, - softmax_scale: float, - freqs_cis: torch.Tensor, - persist: bool = False, -) -> torch.Tensor: - """Batched, CUDA-graph-safe captured-context DSpark draft attention. - - Numerically identical, per request, to :func:`dspark_attention_forward`, but - free of host syncs and data-dependent shapes so it can be captured into a CUDA - graph (the one-engine drafter runs inside the target's graph). The differences - from the scalar path are purely mechanical: - - * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen - request) instead of a python int; RoPE phases are *gathered* per request from - the fixed ``freqs_cis`` table rather than sliced. - * the rolling-window context K/V is written/read through the ``slots`` index - into a shared ``kv_cache`` (``persist=True`` writes through to the caller's - worker-owned buffer; otherwise a clone is used), instead of mutating a - per-request cache in place. - * the attended-position list has the fixed width ``window_size + block_size`` - with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). - - Args: - x: ``[G, block, dim]`` block layer input (per gen request). - main_x: ``[G, 1, hidden]`` projected captured context. - start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows - (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). - slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. - valid_len: optional ``[G]`` count of actually written context entries; - masks holes left when absolute positions are bootstrapped without - receiving the corresponding DSpark rolling-window state. - freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; - must satisfy ``maxlen > start_pos.max() + block_size``. - - Returns: - ``[G, block, dim]`` attention output (residual stream contribution). - """ - g, block, _ = x.shape - if kv_cache.shape[1] != window_size: - raise ValueError( - f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" - ) - rd = rope_head_dim - # Per-request RoPE phases gathered from the fixed table (no host-int slicing). - main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] - blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] - blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] - - # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm_rope_batched(F.linear(main_x, wkv), kv_norm_w, eps, rd, main_freqs) - - # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] - q = _rmsnorm_rope_batched( - q, - kv_norm_w, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - ) - - # Block K/V. - kv = _rmsnorm_rope_batched(F.linear(x, wkv), kv_norm_w, eps, rd, blk_freqs) - - # Write the context K/V into the rolling window at slot start_pos%window_size, - # then attend over [window context | block]. ``persist=True`` writes through to - # the worker-owned buffer (cross-step decode); otherwise clone so single-shot - # callers stay pure. - write_target = kv_cache if persist else kv_cache.clone() - main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) - if ( - valid_len is None - and IS_CUTLASS_DSL_AVAILABLE - and is_fused_dspark_attention_supported( - q, main_kv_flat, kv, write_target, slots, start_pos, attn_sink - ) - ): - # One custom op performs the rolling-cache write/read, validity handling, - # QK, attention-sink online softmax, and PV. In particular it creates no - # topk index, gathered KV, score, or probability tensors. - o = cute_dsl_dspark_attention( - q, - main_kv_flat, - kv, - write_target, - slots, - start_pos, - attn_sink, - softmax_scale, - ) - else: - slot_pos = start_pos % window_size # [G] - write_target[slots, slot_pos] = main_kv_flat - cache_rows = write_target[slots] # [G, window, head_dim] - kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] - topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) - o = dspark_sparse_attn( - q, kv_full, attn_sink, topk, softmax_scale - ) # [G, block, h, head_dim] - o = _rmsnorm_rope_batched( - o, - kv_norm_w, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - apply_rmsnorm=False, - inverse_rope=True, - ) - - # Grouped low-rank O projection. - o = o.reshape(g, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py deleted file mode 100644 index 1b47f4922965..000000000000 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ /dev/null @@ -1,130 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# DSpark draft I/O logic is ported from DeepSeek's DeepSeek-V4-Pro-DSpark -# reference (`inference/model.py`, DSparkBlock.forward_embed / forward_head). -"""DSpark draft I/O: block input and proposal stages. - -This module holds the *framework-agnostic* (pure-torch) input/output stages of -the DSpark draft block, separated from the heavy V4 backbone (MLA + MoE + mHC) so -they can be unit-tested in isolation: - - - ``build_draft_input_ids``: ``[bonus_token, noise, noise, ...]`` block input. - - ``dspark_propose``: given the per-position backbone ``base_logits`` and the - Markov / confidence heads, run the autoregressive Markov refinement to sample - the block tokens and apply the static confidence-threshold truncation. - -The backbone (3 V4 blocks producing ``block_hidden``) lives in the model module; -this file is the part fully specified by the reference and validated against it. -""" - -from typing import Optional - -import torch -from torch import nn - -from .heads import confident_prefix_length - - -def build_draft_input_ids( - bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int -) -> torch.Tensor: - """``[batch] -> [batch, block_size]`` = ``[bonus, noise, noise, ...]``. - - The first position is the verified bonus token (the target's last accepted - token); the rest are the DSpark noise/mask token (id 128799 for V4-Pro). - """ - batch = bonus_token_ids.shape[0] - out = bonus_token_ids.new_full((batch, block_size), int(noise_token_id)) - out[:, 0] = bonus_token_ids - return out - - -def dspark_propose( - base_logits: torch.Tensor, - *, - bonus_token_ids: torch.Tensor, - block_hidden: torch.Tensor, - markov_head: Optional[nn.Module], - confidence_head: Optional[nn.Module], - block_size: int, - temperature: float = 0.0, - confidence_threshold: float = 0.0, - return_logits: bool = False, -) -> tuple: - """Produce DSpark draft tokens for one block (functional-first, static length). - - Args: - base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. - bonus_token_ids: ``[batch]`` the token preceding the first draft position. - block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the - confidence head, and the RNN-head variant). - markov_head / confidence_head: the validated DSpark heads (may be None). - Returns: - draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers - keep the tensor fixed-width for CUDA-graph safety). - num_proposed: ``[batch]`` int32 — how many leading tokens survive the - static confidence-threshold truncation (== block_size when no head / - threshold<=0). - """ - batch = base_logits.shape[0] - # ``draft_logits`` are the per-position distributions the draft token is drawn - # from (markov-corrected when a head is present, else the raw base logits). - # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance - # (1-TV) measurement; the normal path ignores them. - draft_logits = base_logits - if markov_head is not None: - draft_tokens, corrected = markov_head.sample_block_tokens( - base_logits, - first_prev_token_ids=bonus_token_ids, - hidden_states=block_hidden, - temperature=temperature, - ) - draft_logits = corrected - else: - from .heads import greedy_or_sample - - draft_tokens = greedy_or_sample(base_logits, temperature) - - # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. - # The worker always calls with confidence_threshold=0.0, so the block below is - # inert and num_proposed stays == block_size (the full block is proposed). The - # returned num_proposed is intentionally not yet consumed by the speculative - # scheduler/verifier; wiring it through is a follow-up (see PR description). - num_proposed = torch.full( - (batch,), int(block_size), dtype=torch.int32, device=base_logits.device - ) - if confidence_head is not None and confidence_threshold > 0.0: - # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] - prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) - prev_emb = ( - markov_head.get_prev_embeddings(prev_ids) - if (markov_head is not None and getattr(confidence_head, "with_markov", False)) - else None - ) - conf_logits = ( - confidence_head(block_hidden, prev_embeddings=prev_emb) - if prev_emb is not None - else confidence_head(block_hidden) - ) - # Per-request prefix truncation (batch handled row-wise to stay simple; - # functional-first scope typically runs batch=1 for the draft). - for b in range(batch): - num_proposed[b] = confident_prefix_length( - conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold - ) - if return_logits: - return draft_tokens, num_proposed, draft_logits - return draft_tokens, num_proposed diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py deleted file mode 100644 index c49e35fbafa0..000000000000 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ /dev/null @@ -1,254 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# The DSpark Markov/RNN/confidence-head math is ported from DeepSeek's DeepSpec -# reference implementation (https://github.com/deepseek-ai/DeepSpec, MIT License). -"""DSpark draft-network heads (pure-torch, framework-agnostic). - -These modules implement the *sequential refinement* and *acceptance-confidence* -parts of DeepSeek's DSpark speculative-decoding draft network: - - - Markov head: a low-rank token-bigram logit bias ``logits_k += W2(W1[t_{k-1}])`` - applied autoregressively across the ``block_size`` draft positions (the cheap - "sequential" half of DSpark's "semi-parallel" drafting). RNN variant carries - a GRU-style recurrent state across positions. - - Confidence head: predicts a per-position acceptance probability; the cumulative - product over positions estimates prefix-acceptance and is used only to - *truncate* the proposed draft length (NOT to decide acceptance). - -This file deliberately depends on ``torch`` only so it can be unit-tested in -isolation (token-for-token) against the DeepSpec reference. -""" - -from typing import Optional - -import torch -from torch import nn - - -def greedy_or_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: - """Argmax for temperature<=0, else temperature-scaled multinomial. - - Args: - logits: ``[..., vocab]``. - Returns: - token ids with the trailing vocab dim reduced. - """ - if temperature <= 0.0: - return logits.argmax(dim=-1) - probs = torch.softmax(logits.float() / temperature, dim=-1) - flat = probs.reshape(-1, probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1).squeeze(-1) - return sampled.view(probs.shape[:-1]) - - -class VanillaMarkov(nn.Module): - """Low-rank token-bigram logit bias: ``bias = W2(W1[token])``.""" - - markov_head_type = "vanilla" - - def __init__(self, *, vocab_size: int, markov_rank: int): - super().__init__() - self.vocab_size = int(vocab_size) - self.markov_rank = int(markov_rank) - assert self.markov_rank > 0, ( - f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." - ) - self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) - self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) - - def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: - return self.markov_w1(token_ids.long()) - - def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor: - return self.markov_w2(latent_states) - - def compute_step_bias( - self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor] - ) -> torch.Tensor: - del hidden_states - return self.project_bias(self.get_prev_embeddings(token_ids)) - - def apply_step_logits( - self, - logits: torch.Tensor, - *, - token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - ) -> torch.Tensor: - return logits + self.compute_step_bias(token_ids, hidden_states) - - def sample_block_tokens( - self, - base_logits: torch.Tensor, - *, - first_prev_token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - temperature: float = 0.0, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Autoregressive block sampling with the (memoryless) Markov bias. - - Args: - base_logits: ``[batch, block_size, vocab]`` from the backbone+lm_head. - first_prev_token_ids: ``[batch]`` token preceding the first position. - hidden_states: ``[batch, block_size, d]`` (unused by vanilla/gated). - Returns: - sampled_tokens ``[batch, block_size]``, corrected_logits ``[batch, block_size, vocab]``. - """ - batch_size, block_size = base_logits.shape[:2] - if block_size == 0: - empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) - return empty, base_logits - sampled, corrected = [], [] - prev = first_prev_token_ids.long() - for k in range(block_size): - step_hidden = None if hidden_states is None else hidden_states[:, k] - step_logits = self.apply_step_logits( - base_logits[:, k], token_ids=prev, hidden_states=step_hidden - ) - corrected.append(step_logits.unsqueeze(1)) - prev = greedy_or_sample(step_logits, temperature) - sampled.append(prev) - return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) - - -class GatedMarkovHead(VanillaMarkov): - """Markov bias gated by a sigmoid of [hidden, prev_embedding].""" - - markov_head_type = "gated" - - def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): - super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) - self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank) - - def compute_step_bias( - self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor] - ) -> torch.Tensor: - assert hidden_states is not None - prev_emb = self.get_prev_embeddings(token_ids) - gate = torch.sigmoid(self.gate_proj(torch.cat([hidden_states, prev_emb], dim=-1))).to( - dtype=prev_emb.dtype - ) - return self.project_bias(gate * prev_emb) - - -class RNNHead(VanillaMarkov): - """GRU-style head carrying recurrent state across block positions.""" - - markov_head_type = "rnn" - - def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): - super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) - self.hidden_size = int(hidden_size) - # [s_{k-1}; W1[x_{k-1}]; h_k] -> [gate; candidate; output] - self.joint_proj = nn.Linear(2 * markov_rank + hidden_size, 3 * markov_rank) - - def _rnn_step(self, state, prev_embeddings, hidden_states): - z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) - gate_raw, cand_raw, out_raw = self.joint_proj(z).chunk(3, dim=-1) - gate = torch.sigmoid(gate_raw) - candidate = torch.tanh(cand_raw) - new_state = gate * state + (1.0 - gate) * candidate - bias = self.project_bias(torch.tanh(out_raw)) - return new_state, bias - - def sample_block_tokens( - self, - base_logits: torch.Tensor, - *, - first_prev_token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - temperature: float = 0.0, - ) -> tuple[torch.Tensor, torch.Tensor]: - assert hidden_states is not None - batch_size, block_size = base_logits.shape[:2] - if block_size == 0: - empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) - return empty, base_logits - state = torch.zeros( - batch_size, self.markov_rank, device=base_logits.device, dtype=hidden_states.dtype - ) - sampled, corrected = [], [] - prev = first_prev_token_ids.long() - for k in range(block_size): - prev_emb = self.get_prev_embeddings(prev) - state, bias = self._rnn_step(state, prev_emb, hidden_states[:, k]) - step_logits = base_logits[:, k] + bias - corrected.append(step_logits.unsqueeze(1)) - prev = greedy_or_sample(step_logits, temperature) - sampled.append(prev) - return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) - - -def build_markov_head( - *, markov_head_type: str, vocab_size: int, markov_rank: int, hidden_size: int -) -> Optional[nn.Module]: - """Factory mirroring DeepSpec ``build_markov_head``; returns None if rank==0.""" - if int(markov_rank) <= 0: - return None - kind = str(markov_head_type).lower() - if kind == "vanilla": - return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank) - if kind == "gated": - return GatedMarkovHead( - vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size - ) - if kind == "rnn": - return RNNHead(vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size) - raise ValueError(f"Unsupported markov_head_type: {markov_head_type!r}") - - -class DSparkConfidenceHead(nn.Module): - """Per-position acceptance-confidence predictor (DeepSpec AcceptRatePredictor). - - Input features are the backbone hidden state, optionally concatenated with the - Markov head's previous-token embedding. Output is a single logit per position. - """ - - def __init__(self, *, hidden_size: int, markov_rank: int = 0, with_markov: bool = False): - super().__init__() - self.with_markov = bool(with_markov) - input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) - # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the - # confidence score is computed in fp32 (mirrors the DeepSpec reference - # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). - self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) - - def forward( - self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None - ) -> torch.Tensor: - if self.with_markov: - assert prev_embeddings is not None - features = torch.cat([hidden_states, prev_embeddings.to(hidden_states.dtype)], dim=-1) - else: - features = hidden_states - # fp32 matmul for a stable confidence score (mirrors the reference). - return self.proj(features.float()).squeeze(-1) - - -def confident_prefix_length( - confidence_logits: torch.Tensor, *, block_size: int, threshold: float -) -> int: - """First position k where ``sigmoid(confidence_k) < threshold``. - - Returns ``block_size`` when threshold<=0 (no truncation) or all positions - are confident. Assumes batch size 1 (functional-first scope). - """ - if threshold <= 0.0: - return int(block_size) - below = confidence_logits.sigmoid() < threshold - if not bool(below[0].any().item()): - return int(block_size) - return int(torch.nonzero(below[0], as_tuple=False)[0].item()) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py new file mode 100644 index 000000000000..4ea0222e0c52 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -0,0 +1,1323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +from dataclasses import replace +from typing import Dict, Optional + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm.logger import logger + +from ...functional import RotaryScalingType +from ..modules.rotary_embedding import RotaryEmbedding + +try: + from ..custom_ops import flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope +except ImportError: + _flashinfer_rope = None +from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window +from ..speculative.dflash_attention import get_dflash_flash_attention, get_dflash_trtllm_gen_ops +from ..speculative.interface import SpeculativeDecodingMode +from .modeling_utils import get_model_architecture, register_draft_model + + +def dspark_layer_window_size( + use_swa: bool, swa_window: int, layer_types, layer_idx: int +) -> tuple[int, int]: + """flash-attn ``window_size`` for one draft layer of the block decode. + + Sliding-window block decode was introduced by the DSpark drafters + (deepseek-ai/DeepSpec) -- hence the name -- but it is an attention + configuration, not part of the DSpark head set, so it stays in the DFlash + base: the block decode below indexes the resolved windows directly and must + not depend on a subclass-only attribute. + + Those drafters run the draft block through HF attention with + ``sliding_window`` set on 'sliding_attention' layers and is_causal=False. + HF's flash path (transformers/modeling_flash_attention_utils.py) translates + that to ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. + each query attends keys within ``swa_window - 1`` KV-index distance on both + sides. In the DFlash pool layout KV index == token position, so this limits + draft queries to the most recent ``swa_window`` context tokens plus the + (nearby) draft block. Full-attention layers and drafters that do not enable + the window keep flash-attn's default ``(-1, -1)`` (no window). + """ + if not use_swa: + return (-1, -1) + if ( + layer_types is not None + and layer_idx < len(layer_types) + and layer_types[layer_idx] != "sliding_attention" + ): + return (-1, -1) + return (swa_window - 1, swa_window - 1) + + +class DFlashForCausalLM(nn.Module): + """Draft model wrapper for DFlash speculative decoding. + + DFlash uses cross-attention where Q comes from noise/query tokens and K/V + come from the concatenation of target hidden states and noise hidden states. + The target_hidden stays CONSTANT across all layers (no input_layernorm applied). + + Reference: https://arxiv.org/pdf/2602.06036 + """ + + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + """Build the draft model, resolving its architecture from the draft config + (falling back to a model_type-derived name when the checkpoint uses a + custom DFlash architecture label).""" + super().__init__() + + pretrained_cfg = draft_config.pretrained_config + try: + DraftModelClass, _ = get_model_architecture(pretrained_cfg) + except RuntimeError: + model_type = pretrained_cfg.model_type + arch_name = "".join(w.capitalize() for w in model_type.split("_")) + "ForCausalLM" + logger.info( + f"DFlash: architecture {pretrained_cfg.architectures} not found, " + f"falling back to {arch_name} based on model_type={model_type}" + ) + original_archs = pretrained_cfg.architectures + try: + pretrained_cfg.architectures = [arch_name] + DraftModelClass, _ = get_model_architecture(pretrained_cfg) + finally: + pretrained_cfg.architectures = original_archs + + # Remove spec_config to prevent recursive spec-dec initialization + draft_config_no_spec = replace(draft_config, spec_config=None, lm_head_gather_output=False) + + # Weights will be loaded later by ModelLoader.load_draft_weights() + self.draft_model_full = DraftModelClass(draft_config_no_spec) + self.model = self.draft_model_full.model + self.lm_head = self.draft_model_full.lm_head + + # Required by weight mappers + self.model_config = draft_config_no_spec + self.config = draft_config_no_spec.pretrained_config + + # Get mask_token_id from dflash_config + pretrained_config = draft_config.pretrained_config + dflash_config = getattr(pretrained_config, "dflash_config", {}) + self.mask_token_id = dflash_config.get( + "mask_token_id", + getattr(pretrained_config, "mask_token_id", pretrained_config.vocab_size), + ) + + self.target_layer_ids = dflash_config.get("target_layer_ids", None) + self.block_size = dflash_config.get( + "block_size", getattr(pretrained_config, "block_size", None) + ) + self.dflash_attention_backend = dflash_attention_backend + if self.dflash_attention_backend == "VANILLA": + self._dflash_flash_attention = get_dflash_flash_attention() + elif self.dflash_attention_backend == "TRTLLM": + self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() + else: + raise ValueError( + "DFlash attention backend must be VANILLA or TRTLLM, got " + f"{self.dflash_attention_backend!r}." + ) + self._dflash_trtllm_gen_workspace = None + self._dflash_trtllm_gen_counters = None + self.register_buffer("_dflash_batch_indices", None, persistent=False) + self.register_buffer("_dflash_block_offsets", None, persistent=False) + self._dflash_trtllm_gen_device = None + self._dflash_trtllm_gen_sm_count = None + logger.info( + f"DFlash draft model initialized with mask_token_id: {self.mask_token_id}, " + f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}, " + f"attention_backend: {self.dflash_attention_backend}" + ) + + # Sliding-window block decode (use_swa / swa_window_size). Kept in the + # base class rather than in the DSpark subclass because the block decode + # below indexes ``self._layer_windows`` directly: a base-class forward + # must not reach for an attribute only a subclass defines. Drafters that + # leave use_swa unset get all-(-1,-1) windows, i.e. a no-op. + self._use_swa = bool(dflash_config.get("use_swa", False)) + self._swa_window = int(dflash_config.get("swa_window_size", 0) or 0) + if self._use_swa and self._swa_window < 1: + raise ValueError( + "DFlash drafter sets use_swa but swa_window_size=" + f"{dflash_config.get('swa_window_size')} is invalid." + ) + # Per-layer flash-attn window for the block decode, resolved once. + num_draft_layers = getattr(pretrained_config, "num_hidden_layers", 0) + layer_types = getattr(pretrained_config, "layer_types", None) + self._layer_windows = [ + dspark_layer_window_size(self._use_swa, self._swa_window, layer_types, i) + for i in range(num_draft_layers) + ] + + self.logits_processor = None # Set by caller after construction + + # RoPE - lazily initialized from draft model's attention module + self._rope_initialized = False + self._rotary_cos_sin = None + self._is_neox = True + + self._cos_sin_cache_fp32 = None + self._rope_dummy_q = None + + # Lazy-built after weights load (see _build_fused_kv_buffers). + self._fused_kv_weight = None + self._fused_kv_bias = None + self._k_norm_stacked = None + self._k_norm_eps = None + self._num_attn_layers = 0 + self._num_heads = 0 + self._head_dim = 0 + self._num_kv_heads = 0 + self._has_qk_norm = False + self._use_fused_qk_norm_rope = False + # Laguna-specific draft-layer behaviors, disabled by default so generic + # DFlash drafters keep the original contract (no context input_layernorm, + # non-causal block attention). Subclasses opt in. + self._context_input_layernorm = False + self._sliding_layers_causal = False + self._validate_gqa_shape() + self._warn_inferred_attention_windows() + + def _validate_gqa_shape(self): + """Reject a backbone the block decode cannot express. + + The hand-written block decode is GQA-shaped throughout: it splits a + fused ``qkv_proj`` by (num_heads, num_kv_heads, head_dim), fuses K/V + across layers on one uniform head dim, and shares a single RoPE cache. + The registry will happily build a backbone that violates any of that -- + an MLA drafter has no per-head K/V at all -- and without this the + failure surfaces much later inside ``_build_fused_kv_buffers`` or, worse, + as silently mis-sliced weights. Raise at construction instead. + """ + layers = getattr(self.model, "layers", None) + if not layers: + return + for idx, layer in enumerate(layers): + attn = getattr(layer, "self_attn", None) + if attn is None or not hasattr(attn, "qkv_proj"): + raise ValueError( + f"DFlash block decode requires a fused self_attn.qkv_proj on " + f"every draft layer, but layer {idx} of draft backbone " + f"{type(self.config).__name__} has none. Backbones that do not " + "project per-head Q/K/V (e.g. MLA) need their own block decode." + ) + num_kv_heads = layers[0].self_attn.num_key_value_heads + mismatched = [ + idx + for idx, layer in enumerate(layers[1:], start=1) + if layer.self_attn.num_key_value_heads != num_kv_heads + ] + if mismatched: + raise ValueError( + "DFlash fuses draft K/V across layers and needs one uniform " + f"num_key_value_heads, but layers {mismatched} differ from layer 0 " + f"({num_kv_heads})." + ) + + @staticmethod + def _rope_signature(attn): + """Return the effective RoPE configuration used by an attention layer.""" + if attn.rotary_emb is not None: + return ( + attn.rotary_emb.rope_params, + attn.rotary_emb.head_dim, + attn.rotary_emb.is_neox, + ) + if attn.pos_embd_params is not None: + return ( + attn.pos_embd_params.rope, + attn.head_dim, + attn.pos_embd_params.is_neox, + ) + return None + + def _validate_uniform_rope(self): + """Check that all draft layers can safely share one RoPE cache.""" + if len(self.model.layers) == 0: + raise ValueError("DFlash requires at least one draft model layer.") + + signatures = [self._rope_signature(layer.self_attn) for layer in self.model.layers] + + mismatched_layers = [ + layer_idx + for layer_idx, signature in enumerate(signatures[1:], start=1) + if signature != signatures[0] + ] + if mismatched_layers: + layer_types = getattr(self.config, "layer_types", None) + raise ValueError( + "DFlash shares one RoPE cache across draft layers, but layers " + f"{mismatched_layers} have a different effective RoPE " + f"configuration from layer 0. layer_types={layer_types}." + ) + + def _init_rope(self): + """Initialize RoPE from the draft model's attention configuration. + + Reuses the existing RotaryEmbedding infrastructure which correctly + handles all RoPE variants (standard, YaRN, scaled, etc.). + """ + # The flattened context-KV path shares layer 0's RoPE cache. + self._validate_uniform_rope() + attn0 = self.model.layers[0].self_attn + + if attn0.rotary_emb is not None: + self._rotary_cos_sin = attn0.rotary_emb.rotary_cos_sin + self._is_neox = attn0.rotary_emb.is_neox + elif attn0.pos_embd_params is not None: + rope_emb = RotaryEmbedding( + attn0.pos_embd_params.rope, + head_dim=attn0.head_dim, + is_neox=attn0.pos_embd_params.is_neox, + ) + self._rotary_cos_sin = rope_emb.rotary_cos_sin + self._is_neox = rope_emb.is_neox + else: + # Fallback: basic NeoX-style RoPE + config = self.config + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + rope_theta = getattr(config, "rope_theta", 1000000.0) + max_pos = getattr(config, "max_position_embeddings", 32768) + + inv_freq = 1.0 / ( + rope_theta + ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device="cuda") / head_dim) + ) + positions = torch.arange(max_pos, dtype=torch.float32, device="cuda") + freqs = torch.outer(positions, inv_freq) + rope_cos = freqs.cos().to(config.torch_dtype) + rope_sin = freqs.sin().to(config.torch_dtype) + # [max_pos, 2, rot_dim//2] to match RotaryEmbedding format + self._rotary_cos_sin = torch.stack([rope_cos, rope_sin], dim=1) + self._is_neox = True + + self._rope_initialized = True + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project captured target hidden states into the draft hidden space. + + Generic DFlash: fc then hidden_norm. Subclasses (e.g. Laguna) may + normalize the per-aux features first by overriding this method. + """ + hidden_states = hidden_states.to(self.fc.weight.dtype) + return self.hidden_norm(self.fc(hidden_states)) + + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): + """Hook applied to the block-attention output before o_proj. + + No-op for generic DFlash; overridden by drafters that gate (e.g. Laguna). + """ + return attn_output + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Load weights into the DFlash draft model. + + DFlash checkpoints differ from standard HF format: + - Layer weights lack the 'model.' prefix (e.g., 'layers.0...' not 'model.layers.0...') + - Extra DFlash-specific weights: 'fc.weight', 'hidden_norm.weight' + - Missing embed_tokens and lm_head (shared with target model) + """ + # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft + # loader expects split q/k/v (a fused key is silently dropped otherwise). + if any(k.endswith("self_attn.qkv_proj.weight") for k in weights): + for attr in ("num_attention_heads_per_layer", "num_key_value_heads_per_layer"): + per_layer = getattr(self.config, attr, None) + if per_layer is not None and len(set(per_layer)) > 1: + raise ValueError( + "DFlash load_weights() splits the fused qkv_proj using " + "the global head count, but the drafter has heterogeneous " + f"{attr} {sorted(set(per_layer))}; per-layer qkv splitting " + "is required for this checkpoint." + ) + head_dim = getattr( + self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads + ) + num_kv_heads = getattr( + self.config, "num_key_value_heads", self.config.num_attention_heads + ) + q = self.config.num_attention_heads * head_dim + kv = num_kv_heads * head_dim + split = {} + for k, v in weights.items(): + if k.endswith("self_attn.qkv_proj.weight"): + b = k[: -len("qkv_proj.weight")] + split[b + "q_proj.weight"] = v[:q] + split[b + "k_proj.weight"] = v[q : q + kv] + split[b + "v_proj.weight"] = v[q + kv :] + else: + split[k] = v + weights = split + + # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights + remapped = {} + for key, value in weights.items(): + if key in ("fc.weight", "hidden_norm.weight"): + # DFlash-specific projection weights - store directly + remapped[key] = value + elif key == "norm.weight": + remapped["model.norm.weight"] = value + elif not key.startswith("model."): + remapped[f"model.{key}"] = value + else: + remapped[key] = value + + # Load DFlash-specific weights directly + if "fc.weight" in remapped: + self.fc = nn.Linear( + remapped["fc.weight"].shape[1], + remapped["fc.weight"].shape[0], + bias=False, + device="cuda", + dtype=remapped["fc.weight"].dtype, + ) + self.fc.weight.data.copy_(remapped["fc.weight"]) + del remapped["fc.weight"] + + if "hidden_norm.weight" in remapped: + rms_norm_eps = getattr(self.config, "rms_norm_eps", 1e-6) + self.hidden_norm = nn.RMSNorm( + remapped["hidden_norm.weight"].shape[0], + eps=rms_norm_eps, + device="cuda", + elementwise_affine=True, + dtype=remapped["hidden_norm.weight"].dtype, + ) + self.hidden_norm.weight.data.copy_(remapped["hidden_norm.weight"]) + del remapped["hidden_norm.weight"] + + # Load remaining weights into the draft model. + # DFlash checkpoints don't include embed_tokens or lm_head, so allow partial loading + # since those modules won't find matching weights. + self.draft_model_full.load_weights( + weights=remapped, weight_mapper=weight_mapper, allow_partial_loading=True + ) + + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: + """Share embed_tokens and lm_head from the target model.""" + self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens + self.draft_model_full.lm_head = target_model.lm_head + self.lm_head = target_model.lm_head + + def precompute_context_kv( + self, + projected_hidden: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Post-norm / post-RoPE K and V for ALL drafter layers in one fused GEMM. + + Args: + projected_hidden: [N, hidden_size], already fc + hidden_norm'd. + positions: [N] int32/64, RoPE positions for each entry. + Returns: + k: [N, L, nkv, hd] post k_norm and RoPE + v: [N, L, nkv, hd] post split only + """ + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + N = projected_hidden.shape[0] + L = self._num_attn_layers + nkv = self._num_kv_heads + hd = self._head_dim + weight_dtype = self._fused_kv_weight.dtype + if getattr(self, "_input_ln_eps", None) is not None: + ph = projected_hidden.float() + ph = ph * torch.rsqrt(ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) + projected_hidden = ph.to(weight_dtype) + elif projected_hidden.dtype != weight_dtype: + projected_hidden = projected_hidden.to(weight_dtype) + + kv_flat = F.linear(projected_hidden, self._fused_kv_weight, self._fused_kv_bias) + # Per-layer layout [L0_K|L0_V|L1_K|L1_V|...] keeps K and V contiguous + # after the select() splits — no extra copy required. + kv = kv_flat.view(N, L, 2, nkv, hd) + k = kv[:, :, 0].contiguous() + v = kv[:, :, 1].contiguous() + + if self._k_norm_stacked is not None: + # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; + # each layer has its own weight ([L, hd]) but shares eps. + k = F.rms_norm(k, (hd,), eps=self._k_norm_eps) + k = k * self._k_norm_stacked.view(1, L, 1, hd) + + self._fused_rope_inplace(k.view(N * L, nkv * hd), positions, N, L) + return k, v + + def _get_cos_sin_cache(self) -> torch.Tensor: + """Return the flashinfer-style cos/sin cache for the drafter. + + Shape [max_positions, head_dim], fp32 — flashinfer's + apply_rope_with_cos_sin_cache_inplace requires fp32 regardless of + the query/key dtype. + """ + if self._cos_sin_cache_fp32 is not None: + return self._cos_sin_cache_fp32 + if not self._rope_initialized: + self._init_rope() + max_pos = self._rotary_cos_sin.shape[0] + self._cos_sin_cache_fp32 = ( + self._rotary_cos_sin.view(max_pos, -1).to(torch.float32).contiguous() + ) + return self._cos_sin_cache_fp32 + + def _fused_rope_inplace( + self, + k_flat: torch.Tensor, + positions: torch.Tensor, + N: int, + L: int, + ) -> None: + """In-place fused RoPE over [N*L, nkv*hd] K values. + + Layout of k_flat: row (i*L + l) holds layer l of position i, so + positions must be repeat_interleaved by L to match. + """ + positions_int32 = positions.view(-1).to(torch.int32) + if L > 1: + positions_int32 = positions_int32.repeat_interleave(L) + + if _flashinfer_rope is not None: + # flashinfer requires a non-None query tensor; pass a single-head + # scratch so the extra rotate is negligible. + need_rows = k_flat.shape[0] + dummy_q = self._rope_dummy_q + if dummy_q is None or dummy_q.dtype != k_flat.dtype or dummy_q.shape[0] < need_rows: + dummy_q = k_flat.new_empty(need_rows, self._head_dim) + self._rope_dummy_q = dummy_q + _flashinfer_rope( + positions_int32, + dummy_q[:need_rows], + k_flat, + self._head_dim, + self._get_cos_sin_cache(), + self._is_neox, + ) + return + + # Pure-PyTorch fallback (older environments without flashinfer). + cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), dtype=k_flat.dtype) + k_roped = RotaryEmbedding.apply_rotary_pos_emb( + k_flat.view(k_flat.shape[0], -1, self._head_dim), + cos.squeeze(0), + sin.squeeze(0), + unsqueeze_dim=1, + is_neox=self._is_neox, + ) + k_flat.copy_(k_roped.view_as(k_flat)) + + def _build_fused_kv_buffers(self) -> None: + """Stack per-layer KV projection + k_norm weights for a single fused GEMM. + + Must run after weights are loaded. + """ + if self._fused_kv_weight is not None: + return + layers_attn = [layer.self_attn for layer in self.model.layers] + attn0 = layers_attn[0] + q_size = attn0.q_size + kv_size = attn0.kv_size + head_dim = attn0.head_dim + num_heads = attn0.num_heads + num_kv_heads = attn0.num_key_value_heads + # Head counts are read from layer 0 here and in dflash_forward; assert + # uniformity (the target uses per-layer heads, the drafter does not). + for a in layers_attn[1:]: + assert ( + a.q_size == q_size + and a.kv_size == kv_size + and a.head_dim == head_dim + and a.num_heads == num_heads + and a.num_key_value_heads == num_kv_heads + ), ( + "DFlash fused KV requires all drafter layers to share " + "q_size / kv_size / head_dim / num_heads / num_kv_heads." + ) + + has_k_norm = [hasattr(a, "k_norm") for a in layers_attn] + assert all(has_k_norm) or not any(has_k_norm), ( + "DFlash fused KV requires either all or no drafter layers to have k_norm." + ) + + kv_weights = [a.qkv_proj.weight[q_size : q_size + 2 * kv_size] for a in layers_attn] + # Fold each drafter layer's input_layernorm weight into its KV projection + # so context K/V match the query path. vLLM laguna_dflash applies + # layer.input_layernorm to context states before KV; RMSNorm gives + # (x_hat * w) @ Wkv.T == x_hat @ (Wkv * w).T, and the shared 1/rms(x) is + # applied to projected_hidden in precompute_context_kv. + dlayers = self.model.layers + if self._context_input_layernorm and all(hasattr(dl, "input_layernorm") for dl in dlayers): + eps_set = { + getattr( + dl.input_layernorm, + "variance_epsilon", + getattr(self.config, "rms_norm_eps", 1e-6), + ) + for dl in dlayers + } + assert len(eps_set) == 1, ( + "DFlash fused context input_layernorm needs all drafter layers " + f"to share variance_epsilon; got {sorted(eps_set)}" + ) + self._input_ln_eps = eps_set.pop() + folded = [] + for w, dl in zip(kv_weights, dlayers): + scale = dl.input_layernorm.weight.data + if getattr(dl.input_layernorm, "use_gemma", False): + scale = scale + 1 + folded.append(w * scale[None, :].to(w.dtype)) + kv_weights = folded + else: + self._input_ln_eps = None + fused_kv_weight = torch.cat(kv_weights, dim=0).contiguous() + if attn0.qkv_proj.bias is not None: + kv_biases = [a.qkv_proj.bias[q_size : q_size + 2 * kv_size] for a in layers_attn] + self._fused_kv_bias = torch.cat(kv_biases, dim=0).contiguous() + else: + self._fused_kv_bias = None + + if all(has_k_norm): + k_norm0 = layers_attn[0].k_norm + eps = k_norm0.variance_epsilon + eps_set = {a.k_norm.variance_epsilon for a in layers_attn} + assert len(eps_set) == 1, ( + f"DFlash fused k_norm requires all drafter layers to share " + f"variance_epsilon; got {sorted(eps_set)}." + ) + self._k_norm_stacked = torch.stack([a.k_norm.weight.data for a in layers_attn]) + self._k_norm_eps = eps + else: + self._k_norm_stacked = None + self._k_norm_eps = None + self._num_attn_layers = len(layers_attn) + self._num_heads = num_heads + self._head_dim = head_dim + self._num_kv_heads = num_kv_heads + self._fused_kv_weight = fused_kv_weight + + # fused_qk_norm_rope derives YaRN / partial-rotary frequencies on + # the fly, which can disagree with precompute_context_kv's cached + # cos/sin. Only enable it when the drafter uses plain RoPE. + self._has_qk_norm = all(has_k_norm) and all(hasattr(a, "q_norm") for a in layers_attn) + rope_params = getattr(getattr(attn0, "pos_embd_params", None), "rope", None) + scale_type = getattr(rope_params, "scale_type", None) + partial_rotary_factor = getattr( + getattr(attn0, "pretrained_config", None), "partial_rotary_factor", 1.0 + ) + self._use_fused_qk_norm_rope = ( + self._has_qk_norm + and hasattr(attn0, "apply_qk_norm_rope") + and rope_params is not None + and scale_type in (None, RotaryScalingType.none) + and partial_rotary_factor == 1.0 + ) + + logger.debug( + f"DFlash: fused KV weights built for {self._num_attn_layers} layers " + f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})" + ) + + def _get_rope_cos_sin(self, positions, dtype=None): + """Get cos/sin for given positions, suitable for apply_rotary_pos_emb. + + Args: + positions: [B, seq_len] + dtype: target dtype for cos/sin (default: keep original) + Returns: + rope_cos: [B, seq, rot_dim//2] (broadcastable with unsqueeze_dim=1) + rope_sin: [B, seq, rot_dim//2] + """ + if not self._rope_initialized: + self._init_rope() + + # rotary_cos_sin: [max_pos, 2, rot_dim//2] + rope_cache = self._rotary_cos_sin[positions] # [B, seq, 2, rot_dim//2] + rope_cos = rope_cache[..., 0, :] # [B, seq, rot_dim//2] + rope_sin = rope_cache[..., 1, :] + if dtype is not None: + rope_cos = rope_cos.to(dtype) + rope_sin = rope_sin.to(dtype) + return rope_cos, rope_sin + + def _warn_inferred_attention_windows(self) -> None: + """Warn once at initialization when checkpoint metadata enables SWA.""" + if getattr(self.config, "use_sliding_window", None) is not None: + return + + num_hidden_layers = getattr(self.config, "num_hidden_layers", None) + if num_hidden_layers is None: + num_hidden_layers = len(self.model.layers) + layers_by_window = {} + for layer_idx in range(num_hidden_layers): + window = get_layer_attention_window(self.config, layer_idx) + if window is not None: + layers_by_window.setdefault(window, []).append(layer_idx) + + for window, layer_indices in layers_by_window.items(): + logger.warning( + "DFlash inferred pooled-context sliding-window attention from " + f"checkpoint config for draft layers {layer_indices}: " + f"window={window}. Context attention is truncated to {window} " + "tokens for these layers; if the drafter expects full context, " + "acceptance rate may drop. Set use_sliding_window explicitly " + "to confirm or disable windowing." + ) + + def _get_attention_mask_args(self, layer_idx): + """Return FlashAttention causal and local-window arguments for a layer.""" + layer_types = getattr(self.config, "layer_types", None) + is_sliding_layer = False + if layer_types: + layer_type = layer_types[layer_idx % len(layer_types)] + is_sliding_layer = _is_sliding_attention_layer(layer_type) + + sliding_window = get_layer_attention_window(self.config, layer_idx) + is_sliding_layer = is_sliding_layer or sliding_window is not None + if not is_sliding_layer: + return False, (-1, -1) + + causal = self._sliding_layers_causal or sliding_window is not None + if sliding_window is None: + # Legacy drafters without an explicit window preserve their prior + # non-windowed behavior. + return causal, (-1, -1) + # FlashAttention's bounds are inclusive: W tokens are current + W-1 left. + return causal, (sliding_window - 1, 0) + + def _prepare_dflash_trtllm_gen_buffers( + self, + dtype: torch.dtype, + device: torch.device, + max_batch_size: int, + block_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + ) -> None: + trtllm_gen_ops = self._dflash_trtllm_gen_ops + workspace_bytes = trtllm_gen_ops.get_workspace_size( + dtype=dtype, + num_tokens=max_batch_size * block_size, + num_gen_tokens=max_batch_size * block_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + max_num_requests=max_batch_size, + rotary_embedding_dim=0, + fp8_context_fmha=False, + ) + device = torch.device(device) + is_capturing = torch.cuda.is_current_stream_capturing() + if self._dflash_trtllm_gen_device != device: + if is_capturing: + raise RuntimeError( + "DFlash TRTLLM-Gen buffers must be prepared on the current " + "device before CUDA graph capture." + ) + self._dflash_trtllm_gen_device = device + self._dflash_trtllm_gen_sm_count = torch.cuda.get_device_properties( + device + ).multi_processor_count + + workspace = self._dflash_trtllm_gen_workspace + workspace_needs_allocation = ( + workspace is None + or workspace.device != device + or workspace.numel() * workspace.element_size() < workspace_bytes + ) + if workspace_needs_allocation: + if is_capturing: + raise RuntimeError( + "The DFlash TRTLLM-Gen workspace must be allocated at the " + "required size before CUDA graph capture." + ) + self._dflash_trtllm_gen_workspace = torch.empty( + workspace_bytes, dtype=torch.uint8, device=device + ) + + sm_count = self._dflash_trtllm_gen_sm_count + counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( + num_heads, max_batch_size, sm_count + ) + counters = self._dflash_trtllm_gen_counters + counters_need_allocation = ( + counters is None + or counters.device != device + or counters.numel() * counters.element_size() < counter_bytes + ) + if counters_need_allocation: + if is_capturing: + raise RuntimeError( + "The DFlash TRTLLM-Gen counter buffer must be allocated at " + "the required size before CUDA graph capture." + ) + self._dflash_trtllm_gen_counters = torch.zeros( + counter_bytes, dtype=torch.uint8, device=device + ) + + append_batch_indices = self._dflash_batch_indices + block_offsets = self._dflash_block_offsets + static_indices_need_allocation = ( + append_batch_indices is None + or block_offsets is None + or append_batch_indices.device != device + or block_offsets.device != device + or append_batch_indices.size(0) < max_batch_size + or append_batch_indices.size(1) != block_size + or block_offsets.numel() != block_size + ) + if static_indices_need_allocation: + if is_capturing: + raise RuntimeError( + "DFlash TRTLLM-Gen index buffers must be allocated at the " + "required size before CUDA graph capture." + ) + self._dflash_batch_indices = ( + torch.arange(max_batch_size, dtype=torch.int32, device=device) + .view(-1, 1) + .expand(-1, block_size) + .contiguous() + ) + self._dflash_block_offsets = torch.arange(block_size, dtype=torch.int32, device=device) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """DFlash draft forward with cross-attention over a pooled K/V buffer. + + All shapes are fixed so the forward is CUDA-graph compatible. + + Args: + noise_embedding: [B, block_size, hidden_size] + query_positions: [B, block_size] + num_ctx_per_req: [B] — per-batch context length in the pool + ctx_k_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] + ctx_v_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] + ctx_cache_batch_idx: [B] — slot index into the pool per batch entry + Returns: + [B * block_size, hidden_size] + """ + if self.dflash_attention_backend == "TRTLLM": + if ctx_kv_cache is None or ctx_page_table is None: + raise RuntimeError( + "DFlash TRTLLM-Gen requires a paged context cache and page table." + ) + trtllm_gen_ops = self._dflash_trtllm_gen_ops + elif self.dflash_attention_backend == "VANILLA": + flash_attention = self._dflash_flash_attention + else: + raise ValueError( + "DFlash attention backend must be VANILLA or TRTLLM, got " + f"{self.dflash_attention_backend!r}." + ) + + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + + layer0 = self.model.layers[0] + attn0 = layer0.self_attn + q_size = attn0.q_size + kv_size = attn0.kv_size + head_dim = attn0.head_dim + # Uniformity across layers is asserted in _build_fused_kv_buffers (above). + num_heads_per_rank = attn0.num_heads + num_kv_heads_per_rank = attn0.num_key_value_heads + gqa_group_size = num_heads_per_rank // num_kv_heads_per_rank + + has_qk_norm = self._has_qk_norm + is_bf16 = noise_embedding.dtype == torch.bfloat16 + use_fused_qk_norm_rope = self._use_fused_qk_norm_rope and is_bf16 + use_fused_rope = ( + _flashinfer_rope is not None and has_qk_norm and is_bf16 and not use_fused_qk_norm_rope + ) + + B = noise_embedding.shape[0] + block_size = noise_embedding.shape[1] + + hidden_states = noise_embedding # [B, block_size, hidden] + + # Precompute RoPE cos/sin for the pure-PyTorch fallback path only. + # The fused flashinfer path reads self._get_cos_sin_cache() inline. + rope_dtype = hidden_states.dtype + if not use_fused_rope: + q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, dtype=rope_dtype) + _rope = RotaryEmbedding.apply_rotary_pos_emb + + # cache_seqlens (BEFORE append). flash_attn appends block_size + # k/v at cache_seqlens[i]..+block_size for batch i. + cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) + cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) + + if self.dflash_attention_backend == "TRTLLM": + max_batch_size = ctx_page_table.size(0) + self._prepare_dflash_trtllm_gen_buffers( + hidden_states.dtype, + hidden_states.device, + max_batch_size, + block_size, + num_heads_per_rank, + num_kv_heads_per_rank, + head_dim, + ) + block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long()) + pages_per_slot = block_tables.size(1) + page_size = ctx_kv_cache.size(-2) + kv_indices = block_tables.flatten() + kv_indptr = torch.arange( + 0, + (B + 1) * pages_per_slot, + pages_per_slot, + dtype=torch.int32, + device=hidden_states.device, + ) + seq_lens_after = cache_seqlens_i32 + block_size + kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 + batch_indices = self._dflash_batch_indices + append_batch_indices = batch_indices[:B].reshape(-1) + append_positions = ( + (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets) + .reshape(-1) + .contiguous() + ) + + # Flatten query positions once for the fused QK-norm-RoPE kernel. + query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) + + residual = None + + for layer_idx, layer in enumerate(self.model.layers): + attn_mod = layer.self_attn + + # Apply input_layernorm (flatten to 2D for norm, reshape back) + hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if residual is None: + residual = hidden_states.clone() + hs_normed_flat = layer.input_layernorm(hs_flat) + else: + res_flat = residual.reshape(-1, residual.shape[-1]) + hs_normed_flat, res_flat = layer.input_layernorm(hs_flat, res_flat) + residual = res_flat.reshape(B, block_size, -1) + + # QKV projection on normed query tokens (2D) + qkv_query = attn_mod.qkv_proj(hs_normed_flat) # [B*blk, qkv_size] + + if use_fused_qk_norm_rope: + # One kernel does q_norm + k_norm + RoPE in-place on qkv. + # Only safe when the drafter's rope params don't use YaRN / + # long-rope / partial-rotary — otherwise fall back to the + # shared-cache path below. + attn_mod.apply_qk_norm_rope(qkv_query, query_positions_flat_i32) + q_all_2d = qkv_query[:, :q_size] + k_noise_2d = qkv_query[:, q_size : q_size + kv_size] + v_noise_2d = qkv_query[:, q_size + kv_size :] + Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, head_dim) + k_noise_bshd = k_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) + v_noise_bshd = v_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) + elif use_fused_rope: + # Per-head RMSNorm on q/k (returns new contiguous tensors), + # then flashinfer in-place RoPE sharing the same cos/sin cache + # as precompute_context_kv. + q = attn_mod.q_norm(qkv_query[:, :q_size].reshape(-1, head_dim)).view(-1, q_size) + k = attn_mod.k_norm( + qkv_query[:, q_size : q_size + kv_size].reshape(-1, head_dim) + ).view(-1, kv_size) + _flashinfer_rope( + query_positions_flat_i32, + q, + k, + head_dim, + self._get_cos_sin_cache(), + self._is_neox, + ) + Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) + k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, head_dim) + v_noise_bshd = qkv_query[:, q_size + kv_size :].reshape( + B, block_size, num_kv_heads_per_rank, head_dim + ) + else: + qkv_query_3d = qkv_query.reshape(B, block_size, -1) + q_all = qkv_query_3d[..., :q_size] + k_noise_all = qkv_query_3d[..., q_size : q_size + kv_size] + v_noise_all = qkv_query_3d[..., q_size + kv_size :] + if has_qk_norm: + q_for_rope = attn_mod.q_norm(q_all.reshape(-1, head_dim)).reshape( + B, block_size, q_size + ) + k_noise_for_rope = attn_mod.k_norm(k_noise_all.reshape(-1, head_dim)).reshape( + B, block_size, kv_size + ) + else: + q_for_rope = q_all + k_noise_for_rope = k_noise_all + Q = _rope( + q_for_rope.reshape(B, block_size, num_heads_per_rank, head_dim).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox, + ) + k_noise_rope = _rope( + k_noise_for_rope.reshape( + B, block_size, num_kv_heads_per_rank, head_dim + ).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox, + ) + Q_bshd = Q.transpose(1, 2) + k_noise_bshd = k_noise_rope.transpose(1, 2) + v_noise_bshd = v_noise_all.reshape(B, block_size, num_kv_heads_per_rank, head_dim) + + # Per-layer view into the pooled ctx cache. + causal, window_size = self._get_attention_mask_args(layer_idx) + swa_window = ( + self._layer_windows[layer_idx] if layer_idx < len(self._layer_windows) else (-1, -1) + ) + if swa_window != (-1, -1): + window_size = swa_window + if self.dflash_attention_backend == "TRTLLM": + layer_cache = ctx_kv_cache[layer_idx] + trtllm_gen_ops.append_paged_kv_cache( + append_key=k_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + append_value=v_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + batch_indices=append_batch_indices, + positions=append_positions, + paged_kv_cache=layer_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout="HND", + ) + out = torch.empty_like(Q_bshd) + q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) + out_flat = out.reshape(-1, num_heads_per_rank, head_dim) + window_left = window_size[0] + if causal: + trtllm_gen_ops.batch_decode_with_kv_cache( + query=q_flat, + kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), + workspace_buffer=self._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_after, + max_seq_len=pages_per_slot * page_size, + bmm1_scale=head_dim**-0.5, + bmm2_scale=1.0, + window_left=window_left, + out=out_flat, + sinks=None, + enable_pdl=False, + kv_layout="HND", + backend="trtllm-gen", + q_len_per_req=block_size, + max_q_len=None, + cum_seq_lens_q=None, + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + bmm1_scale_log2=None, + multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, + ) + else: + cum_seq_lens_q = torch.arange( + 0, + (B + 1) * block_size, + block_size, + dtype=torch.int32, + device=hidden_states.device, + ) + cum_seq_lens_kv = torch.cat( + ( + torch.zeros(1, dtype=torch.int32, device=hidden_states.device), + seq_lens_after.cumsum(0, dtype=torch.int32), + ) + ) + trtllm_gen_ops.batch_context_with_kv_cache( + query=q_flat, + kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), + workspace_buffer=self._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_after, + max_q_len=block_size, + max_kv_len=pages_per_slot * page_size, + bmm1_scale=head_dim**-0.5, + bmm2_scale=1.0, + batch_size=B, + cum_seq_lens_q=cum_seq_lens_q, + cum_seq_lens_kv=cum_seq_lens_kv, + window_left=window_left, + out=out_flat, + sinks=None, + enable_pdl=False, + kv_layout="HND", + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + causal=False, + multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, + ) + else: # VANILLA, validated before entering the layer loop. + layer_k_cache = ctx_k_cache[:, layer_idx] + layer_v_cache = ctx_v_cache[:, layer_idx] + + # Pack gqa_group_size query heads sharing a KV head into the + # row dimension: [B, blk, h_q, d] -> [B, group*blk, h_kv, d]. + # Each CTA owns a whole query-head group and streams KV head's context once + # instead of gqa_group_size CTAs each re-reading it. + # Exact only while every row of the block attends to the same + # key set, i.e. non-causal, unwindowed layers. Causal or + # windowed layers mask by row, so they stay unpacked. + pack_gqa = gqa_group_size > 1 and not causal and window_size == (-1, -1) + if pack_gqa: + q_grouped = Q_bshd.reshape( + B, block_size, num_kv_heads_per_rank, gqa_group_size, head_dim + ) + q_packed = q_grouped.permute(0, 3, 1, 2, 4) + q_in = q_packed.reshape( + B, gqa_group_size * block_size, num_kv_heads_per_rank, head_dim + ) + else: + q_in = Q_bshd + out = flash_attention( + q=q_in, + k_cache=layer_k_cache, + v_cache=layer_v_cache, + k=k_noise_bshd, + v=v_noise_bshd, + cache_seqlens=cache_seqlens_i32, + cache_batch_idx=cache_batch_idx_i32, + causal=causal, + window_size=window_size, + ) + if pack_gqa: + # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d]. + out = out.view( + B, gqa_group_size, block_size, num_kv_heads_per_rank, head_dim + ).permute(0, 2, 3, 1, 4) + + attn_output = out.reshape(B * block_size, q_size) + + # Per-drafter post-attention gate (no-op for generic DFlash; Laguna + # applies per-head softplus g_proj gating). gate input is the + # input_layernorm output (the attention input). + attn_output = self._post_attention_gate( + attn_output, hs_normed_flat, attn_mod, num_heads_per_rank, head_dim + ) + + # o_proj (flat 2D, handles all-reduce internally) + hidden_out = attn_mod.o_proj(attn_output) + + # Post-attention layernorm + MLP (flat 2D) + res_flat = residual.reshape(-1, residual.shape[-1]) + hidden_out, res_flat = layer.post_attention_layernorm(hidden_out, res_flat) + hidden_out = layer.mlp(hidden_out) + + hidden_states = hidden_out.reshape(B, block_size, -1) + residual = res_flat.reshape(B, block_size, -1) + + # Final norm + hidden_states_out, _ = self.model.norm( + hidden_states.reshape(-1, hidden_states.shape[-1]), + residual.reshape(-1, residual.shape[-1]), + ) + return hidden_states_out + + def forward( + self, + attn_metadata, + input_ids: torch.LongTensor = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + return_context_logits: bool = False, + spec_metadata=None, + hidden_states: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the draft model and return (hidden_states, hidden_states) for the + speculative-decoding contract.""" + hidden_states_out = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + + return hidden_states_out, hidden_states_out + + +class DFlashLagunaForCausalLM(DFlashForCausalLM): + """Laguna DFlash drafter. + + The generic block decode lives in DFlashForCausalLM; this subclass supplies + the Laguna draft-layer specifics: per-head g_proj softplus gating and the + per-aux fc_norm applied to captured target features before fc. + """ + + @staticmethod + def _normalize_config(config: PretrainedConfig) -> None: + """Fill TRT-LLM Laguna defaults missing from dense DFlash drafts.""" + if getattr(config, "num_experts", None) is None: + config.num_experts = 0 + if getattr(config, "mlp_layer_types", None) is None: + config.mlp_layer_types = ["dense"] * config.num_hidden_layers + if getattr(config, "block_size", None) is None: + dflash_config = getattr(config, "dflash_config", {}) + if isinstance(dflash_config, dict): + config.block_size = dflash_config.get("block_size", None) + + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + """Pin the Laguna draft-layer class and enable Laguna-specific behaviors + (context input_layernorm, causal sliding blocks); reject non-per-head + gating.""" + # The checkpoint labels itself with the vLLM name (model_type "llama"); + # remap to the Laguna architecture so TRT-LLM builds the Laguna layers. + draft_config.pretrained_config.architectures = ["LagunaForCausalLM"] + self._normalize_config(draft_config.pretrained_config) + super().__init__( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + self._context_input_layernorm = True + self._sliding_layers_causal = True + gating = getattr(self.config, "gating", True) + if gating not in (True, "per-head"): + raise NotImplementedError( + f"Laguna DFlash drafter supports per-head gating only, got gating={gating!r}" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """Build the per-aux ``fc_norm`` from the drafter's ``aux_hidden_norms.*`` + weights, then defer the remaining weights to the base loader.""" + aux_keys = sorted( + (k for k in weights if k.startswith("aux_hidden_norms.")), + key=lambda k: int(k.split(".")[1]), + ) + if not aux_keys: + raise ValueError("Laguna DFlash checkpoint is missing aux_hidden_norms.* weights") + weights = dict(weights) + eps = getattr(self.config, "rms_norm_eps", 1e-6) + norms = [] + for k in aux_keys: + w = weights.pop(k) + norm = nn.RMSNorm( + w.shape[0], eps=eps, device="cuda", elementwise_affine=True, dtype=w.dtype + ) + norm.weight.data.copy_(w) + norms.append(norm) + self.fc_norm = nn.ModuleList(norms) + super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + def project_target_hidden(self, hidden_states): + """Project captured target features to the draft width: apply the per-aux + ``fc_norm`` to each hidden chunk, then ``fc`` + ``hidden_norm``.""" + hidden_states = hidden_states.to(self.fc.weight.dtype) + fc_norm = getattr(self, "fc_norm", None) + if fc_norm is not None: + chunks = hidden_states.chunk(len(fc_norm), dim=-1) + hidden_states = torch.cat([norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) + return self.hidden_norm(self.fc(hidden_states)) + + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): + """Apply Laguna's per-head softplus output gate (``g_proj``) to the + attention output; a no-op when the layer has no ``g_proj``.""" + g_proj = getattr(attn_mod, "g_proj", None) + if g_proj is None: + return attn_output + gate = F.softplus(g_proj(gate_input).float()).to(attn_output.dtype) + return (attn_output.unflatten(-1, (num_heads, head_dim)) * gate.unsqueeze(-1)).flatten(-2) + + +# Published DSpark drafters spell the head switches four different ways, and a +# reader that knows only one of them degrades silently: the heads are skipped, +# their weights are dropped, and nothing raises. Resolution order matches +# ``TorchLlmArgs.validate_speculative_config`` so the model and the user-visible +# spec_config can never disagree about whether a head is on. +_DSPARK_HEAD_KEY_ALIASES = { + # RadixArk/Kimi-K3-DSpark ships ``enable_confidence_head`` top-level. + "use_confidence_head": ("use_confidence_head", "enable_confidence_head"), +} + + +def resolve_dspark_head_config(pretrained_config, key): + """Resolve one DSpark head switch across every spelling in the wild. + + Looks in ``dspark_config``, then ``dflash_config``, then the top level as + ``dspark_`` and as ````, for each accepted alias of ``key``. + Returns ``None`` when the drafter declares the switch nowhere. + """ + dspark_cfg = getattr(pretrained_config, "dspark_config", None) or {} + dflash_cfg = getattr(pretrained_config, "dflash_config", None) or {} + for name in _DSPARK_HEAD_KEY_ALIASES.get(key, (key,)): + for value in ( + dspark_cfg.get(name), + dflash_cfg.get(name), + getattr(pretrained_config, f"dspark_{name}", None), + getattr(pretrained_config, name, None), + ): + if value is not None: + return value + return None + + +def declares_dspark_heads(pretrained_config) -> bool: + """True when a drafter config asks for the DSpark head set. + + Resolved here rather than in the DSpark module: this module must not import + the DSpark drafters, or the ``modeling_dspark -> modeling_dflash`` + inheritance edge would become a cycle. + """ + return bool( + str(resolve_dspark_head_config(pretrained_config, "projector_type") or "").lower() + == "dspark" + or resolve_dspark_head_config(pretrained_config, "shift_label") + or resolve_dspark_head_config(pretrained_config, "use_confidence_head") + or int(resolve_dspark_head_config(pretrained_config, "markov_rank") or 0) > 0 + ) + + +@register_draft_model(SpeculativeDecodingMode.DFLASH) +def _build_dflash_draft(model_config, draft_config, lm_head, model): + """Build the DFlash drafter. + + Selects the Laguna variant by detecting its architecture in the draft + checkpoint's own config. A drafter that declares the DSpark heads is + rejected rather than silently served without them: DFlash no longer + implements the Markov / confidence / shift_label semantics, so building one + here would degrade the drafter with no error and show up only as a lower + acceptance rate. + """ + if declares_dspark_heads(draft_config.pretrained_config): + raise ValueError( + "This drafter checkpoint declares the DSpark head set " + "(dflash_config with any of projector_type='dspark', shift_label, " + "use_confidence_head, markov_rank > 0), which decoding_type " + "'DFlash' does not implement. Set speculative_config.decoding_type " + "to 'DSpark' to run it." + ) + draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] + dflash_attention_backend = model_config.spec_config.attention_backend + if any("Laguna" in arch for arch in draft_arches): + return DFlashLagunaForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + return DFlashForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 2d556af5e71e..c85dd570db93 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -15,29 +15,73 @@ # # DSpark backbone ported from the DeepSeek-V4-Pro-DSpark reference # (`inference/model.py`: DSparkBlock / Transformer.forward_spec). -"""DeepSeek-V4-Pro DSpark speculative-decoding draft backbone. - -The DSpark draft is ``n_mtp_layers`` (3 for V4-Pro) **full DeepSeek-V4 blocks** -stored under the ``mtp.*`` checkpoint namespace — it reuses the V4 decoder block -(MLA attention + MoE + manifold Hyper-Connections) and adds: - - - **stage 0**: ``main_proj`` (Linear, fp8) + ``main_norm`` (RMSNorm) — projects - the concatenation of captured target-layer hidden states ([58,59,60]) into the - draft's cross-attention context (``main_x``); replaces vanilla-MTP's - enorm/hnorm + e_proj/h_proj single-hidden mixing. - - **last stage**: ``norm`` + ``markov_head`` + ``confidence_head`` + - flat ``hc_head`` — the block-draft output head (see dspark_heads/dspark_draft). - -The per-stage *backbone* forward (block attention whose K/V derive from ``main_x``, -+ MoE + mHC) is brought up and numerically validated against the real fp8 weights -separately; ``forward_embed`` (capture) and ``forward_head`` (block draft) below -are the reference-faithful, unit-validated I/O stages. +# +# The captured-context attention primitives are ported from DeepSeek's DeepSpec +# reference ``inference/kernel.py`` (``sparse_attn``) and ``inference/model.py`` +# (``get_dspark_topk_idxs``). The reference computes these with a TileLang +# kernel; this is a functional-first pure-PyTorch port with the same math +# (index-gather + online softmax + a learnable attention sink that contributes +# only to the softmax denominator). +# +# The draft I/O stages are ported from the same reference +# (`inference/model.py`: DSparkBlock.forward_embed / forward_head). +"""DSpark speculative-decoding drafters. + +``DSpark`` names the speculative-decoding *algorithm*: a parallel block draft +over captured target hidden states, refined by a low-rank Markov head and +scheduled by a confidence head. Every ``decoding_type: DSpark`` drafter is built +here, in one of two flavours that differ only in how the draft is delivered: + +* **embedded** — the draft weights ship inside the DeepSeek-V4-Pro *target* + checkpoint under the ``mtp.*`` namespace and reuse the V4 decoder block, so the + draft inherits the target's EPLB layer namespace and fp8/NVFP4 quantization. + Most of this module is this flavour. +* **standalone** — the drafter ships as its own checkpoint and shares nothing + with the target but the vocabulary and the captured hidden states. Its block + decode is DFlash's, so :class:`GQADSparkForCausalLM` subclasses + ``DFlashForCausalLM`` and adds the Markov head, the confidence head and the + shift_label convention. See "Standalone DSpark drafters" near the bottom. + +The dependency runs one way, ``modeling_dspark -> modeling_dflash``: DFlash is +DSpark minus those three heads, and :mod:`modeling_dflash` must not name a DSpark +class or the edge would become a cycle. + +Four parts live here: + +1. **Draft backbone** — ``n_mtp_layers`` (3 for V4-Pro) full DeepSeek-V4 blocks + (MLA attention + MoE + manifold Hyper-Connections), plus: + + - **stage 0**: ``main_proj`` (Linear, fp8) + ``main_norm`` (RMSNorm) — + projects the concatenation of captured target-layer hidden states + ([58,59,60]) into the draft's cross-attention context (``main_x``); + replaces vanilla-MTP's enorm/hnorm + e_proj/h_proj single-hidden mixing. + - **last stage**: ``norm`` + ``markov_head`` + ``confidence_head`` + flat + ``hc_head`` — the block-draft output head. The heads themselves are shared + with the standalone path and live in :mod:`modeling_speculative`. + +2. **Captured-context attention primitives** — the dense sliding-window MLA the + draft block attends with (``get_dspark_topk_idxs`` / ``dspark_sparse_attn`` + and the rotary helpers). Hardware-agnostic and unit-testable in isolation. + +3. **Block draft I/O** — ``build_draft_input_ids`` (the + ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov + refinement + static confidence truncation). + +4. **Standalone DSpark drafters** — :class:`GQADSparkForCausalLM`, plus the + ``_build_dspark_draft`` dispatch that picks between the two flavours on + deployment form alone. + +The per-stage *backbone* forward (block attention whose K/V derive from +``main_x``, + MoE + mHC) is brought up and numerically validated against the real +fp8 weights separately; ``forward_embed`` (capture) and ``forward_head`` (block +draft) are the reference-faithful, unit-validated I/O stages. """ import copy import json import os import re +from functools import lru_cache from typing import Dict, List, Optional import torch @@ -47,21 +91,14 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo +from ..._utils import is_sm_100f +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..distributed import AllReduceParams from ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead from ..modules.rms_norm import RMSNorm +from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType -from .dspark.attention import ( - _rmsnorm, - _rope_last_dims, - _rope_last_dims_batched, - dspark_attention_forward, - dspark_attention_forward_batched, - precompute_dspark_freqs_cis, -) -from .dspark.draft import build_draft_input_ids, dspark_propose -from .dspark.heads import DSparkConfidenceHead, build_markov_head from .modeling_deepseekv4 import ( DeepseekV4DecoderLayer, DeepseekV4WeightLoader, @@ -71,6 +108,636 @@ _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) +from .modeling_dflash import DFlashForCausalLM, resolve_dspark_head_config +from .modeling_speculative import ( + DSparkConfidenceHead, + build_markov_head, + confident_prefix_length, + dspark_markov_chain_logits, +) +from .modeling_utils import register_draft_model + +if IS_CUTLASS_DSL_AVAILABLE: + from ..custom_ops.dspark_attention_custom_op import ( + cute_dsl_dspark_attention, + is_fused_dspark_attention_supported, + ) + from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope, + is_fused_dspark_rmsnorm_rope_supported, + ) + +# ---------------------------------------------------------------------------- +# Captured-context attention primitives. +# ---------------------------------------------------------------------------- + + +def precompute_dspark_freqs_cis( + rope_head_dim: int, + seqlen: int, + rope_theta: float = 10000.0, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. + + The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the + base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with + ``original_seq_len == 0``). + + Returns: + complex64 tensor ``[seqlen, rope_head_dim // 2]``. + """ + freqs = 1.0 / ( + rope_theta + ** (torch.arange(0, rope_head_dim, 2, dtype=torch.float32, device=device) / rope_head_dim) + ) + t = torch.arange(seqlen, dtype=torch.float32, device=device) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def apply_dspark_rotary( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. + + Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the + last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the + sequence axis, and conjugates for the inverse (de-rotation applied to the + attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or + ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. + """ + orig_dtype = x.dtype + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) + else: + fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) + out = torch.view_as_real(xc * fc).flatten(-2) + return out.to(orig_dtype) + + +def apply_dspark_rotary_batched( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Per-row (batched) variant of :func:`apply_dspark_rotary`. + + Identical math, but ``freqs_cis`` carries a leading batch axis so each row of + ``x`` is rotated by its own per-request phases (the generation draft runs each + request at a different absolute ``start_pos``). ``x`` is the rope-dim slice + only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape + ``[G, s, rd // 2]``. + """ + orig_dtype = x.dtype + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + g, s, half = freqs_cis.shape + if xc.ndim == 3: + fc = freqs_cis.view(g, s, half) + else: + fc = freqs_cis.view(g, s, 1, half) + out = torch.view_as_real(xc * fc).flatten(-2) + return out.to(orig_dtype) + + +@lru_cache(maxsize=64) +def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: + # [min(window, start_pos+1)] context positions in the rolling KV window, + # followed by [block_size] positions for the current block's own K/V (which + # the caller appends to the window at offset ``window_size``). + ctx = torch.arange(min(window_size, start_pos + 1)) + blk = window_size + torch.arange(block_size) + return torch.cat([ctx, blk]).int() + + +def get_dspark_topk_idxs( + window_size: int, + bsz: int, + block_size: int, + start_pos: int, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Per-query attended-position indices for the DSpark draft block. + + Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` + query positions attends to the same set — the ``min(window_size, start_pos+1)`` + most-recent context positions in the rolling KV window, then the + ``block_size`` positions of the current block (stored at offset + ``window_size`` in the concatenated KV). Note this is *non-causal* within the + block (every position sees every block position), matching the reference. + + Args: + window_size: sliding-window length of the captured-context KV cache. + bsz: batch size. + block_size: number of draft positions per request. + start_pos: absolute decode position (must be > 0); bounds the context. + device: device for the returned index tensor. + + Returns: + int32 tensor ``[bsz, block_size, topk]`` with + ``topk = min(window_size, start_pos+1) + block_size``. + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) + return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +def get_dspark_topk_idxs_batched( + window_size: int, + block_size: int, + start_pos: torch.Tensor, + valid_len: torch.Tensor | None = None, +) -> torch.Tensor: + """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. + + Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width + ``min(window_size, start_pos+1) + block_size`` depends on the host int + ``start_pos``), this always returns the **fixed** width ``window_size + + block_size`` and masks the unfilled context slots with ``-1``. The masked + slots are excluded by :func:`dspark_sparse_attn` exactly as if they were + absent while the shape remains CUDA-graph safe. + + Every query attends to the actually written circular-window suffix, followed + by the current-block positions. Without ``valid_len`` this preserves the + legacy ``start_pos``-only behavior. + + Args: + window_size: sliding-window length of the captured-context KV cache. + block_size: number of draft positions per request. + start_pos: ``[G]`` int tensor of per-request absolute decode positions. + valid_len: optional ``[G]`` count of actually written rolling-window + entries. When omitted, preserve the legacy ``start_pos`` mask. + + Returns: + int32 tensor ``[G, block_size, window_size + block_size]``. + """ + device = start_pos.device + g = start_pos.shape[0] + ctx_cols = torch.arange(window_size, device=device) # [win] + if valid_len is None: + valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] + else: + # The valid entries are the contiguous logical suffix ending at + # start_pos, but their physical slots wrap modulo window_size. + valid_len = valid_len.clamp(min=0, max=window_size) + age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) + valid = age < valid_len.unsqueeze(1) + ctx_idx = torch.where( + valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) + ) + blk_idx = window_size + torch.arange(block_size, device=device) # [block] + blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] + row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] + return row.unsqueeze(1).expand(g, block_size, -1).contiguous() + + +def dspark_sparse_attn( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Index-gathered multi-query attention with an attention sink. + + Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For + each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by + ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled + dot-product softmax over them, and adds a per-head learnable *sink* logit that + participates only in the softmax denominator (i.e. an "attend-to-nothing" + option with a zero value vector). KV is shared across query heads (MQA). + + Args: + q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). + kv: ``[b, n, d]`` keys/values (shared across heads). + attn_sink: ``[h]`` per-head sink logits (fp32). + topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). + softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). + + Returns: + ``[b, m, h, d]`` attention output, in ``q.dtype``. + """ + b, m, h, d = q.shape + idx = topk_idxs.long() # [b, m, topk] + valid = idx >= 0 + safe = idx.clamp(min=0) + + # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below + # zeros their softmax probs, so the einsum nullifies them. + kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) + gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() + + # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. + scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + + # Online-softmax max is taken over gathered positions only (the sink is added + # to the denominator afterwards), matching the kernel's reduce order. + smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] + smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) + probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 + sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) + denom = probs.sum(dim=-1) + sink # [b, m, h] + out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) + return out.to(q.dtype) + + +def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" + dtype = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + return (weight.float() * xf).to(dtype) + + +def _rope_last_dims( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + +def _rope_last_dims_batched( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Per-row variant of :func:`_rope_last_dims` (``freqs_cis`` has a batch axis).""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + +def _rmsnorm_rope_batched( + t: torch.Tensor, + weight: torch.Tensor, + eps: float, + rope_head_dim: int, + freqs_cis: torch.Tensor, + *, + num_heads: int = 1, + apply_weight: bool = True, + apply_rmsnorm: bool = True, + inverse_rope: bool = False, +) -> torch.Tensor: + """Fuse DSpark RMSNorm and last-dimension RoPE when supported.""" + if IS_CUTLASS_DSL_AVAILABLE and is_sm_100f(): + freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) + if is_fused_dspark_rmsnorm_rope_supported(t, weight, freqs_real, num_heads, rope_head_dim): + return cute_dsl_dspark_rmsnorm_rope( + t, + weight, + freqs_real, + num_heads, + rope_head_dim, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + ) + + if apply_rmsnorm: + if apply_weight: + t = _rmsnorm(t, weight, eps) + else: + t = t * torch.rsqrt(t.square().mean(-1, keepdim=True) + eps) + elif apply_weight: + t = (t.float() * weight.float()).to(t.dtype) + if rope_head_dim > 0: + t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) + return t + + +def dspark_attention_forward( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: int, + kv_cache: torch.Tensor, + *, + wq_a: torch.Tensor, + q_norm_w: torch.Tensor, + wq_b: torch.Tensor, + wkv: torch.Tensor, + kv_norm_w: torch.Tensor, + wo_a: torch.Tensor, + wo_b: torch.Tensor, + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + persist: bool = False, +) -> torch.Tensor: + """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). + + Functional port of DeepSpec ``DSparkAttention.forward`` for the dense + (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) + with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys + gathered from a rolling captured-context window (``kv_cache``, into which the + projected ``main_x`` context is written at ``start_pos % window_size``) plus the + block's own positions, attention-sink softmax, inverse-RoPE on the output, and a + grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). + + Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / + dequantized projection weights); ``wo_a`` is the raw grouped weight matrix + ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is + ``[b, window_size, head_dim]`` and is updated functionally (cloned). + + Returns: + ``[b, block_size, dim]`` attention output (residual stream contribution). + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + b, block, _ = x.shape + rd = rope_head_dim + main_freqs = freqs_cis[start_pos : start_pos + 1] + blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] + + # Captured-context K/V from main_x (MQA, shared across heads). + main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] + main_kv = _rope_last_dims(main_kv, rd, main_freqs) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) + q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] + # Per-head RMS in the query dtype (matches the reference inline normalization, + # which is NOT the fp32 RMSNorm path). + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + q = _rope_last_dims(q, rd, blk_freqs) + + # Block K/V. + kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] + kv = _rope_last_dims(kv, rd, blk_freqs) + + # Write the context K/V into the rolling window, then attend over + # [window context | block] with the sink. ``persist=True`` writes through + # to the caller's buffer (cross-step decode, worker-owned window); the + # default clones so single-shot callers (golden / unit tests) stay pure. + cache = kv_cache if persist else kv_cache.clone() + cache[:, start_pos % window_size] = main_kv.squeeze(1) + kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] + topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) + o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] + o = _rope_last_dims(o, rd, blk_freqs, inverse=True) + + # Grouped low-rank O projection. + o = o.reshape(b, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) + return F.linear(o.flatten(2), wo_b) + + +def dspark_attention_forward_batched( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + valid_len: torch.Tensor | None = None, + *, + wq_a: torch.Tensor, + q_norm_w: torch.Tensor, + wq_b: torch.Tensor, + wkv: torch.Tensor, + kv_norm_w: torch.Tensor, + wo_a: torch.Tensor, + wo_b: torch.Tensor, + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + persist: bool = False, +) -> torch.Tensor: + """Batched, CUDA-graph-safe captured-context DSpark draft attention. + + Numerically identical, per request, to :func:`dspark_attention_forward`, but + free of host syncs and data-dependent shapes so it can be captured into a CUDA + graph (the one-engine drafter runs inside the target's graph). The differences + from the scalar path are purely mechanical: + + * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen + request) instead of a python int; RoPE phases are *gathered* per request from + the fixed ``freqs_cis`` table rather than sliced. + * the rolling-window context K/V is written/read through the ``slots`` index + into a shared ``kv_cache`` (``persist=True`` writes through to the caller's + worker-owned buffer; otherwise a clone is used), instead of mutating a + per-request cache in place. + * the attended-position list has the fixed width ``window_size + block_size`` + with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). + + Args: + x: ``[G, block, dim]`` block layer input (per gen request). + main_x: ``[G, 1, hidden]`` projected captured context. + start_pos: ``[G]`` int tensor of absolute decode positions (> 0). + kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows + (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). + slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. + valid_len: optional ``[G]`` count of actually written context entries; + masks holes left when absolute positions are bootstrapped without + receiving the corresponding DSpark rolling-window state. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; + must satisfy ``maxlen > start_pos.max() + block_size``. + + Returns: + ``[G, block, dim]`` attention output (residual stream contribution). + """ + g, block, _ = x.shape + if kv_cache.shape[1] != window_size: + raise ValueError( + f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" + ) + rd = rope_head_dim + # Per-request RoPE phases gathered from the fixed table (no host-int slicing). + main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] + blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] + blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] + + # Captured-context K/V from main_x (MQA, shared across heads). + main_kv = _rmsnorm_rope_batched(F.linear(main_x, wkv), kv_norm_w, eps, rd, main_freqs) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) + q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] + q = _rmsnorm_rope_batched( + q, + kv_norm_w, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + ) + + # Block K/V. + kv = _rmsnorm_rope_batched(F.linear(x, wkv), kv_norm_w, eps, rd, blk_freqs) + + # Write the context K/V into the rolling window at slot start_pos%window_size, + # then attend over [window context | block]. ``persist=True`` writes through to + # the worker-owned buffer (cross-step decode); otherwise clone so single-shot + # callers stay pure. + write_target = kv_cache if persist else kv_cache.clone() + main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) + if ( + valid_len is None + and IS_CUTLASS_DSL_AVAILABLE + and is_fused_dspark_attention_supported( + q, main_kv_flat, kv, write_target, slots, start_pos, attn_sink + ) + ): + # One custom op performs the rolling-cache write/read, validity handling, + # QK, attention-sink online softmax, and PV. In particular it creates no + # topk index, gathered KV, score, or probability tensors. + o = cute_dsl_dspark_attention( + q, + main_kv_flat, + kv, + write_target, + slots, + start_pos, + attn_sink, + softmax_scale, + ) + else: + slot_pos = start_pos % window_size # [G] + write_target[slots, slot_pos] = main_kv_flat + cache_rows = write_target[slots] # [G, window, head_dim] + kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] + topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) + o = dspark_sparse_attn( + q, kv_full, attn_sink, topk, softmax_scale + ) # [G, block, h, head_dim] + o = _rmsnorm_rope_batched( + o, + kv_norm_w, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + apply_rmsnorm=False, + inverse_rope=True, + ) + + # Grouped low-rank O projection. + o = o.reshape(g, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) + return F.linear(o.flatten(2), wo_b) + + +# ---------------------------------------------------------------------------- +# Block draft I/O. +# ---------------------------------------------------------------------------- + + +def build_draft_input_ids( + bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int +) -> torch.Tensor: + """``[batch] -> [batch, block_size]`` = ``[bonus, noise, noise, ...]``. + + The first position is the verified bonus token (the target's last accepted + token); the rest are the DSpark noise/mask token (id 128799 for V4-Pro). + """ + batch = bonus_token_ids.shape[0] + out = bonus_token_ids.new_full((batch, block_size), int(noise_token_id)) + out[:, 0] = bonus_token_ids + return out + + +def dspark_propose( + base_logits: torch.Tensor, + *, + bonus_token_ids: torch.Tensor, + block_hidden: torch.Tensor, + markov_head: Optional[nn.Module], + confidence_head: Optional[nn.Module], + block_size: int, + temperature: float = 0.0, + confidence_threshold: float = 0.0, + return_logits: bool = False, +) -> tuple: + """Produce DSpark draft tokens for one block (functional-first, static length). + + Args: + base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. + bonus_token_ids: ``[batch]`` the token preceding the first draft position. + block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the + confidence head, and the RNN-head variant). + markov_head / confidence_head: the validated DSpark heads (may be None). + Returns: + draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers + keep the tensor fixed-width for CUDA-graph safety). + num_proposed: ``[batch]`` int32 — how many leading tokens survive the + static confidence-threshold truncation (== block_size when no head / + threshold<=0). + """ + batch = base_logits.shape[0] + # ``draft_logits`` are the per-position distributions the draft token is drawn + # from (markov-corrected when a head is present, else the raw base logits). + # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance + # (1-TV) measurement; the normal path ignores them. + draft_logits = base_logits + if markov_head is not None: + draft_tokens, corrected = markov_head.sample_block_tokens( + base_logits, + first_prev_token_ids=bonus_token_ids, + hidden_states=block_hidden, + temperature=temperature, + ) + draft_logits = corrected + else: + from .modeling_speculative import greedy_or_sample + + draft_tokens = greedy_or_sample(base_logits, temperature) + + # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. + # The worker always calls with confidence_threshold=0.0, so the block below is + # inert and num_proposed stays == block_size (the full block is proposed). The + # returned num_proposed is intentionally not yet consumed by the speculative + # scheduler/verifier; wiring it through is a follow-up (see PR description). + num_proposed = torch.full( + (batch,), int(block_size), dtype=torch.int32, device=base_logits.device + ) + if confidence_head is not None and confidence_threshold > 0.0: + # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] + prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) + prev_emb = ( + markov_head.get_prev_embeddings(prev_ids) + if (markov_head is not None and getattr(confidence_head, "with_markov", False)) + else None + ) + conf_logits = ( + confidence_head(block_hidden, prev_embeddings=prev_emb) + if prev_emb is not None + else confidence_head(block_hidden) + ) + # Per-request prefix truncation (batch handled row-wise to stay simple; + # functional-first scope typically runs batch=1 for the draft). + for b in range(batch): + num_proposed[b] = confident_prefix_length( + conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold + ) + if return_logits: + return draft_tokens, num_proposed, draft_logits + return draft_tokens, num_proposed + + +# ---------------------------------------------------------------------------- +# Draft backbone (``mtp.*`` stages of DeepSeek-V4 blocks). +# ---------------------------------------------------------------------------- # Matches the draft namespace ``mtp..`` in the V4-Pro-DSpark # checkpoint. Each draft stage is a full DeepSeek-V4 block stored under this @@ -188,7 +855,7 @@ def count_dspark_stages(ckpt_dir: str) -> Optional[int]: def _rename_dspark_stage_subkey(rest: str, routed_scale: str) -> str: - """Map a per-stage checkpoint subkey to the ``DSparkBlock`` param subkey.""" + """Map a per-stage checkpoint subkey to the ``DSv4DSparkBlock`` param subkey.""" if rest == "attn_norm.weight": return "input_layernorm.weight" if rest == "ffn_norm.weight": @@ -208,7 +875,7 @@ def _rename_dspark_stage_subkey(rest: str, routed_scale: str) -> str: if rest.startswith("ffn."): return f"mlp.{_rename_deepseek_v4_ffn_subkey(rest[len('ffn.') :], routed_scale)}" # main_proj.weight, main_norm.weight, norm.weight, markov_head.*, - # confidence_head.* map 1:1 onto the DSparkBlock submodules. + # confidence_head.* map 1:1 onto the DSv4DSparkBlock submodules. return rest @@ -245,7 +912,7 @@ def remap_dspark_draft_keys(weights: Dict, num_stages: int) -> Dict: # dequantize ``wo_a`` (cos 1.0 vs ``wo_a_fp8 * scale``). Always dequantize now. -class DSparkBlock(DeepseekV4DecoderLayer): +class DSv4DSparkBlock(DeepseekV4DecoderLayer): """One DSpark draft stage = a DeepSeek-V4 decoder block + DSpark extras. ``stage_id`` in ``[0, num_stages)``; only stage 0 owns the capture projection @@ -337,7 +1004,7 @@ def has_heads(self) -> bool: return self.stage_id == self.num_stages - 1 -class DSparkDraftModel(nn.Module): +class DSv4DSparkDraftModel(nn.Module): """The ``n_mtp_layers``-stage DSpark draft stacked on a DeepSeek-V4 target. Shares ``embed_tokens`` / ``lm_head`` with the target model. ``forward_embed`` @@ -414,7 +1081,7 @@ def __init__( draft_model_config = self._derive_draft_model_config(model_config, base, self.num_stages) self.mtp_layers = nn.ModuleList( [ - DSparkBlock( + DSv4DSparkBlock( draft_model_config, base + s, aux_stream_dict, @@ -552,7 +1219,7 @@ def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: """Populate ``_dspark_attn`` from an already-loaded in-memory ``weights`` dict (no extra disk I/O); used on the one-engine load path - (``DSparkForCausalLM.load_weights``). Delegates to :meth:`_cache_attn_weights`. + (``DSv4DSparkForCausalLM.load_weights``). Delegates to :meth:`_cache_attn_weights`. """ self._cache_attn_weights(weights) @@ -816,7 +1483,7 @@ def forward_embed( def _forward_stage( self, - stage: "DSparkBlock", + stage: "DSv4DSparkBlock", h: torch.Tensor, main_x: torch.Tensor, start_pos, @@ -1151,12 +1818,12 @@ def forward_head( ) -class DSparkForCausalLM(nn.Module): +class DSv4DSparkForCausalLM(nn.Module): """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). - Wraps :class:`DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) + Wraps :class:`DSv4DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) for the single-engine external-drafter flow: created by ``get_draft_model``, - appended to the target's epilogue, and driven by ``DSparkWorker``. + appended to the target's epilogue, and driven by ``DSv4DSparkWorker``. ``embed_tokens`` / ``lm_head`` are shared with the target model (:meth:`load_weights_from_target_model`). The draft weights live in the SAME @@ -1168,7 +1835,7 @@ class DSparkForCausalLM(nn.Module): def __init__(self, draft_config, aux_stream_dict=None, num_stages=None, block_size=None): super().__init__() - self.dspark_model = DSparkDraftModel( + self.dspark_model = DSv4DSparkDraftModel( draft_config, aux_stream_dict, num_stages=num_stages, @@ -1241,10 +1908,266 @@ def load_weights_from_target_model(self, target_model): self.dspark_model.lm_head = target_model.lm_head +# ---------------------------------------------------------------------------- +# Standalone DSpark drafters. +# +# The other DSpark flavour: the drafter ships as its own checkpoint instead of +# living in the target's ``mtp.*`` namespace, so it shares nothing with the +# target but the vocabulary and the captured hidden states. Its block decode is +# DFlash's -- DSpark is DFlash plus a Markov logit bias, a confidence head and +# the shift_label slot convention -- so these subclass ``DFlashForCausalLM`` and +# add exactly those three. +# ---------------------------------------------------------------------------- + +# Both published drafters (RadixArk/Kimi-K3-DSpark, Inferact/Kimi-K3-DSpark) +# name the head tensors after the submodules that own them: markov_head is a +# VanillaMarkov, confidence_head an AcceptRatePredictor whose linear is ``proj``. +# The bare spellings are kept for drafters exported without that nesting. +_DSPARK_HEAD_WEIGHT_ALIASES = { + "markov_w1.weight": ("markov_head.markov_w1.weight", "markov_w1.weight"), + "markov_w2.weight": ("markov_head.markov_w2.weight", "markov_w2.weight"), + "confidence_proj.weight": ("confidence_head.proj.weight", "confidence_proj.weight"), + "confidence_proj.bias": ("confidence_head.proj.bias", "confidence_proj.bias"), +} + + +class GQADSparkForCausalLM(DFlashForCausalLM): + """DSpark drafter on a GQA-shaped backbone, from a standalone checkpoint. + + Adds the DSpark head set on top of the DFlash block decode: + + - the vanilla Markov intra-block logit bias, applied by ``DSparkWorker`` + through :meth:`apply_markov_chain_logits`; + - the ``shift_label`` output convention (the hidden state at block slot j + predicts draft token j+1, so slot 0 holds the anchor token); + - the confidence head weights. + + Confidence-scheduled verification is not implemented yet: ``confidence_proj`` + is loaded but unused, and drafting always proposes the full K tokens. + + Named for the attention shape, not for a model: the backbone is whatever + the drafter config resolves to through the model registry, and the + inherited block decode works for every GQA family the DFlash drafters + already cover (qwen3, llama, gpt_oss, ...). A per-model subclass would be + empty. The GQA precondition is inherited, not introduced here -- see + ``DFlashForCausalLM._validate_gqa_shape``. An MLA-backboned drafter needs + its own block decode and becomes a sibling, ``MLADSparkForCausalLM``, not a + subclass of this. + + Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. + """ + + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + + cfg = draft_config.pretrained_config + # Defaults on, unlike the DFlash base: the shift_label slot layout is + # part of what DSpark *is*, and both published drafters set + # block_size == max_draft_len, where the DFlash layout (slots 1..K) + # runs one slot past the block and reads the next request's anchor. + # An explicit false still selects the legacy layout. + shift_label = resolve_dspark_head_config(cfg, "shift_label") + self._dspark_shift_label = True if shift_label is None else bool(shift_label) + self._dspark_markov_rank = int(resolve_dspark_head_config(cfg, "markov_rank") or 0) + self._dspark_markov_head_type = str( + resolve_dspark_head_config(cfg, "markov_head_type") or "vanilla" + ).lower() + self._dspark_use_confidence_head = bool( + resolve_dspark_head_config(cfg, "use_confidence_head") or False + ) + # Plain None placeholders rather than nn.Parameter/buffer: the shapes + # ([vocab, rank]) are checkpoint-dependent, so nothing is pre-allocated + # and nothing is constructed in the module's default dtype. Using the + # checkpoint tensor as-is is also what keeps the head in the checkpoint's + # dtype instead of an nn.Module default. load_weights() fills them in; + # consumers treat None as "head absent". + self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) + self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) + self.confidence_proj_weight = None # loaded, unused (follow-up MR) + self.confidence_proj_bias = None + + if self._dspark_markov_rank > 0 and self._dspark_markov_head_type != "vanilla": + raise ValueError( + f"DSpark drafter declares markov_head_type=" + f"'{self._dspark_markov_head_type}'; only 'vanilla' is " + "supported (gated/rnn heads need per-step hidden features)." + ) + # The block decode only supports the non-causal DSpark convention. + # Legacy DFlash drafter configs (e.g. Laguna) also carry a causal field + # and handle it in the legacy decode path, which is why this check lives + # here rather than in the DFlash base. + if resolve_dspark_head_config(cfg, "causal"): + raise ValueError( + "DSpark drafter sets causal=true; the block decode only " + "supports the non-causal DSpark convention." + ) + if self._dspark_use_confidence_head: + logger.warning( + "DSpark drafter declares use_confidence_head; " + "confidence-scheduled verification is not implemented yet " + "(confidence_proj weights are loaded but unused, drafting " + "always proposes the full K tokens)." + ) + + @property + def has_markov_head(self) -> bool: + return self._dspark_markov_rank > 0 and self.markov_w1 is not None + + def apply_markov_chain_logits( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + argmax_fn=None, + vocab_slice: Optional[slice] = None, + ) -> torch.Tensor: + """Apply the vanilla-Markov intra-block bias to block logits. + + No-op (returns ``base_logits`` unchanged) when the checkpoint ships no + Markov head. See :func:`dspark_markov_chain` for the semantics; when + ``base_logits`` is a TP vocab shard the caller must pass this rank's + ``vocab_slice`` (to shard the markov_w2 rows identically) and an + ``argmax_fn`` returning full-vocab token ids -- ``DFlashWorker`` handles + both. + """ + if not self.has_markov_head: + return base_logits + markov_w2 = self.markov_w2 if vocab_slice is None else self.markov_w2[vocab_slice] + return dspark_markov_chain_logits( + base_logits, first_prev_tokens, self.markov_w1, markov_w2, argmax_fn=argmax_fn + ) + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Take the DSpark head weights, then hand the rest to DFlash. + + The head keys are pulled out before the backbone remap: left in, they + would pick up a ``model.`` prefix and be dropped by partial loading. + """ + dspark_weights = {} + consumed = set() + for canonical, aliases in _DSPARK_HEAD_WEIGHT_ALIASES.items(): + for name in aliases: + if name in weights: + dspark_weights[canonical] = weights[name] + consumed.add(name) + break + if consumed: + weights = {k: v for k, v in weights.items() if k not in consumed} + # The inverse of the missing-weights check below. Without it, a config + # whose head switches this build cannot resolve loads the drafter with + # the heads silently dropped -- correct output, lower acceptance. + if self._dspark_markov_rank <= 0 and "markov_w1.weight" in dspark_weights: + raise ValueError( + "DSpark drafter ships markov_w1/markov_w2 but markov_rank resolved to 0. " + "The checkpoint's head switches were not found in dspark_config, " + "dflash_config, or at the top level; loading it would drop the Markov " + "head silently." + ) + if self._dspark_markov_rank > 0: + vocab = self.config.vocab_size + rank = self._dspark_markov_rank + for k in ("markov_w1.weight", "markov_w2.weight"): + if k not in dspark_weights: + raise ValueError( + f"DSpark drafter declares markov_rank=" + f"{self._dspark_markov_rank} but the checkpoint is " + f"missing {k}." + ) + if tuple(dspark_weights[k].shape) != (vocab, rank): + raise ValueError( + f"DSpark {k} has shape " + f"{tuple(dspark_weights[k].shape)}, expected " + f"[vocab, markov_rank] = ({vocab}, {rank})." + ) + self.markov_w1 = dspark_weights["markov_w1.weight"].to("cuda") + self.markov_w2 = dspark_weights["markov_w2.weight"].to("cuda") + if "confidence_proj.weight" in dspark_weights: + self.confidence_proj_weight = dspark_weights["confidence_proj.weight"].to("cuda") + if "confidence_proj.bias" in dspark_weights: + self.confidence_proj_bias = dspark_weights["confidence_proj.bias"].to("cuda") + return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + +def draft_is_embedded_in_target(model_config) -> bool: + """True when the DSpark draft weights live inside the target checkpoint. + + That is the DeepSeek-V4-Pro layout: the draft is ``mtp.*`` inside the target + checkpoint and inherits its block definition, EPLB layer namespace and + quantization. + + The answer comes from ``DSparkDecodingConfig.draft_is_embedded_in_target`` + rather than being re-derived here, because the worker and the spec metadata + have to make the same call from ``_torch/speculative/`` -- which cannot + import this package -- and a builder that disagreed with them would hand + the worker a draft model whose attributes it does not have. + """ + return bool(model_config.spec_config.draft_is_embedded_in_target) + + +@register_draft_model(SpeculativeDecodingMode.DSPARK) +def _build_dspark_draft(model_config, draft_config, lm_head, model): + """Build the DSpark drafter for either flavour. + + Two levels of dispatch: + + 1. Are the draft weights embedded in the target checkpoint? If so this is + the DeepSeek-V4-Pro draft, whose stage count (``n_mtp_layers``) is not in + the HF config and is derived from the ``mtp.*`` namespace. + 2. Otherwise the drafter is standalone, and its own ``model_type`` selects + the backbone-specific class. + + Args: + model_config: the target engine's ``ModelConfig``. + draft_config: the drafter's own ``ModelConfig``. + lm_head: unused; DSpark shares the target's head at weight-load time. + model: the target model, whose aux streams the draft stages reuse. + + Returns: + The draft ``nn.Module`` for this drafter. + """ + if draft_is_embedded_in_target(model_config): + num_stages = count_dspark_stages(model_config.spec_config.speculative_model) + validate_dspark_eplb_layer_base(model_config, draft_config) + return DSv4DSparkForCausalLM( + draft_config, + getattr(model, "aux_stream_dict", None), + num_stages=num_stages, + block_size=model_config.spec_config.block_size, + ) + + # No per-model_type table here. ``DFlashForCausalLM.__init__`` already + # resolves the backbone from the drafter config through the model registry, + # so keying on model_type a second time would only duplicate that dispatch + # and force a new entry for every GQA family that already works. What the + # table really guarded was the block decode's GQA precondition, which is now + # checked where it belongs, in the DFlash base. An MLA-backboned drafter + # (e.g. Inferact/Kimi-K3-DSpark) fails that check with a clear message until + # ``MLADSparkForCausalLM`` lands as a sibling. + return GQADSparkForCausalLM( + draft_config, + dflash_attention_backend=model_config.spec_config.attention_backend, + ) + + __all__ = [ - "DSparkBlock", - "DSparkDraftModel", - "DSparkForCausalLM", + # Embedded (DeepSeek-V4-Pro) flavour. + "DSv4DSparkBlock", + "DSv4DSparkDraftModel", + "DSv4DSparkForCausalLM", + # Standalone flavour. + "GQADSparkForCausalLM", + "draft_is_embedded_in_target", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", + # Captured-context attention primitives. + "get_dspark_topk_idxs", + "get_dspark_topk_idxs_batched", + "dspark_sparse_attn", + "precompute_dspark_freqs_cis", + "apply_dspark_rotary", + "apply_dspark_rotary_batched", + "dspark_attention_forward", + "dspark_attention_forward_batched", + # Block draft I/O. + "build_draft_input_ids", + "dspark_propose", ] diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c22ece009bff..1fd5e8c47c87 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1682,12 +1682,20 @@ def forward( block_residual: torch.Tensor, num_snapshots: int, attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, ) -> Tuple[torch.Tensor, int]: """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). ``block_residual`` is a preallocated snapshot bank in kernel-native ``[K_max, M, H]`` layout. Returns the running prefix sum and the number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. """ prefix_sum = hidden_states valid_block_residual = block_residual[:num_snapshots] @@ -1699,6 +1707,8 @@ def forward( self.self_attention_res_proj, self.self_attention_res_norm, ) + if capture is not None: + capture[0].maybe_capture_hidden_states(capture[1], hidden_states, None) if self.layer_idx % self.attn_res_block_size == 0: block_residual[num_snapshots].copy_(prefix_sum) @@ -1813,19 +1823,52 @@ def forward( hidden_states.shape[1], ) num_snapshots = 0 - for layer in self.layers: + capture_set = ( + getattr(spec_metadata, "_capture_layer_set", None) + if spec_metadata is not None + else None + ) + for i, layer in enumerate(self.layers): + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + capture = None + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + capture = (spec_metadata, self.layers[i - 1].layer_idx) hidden_states, num_snapshots = layer( - hidden_states, block_residual, num_snapshots, attn_metadata + hidden_states, block_residual, num_snapshots, attn_metadata, capture=capture ) - if spec_metadata is not None: - # DFlash hidden-state capture. K3's attn-residual scheme - # already folds the residual into the running prefix sum - # returned by each layer, so unlike Qwen3/Llama we pass the - # full hidden state with residual=None. Whether the drafter - # is trained against this prefix sum or some other tap point - # must be confirmed against the K3 drafter training recipe - # before real weights are used. - spec_metadata.maybe_capture_hidden_states(layer.layer_idx, hidden_states, None) + + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) hidden_states = _apply_attn_res( hidden_states, @@ -1958,16 +2001,20 @@ def __init__(self, model_config: ModelConfig): # separate dense checkpoint (K2.7-Code-DFlash schema) consumed by # the generic DFlashForCausalLM wrapper, and the target only has # to expose per-layer hidden states via maybe_capture_hidden_states - # (see KimiLinearModel.forward). No trained K3 drafter exists yet; - # this path is exercised with synthetic weights - # (examples/kimi_k3/make_synthetic_dflash_drafter.py). + # (see KimiLinearModel.forward). + # - DSpark: the same external-drafter flow with the Markov and + # confidence heads enabled (RadixArk/Kimi-K3-DSpark and friends). + # The target side is identical -- the capture in + # KimiLinearModel.forward is unconditional -- so this gate is the + # only place the mode has to be admitted. # Modes needing draft heads (MTP/Eagle) are blocked until a # draft-head checkpoint exists. assert ( spec_config is None or spec_config.spec_dec_mode.is_sa() or spec_config.spec_dec_mode.is_dflash() - ), "Kimi K3 supports speculative decoding only with SA or DFlash" + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" super().__init__( KimiLinearModel(model_config), model_config, diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index ab8522a54471..d95cf64349c8 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -12,7 +12,7 @@ from tensorrt_llm.logger import logger -from ...functional import PositionEmbeddingType, RotaryScalingType +from ...functional import PositionEmbeddingType from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..model_config import ModelConfig, TConfig @@ -25,25 +25,17 @@ WeightsLoadingConfig) from ..modules.mla import MLA from ..modules.rms_norm import RMSNorm -from ..modules.rotary_embedding import RotaryEmbedding - -try: - from ..custom_ops import \ - flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope -except ImportError: - _flashinfer_rope = None -from ..pyexecutor.config_utils import (_is_sliding_attention_layer, - get_layer_attention_window) from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..speculative import (SpecMetadata, get_spec_worker, should_use_separate_draft_kv_cache) -from ..speculative.dflash_attention import (get_dflash_flash_attention, - get_dflash_trtllm_gen_ops) +from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .checkpoints.base_weight_mapper import BaseWeightMapper from .modeling_auto import AutoModelForCausalLM from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, TModel, - get_model_architecture, register_auto_model) + get_model_architecture, + get_registered_draft_model_builder, + register_auto_model, register_draft_model) _SPECULATIVE_POSITION_HEADROOM = "_speculative_position_headroom" @@ -68,6 +60,406 @@ def _slice_spec_position_ids(position_ids: Optional[torch.Tensor], return position_ids[..., :num_tokens] +# --------------------------------------------------------------------------- +# DSpark draft-network heads, shared by both drafters that implement DSpark: +# the DFlash path (modeling_dflash.py, standalone Kimi K3 style drafters) and +# the DeepSeek-V4-Pro path (modeling_dspark.py, mtp.* stages inside the target +# checkpoint). DFlash is the degenerate case -- DSpark with the Markov and +# confidence heads switched off -- so the math below is the *only* copy; the +# two drafters differ solely in how they store the weights and whether their +# draft lm_head is TP vocab-sharded. +# +# Ported from DeepSeek's DeepSpec reference implementation +# (https://github.com/deepseek-ai/DeepSpec, MIT License). +# --------------------------------------------------------------------------- + + +def greedy_or_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: + """Argmax for temperature<=0, else temperature-scaled multinomial. + + Args: + logits: ``[..., vocab]``. + Returns: + token ids with the trailing vocab dim reduced. + """ + if temperature <= 0.0: + return logits.argmax(dim=-1) + probs = torch.softmax(logits.float() / temperature, dim=-1) + flat = probs.reshape(-1, probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1).squeeze(-1) + return sampled.view(probs.shape[:-1]) + + +def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, + markov_w2: torch.Tensor) -> torch.Tensor: + """Vanilla Markov head logit bias for one intra-block draft step. + + Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ + markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where + markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is + nn.Linear(rank, vocab, bias=False). With both weights stored + [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. + + Args: + prev_tokens: [B] long, previous token per request (draft vocab). + markov_w1: [vocab, rank]. + markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). + Returns: + [B, vocab_or_shard] bias in the markov weights' dtype. + """ + return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) + + +def dspark_markov_chain( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + step_bias_fn, + *, + hidden_states: Optional[torch.Tensor] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """The intra-block Markov refinement loop, shared by every DSpark head. + + Reference: DeepSpec ``VanillaMarkov.sample_block_tokens``. For step i, + ``logits_i += bias(prev_i)`` with ``prev_0`` = the anchor token (the last + accepted token, block slot 0) and ``prev_{i>0}`` = the token drawn from + step i-1's *biased* logits. llama.cpp PR #25173 implements the same chain. + + The chain being sequential is load-bearing and is NOT derivable from a + checkpoint: published DSpark drafters ship only the training-time + ``apply_block_logits``, which biases the whole block in one teacher-forced + pass. SGLang ``srt/models/dspark.py:34-64 run_markov_block`` settles it -- + it steps the same way, feeding back the token drawn from the biased logits. + + Two callers drive this, and a change here has to satisfy both: the embedded + flavour chains it inside the model, which owns its sampler, while the + standalone flavour is driven from the worker through ``next_token_fn`` so + the chain can advance on a TP-gathered global argmax the model cannot + compute on its own. + + Args: + base_logits: [B, K, vocab_or_shard] shared-lm_head logits. + first_prev_tokens: [B] long, anchor token ids (draft vocab). + step_bias_fn: ``(prev_tokens [B], step_hidden or None) -> bias``. A + closure, so a stateful head (RNN) can carry its recurrent state + across positions without a second loop. + hidden_states: [B, K, d] fed to ``step_bias_fn`` one position at a + time; None for the memoryless heads. + next_token_fn: ``([B, vocab_or_shard]) -> [B]`` token ids in the FULL + draft vocab; defaults to a plain argmax. Drafters whose draft + logits are TP vocab-sharded pass a shard-aware argmax here. + cast_bias_to_logits: cast the bias down to ``base_logits.dtype`` + before adding. The DFlash drafter does; the V4-Pro drafter does + not (its Markov weights and its logits already agree), and adding + the cast there would silently narrow its accumulation dtype. + Returns: + sampled_tokens [B, K], corrected_logits [B, K, vocab_or_shard]. + Greedy per-position argmax of the corrected logits reproduces the + reference sampled chain exactly. + """ + batch_size, block_size = base_logits.shape[:2] + if block_size == 0: + empty = torch.empty(batch_size, + 0, + dtype=torch.long, + device=base_logits.device) + return empty, base_logits + sampled, corrected = [], [] + prev = first_prev_tokens.long() + for k in range(block_size): + step_hidden = None if hidden_states is None else hidden_states[:, k] + bias = step_bias_fn(prev, step_hidden) + if cast_bias_to_logits: + bias = bias.to(base_logits.dtype) + step_logits = base_logits[:, k] + bias + corrected.append(step_logits.unsqueeze(1)) + if next_token_fn is None: + prev = torch.argmax(step_logits, dim=-1) + else: + prev = next_token_fn(step_logits).long() + sampled.append(prev) + return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) + + +def dspark_markov_chain_logits(base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + argmax_fn=None) -> torch.Tensor: + """Raw-tensor entry to :func:`dspark_markov_chain`, corrected logits only. + + For drafters that keep the Markov head as plain checkpoint tensors rather + than a :class:`VanillaMarkov` module (the DFlash path). ``markov_w2`` may + already be sliced down to this rank's TP vocab shard, in which case + ``argmax_fn`` must map a shard-local row back to a full-vocab token id. + """ + + def _step_bias(prev_tokens, step_hidden): + del step_hidden + return dspark_markov_step_bias(prev_tokens, markov_w1, markov_w2) + + _, corrected = dspark_markov_chain(base_logits, + first_prev_tokens, + _step_bias, + next_token_fn=argmax_fn, + cast_bias_to_logits=True) + return corrected + + +class VanillaMarkov(nn.Module): + """Low-rank token-bigram logit bias: ``bias = W2(W1[token])``.""" + + markov_head_type = "vanilla" + + def __init__(self, *, vocab_size: int, markov_rank: int): + super().__init__() + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + assert self.markov_rank > 0, ( + f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}.") + self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) + self.markov_w2 = nn.Linear(self.markov_rank, + self.vocab_size, + bias=False) + + def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: + return F.embedding(token_ids.long(), self.markov_w1.weight) + + def project_bias(self, + latent_states: torch.Tensor, + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + w2 = self.markov_w2.weight + if vocab_slice is not None: + w2 = w2[vocab_slice] + return F.linear(latent_states, w2) + + def compute_step_bias(self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + del hidden_states + w2 = self.markov_w2.weight + if vocab_slice is not None: + w2 = w2[vocab_slice] + return dspark_markov_step_bias(token_ids.long(), self.markov_w1.weight, + w2) + + def apply_step_logits( + self, + logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + return logits + self.compute_step_bias(token_ids, hidden_states) + + def sample_block_tokens( + self, + base_logits: torch.Tensor, + *, + first_prev_token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + temperature: float = 0.0, + vocab_slice: Optional[slice] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Autoregressive block sampling with the (memoryless) Markov bias. + + Args: + base_logits: ``[batch, block_size, vocab]`` from backbone+lm_head. + first_prev_token_ids: ``[batch]`` token preceding the 1st position. + hidden_states: ``[batch, block_size, d]`` (unused by vanilla). + vocab_slice / next_token_fn / cast_bias_to_logits: see + :func:`dspark_markov_chain`; only the TP vocab-sharded DFlash + drafter sets them. + Returns: + sampled_tokens ``[batch, block_size]``, + corrected_logits ``[batch, block_size, vocab]``. + """ + + def _step_bias(prev_tokens, step_hidden): + return self.compute_step_bias(prev_tokens, + step_hidden, + vocab_slice=vocab_slice) + + def _sample_step(step_logits): + return greedy_or_sample(step_logits, temperature) + + return dspark_markov_chain( + base_logits, + first_prev_token_ids, + _step_bias, + hidden_states=hidden_states, + next_token_fn=_sample_step + if next_token_fn is None else next_token_fn, + cast_bias_to_logits=cast_bias_to_logits, + ) + + +class GatedMarkovHead(VanillaMarkov): + """Markov bias gated by a sigmoid of [hidden, prev_embedding].""" + + markov_head_type = "gated" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank) + + def compute_step_bias(self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + assert hidden_states is not None + prev_emb = self.get_prev_embeddings(token_ids) + gate = torch.sigmoid( + self.gate_proj(torch.cat([hidden_states, prev_emb], + dim=-1))).to(dtype=prev_emb.dtype) + return self.project_bias(gate * prev_emb, vocab_slice=vocab_slice) + + +class RNNHead(VanillaMarkov): + """GRU-style head carrying recurrent state across block positions.""" + + markov_head_type = "rnn" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.hidden_size = int(hidden_size) + # [s_{k-1}; W1[x_{k-1}]; h_k] -> [gate; candidate; output] + self.joint_proj = nn.Linear(2 * markov_rank + hidden_size, + 3 * markov_rank) + + def _rnn_step(self, + state, + prev_embeddings, + hidden_states, + *, + vocab_slice: Optional[slice] = None): + z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) + gate_raw, cand_raw, out_raw = self.joint_proj(z).chunk(3, dim=-1) + gate = torch.sigmoid(gate_raw) + candidate = torch.tanh(cand_raw) + new_state = gate * state + (1.0 - gate) * candidate + bias = self.project_bias(torch.tanh(out_raw), vocab_slice=vocab_slice) + return new_state, bias + + def sample_block_tokens( + self, + base_logits: torch.Tensor, + *, + first_prev_token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + temperature: float = 0.0, + vocab_slice: Optional[slice] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert hidden_states is not None + state = torch.zeros(base_logits.shape[0], + self.markov_rank, + device=base_logits.device, + dtype=hidden_states.dtype) + + def _step_bias(prev_tokens, step_hidden): + nonlocal state + prev_emb = self.get_prev_embeddings(prev_tokens) + state, bias = self._rnn_step(state, + prev_emb, + step_hidden, + vocab_slice=vocab_slice) + return bias + + def _sample_step(step_logits): + return greedy_or_sample(step_logits, temperature) + + return dspark_markov_chain( + base_logits, + first_prev_token_ids, + _step_bias, + hidden_states=hidden_states, + next_token_fn=_sample_step + if next_token_fn is None else next_token_fn, + cast_bias_to_logits=cast_bias_to_logits, + ) + + +def build_markov_head(*, markov_head_type: str, vocab_size: int, + markov_rank: int, + hidden_size: int) -> Optional[nn.Module]: + """Factory mirroring DeepSpec ``build_markov_head``; None if rank==0.""" + if int(markov_rank) <= 0: + return None + kind = str(markov_head_type).lower() + if kind == "vanilla": + return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank) + if kind == "gated": + return GatedMarkovHead(vocab_size=vocab_size, + markov_rank=markov_rank, + hidden_size=hidden_size) + if kind == "rnn": + return RNNHead(vocab_size=vocab_size, + markov_rank=markov_rank, + hidden_size=hidden_size) + raise ValueError(f"Unsupported markov_head_type: {markov_head_type!r}") + + +class DSparkConfidenceHead(nn.Module): + """Per-position acceptance-confidence predictor (DeepSpec + AcceptRatePredictor). + + Input features are the backbone hidden state, optionally concatenated with + the Markov head's previous-token embedding. Output is a single logit per + position. + """ + + def __init__(self, + *, + hidden_size: int, + markov_rank: int = 0, + with_markov: bool = False): + super().__init__() + self.with_markov = bool(with_markov) + input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) + # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the + # confidence score is computed in fp32 (mirrors the DeepSpec reference + # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). + self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) + + def forward(self, + hidden_states: torch.Tensor, + prev_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor: + if self.with_markov: + assert prev_embeddings is not None + features = torch.cat( + [hidden_states, + prev_embeddings.to(hidden_states.dtype)], + dim=-1) + else: + features = hidden_states + # fp32 matmul for a stable confidence score (mirrors the reference). + return self.proj(features.float()).squeeze(-1) + + +def confident_prefix_length(confidence_logits: torch.Tensor, *, block_size: int, + threshold: float) -> int: + """First position k where ``sigmoid(confidence_k) < threshold``. + + Returns ``block_size`` when threshold<=0 (no truncation) or all positions + are confident. Assumes batch size 1 (functional-first scope). + """ + if threshold <= 0.0: + return int(block_size) + below = confidence_logits.sigmoid() < threshold + if not bool(below[0].any().item()): + return int(block_size) + return int(torch.nonzero(below[0], as_tuple=False)[0].item()) + + class Eagle3Attention(Attention): def __init__( @@ -851,1359 +1243,6 @@ def forward( return hidden_states_out, hidden_states_out -def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, - layer_idx: int) -> tuple[int, int]: - """flash-attn ``window_size`` for one draft layer of the block decode. - - DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF - attention with ``sliding_window`` set on 'sliding_attention' layers and - is_causal=False. HF's flash path - (transformers/modeling_flash_attention_utils.py) translates that to - ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. each - query attends keys within ``swa_window - 1`` KV-index distance on both - sides. In the DFlash pool layout KV index == token position, so this - limits draft queries to the most recent ``swa_window`` context tokens - plus the (nearby) draft block. Full-attention layers and non-dspark - drafters keep flash-attn's default ``(-1, -1)`` (no window). - """ - if not use_swa: - return (-1, -1) - if layer_types is not None and layer_idx < len(layer_types) and \ - layer_types[layer_idx] != 'sliding_attention': - return (-1, -1) - return (swa_window - 1, swa_window - 1) - - -def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, - markov_w2: torch.Tensor) -> torch.Tensor: - """Vanilla Markov head logit bias for one intra-block draft step. - - Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ - markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where - markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is - nn.Linear(rank, vocab, bias=False). With both weights stored - [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. - - Args: - prev_tokens: [B] long, previous token per request (draft vocab). - markov_w1: [vocab, rank]. - markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). - Returns: - [B, vocab_or_shard] bias in the markov weights' dtype. - """ - return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) - - -def dspark_markov_chain_logits( - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - markov_w1: torch.Tensor, - markov_w2: torch.Tensor, - argmax_fn=None, -) -> torch.Tensor: - """Apply the vanilla Markov intra-block bias across a drafted block. - - Reference: DeepSpec ``VanillaMarkov.sample_block_tokens`` at - temperature 0: for step i, ``logits_i += bias(prev_i)`` with - ``prev_0`` = the anchor token (last accepted token, block slot 0) and - ``prev_{i>0}`` = the greedy token from step i-1's *biased* logits. - llama.cpp PR #25173 implements the same greedy chain. - - Args: - base_logits: [B, K, vocab_or_shard] shared-lm_head logits. - first_prev_tokens: [B] long, anchor token ids (draft vocab). - markov_w1 / markov_w2: see :func:`dspark_markov_step_bias`. - argmax_fn: callable([B, vocab_or_shard]) -> [B] token ids in the - full draft vocab; defaults to plain argmax. Workers pass a - TP-aware argmax when the draft logits are vocab-sharded. - Returns: - [B, K, vocab_or_shard] biased logits. Greedy per-position argmax of - the result reproduces the reference sampled chain exactly. - """ - K = base_logits.shape[1] - if K == 0: - return base_logits - prev = first_prev_tokens.long() - steps = [] - for i in range(K): - bias = dspark_markov_step_bias(prev, markov_w1, markov_w2) - step_logits = base_logits[:, i] + bias.to(base_logits.dtype) - steps.append(step_logits) - if argmax_fn is not None: - prev = argmax_fn(step_logits).long() - else: - prev = torch.argmax(step_logits, dim=-1) - return torch.stack(steps, dim=1) - - -class DFlashForCausalLM(nn.Module): - """Draft model wrapper for DFlash speculative decoding. - - DFlash uses cross-attention where Q comes from noise/query tokens and K/V - come from the concatenation of target hidden states and noise hidden states. - The target_hidden stays CONSTANT across all layers (no input_layernorm applied). - - Reference: https://arxiv.org/pdf/2602.06036 - """ - - def __init__(self, - draft_config, - *, - dflash_attention_backend: str = 'VANILLA'): - """Build the draft model, resolving its architecture from the draft config - (falling back to a model_type-derived name when the checkpoint uses a - custom DFlash architecture label).""" - super().__init__() - - pretrained_cfg = draft_config.pretrained_config - try: - DraftModelClass, _ = get_model_architecture(pretrained_cfg) - except RuntimeError: - model_type = pretrained_cfg.model_type - arch_name = "".join(w.capitalize() - for w in model_type.split("_")) + "ForCausalLM" - logger.info( - f"DFlash: architecture {pretrained_cfg.architectures} not found, " - f"falling back to {arch_name} based on model_type={model_type}") - original_archs = pretrained_cfg.architectures - try: - pretrained_cfg.architectures = [arch_name] - DraftModelClass, _ = get_model_architecture(pretrained_cfg) - finally: - pretrained_cfg.architectures = original_archs - - # Remove spec_config to prevent recursive spec-dec initialization - draft_config_no_spec = replace(draft_config, - spec_config=None, - lm_head_gather_output=False) - - # Weights will be loaded later by ModelLoader.load_draft_weights() - self.draft_model_full = DraftModelClass(draft_config_no_spec) - self.model = self.draft_model_full.model - self.lm_head = self.draft_model_full.lm_head - - # Required by weight mappers - self.model_config = draft_config_no_spec - self.config = draft_config_no_spec.pretrained_config - - # Get mask_token_id from dflash_config - pretrained_config = draft_config.pretrained_config - dflash_config = getattr(pretrained_config, 'dflash_config', {}) - self.mask_token_id = dflash_config.get( - 'mask_token_id', - getattr(pretrained_config, 'mask_token_id', - pretrained_config.vocab_size)) - - self.target_layer_ids = dflash_config.get('target_layer_ids', None) - self.block_size = dflash_config.get( - 'block_size', getattr(pretrained_config, 'block_size', None)) - self.dflash_attention_backend = dflash_attention_backend - if self.dflash_attention_backend == 'VANILLA': - self._dflash_flash_attention = get_dflash_flash_attention() - elif self.dflash_attention_backend == 'TRTLLM': - self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() - else: - raise ValueError( - "DFlash attention backend must be VANILLA or TRTLLM, got " - f"{self.dflash_attention_backend!r}.") - self._dflash_trtllm_gen_workspace = None - self._dflash_trtllm_gen_counters = None - self.register_buffer("_dflash_batch_indices", None, persistent=False) - self.register_buffer("_dflash_block_offsets", None, persistent=False) - self._dflash_trtllm_gen_device = None - self._dflash_trtllm_gen_sm_count = None - logger.info( - f"DFlash draft model initialized with mask_token_id: {self.mask_token_id}, " - f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}, " - f"attention_backend: {self.dflash_attention_backend}") - - # DSpark drafters (DFlash + low-rank Markov head + confidence head, - # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- - # independent drafter-forward semantics ARE implemented here: - # - vanilla Markov intra-block logit bias (applied by DFlashWorker - # through apply_markov_chain_logits), - # - sliding-window attention on 'sliding_attention' draft layers - # during the block decode (use_swa / swa_window_size), - # - the shift_label output convention (hidden state at block slot j - # predicts draft token j+1; slot 0 holds the anchor token). - # Confidence-scheduled verification is NOT implemented yet: the - # confidence_proj weights are loaded (for the follow-up MR) but never - # used, and drafting always proposes the full K tokens. - self._dspark_shift_label = bool(dflash_config.get('shift_label', False)) - self._dspark_use_swa = bool(dflash_config.get('use_swa', False)) - self._dspark_swa_window = int( - dflash_config.get('swa_window_size', 0) or 0) - self._dspark_markov_rank = int(dflash_config.get('markov_rank', 0) or 0) - self._dspark_markov_head_type = str( - dflash_config.get('markov_head_type', 'vanilla') - or 'vanilla').lower() - self._dspark_use_confidence_head = bool( - dflash_config.get('use_confidence_head', False)) - # Plain None placeholders rather than nn.Parameter/buffer: most - # DFlash checkpoints don't ship these heads, and their shapes - # ([vocab, rank]) are checkpoint-dependent, so nothing is - # pre-allocated. load_weights() fills them in only when the - # checkpoint ships them; consumers treat None as "head absent". - self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) - self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) - self.confidence_proj_weight = None # loaded, unused (follow-up MR) - self.confidence_proj_bias = None - - if self._dspark_markov_rank > 0 and \ - self._dspark_markov_head_type != 'vanilla': - raise ValueError( - f"DFlash dspark drafter declares markov_head_type=" - f"'{self._dspark_markov_head_type}'; only 'vanilla' is " - "supported (gated/rnn heads need per-step hidden features).") - if self._dspark_use_swa and self._dspark_swa_window < 1: - raise ValueError( - "DFlash dspark drafter sets use_swa but swa_window_size=" - f"{dflash_config.get('swa_window_size')} is invalid.") - # causal=true is only invalid under the dspark convention. Legacy - # DFlash drafter configs (e.g. Laguna) also carry a causal field; - # their causality is handled by the legacy decode path - # (_sliding_layers_causal), so don't reject them here. - is_dspark = (str(dflash_config.get('projector_type', '') - or '').lower() == 'dspark' or self._dspark_shift_label - or self._dspark_use_swa or self._dspark_markov_rank > 0 - or self._dspark_use_confidence_head) - if is_dspark and dflash_config.get('causal'): - raise ValueError( - "DFlash dspark drafter sets causal=true; the block decode " - "only supports the non-causal dspark convention.") - # Per-layer flash-attn window for the block decode, resolved once. - num_draft_layers = getattr(pretrained_config, 'num_hidden_layers', 0) - layer_types = getattr(pretrained_config, 'layer_types', None) - self._dspark_layer_windows = [ - dspark_layer_window_size(self._dspark_use_swa, - self._dspark_swa_window, layer_types, i) - for i in range(num_draft_layers) - ] - if self._dspark_use_confidence_head: - logger.warning( - "DFlash dspark drafter declares use_confidence_head; " - "confidence-scheduled verification is not implemented yet " - "(confidence_proj weights are loaded but unused, drafting " - "always proposes the full K tokens).") - - self.logits_processor = None # Set by caller after construction - - # RoPE - lazily initialized from draft model's attention module - self._rope_initialized = False - self._rotary_cos_sin = None - self._is_neox = True - - self._cos_sin_cache_fp32 = None - self._rope_dummy_q = None - - # Lazy-built after weights load (see _build_fused_kv_buffers). - self._fused_kv_weight = None - self._fused_kv_bias = None - self._k_norm_stacked = None - self._k_norm_eps = None - self._num_attn_layers = 0 - self._num_heads = 0 - self._head_dim = 0 - self._num_kv_heads = 0 - self._has_qk_norm = False - self._use_fused_qk_norm_rope = False - # Laguna-specific draft-layer behaviors, disabled by default so generic - # DFlash drafters keep the original contract (no context input_layernorm, - # non-causal block attention). Subclasses opt in. - self._context_input_layernorm = False - self._sliding_layers_causal = False - self._warn_inferred_attention_windows() - - @staticmethod - def _rope_signature(attn): - """Return the effective RoPE configuration used by an attention layer.""" - if attn.rotary_emb is not None: - return ( - attn.rotary_emb.rope_params, - attn.rotary_emb.head_dim, - attn.rotary_emb.is_neox, - ) - if attn.pos_embd_params is not None: - return ( - attn.pos_embd_params.rope, - attn.head_dim, - attn.pos_embd_params.is_neox, - ) - return None - - def _validate_uniform_rope(self): - """Check that all draft layers can safely share one RoPE cache.""" - if len(self.model.layers) == 0: - raise ValueError("DFlash requires at least one draft model layer.") - - signatures = [ - self._rope_signature(layer.self_attn) for layer in self.model.layers - ] - - mismatched_layers = [ - layer_idx - for layer_idx, signature in enumerate(signatures[1:], start=1) - if signature != signatures[0] - ] - if mismatched_layers: - layer_types = getattr(self.config, 'layer_types', None) - raise ValueError( - "DFlash shares one RoPE cache across draft layers, but layers " - f"{mismatched_layers} have a different effective RoPE " - f"configuration from layer 0. layer_types={layer_types}.") - - def _init_rope(self): - """Initialize RoPE from the draft model's attention configuration. - - Reuses the existing RotaryEmbedding infrastructure which correctly - handles all RoPE variants (standard, YaRN, scaled, etc.). - """ - # The flattened context-KV path shares layer 0's RoPE cache. - self._validate_uniform_rope() - attn0 = self.model.layers[0].self_attn - - if attn0.rotary_emb is not None: - self._rotary_cos_sin = attn0.rotary_emb.rotary_cos_sin - self._is_neox = attn0.rotary_emb.is_neox - elif attn0.pos_embd_params is not None: - rope_emb = RotaryEmbedding( - attn0.pos_embd_params.rope, - head_dim=attn0.head_dim, - is_neox=attn0.pos_embd_params.is_neox, - ) - self._rotary_cos_sin = rope_emb.rotary_cos_sin - self._is_neox = rope_emb.is_neox - else: - # Fallback: basic NeoX-style RoPE - config = self.config - head_dim = getattr(config, 'head_dim', - config.hidden_size // config.num_attention_heads) - rope_theta = getattr(config, 'rope_theta', 1000000.0) - max_pos = getattr(config, 'max_position_embeddings', 32768) - - inv_freq = 1.0 / (rope_theta**(torch.arange( - 0, head_dim, 2, dtype=torch.float32, device='cuda') / head_dim)) - positions = torch.arange(max_pos, - dtype=torch.float32, - device='cuda') - freqs = torch.outer(positions, inv_freq) - rope_cos = freqs.cos().to(config.torch_dtype) - rope_sin = freqs.sin().to(config.torch_dtype) - # [max_pos, 2, rot_dim//2] to match RotaryEmbedding format - self._rotary_cos_sin = torch.stack([rope_cos, rope_sin], dim=1) - self._is_neox = True - - self._rope_initialized = True - - def project_target_hidden(self, - hidden_states: torch.Tensor) -> torch.Tensor: - """Project captured target hidden states into the draft hidden space. - - Generic DFlash: fc then hidden_norm. Subclasses (e.g. Laguna) may - normalize the per-aux features first by overriding this method. - """ - hidden_states = hidden_states.to(self.fc.weight.dtype) - return self.hidden_norm(self.fc(hidden_states)) - - @property - def has_markov_head(self) -> bool: - return self._dspark_markov_rank > 0 and self.markov_w1 is not None - - def apply_markov_chain_logits( - self, - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - argmax_fn=None, - vocab_slice: slice | None = None) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - No-op (returns ``base_logits`` unchanged) for non-dspark drafters. - See :func:`dspark_markov_chain_logits` for the semantics; when - ``base_logits`` is a TP vocab shard, the caller must pass this - rank's ``vocab_slice`` (to shard the markov_w2 rows identically) - and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker - handles both. - """ - if not self.has_markov_head: - return base_logits - markov_w2 = self.markov_w2 if vocab_slice is None else \ - self.markov_w2[vocab_slice] - return dspark_markov_chain_logits(base_logits, - first_prev_tokens, - self.markov_w1, - markov_w2, - argmax_fn=argmax_fn) - - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): - """Hook applied to the block-attention output before o_proj. - - No-op for generic DFlash; overridden by drafters that gate (e.g. Laguna). - """ - return attn_output - - def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): - """Load weights into the DFlash draft model. - - DFlash checkpoints differ from standard HF format: - - Layer weights lack the 'model.' prefix (e.g., 'layers.0...' not 'model.layers.0...') - - Extra DFlash-specific weights: 'fc.weight', 'hidden_norm.weight' - - Missing embed_tokens and lm_head (shared with target model) - """ - # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft - # loader expects split q/k/v (a fused key is silently dropped otherwise). - if any(k.endswith('self_attn.qkv_proj.weight') for k in weights): - for attr in ('num_attention_heads_per_layer', - 'num_key_value_heads_per_layer'): - per_layer = getattr(self.config, attr, None) - if per_layer is not None and len(set(per_layer)) > 1: - raise ValueError( - "DFlash load_weights() splits the fused qkv_proj using " - "the global head count, but the drafter has heterogeneous " - f"{attr} {sorted(set(per_layer))}; per-layer qkv splitting " - "is required for this checkpoint.") - head_dim = getattr( - self.config, 'head_dim', - self.config.hidden_size // self.config.num_attention_heads) - num_kv_heads = getattr(self.config, 'num_key_value_heads', - self.config.num_attention_heads) - q = self.config.num_attention_heads * head_dim - kv = num_kv_heads * head_dim - split = {} - for k, v in weights.items(): - if k.endswith('self_attn.qkv_proj.weight'): - b = k[:-len('qkv_proj.weight')] - split[b + 'q_proj.weight'] = v[:q] - split[b + 'k_proj.weight'] = v[q:q + kv] - split[b + 'v_proj.weight'] = v[q + kv:] - else: - split[k] = v - weights = split - - # DSpark head weights: keep them out of the backbone remap (they'd - # get a 'model.' prefix and be dropped by allow_partial_loading). - # markov_w1/markov_w2 drive the intra-block logit bias; the - # confidence_proj weights are loaded for the confidence-scheduling - # follow-up MR but are not used yet. - dspark_keys = ('markov_w1.weight', 'markov_w2.weight', - 'confidence_proj.weight', 'confidence_proj.bias') - dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} - if dspark_weights: - weights = { - k: v - for k, v in weights.items() if k not in dspark_weights - } - if self._dspark_markov_rank > 0: - vocab = self.config.vocab_size - rank = self._dspark_markov_rank - for k in ('markov_w1.weight', 'markov_w2.weight'): - if k not in dspark_weights: - raise ValueError( - f"DFlash dspark drafter declares markov_rank=" - f"{self._dspark_markov_rank} but the checkpoint is " - f"missing {k}.") - if tuple(dspark_weights[k].shape) != (vocab, rank): - raise ValueError( - f"DFlash dspark {k} has shape " - f"{tuple(dspark_weights[k].shape)}, expected " - f"[vocab, markov_rank] = ({vocab}, {rank}).") - self.markov_w1 = dspark_weights['markov_w1.weight'].to('cuda') - self.markov_w2 = dspark_weights['markov_w2.weight'].to('cuda') - if 'confidence_proj.weight' in dspark_weights: - self.confidence_proj_weight = dspark_weights[ - 'confidence_proj.weight'].to('cuda') - if 'confidence_proj.bias' in dspark_weights: - self.confidence_proj_bias = dspark_weights[ - 'confidence_proj.bias'].to('cuda') - - # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights - remapped = {} - for key, value in weights.items(): - if key in ('fc.weight', 'hidden_norm.weight'): - # DFlash-specific projection weights - store directly - remapped[key] = value - elif key == 'norm.weight': - remapped['model.norm.weight'] = value - elif not key.startswith('model.'): - remapped[f'model.{key}'] = value - else: - remapped[key] = value - - # Load DFlash-specific weights directly - if 'fc.weight' in remapped: - self.fc = nn.Linear(remapped['fc.weight'].shape[1], - remapped['fc.weight'].shape[0], - bias=False, - device='cuda', - dtype=remapped['fc.weight'].dtype) - self.fc.weight.data.copy_(remapped['fc.weight']) - del remapped['fc.weight'] - - if 'hidden_norm.weight' in remapped: - rms_norm_eps = getattr(self.config, 'rms_norm_eps', 1e-6) - self.hidden_norm = nn.RMSNorm( - remapped['hidden_norm.weight'].shape[0], - eps=rms_norm_eps, - device='cuda', - elementwise_affine=True, - dtype=remapped['hidden_norm.weight'].dtype) - self.hidden_norm.weight.data.copy_(remapped['hidden_norm.weight']) - del remapped['hidden_norm.weight'] - - # Load remaining weights into the draft model. - # DFlash checkpoints don't include embed_tokens or lm_head, so allow partial loading - # since those modules won't find matching weights. - self.draft_model_full.load_weights(weights=remapped, - weight_mapper=weight_mapper, - allow_partial_loading=True) - - def load_weights_from_target_model(self, - target_model: torch.nn.Module) -> None: - """Share embed_tokens and lm_head from the target model.""" - self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens - self.draft_model_full.lm_head = target_model.lm_head - self.lm_head = target_model.lm_head - - def precompute_context_kv( - self, - projected_hidden: torch.Tensor, - positions: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Post-norm / post-RoPE K and V for ALL drafter layers in one fused GEMM. - - Args: - projected_hidden: [N, hidden_size], already fc + hidden_norm'd. - positions: [N] int32/64, RoPE positions for each entry. - Returns: - k: [N, L, nkv, hd] post k_norm and RoPE - v: [N, L, nkv, hd] post split only - """ - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - N = projected_hidden.shape[0] - L = self._num_attn_layers - nkv = self._num_kv_heads - hd = self._head_dim - weight_dtype = self._fused_kv_weight.dtype - if getattr(self, '_input_ln_eps', None) is not None: - ph = projected_hidden.float() - ph = ph * torch.rsqrt( - ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) - projected_hidden = ph.to(weight_dtype) - elif projected_hidden.dtype != weight_dtype: - projected_hidden = projected_hidden.to(weight_dtype) - - kv_flat = F.linear(projected_hidden, self._fused_kv_weight, - self._fused_kv_bias) - # Per-layer layout [L0_K|L0_V|L1_K|L1_V|...] keeps K and V contiguous - # after the select() splits — no extra copy required. - kv = kv_flat.view(N, L, 2, nkv, hd) - k = kv[:, :, 0].contiguous() - v = kv[:, :, 1].contiguous() - - if self._k_norm_stacked is not None: - # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; - # each layer has its own weight ([L, hd]) but shares eps. - k = F.rms_norm(k, (hd, ), eps=self._k_norm_eps) - k = k * self._k_norm_stacked.view(1, L, 1, hd) - - self._fused_rope_inplace(k.view(N * L, nkv * hd), positions, N, L) - return k, v - - def _get_cos_sin_cache(self) -> torch.Tensor: - """Return the flashinfer-style cos/sin cache for the drafter. - - Shape [max_positions, head_dim], fp32 — flashinfer's - apply_rope_with_cos_sin_cache_inplace requires fp32 regardless of - the query/key dtype. - """ - if self._cos_sin_cache_fp32 is not None: - return self._cos_sin_cache_fp32 - if not self._rope_initialized: - self._init_rope() - max_pos = self._rotary_cos_sin.shape[0] - self._cos_sin_cache_fp32 = self._rotary_cos_sin.view(max_pos, -1).to( - torch.float32).contiguous() - return self._cos_sin_cache_fp32 - - def _fused_rope_inplace( - self, - k_flat: torch.Tensor, - positions: torch.Tensor, - N: int, - L: int, - ) -> None: - """In-place fused RoPE over [N*L, nkv*hd] K values. - - Layout of k_flat: row (i*L + l) holds layer l of position i, so - positions must be repeat_interleaved by L to match. - """ - positions_int32 = positions.view(-1).to(torch.int32) - if L > 1: - positions_int32 = positions_int32.repeat_interleave(L) - - if _flashinfer_rope is not None: - # flashinfer requires a non-None query tensor; pass a single-head - # scratch so the extra rotate is negligible. - need_rows = k_flat.shape[0] - dummy_q = self._rope_dummy_q - if (dummy_q is None or dummy_q.dtype != k_flat.dtype - or dummy_q.shape[0] < need_rows): - dummy_q = k_flat.new_empty(need_rows, self._head_dim) - self._rope_dummy_q = dummy_q - _flashinfer_rope( - positions_int32, - dummy_q[:need_rows], - k_flat, - self._head_dim, - self._get_cos_sin_cache(), - self._is_neox, - ) - return - - # Pure-PyTorch fallback (older environments without flashinfer). - cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), - dtype=k_flat.dtype) - k_roped = RotaryEmbedding.apply_rotary_pos_emb( - k_flat.view(k_flat.shape[0], -1, self._head_dim), - cos.squeeze(0), - sin.squeeze(0), - unsqueeze_dim=1, - is_neox=self._is_neox, - ) - k_flat.copy_(k_roped.view_as(k_flat)) - - def _build_fused_kv_buffers(self) -> None: - """Stack per-layer KV projection + k_norm weights for a single fused GEMM. - - Must run after weights are loaded. - """ - if self._fused_kv_weight is not None: - return - layers_attn = [layer.self_attn for layer in self.model.layers] - attn0 = layers_attn[0] - q_size = attn0.q_size - kv_size = attn0.kv_size - head_dim = attn0.head_dim - num_heads = attn0.num_heads - num_kv_heads = attn0.num_key_value_heads - # Head counts are read from layer 0 here and in dflash_forward; assert - # uniformity (the target uses per-layer heads, the drafter does not). - for a in layers_attn[1:]: - assert ( - a.q_size == q_size and a.kv_size == kv_size - and a.head_dim == head_dim and a.num_heads == num_heads - and a.num_key_value_heads == num_kv_heads), ( - "DFlash fused KV requires all drafter layers to share " - "q_size / kv_size / head_dim / num_heads / num_kv_heads.") - - has_k_norm = [hasattr(a, 'k_norm') for a in layers_attn] - assert all(has_k_norm) or not any(has_k_norm), ( - "DFlash fused KV requires either all or no drafter layers to have k_norm." - ) - - kv_weights = [ - a.qkv_proj.weight[q_size:q_size + 2 * kv_size] for a in layers_attn - ] - # Fold each drafter layer's input_layernorm weight into its KV projection - # so context K/V match the query path. vLLM laguna_dflash applies - # layer.input_layernorm to context states before KV; RMSNorm gives - # (x_hat * w) @ Wkv.T == x_hat @ (Wkv * w).T, and the shared 1/rms(x) is - # applied to projected_hidden in precompute_context_kv. - dlayers = self.model.layers - if self._context_input_layernorm and all( - hasattr(dl, 'input_layernorm') for dl in dlayers): - eps_set = { - getattr(dl.input_layernorm, 'variance_epsilon', - getattr(self.config, 'rms_norm_eps', 1e-6)) - for dl in dlayers - } - assert len(eps_set) == 1, ( - "DFlash fused context input_layernorm needs all drafter layers " - f"to share variance_epsilon; got {sorted(eps_set)}") - self._input_ln_eps = eps_set.pop() - folded = [] - for w, dl in zip(kv_weights, dlayers): - scale = dl.input_layernorm.weight.data - if getattr(dl.input_layernorm, 'use_gemma', False): - scale = scale + 1 - folded.append(w * scale[None, :].to(w.dtype)) - kv_weights = folded - else: - self._input_ln_eps = None - fused_kv_weight = torch.cat(kv_weights, dim=0).contiguous() - if attn0.qkv_proj.bias is not None: - kv_biases = [ - a.qkv_proj.bias[q_size:q_size + 2 * kv_size] - for a in layers_attn - ] - self._fused_kv_bias = torch.cat(kv_biases, dim=0).contiguous() - else: - self._fused_kv_bias = None - - if all(has_k_norm): - k_norm0 = layers_attn[0].k_norm - eps = k_norm0.variance_epsilon - eps_set = {a.k_norm.variance_epsilon for a in layers_attn} - assert len(eps_set) == 1, ( - f"DFlash fused k_norm requires all drafter layers to share " - f"variance_epsilon; got {sorted(eps_set)}.") - self._k_norm_stacked = torch.stack( - [a.k_norm.weight.data for a in layers_attn]) - self._k_norm_eps = eps - else: - self._k_norm_stacked = None - self._k_norm_eps = None - self._num_attn_layers = len(layers_attn) - self._num_heads = num_heads - self._head_dim = head_dim - self._num_kv_heads = num_kv_heads - self._fused_kv_weight = fused_kv_weight - - # fused_qk_norm_rope derives YaRN / partial-rotary frequencies on - # the fly, which can disagree with precompute_context_kv's cached - # cos/sin. Only enable it when the drafter uses plain RoPE. - self._has_qk_norm = (all(has_k_norm) - and all(hasattr(a, 'q_norm') for a in layers_attn)) - rope_params = getattr(getattr(attn0, 'pos_embd_params', None), 'rope', - None) - scale_type = getattr(rope_params, 'scale_type', None) - partial_rotary_factor = getattr( - getattr(attn0, 'pretrained_config', None), 'partial_rotary_factor', - 1.0) - self._use_fused_qk_norm_rope = (self._has_qk_norm - and hasattr(attn0, 'apply_qk_norm_rope') - and rope_params is not None - and scale_type - in (None, RotaryScalingType.none) - and partial_rotary_factor == 1.0) - - logger.debug( - f"DFlash: fused KV weights built for {self._num_attn_layers} layers " - f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})") - - def _get_rope_cos_sin(self, positions, dtype=None): - """Get cos/sin for given positions, suitable for apply_rotary_pos_emb. - - Args: - positions: [B, seq_len] - dtype: target dtype for cos/sin (default: keep original) - Returns: - rope_cos: [B, seq, rot_dim//2] (broadcastable with unsqueeze_dim=1) - rope_sin: [B, seq, rot_dim//2] - """ - if not self._rope_initialized: - self._init_rope() - - # rotary_cos_sin: [max_pos, 2, rot_dim//2] - rope_cache = self._rotary_cos_sin[positions] # [B, seq, 2, rot_dim//2] - rope_cos = rope_cache[..., 0, :] # [B, seq, rot_dim//2] - rope_sin = rope_cache[..., 1, :] - if dtype is not None: - rope_cos = rope_cos.to(dtype) - rope_sin = rope_sin.to(dtype) - return rope_cos, rope_sin - - def _warn_inferred_attention_windows(self) -> None: - """Warn once at initialization when checkpoint metadata enables SWA.""" - if getattr(self.config, 'use_sliding_window', None) is not None: - return - - num_hidden_layers = getattr(self.config, 'num_hidden_layers', None) - if num_hidden_layers is None: - num_hidden_layers = len(self.model.layers) - layers_by_window = {} - for layer_idx in range(num_hidden_layers): - window = get_layer_attention_window(self.config, layer_idx) - if window is not None: - layers_by_window.setdefault(window, []).append(layer_idx) - - for window, layer_indices in layers_by_window.items(): - logger.warning( - "DFlash inferred pooled-context sliding-window attention from " - f"checkpoint config for draft layers {layer_indices}: " - f"window={window}. Context attention is truncated to {window} " - "tokens for these layers; if the drafter expects full context, " - "acceptance rate may drop. Set use_sliding_window explicitly " - "to confirm or disable windowing.") - - def _get_attention_mask_args(self, layer_idx): - """Return FlashAttention causal and local-window arguments for a layer.""" - layer_types = getattr(self.config, 'layer_types', None) - is_sliding_layer = False - if layer_types: - layer_type = layer_types[layer_idx % len(layer_types)] - is_sliding_layer = _is_sliding_attention_layer(layer_type) - - sliding_window = get_layer_attention_window(self.config, layer_idx) - is_sliding_layer = is_sliding_layer or sliding_window is not None - if not is_sliding_layer: - return False, (-1, -1) - - causal = self._sliding_layers_causal or sliding_window is not None - if sliding_window is None: - # Legacy drafters without an explicit window preserve their prior - # non-windowed behavior. - return causal, (-1, -1) - # FlashAttention's bounds are inclusive: W tokens are current + W-1 left. - return causal, (sliding_window - 1, 0) - - def _prepare_dflash_trtllm_gen_buffers( - self, - dtype: torch.dtype, - device: torch.device, - max_batch_size: int, - block_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - ) -> None: - trtllm_gen_ops = self._dflash_trtllm_gen_ops - workspace_bytes = trtllm_gen_ops.get_workspace_size( - dtype=dtype, - num_tokens=max_batch_size * block_size, - num_gen_tokens=max_batch_size * block_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_dim, - max_num_requests=max_batch_size, - rotary_embedding_dim=0, - fp8_context_fmha=False, - ) - device = torch.device(device) - is_capturing = torch.cuda.is_current_stream_capturing() - if self._dflash_trtllm_gen_device != device: - if is_capturing: - raise RuntimeError( - "DFlash TRTLLM-Gen buffers must be prepared on the current " - "device before CUDA graph capture.") - self._dflash_trtllm_gen_device = device - self._dflash_trtllm_gen_sm_count = ( - torch.cuda.get_device_properties(device).multi_processor_count) - - workspace = self._dflash_trtllm_gen_workspace - workspace_needs_allocation = ( - workspace is None or workspace.device != device - or workspace.numel() * workspace.element_size() < workspace_bytes) - if workspace_needs_allocation: - if is_capturing: - raise RuntimeError( - "The DFlash TRTLLM-Gen workspace must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_trtllm_gen_workspace = torch.empty(workspace_bytes, - dtype=torch.uint8, - device=device) - - sm_count = self._dflash_trtllm_gen_sm_count - counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( - num_heads, max_batch_size, sm_count) - counters = self._dflash_trtllm_gen_counters - counters_need_allocation = (counters is None - or counters.device != device - or counters.numel() * - counters.element_size() < counter_bytes) - if counters_need_allocation: - if is_capturing: - raise RuntimeError( - "The DFlash TRTLLM-Gen counter buffer must be allocated at " - "the required size before CUDA graph capture.") - self._dflash_trtllm_gen_counters = torch.zeros(counter_bytes, - dtype=torch.uint8, - device=device) - - append_batch_indices = self._dflash_batch_indices - block_offsets = self._dflash_block_offsets - static_indices_need_allocation = ( - append_batch_indices is None or block_offsets is None - or append_batch_indices.device != device - or block_offsets.device != device - or append_batch_indices.size(0) < max_batch_size - or append_batch_indices.size(1) != block_size - or block_offsets.numel() != block_size) - if static_indices_need_allocation: - if is_capturing: - raise RuntimeError( - "DFlash TRTLLM-Gen index buffers must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_batch_indices = (torch.arange( - max_batch_size, dtype=torch.int32, - device=device).view(-1, 1).expand(-1, block_size).contiguous()) - self._dflash_block_offsets = torch.arange(block_size, - dtype=torch.int32, - device=device) - - def dflash_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_k_cache: torch.Tensor, - ctx_v_cache: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor] = None, - ctx_page_table: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """DFlash draft forward with cross-attention over a pooled K/V buffer. - - All shapes are fixed so the forward is CUDA-graph compatible. - - Args: - noise_embedding: [B, block_size, hidden_size] - query_positions: [B, block_size] - num_ctx_per_req: [B] — per-batch context length in the pool - ctx_k_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] - ctx_v_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] - ctx_cache_batch_idx: [B] — slot index into the pool per batch entry - Returns: - [B * block_size, hidden_size] - """ - if self.dflash_attention_backend == 'TRTLLM': - if ctx_kv_cache is None or ctx_page_table is None: - raise RuntimeError( - "DFlash TRTLLM-Gen requires a paged context cache and page table." - ) - trtllm_gen_ops = self._dflash_trtllm_gen_ops - elif self.dflash_attention_backend == 'VANILLA': - flash_attention = self._dflash_flash_attention - else: - raise ValueError( - "DFlash attention backend must be VANILLA or TRTLLM, got " - f"{self.dflash_attention_backend!r}.") - - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - - layer0 = self.model.layers[0] - attn0 = layer0.self_attn - q_size = attn0.q_size - kv_size = attn0.kv_size - head_dim = attn0.head_dim - # Uniformity across layers is asserted in _build_fused_kv_buffers (above). - num_heads_per_rank = attn0.num_heads - num_kv_heads_per_rank = attn0.num_key_value_heads - gqa_group_size = num_heads_per_rank // num_kv_heads_per_rank - - has_qk_norm = self._has_qk_norm - is_bf16 = noise_embedding.dtype == torch.bfloat16 - use_fused_qk_norm_rope = self._use_fused_qk_norm_rope and is_bf16 - use_fused_rope = (_flashinfer_rope is not None and has_qk_norm - and is_bf16 and not use_fused_qk_norm_rope) - - B = noise_embedding.shape[0] - block_size = noise_embedding.shape[1] - - hidden_states = noise_embedding # [B, block_size, hidden] - - # Precompute RoPE cos/sin for the pure-PyTorch fallback path only. - # The fused flashinfer path reads self._get_cos_sin_cache() inline. - rope_dtype = hidden_states.dtype - if not use_fused_rope: - q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, - dtype=rope_dtype) - _rope = RotaryEmbedding.apply_rotary_pos_emb - - # cache_seqlens (BEFORE append). flash_attn appends block_size - # k/v at cache_seqlens[i]..+block_size for batch i. - cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) - cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) - - if self.dflash_attention_backend == 'TRTLLM': - max_batch_size = ctx_page_table.size(0) - self._prepare_dflash_trtllm_gen_buffers( - hidden_states.dtype, - hidden_states.device, - max_batch_size, - block_size, - num_heads_per_rank, - num_kv_heads_per_rank, - head_dim, - ) - block_tables = ctx_page_table.index_select( - 0, cache_batch_idx_i32.long()) - pages_per_slot = block_tables.size(1) - page_size = ctx_kv_cache.size(-2) - kv_indices = block_tables.flatten() - kv_indptr = torch.arange( - 0, - (B + 1) * pages_per_slot, - pages_per_slot, - dtype=torch.int32, - device=hidden_states.device, - ) - seq_lens_after = cache_seqlens_i32 + block_size - kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 - batch_indices = self._dflash_batch_indices - append_batch_indices = batch_indices[:B].reshape(-1) - append_positions = ( - cache_seqlens_i32.view(-1, 1) + - self._dflash_block_offsets).reshape(-1).contiguous() - - # Flatten query positions once for the fused QK-norm-RoPE kernel. - query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) - - residual = None - - for layer_idx, layer in enumerate(self.model.layers): - attn_mod = layer.self_attn - - # Apply input_layernorm (flatten to 2D for norm, reshape back) - hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - if residual is None: - residual = hidden_states.clone() - hs_normed_flat = layer.input_layernorm(hs_flat) - else: - res_flat = residual.reshape(-1, residual.shape[-1]) - hs_normed_flat, res_flat = layer.input_layernorm( - hs_flat, res_flat) - residual = res_flat.reshape(B, block_size, -1) - - # QKV projection on normed query tokens (2D) - qkv_query = attn_mod.qkv_proj(hs_normed_flat) # [B*blk, qkv_size] - - if use_fused_qk_norm_rope: - # One kernel does q_norm + k_norm + RoPE in-place on qkv. - # Only safe when the drafter's rope params don't use YaRN / - # long-rope / partial-rotary — otherwise fall back to the - # shared-cache path below. - attn_mod.apply_qk_norm_rope(qkv_query, query_positions_flat_i32) - q_all_2d = qkv_query[:, :q_size] - k_noise_2d = qkv_query[:, q_size:q_size + kv_size] - v_noise_2d = qkv_query[:, q_size + kv_size:] - Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, - head_dim) - k_noise_bshd = k_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - v_noise_bshd = v_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - elif use_fused_rope: - # Per-head RMSNorm on q/k (returns new contiguous tensors), - # then flashinfer in-place RoPE sharing the same cos/sin cache - # as precompute_context_kv. - q = attn_mod.q_norm(qkv_query[:, :q_size].reshape( - -1, head_dim)).view(-1, q_size) - k = attn_mod.k_norm(qkv_query[:, - q_size:q_size + kv_size].reshape( - -1, - head_dim)).view(-1, kv_size) - _flashinfer_rope( - query_positions_flat_i32, - q, - k, - head_dim, - self._get_cos_sin_cache(), - self._is_neox, - ) - Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) - k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, - head_dim) - v_noise_bshd = qkv_query[:, q_size + kv_size:].reshape( - B, block_size, num_kv_heads_per_rank, head_dim) - else: - qkv_query_3d = qkv_query.reshape(B, block_size, -1) - q_all = qkv_query_3d[..., :q_size] - k_noise_all = qkv_query_3d[..., q_size:q_size + kv_size] - v_noise_all = qkv_query_3d[..., q_size + kv_size:] - if has_qk_norm: - q_for_rope = attn_mod.q_norm(q_all.reshape( - -1, head_dim)).reshape(B, block_size, q_size) - k_noise_for_rope = attn_mod.k_norm( - k_noise_all.reshape(-1, head_dim)).reshape( - B, block_size, kv_size) - else: - q_for_rope = q_all - k_noise_for_rope = k_noise_all - Q = _rope(q_for_rope.reshape(B, block_size, num_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) - k_noise_rope = _rope(k_noise_for_rope.reshape( - B, block_size, num_kv_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) - Q_bshd = Q.transpose(1, 2) - k_noise_bshd = k_noise_rope.transpose(1, 2) - v_noise_bshd = v_noise_all.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - - # Per-layer view into the pooled ctx cache. - causal, window_size = self._get_attention_mask_args(layer_idx) - dspark_window = (self._dspark_layer_windows[layer_idx] if layer_idx - < len(self._dspark_layer_windows) else (-1, -1)) - if dspark_window != (-1, -1): - window_size = dspark_window - if self.dflash_attention_backend == 'TRTLLM': - layer_cache = ctx_kv_cache[layer_idx] - trtllm_gen_ops.append_paged_kv_cache( - append_key=k_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), - append_value=v_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), - batch_indices=append_batch_indices, - positions=append_positions, - paged_kv_cache=layer_cache, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_last_page_len=kv_last_page_len, - kv_layout="HND", - ) - out = torch.empty_like(Q_bshd) - q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) - out_flat = out.reshape(-1, num_heads_per_rank, head_dim) - window_left = window_size[0] - if causal: - trtllm_gen_ops.batch_decode_with_kv_cache( - query=q_flat, - kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), - workspace_buffer=self._dflash_trtllm_gen_workspace, - block_tables=block_tables, - seq_lens=seq_lens_after, - max_seq_len=pages_per_slot * page_size, - bmm1_scale=head_dim**-0.5, - bmm2_scale=1.0, - window_left=window_left, - out=out_flat, - sinks=None, - enable_pdl=False, - kv_layout="HND", - backend="trtllm-gen", - q_len_per_req=block_size, - max_q_len=None, - cum_seq_lens_q=None, - kv_cache_sf=None, - uses_shared_paged_kv_idx=True, - bmm1_scale_log2=None, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, - ) - else: - cum_seq_lens_q = torch.arange( - 0, - (B + 1) * block_size, - block_size, - dtype=torch.int32, - device=hidden_states.device, - ) - cum_seq_lens_kv = torch.cat(( - torch.zeros(1, - dtype=torch.int32, - device=hidden_states.device), - seq_lens_after.cumsum(0, dtype=torch.int32), - )) - trtllm_gen_ops.batch_context_with_kv_cache( - query=q_flat, - kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), - workspace_buffer=self._dflash_trtllm_gen_workspace, - block_tables=block_tables, - seq_lens=seq_lens_after, - max_q_len=block_size, - max_kv_len=pages_per_slot * page_size, - bmm1_scale=head_dim**-0.5, - bmm2_scale=1.0, - batch_size=B, - cum_seq_lens_q=cum_seq_lens_q, - cum_seq_lens_kv=cum_seq_lens_kv, - window_left=window_left, - out=out_flat, - sinks=None, - enable_pdl=False, - kv_layout="HND", - kv_cache_sf=None, - uses_shared_paged_kv_idx=True, - causal=False, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, - ) - else: # VANILLA, validated before entering the layer loop. - layer_k_cache = ctx_k_cache[:, layer_idx] - layer_v_cache = ctx_v_cache[:, layer_idx] - - # Pack gqa_group_size query heads sharing a KV head into the - # row dimension: [B, blk, h_q, d] -> [B, group*blk, h_kv, d]. - # Each CTA owns a whole query-head group and streams KV head's context once - # instead of gqa_group_size CTAs each re-reading it. - # Exact only while every row of the block attends to the same - # key set, i.e. non-causal, unwindowed layers. Causal or - # windowed layers mask by row, so they stay unpacked. - pack_gqa = (gqa_group_size > 1 and not causal - and window_size == (-1, -1)) - if pack_gqa: - q_grouped = Q_bshd.reshape(B, block_size, - num_kv_heads_per_rank, - gqa_group_size, head_dim) - q_packed = q_grouped.permute(0, 3, 1, 2, 4) - q_in = q_packed.reshape(B, gqa_group_size * block_size, - num_kv_heads_per_rank, head_dim) - else: - q_in = Q_bshd - out = flash_attention( - q=q_in, - k_cache=layer_k_cache, - v_cache=layer_v_cache, - k=k_noise_bshd, - v=v_noise_bshd, - cache_seqlens=cache_seqlens_i32, - cache_batch_idx=cache_batch_idx_i32, - causal=causal, - window_size=window_size, - ) - if pack_gqa: - # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d]. - out = out.view(B, gqa_group_size, block_size, - num_kv_heads_per_rank, - head_dim).permute(0, 2, 3, 1, 4) - - attn_output = out.reshape(B * block_size, q_size) - - # Per-drafter post-attention gate (no-op for generic DFlash; Laguna - # applies per-head softplus g_proj gating). gate input is the - # input_layernorm output (the attention input). - attn_output = self._post_attention_gate(attn_output, hs_normed_flat, - attn_mod, - num_heads_per_rank, - head_dim) - - # o_proj (flat 2D, handles all-reduce internally) - hidden_out = attn_mod.o_proj(attn_output) - - # Post-attention layernorm + MLP (flat 2D) - res_flat = residual.reshape(-1, residual.shape[-1]) - hidden_out, res_flat = layer.post_attention_layernorm( - hidden_out, res_flat) - hidden_out = layer.mlp(hidden_out) - - hidden_states = hidden_out.reshape(B, block_size, -1) - residual = res_flat.reshape(B, block_size, -1) - - # Final norm - hidden_states_out, _ = self.model.norm( - hidden_states.reshape(-1, hidden_states.shape[-1]), - residual.reshape(-1, residual.shape[-1])) - return hidden_states_out - - def forward( - self, - attn_metadata, - input_ids: torch.LongTensor = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - return_context_logits: bool = False, - spec_metadata=None, - hidden_states: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the draft model and return (hidden_states, hidden_states) for the - speculative-decoding contract.""" - hidden_states_out = self.model( - input_ids=input_ids, - attn_metadata=attn_metadata, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - spec_metadata=spec_metadata, - **kwargs, - ) - - return hidden_states_out, hidden_states_out - - -class DFlashLagunaForCausalLM(DFlashForCausalLM): - """Laguna DFlash drafter. - - The generic block decode lives in DFlashForCausalLM; this subclass supplies - the Laguna draft-layer specifics: per-head g_proj softplus gating and the - per-aux fc_norm applied to captured target features before fc. - """ - - @staticmethod - def _normalize_config(config: PretrainedConfig) -> None: - """Fill TRT-LLM Laguna defaults missing from dense DFlash drafts.""" - if getattr(config, "num_experts", None) is None: - config.num_experts = 0 - if getattr(config, "mlp_layer_types", None) is None: - config.mlp_layer_types = ["dense"] * config.num_hidden_layers - if getattr(config, "block_size", None) is None: - dflash_config = getattr(config, "dflash_config", {}) - if isinstance(dflash_config, dict): - config.block_size = dflash_config.get("block_size", None) - - def __init__(self, - draft_config, - *, - dflash_attention_backend: str = 'VANILLA'): - """Pin the Laguna draft-layer class and enable Laguna-specific behaviors - (context input_layernorm, causal sliding blocks); reject non-per-head - gating.""" - # The checkpoint labels itself with the vLLM name (model_type "llama"); - # remap to the Laguna architecture so TRT-LLM builds the Laguna layers. - draft_config.pretrained_config.architectures = ["LagunaForCausalLM"] - self._normalize_config(draft_config.pretrained_config) - super().__init__( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - self._context_input_layernorm = True - self._sliding_layers_causal = True - gating = getattr(self.config, 'gating', True) - if gating not in (True, 'per-head'): - raise NotImplementedError( - f"Laguna DFlash drafter supports per-head gating only, " - f"got gating={gating!r}") - - def load_weights(self, weights, weight_mapper=None, **kwargs): - """Build the per-aux ``fc_norm`` from the drafter's ``aux_hidden_norms.*`` - weights, then defer the remaining weights to the base loader.""" - aux_keys = sorted( - (k for k in weights if k.startswith('aux_hidden_norms.')), - key=lambda k: int(k.split('.')[1])) - if not aux_keys: - raise ValueError( - "Laguna DFlash checkpoint is missing aux_hidden_norms.* weights" - ) - weights = dict(weights) - eps = getattr(self.config, 'rms_norm_eps', 1e-6) - norms = [] - for k in aux_keys: - w = weights.pop(k) - norm = nn.RMSNorm(w.shape[0], - eps=eps, - device='cuda', - elementwise_affine=True, - dtype=w.dtype) - norm.weight.data.copy_(w) - norms.append(norm) - self.fc_norm = nn.ModuleList(norms) - super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) - - def project_target_hidden(self, hidden_states): - """Project captured target features to the draft width: apply the per-aux - ``fc_norm`` to each hidden chunk, then ``fc`` + ``hidden_norm``.""" - hidden_states = hidden_states.to(self.fc.weight.dtype) - fc_norm = getattr(self, 'fc_norm', None) - if fc_norm is not None: - chunks = hidden_states.chunk(len(fc_norm), dim=-1) - hidden_states = torch.cat( - [norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) - return self.hidden_norm(self.fc(hidden_states)) - - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): - """Apply Laguna's per-head softplus output gate (``g_proj``) to the - attention output; a no-op when the layer has no ``g_proj``.""" - g_proj = getattr(attn_mod, 'g_proj', None) - if g_proj is None: - return attn_output - gate = F.softplus(g_proj(gate_input).float()).to(attn_output.dtype) - return (attn_output.unflatten(-1, (num_heads, head_dim)) * - gate.unsqueeze(-1)).flatten(-2) - - class MTPForCausalLM(nn.Module): def __init__( @@ -2428,86 +1467,110 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens, ) - if spec_config.spec_dec_mode.is_dspark(): + # Only the embedded DSpark draft shares the target's EPLB namespace (its + # stages are target decoder blocks registered into the target's balancer). + # A standalone DSpark drafter is an independent checkpoint, so it falls + # under the "other external drafters" rule above. + if (spec_config.spec_dec_mode.is_dspark() + and spec_config.draft_is_embedded_in_target): kwargs["moe_load_balancer"] = model_config.moe_load_balancer return kwargs +@register_draft_model(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) +def _build_eagle3_one_model_draft(model_config, draft_config, lm_head, model): + """Build the EAGLE3 one-model drafter for the configured draft arch.""" + eagle3_model_arch = model_config.spec_config.eagle3_model_arch + if eagle3_model_arch == "llama3": + # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures + return Eagle3ForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers) + elif eagle3_model_arch == "mistral_large3": + return MistralLarge3EagleForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers, + model.aux_stream_dict) + else: + raise ValueError( + f"Unsupported eagle3 model architecture: {eagle3_model_arch}") + + +@register_draft_model(SpeculativeDecodingMode.MTP) +@register_draft_model(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) +def _build_mtp_one_model_draft(model_config, draft_config, lm_head, model): + """Build the one-model MTP drafter (vanilla MTP and MTP-Eagle share it).""" + return MTPForCausalLM(model_config, + model_config.pretrained_config.num_hidden_layers, + lm_head, model) + + +@register_draft_model(SpeculativeDecodingMode.MTP_EAGLE) +def _build_mtp_eagle_draft(model_config, draft_config, lm_head, model): + """Build the two-model MTP-Eagle drafter.""" + return MTPDraftModelForCausalLM(model_config) + + +@register_draft_model(SpeculativeDecodingMode.PARD) +def _build_pard_draft(model_config, draft_config, lm_head, model): + """Build the PARD drafter.""" + return PARDForCausalLM(draft_config) + + +@register_draft_model(SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL) +def _build_draft_target_one_model_draft(model_config, draft_config, lm_head, + model): + """Build the one-model draft-target drafter from its own checkpoint.""" + # Keep the draft LM head vocab-sharded so greedy draft sampling uses the + # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). + was_frozen = draft_config._frozen + draft_config._frozen = False + draft_config.lm_head_gather_output = False + draft_config._frozen = was_frozen + return AutoModelForCausalLM.from_config(draft_config) + + def get_draft_model(model_config, draft_config, lm_head, model): - """Construct the draft model for the configured speculative-decoding mode - (EAGLE3 / MTP / PARD / DFlash). The DFlash branch selects the Laguna drafter - by detecting its architecture in the draft checkpoint's own config.""" - assert getattr(model_config, 'spec_config', None) is not None - spec_dec_mode = model_config.spec_config.spec_dec_mode - if spec_dec_mode.is_eagle3_one_model(): - if model_config.spec_config.eagle3_model_arch == "llama3": - # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures - return Eagle3ForCausalLM( - draft_config, model_config.pretrained_config.num_hidden_layers) - elif model_config.spec_config.eagle3_model_arch == "mistral_large3": - return MistralLarge3EagleForCausalLM( - draft_config, model_config.pretrained_config.num_hidden_layers, - model.aux_stream_dict) - else: - raise ValueError( - f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}" - ) + """Construct the draft model for the configured speculative-decoding mode. - elif model_config.spec_config.uses_external_draft_model: + Dispatch is registry-based: each mode's builder lives next to the draft + model it constructs and registers itself via ``@register_draft_model``, so + this function never imports a concrete draft implementation (which is what + used to force a lazy import for DSpark, whose provider imports back into + this module through modeling_deepseekv4). + + Args: + model_config: the target engine's ``ModelConfig``, carrying spec_config. + draft_config: the drafter's own ``ModelConfig``, or None when the mode + builds its draft from the target config alone. + lm_head: the target's LM head, shared by the one-model MTP drafter. + model: the target model, for drafters reusing its aux streams. + + Returns: + The draft ``nn.Module`` for this mode. + """ + assert getattr(model_config, 'spec_config', None) is not None + spec_config = model_config.spec_config + spec_dec_mode = spec_config.spec_dec_mode + # An external draft model is loaded straight from its own checkpoint, so it + # has no mode-specific builder to register: this stays an explicit pre-check + # ahead of the registry lookup rather than becoming a registry key. + # + # No mode guard is needed. `uses_external_draft_model` already implies + # `is_mtp_one_model()` (llm_args), which is disjoint from every other mode, + # so this branch cannot divert a drafter that a builder would have claimed. + # Pinned by test_draft_model_registry.py:: + # test_external_draft_model_bypasses_the_registry. + if spec_config.uses_external_draft_model: if draft_config is None: raise ValueError( "MTP speculative decoding with an external draft model requires " "its model config.") return AutoModelForCausalLM.from_config(draft_config) - elif spec_dec_mode.is_mtp_one_model(): - return MTPForCausalLM(model_config, - model_config.pretrained_config.num_hidden_layers, - lm_head, model) - elif spec_dec_mode.is_mtp_eagle(): - return MTPDraftModelForCausalLM(model_config) - elif spec_dec_mode.is_pard(): - return PARDForCausalLM(draft_config) - elif spec_dec_mode.is_dflash(): - draft_arches = getattr(draft_config.pretrained_config, "architectures", - None) or [] - dflash_attention_backend = model_config.spec_config.attention_backend - if any("Laguna" in arch for arch in draft_arches): - return DFlashLagunaForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - return DFlashForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - elif spec_dec_mode.is_dspark(): - # Lazy import to avoid a cycle (modeling_dspark -> modeling_deepseekv4 -> - # modeling_speculative). The DSpark draft reuses the target's aux streams. - # The draft stage count (n_mtp_layers) is not in the HF config, so derive - # it from the checkpoint's mtp.* namespace. - from .modeling_dspark import (DSparkForCausalLM, count_dspark_stages, - validate_dspark_eplb_layer_base) - num_stages = count_dspark_stages( - model_config.spec_config.speculative_model) - validate_dspark_eplb_layer_base(model_config, draft_config) - return DSparkForCausalLM( - draft_config, - getattr(model, "aux_stream_dict", None), - num_stages=num_stages, - block_size=model_config.spec_config.block_size, - ) - elif spec_dec_mode.is_draft_target_one_model(): - # Keep the draft LM head vocab-sharded so greedy draft sampling uses the - # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). - was_frozen = draft_config._frozen - draft_config._frozen = False - draft_config.lm_head_gather_output = False - draft_config._frozen = was_frozen - return AutoModelForCausalLM.from_config(draft_config) - else: + builder = get_registered_draft_model_builder(spec_dec_mode) + if builder is None: raise NotImplementedError( f"get_draft_model does not support speculative decoding mode {spec_dec_mode}." ) + return builder(model_config, draft_config, lm_head, model) class SpecDecOneEngineForCausalLM(DecoderModelForCausalLM[TModel, TConfig], diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1a140ea15e03..16f621397a49 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -8,8 +8,8 @@ import os import time from dataclasses import dataclass -from typing import (Any, Dict, Generic, Iterator, List, Literal, Optional, - Tuple, Type, TypeVar, Union) +from typing import (Any, Callable, Dict, Generic, Iterator, List, Literal, + Optional, Tuple, Type, TypeVar, Union) import torch from torch import nn @@ -33,7 +33,8 @@ from ..modules.logits_processor import LogitsProcessor from ..modules.rms_norm import RMSNorm from ..speculative import SpecMetadata -from ._arch_index import MODEL_ARCH_TO_MODULE, is_builtin_zoo_module +from ._arch_index import (MODEL_ARCH_TO_MODULE, SPEC_MODE_TO_MODULE, + is_builtin_zoo_module) @contextlib.contextmanager @@ -897,6 +898,7 @@ def infer_max_seq_len(self) -> int: MODEL_CLASS_MAPPING = {} +DRAFT_MODEL_BUILDER_MAPPING = {} MODEL_CLASS_VISION_ENCODER_MAPPING = {} MODEL_CLASS_MAPPER_MAPPING = {} MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING = {} @@ -999,6 +1001,120 @@ def get_registered_model_class(model_arch: str) -> Optional[Type[nn.Module]]: return MODEL_CLASS_MAPPING.get(model_arch) +def _is_builtin_draft_model_builder(builder) -> bool: + return is_builtin_zoo_module(getattr(builder, "__module__", "")) + + +# Speculative-decoding modes each decorated builder declared via +# ``register_draft_model``, kept per function (``__dict__``, never inherited). +# Recorded even when the builder loses its ``DRAFT_MODEL_BUILDER_MAPPING`` slot +# to an external registration, so a builder can be mapped back to its modes +# without scanning the mapping by identity (which would silently skip exactly +# those overridden built-ins). +_REGISTERED_SPEC_MODES_ATTR = "_registered_spec_modes" + + +def register_draft_model(mode): + """Register the draft-model builder for a speculative-decoding mode. + + The builder is a plain function + ``(model_config, draft_config, lm_head, model) -> nn.Module`` that owns + everything its mode needs to construct its draft model, so the generic + dispatcher never has to import a concrete draft implementation. Stack the + decorator to serve several modes with one builder (vanilla MTP and + MTP_EAGLE_ONE_MODEL share theirs, mirroring + ``SpeculativeDecodingMode.is_mtp_one_model()``). + + Same registration priority as ``register_auto_model``: built-in builders + only fill empty slots and never overwrite, because under lazy loading a + built-in module may run its decorators *after* an external registration + (e.g. a drafter supplied through ``--custom_module_dirs``) and must not + clobber it. External registrations always overwrite. + + The builder belongs in the ``modeling_*.py`` that defines the draft model, + never in the factory file, and needs an ``SPEC_MODE_TO_MODULE`` row in + ``_arch_index.py`` so it can be found without importing the zoo -- see the + "Adding a speculative decoding mode" notes there. Example (DSpark, whose + stage count is only knowable from the checkpoint):: + + @register_draft_model(SpeculativeDecodingMode.DSPARK) + def _build_dspark_draft(model_config, draft_config, lm_head, model): + num_stages = count_dspark_stages( + model_config.spec_config.speculative_model) + validate_dspark_eplb_layer_base(model_config, draft_config) + return DSv4DSparkForCausalLM( + draft_config, + getattr(model, "aux_stream_dict", None), + num_stages=num_stages, + block_size=model_config.spec_config.block_size, + ) + + Args: + mode: the ``SpeculativeDecodingMode`` member this builder serves. + + Returns: + The decorator binding a builder function to ``mode``. + """ + + def decorator(builder): + modes = builder.__dict__.get(_REGISTERED_SPEC_MODES_ATTR) + if modes is None: + modes = set() + setattr(builder, _REGISTERED_SPEC_MODES_ATTR, modes) + modes.add(mode) + + existing = DRAFT_MODEL_BUILDER_MAPPING.get(mode) + if (existing is not None and existing is not builder + and _is_builtin_draft_model_builder(builder)): + logger.info( + f"Keeping existing draft-model builder " + f"{existing.__module__}.{existing.__qualname__} for " + f"speculative decoding mode {mode.name}; built-in " + f"{builder.__module__}.{builder.__qualname__} not registered.") + return builder + DRAFT_MODEL_BUILDER_MAPPING[mode] = builder + return builder + + return decorator + + +def _ensure_draft_model_registered(mode) -> None: + """Import the module providing ``mode``'s builder, if not yet loaded. + + Mirrors ``_ensure_model_registered``: builders register as an import side + effect, and the model zoo is imported lazily, so this turns a mode into + "the decorator has run". Modes missing from the static index are left to + the caller's normal unsupported-mode handling. + """ + module_name = SPEC_MODE_TO_MODULE.get(mode.name) + if module_name is None: + return + full_name = f"tensorrt_llm._torch.models.{module_name}" + try: + importlib.import_module(full_name) + except ModuleNotFoundError as e: + # Only swallow "the providing module itself is missing" (stale index + # entry); a missing dependency *inside* the module is a real error and + # must not be masked as "unsupported speculative decoding mode". + if e.name != full_name: + raise + logger.warning(f"Lazy import of {module_name} for speculative " + f"decoding mode {mode.name} failed: {e!r}") + + +def get_registered_draft_model_builder(mode) -> Optional[Callable]: + """Resolve ``mode`` to its registered draft-model builder, or ``None``. + + The single entry point for builder lookups: the model zoo is imported + lazily, so this resolves the providing module on demand before reading the + registry. Do not read ``DRAFT_MODEL_BUILDER_MAPPING`` directly — a raw + ``.get()`` silently misses every not-yet-imported provider. + """ + if mode not in DRAFT_MODEL_BUILDER_MAPPING: + _ensure_draft_model_registered(mode) + return DRAFT_MODEL_BUILDER_MAPPING.get(mode) + + def get_registered_vision_encoder( model_arch: str) -> Optional[Tuple[Type[nn.Module], Optional[Type]]]: """Resolve ``model_arch`` to its ``(vision_encoder_cls, vlm_base_model)``. diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 14644ff3e774..b92c1a7c7ac2 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -146,7 +146,13 @@ def is_layer_capture(self, layer_id: int) -> bool: def maybe_capture_hidden_states( self, layer_id: int, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None ) -> None: - """Capture hidden states from a target model layer into the buffer.""" + """Capture hidden states from a target model layer into the buffer. + + The ``residual`` convention is model-specific: Qwen3/Llama-style callers + pass the pre-add pair and this folds them, while K3 hands in an + already-mixed aggregated stream value and passes ``residual=None`` + (see the DSpark tap in ``modeling_kimi_linear.py``). + """ if self.captured_hidden_states is None: return i = self._layer_to_idx.get(layer_id) @@ -239,8 +245,12 @@ def __init__( f"DFlash: acceptance-statistics recording enabled -> {self._accept_stats.path}" ) + # type(self), not a literal: DSparkWorker does not override + # __init__, so a hardcoded name reports the base class and the log + # cannot evidence which worker a run actually used. logger.info( - f"DFlashWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}" + f"{type(self).__name__} initialized with " + f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) @property @@ -314,11 +324,12 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): self.max_draft_len + 1 ) # what the draft forward actually computes - self._compute_block_size = self.max_draft_len + 1 - if self.max_draft_len + 1 > self._resolved_block_size: + self._compute_block_size = self._draft_block_width(draft_model) + if self._compute_block_size > self._resolved_block_size: + slack = self._compute_block_size - self.max_draft_len raise ValueError( f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." - f"Lower max_draft_len to at most {self._resolved_block_size - 1}." + f"Lower max_draft_len to at most {self._resolved_block_size - slack}." f"Current max_draft_len={self.max_draft_len}." ) @@ -698,13 +709,13 @@ def _forward_impl( # Gather K logits per gen request from the block outputs. # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. - # Plain DFlash reads mask slots 1..K; dspark shift_label reads - # slots 0..K-1 (see dflash_draft_slot_ids). + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. block_size = self._compute_block_size - shift_label = getattr(draft_model, "_dspark_shift_label", False) - gen_gather_ids = dflash_draft_slot_ids( - num_gens, block_size, K, shift_label, device="cuda" - ) + gen_gather_ids = self._draft_slot_ids(draft_model, num_gens, block_size, K) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) gen_logits = draft_model.logits_processor( @@ -714,12 +725,9 @@ def _forward_impl( vocab_size = gen_logits.shape[-1] gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - # DSpark Markov head: add the greedy-chained intra-block - # logit bias before sampling (no-op for plain DFlash). - if getattr(draft_model, "has_markov_head", False): - gen_logits = self._apply_dspark_markov_bias( - draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata - ) + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) gen_draft_tokens = self.sample_draft_tokens( gen_logits, @@ -791,60 +799,43 @@ def _forward_impl( "next_new_tokens": next_new_tokens, } - def _apply_dspark_markov_bias( + def _draft_block_width(self, draft_model) -> int: + """Block slots the draft forward must compute for K draft tokens. + + Plain DFlash reads slots 1..K, so slot 0 is pure overhead and the + forward needs K+1. Families whose slot convention differs override + this alongside :meth:`_draft_slot_ids` — the two must agree, since + the width bounds the slots the gather is allowed to name. + """ + return self.max_draft_len + 1 + + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots whose hidden states produce the K draft logits. + + Plain DFlash uses the K2.7 convention: mask slots 1..K. Drafter + families with another convention override this — see + :meth:`DSparkWorker._draft_slot_ids` for the shift_label + variant, which reads slots 0..K-1 instead. + """ + return dflash_draft_slot_ids(num_gens, block_size, num_draft_tokens, False, device="cuda") + + def _refine_block_logits( self, draft_model, gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, + inputs: dict, spec_metadata, ) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): - step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head - logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} - is the greedy token from step i-1's biased logits. Greedy per-position - argmax of the returned logits therefore reproduces the reference - sampled chain; the rejection-sampling path samples from the same - biased distributions (proposal conditioned on the greedy chain). - - Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows - to this rank's contiguous shard and chaining through the TP-aware - global argmax. + """Refine the block logits between the draft forward and sampling. + + Plain DFlash proposes the backbone's logits unchanged. Drafter + families carrying extra heads override this — see + :meth:`DSparkWorker._refine_block_logits` for the Markov + intra-block bias. """ - if self._d2t is not None: - raise NotImplementedError( - "DSpark Markov head requires a shared draft/target vocab " - "(d2t vocab mapping is not supported)." - ) - full_vocab = draft_model.markov_w2.shape[0] - shard = gen_logits.shape[-1] - vocab_slice = None - if shard != full_vocab: - mapping = self.mapping - if ( - mapping is None - or getattr(mapping, "enable_attention_dp", False) - or shard * mapping.tp_size != full_vocab - ): - raise NotImplementedError( - f"DSpark Markov head: draft logits width {shard} does not " - f"match the drafter vocab {full_vocab} and is not a plain " - "TP column shard of it." - ) - vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) - - def argmax_fn(step_logits): - # Full-vocab token ids (TP-aware when sharded); tokens stay in - # draft-vocab space, which is what markov_w1 indexes. - return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() - - return draft_model.apply_markov_chain_logits( - gen_logits, - first_prev_tokens, - argmax_fn=argmax_fn, - vocab_slice=vocab_slice, - ) + return gen_logits def prepare_1st_drafter_inputs( self, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 6510f565fea9..733b44515ea8 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -17,7 +17,7 @@ # hidden states, accept the previous block with standard verification, draft a # new block in one backbone forward), adapted to DSpark's draft model which # produces the whole block (and its confidence-truncated length) inside a single -# ``DSparkDraftModel.forward`` rather than via mask-token cross-attention. +# ``DSv4DSparkDraftModel.forward`` rather than via mask-token cross-attention. from collections import deque from dataclasses import dataclass @@ -30,6 +30,7 @@ from tensorrt_llm.mapping import Mapping from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +from .dflash import DFlashWorker, dflash_draft_slot_ids from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: @@ -44,7 +45,7 @@ class DSparkSpecMetadata(SpecMetadata): the target forward pass. DSpark captures the *mean over the multi-head (mHC) residual streams* at each captured layer (handled by the target-side capture hook), concatenated across layers, and feeds them to the draft - model's ``main_proj`` + ``main_norm`` (inside ``DSparkDraftModel.forward``) + model's ``main_proj`` + ``main_norm`` (inside ``DSv4DSparkDraftModel.forward``) as the captured-context attention input (``main_x``). Mirrors :class:`DFlashSpecMetadata`; the only DSpark-specific detail is that @@ -184,11 +185,11 @@ def get_hidden_states(self, num_tokens: int) -> Optional[torch.Tensor]: ] -class DSparkWorker(SpecWorkerBase): +class DSv4DSparkWorker(SpecWorkerBase): """Worker for DSpark speculative decoding. DSpark drafts a whole block of ``block_size`` tokens in one backbone forward - (``DSparkDraftModel.forward``): it projects the captured target-layer hidden + (``DSv4DSparkDraftModel.forward``): it projects the captured target-layer hidden states (``main_proj`` + ``main_norm``) into the draft's captured-context attention, runs the ``num_stages`` DSpark blocks over a rolling captured window, refines the per-position logits with the Markov head, and predicts a @@ -205,10 +206,18 @@ class DSparkWorker(SpecWorkerBase): The rolling window is kept consistent across the whole decode: it is seeded from the prompt's captured context at prefill and back-filled with the intermediate accepted tokens of a multi-accept step (both via - ``DSparkDraftModel.write_context_windows``), in addition to the per-step bonus + ``DSv4DSparkDraftModel.write_context_windows``), in addition to the per-step bonus write done by the generation path. These affect draft acceptance rate only, not correctness, which the standard target verify guarantees. + Naming: workers are classified by *deployment form*, not by draft + backbone (see :class:`DSparkWorker`). This one is form-specific + because it owns a rolling captured-context window and drives the draft + through attributes only an embedded DeepSeek-V4-Pro draft has -- + ``num_stages``, ``_attn_params``, ``write_context_windows``, + ``write_context_windows_batched`` and ``forward_batched``. A standalone + drafter has none of them and is served by :class:`DSparkWorker`. + Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). """ @@ -245,14 +254,14 @@ def __init__( self._scratch_slot = 0 # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + + # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward_batched`` + # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe # to capture into the target's CUDA graph (DSpark is a one-engine drafter — # its worker forward runs inside that graph, so the draft path MUST be # capture-safe whenever ``cuda_graph_config`` is set). logger.info( - f"DSparkWorker initialized with " + f"DSv4DSparkWorker initialized with " f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) @@ -414,7 +423,7 @@ def _draft_gen_block_batched( (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept back-fill) are gathered as tensors, slots come from the host-built ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSparkDraftModel.forward_batched``. Returns the per-position corrected + ``DSv4DSparkDraftModel.forward_batched``. Returns the per-position corrected block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. Confidence truncation stays disabled — the full block is proposed. @@ -696,3 +705,158 @@ def _forward_impl( "next_draft_tokens": next_draft_tokens, "next_new_tokens": next_new_tokens, } + + +class DSparkWorker(DFlashWorker): + """Worker for a *standalone* DSpark drafter (DFlash lineage). + + DSpark is DFlash plus two extra heads, so the drafting plumbing is + inherited wholesale from :class:`DFlashWorker` -- paged context K/V, + slot management, the mask-token block forward -- and only the two + head-driven policies are overridden here: the block-output slot + convention (``shift_label``) and the Markov intra-block logit bias. + + Mirrors the model side, where ``GQADSparkForCausalLM`` extends + ``DFlashForCausalLM`` with the same two heads. + + Naming: this is the unqualified DSpark worker because a separately + shipped drafter is the ordinary case; :class:`DSv4DSparkWorker` carries + the qualifier because a draft embedded in the target checkpoint is the + special one. Workers are classified by *deployment form*, never by draft + backbone -- so there is no ``Qwen3DSparkWorker``. Note the name meant the + embedded worker before this split; both the rebind and the rename to + ``DSv4DSparkWorker`` land in one commit so the swap reads as a unit. + + A worker is agnostic to the draft backbone: everything backbone-shaped is + supplied by the draft model, which reports its own shapes + (``_num_attn_layers``, ``_num_heads``, ``_num_kv_heads``, ``_head_dim``) + and owns the operators (``_build_fused_kv_buffers``, + ``precompute_context_kv``, ``dflash_forward``, + ``apply_markov_chain_logits``, ``project_target_hidden``). The worker only + allocates against the reported shapes and sequences the calls. An MLA + drafter therefore reuses this class unchanged; its differences (fused-QKV + assumptions, a 576-latent K/V layout) land in its own draft-model + subclass. Naming workers by backbone would produce N classes with + identical bodies. + + Deployment form is the axis the runtime state actually splits on: paged + draft K/V here, a worker-owned rolling window in + :class:`DSv4DSparkWorker`. + """ + + def set_draft_model(self, draft_model) -> None: + """Reject an unsupported vocab mapping here rather than mid-decode. + + ``d2t`` is model-static, so a config mistake should surface at load and + not as a ``NotImplementedError`` raised per decode step, possibly during + CUDA-graph capture. + """ + super().set_draft_model(draft_model) + if self._d2t is not None and getattr(draft_model, "has_markov_head", False): + raise NotImplementedError( + "DSpark Markov head requires a shared draft/target vocab " + "(d2t vocab mapping is not supported); drafter " + f"{type(draft_model).__name__} declares one." + ) + + def _draft_block_width(self, draft_model) -> int: + """Block width under the dspark ``shift_label`` convention. + + shift_label reads slots 0..K-1, so K draft tokens fit in K slots and + the base class' K+1 over-demands by one -- enough to reject a block-7 + checkpoint at max_draft_len=7, which is how both published DSpark + drafters are meant to run. + """ + if getattr(draft_model, "_dspark_shift_label", False): + return self.max_draft_len + return super()._draft_block_width(draft_model) + + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots under the dspark ``shift_label`` convention. + + The drafter checkpoint declares the convention, so it is read off the + draft model rather than assumed: a DSpark drafter trained with the + legacy DFlash slot layout keeps the base class' slots 1..K. + """ + shift_label = getattr(draft_model, "_dspark_shift_label", False) + return dflash_draft_slot_ids( + num_gens, block_size, num_draft_tokens, shift_label, device="cuda" + ) + + def _refine_block_logits( + self, + draft_model, + gen_logits: torch.Tensor, + inputs: dict, + spec_metadata, + ) -> torch.Tensor: + """Add the greedy-chained Markov intra-block bias to the block logits. + + A DSpark drafter checkpoint may omit the Markov head (``markov_rank`` + 0), which loads as a drafter without one; that case falls through to + the unmodified backbone logits. + """ + if not getattr(draft_model, "has_markov_head", False): + return gen_logits + return self._apply_dspark_markov_bias( + draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata + ) + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows + to this rank's contiguous shard and chaining through the TP-aware + global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 313a081b87de..4084fbc3ffc3 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -117,9 +117,14 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool: return False if spec_config._use_shared_kv_cache: return False - # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft - # model does not read the paged draft KV cache managed by attention metadata. - if spec_config.spec_dec_mode.is_dspark(): + # The embedded DSpark draft owns a dedicated rolling-window cache in + # DSv4DSparkWorker and never reads the paged draft KV cache that attention + # metadata manages. A standalone DSpark drafter runs on DSparkWorker + # (DFlash lineage), which does read it, so it keeps the default -- hence a + # form check, not a mode check + # (see DSparkDecodingConfig.draft_is_embedded_in_target). + if (spec_config.spec_dec_mode.is_dspark() + and spec_config.draft_is_embedded_in_target): return False return spec_config._allow_separate_draft_kv_cache diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 488d0412548b..8efbac572925 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -21,7 +21,7 @@ from .dflash import DFlashSpecMetadata, DFlashWorker from .draft_target import (DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .dspark import DSparkSpecMetadata, DSparkWorker +from .dspark import DSparkSpecMetadata, DSparkWorker, DSv4DSparkWorker from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata, MTPEagleWorker) @@ -465,7 +465,13 @@ def _build_spec_metadata(spec_config, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, ) - if spec_config.spec_dec_mode.is_dflash(): + # A standalone DSpark drafter is drafted by DFlashWorker, so it needs the + # DFlash metadata (paged draft KV, DFlash capture buffer). Only the + # embedded DeepSeek-V4-Pro draft uses DSparkSpecMetadata and its rolling + # window. See DSparkDecodingConfig.draft_is_embedded_in_target. + if spec_config.spec_dec_mode.is_dflash() or ( + spec_config.spec_dec_mode.is_dspark() + and not spec_config.draft_is_embedded_in_target): target_layer_ids = getattr(spec_config, 'target_layer_ids', None) return DFlashSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -790,7 +796,16 @@ def get_spec_worker(spec_config, return PARDWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_dflash(): return DFlashWorker(spec_config, mapping, use_separate_draft_kv_cache) + # DSpark splits by deployment form, mirroring the draft-model side. The + # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose rolling-window + # plumbing reads V4-draft-only attributes (num_stages, write_context_windows, + # forward_batched). A standalone drafter is DFlash lineage and is served by + # DSparkWorker, which adds only the Markov bias and the shift_label + # slot convention on top of DFlashWorker. if spec_dec_mode.is_dspark(): + if spec_config.draft_is_embedded_in_target: + return DSv4DSparkWorker(spec_config, mapping, + use_separate_draft_kv_cache) return DSparkWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_sa(): return SAWorker(spec_config, model_config) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index d68aad84b4d7..bfc487767835 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -834,7 +834,8 @@ def __init__(self, output_path: Optional[str] = None, output_dir: Optional[str] = None, post_process_fn: Optional[Callable[[str], str]] = None, - preserve_caller_max_tokens: bool = False): + preserve_caller_max_tokens: bool = False, + num_fewshot: Optional[int] = None): try: import lm_eval except ImportError as e: @@ -889,6 +890,15 @@ def _adjust_config(task_dict, random_seed): else: # NOTE: Few-shot random seed task_obj.set_fewshot_seed(seed=random_seed) + # Caller override of the task yaml's shot count, the same + # call lm-eval's own simple_evaluate makes. Without it a + # 0-shot chat evaluation of a task whose yaml pins 5 shots + # is unreachable, which is the regime a chat-distilled + # speculative drafter has to be measured in. + if num_fewshot is not None: + task_obj.set_config(key="num_fewshot", + value=num_fewshot) + logger.info(f"num_fewshot overridden to {num_fewshot}") adjusted_task_dict[task_name] = task_obj # NOTE: Shuffle dataset @@ -1054,6 +1064,7 @@ def command_harness(cls, ctx, **kwargs): random_seed=kwargs.pop("random_seed", 0), apply_chat_template=kwargs.pop("apply_chat_template", False), fewshot_as_multiturn=kwargs.pop("fewshot_as_multiturn", False), + num_fewshot=kwargs.pop("num_fewshot", None), system_prompt=kwargs.pop("system_prompt", None), is_multimodal=kwargs.pop("is_multimodal", False), chat_template_kwargs=kwargs.pop("chat_template_kwargs", None), @@ -1129,6 +1140,12 @@ def __init__(self, **kwargs): is_flag=True, default=False, help="Apply fewshot as multiturn.") + @click.option("--num_fewshot", + type=int, + default=None, + help="Override the task yaml's shot count. Use 0 with " + "--apply_chat_template for a single-question chat " + "evaluation.") @click.option("--system_prompt", type=str, default=None, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3b95780bd822..a149f7eddcce 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2965,6 +2965,19 @@ class DSparkDecodingConfig(DecodingBaseConfig): decoding_type: Literal["DSpark"] = Field(default="DSpark") + attention_backend: Literal["VANILLA", "TRTLLM"] = Field( + default="VANILLA", + description= + "Attention backend for the pooled-context cross-attention of a " + "standalone DSpark drafter (one shipped as its own checkpoint rather " + "than inside the target's mtp.* namespace). Ignored by the embedded " + "DeepSeek-V4-Pro draft, which uses its own captured-context attention. " + "This is independent of the backend used to construct the drafter's " + "standard attention modules. TRTLLM requires FlashInfer and an NVIDIA " + "Blackwell GPU with SM100 or SM103, and uses generated FMHA kernels " + "with a private paged context cache; VANILLA uses FlashAttention with " + "a contiguous cache.") + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -2982,6 +2995,62 @@ def tokens_per_gen_step(self) -> int: def supports_backend(self, backend: str) -> bool: return backend == "pytorch" + @functools.cached_property + def draft_is_embedded_in_target(self) -> bool: + """True for the embedded (DeepSeek-V4-Pro) flavour of the DSpark draft. + + DSpark ships in two shapes, and they need different runtime plumbing: + + - embedded: the draft is the ``mtp.*`` namespace of the *target* + checkpoint, built from full target decoder blocks, and served by + ``DSv4DSparkWorker`` with its own rolling captured-context window. + - standalone: the draft is its own checkpoint with a registry-resolved + backbone, served by ``DFlashWorker`` and its paged draft KV cache. + + Both are ``decoding_type: DSpark``, so every dispatch that must tell + them apart -- draft-model builder, worker, spec metadata, and the + separate-draft-KV-cache decision -- reads this one flag instead of + re-deriving it. That is what keeps those decisions from drifting apart: + a builder and a worker that disagree produce a draft model whose + attributes the worker does not have. + + The probe is the weight index rather than a config field because the + index is authoritative and cannot be left unset; ``model_type`` is the + fallback for a checkpoint whose index file is absent. Resolution is + memoized here and warmed during ``TorchLlmArgs`` validation, so the + filesystem probe happens once in the main process -- not per rank, and + never at forward or CUDA-graph-capture time. + """ + ckpt_dir = self.speculative_model + if ckpt_dir is None: + return False + ckpt_dir = str(ckpt_dir) + + for name in ("model.safetensors.index.json", + "pytorch_model.bin.index.json"): + index = os.path.join(ckpt_dir, name) + if not os.path.isfile(index): + continue + try: + with open(index, encoding="utf-8") as f: + weight_map = json.load(f).get("weight_map", {}) + except (OSError, ValueError): + break + # An index that parsed is authoritative both ways. Falling through + # to model_type here would classify a standalone V4-shaped drafter + # as embedded, and that only surfaces much later, inside + # count_dspark_stages. + return any(re.match(r"^mtp\.\d+\.", key) for key in weight_map) + + config_path = os.path.join(ckpt_dir, "config.json") + if os.path.isfile(config_path): + try: + with open(config_path, encoding="utf-8") as f: + return json.load(f).get("model_type") == "deepseek_v4" + except (OSError, ValueError): + return False + return False + @functools.cached_property def spec_dec_mode(self): from tensorrt_llm._torch.speculative.interface import \ @@ -5985,34 +6054,55 @@ def validate_speculative_config(self): if not spec_cfg.max_draft_len: raise ValueError("DSpark max_draft_len must be > 0; got " f"{spec_cfg.max_draft_len}") - # The DSpark draft weights live in the ``mtp.*`` namespace of a - # local checkpoint directory; without ``speculative_model`` - # neither the draft weights nor the ``dspark_*`` config - # defaults can be located, and engine construction would fail - # much later with an opaque error. + # Same convention as MTP (see resolve_mtp_checkpoint_source): + # an unset speculative_model means "the draft lives in the + # target checkpoint". For DSpark that is the embedded + # DeepSeek-V4-Pro flavour, whose draft is the target's mtp.* + # namespace. Defaulting rather than rejecting keeps the + # spec-dec API consistent across algorithms; pointing + # speculative_model at the target explicitly is equivalent. if spec_cfg.speculative_model is None: - raise ValueError( - "DSpark requires speculative_config.speculative_model " - "to point at the checkpoint directory containing the " - "mtp.* draft weights (for DeepSeek-V4-Pro-DSpark this " - "is the target checkpoint directory itself).") + spec_cfg.speculative_model = self.model + if not spec_cfg.draft_is_embedded_in_target: + raise ValueError( + "speculative_config.speculative_model is unset, " + "which means 'load the draft from the target " + "checkpoint', but the target has no mtp.* draft " + "weights. Point speculative_model at a standalone " + "DSpark drafter checkpoint directory.") + # Warm the embedded-vs-standalone probe here, while we are in + # the main process and the checkpoint path is known local. The + # flag then travels with the config, so no rank repeats the + # filesystem read and nothing probes at forward time. + _ = spec_cfg.draft_is_embedded_in_target # Resolve target_layer_ids / mask_token_id / block_size / # markov_rank from the draft (or main) model config if not set. - # DSpark ships these as top-level ``dspark_*`` keys in the - # DeepSeek-V4-Pro config.json; also accept a nested - # ``dspark_config`` dict for forward compatibility. + # Three checkpoint spellings are in the wild for the same knobs + # and all are accepted here, because a key the reader misses is + # not an error -- it silently falls back to a default and the + # drafter degrades (a markov_rank read as 0 skips the Markov + # head entirely, costing acceptance with no warning): + # - top-level ``dspark_*`` (DeepSeek-V4-Pro) + # - nested ``dflash_config`` (SpecForge / RadixArk drafters) + # - nested ``dspark_config`` (forward compatibility) + # - plain top-level keys (TorchSpec drafters) draft_config_path = os.path.join(spec_cfg.speculative_model, "config.json") if os.path.exists(draft_config_path): with open(draft_config_path) as f: draft_cfg = json.load(f) - dspark_cfg = draft_cfg.get("dspark_config", {}) + dspark_cfg = draft_cfg.get("dspark_config") or {} + dflash_cfg = draft_cfg.get("dflash_config") or {} def _dspark_get(key, top_level_key): - value = dspark_cfg.get(key) - if value is None: - value = draft_cfg.get(top_level_key) - return value + for source, name in ((dspark_cfg, key), (dflash_cfg, + key), + (draft_cfg, + top_level_key), (draft_cfg, key)): + value = source.get(name) + if value is not None: + return value + return None # The checkpoint's ``dspark_target_layer_ids`` is # authoritative: it fixes both which target hidden states diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 9d19aa27734f..fe695e2210f6 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -39,3 +39,6 @@ TestQwen3_5_4B::test_dflash: TestKimiK3::test_w4a16_mxfp4: ref_al: 1.318 min_al: 1.15 +TestQwen3_8B::test_dspark: + ref_al: 6.259 + min_al: 5.946 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 0d98e20d649b..1a5d5adde3d5 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -192,6 +192,15 @@ Qwen3/Qwen3-8B: accuracy: 87.1114 - spec_dec_algo: DFlash accuracy: 87.1114 + # accuracy 0 means "run the task, do not gate on it" -- the same value + # TRTLLM_ACCURACY_NO_REFERENCE and INTEGRATION_TEST use, and the threshold + # (ref + z_alpha*scale) is then always cleared. test_dspark gates on + # acceptance length instead. Accuracy is not meaningful in its 0-shot chat + # regime: the model stops emitting "#### N", so strict-match is 0 and the + # reported score is the mean of that and flexible-extract (28.92 measured). + - spec_dec_algo: DSpark + extra_acc_spec: zero_shot_al_only + accuracy: 0 - quant_algo: FP8 kv_cache_quant_algo: FP8 accuracy: 87.1114 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index b5cb009adf27..89d4ac7d590f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5061,6 +5061,73 @@ def test_dflash(self): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_hopper + @pytest.mark.parametrize( + "attention_backend", + ["VANILLA", pytest.param("TRTLLM", marks=skip_pre_blackwell)]) + def test_dspark(self, attention_backend): + """Standalone DSpark drafter on Qwen3-8B, 0-shot chat. + + Acceptance length is the gate here. DSpark drafters are distilled on + the target's own chat-mode generations, so the harness default (5-shot + completion) is out of distribution for the drafter and understates + acceptance: 4.42 there against 6.26 here, on the same checkpoints. + GSM8K accuracy is not meaningful in this regime and is not gated -- + see the extra_acc_spec entry in references/gsm8k.yaml. + + Both block-decode backends run: TRTLLM is what deployments use, and + it is the one the acceptance number above was measured on, but it + needs SM100/SM103 so H100 only covers VANILLA. The two differ + numerically (GSM8K 28.9 vs 28.1) while landing the same acceptance + length, so they share one reference. + """ + pytorch_config = dict( + max_batch_size=8, + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(max_batch_size=8, + enable_padding=True), + ) + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.6) + + # DeepSeek's official DSpark head for this target (block_size 7, + # markov_rank 256, confidence head). Head tensors are named after the + # submodules that own them (markov_head.*, confidence_head.proj.*), + # which is the spelling the drafter loader has to resolve. + dspark_model_dir = ( + f"{llm_models_root()}/dspark/dspark_qwen3_8b_block7") + target_model_dir = f"{llm_models_root()}/Qwen3/Qwen3-8B" + + spec_config = DSparkDecodingConfig(max_draft_len=7, + speculative_model=dspark_model_dir, + attention_backend=attention_backend) + + with LLM(model=target_model_dir, + **pytorch_config, + kv_cache_config=kv_cache_config, + max_stats_len=-1, + enable_iter_perf_stats=True, + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + # 0-shot chat, not the harness default 5-shot completion: a DSpark + # drafter is distilled on the target's chat-mode output, so the + # default prompt is out of distribution for it and understates + # acceptance (4.42 vs 6.26 on these same checkpoints). Only the AL + # is gated -- extra_acc_spec selects a gsm8k.yaml entry whose + # reference accuracy is 0, i.e. run the task, do not gate on it. + task.evaluate(llm, + extra_acc_spec="zero_shot_al_only", + extra_evaluator_kwargs=dict( + num_fewshot=0, + apply_chat_template=True, + chat_template_kwargs={"enable_thinking": False}, + )) + acceptance_length = _compute_acceptance_length(llm) + print(f"[AL] test_dspark[{attention_backend}] acceptance_length " + f"= {acceptance_length:.3f}") + assert_acceptance_length("TestQwen3_8B::test_dspark", + acceptance_length) + @skip_pre_blackwell @pytest.mark.parametrize("tp_size,pp_size,ep_size,attention_dp", [(1, 1, 1, False)], diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a4a87cd5f8c0..9e5ef4da23b1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -52,6 +52,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-trtllm-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402 + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark[TRTLLM] # SM100+ only; l0_h100 runs [VANILLA] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_trtllm-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-TRTLLM] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[mxfp8-latency-TRTLLM] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index d4da554eaffb..800c465dd6e6 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -169,6 +169,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark[VANILLA] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 diff --git a/tests/unittest/_torch/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index db9aaa2782a1..bb3236abe243 100644 --- a/tests/unittest/_torch/modeling/test_modeling_speculative.py +++ b/tests/unittest/_torch/modeling/test_modeling_speculative.py @@ -25,8 +25,8 @@ from tensorrt_llm._torch.attention_backend.interface import RopeParams from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM from tensorrt_llm._torch.models.modeling_speculative import ( - DFlashForCausalLM, Eagle3ForCausalLM, SpecDecOneEngineForCausalLM, ) @@ -301,7 +301,7 @@ def test_dflash_attention_mask_args(): assert wrapper._get_attention_mask_args(1) == (False, (-1, -1)) assert wrapper._get_attention_mask_args(2) == (True, (4095, 0)) - with patch("tensorrt_llm._torch.models.modeling_speculative.logger.warning") as warning: + with patch("tensorrt_llm._torch.models.modeling_dflash.logger.warning") as warning: wrapper._warn_inferred_attention_windows() warning.assert_not_called() @@ -344,7 +344,7 @@ def test_dflash_attention_mask_args(): for layer_idx in range(5): assert laguna_wrapper._get_attention_mask_args(layer_idx) == (True, (511, 0)) - with patch("tensorrt_llm._torch.models.modeling_speculative.logger.warning") as warning: + with patch("tensorrt_llm._torch.models.modeling_dflash.logger.warning") as warning: laguna_wrapper._warn_inferred_attention_windows() warning.assert_called_once_with( "DFlash inferred pooled-context sliding-window attention from checkpoint " diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index b73fadde3e5b..40a841f3656f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -27,16 +27,20 @@ import torch import torch.nn.functional as F -import tensorrt_llm._torch.models.dspark.attention as dspark_attention import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( + DSv4DSparkDraftModel, apply_dspark_rotary, dspark_attention_forward, dspark_sparse_attn, get_dspark_topk_idxs, precompute_dspark_freqs_cis, ) -from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel + +# The captured-context attention primitives were folded into +# modeling_dspark; keep the historical alias so monkeypatch targets +# below read unchanged. +dspark_attention = modeling_dspark def test_rmsnorm_rope_fallback_applies_weight_without_rmsnorm(monkeypatch): @@ -118,8 +122,8 @@ def test_rope_table_is_cached_once_per_device(): _freqs_table_cache={}, ) - first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + first = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) assert first.data_ptr() == second.data_ptr() assert len(model._freqs_table_cache) == 1 @@ -162,7 +166,7 @@ def fake_decoder_layer_init( spec_config=None, ) - block = modeling_dspark.DSparkBlock( + block = modeling_dspark.DSv4DSparkBlock( model_config, layer_idx=10, aux_stream_dict={}, @@ -244,7 +248,7 @@ def call(*args, **kwargs): ), ) - actual = DSparkDraftModel._forward_stage( + actual = DSv4DSparkDraftModel._forward_stage( model, stage, h, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 9ada20bf4ef3..95fa963affed 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -27,7 +27,7 @@ import pytest import torch -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( apply_dspark_rotary, apply_dspark_rotary_batched, dspark_attention_forward, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index fed0da3a2803..9d349ef72a2c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -16,8 +16,8 @@ import torch -from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose -from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead, build_markov_head +from tensorrt_llm._torch.models.modeling_dspark import build_draft_input_ids, dspark_propose +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead, build_markov_head VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 NOISE_ID = 199 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py index c59c2c5909ae..4a342e6ee8ef 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py @@ -68,8 +68,12 @@ def _model_config(lb_config=None, num_hidden_layers=NUM_HIDDEN_LAYERS): ) -def _spec_config(mode): - return SimpleNamespace(spec_dec_mode=mode) +def _spec_config(mode, *, embedded=True): + # Only the embedded DSpark draft shares the target's EPLB layer namespace: + # its stages are target decoder blocks registered into the target's + # balancer. A standalone DSpark drafter is an independent checkpoint and is + # treated like the other external drafters. + return SimpleNamespace(spec_dec_mode=mode, draft_is_embedded_in_target=embedded) @pytest.fixture @@ -124,6 +128,17 @@ def test_non_dspark_external_drafters_do_not_inherit_load_balancer(mode): assert "moe_load_balancer" not in kwargs +def test_standalone_dspark_drafter_does_not_inherit_load_balancer(): + # The flavour, not the mode, decides: a standalone DSpark drafter is its own + # checkpoint, so forwarding the target's EPLB config would key its experts + # against a layer namespace that is not the drafter's. + kwargs = external_drafter_config_kwargs( + _model_config(_lb_config(DSPARK_LAYERS)), + _spec_config(SpeculativeDecodingMode.DSPARK, embedded=False), + ) + assert "moe_load_balancer" not in kwargs + + def test_external_drafter_kwargs_are_stable_across_modes(): common = external_drafter_config_kwargs( _model_config(_lb_config(DSPARK_LAYERS)), _spec_config(SpeculativeDecodingMode.PARD) @@ -208,3 +223,47 @@ def test_layer_base_not_checked_without_eplb(): validate_dspark_eplb_layer_base( _model_config(None), _model_config(None, num_hidden_layers=3) ) + + +# --------------------------------------------------------------------------- +# Deployment-form probe. Every DSpark dispatch -- draft model, worker, spec +# metadata, draft-KV decision -- reads this one flag, so a misread does not +# degrade gracefully: it routes a standalone drafter into the V4 worker. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "index_name,weight_map,model_type,expected", + [ + # A parsed index is authoritative in BOTH directions. Falling through + # to model_type on a standalone V4-shaped drafter would classify it as + # embedded, and that only surfaces inside count_dspark_stages. + ("model.safetensors.index.json", {"mtp.0.layers.0.weight": "a"}, "deepseek_v4", True), + ( + "model.safetensors.index.json", + {"layers.0.self_attn.q_proj.weight": "a"}, + "deepseek_v4", + False, + ), + # The bin index is probed too; only the safetensors one used to be. + ("pytorch_model.bin.index.json", {"mtp.1.mlp.weight": "a"}, "qwen3", True), + # No index at all -> the model_type fallback. + (None, None, "deepseek_v4", True), + (None, None, "qwen3", False), + ], + ids=["mtp_index", "standalone_index", "bin_index", "no_index_v4", "no_index_qwen3"], +) +def test_draft_form_probe_reads_the_checkpoint( + tmp_path, index_name, weight_map, model_type, expected +): + import json + + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + if index_name is not None: + (tmp_path / index_name).write_text(json.dumps({"weight_map": weight_map})) + + cfg = DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + + assert cfg.draft_is_embedded_in_target is expected diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py index 16b6a01acd17..5546c1b03f85 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -17,7 +17,7 @@ import pytest import torch -from tensorrt_llm._torch.models.dspark.heads import ( +from tensorrt_llm._torch.models.modeling_speculative import ( DSparkConfidenceHead, RNNHead, VanillaMarkov, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 220aea0656d5..d5d5ec2a0bf2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -16,7 +16,7 @@ Covers the framework-side logic that does NOT need the full draft model: ``DSparkSpecMetadata`` hidden-state capture (incl. the mHC hc-mean reduction) -and ``DSparkWorker`` slot / rolling-KV-window management. The end-to-end block +and ``DSv4DSparkWorker`` slot / rolling-KV-window management. The end-to-end block draft and acceptance path is covered by the DSpark test in ``integration/defs/accuracy/test_llm_api_pytorch.py``. """ @@ -26,7 +26,11 @@ import pytest import torch -from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSparkWorker +from tensorrt_llm._torch.speculative.dspark import ( + DSparkSpecMetadata, + DSparkWorker, + DSv4DSparkWorker, +) from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode pytestmark = pytest.mark.skipif( @@ -101,7 +105,7 @@ def _make_worker(): ) from tensorrt_llm.mapping import Mapping - return DSparkWorker(cfg, Mapping()) + return DSv4DSparkWorker(cfg, Mapping()) def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): @@ -735,3 +739,70 @@ def fake_sample_draft_tokens(gl, sm, bs, *, draft_step): gen_draft = nd[num_contexts:] expected_gen = torch.stack([s[num_contexts:] for s in sampled_per_step], dim=1) assert torch.equal(gen_draft, expected_gen) + + +# --------------------------------------------------------------------------- +# Routing: decoding_type DSpark serves two deployment forms, and the worker, +# the spec metadata and the draft-KV decision must all follow the same flag. +# Mis-routing is not hypothetical: handing a standalone drafter to +# DSv4DSparkWorker raises AttributeError on num_stages at lazy init. +# --------------------------------------------------------------------------- + + +def _routing_config(embedded): + return types.SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + draft_is_embedded_in_target=embedded, + attention_backend="TRTLLM", + confidence_threshold=0.5, + _use_shared_kv_cache=False, + _allow_separate_draft_kv_cache=True, + ) + + +@pytest.mark.parametrize( + "embedded,worker_cls,uses_separate_draft_kv", + [(True, DSv4DSparkWorker, False), (False, DSparkWorker, True)], + ids=["embedded", "standalone"], +) +def test_worker_and_draft_kv_follow_the_draft_form(embedded, worker_cls, uses_separate_draft_kv): + from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache + from tensorrt_llm._torch.speculative.utils import get_spec_worker + from tensorrt_llm.mapping import Mapping + + spec_config = _routing_config(embedded) + + worker = get_spec_worker(spec_config, None, Mapping()) + assert type(worker) is worker_cls + + # The embedded draft owns a rolling window and never reads the paged draft + # KV cache; the standalone one is DFlash lineage and does. + assert should_use_separate_draft_kv_cache(spec_config) is uses_separate_draft_kv + + +def test_dspark_worker_policies_come_from_the_drafter(): + """The two DSpark overrides must read the checkpoint, not hardcode dspark. + + A DSpark drafter trained with the legacy DFlash slot layout, or shipped + without a Markov head, has to get the base-class behaviour. Hardcoding the + dspark answer here would pass every routing test while silently mis-slotting + such a drafter -- and mis-slotting costs acceptance without ever failing. + """ + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + from tensorrt_llm.mapping import Mapping + + worker = DSparkWorker(_routing_config(False), Mapping()) + legacy = types.SimpleNamespace(_dspark_shift_label=False, has_markov_head=False) + + # shift_label off -> the base class' slots 1..K, not the dspark 0..K-1. + ids = worker._draft_slot_ids(legacy, num_gens=2, block_size=5, num_draft_tokens=3) + base_ids = DFlashWorker._draft_slot_ids( + worker, legacy, num_gens=2, block_size=5, num_draft_tokens=3 + ) + assert ids.tolist() == base_ids.tolist() + + # No Markov head -> the backbone logits pass through untouched. + logits = torch.randn(2, 3, 8, device="cuda") + assert worker._refine_block_logits(legacy, logits, {}, None) is logits diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py index ef778488ad8a..51af02a03888 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py @@ -342,3 +342,121 @@ def test_dflash_spec_metadata_capture_prefix_sum_convention(): assert captured.shape == (max_tokens, 2 * hidden_size) torch.testing.assert_close(captured[:, :hidden_size], h1) torch.testing.assert_close(captured[:, hidden_size:], h3) + + +def test_capture_taps_the_next_layers_aggregated_stream(): + """The tap reads layer j+1's mixture, not layer j's prefix sum. + + Worth 4.5pt of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> + 71.4%), and it degrades silently: capturing the prefix sum instead is + still a valid tensor of the right shape, so nothing raises and only + acceptance moves. The two neighbouring tests cover buffer routing and + the forward signature; both still pass if the tap regresses. + + Layers here are stubs that reproduce the one contract the real layer has + with this loop -- compute the attention-side mixture, and when handed a + ``capture`` tuple, record it under the layer id in that tuple. Giving + each stub its own proj/norm weights is what makes the assertions able to + tell "layer j+1's mixture" from "layer j's", and the model's output-side + weights from either. Also covers the no-successor tail branch, which uses + output_attn_res_proj/norm and no layer can produce. + """ + pytest.importorskip("fla") + from types import SimpleNamespace + + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + KimiK3RMSNorm, + KimiLinearModel, + _apply_attn_res, + ) + + torch.manual_seed(0) + num_layers, num_tokens, hidden = 4, 3, 8 + + def _weights(scale): + proj = torch.nn.Linear(hidden, 1, bias=False, dtype=torch.float32) + norm = KimiK3RMSNorm(hidden, eps=1e-6, dtype=torch.float32) + with torch.no_grad(): + proj.weight.copy_(torch.randn(1, hidden) * scale) + norm.weight.copy_(torch.randn(hidden) * scale + 1.0) + return proj, norm + + layer_w = [_weights(0.5 + i) for i in range(num_layers)] + out_proj, out_norm = _weights(9.0) + + seen = {} + + class _StubLayer: + def __init__(self, idx): + self.layer_idx = idx + self.proj, self.norm = layer_w[idx] + + def __call__(self, hidden_states, block_residual, num_snapshots, attn_metadata, capture): + # What the real layer computes on the way into its attention, and + # hands to the tap: the aggregated stream the previous layer's + # consumer sees. + mixture = ( + _apply_attn_res(hidden_states, block_residual[:num_snapshots], self.proj, self.norm) + if num_snapshots > 0 + else hidden_states + ) + if capture is not None: + spec_md, layer_id = capture + spec_md.maybe_capture_hidden_states(layer_id, mixture, None) + block_residual[num_snapshots] = hidden_states * (self.layer_idx + 1) + # A prefix sum that is never equal to the mixture above. + return hidden_states + (self.layer_idx + 1), num_snapshots + 1 + + layers = [_StubLayer(i) for i in range(num_layers)] + last_idx = num_layers - 1 + spec_md = SimpleNamespace( + _capture_layer_set=frozenset({0, last_idx}), + maybe_capture_hidden_states=lambda lid, h, r: seen.__setitem__(lid, h.clone()), + ) + embeds = torch.randn(num_tokens, hidden, dtype=torch.float32) + fake = SimpleNamespace( + embed_tokens=lambda ids: embeds, + layers=layers, + norm=lambda h: h, + output_attn_res_proj=out_proj, + output_attn_res_norm=out_norm, + num_attn_res_snapshots=num_layers, + ) + + KimiLinearModel.forward( + fake, + attn_metadata=SimpleNamespace(num_tokens=num_tokens), + input_ids=torch.zeros(num_tokens, dtype=torch.int32), + spec_metadata=spec_md, + ) + + assert set(seen) == {0, last_idx}, "capture_set not honoured" + + # Replay layer 0 to rebuild what each candidate tensor would have been. + br = torch.empty(num_layers, num_tokens, hidden) + br[0] = embeds * 1 + after_l0 = embeds + 1 + + proj1, norm1 = layer_w[1] + torch.testing.assert_close(seen[0], _apply_attn_res(after_l0, br[:1], proj1, norm1)) + # The two ways this regresses, both silent: + proj0, norm0 = layer_w[0] + assert not torch.allclose(seen[0], _apply_attn_res(after_l0, br[:1], proj0, norm0)), ( + "tap used layer j's own weights instead of its successor's" + ) + assert not torch.allclose(seen[0], after_l0), "tap captured the raw prefix sum" + + # Tail: no successor exists, so it must use the model's output-side + # weights -- not the last layer's, and not the bare prefix sum. + h = embeds + for i in range(num_layers): + br[i] = h * (i + 1) + h = h + (i + 1) + torch.testing.assert_close( + seen[last_idx], _apply_attn_res(h, br[:num_layers], out_proj, out_norm) + ) + proj_last, norm_last = layer_w[last_idx] + assert not torch.allclose( + seen[last_idx], _apply_attn_res(h, br[:num_layers], proj_last, norm_last) + ), "tail used the last layer's weights instead of the output-side ones" + assert not torch.allclose(seen[last_idx], h), "tail captured the raw prefix sum" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index a2cbae410d27..ff3af9a76b90 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -19,13 +19,16 @@ confidence_proj weights load without being used. """ +import re +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM, dspark_layer_window_size +from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM from tensorrt_llm._torch.models.modeling_speculative import ( - DFlashForCausalLM, - dspark_layer_window_size, dspark_markov_chain_logits, dspark_markov_step_bias, ) @@ -205,11 +208,29 @@ def test_swa_window_conventions(): NUM_CAPTURE = 2 -def _tiny_config(dspark: bool): +def _tiny_config(dspark: bool, *, published_spelling: bool = False): + """Tiny drafter config. + + ``published_spelling`` reproduces the two public K3 DSpark checkpoints + (RadixArk, Inferact): the head switches sit at the TOP level, the + confidence flag is ``enable_confidence_head``, and neither ``shift_label`` + nor ``projector_type`` is declared at all -- shift_label rides on the + DSpark default. Reading only ``dflash_config`` resolves markov_rank to 0 + there, which drops the heads without raising. + """ from transformers import Qwen3Config cfg = dict(TINY) dflash = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + if dspark and published_spelling: + cfg.update( + markov_rank=RANK, + markov_head_type="vanilla", + enable_confidence_head=True, + confidence_head_with_markov=True, + ) + cfg["dflash_config"] = dflash + return Qwen3Config.from_dict(cfg) if dspark: dflash.update( projector_type="dspark", @@ -227,7 +248,13 @@ def _tiny_config(dspark: bool): return Qwen3Config.from_dict(cfg) -def _tiny_weights(seed=7): +def _tiny_weights(seed=7, *, published_head_keys=False): + """Tiny drafter weights. + + ``published_head_keys`` names the head tensors the way both public + checkpoints ship them -- after the submodules that own them -- instead of + the bare spellings the DSv4 stage weights use. + """ g = torch.Generator().manual_seed(seed) def rnd(*shape): @@ -235,15 +262,25 @@ def rnd(*shape): h, inter = TINY["hidden_size"], TINY["intermediate_size"] nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) - w = { - "fc.weight": rnd(h, h * NUM_CAPTURE), - "hidden_norm.weight": rnd(h) + 1.0, - "norm.weight": rnd(h) + 1.0, + head = { "markov_w1.weight": rnd(VOCAB, RANK), "markov_w2.weight": rnd(VOCAB, RANK), "confidence_proj.weight": rnd(1, h + RANK), "confidence_proj.bias": rnd(1), } + if published_head_keys: + head = { + "markov_head.markov_w1.weight": head["markov_w1.weight"], + "markov_head.markov_w2.weight": head["markov_w2.weight"], + "confidence_head.proj.weight": head["confidence_proj.weight"], + "confidence_head.proj.bias": head["confidence_proj.bias"], + } + w = { + "fc.weight": rnd(h, h * NUM_CAPTURE), + "hidden_norm.weight": rnd(h) + 1.0, + "norm.weight": rnd(h) + 1.0, + **head, + } for i in range(TINY["num_hidden_layers"]): p = f"layers.{i}." w[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) @@ -260,11 +297,24 @@ def rnd(*shape): return w -def _build_drafter(dspark: bool, weights): +def _build_drafter( + dspark: bool, + weights, + *, + published_spelling: bool = False, + dflash_attention_backend: str = "VANILLA", +): from tensorrt_llm._torch.model_config import ModelConfig - model_config = ModelConfig(pretrained_config=_tiny_config(dspark), attn_backend="TRTLLM") - drafter = DFlashForCausalLM(model_config).to("cuda") + model_config = ModelConfig( + pretrained_config=_tiny_config(dspark, published_spelling=published_spelling), + attn_backend="TRTLLM", + ) + # The DSpark head set lives in the DSpark drafter, not in the DFlash base. + drafter_cls = GQADSparkForCausalLM if dspark else DFlashForCausalLM + drafter = drafter_cls(model_config, dflash_attention_backend=dflash_attention_backend).to( + "cuda" + ) # Drop dspark head tensors for the plain drafter (schema without them). if not dspark: weights = {k: v for k, v in weights.items() if not k.startswith(("markov_", "confidence_"))} @@ -356,9 +406,9 @@ def _has_flash_attn(): def test_dspark_drafter_loads_head_weights_and_parses_config(): weights = _tiny_weights() drafter = _build_drafter(True, weights) - assert drafter._dspark_shift_label and drafter._dspark_use_swa - assert drafter._dspark_swa_window == SWA_WINDOW - assert drafter._dspark_layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 + assert drafter._dspark_shift_label and drafter._use_swa + assert drafter._swa_window == SWA_WINDOW + assert drafter._layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 assert drafter.has_markov_head torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_w1.weight"]) torch.testing.assert_close(drafter.markov_w2.cpu(), weights["markov_w2.weight"]) @@ -369,19 +419,115 @@ def test_dspark_drafter_loads_head_weights_and_parses_config(): assert drafter.confidence_proj_bias is not None +@needs_gpu +def test_published_drafter_spelling_activates_the_heads(): + """Both public K3 DSpark checkpoints load with their heads live. + + They declare the switches at the top level and name the head tensors after + the owning submodules. Reading only ``dflash_config`` and the bare tensor + names resolves markov_rank to 0 and drops markov_w1/w2 on the floor: + correct output, lower acceptance, nothing raised. + """ + weights = _tiny_weights(published_head_keys=True) + drafter = _build_drafter(True, weights, published_spelling=True) + + assert drafter.has_markov_head, "markov weights dropped despite being in the checkpoint" + assert drafter._dspark_use_confidence_head, "enable_confidence_head spelling not resolved" + # Declared nowhere in the published config, so it rides on the DSpark + # default. False would run slots 1..K on a block_size-K drafter and read + # the next request's anchor slot. + assert drafter._dspark_shift_label + torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_head.markov_w1.weight"]) + assert drafter.confidence_proj_bias is not None + + +@needs_gpu +def test_head_weights_without_a_resolvable_rank_raise(): + """The inverse of the missing-weights check. + + A checkpoint that ships markov_w1/w2 while the rank resolves to 0 means the + switches were spelled somewhere this build cannot read. Loading it anyway + would silently cost acceptance, so it is an error. + """ + from tensorrt_llm._torch.model_config import ModelConfig + + # dspark head weights, but a config that declares no head switches at all. + model_config = ModelConfig(pretrained_config=_tiny_config(False), attn_backend="TRTLLM") + drafter = GQADSparkForCausalLM(model_config).to("cuda") + + with pytest.raises(ValueError, match="markov_rank resolved to 0"): + drafter.load_weights(_tiny_weights(published_head_keys=True)) + + +def test_dflash_refuses_a_drafter_that_declares_the_dspark_heads(): + """``decoding_type: DFlash`` must reject a DSpark drafter, not degrade it. + + DFlash does not implement the Markov / confidence / shift_label semantics, + so serving one here would lower the acceptance rate with no error and no + way to attribute it back. + """ + from tensorrt_llm._torch.models import modeling_dflash + + model_config = SimpleNamespace( + spec_config=SimpleNamespace(attention_backend="TRTLLM", speculative_model="/nonexistent") + ) + draft_config = SimpleNamespace(pretrained_config=_tiny_config(True)) + + with pytest.raises(ValueError, match="DSpark"): + modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + +@pytest.mark.parametrize( + "layers,expected", + [ + # MLA-shaped: no per-head q/k/v projection to split at all. + ([SimpleNamespace(self_attn=SimpleNamespace())], "qkv_proj"), + # GQA-shaped but heterogeneous: the cross-layer K/V fusion needs one + # uniform num_key_value_heads. + ( + [ + SimpleNamespace(self_attn=SimpleNamespace(qkv_proj=object(), num_key_value_heads=n)) + for n in (8, 8, 4) + ], + "[2]", + ), + ], + ids=["no_fused_qkv", "mismatched_kv_heads"], +) +def test_block_decode_rejects_a_backbone_it_cannot_express(layers, expected): + """The GQA precondition replaced the old per-model_type whitelist. + + Without it the registry happily builds an unsupported backbone and the + failure surfaces much later inside _build_fused_kv_buffers, or as silently + mis-sliced weights. + """ + drafter = DFlashForCausalLM.__new__(DFlashForCausalLM) + drafter.model = SimpleNamespace(layers=layers) + drafter.config = SimpleNamespace() + + with pytest.raises(ValueError, match=re.escape(expected)): + drafter._validate_gqa_shape() + + @needs_gpu def test_plain_dflash_drafter_keeps_old_gates(): """No-regression: a config WITHOUT dspark fields resolves to the exact - old code path (no window, no markov, mask slots 1..K).""" + old code path (no window, no markov, mask slots 1..K). + + The DFlash base no longer carries the DSpark head set at all, so the + assertions are that those attributes are absent rather than inert. + """ drafter = _build_drafter(False, _tiny_weights()) - assert not drafter._dspark_shift_label - assert not drafter._dspark_use_swa - assert drafter._dspark_layer_windows == [(-1, -1)] * 2 - assert not drafter.has_markov_head - assert drafter.markov_w1 is None and drafter.confidence_proj_weight is None - x = torch.randn(2, 3, VOCAB) - t = torch.zeros(2, dtype=torch.long) - assert drafter.apply_markov_chain_logits(x, t) is x + assert not drafter._use_swa + assert drafter._layer_windows == [(-1, -1)] * 2 + for absent in ( + "_dspark_shift_label", + "has_markov_head", + "markov_w1", + "confidence_proj_weight", + "apply_markov_chain_logits", + ): + assert not hasattr(drafter, absent), f"DFlash base still carries {absent}" @needs_gpu @@ -394,8 +540,8 @@ def test_legacy_causal_dflash_config_constructs(): cfg = _tiny_config(False) cfg.dflash_config = dict(cfg.dflash_config, causal=True) drafter = DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) - assert drafter._dspark_layer_windows == [(-1, -1)] * 2 - assert not drafter.has_markov_head + assert drafter._layer_windows == [(-1, -1)] * 2 + assert not hasattr(drafter, "has_markov_head") @needs_gpu @@ -405,8 +551,8 @@ def test_dspark_causal_config_rejected(): cfg = _tiny_config(True) cfg.dflash_config = dict(cfg.dflash_config, causal=True) - with pytest.raises(ValueError, match="non-causal dspark convention"): - DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + with pytest.raises(ValueError, match="non-causal DSpark convention"): + GQADSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) @needs_gpu @@ -417,8 +563,8 @@ def test_dspark_projector_type_alone_rejects_causal(): cfg = _tiny_config(False) cfg.dflash_config = dict(cfg.dflash_config, projector_type="dspark", causal=True) - with pytest.raises(ValueError, match="non-causal dspark convention"): - DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + with pytest.raises(ValueError, match="non-causal DSpark convention"): + GQADSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) def _run_block_decode(drafter, weights, captured, noise_embed): diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 42a1c69ee2a9..06e35de621ce 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -7,7 +7,7 @@ import torch from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( dspark_sparse_attn, get_dspark_topk_idxs_batched, ) @@ -172,7 +172,7 @@ def test_cute_dsl_dspark_attention_compiles_once_across_batch_sizes(): def test_dspark_attention_forward_batched_fused_matches_fallback(monkeypatch): - import tensorrt_llm._torch.models.dspark.attention as dspark_attention + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention torch.manual_seed(17) device = torch.device("cuda") diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b4b1569ed66d..7c7722c794ec 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -457,13 +457,13 @@ def test_dspark_target_layer_ids_order_mismatch_rejected(tmp_path): @pytest.mark.cpu_only def test_dspark_requires_speculative_model(): - # The DSpark draft weights live in the checkpoint's mtp.* namespace, so an - # unset speculative_model must fail fast at config validation instead of - # raising an opaque TypeError deep inside engine construction. + # Unset speculative_model means "load the draft from the target", as it + # does for MTP. /tmp/dummy_model carries no mtp.* draft weights, so that + # must fail fast at config validation instead of raising an opaque + # TypeError deep inside engine construction. spec_cfg = DSparkDecodingConfig(max_draft_len=5) - with pytest.raises(ValueError, - match="requires speculative_config.speculative_model"): + with pytest.raises(ValueError, match="speculative_model is unset"): TorchLlmArgs( model="/tmp/dummy_model", skip_tokenizer_init=True, diff --git a/tests/unittest/others/test_lazy_model_zoo.py b/tests/unittest/others/test_lazy_model_zoo.py index e82ccf1b8e83..efb50f78dffe 100644 --- a/tests/unittest/others/test_lazy_model_zoo.py +++ b/tests/unittest/others/test_lazy_model_zoo.py @@ -179,6 +179,61 @@ def test_arch_index_matches_decorators(): assert not wrong, f"index points at the wrong module: {wrong}" +def _decorated_draft_model_registrations(): + """AST-scan the modeling files for ``@register_draft_model`` decorators. + + Source scan rather than a registry walk on purpose: the registry is only + populated by importing a provider, and a built-in builder that lost its + slot to an external registration would be missing from it entirely, so a + registry walk would pass while the index is stale. + """ + mode_to_modules = {} + for path in sorted(_MODELS_DIR.glob("*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + name = getattr(func, "id", None) or getattr(func, "attr", None) + if name != "register_draft_model" or not node.args: + continue + arg = node.args[0] + # ``register_draft_model(SpeculativeDecodingMode.DFLASH)`` + if isinstance(arg, ast.Attribute): + mode_to_modules.setdefault(arg.attr, set()).add(path.stem) + return mode_to_modules + + +def test_spec_mode_index_matches_decorators(): + # SPEC_MODE_TO_MODULE is the fourth hand-maintained table in _arch_index, + # and every other one already has a drift test here -- MODEL_ARCH_TO_MODULE + # and MULTIMODAL_MODEL_TYPE_TO_MODULE in test_arch_index_matches_decorators, + # MODEL_CLASS_TO_MODULE in test_class_index_matches_package_all. Keeping the + # set complete is the point: without this, the spec-mode table would be the + # only index whose entries nothing checks. + # + # Drift is not silent, but it is misattributed. A builder added without its + # index entry (or an index entry whose module stopped declaring the mode) + # surfaces as "unsupported speculative decoding mode" to whoever next runs + # that mode -- which reads as "this algorithm is not implemented", not as + # "someone forgot a line in _arch_index". The person who sees it is rarely + # the person who caused it. + from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE + + mode_truth = _decorated_draft_model_registrations() + + missing = set(mode_truth) - set(SPEC_MODE_TO_MODULE) + assert not missing, f"spec modes missing from _arch_index: {missing}" + stale = set(SPEC_MODE_TO_MODULE) - set(mode_truth) + assert not stale, f"stale spec modes in _arch_index: {stale}" + wrong = { + mode: (SPEC_MODE_TO_MODULE[mode], mode_truth[mode]) + for mode in SPEC_MODE_TO_MODULE + if SPEC_MODE_TO_MODULE[mode] not in mode_truth[mode] + } + assert not wrong, f"index points at the wrong module: {wrong}" + + def test_class_index_matches_package_all(): # MODEL_CLASS_TO_MODULE is the one table with no decorator to mirror: it # backs PEP 562 attribute access on the models package. Every name in the