From f8476f03d74a92137c50f52b4196936781c7c1e1 Mon Sep 17 00:00:00 2001
From: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Date: Sat, 20 Jun 2026 11:29:30 +0800
Subject: [PATCH 1/9] Use speculative padding for KV transfer decode bootstrap

Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
---
 .../kv_connector/v1/nixl/base_scheduler.py    |  19 +++
 .../kv_connector/v1/nixl/pull_scheduler.py    |  12 +-
 vllm/v1/core/sched/scheduler.py               | 111 +++++++++++++++++-
 3 files changed, 138 insertions(+), 4 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
index cba81cadd84c..6d551e413266 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
@@ -243,6 +243,25 @@ def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
             ]
         )
 
+    def _clip_blocks_to_num_tokens(
+        self,
+        block_ids: BlockIds,
+        num_tokens: int,
+    ) -> BlockIds:
+        """Drop blocks that cannot contain tokens in ``[0, num_tokens)``."""
+        if len(block_ids) == 0:
+            return block_ids
+
+        clipped: list[list[int]] = []
+        for group, group_spec in zip(block_ids, self.kv_cache_config.kv_cache_groups):
+            if isinstance(group_spec.kv_cache_spec, MambaSpec):
+                clipped.append(list(group))
+                continue
+
+            num_blocks = cdiv(num_tokens, group_spec.kv_cache_spec.block_size)
+            clipped.append(list(group[:num_blocks]))
+        return tuple(clipped)
+
     def set_xfer_handshake_metadata(
         self, metadata: dict[int, KVConnectorHandshakeMetadata]
     ) -> None:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
index f13e21605667..6bd40e6f4bbd 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
@@ -254,14 +254,22 @@ def request_finished(
             self._reqs_need_send[request.request_id] = (
                 time.perf_counter() + request_kv_blocks_ttl
             )
+            remote_num_tokens = request.num_computed_tokens
+            if is_p_node:
+                # The P-side bootstrap request generates one token only to finish
+                # the request and expose prompt KV to D. That generated token is
+                # not part of the D request, so do not expose its KV block.
+                remote_num_tokens = min(remote_num_tokens, request.num_prompt_tokens)
+                block_ids = self._clip_blocks_to_num_tokens(
+                    block_ids, remote_num_tokens
+                )
+
             # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
             # trimming down after allocating for the whole sequence length. Empty
             # blocks are always at the start of the list.
             # Here we "unpad" blocks to send the actual remote blocks to be read.
             block_ids = self.get_sw_clipped_blocks(block_ids)
 
-            remote_num_tokens = request.num_computed_tokens
-
         return delay_free_blocks, dict(
             do_remote_prefill=is_p_node,
             do_remote_decode=is_d_node,
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 25ccf79bc3a7..57b5926c1ecd 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -227,6 +227,13 @@ def __init__(
         speculative_config = vllm_config.speculative_config
         self.use_eagle = False
         self.num_spec_tokens = vllm_config.num_speculative_tokens
+        self.enable_kv_transfer_spec_decode_padding = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "enable_speculative_padding", False
+            )
+            if vllm_config.kv_transfer_config is not None
+            else False
+        )
         self.num_lookahead_tokens = 0
         self.dynamic_sd_lookup: list[int] | None = None
         if speculative_config is not None:
@@ -325,6 +332,57 @@ def __init__(
         # In-flight requests still prefilling (prefill chunks + in-progress
         # async KV loads). Their remaining-block reservation gates async loads.
         self._inflight_prefills: set[Request] = set()
+        self._kv_transfer_spec_decode_padding_req_ids: set[str] = set()
+
+    def _is_kv_transfer_spec_decode_padding_candidate(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> bool:
+        params = request.kv_transfer_params
+        is_remote_prefill = (
+            request.request_id in self._kv_transfer_spec_decode_padding_req_ids
+            or bool(params and params.get("do_remote_prefill"))
+        )
+        speculative_config = self.vllm_config.speculative_config
+        if (
+            not self.enable_kv_transfer_spec_decode_padding
+            or self.num_spec_tokens <= 0
+            or self.dynamic_sd_lookup is not None
+            or (
+                speculative_config is not None
+                and speculative_config.draft_sample_method != "greedy"
+            )
+            or not is_remote_prefill
+            or request.num_output_tokens > 0
+            or num_computed_tokens < request.num_prompt_tokens - 1
+            or self.structured_output_manager.should_advance(request)
+        ):
+            return False
+        return True
+
+    def _maybe_kv_transfer_spec_decode_padding(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> bool:
+        if not self._is_kv_transfer_spec_decode_padding_candidate(
+            request, num_computed_tokens
+        ):
+            return False
+
+        if request.spec_token_ids:
+            return True
+
+        prompt_token_ids = request.prompt_token_ids
+        if not prompt_token_ids:
+            return False
+        pad_token_id = int(prompt_token_ids[0])
+        if pad_token_id < 0:
+            return False
+
+        request.spec_token_ids = [pad_token_id] * self.num_spec_tokens
+        return True
 
     def _mamba_block_aligned_split(
         self,
@@ -777,6 +835,24 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 encoder_inputs_to_schedule = None
                 external_load_encoder_input = []
                 new_encoder_compute_budget = encoder_compute_budget
+                should_kv_transfer_spec_decode_padding = (
+                    self._is_kv_transfer_spec_decode_padding_candidate(
+                        request, num_computed_tokens
+                    )
+                )
+                if (
+                    not load_kv_async
+                    and should_kv_transfer_spec_decode_padding
+                    and num_computed_tokens >= request.num_prompt_tokens
+                ):
+                    num_computed_tokens = max(request.num_prompt_tokens - 1, 0)
+                    request.num_computed_tokens = num_computed_tokens
+                is_kv_transfer_spec_decode_padding = (
+                    not load_kv_async
+                    and self._maybe_kv_transfer_spec_decode_padding(
+                        request, num_computed_tokens
+                    )
+                )
 
                 if load_kv_async:
                     # KVTransfer: loading remote KV, do not allocate for new work.
@@ -792,9 +868,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # 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
+                    num_tokens_to_schedule = (
+                        request.num_tokens_with_spec
+                        if is_kv_transfer_spec_decode_padding
+                        else request.num_tokens
+                    )
+                    num_new_tokens = num_tokens_to_schedule - num_computed_tokens
                     threshold = self.scheduler_config.long_prefill_token_threshold
-                    if 0 < threshold < num_new_tokens:
+                    if (
+                        not is_kv_transfer_spec_decode_padding
+                        and 0 < threshold < num_new_tokens
+                    ):
                         num_new_tokens = threshold
 
                     # chunked prefill has to be enabled explicitly to allow
@@ -807,6 +891,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         # we can stop the scheduling here.
                         break
 
+                    if (
+                        is_kv_transfer_spec_decode_padding
+                        and num_new_tokens > token_budget
+                    ):
+                        break
                     num_new_tokens = min(num_new_tokens, token_budget)
                     assert num_new_tokens > 0
 
@@ -933,6 +1022,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # _update_waiting_for_remote_kv will then cache
                     # only the successfully loaded tokens.
                     request.num_computed_tokens = num_computed_tokens
+                    if should_kv_transfer_spec_decode_padding:
+                        self._kv_transfer_spec_decode_padding_req_ids.add(request_id)
                     self._inflight_prefills.add(request)
                     continue
 
@@ -957,6 +1048,20 @@ 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 is_kv_transfer_spec_decode_padding and request.spec_token_ids:
+                    num_scheduled_spec_tokens = max(
+                        0,
+                        num_new_tokens
+                        + num_computed_tokens
+                        - request.num_tokens
+                        - request.num_output_placeholders,
+                    )
+                    if num_scheduled_spec_tokens:
+                        scheduled_spec_decode_tokens[request_id] = (
+                            request.spec_token_ids[:num_scheduled_spec_tokens]
+                        )
+                    request.spec_token_ids = []
+                    self._kv_transfer_spec_decode_padding_req_ids.discard(request_id)
                 # 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)
@@ -2049,6 +2154,7 @@ def _free_request(
         assert request.is_finished()
 
         self._inflight_prefills.discard(request)
+        self._kv_transfer_spec_decode_padding_req_ids.discard(request.request_id)
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
         self.encoder_cache_manager.free(request)
         request_id = request.request_id
@@ -2358,6 +2464,7 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
         assert self.connector is not None
 
         if request.request_id in self.failed_recving_kv_req_ids:
+            self._kv_transfer_spec_decode_padding_req_ids.discard(request.request_id)
             # Request had KV load failures; num_computed_tokens was already
             # updated in _update_requests_with_invalid_blocks
             if request.num_computed_tokens:

From edf1df49ba3b2e820de8ce78bf67b8a485cdb008 Mon Sep 17 00:00:00 2001
From: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Date: Sun, 21 Jun 2026 14:35:52 +0800
Subject: [PATCH 2/9] cleanup

Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
---
 .../kv_connector/v1/nixl/base_scheduler.py    | 19 -------------------
 .../kv_connector/v1/nixl/pull_scheduler.py    | 11 ++---------
 2 files changed, 2 insertions(+), 28 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
index 6d551e413266..cba81cadd84c 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
@@ -243,25 +243,6 @@ def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
             ]
         )
 
-    def _clip_blocks_to_num_tokens(
-        self,
-        block_ids: BlockIds,
-        num_tokens: int,
-    ) -> BlockIds:
-        """Drop blocks that cannot contain tokens in ``[0, num_tokens)``."""
-        if len(block_ids) == 0:
-            return block_ids
-
-        clipped: list[list[int]] = []
-        for group, group_spec in zip(block_ids, self.kv_cache_config.kv_cache_groups):
-            if isinstance(group_spec.kv_cache_spec, MambaSpec):
-                clipped.append(list(group))
-                continue
-
-            num_blocks = cdiv(num_tokens, group_spec.kv_cache_spec.block_size)
-            clipped.append(list(group[:num_blocks]))
-        return tuple(clipped)
-
     def set_xfer_handshake_metadata(
         self, metadata: dict[int, KVConnectorHandshakeMetadata]
     ) -> None:
diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
index 6bd40e6f4bbd..909a7cafe88c 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
@@ -254,15 +254,6 @@ def request_finished(
             self._reqs_need_send[request.request_id] = (
                 time.perf_counter() + request_kv_blocks_ttl
             )
-            remote_num_tokens = request.num_computed_tokens
-            if is_p_node:
-                # The P-side bootstrap request generates one token only to finish
-                # the request and expose prompt KV to D. That generated token is
-                # not part of the D request, so do not expose its KV block.
-                remote_num_tokens = min(remote_num_tokens, request.num_prompt_tokens)
-                block_ids = self._clip_blocks_to_num_tokens(
-                    block_ids, remote_num_tokens
-                )
 
             # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
             # trimming down after allocating for the whole sequence length. Empty
@@ -270,6 +261,8 @@ def request_finished(
             # Here we "unpad" blocks to send the actual remote blocks to be read.
             block_ids = self.get_sw_clipped_blocks(block_ids)
 
+            remote_num_tokens = request.num_computed_tokens
+
         return delay_free_blocks, dict(
             do_remote_prefill=is_p_node,
             do_remote_decode=is_d_node,

From ed10cd6dd38f006c57380f9897b707970ed279e1 Mon Sep 17 00:00:00 2001
From: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
Date: Sun, 21 Jun 2026 19:43:07 +0800
Subject: [PATCH 3/9] up

Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
---
 .../kv_connector/v1/nixl/pull_scheduler.py        |  1 -
 vllm/v1/core/sched/scheduler.py                   | 15 +++++++--------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
index 909a7cafe88c..f13e21605667 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
@@ -254,7 +254,6 @@ def request_finished(
             self._reqs_need_send[request.request_id] = (
                 time.perf_counter() + request_kv_blocks_ttl
             )
-
             # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
             # trimming down after allocating for the whole sequence length. Empty
             # blocks are always at the start of the list.
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index d4ecdcecdb0a..bfcdca482086 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -357,6 +357,12 @@ def _is_kv_transfer_spec_decode_padding_candidate(
             or not is_remote_prefill
             or request.num_output_tokens > 0
             or num_computed_tokens < request.num_prompt_tokens - 1
+            or (
+                request.num_prompt_tokens
+                + self.num_spec_tokens
+                + self.num_sampled_tokens_per_step
+                > self.max_model_len
+            )
             or self.structured_output_manager.should_advance(request)
         ):
             return False
@@ -375,14 +381,7 @@ def _maybe_kv_transfer_spec_decode_padding(
         if request.spec_token_ids:
             return True
 
-        prompt_token_ids = request.prompt_token_ids
-        if not prompt_token_ids:
-            return False
-        pad_token_id = int(prompt_token_ids[0])
-        if pad_token_id < 0:
-            return False
-
-        request.spec_token_ids = [pad_token_id] * self.num_spec_tokens
+        request.spec_token_ids = [int(request.all_token_ids[0])] * self.num_spec_tokens
         return True
 
     def _mamba_block_aligned_split(

From 2ad774a8e7206591b77bb22d4fce13da8d50c03d Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Mon, 22 Jun 2026 11:55:47 -0700
Subject: [PATCH 4/9] fix/simplify

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 tests/v1/core/test_scheduler.py               |  95 +++++++++
 tests/v1/sample/test_rejection_sampler.py     |  51 +++++
 .../test_rejection_sampler_utils.py           |  37 ++++
 vllm/v1/core/sched/scheduler.py               | 180 ++++++------------
 vllm/v1/sample/rejection_sampler.py           |   6 +-
 .../spec_decode/rejection_sampler_utils.py    |  14 +-
 6 files changed, 258 insertions(+), 125 deletions(-)

diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 6eb5ff5cc44a..e7bea55197ef 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/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
diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py
index 10c4d448f7f6..58aac618a3f6 100644
--- a/tests/v1/sample/test_rejection_sampler.py
+++ b/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}"
+    )
diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py
index 28f6044de877..d294de96b9d5 100644
--- a/tests/v1/spec_decode/test_rejection_sampler_utils.py
+++ b/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()
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index bfcdca482086..9c2fa701f0c6 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -228,13 +228,6 @@ def __init__(
         speculative_config = vllm_config.speculative_config
         self.use_eagle = False
         self.num_spec_tokens = vllm_config.num_speculative_tokens
-        self.enable_kv_transfer_spec_decode_padding = (
-            vllm_config.kv_transfer_config.get_from_extra_config(
-                "enable_speculative_padding", False
-            )
-            if vllm_config.kv_transfer_config is not None
-            else False
-        )
         self.num_lookahead_tokens = 0
         self.dynamic_sd_lookup: list[int] | None = None
         if speculative_config is not None:
@@ -333,56 +326,48 @@ def __init__(
         # In-flight requests still prefilling (prefill chunks + in-progress
         # async KV loads). Their remaining-block reservation gates async loads.
         self._inflight_prefills: set[Request] = set()
-        self._kv_transfer_spec_decode_padding_req_ids: set[str] = set()
 
-    def _is_kv_transfer_spec_decode_padding_candidate(
-        self,
-        request: Request,
-        num_computed_tokens: int,
-    ) -> bool:
-        params = request.kv_transfer_params
-        is_remote_prefill = (
-            request.request_id in self._kv_transfer_spec_decode_padding_req_ids
-            or bool(params and params.get("do_remote_prefill"))
+    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
         )
-        speculative_config = self.vllm_config.speculative_config
-        if (
-            not self.enable_kv_transfer_spec_decode_padding
-            or self.num_spec_tokens <= 0
-            or self.dynamic_sd_lookup is not None
-            or (
-                speculative_config is not None
-                and speculative_config.draft_sample_method != "greedy"
-            )
-            or not is_remote_prefill
-            or request.num_output_tokens > 0
-            or num_computed_tokens < request.num_prompt_tokens - 1
-            or (
-                request.num_prompt_tokens
-                + self.num_spec_tokens
-                + self.num_sampled_tokens_per_step
-                > self.max_model_len
-            )
-            or self.structured_output_manager.should_advance(request)
-        ):
-            return False
-        return True
+        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
 
-    def _maybe_kv_transfer_spec_decode_padding(
-        self,
-        request: Request,
-        num_computed_tokens: int,
-    ) -> bool:
-        if not self._is_kv_transfer_spec_decode_padding_candidate(
-            request, num_computed_tokens
-        ):
-            return False
+        return pad_spec_decode
 
-        if request.spec_token_ids:
-            return True
+    @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
 
-        request.spec_token_ids = [int(request.all_token_ids[0])] * self.num_spec_tokens
-        return True
+        # 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,
@@ -472,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()
@@ -630,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
@@ -638,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:
@@ -835,50 +811,33 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 encoder_inputs_to_schedule = None
                 external_load_encoder_input = []
                 new_encoder_compute_budget = encoder_compute_budget
-                should_kv_transfer_spec_decode_padding = (
-                    self._is_kv_transfer_spec_decode_padding_candidate(
-                        request, num_computed_tokens
-                    )
-                )
-                if (
-                    not load_kv_async
-                    and should_kv_transfer_spec_decode_padding
-                    and num_computed_tokens >= request.num_prompt_tokens
-                ):
-                    num_computed_tokens = max(request.num_prompt_tokens - 1, 0)
-                    request.num_computed_tokens = num_computed_tokens
-                is_kv_transfer_spec_decode_padding = (
-                    not load_kv_async
-                    and self._maybe_kv_transfer_spec_decode_padding(
-                        request, num_computed_tokens
-                    )
-                )
 
+                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_tokens_to_schedule = (
-                        request.num_tokens_with_spec
-                        if is_kv_transfer_spec_decode_padding
-                        else request.num_tokens
-                    )
-                    num_new_tokens = num_tokens_to_schedule - num_computed_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 (
-                        not is_kv_transfer_spec_decode_padding
-                        and 0 < threshold < num_new_tokens
-                    ):
+                    if 0 < threshold < num_new_tokens:
                         num_new_tokens = threshold
 
                     # chunked prefill has to be enabled explicitly to allow
@@ -891,11 +850,6 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         # we can stop the scheduling here.
                         break
 
-                    if (
-                        is_kv_transfer_spec_decode_padding
-                        and num_new_tokens > token_budget
-                    ):
-                        break
                     num_new_tokens = min(num_new_tokens, token_budget)
                     assert num_new_tokens > 0
 
@@ -1022,8 +976,6 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # _update_waiting_for_remote_kv will then cache
                     # only the successfully loaded tokens.
                     request.num_computed_tokens = num_computed_tokens
-                    if should_kv_transfer_spec_decode_padding:
-                        self._kv_transfer_spec_decode_padding_req_ids.add(request_id)
                     self._inflight_prefills.add(request)
                     continue
 
@@ -1048,20 +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 is_kv_transfer_spec_decode_padding and request.spec_token_ids:
-                    num_scheduled_spec_tokens = max(
-                        0,
-                        num_new_tokens
-                        + num_computed_tokens
-                        - request.num_tokens
-                        - request.num_output_placeholders,
+                if request.spec_token_ids:
+                    self._schedule_spec_tokens(
+                        request, num_new_tokens, scheduled_spec_decode_tokens
                     )
-                    if num_scheduled_spec_tokens:
-                        scheduled_spec_decode_tokens[request_id] = (
-                            request.spec_token_ids[:num_scheduled_spec_tokens]
-                        )
-                    request.spec_token_ids = []
-                    self._kv_transfer_spec_decode_padding_req_ids.discard(request_id)
                 # 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)
@@ -2154,7 +2096,6 @@ def _free_request(
         assert request.is_finished()
 
         self._inflight_prefills.discard(request)
-        self._kv_transfer_spec_decode_padding_req_ids.discard(request.request_id)
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
         self.encoder_cache_manager.free(request)
         request_id = request.request_id
@@ -2464,7 +2405,6 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
         assert self.connector is not None
 
         if request.request_id in self.failed_recving_kv_req_ids:
-            self._kv_transfer_spec_decode_padding_req_ids.discard(request.request_id)
             # Request had KV load failures; num_computed_tokens was already
             # updated in _update_requests_with_invalid_blocks
             if request.num_computed_tokens:
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index 8b4d8c9dce70..cf92d1b3acb3 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/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:
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
index 0cfbdf4182be..0bff9de18d23 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -245,7 +245,7 @@ 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(
@@ -253,8 +253,13 @@ def _rejection_kernel(
                     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)

From 5c3f2e5dbf95121663436e14050cbec639e4ce44 Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Tue, 23 Jun 2026 10:16:49 -0700
Subject: [PATCH 5/9] fix DP prefill throttling test

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 tests/v1/core/test_scheduler.py | 55 +++++++++++++++++++++++----------
 1 file changed, 39 insertions(+), 16 deletions(-)

diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index e7bea55197ef..4063c0a03cb9 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -292,34 +292,32 @@ def test_schedule_prefills_gating(has_running: bool):
     assert any(r.req_id == "new0" for r in output.scheduled_new_reqs)
 
 
-def test_throttle_prefills_excludes_remote_kv_resume():
-    """A request resuming after a completed async KV load (num_computed_tokens
-    > 0, e.g. the decode side of P/D disaggregation) must NOT be throttled by
-    the DP prefill cadence: only fresh prefills are deferred. Otherwise the
-    resumed request's first (single-token) step would be needlessly delayed.
+def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int):
+    """Drive a remote-KV request `r2` to the resume point (async load complete)
+    while another request `r1` is already decoding, so the step is throttle-
+    eligible. Returns the scheduler. The connector matches `matched_tokens` of
+    `r2`'s prompt; the rest (if any) is local prefill.
     """
     from tests.v1.kv_connector.unit.utils import create_model_runner_output
 
     BLOCK_SIZE = 16
-    NUM_MATCHED = BLOCK_SIZE * 2
     scheduler = create_scheduler(
         enable_prefix_caching=True,
-        use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED, is_async=True),
+        use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
         block_size=BLOCK_SIZE,
     )
-
-    # Two remote-KV requests with distinct prompts (so r2 gets no local prefix
-    # cache hit from r1, only the connector's external async load).
+    # Distinct prompts so r2 gets no local prefix cache hit from r1, only the
+    # connector's external async load.
     r1, r2 = create_requests(
         num_requests=2,
-        num_tokens=NUM_MATCHED * 2,
+        num_tokens=num_prompt_tokens,
         max_tokens=20,
         block_size=BLOCK_SIZE,
         req_ids=["r1", "r2"],
     )
 
-    # r1: drive through its async KV load and into the running (decode) state,
-    # so that self.running is non-empty for the assertion below.
+    # r1: drive through its async KV load into the running (decode) state, so
+    # self.running is non-empty (which makes the next step throttle-eligible).
     scheduler.add_request(r1)
     _step_until_kv_transfer_finished(scheduler, ["r1"])
     output = scheduler.schedule()  # promote + schedule r1
@@ -337,15 +335,40 @@ def test_throttle_prefills_excludes_remote_kv_resume():
         output, create_model_runner_output([r1], finished_recving={"r2"})
     )
     assert "r2" in scheduler.finished_recving_kv_req_ids
+    return scheduler
+
+
+def test_throttle_prefills_excludes_fully_transferred_remote_kv():
+    """A remote-KV resume whose whole prompt was transferred (no local prefill
+    left, e.g. the decode side of P/D disaggregation) must NOT be throttled by
+    the DP prefill cadence -- its single-token step has no prefill compute to
+    defer, so delaying it would be pointless.
+    """
+    block_size = 16
+    num_prompt = block_size * 2
+    # Fully matched: the whole prompt is loaded remotely.
+    scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt)
 
-    # Throttle prefills. r2's load is complete, so it must be promoted and
-    # scheduled (a resume, not a fresh prefill) even though the running decode
-    # (r1) would otherwise make this a throttled step.
     output = scheduler.schedule(throttle_prefills=True)
     assert "r2" in output.num_scheduled_tokens
     assert "r1" in output.num_scheduled_tokens
 
 
+def test_throttle_prefills_defers_remote_kv_resume_with_local_prefill():
+    """A remote-KV resume with local prefill still to compute (the connector
+    only matched part of the prompt) IS throttled by the DP prefill cadence,
+    like any other request doing local prefill compute this step.
+    """
+    block_size = 16
+    num_prompt = block_size * 4
+    # Half matched: the remaining half is local prefill compute.
+    scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt // 2)
+
+    output = scheduler.schedule(throttle_prefills=True)
+    assert "r2" not in output.num_scheduled_tokens  # deferred (has local prefill)
+    assert "r1" in output.num_scheduled_tokens
+
+
 def test_throttle_defers_inflight_prefill_chunk():
     """DP prefill balancing throttles ALL prefill compute on a throttled step,
     not just new admissions: an in-progress (chunked) prefill already in the

From 6ec4d8b11249f8707e032b63cccff9522ca4063e Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Tue, 23 Jun 2026 11:43:37 -0700
Subject: [PATCH 6/9] Revert DP prefill throttling changes; add comments

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 tests/v1/core/test_scheduler.py               | 55 ++++++-------------
 vllm/v1/core/sched/scheduler.py               |  8 +--
 vllm/v1/sample/rejection_sampler.py           |  2 +
 .../spec_decode/rejection_sampler_utils.py    |  2 +
 4 files changed, 24 insertions(+), 43 deletions(-)

diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 4063c0a03cb9..e7bea55197ef 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -292,32 +292,34 @@ def test_schedule_prefills_gating(has_running: bool):
     assert any(r.req_id == "new0" for r in output.scheduled_new_reqs)
 
 
-def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int):
-    """Drive a remote-KV request `r2` to the resume point (async load complete)
-    while another request `r1` is already decoding, so the step is throttle-
-    eligible. Returns the scheduler. The connector matches `matched_tokens` of
-    `r2`'s prompt; the rest (if any) is local prefill.
+def test_throttle_prefills_excludes_remote_kv_resume():
+    """A request resuming after a completed async KV load (num_computed_tokens
+    > 0, e.g. the decode side of P/D disaggregation) must NOT be throttled by
+    the DP prefill cadence: only fresh prefills are deferred. Otherwise the
+    resumed request's first (single-token) step would be needlessly delayed.
     """
     from tests.v1.kv_connector.unit.utils import create_model_runner_output
 
     BLOCK_SIZE = 16
+    NUM_MATCHED = BLOCK_SIZE * 2
     scheduler = create_scheduler(
         enable_prefix_caching=True,
-        use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
+        use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED, is_async=True),
         block_size=BLOCK_SIZE,
     )
-    # Distinct prompts so r2 gets no local prefix cache hit from r1, only the
-    # connector's external async load.
+
+    # Two remote-KV requests with distinct prompts (so r2 gets no local prefix
+    # cache hit from r1, only the connector's external async load).
     r1, r2 = create_requests(
         num_requests=2,
-        num_tokens=num_prompt_tokens,
+        num_tokens=NUM_MATCHED * 2,
         max_tokens=20,
         block_size=BLOCK_SIZE,
         req_ids=["r1", "r2"],
     )
 
-    # r1: drive through its async KV load into the running (decode) state, so
-    # self.running is non-empty (which makes the next step throttle-eligible).
+    # r1: drive through its async KV load and into the running (decode) state,
+    # so that self.running is non-empty for the assertion below.
     scheduler.add_request(r1)
     _step_until_kv_transfer_finished(scheduler, ["r1"])
     output = scheduler.schedule()  # promote + schedule r1
@@ -335,40 +337,15 @@ def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int):
         output, create_model_runner_output([r1], finished_recving={"r2"})
     )
     assert "r2" in scheduler.finished_recving_kv_req_ids
-    return scheduler
-
-
-def test_throttle_prefills_excludes_fully_transferred_remote_kv():
-    """A remote-KV resume whose whole prompt was transferred (no local prefill
-    left, e.g. the decode side of P/D disaggregation) must NOT be throttled by
-    the DP prefill cadence -- its single-token step has no prefill compute to
-    defer, so delaying it would be pointless.
-    """
-    block_size = 16
-    num_prompt = block_size * 2
-    # Fully matched: the whole prompt is loaded remotely.
-    scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt)
 
+    # Throttle prefills. r2's load is complete, so it must be promoted and
+    # scheduled (a resume, not a fresh prefill) even though the running decode
+    # (r1) would otherwise make this a throttled step.
     output = scheduler.schedule(throttle_prefills=True)
     assert "r2" in output.num_scheduled_tokens
     assert "r1" in output.num_scheduled_tokens
 
 
-def test_throttle_prefills_defers_remote_kv_resume_with_local_prefill():
-    """A remote-KV resume with local prefill still to compute (the connector
-    only matched part of the prompt) IS throttled by the DP prefill cadence,
-    like any other request doing local prefill compute this step.
-    """
-    block_size = 16
-    num_prompt = block_size * 4
-    # Half matched: the remaining half is local prefill compute.
-    scheduler = _setup_remote_kv_resume(num_prompt, matched_tokens=num_prompt // 2)
-
-    output = scheduler.schedule(throttle_prefills=True)
-    assert "r2" not in output.num_scheduled_tokens  # deferred (has local prefill)
-    assert "r1" in output.num_scheduled_tokens
-
-
 def test_throttle_defers_inflight_prefill_chunk():
     """DP prefill balancing throttles ALL prefill compute on a throttled step,
     not just new admissions: an in-progress (chunked) prefill already in the
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 9c2fa701f0c6..f76a46d73073 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -812,14 +812,14 @@ 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 is_prefill:
-                    # DP prefill balancing: defer committing new prefill compute to
-                    # a cadence-aligned step (exclude full kv cache hits).
+                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.
                     break
                 else:
                     # Number of tokens to be scheduled.
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index cf92d1b3acb3..1324191be74c 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -748,6 +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)
+                # -1 is used for padded draft token ids that should be rejected.
                 accepted = (uniform_prob < rate) and draft_token_id >= 0
                 token_id = draft_token_id if accepted else target_argmax_id
                 rejected = not accepted
@@ -806,6 +807,7 @@ def rejection_random_sample_kernel(
             draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
             uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
             if draft_token_id < 0:
+                # -1 is used for padded draft token ids that should be rejected.
                 accepted = False
             elif SYNTHETIC_MODE:
                 rate = tl.load(synthetic_conditional_rates_ptr + pos)
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
index 0bff9de18d23..da3057aec6e7 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -245,6 +245,7 @@ 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)
+                    # -1 is used for padded draft token ids that should be rejected.
                     accepted &= (u < rate) & (draft_sampled >= 0)
                 else:
                     accepted &= target_argmax == draft_sampled
@@ -253,6 +254,7 @@ def _rejection_kernel(
                     draft_sampled if accepted else target_argmax,
                 )
             else:
+                # -1 is used for padded draft token ids that should be rejected.
                 is_valid_draft = draft_sampled >= 0
                 target_logit = tl.load(
                     target_logits_ptr

From 65ba17b4e3ebd183409c1af532a658c17957a57c Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Tue, 23 Jun 2026 19:01:50 -0700
Subject: [PATCH 7/9] simplify, remove rejection sampler fixes

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 tests/v1/sample/test_rejection_sampler.py     | 51 -----------
 .../test_rejection_sampler_utils.py           | 37 --------
 vllm/v1/core/sched/scheduler.py               | 89 +++++++------------
 vllm/v1/sample/rejection_sampler.py           |  8 +-
 .../spec_decode/rejection_sampler_utils.py    | 16 +---
 5 files changed, 39 insertions(+), 162 deletions(-)

diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py
index 58aac618a3f6..10c4d448f7f6 100644
--- a/tests/v1/sample/test_rejection_sampler.py
+++ b/tests/v1/sample/test_rejection_sampler.py
@@ -1150,54 +1150,3 @@ 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}"
-    )
diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py
index d294de96b9d5..28f6044de877 100644
--- a/tests/v1/spec_decode/test_rejection_sampler_utils.py
+++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py
@@ -280,40 +280,3 @@ 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()
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index f76a46d73073..a9a62fcf3cd8 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -327,48 +327,6 @@ 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,
@@ -626,9 +584,21 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
 
             # Speculative decode related.
             if request.spec_token_ids:
-                self._schedule_spec_tokens(
-                    request, num_new_tokens, scheduled_spec_decode_tokens
+                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 = []
 
             # Encoder-related.
             if encoder_inputs_to_schedule:
@@ -811,6 +781,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                 encoder_inputs_to_schedule = None
                 external_load_encoder_input = []
                 new_encoder_compute_budget = encoder_compute_budget
+                pad_spec_decode = False
 
                 if load_kv_async:
                     # KVTransfer: loading remote KV, do not allocate for new work.
@@ -828,13 +799,21 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # 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
+                    # Pad new decode requests to uniform spec decoding size to
+                    # preserve full cudagraph for this step.
+                    if (
+                        (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
+                        and (prefill_scheduled and scheduled_running_reqs)
+                        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
+                    ):
+                        num_new_tokens += self.num_spec_tokens
+                        if num_new_tokens > token_budget:
+                            break
+                        pad_spec_decode = True
 
                     threshold = self.scheduler_config.long_prefill_token_threshold
                     if 0 < threshold < num_new_tokens:
@@ -1000,10 +979,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
-                    )
+                if pad_spec_decode:
+                    scheduled_spec_decode_tokens[request_id] = [
+                        -1
+                    ] * 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)
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index 1324191be74c..8b4d8c9dce70 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -748,8 +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)
-                # -1 is used for padded draft token ids that should be rejected.
-                accepted = (uniform_prob < rate) and draft_token_id >= 0
+                accepted = uniform_prob < rate
                 token_id = draft_token_id if accepted else target_argmax_id
                 rejected = not accepted
             else:
@@ -806,10 +805,7 @@ 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 draft_token_id < 0:
-                # -1 is used for padded draft token ids that should be rejected.
-                accepted = False
-            elif SYNTHETIC_MODE:
+            if SYNTHETIC_MODE:
                 rate = tl.load(synthetic_conditional_rates_ptr + pos)
                 accepted = uniform_prob < rate
             else:
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
index da3057aec6e7..0cfbdf4182be 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
@@ -245,8 +245,7 @@ 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)
-                    # -1 is used for padded draft token ids that should be rejected.
-                    accepted &= (u < rate) & (draft_sampled >= 0)
+                    accepted &= u < rate
                 else:
                     accepted &= target_argmax == draft_sampled
                 tl.store(
@@ -254,14 +253,8 @@ def _rejection_kernel(
                     draft_sampled if accepted else target_argmax,
                 )
             else:
-                # -1 is used for padded draft token ids that should be rejected.
-                is_valid_draft = draft_sampled >= 0
                 target_logit = tl.load(
-                    target_logits_ptr
-                    + logit_idx * target_logits_stride
-                    + draft_sampled,
-                    mask=is_valid_draft,
-                    other=0.0,
+                    target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
                 ).to(tl.float32)
                 target_lse = _compute_global_lse(
                     target_local_max_ptr,
@@ -280,9 +273,7 @@ def _rejection_kernel(
                         draft_logits_ptr
                         + req_state_idx * draft_logits_stride_0
                         + i * draft_logits_stride_1
-                        + draft_sampled,
-                        mask=is_valid_draft,
-                        other=0.0,
+                        + draft_sampled
                     ).to(tl.float32)
                     draft_lse = _compute_global_lse(
                         draft_local_max_ptr,
@@ -305,7 +296,6 @@ 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)

From b742d561d02cd40f1669a71b9ce3870bb46e148d Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Tue, 23 Jun 2026 19:22:04 -0700
Subject: [PATCH 8/9] simplify further

Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com>
Co-authored-by: Zijing Liu <liuzijing2014@gmail.com>
Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 vllm/v1/core/sched/scheduler.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index a9a62fcf3cd8..bb65c61a0fbd 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -803,14 +803,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                     # preserve full cudagraph for this step.
                     if (
                         (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
-                        and (prefill_scheduled and scheduled_running_reqs)
-                        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
+                        and num_new_tokens == 1
+                        and (scheduled_running_reqs and not prefill_scheduled)
                     ):
-                        num_new_tokens += self.num_spec_tokens
+                        num_new_tokens = 1 + self.num_spec_tokens
                         if num_new_tokens > token_budget:
                             break
                         pad_spec_decode = True

From 8aea7b0ece21ab8784cbd2a95e3f40cffb63fde1 Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill123@gmail.com>
Date: Wed, 24 Jun 2026 10:22:05 -0700
Subject: [PATCH 9/9] max_model_len check

Signed-off-by: Nick Hill <nickhill123@gmail.com>
---
 vllm/v1/core/sched/scheduler.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index bb65c61a0fbd..aa19ba7f8798 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -807,7 +807,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
                         and (scheduled_running_reqs and not prefill_scheduled)
                     ):
                         num_new_tokens = 1 + self.num_spec_tokens
-                        if num_new_tokens > token_budget:
+                        if (
+                            num_new_tokens > token_budget
+                            or num_computed_tokens + num_new_tokens > self.max_model_len
+                        ):
+                            # Prefer to not schedule than schedule un-padded here.
                             break
                         pad_spec_decode = True