Specialize SM90 MoE communication and fix backward synchronization

[kolehma8]

Summary

• Specialize SM90 MoE communication for local-only and inter-host execution so local kernels do not carry unnecessary IBGDA code/register pressure, while retaining IBGDA for inter-host communication.
• Fix the SM90 backward hangs: publish the split MLP3 loop bounds and perform the producer/consumer rendezvous before the blocking producer loop; advance shared cursors even when an empty expert range or K slice takes an early continue.
• Add the external non-RDC SM90 module build, packaging, dispatch and compatibility checks, including the all-configuration module and 1/2/4/8/16-PE coverage.
• Extend TMA/routing coverage and ensure a rejected or incompletely initialized external module is unloaded through scoped ownership.

Communication semaphore counts and acquire/release semantics are preserved. The existing SM100-family GPU implementation is unchanged.

MoEBenchmark results: H200 and B300

Protocol: CUDA 12.9, BF16 routed-expert MoE MLP, eight microbatches, 250 ms per-timed-launch cooldown. The summary below uses T=8192 per rank and reports throughput in TFLOP/s per GPU. Sixteen GPUs means two hosts with eight GPUs each. Shared experts and attention are not included in the benchmark FLOP count.

GPU	GPUs	Llama 4 Scout forward	Llama 4 Scout backward	Qwen3-30B-A3B forward	Qwen3-30B-A3B backward
H200	1	511.36	299.54	354.38	199.28
H200	2	461.73	319.30	291.08	222.21
H200	4	473.09	337.02	255.68	217.06
H200	8	461.04	340.12	232.64	209.07
H200	16	318.45	278.54	50.64	74.71
B300	1	977.74	723.75	599.60	376.85
B300	2	939.54	835.58	547.99	396.22
B300	4	1004.79	955.76	517.30	406.17
B300	8	886.71	714.04	467.82	402.50
B300	16	696.53	647.42	172.13	215.10

Representative backend comparisons

Each cell below is forward / backward TFLOP/s per GPU for Qwen3-30B-A3B at T=8192.

GPU	GPUs	TE	DeepEP + TE	Comet	FlashMoE	Liger
H200	8	111.37 / 137.38	100.17 / 158.89	156.37 / N/A	102.37 / N/A	232.64 / 209.07
H200	16	12.70 / 30.63	N/A	N/A	27.44 / N/A	50.64 / 74.71
B300	8	192.26 / 182.09	94.53 / 178.85	N/A	N/A	467.82 / 402.50
B300	16	71.46 / 126.43	N/A	N/A	N/A	172.13 / 215.10

This is not a claim of a uniform win across models. For example, at eight GPUs, Scout backward is 340.12 versus TE's 443.46 TFLOP/s on H200, and 714.04 versus 805.15 on B300. At sixteen B300 GPUs, Scout backward is close to TE: 647.42 versus 653.07 TFLOP/s.

Coverage and limitations

• Seven model families: Llama 4 Scout, Mixtral 8x7B/8x22B, Qwen3-30B-A3B/235B-A22B, and Qwen3.5-35B-A3B/122B-A10B.
• H200 named-model campaign: T=8192, 1/2/4/8/16 GPUs, 238 attempts: 228 numeric results and ten expected expert-divisibility skips.
• B300 named-model campaign: T=1024/2048/4096/8192/16384, 1/2/4/8/16 GPUs, 910 attempts: 870 numeric results and forty expected expert-divisibility skips.
• Mixtral's eight experts cannot be evenly distributed over sixteen expert-parallel ranks.
• DeepEP is included at 2/4/8 GPUs; its sixteen-GPU transport path failed in the preserved mainline baseline and is excluded. Comet is forward-only and single-host in this setup. FlashMoE is forward-only.
• B300 Comet/Flux is unavailable in the pinned source's architecture support. B300 FlashMoE compilation fails against the pinned CUDA 12.9/CCCL interfaces. These cells are N/A, not substituted measurements.
• The final sixteen-GPU Liger results use IBGDA with GPU-local HCA affinity: four usable compute NICs per H200 host and eight per B300 host. Local runs use NVLink; the NIC correction does not change their protocol.
• Kernel-managed RDMA QP counters do not cover IBGDA's direct QPs. Per-rank HCA assignment and transport initialization are retained in provenance, but per-rail traffic volume and aggregate network bandwidth were not measured.

The named-model campaigns above are separate from the completed fixed
eight-GPU sweeps below.

Additional eight-GPU sweeps

• H200 token sweep: seven named-model families at
  T=1024/2048/4096/8192/16384. The final dataset contains 224 new attempts plus
  56 accepted T=8192 rows from the immutable named-model campaign: 280 rows,
  ten LaTeX tables, 63 PNG/PDF/EPS plots, and zero missing cases.
• H200 parameter grid: T=8192, D=I in {2048,4096},
  E in {8,16,32,64,128}, and K in {1,2,4,8} with K<E: 38 configurations and
  304/304 finite-positive attempts. Forward includes TE, DeepEP, Comet,
  FlashMoE and Liger; backward includes TE, DeepEP and Liger. Liger is fastest
  in 38/38 forward rows with a 1.388x geometric-mean speedup over the fastest
  comparison backend. It is fastest in 29/38 backward rows with a 1.103x
  geometric-mean speedup over the fastest comparison backend.
• B300 parameter grid: the same 38 configurations, with TE, DeepEP and
  Liger in both directions: 228/228 finite-positive attempts. Liger is fastest
  in 38/38 forward rows and 38/38 backward rows, with geometric-mean speedups
  of 1.880x forward and 1.406x backward over the fastest comparison backend.
  Both generated tables contain only the three deployed B300 backends; Comet
  and FlashMoE are omitted for the availability reasons documented above.

SM100 regression comparison

A separate, protocol-matched historical comparison uses one microbatch and 250 ms cooldown on both mainline and candidate. Across 330 numeric Liger comparisons, GPU-count/direction geometric-mean throughput ratios range from approximately 0.9955 to 1.0090. The largest individual decrease is 4.81% for eight-GPU Qwen3-30B-A3B at T8192; no individual comparison is more than 5% lower. These are single-run comparisons, not confidence intervals or a claim that every measurement improved.

The candidate's per-host benchmark/core hashes match the SM100 artifacts used for the final campaign. Both sides of the historical sixteen-GPU comparison used the same older four-NIC mapping; those historical values are not mixed into the corrected-affinity, eight-microbatch table above. The comparison shows no broad SM100 throughput loss, while keeping protocol changes and run-to-run variation explicit.

Testing Done

Hardware: NVIDIA H200 (SM90) and NVIDIA B300 SXM6 (SM100-family, sm_103a).

• Complete SM90 native core and all-configuration external cubin built with CUDA 12.9.
• Local and actual two-host forward/backward reference runs completed, including the formerly hanging backward path.
• Named-model campaigns, the H200 token sweep, and the H200/B300 fixed
  eight-GPU parameter grids completed with the coverage above.
• test_cute_build.py and test_export_nonrdc_nvshmem_ptx.py: thirteen tests passed,
  including duplicate-symbol, missing-SETMAXNREG, and identical-path rejection coverage.
• test_moe_bindings.py: six tests passed and the explicit NVSHMEM-runtime placeholder
  remained skipped; changed topology is rejected before either schedule is rewritten.
• test_moe_fwd_bwd_cuda_graph: dense and sparse routing passed with TMA enabled and
  with the IB-capable getmem fallback forced (four cases total, two H200 GPUs).
• Ruff lint and format checks passed for all seven Python files in this change.
• The review follow-up all-configuration SM90 wheel compiled with the existing CUDA
  12.9 Release flags in 44m16s. It removes the dead same-host forward copy path,
  validates cached topology before schedule updates, and fingerprints the selected
  specialization menu.

Repository-wide checklist:

• make test: not rerun; targeted native/CUDA/reference coverage is listed above.
• make checkstyle: targeted non-mutating Ruff checks were used instead of the repository-wide auto-fixing target.
• make test-convergence: not run; this change concerns the native MoE implementation, not the Transformers integration.

[yueyiming2009]

Build/module and forward-path review comments. I found one blocking correctness issue in the same-host X fallback; the other comments are non-blocking questions or test suggestions. I have not yet completed the backward-path review.

[yueyiming2009]

Requesting changes for the blocking same-host forward fallback issue identified inline. When TMA eligibility fails before src_x_ptr[] is fully populated (for example, an otherwise valid SM90 shape with hidden_dim=384, or an early descriptor-encoding failure), do_get<MayUseIb=true> can select the direct peer-copy branch and dereference an unset pointer instead of using the NVSHMEM fallback.

Please populate peer pointers independently of TMA descriptor eligibility, or guard the direct-copy path and retain getmem fallback. Add a same-host multi-GPU regression using a hidden dimension divisible by 8 but not 256, verify numerical parity, and rerun the existing dense/sparse plus TMA on/off targeted tests. The backward synchronization paths reviewed so far look coherent.