[mxfp8 training] fused gated-activation (SwiGLU) + MXFP8 quantization kernel override

[wolfcomos]

Summary

Adds an opt-in override module, torchtitan/overrides/mxfp8_fused_swiglu.py, that runs a
FeedForward / routed-experts MLP entirely in MXFP8 with a fused SwiGLU boundary: one composite
autograd function computes w13 GEMM -> fused SwiGLU + dual MXFP8 quantization -> w2 GEMM
(dense and grouped variants), using torchao's fused gated-activation kernel
(pytorch/ao#4743) for the activation boundary and torchao's MXFP8 grouped/scaled-mm internals
for the GEMMs. Compared to the unfused MXFP8 path, the fused boundary eliminates the standalone
activation-quantization casts (fused mode launches 1 fwd + 1 bwd fused op instead of 6 + 6
standalone cast kernels per MLP) and never materializes the BF16 post-activation tensor.

The integration contained:

• MXFP8FusedSwiGLU(FusedSwiGLU) and MXFP8FusedGroupedExperts(FusedGroupedExperts) subclass
  the existing fused modules (inheriting the w13 parameter, stock-layout checkpoint hooks, and
  init/sharding remaps) and override only forward.
• Two @override factories: mxfp8_fused_swiglu (targets FeedForward.Config) and
  mxfp8_fused_grouped_experts (targets RoutedExperts.Config, which owns both the token
  dispatcher and the inner experts — the factory swaps the dispatcher to the padded variant its
  kernels require, pad_multiple=128).
• Zero changes to existing files' behavior: overrides/fused_swiglu.py,
  components/quantization/mx.py, and the quantization converters are untouched. Enabling the
  override is the only opt-in; no converter or job-config plumbing is added.

Fail-loud contract

There is no silent fallback anywhere in this path. Unsupported configurations raise actionable
errors naming the offending value (unsupported SM at override-application time; non-%128 dims,
non-BF16 dtypes, DTensor activations, and misaligned token counts at call time — data-dependent
token counts use torch._check deferred asserts so routing-dependent shapes work under
torch.compile). The user changes the config; the code never substitutes a different numerics
path than the one configured.

Usage

# dense (llama3 debugmodel example flavor)
--override.imports torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_swiglu

# routed experts (deepseek_v3 debugmodel example flavor; needs EP>=2 so the token
# dispatcher produces the padded expert-major layout the kernels consume)
--override.imports torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_grouped_experts

Example flavors: llama3_debugmodel_mxfp8_fused_swiglu, deepseek_v3_debugmodel_mxfp8_fused_swiglu
(both enable compile; attention/lm_head stay BF16 — composable with MXFP8LinearConverter on
those modules independently).

Numerics and tests

• Shipped: tests/unit_tests/test_mxfp8_swiglu_override.py — CPU-runnable wiring suite
  (config-tree transform, apply_overrides, meta-device build, dispatcher pad_multiple,
  factory fail-loud), validated both with GPUs visible and with CUDA_VISIBLE_DEVICES="".
• GPU numerics are validated in NVIDIA-internal Blackwell CI and will be maintained there
  (upstream CI has no SM100): 25-test suite on GB200 covering bitwise forward parity between the
  fused kernel and standalone-cast quantization, SQNR tracking against the per-GEMM MXFP8 and
  BF16 references, torch.compile with unbacked routing-dependent token counts, dispatcher
  pad-row/tail inertness, and profiler op-count contracts.
• Kernel-level numerics (bitwise forward, one-code-bounded backward at ~6e-7 of elements, root-
  caused to the kernel's correctly-rounded sigmoid + FMA contraction) are documented in
  [mxfp8 training] Add a fused gated-activation (SwiGLU) + MXFP8 quantization kernel pytorch/ao#4743.

[vkuzo]

i'd advocate against falling back to bf16 based on shape, IMO better to honor the user config and let the user change the config if it cannot be executed

[wolfcomos]

Thanks a lot for the review @vkuzo ! I addressed this and it's now updated at pytorch#4257.