Fix PCIe graph lifecycle and preserve safe MoE overlap

[voipmonitor]

Summary

• isolate recycled CUDA graph streams and workspaces from live communication state
• roll back disposable graph channels and coordinate IPC teardown across ranks
• retain the barrier-free native NVFP4 split-decode path
• expose an explicit plan capability for shared-expert auxiliary-stream use
• allow W4A16 decode to use the safe pre-resident overlap contract for 1-7 tokens
• port the concurrency and execution-model regressions to the namespaced sparkinfer APIs

Root cause

The previous stack mixed two separate concerns. Recycled graph resources needed strict ownership and coordinated teardown, while the first safety response disabled W4A16 auxiliary-stream use entirely. That blanket disable prevented resident-grid deadlocks but also removed useful overlap from the common DCP1 A16 decode path.

The corrected contract is narrower: shared-expert work may overlap gate/router work, but the consumer stream must join the auxiliary stream before submitting a resident routed-MoE grid. A positive plan capability therefore does not authorize concurrent resident-grid execution.

Validation

• SparkInfer targeted suite: 76 passed, 6 skipped
• paired vLLM targeted suite: 134 passed
• DCP1 A16 decode A/B:
  • blanket W4A16 disable: 83.565 tok/s mean
  • corrected pre-resident overlap: 86.458 tok/s mean
  • recovery: +3.46%

This supersedes #47 with a clean branch directly based on master. The separate cuBLAS head-major safety work remains in #54.

Summary by CodeRabbit

• New Features

  • Added support for checking auxiliary-stream overlap capabilities in fused MoE execution plans.
  • Improved CUDA Graph capture channel handling, including safer reuse and rollback of temporary channels.
  • Enhanced PCIe communication resource cleanup and coordinated teardown.

• Bug Fixes

  • Prevented stale channel reuse across captures.
  • Improved cleanup reliability when resource-release operations encounter errors.
  • Strengthened NVFP4 split-decode and auxiliary-stream behavior across supported execution modes.

• Tests

  • Added coverage for graph capture, rollback, coordinated cleanup, NVFP4 execution, and auxiliary-stream concurrency.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: aaf34c6f-f14d-41eb-a313-66ce85c71afe

📥 Commits

Reviewing files that changed from the base of the PR and between ec2cd4d and 73c115f.

📒 Files selected for processing (9)

• sparkinfer/comm/pcie/pcie_dcp_a2a.py
• sparkinfer/comm/pcie/pcie_oneshot.py
• sparkinfer/moe/fused_moe/__init__.py
• sparkinfer/moe/fused_moe/_impl.py
• sparkinfer/moe/fused_moe/api.py
• tests/comm/test_pcie_dcp_a2a.py
• tests/comm/test_pcie_oneshot.py
• tests/moe/test_moe_execution_model.py
• tests/moe/test_w4a8_tp_moe.py

---

📝 Walkthrough

Walkthrough

PCIe pools now support phased IPC teardown, channel checkpoint/rollback, and capture-specific channel ownership. Fused MoE adds auxiliary-stream overlap capability detection, updates NVFP4 split-decode gating, exports the capability API, and adds coverage for plan decisions and CUDA Graph replay.

Changes

PCIe channel lifecycle

Layer / File(s)	Summary
Phased IPC channel teardown sparkinfer/comm/pcie/pcie_dcp_a2a.py, sparkinfer/comm/pcie/pcie_oneshot.py	Channel cleanup separates IPC import closure from local IPC export freeing and coordinates cross-rank teardown.
Channel checkpoint and rollback sparkinfer/comm/pcie/pcie_dcp_a2a.py, sparkinfer/comm/pcie/pcie_oneshot.py	Pools retain all created channels, snapshot channel ownership, restore mappings, and close transient channels created after checkpoints.
Capture channel routing sparkinfer/comm/pcie/pcie_dcp_a2a.py, sparkinfer/comm/pcie/pcie_oneshot.py	Nested captures reuse the active capture channel, while multi-channel captures register dedicated graph-owned channels.
Lifecycle behavior validation tests/comm/test_pcie_dcp_a2a.py, tests/comm/test_pcie_oneshot.py	Tests cover distinct capture channels, rollback cleanup, teardown ordering, capture-time rejection, and error-tolerant cleanup.

MoE auxiliary-stream overlap

Layer / File(s)	Summary
Overlap capability and split-decode gating sparkinfer/moe/fused_moe/_impl.py, sparkinfer/moe/fused_moe/api.py, sparkinfer/moe/fused_moe/__init__.py	Native NVFP4 split-decode predicates determine barrier-free execution, and plan_supports_aux_stream_overlap is exported.
Overlap and graph replay validation tests/moe/test_moe_execution_model.py, tests/moe/test_w4a8_tp_moe.py	Tests validate capability decisions, split-decode equivalence, and CUDA Graph replay with auxiliary-stream work.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant MoEPlan
  participant FusedMoEAPI
  participant NVFP4Kernel
  participant AuxStream
  MoEPlan->>FusedMoEAPI: query plan_supports_aux_stream_overlap
  FusedMoEAPI->>NVFP4Kernel: evaluate split-decode eligibility
  NVFP4Kernel->>AuxStream: permit auxiliary-stream overlap
  AuxStream->>NVFP4Kernel: execute concurrent GPU work
  NVFP4Kernel->>MoEPlan: produce replayed output

Loading

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 10.87% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the two main changes: PCIe graph lifecycle fixes and safer MoE overlap behavior.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/namespaced-pcie-moe-concurrency-20260721

---

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.}

[voipmonitor]

Paired vLLM integration: local-inference-lab/vllm#149. Both replacement PRs are based directly on their canonical branches; neither is stacked on a temporary fix branch.

[voipmonitor]

This broad PR has been split by ownership boundary:

• fix(pcie): isolate CUDA graph channel lifecycle #59: PCIe CUDA graph channel lifecycle and coordinated teardown
• fix(moe): expose safe shared-expert overlap plans #60: MoE execution-plan capability and shared-expert overlap safety

Both replacements target master directly, are independently tested, and their combined tree is byte-identical to this PR across all nine changed files. Closing this superseded aggregate PR without merging it.