fix(concurrency): isolate capture channels and barrier-free NVFP4 decode#47
Conversation
|
Warning Review limit reached
Next review available in: 31 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughThe changes add phased PCIe IPC cleanup, channel checkpoint/rollback, capture-channel reuse handling, and unique pool shutdown. They also add bounded NVFP4 auxiliary-stream overlap planning, opt-in split-decode gating, public exports, and CUDA execution coverage. ChangesPCIe capture channel lifecycle
NVFP4 overlap planning
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CaptureContext
participant PCIeOneshotAllReducePool
participant PCIeOneshotAllReduce
CaptureContext->>PCIeOneshotAllReducePool: checkpoint channel state
CaptureContext->>PCIeOneshotAllReducePool: capture(stream_key)
PCIeOneshotAllReducePool->>PCIeOneshotAllReduce: allocate and register channel
CaptureContext->>PCIeOneshotAllReducePool: rollback channel state
PCIeOneshotAllReducePool->>PCIeOneshotAllReduce: close transient IPC resources
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
5d7fe0f to
dc7faa3
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
b12x/distributed/pcie_oneshot.py (1)
871-897: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMissing exception suppression in
_close_ipc_imports/_free_ipc_exports— can deadlock the coordinated rollback and enable double-free on retry.The sibling implementation in
b12x/distributed/pcie_dcp_a2a.py(_close_ipc_imports/_free_ipc_exports, lines 466-494) wraps everydispose/cudaIpcCloseMemHandle/cudaFreecall inwith suppress(Exception). This file's equivalents do not:
- If
cudaIpcCloseMemHandle/cudaFreeraises partway through the loop,self._ipc_exports_freed/self._ipc_imports_closednever get set and_owned_buffersis never cleared, so a later retry would attempt to free already-freed pointers (double free/UB).- Worse,
_coordinated_close_channels()(Lines 150-173) callschannel._close_ipc_imports()/channel._free_ipc_exports()in a loop with no per-channel try/except. An exception on one channel aborts the loop before this rank reaches the nextdist.barrier(), while peer ranks that succeeded will block on that barrier forever — a distributed hang.🐛 Suggested fix (mirrors pcie_dcp_a2a.py)
def _close_ipc_imports(self) -> None: if self._ipc_imports_closed: return self._closed = True if getattr(self, "_ptr", 0): - self._ext.dispose(self._ptr) + with suppress(Exception): + self._ext.dispose(self._ptr) self._ptr = 0 - for shared in self._owned_buffers: - for ptr in shared.remote_ptrs: - if self._ipc is not None: - self._ipc.cudaIpcCloseMemHandle(ptr) + if self._ipc is not None: + for shared in self._owned_buffers: + for ptr in shared.remote_ptrs: + with suppress(Exception): + self._ipc.cudaIpcCloseMemHandle(ptr) self._ipc_imports_closed = True def _free_ipc_exports(self) -> None: if self._ipc_exports_freed: return self._close_ipc_imports() - for shared in self._owned_buffers: - if self._ipc is not None: - self._ipc.cudaFree(shared.local_ptr) + if self._ipc is not None: + for shared in self._owned_buffers: + with suppress(Exception): + self._ipc.cudaFree(shared.local_ptr) self._owned_buffers.clear() self._registered_input_ptrs.clear() self._ipc_exports_freed = True
🧹 Nitpick comments (2)
b12x/distributed/pcie_dcp_a2a.py (1)
774-785: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
close()skips the coordinated IPC teardown used elsewhere in this PR.
rollback_channels()correctly routes through_coordinated_close_channels()(barriers + phased close-imports/free-exports) to avoid a rank freeing its exported buffer before peers unmap it.PCIeDCPA2APool.close()still closes every channel directly viachannel.close()with no barriers, so full-pool shutdown across ranks can hit the exact same cross-rank use-after-free race this PR fixes for capture rollback.♻️ Suggested fix
def close(self) -> None: if self._closed: return self._closed = True - seen: set[int] = set() - for channel in (*self._all_channels, *self._channels.values()): - if id(channel) not in seen: - seen.add(id(channel)) - channel.close() - self._all_channels.clear() + channels = tuple(dict.fromkeys((*self._all_channels, *self._channels.values()))) + _coordinated_close_channels( + channels, exchange_group=self.exchange_group, device=self.device + ) + self._all_channels.clear() self._channels.clear()See consolidated comment for the matching issue in
pcie_oneshot.py.b12x/distributed/pcie_oneshot.py (1)
1241-1252: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
close()skips the coordinated IPC teardown used byrollback_channels.Same root cause as the equivalent gap in
pcie_dcp_a2a.py:PCIeOneshotAllReducePool.close()closes each channel directly (channel.close()) with no barriers, whereasrollback_channels()uses_coordinated_close_channels()specifically to avoid one rank freeing its exported buffer before peers unmap it. Full-pool shutdown reintroduces that same cross-rank race.See consolidated comment for the shared fix across both files.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 58027cac-5e97-4dea-85ce-20ad280ec245
📒 Files selected for processing (8)
b12x/distributed/pcie_dcp_a2a.pyb12x/distributed/pcie_oneshot.pyb12x/integration/__init__.pyb12x/integration/tp_moe.pytests/distributed/test_pcie_dcp_a2a.pytests/distributed/test_pcie_oneshot.pytests/test_moe_execution_model.pytests/test_w4a8_tp_moe.py
🚧 Files skipped from review as they are similar to previous changes (5)
- b12x/integration/init.py
- tests/test_w4a8_tp_moe.py
- b12x/integration/tp_moe.py
- tests/distributed/test_pcie_dcp_a2a.py
- tests/test_moe_execution_model.py
|
Addressed the coordinated-teardown review finding in dee3b5e: every dispose/close/free is exception-isolated, terminal state is always recorded, cleanup remains idempotent, and a regression test verifies all resources are attempted even when every CUDA IPC operation raises. Validation: 32/32 pcie_oneshot tests passed. |
|
Superseded by #58, rebuilt directly on current master with the namespaced sparkinfer APIs, coordinated graph-channel teardown, regression tests, and the corrected W4A16 pre-resident overlap contract. |
1 participant
Summary
B12X_NVFP4_SPLIT_DECODE=0as a diagnostic fallbackRoot causes
Two independent concurrency hazards surfaced under target/draft CUDA graphs and GLM shared-expert overlap. Recycled capture streams could inherit a channel selected for a prior capture owner. Separately, the fused B12X micro MoE launch uses a resident-grid software barrier; co-scheduling unrelated work can prevent every barrier participant from becoming resident.
Disposable profiling graphs add a lifetime constraint: all ranks must destroy graph users and unmap imported CUDA IPC allocations before any rank frees the corresponding exported slab. Rank-local close/free could otherwise leave a peer with a stale mapping. The rollback path now uses a three-barrier teardown protocol and restores only the pre-profile channel map.
The existing W4A8-on-NVFP4 two-phase launch already removed the micro barrier. This PR applies the same split to native NVFP4 A4 and exports the selected launch-plan capability. Dynamic and W4A16 resident launches remain correctly marked unsafe for external overlap.
Validation
barrier -> close imports -> barrier -> free exports -> barrierorder for both poolsgit diff --checkpassedIntegration
Pair with local-inference-lab/vllm#131 for disposable profiling-graph cleanup and local-inference-lab/vllm#132 for stream/workspace isolation and MoE scheduling.
Summary by CodeRabbit