[Core] Avoid mixed batch on spec-dec D-node via padding

tests/v1/core/test_scheduler.py

@@ -1279,6 +1279,101 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks():
     assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens
 
 
+def _model_output(scheduler, output, sampled):
+    """Feed `sampled` (per-request list) back to the scheduler."""
+    req_ids = list(output.num_scheduled_tokens.keys())
+    scheduler.update_from_output(
+        output,
+        ModelRunnerOutput(
+            req_ids=req_ids,
+            req_id_to_index={r: i for i, r in enumerate(req_ids)},
+            sampled_token_ids=sampled,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+
+
+def test_spec_decode_padding_first_decode_step():
+    """A request taking its first decode step (whole prompt already computed via
+    a prefix-cache hit) is padded with placeholder (-1) spec tokens so it enters
+    the worker with the same 1 + num_spec_tokens shape as the other speculative
+    decodes, keeping the batch uniform.
+    """
+    num_spec = 3
+    scheduler = create_scheduler(
+        num_speculative_tokens=num_spec,
+        enable_prefix_caching=True,
+        block_size=16,
+    )
+    # Two identical 33-token prompts: 2 full blocks (32 tokens) get cached, so a
+    # second identical request hits num_computed == num_prompt_tokens - 1.
+    r1, r2 = create_requests(
+        num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
+    )
+
+    # Drive r1 through prefill so its prompt blocks are cached, then give it real
+    # drafts so it is a running speculative decode (1 + num_spec shape).
+    scheduler.add_request(r1)
+    out = scheduler.schedule()
+    assert out.num_scheduled_tokens[r1.request_id] == 33
+    _model_output(scheduler, out, [[100]])
+    scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
+
+    # r2 arrives; its whole prompt is a prefix-cache hit -> first decode step.
+    scheduler.add_request(r2)
+    out = scheduler.schedule()
+
+    # r1 verifies its real drafts.
+    assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
+    # 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
+
+
+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.
+    """
+    num_spec = 3
+    scheduler = create_scheduler(
+        num_speculative_tokens=num_spec,
+        enable_prefix_caching=True,
+        block_size=16,
+        max_num_batched_tokens=64,
+    )
+    # r_warm + r_candidate share a prompt so r_candidate gets a full prefix hit.
+    r_warm, r_candidate = create_requests(
+        num_requests=2, num_tokens=33, same_prompt=True, max_tokens=1
+    )
+    # r_long has a different, long prompt that prefills over multiple chunks.
+    (r_long,) = create_requests(num_requests=1, num_tokens=100, max_tokens=16)
+
+    # Warm the prefix cache with r_warm's prompt (it finishes; blocks stay cached).
+    scheduler.add_request(r_warm)
+    out = scheduler.schedule()
+    assert out.num_scheduled_tokens[r_warm.request_id] == 33
+    _model_output(scheduler, out, [[100]])
+    assert r_warm.request_id in scheduler.finished_req_ids
+
+    # Start r_long; after one chunk it remains a prefill chunk in the running queue.
+    scheduler.add_request(r_long)
+    out = scheduler.schedule()
+    _model_output(scheduler, out, [[]])  # still prefilling, no sampled token
+    assert r_long.is_prefill_chunk
+
+    # r_candidate arrives (prefix-cache hit -> first decode step) alongside the
+    # in-flight prefill chunk.
+    scheduler.add_request(r_candidate)
+    out = scheduler.schedule()
+
+    # The batch has a prefill chunk, so r_candidate is NOT padded.
+    assert r_long.request_id in out.num_scheduled_tokens
+    assert out.num_scheduled_tokens[r_candidate.request_id] == 1
+    assert r_candidate.request_id not in out.scheduled_spec_decode_tokens
+
+
 def test_scheduler_stats_waiting_queues():
     """Test that scheduler stats correctly report waiting and skipped_waiting queues."""
     # Create scheduler with limited capacity so we can have waiting requests

tests/v1/sample/test_rejection_sampler.py

@@ -1150,3 +1150,54 @@ def test_synthetic_all_rejected(all_greedy: bool):
     for row in result:
         assert row[0] != PLACEHOLDER_TOKEN_ID
         assert (row[1:] == PLACEHOLDER_TOKEN_ID).all()
+
+
+def test_placeholder_draft_token_rejected_random(rejection_sampler):
+    """A placeholder/invalid draft id (PLACEHOLDER_TOKEN_ID == -1) must always
+    be rejected in non-greedy sampling, without indexing the probability
+    tensors by the invalid id.
+
+    The scheduler emits ``-1`` for spec positions that must not be verified
+    (structured-output grammar-invalid drafts, and decode-side spec padding).
+    Reproduces the bug where ``rejection_random_sample_kernel`` gathered
+    ``target_probs[..., -1]`` (out of bounds / wrong row), which could
+    spuriously "accept" the placeholder and emit ``-1`` as a real token.
+    """
+    vocab_size = 100
+    # Drafts: two valid tokens (forced-accepted) followed by a placeholder.
+    # The two valid drafts have target prob ~1, so they are accepted and we
+    # reach the placeholder position. The second valid draft id is the last
+    # vocab id, which is exactly the (wrong) cell the buggy kernel reads for
+    # the placeholder position -- so the bug deterministically accepts it.
+    spec_tokens = [[1, vocab_size - 1, PLACEHOLDER_TOKEN_ID]]
+    # Target argmax per draft position: 1, vocab-1, then 7 (recovered token).
+    output_tokens = [[1, vocab_size - 1, 7, 9]]  # 9 is the bonus token
+
+    temperature = torch.ones(1, dtype=torch.float32, device=DEVICE_TYPE)
+    metadata = create_sampling_metadata(
+        all_greedy=False,
+        temperature=temperature,
+        generators={0: torch.Generator(device=DEVICE_TYPE).manual_seed(0)},
+    )
+    logits = create_logits_tensor(output_tokens, vocab_size=vocab_size)
+    bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
+
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
+    output = rejection_sampler(
+        spec_decode_metadata,
+        draft_probs=None,
+        logits=logits,
+        sampling_metadata=metadata,
+    )
+    sampled = output.sampled_token_ids
+
+    # The two valid drafts are accepted.
+    assert sampled[0, 0].item() == 1
+    assert sampled[0, 1].item() == vocab_size - 1
+    # The placeholder draft must be rejected and replaced with a valid recovered
+    # token -- never emitted as -1, and never out of range.
+    recovered = sampled[0, 2].item()
+    assert 0 <= recovered < vocab_size, (
+        f"placeholder draft was not rejected: got {recovered}"
+    )

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -280,3 +280,40 @@ def test_synthetic_rejection_sample(
             f"Step {i}: observed rate {observed_rate:.4f} deviates from "
             f"expected rate {expected_rate:.4f} by more than {deviation_tol}."
         )
+
+
+def test_placeholder_draft_token_rejected():
+    """A placeholder/invalid draft id (-1) must always be rejected without
+    reading the logit tensors out of bounds, for any sampling method.
+
+    The scheduler emits -1 for spec positions that must not be verified
+    (structured-output grammar-invalid drafts, and decode-side spec padding).
+    Reproduces the bug where the random branch gathered target/draft logits at
+    index -1 (out of bounds), which could spuriously accept the placeholder.
+    """
+    torch.manual_seed(0)
+    device = "cuda"
+    num_trials = 64
+    K = 1
+    temperature = 0.6
+
+    target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature
+    draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature
+
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        K,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+    # Replace every draft token with the placeholder id (-1).
+    inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = -1
+
+    sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
+
+    # The placeholder draft is rejected, so exactly one (recovered) token is
+    # emitted per request -- never the placeholder, never out of range.
+    assert torch.equal(num_sampled, torch.ones_like(num_sampled))
+    recovered = sampled[:, 0]
+    assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()

vllm/v1/core/sched/scheduler.py

@@ -327,6 +327,48 @@ def __init__(
         # async KV loads). Their remaining-block reservation gates async loads.
         self._inflight_prefills: set[Request] = set()
 
+    def _pad_spec_decode(self, request: Request, num_computed_tokens: int) -> bool:
+        """Pad a single (first) decode step with placeholder spec tokens (-1),
+        to ensure uniform decode batches that use full cudagraphs.
+        """
+        pad_spec_decode = (
+            (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
+            # Exactly one real token to compute this step. Use num_tokens
+            # rather than num_prompt_tokens to also cover preemption resumes.
+            and num_computed_tokens == request.num_tokens - 1
+            and request.num_tokens
+            + self.num_spec_tokens
+            + self.num_sampled_tokens_per_step
+            <= self.max_model_len
+        )
+        if pad_spec_decode and not request.spec_token_ids:
+            # Pad with the placeholder draft id (-1).
+            request.spec_token_ids = [-1] * self.num_spec_tokens
+
+        return pad_spec_decode
+
+    @staticmethod
+    def _schedule_spec_tokens(
+        request: Request,
+        num_new_tokens: int,
+        scheduled_spec_decode_tokens: dict[str, list[int]],
+    ) -> None:
+        """Record this step's draft tokens for `request` and clear them."""
+        num_scheduled_spec_tokens = (
+            num_new_tokens
+            + request.num_computed_tokens
+            - (request.num_tokens + request.num_output_placeholders)
+        )
+        if num_scheduled_spec_tokens > 0:
+            spec_token_ids = request.spec_token_ids
+            if len(spec_token_ids) > num_scheduled_spec_tokens:
+                spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens]
+            scheduled_spec_decode_tokens[request.request_id] = spec_token_ids
+
+        # New spec tokens will be set in `update_draft_token_ids` before the
+        # next step when applicable.
+        request.spec_token_ids = []
+
     def _mamba_block_aligned_split(
         self,
         request: Request,
@@ -415,6 +457,8 @@ 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]] = {}
+        # Whether the running batch contains any prefill requests.
+        prefill_scheduled = False
 
         # For logging.
         scheduled_timestamp = time.monotonic()
@@ -573,6 +617,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
 
             # Schedule the request.
             scheduled_running_reqs.append(request)
+            prefill_scheduled |= request.is_prefill_chunk
             request_id = request.request_id
             req_to_new_blocks[request_id] = new_blocks
             num_scheduled_tokens[request_id] = num_new_tokens
@@ -581,21 +626,9 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
 
             # Speculative decode related.
             if request.spec_token_ids:
-                num_scheduled_spec_tokens = (
-                    num_new_tokens
-                    + request.num_computed_tokens
-                    - request.num_tokens
-                    - request.num_output_placeholders
+                self._schedule_spec_tokens(
+                    request, num_new_tokens, scheduled_spec_decode_tokens
                 )
-                if num_scheduled_spec_tokens > 0:
-                    spec_token_ids = request.spec_token_ids
-                    if len(spec_token_ids) > num_scheduled_spec_tokens:
-                        spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens]
-                    scheduled_spec_decode_tokens[request.request_id] = spec_token_ids
-
-                # New spec tokens will be set in `update_draft_token_ids` before the
-                # next step when applicable.
-                request.spec_token_ids = []
 
             # Encoder-related.
             if encoder_inputs_to_schedule:
@@ -779,21 +812,30 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 external_load_encoder_input = []
                 new_encoder_compute_budget = encoder_compute_budget
 
+                is_prefill = num_computed_tokens < request.num_tokens - 1
                 if load_kv_async:
                     # KVTransfer: loading remote KV, do not allocate for new work.
                     assert num_external_computed_tokens > 0
                     num_new_tokens = 0
-                elif defer_prefills and request.num_computed_tokens == 0:
-                    # DP prefill balancing: async KV loads (the branch above) are
-                    # allowed to start even on throttled steps, but committing new
-                    # prefill compute is deferred to a cadence-aligned step.
+                elif defer_prefills and is_prefill:
+                    # DP prefill balancing: defer committing new prefill compute to
+                    # a cadence-aligned step (exclude full kv cache hits).
                     break
                 else:
                     # Number of tokens to be scheduled.
                     # We use `request.num_tokens` instead of
                     # `request.num_prompt_tokens` to consider the resumed
                     # requests, which have output tokens.
                     num_new_tokens = request.num_tokens - num_computed_tokens
+
+                    if scheduled_running_reqs and not prefill_scheduled:  # noqa: SIM102
+                        # Pad new decode requests to uniform spec decoding size to
+                        # preserve full cudagraph for this step.
+                        if self._pad_spec_decode(request, num_computed_tokens):
+                            num_new_tokens += len(request.spec_token_ids)
+                            if num_new_tokens > token_budget:
+                                break
+
                     threshold = self.scheduler_config.long_prefill_token_threshold
                     if 0 < threshold < num_new_tokens:
                         num_new_tokens = threshold
@@ -958,6 +1000,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 token_budget -= num_new_tokens
                 request.status = RequestStatus.RUNNING
                 request.num_computed_tokens = num_computed_tokens
+                if request.spec_token_ids:
+                    self._schedule_spec_tokens(
+                        request, num_new_tokens, scheduled_spec_decode_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)

vllm/v1/sample/rejection_sampler.py

@@ -748,7 +748,7 @@ def rejection_greedy_sample_kernel(
             if SYNTHETIC_MODE:
                 uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
                 rate = tl.load(synthetic_conditional_rates_ptr + pos)
-                accepted = uniform_prob < rate
+                accepted = (uniform_prob < rate) and draft_token_id >= 0
                 token_id = draft_token_id if accepted else target_argmax_id
                 rejected = not accepted
             else:
@@ -805,7 +805,9 @@ def rejection_random_sample_kernel(
         if not rejected:
             draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
             uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
-            if SYNTHETIC_MODE:
+            if draft_token_id < 0:
+                accepted = False
+            elif SYNTHETIC_MODE:
                 rate = tl.load(synthetic_conditional_rates_ptr + pos)
                 accepted = uniform_prob < rate
             else:

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

@@ -245,16 +245,21 @@ def _rejection_kernel(
                     pos = tl.load(pos_ptr + logit_idx)
                     u = tl_rand64(seed, pos, includes_zero=False)
                     rate = tl.load(synthetic_conditional_rates_ptr + i)
-                    accepted &= u < rate
+                    accepted &= (u < rate) & (draft_sampled >= 0)
                 else:
                     accepted &= target_argmax == draft_sampled
                 tl.store(
                     sampled_ptr + req_idx * sampled_stride + i,
                     draft_sampled if accepted else target_argmax,
                 )
             else:
+                is_valid_draft = draft_sampled >= 0
                 target_logit = tl.load(
-                    target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
+                    target_logits_ptr
+                    + logit_idx * target_logits_stride
+                    + draft_sampled,
+                    mask=is_valid_draft,
+                    other=0.0,
                 ).to(tl.float32)
                 target_lse = _compute_global_lse(
                     target_local_max_ptr,
@@ -273,7 +278,9 @@ def _rejection_kernel(
                         draft_logits_ptr
                         + req_state_idx * draft_logits_stride_0
                         + i * draft_logits_stride_1
-                        + draft_sampled
+                        + draft_sampled,
+                        mask=is_valid_draft,
+                        other=0.0,
                     ).to(tl.float32)
                     draft_lse = _compute_global_lse(
                         draft_local_max_ptr,
@@ -296,6 +303,7 @@ def _rejection_kernel(
                     # Probability ratio test: p(x) > u * q(x)
                     # Equivalent log form: log_p(x) > log(u) + log_q(x)
                     accepted &= target_log_prob > tl.log(u) + draft_log_prob
+                accepted &= is_valid_draft
                 tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
             rejected_step += accepted
     tl.store(rejected_steps_ptr + req_idx, rejected_step)