From 931e6844989ba9c83386c62f25fb620741d75db0 Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:54:44 -0500 Subject: [PATCH] fix(dcp): support disposable A2A channel profiling --- sparkinfer/comm/pcie/pcie_dcp_a2a.py | 41 +++++++++++ tests/comm/test_pcie_dcp_a2a.py | 102 +++++++++++++++++++++++++++ 2 files changed, 143 insertions(+) diff --git a/sparkinfer/comm/pcie/pcie_dcp_a2a.py b/sparkinfer/comm/pcie/pcie_dcp_a2a.py index 9aa6daed..56d7b27e 100644 --- a/sparkinfer/comm/pcie/pcie_dcp_a2a.py +++ b/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 diff --git a/tests/comm/test_pcie_dcp_a2a.py b/tests/comm/test_pcie_dcp_a2a.py index 991ddc69..e26482c5 100644 --- a/tests/comm/test_pcie_dcp_a2a.py +++ b/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())