From faa01d450015fbcfb84b096d996c611823b1b41e Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 19 Jul 2026 16:06:21 +0000 Subject: [PATCH] fix(pcie): isolate target and draft graph channels --- .../distributed/test_b12x_fused_all_reduce.py | 26 ++++ tests/distributed/test_dcp_a2a.py | 117 ++++++++++++++++++ .../device_communicators/custom_all_reduce.py | 42 ++++--- vllm/distributed/parallel_state.py | 10 ++ vllm/envs.py | 7 +- vllm/v1/attention/ops/dcp_alltoall.py | 24 +++- 6 files changed, 207 insertions(+), 19 deletions(-) diff --git a/tests/distributed/test_b12x_fused_all_reduce.py b/tests/distributed/test_b12x_fused_all_reduce.py index 0fb4a9a312b7..f63823444b4e 100644 --- a/tests/distributed/test_b12x_fused_all_reduce.py +++ b/tests/distributed/test_b12x_fused_all_reduce.py @@ -7,6 +7,7 @@ import torch import torch.distributed as dist +import vllm.envs as envs import vllm.ir.ops from tests.compile.backend import TestBackend from vllm.compilation.passes.fusion import allreduce_rms_fusion @@ -26,6 +27,7 @@ from vllm.distributed import tensor_model_parallel_all_reduce from vllm.distributed.device_communicators.custom_all_reduce import ( CustomAllreduce, + _b12x_pcie_oneshot_limits, get_b12x_pcie_allreduce, ) from vllm.distributed.parallel_state import get_tp_group, graph_capture @@ -131,6 +133,30 @@ def test_b12x_fused_allreduce_zero_cutoff_disables_support() -> None: assert not custom_allreduce.supports_fused_add_rms_norm() +def test_b12x_oneshot_defaults_to_stream_isolation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_PCIE_ONESHOT_SINGLE_CHANNEL", raising=False) + + assert not envs.environment_variables["VLLM_PCIE_ONESHOT_SINGLE_CHANNEL"]() + + +def test_b12x_oneshot_buffer_tracks_dispatch_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE", "64KB") + monkeypatch.setenv( + "VLLM_PCIE_ONESHOT_FUSED_ADD_RMS_NORM_MAX_SIZE", + "84KB", + ) + + assert _b12x_pcie_oneshot_limits() == ( + 64 * 1024, + 84 * 1024, + 84 * 1024, + ) + + def test_b12x_fused_custom_op_dispatch(monkeypatch) -> None: custom_allreduce = MagicMock() custom_allreduce.try_fused_add_rms_norm.return_value = True diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 1b49726a165b..fa244e261524 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -10,6 +10,7 @@ import importlib.util import math +from contextlib import contextmanager from typing import Any import multiprocess as mp @@ -538,6 +539,122 @@ def fake_all_reduce(tensor, *, op, group): assert captured["op"] == dist.ReduceOp.MAX +def test_b12x_pool_uses_independent_stream_channels( + monkeypatch: pytest.MonkeyPatch, +): + from vllm.v1.attention.ops import dcp_alltoall + + captured: dict[str, Any] = {} + + class _FakePool: + @classmethod + def from_exchange_group(cls, **kwargs): + captured.update(kwargs) + return cls() + + def for_stream(self): + captured["warmed"] = True + + group = _FakeCPGroup(2, object()) # type: ignore[arg-type] + monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_POOLS", {}) + monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_DISABLED", set()) + monkeypatch.setattr(dcp_alltoall, "_load_b12x_dcp_a2a_pool", lambda: _FakePool) + monkeypatch.setattr(dcp_alltoall, "_b12x_dcp_init_failed", lambda *args: False) + monkeypatch.setattr( + dcp_alltoall.torch.cuda, + "is_current_stream_capturing", + lambda: False, + ) + + pool = dcp_alltoall._get_b12x_dcp_a2a_pool( + group, # type: ignore[arg-type] + device=torch.device("cuda:0"), + total_heads=64, + head_dim=512, + query_head_dim=576, + max_batch_size=64, + ) + + assert pool is not None + assert captured["single_channel"] is False + assert captured["warmed"] is True + + +def test_b12x_dcp_capture_selects_only_current_group_pools(monkeypatch): + from vllm.v1.attention.ops import dcp_alltoall + + events = [] + + class _FakePool: + def __init__(self, name): + self.name = name + + @contextmanager + def capture(self, *, stream): + events.append(("enter", self.name, stream)) + try: + yield + finally: + events.append(("exit", self.name, stream)) + + device_group = object() + group = _FakeCPGroup(2, device_group) # type: ignore[arg-type] + stream = object() + pools = { + (id(device_group), 0, 64, 512, 576, 64): _FakePool("output"), + (id(device_group), 0, 64, 576, 576, 64): _FakePool("query"), + (id(object()), 0, 64, 512, 576, 64): _FakePool("foreign"), + } + monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_POOLS", pools) + + with dcp_alltoall.capture_b12x_dcp_a2a(group, stream): # type: ignore[arg-type] + events.append(("body", None, stream)) + + assert events == [ + ("enter", "output", stream), + ("enter", "query", stream), + ("body", None, stream), + ("exit", "query", stream), + ("exit", "output", stream), + ] + + +def test_global_graph_capture_enters_b12x_dcp_pool(monkeypatch): + from vllm.distributed import parallel_state + from vllm.v1.attention.ops import dcp_alltoall + + events = [] + + class _FakeGroup: + world_size = 2 + + @contextmanager + def graph_capture(self, context): + yield context + + tp_group = _FakeGroup() + pp_group = _FakeGroup() + dcp_group = _FakeGroup() + stream = object() + context = parallel_state.GraphCaptureContext(stream) # type: ignore[arg-type] + + @contextmanager + def fake_b12x_capture(group, selected_stream): + events.append((group, selected_stream)) + yield + + monkeypatch.setattr(parallel_state, "_DCP", dcp_group) + monkeypatch.setattr(parallel_state, "get_tp_group", lambda: tp_group) + monkeypatch.setattr(parallel_state, "get_pp_group", lambda: pp_group) + monkeypatch.setattr(parallel_state, "get_dcp_group", lambda: dcp_group) + monkeypatch.setattr(dcp_alltoall, "capture_b12x_dcp_a2a", fake_b12x_capture) + + with parallel_state.graph_capture(torch.device("cpu"), context) as actual: + assert actual is context + + assert events == [(dcp_group, stream)] + + @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.") def test_b12x_lse_reduce_honors_token_cap(monkeypatch: pytest.MonkeyPatch): from vllm.v1.attention.ops import dcp_alltoall diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 7d3d48cf3c7e..9ed8284b2c90 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -80,6 +80,20 @@ def _parse_byte_size(value: str) -> int: return int(value) +def _b12x_pcie_oneshot_limits() -> tuple[int, int, int]: + allreduce_max_size = _parse_byte_size(envs.VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE) + fused_max_size = _parse_byte_size( + envs.VLLM_PCIE_ONESHOT_FUSED_ADD_RMS_NORM_MAX_SIZE + ) + if allreduce_max_size < 0 or fused_max_size < 0: + raise ValueError("b12x PCIe oneshot size limits must be non-negative") + # The pool only dispatches tensors within these limits. Reserving the + # scheduler's full prefill tensor capacity per stream wastes hundreds of + # MiB once target and draft graphs use independent channels. + buffer_size = max(allreduce_max_size, fused_max_size, 16) + return allreduce_max_size, fused_max_size, buffer_size + + @lru_cache(maxsize=1) def _load_b12x_pcie_oneshot_pool() -> Any | None: try: @@ -424,9 +438,8 @@ def __init__( "sparkinfer.comm.pcie.OneshotAllReducePool is unavailable." ) return - # The largest allreduce the model can issue (prefill chunk x - # hidden) determines buffer capacity. Dispatch crossovers come - # from envs below so startup does not benchmark the serving GPUs. + # DMA must accommodate the largest scheduled prefill tensor. The + # oneshot pool only needs its much smaller dispatch cutoffs. try: from vllm.config import get_current_vllm_config @@ -446,6 +459,11 @@ def __init__( ) pcie_buffer_size = max_num_batched_tokens * model_hidden_size * 2 pcie_single_channel = envs.VLLM_PCIE_ONESHOT_SINGLE_CHANNEL + ( + self._pcie_allreduce_max_size, + self._pcie_fused_add_rms_norm_max_size, + pcie_oneshot_buffer_size, + ) = _b12x_pcie_oneshot_limits() if self.nccl_group is None: logger.warning( "Custom allreduce is disabled because b12x PCIe oneshot " @@ -453,8 +471,6 @@ def __init__( ) return self.max_size = pcie_buffer_size - self._pcie_allreduce_max_size = 0 - self._pcie_fused_add_rms_norm_max_size = 0 self.rank = rank self.world_size = world_size self.fully_connected = False @@ -464,8 +480,8 @@ def __init__( pcie_runtime = pool_cls.from_exchange_group( exchange_group=self.nccl_group, device=self.device, - eager_buffer_bytes=pcie_buffer_size, - max_size=pcie_buffer_size, + eager_buffer_bytes=pcie_oneshot_buffer_size, + max_size=pcie_oneshot_buffer_size, single_channel=pcie_single_channel, ) pcie_runtime.for_stream() @@ -544,12 +560,6 @@ def __init__( self._pcie_dma = dma logger.debug("b12x PCIe DMA allreduce wire mode: %s", dma.wire_mode) - self._pcie_allreduce_max_size = _parse_byte_size( - envs.VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE - ) - self._pcie_fused_add_rms_norm_max_size = _parse_byte_size( - envs.VLLM_PCIE_ONESHOT_FUSED_ADD_RMS_NORM_MAX_SIZE - ) if rank == 0: logger.info( "Configured b12x PCIe crossovers: " @@ -562,12 +572,14 @@ def __init__( logger.debug( "Using b12x PCIe oneshot allreduce backend " "(world_size=%d, allreduce_max_size=%d, " - "fused_add_rms_norm_max_size=%d, buffer_size=%d, " + "fused_add_rms_norm_max_size=%d, oneshot_buffer_size=%d, " + "dma_buffer_size=%d, " "single_channel=%s).", world_size, self._pcie_allreduce_max_size, self._pcie_fused_add_rms_norm_max_size, - self.max_size, + pcie_oneshot_buffer_size, + pcie_buffer_size, pcie_single_channel, ) return diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 668e78ea03c8..b914700caa4f 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1508,10 +1508,20 @@ def graph_capture( if _DCP is not None and get_dcp_group().world_size > 1 else nullcontext() ) + if _DCP is not None and get_dcp_group().world_size > 1: + # Import locally to avoid making distributed initialization depend on + # attention modules. The helper is a no-op until DCP warmup creates a + # B12X pool for this process group. + from vllm.v1.attention.ops.dcp_alltoall import capture_b12x_dcp_a2a + + maybe_b12x_dcp_capture = capture_b12x_dcp_a2a(get_dcp_group(), context.stream) + else: + maybe_b12x_dcp_capture = nullcontext() with ( get_tp_group().graph_capture(context), get_pp_group().graph_capture(context), maybe_dcp_capture, + maybe_b12x_dcp_capture, ): yield context diff --git a/vllm/envs.py b/vllm/envs.py index a5f870bf875c..358064eaa3da 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -229,7 +229,7 @@ VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE: str = "84KB" VLLM_PCIE_ONESHOT_FUSED_ADD_RMS_NORM_MAX_SIZE: str = "84KB" VLLM_PCIE_ONESHOT_ALLOW_CROSS_NUMA: bool = True - VLLM_PCIE_ONESHOT_SINGLE_CHANNEL: bool = True + VLLM_PCIE_ONESHOT_SINGLE_CHANNEL: bool = False VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 VLLM_XGRAMMAR_CACHE_MB: int = 0 VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5 @@ -1815,9 +1815,10 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_PCIE_ONESHOT_ALLOW_CROSS_NUMA": lambda: ( os.getenv("VLLM_PCIE_ONESHOT_ALLOW_CROSS_NUMA", "1") != "0" ), - # Use a single channel for the b12x PCIe oneshot allreduce. + # Reuse one b12x PCIe oneshot channel across CUDA streams. This saves + # buffers but is unsafe when target and draft graphs replay concurrently. "VLLM_PCIE_ONESHOT_SINGLE_CHANNEL": lambda: ( - os.getenv("VLLM_PCIE_ONESHOT_SINGLE_CHANNEL", "1").strip().lower() + os.getenv("VLLM_PCIE_ONESHOT_SINGLE_CHANNEL", "0").strip().lower() not in ("", "0", "false", "no", "off") ), # Control the workspace buffer size for the FlashInfer backend. diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index dbd852c0647d..dd4c85275952 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -20,6 +20,7 @@ from __future__ import annotations +from contextlib import ExitStack, contextmanager from functools import lru_cache from typing import TYPE_CHECKING, Any @@ -113,7 +114,7 @@ def _get_b12x_dcp_a2a_pool( total_heads=total_heads, head_dim=head_dim, query_head_dim=query_head_dim, - single_channel=True, + single_channel=False, ) pool.for_stream() except Exception as exc: @@ -151,6 +152,27 @@ def _get_b12x_dcp_a2a_pool( return pool +@contextmanager +def capture_b12x_dcp_a2a( + cp_group: GroupCoordinator, + stream: object = None, +): + """Bind each CUDA graph manager to independent B12X DCP channels.""" + group_id = id(cp_group.device_group) + matching_pools = sorted( + ( + (key, pool) + for key, pool in _B12X_DCP_A2A_POOLS.items() + if key[0] == group_id + ), + key=lambda item: item[0][1:], + ) + with ExitStack() as stack: + for _, pool in matching_pools: + stack.enter_context(pool.capture(stream=stream)) + yield + + def _try_b12x_dcp_lse_reduce( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor,