From 446c769482dca3278f9d581d72f9b1a05c0690e0 Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:34:56 +0800 Subject: [PATCH] [Distributed] Add opt-in FlashInfer PCIe IPC all-reduce backend (#53576) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Yongye Zhu Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_comm_ops.py | 1 + .../test_flashinfer_pcie_ipc_all_reduce.py | 146 +++++++++++ .../device_communicators/cuda_communicator.py | 33 +++ .../flashinfer_pcie_ipc_all_reduce.py | 239 ++++++++++++++++++ vllm/distributed/parallel_state.py | 21 +- vllm/envs.py | 7 + vllm/model_executor/warmup/kernel_warmup.py | 7 + 7 files changed, 450 insertions(+), 4 deletions(-) create mode 100644 tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py create mode 100644 vllm/distributed/device_communicators/flashinfer_pcie_ipc_all_reduce.py diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index a7d59ff2f8..b5e552af49 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -396,6 +396,7 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> fi_ar_comm.all_reduce.return_value = output communicator = CudaCommunicator.__new__(CudaCommunicator) communicator.fi_ar_comm = fi_ar_comm + communicator.fi_pcie_ipc_ar_comm = None communicator.pynccl_comm = Mock(world_size=8) communicator.qr_comm = None nccl_selector = Mock(return_value=True) diff --git a/tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py b/tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py new file mode 100644 index 0000000000..c1c62a83b8 --- /dev/null +++ b/tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm.distributed.device_communicators import cuda_communicator +from vllm.distributed.device_communicators import ( + flashinfer_pcie_ipc_all_reduce as pcie_ipc, +) +from vllm.distributed.parallel_state import GroupCoordinator + + +def _uninitialized_comm() -> pcie_ipc.FlashInferPcieIpcAllReduce: + comm = pcie_ipc.FlashInferPcieIpcAllReduce.__new__( + pcie_ipc.FlashInferPcieIpcAllReduce + ) + comm.disabled = False + comm.group = object() + comm.tune_group = object() + comm.device = torch.device("cuda:0") + comm.world_size = 4 + comm.rank = 0 + comm.workspace = None + comm.hidden_dim = 0 + comm.dtype = None + return comm + + +def test_setup_uses_exact_graph_capacity(monkeypatch, tmp_path): + workspace = Mock() + factory = Mock(return_value=workspace) + monkeypatch.setattr( + pcie_ipc, + "flashinfer_comm", + SimpleNamespace(PcieIpcAllReduceWorkspace=factory), + ) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + + comm = _uninitialized_comm() + cache = tmp_path / "tp4.json" + comm.setup( + hidden_dim=4096, + dtype=torch.bfloat16, + capture_sizes=[32, 8, 16, 32], + tune_cache=cache, + ) + + factory.assert_called_once_with( + group=comm.group, + max_numel=32 * 4096, + dtype=torch.bfloat16, + tune_batches=(8, 16, 32), + tune_cache=str(cache), + ) + workspace.tune.assert_called_once_with( + [4096], dtype=torch.bfloat16, tune_group=comm.tune_group + ) + workspace.prepare.assert_called_once_with( + [(8, 4096), (16, 4096), (32, 4096)], dtype=torch.bfloat16 + ) + assert workspace.rebind_stream.call_count == 2 + + +@pytest.mark.parametrize( + ("shape", "dtype", "contiguous", "expected"), + [ + ((16, 4096), torch.bfloat16, True, True), + ((16, 2048), torch.bfloat16, True, False), + ((16, 4096), torch.float16, True, False), + ((16, 4096), torch.bfloat16, False, False), + ], +) +def test_should_use_only_prepared_hidden_shape(shape, dtype, contiguous, expected): + comm = _uninitialized_comm() + comm.hidden_dim = 4096 + comm.dtype = torch.bfloat16 + comm.workspace = Mock() + comm.workspace.supports.return_value = True + inp = SimpleNamespace( + is_cuda=True, + is_contiguous=lambda: contiguous, + dim=lambda: len(shape), + shape=shape, + dtype=dtype, + ) + + assert comm.should_use(inp) is expected + + +def test_capture_rebinds_workspace_even_on_error(monkeypatch): + comm = _uninitialized_comm() + comm.workspace = Mock() + synchronize = Mock() + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + + with pytest.raises(RuntimeError, match="capture failed"), comm.capture(): + raise RuntimeError("capture failed") + + assert synchronize.call_count == 2 + assert comm.workspace.rebind_stream.call_count == 2 + + +def test_cuda_dispatch_prefers_pcie_ipc_before_existing_flashinfer(monkeypatch): + communicator = cuda_communicator.CudaCommunicator.__new__( + cuda_communicator.CudaCommunicator + ) + communicator.pynccl_comm = Mock(world_size=4) + communicator.qr_comm = None + communicator.fi_pcie_ipc_ar_comm = Mock() + communicator.fi_pcie_ipc_ar_comm.should_use.return_value = True + expected = object() + communicator.fi_pcie_ipc_ar_comm.all_reduce.return_value = expected + communicator.fi_ar_comm = Mock() + communicator.aiter_ar_comm = None + communicator.ca_comm = None + communicator.symm_mem_comm = None + monkeypatch.setattr( + cuda_communicator, "should_nccl_symm_mem_allreduce", lambda *_: False + ) + + assert communicator.all_reduce(object()) is expected + communicator.fi_ar_comm.should_use_fi_ar.assert_not_called() + + +def test_group_destroy_releases_communicator_before_process_groups(monkeypatch): + events = [] + coordinator = GroupCoordinator.__new__(GroupCoordinator) + coordinator.device_communicator = SimpleNamespace( + destroy=lambda: events.append("communicator") + ) + coordinator.device_group = object() + coordinator.cpu_group = object() + coordinator.mq_broadcaster = None + monkeypatch.setattr( + torch.distributed, + "destroy_process_group", + lambda _group: events.append("process_group"), + ) + + coordinator.destroy() + + assert events == ["communicator", "process_group", "process_group"] diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index c0f2e0a410..5cef8cba3f 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -52,6 +52,7 @@ class CudaCommunicator(DeviceCommunicatorBase): use_custom_allreduce = False use_torch_symm_mem = False use_flashinfer_allreduce = False + use_flashinfer_pcie_ipc_allreduce = False use_aiter_allreduce = False else: from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE @@ -62,6 +63,10 @@ class CudaCommunicator(DeviceCommunicatorBase): use_flashinfer_allreduce = ( envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT ) + use_flashinfer_pcie_ipc_allreduce = ( + envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC + and not envs.VLLM_BATCH_INVARIANT + ) use_aiter_allreduce = use_custom_allreduce and bool( rocm_aiter_ops.is_custom_all_reduce_enabled() ) @@ -69,6 +74,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.use_custom_allreduce = use_custom_allreduce self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce + self.use_flashinfer_pcie_ipc_allreduce = use_flashinfer_pcie_ipc_allreduce self.use_aiter_allreduce = use_aiter_allreduce # lazy import to avoid documentation build error @@ -78,6 +84,9 @@ class CudaCommunicator(DeviceCommunicatorBase): from vllm.distributed.device_communicators.flashinfer_all_reduce import ( FlashInferAllReduce, ) + from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import ( # noqa: E501 + FlashInferPcieIpcAllReduce, + ) from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, @@ -97,6 +106,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None + self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None self.aiter_ar_comm: AiterCustomAllreduce | None = None if use_torch_symm_mem and current_platform.is_cuda(): @@ -111,6 +121,17 @@ class CudaCommunicator(DeviceCommunicatorBase): device=self.device, ) + if ( + self.use_flashinfer_pcie_ipc_allreduce + and self.world_size > 1 + and self.device_group is not None + ): + self.fi_pcie_ipc_ar_comm = FlashInferPcieIpcAllReduce( + group=self.device_group, + tune_group=self.cpu_group, + device=self.device, + ) + if self.use_aiter_allreduce and self.world_size > 1: self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, @@ -220,6 +241,7 @@ class CudaCommunicator(DeviceCommunicatorBase): depends on the input tensor. """ all_potential_ar_backends = [ + "FLASHINFER_PCIE_IPC", "FLASHINFER", "NCCL_SYMM_MEM", "QUICK_REDUCE", @@ -229,6 +251,11 @@ class CudaCommunicator(DeviceCommunicatorBase): "PYNCCL", ] enabled_ar_backends: list[str] = [] + if ( + self.fi_pcie_ipc_ar_comm is not None + and not self.fi_pcie_ipc_ar_comm.disabled + ): + enabled_ar_backends.append("FLASHINFER_PCIE_IPC") if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: enabled_ar_backends.append("FLASHINFER") # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`: @@ -302,6 +329,9 @@ class CudaCommunicator(DeviceCommunicatorBase): out = qr_comm.quick_all_reduce(input_) assert out is not None return out + fi_pcie_ipc_ar_comm = self.fi_pcie_ipc_ar_comm + if fi_pcie_ipc_ar_comm is not None and fi_pcie_ipc_ar_comm.should_use(input_): + return fi_pcie_ipc_ar_comm.all_reduce(input_) if use_fi_ar: assert fi_ar_comm is not None out = fi_ar_comm.all_reduce(input_) @@ -594,6 +624,9 @@ class CudaCommunicator(DeviceCommunicatorBase): if self.fi_ar_comm is not None: self.fi_ar_comm.destroy() self.fi_ar_comm = None + if self.fi_pcie_ipc_ar_comm is not None: + self.fi_pcie_ipc_ar_comm.destroy() + self.fi_pcie_ipc_ar_comm = None if self.all2all_manager is not None: self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] diff --git a/vllm/distributed/device_communicators/flashinfer_pcie_ipc_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_pcie_ipc_all_reduce.py new file mode 100644 index 0000000000..10e9bf885b --- /dev/null +++ b/vllm/distributed/device_communicators/flashinfer_pcie_ipc_all_reduce.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer PCIe CUDA-IPC all-reduce integration. + +This backend targets small tensor-parallel all-reduces on a single PCIe-only +node. Its workspace has a stricter lifetime and stream contract than the +existing FlashInfer MNNVL/TRT-LLM all-reduce implementation, so it intentionally +lives behind a separate wrapper and an opt-in environment variable. +""" + +from collections.abc import Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_SUPPORTED_WORLD_SIZES = (2, 4, 8) +_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) + +try: + import flashinfer.comm as flashinfer_comm + + _pcie_ipc_available = hasattr(flashinfer_comm, "PcieIpcAllReduceWorkspace") +except ImportError: + flashinfer_comm = None # type: ignore[assignment] + _pcie_ipc_available = False + + +class FlashInferPcieIpcAllReduce: + """vLLM lifecycle wrapper for FlashInfer's PCIe IPC all-reduce.""" + + def __init__( + self, + group: ProcessGroup, + tune_group: ProcessGroup, + device: int | str | torch.device, + ) -> None: + self.disabled = True + self.group = group + self.tune_group = tune_group + self.device = torch.device(device) + self.world_size = dist.get_world_size(group) + self.rank = dist.get_rank(group) + self.workspace: Any | None = None + self.hidden_dim = 0 + self.dtype: torch.dtype | None = None + + if not _pcie_ipc_available: + logger.warning_once( + "FlashInfer PCIe IPC all-reduce was requested but this " + "FlashInfer build does not provide PcieIpcAllReduceWorkspace; " + "falling back to another all-reduce backend." + ) + return + if not current_platform.is_cuda(): + logger.warning_once( + "FlashInfer PCIe IPC all-reduce requires the CUDA platform." + ) + return + if self.world_size not in _SUPPORTED_WORLD_SIZES: + logger.warning_once( + "FlashInfer PCIe IPC all-reduce does not support world_size=%d; " + "supported sizes are %s.", + self.world_size, + _SUPPORTED_WORLD_SIZES, + ) + return + if not all(in_the_same_node_as(tune_group, source_rank=0)): + logger.warning_once( + "FlashInfer PCIe IPC all-reduce requires every rank in the TP " + "group to be on one node." + ) + return + + # Setup is deferred until kernel_warmup, where the model hidden size and + # exact CUDA Graph buckets are known. Until then dispatch falls through. + self.disabled = False + + @property + def initialized(self) -> bool: + return not self.disabled and self.workspace is not None + + def setup( + self, + *, + hidden_dim: int, + dtype: torch.dtype, + capture_sizes: Sequence[int], + tune_cache: Path, + ) -> None: + """Allocate, tune, and prepare the exact graph-capture shapes.""" + if self.disabled or self.workspace is not None: + return + if dtype not in _SUPPORTED_DTYPES: + logger.warning_once( + "FlashInfer PCIe IPC all-reduce does not support dtype=%s; " + "falling back to another all-reduce backend.", + dtype, + ) + self.disabled = True + return + + batches = tuple(sorted({int(size) for size in capture_sizes if size > 0})) + if not batches: + logger.warning_once( + "FlashInfer PCIe IPC all-reduce has no CUDA Graph capture sizes " + "to prepare; falling back to another all-reduce backend." + ) + self.disabled = True + return + + self.hidden_dim = int(hidden_dim) + self.dtype = dtype + max_numel = batches[-1] * self.hidden_dim + workspace = flashinfer_comm.PcieIpcAllReduceWorkspace( + group=self.group, + max_numel=max_numel, + dtype=dtype, + tune_batches=batches, + tune_cache=str(tune_cache), + ) + self.workspace = workspace + + # tune() reuses a complete persisted cache and profiles only cache + # misses. It is deliberately mandatory: FlashInfer's seed policy was + # substantially slower than NCCL for the target TP4 decode workload. + torch.accelerator.synchronize(self.device) + workspace.rebind_stream() + workspace.tune( + [self.hidden_dim], + dtype=dtype, + tune_group=self.tune_group, + ) + workspace.prepare([(batch, self.hidden_dim) for batch in batches], dtype=dtype) + torch.accelerator.synchronize(self.device) + workspace.rebind_stream() + logger.info_once( + "Initialized FlashInfer PCIe IPC all-reduce for TP%d, hidden_dim=%d, " + "dtype=%s, max_tokens=%d.", + self.world_size, + self.hidden_dim, + dtype, + batches[-1], + ) + + def should_use(self, inp: torch.Tensor) -> bool: + workspace = self.workspace + return bool( + not self.disabled + and workspace is not None + and inp.is_cuda + and inp.is_contiguous() + and inp.dim() == 2 + and inp.shape[1] == self.hidden_dim + and inp.dtype == self.dtype + and workspace.supports(inp) + ) + + def all_reduce(self, inp: torch.Tensor) -> torch.Tensor: + assert self.workspace is not None + return self.workspace.all_reduce(inp) + + @contextmanager + def capture(self): + """Move the single-stream workspace to and from the capture stream.""" + workspace = self.workspace + if workspace is None: + yield + return + + # The caller orders graph capture against normal execution. A device + # sync makes that assertion explicit before relaxing FlashInfer's + # single-stream guard. Captured launches themselves are exempt from the + # guard, while capture warmups execute and bind to the capture stream. + torch.accelerator.synchronize(self.device) + workspace.rebind_stream() + try: + yield + finally: + torch.accelerator.synchronize(self.device) + workspace.rebind_stream() + + def destroy(self) -> None: + workspace = self.workspace + if workspace is not None: + workspace.destroy() + self.workspace = None + + +def warmup_flashinfer_pcie_ipc_allreduce(worker: "Worker") -> None: + """Initialize the TP PCIe IPC backend immediately before graph capture.""" + from vllm.distributed.device_communicators.cuda_communicator import ( + CudaCommunicator, + ) + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + resolve_flashinfer_autotune_file, + ) + + tp_group = get_tp_group() + communicator = tp_group.device_communicator + if not isinstance(communicator, CudaCommunicator): + return + pcie_comm = communicator.fi_pcie_ipc_ar_comm + if pcie_comm is None or pcie_comm.disabled: + return + if worker.vllm_config.parallel_config.use_ubatching: + logger.warning_once( + "FlashInfer PCIe IPC all-reduce does not yet support DBO or " + "multi-ubatch execution; falling back to another backend." + ) + pcie_comm.disabled = True + return + + capture_sizes = worker.vllm_config.compilation_config.cudagraph_capture_sizes + if not capture_sizes: + return + + base_cache = resolve_flashinfer_autotune_file(worker.model_runner) + ranks = "-".join(str(rank) for rank in tp_group.ranks) + tune_cache = base_cache.with_name(f"pcie_ipc_allreduce_tp_{ranks}.json") + pcie_comm.setup( + hidden_dim=worker.model_config.get_hidden_size(), + dtype=worker.model_config.dtype, + capture_sizes=capture_sizes, + tune_cache=tune_cache, + ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 0291a3d0ae..c3d8e70358 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -29,7 +29,7 @@ import pickle import weakref from collections import deque, namedtuple from collections.abc import Callable -from contextlib import contextmanager, nullcontext +from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from datetime import timedelta from multiprocessing import shared_memory @@ -666,6 +666,7 @@ class GroupCoordinator: # only cuda uses this function, # so we don't abstract it into the base class maybe_ca_context = nullcontext() + maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext() maybe_aiter_context = nullcontext() from vllm.distributed.device_communicators.cuda_communicator import ( CudaCommunicator, @@ -682,6 +683,10 @@ class GroupCoordinator: ca_comm = self.device_communicator.ca_comm if ca_comm is not None: maybe_ca_context = ca_comm.capture() # type: ignore + if isinstance(self.device_communicator, CudaCommunicator): + fi_pcie_ipc_ar_comm = self.device_communicator.fi_pcie_ipc_ar_comm + if fi_pcie_ipc_ar_comm is not None: + maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture() from vllm._aiter_ops import rocm_aiter_ops @@ -696,7 +701,12 @@ class GroupCoordinator: if curr_stream != stream: stream.wait_stream(curr_stream) - with torch.cuda.stream(stream), maybe_ca_context, maybe_aiter_context: + with ( + torch.cuda.stream(stream), + maybe_ca_context, + maybe_fi_pcie_ipc_context, + maybe_aiter_context, + ): yield graph_capture_context def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: @@ -1360,14 +1370,17 @@ class GroupCoordinator: return self.device_communicator.recv(size, dtype, src) def destroy(self): + # Device communicators can own collective workspaces whose teardown + # uses these process groups (for example FlashInfer PCIe IPC barriers). + # Release them before destroying the groups they depend on. + if self.device_communicator is not None: + self.device_communicator.destroy() if hasattr(self, "device_group"): torch.distributed.destroy_process_group(self.device_group) del self.device_group if hasattr(self, "cpu_group"): torch.distributed.destroy_process_group(self.cpu_group) del self.cpu_group - if self.device_communicator is not None: - self.device_communicator.destroy() if self.mq_broadcaster is not None: self.mq_broadcaster = None diff --git a/vllm/envs.py b/vllm/envs.py index 786864b912..f2e24c576e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -264,6 +264,7 @@ if TYPE_CHECKING: VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True VLLM_ALLREDUCE_USE_FLASHINFER: bool = True + VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC: bool = False VLLM_TUNED_CONFIG_FOLDER: str | None = None VLLM_ENABLE_STARTUP_PLAN: bool = False VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() @@ -1864,6 +1865,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool( int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "1")) ), + # Whether to use FlashInfer's single-node PCIe CUDA-IPC all-reduce. + # This backend has a strict single-stream contract and is opt-in while its + # integration is being qualified. + "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": lambda: bool( + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC", "0")) + ), # Experimental: use this to enable MCP tool calling for non harmony models "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool( int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0")) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 42b3746e51..b406ec3873 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -13,6 +13,9 @@ from typing import TYPE_CHECKING import torch import vllm.envs as envs +from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import ( + warmup_flashinfer_pcie_ipc_allreduce, +) from vllm.logger import init_logger from vllm.model_executor.warmup.b12x_warmup import b12x_warmup from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup @@ -182,6 +185,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): minimax_m3_msa_warmup(worker) + # Allocate the exact decode-sized workspace, autotune cache misses, and + # resolve every CUDA Graph bucket before capture begins. + warmup_flashinfer_pcie_ipc_allreduce(worker) + enable_flashinfer_autotune = ( worker.vllm_config.kernel_config.enable_flashinfer_autotune )