Add swiglu + MXFP8 fused CuTe DSL 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 https://github.com/wolfcomos/torchtitan/pull/2/changes (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          105.38     62.53  1.69x             1.3
(128, 128)      forward      colwise          115.47     65.70  1.76x             1.3
(128, 128)      forward      both             179.42     73.86  2.43x             1.3
(128, 128)      backward     rowwise          110.19     65.54  1.68x             2.0
(128, 128)      backward     colwise          116.06     66.93  1.73x             2.0
(128, 128)      backward     both             199.01     79.81  2.49x             2.1
(4096, 2048)    forward      rowwise          112.58     64.61  1.74x           653.3
(4096, 2048)    forward      colwise          113.76     65.79  1.73x           641.5
(4096, 2048)    forward      both             172.75     73.07  2.36x           696.0
(4096, 2048)    backward     rowwise          103.12     65.86  1.57x          1027.0
(4096, 2048)    backward     colwise          115.60     68.14  1.70x           992.5
(4096, 2048)    backward     both             176.45     75.60  2.33x          1123.5
(4096, 7168)    forward      rowwise          111.23     71.71  1.55x          2059.9
(4096, 7168)    forward      colwise          110.19     72.06  1.53x          2049.8
(4096, 7168)    forward      both             177.20     80.26  2.21x          2217.9
(4096, 7168)    backward     rowwise          118.11     72.06  1.64x          3284.8
(4096, 7168)    backward     colwise          120.19     75.10  1.60x          3151.8
(4096, 7168)    backward     both             181.25     80.22  2.26x          3705.5
(16384, 7168)   forward      rowwise          164.86    110.59  1.49x          5342.8
(16384, 7168)   forward      colwise          164.26    120.83  1.36x          4890.0
(16384, 7168)   forward      both             219.20    143.36  1.53x          4966.4
(16384, 7168)   backward     rowwise          406.53    155.65  2.61x          6083.4
(16384, 7168)   backward     colwise          407.55    196.61  2.07x          4816.0
(16384, 7168)   backward     both             514.43    219.14  2.35x          5426.2
(131072, 8192)  forward      rowwise         1372.16    904.19  1.52x          5974.7
(131072, 8192)  forward      colwise         1387.52   1022.98  1.36x          5280.9
(131072, 8192)  forward      both            1849.09   1216.08  1.52x          5352.9
(131072, 8192)  backward     rowwise         3676.03   1321.98  2.78x          6548.5
(131072, 8192)  backward     colwise         3708.94   1710.08  2.17x          5062.4
(131072, 8192)  backward     both            4624.83   1859.09  2.49x          5847.8

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 pytorch/ao#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_gated_act_mxfp8_unified.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.