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 6, 04:30 UTC
[MRV2] Consider spec decoding in warmup (#37812)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Woosuk Kwon committed
Mar 23, 2026 at 10:45 UTC
ffb5b32b5f2a37ea58261747aef9e5f3907b9941
1 file changed
+16
-5
vllm/v1/worker/gpu/warmup.py
+16
-5
@@ -35,7 +35,9 @@ def warmup_kernels(
35
"""
36
prompt_token_ids = [0, 1]
37
prompt_len = len(prompt_token_ids)
38
- decode_len = prompt_len + 1 # After prefill, one decode token is added.
38
+ num_spec_steps = model_runner.num_speculative_steps
39
+ # After prefill, decode generates 1 verified + num_spec_steps draft tokens.
40
+ decode_len = prompt_len + 1 + num_spec_steps
41
42
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
43
num_kv_cache_groups = len(kv_cache_groups)
@@ -51,7 +53,8 @@ def warmup_kernels(
53
54
num_reqs = min(
55
model_runner.scheduler_config.max_num_seqs,
54
- model_runner.scheduler_config.max_num_batched_tokens // prompt_len,
56
+ model_runner.scheduler_config.max_num_batched_tokens
57
+ // max(prompt_len, 1 + num_spec_steps),
58
# Reserve block 0 (null block) and ensure we have enough blocks.
59
max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
60
)
@@ -111,7 +114,7 @@ def warmup_kernels(
114
115
worker_sample_tokens(grammar_output)
116
114
- # Step 2: Decode all requests with 1 token each.
117
+ # Step 2: Decode all requests with 1 + num_spec_steps tokens each.
118
cached_req_data = CachedRequestData.make_empty()
119
cached_req_data.req_ids = list(req_ids)
120
cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
@@ -124,8 +127,16 @@ def warmup_kernels(
127
128
decode_output = SchedulerOutput.make_empty()
129
decode_output.scheduled_cached_reqs = cached_req_data
127
- decode_output.num_scheduled_tokens = {rid: 1 for rid in req_ids}
128
- decode_output.total_num_scheduled_tokens = num_reqs
130
+ decode_output.num_scheduled_tokens = {
131
+ req_id: 1 + num_spec_steps for req_id in req_ids
132
+ }
133
+ if num_spec_steps > 0:
134
+ decode_output.scheduled_spec_decode_tokens = {
135
+ req_id: [0] * num_spec_steps for req_id in req_ids
136
+ }
137
+ decode_output.total_num_scheduled_tokens = sum(
138
+ decode_output.num_scheduled_tokens.values()
139
+ )
140
decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
141
142
worker_execute_model(decode_output)