[Spec Decode] Tighten confidence graph allocation

Compute current confidence ordering only on TP rank zero, retain two request-count CUDA graph tiers, and use automatic SPS in the confidence evaluation config.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml

@@ -18,5 +18,5 @@ server_args: >-
   --speculative-config '{"method":"dspark",
   "model":"deepseek-ai/DeepSeek-V4-Flash-DSpark",
   "attention_backend":"FLASH_ATTN","num_speculative_tokens":7,
-  "draft_sample_method":"probabilistic","dspark_confidence_threshold":0.0,
-  "dspark_budget_frac":0.5,"confidence_based_verification":"auto"}'
+  "draft_sample_method":"probabilistic","dspark_sps_curve":"auto",
+  "confidence_based_verification":"auto"}'

tests/test_config.py

@@ -1576,11 +1576,9 @@ def test_dspark_confidence_config_validation():
     config = SpeculativeConfig(
         method="ngram",
         num_speculative_tokens=1,
-        dspark_confidence_threshold=0.25,
         dspark_budget_frac=0.5,
         confidence_based_verification="mask",
     )
-    assert config.dspark_confidence_threshold == 0.25
     assert config.dspark_budget_frac == 0.5
     assert config.confidence_based_verification == "mask"
     with pytest.raises(ValidationError, match="confidence_based_verification"):
@@ -1594,8 +1592,6 @@ def test_dspark_confidence_config_validation():
 @pytest.mark.parametrize(
     ("field", "value"),
     [
-        ("dspark_confidence_threshold", -0.1),
-        ("dspark_confidence_threshold", 1.1),
         ("dspark_budget_frac", 0.0),
         ("dspark_budget_frac", 1.1),
     ],

tests/v1/spec_decode/test_dspark_confidence.py

@@ -45,7 +45,6 @@
 
 def _spec_config(**overrides: Any) -> SimpleNamespace:
     config = SimpleNamespace(
-        dspark_confidence_threshold=0.0,
         dspark_budget_frac=1.0,
         dspark_sps_curve=None,
         dspark_online_sts=False,
@@ -62,6 +61,7 @@ def _logit(probs: np.ndarray) -> np.ndarray:
 @pytest.fixture(autouse=True)
 def _single_rank_tp_group(monkeypatch):
     group = SimpleNamespace(
+        rank_in_group=0,
         broadcast=lambda tensor, src=0: tensor,
         broadcast_object=lambda value, src=0: value,
     )
@@ -196,22 +196,6 @@ def test_varlen_indexer_indices_use_device_lengths():
     assert indices.cpu().tolist() == [0, 1, 1, 2, 3]
 
 
-def test_select_budget_uses_global_prefix_order():
-    survival = compute_prefix_survival(
-        _logit(
-            np.array(
-                [
-                    [0.90, 0.90, 0.90],
-                    [0.95, 0.10, 0.99],
-                    [0.70, 0.70, 0.70],
-                ]
-            )
-        )
-    )
-    budget = select_draft_token_budget(survival, 3, 3, 3, min_survival=0.75)
-    assert budget == 3
-
-
 def test_select_budget_applies_fraction_globally():
     survival = compute_prefix_survival(_logit(np.array([[0.90, 0.80], [0.80, 0.80]])))
     budget = select_draft_token_budget(survival, 2, 2, 2, budget_frac=0.5)
@@ -362,6 +346,7 @@ def test_capacity_manager_assigns_scalar_budget_from_gpu_scores(monkeypatch):
     monkeypatch.setattr(
         "vllm.v1.worker.gpu.spec_decode.confidence.get_tp_group",
         lambda: SimpleNamespace(
+            rank_in_group=0,
             broadcast=lambda tensor, src=0: broadcasts.append(tensor),
             broadcast_object=lambda value, src=0: value,
         ),
@@ -679,7 +664,6 @@ def test_varlen_cudagraph_capture_adds_full_desc():
 
     manager._init_candidates()
 
-    assert any(
-        desc.max_req_tokens == manager.decode_query_len
-        for desc in manager._capture_descs[CUDAGraphMode.FULL]
-    )
+    descs = manager._capture_descs[CUDAGraphMode.FULL]
+    assert [desc.num_reqs for desc in descs] == [3, 5]
+    assert all(desc.max_req_tokens == manager.decode_query_len for desc in descs)

vllm/config/speculative.py

@@ -232,11 +232,6 @@ class SpeculativeConfig:
     synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'.
     Mutually exclusive with synthetic_acceptance_rates."""
 
-    dspark_confidence_threshold: float = 0.0
-    """Minimum DSpark cumulative prefix-survival probability for keeping a
-    per-request draft prefix. Set to 0.0 to use budget-based global top-k
-    allocation."""
-
     dspark_budget_frac: float = 1.0
     """Fraction of the full per-request draft-token budget available to the
     DSpark global prefix allocator."""
@@ -1217,11 +1212,6 @@ def _verify_args(self) -> Self:
                 "are only valid with rejection_sample_method='synthetic'."
             )
 
-        if not 0.0 <= self.dspark_confidence_threshold <= 1.0:
-            raise ValueError(
-                "dspark_confidence_threshold must be in [0, 1], got "
-                f"{self.dspark_confidence_threshold}."
-            )
         if not 0.0 < self.dspark_budget_frac <= 1.0:
             raise ValueError(
                 f"dspark_budget_frac must be in (0, 1], got {self.dspark_budget_frac}."
@@ -1254,17 +1244,10 @@ def _verify_args(self) -> Self:
         if (
             self.method == "dspark"
             and self.confidence_based_verification in ("auto", "mask")
-            and (
-                self.dspark_confidence_threshold > 0.0
-                or self.dspark_budget_frac < 1.0
-                or self.dspark_sps_curve is not None
-            )
+            and (self.dspark_budget_frac < 1.0 or self.dspark_sps_curve is not None)
             and "VLLM_MOE_SKIP_PADDING" not in os.environ
         ):
-            # The masked fallback keeps pruned verify rows in the batch as
-            # padding; pruning only saves work if MoE kernels skip those rows.
-            # Set here (frontend) so spawned workers inherit it before their
-            # env caches freeze. Set VLLM_MOE_SKIP_PADDING=0 to override.
+            # Set before spawning workers so masked rows can skip MoE work.
             logger.info(
                 "Confidence-based verification: defaulting "
                 "VLLM_MOE_SKIP_PADDING=1 so MoE kernels can skip masked "

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -232,11 +232,7 @@ def decode_descs(
                 if num_tokens > max_decode_tokens or num_tokens > max_cg_capture_size:
                     return
                 max_requests = min(num_tokens, self.max_num_reqs)
-                request_counts = {
-                    (max_requests + 1) // 2,
-                    (3 * max_requests + 3) // 4,
-                    max_requests,
-                }
+                request_counts = {(max_requests + 1) // 2, max_requests}
                 for num_reqs in sorted(request_counts):
                     if num_reqs * self.decode_query_len < num_tokens:
                         continue

vllm/v1/worker/gpu/model_runner.py

@@ -1048,7 +1048,6 @@ def prepare_inputs(
         if self.spec_decode_confidence_manager is not None:
             self.spec_decode_confidence_manager.trim_batch(input_batch, draft_tokens)
         if self.use_dcp:
-            # Prepare dcp local seq_lens.
             prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
                 self.input_buffers.seq_lens,
@@ -1061,8 +1060,7 @@ def prepare_inputs(
                 : input_batch.num_reqs_after_padding
             ]
         if uses_padding_mask:
-            # Mark trailing cudagraph-padding rows so kernels can skip work for
-            # them when supported.
+            # Mark trailing cudagraph padding.
             self.input_buffers.is_padding[
                 input_batch.num_tokens : num_tokens_after_padding
             ].fill_(True)

vllm/v1/worker/gpu/spec_decode/confidence.py

@@ -222,7 +222,6 @@ def select_draft_token_budget(
     num_required_target_tokens: int,
     budget_frac: float = 1.0,
     sps_table: np.ndarray | None = None,
-    min_survival: float = 0.0,
 ) -> int:
     """Choose the scalar draft budget from stale prefix scores.
 
@@ -232,9 +231,6 @@ def select_draft_token_budget(
     """
     flat = survival.reshape(-1)
     scores = flat[np.isfinite(flat)]
-    if min_survival > 0.0:
-        return int((scores >= min_survival).sum())
-
     scores = np.sort(scores)[::-1]
     total = scores.size
     max_admissions = min(int(total * budget_frac) + 1, total)
@@ -314,7 +310,6 @@ def __init__(
         self.copy_stream = torch.cuda.Stream(device)
         self.varlen_spec_decode = False
 
-        self.min_survival_probability = speculative_config.dspark_confidence_threshold
         self.capacity_budget_frac = speculative_config.dspark_budget_frac
         sps_curve = speculative_config.dspark_sps_curve
         self.wants_auto_sps_curve = sps_curve == "auto"
@@ -740,7 +735,6 @@ def _plan_draft_token_budget(
             num_required_target_tokens=int(num_required_target_tokens_per_req.sum()),
             budget_frac=self.capacity_budget_frac,
             sps_table=self.sps_table_np,
-            min_survival=self.min_survival_probability,
         )
         budget = min(budget, int(valid_draft_tokens_per_req.sum()))
         return int(get_tp_group().broadcast_object(budget, src=0))
@@ -764,45 +758,47 @@ def _assign_draft_token_budget(
                 capacities.clamp_(max=self.forced_capacity)
             return capacities
 
-        temperatures: np.ndarray | float = 1.0
-        if self.online_sts is not None:
-            temperatures = self.online_sts.temperatures
-        self._temperatures_np[:] = temperatures
-        self._temperatures.copy_(self._temperatures_cpu, non_blocking=True)
-        _compute_prefix_survival_kernel[(num_reqs,)](
-            self._confidence_logits,
-            self._confidence_logits.stride(0),
-            input_batch.idx_mapping,
-            self._valid_draft_tokens,
-            self._temperatures,
-            self._survival,
-            NUM_STEPS=self.num_speculative_steps,
-        )
-        num_candidates = num_reqs * self.num_speculative_steps
-        torch.sort(
-            self._survival[:num_reqs].flatten(),
-            descending=True,
-            stable=True,
-            out=(
-                self._sorted_survival[:num_candidates],
-                self._sorted_indices[:num_candidates],
-            ),
-        )
-        capacities.zero_()
-        if draft_token_budget > 0:
-            admitted = self._sorted_indices[:draft_token_budget]
-            torch.div(
-                admitted,
-                self.num_speculative_steps,
-                rounding_mode="floor",
-                out=admitted,
+        tp_group = get_tp_group()
+        if tp_group.rank_in_group == 0:
+            temperatures: np.ndarray | float = 1.0
+            if self.online_sts is not None:
+                temperatures = self.online_sts.temperatures
+            self._temperatures_np[:] = temperatures
+            self._temperatures.copy_(self._temperatures_cpu, non_blocking=True)
+            _compute_prefix_survival_kernel[(num_reqs,)](
+                self._confidence_logits,
+                self._confidence_logits.stride(0),
+                input_batch.idx_mapping,
+                self._valid_draft_tokens,
+                self._temperatures,
+                self._survival,
+                NUM_STEPS=self.num_speculative_steps,
             )
-            capacities.scatter_add_(
-                0,
-                admitted,
-                self._capacity_ones[:draft_token_budget],
+            num_candidates = num_reqs * self.num_speculative_steps
+            torch.sort(
+                self._survival[:num_reqs].flatten(),
+                descending=True,
+                stable=True,
+                out=(
+                    self._sorted_survival[:num_candidates],
+                    self._sorted_indices[:num_candidates],
+                ),
             )
-        get_tp_group().broadcast(capacities, src=0)
+            capacities.zero_()
+            if draft_token_budget > 0:
+                admitted = self._sorted_indices[:draft_token_budget]
+                torch.div(
+                    admitted,
+                    self.num_speculative_steps,
+                    rounding_mode="floor",
+                    out=admitted,
+                )
+                capacities.scatter_add_(
+                    0,
+                    admitted,
+                    self._capacity_ones[:draft_token_budget],
+                )
+        tp_group.broadcast(capacities, src=0)
         return capacities
 
     def warmup(self, input_buffers: "InputBuffers") -> None:
@@ -1116,7 +1112,7 @@ def __init__(self, *args: Any, **kwargs: Any):
             self.wants_auto_sps_curve = False
             logger.info_once(
                 "DSpark auto SPS profiling requires compact verification; "
-                "masked verification will use the configured threshold or budget."
+                "masked verification will use the configured budget."
             )
 
     def _prepare_forward_skip_mask(

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

@@ -80,8 +80,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.use_confidence_based_verification = (
             self.speculative_config.confidence_based_verification not in ("none", "off")
             and (
-                self.speculative_config.dspark_confidence_threshold > 0.0
-                or self.speculative_config.dspark_budget_frac < 1.0
+                self.speculative_config.dspark_budget_frac < 1.0
                 or sps_curve is not None
             )
         )

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

@@ -53,8 +53,7 @@ def __init__(
             spec_config.method == "dspark"
             and spec_config.confidence_based_verification not in ("none", "off")
             and (
-                spec_config.dspark_confidence_threshold > 0.0
-                or spec_config.dspark_budget_frac < 1.0
+                spec_config.dspark_budget_frac < 1.0
                 or spec_config.dspark_sps_curve is not None
             )
         )