@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 446c769482

[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>

Gabriel Wu committed Sep 1, 2026 at 11:34 UTC 446c769482dca3278f9d581d72f9b1a05c0690e0
7 files changed +450 -4
tests/distributed/test_comm_ops.py
+1
@@ -396,6 +396,7 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) ->
396 fi_ar_comm.all_reduce.return_value = output
397 communicator = CudaCommunicator.__new__(CudaCommunicator)
398 communicator.fi_ar_comm = fi_ar_comm
399 + communicator.fi_pcie_ipc_ar_comm = None
400 communicator.pynccl_comm = Mock(world_size=8)
401 communicator.qr_comm = None
402 nccl_selector = Mock(return_value=True)
tests/distributed/test_flashinfer_pcie_ipc_all_reduce.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from unittest.mock import Mock
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.distributed.device_communicators import cuda_communicator
11 +from vllm.distributed.device_communicators import (
12 + flashinfer_pcie_ipc_all_reduce as pcie_ipc,
13 +)
14 +from vllm.distributed.parallel_state import GroupCoordinator
15 +
16 +
17 +def _uninitialized_comm() -> pcie_ipc.FlashInferPcieIpcAllReduce:
18 + comm = pcie_ipc.FlashInferPcieIpcAllReduce.__new__(
19 + pcie_ipc.FlashInferPcieIpcAllReduce
20 + )
21 + comm.disabled = False
22 + comm.group = object()
23 + comm.tune_group = object()
24 + comm.device = torch.device("cuda:0")
25 + comm.world_size = 4
26 + comm.rank = 0
27 + comm.workspace = None
28 + comm.hidden_dim = 0
29 + comm.dtype = None
30 + return comm
31 +
32 +
33 +def test_setup_uses_exact_graph_capacity(monkeypatch, tmp_path):
34 + workspace = Mock()
35 + factory = Mock(return_value=workspace)
36 + monkeypatch.setattr(
37 + pcie_ipc,
38 + "flashinfer_comm",
39 + SimpleNamespace(PcieIpcAllReduceWorkspace=factory),
40 + )
41 + monkeypatch.setattr(torch.accelerator, "synchronize", Mock())
42 +
43 + comm = _uninitialized_comm()
44 + cache = tmp_path / "tp4.json"
45 + comm.setup(
46 + hidden_dim=4096,
47 + dtype=torch.bfloat16,
48 + capture_sizes=[32, 8, 16, 32],
49 + tune_cache=cache,
50 + )
51 +
52 + factory.assert_called_once_with(
53 + group=comm.group,
54 + max_numel=32 * 4096,
55 + dtype=torch.bfloat16,
56 + tune_batches=(8, 16, 32),
57 + tune_cache=str(cache),
58 + )
59 + workspace.tune.assert_called_once_with(
60 + [4096], dtype=torch.bfloat16, tune_group=comm.tune_group
61 + )
62 + workspace.prepare.assert_called_once_with(
63 + [(8, 4096), (16, 4096), (32, 4096)], dtype=torch.bfloat16
64 + )
65 + assert workspace.rebind_stream.call_count == 2
66 +
67 +
68 +@pytest.mark.parametrize(
69 + ("shape", "dtype", "contiguous", "expected"),
70 + [
71 + ((16, 4096), torch.bfloat16, True, True),
72 + ((16, 2048), torch.bfloat16, True, False),
73 + ((16, 4096), torch.float16, True, False),
74 + ((16, 4096), torch.bfloat16, False, False),
75 + ],
76 +)
77 +def test_should_use_only_prepared_hidden_shape(shape, dtype, contiguous, expected):
78 + comm = _uninitialized_comm()
79 + comm.hidden_dim = 4096
80 + comm.dtype = torch.bfloat16
81 + comm.workspace = Mock()
82 + comm.workspace.supports.return_value = True
83 + inp = SimpleNamespace(
84 + is_cuda=True,
85 + is_contiguous=lambda: contiguous,
86 + dim=lambda: len(shape),
87 + shape=shape,
88 + dtype=dtype,
89 + )
90 +
91 + assert comm.should_use(inp) is expected
92 +
93 +
94 +def test_capture_rebinds_workspace_even_on_error(monkeypatch):
95 + comm = _uninitialized_comm()
96 + comm.workspace = Mock()
97 + synchronize = Mock()
98 + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize)
99 +
100 + with pytest.raises(RuntimeError, match="capture failed"), comm.capture():
101 + raise RuntimeError("capture failed")
102 +
103 + assert synchronize.call_count == 2
104 + assert comm.workspace.rebind_stream.call_count == 2
105 +
106 +
107 +def test_cuda_dispatch_prefers_pcie_ipc_before_existing_flashinfer(monkeypatch):
108 + communicator = cuda_communicator.CudaCommunicator.__new__(
109 + cuda_communicator.CudaCommunicator
110 + )
111 + communicator.pynccl_comm = Mock(world_size=4)
112 + communicator.qr_comm = None
113 + communicator.fi_pcie_ipc_ar_comm = Mock()
114 + communicator.fi_pcie_ipc_ar_comm.should_use.return_value = True
115 + expected = object()
116 + communicator.fi_pcie_ipc_ar_comm.all_reduce.return_value = expected
117 + communicator.fi_ar_comm = Mock()
118 + communicator.aiter_ar_comm = None
119 + communicator.ca_comm = None
120 + communicator.symm_mem_comm = None
121 + monkeypatch.setattr(
122 + cuda_communicator, "should_nccl_symm_mem_allreduce", lambda *_: False
123 + )
124 +
125 + assert communicator.all_reduce(object()) is expected
126 + communicator.fi_ar_comm.should_use_fi_ar.assert_not_called()
127 +
128 +
129 +def test_group_destroy_releases_communicator_before_process_groups(monkeypatch):
130 + events = []
131 + coordinator = GroupCoordinator.__new__(GroupCoordinator)
132 + coordinator.device_communicator = SimpleNamespace(
133 + destroy=lambda: events.append("communicator")
134 + )
135 + coordinator.device_group = object()
136 + coordinator.cpu_group = object()
137 + coordinator.mq_broadcaster = None
138 + monkeypatch.setattr(
139 + torch.distributed,
140 + "destroy_process_group",
141 + lambda _group: events.append("process_group"),
142 + )
143 +
144 + coordinator.destroy()
145 +
146 + assert events == ["communicator", "process_group", "process_group"]
vllm/distributed/device_communicators/cuda_communicator.py
+33
@@ -52,6 +52,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
52 use_custom_allreduce = False
53 use_torch_symm_mem = False
54 use_flashinfer_allreduce = False
55 + use_flashinfer_pcie_ipc_allreduce = False
56 use_aiter_allreduce = False
57 else:
58 from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
@@ -62,6 +63,10 @@ class CudaCommunicator(DeviceCommunicatorBase):
63 use_flashinfer_allreduce = (
64 envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT
65 )
66 + use_flashinfer_pcie_ipc_allreduce = (
67 + envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC
68 + and not envs.VLLM_BATCH_INVARIANT
69 + )
70 use_aiter_allreduce = use_custom_allreduce and bool(
71 rocm_aiter_ops.is_custom_all_reduce_enabled()
72 )
@@ -69,6 +74,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
74 self.use_custom_allreduce = use_custom_allreduce
75 self.use_torch_symm_mem = use_torch_symm_mem
76 self.use_flashinfer_allreduce = use_flashinfer_allreduce
77 + self.use_flashinfer_pcie_ipc_allreduce = use_flashinfer_pcie_ipc_allreduce
78 self.use_aiter_allreduce = use_aiter_allreduce
79
80 # lazy import to avoid documentation build error
@@ -78,6 +84,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
84 from vllm.distributed.device_communicators.flashinfer_all_reduce import (
85 FlashInferAllReduce,
86 )
87 + from vllm.distributed.device_communicators.flashinfer_pcie_ipc_all_reduce import ( # noqa: E501
88 + FlashInferPcieIpcAllReduce,
89 + )
90 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
91 from vllm.distributed.device_communicators.quick_all_reduce import (
92 QuickAllReduce,
@@ -97,6 +106,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
106 self.qr_comm: QuickAllReduce | None = None
107 self.symm_mem_comm: SymmMemCommunicator | None = None
108 self.fi_ar_comm: FlashInferAllReduce | None = None
109 + self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None
110 self.aiter_ar_comm: AiterCustomAllreduce | None = None
111
112 if use_torch_symm_mem and current_platform.is_cuda():
@@ -111,6 +121,17 @@ class CudaCommunicator(DeviceCommunicatorBase):
121 device=self.device,
122 )
123
124 + if (
125 + self.use_flashinfer_pcie_ipc_allreduce
126 + and self.world_size > 1
127 + and self.device_group is not None
128 + ):
129 + self.fi_pcie_ipc_ar_comm = FlashInferPcieIpcAllReduce(
130 + group=self.device_group,
131 + tune_group=self.cpu_group,
132 + device=self.device,
133 + )
134 +
135 if self.use_aiter_allreduce and self.world_size > 1:
136 self.aiter_ar_comm = AiterCustomAllreduce(
137 group=self.cpu_group,
@@ -220,6 +241,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
241 depends on the input tensor.
242 """
243 all_potential_ar_backends = [
244 + "FLASHINFER_PCIE_IPC",
245 "FLASHINFER",
246 "NCCL_SYMM_MEM",
247 "QUICK_REDUCE",
@@ -229,6 +251,11 @@ class CudaCommunicator(DeviceCommunicatorBase):
251 "PYNCCL",
252 ]
253 enabled_ar_backends: list[str] = []
254 + if (
255 + self.fi_pcie_ipc_ar_comm is not None
256 + and not self.fi_pcie_ipc_ar_comm.disabled
257 + ):
258 + enabled_ar_backends.append("FLASHINFER_PCIE_IPC")
259 if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
260 enabled_ar_backends.append("FLASHINFER")
261 # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`:
@@ -302,6 +329,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
329 out = qr_comm.quick_all_reduce(input_)
330 assert out is not None
331 return out
332 + fi_pcie_ipc_ar_comm = self.fi_pcie_ipc_ar_comm
333 + if fi_pcie_ipc_ar_comm is not None and fi_pcie_ipc_ar_comm.should_use(input_):
334 + return fi_pcie_ipc_ar_comm.all_reduce(input_)
335 if use_fi_ar:
336 assert fi_ar_comm is not None
337 out = fi_ar_comm.all_reduce(input_)
@@ -594,6 +624,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
624 if self.fi_ar_comm is not None:
625 self.fi_ar_comm.destroy()
626 self.fi_ar_comm = None
627 + if self.fi_pcie_ipc_ar_comm is not None:
628 + self.fi_pcie_ipc_ar_comm.destroy()
629 + self.fi_pcie_ipc_ar_comm = None
630 if self.all2all_manager is not None:
631 self.all2all_manager.destroy()
632 self.all2all_manager = None # type: ignore[assignment]
vllm/distributed/device_communicators/flashinfer_pcie_ipc_all_reduce.py new
+239
@@ -0,0 +1,239 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FlashInfer PCIe CUDA-IPC all-reduce integration.
4 +
5 +This backend targets small tensor-parallel all-reduces on a single PCIe-only
6 +node. Its workspace has a stricter lifetime and stream contract than the
7 +existing FlashInfer MNNVL/TRT-LLM all-reduce implementation, so it intentionally
8 +lives behind a separate wrapper and an opt-in environment variable.
9 +"""
10 +
11 +from collections.abc import Sequence
12 +from contextlib import contextmanager
13 +from pathlib import Path
14 +from typing import TYPE_CHECKING, Any
15 +
16 +import torch
17 +import torch.distributed as dist
18 +from torch.distributed import ProcessGroup
19 +
20 +from vllm.distributed.parallel_state import in_the_same_node_as
21 +from vllm.logger import init_logger
22 +from vllm.platforms import current_platform
23 +
24 +if TYPE_CHECKING:
25 + from vllm.v1.worker.gpu_worker import Worker
26 +
27 +logger = init_logger(__name__)
28 +
29 +_SUPPORTED_WORLD_SIZES = (2, 4, 8)
30 +_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
31 +
32 +try:
33 + import flashinfer.comm as flashinfer_comm
34 +
35 + _pcie_ipc_available = hasattr(flashinfer_comm, "PcieIpcAllReduceWorkspace")
36 +except ImportError:
37 + flashinfer_comm = None # type: ignore[assignment]
38 + _pcie_ipc_available = False
39 +
40 +
41 +class FlashInferPcieIpcAllReduce:
42 + """vLLM lifecycle wrapper for FlashInfer's PCIe IPC all-reduce."""
43 +
44 + def __init__(
45 + self,
46 + group: ProcessGroup,
47 + tune_group: ProcessGroup,
48 + device: int | str | torch.device,
49 + ) -> None:
50 + self.disabled = True
51 + self.group = group
52 + self.tune_group = tune_group
53 + self.device = torch.device(device)
54 + self.world_size = dist.get_world_size(group)
55 + self.rank = dist.get_rank(group)
56 + self.workspace: Any | None = None
57 + self.hidden_dim = 0
58 + self.dtype: torch.dtype | None = None
59 +
60 + if not _pcie_ipc_available:
61 + logger.warning_once(
62 + "FlashInfer PCIe IPC all-reduce was requested but this "
63 + "FlashInfer build does not provide PcieIpcAllReduceWorkspace; "
64 + "falling back to another all-reduce backend."
65 + )
66 + return
67 + if not current_platform.is_cuda():
68 + logger.warning_once(
69 + "FlashInfer PCIe IPC all-reduce requires the CUDA platform."
70 + )
71 + return
72 + if self.world_size not in _SUPPORTED_WORLD_SIZES:
73 + logger.warning_once(
74 + "FlashInfer PCIe IPC all-reduce does not support world_size=%d; "
75 + "supported sizes are %s.",
76 + self.world_size,
77 + _SUPPORTED_WORLD_SIZES,
78 + )
79 + return
80 + if not all(in_the_same_node_as(tune_group, source_rank=0)):
81 + logger.warning_once(
82 + "FlashInfer PCIe IPC all-reduce requires every rank in the TP "
83 + "group to be on one node."
84 + )
85 + return
86 +
87 + # Setup is deferred until kernel_warmup, where the model hidden size and
88 + # exact CUDA Graph buckets are known. Until then dispatch falls through.
89 + self.disabled = False
90 +
91 + @property
92 + def initialized(self) -> bool:
93 + return not self.disabled and self.workspace is not None
94 +
95 + def setup(
96 + self,
97 + *,
98 + hidden_dim: int,
99 + dtype: torch.dtype,
100 + capture_sizes: Sequence[int],
101 + tune_cache: Path,
102 + ) -> None:
103 + """Allocate, tune, and prepare the exact graph-capture shapes."""
104 + if self.disabled or self.workspace is not None:
105 + return
106 + if dtype not in _SUPPORTED_DTYPES:
107 + logger.warning_once(
108 + "FlashInfer PCIe IPC all-reduce does not support dtype=%s; "
109 + "falling back to another all-reduce backend.",
110 + dtype,
111 + )
112 + self.disabled = True
113 + return
114 +
115 + batches = tuple(sorted({int(size) for size in capture_sizes if size > 0}))
116 + if not batches:
117 + logger.warning_once(
118 + "FlashInfer PCIe IPC all-reduce has no CUDA Graph capture sizes "
119 + "to prepare; falling back to another all-reduce backend."
120 + )
121 + self.disabled = True
122 + return
123 +
124 + self.hidden_dim = int(hidden_dim)
125 + self.dtype = dtype
126 + max_numel = batches[-1] * self.hidden_dim
127 + workspace = flashinfer_comm.PcieIpcAllReduceWorkspace(
128 + group=self.group,
129 + max_numel=max_numel,
130 + dtype=dtype,
131 + tune_batches=batches,
132 + tune_cache=str(tune_cache),
133 + )
134 + self.workspace = workspace
135 +
136 + # tune() reuses a complete persisted cache and profiles only cache
137 + # misses. It is deliberately mandatory: FlashInfer's seed policy was
138 + # substantially slower than NCCL for the target TP4 decode workload.
139 + torch.accelerator.synchronize(self.device)
140 + workspace.rebind_stream()
141 + workspace.tune(
142 + [self.hidden_dim],
143 + dtype=dtype,
144 + tune_group=self.tune_group,
145 + )
146 + workspace.prepare([(batch, self.hidden_dim) for batch in batches], dtype=dtype)
147 + torch.accelerator.synchronize(self.device)
148 + workspace.rebind_stream()
149 + logger.info_once(
150 + "Initialized FlashInfer PCIe IPC all-reduce for TP%d, hidden_dim=%d, "
151 + "dtype=%s, max_tokens=%d.",
152 + self.world_size,
153 + self.hidden_dim,
154 + dtype,
155 + batches[-1],
156 + )
157 +
158 + def should_use(self, inp: torch.Tensor) -> bool:
159 + workspace = self.workspace
160 + return bool(
161 + not self.disabled
162 + and workspace is not None
163 + and inp.is_cuda
164 + and inp.is_contiguous()
165 + and inp.dim() == 2
166 + and inp.shape[1] == self.hidden_dim
167 + and inp.dtype == self.dtype
168 + and workspace.supports(inp)
169 + )
170 +
171 + def all_reduce(self, inp: torch.Tensor) -> torch.Tensor:
172 + assert self.workspace is not None
173 + return self.workspace.all_reduce(inp)
174 +
175 + @contextmanager
176 + def capture(self):
177 + """Move the single-stream workspace to and from the capture stream."""
178 + workspace = self.workspace
179 + if workspace is None:
180 + yield
181 + return
182 +
183 + # The caller orders graph capture against normal execution. A device
184 + # sync makes that assertion explicit before relaxing FlashInfer's
185 + # single-stream guard. Captured launches themselves are exempt from the
186 + # guard, while capture warmups execute and bind to the capture stream.
187 + torch.accelerator.synchronize(self.device)
188 + workspace.rebind_stream()
189 + try:
190 + yield
191 + finally:
192 + torch.accelerator.synchronize(self.device)
193 + workspace.rebind_stream()
194 +
195 + def destroy(self) -> None:
196 + workspace = self.workspace
197 + if workspace is not None:
198 + workspace.destroy()
199 + self.workspace = None
200 +
201 +
202 +def warmup_flashinfer_pcie_ipc_allreduce(worker: "Worker") -> None:
203 + """Initialize the TP PCIe IPC backend immediately before graph capture."""
204 + from vllm.distributed.device_communicators.cuda_communicator import (
205 + CudaCommunicator,
206 + )
207 + from vllm.distributed.parallel_state import get_tp_group
208 + from vllm.model_executor.warmup.flashinfer_autotune_cache import (
209 + resolve_flashinfer_autotune_file,
210 + )
211 +
212 + tp_group = get_tp_group()
213 + communicator = tp_group.device_communicator
214 + if not isinstance(communicator, CudaCommunicator):
215 + return
216 + pcie_comm = communicator.fi_pcie_ipc_ar_comm
217 + if pcie_comm is None or pcie_comm.disabled:
218 + return
219 + if worker.vllm_config.parallel_config.use_ubatching:
220 + logger.warning_once(
221 + "FlashInfer PCIe IPC all-reduce does not yet support DBO or "
222 + "multi-ubatch execution; falling back to another backend."
223 + )
224 + pcie_comm.disabled = True
225 + return
226 +
227 + capture_sizes = worker.vllm_config.compilation_config.cudagraph_capture_sizes
228 + if not capture_sizes:
229 + return
230 +
231 + base_cache = resolve_flashinfer_autotune_file(worker.model_runner)
232 + ranks = "-".join(str(rank) for rank in tp_group.ranks)
233 + tune_cache = base_cache.with_name(f"pcie_ipc_allreduce_tp_{ranks}.json")
234 + pcie_comm.setup(
235 + hidden_dim=worker.model_config.get_hidden_size(),
236 + dtype=worker.model_config.dtype,
237 + capture_sizes=capture_sizes,
238 + tune_cache=tune_cache,
239 + )
vllm/distributed/parallel_state.py
+17 -4
@@ -29,7 +29,7 @@ import pickle
29 import weakref
30 from collections import deque, namedtuple
31 from collections.abc import Callable
32 -from contextlib import contextmanager, nullcontext
32 +from contextlib import AbstractContextManager, contextmanager, nullcontext
33 from dataclasses import dataclass
34 from datetime import timedelta
35 from multiprocessing import shared_memory
@@ -666,6 +666,7 @@ class GroupCoordinator:
666 # only cuda uses this function,
667 # so we don't abstract it into the base class
668 maybe_ca_context = nullcontext()
669 + maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext()
670 maybe_aiter_context = nullcontext()
671 from vllm.distributed.device_communicators.cuda_communicator import (
672 CudaCommunicator,
@@ -682,6 +683,10 @@ class GroupCoordinator:
683 ca_comm = self.device_communicator.ca_comm
684 if ca_comm is not None:
685 maybe_ca_context = ca_comm.capture() # type: ignore
686 + if isinstance(self.device_communicator, CudaCommunicator):
687 + fi_pcie_ipc_ar_comm = self.device_communicator.fi_pcie_ipc_ar_comm
688 + if fi_pcie_ipc_ar_comm is not None:
689 + maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture()
690
691 from vllm._aiter_ops import rocm_aiter_ops
692
@@ -696,7 +701,12 @@ class GroupCoordinator:
701 if curr_stream != stream:
702 stream.wait_stream(curr_stream)
703
699 - with torch.cuda.stream(stream), maybe_ca_context, maybe_aiter_context:
704 + with (
705 + torch.cuda.stream(stream),
706 + maybe_ca_context,
707 + maybe_fi_pcie_ipc_context,
708 + maybe_aiter_context,
709 + ):
710 yield graph_capture_context
711
712 def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
@@ -1360,14 +1370,17 @@ class GroupCoordinator:
1370 return self.device_communicator.recv(size, dtype, src)
1371
1372 def destroy(self):
1373 + # Device communicators can own collective workspaces whose teardown
1374 + # uses these process groups (for example FlashInfer PCIe IPC barriers).
1375 + # Release them before destroying the groups they depend on.
1376 + if self.device_communicator is not None:
1377 + self.device_communicator.destroy()
1378 if hasattr(self, "device_group"):
1379 torch.distributed.destroy_process_group(self.device_group)
1380 del self.device_group
1381 if hasattr(self, "cpu_group"):
1382 torch.distributed.destroy_process_group(self.cpu_group)
1383 del self.cpu_group
1369 - if self.device_communicator is not None:
1370 - self.device_communicator.destroy()
1384 if self.mq_broadcaster is not None:
1385 self.mq_broadcaster = None
1386
vllm/envs.py
+7
@@ -264,6 +264,7 @@ if TYPE_CHECKING:
264 VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
265 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
266 VLLM_ALLREDUCE_USE_FLASHINFER: bool = True
267 + VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC: bool = False
268 VLLM_TUNED_CONFIG_FOLDER: str | None = None
269 VLLM_ENABLE_STARTUP_PLAN: bool = False
270 VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
@@ -1864,6 +1865,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1865 "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
1866 int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "1"))
1867 ),
1868 + # Whether to use FlashInfer's single-node PCIe CUDA-IPC all-reduce.
1869 + # This backend has a strict single-stream contract and is opt-in while its
1870 + # integration is being qualified.
1871 + "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": lambda: bool(
1872 + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC", "0"))
1873 + ),
1874 # Experimental: use this to enable MCP tool calling for non harmony models
1875 "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
1876 int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))
vllm/model_executor/warmup/kernel_warmup.py
+7
@@ -13,6 +13,9 @@ 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 +)
19 from vllm.logger import init_logger
20 from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
21 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
@@ -182,6 +185,10 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
185
186 minimax_m3_msa_warmup(worker)
187
188 + # Allocate the exact decode-sized workspace, autotune cache misses, and
189 + # resolve every CUDA Graph bucket before capture begins.
190 + warmup_flashinfer_pcie_ipc_allreduce(worker)
191 +
192 enable_flashinfer_autotune = (
193 worker.vllm_config.kernel_config.enable_flashinfer_autotune
194 )