fix(dcp): support disposable A2A channel profiling

sparkinfer/comm/pcie/pcie_dcp_a2a.py

@@ -480,6 +480,12 @@ def __del__(self) -> None:
             self.close()
 
 
+@dataclass(frozen=True)
+class _ChannelCheckpoint:
+    pool_id: int
+    channels: tuple[tuple[int, PCIeDCPA2A], ...]
+
+
 class PCIeDCPA2APool:
     """Create an independent DCP collective channel for each CUDA stream."""
 
@@ -685,6 +691,41 @@ def capture(self, stream: object = None):
             if popped is not channel:
                 raise RuntimeError("PCIe DCP A2A capture channel stack corrupted")
 
+    def checkpoint_channels(self) -> _ChannelCheckpoint:
+        """Snapshot channels before a disposable CUDA graph capture."""
+        if self._closed:
+            raise RuntimeError("PCIeDCPA2APool is closed")
+        if self._capture_channel_stack:
+            raise RuntimeError("cannot checkpoint channels during CUDA graph capture")
+        return _ChannelCheckpoint(
+            pool_id=id(self),
+            channels=tuple(self._channels.items()),
+        )
+
+    def rollback_channels(self, checkpoint: _ChannelCheckpoint) -> None:
+        """Release channels created after ``checkpoint`` and restore aliases."""
+        if self._closed:
+            raise RuntimeError("PCIeDCPA2APool is closed")
+        if self._capture_channel_stack:
+            raise RuntimeError("cannot roll back channels during CUDA graph capture")
+        if checkpoint.pool_id != id(self):
+            raise ValueError("channel checkpoint belongs to a different pool")
+
+        saved_channels = dict(checkpoint.channels)
+        saved_channel_ids = {id(channel) for channel in saved_channels.values()}
+        new_channels: list[PCIeDCPA2A] = []
+        seen: set[int] = set()
+        for channel in self._channels.values():
+            channel_id = id(channel)
+            if channel_id not in saved_channel_ids and channel_id not in seen:
+                seen.add(channel_id)
+                new_channels.append(channel)
+
+        self._channels.clear()
+        self._channels.update(saved_channels)
+        for channel in new_channels:
+            channel.close()
+
     def close(self) -> None:
         if self._closed:
             return

tests/comm/test_pcie_dcp_a2a.py

@@ -5,6 +5,7 @@
 
 from sparkinfer.comm.pcie.pcie_dcp_a2a import (
     PCIeDCPA2A,
+    PCIeDCPA2APool,
     _staging_layout,
     lse_reduce_scatter_reference,
 )
@@ -202,3 +203,104 @@ 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_rollback_closes_only_channels_created_after_checkpoint(monkeypatch):
+    created = []
+    current_stream = [7]
+    capturing = [False]
+
+    def make_channel(stream_key):
+        ext = _FakeExt()
+        runtime = _make_runtime(ext)
+        created.append((stream_key, runtime, ext))
+        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(
+        "sparkinfer.comm.pcie.pcie_dcp_a2a._current_stream_key",
+        lambda device, stream=None: (
+            current_stream[0] if stream is None else int(stream)
+        ),
+    )
+    monkeypatch.setattr(
+        "sparkinfer.comm.pcie.pcie_dcp_a2a._is_current_stream_capturing",
+        lambda device: capturing[0],
+    )
+
+    target_channel = pool.for_stream(7)
+    checkpoint = pool.checkpoint_channels()
+
+    with pool.capture(8) as draft_channel:
+        capturing[0] = True
+        current_stream[0] = 80
+        assert pool.for_stream() is draft_channel
+        capturing[0] = False
+
+    pool.rollback_channels(checkpoint)
+
+    assert pool._channels == {7: target_channel}
+    assert created[0][2].dispose_calls == []
+    assert created[1][2].dispose_calls == [1234]
+
+
+def test_pool_rollback_removes_new_alias_without_closing_saved_channel(monkeypatch):
+    ext = _FakeExt()
+    pool = PCIeDCPA2APool(
+        rank=0,
+        world_size=2,
+        device=torch.device("cpu"),
+        max_batch_size=4,
+        total_heads=32,
+        head_dim=64,
+        channel_factory=lambda stream_key: _make_runtime(ext),
+    )
+    monkeypatch.setattr(
+        "sparkinfer.comm.pcie.pcie_dcp_a2a._current_stream_key",
+        lambda device, stream=None: 70 if stream is None else int(stream),
+    )
+    capturing = [False]
+    monkeypatch.setattr(
+        "sparkinfer.comm.pcie.pcie_dcp_a2a._is_current_stream_capturing",
+        lambda device: capturing[0],
+    )
+
+    channel = pool.for_stream(7)
+    checkpoint = pool.checkpoint_channels()
+    with pool.capture(7):
+        capturing[0] = True
+        assert pool.for_stream() is channel
+        capturing[0] = False
+
+    assert pool._channels == {7: channel, 70: channel}
+    pool.rollback_channels(checkpoint)
+
+    assert pool._channels == {7: channel}
+    assert ext.dispose_calls == []
+
+
+def test_pool_rejects_checkpoint_from_another_pool():
+    def make_pool():
+        return PCIeDCPA2APool(
+            rank=0,
+            world_size=2,
+            device=torch.device("cpu"),
+            max_batch_size=4,
+            total_heads=32,
+            head_dim=64,
+            channel_factory=lambda stream_key: _make_runtime(),
+        )
+
+    first = make_pool()
+    second = make_pool()
+
+    with pytest.raises(ValueError, match="different pool"):
+        second.rollback_channels(first.checkpoint_channels())