[Bugfix][Spec Decode] Skip uniform spec-decode padding for diffusion models (#47464)
Signed-off-by: kl527 <kl527@cornell.edu> Signed-off-by: Kyung Sub Lee (Daniel) <66861800+kl527@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Kyung Sub Lee (Daniel) committed
Jul 7, 2026 at 04:25 UTC
c5b66233b26824f463fcca8af0779ff9e9d7930c
2 files changed
+37
tests/v1/core/test_scheduler.py
+35
@@ -1356,6 +1356,41 @@ def test_spec_decode_padding_first_decode_step():
1356
assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec
1357
1358
1359
+def test_spec_decode_padding_skipped_for_diffusion():
1360
+ """Diffusion spec tokens are the fixed-size denoising canvas, not
1361
+ rejectable drafts: a first-decode-step request must keep its 1-token span
1362
+ instead of being padded to 1 + num_spec_tokens, which would overflow the
1363
+ canvas.
1364
+ """
1365
+ num_spec = 3
1366
+ scheduler = create_scheduler(
1367
+ num_speculative_tokens=num_spec,
1368
+ enable_prefix_caching=True,
1369
+ block_size=16,
1370
+ )
1371
+ # Diffusion schedulers initialize this to 0 (model_config.is_diffusion).
1372
+ scheduler.num_sampled_tokens_per_step = 0
1373
+ r1, r2 = create_requests(
1374
+ num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
1375
+ )
1376
+
1377
+ scheduler.add_request(r1)
1378
+ out = scheduler.schedule()
1379
+ assert out.num_scheduled_tokens[r1.request_id] == 33
1380
+ _model_output(scheduler, out, [[100]])
1381
+ scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
1382
+
1383
+ # r2 arrives; its whole prompt is a prefix-cache hit -> needs exactly
1384
+ # 1 token while r1 is a running speculative decode.
1385
+ scheduler.add_request(r2)
1386
+ out = scheduler.schedule()
1387
+
1388
+ assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
1389
+ # r2 keeps its true 1-token span; no placeholder drafts are attached.
1390
+ assert out.num_scheduled_tokens[r2.request_id] == 1
1391
+ assert r2.request_id not in out.scheduled_spec_decode_tokens
1392
+
1393
+
1394
def test_spec_decode_padding_skipped_with_prefill_in_batch():
1395
"""Padding is skipped when the batch contains a prefill chunk: the batch is
1396
already mixed/non-uniform, so padding a new decode request buys nothing.
vllm/v1/core/sched/scheduler.py
+2
@@ -811,8 +811,10 @@ class Scheduler(SchedulerInterface):
811
812
# Pad new decode requests to uniform spec decoding size to
813
# preserve full cudagraph for this step.
814
+ # Not for diffusion where draft tokens can't be padded.
815
if (
816
(self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)
817
+ and self.num_sampled_tokens_per_step > 0
818
and num_new_tokens == 1
819
and (scheduled_running_reqs and not prefill_scheduled)
820
):