fix(pcie): isolate nested CUDA graph capture channels

[voipmonitor]

Summary

• bind torch/Inductor nested CUDA capture streams to the B12X channel selected by their enclosing graph capture
• keep target and MTP/draft graph managers on distinct PCIe oneshot signal and staging buffers
• apply the same capture ownership model to the PCIe DCP A2A pool
• preserve legacy fallback behavior for captures entered without a pool capture context
• add CPU regression coverage for separate target/draft outer and nested stream pairs in both pools

Root cause

A communication channel owns mutable signal counters and double-buffered IPC staging memory. The old nested-capture fallback selected the first existing channel because allocating during CUDA capture is illegal. Independent target and MTP graph managers could therefore capture one shared channel and replay it concurrently, corrupting its state.

The pools now maintain an enclosing capture-channel stack. A torch/Inductor-owned nested capture stream aliases only the channel selected by its own outer graph manager. Target and draft captures consequently remain isolated without adding a runtime copy or synchronization.

Validation

• tests/distributed/test_pcie_oneshot.py and tests/distributed/test_pcie_dcp_a2a.py: 32 passed
• Ruff and git diff --check: passed
• GLM-5.2 TP6/DCP6/MTP3 with vLLM multichannel wiring survived the previously deterministic CC4 failure for both original and online-MXFP8 configurations

Integration

• Pair with [GG] fix(concurrency): isolate graph state and gate B12X MoE overlap vllm#132, which changes the unsafe oneshot default, right-sizes per-stream buffers, makes DCP A2A multichannel, and enters these pool contexts during graph capture.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 3c1071d7-7469-4f02-8bb3-bb76f43904bb

📥 Commits

Reviewing files that changed from the base of the PR and between 21e508e and 2d710b8.

📒 Files selected for processing (4)

• b12x/distributed/pcie_dcp_a2a.py
• b12x/distributed/pcie_oneshot.py
• tests/distributed/test_pcie_dcp_a2a.py
• tests/distributed/test_pcie_oneshot.py

---

Walkthrough

PCIeDCPA2APool and PCIeOneshotAllReducePool now preserve channel selection across nested CUDA graph captures, register reused channels under current stream keys, validate capture-stack integrity, and test separate and nested capture scenarios.

Changes

PCIe capture channel reuse

Layer / File(s)	Summary
Pool capture channel behavior b12x/distributed/pcie_dcp_a2a.py, b12x/distributed/pcie_oneshot.py	Both pools track active capture channels, reuse the enclosing channel during capture, register it under the current stream key, and validate stack push/pop consistency.
Capture channel validation tests/distributed/test_pcie_dcp_a2a.py, tests/distributed/test_pcie_oneshot.py	Tests verify distinct channels for separate captures and reuse of the outer channel during nested captures.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant PCIePool
  participant ChannelFactory
  Caller->>PCIePool: enter capture(stream)
  PCIePool->>ChannelFactory: create or retrieve channel
  PCIePool->>PCIePool: push channel onto capture stack
  Caller->>PCIePool: for_stream()
  PCIePool->>PCIePool: reuse top captured channel
  PCIePool-->>Caller: return captured channel
  Caller->>PCIePool: exit capture
  PCIePool->>PCIePool: pop and validate channel

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: isolating nested CUDA graph capture channels for PCIe oneshot behavior.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}