From e9f331d72e90f34614363101528afe6c6fcdf7c5 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 23 Apr 2026 18:33:26 -0700 Subject: [PATCH] [MRV2] Ensure warmup covers prefill path (#40746) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/warmup.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 026b6a7d7eb..83d87c74a4a 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -29,13 +29,16 @@ def warmup_kernels( triton kernels. We must call the provided worker's execute_model for pipeline parallel coordination. - The first iteration simulates a prefill with requests of 2 prompt - tokens each. The second iteration simulates a decode step with all - requests generating 1 token each. + The first iteration simulates a prefill with requests of + 2 + num_spec_steps prompt tokens each. The second iteration simulates + a decode step with all requests generating 1 + num_spec_steps tokens. """ - prompt_token_ids = [0, 1] - prompt_len = len(prompt_token_ids) num_spec_steps = model_runner.num_speculative_steps + # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request + # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing + # it from being misclassified as a uniform decode batch. + prompt_len = 2 + num_spec_steps + prompt_token_ids = list(range(prompt_len)) # After prefill, decode generates 1 verified + num_spec_steps draft tokens. decode_len = prompt_len + 1 + num_spec_steps @@ -76,7 +79,7 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) - # Step 1: Prefill all requests with 2 prompt tokens each. + # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each. new_reqs = [ NewRequestData.from_request( Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),