From 3bb55e72a4b4853e3b2d0e5414ca74b9edcab984 Mon Sep 17 00:00:00 2001 From: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> Date: Fri, 9 Oct 2026 14:54:35 -0700 Subject: [PATCH] [TRTLLM-14620][feat] Integrate calibrated DSpark confidence scoring Signed-off-by: Zhaoyuan He <195028209+zhaoyuanh-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_dspark.py | 99 +++--- .../_torch/models/modeling_dspark_v41.py | 1 + .../_torch/models/modeling_speculative.py | 92 ++++-- tensorrt_llm/_torch/speculative/dspark.py | 254 ++++++++++++++- .../_torch/speculative/dspark_schedule.py | 8 + tensorrt_llm/llmapi/llm_args.py | 18 +- .../test_dspark_confidence_routing.py | 106 +++++++ .../hw_agnostic/test_dspark_draft.py | 31 +- .../hw_agnostic/test_dspark_heads.py | 63 +++- .../hw_agnostic/test_dspark_sts.py | 80 +++++ .../hw_agnostic/test_dspark_v41.py | 11 +- .../hw_agnostic/test_dspark_worker.py | 83 ++++- .../test_dspark_worker_confidence_state.py | 300 ++++++++++++++++++ 13 files changed, 1017 insertions(+), 129 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_routing.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 9ef7a91a1f71..030b34272450 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -123,7 +123,6 @@ class or the edge would become a cycle. DSparkConfidenceHead, SpecDecOneEngineForCausalLM, build_markov_head, - confident_prefix_length, dspark_markov_chain_logits, ) from .modeling_utils import DecoderModel, register_auto_model, register_draft_model @@ -768,29 +767,35 @@ def dspark_propose( confidence_head: Optional[nn.Module], block_size: int, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, + confidence_hidden: Optional[torch.Tensor] = None, ) -> tuple: """Produce DSpark draft tokens for one block (functional-first, static length). + Always proposes the full block: confidence scoring does not shorten the + parallel draft pass. The verification scheduler consumes the scores and + decides how many draft tokens are worth sending to the target. + Args: base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. bonus_token_ids: ``[batch]`` the token preceding the first draft position. - block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the - confidence head, and the RNN-head variant). + block_hidden: ``[batch, block_size, hidden]`` proposal-head hidden. + confidence_hidden: Optional separate pre-normalization confidence input. + Defaults to ``block_hidden`` for callers without a separate input. markov_head / confidence_head: the validated DSpark heads (may be None). + return_confidence: Compute fixed-shape per-position confidence logits. + This is a run-constant flag so CUDA-graph execution cannot diverge. Returns: draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers keep the tensor fixed-width for CUDA-graph safety). - num_proposed: ``[batch]`` int32 — how many leading tokens survive the - static confidence-threshold truncation (== block_size when no head / - threshold<=0). + confidence: ``[batch, block_size]`` fp32 raw confidence logits, or None + when scoring is disabled. Apply the confidence head's STS + calibration before computing prefix survival. """ - batch = base_logits.shape[0] # ``draft_logits`` are the per-position distributions the draft token is drawn # from (markov-corrected when a head is present, else the raw base logits). - # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance - # (1-TV) measurement; the normal path ignores them. + # Surfaced under ``return_logits`` for rejection sampling. draft_logits = base_logits if markov_head is not None: draft_tokens, corrected = markov_head.sample_block_tokens( @@ -805,36 +810,26 @@ def dspark_propose( draft_tokens = greedy_or_sample(base_logits, temperature) - # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. - # The worker always calls with confidence_threshold=0.0, so the block below is - # inert and num_proposed stays == block_size (the full block is proposed). The - # returned num_proposed is intentionally not yet consumed by the speculative - # scheduler/verifier; wiring it through is a follow-up (see PR description). - num_proposed = torch.full( - (batch,), int(block_size), dtype=torch.int32, device=base_logits.device - ) - if confidence_head is not None and confidence_threshold > 0.0: + # Branch-free, fixed-shape confidence scoring keeps the full draft path + # CUDA-graph capturable. + confidence = None + if return_confidence and confidence_head is not None: + scoring_hidden = block_hidden if confidence_hidden is None else confidence_hidden # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) prev_emb = ( markov_head.get_prev_embeddings(prev_ids) - if (markov_head is not None and getattr(confidence_head, "with_markov", False)) + if (markov_head is not None and confidence_head.with_markov) else None ) - conf_logits = ( - confidence_head(block_hidden, prev_embeddings=prev_emb) + confidence = ( + confidence_head(scoring_hidden, prev_embeddings=prev_emb) if prev_emb is not None - else confidence_head(block_hidden) + else confidence_head(scoring_hidden) ) - # Per-request prefix truncation (batch handled row-wise to stay simple; - # functional-first scope typically runs batch=1 for the draft). - for b in range(batch): - num_proposed[b] = confident_prefix_length( - conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold - ) if return_logits: - return draft_tokens, num_proposed, draft_logits - return draft_tokens, num_proposed + return draft_tokens, confidence, draft_logits + return draft_tokens, confidence # ---------------------------------------------------------------------------- @@ -1036,6 +1031,7 @@ def __init__( stage_id: int, num_stages: int, num_capture_layers: int, + block_size: int = 0, ): # The inherited attention uses a draft-local layer index, while the # decoder layer keeps its model-level index for weights and captures. @@ -1102,6 +1098,8 @@ def __init__( # markov_rank <= 0); otherwise dspark_propose passes no # prev_embeddings and DSparkConfidenceHead.forward would assert. with_markov=self.markov_rank > 0, + # Sizes the per-position STS calibration table once, up front. + block_size=int(block_size), ) @property @@ -1229,6 +1227,7 @@ def __init__( stage_id=s, num_stages=self.num_stages, num_capture_layers=self.num_capture_layers, + block_size=self.block_size, ) for s in range(self.num_stages) ] @@ -1802,7 +1801,7 @@ def forward( *, kv_windows: Optional[torch.Tensor] = None, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -1822,7 +1821,7 @@ def forward( worker; updated in place each call. ``None`` allocates fresh zero windows (single-shot golden / test path). Returns: - ``(draft_tokens [T, block], num_proposed [T])`` from ``forward_head``. + ``(draft_tokens [T, block], confidence [T, block] or None)``. """ assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: @@ -1857,7 +1856,7 @@ def forward( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, pre_mix=pre_mix, ) @@ -1872,7 +1871,7 @@ def forward_batched( slots: torch.Tensor, valid_len: Optional[torch.Tensor] = None, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -1884,9 +1883,9 @@ def forward_batched( graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per gen request); the rolling captured-context windows are written/read through ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are - gathered from a fixed table. ``forward_head`` is run with - ``confidence_threshold == 0`` (the worker proposes the full block), which is - the graph-safe branch of :func:`dspark_propose`. + gathered from a fixed table. ``forward_head`` always proposes the full + block; ``return_confidence`` only adds a fixed-shape scoring matmul, so + both settings stay capturable. Args: main_hidden: ``[G, num_capture * hidden]`` captured target context. @@ -1897,7 +1896,7 @@ def forward_batched( slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. valid_len: ``[G]`` count of actually written rolling-window entries. Returns: - ``(draft_tokens [G, block], num_proposed [G])`` from ``forward_head``. + ``(draft_tokens [G, block], confidence [G, block] or None)``. """ if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( @@ -1933,7 +1932,7 @@ def forward_batched( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, pre_mix=pre_mix, ) @@ -1995,15 +1994,15 @@ def forward_head( bonus_token_ids: torch.Tensor, *, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, pre_mix: Optional[torch.Tensor] = None, ) -> tuple: """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. - Returns (draft_tokens [*, block], num_proposed [*]); with ``return_logits`` - also returns the per-position draft logits [*, block, vocab] (§7.9 1-TV). + Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` + also returns the per-position draft logits [*, block, vocab]. """ last = self.mtp_layers[-1] if self.uses_lagged_pre: @@ -2012,6 +2011,7 @@ def forward_head( h = last.hc_ffn.collapse(block_hidden, pre_mix) else: h = last.hc_head(block_hidden) + confidence_hidden = h if return_confidence and last.confidence_head is not None else None normalized = last.norm(h) base_logits = self.lm_head(normalized) proposal_hidden = normalized if self.confidence_uses_normalized_hidden else h @@ -2023,8 +2023,9 @@ def forward_head( confidence_head=last.confidence_head, block_size=self.block_size, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, + confidence_hidden=confidence_hidden, ) @@ -2101,6 +2102,15 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_ main_hidden, positions, slots, mask, kv_windows ) + def _prepare_confidence_head_for_load(self, weights: Dict) -> None: + """Match the projection's optional bias to the checkpoint before loading.""" + head = self.dspark_model.mtp_layers[-1].confidence_head + bias_key = f"mtp_layers.{self.num_stages - 1}.confidence_head.proj.bias" + if head is not None and bias_key in weights and head.proj.bias is None: + # This runs before weight loading and graph capture. Empty storage + # avoids consuming RNG or reinitializing the existing projection. + head.proj.bias = nn.Parameter(head.proj.weight.new_empty(head.proj.out_features)) + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Load the ``mtp.*`` draft weights from the (full) checkpoint dict. @@ -2109,6 +2119,7 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): remap (``remap_dspark_draft_keys``) onto the shared V4 weight loader. """ remapped = remap_dspark_draft_keys(weights, num_stages=self.num_stages) + self._prepare_confidence_head_for_load(remapped) logger.info( f"[DSpark] loading {len(remapped)} draft params across {self.num_stages} stages" ) diff --git a/tensorrt_llm/_torch/models/modeling_dspark_v41.py b/tensorrt_llm/_torch/models/modeling_dspark_v41.py index 9d4943690dcf..32dd58fbf150 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark_v41.py +++ b/tensorrt_llm/_torch/models/modeling_dspark_v41.py @@ -298,6 +298,7 @@ def load_weights(self, weights: dict[str, torch.Tensor], weight_mapper=None, **k confidence_prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head." if not any(name.startswith(confidence_prefix) for name in remapped): last.confidence_head = None + self._prepare_confidence_head_for_load(remapped) DeepseekV41WeightLoader(self.dspark_model).load_weights(remapped) unexpected, unfed = _v41_load_coverage(self.dspark_model, remapped) if unexpected or unfed: diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index f42f9ec727f8..569ea56dced9 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -501,30 +501,37 @@ def build_markov_head(*, markov_head_type: str, vocab_size: int, class DSparkConfidenceHead(nn.Module): - """Per-position acceptance-confidence predictor (DeepSpec - AcceptRatePredictor). + """Per-position acceptance-confidence predictor. - Input features are the backbone hidden state, optionally concatenated with - the Markov head's previous-token embedding. Output is a single logit per - position. + Emits one raw logit per position. :meth:`apply_sts` converts logits to + calibrated conditional acceptance probabilities. """ def __init__(self, *, hidden_size: int, markov_rank: int = 0, - with_markov: bool = False): + with_markov: bool = False, + bias: bool = False, + block_size: int = 0): super().__init__() self.with_markov = bool(with_markov) input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) - # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the - # confidence score is computed in fp32 (mirrors the DeepSpec reference - # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). - self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) + self.proj = nn.Linear(input_dim, + 1, + bias=bool(bias), + dtype=torch.float32) + self.register_buffer( + "sts_temperatures", + torch.ones(max(int(block_size), 1), dtype=torch.float32), + persistent=False, + ) + self._sts_temperatures_host: Optional[torch.Tensor] = None def forward(self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor: + """Return raw, uncalibrated logits with shape ``[..., block]``.""" if self.with_markov: assert prev_embeddings is not None features = torch.cat( @@ -536,20 +543,57 @@ def forward(self, # fp32 matmul for a stable confidence score (mirrors the reference). return self.proj(features.float()).squeeze(-1) - -def confident_prefix_length(confidence_logits: torch.Tensor, *, block_size: int, - threshold: float) -> int: - """First position k where ``sigmoid(confidence_k) < threshold``. - - Returns ``block_size`` when threshold<=0 (no truncation) or all positions - are confident. Assumes batch size 1 (functional-first scope). - """ - if threshold <= 0.0: - return int(block_size) - below = confidence_logits.sigmoid() < threshold - if not bool(below[0].any().item()): - return int(block_size) - return int(torch.nonzero(below[0], as_tuple=False)[0].item()) + def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: + """Convert raw logits to calibrated acceptance probabilities.""" + temperatures = self.sts_temperatures + if temperatures.device != confidence_logits.device: + if confidence_logits.device.type == "cpu": + temperatures = self._host_sts_temperatures() + else: + temperatures = temperatures.to(confidence_logits.device) + return torch.sigmoid(confidence_logits.float() / temperatures) + + def _host_sts_temperatures(self) -> torch.Tensor: + """Return a lazily materialized CPU mirror of the STS table.""" + cached = self._sts_temperatures_host + if cached is None: + cached = self.sts_temperatures.detach().to("cpu") + self._sts_temperatures_host = cached + return cached + + @torch.no_grad() + def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: + """Update the fixed-address STS table in place.""" + flat = temperatures.reshape(-1).to(device=self.sts_temperatures.device, + dtype=torch.float32) + if flat.numel() != self.sts_temperatures.numel(): + raise ValueError( + f"STS temperature table has {flat.numel()} entries but the confidence " + f"head expects {self.sts_temperatures.numel()} (one per block position)" + ) + if not bool(torch.all(flat > 0.0)): + raise ValueError("STS temperatures must be strictly positive") + self.sts_temperatures.copy_(flat) + self._sts_temperatures_host = None + + def load_weights(self, weights: list) -> None: + """Load every confidence-head checkpoint parameter strictly.""" + (module_weights, ) = weights + expected = dict(self.named_parameters()) + unexpected = [key for key in module_weights if key not in expected] + missing = [key for key in expected if key not in module_weights] + if unexpected: + raise ValueError( + f"DSpark confidence head checkpoint has unsupported key(s) {sorted(unexpected)}; " + f"the module exposes {sorted(expected)}. A checkpoint bias requires constructing " + "DSparkConfidenceHead(bias=True); silently dropping it would bias every score." + ) + if missing: + raise ValueError( + f"DSpark confidence head checkpoint is missing key(s) {sorted(missing)}" + ) + for name, param in expected.items(): + param.data.copy_(module_weights[name][:]) class Eagle3Attention(Attention): diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index a183d72d6d05..f1273926d683 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -16,9 +16,10 @@ # DSpark worker / metadata mirror the DFlash plumbing (capture target-layer # hidden states, accept the previous block with standard verification, draft a # new block in one backbone forward), adapted to DSpark's draft model which -# produces the whole block (and its confidence-truncated length) inside a single +# produces the whole block and fixed-shape confidence scores inside a single # ``DSv4DSparkDraftModel.forward`` rather than via mask-token cross-attention. +import json from collections import deque from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional @@ -31,11 +32,45 @@ from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from .dflash import DFlashWorker, dflash_draft_slot_ids +from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig +# Unscored-slot fill: sigmoid saturates to 1.0, so an unwritten row schedules +# as "verify the whole block" (fail-safe). +_NEUTRAL_CONFIDENCE_LOGIT = NEUTRAL_CONFIDENCE_LOGIT + + +def load_sts_temperatures_from_path(path: str) -> list[float]: + """Read per-position calibration temperatures from an STS artifact.""" + with open(path, encoding="utf-8") as handle: + payload = json.load(handle) + if not isinstance(payload, dict): + raise ValueError(f"{path} must contain a JSON object") + for key in ("sts_temperatures", "temperatures"): + if key in payload: + temps = payload[key] + break + else: + raise KeyError( + f"{path} has neither 'sts_temperatures' nor 'temperatures'; " + f"found keys {sorted(payload)}" + ) + if not isinstance(temps, list) or not temps: + raise ValueError(f"{path} must contain a non-empty temperature list") + return [float(t) for t in temps] + + +def resolve_confidence_head(draft_model: torch.nn.Module) -> Optional[torch.nn.Module]: + """Return the last-stage confidence head on a bare or wrapped drafter.""" + inner = getattr(draft_model, "dspark_model", draft_model) + stages = getattr(inner, "mtp_layers", None) + if not stages: + return None + return getattr(stages[-1], "confidence_head", None) + def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: """Return the number of RoPE entries needed by the DSv4 block drafter. @@ -304,7 +339,7 @@ class DSv4DSparkWorker(SpecWorkerBase): states (``main_proj`` + ``main_norm``) into the draft's captured-context attention, runs the ``num_stages`` DSpark blocks over a rolling captured window, refines the per-position logits with the Markov head, and predicts a - per-position acceptance confidence used to truncate the proposed prefix. + per-position acceptance confidence used by the verification scheduler. Unlike DFlash, the draft does NOT use the paged KV cache or mask-token cross-attention: its attention K/V come from the worker-owned rolling window @@ -369,6 +404,21 @@ def __init__( # never handed out through ``_free_slots``. self._scratch_slot = 0 + # ``return_confidence`` is read once here (never per step) so the + # captured draft graph cannot diverge. ``_confidence_logits`` is the + # slot-indexed handoff to the verification scheduler. + self.return_confidence = bool(getattr(spec_config, "enable_confidence_scheduling", False)) + self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+2, block] + # Draft-pass stamps for the buffer above; allocated in ``_lazy_init``. + self._confidence_stamp: Optional[torch.Tensor] = None + self._draft_seq_host = 0 + self._draft_seq_cuda: Optional[torch.Tensor] = None + # Row index of the permanently-neutral confidence row; see ``_lazy_init``. + self._neutral_conf_row = 0 + # Host-side planner turning the confidence snapshot into this + # iteration's draft length; built in ``_lazy_init``. + self.verify_planner = None + # The generation draft path is the batched, host-sync-free # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward_batched`` + # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe @@ -381,6 +431,66 @@ def __init__( f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) + def staged_confidence_buffer(self) -> Optional[torch.Tensor]: + """The whole slot-indexed confidence buffer, or None when disabled. + + The whole buffer, not the current batch's rows: the planner reads it + back by slot across the one-iteration lag. + """ + return self._confidence_logits + + def confidence_stamp_buffer(self) -> Optional[torch.Tensor]: + """The slot-indexed draft-pass stamps, or None before ``_lazy_init``.""" + return self._confidence_stamp + + def bump_draft_seq(self) -> int: + """Advance the draft-pass sequence and return the PREVIOUS value. + + Must be called once per step from the executor's host path before the + step's forward is enqueued; the returned previous value is the + sequence of the last completed draft. + """ + prev = self._draft_seq_host + self._draft_seq_host += 1 + if self._draft_seq_cuda is not None: + self._draft_seq_cuda.fill_(self._draft_seq_host) + return prev + + def confidence_row_for(self, req_id: int) -> int: + """Buffer row holding ``req_id``'s confidence, host-side. + + Falls back to the permanently-neutral row ("verify the full block" -- + fail-safe). Keyed by slot, not batch position, so the one-iteration- + lagged snapshot survives batch reshuffles. + """ + if self._confidence_logits is None: + return self._neutral_conf_row + return self._req_to_slot.get(req_id, self._neutral_conf_row) + + def verified_draft_seq_cuda(self) -> Optional[torch.Tensor]: + """The verified draft sequence as a device scalar for the prologue. + + ``_draft_seq_cuda`` already holds the CURRENT step's sequence by + launch time (``bump_draft_seq`` fills it before the forward is + enqueued), so the pass that stamped the block being verified is one + less. Passing the buffer itself would make every fresh row read as + stale and silently degrade the whole batch to neutral. + """ + if self._draft_seq_cuda is None: + return None + return self._draft_seq_cuda - 1 + + def batch_slot_view(self, num_rows: int) -> Optional[torch.Tensor]: + """The batch-position -> confidence-row map staged this step. + + ``[num_rows]`` int64 on device; rows past the real batch point at the + scratch row (their survival is zeroed by the prologue's ``num_real`` + mask, so the content never matters). + """ + if self._batch_to_slot is None: + return None + return self._batch_to_slot[:num_rows] + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len @@ -470,6 +580,27 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._batch_to_slot = torch.full( (max_batch,), self._scratch_slot, dtype=torch.long, device="cuda" ) + if self.return_confidence: + # Neutral fill = large positive logit (sigmoid ~ 1.0): + # unscored slots fail safe by verifying the full block. The + # neutral row sits one past the scratch row and is never used + # as a draft-output destination. + self._neutral_conf_row = num_rows + self._confidence_logits = torch.full( + (num_rows + 1, block_size), + _NEUTRAL_CONFIDENCE_LOGIT, + dtype=torch.float32, + device="cuda", + ) + self._confidence_stamp = torch.zeros( + (num_rows + 1,), dtype=torch.int32, device="cuda" + ) + self._draft_seq_host = 0 + # The executor updates this scalar outside inference mode; + # inference tensors reject that in-place host-side update. + with torch.inference_mode(False): + self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, device="cuda") + self._build_verify_planner(draft_model, block_size) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} logger.info( @@ -508,6 +639,59 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: ) self._attention_warmup_attempted = True + def _build_verify_planner(self, draft_model, block_size: int) -> None: + """Construct the host-side verify planner (once, at lazy init). + + ModelEngine owns and authenticates the deployment cost table. The + worker creates only the policy state and calibration, then receives + that exact runtime object before its first decision. + """ + from .dspark_schedule import DSparkScheduleConfig + from .dspark_verify import DSparkVerifyPlanner + + cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) + + table_path = self.spec_config.confidence_sps_table_path + # ModelEngine owns the actual CUDA graph ladder, so it is the only + # component that may read, authenticate, and construct the deployment + # cost object. Reopening this path here creates a TOCTOU window in + # which graph capture and scheduling can observe different files. + logger.info( + f"DSpark: deferring exact SPS table {table_path} to the validated " + "ModelEngine runtime object" + ) + + # Calibration lives on the draft's confidence head; fall back to a plain + # sigmoid when the head is absent (no confidence to calibrate anyway). + head = resolve_confidence_head(draft_model) + apply_calibration = head.apply_sts if head is not None else None + sts_path = self.spec_config.confidence_sts_path + if sts_path and head is None: + # Refuse rather than degrade: a configured calibration that cannot + # be attached is a wiring bug, never a preference. + raise ValueError( + f"confidence_sts_path is set but no confidence head was found " + f"on {type(draft_model).__name__}; calibration cannot be " + f"applied, so the planner would schedule on raw sigmoid " + f"survivals. This is a model-wiring bug, not a config choice." + ) + if sts_path and head is not None: + # Accepts either key spelling ("sts_temperatures" here, + # "temperatures" in SGLang); the vectors are interchangeable. + temps = load_sts_temperatures_from_path(sts_path) + head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) + logger.info(f"DSpark: loaded STS calibration from {sts_path}") + + self.verify_planner = DSparkVerifyPlanner( + cfg=cfg, + apply_calibration=apply_calibration, + device_windows=bool(self.spec_config.enable_fused_confidence_scheduler), + ) + logger.info( + f"DSpark verify planner: max_verify_len={self.verify_planner.max_verify_len}, " + f"device_windows={self.verify_planner.device_windows}" + ) + def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: @@ -529,6 +713,11 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._valid_len[slot] = 0 self._position_initialized[slot] = False self._kv_windows[slot].zero_() + if self._confidence_logits is not None: + # A recycled slot still holds the previous occupant's scores; + # reset to neutral so the first scheduling decision verifies + # the full block. + self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] def _seed_context_windows( @@ -633,9 +822,16 @@ def _draft_gen_block_batched( ``DSv4DSparkDraftModel.forward_batched``. Returns the per-position corrected block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. - Confidence truncation stays disabled — the full block is proposed. + Confidence truncation stays disabled — the full block is proposed. When + enabled, fixed-shape confidence scores are stored separately for + verification scheduling; they never truncate the proposal. """ num_gens = batch_size - num_contexts + # K is the draft block width (fixed by the checkpoint); Kp1 is the + # target's per-gen-request token stride in ``captured`` + # (runtime_draft_len + 1, which shrinks when verification is trimmed). + # Striding by max_draft_len instead reads into the next request's + # hidden states. K = self.max_draft_len device = accepted_tokens.device @@ -665,9 +861,22 @@ def _draft_gen_block_batched( accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long() ) # [G] - arange_g = torch.arange(num_gens, device=device) - base = gen_start + arange_g * target_width # [G] - main_hidden = captured[base + gidx] # [G, ncap*hidden] + # Start offset of each request's processed tokens inside ``captured``. + # Ragged verification uses the same qo_indptr as the packed input; + # uniform and bootstrap iterations use the accepted-row width. + qo_indptr = spec_metadata.qo_indptr + if spec_metadata.verify_lens is not None and qo_indptr is not None: + base = gen_start + qo_indptr[:num_gens].to(device=device, dtype=torch.long) + else: + arange_g = torch.arange(num_gens, device=device) + base = gen_start + arange_g * target_width # [G] + # Graph-padding rows own only scratch state and may lie beyond the + # captured token extent. Route only those rows to a safe placeholder; + # an invalid real-request offset must not silently read another row. + is_scratch = slots == self._scratch_slot + base = torch.where(is_scratch, torch.zeros_like(base), base) + bonus_row = torch.where(is_scratch, torch.zeros_like(base), base + gidx) + main_hidden = captured.index_select(0, bonus_row) # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens # (everything but the bonus) into the rolling window — same frames as the @@ -676,7 +885,7 @@ def _draft_gen_block_batched( # slot has no absolute decode position. Bootstrap it once from the first # target input position; locally-prefilled and existing slots keep their # monotonically advanced position. - input_positions = position_ids.reshape(-1)[base].long() + input_positions = position_ids.reshape(-1).index_select(0, base).long() old, start_pos = self._advance_generation_state(slots, nacc, input_positions) j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -692,7 +901,7 @@ def _draft_gen_block_batched( # Surface the per-position corrected block logits ([num_gens, K, vocab]) # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, _num_proposed, block_logits = draft_model.forward_batched( + _toks, confidence, block_logits = draft_model.forward_batched( main_hidden, bonus, start_pos, @@ -700,10 +909,20 @@ def _draft_gen_block_batched( slots=slots, valid_len=self._valid_len[slots], temperature=0.0, - confidence_threshold=0.0, + return_confidence=self.return_confidence, return_logits=True, all_rank_num_tokens=all_rank_num_tokens, ) + # Stash the [G, K] confidence for the verification scheduler. + # Slot-indexed (survives batch reshuffles), written in place (a + # captured graph keeps the same storage), and scattered through + # ``slots`` so a replay lands on the current batch's slots. + if confidence is not None: + # Padding may write the same scratch slot repeatedly; its content + # is undefined and must never be consumed as a real confidence row. + self._confidence_logits[slots] = confidence.detach() + # Same slots, same capture region: content-only update, graph-safe. + self._confidence_stamp[slots] = self._draft_seq_cuda return block_logits def _sample_draft_tokens_guided( @@ -782,9 +1001,10 @@ def _forward_impl( # NOT yet capturing) run synthetic gen batches that would otherwise advance # the persistent rolling-window state. Snapshot and restore it so warmup is # side-effect-free. (During the capture pass itself the stream IS capturing, - # so we skip the save/restore and let the ops be recorded; real requests - # reset their slot's window+ctx_len at prefill, wiping any capture-time - # mutation.) + # so we skip the save/restore and let the ops be recorded. Capture-time + # mutation stays on dummy/scratch state: a locally-prefilled request + # resets its slot at position zero, while a disaggregated generation-only + # request receives a freshly zeroed slot.) is_warmup = ( getattr(spec_metadata, "is_cuda_graph", False) and not torch.cuda.is_current_stream_capturing() @@ -794,6 +1014,12 @@ def _forward_impl( saved_valid_len = self._valid_len.clone() saved_position_initialized = self._position_initialized.clone() saved_windows = self._kv_windows.clone() + # The confidence buffer is persistent state too; warmup scores + # must not reach the verification planner. + saved_confidence = ( + None if self._confidence_logits is None else self._confidence_logits.clone() + ) + saved_stamp = None if self._confidence_stamp is None else self._confidence_stamp.clone() # Assign / reset window slots for context (prefill) requests and seed each # request's rolling KV window from its prompt's captured context, so the @@ -908,6 +1134,10 @@ def _forward_impl( self._valid_len.copy_(saved_valid_len) self._position_initialized.copy_(saved_position_initialized) self._kv_windows.copy_(saved_windows) + if saved_confidence is not None: + self._confidence_logits.copy_(saved_confidence) + if saved_stamp is not None: + self._confidence_stamp.copy_(saved_stamp) outputs = { "logits": raw_logits, diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index 3a12d6d27bf1..e9bd692b6259 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -30,6 +30,8 @@ __all__ = [ "DSparkScheduleConfig", + "HOST_POLICY_WINDOWS_SNAPSHOT_OUTPUT", + "NATIVE_UNIFORM_VERIFY_OUTPUT", "compute_survival", "schedule_verify_lens_topk", ] @@ -41,6 +43,12 @@ # planner (counts it) cannot import each other. NEUTRAL_CONFIDENCE_LOGIT = 30.0 +# The DSpark forward publishes one of these markers when it cannot return a +# batch-aligned verify-length tensor. The sampler uses the marker to preserve +# the current iteration's policy across overlap with the next iteration. +HOST_POLICY_WINDOWS_SNAPSHOT_OUTPUT = "host_policy_windows_snapshot" +NATIVE_UNIFORM_VERIFY_OUTPUT = "native_uniform_verify" + @dataclass(frozen=True) class DSparkScheduleConfig: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1f4924b798c8..2626af1ebf23 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3469,14 +3469,16 @@ class DSparkDecodingConfig(DecodingBaseConfig): layers as cross-attention context and drafts a whole block in a single backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN - head) and predicts an acceptance-confidence per position to truncate the - proposed prefix. + head). A confidence head can score each draft position for verification + scheduling without truncating the proposed block. This configuration selects + native full-block drafting and does not expose confidence-scheduling or + calibration controls. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. - Semi-parallel: one block backbone forward + cheap sequential head refine. - - Confidence head: truncates the proposed draft length (NOT the accept rule; - acceptance stays standard target verification, preserving greedy parity). + - Full-block proposals: confidence scoring does not change the draft tokens + or the standard target-verification acceptance rule. Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). """ @@ -3512,14 +3514,6 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") - # NOTE: confidence-based dynamic drafting (the draft model's confidence head - # that truncates the proposed block) is NOT enabled in this PR. The user-facing - # ``enable_confidence_head`` / ``confidence_threshold`` knobs are intentionally - # omitted and will be added when the feature is actually wired into the - # speculative scheduling/verification path. The confidence head module and its - # internal plumbing remain as scaffolding (see DSparkConfidenceHead / - # dspark_propose). - decoding_type: Literal["DSpark"] = Field(default="DSpark") attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_routing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_routing.py new file mode 100644 index 000000000000..969430b71807 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_routing.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU numerical regressions for confidence input routing, not model accuracy.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_dspark import DSv4DSparkDraftModel, dspark_propose +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead, build_markov_head + + +@pytest.mark.parametrize("physical_k", range(1, 9)) +@pytest.mark.parametrize( + "head_type,with_markov", + [ + ("none", False), + ("vanilla", False), + ("vanilla", True), + ("gated", False), + ("gated", True), + ("rnn", False), + ("rnn", True), + ], +) +@torch.no_grad() +def test_confidence_uses_prenorm_without_changing_proposals(physical_k, head_type, with_markov): + torch.manual_seed(18000 + physical_k) + hidden_size, rank, vocab_size, batch = 16, 8, 31, 3 + markov = build_markov_head( + markov_head_type=head_type, + vocab_size=vocab_size, + markov_rank=0 if head_type == "none" else rank, + hidden_size=hidden_size, + ) + confidence = DSparkConfidenceHead( + hidden_size=hidden_size, + markov_rank=rank, + with_markov=with_markov, + block_size=physical_k, + ) + norm = torch.nn.RMSNorm(hidden_size, eps=1e-6) + norm.weight.copy_(torch.linspace(0.25, 2.5, hidden_size)) + last = SimpleNamespace( + hc_head=lambda h: h.mean(dim=-2), + norm=norm, + markov_head=markov, + confidence_head=confidence, + ) + model = SimpleNamespace( + mtp_layers=[last], + lm_head=torch.nn.Linear(hidden_size, vocab_size, bias=False), + block_size=physical_k, + uses_lagged_pre=False, + confidence_uses_normalized_hidden=True, + ) + hidden = torch.randn(batch, physical_k, 4, hidden_size) * 3 + bonus = torch.arange(batch) + 2 + pre_norm = last.hc_head(hidden) + post_norm = norm(pre_norm) + + for temperature in (0.0, 0.8): + outputs, states = [], [] + for enabled in (False, True): + torch.manual_seed(912) + outputs.append( + DSv4DSparkDraftModel.forward_head( + model, + hidden, + bonus, + return_confidence=enabled, + return_logits=True, + temperature=temperature, + ) + ) + states.append(torch.get_rng_state()) + off, on = outputs + assert off[1] is None + assert torch.equal(off[0], on[0]) + assert torch.equal(off[2], on[2]) + assert torch.equal(states[0], states[1]) + # The legacy functional caller keeps its proposal-head contract. + torch.manual_seed(912) + legacy = dspark_propose( + model.lm_head(post_norm), + bonus_token_ids=bonus, + block_hidden=post_norm, + markov_head=markov, + confidence_head=confidence, + block_size=physical_k, + temperature=temperature, + return_confidence=True, + return_logits=True, + ) + assert torch.equal(on[0], legacy[0]) + assert torch.equal(on[2], legacy[2]) + prev = torch.cat([bonus.unsqueeze(1), on[0][:, :-1]], dim=1) + embeddings = markov.get_prev_embeddings(prev) if with_markov else None + expected = confidence(pre_norm, prev_embeddings=embeddings) + torch.testing.assert_close(on[1], expected, rtol=0, atol=0) + torch.testing.assert_close( + legacy[1], confidence(post_norm, prev_embeddings=embeddings), rtol=0, atol=0 + ) + # This fixture must expose the previous post-norm scoring defect. + assert not torch.allclose(on[1], legacy[1], rtol=1e-6, atol=1e-6) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index 9d349ef72a2c..ee275aff985d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -40,7 +40,7 @@ def test_dspark_propose_full_block_no_confidence(): bonus = torch.randint(0, VOCAB, (B,)) hid = torch.randn(B, BLK, HID) with torch.no_grad(): - tokens, num = dspark_propose( + tokens, confidence = dspark_propose( base, bonus_token_ids=bonus, block_hidden=hid, @@ -49,8 +49,7 @@ def test_dspark_propose_full_block_no_confidence(): block_size=BLK, ) assert tokens.shape == (B, BLK) - # No confidence head -> propose the full block. - assert torch.all(num == BLK) + assert confidence is None # Tokens match the markov head's own greedy block sampling. ref_tokens, _ = markov.sample_block_tokens( base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 @@ -58,42 +57,42 @@ def test_dspark_propose_full_block_no_confidence(): assert torch.equal(tokens, ref_tokens) -def test_dspark_propose_confidence_truncates(): +def test_dspark_propose_scores_without_shortening_the_block(): torch.manual_seed(1) markov = build_markov_head( markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID ).eval() - conf = DSparkConfidenceHead(hidden_size=HID).eval() + conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() # The confidence proj is bias-free, so drive the logit via a constant weight # against a constant hidden: logit = weight_val * HID per position. base = torch.randn(1, BLK, VOCAB) bonus = torch.randint(0, VOCAB, (1,)) hid = torch.ones(1, BLK, HID) with torch.no_grad(): - conf.proj.weight.fill_(5.0 / HID) # logit ~ 5 -> sigmoid ~ 0.993, all confident + conf.proj.weight.fill_(-5.0 / HID) with torch.no_grad(): - _, num = dspark_propose( + tokens, confidence = dspark_propose( base, bonus_token_ids=bonus, block_hidden=hid, markov_head=markov, confidence_head=conf, block_size=BLK, - confidence_threshold=0.5, + return_confidence=True, ) - # All-confident -> full block proposed. - assert int(num[0]) == BLK - # Now make the head output low confidence everywhere -> truncate to 0... but - # confident_prefix_length returns first sub-threshold index (0 here). + assert tokens.shape == (1, BLK) + assert confidence.shape == (1, BLK) + assert torch.all(confidence < 0.0) + assert torch.all(conf.apply_sts(confidence) < 0.5) + with torch.no_grad(): - conf.proj.weight.fill_(-5.0 / HID) # logit ~ -5 -> sigmoid ~ 0.0067 < 0.5 - _, num2 = dspark_propose( + _, confidence = dspark_propose( base, bonus_token_ids=bonus, block_hidden=hid, markov_head=markov, confidence_head=conf, block_size=BLK, - confidence_threshold=0.5, + return_confidence=False, ) - assert int(num2[0]) == 0 + assert confidence is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py index 036152f00b08..f1caa75022b2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -22,7 +22,6 @@ RNNHead, VanillaMarkov, build_markov_head, - confident_prefix_length, dspark_markov_step_bias, markov_prev_embeddings, ) @@ -95,18 +94,56 @@ def test_build_markov_head_rank_zero_returns_none(): ) -def test_confidence_head_and_prefix_truncation(): - head = DSparkConfidenceHead(hidden_size=HID) - conf = head(torch.randn(B, BLK, HID)) - assert conf.shape == (B, BLK) - # threshold 0 disables truncation. - assert confident_prefix_length(conf, block_size=BLK, threshold=0.0) == BLK - # First sub-threshold position truncates the prefix. - logits = torch.tensor([[10.0, 10.0, -10.0, 10.0, 10.0]]) - assert confident_prefix_length(logits, block_size=BLK, threshold=0.5) == 2 - # All-confident -> full block. - logits_hi = torch.full((1, BLK), 10.0) - assert confident_prefix_length(logits_hi, block_size=BLK, threshold=0.5) == BLK +def test_confidence_head_emits_raw_logits_and_applies_sts(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + logits = head(torch.randn(B, BLK, HID)) + assert logits.shape == (B, BLK) + assert logits.dtype == torch.float32 + + temperatures = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) + head.load_sts_temperatures(temperatures) + raw = torch.full((1, BLK), 4.0) + assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw / temperatures)) + + +@pytest.mark.parametrize( + ("temperatures", "message"), + [ + (torch.ones(BLK + 1), "one per block position"), + (torch.zeros(BLK), "strictly positive"), + ], +) +def test_confidence_head_rejects_invalid_sts_tables(temperatures, message): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + with pytest.raises(ValueError, match=message): + head.load_sts_temperatures(temperatures) + + +def test_confidence_head_updates_sts_in_place(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + storage = head.sts_temperatures + + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + + assert head.sts_temperatures is storage + assert "sts_temperatures" not in head.state_dict() + + +def test_confidence_head_loads_checkpoint_weights_strictly(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + weight = torch.randn(1, HID) + head.load_weights([{"proj.weight": weight}]) + assert torch.equal(head.proj.weight, weight) + + with pytest.raises(ValueError, match="bias=True"): + head.load_weights([{"proj.weight": weight, "proj.bias": torch.zeros(1)}]) + with pytest.raises(ValueError, match="missing key"): + head.load_weights([{}]) + + biased = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) + bias = torch.randn(1) + biased.load_weights([{"proj.weight": weight, "proj.bias": bias}]) + assert torch.equal(biased.proj.bias, bias) def test_confidence_head_with_markov_concat_dim(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py new file mode 100644 index 000000000000..04ce5bc4a276 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""STS calibration loading and confidence-head resolution tests.""" + +import json + +import pytest + +from tensorrt_llm._torch.speculative.dspark import ( + load_sts_temperatures_from_path, + resolve_confidence_head, +) + + +@pytest.mark.parametrize("key", ["sts_temperatures", "temperatures"]) +def test_loads_native_and_sglang_temperature_keys(tmp_path, key): + path = tmp_path / "sts.json" + path.write_text(json.dumps({key: [0.75, 1.0, 1.25]}), encoding="utf-8") + + assert load_sts_temperatures_from_path(str(path)) == [0.75, 1.0, 1.25] + + +@pytest.mark.parametrize( + ("payload", "error"), + [({}, KeyError), ({"temperatures": []}, ValueError)], +) +def test_rejects_missing_or_empty_temperatures(tmp_path, payload, error): + path = tmp_path / "sts.json" + path.write_text(json.dumps(payload), encoding="utf-8") + + with pytest.raises(error): + load_sts_temperatures_from_path(str(path)) + + +@pytest.mark.parametrize("payload", [None, [], "temperatures", 1]) +def test_rejects_non_object_payload(tmp_path, payload): + path = tmp_path / "sts.json" + path.write_text(json.dumps(payload), encoding="utf-8") + + with pytest.raises(ValueError, match="must contain a JSON object"): + load_sts_temperatures_from_path(str(path)) + + +@pytest.mark.parametrize("key", ["sts_temperatures", "temperatures"]) +@pytest.mark.parametrize("temperatures", ["123", {"1": 0.75, "2": 1.0, "3": 1.25}, 1.0, None]) +def test_rejects_non_list_temperatures(tmp_path, key, temperatures): + path = tmp_path / "sts.json" + path.write_text(json.dumps({key: temperatures}), encoding="utf-8") + + with pytest.raises(ValueError, match="must contain a non-empty temperature list"): + load_sts_temperatures_from_path(str(path)) + + +class _Head: + pass + + +class _Bare: + def __init__(self, head): + stage = type("Stage", (), {})() + if head is not None: + stage.confidence_head = head + self.mtp_layers = [object(), object(), stage] + + +class _Wrapper: + def __init__(self, head): + self.dspark_model = _Bare(head) + + +@pytest.mark.parametrize("wrap", [lambda head: _Bare(head), lambda head: _Wrapper(head)]) +def test_resolves_known_model_layouts(wrap): + head = _Head() + + assert resolve_confidence_head(wrap(head)) is head + assert resolve_confidence_head(wrap(None)) is None + + +def test_unknown_model_layout_has_no_confidence_head(): + assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index 6832cf28304d..8fd59c23d4f6 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -136,7 +136,7 @@ def test_v41_remap_rejects_orphan_draft_scale(): @pytest.mark.parametrize("v41", [False, True]) -def test_dspark_confidence_and_lm_head_use_variant_correct_hidden_states(v41): +def test_dspark_confidence_uses_prenorm_and_lm_head_uses_normalized_hidden_states(v41): cls = DSv41DSparkDraftModel if v41 else modeling_dspark.DSv4DSparkDraftModel model = cls.__new__(cls) nn.Module.__init__(model) @@ -154,16 +154,17 @@ def test_dspark_confidence_and_lm_head_use_variant_correct_hidden_states(v41): confidence_head=confidence, ) ] - _, num_proposed, logits = model.forward_head( + draft_tokens, confidence_scores, logits = model.forward_head( torch.empty(1, 2, 4, 2), torch.tensor([1]), pre_mix=torch.empty(1, 2, 4, 1), - confidence_threshold=0.5, + return_confidence=True, return_logits=True, ) torch.testing.assert_close(logits, normalized) - assert num_proposed.tolist() == [2] - assert confidence.call_args.args[0] is (collapsed if v41 else normalized) + assert draft_tokens.shape == (1, 2) + torch.testing.assert_close(confidence_scores, torch.full((1, 2), 10.0)) + assert confidence.call_args.args[0] is collapsed def _draft_routing(bias, *, hashed=False, table=None, top_k=3): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index e0406f881b54..40786627b200 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -97,11 +97,14 @@ def test_metadata_prepare_batch_indices(): assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] -def _make_worker(): +def _make_worker(enable_confidence_scheduling=False): cfg = types.SimpleNamespace( max_draft_len=5, spec_dec_mode=SpeculativeDecodingMode.DSPARK, - confidence_threshold=0.5, + enable_confidence_scheduling=enable_confidence_scheduling, + enable_fused_confidence_scheduler=False, + confidence_sps_table_path=None, + confidence_sts_path=None, ) from tensorrt_llm.mapping import Mapping @@ -585,7 +588,11 @@ def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): target rows by their runtime width rather than the configured K+1 width.""" worker = _make_worker() draft_model = _RecordingDraftModel() - metadata = types.SimpleNamespace(max_num_requests=2) + metadata = types.SimpleNamespace( + max_num_requests=2, + qo_indptr=None, + verify_lens=None, + ) worker._lazy_init(draft_model, metadata) slots = [worker._assign_slot(1000, reset=False), worker._assign_slot(1001, reset=False)] @@ -915,6 +922,7 @@ def _routing_config(embedded): draft_is_embedded_in_target=embedded, attention_backend="TRTLLM", confidence_threshold=0.5, + enable_confidence_scheduling=False, _use_shared_kv_cache=False, _allow_separate_draft_kv_cache=True, ) @@ -1105,3 +1113,72 @@ def test_prepare_resets_the_scratch_row_every_iteration(): assert int(worker._ctx_len[live]) == 17 assert int(worker._valid_len[live]) == 5 assert bool(worker._position_initialized[live]) + + +def test_ragged_draft_gathers_each_requests_hidden_state(): + worker = _make_worker() + draft_model = _RecordingDraftModel() + metadata = types.SimpleNamespace( + max_num_requests=2, + qo_indptr=torch.tensor([0, 4, 6], dtype=torch.int32, device="cuda"), + verify_lens=torch.tensor([4, 2], dtype=torch.int32, device="cuda"), + ) + worker._lazy_init(draft_model, metadata) + worker._batch_to_slot[:2] = torch.tensor([0, 1], device="cuda") + captured = torch.stack( + [torch.full((HIDDEN * NCAP,), float(index), device="cuda") for index in range(6)] + ) + metadata.get_hidden_states = lambda _num_tokens: captured + accepted = torch.zeros((2, 6), dtype=torch.int32, device="cuda") + + worker._draft_gen_block_batched( + draft_model, + metadata, + types.SimpleNamespace(num_ctx_tokens=0), + accepted, + torch.ones(2, dtype=torch.int32, device="cuda"), + num_contexts=0, + batch_size=2, + total_target_tokens=6, + position_ids=torch.arange(6, device="cuda").unsqueeze(0), + ) + + assert draft_model.forward_calls[-1]["main_hidden"][:, 0].tolist() == [0.0, 4.0] + + +def test_confidence_rows_are_slot_indexed_and_fail_open(): + worker = _make_worker(enable_confidence_scheduling=True) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + + buffer = worker.staged_confidence_buffer() + assert buffer is not None + neutral = worker.confidence_row_for(12345) + assert neutral == worker._neutral_conf_row + assert neutral > worker._scratch_slot + assert torch.sigmoid(buffer[neutral]).min().item() > 0.999 + + worker._req_to_slot = {100: 5, 101: 2, 102: 7} + for slot in range(worker._scratch_slot): + buffer[slot] = float(slot) + rows = [worker.confidence_row_for(request_id) for request_id in (100, 101, 102)] + assert rows == [5, 2, 7] + assert buffer[rows, 0].tolist() == [5.0, 2.0, 7.0] + + +def test_recycled_confidence_slot_is_reset_to_neutral(): + worker = _make_worker(enable_confidence_scheduling=True) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + worker._confidence_logits[: worker._scratch_slot] = -20.0 + + slot = worker._assign_slot(100, reset=False) + + assert torch.sigmoid(worker._confidence_logits[slot]).min().item() > 0.999 + + +def test_confidence_state_is_absent_when_scheduling_is_disabled(): + worker = _make_worker(enable_confidence_scheduling=False) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + + assert worker.staged_confidence_buffer() is None + assert worker.confidence_stamp_buffer() is None + assert worker.verify_planner is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py new file mode 100644 index 000000000000..d94e995df727 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py @@ -0,0 +1,300 @@ +# 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. +"""CPU regression tests for the embedded worker's confidence-state contract.""" + +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models import modeling_dspark +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead +from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker + +pytestmark = pytest.mark.cpu_only + + +def _worker(num_rows=2): + worker = DSv4DSparkWorker.__new__(DSv4DSparkWorker) + worker.spec_config = SimpleNamespace(max_draft_len=3) + worker.return_confidence = True + worker._scratch_slot = num_rows + worker._batch_to_slot = torch.arange(num_rows) + worker._ctx_len = torch.zeros(num_rows + 1, dtype=torch.long) + worker._valid_len = torch.zeros(num_rows + 1, dtype=torch.long) + worker._position_initialized = torch.zeros(num_rows + 1, dtype=torch.bool) + worker._kv_windows = torch.zeros(num_rows + 1, 1, 4, 2) + worker._win = 4 + worker._position_cap = None + worker._confidence_logits = torch.full((num_rows + 2, 3), 30.0) + worker._confidence_stamp = torch.zeros(num_rows + 2, dtype=torch.int32) + worker._draft_seq_host = 0 + worker._draft_seq_cuda = torch.zeros((), dtype=torch.int32) + worker._output_verify_lens = None + return worker + + +def _draft_model(): + def forward(main_hidden, bonus, start_pos, **kwargs): + batch = main_hidden.shape[0] + confidence = torch.full((batch, 3), -3.0) if kwargs["return_confidence"] else None + return None, confidence, torch.zeros(batch, 3, 4) + + return SimpleNamespace( + block_size=3, + forward_batched=Mock(side_effect=forward), + write_context_windows_batched=Mock(), + ) + + +def _draft( + worker, + model, + captured, + *, + num_contexts=0, + gen_start=0, + lengths=None, + offsets=None, + accepted_counts=None, + positions=None, +): + num_gens = len(accepted_counts) if accepted_counts is not None else 2 + accepted_counts = [1] * num_gens if accepted_counts is None else accepted_counts + metadata = SimpleNamespace( + qo_indptr=None if offsets is None else torch.tensor(offsets, dtype=torch.int32), + verify_lens=None if lengths is None else torch.tensor(lengths, dtype=torch.int32), + get_hidden_states=lambda _num_tokens: captured, + ) + accepted = torch.zeros(num_contexts + num_gens, 4, dtype=torch.int32) + counts = torch.tensor([1] * num_contexts + accepted_counts, dtype=torch.int32) + positions = torch.arange(len(captured)) if positions is None else positions + return worker._draft_gen_block_batched( + model, + metadata, + SimpleNamespace(num_ctx_tokens=gen_start), + accepted, + counts, + num_contexts, + num_contexts + num_gens, + len(captured), + positions, + ) + + +@pytest.mark.parametrize("return_confidence", [False, True]) +def test_mixed_ragged_draft_uses_generation_offsets_for_hidden_and_positions(return_confidence): + worker = _worker(3) + worker.return_confidence = return_confidence + model = _draft_model() + captured = torch.arange(9).float().unsqueeze(1).repeat(1, 2) + positions = torch.tensor([0, 1, 2, 100, 101, 102, 103, 200, 201]) + + _draft( + worker, + model, + captured, + num_contexts=1, + gen_start=3, + lengths=[4, 2], + offsets=[0, 4, 6], + accepted_counts=[3, 2], + positions=positions, + ) + + args = model.forward_batched.call_args.args + assert args[0][:, 0].tolist() == [5.0, 8.0] + assert args[2].tolist() == [103, 202] + backfill = model.write_context_windows_batched.call_args.args + assert backfill[0][:, :, 0][backfill[3]].tolist() == [3.0, 4.0, 7.0] + assert backfill[3].tolist() == [[True, True, False], [True, False, False]] + + +@pytest.mark.parametrize("offset, accepted_count", [(0, 3), (-1, 1)]) +def test_real_bonus_row_out_of_range_is_not_clamped_to_another_request(offset, accepted_count): + worker = _worker(1) + with pytest.raises(IndexError): + _draft( + worker, + _draft_model(), + torch.zeros(2, 2), + lengths=[3], + offsets=[offset, offset + 3], + accepted_counts=[accepted_count], + ) + + +def test_padding_offsets_do_not_index_past_captured_hidden_or_positions(): + worker = _worker(2) + worker._batch_to_slot[1] = worker._scratch_slot + model = _draft_model() + + _draft( + worker, + model, + torch.tensor([[1.0, 1.0], [2.0, 2.0]]), + lengths=[2, 2], + offsets=[0, 3, 5], + accepted_counts=[2, 2], + ) + + assert model.forward_batched.call_args.args[0][0].tolist() == [2.0, 2.0] + assert worker._confidence_logits[0].tolist() == [-3.0, -3.0, -3.0] + assert worker._confidence_logits[worker._scratch_slot + 1].tolist() == [30.0, 30.0, 30.0] + + +def test_draft_stamp_matches_previous_sequence_after_next_step_bump(): + worker = _worker(2) + assert worker.bump_draft_seq() == 0 + _draft(worker, _draft_model(), torch.zeros(8, 2)) + assert worker.confidence_stamp_buffer()[:2].tolist() == [1, 1] + assert worker.bump_draft_seq() == 1 + assert worker.verified_draft_seq_cuda().item() == 1 + assert torch.equal( + worker.confidence_stamp_buffer()[:2], worker.verified_draft_seq_cuda().expand(2) + ) + assert worker._confidence_stamp[worker._scratch_slot + 1].item() == 0 + + +def test_warmup_restores_confidence_logits_and_their_matching_stamps(monkeypatch): + worker = _worker(1) + worker.guided_decoder = None + worker._lazy_init = Mock() + worker._execute_guided_decoder_if_present = Mock() + accepted = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32) + counts = torch.tensor([2], dtype=torch.int32) + worker.sample_and_accept_draft_tokens = Mock(return_value=(accepted, counts)) + worker.sample_draft_tokens = Mock(return_value=torch.zeros(1, 3, dtype=torch.int32)) + worker.write_context_onehot_draft_probs = Mock() + worker._prepare_next_new_tokens = Mock(return_value=accepted) + worker._confidence_logits[0] = torch.tensor([4.0, 5.0, 6.0]) + worker._confidence_stamp[0] = 7 + worker._draft_seq_cuda.fill_(8) + saved_logits = worker._confidence_logits.clone() + saved_stamps = worker._confidence_stamp.clone() + saved_addresses = (worker._confidence_logits.data_ptr(), worker._confidence_stamp.data_ptr()) + metadata = SimpleNamespace( + is_cuda_graph=True, + qo_indptr=None, + verify_lens=None, + get_hidden_states=lambda _num_tokens: torch.zeros(4, 2), + draft_probs_last_dim=4, + batch_indices_cuda=torch.tensor([0]), + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + worker._forward_impl( + torch.tensor([1, 2, 3, 4]), + torch.arange(4), + None, + torch.zeros(4, 4), + SimpleNamespace(num_seqs=1, num_contexts=0, num_ctx_tokens=0), + metadata, + _draft_model(), + ) + + assert torch.equal(worker._confidence_logits, saved_logits) + assert torch.equal(worker._confidence_stamp, saved_stamps) + assert saved_addresses == ( + worker._confidence_logits.data_ptr(), + worker._confidence_stamp.data_ptr(), + ) + + +def test_configured_sts_is_loaded_on_the_resolved_head_and_planner(tmp_path): + worker = _worker(1) + path = tmp_path / "sts.json" + path.write_text(json.dumps({"sts_temperatures": [0.5, 1.0, 2.0]}), encoding="utf-8") + worker.spec_config.confidence_sts_path = str(path) + worker.spec_config.confidence_sps_table_path = None + worker.spec_config.enable_fused_confidence_scheduler = False + head = DSparkConfidenceHead(hidden_size=2, block_size=3) + model = SimpleNamespace( + dspark_model=SimpleNamespace(mtp_layers=[SimpleNamespace(confidence_head=head)]) + ) + temperatures = head.sts_temperatures + raw = torch.tensor([[1.0, 2.0, 3.0]]) + + worker._build_verify_planner(model, block_size=3) + + assert head.sts_temperatures is temperatures + assert torch.equal(temperatures, torch.tensor([0.5, 1.0, 2.0])) + assert worker.verify_planner.apply_calibration.__self__ is head + assert torch.allclose( + worker.verify_planner.apply_calibration(raw), torch.sigmoid(raw / temperatures) + ) + assert worker.verify_planner.max_verify_len == 3 + assert worker.verify_planner.cost_table is None + + +def test_configured_sts_without_a_confidence_head_is_rejected_before_loading(tmp_path): + worker = _worker(1) + worker.spec_config.confidence_sts_path = str(tmp_path / "must-not-be-opened.json") + worker.spec_config.confidence_sps_table_path = None + worker.spec_config.enable_fused_confidence_scheduler = False + + with pytest.raises( + ValueError, match="confidence_sts_path is set but no confidence head was found" + ): + worker._build_verify_planner(SimpleNamespace(), block_size=3) + + assert not hasattr(worker, "verify_planner") + + +@pytest.mark.parametrize("bias", [False, True]) +def test_embedded_checkpoint_load_preserves_optional_confidence_bias(monkeypatch, bias): + draft = modeling_dspark.DSv4DSparkDraftModel.__new__(modeling_dspark.DSv4DSparkDraftModel) + nn.Module.__init__(draft) + stage = nn.Module() + stage.confidence_head = DSparkConfidenceHead(hidden_size=2, block_size=2) + draft.mtp_layers = nn.ModuleList([stage]) + draft.num_stages = 1 + draft.post_load_weights = Mock() + draft.cache_attn_weights_from_state_dict = Mock() + wrapper = modeling_dspark.DSv4DSparkForCausalLM.__new__(modeling_dspark.DSv4DSparkForCausalLM) + nn.Module.__init__(wrapper) + wrapper.dspark_model = draft + wrapper.num_stages = 1 + checkpoint = {"proj.weight": torch.tensor([[2.0, 3.0]])} + if bias: + checkpoint["proj.bias"] = torch.tensor([4.0]) + prefix = "mtp_layers.0.confidence_head." + remapped = {prefix + name: value for name, value in checkpoint.items()} + expected_weights = {name: value.clone() for name, value in remapped.items()} + monkeypatch.setattr(modeling_dspark, "remap_dspark_draft_keys", Mock(return_value=remapped)) + + def load_received(weights: dict[str, torch.Tensor]) -> None: + assert set(weights) == set(expected_weights) + for name, value in expected_weights.items(): + assert torch.equal(weights[name], value) + head_weights = {name.removeprefix(prefix): value for name, value in weights.items()} + stage.confidence_head.load_weights([head_weights]) + + loader = SimpleNamespace(load_weights=Mock(side_effect=load_received)) + monkeypatch.setattr(modeling_dspark, "DeepseekV4WeightLoader", lambda _draft: loader) + weight_address = stage.confidence_head.proj.weight.data_ptr() + rng = torch.get_rng_state() + + wrapper.load_weights({}) + + loader.load_weights.assert_called_once() + assert stage.confidence_head.proj.weight.data_ptr() == weight_address + assert torch.equal(torch.get_rng_state(), rng) + assert (stage.confidence_head.proj.bias is not None) is bias + expected = 5.0 + (4.0 if bias else 0.0) + assert stage.confidence_head(torch.ones(1, 2)).item() == expected