[SpecDecode] Preserve reject-only residual verifier padding

tests/v1/core/test_scheduler.py

@@ -1724,6 +1724,39 @@ def test_spec_decode_padding_first_decode_step():
     # r2 is padded to the 1 + num_spec shape with placeholder (-1) drafts.
     assert out.num_scheduled_tokens[r2.request_id] == 1 + num_spec
     assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec
+    assert out.num_invalid_spec_tokens == {r2.request_id: num_spec}
+
+
+def test_spec_decode_padding_follows_mamba_alignment():
+    """Mamba alignment must not clip reject-only verifier padding."""
+    num_spec = 3
+    scheduler = create_scheduler(
+        num_speculative_tokens=num_spec,
+        enable_prefix_caching=True,
+        block_size=16,
+        use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
+    )
+    scheduler.need_mamba_block_aligned_split = True
+    running, prompt_tail = create_requests(num_requests=2, num_tokens=14)
+    scheduler.connector.get_num_new_matched_tokens = Mock(
+        side_effect=lambda request, _: (
+            (13, False) if request is prompt_tail else (0, False)
+        )
+    )
+
+    scheduler.add_request(running)
+    out = scheduler.schedule()
+    _model_output(scheduler, out, [[100]])
+    scheduler.update_draft_token_ids(DraftTokenIds([running.request_id], [[1, 2, 3]]))
+
+    scheduler.add_request(prompt_tail)
+    out = scheduler.schedule()
+
+    assert out.num_scheduled_tokens == {
+        running.request_id: 1 + num_spec,
+        prompt_tail.request_id: 1 + num_spec,
+    }
+    assert out.scheduled_spec_decode_tokens[prompt_tail.request_id] == [-1] * num_spec
 
 
 def test_spec_decode_padding_skipped_for_diffusion():
@@ -1761,6 +1794,33 @@ def test_spec_decode_padding_skipped_for_diffusion():
     assert r2.request_id not in out.scheduled_spec_decode_tokens
 
 
+def test_spec_decode_padding_skipped_for_adaptive_verification():
+    num_spec = 3
+    scheduler = create_scheduler(
+        num_speculative_tokens=num_spec,
+        enable_prefix_caching=True,
+        block_size=16,
+    )
+    speculative_config = scheduler.vllm_config.speculative_config
+    assert speculative_config is not None
+    speculative_config.enable_adaptive_verification = True
+    r1, r2 = create_requests(
+        num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
+    )
+
+    scheduler.add_request(r1)
+    out = scheduler.schedule()
+    _model_output(scheduler, out, [[100]])
+    scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
+
+    scheduler.add_request(r2)
+    out = scheduler.schedule()
+
+    assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
+    assert out.num_scheduled_tokens[r2.request_id] == 1
+    assert r2.request_id not in out.scheduled_spec_decode_tokens
+
+
 def test_spec_decode_padding_skipped_with_prefill_in_batch():
     """Padding is skipped when the batch contains a prefill chunk: the batch is
     already mixed/non-uniform, so padding a new decode request buys nothing.

tests/v1/spec_decode/test_dynamic_sd_cug.py

@@ -15,7 +15,9 @@
     SchedulerConfig,
     VllmConfig,
 )
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
+from vllm.v1.worker.gpu import model_runner as gpu_model_runner
 from vllm.v1.worker.utils import get_uniform_decode_token_count
 
 pytestmark = pytest.mark.cpu_test
@@ -84,6 +86,42 @@ def _create_vllm_config_for_dsd(
     return vllm_config
 
 
+def test_model_runner_dispatches_full_graph_only_for_verifier_topology(
+    monkeypatch,
+):
+    dispatched_uniform_tokens = []
+
+    def fake_dispatch(
+        _manager, _num_reqs, _num_tokens, uniform_token_count, *_args, **_kwargs
+    ):
+        dispatched_uniform_tokens.append(uniform_token_count)
+        return SimpleNamespace(num_tokens=0), None
+
+    monkeypatch.setattr(gpu_model_runner, "dispatch_cg_and_sync_dp", fake_dispatch)
+    runner = object.__new__(gpu_model_runner.GPUModelRunner)
+    runner.speculative_config = SimpleNamespace()
+    runner.model_state = SimpleNamespace(num_new_sampled_tokens_per_step=1)
+    runner.lora_config = None
+    runner.is_encoder_decoder = False
+    runner.cudagraph_manager = None
+    runner.dp_size = 1
+    runner.dp_rank = 0
+    runner.kv_connector = SimpleNamespace(no_forward=lambda _output: None)
+    runner._merge_ec_connector_no_forward = lambda *_args: None
+
+    def execute(scheduled_drafts):
+        output = SchedulerOutput.make_empty()
+        output.num_scheduled_tokens = {"r0": 8, "r1": 8}
+        output.total_num_scheduled_tokens = 16
+        output.scheduled_spec_decode_tokens = scheduled_drafts
+        runner.execute_model(output, dummy_run=True)
+
+    execute({"r0": [1] * 7, "r1": [2] * 7})
+    execute({})
+    execute({"r0": [1] * 7})
+    assert dispatched_uniform_tokens == [8, None, None]
+
+
 def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch):
     """Dynamic SD should create FULL decode candidates for every k in [1, K+1].
 

tests/v1/worker/test_gpu_rejection_sampler_chunking.py

@@ -10,6 +10,9 @@
 
 from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
 from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.spec_decode import (
+    rejection_sampler as rejection_sampler_module,
+)
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
     RejectionSampler,
     _iter_request_chunks,
@@ -26,6 +29,47 @@ def test_iter_request_chunks_preserves_request_boundaries():
     ]
 
 
+def test_rejection_sampler_restores_synthetic_draft_sentinels(monkeypatch):
+    input_ids = torch.tensor([10, 13, 0, 0], dtype=torch.int32)
+    input_batch = SimpleNamespace(
+        input_ids=input_ids,
+        logits_indices=torch.arange(4),
+        num_invalid_spec_tokens={"req": 2},
+        req_ids=["req"],
+        cu_num_logits_np=np.array([0, 4], dtype=np.int32),
+        positions=torch.arange(4),
+        idx_mapping_np=np.array([0], dtype=np.int32),
+        seq_lens=torch.tensor([4], dtype=torch.int32),
+        cu_num_logits=torch.tensor([0, 4], dtype=torch.int32),
+        idx_mapping=torch.tensor([0], dtype=torch.int32),
+    )
+    sampler = object.__new__(RejectionSampler)
+    sampler.enable_adaptive_verification = False
+    sampler.sampler = SimpleNamespace(
+        compute_nans=False,
+        sampling_states=SimpleNamespace(max_num_logprobs=lambda _indices: 0),
+        req_states=SimpleNamespace(
+            prefill_len=SimpleNamespace(gpu=torch.tensor([0], dtype=torch.int32))
+        ),
+    )
+
+    def fake_verify_in_chunks(
+        _self, _logits, _input_batch, _draft_logits, draft_sampled, *_args
+    ):
+        assert draft_sampled.tolist() == [10, 13, -1, -1]
+        return torch.tensor([[1]]), torch.tensor([1], dtype=torch.int32), None
+
+    sampler._verify_in_chunks = MethodType(fake_verify_in_chunks, sampler)
+    monkeypatch.setattr(
+        rejection_sampler_module,
+        "get_num_sampled_and_rejected",
+        lambda num_sampled, *_args: (num_sampled, torch.zeros_like(num_sampled)),
+    )
+
+    sampler(torch.zeros((4, 4)), input_batch)
+    assert input_ids.tolist() == [10, 13, 0, 0]
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
 @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))
 def test_chunked_scores_match_full_batch(logprobs_mode: str):

vllm/v1/core/sched/output.py

@@ -256,7 +256,7 @@ class SchedulerOutput:
     # need to perform grammar bitmask computation.
     pending_structured_output_tokens: bool = False
 
-    # Used for adjusting acceptance rate calculation.
+    # True reject-only draft padding, also excluded from acceptance accounting.
     num_invalid_spec_tokens: dict[str, int] | None = None
 
     # KV Cache Connector metadata.

vllm/v1/core/sched/scheduler.py

@@ -530,6 +530,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         encoder_compute_budget = self.max_num_encoder_input_tokens
         # Spec decode-related.
         scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+        num_invalid_spec_tokens: dict[str, int] = {}
         # Whether the running batch contains any prefill requests.
         prefill_scheduled = False
 
@@ -960,6 +961,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # Not for diffusion where draft tokens can't be padded.
                     if (
                         (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
+                        and not (spec is not None and spec.enable_adaptive_verification)
                         and self.num_sampled_tokens_per_step > 0
                         and num_new_tokens == 1
                         and (scheduled_running_reqs and not prefill_scheduled)
@@ -975,7 +977,6 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                             if padded_num_tokens > request_token_budget:
                                 # Prefer to not schedule than schedule un-padded.
                                 break
-                            num_new_tokens = padded_num_tokens
                             pad_spec_decode = True
 
                     threshold = self.scheduler_config.long_prefill_token_threshold
@@ -1031,6 +1032,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         # The request cannot be scheduled.
                         break
 
+                    if pad_spec_decode:
+                        assert num_new_tokens == 1
+                        num_new_tokens += self.num_spec_tokens
+
                 # During async KV load, no forward pass is run yet.
                 # Allocate speculative lookahead slots later to avoid
                 # mismatching local and remote block counts.
@@ -1167,6 +1172,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     scheduled_spec_decode_tokens[request_id] = [
                         -1
                     ] * self.num_spec_tokens
+                    num_invalid_spec_tokens[request_id] = self.num_spec_tokens
                 # Only track requests that will still be prefilling after this chunk.
                 if num_computed_tokens + num_new_tokens < request.num_tokens:
                     self._inflight_prefills.add(request)
@@ -1321,6 +1327,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
             kv_cache_block_copies=pending_kv_cache_block_copies,
             partial_tail_offloads=pending_partial_tail_offloads,
             num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
+            num_invalid_spec_tokens=num_invalid_spec_tokens or None,
             ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(),
         )
 
@@ -2309,7 +2316,7 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
     def update_draft_token_ids_in_output(
         self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput
     ) -> None:
-        num_invalid_spec_tokens: dict[str, int] = {}
+        num_invalid_spec_tokens = dict(scheduler_output.num_invalid_spec_tokens or {})
 
         sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
         for req_id, spec_token_ids in zip(
@@ -2337,11 +2344,13 @@ def update_draft_token_ids_in_output(
             num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
             if num_invalid_tokens:
                 spec_token_ids.extend([-1] * num_invalid_tokens)
-                num_invalid_spec_tokens[req_id] = num_invalid_tokens
+                num_invalid_spec_tokens[req_id] = max(
+                    num_invalid_spec_tokens.get(req_id, 0), num_invalid_tokens
+                )
 
             sched_spec_tokens[req_id] = spec_token_ids
 
-        scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens
+        scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens or None
 
     def get_request_counts(self) -> tuple[int, int]:
         """Returns (num_running_reqs, num_waiting_reqs)."""

vllm/v1/worker/gpu/input_batch.py

@@ -107,6 +107,9 @@ class InputBatch:
     # [num_reqs] per-request prompt length, only populated for R-SWA.
     prompt_lens: torch.Tensor | None
 
+    # Per-request reject-only draft suffix lengths.
+    num_invalid_spec_tokens: dict[str, int] | None = None
+
     # Longest query the batch may contain. Set when a cudagraph descriptor promises
     # a query length this batch's own split does not reach, so attention metadata
     # stays valid for every replay the graph serves.

vllm/v1/worker/gpu/model_runner.py

@@ -713,6 +713,12 @@ def _dummy_run(
         dummy_scheduler_output = SchedulerOutput.make_empty()
         dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
         dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens
+        if uniform_decode and self.speculative_config is not None:
+            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
+            dummy_scheduler_output.scheduled_spec_decode_tokens = {
+                req_id: [-1] * (n - num_bonus_tokens)
+                for req_id, n in num_scheduled_tokens.items()
+            }
 
         # Disable any use of KVConnector for dummy runs.
         self.kv_connector.set_disabled(True)
@@ -1333,6 +1339,7 @@ def prepare_inputs(
             cu_num_logits_np=cu_num_logits_np,
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
             prompt_lens=prompt_lens,
+            num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
             max_query_len=(
                 int(num_scheduled_tokens_upper_bound.max())
                 if adaptive_verification is not None
@@ -1547,6 +1554,24 @@ def execute_model(
             # cross-attention cache with dynamic encoder outputs.
             skip_compiled = True
 
+        if self.speculative_config is not None:
+            # A verifier FULL graph captures model-wide branch topology, not
+            # only its tensor shape. Do not let a same-shape short prefill
+            # select that graph.
+            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
+            draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+            is_target_verifier_batch = bool(draft_tokens) and all(
+                num_tokens > num_bonus_tokens
+                and len(draft_tokens.get(req_id, ())) == num_tokens - num_bonus_tokens
+                for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items()
+            )
+            if (
+                uniform_tok_count is not None
+                and uniform_tok_count > num_bonus_tokens
+                and not is_target_verifier_batch
+            ):
+                uniform_tok_count = None
+
         batch_desc, dp_sync = dispatch_cg_and_sync_dp(
             self.cudagraph_manager,
             num_reqs,

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

@@ -268,6 +268,21 @@ def __call__(
         num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
 
         draft_sampled = input_batch.input_ids[input_batch.logits_indices]
+        if not self.enable_adaptive_verification and (
+            invalid_counts := input_batch.num_invalid_spec_tokens
+        ):
+            invalid_indices: list[int] = []
+            for req_id, start, end in zip(
+                input_batch.req_ids,
+                input_batch.cu_num_logits_np[:-1],
+                input_batch.cu_num_logits_np[1:],
+                strict=True,
+            ):
+                num_invalid = invalid_counts.get(req_id, 0)
+                assert 0 <= num_invalid <= end - start
+                invalid_indices.extend(range(int(end) - num_invalid, int(end)))
+            if invalid_indices:
+                draft_sampled[invalid_indices] = -1
         pos = input_batch.positions[input_batch.logits_indices]
 
         max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(