diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 1b47f4922965..0988e05e4a0e 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -38,6 +38,20 @@ 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. + + ``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) + 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..c86c2b494562 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 @@ -375,14 +383,7 @@ def __init__( 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 +1147,549 @@ def forward_head( ) -class DSparkForCausalLM(nn.Module): - """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). +# --------------------------------------------------------------------------- # +# Qwen3 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``. +_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 + + 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)) + 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. + 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. + """ + 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``. + + 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. 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. 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 `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( + ( + 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 = 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 +1719,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 +1779,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_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index a991e266965d..9101149143b6 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1912,11 +1912,22 @@ def get_draft_model(model_config, draft_config, lm_head, model): return DFlashForCausalLM(draft_config) elif spec_dec_mode.is_dspark(): # Lazy import to avoid a cycle (modeling_dspark -> modeling_deepseekv4 -> - # modeling_speculative). The DSpark draft reuses the target's aux streams. - # The draft stage count (n_mtp_layers) is not in the HF config, so derive - # it from the checkpoint's mtp.* namespace. - from .modeling_dspark import (DSparkForCausalLM, count_dspark_stages, + # modeling_speculative). + from .modeling_dspark import (DSparkForCausalLM, Qwen3DSparkForCausalLM, + count_dspark_stages, validate_dspark_eplb_layer_base) + + # 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, + # 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) num_stages = count_dspark_stages( model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) @@ -1925,6 +1936,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/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..fa162e4a92e7 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_qwen3.py @@ -0,0 +1,469 @@ +# 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.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 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): + # 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) + 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, + # 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] + 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 - 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(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") + 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_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 + 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())