[Bugfix][V1] Honor enable_jit_warmup for V2 kernel warmup (#55146)
Co-authored-by: mgoin <mgoin64@gmail.com>
pengyihang committed
Sep 23, 2026 at 03:34 UTC
415e11c6f1256e8e007ee88f2350049c51f11a56
1 file changed
+6
-3
vllm/v1/worker/gpu_worker.py
+6
-3
@@ -838,9 +838,12 @@ class Worker(WorkerBase):
838
# cuda graph capture.
839
kernel_warmup(self)
840
841
- if self.use_v2_model_runner:
842
- # A workspace resize after capture frees what the graphs point at.
843
- warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
841
+ if self.use_v2_model_runner: # noqa: SIM102
842
+ if self.vllm_config.kernel_config.enable_jit_warmup:
843
+ # A workspace resize after capture frees what the graphs point at.
844
+ warmup_kernels(
845
+ self.model_runner, self.execute_model, self.sample_tokens
846
+ )
847
848
cuda_graph_memory_bytes = 0
849
if not self.model_config.enforce_eager: