[GG] fix(pcie): isolate target and draft graph channels

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

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

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,15 +459,18 @@ 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 "
                     "allreduce requires a CUDA/NCCL device process group."
                 )
                 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

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
 

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.

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,