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

[wolfcomos]

Summary

This PR adds one CuTe DSL kernel that fuses the swiglu with the RCEIL MXFP8 cast: gate/up are read once, the bf16 activation tensor is never materialized, and rowwise (1x32), colwise (32x1), or both scales come out
of a single pass. This kernel is bit-compatible with the standalone mxfp8 cute dsl quantizers. The consuming
composite autograd op lives in torchtitan branch pytorch/torchtitan#4257 (will refactor for upstream torchtitan PR soon), this torchao PR is kernels-only.

Performance

GB200, torch 2.14 nightly, CUDA 13.4, public DSL 4.7.0. Baseline is an already-fused-activation arm: torch.compile-fused SwiGLU (one static Triton kernel per shape) + the standalone CuTe DSL quantizers

Eager:

input_shape     direction    scales    baseline_us    fused_us  speedup    fused_gbps
--------------  -----------  --------  -------------  --------  ---------  ----------
(128, 128)      forward      rowwise          127.71     68.67  1.86x             1.2
(128, 128)      forward      colwise          123.01     71.49  1.72x             1.2
(128, 128)      forward      both             188.99     82.77  2.28x             1.2
(128, 128)      backward     rowwise          132.58     75.18  1.76x             1.8
(128, 128)      backward     colwise          131.26     75.22  1.75x             1.8
(128, 128)      backward     both             209.95    112.14  1.87x             1.5
(4096, 2048)    forward      rowwise          126.62     71.84  1.76x           587.5
(4096, 2048)    forward      colwise          125.63     71.71  1.75x           588.5
(4096, 2048)    forward      both             192.37     79.04  2.43x           643.4
(4096, 2048)    backward     rowwise          133.58     77.41  1.73x           873.7
(4096, 2048)    backward     colwise          132.51     75.44  1.76x           896.5
(4096, 2048)    backward     both             199.94     85.38  2.34x           994.8
(4096, 7168)    forward      rowwise          122.42     83.33  1.47x          1772.7
(4096, 7168)    forward      colwise          134.58     87.06  1.55x          1696.8
(4096, 7168)    forward      both             207.52     85.38  2.43x          2084.8
(4096, 7168)    backward     rowwise          123.26     80.19  1.54x          2951.9
(4096, 7168)    backward     colwise           133.5     80.26  1.66x          2949.5
(4096, 7168)    backward     both             197.63     87.42  2.26x          3400.3
(16384, 7168)   forward      rowwise           164.8    110.24  1.49x          5359.9
(16384, 7168)   forward      colwise          162.85    116.74  1.40x          5061.6
(16384, 7168)   forward      both             219.14    135.17  1.62x          5267.4
(16384, 7168)   backward     rowwise           405.5    155.65  2.61x          6083.4
(16384, 7168)   backward     colwise          406.78    165.89  2.45x          5707.9
(16384, 7168)   backward     both             513.47    219.17  2.34x          5425.5
(131072, 8192)  forward      rowwise         1371.14    902.11  1.52x          5988.5
(131072, 8192)  forward      colwise         1385.47    978.94  1.42x          5518.5
(131072, 8192)  forward      both            1846.53   1167.33  1.58x          5576.5
(131072, 8192)  backward     rowwise         3670.02    1356.8  2.70x          6380.5
(131072, 8192)  backward     colwise         3702.82   1460.22  2.54x          5928.6
(131072, 8192)  backward     both            4616.99   1891.34  2.44x          5748.1

This host had application clocks pinned to 2062 Mhz.

End-to-end in torchtitan (DeepSeek-V3 16B, 4x GB200, EP=4, TP=1, batch 4, seq 4096, MXFP8 MoE recipe, 50 steps), three arms — the stock MXFP8 path, the fused-w13 MoE composite running the unfused activation + standalone quantizers, and the same composite backed by this PR's kernels:

arm                             tflops(11-50)  vs stock  median step ms  loss@50  peak GiB
------------------------------  -------------  --------  --------------  -------  --------
stock MXFP8                            381.05        --          777.84  2.78596     96.15
w13 composite, unfused quant           379.31    -0.46%          782.05  2.78251     99.47
w13 composite, fused (this PR)         390.64    +2.52%          763.35  2.78362     99.59

Net +2.5% throughput (+3.0% vs the structurally identical unfused composite the kernel-only effect), median step time -1.8%, loss parity (max spread 3.5e-3 at step 50), and the tightest step-time tail (max step 869 ms vs 1081 stock).

The same A/B under the BF16-wgrad recipe (mxfp8_rceil_wgrad_with_hp: the kernel emits rowwise-only casts, and the composite recomputes the activation in BF16 at the autograd boundary for the two wgrad GEMMs — a boundary with less redundant traffic for fusion to remove):

arm                             tflops(11-50)  vs unfused  median step ms  loss@50  peak GiB
------------------------------  -------------  ----------  --------------  -------  --------
w13 composite, unfused                 382.00          --          776.61  2.77950     99.47
w13 composite, fused (this PR)         386.48      +1.17%          771.34  2.78439     99.41

The fused kernel still wins here (+1.2% throughput, -0.7% median step, kernel-only), with the margin narrowed as expected for the cheaper boundary.

API

from torchao.prototype.moe_training.kernels.mxfp8 import (
    gated_act_mxfp8_forward,   # (gated_input [M, 2K] bf16, *, rowwise, colwise)
    gated_act_mxfp8_backward,  # (grad_h [M, K] bf16, gated_input [M, 2K] bf16, *, rowwise, colwise)
)
# both return (output_rowwise, output_colwise, scales_rowwise, scales_colwise)

gated_input packs [gate | up] — the layout a fused w13 projection produces; backward returns concatenated [dGate | dUp], so there is no torch.cat on the hot path. Disabled directions come back zero-sized (fixed arity). Registered as torch custom ops with fake impls; torch.compile(fullgraph=True) works. One kernel source: direction,
layout mode, and the activation (a module-level Constexpr policy function, _silu_pair) are compile-time parameters; shapes are runtime, so one compilation per mode serves every shape.

Numerics and tests

• E8M0 scales and forward E4M3 data are bitwise identical to the eager activation followed by the standalone quantizers — including NaN/Inf, zero, subnormal-tiny, and E4M3-boundary amax regimes under the Refine MXFP8 CUDA quantization semantics #4725 contract.
• Backward E4M3 data may differ by at most one code in ≤1e-5 of elements (measured 5.7e-7): the kernel's PTX sigmoid and FMA contraction have no bit-exact eager equivalent without unfusing.
• 80 cases, ~14 s on GB200: {fwd,bwd} × {rowwise,colwise,both} × edge-value patterns and a shape sweep; compile-fullgraph and cross-mode consistency checked bitwise; every launcher error path exercised (INT32 bound probed allocation-free via FakeTensorMode; pointer-alignment and geometry-override guards, each with a bitwise positive case); opt-in largest-legal-shape check against the standalone quantizer. Upstream's cutedsl test_mxfp8_dsl_special_value_semantics cases pass in the same environment.

Reproducing

pytest test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py -q

# MXFP8_BENCH_VALIDATE=1 checks every configuration against eager before timing;
# --shape/--direction/--scales run a single cell
MXFP8_BENCH_VALIDATE=1 python \
  benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_gated_act_mxfp8.py [--compile]

Follow-ups

1. SiLU is the only activation exposed today, plan to add more activations in future PR.

2. the TorchAO op itself has no register_sharding or automatic fallback. The external TorchTitan composite must choose the unfused path before invoking it, so there might be kernel sharding register work need to do for torchtitan integration.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4743

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

do we have more context on what exactly is not matching?

[wolfcomos]

Thanks for the reviews! The sigmoid fast math caused backward not bitwise exact: https://github.com/wolfcomos/ao/blob/61eb0939b8fc4bf19ab851a023523e0c7d0c70cb/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_gated_act_mxfp8.py#L308-L338.

Bitwise exact: forward qdata, forward scales, backward dUp, and backward scales (0 scale-byte diffs across 8×10⁸ elements). The entire MXFP8 quantization contributes zero delta. The mismatch (472 of 805M backward codes, 5.9e-7, all ±1 code, all in the dGate half) is from pre-quantization fp32 rounding. Sigmoid kernel uses rcp.rn(1 + ex2.approx(−x·log2e)), which differs from fp32 torch.sigmoid at 795 of 65,280 finite bf16 codes. And FMA contraction in d_silu computes dact = fma(silu, 1−s, s) in one rounding, eager does the multiply and add as two roundings.

[vkuzo]

I think this looks reasonable. Can we get the torchtitan PR in close-to-final form and then review both together?

[wolfcomos]

Thanks a lot for the reviews! Torchtitan side integration PR is now live at pytorch/torchtitan#4257. cc @vkuzo @syed-ahmed .