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 7, 02:45 UTC
[Model Runner V2] Fix warmup for pipeline parallel (#36280)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 6, 2026 at 16:58 UTC
b3546865247d5f61025b6fa256fe08c2843f6ea0
2 files changed
+16
-8
vllm/v1/worker/gpu/warmup.py
+15
-7
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Callable
5
+from typing import Any
6
+
7
import numpy as np
8
import torch
9
@@ -17,9 +20,14 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner
20
21
22
@torch.inference_mode()
20
-def warmup_kernels(model_runner: GPUModelRunner) -> None:
23
+def warmup_kernels(
24
+ model_runner: GPUModelRunner,
25
+ worker_execute_model: Callable[[SchedulerOutput], Any],
26
+ worker_sample_tokens: Callable[[GrammarOutput | None], Any],
27
+) -> None:
28
"""Run two execute_model + sample_tokens iterations to JIT compile
22
- triton 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
@@ -83,7 +91,7 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
91
92
# Disable KV connector for warmup run.
93
model_runner.kv_connector.set_disabled(True)
86
- model_runner.execute_model(prefill_output)
94
+ worker_execute_model(prefill_output)
95
96
if not model_runner.is_pooling_model:
97
# Warm up sampler and perform a decode step for non-pooling models.
@@ -101,7 +109,7 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
109
structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask
110
)
111
104
- model_runner.sample_tokens(grammar_output)
112
+ worker_sample_tokens(grammar_output)
113
114
# Step 2: Decode all requests with 1 token each.
115
cached_req_data = CachedRequestData.make_empty()
@@ -120,12 +128,12 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:
128
decode_output.total_num_scheduled_tokens = num_reqs
129
decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
130
123
- model_runner.execute_model(decode_output)
124
- model_runner.sample_tokens(None)
131
+ worker_execute_model(decode_output)
132
+ worker_sample_tokens(None)
133
134
# Clean up - process finish_req_ids.
135
cleanup_output = SchedulerOutput.make_empty()
136
cleanup_output.finished_req_ids = set(req_ids)
129
- model_runner.execute_model(cleanup_output)
137
+ worker_execute_model(cleanup_output)
138
model_runner.kv_connector.set_disabled(False)
139
torch.accelerator.synchronize()
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -584,7 +584,7 @@ class Worker(WorkerBase):
584
585
if self.use_v2_model_runner:
586
# V2: Run full execute_model + sample_tokens to JIT compile triton kernels.
587
- warmup_kernels(self.model_runner)
587
+ warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
588
elif get_pp_group().is_last_rank:
589
# V1: Warm up sampler and preallocate memory buffer for logits and other
590
# sampling related tensors of max possible shape to avoid memory