Wire the fully-fused grouped-GEMM MXFP8 plan behind fusion_plan

[wolfcomos]

Stacked on swiglu-mxfp8-upstream (pytorch#4257) — this diff is exactly the follow-up wiring change: +1087/−29 across the override module, the single test file, and one registry flavor. Kernel dependency: pytorch/ao#4820 (the four fused grouped-GEMM ops); until it merges, selecting the plan raises with the dependency named. On approval this becomes the separate upstream pytorch/torchtitan PR, opened after pytorch#4257 merges and ao#4820 lands.

What it adds

MXFP8FusedGroupedMLP.Config gains fusion_plan: Literal["swiglu", "grouped_gemm_swiglu"] = "swiglu":

• "swiglu" (default): unchanged — the existing composite, pad_multiple=128.
• "grouped_gemm_swiglu": the whole routed-experts MLP runs through [mxfp8 training] Add fused grouped GEMM + SwiGLU + MXFP8 quantization ops for MoE routed experts pytorch/ao#4820's four ops (FC1 grouped GEMM + SwiGLU + dual MXFP8 quantization in one launch, FC2 GEMM, fused dSwiGLU backward, wgrad). Parameters stay stock — the forward packs w1_EFD/w3_EFD into the kernels' 32-block [E, 2F, D] order (byte-verified against the proven mapping, with an unpack round-trip test). The factory installs the pad_multiple=256 dispatcher the kernels' hard per-expert alignment requires (update-in-place for an existing TorchAO dispatcher, swap for AllToAll, refuse otherwise), and raises actionably on unknown plans, fuse_activation=False under this plan, or missing ops.

Because both plans share the stock-parameter structure from pytorch#4257, the plan differs only in its forward pack + composite: checkpoints, fresh init, and sharding are identical across plans and identical to stock modules. Only 3 executable pre-existing lines changed in the module (import, dispatcher call re-indent, derive kwarg); the swiglu Functions/helpers/forward are byte-untouched.

Evidence

• Default-plan inertness, checkpoint-paired bitwise: the unchanged deepseek_v3_debugmodel_mxfp8_fused_mlp flavor run from a shared seed checkpoint at the base commit vs this commit (NGPU=2 EP=2, 5 deterministic steps) — every printed loss and grad-norm equal (8.12229/3.6271 → 4.52541/2.4915, both arms).
• Fully-fused smoke (deepseek_v3_debugmodel_mxfp8_grouped_gemm_mlp, fresh init, EP=2): plan active on all 5 MoE layers, loss 8.036 → 4.682.
• Tests: 21 in the one file — 14 wiring (CPU-executable; mocked capability/availability) + 7 GPU numerics gated on SM100 + ops: composite vs a quantized-unfused reference with calibrated SQNR bands and a NaN-poisoned-tail negative control, zero-token-expert exact-zero grads, zero-routed-tokens backward, torch.compile vs eager bitwise on outputs and all three weight grads, rejection matrix, checkpoint-keys == stock, and the 32-block pack byte-verification. Reference helpers and thresholds are identical to the previously gate-proven suite.
• Upstream-CI simulation (no GPU, stock PyPI torchao 0.18.0, no dev trees): 21 collected cleanly, all 14 wiring tests execute green, 7 numerics skip with a clear reason.
• flake8/py_compile clean; overrides/fused_swiglu.py remains byte-untouched.

Performance

Measured end-to-end with these kernels at DSv3-16B EP=4 (details and pad-tax table in pytorch/ao#4820): +2.3% tokens/s over the pad-matched unfused MXFP8 baseline at bs4, +5.0% at bs8, with paired-seed loss-parity gates and a 3-seed × 500-step convergence protocol.

🤖 Generated with Claude Code