From 2d710b88205bd6778d689dd3a8998e9c30a6fef0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 19 Jul 2026 15:12:03 +0000 Subject: [PATCH] fix(pcie): isolate nested CUDA graph capture channels --- b12x/distributed/pcie_dcp_a2a.py | 22 +++++++++++- b12x/distributed/pcie_oneshot.py | 26 ++++++++++---- tests/distributed/test_pcie_dcp_a2a.py | 49 ++++++++++++++++++++++++++ tests/distributed/test_pcie_oneshot.py | 45 +++++++++++++++++++++++ 4 files changed, 135 insertions(+), 7 deletions(-) diff --git a/b12x/distributed/pcie_dcp_a2a.py b/b12x/distributed/pcie_dcp_a2a.py index 4c762b82..8b4a94f1 100644 --- a/b12x/distributed/pcie_dcp_a2a.py +++ b/b12x/distributed/pcie_dcp_a2a.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from contextlib import suppress +from contextlib import contextmanager, suppress from dataclasses import dataclass from functools import lru_cache from pathlib import Path @@ -512,6 +512,7 @@ def __init__( self.single_channel = bool(single_channel) self._channel_factory = channel_factory self._channels: dict[int, PCIeDCPA2A] = {} + self._capture_channel_stack: list[PCIeDCPA2A] = [] self._closed = False if channel_factory is None and exchange_group is None: raise ValueError("exchange_group is required unless channel_factory is set") @@ -597,6 +598,13 @@ def for_stream(self, stream: object = None) -> PCIeDCPA2A: if channel is not None: return channel if _is_current_stream_capturing(self.device): + # Nested piecewise captures use a torch-owned stream key that is + # unavailable before capture begins. Reuse the channel selected by + # the enclosing graph manager, never one owned by another graph. + if self._capture_channel_stack: + channel = self._capture_channel_stack[-1] + self._channels[key] = channel + return channel if self._channels: channel = next(iter(self._channels.values())) self._channels[key] = channel @@ -665,6 +673,18 @@ def all_gather_heads( block_limit=block_limit, ) + @contextmanager + def capture(self, stream: object = None): + """Bind nested CUDA captures to the enclosing stream's channel.""" + channel = self.for_stream(stream) + self._capture_channel_stack.append(channel) + try: + yield channel + finally: + popped = self._capture_channel_stack.pop() + if popped is not channel: + raise RuntimeError("PCIe DCP A2A capture channel stack corrupted") + def close(self) -> None: if self._closed: return diff --git a/b12x/distributed/pcie_oneshot.py b/b12x/distributed/pcie_oneshot.py index e6f11558..390d8938 100644 --- a/b12x/distributed/pcie_oneshot.py +++ b/b12x/distributed/pcie_oneshot.py @@ -913,6 +913,7 @@ def __init__( self.single_channel = bool(single_channel) self._channel_factory = channel_factory self._channels: dict[int, PCIeOneshotAllReduce] = {} + self._capture_channel_stack: list[PCIeOneshotAllReduce] = [] self._closed = False self._ipc = ipc @@ -1045,11 +1046,18 @@ def for_stream(self, stream: object = None) -> PCIeOneshotAllReduce: if _is_current_stream_capturing(self.device): # Piecewise / inductor CUDA graphs (MTP, spec-decode) capture on a # torch-owned stream that we cannot pre-register before capture - # starts. Allocating a fresh channel here is impossible (it would - # touch CUDA APIs that are illegal mid-capture), so reuse an - # existing channel: the signal/IPC buffers are stream-agnostic and - # the all-reduce is recorded on the current (capturing) stream and - # replays on whatever stream the caller uses. + # starts. Reuse the channel selected by the enclosing vLLM graph + # capture, not an arbitrary channel from another graph manager. + # Target and draft graph managers can replay independently; if + # their nested captures share signal/staging buffers, they race. + if self._capture_channel_stack: + channel = self._capture_channel_stack[-1] + self._channels[channel_key] = channel + return channel + # Preserve compatibility for callers that enter CUDA capture + # without the pool.capture() context. They must already have a + # channel because allocating CUDA IPC storage mid-capture is + # illegal. if self._channels: channel = next(iter(self._channels.values())) self._channels[channel_key] = channel @@ -1111,7 +1119,13 @@ def run() -> tuple[torch.Tensor, torch.Tensor]: def capture(self, stream: object = None): channel = self.for_stream(stream) with channel.capture(stream=stream): - yield channel + self._capture_channel_stack.append(channel) + try: + yield channel + finally: + popped = self._capture_channel_stack.pop() + if popped is not channel: + raise RuntimeError("PCIe oneshot capture channel stack corrupted") def close(self) -> None: if self._closed: diff --git a/tests/distributed/test_pcie_dcp_a2a.py b/tests/distributed/test_pcie_dcp_a2a.py index c5bcb936..6346e3e9 100644 --- a/tests/distributed/test_pcie_dcp_a2a.py +++ b/tests/distributed/test_pcie_dcp_a2a.py @@ -5,6 +5,7 @@ from b12x.distributed.pcie_dcp_a2a import ( PCIeDCPA2A, + PCIeDCPA2APool, _staging_layout, lse_reduce_scatter_reference, ) @@ -202,3 +203,51 @@ def test_runtime_rejects_shape_dtype_and_capacity_mismatches(): runtime.all_gather_heads(good_output[:, :8]) with pytest.raises(ValueError, match="exceeds configured capacity"): runtime.all_gather_heads(torch.zeros(5, 16, 64, dtype=torch.bfloat16)) + + +def test_pool_uses_distinct_channels_for_target_and_draft_captures(monkeypatch): + created = [] + current_stream = [7] + capturing = [False] + + def make_channel(stream_key): + runtime = _make_runtime() + created.append((stream_key, runtime)) + return runtime + + pool = PCIeDCPA2APool( + rank=0, + world_size=2, + device=torch.device("cpu"), + max_batch_size=4, + total_heads=32, + head_dim=64, + channel_factory=make_channel, + ) + monkeypatch.setattr( + "b12x.distributed.pcie_dcp_a2a._current_stream_key", + lambda device, stream=None: ( + current_stream[0] if stream is None else int(stream) + ), + ) + monkeypatch.setattr( + "b12x.distributed.pcie_dcp_a2a._is_current_stream_capturing", + lambda device: capturing[0], + ) + + with pool.capture(7) as target_channel: + capturing[0] = True + current_stream[0] = 70 + assert pool.for_stream() is target_channel + capturing[0] = False + + with pool.capture(8) as draft_channel: + capturing[0] = True + current_stream[0] = 80 + assert pool.for_stream() is draft_channel + capturing[0] = False + + assert target_channel is not draft_channel + assert pool._channels[70] is target_channel + assert pool._channels[80] is draft_channel + assert [entry[0] for entry in created] == [7, 8] diff --git a/tests/distributed/test_pcie_oneshot.py b/tests/distributed/test_pcie_oneshot.py index ed009846..0906dd9b 100644 --- a/tests/distributed/test_pcie_oneshot.py +++ b/tests/distributed/test_pcie_oneshot.py @@ -660,3 +660,48 @@ def test_pool_requires_precreated_channel_during_capture(monkeypatch): pool._channels[7] = _make_runtime(eager=True) assert pool.for_stream() is pool._channels[7] + + +def test_nested_capture_reuses_its_outer_channel(monkeypatch): + created = [] + current_stream = [7] + capturing = [False] + + def make_channel(stream_key): + runtime = _make_runtime(eager=True) + created.append((stream_key, runtime)) + return runtime + + pool = PCIeOneshotAllReducePool( + rank=0, + world_size=2, + device=torch.device("cpu"), + channel_factory=make_channel, + ) + monkeypatch.setattr( + "b12x.distributed.pcie_oneshot._current_stream_key", + lambda device, stream=None: ( + current_stream[0] if stream is None else int(stream) + ), + ) + monkeypatch.setattr( + "b12x.distributed.pcie_oneshot._is_current_stream_capturing", + lambda device: capturing[0], + ) + + with pool.capture(7) as target_channel: + capturing[0] = True + current_stream[0] = 70 + assert pool.for_stream() is target_channel + capturing[0] = False + + with pool.capture(8) as draft_channel: + capturing[0] = True + current_stream[0] = 80 + assert pool.for_stream() is draft_channel + capturing[0] = False + + assert target_channel is not draft_channel + assert pool._channels[70] is target_channel + assert pool._channels[80] is draft_channel + assert [entry[0] for entry in created] == [7, 8]