A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 17:30 UTC
[MRV2] Ensure warmup covers prefill path (#40746)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Apr 23, 2026 at 18:33 UTC
e9f331d72e90f34614363101528afe6c6fcdf7c5
1 file changed
+9
-6
vllm/v1/worker/gpu/warmup.py
+9
-6
@@ -29,13 +29,16 @@ def warmup_kernels(
29
triton kernels. We must call the provided worker's execute_model for
30
pipeline parallel coordination.
31
32
- The first iteration simulates a prefill with requests of 2 prompt
33
- tokens each. The second iteration simulates a decode step with all
34
- requests generating 1 token each.
32
+ The first iteration simulates a prefill with requests of
33
+ 2 + num_spec_steps prompt tokens each. The second iteration simulates
34
+ a decode step with all requests generating 1 + num_spec_steps tokens.
35
"""
36
- prompt_token_ids = [0, 1]
37
- prompt_len = len(prompt_token_ids)
36
num_spec_steps = model_runner.num_speculative_steps
37
+ # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request
38
+ # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing
39
+ # it from being misclassified as a uniform decode batch.
40
+ prompt_len = 2 + num_spec_steps
41
+ prompt_token_ids = list(range(prompt_len))
42
# After prefill, decode generates 1 verified + num_spec_steps draft tokens.
43
decode_len = prompt_len + 1 + num_spec_steps
44
@@ -76,7 +79,7 @@ def warmup_kernels(
79
nonlocal next_block_id
80
return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
81
79
- # Step 1: Prefill all requests with 2 prompt tokens each.
82
+ # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each.
83
new_reqs = [
84
NewRequestData.from_request(
85
Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),