diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 004adf4a67ba..d0168c9a9352 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1302,6 +1302,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/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 55e40b204366..ab9fd5e34331 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -415,6 +415,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 +575,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 @@ -778,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. @@ -793,6 +797,23 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # `request.num_prompt_tokens` to consider the resumed # requests, which have output tokens. num_new_tokens = request.num_tokens - num_computed_tokens + + # 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 num_new_tokens == 1 + and (scheduled_running_reqs and not prefill_scheduled) + ): + num_new_tokens = 1 + self.num_spec_tokens + 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 + threshold = self.scheduler_config.long_prefill_token_threshold if 0 < threshold < num_new_tokens: num_new_tokens = threshold @@ -957,6 +978,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 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)