[Blackwell] Reserve SMs for DeepGEMM MegaMoE grid barriers

[weireweire]

Summary

• Reserve two SMs by default for Blackwell DeepGEMM MegaMoE launches.
• Scope the process-wide SM override to the actual MegaMoE kernel calls.
• Keep symmetric-buffer caching independent of the launch SM count.
• Keep Hopper/SM90 behavior unchanged and allow the reserve to be tuned or disabled.

Root cause

Blackwell MegaMoE uses an even clustered grid with a whole-grid software barrier. If the launch uses every available SM while another CUDA stream temporarily occupies an SM, not all clusters can become resident. Resident clusters then wait at the barrier for clusters that cannot launch, eventually producing a DeepGEMM grid-sync timeout and a subsequent CUDA launch failure.

Fix

For SM100 and newer, derive the MegaMoE launch size from DeepGEMM's current SM count and reserve two SMs by default. Round the active count down to an even number so it remains compatible with the two-CTA cluster.

Apply the override only around the DeepSeek and Kimi K3 deep_gemm.fp8_fp4_mega_moe calls, then restore DeepGEMM's process-wide setting even when launch setup raises. Gate computation, top-k selection, symmetric-buffer setup, and pre-dispatch continue to use the original setting. The symmetric-buffer cache key remains based only on parameters that affect buffer size or layout.

The separate SM90 MegaMoE implementation remains unchanged. This supersedes the closed #30592 with a smaller diff and explicit architecture gating.

Fixes #30399.

Validation

Validation	Full SMs / reserve=0	Reserve=2
CPU unit tests	—	Passed: default reserve, even rounding, zero/excessive reserve, exception restoration, and SM90 no-op
GB300 side-stream repro	MegaMoE remained blocked while the side stream was active	MegaMoE completed while the side stream was still active
Kernel performance, 100 iterations	0.059546 ms	0.059393 ms
Overlap-scheduler E2E	Reproduced DeepGEMM grid-sync timeout	12,748/12,748 warmup requests; 11,052 profiled requests; 0 errors; ~126.7K tok/s/GPU
CUDA Graph/SBO E2E	Reproduced liveness failure	12,345 requests; 0 errors; ~69.5K tok/s/GPU
Static/CI checks	—	compileall, pre-commit, lint, build, and Base-A CPU CI passed

---

CI States

Latest PR Test (Base): ⏳ Run #33208972574
Latest PR Test (Extra): ⏳ Run #33208972664
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33208972751

[weireweire]

/tag-and-rerun-ci

[BBuf]

Could we scope this SM override to the actual deep_gemm.fp8_fp4_mega_moe call?

set_num_sms is process-wide, and the current context also covers the gate, top-k, buffer setup, and pre-dispatch. Any other DeepGEMM call in this path would inherit the reduced SM count, which seems broader than this fix requires.

[weireweire]

make sense, changed to only guard fp8_fp4_mega_moe

[BBuf]

Does the symmetric buffer actually depend on num_sms?

In sgl-deep-gemm 0.1.5.post3, its size and layout are derived from the rank, token, and model dimensions, not the active SM count. If there is no hidden dependency, can we remove this cache-key field and the related argument plumbing? Otherwise changing the reserve can allocate and rendezvous another large symmetric buffer unnecessarily.

[weireweire]

thanks, this num_sms won't change after setting anyway, changed.

[BBuf]

The fix direction makes sense. Before merging, could we carry the Blackwell validation from #30592 over to the current PR head?

At minimum, I would like to see:

• the side-stream repro with full SMs versus a 2-SM reserve;
• an overlap-scheduler E2E run confirming that the timeout is gone;
• a throughput A/B for reserve=0 versus reserve=2;
• CUDA Graph/SBO coverage.

The current PR only reports CPU unit tests, while this changes a GPU liveness condition and the default MegaMoE launch size.

[weireweire]

@BBuf thanks for review, test is updated on pr decripetion.

[BBuf]

LGTM