[TRTLLM-14620][feat] Integrate calibrated DSpark confidence scoring

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"
         )

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:

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):