fix(pcie): isolate nested CUDA graph capture channels

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

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:

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]

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]