diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 52f16c84ffe9..7caf1c6ecc28 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -529,6 +529,10 @@ def prepare_for_indexer_k_cache(self): self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True, ) + # Columns beyond each sequence's allocated indexer blocks contain BAD_PAGE_INDEX (-1). + # CUDA-graph padded token slots may still compute scatter addresses from those columns + # before being ignored, so map them to block 0, matching the base DSA metadata path. + self.indexer_k_cache_block_offsets.clamp_(min=0) def prepare_for_block_tables(self): """Prepare block tables for sliding-window and compressed attention.""" diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 29b8b8fd390e..fedd19ad0aaa 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -956,9 +956,14 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): pretrained_config, 'compress_ratios', None) num_base_layers = pretrained_config.num_hidden_layers spec_config = kwargs.get('spec_config', None) - if (spec_config is not None - and getattr(spec_config, 'num_nextn_predict_layers', - None) is None): + # ``num_nextn_predict_layers`` is MTP-specific (only read on + # the is_mtp_one_model path). Only set it on configs that + # actually declare the field; other DeepSeek-V4 spec modes + # (e.g. DSpark, which carries its own draft stage count) do + # not, and a blind setattr would fail pydantic validation. + if (spec_config is not None and 'num_nextn_predict_layers' + in type(spec_config).model_fields + and spec_config.num_nextn_predict_layers is None): spec_config.num_nextn_predict_layers = getattr( pretrained_config, 'num_nextn_predict_layers', 1) mtp_enabled = (spec_config is not None and @@ -990,6 +995,22 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): if window_size is None: window_size = pretrained_config.sliding_window + # DeepSeek-V4 needs explicit per-layer compress ratios. They + # must come from the checkpoint config or a user override; we + # intentionally do not synthesize a default list (it would + # silently change sparse-attention semantics). Fail fast with + # an actionable message instead of letting the normalization + # below raise an opaque TypeError on None. + if compress_ratios is None: + raise ValueError( + "DeepSeek-V4 requires per-layer `compress_ratios`, " + "but none were found in the checkpoint config and " + "none were provided via `sparse_attention_config`. " + "Set `compress_ratios` in the model's config.json, or " + "pass `sparse_attention_config=" + "DeepSeekV4SparseAttentionConfig(compress_ratios=[...])`" + " in --extra_llm_api_options.") + # Normalize checkpoint-facing ratio 0 (SWA-only/uncompressed) # to 1 internally so cache allocation math works. The # external config keeps the original semantics. diff --git a/tensorrt_llm/_torch/models/dspark/__init__.py b/tensorrt_llm/_torch/models/dspark/__init__.py new file mode 100644 index 000000000000..b33d7553d877 --- /dev/null +++ b/tensorrt_llm/_torch/models/dspark/__init__.py @@ -0,0 +1,16 @@ +# 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 new file mode 100644 index 000000000000..7f63a40edfbe --- /dev/null +++ b/tensorrt_llm/_torch/models/dspark/attention.py @@ -0,0 +1,462 @@ +# 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 + +__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, +) -> 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, so the result is numerically identical to gathering only the + ``min(window_size, start_pos+1)`` valid context positions — but the shape no + longer depends on the data, which is what CUDA-graph capture requires. + + Every query position attends to the same set: context window slots + ``0..window_size-1`` (slot ``c`` valid iff ``c <= start_pos[g]``, i.e. it has + been written) followed by the ``block_size`` block positions at offset + ``window_size`` (always valid). + + 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. + + 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] + # Context slot c holds a written key iff c <= start_pos (slots 0..start_pos + # filled; for start_pos >= window_size-1 the whole rolling window is filled). + valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] + 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 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, + *, + 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. + 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 + 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(F.linear(main_x, wkv), kv_norm_w, eps) # [G, 1, head_dim] + main_kv = _rope_last_dims_batched(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)) # [G, block, h, head_dim] + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + q = _rope_last_dims_batched(q, rd, blk_freqs) + + # Block K/V. + kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [G, block, head_dim] + kv = _rope_last_dims_batched(kv, 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. Indexed scatter/gather by (slots, slot_pos) is graph-safe. + write_target = kv_cache if persist else kv_cache.clone() + slot_pos = start_pos % window_size # [G] + write_target[slots, slot_pos] = main_kv.squeeze(1).to(write_target.dtype) + 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) + o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [G, block, h, head_dim] + o = _rope_last_dims_batched(o, rd, blk_freqs, inverse=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 new file mode 100644 index 000000000000..1b47f4922965 --- /dev/null +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -0,0 +1,130 @@ +# 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 new file mode 100644 index 000000000000..c49e35fbafa0 --- /dev/null +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -0,0 +1,254 @@ +# 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_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 2fb611ecb999..3d65b1ee7595 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -228,6 +228,66 @@ def moe_reduce_add_shared_output(routed_output, shared_output): } +def _get_deepseek_v4_routed_moe_scale_name(weights: Dict, key_prefix: str) -> str: + """Return the model scale suffix for routed-expert checkpoint tensors.""" + for key, value in weights.items(): + if ( + key.startswith(key_prefix) + and ".ffn.experts." in key + and key.endswith(".weight") + and getattr(value, "ndim", 0) == 2 + and getattr(value, "dtype", None) in (torch.int8, torch.uint8) + ): + return "weight_scale" + return "weight_scale_inv" + + +def _rename_deepseek_v4_attn_subkey(rest: str) -> str: + """Rename a DeepSeek-V4 attention checkpoint subkey.""" + if rest == "attn_sink": + return "attn_sink" + if rest == "wo_a.weight": + return "o_a_proj" + if rest == "wo_a.scale": + return "o_a_proj.weight_scale_inv" + if rest.startswith("compressor.") or rest.startswith("indexer."): + return rest.replace(".scale", ".weight_scale_inv") + head, sep, tail = rest.partition(".") + new_head = _ATTN_PARAM_RENAME.get(head, head) + if tail == "scale": + tail = "weight_scale_inv" + return f"{new_head}.{tail}" if sep else new_head + + +def _rename_deepseek_v4_ffn_subkey(rest: str, routed_moe_scale_name: str) -> str: + """Rename a DeepSeek-V4 FFN checkpoint subkey.""" + if rest == "gate.bias": + return "gate.e_score_correction_bias" + if rest.startswith("experts.") and rest.endswith(".scale"): + return f"{rest[: -len('.scale')]}.{routed_moe_scale_name}" + rest = rest.replace(".scale", ".weight_scale_inv") + if rest.startswith("shared_experts."): + parts = rest.split(".") + if len(parts) >= 2 and parts[1] in _SHARED_EXPERT_RENAME: + parts[1] = _SHARED_EXPERT_RENAME[parts[1]] + rest = ".".join(parts) + return rest + + +def _maybe_view_deepseek_v4_routed_moe_tensor( + model_key: str, tensor: torch.Tensor, routed_moe_scale_name: str +) -> torch.Tensor: + """Expose packed MXFP4 routed-expert tensors through their uint8 view.""" + if ( + routed_moe_scale_name == "weight_scale" + and ".mlp.experts." in model_key + and (model_key.endswith(".weight") or model_key.endswith(".weight_scale")) + and tensor.dtype != torch.uint8 + ): + return tensor.view(torch.uint8) + return tensor + + def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: """Resolve the DeepSeek-V4 fused HC boundary-fusion knob.""" env = os.environ.get("TRTLLM_MHC_ENABLE_FUSED_HC") @@ -315,66 +375,7 @@ def _remap_deepseek_v4_checkpoint_keys( carries it but matches the main head, so we let the main head win. """ mtp_layer_prefix = f"model.layers.{num_hidden_layers}" - routed_moe_scale_name = "weight_scale_inv" - for key, value in weights.items(): - if ( - key.startswith("layers.") - and ".ffn.experts." in key - and key.endswith(".weight") - and getattr(value, "ndim", 0) == 2 - and value.dtype in (torch.int8, torch.uint8) - ): - routed_moe_scale_name = "weight_scale" - break - - def _rename_attn_subkey(rest: str) -> Optional[str]: - # rest examples: "wq_a.weight", "wq_a.scale", "wo_a.weight", - # "attn_sink", "compressor.wkv.weight", "indexer.wq_b.scale", - # "kv_norm.weight" - # ``attn_sink`` is loaded by the ``mqa`` branch in the per-module - # loader, which reads it under the parent ``self_attn.attn_sink`` - # key. Pass through unchanged. - if rest == "attn_sink": - return "attn_sink" - # `wo_a` is an nn.Parameter on the model side (not a Linear), so - # `wo_a.weight` carries the value directly into `o_a_proj` without - # a trailing ``.weight``. Retain `.scale` so the loader can dequantize - # FP8 block-scaled checkpoints before assigning the bf16 parameter. - if rest == "wo_a.weight": - return "o_a_proj" - if rest == "wo_a.scale": - return "o_a_proj.weight_scale_inv" - # Compressor / indexer paths — pass through with .scale rename, plus - # wkv+wgate fusion handled separately below. - if rest.startswith("compressor.") or rest.startswith("indexer."): - return rest.replace(".scale", ".weight_scale_inv") - head, sep, tail = rest.partition(".") - new_head = _ATTN_PARAM_RENAME.get(head, head) - if tail == "scale": - tail = "weight_scale_inv" - return f"{new_head}.{tail}" if sep else new_head - - def _rename_ffn_subkey(rest: str) -> str: - # Examples: - # gate.weight / gate.tid2eid → gate.weight / gate.tid2eid - # gate.bias → gate.e_score_correction_bias - # experts... → experts... - # shared_experts.. → shared_experts._proj. - if rest == "gate.bias": - return "gate.e_score_correction_bias" - if rest.startswith("experts.") and rest.endswith(".scale"): - return f"{rest[: -len('.scale')]}.{routed_moe_scale_name}" - rest = rest.replace(".scale", ".weight_scale_inv") - # Non-hashed layers carry the routing logit bias as `gate.bias`; the - # model wires it through `DeepseekV4Gate.e_score_correction_bias`. - if rest == "gate.bias": - return "gate.e_score_correction_bias" - if rest.startswith("shared_experts."): - parts = rest.split(".") - if len(parts) >= 2 and parts[1] in _SHARED_EXPERT_RENAME: - parts[1] = _SHARED_EXPERT_RENAME[parts[1]] - rest = ".".join(parts) - return rest + routed_moe_scale_name = _get_deepseek_v4_routed_moe_scale_name(weights, "layers.") def _rename_layer_subkey(rest: str) -> Optional[str]: # rest examples: "attn_norm.weight", "ffn_norm.weight", @@ -390,10 +391,10 @@ def _rename_layer_subkey(rest: str) -> Optional[str]: if rest.startswith("hc_attn_") or rest.startswith("hc_ffn_"): return rest if rest.startswith("attn."): - new_sub = _rename_attn_subkey(rest[len("attn.") :]) - return None if new_sub is None else f"self_attn.{new_sub}" + return f"self_attn.{_rename_deepseek_v4_attn_subkey(rest[len('attn.') :])}" if rest.startswith("ffn."): - return f"mlp.{_rename_ffn_subkey(rest[len('ffn.') :])}" + new_sub = _rename_deepseek_v4_ffn_subkey(rest[len("ffn.") :], routed_moe_scale_name) + return f"mlp.{new_sub}" return rest out: Dict[str, torch.Tensor] = {} @@ -417,13 +418,7 @@ def _emit_or_collect(model_key: str, tensor: torch.Tensor): part = "wkv" if model_key.endswith(".wkv.weight") else "wgate" _record_compressor_part(model_key, part, tensor) return - if ( - routed_moe_scale_name == "weight_scale" - and ".mlp.experts." in model_key - and (model_key.endswith(".weight") or model_key.endswith(".weight_scale")) - and tensor.dtype != torch.uint8 - ): - tensor = tensor.view(torch.uint8) + tensor = _maybe_view_deepseek_v4_routed_moe_tensor(model_key, tensor, routed_moe_scale_name) out[model_key] = tensor for k, v in weights.items(): @@ -1703,7 +1698,7 @@ def __init__( model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, + attention_layer_idx: Optional[int] = None, mapping_with_cp: Optional[Mapping] = None, disable_post_moe_fusion: bool = False, ): @@ -1732,14 +1727,12 @@ def __init__( post_mult_value=2.0, ) - layer_idx_for_attention = layer_idx - if is_separate_draft_engine: - # KVCacheManager only support 1 layer for separate draft engine - layer_idx_for_attention = layer_idx - model_config.pretrained_config.num_hidden_layers + if attention_layer_idx is None: + attention_layer_idx = layer_idx self.self_attn = DeepseekV4Attention( model_config, - layer_idx=layer_idx_for_attention, + layer_idx=attention_layer_idx, aux_stream=aux_stream_dict[AuxStreamType.Attention], reduce_output=not self.enable_attention_dp and self.mapping.tp_size > 1, ) @@ -1960,8 +1953,18 @@ def forward( # No engram concern here because engram only fires at layer entry. # When enable_fused_hc=False, fall back to the unfused chain. # ------------------------------------------------------------------- - if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + capture_this_layer = spec_metadata is not None and spec_metadata.is_layer_capture( + self.layer_idx + ) + is_dspark_capture = capture_this_layer and spec_metadata.spec_dec_mode.is_dspark() + if capture_this_layer: self.fusion_config.POST_MOE_FUSION = False + if is_dspark_capture: + # DSpark captures the FULL post-mapped mHC residual stream (reference + # `h.mean(dim=2)`), which is only materialized after + # hc_ffn.post_mapping -- so post_mapping must resolve in-layer rather + # than being deferred into the next layer's fused_hc. + self.defer_post_mapping = False if self.enable_fused_hc: residual, post_mix, comb_mix, layer_input = self.hc_ffn.fused_hc( x_prev=x_attn, @@ -2003,6 +2006,17 @@ def forward( post_layer_mix=post_mix, comb_res_mix=comb_mix, ) + if is_dspark_capture: + # Capture the full mHC residual stream [N, hc_mult*hidden]; the DSpark + # metadata means over the hc streams (reference `h.mean(dim=2)`) to + # form the draft's captured context (``main_x``). This is the correct + # representation -- NOT the pre-post_mapping MoE delta that the generic + # capture in forward_MoE records for other spec modes. + spec_metadata.maybe_capture_hidden_states( + self.layer_idx, + resolved_residual.reshape(resolved_residual.shape[0], -1), + None, + ) return HCState.resolved(resolved_residual) def _entry_boundary(self, hc_state, engram_embeddings, has_engram): @@ -2131,7 +2145,14 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize, input_ids): fc2_output, all_reduce_params=moe_all_reduce_params ) else: - if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + # DSpark captures the post-mapped mHC residual stream after this layer + # (done in the decoder-layer forward), not the pre-post_mapping MoE + # output recorded here for other spec modes. + if ( + spec_metadata is not None + and spec_metadata.is_layer_capture(self.layer_idx) + and not spec_metadata.spec_dec_mode.is_dspark() + ): spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, None) return hidden_states @@ -2143,13 +2164,13 @@ def __init__( model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, + attention_layer_idx: Optional[int] = None, ): super().__init__( model_config, layer_idx, aux_stream_dict, - is_separate_draft_engine, + attention_layer_idx=attention_layer_idx, disable_post_moe_fusion=True, ) config = model_config.pretrained_config diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py new file mode 100644 index 000000000000..2ef467061a69 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -0,0 +1,1100 @@ +# 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 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. +""" + +import copy +import json +import os +import re +from typing import Dict, List, Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm.logger import logger + +from ..distributed import AllReduceParams +from ..modules.linear import Linear +from ..modules.mhc.hyper_connection import HCHead +from ..modules.rms_norm import RMSNorm +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, + _get_deepseek_v4_routed_moe_scale_name, + _maybe_view_deepseek_v4_routed_moe_tensor, + _rename_deepseek_v4_attn_subkey, + _rename_deepseek_v4_ffn_subkey, +) + +# Matches the draft namespace ``mtp..`` in the V4-Pro-DSpark +# checkpoint. Each draft stage is a full DeepSeek-V4 block stored under this +# prefix; the main model's keys (``layers.*``, ``embed.weight``, ``head.weight``, +# top-level ``norm.weight`` / ``hc_head_*``) are loaded by the target model. +_DSPARK_MTP_RE = re.compile(r"^mtp\.(\d+)\.(.+)$") + + +def count_dspark_stages(ckpt_dir: str) -> Optional[int]: + """Count the DSpark draft stages (``mtp.{s}.*``) in a checkpoint index. + + The HF ``config.json`` does not expose ``n_mtp_layers`` (only the reference + ``inference/config.json`` does), so the authoritative draft stage count is + the number of distinct ``mtp.`` prefixes in the weight index. Returns + ``None`` if the index is missing or has no ``mtp.*`` keys (caller falls back + to the config-derived default). + """ + index = os.path.join(ckpt_dir, "model.safetensors.index.json") + if not os.path.isfile(index): + return None + with open(index, encoding="utf-8") as f: + weight_map = json.load(f).get("weight_map", {}) + stages = {int(m.group(1)) for k in weight_map if (m := _DSPARK_MTP_RE.match(k))} + return (max(stages) + 1) if stages else None + + +def _rename_dspark_stage_subkey(rest: str, routed_scale: str) -> str: + """Map a per-stage checkpoint subkey to the ``DSparkBlock`` param subkey.""" + if rest == "attn_norm.weight": + return "input_layernorm.weight" + if rest == "ffn_norm.weight": + return "post_attention_layernorm.weight" + # Flat manifold-Hyper-Connections / draft-head weights are loaded via + # ``load_flat_hc_weights`` (keyed by the parent module stem), so pass the + # flat-underscore form through unchanged: + # hc_attn_* / hc_ffn_* -> mHC on every block + # hc_head_* -> HCHead on the last stage + if rest.startswith(("hc_attn_", "hc_ffn_", "hc_head_")): + return rest + # DSpark capture projection (stage 0): fp8 Linear .scale -> .weight_scale_inv. + if rest == "main_proj.scale": + return "main_proj.weight_scale_inv" + if rest.startswith("attn."): + return f"self_attn.{_rename_deepseek_v4_attn_subkey(rest[len('attn.') :])}" + 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. + return rest + + +def remap_dspark_draft_keys(weights: Dict, num_stages: int) -> Dict: + """Convert checkpoint ``mtp.{s}.*`` keys to ``mtp_layers.{s}.*`` model keys. + + Only the draft namespace is consumed (stages ``[0, num_stages)``); shared + ``embed_tokens`` / ``lm_head`` and other top-level keys belong to the target + model and are skipped here. The routed-expert scale suffix mirrors the V4 + loader: ``weight_scale`` for the packed MXFP4 layout, else ``weight_scale_inv``. + """ + routed_scale = _get_deepseek_v4_routed_moe_scale_name(weights, "mtp.") + out: Dict[str, torch.Tensor] = {} + for k, v in weights.items(): + m = _DSPARK_MTP_RE.match(k) + if not m: + continue + stage = int(m.group(1)) + if stage >= num_stages: + continue + sub = _rename_dspark_stage_subkey(m.group(2), routed_scale) + model_key = f"mtp_layers.{stage}.{sub}" + v = _maybe_view_deepseek_v4_routed_moe_tensor(model_key, v, routed_scale) + out[model_key] = v + return out + + +# The checkpoint stores +# ``mtp.{s}.attn.wo_a`` as fp8_e4m3 + a UE8M0 128x128 block scale (verified), and +# the reference (`inference/model.py`, ``self.wo_a`` is a bf16 ColumnParallelLinear +# loaded from the fp8 ckpt) uses the DEQUANTIZED bf16 ``wo_a`` (== ``wo_a_fp8 * +# scale`` ~ absmean 0.065). The bf16 captured-context path historically skipped +# this dequant (raw fp8-cast-to-bf16, ~993x too large); the correct behavior is to +# dequantize ``wo_a`` (cos 1.0 vs ``wo_a_fp8 * scale``). Always dequantize now. + + +class DSparkBlock(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 + and only the last stage owns the draft heads, matching the ``mtp.*`` schema. + """ + + def __init__( + self, + model_config, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + *, + stage_id: int, + num_stages: int, + num_capture_layers: int, + ): + # The inherited attention uses a draft-local layer index, while the + # decoder layer keeps its model-level index for weights and captures. + super().__init__( + model_config, + layer_idx, + aux_stream_dict, + attention_layer_idx=stage_id, + disable_post_moe_fusion=True, + ) + config = model_config.pretrained_config + spec_cfg = getattr(model_config, "spec_config", None) + self.stage_id = int(stage_id) + self.num_stages = int(num_stages) + # mask_token_id is a user override on the speculative_config; None means + # fall back to the draft checkpoint's dspark_noise_token_id. + mask_token_id = getattr(spec_cfg, "mask_token_id", None) + self.noise_token_id = int( + mask_token_id + if mask_token_id is not None + else getattr(config, "dspark_noise_token_id", config.vocab_size) + ) + self.markov_rank = int(getattr(config, "dspark_markov_rank", 0)) + self.hc_mult = config.hc_mult + # markov_head_type is a user override on the speculative_config; None + # means fall back to the draft checkpoint's dspark_markov_head_type. + markov_head_type = getattr(spec_cfg, "markov_head_type", None) + if markov_head_type is None: + markov_head_type = getattr(config, "dspark_markov_head_type", "vanilla") + self.markov_head_type = markov_head_type + + # Stage 0: capture projection of the concatenated target-layer hiddens. + if self.has_capture: + self.main_proj = Linear( + config.hidden_size * num_capture_layers, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.main_norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + # Last stage: the block-draft output heads + mHC head + final norm. + if self.has_heads: + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.hc_head = HCHead(config.hc_mult, config.hidden_size) + self.markov_head = build_markov_head( + markov_head_type=self.markov_head_type, + vocab_size=config.vocab_size, + markov_rank=self.markov_rank, + hidden_size=config.hidden_size, + ) + self.confidence_head = DSparkConfidenceHead( + hidden_size=config.hidden_size, + markov_rank=self.markov_rank, + # Only concat the Markov prev-token embedding when a Markov head + # actually exists (build_markov_head returns None for + # markov_rank <= 0); otherwise dspark_propose passes no + # prev_embeddings and DSparkConfidenceHead.forward would assert. + with_markov=self.markov_rank > 0, + ) + + @property + def has_capture(self) -> bool: + return self.stage_id == 0 + + @property + def has_heads(self) -> bool: + return self.stage_id == self.num_stages - 1 + + +class DSparkDraftModel(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`` + builds the block input from the captured context; the per-stage backbone runs + the 3 blocks; ``forward_head`` produces the block draft tokens + confidence. + """ + + def __init__( + self, + model_config, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + num_stages: Optional[int] = None, + block_size: Optional[int] = None, + ): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.config = config + # The DSpark stage count is NOT the HF ``num_nextn_predict_layers`` (=1). + # It is ``n_mtp_layers`` (3 for V4-Pro), which lives in the draft + # sub-checkpoint config (inference/config.json) and is reflected by the + # ``mtp.{0..n-1}.*`` weight namespace. Resolve it from (in priority): + # an explicit override, the spec config's ``num_draft_layers``, a + # pretrained-config ``n_mtp_layers``, else fall back to nextn. + spec_cfg = getattr(model_config, "spec_config", None) + self.num_stages = int( + num_stages + if num_stages is not None + else getattr(spec_cfg, "num_draft_layers", None) + or getattr(config, "n_mtp_layers", None) + or config.num_nextn_predict_layers + ) + # Production passes the validated speculative-config value explicitly; + # direct construction falls back to the checkpoint's trained block size. + self.block_size = int( + block_size if block_size is not None else getattr(config, "dspark_block_size", 5) + ) + # mask_token_id is a user override on the speculative_config; None means + # fall back to the draft checkpoint's dspark_noise_token_id. + mask_token_id = getattr(spec_cfg, "mask_token_id", None) + self.noise_token_id = int( + mask_token_id + if mask_token_id is not None + else getattr(config, "dspark_noise_token_id", config.vocab_size) + ) + self.hc_mult = config.hc_mult + target_layer_ids = getattr(config, "dspark_target_layer_ids", []) + self.num_capture_layers = len(target_layer_ids) + base = config.num_hidden_layers + # Derive a draft-only model_config (a shallow copy so the shared config + # and the target model are untouched) carrying two draft-specific fixes: + # + # 1. compress_ratios SLICE — the draft runs as a separate engine, so the + # inherited DeepSeek-V4 block remaps each block's layer_idx to a + # draft-local index in [0, num_stages) (the 1-layer-style draft KV + # cache). Sparse-attention compress_ratios / RoPE are indexed by that + # draft-local id, so they must be the draft slice + # (compress_ratios[base : base + num_stages]); otherwise indices + # 0..n-1 resolve to the first *main* layers' sparse ratios — building + # a compressor the DSpark draft lacks and selecting YaRN over the + # dense path. For V4-Pro the draft slice is [1, 1, 1] (dense). + # + # 2. quant_config_dict EXTENSION — the checkpoint's per-module quant map + # only enumerates the base layers, so the draft layers' routed + # experts fall back to the global fp8 config and build fp8-shaped + # buffers. The draft experts are physically MXFP4 (same as the main + # MoE layers), so copy a main MoE layer's experts quant onto the + # draft layer keys. + draft_model_config = self._derive_draft_model_config(model_config, base, self.num_stages) + self.mtp_layers = nn.ModuleList( + [ + DSparkBlock( + draft_model_config, + base + s, + aux_stream_dict, + stage_id=s, + num_stages=self.num_stages, + num_capture_layers=self.num_capture_layers, + ) + for s in range(self.num_stages) + ] + ) + # Shared with target; wired by the spec wrapper after construction. + self.embed_tokens: Optional[nn.Module] = None + self.lm_head: Optional[nn.Module] = None + + # Scalar attention params for the captured-context draft attention. These + # are the dense (compress_ratio == 0) DSparkAttention constants — see the + # reference ``inference/model.py`` ``Attention.__init__``. ``head_dim`` is + # the MLA latent (MQA) dim; ``softmax_scale = head_dim ** -0.5``; the dense + # draft disables YaRN and uses the base ``rope_theta``. + self._attn_params = dict( + n_heads=int(config.num_attention_heads), + head_dim=int( + getattr(config, "head_dim", config.kv_lora_rank + config.qk_rope_head_dim) + ), + rope_head_dim=int(config.qk_rope_head_dim), + n_groups=int(config.o_groups), + o_lora_rank=int(config.o_lora_rank), + window_size=int(getattr(config, "window_size", 128)), + eps=float(config.rms_norm_eps), + ) + self._attn_params["softmax_scale"] = self._attn_params["head_dim"] ** -0.5 + self._rope_theta = float(getattr(config, "rope_theta", 10000.0)) + # Fixed-cap plain-RoPE table shared by the eager and CUDA-graph-safe + # batched paths. It is built once per device and gathered/sliced by the + # runtime decode positions, so the cache does not grow with sequence + # length and the batched consuming op's shape remains static. + self._freqs_cap = ( + int(getattr(config, "max_position_embeddings", 163840)) + self.block_size + 2 + ) + self._freqs_table_cache: Dict = {} + + def post_load_weights(self) -> None: + """Run the one-shot post-load transforms for the draft's quant linears. + + The fp8 UE8M0 linears we invoke as modules (``main_proj``, shared experts, + the heads) need ``resmooth_to_fp8_e8m0`` + ``transform_sf_into_required_layout`` + before the first forward, or the kernel reads raw scales and emits NaNs. + ``Linear.transform_weights`` is idempotent; the routed-expert MoE packs + itself in its own ``load_weights``. + + The bf16 captured-context attention does NOT use the MLA module's forward — + it runs ``dspark_attention_forward`` on dequantized bf16 weights cached via + :meth:`cache_attn_weights_from_checkpoint` — so the MLA projection linears are + skipped here (they would otherwise be transformed into the deep_gemm layout we + don't consume). + """ + attn_linear_ids = set() + for stage in self.mtp_layers: + for m in stage.self_attn.modules(): + if isinstance(m, Linear): + attn_linear_ids.add(id(m)) + + for module in self.modules(): + if isinstance(module, Linear) and id(module) not in attn_linear_ids: + module.transform_weights() + + @staticmethod + def _block_dequant(w_fp8: torch.Tensor, scale: torch.Tensor, block: int = 128) -> torch.Tensor: + """DeepSeek ``block``×``block`` block-scale dequant → bf16: ``real = fp8 * scale``. + + ``scale`` (possibly UE8M0) is broadcast over each ``block``×``block`` tile. + Pure-torch (matches the golden-validated reference dequant), robust to the + e8m0 scale dtype. + """ + wf = w_fp8.float() + out, inn = wf.shape + s = scale.float() + s_full = s.repeat_interleave(block, 0)[:out].repeat_interleave(block, 1)[:, :inn] + return (wf * s_full).to(torch.bfloat16) + + def _cache_attn_weights(self, src: Dict) -> None: + """Populate each stage's ``_dspark_attn`` from a dict of raw ``mtp.{s}.attn.*`` + tensors (source-agnostic core shared by the two public entry points). + + The captured-context attention runs the validated ``dspark_attention_forward`` + free function on reference-layout bf16 weights dequantized here. Sourcing the + separate ``wq_a``/``wkv`` (plain 128×128 block scale) sidesteps the TRT-LLM + ``MLA`` module's fused + interleaved fp8 storage (``kv_a_proj_with_mqa`` fuses + ``q_a``+``kv`` and stores the scale interleaved). This mirrors the + golden-validated dequant exactly. + """ + for s, stage in enumerate(self.mtp_layers): + pref = f"mtp.{s}.attn." + dev = stage.input_layernorm.weight.device + + def deq(name: str, fp8: bool) -> torch.Tensor: + w = src[f"{pref}{name}.weight"].to(dev) + if fp8: + return self._block_dequant(w, src[f"{pref}{name}.scale"].to(dev)) + return w.to(torch.bfloat16) + + stage._dspark_attn = dict( + wq_a=deq("wq_a", True), + q_norm_w=src[f"{pref}q_norm.weight"].to(dev).to(torch.bfloat16), + wq_b=deq("wq_b", True), + wkv=deq("wkv", True), + kv_norm_w=src[f"{pref}kv_norm.weight"].to(dev).to(torch.bfloat16), + # wo_a IS fp8+scale in the checkpoint (verified); always dequant. + wo_a=deq("wo_a", True), + wo_b=deq("wo_b", True), + attn_sink=src[f"{pref}attn_sink"].to(dev).float(), + ) + + def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str, str]) -> None: + """Populate ``_dspark_attn`` by reading the ``mtp.{s}.attn.*`` tensors from the + checkpoint shards on disk, then dequantizing via :meth:`_cache_attn_weights`. + + TODO(step 3): source these from the loaded ``MLA`` modules instead, once the + fused/interleaved fp8 scale layout is decoded, to drop the checkpoint I/O. + """ + from safetensors import safe_open + + prefixes = tuple(f"mtp.{s}.attn." for s in range(len(self.mtp_layers))) + shards: Dict[str, list] = {} + for k in weight_map: + if k.startswith(prefixes): + shards.setdefault(weight_map[k], []).append(k) + raw: Dict[str, torch.Tensor] = {} + for shard, ks in shards.items(): + with safe_open(os.path.join(ckpt_dir, shard), framework="pt", device="cpu") as f: + for k in ks: + raw[k] = f.get_tensor(k) + self._cache_attn_weights(raw) + + 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`. + """ + self._cache_attn_weights(weights) + + def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: + """Return the fixed-size plain-RoPE table cached for ``device``.""" + key = str(device) + cached = self._freqs_table_cache.get(key) + if cached is None: + cached = precompute_dspark_freqs_cis( + self._attn_params["rope_head_dim"], + self._freqs_cap, + rope_theta=self._rope_theta, + device=device, + ) + self._freqs_table_cache[key] = cached + return cached + + @classmethod + def _derive_draft_model_config(cls, model_config, base: int, num_stages: int): + """Return a draft-only ``model_config`` copy with draft-specific fixes. + + Applies (1) the ``compress_ratios`` draft slice and (2) the + ``quant_config_dict`` MXFP4 extension for the draft layers' routed + experts. A single shallow copy is made (and only when something needs to + change) so the shared ``model_config`` and the target model are untouched. + + The draft MoE backend is **inherited** from the target's + ``model_config.moe_backend`` (carried by the shallow copy) — not pinned — + matching every other drafter (the MTP module reuses the V4 decoder layer, + whose MoE is built with ``moe_backend=model_config.moe_backend``; separate + Eagle3/DFlash drafts resolve it from their own config the same way). The + draft ``mtp.*`` stages are full V4 blocks, so they share the target's + MXFP4 ``n_routed_experts=384`` / ``n_group=8`` (= 48 experts/group) layout + and therefore the same backend constraints: pick a backend that supports + it (CUTLASS today, DeepGEMM megaMoE once available) on the target and the + draft follows. Note the TRTLLM-Gen ``blockScaleMoe`` routing kernel asserts + ``experts/group <= 32`` (warp size), so it is incompatible with this layout + for both the target and the draft. + """ + new_sa = cls._draft_sparse_config(model_config, base, num_stages) + new_qcd = cls._draft_quant_config_dict(model_config, base, num_stages) + if new_sa is None and new_qcd is None: + return model_config + draft_cfg = copy.copy(model_config) + # ModelConfig is a frozen dataclass; bypass the guard for these fields. + if new_sa is not None: + object.__setattr__(draft_cfg, "sparse_attention_config", new_sa) + if new_qcd is not None: + object.__setattr__(draft_cfg, "quant_config_dict", new_qcd) + return draft_cfg + + @staticmethod + def _draft_sparse_config(model_config, base: int, num_stages: int): + """Sparse-attention config sliced to the draft layers, or None if N/A. + + The inherited block remaps ``layer_idx`` to a draft-local index, so the + sparse config must expose the draft layers' per-layer ratios at indices + ``[0, num_stages)``. + """ + sa = getattr(model_config, "sparse_attention_config", None) + compress_ratios = getattr(sa, "compress_ratios", None) if sa is not None else None + if not compress_ratios or len(compress_ratios) < base + num_stages: + return None + draft_ratios = list(compress_ratios)[base : base + num_stages] + # Already draft-local (e.g. a draft-only checkpoint config); no slice. + if ( + draft_ratios == list(compress_ratios)[:num_stages] + and len(compress_ratios) == num_stages + ): + return None + return sa.model_copy(update={"compress_ratios": draft_ratios}) + + @staticmethod + def _draft_quant_config_dict(model_config, base: int, num_stages: int): + """quant_config_dict extended to cover draft-layer experts, or None. + + The checkpoint's per-module quant map only enumerates the base layers, so + ``model.layers.{base+s}.mlp.experts`` would fall back to the global fp8 + config and build fp8-shaped expert buffers. The draft routed experts are + physically MXFP4 (identical to the main MoE layers), so copy a + representative main MoE layer's experts quant onto the draft layer keys. + """ + qcd = getattr(model_config, "quant_config_dict", None) + if not qcd: + return None + src = next( + ( + qcd[f"model.layers.{li}.mlp.experts"] + for li in range(base) + if f"model.layers.{li}.mlp.experts" in qcd + ), + None, + ) + if src is None: + return None + new_qcd = dict(qcd) + changed = False + for s in range(num_stages): + key = f"model.layers.{base + s}.mlp.experts" + if new_qcd.get(key) is not src: + new_qcd[key] = src + changed = True + return new_qcd if changed else None + + @torch.inference_mode() + def write_context_windows( + self, + main_hidden: torch.Tensor, + positions: torch.Tensor, + stage_windows: torch.Tensor, + ) -> None: + """Write captured-context ``main_kv`` into the rolling per-stage KV windows. + + Replicates exactly the per-position context write that + :func:`dspark_attention_forward` performs each generation step + (``main_kv = RoPE_pos(rmsnorm(wkv @ main_x))`` written at + ``window[pos % window_size]``), but for an arbitrary set of + ``(captured-hidden, position)`` pairs. Used to (a) seed a request's + window from its prompt at prefill and (b) back-fill the intermediate + accepted tokens of a multi-accept step — both of which the per-step + generation path would otherwise leave as holes, starving the draft + attention of context (acceptance-rate only; verified decoding keeps + output correctness regardless). + + Args: + main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. + positions: ``[M]`` absolute window positions (used for BOTH the RoPE + phase and the slot ``pos % window_size``). By the generation-path + convention this is ``committed_position + 1``. Must hold at most + ``window_size`` entries with distinct slots (the caller passes a + contiguous, deduplicated range) so the scatter is well defined. + stage_windows: ``[num_stages, window_size, head_dim]`` window for one + request's slot; updated in place. + """ + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + return + M = int(main_hidden.shape[0]) + if M == 0: + return + win = int(self._attn_params["window_size"]) + rd = int(self._attn_params["rope_head_dim"]) + eps = float(self._attn_params["eps"]) + positions = positions.to(main_hidden.device).long() + # main_x is stage-invariant (stage 0's projection), matching forward_embed. + stage0 = self.mtp_layers[0] + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [M, hidden] + freqs = self._dspark_freqs_table(main_x.device)[positions] + slots = positions % win # [M] + mx = main_x.unsqueeze(0) # [1, M, hidden] for the per-position RoPE layout + for s, stage in enumerate(self.mtp_layers): + a = stage._dspark_attn + kv = _rmsnorm(F.linear(mx, a["wkv"]), a["kv_norm_w"], eps) # [1, M, head_dim] + kv = _rope_last_dims(kv, rd, freqs) # [1, M, head_dim] + stage_windows[s, slots] = kv[0].to(stage_windows.dtype) + + def write_context_windows_batched( + self, + main_hidden: torch.Tensor, + positions: torch.Tensor, + slots: torch.Tensor, + mask: torch.Tensor, + kv_windows: torch.Tensor, + ) -> None: + """CUDA-graph-safe batched + masked variant of :meth:`write_context_windows`. + + Back-fills the intermediate accepted tokens of a multi-accept step into the + rolling per-stage KV windows for ALL gen requests at once, with a fixed + ``[G, M]`` shape (``M = max interim per request``) and a validity mask + (invalid entries are no-ops via a read-modify-write), so nothing depends on + the per-request accept count. Same per-position math as the scalar version + (``RoPE_pos(rmsnorm(wkv @ main_x))`` written at ``window[pos % window]``), + but indexed/scattered through ``slots`` into the shared persistent buffer. + + Args: + main_hidden: ``[G, M, num_capture * hidden]`` captured target hiddens + (rows beyond a request's interim count are masked). + positions: ``[G, M]`` absolute window positions (RoPE phase + slot). + slots: ``[G]`` row index of each request into ``kv_windows``. + mask: ``[G, M]`` bool — which ``(g, m)`` entries are real interim writes. + kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent buffer; + updated in place. + """ + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + return + G, M = positions.shape + if G == 0 or M == 0: + return + win = int(self._attn_params["window_size"]) + rd = int(self._attn_params["rope_head_dim"]) + eps = float(self._attn_params["eps"]) + positions = positions.long() + slots = slots.long() + freqs = self._dspark_freqs_table(main_hidden.device)[positions] # [G, M, rd//2] + cols = positions % win # [G, M] + rows = slots[:, None].expand(-1, M) # [G, M] + mask3 = mask.unsqueeze(-1) # [G, M, 1] + stage0 = self.mtp_layers[0] + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [G, M, hidden] + for s, stage in enumerate(self.mtp_layers): + a = stage._dspark_attn + kv = _rmsnorm(F.linear(main_x, a["wkv"]), a["kv_norm_w"], eps) # [G, M, head_dim] + kv = _rope_last_dims_batched(kv, rd, freqs) # [G, M, head_dim] + win_s = kv_windows[:, s] # [N, win, head_dim] view onto the base buffer + # Read-modify-write so masked-out (g, m) entries keep their current + # value — a graph-safe masked scatter (no dynamic-shape compaction). + cur = win_s[rows, cols] # [G, M, head_dim] + win_s[rows, cols] = torch.where(mask3, kv.to(win_s.dtype), cur) + + def forward_embed( + self, main_hidden: torch.Tensor, bonus_token_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the draft block input and the cross-attention context. + + Args: + main_hidden: ``[num_tokens, num_capture * hidden]`` captured context. + bonus_token_ids: ``[num_tokens]`` last accepted token per request. + Returns: + x: hc-expanded block embeddings ``[num_tokens, block_size, hc_mult, hidden]`` + main_x: projected context ``[num_tokens, hidden]`` + draft_ids: block input token ids ``[num_tokens, block_size]`` (the callers + reuse these as the per-position MoE routing ids, so we return them + rather than rebuild). + """ + stage0 = self.mtp_layers[0] + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) + draft_ids = build_draft_input_ids( + bonus_token_ids, block_size=self.block_size, noise_token_id=self.noise_token_id + ) + x = self.embed_tokens(draft_ids) + x = x.unsqueeze(-2).repeat(1, 1, self.hc_mult, 1) + return x, main_x, draft_ids + + def _forward_stage( + self, + stage: "DSparkBlock", + h: torch.Tensor, + main_x: torch.Tensor, + start_pos, + freqs_cis: torch.Tensor, + moe_input_ids: torch.Tensor, + stage_window: Optional[torch.Tensor] = None, + slots: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> torch.Tensor: + """One DSpark stage = reference ``Block.forward`` with captured-context attn. + + ``h`` is the mHC residual stream ``[T, block, hc_mult, hidden]``. The mHC + ``pre_mapping``/``post_mapping`` preserve the leading ``[T, block]`` dims; + the captured-context attention and MoE run on the collapsed token axis. + Mirrors the reference (unfused) mHC boundaries exactly: + ``hc_pre → attn_norm → DSparkAttention → hc_post`` then + ``hc_pre → ffn_norm → MoE → hc_post``. + + ``stage_window`` is this stage's persistent rolling captured-context KV + window ``[T, window_size, head_dim]`` owned by the worker across decode + steps; the attention writes the current ``main_kv`` into it in place. When + ``None`` (golden / single-shot) a zero window is allocated per call. + + When ``slots`` (a ``[G]`` int tensor) is given, the CUDA-graph-safe batched + attention (:func:`dspark_attention_forward_batched`) is used instead: it + takes ``start_pos`` as a ``[G]`` tensor and writes/reads ``stage_window`` + (then shaped ``[N, window_size, head_dim]``) through the ``slots`` index. + """ + T, block, _, hidden = h.shape + + # --- attention sub-block (captured-context, not paged-KV MLA) --- + residual = h + post_mix, comb_mix, layer_input = stage.hc_attn.pre_mapping(residual) + layer_input = stage.input_layernorm(layer_input) # [T, block, hidden] + # Rolling-window cache: persist through the worker-owned ``stage_window`` + # for cross-step decode, else a fresh zero window for a single block. + persist = stage_window is not None + kv_cache = ( + stage_window + if persist + else torch.zeros( + T, + self._attn_params["window_size"], + self._attn_params["head_dim"], + dtype=torch.bfloat16, + device=h.device, + ) + ) + if slots is not None: + # Batched, CUDA-graph-safe path (start_pos is a [G] tensor; window is + # written/read through ``slots``). + attn = dspark_attention_forward_batched( + layer_input, + main_x, + start_pos, + kv_cache, + slots, + freqs_cis=freqs_cis, + persist=True, + **stage._dspark_attn, + **self._attn_params, + ) + else: + attn = dspark_attention_forward( + layer_input, + main_x, + start_pos, + kv_cache, + freqs_cis=freqs_cis, + persist=persist, + **stage._dspark_attn, + **self._attn_params, + ) + if stage.enable_fused_hc: + residual, post_mix, comb_mix, layer_input = stage.hc_ffn.fused_hc( + x_prev=attn, + residual_prev=residual, + post_mix_prev=post_mix, + comb_mix_prev=comb_mix, + norm_weight=stage.post_attention_layernorm.weight, + norm_eps=stage.post_attention_layernorm.variance_epsilon, + ) + else: + residual = stage.hc_attn.post_mapping( + x=attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + post_mix, comb_mix, layer_input = stage.hc_ffn.pre_mapping(residual) + layer_input = stage.post_attention_layernorm(layer_input) + num_tokens = T * block + # FUSED_COMM MoE backends (DeepGEMM MegaMoE) size their in-kernel + # NVLink-barrier chunk loop from ``max(all_rank_num_tokens)`` and index + # the local slice by ``moe_ep_rank``, so every EP rank must pass the same + # globally-gathered per-rank list (here: gen tokens = num_gens * block per + # rank). Passing only the local ``[num_tokens]`` desyncs the phase-flip + # barrier across ranks (hang / "unspecified launch failure"). Fall back to + # the local count for single-rank / non-ADP runs where no list is threaded. + moe_all_rank_num_tokens = ( + all_rank_num_tokens if all_rank_num_tokens is not None else [num_tokens] + ) + moe_out = stage.mlp( + layer_input.reshape(num_tokens, hidden), + input_ids=moe_input_ids, + all_rank_num_tokens=moe_all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ).reshape(T, block, hidden) + h = stage.hc_ffn.post_mapping( + x=moe_out, residual=residual, post_layer_mix=post_mix, comb_res_mix=comb_mix + ) + return h + + def forward( + self, + main_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + start_pos: int, + *, + kv_windows: Optional[torch.Tensor] = None, + temperature: float = 0.0, + confidence_threshold: float = 0.0, + return_logits: bool = False, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> tuple: + """Full block-draft forward: chain the ``num_stages`` DSpark stages. + + Mirrors the reference ``Transformer.forward_spec`` (generation path, + ``start_pos > 0``): ``forward_embed`` builds the block input + captured + context, each stage runs the captured-context backbone, and + ``forward_head`` emits the block draft tokens + per-position confidence. + + Args: + main_hidden: ``[T, num_capture * hidden]`` captured target context. + bonus_token_ids: ``[T]`` last accepted token per request. + start_pos: absolute decode position (must be > 0). + kv_windows: optional persistent per-stage rolling captured-context KV + windows ``[T, num_stages, window_size, head_dim]`` owned by the + worker; updated in place each call. ``None`` allocates fresh zero + windows (single-shot golden / test path). + Returns: + ``(draft_tokens [T, block], num_proposed [T])`` from ``forward_head``. + """ + assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + raise RuntimeError( + "DSpark attention weights not cached; call " + "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + ) + x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) + main_x = main_x.unsqueeze(1) # [T, 1, hidden] for the MQA K/V projection + freqs_cis = self._dspark_freqs_table(x.device) + moe_input_ids = draft_ids.reshape(-1) + + h = x + for s, stage in enumerate(self.mtp_layers): + stage_window = kv_windows[:, s] if kv_windows is not None else None + h = self._forward_stage( + stage, + h, + main_x, + start_pos, + freqs_cis, + moe_input_ids, + stage_window, + all_rank_num_tokens=all_rank_num_tokens, + ) + + return self.forward_head( + h, + bonus_token_ids, + temperature=temperature, + confidence_threshold=confidence_threshold, + return_logits=return_logits, + ) + + def forward_batched( + self, + main_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + start_pos: torch.Tensor, + *, + kv_windows: torch.Tensor, + slots: torch.Tensor, + temperature: float = 0.0, + confidence_threshold: float = 0.0, + return_logits: bool = False, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> tuple: + """CUDA-graph-safe batched block-draft forward (all gen requests at once). + + Same computation as :meth:`forward`, but every host-int / data-dependent + operation is tensorized so the whole path can be captured into the target's + CUDA graph (DSpark is a one-engine drafter — its worker runs inside the + graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per + gen request); the rolling captured-context windows are written/read through + ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are + gathered from a fixed table. ``forward_head`` is run with + ``confidence_threshold == 0`` (the worker proposes the full block), which is + the graph-safe branch of :func:`dspark_propose`. + + Args: + main_hidden: ``[G, num_capture * hidden]`` captured target context. + bonus_token_ids: ``[G]`` last accepted token per gen request. + start_pos: ``[G]`` int tensor of absolute decode positions (> 0). + kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent rolling + windows; written in place through ``slots``. + slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. + Returns: + ``(draft_tokens [G, block], num_proposed [G])`` from ``forward_head``. + """ + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + raise RuntimeError( + "DSpark attention weights not cached; call " + "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + ) + x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) + main_x = main_x.unsqueeze(1) # [G, 1, hidden] for the MQA K/V projection + freqs_cis = self._dspark_freqs_table(x.device) + moe_input_ids = draft_ids.reshape(-1) + + h = x + for s, stage in enumerate(self.mtp_layers): + stage_window = kv_windows[:, s] # [N, window_size, head_dim] + h = self._forward_stage( + stage, + h, + main_x, + start_pos, + freqs_cis, + moe_input_ids, + stage_window, + slots, + all_rank_num_tokens=all_rank_num_tokens, + ) + + return self.forward_head( + h, + bonus_token_ids, + temperature=temperature, + confidence_threshold=confidence_threshold, + return_logits=return_logits, + ) + + def run_moe_lockstep_noop( + self, all_rank_num_tokens: Optional[List[int]], device: torch.device + ) -> None: + """Cross the FUSED_COMM MoE NVLink barrier the same number of times as + gen-bearing ranks, for an EP rank whose local draft batch is empty. + + DeepGEMM MegaMoE (``scheduler_kind == FUSED_COMM``) synchronizes EP ranks + with an in-kernel phase-flip NVLink barrier that flips on every kernel + call, so every rank must invoke the MoE the same number of times or the + barrier desyncs (hang / "unspecified launch failure"). In the DSpark + draft only the MoE carries a cross-rank barrier (the captured-context + attention and the markov/confidence heads are per-rank), so a rank with + zero local generation requests replays just the per-stage MoE call with a + single 1-row dummy (its entry in ``all_rank_num_tokens`` is ``1``). The + scheduler runs its ``max``-derived chunk count, slicing this rank to the + 1 dummy row and zero-padding the remaining chunks, keeping the barrier + lockstep. No-op when there is no cross-rank work (single-rank / non-ADP, + or every rank is empty). + """ + if all_rank_num_tokens is None or max(all_rank_num_tokens) == 0: + return + hidden = self.config.hidden_size + # Use a 1-row dummy, NOT a 0-row tensor: DeepseekV4MoE's router / + # shared-expert dense GEMMs reject a 0-row input (cuBLAS + # CUBLAS_STATUS_INVALID_VALUE). The paired ``all_rank_num_tokens`` encodes + # 1 for this rank, so the FUSED_COMM scheduler slices to this 1 dummy row + # and still launches ``num_chunks`` cross-rank barrier crossings in + # lockstep with the gen-bearing ranks. + dummy_x = torch.zeros((1, hidden), dtype=torch.bfloat16, device=device) + dummy_ids = torch.zeros((1,), dtype=torch.long, device=device) + for stage in self.mtp_layers: + stage.mlp( + dummy_x, + input_ids=dummy_ids, + all_rank_num_tokens=all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=False), + do_finalize=True, + ) + + def forward_head( + self, + block_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + *, + temperature: float = 0.0, + confidence_threshold: float = 0.0, + return_logits: bool = False, + ) -> tuple: + """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. + + ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. + Returns (draft_tokens [*, block], num_proposed [*]); with ``return_logits`` + also returns the per-position draft logits [*, block, vocab] (§7.9 1-TV). + """ + last = self.mtp_layers[-1] + h = last.hc_head(block_hidden) + h = last.norm(h) + base_logits = self.lm_head(h) + return dspark_propose( + base_logits, + bonus_token_ids=bonus_token_ids, + block_hidden=h, + markov_head=last.markov_head, + confidence_head=last.confidence_head, + block_size=self.block_size, + temperature=temperature, + confidence_threshold=confidence_threshold, + return_logits=return_logits, + ) + + +class DSparkForCausalLM(nn.Module): + """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). + + Wraps :class:`DSparkDraftModel` (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``. + + ``embed_tokens`` / ``lm_head`` are shared with the target model + (:meth:`load_weights_from_target_model`). The draft weights live in the SAME + checkpoint under ``mtp.*``; :meth:`load_weights` remaps them + (``remap_dspark_draft_keys``), loads via ``DeepseekV4WeightLoader``, runs the + fp8 ``post_load_weights`` transforms, and caches the bf16 captured-context + attention weights from the in-memory state dict. + """ + + def __init__(self, draft_config, aux_stream_dict=None, num_stages=None, block_size=None): + super().__init__() + self.dspark_model = DSparkDraftModel( + draft_config, + aux_stream_dict, + num_stages=num_stages, + block_size=block_size, + ) + # Generic handles expected by the loader / weight mappers. + self.model = self.dspark_model + self.model_config = draft_config + self.config = draft_config.pretrained_config + # Worker-facing interface (the worker receives this wrapper as + # ``draft_model`` and calls forward()/reads these properties and scalars). + self.num_stages = self.dspark_model.num_stages + self._attn_params = self.dspark_model._attn_params + self.lm_head = None # shared from the target (load_weights_from_target_model) + self.logits_processor = None # set by the caller after construction + + @property + def block_size(self): + return self.dspark_model.block_size + + @property + def embed_tokens(self): + return self.dspark_model.embed_tokens + + def forward(self, main_hidden, bonus_token_ids, start_pos, **kwargs): + return self.dspark_model.forward(main_hidden, bonus_token_ids, start_pos, **kwargs) + + def forward_batched(self, main_hidden, bonus_token_ids, start_pos, **kwargs): + """CUDA-graph-safe batched draft forward (delegates to the draft model).""" + return self.dspark_model.forward_batched(main_hidden, bonus_token_ids, start_pos, **kwargs) + + def run_moe_lockstep_noop(self, all_rank_num_tokens, device): + """Empty-batch MoE barrier lockstep (delegates to the draft model).""" + return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) + + def write_context_windows(self, main_hidden, positions, stage_windows): + """Seed / back-fill the rolling KV windows (delegates to the draft model).""" + return self.dspark_model.write_context_windows(main_hidden, positions, stage_windows) + + def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_windows): + """Batched + masked window back-fill (delegates to the draft model).""" + return self.dspark_model.write_context_windows_batched( + main_hidden, positions, slots, mask, kv_windows + ) + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. + + ``weight_mapper`` is accepted for interface parity with the draft-weight + loader but unused: DSpark does its own ``mtp.{s}.* -> mtp_layers.{s}.*`` + remap (``remap_dspark_draft_keys``) onto the shared V4 weight loader. + """ + remapped = remap_dspark_draft_keys(weights, num_stages=self.num_stages) + logger.info( + f"[DSpark] loading {len(remapped)} draft params across {self.num_stages} stages" + ) + DeepseekV4WeightLoader(self.dspark_model).load_weights(remapped) + self.dspark_model.post_load_weights() + # bf16 captured-context attention path: dequantize the raw mtp.{s}.attn.* + # tensors for ``dspark_attention_forward``. + self.dspark_model.cache_attn_weights_from_state_dict(weights) + logger.info("[DSpark] draft weight load complete") + + def load_weights_from_target_model(self, target_model): + """Share the target's embed_tokens / lm_head (DSpark has neither).""" + if self.dspark_model.embed_tokens is None: + self.dspark_model.embed_tokens = target_model.model.embed_tokens + if self.lm_head is None: + self.lm_head = target_model.lm_head + self.dspark_model.lm_head = target_model.lm_head + + +__all__ = ["DSparkBlock", "DSparkDraftModel", "DSparkForCausalLM"] diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 88fb8979fa60..a52534de4f50 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1863,6 +1863,20 @@ def get_draft_model(model_config, draft_config, lm_head, model): if any("Laguna" in arch for arch in draft_arches): return DFlashLagunaForCausalLM(draft_config) return DFlashForCausalLM(draft_config) + 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 + num_stages = count_dspark_stages( + model_config.spec_config.speculative_model) + 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). @@ -2057,7 +2071,8 @@ def load_draft_weights(self, if self.spec_config and ( not self.spec_config.spec_dec_mode.is_external_drafter() - or self.spec_config.spec_dec_mode.is_dflash()): + or self.spec_config.spec_dec_mode.is_dflash() + or self.spec_config.spec_dec_mode.is_dspark()): self.draft_model.load_weights_from_target_model(self) def set_guided_decoder(self, diff --git a/tensorrt_llm/_torch/modules/mhc/hyper_connection.py b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py index 73f2b59c6efe..477ea75ddd66 100644 --- a/tensorrt_llm/_torch/modules/mhc/hyper_connection.py +++ b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py @@ -276,7 +276,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if not _cuda_available: raise RuntimeError("CUDA MHC kernels not available") dtype = x.dtype - x_bf16 = x.to(torch.bfloat16).contiguous() + # The CUDA head consumes a flat ``[M, mult, hidden]`` tensor. Preserve any + # leading dims (e.g. ``[batch, block, mult, hidden]`` from the DSpark draft + # block) by collapsing to a single token axis and restoring afterwards; + # the RMS-norm + weighted sum is independent per token, so this matches the + # reference ``hc_head`` which keeps the leading dims. + lead = x.shape[:-2] + x_bf16 = x.reshape(-1, self.mult, self.hidden_size).to(torch.bfloat16).contiguous() y = mhc_hc_head_cuda( x_bf16, self.fn, @@ -287,7 +293,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: norm_eps=self.norm_eps, eps=self.eps, ) - return y.to(dtype) + return y.reshape(*lead, self.hidden_size).to(dtype) def skip_forward(self, x: torch.Tensor) -> torch.Tensor: """Skip HCHead computation for pipeline parallelism on non-last ranks.""" diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index e2f708525af2..63f9a1bd65d3 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -495,6 +495,15 @@ def __init__( if config is not None: if "mla_layers" not in config.extra_attrs: config.extra_attrs["mla_layers"] = {} + suffix = 0 + # ``layer_idx`` is local to an attention stack, while this registry is shared + # by target and draft modules in one-model speculative decoding. Keep the first + # registration under ``""`` and suffix later collisions as + # ``"_"`` so custom ops resolve the originating module without + # overwriting another stack's weak reference. + while self.layer_idx_str in config.extra_attrs["mla_layers"]: + self.layer_idx_str = str(layer_idx) + f"_{suffix}" + suffix += 1 config.extra_attrs["mla_layers"][self.layer_idx_str] = weakref.ref(self) self.register_to_config = True diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d83f51a1bbe3..86ebe0492451 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2592,14 +2592,26 @@ def _get_all_rank_ctx_requests(self, num_ctx_requests: int): return None def _set_spec_metadata_all_rank_num_tokens( - self, spec_metadata: SpecMetadata, + self, + spec_metadata: SpecMetadata, spec_all_rank_num_tokens: List[int], - all_rank_num_seqs: List[int]) -> None: + all_rank_num_seqs: List[int], + all_rank_num_gens: Optional[List[int]] = None) -> None: # Eagle3 / MTP-eagle one-model use subseq_all_rank_num_tokens for # draft loop iterations i>0 (per-sequence counts, since each # sequence contributes one token per iteration). spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens spec_metadata.all_rank_num_seqs = all_rank_num_seqs + # DSpark can draft only after the target processes the current bonus token, + # because it consumes captured target-layer hidden states for that token. + # Prefill computes hidden states for prompt tokens; the first generated token + # is sampled from the last prompt logits and has not itself passed through the + # target layers. Thus context requests seed the rolling window but do not run + # the draft. On mixed steps, num_seqs therefore over-counts the draft MoE + # workload; gen-only per-rank counts keep the FUSED_COMM (DeepGEMM MegaMoE) + # chunk loop identical across EP ranks. + if all_rank_num_gens is not None: + spec_metadata.all_rank_num_gens = all_rank_num_gens if (spec_metadata.spec_dec_mode.is_mtp_eagle_one_model() or spec_metadata.spec_dec_mode.is_eagle3_one_model()): spec_metadata.subseq_all_rank_num_tokens = all_rank_num_seqs @@ -2928,12 +2940,14 @@ def _prepare_incremental_update_metadata( # Handle distributed spec metadata if enable_attention_dp: sequence_lengths = spec_metadata.seq_lens - all_rank_num_tokens = self.dist.tp_cp_allgather( - [spec_metadata.num_tokens, - len(sequence_lengths)]) + all_rank_num_tokens = self.dist.tp_cp_allgather([ + spec_metadata.num_tokens, + len(sequence_lengths), attn_metadata.num_generations + ]) self._set_spec_metadata_all_rank_num_tokens( spec_metadata, [item[0] for item in all_rank_num_tokens], - [item[1] for item in all_rank_num_tokens]) + [item[1] for item in all_rank_num_tokens], + [item[2] for item in all_rank_num_tokens]) # Set iteration states - batch dictionary updates self.iter_states.update({ @@ -4503,12 +4517,14 @@ def previous_seq_slots_device(): inputs['spec_metadata'] = spec_metadata if self.enable_attention_dp: - all_rank_num_tokens = self.dist.tp_cp_allgather( - [spec_metadata.num_tokens, - len(sequence_lengths)]) + all_rank_num_tokens = self.dist.tp_cp_allgather([ + spec_metadata.num_tokens, + len(sequence_lengths), spec_metadata.num_generations + ]) self._set_spec_metadata_all_rank_num_tokens( spec_metadata, [item[0] for item in all_rank_num_tokens], - [item[1] for item in all_rank_num_tokens]) + [item[1] for item in all_rank_num_tokens], + [item[2] for item in all_rank_num_tokens]) if mm_token_indices is not None: self._ship_multimodal_indices( @@ -4747,14 +4763,15 @@ def _prepare_tp_inputs_no_cache( if spec_metadata is not None: all_rank_num_tokens = self.dist.tp_cp_allgather([ attn_metadata.num_tokens, spec_metadata.num_tokens, - len(sequence_lengths) + len(sequence_lengths), spec_metadata.num_generations ]) attn_metadata.all_rank_num_tokens = [ item[0] for item in all_rank_num_tokens ] self._set_spec_metadata_all_rank_num_tokens( spec_metadata, [item[1] for item in all_rank_num_tokens], - [item[2] for item in all_rank_num_tokens]) + [item[2] for item in all_rank_num_tokens], + [item[3] for item in all_rank_num_tokens]) else: all_rank_num_tokens = self.dist.tp_cp_allgather( attn_metadata.num_tokens) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py new file mode 100644 index 000000000000..f2597f5e9ef3 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -0,0 +1,564 @@ +# 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 worker / metadata mirror the DFlash plumbing (capture target-layer +# 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. + +from collections import deque +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional + +import torch + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +from .interface import SpecMetadata, SpecWorkerBase + +if TYPE_CHECKING: + from ...llmapi.llm_args import DSparkDecodingConfig + + +@dataclass +class DSparkSpecMetadata(SpecMetadata): + """Metadata for DSpark speculative decoding. + + Captures hidden states from the target model's ``layers_to_capture`` during + 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``) + as the captured-context attention input (``main_x``). + + Mirrors :class:`DFlashSpecMetadata`; the only DSpark-specific detail is that + the per-layer captured width is the model hidden size (post hc-mean), so the + buffer is ``[max_num_tokens, hidden_size * num_capture_layers]``. + """ + + batch_indices_cuda: Optional[torch.Tensor] = None + + # Hidden state capture fields + layers_to_capture: Optional[List[int]] = None + hidden_size: int = 0 + max_num_tokens: int = 0 + dtype: torch.dtype = torch.bfloat16 + captured_hidden_states: Optional[torch.Tensor] = None + + def __post_init__(self): + self.batch_indices_cuda = torch.empty( + [self.max_num_requests], + dtype=torch.int, + device="cuda", + ) + + self.is_spec_dec_tree = False + self.is_spec_dec_dynamic_tree = False + + # Set up hidden state capture buffer + if self.layers_to_capture is not None and len(self.layers_to_capture) > 0: + self.layers_to_capture = sorted(list(self.layers_to_capture)) + self.num_capture_layers = len(self.layers_to_capture) + # O(1) lookups for is_layer_capture() and maybe_capture_hidden_states() + self._capture_layer_set = frozenset(self.layers_to_capture) + self._layer_to_idx = {lid: i for i, lid in enumerate(self.layers_to_capture)} + self.captured_hidden_states = torch.empty( + (self.max_num_tokens, self.hidden_size * self.num_capture_layers), + dtype=self.dtype, + device="cuda", + ) + logger.info( + f"DSpark: capturing hidden states from layers {self.layers_to_capture}, " + f"buffer shape {self.captured_hidden_states.shape}" + ) + else: + self.num_capture_layers = 0 + self._capture_layer_set = frozenset() + self._layer_to_idx = {} + + def prepare(self): + assert self.request_ids is not None + num_seqs = len(self.request_ids) + batch_indices = torch.arange( + num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned() + ) + self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + + # CUDA-graph-safe path: maintain the request->slot mapping on the host + # (outside the captured region) and mirror it into ``_batch_to_slot`` so the + # captured gen forward can index the rolling windows by tensor. Mirrors + # ``DFlashSpecMetadata.prepare`` (dflash.py:96-113). + worker = getattr(self, "_dspark_worker", None) + if worker is not None and worker._win_inited: + current = set(self.request_ids) + for rid in list(worker._req_to_slot.keys()): + if rid not in current: + slot = worker._req_to_slot.pop(rid) + worker._ctx_len[slot] = 0 + worker._kv_windows[slot].zero_() + worker._free_slots.append(slot) + # Unknown request IDs (e.g. synthetic warmup requests) default to slot 0. + mapping = torch.tensor( + [worker._req_to_slot.get(rid, 0) for rid in self.request_ids], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ) + worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) + + def is_layer_capture(self, layer_id: int) -> bool: + return layer_id in self._capture_layer_set + + 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. + + DeepSeek-V4 keeps the multi-head (mHC) residual stream flattened as + ``[num_tokens, hc_mult * hidden]``; DSpark captures the *mean over the hc + streams* (reference ``h.mean(dim=2)`` with ``h`` shaped + ``[*, hc_mult, hidden]``). We reduce here so the V4 decoder layer's + existing capture call is unchanged. A ``[num_tokens, hidden]`` input + (already reduced / non-mHC) is stored as-is. + """ + if self.captured_hidden_states is None: + return + i = self._layer_to_idx.get(layer_id) + if i is not None: + num_tokens = hidden_states.shape[0] + to_save = hidden_states + residual if residual is not None else hidden_states + # mHC residual -> mean over the hc_mult streams. + if to_save.shape[-1] != self.hidden_size: + hc_mult = to_save.shape[-1] // self.hidden_size + to_save = to_save.reshape(num_tokens, hc_mult, self.hidden_size).mean(dim=1) + self.captured_hidden_states[ + :num_tokens, i * self.hidden_size : (i + 1) * self.hidden_size + ].copy_(to_save, non_blocking=True) + + def get_hidden_states(self, num_tokens: int) -> Optional[torch.Tensor]: + """Get captured hidden states (all layers concatenated).""" + if self.captured_hidden_states is None: + return None + return self.captured_hidden_states[ + :num_tokens, : self.hidden_size * self.num_capture_layers + ] + + +class DSparkWorker(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 + 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 + per-position acceptance confidence used to truncate the proposed prefix. + + Unlike DFlash, the draft does NOT use the paged KV cache or mask-token + cross-attention: its attention K/V come from the worker-owned rolling window + of projected captured context (one ``main_kv`` per decode step, per stage). + Acceptance of the previous block goes through the unified + :meth:`SpecWorkerBase.sample_and_accept_draft_tokens` (strict target-verify, + or rejection sampling for a non-greedy batch), so greedy parity with no-spec + is preserved regardless of draft quality. + + 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 + write done by the generation path. These affect draft acceptance rate only, + not correctness, which the standard target verify guarantees. + + Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). + """ + + def __init__( + self, + spec_config: "DSparkDecodingConfig", + mapping: Mapping, + use_separate_draft_kv_cache: bool = False, + ): + super().__init__(use_separate_draft_kv_cache) + self.spec_config = spec_config + self.mapping = mapping + + # Per-slot rolling captured-context KV windows, built lazily on the + # first forward (fixed-size for slot-indexed reads/writes). + self._win_inited = False + self._kv_windows: Optional[torch.Tensor] = None # [max_batch, num_stages, win, hd] + self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position + self._win = 0 + + # Slot management. ``_req_to_slot`` (python dict) + ``_free_slots`` are the + # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the + # CUDA mirror (request-order -> slot) read by the CUDA-graph-safe batched + # gen path (set on the host in prepare(), so the captured forward indexes + # the rolling windows through a tensor instead of a python dict lookup). + self._req_to_slot = {} # request_id -> slot index + self._free_slots = deque() # available slot indices + self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda + + # The generation draft path is the batched, host-sync-free + # ``_draft_gen_block_batched`` + ``DSparkDraftModel.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"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" + ) + + @property + def max_draft_len(self) -> int: + return self.spec_config.max_draft_len + + def _lazy_init(self, draft_model, spec_metadata) -> None: + block_size = int(draft_model.block_size) + if block_size != self.max_draft_len: + raise ValueError( + "DSpark draft model block_size must equal worker max_draft_len; " + f"got block_size={block_size} and max_draft_len={self.max_draft_len}" + ) + + if self._win_inited: + return + max_batch = spec_metadata.max_num_requests + num_stages = draft_model.num_stages + self._win = int(draft_model._attn_params["window_size"]) + head_dim = int(draft_model._attn_params["head_dim"]) + + self._kv_windows = torch.zeros( + (max_batch, num_stages, self._win, head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + self._ctx_len = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._free_slots = deque(range(max_batch)) + self._req_to_slot = {} + self._win_inited = True + logger.info( + f"DSpark: allocated rolling KV windows " + f"[{max_batch}, {num_stages}, {self._win}, {head_dim}]" + ) + + def _assign_slot(self, req_id: int, reset: bool) -> int: + """Get (or refresh) the slot for a request; reset clears its window.""" + if reset and req_id in self._req_to_slot: + old = self._req_to_slot.pop(req_id) + self._ctx_len[old] = 0 + self._kv_windows[old].zero_() + self._free_slots.append(old) + if req_id not in self._req_to_slot: + if not self._free_slots: + raise RuntimeError( + "DSpark has no free rolling-window slots for request " + f"{req_id}; increase max_num_requests" + ) + slot = self._free_slots.popleft() + self._req_to_slot[req_id] = slot + self._ctx_len[slot] = 0 + self._kv_windows[slot].zero_() + return self._req_to_slot[req_id] + + def _seed_context_windows( + self, + draft_model, + spec_metadata: "DSparkSpecMetadata", + attn_metadata, + position_ids: torch.Tensor, + total_target_tokens: int, + ) -> None: + """Seed context chunks using their absolute positions. + + A request can arrive in multiple prefill chunks. Only its first chunk + starts at position zero and resets the persistent rolling window; + continuation chunks append to the same request slot. + """ + captured = spec_metadata.get_hidden_states(total_target_tokens) + flat_position_ids = position_ids.reshape(-1) + context_offset = 0 + for i in range(attn_metadata.num_contexts): + chunk_len = int(attn_metadata._seq_lens[i]) + chunk_positions = flat_position_ids[context_offset : context_offset + chunk_len].long() + if chunk_len == 0: + context_offset += chunk_len + continue + + req_id = spec_metadata.request_ids[i] + first_position = int(chunk_positions[0].item()) + slot = self._assign_slot(req_id, reset=first_position == 0) + self._ctx_len[slot] = chunk_positions[-1] + 1 + + if captured is not None: + keep = min(self._win, chunk_len) + hidden = captured[context_offset + chunk_len - keep : context_offset + chunk_len] + # A prompt token at absolute position p is stored in frame p+1, + # matching the generation path's start_pos convention. + window_positions = chunk_positions[-keep:] + 1 + draft_model.write_context_windows(hidden, window_positions, self._kv_windows[slot]) + context_offset += chunk_len + + def _draft_gen_block_batched( + self, + draft_model, + spec_metadata: "DSparkSpecMetadata", + attn_metadata, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + total_target_tokens: int, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> torch.Tensor: + """CUDA-graph-safe batched gen draft (all gen requests in one forward). + + Free of host syncs and data-dependent shapes: per-request quantities + (``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 + 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. + """ + num_gens = batch_size - num_contexts + K = self.max_draft_len + Kp1 = K + 1 + device = accepted_tokens.device + + if num_gens == 0: + return None + captured = spec_metadata.get_hidden_states(total_target_tokens) + if captured is None: + return None + + # gen-only graph batches have num_ctx_tokens == 0; mixed eager batches put + # the gen tokens after the context tokens. + gen_start = attn_metadata.num_ctx_tokens + slots = self._batch_to_slot[num_contexts:batch_size] # [G] + nacc = num_accepted_tokens[num_contexts:batch_size].long() # [G] + gidx = nacc - 1 # [G] index of the bonus within each verified prefix + + # Bonus token = last accepted token of the verified prefix. + bonus = ( + accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long() + ) # [G] + + # Captured target hidden at the bonus position within each request's Kp1 + # processed tokens. + arange_g = torch.arange(num_gens, device=device) + base = gen_start + arange_g * Kp1 # [G] + main_hidden = captured[base + gidx] # [G, ncap*hidden] + + # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens + # (everything but the bonus) into the rolling window — same frames as the + # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. + old = self._ctx_len[slots] # [G] pre-increment decode position + j = torch.arange(K, device=device) # [K] + interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] + interim_pos = old.unsqueeze(1) + 1 + j.unsqueeze(0) # [G, K] + interim_base = (base.unsqueeze(1) + j.unsqueeze(0)).clamp( + min=0, max=captured.shape[0] - 1 + ) # [G, K] (clamped; invalid entries are masked out anyway) + interim_hidden = captured[interim_base] # [G, K, ncap*hidden] + draft_model.write_context_windows_batched( + interim_hidden, interim_pos, slots, interim_valid, self._kv_windows + ) + + # Advance the decode position by the accepted count; start_pos (= post- + # increment ctx_len) matches the eager path's frame value. + start_pos = old + nacc # [G] + self._ctx_len[slots] = start_pos + + # Surface the per-position corrected block logits ([num_gens, K, vocab]) + # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) + # sampling + TP gather + draft_probs scatter, rather than argmaxing here. + _toks, _num_proposed, block_logits = draft_model.forward_batched( + main_hidden, + bonus, + start_pos, + kv_windows=self._kv_windows, + slots=slots, + temperature=0.0, + confidence_threshold=0.0, + return_logits=True, + all_rank_num_tokens=all_rank_num_tokens, + ) + return block_logits + + def forward( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + raw_logits = logits + K = self.max_draft_len + + self._lazy_init(draft_model, spec_metadata) + # Backref so DSparkSpecMetadata.prepare() can maintain the host slot map + # and mirror it into _batch_to_slot for the CUDA-graph-safe gen path. + spec_metadata._dspark_worker = self + self._execute_guided_decoder_if_present(logits) + + # Target-verify acceptance via the unified SpecWorkerBase entry: it + # reshapes the stored draft tokens (default (num_gens, runtime_draft_len) + # hook), then routes to strict or rejection sampling. Greedy parity with + # the previous hand-rolled path is preserved (rejection only engages for a + # non-greedy batch with valid draft_probs). + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + total_target_tokens = input_ids.shape[0] + + # CUDA-graph warmup guard: the warmup forwards (is_cuda_graph set, stream + # NOT yet capturing) run synthetic gen batches that would otherwise advance + # the persistent rolling-window state. Snapshot and restore it so warmup is + # side-effect-free. (During the capture pass itself the stream IS capturing, + # so we skip the save/restore and let the ops be recorded; real requests + # reset their slot's window+ctx_len at prefill, wiping any capture-time + # mutation.) + is_warmup = ( + getattr(spec_metadata, "is_cuda_graph", False) + and not torch.cuda.is_current_stream_capturing() + ) + if is_warmup: + saved_ctx_len = self._ctx_len.clone() + saved_windows = self._kv_windows.clone() + + # Assign / reset window slots for context (prefill) requests and seed each + # request's rolling KV window from its prompt's captured context, so the + # first generation step drafts against real context instead of an all-zero + # window (acceptance-rate only; verified decoding keeps output correct). + if num_contexts > 0: + self._seed_context_windows( + draft_model, + spec_metadata, + attn_metadata, + position_ids, + total_target_tokens, + ) + + # FUSED_COMM MoE backends (DeepGEMM MegaMoE) synchronize EP ranks with an + # in-kernel phase-flip NVLink barrier that flips on every kernel call, so + # every rank must invoke the draft MoE the same number of times and with + # the same globally-gathered per-rank token list, or the barrier desyncs + # (hang / "unspecified launch failure"). The draft runs over generation + # requests only, each expanded to ``block`` positions, so the per-rank + # draft-MoE token count is ``num_gens * block``. ``all_rank_num_gens`` is + # gathered at metadata-prep time (model_engine, outside any CUDA-graph + # capture region); it is None for non-ADP / single-rank runs, where the + # local ``[num_tokens]`` fallback in ``_forward_stage`` is correct. + block = int(draft_model.block_size) + all_rank_num_gens = getattr(spec_metadata, "all_rank_num_gens", None) + # A rank with zero local gen requests still has to cross the draft MoE's + # cross-rank barrier, but DeepseekV4MoE's router / shared-expert dense + # GEMMs reject a 0-row input (cuBLAS CUBLAS_STATUS_INVALID_VALUE), so such + # a rank runs a single 1-row dummy through the MoE (like ADP padding). + # Encode that as ``1`` in the globally-shared per-rank token list so every + # rank agrees on the FUSED_COMM chunk count and per-rank slice. + all_rank_draft_tokens = ( + [max(1, int(g) * block) for g in all_rank_num_gens] + if all_rank_num_gens is not None + else None + ) + global_has_gen = ( + max(all_rank_num_gens) > 0 if all_rank_num_gens is not None else num_gens > 0 + ) + + if num_gens > 0: + # The batched gen-block draft returns the per-position corrected block + # logits [num_gens, K, vocab] and is CUDA-graph-safe. + gen_logits = self._draft_gen_block_batched( + draft_model, + spec_metadata, + attn_metadata, + accepted_tokens, + num_accepted_tokens, + num_contexts, + batch_size, + total_target_tokens, + all_rank_num_tokens=all_rank_draft_tokens, + ) + if gen_logits is not None: + # SpecWorkerBase samples the draft tokens (greedy argmax, or + # rejection sampling for a non-greedy batch), performs the TP + # gather, and scatters the proposal distribution into draft_probs. + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, spec_metadata, batch_size, num_contexts=num_contexts + ) + # The context one-hot must match the width the gen scatter just + # published to draft_probs, NOT gen_logits.shape[-1]: under TP the + # draft logits are vocab-sharded and sample_draft_tokens gathers + # them to full vocab before scattering, so the pre-gather shard + # width would leave stale columns and corrupt rejection. + gen_vocab = spec_metadata.draft_probs_last_dim + else: + gen_draft_tokens = torch.zeros((num_gens, K), dtype=torch.int32, device="cuda") + gen_vocab = None + else: + # No local generation requests: if any peer EP rank has some, we must + # still cross the draft MoE's cross-rank barrier the same number of + # times (zero-token) so a FUSED_COMM phase-flip barrier stays lockstep. + if global_has_gen: + draft_model.run_moe_lockstep_noop(all_rank_draft_tokens, accepted_tokens.device) + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + gen_vocab = None + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a legal one-hot so they are + # a valid distribution when they become gen requests next iteration. + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + + if num_contexts > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + else: + next_draft_tokens = gen_draft_tokens + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + if is_warmup: + self._ctx_len.copy_(saved_ctx_len) + self._kv_windows.copy_(saved_windows) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 48ae77f716e8..5ad8f374d4ce 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -111,6 +111,10 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool: return False if not spec_config.spec_dec_mode.use_one_engine(): 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(): + return False return spec_config._allow_separate_draft_kv_cache @@ -276,6 +280,7 @@ class SpeculativeDecodingMode(IntEnum): SAVE_HIDDEN_STATES = auto() PARD = auto() DFLASH = auto() + DSPARK = auto() NONE = auto() AUTO = auto() @@ -310,8 +315,11 @@ def is_pard(self): def is_dflash(self): return self == SpeculativeDecodingMode.DFLASH + def is_dspark(self): + return self == SpeculativeDecodingMode.DSPARK + def is_parallel_draft(self): - return self.is_pard() or self.is_dflash() + return self.is_pard() or self.is_dflash() or self.is_dspark() def is_ngram(self): return self == SpeculativeDecodingMode.NGRAM @@ -476,6 +484,11 @@ class SpecMetadata: # The number of sequences for speculative model/layer of different rank all_rank_num_seqs: Optional[List[int]] = None + # The number of generation requests for the speculative model/layer of each + # rank (num_seqs - num_contexts). Used by external drafters (e.g. DSpark) + # whose draft forward processes only generation requests and must size a + # FUSED_COMM MoE (DeepGEMM MegaMoE) chunk loop identically across EP ranks. + all_rank_num_gens: Optional[List[int]] = None # The number of extra kv tokens # Some speculative decoding methods need to use different kv lengths for the # draft/target layers. But KVCacheManager can only support kv caches with the diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 4f9c5af8846a..0c04e9829608 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -20,6 +20,7 @@ from .draft_target import (DraftTargetOneModelSampler, DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) +from .dspark import DSparkSpecMetadata, DSparkWorker from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, @@ -203,6 +204,21 @@ def get_spec_metadata(spec_config, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, ) + if spec_config.spec_dec_mode.is_dspark(): + target_layer_ids = getattr(spec_config, 'target_layer_ids', None) + return DSparkSpecMetadata( + max_draft_len=spec_config.max_draft_len, + max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, + spec_dec_mode=spec_config.spec_dec_mode, + max_num_requests=max_num_requests, + layers_to_capture=target_layer_ids, + hidden_size=model_config.hidden_size, + max_num_tokens=max_num_tokens, + dtype=model_config.torch_dtype, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, + ) if spec_config.spec_dec_mode.is_draft_target_one_model(): return DraftTargetOneModelSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -465,6 +481,8 @@ 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) + if spec_dec_mode.is_dspark(): + return DSparkWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_sa(): return SAWorker(spec_config, model_config) if spec_dec_mode.is_draft_target_one_model(): diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 019b97913d38..e4ec7ac25137 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -12,12 +12,13 @@ CudaGraphConfig, DecodeCudaGraphConfig, DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, - DraftTargetDecodingConfig, DynamicBatchConfig, - Eagle3DecodingConfig, EagleDecodingConfig, - EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, - MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, NGramDecodingConfig, + DraftTargetDecodingConfig, DSparkDecodingConfig, + DynamicBatchConfig, Eagle3DecodingConfig, + EagleDecodingConfig, EncodeCudaGraphConfig, + ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, + LookaheadDecodingConfig, MedusaDecodingConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, + MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -67,6 +68,7 @@ 'NGramDecodingConfig', 'PARDDecodingConfig', 'DFlashDecodingConfig', + 'DSparkDecodingConfig', 'SADecodingConfig', 'SAEnhancerConfig', 'UserProvidedDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 57c3de25dcfe..65054a175f00 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2629,6 +2629,91 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.DFLASH +class DSparkDecodingConfig(DecodingBaseConfig): + """Configuration for DSpark speculative decoding. + + DSpark (DeepSeek) is a target-dependent, "semi-parallel" speculative + decoding method. Like DFlash it captures hidden states from several target + layers as cross-attention context and drafts a whole block in a single + backbone forward, but it additionally refines the per-position draft logits + with a lightweight sequential head (a low-rank Markov head, optionally an RNN + head) and predicts an acceptance-confidence per position to truncate the + proposed prefix. + + Key features: + - Target-dependent: captures hidden states from ``target_layer_ids``. + - Semi-parallel: one block backbone forward + cheap sequential head refine. + - Confidence head: truncates the proposed draft length (NOT the accept rule; + acceptance stays standard target verification, preserving greedy parity). + + Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). + """ + mask_token_id: Optional[int] = Field( + default=None, + description= + "Token ID used as the mask/noise token for parallel draft prediction. " + "If None, read from the draft model config (dspark_noise_token_id).") + + target_layer_ids: Optional[List[int]] = Field( + default=None, + description= + "Target model layer indices whose hidden states are captured for " + "cross-attention in the draft model. If None, read from the draft model " + "config (dspark_target_layer_ids).") + + block_size: Optional[PositiveInt] = Field( + default=None, + description= + "Number of draft positions produced per block. If None, read from the " + "draft model config (dspark_block_size). Should equal max_draft_len.") + + markov_rank: Optional[int] = Field( + default=None, + description= + "Low-rank dimension of the Markov head logit-bias. If None, read from " + "the draft model config (dspark_markov_rank). 0 disables the head.") + + markov_head_type: Optional[Literal["vanilla", "gated", "rnn"]] = Field( + default=None, + description= + "Type of the sequential refinement head used within a block. If None, " + "read from the draft model config (dspark_markov_head_type), " + "defaulting to \"vanilla\".") + + # NOTE: confidence-based dynamic drafting (the draft model's confidence head + # that truncates the proposed block) is NOT enabled in this PR. The user-facing + # ``enable_confidence_head`` / ``confidence_threshold`` knobs are intentionally + # omitted and will be added when the feature is actually wired into the + # speculative scheduling/verification path. The confidence head module and its + # internal plumbing remain as scaffolding (see DSparkConfidenceHead / + # dspark_propose). + + decoding_type: Literal["DSpark"] = Field(default="DSpark") + + @model_validator(mode="after") + def set_max_total_draft_tokens(self): + self.max_total_draft_tokens = self.max_draft_len + return self + + @property + def tokens_per_gen_step(self) -> int: + """DSpark needs K+1 tokens per gen request (K drafts + 1 bonus). + + The draft produces its own mask queries internally; passing mask + fillers through the target is pure wasted work at large batch size. + """ + return self.max_draft_len + 1 + + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + + @functools.cached_property + def spec_dec_mode(self): + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + return TorchSpeculativeDecodingMode.DSPARK + + class AutoDecodingConfig(DecodingBaseConfig): """Configuration for auto speculative decoding. @@ -3295,6 +3380,7 @@ def supports_backend(self, backend: str) -> bool: SaveHiddenStatesDecodingConfig, PARDDecodingConfig, DFlashDecodingConfig, + DSparkDecodingConfig, AutoDecodingConfig, ], Field(discriminator="decoding_type"), @@ -5105,6 +5191,7 @@ def validate_speculative_config(self): or spec_mode == TorchSpeculativeDecodingMode.MTP or spec_mode.is_mtp_eagle_one_model() or spec_mode.is_pard() or spec_mode.is_dflash() + or spec_mode.is_dspark() or spec_mode.is_draft_target_one_model()) # Combinations that break the proposal-distribution invariant. @@ -5114,6 +5201,7 @@ def validate_speculative_config(self): is_new_rejection_method = ( spec_mode == TorchSpeculativeDecodingMode.MTP or spec_mode.is_pard() or spec_mode.is_dflash() + or spec_mode.is_dspark() or spec_mode.is_draft_target_one_model()) # Plain tensor parallelism is supported (the draft path # all-gathers vocab-sharded draft logits before rejection, see @@ -5231,6 +5319,92 @@ def validate_speculative_config(self): if mask_id is not None: self.speculative_config.mask_token_id = mask_id + if isinstance(self.speculative_config, DSparkDecodingConfig): + spec_cfg = self.speculative_config + 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. + 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).") + # 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. + 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", {}) + + 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 + + # The checkpoint's ``dspark_target_layer_ids`` is + # authoritative: it fixes both which target hidden states + # were captured during draft training and the input width of + # ``main_proj`` (hidden_size * num_capture_layers). If the + # user leaves ``target_layer_ids`` unset we copy it verbatim + # (order preserved, since the projection columns are + # order-dependent). An explicit override that does not match + # the checkpoint list exactly would either mismatch the + # projection shape at runtime (different count) or feed the + # draft hidden states it was not trained on (same count, + # different layers), so reject it during validation. + ckpt_layer_ids = _dspark_get("target_layer_ids", + "dspark_target_layer_ids") + if spec_cfg.target_layer_ids is None: + if ckpt_layer_ids is not None: + spec_cfg.target_layer_ids = list(ckpt_layer_ids) + elif ckpt_layer_ids is not None and list( + spec_cfg.target_layer_ids) != list(ckpt_layer_ids): + raise ValueError( + "DSpark target_layer_ids must match the checkpoint's " + "dspark_target_layer_ids exactly (the draft " + "projection weights are trained for that specific " + f"layer set); got override {spec_cfg.target_layer_ids} " + f"but the checkpoint specifies {list(ckpt_layer_ids)}. " + "Leave target_layer_ids unset to use the checkpoint " + "value.") + if spec_cfg.mask_token_id is None: + mask_id = _dspark_get("mask_token_id", + "dspark_noise_token_id") + if mask_id is not None: + spec_cfg.mask_token_id = mask_id + if spec_cfg.block_size is None: + block_size = _dspark_get("block_size", + "dspark_block_size") + if block_size is not None: + spec_cfg.block_size = int(block_size) + if spec_cfg.markov_rank is None: + markov_rank = _dspark_get("markov_rank", + "dspark_markov_rank") + if markov_rank is not None: + spec_cfg.markov_rank = markov_rank + if spec_cfg.block_size is None: + spec_cfg.block_size = spec_cfg.max_draft_len + if spec_cfg.block_size <= 0: + raise ValueError( + "DSpark block_size must be greater than zero") + if spec_cfg.block_size != spec_cfg.max_draft_len: + raise ValueError( + "DSpark block_size must equal max_draft_len; got " + f"block_size={spec_cfg.block_size} and " + f"max_draft_len={spec_cfg.max_draft_len}") + if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size if pool_size is not None and self.max_batch_size is not None: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ca6e762490cf..28760a844f9d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1607,10 +1607,18 @@ "kind": "value", "path": "speculative_config.begin_thinking_phase_token" }, + { + "allowed_values": [], + "annotation": "Optional[Annotated[int, Gt(gt=0)]]", + "converter": "", + "kind": "value", + "path": "speculative_config.block_size" + }, { "allowed_values": [ "AUTO", "DFlash", + "DSpark", "Draft_Target", "Eagle3", "Eagle", @@ -1715,6 +1723,24 @@ "kind": "value", "path": "speculative_config.is_use_oldest" }, + { + "allowed_values": [ + "vanilla", + "gated", + "rnn" + ], + "annotation": "Optional[Literal['vanilla', 'gated', 'rnn']]", + "converter": "", + "kind": "categorical", + "path": "speculative_config.markov_head_type" + }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "speculative_config.markov_rank" + }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index facfca89379b..65202061ad8f 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -155,6 +155,12 @@ deepseek-ai/DeepSeek-V4-Pro: kv_cache_quant_algo: FP8 spec_dec_algo: MTP accuracy: 96.0 + # Full GSM8K DEP8 DSpark path: TP=8, EP=8, attention DP, MegaMoE DeepGEMM, + # and max_draft_len=5. Measured 96.475 on 1319 samples; floor set to 96.0 + # for run-to-run margin. + - quant_algo: FP8_BLOCK_SCALES + spec_dec_algo: DSpark + accuracy: 96.0 Qwen3/Qwen3-4B: - spec_dec_algo: Eagle accuracy: 85.823 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2305c0852664..550ff8641d80 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -30,11 +30,12 @@ # isort: off from tensorrt_llm.llmapi import ( AttentionDpConfig, CudaGraphConfig, DeepSeekSparseAttentionConfig, - DFlashDecodingConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, - KvCacheConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, - NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, - SADecodingConfig, SamplingParams, SchedulerConfig, - SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) + DFlashDecodingConfig, DSparkDecodingConfig, DraftTargetDecodingConfig, + Eagle3DecodingConfig, KvCacheConfig, MiniMaxM3SparseAttentionConfig, + MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, + RocketSparseAttentionConfig, SADecodingConfig, SamplingParams, + SchedulerConfig, SkipSoftmaxAttentionConfig, SAEnhancerConfig, + TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -3882,6 +3883,12 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): mtp_nextn=mtp_nextn) +_DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT = ( + "Solve the problem carefully. End your response with a final line exactly " + "in the form #### , using the simplest numeric form without units " + "or trailing zeros.") + + @pytest.mark.timeout(14400) @pytest.mark.skip_less_device(8) @pytest.mark.skip_less_device_memory(140000) @@ -3912,6 +3919,52 @@ def test_gsm8k_full_accuracy(self): f"{acc_params.ref_accuracy:.3f}") +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4ProDSpark(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Pro" + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Pro-DSpark" + EXTRA_EVALUATOR_KWARGS = dict( + apply_chat_template=True, + system_prompt=_DEEPSEEK_V4_GSM8K_SYSTEM_PROMPT, + ) + + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_dep8_megamoe_deepgemm(self): + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + spec_config = DSparkDecodingConfig(max_draft_len=5, + speculative_model=self.MODEL_PATH) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, + disable_overlap_scheduler=True, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + acc_params = task.get_hypothesis_testing_params( + dtype=llm.args.dtype, + quant_algo=llm.args.quant_config.quant_algo, + kv_cache_quant_algo=llm.args.quant_config.kv_cache_quant_algo, + spec_dec_algo=llm.args.speculative_config.decoding_type) + assert acc_params.num_samples == GSM8K.NUM_SAMPLES + with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): + score = task.evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} is below recorded reference " + f"{acc_params.ref_accuracy:.3f}") + + @pytest.mark.timeout(14400) @pytest.mark.skip_less_device_memory(140000) @skip_pre_blackwell diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d5d316518f22..4a76ae952966 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -174,6 +174,7 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8] TIMEOUT (60) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm TIMEOUT (240) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_on-cutedsl] TIMEOUT (60) - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] TIMEOUT (60) diff --git a/tests/unittest/_torch/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py new file mode 100644 index 000000000000..908069a02560 --- /dev/null +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -0,0 +1,82 @@ +# 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. + +from unittest.mock import patch + +import torch +from torch import nn + +from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.mla import MLA +from tensorrt_llm.functional import PositionEmbeddingType + + +class _FakeAttention(nn.Module): + def support_fused_rope(self) -> bool: + return True + + def update_quant_config(self, _quant_config: object) -> None: + pass + + +def _make_mla(config: ModelConfig) -> MLA: + position_embedding = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=2, max_positions=8), + ) + return MLA( + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + qk_nope_head_dim=2, + qk_rope_head_dim=2, + v_head_dim=2, + q_lora_rank=4, + kv_lora_rank=4, + predicted_tokens_per_seq=1, + max_position_embeddings=8, + bias=False, + pos_embd_params=position_embedding, + layer_idx=0, + dtype=torch.bfloat16, + config=config, + o_lora_rank=2, + ) + + +def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: + target_config = ModelConfig(skip_create_weights_in_init=True) + draft_config = ModelConfig(skip_create_weights_in_init=True) + next_config = ModelConfig(skip_create_weights_in_init=True) + draft_config.extra_attrs = target_config.extra_attrs + next_config.extra_attrs = target_config.extra_attrs + + with patch( + "tensorrt_llm._torch.modules.mla.create_attention", + side_effect=lambda *args, **kwargs: _FakeAttention(), + ): + target_mla = _make_mla(target_config) + draft_mla = _make_mla(draft_config) + next_mla = _make_mla(next_config) + + assert target_mla.layer_idx == draft_mla.layer_idx == next_mla.layer_idx == 0 + assert target_mla.layer_idx_str == "0" + assert draft_mla.layer_idx_str == "0_0" + assert next_mla.layer_idx_str == "0_1" + registry = target_config.extra_attrs["mla_layers"] + assert registry["0"]() is target_mla + assert registry["0_0"]() is draft_mla + assert registry["0_1"]() is next_mla diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py new file mode 100644 index 000000000000..dccfbfe2e562 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -0,0 +1,447 @@ +# 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. +"""Unit tests for the DSpark captured-context attention primitives (CPU). + +The two primitives are validated against fully independent computations: +``dspark_sparse_attn`` vs (a) a per-element loop reference of the kernel formula +and (b) ``torch.nn.functional.scaled_dot_product_attention`` for the no-sink, +all-valid case; ``get_dspark_topk_idxs`` vs the reference index formula. +""" + +import types +from unittest.mock import Mock + +import pytest +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark +from tensorrt_llm._torch.models.dspark.attention import ( + 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 + + +def test_rope_table_is_cached_once_per_device(): + model = types.SimpleNamespace( + _attn_params={"rope_head_dim": 16}, + _freqs_cap=64, + _rope_theta=10000.0, + _freqs_table_cache={}, + ) + + first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + + assert first.data_ptr() == second.data_ptr() + assert len(model._freqs_table_cache) == 1 + positions = torch.tensor([1, 17, 63]) + expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) + torch.testing.assert_close(first[positions], expected[positions]) + + +def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): + captured = {} + + def fake_decoder_layer_init( + self, + model_config, + layer_idx, + aux_stream_dict, + attention_layer_idx=None, + mapping_with_cp=None, + disable_post_moe_fusion=False, + ): + torch.nn.Module.__init__(self) + self.model_config = model_config + self.config = model_config.pretrained_config + self.layer_idx = layer_idx + captured.update( + layer_idx=layer_idx, + attention_layer_idx=attention_layer_idx, + aux_stream_dict=aux_stream_dict, + mapping_with_cp=mapping_with_cp, + disable_post_moe_fusion=disable_post_moe_fusion, + ) + + monkeypatch.setattr( + modeling_dspark.DeepseekV4DecoderLayer, + "__init__", + fake_decoder_layer_init, + ) + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), + spec_config=None, + ) + + block = modeling_dspark.DSparkBlock( + model_config, + layer_idx=10, + aux_stream_dict={}, + stage_id=1, + num_stages=3, + num_capture_layers=0, + ) + + assert block.layer_idx == captured["layer_idx"] == 10 + assert captured["attention_layer_idx"] == block.stage_id == 1 + assert captured["disable_post_moe_fusion"] is True + + +@pytest.mark.parametrize("enable_fused_hc", [True, False]) +def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): + """The draft stage must use the inherited fused-HC rollback setting.""" + torch.manual_seed(71) + num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 + h = torch.randn(num_requests, block_size, hc_mult, hidden_size) + attention_input = torch.randn(num_requests, block_size, hidden_size) + attention_output = torch.randn_like(attention_input) + mid_residual = torch.randn_like(h) + attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) + attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) + ffn_post_mix = torch.randn_like(attention_post_mix) + ffn_comb_mix = torch.randn_like(attention_comb_mix) + raw_ffn_input = torch.randn_like(attention_input) + normed_ffn_input = torch.randn_like(attention_input) + moe_output = torch.randn(num_requests * block_size, hidden_size) + final_h = torch.randn_like(h) + events = [] + + def record(name, result): + def call(*args, **kwargs): + events.append(name) + return result + + return call + + monkeypatch.setattr( + modeling_dspark, + "dspark_attention_forward", + Mock(return_value=attention_output), + ) + + hc_attn = types.SimpleNamespace( + pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), + post_mapping=Mock(side_effect=record("attention_post", mid_residual)), + ) + hc_ffn = types.SimpleNamespace( + fused_hc=Mock( + side_effect=record( + "fused", + (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), + ) + ), + pre_mapping=Mock( + side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) + ), + post_mapping=Mock(side_effect=record("ffn_post", final_h)), + ) + post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) + post_attention_layernorm.weight = torch.ones(hidden_size) + post_attention_layernorm.variance_epsilon = 1e-6 + stage = types.SimpleNamespace( + enable_fused_hc=enable_fused_hc, + hc_attn=hc_attn, + hc_ffn=hc_ffn, + input_layernorm=Mock(side_effect=lambda tensor: tensor), + post_attention_layernorm=post_attention_layernorm, + mlp=Mock(return_value=moe_output), + _dspark_attn={}, + ) + model = types.SimpleNamespace( + use_real_mla=False, + _attn_params={"window_size": 2, "head_dim": 1}, + ) + + actual = DSparkDraftModel._forward_stage( + model, + stage, + h, + torch.randn(num_requests, hidden_size), + 1, + torch.empty(0), + torch.zeros(num_requests, block_size, dtype=torch.long), + ) + + assert actual is final_h + torch.testing.assert_close( + stage.mlp.call_args.args[0], + normed_ffn_input.reshape(num_requests * block_size, hidden_size), + ) + if enable_fused_hc: + assert events == ["fused", "ffn_post"] + hc_ffn.fused_hc.assert_called_once() + hc_attn.post_mapping.assert_not_called() + hc_ffn.pre_mapping.assert_not_called() + post_attention_layernorm.assert_not_called() + fused_kwargs = hc_ffn.fused_hc.call_args.kwargs + assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight + assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon + else: + assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] + hc_ffn.fused_hc.assert_not_called() + hc_attn.post_mapping.assert_called_once() + hc_ffn.pre_mapping.assert_called_once_with(mid_residual) + post_attention_layernorm.assert_called_once_with(raw_ffn_input) + + +def _ref_precompute_freqs_cis(dim, seqlen, base): + """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): + """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" + 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)) + return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) + + +def _loop_reference(q, kv, attn_sink, topk_idxs, scale): + """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" + b, m, h, d = q.shape + out = torch.zeros(b, m, h, d, dtype=torch.float32) + qf, kvf, sink = q.float(), kv.float(), attn_sink.float() + for bi in range(b): + for mi in range(m): + idxs = topk_idxs[bi, mi].tolist() + for hi in range(h): + scores, vecs = [], [] + for j in idxs: + if j < 0: + continue + k = kvf[bi, j] + scores.append(torch.dot(qf[bi, mi, hi], k) * scale) + vecs.append(k) + if not scores: + continue + s = torch.stack(scores) + smax = s.max() + p = torch.exp(s - smax) + denom = p.sum() + torch.exp(sink[hi] - smax) + num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) + out[bi, mi, hi] = num / denom + return out + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_sparse_attn_matches_loop_reference(seed): + torch.manual_seed(seed) + b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, n, d) + attn_sink = torch.randn(h) + idx = torch.stack( + [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] + ).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + ref = _loop_reference(q, kv, attn_sink, idx, scale) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_no_sink_matches_sdpa(): + """With a -inf sink (no sink mass) and all-valid contiguous indices, the + primitive must equal standard scaled-dot-product attention over the gathered + KV — an independent implementation.""" + torch.manual_seed(0) + b, m, h, d, topk = 2, 4, 5, 16, 9 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) # n == topk, attend to all + attn_sink = torch.full((h,), float("-inf")) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + + # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). + qh = q.permute(0, 2, 1, 3) + kvh = kv.unsqueeze(1).expand(b, h, topk, d) + ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_sink_reduces_mass(): + """A finite sink must strictly shrink the attention output magnitude vs an + infinitely-negative (disabled) sink, because it adds denominator mass only.""" + torch.manual_seed(0) + b, m, h, d, topk = 1, 2, 2, 16, 6 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + no_sink = dspark_sparse_attn(q, kv, torch.full((h,), float("-inf")), idx, scale) + with_sink = dspark_sparse_attn(q, kv, torch.zeros(h), idx, scale) + assert with_sink.abs().sum() < no_sink.abs().sum() + + +def test_sparse_attn_masked_indices_excluded(): + """An index of -1 must be excluded exactly (equiv. to dropping that column).""" + torch.manual_seed(0) + b, m, h, d = 1, 1, 2, 16 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, 5, d) + sink = torch.full((h,), float("-inf")) + scale = d**-0.5 + full = torch.tensor([[[0, 1, 2, 3]]]).int() + masked = torch.tensor([[[0, 1, 2, -1]]]).int() + drop3 = torch.tensor([[[0, 1, 2]]]).int() + got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) + got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) + torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) + # And masking genuinely changes the result vs attending to position 3. + got_full = dspark_sparse_attn(q, kv, sink, full, scale) + assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize( + "start_pos,window,block", [(1, 128, 5), (3, 128, 5), (10, 4, 5), (200, 128, 6)] +) +def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): + bsz = 3 + got = get_dspark_topk_idxs(window, bsz, block, start_pos) + # Reference formula (DeepSpec get_dspark_topk_idxs). + ctx = torch.arange(min(window, start_pos + 1)) + blk = window + torch.arange(block) + ref_row = torch.cat([ctx, blk]).int() + assert got.shape == (bsz, block, ref_row.numel()) + for bi in range(bsz): + for mi in range(block): + torch.testing.assert_close(got[bi, mi], ref_row) + + +def test_get_dspark_topk_idxs_requires_generation(): + with pytest.raises(AssertionError): + get_dspark_topk_idxs(128, 1, 5, 0) + + +@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) +def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): + got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) + ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_matches_reference(ndim): + torch.manual_seed(0) + b, s, h, rd = 2, 5, 4, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + got = apply_dspark_rotary(x, fc) + ref = _ref_apply_rotary_emb(x, fc) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_inverse_roundtrip(ndim): + """De-rotation (inverse) must undo the forward rotation (property test).""" + torch.manual_seed(1) + b, s, h, rd = 2, 6, 3, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) + torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) + + +def _make_attn_inputs(seed=0): + """Small synthetic DSpark attention inputs/weights (CPU bf16).""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block, start_pos = 8, 3, 5 + b = 2 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + start_pos=start_pos, + b=b, + eps=1e-6, + softmax_scale=head_dim**-0.5, + ) + bf = torch.bfloat16 + g["x"] = torch.randn(b, block, dim, dtype=bf) + g["main_x"] = torch.randn(b, 1, dim, dtype=bf) + g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) + g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 + g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 + g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 + g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 + g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 + g["q_norm"] = torch.ones(q_lora) + g["kv_norm"] = torch.ones(head_dim) + g["attn_sink"] = torch.randn(n_heads) + g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) + return g + + +def _run(g): + return dspark_attention_forward( + g["x"], + g["main_x"], + g["start_pos"], + g["kv_cache0"], + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def test_attention_forward_shape_and_determinism(): + g = _make_attn_inputs() + o = _run(g) + assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) + assert torch.isfinite(o.float()).all() + torch.testing.assert_close(o, _run(g)) # deterministic + + +def test_attention_forward_does_not_mutate_kv_cache(): + """The rolling window write must be functional (cache cloned, not mutated).""" + g = _make_attn_inputs() + before = g["kv_cache0"].clone() + _run(g) + torch.testing.assert_close(g["kv_cache0"], before) 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 new file mode 100644 index 000000000000..f66a37834fab --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -0,0 +1,252 @@ +# 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. +"""CUDA-graph-safety unit tests for the DSpark batched draft-attention path. + +The load-bearing invariant for the DSpark batched draft path (the default, used +whenever ``cuda_graph_config`` is set, since the one-engine drafter is captured in +the target's graph) is that the batched, sync-free primitives are **numerically +identical, per request**, to the validated scalar path — only the host-int +``start_pos`` and the per-request window indexing are tensorized. These tests assert +that equivalence on CPU (so they run in pre-merge CI without a GPU), plus a +GPU-gated capture+replay smoke test that proves the batched attention is actually +graph-capturable. +""" + +import pytest +import torch + +from tensorrt_llm._torch.models.dspark.attention import ( + apply_dspark_rotary, + apply_dspark_rotary_batched, + dspark_attention_forward, + dspark_attention_forward_batched, + get_dspark_topk_idxs, + get_dspark_topk_idxs_batched, + precompute_dspark_freqs_cis, +) + + +def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): + """Per-request DSpark attention inputs/weights (shared weights, distinct pos). + + Mirrors ``test_dspark_attention._make_attn_inputs`` but builds ``G`` requests + each with its own ``start_pos`` (small => partial context, large => full + rolling window) and its own pre-seeded window, so the batched-vs-scalar + comparison exercises the per-request RoPE gather + windowed context read. + """ + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block = 8, 3 + G = len(start_positions) + bf = torch.bfloat16 + # A single fixed RoPE table covering every request's positions (both paths + # index/gather the same values, so freqs are identical across paths). + maxlen = max(start_positions) + 1 + block + 4 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + G=G, + start_positions=list(start_positions), + eps=1e-6, + softmax_scale=head_dim**-0.5, + x=torch.randn(G, block, dim, dtype=bf), + main_x=torch.randn(G, 1, dim, dtype=bf), + # Distinct, non-zero seeded window per request to exercise context reads. + kv_cache=torch.randn(G, window, head_dim, dtype=bf), + wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, + wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, + wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, + wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, + wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, + q_norm=torch.ones(q_lora), + kv_norm=torch.ones(head_dim), + attn_sink=torch.randn(n_heads), + freqs=precompute_dspark_freqs_cis(rd, maxlen), + ) + return g + + +def _attn_kwargs(g): + return dict( + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def _scalar_reference(g): + """Run the validated scalar attention once per request and stack the outputs.""" + outs = [] + for i, sp in enumerate(g["start_positions"]): + out_i = dspark_attention_forward( + g["x"][i : i + 1], + g["main_x"][i : i + 1], + int(sp), + g["kv_cache"][i : i + 1].clone(), + **_attn_kwargs(g), + ) + outs.append(out_i) + return torch.cat(outs, dim=0) + + +def _batched(g, persist=False): + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + return dspark_attention_forward_batched( + g["x"], + g["main_x"], + start_pos, + g["kv_cache"].clone(), + slots, + persist=persist, + **_attn_kwargs(g), + ) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_batched_attention_matches_scalar_per_request(seed): + """Batched attention == per-request scalar attention at distinct start_pos. + + This is the invariant that lets the batched path replace the per-request loop + under CUDA graphs without changing draft quality / greedy parity. + """ + g = _make_batched_inputs(seed=seed) + ref = _scalar_reference(g) + got = _batched(g) + assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + + +def test_batched_attention_persist_writes_through_window(): + """persist=True writes main_kv into the shared window at start_pos%window.""" + g = _make_batched_inputs(seed=3) + G, win = g["G"], g["window"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + cache = g["kv_cache"].clone() + before = cache.clone() + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) + ) + # Exactly the start_pos%win row of each request changed. + for i, sp in enumerate(g["start_positions"]): + changed = (cache[i] != before[i]).any(dim=-1) + expected = torch.zeros(win, dtype=torch.bool) + expected[sp % win] = True + assert torch.equal(changed, expected), f"req {i}: wrong window row written" + + +def test_batched_attention_no_persist_keeps_window(): + """persist=False must not mutate the caller's window (functional).""" + g = _make_batched_inputs(seed=4) + before = g["kv_cache"].clone() + _batched(g, persist=False) + torch.testing.assert_close(g["kv_cache"], before) + + +@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) +def test_batched_topk_matches_scalar(start_positions): + """Fixed-size masked batched topk == scalar topk per request (valid slots).""" + window, block = 8, 3 + start_pos = torch.tensor(start_positions, dtype=torch.long) + batched = get_dspark_topk_idxs_batched(window, block, start_pos) + assert tuple(batched.shape) == (len(start_positions), block, window + block) + for i, sp in enumerate(start_positions): + scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] + # The batched row drops the -1-masked context slots to recover the scalar + # (variable-width) index list; both must then be identical. + for m in range(block): + valid = batched[i, m][batched[i, m] >= 0] + torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_batched_rotary_matches_scalar_per_row(ndim): + """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" + torch.manual_seed(0) + G, s, h, rd = 3, 4, 2, 8 + x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) + table = precompute_dspark_freqs_cis(rd, 64) + # Per-row absolute start positions -> per-row freq windows. + starts = [1, 9, 30] + per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] + got = apply_dspark_rotary_batched(x, per_row) + for i, sp in enumerate(starts): + ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) + torch.testing.assert_close(got[i : i + 1], ref_i) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_batched_attention_cuda_graph_capture_replay(): + """The batched attention captures + replays and matches eager output. + + Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). + """ + g = _make_batched_inputs(seed=0) + dev = "cuda" + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) + slots = torch.arange(G, dtype=torch.long, device=dev) + # Static input tensors the graph reads/writes. + x = g["x"].to(dev) + main_x = g["main_x"].to(dev) + cache = g["kv_cache"].to(dev) + kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} + + def run(persist): + return dspark_attention_forward_batched( + x, main_x, start_pos, cache, slots, persist=persist, **kw + ) + + eager = run(persist=False) + + # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so + # the comparison isn't perturbed by the window write-through. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run(persist=False) + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = run(persist=False) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py new file mode 100644 index 000000000000..fed0da3a2803 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -0,0 +1,99 @@ +# 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. +"""Unit tests for the DSpark draft I/O proposal stage.""" + +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 + +VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 +NOISE_ID = 199 + + +def test_build_draft_input_ids(): + bonus = torch.tensor([7, 9]) + ids = build_draft_input_ids(bonus, block_size=BLK, noise_token_id=NOISE_ID) + assert ids.shape == (B, BLK) + assert torch.equal(ids[:, 0], bonus) + assert torch.all(ids[:, 1:] == NOISE_ID) + + +def test_dspark_propose_full_block_no_confidence(): + torch.manual_seed(0) + markov = build_markov_head( + markov_head_type="rnn", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(B, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (B,)) + hid = torch.randn(B, BLK, HID) + with torch.no_grad(): + tokens, num = dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=None, + block_size=BLK, + ) + assert tokens.shape == (B, BLK) + # No confidence head -> propose the full block. + assert torch.all(num == BLK) + # Tokens match the markov head's own greedy block sampling. + ref_tokens, _ = markov.sample_block_tokens( + base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 + ) + assert torch.equal(tokens, ref_tokens) + + +def test_dspark_propose_confidence_truncates(): + torch.manual_seed(1) + markov = build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + conf = DSparkConfidenceHead(hidden_size=HID).eval() + # The confidence proj is bias-free, so drive the logit via a constant weight + # against a constant hidden: logit = weight_val * HID per position. + base = torch.randn(1, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (1,)) + hid = torch.ones(1, BLK, HID) + with torch.no_grad(): + conf.proj.weight.fill_(5.0 / HID) # logit ~ 5 -> sigmoid ~ 0.993, all confident + with torch.no_grad(): + _, num = dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=conf, + block_size=BLK, + confidence_threshold=0.5, + ) + # All-confident -> full block proposed. + assert int(num[0]) == BLK + # Now make the head output low confidence everywhere -> truncate to 0... but + # confident_prefix_length returns first sub-threshold index (0 here). + with torch.no_grad(): + conf.proj.weight.fill_(-5.0 / HID) # logit ~ -5 -> sigmoid ~ 0.0067 < 0.5 + _, num2 = dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=conf, + block_size=BLK, + confidence_threshold=0.5, + ) + assert int(num2[0]) == 0 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py new file mode 100644 index 000000000000..16b6a01acd17 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -0,0 +1,117 @@ +# 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. +"""Unit tests for the DSpark draft-network heads (hardware-agnostic, CPU).""" + +import pytest +import torch + +from tensorrt_llm._torch.models.dspark.heads import ( + DSparkConfidenceHead, + RNNHead, + VanillaMarkov, + build_markov_head, + confident_prefix_length, +) + +VOCAB, RANK, HID, B, BLK = 257, 16, 32, 3, 5 + + +@pytest.mark.parametrize("head_type", ["vanilla", "gated", "rnn"]) +def test_markov_block_sampling_shapes_and_determinism(head_type): + torch.manual_seed(0) + head = build_markov_head( + markov_head_type=head_type, vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (B,)) + hid = torch.randn(B, BLK, HID) + with torch.no_grad(): + tok, logits = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + tok2, _ = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + assert tok.shape == (B, BLK) + assert logits.shape == (B, BLK, VOCAB) + # Greedy is deterministic. + assert torch.equal(tok, tok2) + # Each sampled token is the argmax of its (bias-corrected) step logits. + assert torch.equal(tok, logits.argmax(dim=-1)) + + +def test_markov_bias_is_additive_low_rank(): + # bias = W2(W1[token]); the corrected first-step logits == base + bias. + torch.manual_seed(1) + head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() + base = torch.randn(B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (B,)) + with torch.no_grad(): + _, corrected = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 + ) + expected0 = base[:, 0] + head.markov_w2(head.markov_w1(first)) + assert torch.allclose(corrected[:, 0], expected0, atol=1e-5) + + +def test_rnn_state_carries_across_positions(): + torch.manual_seed(2) + head = RNNHead(vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID).eval() + initial_state = torch.zeros(1, RANK) + prev_embedding = head.get_prev_embeddings(torch.zeros(1, dtype=torch.long)) + prefix_hidden = torch.randn(1, HID) + current_hidden = torch.randn(1, HID) + + with torch.no_grad(): + state_a, _ = head._rnn_step(initial_state, prev_embedding, prefix_hidden) + state_b, _ = head._rnn_step(initial_state, prev_embedding, -prefix_hidden) + _, bias_a = head._rnn_step(state_a, prev_embedding, current_hidden) + _, bias_b = head._rnn_step(state_b, prev_embedding, current_hidden) + + assert not torch.allclose(state_a, state_b) + assert not torch.allclose(bias_a, bias_b) + + +def test_build_markov_head_rank_zero_returns_none(): + assert ( + build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=0, hidden_size=HID + ) + is None + ) + + +def test_confidence_head_and_prefix_truncation(): + head = DSparkConfidenceHead(hidden_size=HID) + conf = head(torch.randn(B, BLK, HID)) + assert conf.shape == (B, BLK) + # threshold 0 disables truncation. + assert confident_prefix_length(conf, block_size=BLK, threshold=0.0) == BLK + # First sub-threshold position truncates the prefix. + logits = torch.tensor([[10.0, 10.0, -10.0, 10.0, 10.0]]) + assert confident_prefix_length(logits, block_size=BLK, threshold=0.5) == 2 + # All-confident -> full block. + logits_hi = torch.full((1, BLK), 10.0) + assert confident_prefix_length(logits_hi, block_size=BLK, threshold=0.5) == BLK + + +def test_confidence_head_with_markov_concat_dim(): + head = DSparkConfidenceHead(hidden_size=HID, markov_rank=RANK, with_markov=True) + hid = torch.randn(B, BLK, HID) + prev_emb = torch.randn(B, BLK, RANK) + out = head(hid, prev_embeddings=prev_emb) + assert out.shape == (B, BLK) + with pytest.raises(AssertionError): + head(hid) # with_markov requires prev_embeddings diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py new file mode 100644 index 000000000000..34ee08006ee9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -0,0 +1,356 @@ +# 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. +"""GPU unit tests for the DSpark worker and metadata plumbing. + +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 +draft and acceptance path is covered by the DSpark test in +``integration/defs/accuracy/test_llm_api_pytorch.py``. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="DSpark metadata/worker allocate CUDA buffers" +) + +HIDDEN = 128 +NCAP = 3 +HC_MULT = 4 + + +def _make_metadata(max_num_requests=8, max_num_tokens=64, layers=(58, 59, 60)): + return DSparkSpecMetadata( + max_draft_len=5, + max_total_draft_tokens=5, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + max_num_requests=max_num_requests, + layers_to_capture=list(layers), + hidden_size=HIDDEN, + max_num_tokens=max_num_tokens, + dtype=torch.bfloat16, + ) + + +def test_metadata_buffer_and_layer_lookup(): + meta = _make_metadata() + assert meta.num_capture_layers == NCAP + assert meta.captured_hidden_states.shape == (64, HIDDEN * NCAP) + # sorted, O(1) membership + assert meta.is_layer_capture(58) and meta.is_layer_capture(60) + assert not meta.is_layer_capture(0) and not meta.is_layer_capture(61) + + +def test_metadata_capture_plain_hidden(): + """A [num_tokens, hidden] capture is stored at the layer's slice as-is.""" + meta = _make_metadata() + hs = torch.randn(4, HIDDEN, device="cuda", dtype=torch.bfloat16) + meta.maybe_capture_hidden_states(59, hs) # layer 59 -> capture index 1 + got = meta.get_hidden_states(4) + assert torch.equal(got[:, HIDDEN : 2 * HIDDEN], hs) + + +def test_metadata_capture_hc_mean_reduction(): + """A flattened mHC residual [N, hc_mult*hidden] is reduced by mean over hc.""" + meta = _make_metadata() + mhc = torch.randn(4, HC_MULT * HIDDEN, device="cuda", dtype=torch.bfloat16) + meta.maybe_capture_hidden_states(58, mhc) # layer 58 -> capture index 0 + expected = mhc.reshape(4, HC_MULT, HIDDEN).mean(dim=1) + got = meta.get_hidden_states(4) + assert torch.equal(got[:, 0:HIDDEN], expected) + + +def test_metadata_no_capture_for_unlisted_layer(): + meta = _make_metadata() + meta.captured_hidden_states.zero_() + meta.maybe_capture_hidden_states(10, torch.randn(4, HIDDEN, device="cuda")) + assert torch.count_nonzero(meta.get_hidden_states(4)) == 0 + + +def test_metadata_prepare_batch_indices(): + meta = _make_metadata() + meta.request_ids = [7, 3, 5] + meta.prepare() + assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] + + +def _make_worker(): + cfg = types.SimpleNamespace( + max_draft_len=5, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + confidence_threshold=0.5, + ) + from tensorrt_llm.mapping import Mapping + + return DSparkWorker(cfg, Mapping()) + + +def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): + return types.SimpleNamespace( + num_stages=num_stages, + block_size=5, + _attn_params={"window_size": window_size, "head_dim": head_dim}, + ) + + +def test_worker_lazy_init_window_buffers(): + worker = _make_worker() + dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) + meta = _make_metadata(max_num_requests=8) + worker._lazy_init(dm, meta) + assert worker._kv_windows.shape == (8, 3, 128, 64) + assert worker._ctx_len.shape == (8,) + assert list(worker._free_slots) == list(range(8)) + assert worker._batch_to_slot is not None + assert worker._batch_to_slot.shape == (8,) + assert worker._batch_to_slot.device.type == "cuda" + # idempotent + buf_id = id(worker._kv_windows) + worker._lazy_init(dm, meta) + assert id(worker._kv_windows) == buf_id + + +def test_worker_rejects_mismatched_block_size(): + worker = _make_worker() + draft_model = _fake_draft_model() + draft_model.block_size = 4 + + with pytest.raises(ValueError, match="block_size must equal worker max_draft_len"): + worker._lazy_init(draft_model, _make_metadata()) + + +def test_worker_slot_assignment_and_reset(): + worker = _make_worker() + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=4)) + + s0 = worker._assign_slot(100, reset=False) + s1 = worker._assign_slot(101, reset=False) + assert s0 != s1 + # same request id -> same slot (no reset) + assert worker._assign_slot(100, reset=False) == s0 + + # mark a position, then reset -> slot freed + window/pos cleared + worker._ctx_len[s0] = 42 + worker._kv_windows[s0].fill_(1.0) + s0b = worker._assign_slot(100, reset=True) + assert int(worker._ctx_len[s0b]) == 0 + assert float(worker._kv_windows[s0b].abs().sum()) == 0.0 + + +def test_worker_slot_exhaustion_preserves_live_request(): + worker = _make_worker() + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=1)) + + slot = worker._assign_slot(100, reset=False) + worker._ctx_len[slot] = 42 + worker._kv_windows[slot].fill_(1.0) + + with pytest.raises(RuntimeError, match="no free rolling-window slots"): + worker._assign_slot(101, reset=False) + + assert worker._req_to_slot == {100: slot} + assert int(worker._ctx_len[slot]) == 42 + assert torch.all(worker._kv_windows[slot] == 1.0) + + +def test_seed_context_windows_preserves_state_across_prefill_chunks(): + class DraftModel: + num_stages = 1 + block_size = 5 + _attn_params = {"window_size": 8, "head_dim": 4} + + def __init__(self): + self.written_positions = [] + + def write_context_windows(self, hidden, positions, windows): + self.written_positions.append(positions.clone()) + windows.add_(1) + + worker = _make_worker() + draft_model = DraftModel() + metadata = types.SimpleNamespace( + max_num_requests=1, + request_ids=[100], + get_hidden_states=lambda _num_tokens: torch.zeros( + 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ), + ) + worker._lazy_init(draft_model, metadata) + + first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) + worker._seed_context_windows( + draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 + ) + slot = worker._req_to_slot[100] + assert int(worker._ctx_len[slot]) == 3 + + metadata.get_hidden_states = lambda _num_tokens: torch.zeros( + 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ) + second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) + worker._seed_context_windows( + draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 + ) + + assert int(worker._ctx_len[slot]) == 5 + assert [positions.tolist() for positions in draft_model.written_positions] == [ + [1, 2, 3], + [4, 5], + ] + assert torch.all(worker._kv_windows[slot] == 2.0) + + +def test_prepare_builds_batch_to_slot_on_batched_path(): + """prepare() mirrors the host slot map into _batch_to_slot (default batched path).""" + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) # batched is the default + meta._dspark_worker = worker + + # Assign slots for two requests (as the prefill path would). + sa = worker._assign_slot(100, reset=True) + sb = worker._assign_slot(101, reset=True) + + meta.request_ids = [101, 100] + meta.prepare() + # Mirror reflects request-order -> slot. + assert worker._batch_to_slot[:2].tolist() == [sb, sa] + + +def test_prepare_frees_stale_slots_on_batched_path(): + """A request that drops out of the batch returns its slot to the free pool.""" + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + + sa = worker._assign_slot(100, reset=True) + worker._assign_slot(101, reset=True) + worker._ctx_len[sa] = 17 + + # Only request 101 survives; 100's slot must be freed + cleared. + meta.request_ids = [101] + meta.prepare() + assert 100 not in worker._req_to_slot + assert sa in worker._free_slots + assert int(worker._ctx_len[sa]) == 0 + + +def test_forward_mixed_batch_routes_through_base_entries(monkeypatch): + """Mixed (context + gen) batch: ``forward`` must route acceptance and + production through the unified ``SpecWorkerBase`` entries, one-hot-fill the + context requests' draft-prob rows, and assemble + ``next_draft_tokens = [ctx zeros ; gen argmax]``. + + Spies replace the base sampling entries and the heavy sub-calls (context + seeding, per-request draft backbone) so this exercises the worker's + context/gen orchestration — the exact surface the #15775 refactor changed — + without a real draft model or MPI. + """ + worker = _make_worker() + worker.guided_decoder = None + dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) + + K = worker.max_draft_len + vocab = 16 + num_contexts, num_gens = 2, 3 + batch_size = num_contexts + num_gens + + meta = _make_metadata(max_num_requests=8) + meta.request_ids = [10, 11, 20, 21, 22] # 2 context + 3 gen + meta.prepare() + + attn_metadata = types.SimpleNamespace( + num_seqs=batch_size, + num_contexts=num_contexts, + num_ctx_tokens=0, + num_tokens=batch_size, + ) + + # Acceptance: return a fixed verified prefix (one accepted token per request). + accepted = torch.arange(batch_size * (K + 1), dtype=torch.int32, device="cuda").reshape( + batch_size, K + 1 + ) + num_accepted = torch.ones(batch_size, dtype=torch.int32, device="cuda") + accept_calls = {} + + def fake_accept(logits, am, sm): + accept_calls["args"] = (am, sm) + return accepted, num_accepted + + monkeypatch.setattr(worker, "sample_and_accept_draft_tokens", fake_accept) + # Context-window seeding is covered by its own test; stub it out here. + monkeypatch.setattr(worker, "_seed_context_windows", lambda *a, **k: None) + + # The gen-block helper now returns the corrected block logits [num_gens,K,vocab]. + gen_logits = torch.randn(num_gens, K, vocab, device="cuda") + monkeypatch.setattr(worker, "_draft_gen_block_batched", lambda *a, **k: gen_logits) + + sdt_calls = {} + # The gen scatter publishes the FULL (post-TP-gather) vocab width, which is + # wider than the sharded gen_logits width (`vocab`). The worker must pass this + # published width (draft_probs_last_dim) to write_context_onehot_draft_probs, + # NOT gen_logits.shape[-1]. + FULL_VOCAB = 97 + + def fake_sample_draft_tokens(gl, sm, bs, *, num_contexts): + sdt_calls["logits"] = gl + sdt_calls["batch_size"] = bs + sdt_calls["num_contexts"] = num_contexts + sm.draft_probs_last_dim = FULL_VOCAB # simulate the full-vocab scatter + return gl.argmax(dim=-1).to(torch.int32) + + monkeypatch.setattr(worker, "sample_draft_tokens", fake_sample_draft_tokens) + + onehot_calls = {} + monkeypatch.setattr( + worker, + "write_context_onehot_draft_probs", + lambda sm, nc, ng, k, gv: onehot_calls.update(nc=nc, ng=ng, k=k, gv=gv), + ) + + input_ids = torch.zeros(batch_size, dtype=torch.long, device="cuda") + position_ids = torch.zeros(batch_size, dtype=torch.long, device="cuda") + hidden = torch.zeros(batch_size, HIDDEN, device="cuda", dtype=torch.bfloat16) + logits = torch.zeros(batch_size, vocab, device="cuda") + + out = worker.forward(input_ids, position_ids, hidden, logits, attn_metadata, meta, dm) + + # Acceptance went through the unified entry with the right metadata objects. + assert accept_calls["args"] == (attn_metadata, meta) + # Production fed the [num_gens, K, vocab] block logits to the base sampler, + # with num_contexts so it slices the gen segment. + assert sdt_calls["num_contexts"] == num_contexts + assert sdt_calls["logits"].shape == (num_gens, K, vocab) + # Context rows one-hot-filled with the *scatter* width (draft_probs_last_dim, + # FULL_VOCAB), not the sharded gen_logits width (vocab). + assert onehot_calls == {"nc": num_contexts, "ng": num_gens, "k": K, "gv": FULL_VOCAB} + + # next_draft_tokens = [context zeros ; gen argmax]; gen subset is not polluted + # by the context rows. + nd = out["next_draft_tokens"] + assert nd.shape == (batch_size, K) + assert torch.all(nd[:num_contexts] == 0) + assert torch.equal(nd[num_contexts:], gen_logits.argmax(dim=-1).to(torch.int32)) + # Verified tokens are surfaced unchanged. + assert torch.equal(out["new_tokens"], accepted) + assert torch.equal(out["new_tokens_lens"], num_accepted) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 486c150f2040..775f3e7b6be1 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -170,6 +170,34 @@ def test_deepseek_v4_base_checkpoint_detection( assert ModelConfig._is_deepseek_v4_base_checkpoint(str(tmp_path)) is expected_is_base +def test_deepseek_v4_missing_compress_ratios_raises(tmp_path, monkeypatch): + """DeepSeek-V4 load must fail fast with a clear error when neither the + checkpoint config nor a user ``sparse_attention_config`` provides + ``compress_ratios``. + + Regression test: previously ``compress_ratios`` could stay ``None`` and the + internal normalization comprehension raised an opaque + ``TypeError: 'NoneType' object is not iterable`` mid-load. + """ + from tensorrt_llm._torch import model_config as model_config_module + from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config + + pretrained_config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + compress_ratios=None, + num_hidden_layers=4, + ) + + # Avoid touching the filesystem for the HF config load; the empty tmp_path + # makes the real ``_is_deepseek_v4_base_checkpoint`` probe return False. + monkeypatch.setattr( + model_config_module, "load_pretrained_config", lambda *args, **kwargs: pretrained_config + ) + + with pytest.raises(ValueError, match="compress_ratios"): + ModelConfig.from_pretrained(str(tmp_path)) + + def test_model_config_sets_is_encoder_decoder_from_pretrained_config(): model_config = ModelConfig( pretrained_config=make_pretrained_config( diff --git a/tests/unittest/api_stability/references_committed/llm.yaml b/tests/unittest/api_stability/references_committed/llm.yaml index cf6f04373a90..8f36fe11823e 100644 --- a/tests/unittest/api_stability/references_committed/llm.yaml +++ b/tests/unittest/api_stability/references_committed/llm.yaml @@ -59,7 +59,7 @@ methods: default: null # Speculative decoding speculative_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.DraftTargetDecodingConfig, tensorrt_llm.llmapi.llm_args.EagleDecodingConfig, tensorrt_llm.llmapi.llm_args.Eagle3DecodingConfig, tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_args.MedusaDecodingConfig, tensorrt_llm.llmapi.llm_args.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, tensorrt_llm.llmapi.llm_args.SADecodingConfig, tensorrt_llm.llmapi.llm_args.UserProvidedDecodingConfig, tensorrt_llm.llmapi.llm_args.SaveHiddenStatesDecodingConfig, tensorrt_llm.llmapi.llm_args.PARDDecodingConfig, tensorrt_llm.llmapi.llm_args.DFlashDecodingConfig, tensorrt_llm.llmapi.llm_args.AutoDecodingConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.DraftTargetDecodingConfig, tensorrt_llm.llmapi.llm_args.EagleDecodingConfig, tensorrt_llm.llmapi.llm_args.Eagle3DecodingConfig, tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_args.MedusaDecodingConfig, tensorrt_llm.llmapi.llm_args.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, tensorrt_llm.llmapi.llm_args.SADecodingConfig, tensorrt_llm.llmapi.llm_args.UserProvidedDecodingConfig, tensorrt_llm.llmapi.llm_args.SaveHiddenStatesDecodingConfig, tensorrt_llm.llmapi.llm_args.PARDDecodingConfig, tensorrt_llm.llmapi.llm_args.DFlashDecodingConfig, tensorrt_llm.llmapi.llm_args.DSparkDecodingConfig, tensorrt_llm.llmapi.llm_args.AutoDecodingConfig, NoneType] default: null # generation constraints max_batch_size: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b76bdea0810b..3365ce6685ab 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -34,6 +34,7 @@ DecodeCudaGraphConfig, DecodingBaseConfig, DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, @@ -241,6 +242,145 @@ def _capture_warning(msg, *args, **kwargs): for m in warnings_seen) +def test_dspark_block_size_resolved_from_checkpoint(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 5}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path)) + + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + assert args.speculative_config.block_size == 5 + + +def test_dspark_block_size_must_match_max_draft_len(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 4}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path)) + + with pytest.raises(ValueError, match="block_size must equal max_draft_len"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + +def test_dspark_target_layer_ids_resolved_from_checkpoint(tmp_path): + # When the user leaves target_layer_ids unset, the checkpoint's ordered + # dspark_target_layer_ids must be adopted verbatim. + (tmp_path / "config.json").write_text( + '{"dspark_block_size": 5, "dspark_target_layer_ids": [3, 1, 2]}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path)) + + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + # Order is preserved (projection columns are order-dependent). + assert args.speculative_config.target_layer_ids == [3, 1, 2] + + +def test_dspark_target_layer_ids_matching_override_accepted(tmp_path): + # An explicit override that matches the checkpoint list exactly is fine. + (tmp_path / "config.json").write_text( + '{"dspark_block_size": 5, "dspark_target_layer_ids": [1, 2, 3]}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path), + target_layer_ids=[1, 2, 3]) + + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + assert args.speculative_config.target_layer_ids == [1, 2, 3] + + +def test_dspark_target_layer_ids_mismatched_count_rejected(tmp_path): + # A different number of layers would mismatch main_proj.in_features. + (tmp_path / "config.json").write_text( + '{"dspark_block_size": 5, "dspark_target_layer_ids": [1, 2, 3]}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path), + target_layer_ids=[1, 2]) + + with pytest.raises(ValueError, match="must match the checkpoint"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + +def test_dspark_target_layer_ids_same_count_different_layers_rejected(tmp_path): + # Same count but different layers: shapes line up, but the draft would see + # hidden states it was not trained on, so this must be rejected too. + (tmp_path / "config.json").write_text( + '{"dspark_block_size": 5, "dspark_target_layer_ids": [1, 2, 3]}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path), + target_layer_ids=[1, 2, 4]) + + with pytest.raises(ValueError, match="must match the checkpoint"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + +def test_dspark_target_layer_ids_order_mismatch_rejected(tmp_path): + # Same set but different order: projection columns are order-dependent, so a + # reordered override must be rejected rather than silently accepted. + (tmp_path / "config.json").write_text( + '{"dspark_block_size": 5, "dspark_target_layer_ids": [1, 2, 3]}') + spec_cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model=str(tmp_path), + target_layer_ids=[3, 2, 1]) + + with pytest.raises(ValueError, match="must match the checkpoint"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + +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. + spec_cfg = DSparkDecodingConfig(max_draft_len=5) + + with pytest.raises(ValueError, + match="requires speculative_config.speculative_model"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + +def test_dspark_requires_positive_max_draft_len(tmp_path): + (tmp_path / "config.json").write_text('{"dspark_block_size": 5}') + spec_cfg = DSparkDecodingConfig(speculative_model=str(tmp_path)) + + with pytest.raises(ValueError, match="max_draft_len must be > 0"): + TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=spec_cfg, + ) + + def test_post_processor_hook_rejected_with_skip_tokenizer_init(): """post_processor_hook + skip_tokenizer_init must fail fast.