[Bugfix][PD] Pad resumed speculative decode requests (#55126)
Signed-off-by: Ziming Huang <48115868+ZeldaHuang@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Ziming Huang committed
Sep 4, 2026 at 11:05 UTC
a26b71d868102dea759d469eaf6f08be0765d831
2 files changed
+18
-1
tests/v1/core/test_scheduler.py
+16
@@ -1828,6 +1828,22 @@ def test_spec_decode_padding_first_decode_step():
1828
assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec
1829
1830
1831
+def test_spec_decode_padding_resumed_request_without_running_requests():
1832
+ """Pad a synchronously resumed request even without local running work."""
1833
+ num_spec = 3
1834
+ scheduler = create_scheduler(
1835
+ num_speculative_tokens=num_spec,
1836
+ use_kv_connector=mock_kv(matched_tokens=32, is_async=False),
1837
+ )
1838
+ (resumed,) = create_requests(num_requests=1, num_tokens=33, max_tokens=16)
1839
+
1840
+ scheduler.add_request(resumed)
1841
+ out = scheduler.schedule()
1842
+
1843
+ assert out.num_scheduled_tokens[resumed.request_id] == 1 + num_spec
1844
+ assert out.scheduled_spec_decode_tokens[resumed.request_id] == [-1] * num_spec
1845
+
1846
+
1847
def test_spec_decode_padding_skipped_for_diffusion():
1848
"""Diffusion spec tokens are the fixed-size denoising canvas, not
1849
rejectable drafts: a first-decode-step request must keep its 1-token span
vllm/v1/core/sched/scheduler.py
+2
-1
@@ -973,7 +973,8 @@ class Scheduler(SchedulerInterface):
973
(self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
974
and self.num_sampled_tokens_per_step > 0
975
and num_new_tokens == 1
976
- and (scheduled_running_reqs and not prefill_scheduled)
976
+ and not prefill_scheduled
977
+ and (scheduled_running_reqs or num_computed_tokens > 0)
978
):
979
padded_num_tokens = 1 + self.num_spec_tokens
980
# Pad only when there is room for the sampled token(s).