From 40289c86d092d0a94681210d676eb34a94237e68 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Thu, 23 Jul 2026 00:19:08 +0000 Subject: [PATCH 1/3] [None][feat] Add Qwen3-based DSpark drafter (DeepSpec dense checkpoints) Support the DeepSpec-released dense DSpark drafters for Qwen3 targets (e.g. deepseek-ai/dspark_qwen3_8b_block7) alongside the existing DeepSeek-V4 mtp.*-namespace drafter: - modeling_dspark_qwen3.py: Qwen3DSparkDraftModel/Qwen3DSparkForCausalLM, a pure-torch dense GQA draft backbone (fc+hidden_norm context projection, Qwen3 layers with captured-context K/V and bidirectional block attention, Markov head refinement) implementing the same worker-facing protocol as DSparkDraftModel, so DSparkWorker / DSparkSpecMetadata / CUDA-graph plumbing are reused unchanged. The worker rolling buffer holds per-layer context K/V as a ring over the last TRTLLM_DSPARK_QWEN3_CTX_WINDOW (default 2048) committed positions. - get_draft_model: dispatch on the drafter checkpoint's Qwen3DSparkModel architecture. - llm_args DSpark validation: also resolve unprefixed top-level config keys (block_size / target_layer_ids / mask_token_id / markov_rank) used by the DeepSpec drafter checkpoints (no schema change; golden manifest unchanged). - tests: golden tests vs a torch-only port of the DeepSpec reference (worker frame conventions across multi-step decode, batched-vs-singleton parity, ring wraparound). Validation: golden unit tests vs a torch-only DeepSpec reference port; end-to-end + GSM8K/MATH-500/HumanEval benchmarks (aiperf, 1xB300, conc 1..64) were run on the v1.3.0rc22 backport of this change (branch dspark-qwen3-rc22, same diff) inside the 1.3.0rc22 release container: 3.3-3.7x per-user speedup over vanilla at concurrency 1 (~750 tok/s/user), 6.17 avg decoded tokens/iter on GSM8K (block=7). Signed-off-by: chungen04 --- .../_torch/models/modeling_dspark_qwen3.py | 602 ++++++++++++++++++ .../_torch/models/modeling_speculative.py | 9 + tensorrt_llm/llmapi/llm_args.py | 5 + .../hw_agnostic/test_dspark_qwen3.py | 424 ++++++++++++ 4 files changed, 1040 insertions(+) create mode 100644 tensorrt_llm/_torch/models/modeling_dspark_qwen3.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py diff --git a/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py b/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py new file mode 100644 index 000000000000..cf0eb2a19c68 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py @@ -0,0 +1,602 @@ +# 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. +"""Qwen3-based DSpark draft model (DeepSeek DeepSpec ``Qwen3DSparkModel``). + +Unlike the DeepSeek-V4 DSpark drafter (``modeling_dspark.py``) — whose draft +weights live in the target checkpoint's ``mtp.*`` namespace and whose stages +are full V4 blocks (MLA + MoE + mHC) — the DeepSpec-released Qwen3 drafters +(e.g. ``deepseek-ai/dspark_qwen3_8b_block7``) are separate dense bf16 +checkpoints with a flat namespace: + + - ``fc`` + ``hidden_norm``: project the concatenated captured target-layer + hidden states (``hidden_size * num_capture_layers``) into the draft hidden + size — the layer-invariant *context stream* (``main_x``). + - ``layers.{0..n-1}``: standard Qwen3 GQA decoder layers (q/k per-head + RMSNorm, RoPE, gated-SiLU MLP), except attention keys/values come from + BOTH the projected context (one row per committed token) and the draft + block itself, and the ``block_size`` draft queries attend bidirectionally + within the block. + - ``norm`` + (target-shared) ``lm_head`` -> :func:`dspark_propose` (Markov + head refinement; the confidence head is inert scaffolding, as in the V4 + path). + +The worker-facing protocol (``write_context_windows`` / +``write_context_windows_batched`` / ``forward_batched`` / ``_attn_params`` / +``num_stages`` / ``run_moe_lockstep_noop``) matches :class:`DSparkDraftModel`, +so :class:`DSparkWorker` and :class:`DSparkSpecMetadata` drive both drafters +unchanged. The worker-owned rolling buffer here holds per-layer context K/V +(``[max_batch, num_layers, window, 2 * num_kv_heads * head_dim]``, K then V, +K stored RoPE'd/k-normed). Frame convention: the worker passes window frames +``f = absolute_position + 1``; this model stores the context row of the token +at absolute position ``p`` at ring index ``p % window`` with RoPE phase ``p``, +matching the DeepSpec reference (its draft ``DynamicCache`` holds the context +K/V of positions ``0..start-1`` in order; the ring is an O(1)-memory window +over the most recent ``window`` positions — acceptance-rate only, standard +target verification keeps outputs correct regardless). + +Reference: https://github.com/deepseek-ai/DeepSpec +(``deepspec/modeling/dspark/qwen3/modeling.py`` and +``deepspec/eval/dspark/draft_ops.py``). +""" + +import os +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 .dspark.draft import build_draft_input_ids, dspark_propose +from .dspark.heads import build_markov_head + +# Ring-window length for the per-layer context K/V cache (per request slot). +# The DeepSpec reference attends over the FULL committed context; a ring +# window bounds the worker buffer at +# ``max_batch * num_layers * window * 2 * kv_dim`` bytes while remaining +# exactly full-context for sequences up to ``window`` tokens. +_DEFAULT_CTX_WINDOW = 2048 +_CTX_WINDOW_ENV = "TRTLLM_DSPARK_QWEN3_CTX_WINDOW" + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """HF-style rotary embedding on ``[..., num_heads, head_dim]``. + + ``cos``/``sin`` are ``[..., head_dim]`` (one row per position) and are + broadcast over the heads axis. + """ + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + return x * cos + _rotate_half(x) * sin + + +class _RMSNorm(nn.Module): + """Qwen3RMSNorm: fp32 normalize, cast back, then scale by the bf16 weight.""" + + def __init__(self, dim: int, eps: float): + super().__init__() + self.weight = nn.Parameter(torch.empty(dim)) + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dt = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return self.weight * xf.to(dt) + + +class _Attention(nn.Module): + """Weight container for one draft layer's GQA attention (names match ckpt).""" + + def __init__(self, hidden: int, n_heads: int, n_kv_heads: int, head_dim: int, eps: float): + super().__init__() + self.q_proj = nn.Linear(hidden, n_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) + self.o_proj = nn.Linear(n_heads * head_dim, hidden, bias=False) + self.q_norm = _RMSNorm(head_dim, eps) + self.k_norm = _RMSNorm(head_dim, eps) + + +class _MLP(nn.Module): + def __init__(self, hidden: int, intermediate: int): + super().__init__() + self.gate_proj = nn.Linear(hidden, intermediate, bias=False) + self.up_proj = nn.Linear(hidden, intermediate, bias=False) + self.down_proj = nn.Linear(intermediate, hidden, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _DecoderLayer(nn.Module): + def __init__( + self, + hidden: int, + intermediate: int, + n_heads: int, + n_kv_heads: int, + head_dim: int, + eps: float, + ): + super().__init__() + self.self_attn = _Attention(hidden, n_heads, n_kv_heads, head_dim, eps) + self.mlp = _MLP(hidden, intermediate) + self.input_layernorm = _RMSNorm(hidden, eps) + self.post_attention_layernorm = _RMSNorm(hidden, eps) + + +class _ConfidenceHead(nn.Module): + """DeepSpec ``AcceptRatePredictor`` (fp32 linear WITH bias for the Qwen3 + drafters). Inert scaffolding — the worker always passes + ``confidence_threshold=0.0`` — kept for checkpoint completeness.""" + + def __init__(self, input_dim: int, with_markov: bool): + super().__init__() + self.with_markov = bool(with_markov) + self.proj = nn.Linear(int(input_dim), 1, bias=True, 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 + return self.proj(features.float()).squeeze(-1) + + +class Qwen3DSparkDraftModel(nn.Module): + """The dense Qwen3 DSpark draft backbone + heads (DeepSpec reference math). + + Shares ``embed_tokens`` / ``lm_head`` with the target model (the checkpoint + carries frozen copies of both; they are identical to the target's, so the + shared modules are used and the copies skipped at load). + """ + + def __init__(self, model_config, block_size: Optional[int] = None): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.config = config + + self.hidden_size = int(config.hidden_size) + self.num_heads = int(config.num_attention_heads) + self.num_kv_heads = int(config.num_key_value_heads) + self.head_dim = int(getattr(config, "head_dim", None) or self.hidden_size // self.num_heads) + self.kv_dim = self.num_kv_heads * self.head_dim + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.softmax_scale = self.head_dim**-0.5 + eps = float(config.rms_norm_eps) + num_layers = int(config.num_hidden_layers) + # Worker-facing stage count (== draft layer count for this drafter). + self.num_stages = num_layers + + spec_cfg = getattr(model_config, "spec_config", None) + self.block_size = int( + block_size if block_size is not None else getattr(config, "block_size", 7) + ) + 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, "mask_token_id", config.vocab_size) + ) + + target_layer_ids = list(getattr(config, "target_layer_ids", []) or []) + self.num_capture_layers = len(target_layer_ids) + assert self.num_capture_layers > 0, ( + "Qwen3 DSpark drafter config must provide target_layer_ids" + ) + + # Plain-RoPE parameters. transformers>=5 nests rope_theta under + # rope_parameters; older versions keep the flat attribute. + rope_params = getattr(config, "rope_parameters", None) or {} + self._rope_theta = float( + rope_params.get("rope_theta", None) or getattr(config, "rope_theta", 1000000.0) + ) + max_pos = int(getattr(config, "max_position_embeddings", 40960)) + self._freqs_cap = max_pos + self.block_size + 2 + self._rope_cache: Dict[str, tuple] = {} + + # Ring-window length for the worker-owned context K/V buffer. + window = int(os.environ.get(_CTX_WINDOW_ENV, _DEFAULT_CTX_WINDOW)) + window = max(self.block_size + 2, min(window, max_pos)) + # Worker allocation contract (DSparkWorker._lazy_init): + # kv_windows = [max_batch, num_stages, window_size, head_dim], where + # "head_dim" is the per-position row width — here K||V flattened. + self._attn_params = dict( + window_size=window, + head_dim=2 * self.kv_dim, + ) + + # Build the weight modules on the meta device; load_weights() assigns + # the real cuda tensors (torch load_state_dict(assign=True)). + with torch.device("meta"): + self.fc = nn.Linear( + self.num_capture_layers * self.hidden_size, self.hidden_size, bias=False + ) + self.hidden_norm = _RMSNorm(self.hidden_size, eps) + self.layers = nn.ModuleList( + [ + _DecoderLayer( + self.hidden_size, + int(config.intermediate_size), + self.num_heads, + self.num_kv_heads, + self.head_dim, + eps, + ) + for _ in range(num_layers) + ] + ) + self.norm = _RMSNorm(self.hidden_size, eps) + self.markov_head = build_markov_head( + markov_head_type=str(getattr(config, "markov_head_type", "vanilla")), + vocab_size=int(config.vocab_size), + markov_rank=int(getattr(config, "markov_rank", 0)), + hidden_size=self.hidden_size, + ) + self.confidence_head = None + if bool(getattr(config, "enable_confidence_head", False)): + with_markov = bool(getattr(config, "confidence_head_with_markov", False)) + input_dim = self.hidden_size + (int(config.markov_rank) if with_markov else 0) + self.confidence_head = _ConfidenceHead(input_dim, with_markov) + + # Shared with target; wired by the spec wrapper after construction. + self.embed_tokens: Optional[nn.Module] = None + self.lm_head: Optional[nn.Module] = None + + # ------------------------------------------------------------------ RoPE + + def _cos_sin_tables(self, device: torch.device): + key = str(device) + cached = self._rope_cache.get(key) + if cached is None: + inv_freq = 1.0 / ( + self._rope_theta + ** ( + torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) + / self.head_dim + ) + ) + t = torch.arange(self._freqs_cap, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cached = (emb.cos(), emb.sin()) + self._rope_cache[key] = cached + return cached + + def _gather_cos_sin(self, positions: torch.Tensor, dtype: torch.dtype): + cos_t, sin_t = self._cos_sin_tables(positions.device) + # Clamp for graph-safety: masked-out entries may carry arbitrary + # (already clamped by the worker) positions. + p = positions.long().clamp(min=0, max=self._freqs_cap - 1) + return cos_t[p].to(dtype), sin_t[p].to(dtype) + + # ---------------------------------------------------------- context K/V + + def _project_ctx(self, main_hidden: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(captured))`` — the layer-invariant context stream.""" + return self.hidden_norm(self.fc(main_hidden)) + + def _ctx_kv( + self, layer: _DecoderLayer, main_x: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + """Per-layer context K/V rows ``[..., 2*kv_dim]`` (K RoPE'd/normed || V).""" + a = layer.self_attn + shp = main_x.shape[:-1] + k = a.k_proj(main_x).view(*shp, self.num_kv_heads, self.head_dim) + k = a.k_norm(k) + cos, sin = self._gather_cos_sin(positions, k.dtype) + k = _apply_rope(k, cos, sin) + v = a.v_proj(main_x) + return torch.cat([k.reshape(*shp, self.kv_dim), v], dim=-1) + + @torch.inference_mode() + def write_context_windows( + self, main_hidden: torch.Tensor, positions: torch.Tensor, stage_windows: torch.Tensor + ) -> None: + """Seed one request's ring windows from captured context (prefill path). + + Args: + main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. + positions: ``[M]`` window frames (= absolute position + 1, the + worker's generation-path convention). + stage_windows: ``[num_layers, window, 2*kv_dim]``, updated in place. + """ + M = int(main_hidden.shape[0]) + if M == 0: + return + win = int(self._attn_params["window_size"]) + p = positions.to(main_hidden.device).long() - 1 + cols = p % win + main_x = self._project_ctx(main_hidden) + for li, layer in enumerate(self.layers): + kv = self._ctx_kv(layer, main_x, p) + stage_windows[li, cols] = kv.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 masked back-fill of interim accepted tokens. + + Args: + main_hidden: ``[G, M, num_capture * hidden]``. + positions: ``[G, M]`` window frames (masked entries arbitrary). + slots: ``[G]`` request rows into ``kv_windows``. + mask: ``[G, M]`` bool validity. + kv_windows: ``[N, num_layers, window, 2*kv_dim]``, in place. + """ + G, M = positions.shape + if G == 0 or M == 0: + return + win = int(self._attn_params["window_size"]) + p = positions.long() - 1 + cols = p % win + rows = slots.long()[:, None].expand(-1, M) + mask3 = mask.unsqueeze(-1) + main_x = self._project_ctx(main_hidden) + for li, layer in enumerate(self.layers): + kv = self._ctx_kv(layer, main_x, p) # [G, M, 2*kv_dim] + win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view + cur = win_l[rows, cols] + win_l[rows, cols] = torch.where(mask3, kv.to(win_l.dtype), cur) + + # ------------------------------------------------------------- backbone + + 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 (all gen requests at once). + + Mirrors DeepSpec ``forward_dspark_draft_block`` + ``build_dspark_proposal``: + write the newest committed token's context K/V (position ``start_pos-1``), + run the ``block_size`` draft queries ``[bonus, mask, ...]`` at positions + ``start_pos + [0..block)`` through the Qwen3 layers attending to the ring + context + the block (bidirectional), then Markov-refine the block logits. + + Args: + main_hidden: ``[G, num_capture * hidden]`` captured hidden of the + newest committed context token per request. + bonus_token_ids: ``[G]`` last accepted token per request. + start_pos: ``[G]`` absolute decode position of the bonus token. + kv_windows: ``[N, num_layers, window, 2*kv_dim]`` persistent buffer. + slots: ``[G]`` request rows into ``kv_windows``. + Returns: + ``(draft_tokens [G, block], num_proposed [G])`` and, with + ``return_logits``, the corrected block logits ``[G, block, vocab]``. + """ + del all_rank_num_tokens # dense drafter: no cross-rank MoE lockstep + G = int(main_hidden.shape[0]) + B = self.block_size + win = int(self._attn_params["window_size"]) + device = main_hidden.device + start_pos = start_pos.long() + + main_x = self._project_ctx(main_hidden) # [G, hidden] + p0 = start_pos - 1 + cols0 = p0 % win + + draft_ids = build_draft_input_ids( + bonus_token_ids, block_size=B, noise_token_id=self.noise_token_id + ) + h = self.embed_tokens(draft_ids) # [G, B, hidden] + + pos_q = start_pos.unsqueeze(1) + torch.arange(B, device=device) # [G, B] + cos_q, sin_q = self._gather_cos_sin(pos_q, h.dtype) + + # Bool attention mask [G, 1, B, win + B]: ring rows valid iff their + # index < min(start_pos, win) (ring holds the last `win` positions); + # the block attends to itself bidirectionally. + ctx_valid = ( + torch.arange(win, device=device)[None, :] < start_pos.clamp(max=win)[:, None] + ) # [G, win] + attn_mask = torch.cat( + [ + ctx_valid[:, None, None, :].expand(G, 1, B, win), + torch.ones(G, 1, B, B, dtype=torch.bool, device=device), + ], + dim=-1, + ) + + for li, layer in enumerate(self.layers): + a = layer.self_attn + win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view + # Write the newest committed token's context row, then read the ring. + kv0 = self._ctx_kv(layer, main_x, p0) # [G, 2*kv_dim] + win_l[slots, cols0] = kv0.to(win_l.dtype) + ctx = win_l[slots] # [G, win, 2*kv_dim] + k_ctx = ctx[..., : self.kv_dim].view(G, win, self.num_kv_heads, self.head_dim) + v_ctx = ctx[..., self.kv_dim :].view(G, win, self.num_kv_heads, self.head_dim) + + residual = h + x = layer.input_layernorm(h) + q = a.q_norm(a.q_proj(x).view(G, B, self.num_heads, self.head_dim)) + q = _apply_rope(q, cos_q, sin_q) + k_blk = a.k_norm(a.k_proj(x).view(G, B, self.num_kv_heads, self.head_dim)) + k_blk = _apply_rope(k_blk, cos_q, sin_q) + v_blk = a.v_proj(x).view(G, B, self.num_kv_heads, self.head_dim) + + k = torch.cat([k_ctx.to(h.dtype), k_blk], dim=1).transpose(1, 2) + v = torch.cat([v_ctx.to(h.dtype), v_blk], dim=1).transpose(1, 2) + k = k.repeat_interleave(self.num_kv_groups, dim=1) + v = v.repeat_interleave(self.num_kv_groups, dim=1) + o = F.scaled_dot_product_attention( + q.transpose(1, 2), k, v, attn_mask=attn_mask, scale=self.softmax_scale + ) + o = o.transpose(1, 2).reshape(G, B, self.num_heads * self.head_dim) + h = residual + a.o_proj(o) + + residual = h + h = residual + layer.mlp(layer.post_attention_layernorm(h)) + + h = self.norm(h) + base_logits = self.lm_head(h) + return dspark_propose( + base_logits, + bonus_token_ids=bonus_token_ids, + block_hidden=h, + markov_head=self.markov_head, + confidence_head=self.confidence_head, + block_size=B, + temperature=temperature, + confidence_threshold=confidence_threshold, + return_logits=return_logits, + ) + + def forward( + self, + main_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + start_pos, + *, + kv_windows: Optional[torch.Tensor] = None, + **kwargs, + ) -> tuple: + """Eager/single-shot convenience wrapper over :meth:`forward_batched`.""" + T = int(main_hidden.shape[0]) + device = main_hidden.device + if not torch.is_tensor(start_pos): + start_pos = torch.full((T,), int(start_pos), dtype=torch.long, device=device) + if kv_windows is None: + kv_windows = torch.zeros( + ( + T, + self.num_stages, + self._attn_params["window_size"], + self._attn_params["head_dim"], + ), + dtype=torch.bfloat16, + device=device, + ) + slots = torch.arange(T, device=device) + return self.forward_batched( + main_hidden, bonus_token_ids, start_pos, kv_windows=kv_windows, slots=slots, **kwargs + ) + + def run_moe_lockstep_noop(self, all_rank_num_tokens, device) -> None: + """Dense drafter: no cross-rank MoE barrier to keep in lockstep.""" + return None + + # --------------------------------------------------------------- loading + + def load_weights(self, weights: Dict) -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + sd = {} + for k, v in weights.items(): + # embed_tokens / lm_head are frozen copies of the target's; the + # shared target modules are used instead (load_weights_from_target_model). + if k.startswith(("embed_tokens.", "lm_head.")): + continue + t = v.to(device) + t = t.float() if k.startswith("confidence_head.") else t.to(torch.bfloat16) + sd[k] = t + self.load_state_dict(sd, strict=True, assign=True) + logger.info( + f"[DSpark-Qwen3] loaded {len(sd)} draft params " + f"({self.num_stages} layers, block_size={self.block_size}, " + f"ctx_window={self._attn_params['window_size']})" + ) + + +class Qwen3DSparkForCausalLM(nn.Module): + """One-engine draft wrapper for the Qwen3 DSpark drafter. + + Mirrors :class:`DSparkForCausalLM`: created by ``get_draft_model``, + appended to the target's epilogue, and driven by ``DSparkWorker`` through + the shared draft-model protocol. ``embed_tokens`` / ``lm_head`` are shared + with the target model (the drafter checkpoint's copies are identical + frozen snapshots and are skipped at load). + """ + + def __init__(self, draft_config, block_size: Optional[int] = None): + super().__init__() + self.dspark_model = Qwen3DSparkDraftModel(draft_config, 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. + self.num_stages = self.dspark_model.num_stages + self._attn_params = self.dspark_model._attn_params + self.lm_head = None # shared from the target + 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): + 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): + return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) + + def write_context_windows(self, main_hidden, positions, stage_windows): + 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): + 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 flat-namespace Qwen3 DSpark drafter checkpoint. + + ``weight_mapper`` is accepted for loader-interface parity but unused — + the checkpoint names map 1:1 onto the module tree. + """ + self.dspark_model.load_weights(weights) + + def load_weights_from_target_model(self, target_model): + """Share the target's embed_tokens / lm_head (identical frozen copies).""" + 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__ = ["Qwen3DSparkDraftModel", "Qwen3DSparkForCausalLM"] diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index a991e266965d..b1d4a86b1b28 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1911,6 +1911,15 @@ def get_draft_model(model_config, draft_config, lm_head, model): return DFlashLagunaForCausalLM(draft_config) return DFlashForCausalLM(draft_config) elif spec_dec_mode.is_dspark(): + # DeepSpec-released dense drafters (e.g. dspark_qwen3_8b_block7) are + # separate checkpoints that declare their own architecture; the + # DeepSeek-V4 drafter lives in the target checkpoint's mtp.* namespace. + draft_arches = getattr(draft_config.pretrained_config, "architectures", + None) or [] + if any("Qwen3DSpark" in arch for arch in draft_arches): + from .modeling_dspark_qwen3 import Qwen3DSparkForCausalLM + return Qwen3DSparkForCausalLM( + draft_config, block_size=model_config.spec_config.block_size) # 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 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 00f2d2c4f678..dc8e266510cf 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5731,6 +5731,11 @@ def _dspark_get(key, top_level_key): value = dspark_cfg.get(key) if value is None: value = draft_cfg.get(top_level_key) + if value is None: + # DeepSpec-released dense drafter checkpoints + # (e.g. Qwen3DSparkModel) use unprefixed top-level + # keys in their own config.json. + value = draft_cfg.get(key) return value # The checkpoint's ``dspark_target_layer_ids`` is diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py new file mode 100644 index 000000000000..5f4bfe67c9f1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py @@ -0,0 +1,424 @@ +# 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. +"""Golden tests for the Qwen3 DSpark drafter (modeling_dspark_qwen3.py). + +The reference implementation below is a line-for-line torch-only port of the +DeepSpec ``Qwen3DSparkModel`` inference path +(``deepspec/modeling/dspark/qwen3/modeling.py`` `_forward_backbone` + +``deepspec/eval/dspark/draft_ops.py`` `forward_dspark_draft_block` / +`build_dspark_proposal`): full-context draft attention over +``[ctx_kv_cache, block]`` with bidirectional block attention, HF-style RoPE, +per-head q/k RMSNorm, and greedy Markov-chained block sampling. + +The tests drive ``Qwen3DSparkDraftModel`` exactly the way ``DSparkWorker`` +does — seeding via ``write_context_windows`` (frames = position + 1), +back-filling interims via ``write_context_windows_batched``, drafting via +``forward_batched`` — across multiple decode steps, and assert token-exact / +logit-close agreement with the reference. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.models.modeling_dspark_qwen3 import Qwen3DSparkDraftModel, _apply_rope + +VOCAB = 97 +HID = 32 +INTER = 48 +N_HEADS = 4 +N_KV_HEADS = 2 +HEAD_DIM = 8 +N_LAYERS = 2 +N_CAP = 3 +BLOCK = 3 +MASK_ID = 90 +RANK = 16 +THETA = 10000.0 +MAX_POS = 512 +EPS = 1e-6 + +DTYPE = torch.bfloat16 + + +def _make_config(): + pretrained = SimpleNamespace( + architectures=["Qwen3DSparkModel"], + hidden_size=HID, + intermediate_size=INTER, + num_attention_heads=N_HEADS, + num_key_value_heads=N_KV_HEADS, + head_dim=HEAD_DIM, + num_hidden_layers=N_LAYERS, + vocab_size=VOCAB, + rms_norm_eps=EPS, + rope_parameters={"rope_theta": THETA}, + max_position_embeddings=MAX_POS, + block_size=BLOCK, + mask_token_id=MASK_ID, + target_layer_ids=[0, 2, 4], + markov_rank=RANK, + markov_head_type="vanilla", + enable_confidence_head=True, + confidence_head_with_markov=True, + num_anchors=8, + ) + return SimpleNamespace(pretrained_config=pretrained, spec_config=None) + + +def _rand_weights(gen): + def r(*shape, scale=0.05): + return (torch.randn(*shape, generator=gen) * scale).to(DTYPE) + + w = { + "fc.weight": r(HID, N_CAP * HID), + "hidden_norm.weight": 1.0 + r(HID), + "norm.weight": 1.0 + r(HID), + "markov_head.markov_w1.weight": r(VOCAB, RANK), + "markov_head.markov_w2.weight": r(VOCAB, RANK), + "confidence_head.proj.weight": r(1, HID + RANK), + "confidence_head.proj.bias": r(1), + # frozen copies (skipped by the model; shared modules used instead) + "embed_tokens.weight": r(VOCAB, HID, scale=0.5), + "lm_head.weight": r(VOCAB, HID, scale=0.5), + } + for i in range(N_LAYERS): + p = f"layers.{i}." + w[p + "self_attn.q_proj.weight"] = r(N_HEADS * HEAD_DIM, HID) + w[p + "self_attn.k_proj.weight"] = r(N_KV_HEADS * HEAD_DIM, HID) + w[p + "self_attn.v_proj.weight"] = r(N_KV_HEADS * HEAD_DIM, HID) + w[p + "self_attn.o_proj.weight"] = r(HID, N_HEADS * HEAD_DIM) + w[p + "self_attn.q_norm.weight"] = 1.0 + r(HEAD_DIM) + w[p + "self_attn.k_norm.weight"] = 1.0 + r(HEAD_DIM) + w[p + "mlp.gate_proj.weight"] = r(INTER, HID) + w[p + "mlp.up_proj.weight"] = r(INTER, HID) + w[p + "mlp.down_proj.weight"] = r(HID, INTER) + w[p + "input_layernorm.weight"] = 1.0 + r(HID) + w[p + "post_attention_layernorm.weight"] = 1.0 + r(HID) + return w + + +def _build_model(weights): + model = Qwen3DSparkDraftModel(_make_config()) + model.load_weights(weights) + device = model.fc.weight.device + embed = torch.nn.Embedding(VOCAB, HID) + embed.weight.data = weights["embed_tokens.weight"].to(device) + lm_head = torch.nn.Linear(HID, VOCAB, bias=False) + lm_head.weight.data = weights["lm_head.weight"].to(device) + model.embed_tokens = embed.to(device) + model.lm_head = lm_head.to(device) + return model, device + + +# -------------------------------------------------------------------------- +# Reference: DeepSpec Qwen3DSparkModel inference path (torch-only port) +# -------------------------------------------------------------------------- + + +class _Ref: + """Full-context reference drafter operating on one request.""" + + def __init__(self, w, device): + self.w = {k: v.to(device) for k, v in w.items()} + self.device = device + inv = 1.0 / ( + THETA ** (torch.arange(0, HEAD_DIM, 2, device=device, dtype=torch.float32) / HEAD_DIM) + ) + t = torch.arange(MAX_POS, device=device, dtype=torch.float32) + emb = torch.cat([torch.outer(t, inv)] * 2, dim=-1) + self.cos, self.sin = emb.cos(), emb.sin() + # per-request context stream: [T, HID] projected hiddens, in order + self.ctx_x = torch.zeros(0, HID, dtype=DTYPE, device=device) + + def _norm(self, x, wname): + wt = self.w[wname] + dt = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + EPS) + return wt * xf.to(dt) + + def append_ctx(self, captured): + """captured: [M, N_CAP*HID] raw target hiddens (positions in order).""" + x = F.linear(captured.to(self.device), self.w["fc.weight"]) + x = self._norm(x, "hidden_norm.weight") + self.ctx_x = torch.cat([self.ctx_x, x.to(DTYPE)], dim=0) + + def draft(self, bonus_id, start_pos): + assert self.ctx_x.shape[0] == start_pos + w = self.w + ids = torch.full((BLOCK,), MASK_ID, dtype=torch.long, device=self.device) + ids[0] = bonus_id + h = F.embedding(ids, w["embed_tokens.weight"]) # [B, HID] + pos_q = start_pos + torch.arange(BLOCK, device=self.device) + pos_c = torch.arange(start_pos, device=self.device) + for i in range(N_LAYERS): + p = f"layers.{i}." + x = self._norm(h, p + "input_layernorm.weight") + q = F.linear(x, w[p + "self_attn.q_proj.weight"]).view(BLOCK, N_HEADS, HEAD_DIM) + q = self._norm(q, p + "self_attn.q_norm.weight") + q = _apply_rope(q, self.cos[pos_q].to(DTYPE), self.sin[pos_q].to(DTYPE)) + src = torch.cat([self.ctx_x, x], dim=0) # [T+B, HID] + pos_k = torch.cat([pos_c, pos_q]) + k = F.linear(src, w[p + "self_attn.k_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) + k = self._norm(k, p + "self_attn.k_norm.weight") + k = _apply_rope(k, self.cos[pos_k].to(DTYPE), self.sin[pos_k].to(DTYPE)) + v = F.linear(src, w[p + "self_attn.v_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) + rep = N_HEADS // N_KV_HEADS + kk = k.transpose(0, 1).repeat_interleave(rep, dim=0) + vv = v.transpose(0, 1).repeat_interleave(rep, dim=0) + o = F.scaled_dot_product_attention(q.transpose(0, 1), kk, vv, scale=HEAD_DIM**-0.5) + o = o.transpose(0, 1).reshape(BLOCK, N_HEADS * HEAD_DIM) + h = h + F.linear(o, w[p + "self_attn.o_proj.weight"]) + x = self._norm(h, p + "post_attention_layernorm.weight") + mlp = F.linear( + F.silu(F.linear(x, w[p + "mlp.gate_proj.weight"])) + * F.linear(x, w[p + "mlp.up_proj.weight"]), + w[p + "mlp.down_proj.weight"], + ) + h = h + mlp + h = self._norm(h, "norm.weight") + base = F.linear(h, w["lm_head.weight"]) # [B, VOCAB] + # Greedy Markov-chained sampling (VanillaMarkov). + toks, logits = [], [] + prev = bonus_id.view(1).long() + for kstep in range(BLOCK): + bias = F.linear( + F.embedding(prev, w["markov_head.markov_w1.weight"]), + w["markov_head.markov_w2.weight"], + ) + step = base[kstep : kstep + 1] + bias + logits.append(step) + prev = step.argmax(dim=-1) + toks.append(prev) + return torch.cat(toks), torch.cat(logits) + + +@pytest.fixture(scope="module") +def setup(): + gen = torch.Generator().manual_seed(1234) + weights = _rand_weights(gen) + model, device = _build_model(weights) + return weights, model, device + + +def _rand_hidden(gen, n): + return (torch.randn(n, N_CAP * HID, generator=gen) * 0.3).to(DTYPE) + + +def test_worker_protocol_golden(setup): + """Drive the model exactly like DSparkWorker across 3 decode steps.""" + weights, model, device = setup + gen = torch.Generator().manual_seed(7) + + prompt_len = 6 + max_batch = 4 + win = model._attn_params["window_size"] + kv_windows = torch.zeros( + max_batch, model.num_stages, win, model._attn_params["head_dim"], dtype=DTYPE, device=device + ) + slot = 1 + ref = _Ref(weights, device) + + # ---- prefill: the worker seeds ALL prompt positions (frames = pos+1). + # The target processed positions 0..L-1, so the ring holds 0..L-1 and + # start_pos == L for the first draft; each gen step's main_hidden is the + # newest committed token's captured hidden (position start_pos-1). + prompt_hidden = _rand_hidden(gen, prompt_len).to(device) + positions = torch.arange(prompt_len, device=device) + 1 + model.write_context_windows(prompt_hidden, positions, kv_windows[slot]) + ref.append_ctx(prompt_hidden) + + start_pos = prompt_len + bonus_val = 13 + + for step in range(3): + # The worker's gen step: target verified some tokens; captured hiddens + # for the (nacc) newly committed context tokens exist. nacc-1 interim + # rows are back-filled batched; the newest one rides in main_hidden. + nacc = [1, 3, 2][step] # 1 = no interim (only the bonus position) + new_hidden = _rand_hidden(gen, nacc).to(device) + + if nacc > 1: + interim = new_hidden[: nacc - 1].unsqueeze(0) # [1, nacc-1, ...] + # frames old+1+j for j in 0..nacc-2 (old = pre-step start_pos) + interim_pos = (start_pos + 1 + torch.arange(nacc - 1, device=device)).unsqueeze(0) + mask = torch.ones(1, nacc - 1, dtype=torch.bool, device=device) + model.write_context_windows_batched( + interim, interim_pos, torch.tensor([slot], device=device), mask, kv_windows + ) + + ref.append_ctx(new_hidden) + start_pos += nacc + bonus = torch.tensor([bonus_val], device=device) + + got_toks, got_n, got_logits = model.forward_batched( + new_hidden[-1:].reshape(1, -1), + bonus, + torch.tensor([start_pos], device=device), + kv_windows=kv_windows, + slots=torch.tensor([slot], device=device), + return_logits=True, + ) + exp_toks, exp_logits = ref.draft(bonus[0], start_pos) + + assert got_n.item() == BLOCK + torch.testing.assert_close(got_logits[0].float(), exp_logits.float(), atol=0.05, rtol=0.05) + assert torch.equal(got_toks[0].long().cpu(), exp_toks.long().cpu()), ( + f"step {step}: draft tokens diverge from reference" + ) + bonus_val = (bonus_val * 7 + 3) % VOCAB + + +def test_batched_matches_eager_singletons(setup): + """One batched call over G requests == G independent eager calls.""" + weights, model, device = setup + gen = torch.Generator().manual_seed(21) + G = 3 + win = model._attn_params["window_size"] + kv_windows = torch.zeros( + G, model.num_stages, win, model._attn_params["head_dim"], dtype=DTYPE, device=device + ) + ctx_lens = [4, 7, 5] + for g in range(G): + h = _rand_hidden(gen, ctx_lens[g]).to(device) + pos = torch.arange(ctx_lens[g], device=device) + 1 + model.write_context_windows(h, pos, kv_windows[g]) + + main = _rand_hidden(gen, G).to(device) + bonus = torch.tensor([11, 22, 33], device=device) + start = torch.tensor([c + 1 for c in ctx_lens], device=device) + # per-request singleton calls on cloned windows + exp_toks, exp_logits = [], [] + for g in range(G): + wins = kv_windows[g : g + 1].clone() + t, _, lg = model.forward_batched( + main[g : g + 1], + bonus[g : g + 1], + start[g : g + 1], + kv_windows=wins, + slots=torch.tensor([0], device=device), + return_logits=True, + ) + exp_toks.append(t) + exp_logits.append(lg) + got_toks, _, got_logits = model.forward_batched( + main, + bonus, + start, + kv_windows=kv_windows, + slots=torch.arange(G, device=device), + return_logits=True, + ) + assert torch.equal(got_toks, torch.cat(exp_toks, dim=0)) + torch.testing.assert_close(got_logits, torch.cat(exp_logits, dim=0)) + + +def test_ring_window_wraparound(setup, monkeypatch): + """start_pos > window: the ring holds the last `win` positions and the + draft attends to all of them (mask all-valid).""" + weights, _, device = setup + monkeypatch.setenv("TRTLLM_DSPARK_QWEN3_CTX_WINDOW", "32") + model, device = _build_model(weights) + gen = torch.Generator().manual_seed(5) + win = model._attn_params["window_size"] + assert win == 32 + total = win + 9 # forces wraparound + kv_windows = torch.zeros( + 1, model.num_stages, win, model._attn_params["head_dim"], dtype=DTYPE, device=device + ) + h = _rand_hidden(gen, total).to(device) + pos = torch.arange(total, device=device) + 1 + # Seed in two chunks like chunked prefill (worker keeps last min(win, len)). + model.write_context_windows(h[:win], pos[:win], kv_windows[0]) + model.write_context_windows(h[win:], pos[win:], kv_windows[0]) + + # Reference limited to the last `win` context positions. + ref = _Ref(weights, device) + ref.append_ctx(h) + ref.ctx_x = ref.ctx_x[-win:] + + bonus = torch.tensor([42], device=device) + start = torch.tensor([total + 1], device=device) + main = _rand_hidden(gen, 1).to(device) + # keep ref in sync: main_hidden row is position `total` (= start-1) + ref.append_ctx(main) + ref.ctx_x = ref.ctx_x[-win:] + + got_toks, _, got_logits = model.forward_batched( + main, + bonus, + start, + kv_windows=kv_windows, + slots=torch.tensor([0], device=device), + return_logits=True, + ) + + # Reference drafts with positions: ctx = last `win` absolute positions. + class _WrapRef(_Ref): + pass + + wref = _WrapRef(weights, device) + wref.ctx_x = ref.ctx_x + # override position bookkeeping: ctx positions are total+1-win .. total + w = wref.w + ids = torch.full((BLOCK,), MASK_ID, dtype=torch.long, device=device) + ids[0] = bonus[0] + hh = F.embedding(ids, w["embed_tokens.weight"]) + pos_q = start[0] + torch.arange(BLOCK, device=device) + pos_c = torch.arange(start[0] - win, start[0], device=device) + for i in range(N_LAYERS): + p = f"layers.{i}." + x = wref._norm(hh, p + "input_layernorm.weight") + q = F.linear(x, w[p + "self_attn.q_proj.weight"]).view(BLOCK, N_HEADS, HEAD_DIM) + q = wref._norm(q, p + "self_attn.q_norm.weight") + q = _apply_rope(q, wref.cos[pos_q].to(DTYPE), wref.sin[pos_q].to(DTYPE)) + src = torch.cat([wref.ctx_x, x], dim=0) + pos_k = torch.cat([pos_c, pos_q]) + k = F.linear(src, w[p + "self_attn.k_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) + k = wref._norm(k, p + "self_attn.k_norm.weight") + k = _apply_rope(k, wref.cos[pos_k].to(DTYPE), wref.sin[pos_k].to(DTYPE)) + v = F.linear(src, w[p + "self_attn.v_proj.weight"]).view(-1, N_KV_HEADS, HEAD_DIM) + rep = N_HEADS // N_KV_HEADS + kk = k.transpose(0, 1).repeat_interleave(rep, dim=0) + vv = v.transpose(0, 1).repeat_interleave(rep, dim=0) + o = F.scaled_dot_product_attention(q.transpose(0, 1), kk, vv, scale=HEAD_DIM**-0.5) + o = o.transpose(0, 1).reshape(BLOCK, N_HEADS * HEAD_DIM) + hh = hh + F.linear(o, w[p + "self_attn.o_proj.weight"]) + x = wref._norm(hh, p + "post_attention_layernorm.weight") + mlp = F.linear( + F.silu(F.linear(x, w[p + "mlp.gate_proj.weight"])) + * F.linear(x, w[p + "mlp.up_proj.weight"]), + w[p + "mlp.down_proj.weight"], + ) + hh = hh + mlp + hh = wref._norm(hh, "norm.weight") + base = F.linear(hh, w["lm_head.weight"]) + toks, logits = [], [] + prev = bonus.long() + for kstep in range(BLOCK): + bias = F.linear( + F.embedding(prev, w["markov_head.markov_w1.weight"]), w["markov_head.markov_w2.weight"] + ) + step = base[kstep : kstep + 1] + bias + logits.append(step) + prev = step.argmax(dim=-1) + toks.append(prev) + torch.testing.assert_close( + got_logits[0].float(), torch.cat(logits).float(), atol=0.05, rtol=0.05 + ) + assert torch.equal(got_toks[0].long(), torch.cat(toks).long()) From a23949d86318cff6063324fac3033a4609f75692 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Wed, 5 Aug 2026 18:15:24 +0000 Subject: [PATCH 2/3] [None][refactor] blend Qwen3-based speculator support into existing framework Fold modeling_dspark_qwen3.py into the shared DSpark modeling path instead of keeping a parallel Qwen3-specific implementation, so drafter construction, heads and worker plumbing go through one code path. Signed-off-by: chungen04 --- tensorrt_llm/_torch/models/dspark/draft.py | 18 + tensorrt_llm/_torch/models/dspark/heads.py | 14 +- tensorrt_llm/_torch/models/modeling_dspark.py | 703 ++++++++++++++++-- .../_torch/models/modeling_dspark_qwen3.py | 602 --------------- .../_torch/models/modeling_speculative.py | 24 +- .../hw_agnostic/test_dspark_qwen3.py | 57 +- 6 files changed, 751 insertions(+), 667 deletions(-) delete mode 100644 tensorrt_llm/_torch/models/modeling_dspark_qwen3.py diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 1b47f4922965..aa76ffd5a725 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -38,6 +38,24 @@ from .heads import confident_prefix_length +def resolve_noise_token_id(mask_token_id: Optional[int], config, ckpt_attr: str) -> int: + """Resolve the DSpark noise/mask token id (shared by every drafter). + + ``mask_token_id`` is the speculative config's value — the user's override, + or the drafter checkpoint's id copied in by ``DSparkDecodingConfig`` + validation. ``None`` (direct construction, e.g. tests) falls back to the + drafter checkpoint's ``ckpt_attr`` (``dspark_noise_token_id`` for the V4 + ``mtp.*`` drafter, ``mask_token_id`` for the DeepSpec Qwen3 checkpoints), + else ``vocab_size``. A config that declares the key with a null value is + treated as absent. + """ + if mask_token_id is None: + mask_token_id = getattr(config, ckpt_attr, None) + if mask_token_id is None: + mask_token_id = config.vocab_size + return int(mask_token_id) + + def build_draft_input_ids( bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index c49e35fbafa0..642c6b2f8b6a 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -217,14 +217,22 @@ class DSparkConfidenceHead(nn.Module): 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): + def __init__( + self, + *, + hidden_size: int, + markov_rank: int = 0, + with_markov: bool = False, + bias: 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 + # The V4-Pro checkpoint stores ``proj`` as a bias-free bf16 weight; the + # DeepSpec Qwen3 drafter checkpoints carry a bias. Either way 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) + self.proj = nn.Linear(input_dim, 1, bias=bias, dtype=torch.float32) def forward( self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 290732a73037..bfbd989ec48f 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -15,9 +15,21 @@ # # 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. +"""DSpark speculative-decoding draft backbones. -The DSpark draft is ``n_mtp_layers`` (3 for V4-Pro) **full DeepSeek-V4 blocks** +This module is the home of both DSpark drafters: + + - the DeepSeek-V4-Pro drafter (``DSparkDraftModel``), whose weights live in + the target checkpoint's ``mtp.*`` namespace; + - the DeepSpec-released dense Qwen3 drafter (``Qwen3DSparkDraftModel``), + a separate flat-namespace checkpoint. + +Both implement the same worker-facing protocol (``DSparkForCausalLMBase``) and +share the head / propose / input helpers in the ``dspark/`` package, so +``DSparkWorker`` / ``DSparkSpecMetadata`` / CUDA-graph plumbing drive them +unchanged. + +The DeepSeek-V4 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: @@ -60,7 +72,7 @@ dspark_attention_forward_batched, precompute_dspark_freqs_cis, ) -from .dspark.draft import build_draft_input_ids, dspark_propose +from .dspark.draft import build_draft_input_ids, dspark_propose, resolve_noise_token_id from .dspark.heads import DSparkConfidenceHead, build_markov_head from .modeling_deepseekv4 import ( DeepseekV4DecoderLayer, @@ -275,13 +287,8 @@ def __init__( 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.noise_token_id = resolve_noise_token_id( + getattr(spec_cfg, "mask_token_id", None), config, "dspark_noise_token_id" ) self.markov_rank = int(getattr(config, "dspark_markov_rank", 0)) self.hc_mult = config.hc_mult @@ -351,6 +358,7 @@ def __init__( aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], num_stages: Optional[int] = None, block_size: Optional[int] = None, + mask_token_id: Optional[int] = None, ): super().__init__() config = model_config.pretrained_config @@ -370,19 +378,14 @@ def __init__( 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. + # Production passes the validated speculative-config values explicitly + # (the drafter's own ModelConfig carries spec_config=None to avoid + # recursive spec-dec, so they cannot be read back off ``spec_cfg``); + # direct construction falls back to the checkpoint's trained values. 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.noise_token_id = resolve_noise_token_id(mask_token_id, config, "dspark_noise_token_id") self.hc_mult = config.hc_mult target_layer_ids = getattr(config, "dspark_target_layer_ids", []) self.num_capture_layers = len(target_layer_ids) @@ -1146,40 +1149,578 @@ def forward_head( ) -class DSparkForCausalLM(nn.Module): - """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). +# --------------------------------------------------------------------------- # +# Qwen3 (DeepSpec) DSpark drafter +# --------------------------------------------------------------------------- # - 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``. +# Ring-window length for the per-layer context K/V cache (per request slot). +# The DeepSpec reference attends over the FULL committed context; a ring +# window bounds the worker buffer at +# ``max_batch * num_layers * window * 2 * kv_dim`` bytes while remaining +# exactly full-context for sequences up to ``window`` tokens. +_DEFAULT_CTX_WINDOW = 2048 +_CTX_WINDOW_ENV = "TRTLLM_DSPARK_QWEN3_CTX_WINDOW" - ``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 _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """HF-style rotary embedding on ``[..., num_heads, head_dim]``. + + ``cos``/``sin`` are ``[..., head_dim]`` (one row per position) and are + broadcast over the heads axis. """ + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + return x * cos + _rotate_half(x) * sin - def __init__(self, draft_config, aux_stream_dict=None, num_stages=None, block_size=None): + +class _Qwen3RMSNorm(nn.Module): + """Qwen3RMSNorm: fp32 normalize, cast back, then scale by the bf16 weight.""" + + def __init__(self, dim: int, eps: float): super().__init__() - self.dspark_model = DSparkDraftModel( - draft_config, - aux_stream_dict, - num_stages=num_stages, - block_size=block_size, + self.weight = nn.Parameter(torch.empty(dim)) + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dt = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return self.weight * xf.to(dt) + + +class _Qwen3Attention(nn.Module): + """Weight container for one draft layer's GQA attention (names match ckpt).""" + + def __init__(self, hidden: int, n_heads: int, n_kv_heads: int, head_dim: int, eps: float): + super().__init__() + self.q_proj = nn.Linear(hidden, n_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) + self.o_proj = nn.Linear(n_heads * head_dim, hidden, bias=False) + self.q_norm = _Qwen3RMSNorm(head_dim, eps) + self.k_norm = _Qwen3RMSNorm(head_dim, eps) + + +class _Qwen3MLP(nn.Module): + def __init__(self, hidden: int, intermediate: int): + super().__init__() + self.gate_proj = nn.Linear(hidden, intermediate, bias=False) + self.up_proj = nn.Linear(hidden, intermediate, bias=False) + self.down_proj = nn.Linear(intermediate, hidden, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _Qwen3DecoderLayer(nn.Module): + def __init__( + self, + hidden: int, + intermediate: int, + n_heads: int, + n_kv_heads: int, + head_dim: int, + eps: float, + ): + super().__init__() + self.self_attn = _Qwen3Attention(hidden, n_heads, n_kv_heads, head_dim, eps) + self.mlp = _Qwen3MLP(hidden, intermediate) + self.input_layernorm = _Qwen3RMSNorm(hidden, eps) + self.post_attention_layernorm = _Qwen3RMSNorm(hidden, eps) + + +class Qwen3DSparkDraftModel(nn.Module): + """Qwen3-based DSpark draft backbone + heads (DeepSpec ``Qwen3DSparkModel``). + + Unlike the DeepSeek-V4 drafter above — whose draft weights live in the + target checkpoint's ``mtp.*`` namespace and whose stages are full V4 blocks + (MLA + MoE + mHC) — the DeepSpec-released Qwen3 drafters (e.g. + ``deepseek-ai/dspark_qwen3_8b_block7``) are separate dense bf16 checkpoints + with a flat namespace: + + - ``fc`` + ``hidden_norm``: project the concatenated captured + target-layer hidden states (``hidden_size * num_capture_layers``) into + the draft hidden size — the layer-invariant *context stream* + (``main_x``). + - ``layers.{0..n-1}``: standard Qwen3 GQA decoder layers (q/k per-head + RMSNorm, RoPE, gated-SiLU MLP), except attention keys/values come from + BOTH the projected context (one row per committed token) and the draft + block itself, and the ``block_size`` draft queries attend + bidirectionally within the block. + - ``norm`` + (target-shared) ``lm_head`` -> :func:`dspark_propose` + (Markov head refinement; the confidence head is inert scaffolding, as + in the V4 path). + + The worker-facing protocol matches :class:`DSparkDraftModel`, so + ``DSparkWorker`` and ``DSparkSpecMetadata`` drive both drafters unchanged. + The worker-owned rolling buffer here holds per-layer context K/V + (``[max_batch, num_layers, window, 2 * num_kv_heads * head_dim]``, K then + V, K stored RoPE'd/k-normed). Frame convention: the worker passes window + frames ``f = absolute_position + 1``; this model stores the context row of + the token at absolute position ``p`` at ring index ``p % window`` with + RoPE phase ``p``, matching the DeepSpec reference (its draft + ``DynamicCache`` holds the context K/V of positions ``0..start-1`` in + order; the ring is an O(1)-memory window over the most recent ``window`` + positions — acceptance-rate only, standard target verification keeps + outputs correct regardless). + + ``embed_tokens`` / ``lm_head`` are shared with the target model (the + checkpoint carries frozen copies of both; they are identical to the + target's, so the shared modules are used and the copies skipped at load). + + Reference: https://github.com/deepseek-ai/DeepSpec + (``deepspec/modeling/dspark/qwen3/modeling.py`` and + ``deepspec/eval/dspark/draft_ops.py``). + """ + + def __init__( + self, + model_config, + block_size: Optional[int] = None, + mask_token_id: Optional[int] = None, + ): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.config = config + + self.hidden_size = int(config.hidden_size) + self.num_heads = int(config.num_attention_heads) + self.num_kv_heads = int(config.num_key_value_heads) + self.head_dim = int(getattr(config, "head_dim", None) or self.hidden_size // self.num_heads) + self.kv_dim = self.num_kv_heads * self.head_dim + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.softmax_scale = self.head_dim**-0.5 + eps = float(config.rms_norm_eps) + num_layers = int(config.num_hidden_layers) + # Worker-facing stage count (== draft layer count for this drafter). + self.num_stages = num_layers + + # Production passes the validated speculative-config values explicitly + # (the drafter's own ModelConfig carries spec_config=None to avoid + # recursive spec-dec, so they cannot be read back off it); direct + # construction falls back to the checkpoint's trained values. + self.block_size = int( + block_size if block_size is not None else getattr(config, "block_size", 7) + ) + self.noise_token_id = resolve_noise_token_id(mask_token_id, config, "mask_token_id") + + target_layer_ids = list(getattr(config, "target_layer_ids", []) or []) + self.num_capture_layers = len(target_layer_ids) + assert self.num_capture_layers > 0, ( + "Qwen3 DSpark drafter config must provide target_layer_ids" + ) + + # Plain-RoPE parameters. transformers>=5 nests rope_theta under + # rope_parameters; older versions keep the flat attribute. + rope_params = getattr(config, "rope_parameters", None) or {} + self._rope_theta = float( + rope_params.get("rope_theta", None) or getattr(config, "rope_theta", 1000000.0) + ) + max_pos = int(getattr(config, "max_position_embeddings", 40960)) + self._freqs_cap = max_pos + self.block_size + 2 + self._build_rope_tables() + + # Ring-window length for the worker-owned context K/V buffer. + window = int(os.environ.get(_CTX_WINDOW_ENV, _DEFAULT_CTX_WINDOW)) + window = max(self.block_size + 2, min(window, max_pos)) + # Worker allocation contract (DSparkWorker._lazy_init): + # kv_windows = [max_batch, num_stages, window_size, head_dim], where + # "head_dim" is the per-position row width — here K||V flattened. + self._attn_params = dict( + window_size=window, + head_dim=2 * self.kv_dim, + ) + + self.markov_rank = int(getattr(config, "markov_rank", 0)) + + # Weight modules are built normally: the model loader wraps construction + # in ``MetaInitMode`` and materializes the parameters itself, and + # load_weights() then assigns the checkpoint tensors + # (torch load_state_dict(assign=True)). Forcing the meta device here + # instead would leave the drafter on meta whenever the loader falls back + # to regular init. + self.fc = nn.Linear( + self.num_capture_layers * self.hidden_size, self.hidden_size, bias=False + ) + self.hidden_norm = _Qwen3RMSNorm(self.hidden_size, eps) + self.layers = nn.ModuleList( + [ + _Qwen3DecoderLayer( + self.hidden_size, + int(config.intermediate_size), + self.num_heads, + self.num_kv_heads, + self.head_dim, + eps, + ) + for _ in range(num_layers) + ] + ) + self.norm = _Qwen3RMSNorm(self.hidden_size, eps) + self.markov_head = build_markov_head( + markov_head_type=str(getattr(config, "markov_head_type", "vanilla")), + vocab_size=int(config.vocab_size), + markov_rank=self.markov_rank, + hidden_size=self.hidden_size, + ) + self.confidence_head = None + if bool(getattr(config, "enable_confidence_head", False)): + # Inert scaffolding — the worker always passes + # ``confidence_threshold=0.0`` — kept for checkpoint + # completeness. The DeepSpec Qwen3 drafters store the + # confidence ``proj`` WITH a bias (the V4 one is bias-free). + self.confidence_head = DSparkConfidenceHead( + hidden_size=self.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=( + bool(getattr(config, "confidence_head_with_markov", False)) + and self.markov_rank > 0 + ), + bias=True, + ) + + # Shared with target; wired by the spec wrapper after construction. + self.embed_tokens: Optional[nn.Module] = None + self.lm_head: Optional[nn.Module] = None + + # ------------------------------------------------------------------ RoPE + + def _build_rope_tables(self) -> None: + """Precompute the cos/sin tables as non-persistent buffers. + + Built eagerly at construction rather than lazily on first use: the + drafter runs inside the target's CUDA graph, and an allocation made + during capture would live in the graph's private memory pool. Buffers + ride along with the module's ``.to(device)`` and stay out of + ``state_dict`` (``persistent=False``), so the strict checkpoint load is + unaffected. + """ + inv_freq = 1.0 / ( + self._rope_theta + ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim) + ) + t = torch.arange(self._freqs_cap, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_rope_cos", emb.cos(), persistent=False) + self.register_buffer("_rope_sin", emb.sin(), persistent=False) + + def _gather_cos_sin(self, positions: torch.Tensor, dtype: torch.dtype): + # Clamp for graph-safety: masked-out entries may carry arbitrary + # (already clamped by the worker) positions. + p = positions.long().clamp(min=0, max=self._freqs_cap - 1) + return self._rope_cos[p].to(dtype), self._rope_sin[p].to(dtype) + + # ---------------------------------------------------------- context K/V + + def _project_ctx(self, main_hidden: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(captured))`` — the layer-invariant context stream.""" + return self.hidden_norm(self.fc(main_hidden)) + + def _ctx_kv( + self, layer: _Qwen3DecoderLayer, main_x: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + """Per-layer context K/V rows ``[..., 2*kv_dim]`` (K RoPE'd/normed || V).""" + a = layer.self_attn + shp = main_x.shape[:-1] + k = a.k_proj(main_x).view(*shp, self.num_kv_heads, self.head_dim) + k = a.k_norm(k) + cos, sin = self._gather_cos_sin(positions, k.dtype) + k = _apply_rope(k, cos, sin) + v = a.v_proj(main_x) + return torch.cat([k.reshape(*shp, self.kv_dim), v], dim=-1) + + @torch.inference_mode() + def write_context_windows( + self, main_hidden: torch.Tensor, positions: torch.Tensor, stage_windows: torch.Tensor + ) -> None: + """Seed one request's ring windows from captured context (prefill path). + + Args: + main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. + positions: ``[M]`` window frames (= absolute position + 1, the + worker's generation-path convention). + stage_windows: ``[num_layers, window, 2*kv_dim]``, updated in place. + """ + M = int(main_hidden.shape[0]) + if M == 0: + return + win = int(self._attn_params["window_size"]) + p = positions.to(main_hidden.device).long() - 1 + cols = p % win + main_x = self._project_ctx(main_hidden) + for li, layer in enumerate(self.layers): + kv = self._ctx_kv(layer, main_x, p) + stage_windows[li, cols] = kv.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 masked back-fill of interim accepted tokens. + + Args: + main_hidden: ``[G, M, num_capture * hidden]``. + positions: ``[G, M]`` window frames (masked entries arbitrary). + slots: ``[G]`` request rows into ``kv_windows``. + mask: ``[G, M]`` bool validity. + kv_windows: ``[N, num_layers, window, 2*kv_dim]``, in place. + """ + G, M = positions.shape + if G == 0 or M == 0: + return + win = int(self._attn_params["window_size"]) + p = positions.long() - 1 + cols = p % win + rows = slots.long()[:, None].expand(-1, M) + mask3 = mask.unsqueeze(-1) + main_x = self._project_ctx(main_hidden) + for li, layer in enumerate(self.layers): + kv = self._ctx_kv(layer, main_x, p) # [G, M, 2*kv_dim] + win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view + cur = win_l[rows, cols] + win_l[rows, cols] = torch.where(mask3, kv.to(win_l.dtype), cur) + + # ------------------------------------------------------------- backbone + + 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 (all gen requests at once). + + Mirrors DeepSpec ``forward_dspark_draft_block`` + ``build_dspark_proposal``: + write the newest committed token's context K/V (position ``start_pos-1``), + run the ``block_size`` draft queries ``[bonus, mask, ...]`` at positions + ``start_pos + [0..block)`` through the Qwen3 layers attending to the ring + context + the block (bidirectional), then Markov-refine the block logits. + + Args: + main_hidden: ``[G, num_capture * hidden]`` captured hidden of the + newest committed context token per request. + bonus_token_ids: ``[G]`` last accepted token per request. + start_pos: ``[G]`` absolute decode position of the bonus token. + kv_windows: ``[N, num_layers, window, 2*kv_dim]`` persistent buffer. + slots: ``[G]`` request rows into ``kv_windows``. + Returns: + ``(draft_tokens [G, block], num_proposed [G])`` and, with + ``return_logits``, the corrected block logits ``[G, block, vocab]``. + """ + del all_rank_num_tokens # dense drafter: no cross-rank MoE lockstep + G = int(main_hidden.shape[0]) + B = self.block_size + win = int(self._attn_params["window_size"]) + device = main_hidden.device + start_pos = start_pos.long() + + main_x = self._project_ctx(main_hidden) # [G, hidden] + p0 = start_pos - 1 + cols0 = p0 % win + + draft_ids = build_draft_input_ids( + bonus_token_ids, block_size=B, noise_token_id=self.noise_token_id + ) + h = self.embed_tokens(draft_ids) # [G, B, hidden] + + pos_q = start_pos.unsqueeze(1) + torch.arange(B, device=device) # [G, B] + cos_q, sin_q = self._gather_cos_sin(pos_q, h.dtype) + + # Bool attention mask [G, 1, B, win + B]: the block attends to itself + # bidirectionally, and to the ring rows that actually hold a context + # row. Ring validity is NOT ``index < min(start_pos, win)``: the worker + # zeroes a slot's window when it assigns it and only seeds the positions + # the target actually processed, so under KV-cache prefix reuse (or any + # partially seeded slot) the rows below the request's first seeded + # position hold nothing, and attending to them would mix all-zero K/V + # into the softmax. Validity is read off the buffer instead — an + # unwritten row is exactly zero, and every position from the first + # seeded one on is written exactly once, so "non-zero" == "live". All + # layers are written together, so layer 0 answers for the whole stack. + blk_valid = torch.ones(G, 1, B, B, dtype=torch.bool, device=device) + attn_mask = None + + for li, layer in enumerate(self.layers): + a = layer.self_attn + win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view + # Write the newest committed token's context row, then read the ring. + kv0 = self._ctx_kv(layer, main_x, p0) # [G, 2*kv_dim] + win_l[slots, cols0] = kv0.to(win_l.dtype) + ctx = win_l[slots] # [G, win, 2*kv_dim] + if attn_mask is None: # layer 0, after its own context write + ctx_valid = (ctx != 0).any(dim=-1) # [G, win] + attn_mask = torch.cat( + [ctx_valid[:, None, None, :].expand(G, 1, B, win), blk_valid], dim=-1 + ) + k_ctx = ctx[..., : self.kv_dim].view(G, win, self.num_kv_heads, self.head_dim) + v_ctx = ctx[..., self.kv_dim :].view(G, win, self.num_kv_heads, self.head_dim) + + residual = h + x = layer.input_layernorm(h) + q = a.q_norm(a.q_proj(x).view(G, B, self.num_heads, self.head_dim)) + q = _apply_rope(q, cos_q, sin_q) + k_blk = a.k_norm(a.k_proj(x).view(G, B, self.num_kv_heads, self.head_dim)) + k_blk = _apply_rope(k_blk, cos_q, sin_q) + v_blk = a.v_proj(x).view(G, B, self.num_kv_heads, self.head_dim) + + k = torch.cat([k_ctx.to(h.dtype), k_blk], dim=1).transpose(1, 2) + v = torch.cat([v_ctx.to(h.dtype), v_blk], dim=1).transpose(1, 2) + k = k.repeat_interleave(self.num_kv_groups, dim=1) + v = v.repeat_interleave(self.num_kv_groups, dim=1) + o = F.scaled_dot_product_attention( + q.transpose(1, 2), k, v, attn_mask=attn_mask, scale=self.softmax_scale + ) + o = o.transpose(1, 2).reshape(G, B, self.num_heads * self.head_dim) + h = residual + a.o_proj(o) + + residual = h + h = residual + layer.mlp(layer.post_attention_layernorm(h)) + + h = self.norm(h) + base_logits = self.lm_head(h) + return dspark_propose( + base_logits, + bonus_token_ids=bonus_token_ids, + block_hidden=h, + markov_head=self.markov_head, + confidence_head=self.confidence_head, + block_size=B, + temperature=temperature, + confidence_threshold=confidence_threshold, + return_logits=return_logits, ) - # Generic handles expected by the loader / weight mappers. - self.model = self.dspark_model + + def forward( + self, + main_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + start_pos, + *, + kv_windows: Optional[torch.Tensor] = None, + **kwargs, + ) -> tuple: + """Eager/single-shot convenience wrapper over :meth:`forward_batched`. + + With no ``kv_windows`` the request drafts against an empty ring (only + the newest committed token's row is written, and the all-zero rows are + masked out); callers that want the prompt's context must seed a buffer + with :meth:`write_context_windows` and pass it, as ``DSparkWorker`` does. + """ + T = int(main_hidden.shape[0]) + device = main_hidden.device + if not torch.is_tensor(start_pos): + start_pos = torch.full((T,), int(start_pos), dtype=torch.long, device=device) + if kv_windows is None: + kv_windows = torch.zeros( + ( + T, + self.num_stages, + self._attn_params["window_size"], + self._attn_params["head_dim"], + ), + dtype=torch.bfloat16, + device=device, + ) + slots = torch.arange(T, device=device) + return self.forward_batched( + main_hidden, bonus_token_ids, start_pos, kv_windows=kv_windows, slots=slots, **kwargs + ) + + def run_moe_lockstep_noop(self, all_rank_num_tokens, device) -> None: + """Dense drafter: no cross-rank MoE barrier to keep in lockstep.""" + return None + + # --------------------------------------------------------------- loading + + def load_weights(self, weights: Dict) -> None: + # The loader has already placed the (empty) parameters on their target + # device; follow them rather than assuming the current cuda device. + device = self.fc.weight.device + if device.type == "meta": + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + sd = {} + for k, v in weights.items(): + # embed_tokens / lm_head are frozen copies of the target's; the + # shared target modules are used instead (load_weights_from_target_model). + if k.startswith(("embed_tokens.", "lm_head.")): + continue + t = v.to(device) + t = t.float() if k.startswith("confidence_head.") else t.to(torch.bfloat16) + sd[k] = t + self.load_state_dict(sd, strict=True, assign=True) + # assign=True rebinds the parameters only; move the RoPE buffers along. + self.to(device) + logger.info( + f"[DSpark-Qwen3] loaded {len(sd)} draft params " + f"({self.num_stages} layers, block_size={self.block_size}, " + f"ctx_window={self._attn_params['window_size']})" + ) + + +# --------------------------------------------------------------------------- # +# One-engine draft wrappers +# --------------------------------------------------------------------------- # + + +class DSparkForCausalLMBase(nn.Module): + """Shared one-engine draft wrapper for the DSpark drafters. + + Every DSpark drafter is created by ``get_draft_model``, appended to the + target's epilogue, and driven by ``DSparkWorker`` through the same + protocol: ``forward`` / ``forward_batched`` / ``write_context_windows`` / + ``write_context_windows_batched`` / ``run_moe_lockstep_noop`` plus the + ``num_stages`` / ``_attn_params`` / ``block_size`` worker-allocation + scalars. This base class pins that surface in one place; a drafter + subclass constructs its backbone (``dspark_model``) and implements + ``load_weights``. + + ``embed_tokens`` / ``lm_head`` are shared with the target model + (:meth:`load_weights_from_target_model`). + """ + + def __init__(self, dspark_model: nn.Module, draft_config): + super().__init__() + self.dspark_model = 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.num_stages = dspark_model.num_stages + self._attn_params = 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 model(self): + """Generic handle expected by the loader / weight mappers. + + A property, not a second attribute binding: registering the same + nn.Module twice would emit every draft tensor twice in ``state_dict``. + """ + return self.dspark_model + @property def block_size(self): return self.dspark_model.block_size @@ -1209,6 +1750,48 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_ main_hidden, positions, slots, mask, kv_windows ) + def load_weights_from_target_model(self, target_model): + """Share the target's embed_tokens / lm_head with the drafter.""" + 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 + + +class DSparkForCausalLM(DSparkForCausalLMBase): + """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 — see + :class:`DSparkForCausalLMBase` for the worker-facing surface. + + 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, + mask_token_id=None, + ): + super().__init__( + DSparkDraftModel( + draft_config, + aux_stream_dict, + num_stages=num_stages, + block_size=block_size, + mask_token_id=mask_token_id, + ), + draft_config, + ) + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. @@ -1227,19 +1810,43 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): 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 + +class Qwen3DSparkForCausalLM(DSparkForCausalLMBase): + """One-engine draft wrapper for the Qwen3 DSpark drafter. + + Wraps :class:`Qwen3DSparkDraftModel` — see :class:`DSparkForCausalLMBase` + for the worker-facing surface. ``embed_tokens`` / ``lm_head`` are shared + with the target model (the drafter checkpoint's copies are identical + frozen snapshots and are skipped at load). + """ + + def __init__( + self, + draft_config, + block_size: Optional[int] = None, + mask_token_id: Optional[int] = None, + ): + super().__init__( + Qwen3DSparkDraftModel(draft_config, block_size=block_size, mask_token_id=mask_token_id), + draft_config, + ) + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Load the flat-namespace Qwen3 DSpark drafter checkpoint. + + ``weight_mapper`` is accepted for loader-interface parity but unused — + the checkpoint names map 1:1 onto the module tree. + """ + self.dspark_model.load_weights(weights) __all__ = [ "DSparkBlock", "DSparkDraftModel", "DSparkForCausalLM", + "DSparkForCausalLMBase", + "Qwen3DSparkDraftModel", + "Qwen3DSparkForCausalLM", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", ] diff --git a/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py b/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py deleted file mode 100644 index cf0eb2a19c68..000000000000 --- a/tensorrt_llm/_torch/models/modeling_dspark_qwen3.py +++ /dev/null @@ -1,602 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3-based DSpark draft model (DeepSeek DeepSpec ``Qwen3DSparkModel``). - -Unlike the DeepSeek-V4 DSpark drafter (``modeling_dspark.py``) — whose draft -weights live in the target checkpoint's ``mtp.*`` namespace and whose stages -are full V4 blocks (MLA + MoE + mHC) — the DeepSpec-released Qwen3 drafters -(e.g. ``deepseek-ai/dspark_qwen3_8b_block7``) are separate dense bf16 -checkpoints with a flat namespace: - - - ``fc`` + ``hidden_norm``: project the concatenated captured target-layer - hidden states (``hidden_size * num_capture_layers``) into the draft hidden - size — the layer-invariant *context stream* (``main_x``). - - ``layers.{0..n-1}``: standard Qwen3 GQA decoder layers (q/k per-head - RMSNorm, RoPE, gated-SiLU MLP), except attention keys/values come from - BOTH the projected context (one row per committed token) and the draft - block itself, and the ``block_size`` draft queries attend bidirectionally - within the block. - - ``norm`` + (target-shared) ``lm_head`` -> :func:`dspark_propose` (Markov - head refinement; the confidence head is inert scaffolding, as in the V4 - path). - -The worker-facing protocol (``write_context_windows`` / -``write_context_windows_batched`` / ``forward_batched`` / ``_attn_params`` / -``num_stages`` / ``run_moe_lockstep_noop``) matches :class:`DSparkDraftModel`, -so :class:`DSparkWorker` and :class:`DSparkSpecMetadata` drive both drafters -unchanged. The worker-owned rolling buffer here holds per-layer context K/V -(``[max_batch, num_layers, window, 2 * num_kv_heads * head_dim]``, K then V, -K stored RoPE'd/k-normed). Frame convention: the worker passes window frames -``f = absolute_position + 1``; this model stores the context row of the token -at absolute position ``p`` at ring index ``p % window`` with RoPE phase ``p``, -matching the DeepSpec reference (its draft ``DynamicCache`` holds the context -K/V of positions ``0..start-1`` in order; the ring is an O(1)-memory window -over the most recent ``window`` positions — acceptance-rate only, standard -target verification keeps outputs correct regardless). - -Reference: https://github.com/deepseek-ai/DeepSpec -(``deepspec/modeling/dspark/qwen3/modeling.py`` and -``deepspec/eval/dspark/draft_ops.py``). -""" - -import os -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 .dspark.draft import build_draft_input_ids, dspark_propose -from .dspark.heads import build_markov_head - -# Ring-window length for the per-layer context K/V cache (per request slot). -# The DeepSpec reference attends over the FULL committed context; a ring -# window bounds the worker buffer at -# ``max_batch * num_layers * window * 2 * kv_dim`` bytes while remaining -# exactly full-context for sequences up to ``window`` tokens. -_DEFAULT_CTX_WINDOW = 2048 -_CTX_WINDOW_ENV = "TRTLLM_DSPARK_QWEN3_CTX_WINDOW" - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - - -def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - """HF-style rotary embedding on ``[..., num_heads, head_dim]``. - - ``cos``/``sin`` are ``[..., head_dim]`` (one row per position) and are - broadcast over the heads axis. - """ - cos = cos.unsqueeze(-2) - sin = sin.unsqueeze(-2) - return x * cos + _rotate_half(x) * sin - - -class _RMSNorm(nn.Module): - """Qwen3RMSNorm: fp32 normalize, cast back, then scale by the bf16 weight.""" - - def __init__(self, dim: int, eps: float): - super().__init__() - self.weight = nn.Parameter(torch.empty(dim)) - self.eps = float(eps) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - dt = x.dtype - xf = x.float() - xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) - return self.weight * xf.to(dt) - - -class _Attention(nn.Module): - """Weight container for one draft layer's GQA attention (names match ckpt).""" - - def __init__(self, hidden: int, n_heads: int, n_kv_heads: int, head_dim: int, eps: float): - super().__init__() - self.q_proj = nn.Linear(hidden, n_heads * head_dim, bias=False) - self.k_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) - self.v_proj = nn.Linear(hidden, n_kv_heads * head_dim, bias=False) - self.o_proj = nn.Linear(n_heads * head_dim, hidden, bias=False) - self.q_norm = _RMSNorm(head_dim, eps) - self.k_norm = _RMSNorm(head_dim, eps) - - -class _MLP(nn.Module): - def __init__(self, hidden: int, intermediate: int): - super().__init__() - self.gate_proj = nn.Linear(hidden, intermediate, bias=False) - self.up_proj = nn.Linear(hidden, intermediate, bias=False) - self.down_proj = nn.Linear(intermediate, hidden, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) - - -class _DecoderLayer(nn.Module): - def __init__( - self, - hidden: int, - intermediate: int, - n_heads: int, - n_kv_heads: int, - head_dim: int, - eps: float, - ): - super().__init__() - self.self_attn = _Attention(hidden, n_heads, n_kv_heads, head_dim, eps) - self.mlp = _MLP(hidden, intermediate) - self.input_layernorm = _RMSNorm(hidden, eps) - self.post_attention_layernorm = _RMSNorm(hidden, eps) - - -class _ConfidenceHead(nn.Module): - """DeepSpec ``AcceptRatePredictor`` (fp32 linear WITH bias for the Qwen3 - drafters). Inert scaffolding — the worker always passes - ``confidence_threshold=0.0`` — kept for checkpoint completeness.""" - - def __init__(self, input_dim: int, with_markov: bool): - super().__init__() - self.with_markov = bool(with_markov) - self.proj = nn.Linear(int(input_dim), 1, bias=True, 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 - return self.proj(features.float()).squeeze(-1) - - -class Qwen3DSparkDraftModel(nn.Module): - """The dense Qwen3 DSpark draft backbone + heads (DeepSpec reference math). - - Shares ``embed_tokens`` / ``lm_head`` with the target model (the checkpoint - carries frozen copies of both; they are identical to the target's, so the - shared modules are used and the copies skipped at load). - """ - - def __init__(self, model_config, block_size: Optional[int] = None): - super().__init__() - config = model_config.pretrained_config - self.model_config = model_config - self.config = config - - self.hidden_size = int(config.hidden_size) - self.num_heads = int(config.num_attention_heads) - self.num_kv_heads = int(config.num_key_value_heads) - self.head_dim = int(getattr(config, "head_dim", None) or self.hidden_size // self.num_heads) - self.kv_dim = self.num_kv_heads * self.head_dim - self.num_kv_groups = self.num_heads // self.num_kv_heads - self.softmax_scale = self.head_dim**-0.5 - eps = float(config.rms_norm_eps) - num_layers = int(config.num_hidden_layers) - # Worker-facing stage count (== draft layer count for this drafter). - self.num_stages = num_layers - - spec_cfg = getattr(model_config, "spec_config", None) - self.block_size = int( - block_size if block_size is not None else getattr(config, "block_size", 7) - ) - 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, "mask_token_id", config.vocab_size) - ) - - target_layer_ids = list(getattr(config, "target_layer_ids", []) or []) - self.num_capture_layers = len(target_layer_ids) - assert self.num_capture_layers > 0, ( - "Qwen3 DSpark drafter config must provide target_layer_ids" - ) - - # Plain-RoPE parameters. transformers>=5 nests rope_theta under - # rope_parameters; older versions keep the flat attribute. - rope_params = getattr(config, "rope_parameters", None) or {} - self._rope_theta = float( - rope_params.get("rope_theta", None) or getattr(config, "rope_theta", 1000000.0) - ) - max_pos = int(getattr(config, "max_position_embeddings", 40960)) - self._freqs_cap = max_pos + self.block_size + 2 - self._rope_cache: Dict[str, tuple] = {} - - # Ring-window length for the worker-owned context K/V buffer. - window = int(os.environ.get(_CTX_WINDOW_ENV, _DEFAULT_CTX_WINDOW)) - window = max(self.block_size + 2, min(window, max_pos)) - # Worker allocation contract (DSparkWorker._lazy_init): - # kv_windows = [max_batch, num_stages, window_size, head_dim], where - # "head_dim" is the per-position row width — here K||V flattened. - self._attn_params = dict( - window_size=window, - head_dim=2 * self.kv_dim, - ) - - # Build the weight modules on the meta device; load_weights() assigns - # the real cuda tensors (torch load_state_dict(assign=True)). - with torch.device("meta"): - self.fc = nn.Linear( - self.num_capture_layers * self.hidden_size, self.hidden_size, bias=False - ) - self.hidden_norm = _RMSNorm(self.hidden_size, eps) - self.layers = nn.ModuleList( - [ - _DecoderLayer( - self.hidden_size, - int(config.intermediate_size), - self.num_heads, - self.num_kv_heads, - self.head_dim, - eps, - ) - for _ in range(num_layers) - ] - ) - self.norm = _RMSNorm(self.hidden_size, eps) - self.markov_head = build_markov_head( - markov_head_type=str(getattr(config, "markov_head_type", "vanilla")), - vocab_size=int(config.vocab_size), - markov_rank=int(getattr(config, "markov_rank", 0)), - hidden_size=self.hidden_size, - ) - self.confidence_head = None - if bool(getattr(config, "enable_confidence_head", False)): - with_markov = bool(getattr(config, "confidence_head_with_markov", False)) - input_dim = self.hidden_size + (int(config.markov_rank) if with_markov else 0) - self.confidence_head = _ConfidenceHead(input_dim, with_markov) - - # Shared with target; wired by the spec wrapper after construction. - self.embed_tokens: Optional[nn.Module] = None - self.lm_head: Optional[nn.Module] = None - - # ------------------------------------------------------------------ RoPE - - def _cos_sin_tables(self, device: torch.device): - key = str(device) - cached = self._rope_cache.get(key) - if cached is None: - inv_freq = 1.0 / ( - self._rope_theta - ** ( - torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) - / self.head_dim - ) - ) - t = torch.arange(self._freqs_cap, device=device, dtype=torch.float32) - freqs = torch.outer(t, inv_freq) - emb = torch.cat((freqs, freqs), dim=-1) - cached = (emb.cos(), emb.sin()) - self._rope_cache[key] = cached - return cached - - def _gather_cos_sin(self, positions: torch.Tensor, dtype: torch.dtype): - cos_t, sin_t = self._cos_sin_tables(positions.device) - # Clamp for graph-safety: masked-out entries may carry arbitrary - # (already clamped by the worker) positions. - p = positions.long().clamp(min=0, max=self._freqs_cap - 1) - return cos_t[p].to(dtype), sin_t[p].to(dtype) - - # ---------------------------------------------------------- context K/V - - def _project_ctx(self, main_hidden: torch.Tensor) -> torch.Tensor: - """``hidden_norm(fc(captured))`` — the layer-invariant context stream.""" - return self.hidden_norm(self.fc(main_hidden)) - - def _ctx_kv( - self, layer: _DecoderLayer, main_x: torch.Tensor, positions: torch.Tensor - ) -> torch.Tensor: - """Per-layer context K/V rows ``[..., 2*kv_dim]`` (K RoPE'd/normed || V).""" - a = layer.self_attn - shp = main_x.shape[:-1] - k = a.k_proj(main_x).view(*shp, self.num_kv_heads, self.head_dim) - k = a.k_norm(k) - cos, sin = self._gather_cos_sin(positions, k.dtype) - k = _apply_rope(k, cos, sin) - v = a.v_proj(main_x) - return torch.cat([k.reshape(*shp, self.kv_dim), v], dim=-1) - - @torch.inference_mode() - def write_context_windows( - self, main_hidden: torch.Tensor, positions: torch.Tensor, stage_windows: torch.Tensor - ) -> None: - """Seed one request's ring windows from captured context (prefill path). - - Args: - main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. - positions: ``[M]`` window frames (= absolute position + 1, the - worker's generation-path convention). - stage_windows: ``[num_layers, window, 2*kv_dim]``, updated in place. - """ - M = int(main_hidden.shape[0]) - if M == 0: - return - win = int(self._attn_params["window_size"]) - p = positions.to(main_hidden.device).long() - 1 - cols = p % win - main_x = self._project_ctx(main_hidden) - for li, layer in enumerate(self.layers): - kv = self._ctx_kv(layer, main_x, p) - stage_windows[li, cols] = kv.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 masked back-fill of interim accepted tokens. - - Args: - main_hidden: ``[G, M, num_capture * hidden]``. - positions: ``[G, M]`` window frames (masked entries arbitrary). - slots: ``[G]`` request rows into ``kv_windows``. - mask: ``[G, M]`` bool validity. - kv_windows: ``[N, num_layers, window, 2*kv_dim]``, in place. - """ - G, M = positions.shape - if G == 0 or M == 0: - return - win = int(self._attn_params["window_size"]) - p = positions.long() - 1 - cols = p % win - rows = slots.long()[:, None].expand(-1, M) - mask3 = mask.unsqueeze(-1) - main_x = self._project_ctx(main_hidden) - for li, layer in enumerate(self.layers): - kv = self._ctx_kv(layer, main_x, p) # [G, M, 2*kv_dim] - win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view - cur = win_l[rows, cols] - win_l[rows, cols] = torch.where(mask3, kv.to(win_l.dtype), cur) - - # ------------------------------------------------------------- backbone - - 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 (all gen requests at once). - - Mirrors DeepSpec ``forward_dspark_draft_block`` + ``build_dspark_proposal``: - write the newest committed token's context K/V (position ``start_pos-1``), - run the ``block_size`` draft queries ``[bonus, mask, ...]`` at positions - ``start_pos + [0..block)`` through the Qwen3 layers attending to the ring - context + the block (bidirectional), then Markov-refine the block logits. - - Args: - main_hidden: ``[G, num_capture * hidden]`` captured hidden of the - newest committed context token per request. - bonus_token_ids: ``[G]`` last accepted token per request. - start_pos: ``[G]`` absolute decode position of the bonus token. - kv_windows: ``[N, num_layers, window, 2*kv_dim]`` persistent buffer. - slots: ``[G]`` request rows into ``kv_windows``. - Returns: - ``(draft_tokens [G, block], num_proposed [G])`` and, with - ``return_logits``, the corrected block logits ``[G, block, vocab]``. - """ - del all_rank_num_tokens # dense drafter: no cross-rank MoE lockstep - G = int(main_hidden.shape[0]) - B = self.block_size - win = int(self._attn_params["window_size"]) - device = main_hidden.device - start_pos = start_pos.long() - - main_x = self._project_ctx(main_hidden) # [G, hidden] - p0 = start_pos - 1 - cols0 = p0 % win - - draft_ids = build_draft_input_ids( - bonus_token_ids, block_size=B, noise_token_id=self.noise_token_id - ) - h = self.embed_tokens(draft_ids) # [G, B, hidden] - - pos_q = start_pos.unsqueeze(1) + torch.arange(B, device=device) # [G, B] - cos_q, sin_q = self._gather_cos_sin(pos_q, h.dtype) - - # Bool attention mask [G, 1, B, win + B]: ring rows valid iff their - # index < min(start_pos, win) (ring holds the last `win` positions); - # the block attends to itself bidirectionally. - ctx_valid = ( - torch.arange(win, device=device)[None, :] < start_pos.clamp(max=win)[:, None] - ) # [G, win] - attn_mask = torch.cat( - [ - ctx_valid[:, None, None, :].expand(G, 1, B, win), - torch.ones(G, 1, B, B, dtype=torch.bool, device=device), - ], - dim=-1, - ) - - for li, layer in enumerate(self.layers): - a = layer.self_attn - win_l = kv_windows[:, li] # [N, win, 2*kv_dim] view - # Write the newest committed token's context row, then read the ring. - kv0 = self._ctx_kv(layer, main_x, p0) # [G, 2*kv_dim] - win_l[slots, cols0] = kv0.to(win_l.dtype) - ctx = win_l[slots] # [G, win, 2*kv_dim] - k_ctx = ctx[..., : self.kv_dim].view(G, win, self.num_kv_heads, self.head_dim) - v_ctx = ctx[..., self.kv_dim :].view(G, win, self.num_kv_heads, self.head_dim) - - residual = h - x = layer.input_layernorm(h) - q = a.q_norm(a.q_proj(x).view(G, B, self.num_heads, self.head_dim)) - q = _apply_rope(q, cos_q, sin_q) - k_blk = a.k_norm(a.k_proj(x).view(G, B, self.num_kv_heads, self.head_dim)) - k_blk = _apply_rope(k_blk, cos_q, sin_q) - v_blk = a.v_proj(x).view(G, B, self.num_kv_heads, self.head_dim) - - k = torch.cat([k_ctx.to(h.dtype), k_blk], dim=1).transpose(1, 2) - v = torch.cat([v_ctx.to(h.dtype), v_blk], dim=1).transpose(1, 2) - k = k.repeat_interleave(self.num_kv_groups, dim=1) - v = v.repeat_interleave(self.num_kv_groups, dim=1) - o = F.scaled_dot_product_attention( - q.transpose(1, 2), k, v, attn_mask=attn_mask, scale=self.softmax_scale - ) - o = o.transpose(1, 2).reshape(G, B, self.num_heads * self.head_dim) - h = residual + a.o_proj(o) - - residual = h - h = residual + layer.mlp(layer.post_attention_layernorm(h)) - - h = self.norm(h) - base_logits = self.lm_head(h) - return dspark_propose( - base_logits, - bonus_token_ids=bonus_token_ids, - block_hidden=h, - markov_head=self.markov_head, - confidence_head=self.confidence_head, - block_size=B, - temperature=temperature, - confidence_threshold=confidence_threshold, - return_logits=return_logits, - ) - - def forward( - self, - main_hidden: torch.Tensor, - bonus_token_ids: torch.Tensor, - start_pos, - *, - kv_windows: Optional[torch.Tensor] = None, - **kwargs, - ) -> tuple: - """Eager/single-shot convenience wrapper over :meth:`forward_batched`.""" - T = int(main_hidden.shape[0]) - device = main_hidden.device - if not torch.is_tensor(start_pos): - start_pos = torch.full((T,), int(start_pos), dtype=torch.long, device=device) - if kv_windows is None: - kv_windows = torch.zeros( - ( - T, - self.num_stages, - self._attn_params["window_size"], - self._attn_params["head_dim"], - ), - dtype=torch.bfloat16, - device=device, - ) - slots = torch.arange(T, device=device) - return self.forward_batched( - main_hidden, bonus_token_ids, start_pos, kv_windows=kv_windows, slots=slots, **kwargs - ) - - def run_moe_lockstep_noop(self, all_rank_num_tokens, device) -> None: - """Dense drafter: no cross-rank MoE barrier to keep in lockstep.""" - return None - - # --------------------------------------------------------------- loading - - def load_weights(self, weights: Dict) -> None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - sd = {} - for k, v in weights.items(): - # embed_tokens / lm_head are frozen copies of the target's; the - # shared target modules are used instead (load_weights_from_target_model). - if k.startswith(("embed_tokens.", "lm_head.")): - continue - t = v.to(device) - t = t.float() if k.startswith("confidence_head.") else t.to(torch.bfloat16) - sd[k] = t - self.load_state_dict(sd, strict=True, assign=True) - logger.info( - f"[DSpark-Qwen3] loaded {len(sd)} draft params " - f"({self.num_stages} layers, block_size={self.block_size}, " - f"ctx_window={self._attn_params['window_size']})" - ) - - -class Qwen3DSparkForCausalLM(nn.Module): - """One-engine draft wrapper for the Qwen3 DSpark drafter. - - Mirrors :class:`DSparkForCausalLM`: created by ``get_draft_model``, - appended to the target's epilogue, and driven by ``DSparkWorker`` through - the shared draft-model protocol. ``embed_tokens`` / ``lm_head`` are shared - with the target model (the drafter checkpoint's copies are identical - frozen snapshots and are skipped at load). - """ - - def __init__(self, draft_config, block_size: Optional[int] = None): - super().__init__() - self.dspark_model = Qwen3DSparkDraftModel(draft_config, 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. - self.num_stages = self.dspark_model.num_stages - self._attn_params = self.dspark_model._attn_params - self.lm_head = None # shared from the target - 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): - 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): - return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) - - def write_context_windows(self, main_hidden, positions, stage_windows): - 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): - 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 flat-namespace Qwen3 DSpark drafter checkpoint. - - ``weight_mapper`` is accepted for loader-interface parity but unused — - the checkpoint names map 1:1 onto the module tree. - """ - self.dspark_model.load_weights(weights) - - def load_weights_from_target_model(self, target_model): - """Share the target's embed_tokens / lm_head (identical frozen copies).""" - 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__ = ["Qwen3DSparkDraftModel", "Qwen3DSparkForCausalLM"] diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index b1d4a86b1b28..e3713395508b 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1911,21 +1911,28 @@ def get_draft_model(model_config, draft_config, lm_head, model): 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). + from .modeling_dspark import (DSparkForCausalLM, Qwen3DSparkForCausalLM, + count_dspark_stages, + validate_dspark_eplb_layer_base) + # DeepSpec-released dense drafters (e.g. dspark_qwen3_8b_block7) are # separate checkpoints that declare their own architecture; the # DeepSeek-V4 drafter lives in the target checkpoint's mtp.* namespace. draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] + # The drafter's own ModelConfig carries spec_config=None (no recursive + # spec-dec), so the validated speculative-config values are passed in + # explicitly here. if any("Qwen3DSpark" in arch for arch in draft_arches): - from .modeling_dspark_qwen3 import Qwen3DSparkForCausalLM return Qwen3DSparkForCausalLM( - draft_config, block_size=model_config.spec_config.block_size) - # Lazy import to avoid a cycle (modeling_dspark -> modeling_deepseekv4 -> - # modeling_speculative). The DSpark draft reuses the target's aux streams. - # The draft stage count (n_mtp_layers) is not in the HF config, so derive - # it from the checkpoint's mtp.* namespace. - from .modeling_dspark import (DSparkForCausalLM, count_dspark_stages, - validate_dspark_eplb_layer_base) + draft_config, + block_size=model_config.spec_config.block_size, + mask_token_id=model_config.spec_config.mask_token_id) + # 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. num_stages = count_dspark_stages( model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) @@ -1934,6 +1941,7 @@ def get_draft_model(model_config, draft_config, lm_head, model): getattr(model, "aux_stream_dict", None), num_stages=num_stages, block_size=model_config.spec_config.block_size, + mask_token_id=model_config.spec_config.mask_token_id, ) elif spec_dec_mode.is_draft_target_one_model(): # Keep the draft LM head vocab-sharded so greedy draft sampling uses the diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py index 5f4bfe67c9f1..fa162e4a92e7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py @@ -12,7 +12,7 @@ # 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. -"""Golden tests for the Qwen3 DSpark drafter (modeling_dspark_qwen3.py). +"""Golden tests for the Qwen3 DSpark drafter (modeling_dspark.py). The reference implementation below is a line-for-line torch-only port of the DeepSpec ``Qwen3DSparkModel`` inference path @@ -35,7 +35,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.models.modeling_dspark_qwen3 import Qwen3DSparkDraftModel, _apply_rope +from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkDraftModel, _apply_rope VOCAB = 97 HID = 32 @@ -113,7 +113,10 @@ def r(*shape, scale=0.05): def _build_model(weights): - model = Qwen3DSparkDraftModel(_make_config()) + # Construct under a device context like the model loader does; load_weights + # then follows the modules (params and the RoPE buffers) onto that device. + with torch.device("cuda" if torch.cuda.is_available() else "cpu"): + model = Qwen3DSparkDraftModel(_make_config()) model.load_weights(weights) device = model.fc.weight.device embed = torch.nn.Embedding(VOCAB, HID) @@ -142,8 +145,11 @@ def __init__(self, w, device): t = torch.arange(MAX_POS, device=device, dtype=torch.float32) emb = torch.cat([torch.outer(t, inv)] * 2, dim=-1) self.cos, self.sin = emb.cos(), emb.sin() - # per-request context stream: [T, HID] projected hiddens, in order + # per-request context stream: [T, HID] projected hiddens, in order, + # starting at absolute position ``first_pos`` (non-zero when the target + # skipped a KV-cache-reused prefix, so no context exists below it). self.ctx_x = torch.zeros(0, HID, dtype=DTYPE, device=device) + self.first_pos = 0 def _norm(self, x, wname): wt = self.w[wname] @@ -159,13 +165,13 @@ def append_ctx(self, captured): self.ctx_x = torch.cat([self.ctx_x, x.to(DTYPE)], dim=0) def draft(self, bonus_id, start_pos): - assert self.ctx_x.shape[0] == start_pos + assert self.ctx_x.shape[0] == start_pos - self.first_pos w = self.w ids = torch.full((BLOCK,), MASK_ID, dtype=torch.long, device=self.device) ids[0] = bonus_id h = F.embedding(ids, w["embed_tokens.weight"]) # [B, HID] pos_q = start_pos + torch.arange(BLOCK, device=self.device) - pos_c = torch.arange(start_pos, device=self.device) + pos_c = torch.arange(self.first_pos, start_pos, device=self.device) for i in range(N_LAYERS): p = f"layers.{i}." x = self._norm(h, p + "input_layernorm.weight") @@ -284,6 +290,45 @@ def test_worker_protocol_golden(setup): bonus_val = (bonus_val * 7 + 3) % VOCAB +def test_prefix_reuse_masks_unseeded_rows(setup): + """KV-cache prefix reuse: only the uncached suffix reaches the drafter. + + The worker seeds the ring for the positions the target actually processed, + so the rows below the first seeded position hold nothing. The draft must + attend to the seeded rows only — treating the whole ``0..start_pos`` range + as live would mix all-zero K/V rows into the softmax. + """ + weights, model, device = setup + gen = torch.Generator().manual_seed(33) + win = model._attn_params["window_size"] + kv_windows = torch.zeros( + 1, model.num_stages, win, model._attn_params["head_dim"], dtype=DTYPE, device=device + ) + + reused, suffix = 5, 4 # positions 0..4 came from the KV cache; 5..8 ran + seeded = _rand_hidden(gen, suffix).to(device) + positions = torch.arange(reused, reused + suffix, device=device) + 1 + model.write_context_windows(seeded, positions, kv_windows[0]) + + ref = _Ref(weights, device) + ref.first_pos = reused + ref.append_ctx(seeded) + + start_pos = reused + suffix + bonus = torch.tensor([41], device=device) + got_toks, _, got_logits = model.forward_batched( + seeded[-1:].reshape(1, -1), + bonus, + torch.tensor([start_pos], device=device), + kv_windows=kv_windows, + slots=torch.tensor([0], device=device), + return_logits=True, + ) + exp_toks, exp_logits = ref.draft(bonus[0], start_pos) + torch.testing.assert_close(got_logits[0].float(), exp_logits.float(), atol=0.05, rtol=0.05) + assert torch.equal(got_toks[0].long().cpu(), exp_toks.long().cpu()) + + def test_batched_matches_eager_singletons(setup): """One batched call over G requests == G independent eager calls.""" weights, model, device = setup From f301332b4966a9a74d055255cbbf4183e645ee2b Mon Sep 17 00:00:00 2001 From: chungen04 Date: Thu, 6 Aug 2026 00:09:06 +0000 Subject: [PATCH 3/3] trim comments Signed-off-by: chungen04 --- tensorrt_llm/_torch/models/dspark/draft.py | 14 ++--- tensorrt_llm/_torch/models/modeling_dspark.py | 55 ++++--------------- .../_torch/models/modeling_speculative.py | 11 +--- 3 files changed, 20 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index aa76ffd5a725..0988e05e4a0e 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -39,15 +39,11 @@ def resolve_noise_token_id(mask_token_id: Optional[int], config, ckpt_attr: str) -> int: - """Resolve the DSpark noise/mask token id (shared by every drafter). - - ``mask_token_id`` is the speculative config's value — the user's override, - or the drafter checkpoint's id copied in by ``DSparkDecodingConfig`` - validation. ``None`` (direct construction, e.g. tests) falls back to the - drafter checkpoint's ``ckpt_attr`` (``dspark_noise_token_id`` for the V4 - ``mtp.*`` drafter, ``mask_token_id`` for the DeepSpec Qwen3 checkpoints), - else ``vocab_size``. A config that declares the key with a null value is - treated as absent. + """Resolve the DSpark noise/mask token id. + + ``mask_token_id`` is the speculative config's value, either indicated in ``DSparkDecodingConfig`` + validation. ``None`` falls back to the drafter checkpoint's ``ckpt_attr``, + else ``vocab_size``. """ if mask_token_id is None: mask_token_id = getattr(config, ckpt_attr, None) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index bfbd989ec48f..c86c2b494562 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -378,10 +378,8 @@ def __init__( or getattr(config, "n_mtp_layers", None) or config.num_nextn_predict_layers ) - # Production passes the validated speculative-config values explicitly - # (the drafter's own ModelConfig carries spec_config=None to avoid - # recursive spec-dec, so they cannot be read back off ``spec_cfg``); - # direct construction falls back to the checkpoint's trained values. + # 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) ) @@ -1150,14 +1148,9 @@ def forward_head( # --------------------------------------------------------------------------- # -# Qwen3 (DeepSpec) DSpark drafter +# Qwen3 DSpark drafter # --------------------------------------------------------------------------- # -# Ring-window length for the per-layer context K/V cache (per request slot). -# The DeepSpec reference attends over the FULL committed context; a ring -# window bounds the worker buffer at -# ``max_batch * num_layers * window * 2 * kv_dim`` bytes while remaining -# exactly full-context for sequences up to ``window`` tokens. _DEFAULT_CTX_WINDOW = 2048 _CTX_WINDOW_ENV = "TRTLLM_DSPARK_QWEN3_CTX_WINDOW" @@ -1301,10 +1294,6 @@ def __init__( # Worker-facing stage count (== draft layer count for this drafter). self.num_stages = num_layers - # Production passes the validated speculative-config values explicitly - # (the drafter's own ModelConfig carries spec_config=None to avoid - # recursive spec-dec, so they cannot be read back off it); direct - # construction falls back to the checkpoint's trained values. self.block_size = int( block_size if block_size is not None else getattr(config, "block_size", 7) ) @@ -1338,13 +1327,6 @@ def __init__( ) self.markov_rank = int(getattr(config, "markov_rank", 0)) - - # Weight modules are built normally: the model loader wraps construction - # in ``MetaInitMode`` and materializes the parameters itself, and - # load_weights() then assigns the checkpoint tensors - # (torch load_state_dict(assign=True)). Forcing the meta device here - # instead would leave the drafter on meta whenever the loader falls back - # to regular init. self.fc = nn.Linear( self.num_capture_layers * self.hidden_size, self.hidden_size, bias=False ) @@ -1372,9 +1354,7 @@ def __init__( self.confidence_head = None if bool(getattr(config, "enable_confidence_head", False)): # Inert scaffolding — the worker always passes - # ``confidence_threshold=0.0`` — kept for checkpoint - # completeness. The DeepSpec Qwen3 drafters store the - # confidence ``proj`` WITH a bias (the V4 one is bias-free). + # ``confidence_threshold=0.0`` — kept for checkpoint completeness. self.confidence_head = DSparkConfidenceHead( hidden_size=self.hidden_size, markov_rank=self.markov_rank, @@ -1398,12 +1378,7 @@ def __init__( def _build_rope_tables(self) -> None: """Precompute the cos/sin tables as non-persistent buffers. - Built eagerly at construction rather than lazily on first use: the - drafter runs inside the target's CUDA graph, and an allocation made - during capture would live in the graph's private memory pool. Buffers - ride along with the module's ``.to(device)`` and stay out of - ``state_dict`` (``persistent=False``), so the strict checkpoint load is - unaffected. + Built eagerly at construction rather than lazily on first use. """ inv_freq = 1.0 / ( self._rope_theta @@ -1512,11 +1487,7 @@ def forward_batched( ) -> tuple: """CUDA-graph-safe batched block draft (all gen requests at once). - Mirrors DeepSpec ``forward_dspark_draft_block`` + ``build_dspark_proposal``: - write the newest committed token's context K/V (position ``start_pos-1``), - run the ``block_size`` draft queries ``[bonus, mask, ...]`` at positions - ``start_pos + [0..block)`` through the Qwen3 layers attending to the ring - context + the block (bidirectional), then Markov-refine the block logits. + Mirrors DeepSpec ``forward_dspark_draft_block`` + ``build_dspark_proposal``. Args: main_hidden: ``[G, num_capture * hidden]`` captured hidden of the @@ -1550,14 +1521,11 @@ def forward_batched( # Bool attention mask [G, 1, B, win + B]: the block attends to itself # bidirectionally, and to the ring rows that actually hold a context - # row. Ring validity is NOT ``index < min(start_pos, win)``: the worker - # zeroes a slot's window when it assigns it and only seeds the positions + # row. The worker zeroes a slot's window when it assigns it and only seeds the positions # the target actually processed, so under KV-cache prefix reuse (or any # partially seeded slot) the rows below the request's first seeded # position hold nothing, and attending to them would mix all-zero K/V - # into the softmax. Validity is read off the buffer instead — an - # unwritten row is exactly zero, and every position from the first - # seeded one on is written exactly once, so "non-zero" == "live". All + # into the softmax. Validity is read off the buffer instead. All # layers are written together, so layer 0 answers for the whole stack. blk_valid = torch.ones(G, 1, B, B, dtype=torch.bool, device=device) attn_mask = None @@ -1626,12 +1594,15 @@ def forward( With no ``kv_windows`` the request drafts against an empty ring (only the newest committed token's row is written, and the all-zero rows are masked out); callers that want the prompt's context must seed a buffer - with :meth:`write_context_windows` and pass it, as ``DSparkWorker`` does. + with `write_context_windows` and pass it, as ``DSparkWorker`` does. """ T = int(main_hidden.shape[0]) device = main_hidden.device if not torch.is_tensor(start_pos): start_pos = torch.full((T,), int(start_pos), dtype=torch.long, device=device) + # Same precondition as DSparkDraftModel.forward + # Checked here (host sync) because this is the eager single-shot path. + assert bool((start_pos > 0).all()), "DSpark draft runs at generation (start_pos > 0)" if kv_windows is None: kv_windows = torch.zeros( ( @@ -1655,8 +1626,6 @@ def run_moe_lockstep_noop(self, all_rank_num_tokens, device) -> None: # --------------------------------------------------------------- loading def load_weights(self, weights: Dict) -> None: - # The loader has already placed the (empty) parameters on their target - # device; follow them rather than assuming the current cuda device. device = self.fc.weight.device if device.type == "meta": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index e3713395508b..9101149143b6 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1917,22 +1917,17 @@ def get_draft_model(model_config, draft_config, lm_head, model): count_dspark_stages, validate_dspark_eplb_layer_base) - # DeepSpec-released dense drafters (e.g. dspark_qwen3_8b_block7) are - # separate checkpoints that declare their own architecture; the + # Dense drafters (e.g. dspark_qwen3_8b_block7) are separate checkpoints. The # DeepSeek-V4 drafter lives in the target checkpoint's mtp.* namespace. draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] - # The drafter's own ModelConfig carries spec_config=None (no recursive - # spec-dec), so the validated speculative-config values are passed in - # explicitly here. + # The drafter's own ModelConfig carries spec_config=None, + # so the validated speculative-config values are passed in explicitly here. if any("Qwen3DSpark" in arch for arch in draft_arches): return Qwen3DSparkForCausalLM( draft_config, block_size=model_config.spec_config.block_size, mask_token_id=model_config.spec_config.mask_token_id) - # 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. num_stages = count_dspark_stages( model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config)