[Spec Decode] Tighten confidence allocation

Keep draft and verifier costs as separate one-dimensional tables, compute current confidence ordering only on TP rank zero, reduce manager staging state, and retain two request-count CUDA graph tiers.

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

@@ -32,7 +32,7 @@
 from vllm.v1.worker.gpu.spec_decode.confidence import (  # noqa: E402
     MaskedConfidenceManager,
     VarlenConfidenceManager,
-    build_sps_table_from_costs,
+    build_cost_tables_from_curves,
     compute_prefix_survival,
     make_confidence_manager,
     select_draft_token_budget,
@@ -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,
     )
@@ -95,9 +95,7 @@ def test_auto_verification_selects_supported_manager(cg_support, expected_type):
     manager = make_confidence_manager(
         "auto",
         attn_cg_support,
-        max_num_tokens=8,
         req_states=req_states,
-        device=device,
         speculative_config=_spec_config(),
     )
 
@@ -107,7 +105,6 @@ def test_auto_verification_selects_supported_manager(cg_support, expected_type):
 def test_masked_verification_does_not_auto_profile_sps():
     device = torch.device("cuda")
     manager = MaskedConfidenceManager(
-        max_num_tokens=8,
         req_states=RequestState(
             max_num_reqs=2,
             max_model_len=4,
@@ -116,12 +113,12 @@ def test_masked_verification_does_not_auto_profile_sps():
             vocab_size=16,
             device=device,
         ),
-        device=device,
         speculative_config=_spec_config(dspark_sps_curve="auto"),
     )
 
     assert not manager.wants_auto_sps_curve
-    assert manager.sps_table_np is None
+    assert manager.draft_cost_ms is None
+    assert manager.verify_and_sample_cost_ms is None
 
 
 def test_prepare_pos_seq_lens_clears_active_padding():
@@ -196,32 +193,16 @@ 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)
-    assert budget == 3
+    assert budget == 2
 
 
 def test_select_budget_is_hard_cap_under_ties():
     survival = compute_prefix_survival(_logit(np.array([[0.90, 0.80], [0.90, 0.80]])))
     budget = select_draft_token_budget(survival, 2, 2, 2, budget_frac=0.25)
-    assert budget == 2
+    assert budget == 1
 
 
 def test_select_budget_is_hard_cap_for_saturated_scores():
@@ -232,7 +213,7 @@ def test_select_budget_is_hard_cap_for_saturated_scores():
     """
     survival = compute_prefix_survival(np.full((4, 7), 40.0))
     budget = select_draft_token_budget(survival, 4, 4, 4, budget_frac=0.5)
-    assert budget == int(4 * 7 * 0.5) + 1
+    assert budget == int(4 * 7 * 0.5)
 
 
 def test_select_budget_never_admits_zero_survival():
@@ -256,46 +237,45 @@ def test_select_budget_uses_sps_argmax():
     # SPS drops sharply after B=4 so theta peaks at k=2:
     #   k=0: 2.00*1.00, k=1: 2.90*0.95=2.755, k=2: 3.62*0.90=3.258,
     #   k=3: 4.22*0.20=0.844, k=4: 4.52*0.10=0.452.
-    sps_table = np.broadcast_to(
-        np.array([1.0, 1.0, 1.0, 0.95, 0.90, 0.20, 0.10]),
-        (3, 7),
-    )
+    rates = np.array([1.0, 1.0, 1.0, 0.95, 0.90, 0.20, 0.10])
     budget = select_draft_token_budget(
         survival,
         num_batch_requests=2,
         num_sampling_requests=2,
         num_required_target_tokens=2,
-        sps_table=sps_table,
+        draft_cost_ms=np.zeros(3),
+        verify_and_sample_cost_ms=1000.0 / rates,
     )
     assert budget == 2
 
 
 def test_select_budget_counts_only_requests_that_sample():
     survival = np.array([[0.9, 0.8], [0.7, 0.6]])
-    sps_table = np.ones((3, 9))
-    sps_table[:, 5:] = 0.1
+    rates = np.ones(9)
+    rates[5:] = 0.1
 
     budget = select_draft_token_budget(
         survival,
         num_batch_requests=2,
         num_sampling_requests=1,
         num_required_target_tokens=4,
-        sps_table=sps_table,
+        draft_cost_ms=np.zeros(3),
+        verify_and_sample_cost_ms=1000.0 / rates,
     )
 
     assert budget == 0
 
 
-def test_sps_table_combines_direct_draft_and_verification_costs():
+def test_cost_tables_keep_draft_and_verification_costs_separate():
     draft_curve = [(1, 0.2), (4, 0.8), (8, 1.6)]
     verify_curve = [(1, 2.0), (8, 3.0), (64, 5.0)]
 
-    table = build_sps_table_from_costs(
+    draft_cost_ms, verify_and_sample_cost_ms = build_cost_tables_from_curves(
         draft_curve, verify_curve, max_num_reqs=8, max_batch_tokens=64
     )
 
-    assert 1000.0 / table[1, 8] == pytest.approx(3.2)
-    assert 1000.0 / table[8, 8] == pytest.approx(4.6)
+    assert draft_cost_ms[1] + verify_and_sample_cost_ms[8] == pytest.approx(3.2)
+    assert draft_cost_ms[8] + verify_and_sample_cost_ms[8] == pytest.approx(4.6)
 
 
 def test_select_budget_temperature_desaturates_zeros():
@@ -312,9 +292,8 @@ def test_select_budget_temperature_desaturates_zeros():
     )
     # T=1: exact-zero survival past position 2 truncates both requests.
     assert budget_t1 == 4
-    # T=80: sigmoid(-10) > 0, so the budget (int(10*0.9)+1 = 10 admissions,
-    # capped at 5 per request) is fully spent.
-    assert budget_t80 == 10
+    # T=80: sigmoid(-10) > 0, so 90% of the candidates are admitted.
+    assert budget_t80 == 9
 
 
 def test_online_sts_fits_order_preserving_temperatures():
@@ -362,6 +341,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,
         ),
@@ -377,9 +357,7 @@ def test_capacity_manager_assigns_scalar_budget_from_gpu_scores(monkeypatch):
     )
     req_states.req_id_to_index = {"req0": 0, "req1": 1}
     handler = VarlenConfidenceManager(
-        max_num_tokens=16,
         req_states=req_states,
-        device=device,
         speculative_config=_spec_config(),
     )
     input_batch: Any = SimpleNamespace(
@@ -433,9 +411,7 @@ def test_varlen_capacity_manager_compacts_verifier_batch():
     )
     req_states.req_id_to_index = {"req0": 0, "req1": 1}
     handler = VarlenConfidenceManager(
-        max_num_tokens=16,
         req_states=req_states,
-        device=device,
         speculative_config=_spec_config(dspark_budget_frac=0.34),
     )
     handler._confidence_logits[:2].copy_(
@@ -541,9 +517,7 @@ def test_masked_capacity_manager_marks_pruned_tokens_for_forward_and_sampler():
     )
     req_states.req_id_to_index = {"req0": 0, "req1": 1}
     handler = MaskedConfidenceManager(
-        max_num_tokens=16,
         req_states=req_states,
-        device=device,
         speculative_config=_spec_config(dspark_budget_frac=0.34),
     )
     handler._confidence_logits[:2].copy_(
@@ -679,7 +653,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

@@ -454,9 +454,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
             self.spec_decode_confidence_manager = make_confidence_manager(
                 self.speculative_config.confidence_based_verification,
                 attn_cg_support,
-                self.max_num_tokens,
                 self.req_states,
-                self.device,
                 self.speculative_config,
             )
         self.block_tables = BlockTables(
@@ -1048,7 +1046,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 +1058,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)