Add SM103 FP8 block128 MegaMoE training backend

[renfeichen-fw]

What changed

• add an SM103-only FP8 E4M3 activation × FP8 E4M3 block-128 weight MegaMoE training backend
• own routed quantization, expert-parallel all-to-all transport, grouped W13/SwiGLU/W2 compute, fixed post-down combine, and the complete routed backward
• return BF16 input/master-weight gradients and exact FP32 route-score gradients while retaining GLM's canonical gate/up master layout
• embed the full source commit as deep_gemm.__git_commit__ and expose a non-launching, fail-closed capability manifest
• add reference, boundary, empty/skewed-route, two-rank distributed, provenance, and structured performance coverage

Why

FireTitan's GLM-5.2 FP8-block128 training path needs a single routed backend that preserves E4M3 weights and FP32 128 × 128 inverse scales. Existing MegaMoE paths use FP4 or BF16 weights and do not provide this forward/backward contract.

This backend deliberately targets compute capability 10.3 only. The extension contains only an sm_103a cubin and rejects every other runtime capability; there is no fallback.

Validation

• installed-wheel non-launching capability/provenance test: 1 passed
• installed-wheel SM103 primitive and single-rank forward/backward suite: 22 passed
• installed-wheel two-rank uneven/cross-rank/skew/empty-expert suite: 1 passed
• cuobjdump --list-elf: exactly _C.cpython-312-x86_64-linux-gnu.1.sm_103a.cubin
• full-width structured smoke (D=6144, H=2048, top-8, 15,625 tokens, 16 local experts): median complete forward/backward 51.741 ms on L20D/SM103

The FireTitan integration will pin exact commit a0cba68c6e1701023ce93e3eaef5ad08d9b5d23f and validate the installed extension SHA separately.

---

Note

High Risk
Large new training-critical CUDA surface (routing, EP transport, autograd) with strict SM103-only behavior and tight coupling to persistent Mega MoE kernels; mistakes affect numerical correctness and distributed training.

Overview
Adds an SM103-only (cc 10.3, no fallback) FP8 E4M3 / 128×128 block-scale routed MoE training path: Python exposes fp8_block128_mega_moe, capability introspection, and W13 layout helpers; the native extension registers a large sm103_fp8_block128 API surface and embeds deep_gemm.__git_commit__.

The new csrc/sm103_fp8_block128.cu stack owns persistent workspace layout, input prep, EP2/EP16 persistent forward/backward (reusing SM100 Mega MoE device code with FP8 block128 / canonical W13 wiring), grouped CUTLASS blockwise GEMMs, SwiGLU/route/combine helpers, and grouped BF16 wgrad with 128-row padding handling.

Smaller cross-cutting tweaks: README documents SM103 and the backend; grouped BF16 k-GEMM allows optional c; UE8M0 pack accepts column-major scale tensors when using PSUM layout; SM100 fp8_fp4_mega_moe JIT launch gains token source metadata, runtime ring sizes, and bitwise-stable activation-clamp constants; shared enums add route/combine modes.

^{Reviewed by Cursor Bugbot for commit 5018de9. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Benchmark skips transport setup

Medium Severity

The structured performance script runs fp8_block128_mega_moe under torchrun with a multi-rank process group but never calls _configure_fp8_block128_mega_moe_transport. The DeepEP path requires a pre-registered context token capacity, so multi-rank runs fail at buffer lookup instead of benchmarking.

 

^{Reviewed by Cursor Bugbot for commit 5b2aa57. Configure here.}

[cursor]

Route quantize skips power-of-two scales

High Severity

The new sm103_route_scale_quantize_group128_kernel and sm103_expanded_route_scale_quantize_group128_kernel store block scales as raw amax / 448, while sm103_quantize_bf16_e4m3_group128_kernel, sm103_swiglu_quantize_group128_kernel, and persistent input prep round scales up to the next power of two. The capability manifest advertises fp32_power_of_two scales, so backward route quantization can disagree with forward and MMA scale expectations.

Additional Locations (1)

• csrc/sm103_fp8_block128.cu#L848-L874

 

^{Reviewed by Cursor Bugbot for commit c84bbed. Configure here.}

[cursor]

Prepare forward rank count mismatch

Medium Severity

sm103_fp8_block128_prepare_persistent_inputs sizes the symmetric workspace from its num_ranks argument, while sm103_fp8_block128_persistent_forward and backward rebuild PersistentWorkspaceLayout from buffer_ptrs.size(). If those counts differ, preparation writes quantized inputs and route metadata at the wrong offsets while the persistent kernel interprets the same buffer with a different layout, corrupting routing and activations without a hard error when the buffer is large enough.

Additional Locations (1)

• csrc/sm103_fp8_block128.cu#L2010-L2014

 

^{Reviewed by Cursor Bugbot for commit 9c63f41. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 4 total unresolved issues (including 3 from previous reviews).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 06246eb. Configure here.}

[cursor]

Empty expert PSUM traps GPU

High Severity

In sm103_prepare_expanded_wgrad_metadata_kernel, each expert’s valid row range starts at the 128-aligned successor of the previous PSUM end. When an expert receives no tokens, its PSUM end equals the prior end; if that value is not already 128-aligned, the computed start exceeds end, the kernel hits asm trap, and grouped_bf16_wgrad_expanded can hard-fail on common empty-expert routing.

 

^{Reviewed by Cursor Bugbot for commit 06246eb. Configure here.}