[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 <zyy1102000@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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)

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"]

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]

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,
+    )

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
 

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"))

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
     )