Add the fully-fused MXFP8 grouped-MLP override

[wolfcomos]

Stacked on swiglu-mxfp8-upstream (pytorch#4257). Adds the fully-fused sibling of the SwiGLU-fused override: the whole routed-experts MLP — FC1 grouped GEMM + SwiGLU + dual MXFP8 quantization, FC2, fused dSwiGLU backward, and wgrad — runs through four torchao grouped-GEMM custom ops (mxfp8_grouped_gemm_swiglu_fwd, mxfp8_grouped_gemm, mxfp8_grouped_gemm_dswiglu_bwd, mxfp8_grouped_gemm_wgrad).

Purely additive: one commit, 3 files, +1409/−0. No line of the base branch is modified; the only touched existing file is the DSv3 config registry, which gains one appended debugmodel flavor.

What's in it

• torchtitan/overrides/mxfp8_grouped_mlp.py — self-contained opt-in override. MXFP8FusedGroupedExperts holds w13 natively in [E, 2F, D] 32-block gate/up order (the kernels' layout), with DTensor-RNG-preserving param init and save/load hooks that keep checkpoints in stock w1/w3 format. The factory mxfp8_grouped_experts targets stock RoutedExperts.Config, swaps the token dispatcher to pad_multiple=256 (hard kernel alignment contract: per-expert row counts must be multiples of 256), and fails loud on any non-stock or already-converted config — it and the SwiGLU-fused override reject each other's configs bidirectionally.
• tests/unit_tests/test_mxfp8_grouped_mlp.py — 17 tests: composite numerics vs a quantized-unfused reference with negative controls, checkpoint round-trips both directions, dispatcher/factory wiring, rejection matrix, zero-token backward, compile. Skips cleanly (not ImportError) on hosts without SM100 or without the torchao ops.
• deepseek_v3_debugmodel_mxfp8_grouped_mlp config flavor.

torchao dependency

The four ops live on the torchao side (wolfcomos/ao#5, torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py). Availability is probed at import; when the ops are absent the factory raises with the missing module and requirements named, and the tests skip.

Validation

• Ported suite 17/17 on GB200; the base branch's test_mxfp8_swiglu_override.py unaffected (7/7 with GPU and with CUDA_VISIBLE_DEVICES="").
• Fresh-init NGPU=2 EP=2 debugmodel smoke: override active on all 5 MoE layers, loss 8.12 → 4.54 over 5 steps.
• Evidence lineage (same module bytes, measured on the fork's benchmark branch): full parity gate ladder (paired-seed loss deltas in-band, trace-verified kernel engagement, 3-seed × 500-step convergence) and 16B EP4 A/B: +2.3% tps over the pad-matched MXFP8 baseline at bs4, +5.0% at bs8.

Replaces the closed draft #4 (cudnn-grouped-mlp), which carried the benchmark arms and evidence record.

🤖 Generated with Claude Code

[wolfcomos]

Folded into the base branch: swiglu-mxfp8-upstream@89fcb0feb unifies the grouped override behind a single MXFP8FusedGroupedExperts with fusion_plan={swiglu, grouped_gemm_swiglu} (own params on stock GroupedExperts, plan-native w13 layout, per-plan dispatcher padding, stock w1/w3 checkpoints both ways). The swiglu plan is bitwise-verified against the pre-unification head; this PR's module and tests are carried over inside that commit.