From 50804c614432edd56f5a88dee49e2ae871751c58 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 23:11:16 -0700 Subject: [PATCH] Add the fully-fused MXFP8 grouped-MLP override Stacked on the SwiGLU-override branch: a self-contained override module that runs the whole routed-expert MLP (grouped GEMM + SwiGLU + quant) through fused torchao ops, its unit-test suite, and one deepseek_v3 debugmodel flavor activating it. The ops are gated on torchao availability; without them the tests skip and the factory fails loud. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_mxfp8_grouped_mlp.py | 780 ++++++++++++++++++ .../models/deepseek_v3/config_registry.py | 33 + torchtitan/overrides/mxfp8_grouped_mlp.py | 596 +++++++++++++ 3 files changed, 1409 insertions(+) create mode 100644 tests/unit_tests/test_mxfp8_grouped_mlp.py create mode 100644 torchtitan/overrides/mxfp8_grouped_mlp.py diff --git a/tests/unit_tests/test_mxfp8_grouped_mlp.py b/tests/unit_tests/test_mxfp8_grouped_mlp.py new file mode 100644 index 0000000000..1209040df8 --- /dev/null +++ b/tests/unit_tests/test_mxfp8_grouped_mlp.py @@ -0,0 +1,780 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the cuDNN-frontend MXFP8 fused grouped-MLP override +(mxfp8_grouped_mlp.py). + +The five test groups: + +1. Composite numerics: forward+backward vs an independent quantized-unfused + reference built here from standalone RCEIL casts (``to_mx``), raw + ``torch._scaled_grouped_mm``, and first-principles eager SwiGLU + forward/backward -- deliberately NOT the override module's own cast + helpers. Shapes include a D != F case, zero-token experts (asserting + exactly-zero param grads), an all-experts-empty R == 0 case (the ops' + documented early-outs), and a strict inactive tail (A < R) filled with + deliberate garbage including NaN (a NaN-poisoned-inactive-tail check -- + critical because every e2e gate is force-balanced and never exercises + ragged routing). +2. Composition: the named fused config through the real ``apply_overrides`` + pipeline. The override is self-contained: stock experts in, fused experts + plus the factory-installed pad_multiple=256 TorchAO dispatcher out, no + converter involved -- and fail-loud: converter-quantized experts raise + instead of falling back. Activation evidence is module/config TYPE only. +3. Autograd/AC: under the real SelectiveAC policy the composite forward runs + exactly twice (save-from-recompute), gradients match no-AC bitwise, and + the by-reference parameter save survives optimizer.step into the next + step. +4. State dict / param layout: the 32-block ``w13 [E, 2F, D]`` hooks + round-trip the stock ``w1_EFD``/``w3_EFD`` checkpoint layout, and the + param-init remap initializes the gate/up 32-row blocks with their own + initializers. +5. Trace shape: one fwd+bwd calls exactly 1x fwd op, 2x mm op (FC2 + FC1 + dgrad), 1x bwd op, 2x wgrad op, counted via ``torchao::`` op names ONLY. + +Every per-expert row count in this file is a 256-multiple: the cuDNN FE +kernels hard-code FIX_PAD_SIZE=256 and 128-multiple-only splits corrupt +silently and NONDETERMINISTICALLY (the corruption locus migrates between +identical-input reruns, so no passing run proves such a split safe) -- there +is deliberately no "sub-256 splits still work" fixture. +""" + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F + +if not (torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)): + pytest.skip("Requires CUDA SM 10.0 (Blackwell)", allow_module_level=True) + +# The override module's own availability flag is the single source of truth +# for whether the torchao fused grouped-MLP ops exist AND their cudnn-frontend +# kernels are usable; skipping on it (with its reason) instead of a bare +# try/except keeps a broken environment loud rather than silently skipped. +try: + from torchtitan.overrides.mxfp8_grouped_mlp import ( + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE, + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON, + ) +except ImportError as e: # torchao (or a transitive dep) not installed + pytest.skip( + f"torchao is required for the MXFP8 grouped-MLP override: {e}", + allow_module_level=True, + ) + +if not _TORCHAO_GROUPED_MLP_OPS_AVAILABLE: + pytest.skip( + "torchao fused grouped-MLP ops unavailable: " + f"{_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON}", + allow_module_level=True, + ) + +from torch.profiler import ProfilerActivity, profile +from torchao.prototype.mx_formats.config import ScaleCalculationMode +from torchao.prototype.mx_formats.mx_tensor import to_mx +from torchao.prototype.mx_formats.utils import to_blocked +from torchao.quantization.utils import compute_error + +from torchtitan.config import apply_overrides, derive +from torchtitan.distributed.activation_checkpoint import SelectiveAC +from torchtitan.models.common.moe import GroupedExperts, RoutedExperts +from torchtitan.models.common.token_dispatcher import TorchAOTokenDispatcher +from torchtitan.models.deepseek_v3.config_registry import ( + deepseek_v3_debugmodel, + deepseek_v3_debugmodel_hybridep, + deepseek_v3_debugmodel_mxfp8, + deepseek_v3_debugmodel_mxfp8_grouped_mlp, +) +from torchtitan.overrides.mxfp8_grouped_mlp import ( + MXFP8FusedGroupedExperts, + _make_w13_init, + mxfp8_fused_grouped_mlp, + mxfp8_grouped_experts, +) + +# The activation string is part of the frozen override interface. +_OVERRIDE_TARGET = "torchtitan.overrides.mxfp8_grouped_mlp.mxfp8_grouped_experts" + +_OP_FWD = "torchao::mxfp8_grouped_gemm_swiglu_fwd" +_OP_MM = "torchao::mxfp8_grouped_gemm" +_OP_BWD = "torchao::mxfp8_grouped_gemm_dswiglu_bwd" +_OP_WGRAD = "torchao::mxfp8_grouped_gemm_wgrad" +_ALL_OPS = (_OP_FWD, _OP_MM, _OP_BWD, _OP_WGRAD) + +_BLOCK = 32 +_E4M3 = torch.float8_e4m3fn +_RCEIL = ScaleCalculationMode.RCEIL + +# --------------------------------------------------------------------------- +# Test-1 tolerances: derived from the measured variability of the unfused +# lane itself, never copied from kernel-level (ao op suite) gates. +# Calibration method (rerunnable in-place after any torch/torchao numerics +# change): run THIS file's reference against (a) an alternative reduction +# order of the SAME unfused math (per-expert fp32 dequant-loop GEMMs +# consuming the reference's own BF16 z/h/dz boundary values) and (b) the +# fp32 eager MLP, at both parametrized shapes. Measured 2026-08-18 on GB200 +# (torch 2.14.0a0 nightly, app clocks 2062 MHz): +# +# output ref-vs-alt-order (dB) ref-vs-fp32 (dB) +# y inf / 128.16 23.70 / 23.66 +# dx 112.30 / 118.86 23.68 / 23.68 +# dw13 inf / inf 23.69 / 23.68 +# dw2 inf / inf 23.65 / 23.66 +# +# Identical RCEIL quantization boundaries make the two unfused lanes +# near-bitwise (their mutual floor is 112.3 dB) — but the COMPOSITE cannot +# reach that floor against any bf16-z reference: the cuDNN GLU/dGLU kernels +# evaluate SwiGLU/dSwiGLU from their in-kernel FP32 accumulators (h is +# quantized from f32 silu(z_f32)*up_f32; dz from f32 dh), while this +# reference — and the real unfused MXFP8 baseline it models — round z and dh +# to BF16 first. That one-boundary difference dominates every +# composite output; measured composite-vs-reference band (same host/session +# as the table above): +# +# output debugmodel_tail_zero_expert asym_d_ne_f_tail +# y 35.39 dB 35.32 dB +# dx 35.64 dB 35.70 dB +# dw13 35.68 dB 35.64 dB +# dw2 35.32 dB 35.30 dB +# +# The gate sits 5.3 dB below the measured band floor (35.30) and 6.3 dB +# ABOVE the 23.65-23.70 dB "independent-but-correct lane" level (what fp32 +# itself scores against this reference), so it still discriminates "shares +# every quantization boundary except the kernel-internal h/dh rounds" from +# "merely correct" — and a real dataflow/layout/offsets bug lands near 0 dB. +# Kernel-level exactness (60-160 dB) is enforced separately by the ao op +# suite against kernel-native references. +_SQNR_VS_REFERENCE_DB = 30.0 +# Secondary tracking gate vs the FP32 eager MLP: measured 23.65-23.70 dB for +# every output at both shapes; 2.6 dB of seed headroom. Catches a blind spot +# shared by both MXFP8 lanes. +_SQNR_VS_FP32_DB = 21.0 +# --------------------------------------------------------------------------- + +# Fixture shapes: per-expert row counts are 256-multiples (the dispatcher's +# pad_multiple=256 ABI guarantee), zero-token experts are legal anywhere, and +# `tail` allocates inactive rows past offsets[-1] (A < R). D != F in the +# second case so a wrong-axis weight cast cannot cancel. +_CASES = { + "debugmodel_tail_zero_expert": dict( + d=256, f=256, sizes=[256, 0, 256, 512, 0, 256], tail=256, seed=0 + ), + "asym_d_ne_f_tail": dict(d=256, f=512, sizes=[256, 0, 512], tail=256, seed=1), +} + + +def _blk_view(w13: torch.Tensor): + """[E, 2F, D] 32-block order -> view [E, F//32, 2, 32, D] with the + gate/up axis at dim 2.""" + e, two_f, d = w13.shape + return w13.view(e, two_f // 64, 2, 32, d) + + +def _to_blk(w1: torch.Tensor, w3: torch.Tensor) -> torch.Tensor: + """Stock [E, F, D] pairs -> 32-block [E, 2F, D].""" + e, f, d = w1.shape + return ( + torch.stack([w1, w3], dim=2) + .view(e, f // 32, 32, 2, d) + .permute(0, 1, 3, 2, 4) + .reshape(e, 2 * f, d) + ) + + +def _zsplit(z: torch.Tensor, f: int): + """[R, 2F] in 32-block order -> (gate [R, F], up [R, F]).""" + r = z.shape[0] + v = z.view(r, f // 32, 2, 32) + return v[:, :, 0, :].reshape(r, f), v[:, :, 1, :].reshape(r, f) + + +def _zmerge(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """(gate [R, F], up [R, F]) -> [R, 2F] in 32-block order.""" + r, f = gate.shape + out = torch.empty(r, 2 * f, dtype=gate.dtype, device=gate.device) + v = out.view(r, f // 32, 2, 32) + v[:, :, 0, :] = gate.view(r, f // 32, 32) + v[:, :, 1, :] = up.view(r, f // 32, 32) + return out + + +def _make_case(*, d, f, sizes, tail, seed=0): + """Dispatcher-shaped fixture: expert-major x [R, D] with per-expert row + counts ``sizes`` (256-multiples), offsets = inclusive cumsum, and a + strict inactive tail of ``tail`` rows. Tail rows carry large deliberate + garbage plus NaN (a NaN-poisoning attack on the undefined inactive tail): + producers do not define them, kernels must never let them contaminate + active rows, and y/dx comparisons mask them because the mm op leaves + output tail rows unwritten.""" + g = len(sizes) + a = sum(sizes) + r = a + tail + torch.manual_seed(seed) + offsets = torch.tensor( + [sum(sizes[: i + 1]) for i in range(g)], device="cuda", dtype=torch.int32 + ) + x = torch.randn(r, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + dy = torch.randn(r, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w1 = torch.randn(g, f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w3 = torch.randn(g, f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w13 = _to_blk(w1, w3) + w2 = torch.randn(g, d, f, device="cuda", dtype=torch.bfloat16) / d**0.5 + if tail: + x[a:] = 12345.0 + dy[a:] = -6789.0 + x[a : a + tail // 2] = float("nan") + dy[a : a + tail // 2] = float("nan") + return dict(x=x, dy=dy, w13=w13, w2=w2, offsets=offsets, sizes=sizes, a=a, r=r) + + +# --------------------------------------------------------------------------- +# Independent quantized-unfused reference. Standalone torchao RCEIL casts +# (``to_mx``) + raw ``torch._scaled_grouped_mm`` + eager SwiGLU +# forward/backward (the BF16 round of z precedes SwiGLU; the BF16 rounds of +# h/dz precede their quantizers); wgrads are colwise quant-dequant + fp32 +# matmul per expert. Built from first principles, sharing no code with the +# module under test. Operates directly on the 32-block ``w13 [G, 2F, D]``: +# rowwise quantization is per-row (row order is irrelevant) and colwise +# 32-blocks along 2F are pure-gate or pure-up in this order, so quantization +# boundaries match the composite exactly. +# --------------------------------------------------------------------------- + + +def _rceil_rowwise(t): + scale, q = to_mx(t, _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q, to_blocked(scale) + + +def _rceil_rowwise_3d(w): + qs, sfs = zip(*(_rceil_rowwise(w[g]) for g in range(w.shape[0]))) + return torch.stack(list(qs)), torch.stack(list(sfs)) + + +def _rceil_colwise_3d(w): + """[G, N, K] -> qdata [G, N, K] stride (N*K, 1, N) quantized along N + + per-group blocked scales (the ``mat2`` of a dgrad ``_scaled_grouped_mm``).""" + qs, sfs = zip(*(_rceil_rowwise(w[g].t().contiguous()) for g in range(w.shape[0]))) + return torch.stack(list(qs)).transpose(-2, -1), torch.stack(list(sfs)) + + +def _dequant(q, scale): + m, k = q.shape + return ( + q.float().view(m, k // _BLOCK, _BLOCK) + * scale.to(torch.float32).view(m, k // _BLOCK, 1) + ).view(m, k) + + +def _quant_dequant_colwise(t): + """[m, N] bf16 -> fp32 [N, m]: RCEIL-quantize along the row axis (32x1) + and dequantize. Per-expert slices quantize identically to the whole + matrix because 256-multiple group sizes keep every 32-value block inside + one group.""" + scale, q = to_mx(t.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + return _dequant(q, scale) + + +def _wgrad_expert(a, b): + """Normative wgrad for one expert: dequant(a_col).T @ dequant(b_col), + fp32 accumulation, one BF16 round. a [m, N], b [m, K] -> [N, K].""" + return (_quant_dequant_colwise(a) @ _quant_dequant_colwise(b).t()).to( + torch.bfloat16 + ) + + +def _reference_forward_backward(x, w13, w2, dy, offsets, sizes, a): + """Returns (y, dx, dw13 [G, 2F, D] 32-block order, dw2 [G, D, F]). y/dx + tail rows [a:] are defined as zero here (the real ops leave them + unwritten; callers mask them out of every comparison).""" + r, d = x.shape + g, two_f = w13.shape[0], w13.shape[1] + f = two_f // 2 + + # FC1 forward, then eager SwiGLU on the BF16-rounded z. The gate/up split + # follows the 32-block column order z inherits from the w13 row order. + x_q, x_sf = _rceil_rowwise(x) + w13_row_q, w13_row_sf = _rceil_rowwise_3d(w13) + z = torch._scaled_grouped_mm( + x_q, + w13_row_q.transpose(-2, -1), + x_sf.reshape(r, -1), + w13_row_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + z[a:] = 0 + gate_bf16, up_bf16 = _zsplit(z, f) + gate = gate_bf16.float() + up = up_bf16.float() + h = (F.silu(gate) * up).to(torch.bfloat16) + + # FC2 forward. + h_q, h_sf = _rceil_rowwise(h) + w2_row_q, w2_row_sf = _rceil_rowwise_3d(w2) + y = torch._scaled_grouped_mm( + h_q, + w2_row_q.transpose(-2, -1), + h_sf.reshape(r, -1), + w2_row_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + y[a:] = 0 + + # FC2 dgrad, then eager dSwiGLU on the BF16-rounded dh. + dy_q, dy_sf = _rceil_rowwise(dy) + w2_col_q, w2_col_sf = _rceil_colwise_3d(w2) + dh = torch._scaled_grouped_mm( + dy_q, + w2_col_q, + dy_sf.reshape(r, -1), + w2_col_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + dh[a:] = 0 + sig = torch.sigmoid(gate) + silu_g = gate * sig + dsilu = sig * (1.0 + gate * (1.0 - sig)) + dhf = dh.float() + dgate = (dhf * up * dsilu).to(torch.bfloat16) + dup = (dhf * silu_g).to(torch.bfloat16) + dz = _zmerge(dgate, dup) + + # FC1 dgrad. + dz_q, dz_sf = _rceil_rowwise(dz) + w13_col_q, w13_col_sf = _rceil_colwise_3d(w13) + dx = torch._scaled_grouped_mm( + dz_q, + w13_col_q, + dz_sf.reshape(r, -1), + w13_col_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + dx[a:] = 0 + + # Wgrads over active rows only; zero-token experts stay all-zero. + dw13 = torch.zeros(g, two_f, d, device=x.device, dtype=torch.bfloat16) + dw2 = torch.zeros(g, d, f, device=x.device, dtype=torch.bfloat16) + prev = 0 + for gi in range(g): + end = int(offsets[gi]) + if end > prev: + dw13[gi] = _wgrad_expert(dz[prev:end], x[prev:end]) + dw2[gi] = _wgrad_expert(dy[prev:end], h[prev:end]) + prev = end + return y, dx, dw13, dw2 + + +def _fp32_reference(x, w13, w2, dy, offsets, a): + """FP32 eager autograd MLP over the active rows. Returns + (y, dx, dw13 [G, 2F, D] 32-block order, dw2).""" + g, two_f = w13.shape[0], w13.shape[1] + f = two_f // 2 + x32 = x[:a].float().detach().requires_grad_(True) + w13_32 = w13.float().detach().requires_grad_(True) + w2_32 = w2.float().detach().requires_grad_(True) + v = _blk_view(w13_32) + outs, prev = [], 0 + for gi in range(g): + end = int(offsets[gi]) + w1_g = v[gi, :, 0].reshape(f, x.shape[1]) + w3_g = v[gi, :, 1].reshape(f, x.shape[1]) + gate = x32[prev:end] @ w1_g.t() + up = x32[prev:end] @ w3_g.t() + h = F.silu(gate) * up + outs.append(h @ w2_32[gi].t()) + prev = end + y_ref = torch.cat(outs, dim=0) + y_ref.backward(dy[:a].float()) + return y_ref, x32.grad, w13_32.grad, w2_32.grad + + +# --------------------------------------------------------------------------- +# 1. Composite numerics +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("case", sorted(_CASES)) +def test_composite_matches_quantized_unfused_reference(case): + fx = _make_case(**_CASES[case]) + a, r, d = fx["a"], fx["r"], fx["x"].shape[1] + + x = fx["x"].clone().detach().requires_grad_(True) + w13 = fx["w13"].clone().detach().requires_grad_(True) + w2 = fx["w2"].clone().detach().requires_grad_(True) + y = mxfp8_fused_grouped_mlp(x, w13, w2, fx["offsets"]) + assert y.shape == (r, d) + assert y.dtype == torch.bfloat16 + y.backward(fx["dy"]) + + ref_y, ref_dx, ref_dw13, ref_dw2 = _reference_forward_backward( + fx["x"], fx["w13"], fx["w2"], fx["dy"], fx["offsets"], fx["sizes"], a + ) + fp32 = _fp32_reference(fx["x"], fx["w13"], fx["w2"], fx["dy"], fx["offsets"], a) + + # Zero-token experts must produce exactly-zero weight gradients through + # the autograd path (the wgrad op writes empty-group outputs as zero). + for gi, m in enumerate(fx["sizes"]): + if m == 0: + assert w13.grad[gi].abs().max().item() == 0.0 + assert w2.grad[gi].abs().max().item() == 0.0 + + # y/dx are compared over active rows only: both lanes leave the inactive + # tail [A, R) unwritten, and the garbage+NaN planted in the x/dy tails + # must not move (or NaN-poison) any active output. + for name, got, ref, hp in [ + ("y", y[:a], ref_y[:a], fp32[0]), + ("dx", x.grad[:a], ref_dx[:a], fp32[1]), + ("dw13", w13.grad, ref_dw13, fp32[2]), + ("dw2", w2.grad, ref_dw2, fp32[3]), + ]: + assert torch.isfinite(got).all(), f"{name} contains non-finite values" + sqnr = compute_error(ref.float(), got.float()) + assert sqnr >= _SQNR_VS_REFERENCE_DB, ( + f"{name} SQNR vs quantized-unfused reference {sqnr} < " + f"{_SQNR_VS_REFERENCE_DB}" + ) + sqnr_hp = compute_error(hp.float(), got.float()) + assert sqnr_hp >= _SQNR_VS_FP32_DB, ( + f"{name} SQNR vs fp32 {sqnr_hp} < {_SQNR_VS_FP32_DB}" + ) + + +def test_composite_zero_routed_tokens(): + # A local expert set receiving zero routed tokens (R == 0, every offset + # 0): the titan cast chain must produce empty quantized operands and the + # ops' documented R == 0 early-outs must return empty y/dx and + # exactly-zero weight grads without error. + d, f, g = 256, 256, 3 + torch.manual_seed(8) + x = torch.zeros(0, d, device="cuda", dtype=torch.bfloat16, requires_grad=True) + w13 = ( + torch.randn(g, 2 * f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + ).requires_grad_(True) + w2 = ( + torch.randn(g, d, f, device="cuda", dtype=torch.bfloat16) / d**0.5 + ).requires_grad_(True) + offsets = torch.zeros(g, device="cuda", dtype=torch.int32) + + y = mxfp8_fused_grouped_mlp(x, w13, w2, offsets) + assert y.shape == (0, d) + assert y.dtype == torch.bfloat16 + y.backward(torch.zeros_like(y)) + + assert x.grad is not None and x.grad.shape == (0, d) + assert w13.grad is not None and w13.grad.abs().max().item() == 0.0 + assert w2.grad is not None and w2.grad.abs().max().item() == 0.0 + + +# --------------------------------------------------------------------------- +# 2. Composition (config-time factory gating through the real pipeline) +# --------------------------------------------------------------------------- + + +def _prepare(config): + """Mimic the trainer's pre-override step (sharding fill).""" + config.model_spec.model.update_from_config(config=config) + return config + + +def _routed_experts_nodes(config): + return list(config.traverse(RoutedExperts.Config)) + + +def test_composition_fired_on_named_config(): + config = deepseek_v3_debugmodel_mxfp8_grouped_mlp() + # The activation string is part of the frozen interface. + assert _OVERRIDE_TARGET in config.override.imports + _prepare(config) + # Self-contained: before the apply, the experts are STOCK (no grouped + # converter ran on this flavor) and the dispatcher is not yet padded. + for _fqn, cfg, _parent, _attr in _routed_experts_nodes(config): + assert type(cfg.inner_experts) is GroupedExperts.Config + assert not isinstance(cfg.token_dispatcher, TorchAOTokenDispatcher.Config) + replacements = apply_overrides(config.override, config) + assert replacements + + nodes = _routed_experts_nodes(config) + assert nodes + for _fqn, cfg, _parent, _attr in nodes: + assert type(cfg.inner_experts) is MXFP8FusedGroupedExperts.Config + # The factory installed the padded dispatcher itself. + assert isinstance(cfg.token_dispatcher, TorchAOTokenDispatcher.Config) + assert cfg.token_dispatcher.pad_multiple == 256 + + # Module TYPE is the accepted activation evidence (never log text). + with torch.device("meta"): + experts = nodes[0][1].inner_experts.build() + assert type(experts) is MXFP8FusedGroupedExperts + + +def test_composition_fires_on_stock_config(): + # No converter anywhere: the override alone opts the model in, dispatcher + # swap included. + config = _prepare(deepseek_v3_debugmodel()) + config.override.imports.append(_OVERRIDE_TARGET) + apply_overrides(config.override, config) + + nodes = _routed_experts_nodes(config) + assert nodes + for _fqn, cfg, _parent, _attr in nodes: + assert type(cfg.inner_experts) is MXFP8FusedGroupedExperts.Config + assert isinstance(cfg.token_dispatcher, TorchAOTokenDispatcher.Config) + assert cfg.token_dispatcher.pad_multiple == 256 + + +def test_composition_raises_on_converter_quantized_experts(): + # The composite quantizes every grouped GEMM itself; layering it on the + # MXFP8 grouped-experts converter's output is a config error that must + # raise -- never a silent fallback to the converter's unfused path. + config = _prepare(deepseek_v3_debugmodel_mxfp8()) + config.override.imports.append(_OVERRIDE_TARGET) + with pytest.raises(ValueError, match="grouped-experts converter"): + apply_overrides(config.override, config) + + +def test_composition_raises_on_non_stock_routed_experts_subclass(): + # A RoutedExperts.Config SUBCLASS must raise, not no-op (the decorator no + # longer carries exact=True, so subclass nodes are claimed and gated). + @dataclass(kw_only=True) + class _DerivedRoutedExpertsConfig(RoutedExperts.Config): + pass + + config = _prepare(deepseek_v3_debugmodel()) + node = _routed_experts_nodes(config)[0][1] + with pytest.raises(ValueError, match="stock RoutedExperts"): + mxfp8_grouped_experts(derive(node, _DerivedRoutedExpertsConfig)) + + +def test_composition_raises_on_unsupported_dims(): + config = _prepare(deepseek_v3_debugmodel()) + node = _routed_experts_nodes(config)[0][1] + node.inner_experts.hidden_dim = 100 # not a 128-multiple + with pytest.raises(ValueError, match="is_supported"): + mxfp8_grouped_experts(node) + + +def test_composition_raises_on_non_alltoall_dispatcher(): + # hybridep's padded dispatcher is not validated for the cuDNN FE 256-row + # contract; the factory must refuse it rather than swap or accept it. + # HybridEP's own config gate requires EP>1; satisfy it so the factory + # refusal (not the dispatcher validation) is what's under test. + config = deepseek_v3_debugmodel_hybridep() + config.parallelism.expert_parallel_degree = 2 + _prepare(config) + node = _routed_experts_nodes(config)[0][1] + with pytest.raises(ValueError, match="TorchAO padded"): + mxfp8_grouped_experts(node) + + +def test_composition_raises_when_ops_unavailable(monkeypatch): + import torchtitan.overrides.mxfp8_grouped_mlp as override_module + + monkeypatch.setattr(override_module, "_TORCHAO_GROUPED_MLP_OPS_AVAILABLE", False) + monkeypatch.setattr( + override_module, + "_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON", + "unavailable for the test", + ) + config = _prepare(deepseek_v3_debugmodel()) + node = _routed_experts_nodes(config)[0][1] + with pytest.raises(ValueError, match="unavailable for the test"): + override_module.mxfp8_grouped_experts(node) + + +# --------------------------------------------------------------------------- +# 3/5. Module-level fixtures (SelectiveAC + trace shape) +# --------------------------------------------------------------------------- + +_MOD_D, _MOD_F, _MOD_E = 256, 256, 4 +_MOD_SIZES = [256, 512, 256, 256] + + +def _build_module(seed): + torch.manual_seed(seed) + module = MXFP8FusedGroupedExperts.Config( + dim=_MOD_D, hidden_dim=_MOD_F, num_experts=_MOD_E + ).build() + module = module.to("cuda") + with torch.no_grad(): + # fp32 master weights: the .bfloat16() casts stay outside the Function + # so autograd routes bf16 grads back to fp32 params. + module.w13.normal_(0.0, _MOD_D**-0.5) + module.w2_EDF.normal_(0.0, _MOD_D**-0.5) + return module + + +def _module_inputs(seed=0): + torch.manual_seed(seed) + r = sum(_MOD_SIZES) + x = torch.randn(r, _MOD_D, device="cuda", dtype=torch.bfloat16) / _MOD_D**0.5 + dy = torch.randn(r, _MOD_D, device="cuda", dtype=torch.bfloat16) / _MOD_D**0.5 + num_tokens = torch.tensor(_MOD_SIZES, device="cuda") + return x, dy, num_tokens + + +def _run_module(module, x, dy, num_tokens): + x = x.clone().detach().requires_grad_(True) + y = module(x, num_tokens) + y.backward(dy) + return y.detach(), x.grad + + +def _torchao_op_counts(prof): + # Count ONLY the torchao:: custom-op events: aten CPU events double-count + # under SAC recompute, so they are never used for launch evidence. + counts = {} + for evt in prof.key_averages(): + if evt.key in _ALL_OPS: + counts[evt.key] = counts.get(evt.key, 0) + evt.count + return counts + + +def test_selective_ac_recompute_count_and_bitwise_grads(): + x, dy, num_tokens = _module_inputs() + + ref = _build_module(seed=1) + y_ref, dx_ref = _run_module(ref, x, dy, num_tokens) + + acm = _build_module(seed=2) + with torch.no_grad(): + acm.w13.copy_(ref.w13) + acm.w2_EDF.copy_(ref.w2_EDF) + wrapped = SelectiveAC(SelectiveAC.Config())._wrap_block(acm) + + # Warm up kernel JIT outside the profiled region. + _run_module(wrapped, x, dy, num_tokens) + acm.zero_grad(set_to_none=True) + + with profile(activities=[ProfilerActivity.CPU]) as prof: + y_ac, dx_ac = _run_module(wrapped, x, dy, num_tokens) + counts = _torchao_op_counts(prof) + + # The composite forward runs exactly twice under SelectiveAC (original + + # recompute; saves come from the recompute pass), backward once. The mm + # op runs 2x in the two forwards (FC2) + 1x in backward (FC1 dgrad). + assert counts.get(_OP_FWD, 0) == 2, counts + assert counts.get(_OP_MM, 0) == 3, counts + assert counts.get(_OP_BWD, 0) == 1, counts + assert counts.get(_OP_WGRAD, 0) == 2, counts + + # Deterministic kernels + save-from-recompute => bitwise-identical results. + assert torch.equal(y_ac, y_ref) + assert torch.equal(dx_ac, dx_ref) + assert torch.equal(acm.w13.grad, ref.w13.grad) + assert torch.equal(acm.w2_EDF.grad, ref.w2_EDF.grad) + + +def test_param_ref_save_survives_optimizer_step(): + # The Function saves w13/w2 by reference; the same-step backward precedes + # the optimizer update, so step -> next fwd+bwd must work. + x, dy, num_tokens = _module_inputs() + module = _build_module(seed=3) + wrapped = SelectiveAC(SelectiveAC.Config())._wrap_block(module) + optimizer = torch.optim.SGD(module.parameters(), lr=1e-3) + + _run_module(wrapped, x, dy, num_tokens) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + _, dx = _run_module(wrapped, x, dy, num_tokens) + assert dx is not None + assert torch.isfinite(dx).all() + assert module.w13.grad is not None + assert torch.isfinite(module.w13.grad).all() + + +# --------------------------------------------------------------------------- +# 4. State dict / param layout (32-block checkpoint hooks + init remap) +# --------------------------------------------------------------------------- + + +def test_state_dict_round_trips_stock_layout(): + src = _build_module(seed=4) + sd = src.state_dict() + + # Saved in the stock GroupedExperts layout, not as the fused parameter. + assert set(sd) == {"w1_EFD", "w3_EFD", "w2_EDF"} + v = _blk_view(src.w13) + f = _MOD_F + assert torch.equal(sd["w1_EFD"], v[:, :, 0].reshape(_MOD_E, f, _MOD_D)) + assert torch.equal(sd["w3_EFD"], v[:, :, 1].reshape(_MOD_E, f, _MOD_D)) + + dst = _build_module(seed=5) + dst.load_state_dict(sd) + assert torch.equal(dst.w13, src.w13) + assert torch.equal(dst.w2_EDF, src.w2_EDF) + + +def test_state_dict_loads_stock_grouped_experts_checkpoint(): + stock = GroupedExperts.Config( + dim=_MOD_D, hidden_dim=_MOD_F, num_experts=_MOD_E + ).build() + with torch.no_grad(): + for param in stock.parameters(): + param.normal_() + + fused = _build_module(seed=6).cpu() + fused.load_state_dict(stock.state_dict()) + assert torch.equal(fused.w13, _to_blk(stock.w1_EFD, stock.w3_EFD)) + assert torch.equal(fused.w2_EDF, stock.w2_EDF) + + +def test_remap_round_trip_identity(): + # elem -> 32-block -> elem is the identity, and the 32-block rows are the + # documented [gate_0..31 | up_0..31 | ...] pattern. + e, f, d = 2, 64, 8 + w1 = torch.arange(e * f * d, dtype=torch.float32).reshape(e, f, d) + w3 = -torch.arange(e * f * d, dtype=torch.float32).reshape(e, f, d) + blk = _to_blk(w1, w3) + assert torch.equal(blk[:, 0:32], w1[:, 0:32]) # first 32 gate rows + assert torch.equal(blk[:, 32:64], w3[:, 0:32]) # then their up rows + assert torch.equal(blk[:, 64:96], w1[:, 32:64]) + v = _blk_view(blk) + assert torch.equal(v[:, :, 0].reshape(e, f, d), w1) + assert torch.equal(v[:, :, 1].reshape(e, f, d), w3) + + +def test_param_init_remap_initializes_gate_and_up_blocks(): + t = torch.empty(2, 2 * 64, 8) + init = _make_w13_init( + lambda w: torch.nn.init.constant_(w, 1.0), + lambda w: torch.nn.init.constant_(w, 2.0), + ) + init(t) + v = _blk_view(t) + assert (v[:, :, 0] == 1.0).all() # gate blocks + assert (v[:, :, 1] == 2.0).all() # up blocks + # Alternating 32-row pattern in the flat layout. + assert (t[:, 0:32] == 1.0).all() + assert (t[:, 32:64] == 2.0).all() + assert (t[:, 64:96] == 1.0).all() + + +# --------------------------------------------------------------------------- +# 5. Trace shape +# --------------------------------------------------------------------------- + + +def test_trace_counts_one_fwd_two_mm_one_bwd_two_wgrad(): + x, dy, num_tokens = _module_inputs() + module = _build_module(seed=7) + + # Warm up kernel JIT outside the profiled region. + _run_module(module, x, dy, num_tokens) + module.zero_grad(set_to_none=True) + + with profile(activities=[ProfilerActivity.CPU]) as prof: + _run_module(module, x, dy, num_tokens) + counts = _torchao_op_counts(prof) + + assert counts.get(_OP_FWD, 0) == 1, counts + assert counts.get(_OP_MM, 0) == 2, counts + assert counts.get(_OP_BWD, 0) == 1, counts + assert counts.get(_OP_WGRAD, 0) == 2, counts diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index eb26e320d8..d478c9cf94 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -128,6 +128,39 @@ def deepseek_v3_debugmodel_mxfp8_fused_swiglu() -> Trainer.Config: return config +def deepseek_v3_debugmodel_mxfp8_grouped_mlp() -> Trainer.Config: + # Routed experts via the self-contained fully-fused MXFP8 grouped-MLP + # override (torchtitan/overrides/mxfp8_grouped_mlp.py), which owns the + # whole routed-expert path — the fused grouped GEMM + SwiGLU + quant + # composite plus the pad_multiple=256 TorchAO dispatcher its factory + # installs — so no grouped-experts converter runs here; only the dense + # MXFP8Linear swap, keeping the non-expert layers identical to the + # unfused MXFP8 flavor (checkpoints stay interchangeable through the + # override's stock-layout state-dict hooks). + # disable_cuda_graphs is required by the TorchAOTokenDispatcher under + # EP>1; moe_force_load_balance keeps the per-rank routed row count fixed + # (a constant R avoids fused-kernel JIT churn per newly seen R). + config = deepseek_v3_debugmodel() + model_compile_enabled = ( + config.compile.enable and "model" in config.compile.components + ) + config.model_spec = model_registry( + "debugmodel", + converters=[ + MXFP8LinearConverter.Config( + model_compile_enabled=model_compile_enabled, + fqns=["attention", "shared_experts", "feed_forward"], + ), + ], + ) + config.override.imports.append( + "torchtitan.overrides.mxfp8_grouped_mlp.mxfp8_grouped_experts" + ) + config.training.disable_cuda_graphs = True + config.debug.moe_force_load_balance = True + return config + + def deepseek_v3_debugmodel_hybridep() -> Trainer.Config: config = deepseek_v3_debugmodel() config.model_spec = model_registry( diff --git a/torchtitan/overrides/mxfp8_grouped_mlp.py b/torchtitan/overrides/mxfp8_grouped_mlp.py new file mode 100644 index 0000000000..66ebbb073d --- /dev/null +++ b/torchtitan/overrides/mxfp8_grouped_mlp.py @@ -0,0 +1,596 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyrefly: ignore-errors + +"""Fused MXFP8 grouped-MLP routed experts over the torchao cuDNN-frontend ops. + +The composite autograd Function runs the full routed-expert SwiGLU MLP with +four torchao custom ops that wrap the ``cudnn.grouped_gemm_*_wrapper_sm100`` +CuTe DSL kernels (cuDNN frontend >= 1.27, no TransformerEngine involved): + +* forward: casts -> ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` (FC1 grouped + GEMM + SwiGLU + rowwise/colwise MXFP8 RCEIL quantization of the activation) + -> ``torchao::mxfp8_grouped_gemm`` FC2. +* backward: casts -> ``torchao::mxfp8_grouped_gemm_dswiglu_bwd`` (FC2 dgrad + + dSwiGLU + dual quantization) -> ``torchao::mxfp8_grouped_gemm`` FC1 + dgrad -> ``torchao::mxfp8_grouped_gemm_wgrad`` twice. + +Every cast in this module is non-CuTe (triton/CUDA-extension torchao kernels); +only the cudnn package's own kernels are CuTe DSL. This module must never call +into ``cute_utils``. + +ABI preconditions (metadata-validated by the ops; offset VALUES are the +dispatcher's contract): ``D`` and ``F`` are positive multiples of 128, the +allocated row count is a multiple of 256, and every per-expert row count +``m[g]`` is a nonnegative multiple of **256** — the cuDNN FE kernels hard-code +``FIX_PAD_SIZE = 256``, and per-expert splits that are only 128-multiples +corrupt the chain SILENTLY and NONDETERMINISTICALLY (the corruption locus +migrates between identical-input reruns; no smoke test can prove such a +configuration safe). The %256 guarantee comes from a ``TorchAOTokenDispatcher`` +with ``pad_multiple=256``, which the override factory installs itself — hence +it targets ``RoutedExperts.Config`` (the one node owning both the dispatcher +and the inner experts configs). + +The fused ``w13`` parameter is stored ``[E, 2F, D]`` in the cuDNN 32-block GLU +row order ``[gate_0..31 | up_0..31 | gate_32..63 | up_32..63 | ...]`` so that +no per-step layout remap exists anywhere: the FC1 kernel consumes the rows +as-is and the wgrad op emits ``dw13`` directly in parameter order. Checkpoints +still save/load the stock ``w1_EFD``/``w3_EFD`` layout through this module's +state-dict hooks. + +Activate with ``--override.imports +torchtitan.overrides.mxfp8_grouped_mlp.mxfp8_grouped_experts`` on a STOCK +``RoutedExperts.Config``: the override is self-contained. Its factory swaps +the token dispatcher for the padded TorchAO variant itself and needs no +quantization converter (the composite quantizes every grouped GEMM). There +is no silent fallback: configurations the kernels cannot execute (missing +torchao ops, non-SM100 hardware, converter-quantized or otherwise non-stock +experts, dims violating the 128-alignment contract) raise an actionable +error at config-application time rather than training the unfused path +silently. +""" + +from dataclasses import dataclass + +import torch +from torch.distributed.tensor import DTensor +from torchao.prototype.moe_training.kernels.mxfp8 import ( + triton_mx_block_rearrange_2d_K_groups, + triton_mx_block_rearrange_per_group_3d, +) + +# Importing the wrapper module registers the four torchao:: custom ops; the +# cudnn package is only imported lazily inside the op bodies at first launch. +# The module is newer than several torchao releases, so its absence must +# surface as the factory's actionable config-time error (with this reason), +# not an ImportError at override-import time. Tests skip on this flag too. +try: + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + _mxfp8_grouped_mlp_kernels_available, + is_supported, + ) +except ImportError as _exc: + is_supported = None + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE = False + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON = ( + "the installed torchao has no torchao.prototype.moe_training.kernels." + f"mxfp8.cutedsl_grouped_mlp module (a torchao build that ships the " + f"fused grouped-MLP custom ops is required): {_exc}" + ) +else: + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE = bool(_mxfp8_grouped_mlp_kernels_available) + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON = ( + "" + if _TORCHAO_GROUPED_MLP_OPS_AVAILABLE + else ( + "torchao cuDNN-frontend grouped-MLP ops are unavailable in this " + "environment (needs the cudnn python package >= 1.27 with the " + "grouped_gemm_*_wrapper_sm100 kernels)." + ) + ) +from torchao.prototype.mx_formats.config import ( + MXFP8Dim1CastKernelChoice, + ScaleCalculationMode, +) +from torchao.prototype.mx_formats.kernels import triton_to_mxfp8_dim0 +from torchao.prototype.mx_formats.utils import ( + _to_mxfp8_dim1_kernel_wrapper, + to_blocked, +) +from torchao.quantization.quantize_.common import KernelPreference + +from torchtitan.components.quantization.utils import swap_token_dispatcher +from torchtitan.config import derive, override +from torchtitan.models.common.moe import GroupedExperts, RoutedExperts +from torchtitan.models.common.token_dispatcher import ( + AllToAllTokenDispatcher, + TorchAOTokenDispatcher, +) +from torchtitan.overrides.fused_swiglu import _fuse_w13_grouped_experts_sharding + +__all__ = [ + "MXFP8FusedGroupedExperts", + "mxfp8_fused_grouped_mlp", + "mxfp8_grouped_experts", +] + +_BLOCK_SIZE = 32 +_SCALING_MODE = "rceil" +# Per-expert row groups must be 256-multiples (cuDNN FE FIX_PAD_SIZE); feature +# dims must be 128-multiples (blocked-scale tiles). The row guarantee is the +# dispatcher's pad_multiple; the ops re-validate R % 256 statically. +_ROW_ALIGNMENT = 256 +_DIM_ALIGNMENT = 128 + + +def _cast_rowwise(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """1x32 rowwise RCEIL cast: row-major qdata + whole-matrix blocked scales + (identical to the per-group concatenation because every per-expert row + count is a 256-multiple, so 128-row scale tiles never straddle groups).""" + qdata, scales = triton_to_mxfp8_dim0(t, _BLOCK_SIZE, _SCALING_MODE) + return qdata, to_blocked(scales) + + +def _cast_weight_rowwise_3d(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[G, N, K]`` quantized along K: contiguous qdata + per-group blocked + scales for logical ``(N, K/32)`` per group — the rowwise ``b`` operand of + the fwd/mm ops.""" + qdata, scales = triton_to_mxfp8_dim0(w, _BLOCK_SIZE, _SCALING_MODE) + return qdata, triton_mx_block_rearrange_per_group_3d(scales) + + +def _cast_weight_colwise_3d(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[G, N, K]`` quantized along N: k-major per-group qdata + + per-group blocked scales for logical ``(K, N/32)``. + + Batched: ONE (32x1 RCEIL) cast of the flat ``[G*N, K]`` view along dim0 + + ONE ``K_groups`` swizzle with uniform scale-column offsets. Exact + because N is a 256-multiple, so 32-row quantization blocks never + straddle groups, and every group's N/32 scale columns are 4-multiples, + so the swizzle packs the same per-group ``to_blocked`` bytes densely + from the buffer start. qdata, scales, AND downstream op outputs are + BITWISE-equal to a naive per-group ``triton_to_mxfp8_dim1`` + + ``to_blocked`` loop (measured on GB200) at ~4x less time and ~6*G + fewer launches per weight. + + The cast's native ``[G, N, K]`` view carries an interleaved batch + stride ``(N, 1, G*N)``, which the cudnn wrappers reject (B must be + per-group-contiguous, k- or n-major); one fp8 repack to k-major — the + same major the rowwise casts pass — restores an accepted layout. + """ + g, n, k = w.shape + mx = _to_mxfp8_dim1_kernel_wrapper( + w.reshape(g * n, k), + _BLOCK_SIZE, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=w.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + ) + scale_offsets = ( + torch.arange(1, g + 1, device=w.device, dtype=torch.int32) + * (n // _BLOCK_SIZE) + ) + # Same pow2 quirk as _cast_colwise_grouped: the K_groups swizzle's + # tl.arange needs a power-of-2 group count; repeated end-offsets are + # zero-size groups the kernel skips. + g_pow2 = 1 << (g - 1).bit_length() + if g_pow2 != g: + scale_offsets = torch.cat( + [scale_offsets, scale_offsets[-1:].expand(g_pow2 - g)] + ) + col_scales = triton_mx_block_rearrange_2d_K_groups(mx.scale, scale_offsets) + k_pad = -(-k // 128) * 128 + flat = col_scales.reshape(-1)[: k_pad * (g * n // _BLOCK_SIZE)] + qdata = mx.qdata.view(k, g, n).permute(1, 2, 0).contiguous() + return qdata, flat.view(g, -1) + + +def _cast_colwise_grouped( + t: torch.Tensor, offsets: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Ragged colwise (32x1) RCEIL cast of ``[R, N]`` for the wgrad operands: + torchao-native qdata (``[R, N]``-logical, ``(1, R)`` strides — the cudnn + wgrad kernel accepts this major directly, verified on GB200) + PER-GROUP + blocked scales via ``triton_mx_block_rearrange_2d_K_groups``. + + Quantizing the whole ragged tensor in one launch is safe ONLY because + every per-expert row count is a 256-multiple, so 32-row quantization + blocks never straddle an expert boundary. + + The K_groups swizzle's ``tl.arange(0, num_groups)`` requires a power-of-2 + bound, so the scale-column offsets are padded to the next power of 2 by + repeating the final offset — repeated end-offsets are zero-size groups the + kernel skips. Its output also carries 4 trailing padding columns per + group slot (d2h-sync avoidance); with 256-multiple groups the real blocks + pack densely from the start of the buffer (total real content = + ``round_up(N, 128) * offsets[-1]/32`` elements <= ``... * R/32``), so the + flat buffer is statically sliced to ``round_up(N, 128) * R/32`` — the op's + documented maximum — without any device sync; the wgrad kernel never reads + past the offsets-bounded span. + """ + mx = _to_mxfp8_dim1_kernel_wrapper( + t, + _BLOCK_SIZE, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=t.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + ) + scale_offsets = offsets // _BLOCK_SIZE + g = scale_offsets.shape[0] + g_pow2 = 1 << (g - 1).bit_length() + if g_pow2 != g: + scale_offsets = torch.cat( + [scale_offsets, scale_offsets[-1:].expand(g_pow2 - g)] + ) + col_scales = triton_mx_block_rearrange_2d_K_groups(mx.scale, scale_offsets) + r, n = t.shape + n_pad = -(-n // 128) * 128 + # mx.qdata is [N, R]-shaped; .t() presents the op's [R, N]-logical view. + return mx.qdata.t(), col_scales.reshape(-1)[: n_pad * (r // _BLOCK_SIZE)] + + +class _MXFP8GroupedMLP(torch.autograd.Function): + """Composite MXFP8 grouped SwiGLU MLP over the four cudnn-FE ops. + + All inputs are plain BF16 CUDA tensors (the module prologue casts and + un-DTensors them); ``dy`` arrives as contiguous BF16 ``[R, D]``. ``w13`` + is ``[G, 2F, D]`` in 32-block GLU row order. All backward-only casts are + lazy: forward quantizes only what forward consumes (the rowwise views); + backward requantizes the colwise weight views and the colwise ``x`` from + the saved BF16 references — safe because the same-step backward always + precedes the optimizer update (an update in between trips the autograd + version counter), and cheaper under per-op SAC because the forward (and + thus any forward-side cast) re-runs in the recompute pass. + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + offsets: torch.Tensor, + ) -> torch.Tensor: + x_row_q, x_row_sf = _cast_rowwise(x) + w13_row_q, w13_row_sf = _cast_weight_rowwise_3d(w13) + z, h_row_q, h_row_sf, h_col_q, h_col_sf = ( + torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_row_q, + x_row_sf, + w13_row_q, + w13_row_sf.reshape(-1), + offsets, + ) + ) + w2_row_q, w2_row_sf = _cast_weight_rowwise_3d(w2) + # FC2 forward: b [G, N=D, K=F] rowwise (quantized along F = the + # contraction), row-major as cast. + y = torch.ops.torchao.mxfp8_grouped_gemm( + h_row_q, + h_row_sf, + w2_row_q, + w2_row_sf.reshape(-1), + offsets, + ) + # x is saved BF16; its colwise cast is deferred to backward. Under + # per-op SAC the whole forward re-runs in the recompute pass, so a + # forward-side cast would execute twice per step for one consumer + # (the FC1 wgrad) — deferring makes it run exactly once. Safe for the + # same reason the weight casts are lazy: the same-step backward always + # precedes the optimizer update. + ctx.save_for_backward(z, h_col_q, h_col_sf, x, offsets, w13, w2) + return y + + @staticmethod + def backward(ctx, dy: torch.Tensor): + z, h_col_q, h_col_sf, x, offsets, w13, w2 = ctx.saved_tensors + if x.shape[0] == 0: + # A rank whose local experts received zero routed tokens: every + # grad is zero by construction, and torchao's CUDA colwise cast + # rejects 0-row inputs, so skip the cast/GEMM chain outright. + # (The forward needs no such guard: its rowwise casts accept 0 + # rows and the ops early-return at R == 0.) + return torch.empty_like(x), torch.zeros_like(w13), torch.zeros_like(w2), None + # The casts assert contiguity; dy is contiguous today (BF16 [R, D] + # stride (D, 1)) but that is a live invariant, not a given. + dy = dy.contiguous() + dy_row_q, dy_row_sf = _cast_rowwise(dy) + # w2 colwise (quantized along D = the dgrad contraction): the bwd op's + # ABI takes the [G, D, F]-logical cast output as-is. + w2_col_q, w2_col_sf = _cast_weight_colwise_3d(w2) + dz_row_q, dz_row_sf, dz_col_q, dz_col_sf = ( + torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + dy_row_q, + dy_row_sf, + w2_col_q, + w2_col_sf.reshape(-1), + z, + offsets, + ) + ) + # FC1 dgrad: b [G, N=D, K=2F] quantized along 2F. The colwise cast + # yields [G, 2F, D]; the mm op's b orientation is [G, N, K], so the + # call site transposes (unlike ``torch._scaled_grouped_mm``, whose + # [G, K, N] mat2 convention would take the cast output as-is). + w13_col_q, w13_col_sf = _cast_weight_colwise_3d(w13) + dx = torch.ops.torchao.mxfp8_grouped_gemm( + dz_row_q, + dz_row_sf, + w13_col_q.transpose(-2, -1), + w13_col_sf.reshape(-1), + offsets, + ) + dy_col_q, dy_col_sf = _cast_colwise_grouped(dy, offsets) + x_col_q, x_col_sf = _cast_colwise_grouped(x, offsets) + # dw2 [G, D, F] = dy^T @ h per expert; dw13 [G, 2F, D] = dz^T @ x per + # expert, landing directly in the 32-block parameter order. + dw2 = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, h_col_q, h_col_sf, offsets + ) + dw13 = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dz_col_q, dz_col_sf, x_col_q, x_col_sf, offsets + ) + return dx, dw13, dw2, None + + +def mxfp8_fused_grouped_mlp( + x: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + """Fused MXFP8 grouped SwiGLU MLP: ``x [R, D] -> y [R, D]`` (BF16). + + Args: + x: BF16 ``[R, D]`` expert-major padded rows; every per-expert group + is a multiple of 256 rows. + w13: BF16 ``[G, 2F, D]`` fused gate/up weight in 32-block GLU row + order (32 gate rows, then the same features' 32 up rows, ...). + w2: BF16 ``[G, D, F]`` down-projection weight. + offsets: int32 CUDA ``[G]`` exclusive per-expert end rows, + ``offsets[-1] <= R``. Rows past ``offsets[-1]`` of ``y`` (and of + ``dx`` in backward) are left UNWRITTEN. + """ + return _MXFP8GroupedMLP.apply(x, w13, w2, offsets) + + +def _stock_pair_to_w13(w1: torch.Tensor, w3: torch.Tensor) -> torch.Tensor: + """Stock ``w1_EFD``/``w3_EFD`` ``[E, F, D]`` pair -> the 32-block GLU + row-ordered ``w13 [E, 2F, D]`` (the inverse of ``_split_w13_on_save``'s + view).""" + e, f, d = w1.shape + return ( + torch.stack([w1, w3], dim=2) # [E, F, 2, D] + .view(e, f // 32, 32, 2, d) + .permute(0, 1, 3, 2, 4) + .reshape(e, 2 * f, d) + ) + + +def _make_w13_init(gate_init, up_init): + """Initializer for the 32-block-ordered ``w13 [E, 2F, D]`` from the stock + per-half initializers. + + Each half is initialized IN PLACE through a strided sub-view of ``w13`` + (the same ``(E, F/32, 2, 32, D)`` view the save hook uses), mirroring + ``_make_fused_gate_up_init``: initializing the parameter itself keeps + DTensor init semantics under parallelism — shard-distinct, globally + consistent draws through the DTensor RNG tracker. (Plain-tensor + temporaries would draw IDENTICAL values on every rank — torchtitan + seeds all non-PP ranks the same — silently duplicating experts across + EP/FSDP shards; and a plain ``copy_`` into a DTensor raises.) The cost: + fan-computing initializers (e.g. ``nn.init.xavier_uniform_``) would see + the blocked 4-D sub-view geometry instead of stock ``(E, F, D)``; every + in-tree ``w1_EFD``/``w3_EFD`` initializer is a fixed-std + ``trunc_normal_``, which is shape-agnostic.""" + + def _init(t: torch.Tensor) -> None: + e, two_f, d = t.shape + v = t.view(e, two_f // 64, 2, 32, d) + gate_init(v[:, :, 0]) # gate (stock w1) half + up_init(v[:, :, 1]) # up (stock w3) half + + return _init + + +def _w13_grouped_experts_param_init(param_init: dict | None) -> dict | None: + """Remap ``w1_EFD`` / ``w3_EFD`` initializers onto the 32-block ``w13``. + + Other entries (e.g. ``w2_EDF``) are kept as-is. + """ + if param_init is None: + return None + w1_init = param_init.get("w1_EFD") + w3_init = param_init.get("w3_EFD") + fused = {k: v for k, v in param_init.items() if k not in ("w1_EFD", "w3_EFD")} + if w1_init is not None and w3_init is not None: + fused["w13"] = _make_w13_init(w1_init, w3_init) + return fused or None + + +class MXFP8FusedGroupedExperts(GroupedExperts): + """Routed experts computed by the cudnn-FE MXFP8 grouped-MLP composite. + + ``w13`` has shape ``(num_experts, 2*hidden_dim, dim)`` in the cuDNN + 32-block GLU row order (NOT the FusedGroupedExperts ``(E, F, 2, D)`` + element interleave); checkpoints save/load the stock ``w1_EFD``/``w3_EFD`` + layout through this class's own hooks. The layout is fixed at + parameter-registration time so the training loop performs zero per-step + remaps and the wgrad op emits ``dw13`` directly in parameter order. + """ + + @dataclass(kw_only=True, slots=True) + class Config(GroupedExperts.Config): + # No new fields in v1. derive() carries param_init / sharding_config / + # dim / hidden_dim / num_experts from the stock config by name; + # any future knob must be re-declared here or derive() drops it. + pass + + def __init__(self, config: Config): + super().__init__(config) + + # delete separate w1/w3 and fuse in 32-block GLU row order + del self.w1_EFD + del self.w3_EFD + self.w13 = torch.nn.Parameter( + torch.empty(config.num_experts, 2 * config.hidden_dim, config.dim) + ) + + self.register_state_dict_post_hook(self._split_w13_on_save) + self.register_load_state_dict_pre_hook(self._merge_w13_on_load) + + def forward( + self, + x_RD: torch.Tensor, + num_tokens_per_expert_E: torch.Tensor, + ) -> torch.Tensor: + if isinstance(self.w13, DTensor): + w13 = self.w13.to_local() + assert isinstance(self.w2_EDF, DTensor) + w2_EDF = self.w2_EDF.to_local() + else: + w13 = self.w13 + w2_EDF = self.w2_EDF + + # The factory gate can only validate the GLOBAL dims (the config + # carries sharding placements, not mesh degrees), so the local shard + # dims are validated here at first call: under dense tensor + # parallelism (expert_parallel_degree=1, tensor_parallel_degree>1) + # w13/w2 are Shard(1)/Shard(2)-split on hidden_dim, and a TP degree + # with hidden_dim/tp not a 128-multiple would otherwise fail deep + # inside the first fused op launch. + local_f = w13.shape[1] // 2 + local_d = w13.shape[2] + if local_f % _DIM_ALIGNMENT != 0 or local_d % _DIM_ALIGNMENT != 0: + raise ValueError( + f"mxfp8_grouped_experts: the LOCAL expert shard dims " + f"(D={local_d}, F={local_f} from w13 of local shape " + f"{tuple(w13.shape)}) must be positive multiples of " + f"{_DIM_ALIGNMENT} for the fused cuDNN-FE grouped-MLP " + "kernels. This typically means tensor parallelism split " + "hidden_dim into a non-128-multiple shard; choose a " + "tensor_parallel_degree such that hidden_dim / tp stays a " + f"multiple of {_DIM_ALIGNMENT}, or drop the " + "mxfp8_grouped_experts override import to use the unfused " + "MXFP8 path." + ) + + offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) + # The .bfloat16() casts stay OUTSIDE the Function so autograd handles + # high-precision master-weight configs and dy reaches backward() BF16. + y_RD = _MXFP8GroupedMLP.apply( + x_RD.bfloat16(), w13.bfloat16(), w2_EDF.bfloat16(), offsets_E + ) + return y_RD.type_as(x_RD) + + @staticmethod + def _split_w13_on_save(module, state_dict, prefix, local_metadata) -> None: + """Save the 32-block fused weight as stock ``w1_EFD`` / ``w3_EFD``.""" + w13 = state_dict.pop(f"{prefix}w13") + e, two_f, d = w13.shape + f = two_f // 2 + v = w13.view(e, f // 32, 2, 32, d) + state_dict[f"{prefix}w1_EFD"] = v[:, :, 0].reshape(e, f, d) + state_dict[f"{prefix}w3_EFD"] = v[:, :, 1].reshape(e, f, d) + + @staticmethod + def _merge_w13_on_load(module, state_dict, prefix, *args) -> None: + """Combine stock ``w1_EFD`` / ``w3_EFD`` into the 32-block ``w13``.""" + w1_key, w3_key = f"{prefix}w1_EFD", f"{prefix}w3_EFD" + if w1_key in state_dict and w3_key in state_dict: + state_dict[f"{prefix}w13"] = _stock_pair_to_w13( + state_dict.pop(w1_key), state_dict.pop(w3_key) + ) + + +@override( + target=RoutedExperts.Config, + description="Fully-fused MXFP8 grouped MLP (grouped GEMM + SwiGLU + quant)", +) +def mxfp8_grouped_experts(cfg: RoutedExperts.Config) -> RoutedExperts.Config: + """Swap stock grouped experts for the cudnn-FE fused MXFP8 composite. + + Targets ``RoutedExperts.Config`` because the composite constrains BOTH + children: the ``inner_experts`` (replaced with the fused module) and the + ``token_dispatcher``, which the factory swaps for the padded TorchAO + variant that guarantees the ABI's ``m[g] % 256``. Self-contained and + fail-loud: no converter is involved, and any configuration the kernels + cannot execute raises at config-application time instead of silently + training the unfused path. + """ + if not ( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + ): + raise ValueError( + "mxfp8_grouped_experts requires CUDA device capability exactly " + "(10, 0) (the torchao ops wrap cudnn grouped_gemm_*_wrapper_sm100 " + "kernels); remove the override or run on supported hardware." + ) + if not _TORCHAO_GROUPED_MLP_OPS_AVAILABLE: + raise ValueError( + f"mxfp8_grouped_experts: {_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON}" + ) + if type(cfg) is not RoutedExperts.Config: + raise ValueError( + "mxfp8_grouped_experts targets the stock RoutedExperts.Config, " + f"got {type(cfg).__qualname__}; narrow this override's fqns or " + "remove the conflicting override." + ) + experts = cfg.inner_experts + if type(experts) is not GroupedExperts.Config: + raise ValueError( + "mxfp8_grouped_experts requires the stock GroupedExperts.Config, " + f"but inner_experts is {type(experts).__qualname__}. The " + "composite quantizes every grouped GEMM itself — do not combine " + "it with the MXFP8 grouped-experts converter." + ) + # GLOBAL dims only: the config carries sharding placements (e.g. the TP + # Shard(1) on hidden_dim) but not mesh degrees, so the per-rank shard dims + # cannot be computed here. MXFP8FusedGroupedExperts.forward re-validates + # the LOCAL dims at first call and raises with the config fix. + if not is_supported(experts.dim, experts.hidden_dim): + raise ValueError( + f"mxfp8_grouped_experts: is_supported(D={experts.dim}, " + f"F={experts.hidden_dim}) is False; both dims must be positive " + f"multiples of {_DIM_ALIGNMENT}." + ) + + # The %256 padding contract is the factory's own work (no converter + # involved): the cuDNN FE kernels require per-expert groups padded to + # 256 (FIX_PAD_SIZE) — 128-multiple-only splits corrupt silently and + # nondeterministically. + dispatcher = cfg.token_dispatcher + if isinstance(dispatcher, TorchAOTokenDispatcher.Config): + dispatcher.pad_multiple = _ROW_ALIGNMENT + elif isinstance(dispatcher, AllToAllTokenDispatcher.Config): + swap_token_dispatcher(cfg, pad_multiple=_ROW_ALIGNMENT) + else: + raise ValueError( + f"mxfp8_grouped_experts: token_dispatcher is " + f"{type(dispatcher).__qualname__}; only the TorchAO padded " + "dispatcher (swapped in from the stock all-to-all) is validated " + "for the per-expert 256-row contract." + ) + + fused = derive(experts, MXFP8FusedGroupedExperts.Config) + # The w1_EFD/w3_EFD -> w13 param-init remap is factory work (it is not + # inherited through derive()); the sharding remap simply carries w1's + # layout onto w13 (both rank-3, same shard axes), never a gate: + # deepseek EP=1 configs carry TP Shard(1) on w1_EFD unconditionally, so + # raising on a fused-dim Shard would reject exactly the single-GPU debug + # mode. + fused.param_init = _w13_grouped_experts_param_init(fused.param_init) + if fused.sharding_config is not None: + fused.sharding_config = _fuse_w13_grouped_experts_sharding( + fused.sharding_config + ) + cfg.inner_experts = fused + return cfg