zhaoyuanh-nvidia · Nov 3, 2026 · Nov 2, 2026
@@ -36,6 +36,7 @@ for how to build TensorRT LLM from source and start a TRT-LLM Docker container.
     - [Sparse attention](#sparse-attention)
     - [KV cache](#kv-cache)
     - [Quantized checkpoints](#quantized-checkpoints)
+  - [DSpark confidence scheduling (prototype)](#dspark-confidence-scheduling-prototype)
   - [Notes and Troubleshooting](#notes-and-troubleshooting)
 
 
@@ -472,6 +473,73 @@ MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-e
 configuration.
 
 
+## DSpark confidence scheduling (prototype)
+
+DSpark always drafts a block of `K` tokens. With confidence scheduling on, a trained confidence head
+scores each drafted position, and in decode-only steps the runtime may verify fewer than all
+`G * (K + 1)` rows (`G` = requests per rank, `V` = verified rows). It picks `V` from a **cost table**
+of measured step times, and only when the predicted accepted tokens per millisecond beat full
+verification by `confidence_min_predicted_gain`. Target verification and the acceptance rule are
+unchanged; mixed prefill/decode steps always verify every row. The feature is off by default.
+
+### Enable it
+
+```yaml
+cuda_graph_config:
+  enable_padding: true
+  batch_sizes: [16, 32, 64, 128]        # rank-local G; the cost table must cover every entry
+speculative_config:
+  decoding_type: DSpark
+  max_draft_len: 5
+  enable_confidence_scheduling: true
+  confidence_cost_table_path: /path/to/dspark-costs.json
+  confidence_calibration_path: /path/to/calibration.json   # optional
+  confidence_min_predicted_gain: 0.04                      # default
+```
+
+A static baseline can reuse the same file with `enable_confidence_scheduling: false`; the other
+confidence fields are then ignored. `max_draft_len: 1` has no smaller tier and always runs static.
+
+### Get a cost table
+
+Step costs depend on the GPU, parallel layout and kernels, so generate the table on the system
+and with the runtime YAML you serve with. One file can hold several `K`; the runtime uses the entry
+for `max_draft_len`:
+
+```bash
+python -m tensorrt_llm.tools.dspark_profile_costs \
+  --model <YOUR_MODEL_DIR> --config runtime.yaml --k 3 5 --output dspark-costs.json
+```
+
+The tool loads the model once per `K`. For every captured `G` it runs a short wave of requests
+(another one while a budget has too few samples), cycles the full and four compact budgets through
+steady decode steps, and records the median whole-step time (draft and verify included). Running it
+again adds or replaces the profiled `K` and `G` rows in the same file, so `--k 4` or
+`--graph-batch-sizes 512` extends an existing table. `--graph-batch-sizes`, `--tier-fractions` and
+`--dataset` override the defaults; see `--help`.
+
+At startup the runtime checks the table against the engine and fails with the command above when:
+
+- the file is missing, or has no entry for `max_draft_len`;
+- a captured `cuda_graph_config.batch_sizes` entry has no row (extra rows are ignored);
+- it was measured on another GPU model, world size, TP/EP size or attention-DP setting.
+
+A table from another TensorRT LLM version only warns. Tables in the earlier single-`K` format
+(`schema_version: 2`) still load.
+
+### Calibration
+
+`confidence_calibration_path` points to per-position temperatures (`{"sts_temperatures": [...]}`).
+A vector longer than `K` is truncated, so one file fitted at the checkpoint's block size serves
+every `K`. Without it the raw sigmoid confidence is used.
+
+### Limitations
+
+Confidence scheduling requires the embedded DeepSeek-V4 DSpark drafter, padded CUDA graphs,
+PP=1 and CP=1. Guided decoding, LoRA, per-request logits post-processors, per-layer dynamic EPLB,
+heuristic sparse top-k, the acceptance-rate gate (`acceptance_rate_window_size` /
+`acceptance_rate_threshold`) and non-greedy rejection draft sampling are rejected.
+
 ## Notes and Troubleshooting
 
 - `DeepSeek-V4 on Hopper requires kv_cache_config.dtype='fp8_ds_mla'`: remove an explicit BF16/FP8

@@ -2111,13 +2111,28 @@ def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_
         )
 
     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:
+        """Prepare and validate the embedded confidence namespace before loading."""
+        last = self.dspark_model.mtp_layers[-1]
+        head = last.confidence_head
+        prefix = f"mtp_layers.{self.num_stages - 1}.confidence_head."
+        module_weights = {
+            name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix)
+        }
+        if head is not None and not module_weights:
+            # A checkpoint without a trained confidence head still serves static
+            # DSpark; drop the untrained head instead of rejecting the load.
+            last.confidence_head = head = None
+        if head is None:
+            if module_weights:
+                raise ValueError("DSpark confidence head checkpoint has no constructed head")
+            return
+        if "proj.bias" in module_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))
+        # The generic V4 loader skips this parameter-free head container and
+        # visits its Linear child, so strict validation must run at this seam.
+        head._validate_checkpoint_weights(module_weights)
 
     def load_weights(self, weights: Dict, weight_mapper=None, **kwargs):
         """Load the ``mtp.*`` draft weights from the (full) checkpoint dict.

@@ -295,9 +295,6 @@ def load_weights(self, weights: dict[str, torch.Tensor], weight_mapper=None, **k
         missing = [name for name in required if name not in remapped]
         if missing:
             raise ValueError(f"V4.1 DSpark checkpoint is missing draft head weights: {missing}")
-        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)

@@ -555,6 +555,13 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None:
     def load_weights(self, weights: list) -> None:
         """Load every confidence-head checkpoint parameter strictly."""
         (module_weights, ) = weights
+        self._validate_checkpoint_weights(module_weights)
+        for name, param in self.named_parameters():
+            param.data.copy_(module_weights[name][:])
+
+    def _validate_checkpoint_weights(
+            self, module_weights: dict[str, torch.Tensor]) -> None:
+        """Validate a head namespace before generic embedded parameter loading."""
         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]
@@ -569,7 +576,11 @@ def load_weights(self, weights: list) -> None:
                 f"DSpark confidence head checkpoint is missing key(s) {sorted(missing)}"
             )
         for name, param in expected.items():
-            param.data.copy_(module_weights[name][:])
+            actual_shape = tuple(module_weights[name].shape)
+            if actual_shape != tuple(param.shape):
+                raise ValueError(
+                    f"DSpark confidence head checkpoint {name} has shape {actual_shape}; "
+                    f"expected {tuple(param.shape)}")
 
 
 class Eagle3Attention(Attention):

@@ -89,6 +89,8 @@
 from .seq_slot_manager import SeqSlotManager
 
 GB = 1 << 30
+_DEFAULT_INDEX_MAPPER_RESERVED_SLOTS = 1
+_CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS = 2
 
 # Page sizes of a standalone drafter's separate KV pool (DFlash, standalone
 # DSpark). The pool keeps the target's kv_cache_config.tokens_per_block when it
@@ -107,6 +109,15 @@
 _HELIX_AGG_FREE_GPU_MEMORY_FRACTION = 0.6
 
 
+def _get_num_reserved_index_slots(spec_config: Optional[SpeculativeConfig],
+                                  estimating_kv_cache: bool) -> int:
+    """Return IndexMapper headroom required by the execution policy."""
+    if getattr(spec_config, "enable_confidence_scheduling", False):
+        return (0 if estimating_kv_cache else
+                _CONFIDENCE_INDEX_MAPPER_RESERVED_SLOTS)
+    return _DEFAULT_INDEX_MAPPER_RESERVED_SLOTS
+
+
 def ceil_div(a: int, b: int) -> int:
     return (a + b - 1) // b
 
@@ -3524,6 +3535,9 @@ def _create_kv_cache_manager(
         manager_extra_kwargs["max_copy_beam_width"] = max_copy_beam_width
         manager_extra_kwargs[
             "cold_page_codec_provider"] = cold_page_codec_provider
+        if getattr(spec_config, "enable_confidence_scheduling", False):
+            manager_extra_kwargs["num_reserved_index_slots"] = (
+                _get_num_reserved_index_slots(spec_config, estimating_kv_cache))
         manager_extra_kwargs["kv_events_config"] = kv_events_config
         manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse
         manager_extra_kwargs[

@@ -1598,8 +1598,29 @@ def __init__(self, config: CUDAGraphRunnerConfig) -> None:
         # The decoder publishes the fitted bucket before graph lookup.
         self.agreed_ragged_bucket: Optional[int] = None
         self.adp_shape_agreement: Optional[ADPShapeAgreement] = None
+        self._policy_iteration: Optional[int] = None
         self.secondary_padding_dummy_requests: Dict[int, LlmRequest] = {}
 
+    def set_policy_iteration(self, iteration: int) -> None:
+        """Advance the executor-owned epoch before publishing this step's vote."""
+        iteration = int(iteration)
+        if getattr(self, "_policy_iteration", None) != iteration:
+            self.adp_shape_agreement = None
+        self._policy_iteration = iteration
+
+    @contextlib.contextmanager
+    def pad_batch(self,
+                  scheduled_requests: ScheduledRequests,
+                  resource_manager: ResourceManager,
+                  runtime_draft_len: int = 0):
+        """Consume the shape vote once, including an exceptional forward exit."""
+        try:
+            with super().pad_batch(scheduled_requests, resource_manager,
+                                   runtime_draft_len) as padded_requests:
+                yield padded_requests
+        finally:
+            self.adp_shape_agreement = None
+
     def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]:
         """Return the agreed token bucket only when every row has a window."""
         # A window-less batch must not replay stale token-major row maps.
@@ -1626,7 +1647,8 @@ def _agree_graph_shape(
         reuse_agreement = bool(
             agreement is not None and agreement.reuse_graph_shape
             and not self.enable_in_graph_sampling
-            and agreement.matches(batch, agreement.iteration, padded=True))
+            and getattr(self, "_policy_iteration", None) is not None
+            and agreement.matches(batch, self._policy_iteration, padded=True))
         if reuse_agreement:
             if (not can_run_cuda_graph
                     or batch.batch_size != agreement.graph_batch_size or
@@ -1675,8 +1697,9 @@ def _get_num_tokens_for_key(self, key: KeyType) -> int:
     def _reusable_padding_agreement(
             self, batch: ScheduledRequests) -> Optional[ADPShapeAgreement]:
         agreement = self.adp_shape_agreement
-        if (agreement is not None and agreement.reuse_graph_shape and
-                agreement.matches(batch, agreement.iteration, padded=False)):
+        if (agreement is not None and agreement.reuse_graph_shape
+                and getattr(self, "_policy_iteration", None) is not None and
+                agreement.matches(batch, self._policy_iteration, padded=False)):
             return agreement
         return None
 
@@ -1760,6 +1783,7 @@ def reset_padding_dummies(self) -> None:
         self.secondary_padding_dummy_requests = {}
         self.ragged_zero_real_high_rows = 0
         self.adp_shape_agreement = None
+        self._policy_iteration = None
 
     def _on_padding_dummy_release(self) -> None:
         # Releasing a dummy invalidates the agreement that certified it.
@@ -1768,7 +1792,7 @@ def _on_padding_dummy_release(self) -> None:
     def _preallocate_extra_padding_dummy(self,
                                          resource_manager: ResourceManager,
                                          draft_len: int) -> None:
-        if not getattr(self.spec_config, "confidence_sps_table_path", None):
+        if not getattr(self.spec_config, "confidence_cost_table_path", None):
             return
         if self._get_or_create_padding_dummy(resource_manager,
                                              draft_len,

@@ -56,7 +56,7 @@
 from .config import DecoderRunnerConfig
 from .dummy_requests import DummyRequestFactory
 from .inputs import DecoderInputPreparer
-from .metadata import DecoderMetadata, resolve_mamba_metadata_cls
+from .metadata import DecoderMetadata, resolve_mamba_metadata_cls, supports_advanced_draft_sampling
 from .warmup import WarmupAgreement
 
 # Context requests per warmup batch whose breakable prefill graph breaks get
@@ -924,7 +924,8 @@ def _capture_variant(
             # Dummy warmup requests carry no sampling params, so the greedy pass
             # needs no override while the advanced one has to force the flag.
             _capture_variant("greedy")
-            _capture_variant("advanced sampling", force_non_greedy=True)
+            if supports_advanced_draft_sampling(self._config.spec_config):
+                _capture_variant("advanced sampling", force_non_greedy=True)
             if temperature_only_graphs:
                 # Temperature-only batches replay the vocabulary-sharded Gumbel-max
                 # target sampling (the advanced pass records the stock sampler: its

@@ -417,16 +417,26 @@ def __init__(
         self._dspark_sps_cost_table = None
         self._dspark_exact_candidate_cells = ()
         self._dspark_exact_identity_words = (0,) * 8
-        table_path = getattr(spec_config, "confidence_sps_table_path", None)
+        table_path = getattr(spec_config, "confidence_cost_table_path", None)
         if self._dspark_confidence_enabled and table_path:
-            from tensorrt_llm._torch.speculative.dspark_planner import load_runtime_sps_cost_table
+            from tensorrt_llm._torch.speculative.dspark_planner import (
+                CostTableSetup,
+                load_runtime_sps_cost_table,
+            )
+            from tensorrt_llm.version import __version__ as trtllm_version
 
-            self._dspark_sps_cost_table, _ = load_runtime_sps_cost_table(
+            self._dspark_sps_cost_table = load_runtime_sps_cost_table(
                 table_path,
                 graph_batch_sizes=self._config.cuda_graph_batch_sizes,
                 max_draft_len=int(spec_config.max_draft_len),
-                live_engine_fingerprint_path=getattr(
-                    spec_config, "confidence_sps_live_fingerprint_path", None
+                minimum_predicted_gain=float(spec_config.confidence_min_predicted_gain),
+                live_setup=CostTableSetup(
+                    gpu=torch.cuda.get_device_name(self._buffers.input_ids_cuda.device),
+                    world_size=int(self.mapping.world_size),
+                    tensor_parallel_size=int(self.mapping.tp_size),
+                    moe_expert_parallel_size=int(self.mapping.moe_ep_size),
+                    enable_attention_dp=bool(self._config.enable_attention_dp),
+                    tensorrt_llm_version=trtllm_version,
                 ),
             )
             self._dspark_exact_candidate_cells = self._dspark_sps_cost_table.candidate_cells()

@@ -41,6 +41,7 @@
     from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager
     from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
     from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager
+    from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
     from .inputs import DecoderInputBuffers
 
@@ -372,7 +373,11 @@ def update_spec_metadata(
 
 
 def sync_group_all_greedy_sample(
-    spec_metadata: SpecMetadata, mapping: Mapping, dist: Distributed | None
+    spec_metadata: SpecMetadata,
+    mapping: Mapping,
+    dist: Distributed | None,
+    *,
+    confidence_enabled: bool = False,
 ) -> None:
     """All-gather the per-rank greedy flags and store the group AND.
 
@@ -387,20 +392,47 @@ def sync_group_all_greedy_sample(
     ``all_rank_num_tokens`` exchange, which runs in ``prepare_inputs`` --
     after the graph key, too late for the key to see the synced value.
     """
-    # enable_lm_head_tp_in_adp implies enable_attention_dp (asserted in
-    # Mapping.__init__), so ADP needs no separate check here.
-    if not (mapping.enable_lm_head_tp_in_adp and spec_metadata.use_rejection_sampling):
+    # Confidence needs the same group-uniform rejection decision even without
+    # LM-head TP. Feature-off execution retains the existing gate.
+    if not (
+        spec_metadata.use_rejection_sampling
+        and (mapping.enable_lm_head_tp_in_adp or confidence_enabled)
+    ):
         return
     local_flag = bool(spec_metadata.is_all_greedy_sample)
-    all_flags = dist.tp_allgather_int64([local_flag])[:, 0]
-    spec_metadata.group_all_greedy_sample = bool(all_flags.all())
+    if mapping.tp_size > 1:
+        all_flags = dist.tp_allgather_int64([local_flag])[:, 0]
+        spec_metadata.group_all_greedy_sample = bool(all_flags.all())
+    else:
+        spec_metadata.group_all_greedy_sample = local_flag
     # Also overwrite the live flag directly: this iteration's scan already
     # ran (update_is_all_greedy_sample just returned) and the CUDA graph
     # key reads the flag next -- the stored override only takes effect on
     # the NEXT rescan (populate), which is after key selection.
     spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample
 
 
+def supports_advanced_draft_sampling(spec_config: DecodingBaseConfig | None) -> bool:
+    """Whether the engine may capture an advanced draft-sampling variant."""
+    return not (
+        getattr(spec_config, "enable_confidence_scheduling", False)
+        and spec_config.use_rejection_sampling
+    )
+
+
+def validate_dspark_draft_sampling(
+    spec_metadata: SpecMetadata,
+    *,
+    confidence_enabled: bool,
+) -> None:
+    """Reject unsupported group-uniform sampling before graph lookup."""
+    if confidence_enabled and spec_metadata.wants_advanced_draft_sampling:
+        raise ValueError(
+            "DSpark confidence scheduling does not support advanced rejection "
+            "draft sampling: confidence is conditioned on the greedy proposal prefix"
+        )
+
+
 def set_spec_metadata_all_rank_num_tokens(
     spec_metadata: SpecMetadata,
     spec_all_rank_num_tokens: list[int],