[Bugfix] Lazy-import FlashInfer PCIe IPC all-reduce in kernel_warmup (#54869)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>
Gabriel Wu committed
Sep 2, 2026 at 12:31 UTC
396c5a563223244fef5946caa1be10256eca7189
1 file changed
+8
-4
vllm/model_executor/warmup/kernel_warmup.py
+8
-4
@@ -13,9 +13,6 @@ from typing import TYPE_CHECKING
13
import torch
14
15
import vllm.envs as envs
16
-from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import (
17
- warmup_flashinfer_pcie_ipc_allreduce,
18
-)
16
from vllm.logger import init_logger
17
from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
18
from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
@@ -200,7 +197,14 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
197
198
# Allocate the exact decode-sized workspace, autotune cache misses, and
199
# resolve every CUDA Graph bucket before capture begins.
203
- warmup_flashinfer_pcie_ipc_allreduce(worker)
200
+ # Lazy import: flashinfer_pcie_ipc_all_reduce imports flashinfer.comm,
201
+ # which initializes CUDA at import time and must not run at module scope.
202
+ if envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC:
203
+ from vllm.distributed.device_communicators import (
204
+ flashinfer_pcie_ipc_all_reduce,
205
+ )
206
+
207
+ flashinfer_pcie_ipc_all_reduce.warmup_flashinfer_pcie_ipc_allreduce(worker)
208
209
enable_flashinfer_autotune = (
210
worker.vllm_config.kernel_config.enable_flashinfer_autotune