fix(moe): expose safe shared-expert overlap plans

[voipmonitor]

Summary

• expose an explicit execution-plan capability for shared-expert auxiliary-stream scheduling
• retain bounded native NVFP4 micro-decode overlap
• reject auxiliary overlap for resident split-decode plans
• expose the safe W4A16 pre-resident capability for 1-7 token decode
• keep all PCIe graph/channel lifecycle changes out of this PR

Why

Resident-grid MoE kernels can deadlock when auxiliary work occupies an SM required by a grid-wide barrier. The caller must query the concrete launch plan instead of treating every SparkInfer kernel as equally overlap-safe.

This PR contains only the MoE execution-model half of #58 and is based directly on master.

Validation

• tests/moe/test_moe_execution_model.py
• tests/moe/test_w4a8_tp_moe.py
• CUDA image run: 35 passed, 6 skipped
• CPU-only run: 28 passed, 13 skipped
• Ruff lint and format checks passed
• split integration tree is byte-identical to the MoE files in Fix PCIe graph lifecycle and preserve safe MoE overlap #58

Summary by CodeRabbit

• New Features
  • Added support for checking whether auxiliary-stream overlap is safe for MoE execution plans.
  • Added optional NVFP4 split-decode behavior, configurable through an environment setting.
• Bug Fixes
  • Improved stream-overlap safety for NVFP4 and W4A16 execution scenarios.
  • Verified consistent outputs when toggling NVFP4 split decode and during CUDA graph replay with auxiliary stream work.
• Tests
  • Expanded coverage across quantization modes, token sizes, and CUDA stream configurations.

[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: 5dab281d-c36b-427a-b8d5-1d9ef8bb677d

📥 Commits

Reviewing files that changed from the base of the PR and between ec2cd4d and 52379f0.

📒 Files selected for processing (5)

• sparkinfer/moe/fused_moe/__init__.py
• sparkinfer/moe/fused_moe/_impl.py
• sparkinfer/moe/fused_moe/api.py
• tests/moe/test_moe_execution_model.py
• tests/moe/test_w4a8_tp_moe.py

---

📝 Walkthrough

Walkthrough

The change adds environment-gated native NVFP4 split decode, reports auxiliary-stream overlap support for TP MoE plans, exposes the helper through the fused MoE API, and adds plan-level and CUDA graph replay tests.

Changes

NVFP4 auxiliary-stream overlap

Layer / File(s)	Summary
Split-decode gating and overlap support sparkinfer/moe/fused_moe/_impl.py	Centralized native NVFP4 predicates, gated barrier-free split decode with SPARKINFER_NVFP4_SPLIT_DECODE, and added TP MoE plan overlap detection.
Public API registration sparkinfer/moe/fused_moe/__init__.py, sparkinfer/moe/fused_moe/api.py	Registered and exported plan_supports_aux_stream_overlap for runtime and static-analysis access.
Plan and CUDA validation tests/moe/test_moe_execution_model.py, tests/moe/test_w4a8_tp_moe.py	Added coverage for overlap decisions, split-decode output equivalence, and CUDA graph replay with auxiliary-stream matrix work.

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

Suggested reviewers: lukealonso

Sequence Diagram(s)

sequenceDiagram
  participant Test
  participant plan_supports_aux_stream_overlap
  participant tp_moe_plan_supports_aux_stream_overlap
  participant sparkinfer_moe_fp4
  participant AuxiliaryStream
  Test->>plan_supports_aux_stream_overlap: evaluate TP MoE plan
  plan_supports_aux_stream_overlap->>tp_moe_plan_supports_aux_stream_overlap: delegate overlap check
  tp_moe_plan_supports_aux_stream_overlap-->>Test: return supported or rejected
  Test->>sparkinfer_moe_fp4: replay captured NVFP4 graph
  Test->>AuxiliaryStream: run concurrent torch.mm work
  sparkinfer_moe_fp4-->>Test: produce replay output

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 matches the main change: exposing safe overlap planning for MoE shared-expert execution.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/moe-shared-expert-overlap-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]

Closing after revalidating the original root-cause evidence.

The claimed shared-expert overlap failure was not reproduced. On both 2829061 and current 695c011, M=1/4/7 pass serial eager, serial CUDA graph, routed-first overlap, and aux-first overlap. M=8 differs even in pure serial eager execution with no graph and no auxiliary stream, while heavy auxiliary work stays within the same variation range. The exact M=4 stress workload also passes on both commits, and no deadlock was observed.

Therefore the M<=7 capability boundary and caller-facing overlap policy were inferred from a pre-existing M=8 numerical/race behavior rather than an overlap-specific failure. Commit 695c011 now expresses the all-CTA admission requirement at the kernel launch through a cooperative launch, which is the appropriate layer for that contract.

The M=8 behavior is being investigated separately against real checkpoint weights and an independent oracle.

[voipmonitor]

Final M=8 result from a real GLM-5.2 NVFP4 layer (K=6144, local N=256, E=256, top-k=8): this case does not execute fused micro. The production planner selects micro for M=7 (56 routed rows) and dynamic for M=8 (64 routed rows, at the strict cutover).

M=7 was bitwise identical across all repeats and matched the independent oracle at cosine 0.9999105. M=8 used the default BF16 atomic scatter of the dynamic kernel; its expected reduction-order variation was small (pairwise cosine about 0.999977, max absolute difference 0.000366) and was present in serial execution. Setting SPARKINFER_DYNAMIC_DETERMINISTIC_OUTPUT=1 made all eight M=8 repeats bitwise identical while retaining the same oracle accuracy (cosine 0.9995970, max absolute error 0.0006714).

Therefore there is no remaining evidence of a shared-expert auxiliary-stream race or corruption in this report. The M=8 signal came from comparing different kernel families and requiring exact equality from the dynamic atomic-output path. PR #71 adds a CPU planner regression test that pins the real GLM M=7/M=8 dispatch boundary.