@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d490b98162

[Core] Avoid mixed length specdec batches via padding (#45237)

Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com>

Yiliu Dong committed Jun 25, 2026 at 23:34 UTC d490b98162422922230c8b91757c58077e24682f

2 files changed +120

tests/v1/core/test_scheduler.py

+95

@@ -1302,6 +1302,101 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks():

1302 assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens

1303

1304

1305 +def _model_output(scheduler, output, sampled):

1306 + """Feed `sampled` (per-request list) back to the scheduler."""

1307 + req_ids = list(output.num_scheduled_tokens.keys())

1308 + scheduler.update_from_output(

1309 + output,

1310 + ModelRunnerOutput(

1311 + req_ids=req_ids,

1312 + req_id_to_index={r: i for i, r in enumerate(req_ids)},

1313 + sampled_token_ids=sampled,

1314 + logprobs=None,

1315 + prompt_logprobs_dict={},

1316 + pooler_output=[],

1317 + ),

1318 + )

1319 +

1320 +

1321 +def test_spec_decode_padding_first_decode_step():

1322 + """A request taking its first decode step (whole prompt already computed via

1323 + a prefix-cache hit) is padded with placeholder (-1) spec tokens so it enters

1324 + the worker with the same 1 + num_spec_tokens shape as the other speculative

1325 + decodes, keeping the batch uniform.

1326 + """

1327 + num_spec = 3

1328 + scheduler = create_scheduler(

1329 + num_speculative_tokens=num_spec,

1330 + enable_prefix_caching=True,

1331 + block_size=16,

1332 + )

1333 + # Two identical 33-token prompts: 2 full blocks (32 tokens) get cached, so a

1334 + # second identical request hits num_computed == num_prompt_tokens - 1.

1335 + r1, r2 = create_requests(

1336 + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16

1337 + )

1338 +

1339 + # Drive r1 through prefill so its prompt blocks are cached, then give it real

1340 + # drafts so it is a running speculative decode (1 + num_spec shape).

1341 + scheduler.add_request(r1)

1342 + out = scheduler.schedule()

1343 + assert out.num_scheduled_tokens[r1.request_id] == 33

1344 + _model_output(scheduler, out, [[100]])

1345 + scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))

1346 +

1347 + # r2 arrives; its whole prompt is a prefix-cache hit -> first decode step.

1348 + scheduler.add_request(r2)

1349 + out = scheduler.schedule()

1350 +

1351 + # r1 verifies its real drafts.

1352 + assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]

1353 + # r2 is padded to the 1 + num_spec shape with placeholder (-1) drafts.

1354 + assert out.num_scheduled_tokens[r2.request_id] == 1 + num_spec

1355 + assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec

1356 +

1357 +

1358 +def test_spec_decode_padding_skipped_with_prefill_in_batch():

1359 + """Padding is skipped when the batch contains a prefill chunk: the batch is

1360 + already mixed/non-uniform, so padding a new decode request buys nothing.

1361 + """

1362 + num_spec = 3

1363 + scheduler = create_scheduler(

1364 + num_speculative_tokens=num_spec,

1365 + enable_prefix_caching=True,

1366 + block_size=16,

1367 + max_num_batched_tokens=64,

1368 + )

1369 + # r_warm + r_candidate share a prompt so r_candidate gets a full prefix hit.

1370 + r_warm, r_candidate = create_requests(

1371 + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=1

1372 + )

1373 + # r_long has a different, long prompt that prefills over multiple chunks.

1374 + (r_long,) = create_requests(num_requests=1, num_tokens=100, max_tokens=16)

1375 +

1376 + # Warm the prefix cache with r_warm's prompt (it finishes; blocks stay cached).

1377 + scheduler.add_request(r_warm)

1378 + out = scheduler.schedule()

1379 + assert out.num_scheduled_tokens[r_warm.request_id] == 33

1380 + _model_output(scheduler, out, [[100]])

1381 + assert r_warm.request_id in scheduler.finished_req_ids

1382 +

1383 + # Start r_long; after one chunk it remains a prefill chunk in the running queue.

1384 + scheduler.add_request(r_long)

1385 + out = scheduler.schedule()

1386 + _model_output(scheduler, out, [[]]) # still prefilling, no sampled token

1387 + assert r_long.is_prefill_chunk

1388 +

1389 + # r_candidate arrives (prefix-cache hit -> first decode step) alongside the

1390 + # in-flight prefill chunk.

1391 + scheduler.add_request(r_candidate)

1392 + out = scheduler.schedule()

1393 +

1394 + # The batch has a prefill chunk, so r_candidate is NOT padded.

1395 + assert r_long.request_id in out.num_scheduled_tokens

1396 + assert out.num_scheduled_tokens[r_candidate.request_id] == 1

1397 + assert r_candidate.request_id not in out.scheduled_spec_decode_tokens

1398 +

1399 +

1400 def test_scheduler_stats_waiting_queues():

1401 """Test that scheduler stats correctly report waiting and skipped_waiting queues."""

1402 # Create scheduler with limited capacity so we can have waiting requests

vllm/v1/core/sched/scheduler.py

+25

@@ -415,6 +415,8 @@ class Scheduler(SchedulerInterface):

415 encoder_compute_budget = self.max_num_encoder_input_tokens

416 # Spec decode-related.

417 scheduled_spec_decode_tokens: dict[str, list[int]] = {}

418 + # Whether the running batch contains any prefill requests.

419 + prefill_scheduled = False

420

421 # For logging.

422 scheduled_timestamp = time.monotonic()

@@ -573,6 +575,7 @@ class Scheduler(SchedulerInterface):

575

576 # Schedule the request.

577 scheduled_running_reqs.append(request)

578 + prefill_scheduled |= request.is_prefill_chunk

579 request_id = request.request_id

580 req_to_new_blocks[request_id] = new_blocks

581 num_scheduled_tokens[request_id] = num_new_tokens

@@ -778,6 +781,7 @@ class Scheduler(SchedulerInterface):

781 encoder_inputs_to_schedule = None

782 external_load_encoder_input = []

783 new_encoder_compute_budget = encoder_compute_budget

784 + pad_spec_decode = False

785

786 if load_kv_async:

787 # KVTransfer: loading remote KV, do not allocate for new work.

@@ -793,6 +797,23 @@ class Scheduler(SchedulerInterface):

797 # `request.num_prompt_tokens` to consider the resumed

798 # requests, which have output tokens.

799 num_new_tokens = request.num_tokens - num_computed_tokens

800 +

801 + # Pad new decode requests to uniform spec decoding size to

802 + # preserve full cudagraph for this step.

803 + if (

804 + (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None)

805 + and num_new_tokens == 1

806 + and (scheduled_running_reqs and not prefill_scheduled)

807 + ):

808 + num_new_tokens = 1 + self.num_spec_tokens

809 + if (

810 + num_new_tokens > token_budget

811 + or num_computed_tokens + num_new_tokens > self.max_model_len

812 + ):

813 + # Prefer to not schedule than schedule un-padded here.

814 + break

815 + pad_spec_decode = True

816 +

817 threshold = self.scheduler_config.long_prefill_token_threshold

818 if 0 < threshold < num_new_tokens:

819 num_new_tokens = threshold

@@ -957,6 +978,10 @@ class Scheduler(SchedulerInterface):

978 token_budget -= num_new_tokens

979 request.status = RequestStatus.RUNNING

980 request.num_computed_tokens = num_computed_tokens

981 + if pad_spec_decode:

982 + scheduled_spec_decode_tokens[request_id] = [

983 + -1

984 + ] * self.num_spec_tokens

985 # Only track requests that will still be prefilling after this chunk.

986 if num_computed_tokens + num_new_tokens < request.num_tokens:

987 self._inflight_prefills.add(request)