diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mlp.py b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py new file mode 100644 index 0000000000..f53c71d174 --- /dev/null +++ b/test/prototype/moe_training/test_mxfp8_grouped_mlp.py @@ -0,0 +1,979 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the MXFP8 fused grouped-MLP custom ops. + +The four ops wrap the cudnn-frontend package's CuTe DSL grouped-GEMM kernels +(``cudnn.grouped_gemm_{glu,quant,dglu,wgrad}_wrapper_sm100``). + +Every numerics gate is DERIVED at test time, never hard-coded: + +* ``refA`` (GEMM-exactness) gates come from the variability between two + legitimate evaluations of the same dequantized-operand reference that differ + only in FP32 reduction order (whole-K vs chunked-K), minus a 12 dB margin, + capped at 60 dB. Measured bands on GB200: 63-75 dB at the debug shapes + (gates land at the 51-60 dB cap region); op outputs measure 85-160 dB. +* ``refB`` (independent-chain) gates come from the SQNR of a quantized-unfused + evaluation against the exact FP32 chain computed from the ORIGINAL BF16 + inputs with no quantization at all, minus a 6 dB margin. This reference + shares NO layout helper with the op inputs, so a self-consistent layout bug + (wrong scale blocking built and decoded the same wrong way) cannot pass it. + Measured band: ~30-40 dB at these shapes (pure MXFP8 requantization error). + +Layout vocabulary (probe-derived): columnwise operands are accepted in ANY +major -- "rowmajor" here means un-transposed ``[R, N]`` row-major bytes (also +the layout the fwd/bwd ops emit for their columnwise outputs) and "native" +means the transposed-memory layout torchao's dim1 quantizers produce. +Columnwise scale buffers are PER-GROUP blocked (each expert's block +``to_blocked``-ed independently, concatenated); whole-matrix blocking has the +same byte count and is silently wrong -- a dedicated negative control asserts +the gap. +""" + +import pytest +import torch +import torch.nn.functional as F +from torch._subclasses.fake_tensor import FakeTensorMode + +from torchao.utils import is_sm_version + +# Exactly SM 10.0, matching the ops module's availability gate: the wrapped +# cudnn kernels are *_sm100-specific. +if not (torch.cuda.is_available() and is_sm_version(10, 0)): + pytest.skip( + "MXFP8 fused grouped MLP requires CUDA SM100", + allow_module_level=True, + ) + +try: + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + _mxfp8_grouped_mlp_kernels_available, + _mxfp8_grouped_mlp_unavailable_reason, + is_supported, + validate_group_offsets, + ) +except ImportError: + pytest.skip( + "installed torchao does not provide the cutedsl_grouped_mlp module", + allow_module_level=True, + ) + +if not _mxfp8_grouped_mlp_kernels_available: + pytest.skip( + f"cudnn-frontend grouped-GEMM wrappers unavailable: " + f"{_mxfp8_grouped_mlp_unavailable_reason}", + allow_module_level=True, + ) + +from torchao.float8.float8_utils import compute_error +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 from_blocked, to_blocked + +_E4M3 = torch.float8_e4m3fn +_E8M0 = torch.float8_e8m0fnu +_BLOCK = 32 +_RCEIL = ScaleCalculationMode.RCEIL + +_OPS = torch.ops.torchao + + +# --------------------------------------------------------------------------- +# Pure-torch quantization / dequantization helpers. +# --------------------------------------------------------------------------- + + +def _bytes(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _e8m0_to_f64(s: torch.Tensor) -> torch.Tensor: + u = s.view(torch.uint8).to(torch.int32) + out = torch.exp2((u - 127).to(torch.float64)) + return torch.where(u == 255, torch.full_like(out, float("nan")), out) + + +def _quant_rowwise(x: torch.Tensor): + """[M, K] -> (qdata [M, K] e4m3 row-major, flat blocked scales).""" + s, q = to_mx(x, _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q, to_blocked(s.view(_E8M0)).view(_E8M0) + + +def _quant_colwise(x: torch.Tensor, native: bool): + """[M, K] quantized along M in 32-blocks. + + native=False: un-transposed row-major [M, K] bytes ("rowmajor"). + native=True: the dim1-quantizer layout, [M, K] logical with (1, M) strides. + Scales: flat blocked of the transposed [K, M/32] scale matrix (one group). + """ + M, K = x.shape + if M == 0: + return ( + torch.empty(0, K, dtype=_E4M3, device=x.device), + torch.empty(0, dtype=_E8M0, device=x.device), + ) + s_t, q_t = to_mx(x.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + q = q_t.t() if native else q_t.t().contiguous() + return q, to_blocked(s_t.view(_E8M0)).view(_E8M0) + + +def _cat8(ts, dim=0): + dt = ts[0].dtype + return torch.cat([t.view(torch.uint8) for t in ts], dim).view(dt) + + +def _quant_colwise_grouped(x: torch.Tensor, sizes, native: bool): + """Ragged [R, K]: per-group colwise quantization, per-group blocked scales.""" + qs, sfs = [], [] + off = 0 + for m in sizes: + q, sf = _quant_colwise(x[off : off + m], native=False) + qs.append(q) + sfs.append(sf.reshape(-1)) + off += m + q = _cat8(qs, 0) + if native: + q = q.t().contiguous().t() # values identical; (1, R) strides + return q, _cat8(sfs) + + +def _quant_weight_rowwise(w: torch.Tensor): + """[G, N, K] quantized along K -> (contiguous stack, per-group blocked).""" + qs, sfs = [], [] + for g in range(w.shape[0]): + q, sf = _quant_rowwise(w[g]) + qs.append(q.view(torch.uint8)) + sfs.append(sf.reshape(-1)) + return torch.stack(qs).view(_E4M3), _cat8(sfs) + + +def _quant_weight_colwise(w: torch.Tensor, native: bool = False): + """[G, N, K] quantized along N. + + native=False: contiguous row-major [G, N, K] bytes. + native=True: the dim1-quantizer memory-transposed major -- [G, N, K] + logical with per-group (1, N) strides (values identical). + """ + qs, sfs = [], [] + for g in range(w.shape[0]): + q, sf = _quant_colwise(w[g], native=False) + qs.append(q.view(torch.uint8)) + sfs.append(sf.reshape(-1)) + q = torch.stack(qs).view(_E4M3) + if native: + q = q.transpose(-2, -1).contiguous().transpose(-2, -1) + return q, _cat8(sfs) + + +def _dequant_rowwise(q: torch.Tensor, sf_flat: torch.Tensor): + M, K = q.shape + s = _e8m0_to_f64(from_blocked(sf_flat.view(_E8M0), M, K // _BLOCK)) + return (q.to(torch.float64) * s.repeat_interleave(_BLOCK, dim=1)).to(torch.float32) + + +def _dequant_colwise_grouped(q: torch.Tensor, sf_flat: torch.Tensor, sizes, K: int): + """q [R, K] logical (any major), per-group blocked scales -> f32 [R, K].""" + R = q.shape[0] + out = torch.empty(R, K, dtype=torch.float32, device=q.device) + off, soff = 0, 0 + for m in sizes: + if m == 0: + continue + n = K * (m // _BLOCK) + s_t = _e8m0_to_f64( + from_blocked(sf_flat[soff : soff + n].view(_E8M0), K, m // _BLOCK) + ) + block = q[off : off + m].to(torch.float64) + out[off : off + m] = (block * s_t.t().repeat_interleave(_BLOCK, dim=0)).to( + torch.float32 + ) + off += m + soff += n + return out + + +def _mk_offsets(sizes, device): + return torch.cumsum( + torch.tensor(sizes, dtype=torch.int64, device=device), dim=0 + ).to(torch.int32) + + +def _zsplit(z: torch.Tensor, hidden: int): + """32-block interleaved [R, 2F] -> (gate [R, F], up [R, F]).""" + v = z.view(z.shape[0], hidden // _BLOCK, 2, _BLOCK) + return ( + v[:, :, 0, :].reshape(z.shape[0], hidden), + v[:, :, 1, :].reshape(z.shape[0], hidden), + ) + + +def _to_32block(w13_elem: torch.Tensor) -> torch.Tensor: + """Element-interleaved [G, F, 2, D] -> 32-block GLU order [G, 2F, D].""" + G, hidden, _, D = w13_elem.shape + return ( + w13_elem.view(G, hidden // _BLOCK, _BLOCK, 2, D) + .permute(0, 1, 3, 2, 4) + .reshape(G, 2 * hidden, D) + .contiguous() + ) + + +def _grouped_matmul(a_f32, b_f32_per_group, sizes, transpose_b: bool, chunks: int = 1): + """Per-group f32 matmul with an optional chunked-K reduction order.""" + R = a_f32.shape[0] + N = b_f32_per_group[0].shape[0 if transpose_b else 1] + out = torch.zeros(R, N, dtype=torch.float32, device=a_f32.device) + off = 0 + for g, m in enumerate(sizes): + b = b_f32_per_group[g] + bt = b.t() if transpose_b else b + if chunks == 1: + out[off : off + m] = a_f32[off : off + m] @ bt + else: + K = a_f32.shape[1] + step = K // chunks + acc = torch.zeros(m, N, dtype=torch.float32, device=a_f32.device) + for c in range(chunks): + lo, hi = c * step, K if c == chunks - 1 else (c + 1) * step + acc += a_f32[off : off + m, lo:hi] @ bt[lo:hi] + out[off : off + m] = acc + off += m + return out + + +def _refA_gate(ref_whole: torch.Tensor, ref_chunked: torch.Tensor) -> float: + """GEMM-exactness gate from the reduction-order variability band - 12 dB.""" + band = compute_error(ref_whole.bfloat16(), ref_chunked.bfloat16()).item() + return min(band - 12.0, 60.0) + + +def _dswiglu(dh, gate, up): + s = torch.sigmoid(gate) + return dh * up * (s * (1 + gate * (1 - s))), dh * F.silu(gate) + + +# --------------------------------------------------------------------------- +# Case construction: quantize everything once per case, with references. +# --------------------------------------------------------------------------- + +_CASES = { + # name: (D, F, sizes) + "dbg_zero_token": (256, 256, [256, 0, 512, 256]), + "dnef": (256, 384, [512, 256]), + "g1": (256, 256, [512]), + "16b": (2048, 1408, [256] * 8), +} + + +def _build_case(D, hidden, sizes, device="cuda", seed=0): + torch.manual_seed(seed) + G = len(sizes) + R = sum(sizes) + c = {} + c["sizes"], c["G"], c["R"], c["D"], c["F"] = sizes, G, R, D, hidden + c["offsets"] = _mk_offsets(sizes, device) + c["x"] = torch.randn(R, D, dtype=torch.bfloat16, device=device) * 0.5 + w13_elem = torch.randn(G, hidden, 2, D, dtype=torch.bfloat16, device=device) * 0.02 + c["w13"] = _to_32block(w13_elem) + c["w2"] = torch.randn(G, D, hidden, dtype=torch.bfloat16, device=device) * 0.02 + c["dy"] = torch.randn(R, D, dtype=torch.bfloat16, device=device) * 0.5 + + c["x_q"], c["x_sf"] = _quant_rowwise(c["x"]) + c["w13_q"], c["w13_sf"] = _quant_weight_rowwise(c["w13"]) + c["w2_q"], c["w2_sf"] = _quant_weight_rowwise(c["w2"]) + c["w13c_q"], c["w13c_sf"] = _quant_weight_colwise(c["w13"]) + c["w2c_q"], c["w2c_sf"] = _quant_weight_colwise(c["w2"]) + c["dy_q"], c["dy_sf"] = _quant_rowwise(c["dy"]) + c["x_colq"], c["x_col_sf"] = _quant_colwise_grouped(c["x"], sizes, native=True) + c["dy_colq"], c["dy_col_sf"] = _quant_colwise_grouped(c["dy"], sizes, native=True) + + c["x_deq"] = _dequant_rowwise(c["x_q"], c["x_sf"]) + c["w13_deq"] = [ + _dequant_rowwise(c["w13_q"][g], c["w13_sf"].view(G, -1)[g]) for g in range(G) + ] + c["w2_deq"] = [ + _dequant_rowwise(c["w2_q"][g], c["w2_sf"].view(G, -1)[g]) for g in range(G) + ] + return c + + +@pytest.fixture(scope="module") +def dbg(): + return _build_case(*_CASES["dbg_zero_token"]) + + +# --------------------------------------------------------------------------- +# Registration / availability / fakes (no GPU launch). +# --------------------------------------------------------------------------- + + +def test_ops_registered(): + for name in ( + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_wgrad", + ): + assert hasattr(_OPS, name), f"torchao::{name} is not registered" + + +def test_is_supported(): + assert is_supported(2048, 1408) + assert is_supported(256, 256) + assert not is_supported(192, 256) + assert not is_supported(256, 64) + assert not is_supported(0, 256) + + +def _fake_chain_shapes(R=512, D=256, hidden=256, G=2): + N1 = 2 * hidden + dev = "cuda" + x_q = torch.empty(R, D, dtype=_E4M3, device=dev) + x_sf = torch.empty(R * D // _BLOCK, dtype=_E8M0, device=dev) + w13_q = torch.empty(G, N1, D, dtype=_E4M3, device=dev) + w13_sf = torch.empty(G * N1 * D // _BLOCK, dtype=_E8M0, device=dev) + offsets = torch.empty(G, dtype=torch.int32, device=dev) + outs = {} + outs["fwd"] = _OPS.mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets) + w2_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) + w2_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) + outs["mm"] = _OPS.mxfp8_grouped_gemm( + outs["fwd"][1], outs["fwd"][2], w2_q, w2_sf, offsets + ) + dy_q = torch.empty(R, D, dtype=_E4M3, device=dev) + dy_sf = torch.empty(R * D // _BLOCK, dtype=_E8M0, device=dev) + w2c_q = torch.empty(G, D, hidden, dtype=_E4M3, device=dev) + w2c_sf = torch.empty(G * D * hidden // _BLOCK, dtype=_E8M0, device=dev) + outs["bwd"] = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + dy_q, dy_sf, w2c_q, w2c_sf, outs["fwd"][0], offsets + ) + outs["wgrad"] = _OPS.mxfp8_grouped_gemm_wgrad( + torch.empty(R, D, dtype=_E4M3, device=dev), + torch.empty(D * R // _BLOCK, dtype=_E8M0, device=dev), + torch.empty(R, hidden, dtype=_E4M3, device=dev), + torch.empty( + ((hidden + 127) // 128 * 128) * R // _BLOCK, dtype=_E8M0, device=dev + ), + offsets, + ) + return outs + + +def test_fake_contracts_match_specs(): + """All four fakes produce the documented shapes/dtypes/contiguity.""" + with FakeTensorMode(): + outs = _fake_chain_shapes() + R, D, hidden, N1 = 512, 256, 256, 512 + z, hq, hsf, hcq, hcsf = outs["fwd"] + assert tuple(z.shape) == (R, N1) and z.dtype == torch.bfloat16 + assert tuple(hq.shape) == (R, hidden) and hq.dtype == _E4M3 + assert hsf.numel() == R * hidden // _BLOCK and hsf.dtype == _E8M0 + assert tuple(hcq.shape) == (R, hidden) and hcq.dtype == _E4M3 + assert hcsf.numel() == hidden * R // _BLOCK + assert all(t.is_contiguous() for t in outs["fwd"]) + y = outs["mm"] + assert tuple(y.shape) == (R, D) and y.dtype == torch.bfloat16 and y.is_contiguous() + dz_q, dz_sf, dzc_q, dzc_sf = outs["bwd"] + assert tuple(dz_q.shape) == (R, N1) and dz_sf.numel() == R * N1 // _BLOCK + assert tuple(dzc_q.shape) == (R, N1) and dzc_sf.numel() == N1 * R // _BLOCK + dw = outs["wgrad"] + assert tuple(dw.shape) == (2, D, hidden) and dw.dtype == torch.bfloat16 + + +# --------------------------------------------------------------------------- +# Full-chain numerics: two references per stage, derived gates. +# --------------------------------------------------------------------------- + + +def _run_chain(c): + """fwd -> FC2 mm -> bwd -> FC1-dgrad mm -> wgrad x2, production layouts.""" + r = {} + r["z"], r["h_q"], r["h_sf"], r["h_colq"], r["h_col_sf"] = ( + _OPS.mxfp8_grouped_gemm_swiglu_fwd( + c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] + ) + ) + r["y"] = _OPS.mxfp8_grouped_gemm( + r["h_q"], r["h_sf"], c["w2_q"], c["w2_sf"], c["offsets"] + ) + r["dz_q"], r["dz_sf"], r["dz_colq"], r["dz_col_sf"] = ( + _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + c["dy_q"], c["dy_sf"], c["w2c_q"], c["w2c_sf"], r["z"], c["offsets"] + ) + ) + # FC1 dgrad: colwise weight cast enters the mm op TRANSPOSED into + # [G, N=D, K=2F]. + r["dx"] = _OPS.mxfp8_grouped_gemm( + r["dz_q"], + r["dz_sf"], + c["w13c_q"].transpose(-2, -1), + c["w13c_sf"], + c["offsets"], + ) + # Production wgrad layout mixes: native dy x kernel h; kernel dz x native x. + r["dw2"] = _OPS.mxfp8_grouped_gemm_wgrad( + c["dy_colq"], c["dy_col_sf"], r["h_colq"], r["h_col_sf"], c["offsets"] + ) + r["dw13"] = _OPS.mxfp8_grouped_gemm_wgrad( + r["dz_colq"], r["dz_col_sf"], c["x_colq"], c["x_col_sf"], c["offsets"] + ) + return r + + +@pytest.mark.parametrize("case", list(_CASES)) +def test_chain_numerics(case): + D, hidden, sizes = _CASES[case] + c = _build_case(D, hidden, sizes) + G, R = c["G"], c["R"] + r = _run_chain(c) + + # --- z: refA (dequantized operands, two reduction orders) + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], sizes, transpose_b=True) + z_ref2 = _grouped_matmul( + c["x_deq"], c["w13_deq"], sizes, transpose_b=True, chunks=4 + ) + gate_a = _refA_gate(z_ref, z_ref2) + z_db = compute_error(z_ref.bfloat16(), r["z"]).item() + assert z_db >= gate_a, f"z {z_db:.1f} dB < derived refA gate {gate_a:.1f}" + + # --- z: refB (exact chain from ORIGINAL bf16 tensors; no quant helpers) + z_exact = _grouped_matmul( + c["x"].float(), [w.float() for w in c["w13"]], sizes, transpose_b=True + ) + band_b = compute_error(z_exact, z_ref).item() # quantization band + z_db_b = compute_error(z_exact.bfloat16(), r["z"]).item() + assert z_db_b >= band_b - 6.0, ( + f"z vs independent exact chain {z_db_b:.1f} dB < band {band_b:.1f} - 6" + ) + + # --- h (both quantized orientations) vs silu ref from the KERNEL's z + gate_f, up_f = _zsplit(r["z"].float(), hidden) + h_ref = F.silu(gate_f) * up_f + band_h = compute_error( + h_ref, _dequant_rowwise(*_quant_rowwise(h_ref.bfloat16())) + ).item() + h_deq = _dequant_rowwise(r["h_q"], r["h_sf"]) + h_db = compute_error(h_ref, h_deq).item() + assert h_db >= band_h - 6.0, f"h {h_db:.1f} dB < requant band {band_h:.1f} - 6" + h_col_deq = _dequant_colwise_grouped(r["h_colq"], r["h_col_sf"], sizes, hidden) + h_col_db = compute_error(h_ref, h_col_deq).item() + assert h_col_db >= band_h - 6.0, ( + f"h_col {h_col_db:.1f} dB < requant band {band_h:.1f} - 6 " + "(a whole-matrix-vs-per-group scale layout bug lands at 2-5 dB)" + ) + + # --- y: refA from the op's own quantized h + refB independent chain + w2_deq = c["w2_deq"] + y_ref = _grouped_matmul(h_deq, w2_deq, sizes, transpose_b=True) + y_ref2 = _grouped_matmul(h_deq, w2_deq, sizes, transpose_b=True, chunks=4) + y_gate = _refA_gate(y_ref, y_ref2) + y_db = compute_error(y_ref.bfloat16(), r["y"]).item() + assert y_db >= y_gate, f"y {y_db:.1f} dB < derived refA gate {y_gate:.1f}" + # refB for y: the whole forward computed from ORIGINAL bf16 tensors only. + gate_x, up_x = _zsplit(z_exact, hidden) + y_exact = _grouped_matmul( + F.silu(gate_x) * up_x, [w.float() for w in c["w2"]], sizes, transpose_b=True + ) + y_band_b = compute_error(y_exact, y_ref).item() + y_db_b = compute_error(y_exact.bfloat16(), r["y"]).item() + assert y_db_b >= y_band_b - 6.0, ( + f"y vs independent chain {y_db_b:.1f} dB < band {y_band_b:.1f} - 6" + ) + + # --- dz vs closed-form dSwiGLU from the kernel's z + dy_deq = _dequant_rowwise(c["dy_q"], c["dy_sf"]) + w2c_deq = [ + _dequant_colwise_grouped(c["w2c_q"][g], c["w2c_sf"].view(G, -1)[g], [D], hidden) + for g in range(G) + ] + dh_ref = _grouped_matmul(dy_deq, w2c_deq, sizes, transpose_b=False) + dgate, dup = _dswiglu(dh_ref, gate_f, up_f) + dz_ref = torch.empty(R, 2 * hidden, dtype=torch.float32, device="cuda") + v = dz_ref.view(R, hidden // _BLOCK, 2, _BLOCK) + v[:, :, 0, :] = dgate.view(R, hidden // _BLOCK, _BLOCK) + v[:, :, 1, :] = dup.view(R, hidden // _BLOCK, _BLOCK) + band_dz = compute_error( + dz_ref, _dequant_rowwise(*_quant_rowwise(dz_ref.bfloat16())) + ).item() + dz_deq = _dequant_rowwise(r["dz_q"], r["dz_sf"]) + dz_db = compute_error(dz_ref, dz_deq).item() + assert dz_db >= band_dz - 6.0, f"dz {dz_db:.1f} dB < band {band_dz:.1f} - 6" + + # --- dx refA + w13c_deq = [ + _dequant_colwise_grouped( + c["w13c_q"][g], c["w13c_sf"].view(G, -1)[g], [2 * hidden], D + ) + for g in range(G) + ] + dx_ref = _grouped_matmul(dz_deq, w13c_deq, sizes, transpose_b=False) + dx_ref2 = _grouped_matmul(dz_deq, w13c_deq, sizes, transpose_b=False, chunks=4) + dx_gate = _refA_gate(dx_ref, dx_ref2) + dx_db = compute_error(dx_ref.bfloat16(), r["dx"]).item() + assert dx_db >= dx_gate, f"dx {dx_db:.1f} dB < derived refA gate {dx_gate:.1f}" + + # --- wgrads refA (production layout mixes) + dy_col_deq = _dequant_colwise_grouped(c["dy_colq"], c["dy_col_sf"], sizes, D) + dz_col_deq = _dequant_colwise_grouped( + r["dz_colq"], r["dz_col_sf"], sizes, 2 * hidden + ) + x_col_deq = _dequant_colwise_grouped(c["x_colq"], c["x_col_sf"], sizes, D) + off = 0 + dw2_ref = torch.zeros(G, D, hidden, dtype=torch.float32, device="cuda") + dw13_ref = torch.zeros(G, 2 * hidden, D, dtype=torch.float32, device="cuda") + for g, m in enumerate(sizes): + dw2_ref[g] = dy_col_deq[off : off + m].t() @ h_col_deq[off : off + m] + dw13_ref[g] = dz_col_deq[off : off + m].t() @ x_col_deq[off : off + m] + off += m + dw2_db = compute_error(dw2_ref.bfloat16(), r["dw2"]).item() + dw13_db = compute_error(dw13_ref.bfloat16(), r["dw13"]).item() + assert dw2_db >= 50.0, f"dw2 {dw2_db:.1f} dB < 50 (probe level: 98-155)" + assert dw13_db >= 50.0, f"dw13 {dw13_db:.1f} dB < 50 (probe level: 91-160)" + + # zero-token experts must come back written as exact zeros + for g, m in enumerate(sizes): + if m == 0: + assert (r["dw2"][g] == 0).all() and (r["dw13"][g] == 0).all(), ( + f"zero-token expert {g} weight gradients must be exactly zero" + ) + + +# --------------------------------------------------------------------------- +# Wgrad stride matrix: both operands in each major, all four combinations. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("a_native", [False, True], ids=["aRM", "aNat"]) +@pytest.mark.parametrize("b_native", [False, True], ids=["bRM", "bNat"]) +def test_wgrad_stride_matrix(dbg, a_native, b_native): + c = dbg + sizes, D = c["sizes"], c["D"] + dy_q, dy_sf = _quant_colwise_grouped(c["dy"], sizes, native=a_native) + x_q, x_sf = _quant_colwise_grouped(c["x"], sizes, native=b_native) + dw = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf, x_q, x_sf, c["offsets"]) + dy_deq = _dequant_colwise_grouped(dy_q, dy_sf, sizes, D) + x_deq = _dequant_colwise_grouped(x_q, x_sf, sizes, D) + ref = torch.zeros(c["G"], D, D, dtype=torch.float32, device="cuda") + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_deq[off : off + m].t() @ x_deq[off : off + m] + off += m + db = compute_error(ref.bfloat16(), dw).item() + assert db >= 50.0, f"wgrad[{a_native=} {b_native=}] {db:.1f} dB < 50" + + +def test_native_weight_major_mm_bwd(dbg): + """Ops 2 and 3 accept the production dim1-native (memory-transposed) + colwise weight major. Both majors carry identical logical values, so each + native arm must agree with the rowmajor arm far above any + reduction-order band.""" + c = dbg + r = _run_chain(c) + # Op 3: dim1-native w2 colwise major. + w2c_nat, w2c_nat_sf = _quant_weight_colwise(c["w2"], native=True) + assert not w2c_nat.is_contiguous() + assert torch.equal(_bytes(w2c_nat), _bytes(c["w2c_q"])) + dz_q, dz_sf, _, _ = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + c["dy_q"], c["dy_sf"], w2c_nat, w2c_nat_sf, r["z"], c["offsets"] + ) + db = compute_error( + _dequant_rowwise(r["dz_q"], r["dz_sf"]), _dequant_rowwise(dz_q, dz_sf) + ).item() + assert db >= 50.0, f"native-major w2 dz vs rowmajor arm: {db:.1f} dB < 50" + # Op 2 (FC1 dgrad): dim1-native w13 colwise major, transposed into + # [G, N=D, K=2F] exactly like the production call. + w13c_nat, w13c_nat_sf = _quant_weight_colwise(c["w13"], native=True) + dx_nat = _OPS.mxfp8_grouped_gemm( + r["dz_q"], r["dz_sf"], w13c_nat.transpose(-2, -1), w13c_nat_sf, c["offsets"] + ) + db = compute_error(r["dx"].float(), dx_nat.float()).item() + assert db >= 50.0, f"native-major w13 dx vs rowmajor arm: {db:.1f} dB < 50" + + +# --------------------------------------------------------------------------- +# A < R strict tail with planted garbage. +# --------------------------------------------------------------------------- + + +def test_tail_a_lt_r_poisoned(): + D = hidden = 256 + sizes = [256, 0, 512, 256] + A, R = sum(sizes), 1280 + torch.manual_seed(3) + dev = "cuda" + offsets = _mk_offsets(sizes, dev) + x = torch.randn(R, D, dtype=torch.bfloat16, device=dev) * 0.5 + dy = torch.randn(R, D, dtype=torch.bfloat16, device=dev) * 0.5 + x[A:] = float("nan") + dy[A::2] = float("inf") + dy[A + 1 :: 2] = float("nan") + w13 = _to_32block( + torch.randn(4, hidden, 2, D, dtype=torch.bfloat16, device=dev) * 0.02 + ) + w2 = torch.randn(4, D, hidden, dtype=torch.bfloat16, device=dev) * 0.02 + + x_q, x_sf = _quant_rowwise(x) + dy_q, dy_sf = _quant_rowwise(dy) + w13_q, w13_sf = _quant_weight_rowwise(w13) + w2_q, w2_sf = _quant_weight_rowwise(w2) + w2c_q, w2c_sf = _quant_weight_colwise(w2) + + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + assert not z[:A].isnan().any(), "active z rows contaminated by the poisoned tail" + y = _OPS.mxfp8_grouped_gemm(h_q, h_sf, w2_q, w2_sf, offsets) + assert not y[:A].isnan().any(), "active y rows contaminated" + dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + dy_q, dy_sf, w2c_q, w2c_sf, z, offsets + ) + w13c_q, w13c_sf = _quant_weight_colwise(w13) + dx = _OPS.mxfp8_grouped_gemm( + dz_q, dz_sf, w13c_q.transpose(-2, -1), w13c_sf, offsets + ) + assert not dx[:A].isnan().any(), "active dx rows contaminated" + + # wgrad: colwise scales cover only the routed A rows; the qdata tail is + # additionally poisoned with NaN bytes and must never be read. + dy_colq, dy_col_sf = _quant_colwise_grouped(dy[:A], sizes, native=True) + dy_colq_full = _cat8( + [ + dy_colq.contiguous(), + torch.full((R - A, D), 0x7F, dtype=torch.uint8, device=dev).view(_E4M3), + ], + 0, + ) + dw2 = _OPS.mxfp8_grouped_gemm_wgrad( + dy_colq_full, dy_col_sf, h_colq, h_col_sf, offsets + ) + assert not dw2.isnan().any(), "wgrad read the NaN-poisoned inactive tail" + dy_col_deq = _dequant_colwise_grouped(dy_colq, dy_col_sf, sizes, D) + h_col_deq = _dequant_colwise_grouped(h_colq[:A], h_col_sf, sizes, hidden) + ref = torch.zeros(4, D, hidden, dtype=torch.float32, device=dev) + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_col_deq[off : off + m].t() @ h_col_deq[off : off + m] + off += m + db = compute_error(ref.bfloat16(), dw2).item() + assert db >= 50.0, f"tail-poisoned dw2 {db:.1f} dB < 50" + + +# --------------------------------------------------------------------------- +# Determinism, compile, R == 0. +# --------------------------------------------------------------------------- + + +def test_determinism_all_ops_bitwise(dbg): + c = dbg + r1 = _run_chain(c) + r2 = _run_chain(c) + for key in r1: + assert torch.equal(_bytes(r1[key]), _bytes(r2[key])), ( + f"{key} is not bitwise deterministic across identical launches" + ) + + +def test_compile_fullgraph_bitwise(dbg): + c = dbg + + def fwd_then_mm(x_q, x_sf, w13_q, w13_sf, w2_q, w2_sf, offsets): + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + y = _OPS.mxfp8_grouped_gemm(h_q, h_sf, w2_q, w2_sf, offsets) + return z, h_q, y + + eager = fwd_then_mm( + c["x_q"], + c["x_sf"], + c["w13_q"], + c["w13_sf"], + c["w2_q"], + c["w2_sf"], + c["offsets"], + ) + compiled = torch.compile(fwd_then_mm, fullgraph=True)( + c["x_q"], + c["x_sf"], + c["w13_q"], + c["w13_sf"], + c["w2_q"], + c["w2_sf"], + c["offsets"], + ) + for e, co, name in zip(eager, compiled, ("z", "h_q", "y")): + assert torch.equal(_bytes(e), _bytes(co)), f"compiled {name} != eager" + + +def test_r0_all_ops(): + dev = "cuda" + D = hidden = 256 + offsets = torch.zeros(2, dtype=torch.int32, device=dev) + z, h_q, h_sf, h_colq, h_col_sf = _OPS.mxfp8_grouped_gemm_swiglu_fwd( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, 2 * hidden, D, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * 2 * hidden * D // _BLOCK, dtype=_E8M0, device=dev), + offsets, + ) + assert z.shape == (0, 2 * hidden) and h_q.shape == (0, hidden) + assert h_sf.numel() == 0 and h_col_sf.numel() == 0 + y = _OPS.mxfp8_grouped_gemm( + torch.empty(0, hidden, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, D, hidden, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * D * hidden // _BLOCK, dtype=_E8M0, device=dev), + offsets, + ) + assert y.shape == (0, D) and y.dtype == torch.bfloat16 + dz_q, dz_sf, dz_colq, dz_col_sf = _OPS.mxfp8_grouped_gemm_dswiglu_bwd( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.zeros(2, D, hidden, dtype=torch.uint8, device=dev).view(_E4M3), + torch.empty(2 * hidden * D // _BLOCK, dtype=_E8M0, device=dev), + torch.empty(0, 2 * hidden, dtype=torch.bfloat16, device=dev), + offsets, + ) + assert dz_q.shape == (0, 2 * hidden) and dz_sf.numel() == 0 + assert dz_colq.shape == (0, 2 * hidden) and dz_col_sf.numel() == 0 + dw = _OPS.mxfp8_grouped_gemm_wgrad( + torch.empty(0, D, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + torch.empty(0, hidden, dtype=_E4M3, device=dev), + torch.empty(0, dtype=_E8M0, device=dev), + offsets, + ) + assert dw.shape == (2, D, hidden) and (dw == 0).all() + + +# --------------------------------------------------------------------------- +# Negative controls: each sabotage must fail the numerics gates decisively. +# --------------------------------------------------------------------------- + + +def test_negative_control_whole_matrix_colwise_scales(dbg): + """Whole-matrix to_blocked colwise scales: same bytes, silently wrong order.""" + c = dbg + sizes, D, G = c["sizes"], c["D"], c["G"] + dy_q, dy_sf_pg = _quant_colwise_grouped(c["dy"], sizes, native=False) + x_q, x_sf_pg = _quant_colwise_grouped(c["x"], sizes, native=False) + # Rebuild the SAME logical scales in whole-matrix blocked order. + s_t, _ = to_mx(c["dy"].t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + dy_sf_wm = to_blocked(s_t.view(_E8M0)).view(_E8M0) + assert dy_sf_wm.numel() == dy_sf_pg.numel() + good = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf_pg, x_q, x_sf_pg, c["offsets"]) + bad = _OPS.mxfp8_grouped_gemm_wgrad(dy_q, dy_sf_wm, x_q, x_sf_pg, c["offsets"]) + dy_deq = _dequant_colwise_grouped(dy_q, dy_sf_pg, sizes, D) + x_deq = _dequant_colwise_grouped(x_q, x_sf_pg, sizes, D) + ref = torch.zeros(G, D, D, dtype=torch.float32, device="cuda") + off = 0 + for g, m in enumerate(sizes): + ref[g] = dy_deq[off : off + m].t() @ x_deq[off : off + m] + off += m + good_db = compute_error(ref.bfloat16(), good).item() + bad_db = compute_error(ref.bfloat16(), bad).item() + assert good_db >= 50.0 + assert bad_db < 25.0, ( + f"whole-matrix colwise scales scored {bad_db:.1f} dB -- the negative " + f"control lost its teeth (good arm: {good_db:.1f})" + ) + + +def test_negative_control_gate_up_swap(dbg): + """Swapping the gate/up 32-blocks must collapse h against the correct ref.""" + c = dbg + hidden, G, D = c["F"], c["G"], c["D"] + w13_sw = ( + c["w13"] + .view(G, hidden // _BLOCK, 2, _BLOCK, D) + .flip(2) + .reshape(G, 2 * hidden, D) + .contiguous() + ) + w13_sw_q, w13_sw_sf = _quant_weight_rowwise(w13_sw) + _, h_q, h_sf, _, _ = _OPS.mxfp8_grouped_gemm_swiglu_fwd( + c["x_q"], c["x_sf"], w13_sw_q, w13_sw_sf, c["offsets"] + ) + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) + gate_f, up_f = _zsplit(z_ref, hidden) + h_ref = F.silu(gate_f) * up_f + good_db = compute_error( + h_ref, + _dequant_rowwise( + *( + _OPS.mxfp8_grouped_gemm_swiglu_fwd( + c["x_q"], c["x_sf"], c["w13_q"], c["w13_sf"], c["offsets"] + )[1:3] + ) + ), + ).item() + bad_db = compute_error(h_ref, _dequant_rowwise(h_q, h_sf)).item() + assert bad_db < good_db - 10.0, ( + f"gate/up swap only moved h from {good_db:.1f} to {bad_db:.1f} dB -- " + "the 32-block order convention is not actually being exercised" + ) + + +def test_negative_control_scale_byte_flip(dbg): + """One +2-code E8M0 flip (x4) in the weight scales must break refA.""" + c = dbg + sf_bad = c["w13_sf"].view(torch.uint8).clone() + sf_bad[sf_bad.numel() // 2] += 2 + z_bad = _OPS.mxfp8_grouped_gemm_swiglu_fwd( + c["x_q"], c["x_sf"], c["w13_q"], sf_bad, c["offsets"] + )[0] + z_ref = _grouped_matmul(c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True) + z_ref2 = _grouped_matmul( + c["x_deq"], c["w13_deq"], c["sizes"], transpose_b=True, chunks=4 + ) + gate = _refA_gate(z_ref, z_ref2) + bad_db = compute_error(z_ref.bfloat16(), z_bad).item() + assert bad_db < gate, ( + f"single scale-byte flip still passes refA ({bad_db:.1f} >= {gate:.1f} dB)" + ) + + +def test_kernel_scale_mode_is_rceil(dbg): + """The fwd op's h scale bytes must match RCEIL, and not FLOOR, quantization.""" + c = dbg + r = _run_chain(c) + gate_f, up_f = _zsplit(r["z"].float(), c["F"]) + h_ref = (F.silu(gate_f) * up_f).bfloat16() + s_rceil, _ = to_mx(h_ref, _E4M3, _BLOCK, scaling_mode=_RCEIL) + s_floor, _ = to_mx(h_ref, _E4M3, _BLOCK, scaling_mode=ScaleCalculationMode.FLOOR) + got = from_blocked(r["h_sf"].view(_E8M0), c["R"], c["F"] // _BLOCK).view( + torch.uint8 + ) + rceil_frac = (got == s_rceil.view(torch.uint8)).float().mean().item() + floor_frac = (got == s_floor.view(torch.uint8)).float().mean().item() + assert rceil_frac > 0.98, f"h scales match RCEIL on only {rceil_frac:.3f}" + assert rceil_frac > floor_frac + 0.1, ( + f"RCEIL ({rceil_frac:.3f}) does not dominate FLOOR ({floor_frac:.3f})" + ) + + +# --------------------------------------------------------------------------- +# Validation: rejection matrix and the opt-in offsets path. +# --------------------------------------------------------------------------- + + +def _valid_fwd_args(device="cuda", R=512, D=256, hidden=256, G=2): + N1 = 2 * hidden + return dict( + x_q=torch.zeros(R, D, dtype=_E4M3, device=device), + x_sf=torch.zeros(R * D // _BLOCK, dtype=_E8M0, device=device), + w13_q=torch.zeros(G, N1, D, dtype=_E4M3, device=device), + w13_sf=torch.zeros(G * N1 * D // _BLOCK, dtype=_E8M0, device=device), + offsets=torch.tensor([256, 512], dtype=torch.int32, device=device), + ) + + +_NEGATIVES = [ + ("x_q_dtype", lambda a: a.update(x_q=a["x_q"].view(torch.int8)), "float8_e4m3fn"), + ("x_q_cpu", lambda a: a.update(x_q=a["x_q"].cpu()), "CUDA"), + ( + "r_not_256", + lambda a: a.update( + x_q=torch.zeros(384, 256, dtype=_E4M3, device="cuda"), + x_sf=torch.zeros(384 * 8, dtype=_E8M0, device="cuda"), + ), + "multiple of 256", + ), + ( + "d_192", + lambda a: a.update( + x_q=torch.zeros(512, 192, dtype=_E4M3, device="cuda"), + x_sf=torch.zeros(512 * 6, dtype=_E8M0, device="cuda"), + w13_q=torch.zeros(2, 512, 192, dtype=_E4M3, device="cuda"), + w13_sf=torch.zeros(2 * 512 * 6, dtype=_E8M0, device="cuda"), + ), + "multiple of 128", + ), + ( + "offsets_i64", + lambda a: a.update(offsets=a["offsets"].to(torch.int64)), + "int32", + ), + ( + "offsets_wrong_len", + lambda a: a.update(offsets=a["offsets"][:1]), + "one entry per local expert", + ), + ( + "g0", + lambda a: a.update( + w13_q=torch.zeros(0, 512, 256, dtype=_E4M3, device="cuda"), + w13_sf=torch.zeros(0, dtype=_E8M0, device="cuda"), + offsets=torch.zeros(0, dtype=torch.int32, device="cuda"), + ), + "at least one expert group", + ), + ( + "x_sf_short", + lambda a: a.update(x_sf=a["x_sf"][:-8].clone()), + "blocked scale bytes", + ), + ( + "w13_stride", + lambda a: a.update( + w13_q=a["w13_q"].transpose(-2, -1).contiguous().transpose(-2, -1) + ), + "stride", + ), +] + + +@pytest.mark.parametrize("case", _NEGATIVES, ids=[c[0] for c in _NEGATIVES]) +def test_validation_negatives(case): + _name, mutate, needle = case + args = _valid_fwd_args() + mutate(args) + with pytest.raises(ValueError) as exc_info: + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**args) + assert needle.lower() in str(exc_info.value).lower(), ( + f"rejection message {str(exc_info.value)!r} does not name the defect " + f"({needle!r})" + ) + + +def test_optin_offsets_validation(monkeypatch): + """Offset VALUES are checked only under TORCHAO_MXFP8_VALIDATE_OFFSETS=1. + + The default-build non-rejection of a 128-row group is asserted on the + validator directly: launching a kernel with misaligned offsets is the + exact out-of-contract config the module documents as corrupting silently + and nondeterministically, so this test must never perform that launch. + The opt-in rejections DO go through the ops, which raise before any + launch. + """ + args = _valid_fwd_args() + bad_offsets = torch.tensor([128, 512], dtype=torch.int32, device="cuda") + + # Default build: metadata-only, misaligned VALUES are not (and cannot be) + # caught without a D2H sync. + monkeypatch.delenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", raising=False) + validate_group_offsets( + bad_offsets, num_groups=2, allocated_rows=512, device=bad_offsets.device + ) + + monkeypatch.setenv("TORCHAO_MXFP8_VALIDATE_OFFSETS", "1") + bad = dict(args, offsets=bad_offsets) + with pytest.raises(ValueError, match="FIX_PAD_SIZE"): + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**bad) + + dec = dict(args, offsets=torch.tensor([512, 256], dtype=torch.int32, device="cuda")) + with pytest.raises(ValueError, match="nondecreasing"): + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**dec) + + over = dict( + args, offsets=torch.tensor([256, 768], dtype=torch.int32, device="cuda") + ) + with pytest.raises(ValueError, match="exceeds the allocated row count"): + _OPS.mxfp8_grouped_gemm_swiglu_fwd(**over) + + # The opt-in check must not break fake tracing (no values to read). + with FakeTensorMode(): + _fake_chain_shapes() diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py new file mode 100644 index 0000000000..1c9abd34e5 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py @@ -0,0 +1,1254 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +"""MXFP8 routed-expert grouped-MLP ops over the cuDNN-frontend CuTe DSL kernels. + +Four custom ops, each one launch of a ``cudnn.grouped_gemm_*_wrapper_sm100`` +kernel from the standalone cudnn-frontend python package (>= 1.27, Blackwell +SM 10.0 exactly -- the wrappers are sm100-specific; no TransformerEngine +dependency); the matching public wrappers live at the bottom of this module: + +* :func:`mxfp8_grouped_gemm_swiglu_fwd` -- FC1 ragged grouped GEMM + SwiGLU + + rowwise 1x32 AND columnwise 32x1 MXFP8 RCEIL quantization + BF16 pre-GLU. +* :func:`mxfp8_grouped_gemm` -- ragged grouped GEMM on + prequantized operands to BF16 (FC2 forward and FC1 dgrad). +* :func:`mxfp8_grouped_gemm_dswiglu_bwd` -- FC2 dgrad + dSwiGLU + dual MXFP8 + quantization of the FC1 gradient. +* :func:`mxfp8_grouped_gemm_wgrad` -- ragged-reduction grouped weight + gradient (dense output mode; called once for FC1 and once for FC2). + +CONTRACT: every per-expert row count and the allocated row count must be +multiples of **256** -- the cuDNN FE kernels hard-code ``FIX_PAD_SIZE = 256``, +and groups that are only 128-row aligned corrupt results SILENTLY and +NONDETERMINISTICALLY (the corruption locus migrates between identical-input +reruns; no smoke test can prove a misaligned config safe). Use a token +dispatcher with ``pad_multiple=256``. Enforcement is two-tier: metadata-only +checks always run (memoized per signature, FakeTensor-safe, back +``register_fake`` so torch.compile rejects at capture time); the offset +VALUES (nondecreasing, per-expert %256, ``offsets[-1] <= R``) are checked +only under ``TORCHAO_MXFP8_VALIDATE_OFFSETS=1`` because reading them forces a +D2H sync. Checks raise ValueError, never assert, so ``python -O`` cannot +strip them. + +All scale arguments are FLAT blocked E8M0 buffers (uint8 or float8_e8m0fnu); +the ops build the kernel-native 6-D / 2-D views internally with probe-proven +recipes. The FC1 weight is E4M3 ``[G, 2F, D]`` with rows in the cuDNN +32-block GLU order ``[gate0(32) | up0(32) | gate1(32) | ...]`` (gate = the +SiLU'd operand). ``offsets`` is int32 CUDA ``[G]`` exclusive-end rows. Rows +in ``[offsets[-1], R)``: caller-allocated outputs (the grouped-mm result and +the weight gradients) keep their tails untouched, while kernel-allocated +outputs (z, h, dz and their scales) carry garbage tails that are +read-forbidden -- both behaviors probe-verified with NaN-poisoned tails. + +Importing this module registers the four ``torchao::`` custom ops; the +``cudnn`` package itself is imported lazily inside the op bodies at first +real launch. :func:`is_supported` is the static shape predicate to call +before selecting this family. +""" + +import importlib.util +import os +from typing import Optional, Tuple + +import torch + +__all__ = [ + "DIM_ALIGNMENT", + "ROW_GROUP_ALIGNMENT", + "SCALE_BLOCK_SIZE", + "is_supported", + "mxfp8_grouped_gemm", + "mxfp8_grouped_gemm_dswiglu_bwd", + "mxfp8_grouped_gemm_swiglu_fwd", + "mxfp8_grouped_gemm_wgrad", +] + +# MXFP8 scaling block: 32 values share one E8M0 scale. +SCALE_BLOCK_SIZE = 32 +# tcgen05 blocked scale tile: 128 rows x 4 columns, 512 bytes. +SCALE_TILE_ROWS = 128 +SCALE_TILE_COLS = 4 +# Feature-dimension granularity (D and F). +DIM_ALIGNMENT = 128 +# Row-count granularity: per-expert groups AND the allocated row count (the +# cuDNN FE kernels' FIX_PAD_SIZE). +ROW_GROUP_ALIGNMENT = 256 +# Byte alignment for TMA/vectorized accesses. +_PTR_ALIGNMENT = 16 + +_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) + +_E4M3 = torch.float8_e4m3fn +_E8M0 = torch.float8_e8m0fnu +_BLOCK = SCALE_BLOCK_SIZE + + +# -------------------------------------------------------------------------- +# Availability probe and the static shape predicate. +# -------------------------------------------------------------------------- + +_REQUIRED_WRAPPERS = ( + "grouped_gemm_glu_wrapper_sm100", + "grouped_gemm_quant_wrapper_sm100", + "grouped_gemm_dglu_wrapper_sm100", + "grouped_gemm_wgrad_wrapper_sm100", +) +# 1.27 is required: earlier frontends reject prob_tensor=None. +_MIN_FE_VERSION = (1, 27) + + +def _fe_version_tuple(version: str) -> tuple: + """Numeric prefix as a tuple ('1.27.0' -> (1, 27, 0)); never compare + version STRINGS ('1.100' < '1.27' lexicographically).""" + parts = [] + for piece in version.split("."): + digits = "" + for ch in piece: + if not ch.isdigit(): + break + digits += ch + if not digits: + break + parts.append(int(digits)) + return tuple(parts) + + +def _is_sm100() -> bool: + # Exactly capability (10, 0): the cudnn wrappers are *_sm100-specific and + # unproven on other SM 10.x parts. + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + + +def _probe_cudnn_frontend() -> str: + """Empty string when usable; else the reason it is not.""" + if importlib.util.find_spec("cudnn") is None: + return "the cudnn-frontend python package ('cudnn') is not installed" + try: + import cudnn + except Exception as exc: # pragma: no cover - environment-specific + return f"'import cudnn' failed: {exc!r}" + version = getattr(cudnn, "__version__", "0") + if _fe_version_tuple(version) < _MIN_FE_VERSION: + return ( + f"cudnn-frontend {version} is too old; >= " + f"{'.'.join(map(str, _MIN_FE_VERSION))} is required " + "(prob_tensor=None support)" + ) + missing = [name for name in _REQUIRED_WRAPPERS if not hasattr(cudnn, name)] + if missing: + return "cudnn-frontend lacks required wrappers: " + ", ".join(missing) + return "" + + +_mxfp8_grouped_mlp_unavailable_reason = ( + _probe_cudnn_frontend() + if _is_sm100() + else ( + "requires an SM 10.0 (Blackwell) GPU; the cudnn wrappers are sm100-specific" + if torch.cuda.is_available() + else "CUDA is not available" + ) +) +_mxfp8_grouped_mlp_kernels_available = _mxfp8_grouped_mlp_unavailable_reason == "" + + +def _require_available() -> None: + if not _mxfp8_grouped_mlp_kernels_available: + raise NotImplementedError( + "cuDNN-frontend MXFP8 grouped-MLP kernels are unavailable: " + + _mxfp8_grouped_mlp_unavailable_reason + ) + + +def is_supported(model_dim: int, hidden_dim: int) -> bool: + """True when D and F are positive multiples of 128. Integration code must + ALSO guarantee the runtime row contract (per-expert groups and the row + allocation padded to multiples of 256): row counts live in device memory + and are not checkable here. Environment availability is a separate + concern (``_mxfp8_grouped_mlp_kernels_available``).""" + return ( + model_dim > 0 + and hidden_dim > 0 + and model_dim % DIM_ALIGNMENT == 0 + and hidden_dim % DIM_ALIGNMENT == 0 + ) + + +# -------------------------------------------------------------------------- +# Metadata validation helpers (see the module docstring for the two tiers). +# -------------------------------------------------------------------------- + + +def _round_up(x: int, to: int) -> int: + return ((x + to - 1) // to) * to + + +def blocked_scale_numel(rows: int, cols: int) -> int: + """Blocked-buffer element count for a logical [rows, cols] scale matrix + (``cols`` counts scale values: the reduced dimension divided by 32).""" + return _round_up(rows, SCALE_TILE_ROWS) * _round_up(cols, SCALE_TILE_COLS) + + +def host_offsets_validation_enabled() -> bool: + """Opt-in offset-VALUES validation; off by default (forces a D2H sync).""" + return os.environ.get("TORCHAO_MXFP8_VALIDATE_OFFSETS", "0") == "1" + + +def _is_fake(tensor: torch.Tensor) -> bool: + """True for meta/fake tensors (no usable data pointer or values).""" + if tensor.device.type == "meta": + return True + try: + from torch._subclasses.fake_tensor import FakeTensor + except ImportError: + return False + return isinstance(tensor, FakeTensor) + + +def validate_group_offsets( + offsets: torch.Tensor, + *, + num_groups: int, + allocated_rows: int, + device: Optional[torch.device] = None, + name: str = "offsets", +) -> None: + """Metadata always; VALUES only when opted in and the tensor is real.""" + if not isinstance(offsets, torch.Tensor): + raise ValueError(f"{name} must be a torch.Tensor, got {type(offsets)}") + if num_groups < 1: + raise ValueError( + f"{name} must describe at least one expert group, got G={num_groups}" + ) + if offsets.dtype != torch.int32: + raise ValueError(f"{name} must be int32, got {offsets.dtype}") + if not offsets.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device {offsets.device}") + if device is not None and offsets.device != device: + raise ValueError( + f"{name} must be on {device}, got {offsets.device}; all operands and " + "destinations must share one CUDA device" + ) + if offsets.ndim != 1: + raise ValueError(f"{name} must be 1D, got shape {tuple(offsets.shape)}") + if offsets.numel() != num_groups: + raise ValueError( + f"{name} must have one entry per local expert: expected {num_groups}, " + f"got {offsets.numel()}" + ) + if not offsets.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {offsets.stride()}") + + if not host_offsets_validation_enabled() or _is_fake(offsets): + return + + values = offsets.tolist() # d2h sync; opt-in debugging path only + previous = 0 + for group, end in enumerate(values): + if end < previous: + raise ValueError( + f"{name} must be nondecreasing, but entry {group} is {end} " + f"after {previous}" + ) + size = end - previous + if size % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"per-expert row counts must be multiples of {ROW_GROUP_ALIGNMENT} " + f"(cuDNN FE FIX_PAD_SIZE; sub-256 groups corrupt results " + f"nondeterministically): expert {group} has {size} rows " + f"(offsets {previous} -> {end})" + ) + previous = end + if previous > allocated_rows: + raise ValueError( + f"{name}[-1] ({previous}) exceeds the allocated row count " + f"({allocated_rows})" + ) + + +def _check_pointer_alignment(tensor: torch.Tensor, *, name: str) -> None: + """16-byte data_ptr gate (TMA/vectorized accesses); fakes have no pointer.""" + if _is_fake(tensor): + return + if tensor.data_ptr() % _PTR_ALIGNMENT != 0: + raise ValueError( + f"{name} must be {_PTR_ALIGNMENT}-byte aligned, but its data " + f"pointer is {tensor.data_ptr() % _PTR_ALIGNMENT} bytes past an " + "aligned address. A contiguous view with a nonzero storage " + "offset can violate this." + ) + + +def validate_operand( + tensor: torch.Tensor, + *, + name: str, + shape: tuple, + dtype: torch.dtype, + device: torch.device, + stride: Optional[tuple] = None, + check_pointer_alignment: bool = True, +) -> None: + """dtype/shape/device, optional EXACT stride (None = any: the wrappers + consume both majors, every composite combination probe-proven), pointer + alignment. Metadata gates run before the ``data_ptr()`` gate so + FakeTensor tracing exercises the same checks.""" + if tensor.dtype != dtype: + raise ValueError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != tuple(shape): + raise ValueError( + f"{name} must have shape {tuple(shape)}, got {tuple(tensor.shape)}" + ) + if stride is not None and tuple(tensor.stride()) != tuple(stride): + raise ValueError( + f"{name} must have stride {tuple(stride)}, got {tuple(tensor.stride())}. " + "This layout is part of the ABI; a values-equal tensor with a " + "different stride is not interchangeable." + ) + if tensor.device != device: + raise ValueError( + f"{name} must be on {device}, got {tensor.device}; all operands and " + "destinations must share one CUDA device" + ) + if check_pointer_alignment: + _check_pointer_alignment(tensor, name=name) + + +def validate_blocked_scales( + scales: torch.Tensor, + *, + name: str, + logical_rows: int, + logical_cols: int, + device: torch.device, + groups: int = 1, +) -> None: + """Flat blocked E8M0 buffer with a statically known size; ``groups > 1`` + means per-expert blocks concatenated.""" + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + expected = groups * blocked_scale_numel(logical_rows, logical_cols) + if scales.numel() != expected: + raise ValueError( + f"{name} must hold {expected} blocked scale bytes for a logical " + f"[{logical_rows}, {logical_cols}] scale matrix" + + (f" across {groups} experts" if groups > 1 else "") + + f", got {scales.numel()}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + _check_pointer_alignment(scales, name=name) + + +def validate_ragged_colwise_scales( + scales: torch.Tensor, + *, + name: str, + features: int, + allocated_rows: int, + device: torch.device, +) -> None: + """Per-group columnwise scale buffer sized by ``offsets[-1]`` -- a device + value -- so only dtype/device/contiguity, granule divisibility, and the + allocated-rows maximum are host-checkable (an ``offsets[-1] < R`` buffer + legitimately covers fewer scale columns, probe-verified).""" + if scales.dtype not in _SCALE_DTYPES: + raise ValueError( + f"{name} must be uint8 or float8_e8m0fnu (raw E8M0 bytes), " + f"got {scales.dtype}" + ) + if not scales.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {scales.stride()}") + if scales.device != device: + raise ValueError(f"{name} must be on {device}, got {scales.device}") + rows_pad = _round_up(features, SCALE_TILE_ROWS) + # Each 256-row group contributes features_pad * (group_rows/32) bytes and + # group_rows/32 is a multiple of 8. + granule = rows_pad * (ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE) + if scales.numel() % granule != 0: + raise ValueError( + f"{name} numel {scales.numel()} is not a multiple of {granule} " + f"(= round_up({features},128) x {ROW_GROUP_ALIGNMENT // SCALE_BLOCK_SIZE} " + "scale columns per 256-row group)" + ) + max_numel = rows_pad * (allocated_rows // SCALE_BLOCK_SIZE) + if scales.numel() > max_numel: + raise ValueError( + f"{name} numel {scales.numel()} exceeds the maximum {max_numel} implied " + f"by the allocated row count {allocated_rows}" + ) + _check_pointer_alignment(scales, name=name) + + +def validate_feature_dims( + *, + model_dim: int, + hidden_dim: int, + model_dim_name: str = "model dimension D", + hidden_dim_name: str = "routed-expert hidden dimension F", +) -> None: + """The name arguments let mm/wgrad call sites report their generic N/K + dims instead of the fwd/bwd ops' D/F.""" + if model_dim <= 0 or model_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"{model_dim_name} must be a positive multiple of {DIM_ALIGNMENT}, " + f"got {model_dim}" + ) + if hidden_dim <= 0 or hidden_dim % DIM_ALIGNMENT != 0: + raise ValueError( + f"{hidden_dim_name} must be a positive multiple of " + f"{DIM_ALIGNMENT}, got {hidden_dim}" + ) + + +def validate_allocated_rows(rows: int, *, name: str = "R") -> None: + """%256 (may be zero): the allocation must be reachable by a legal offsets + vector plus an inactive tail, and a non-256 allocation also breaks the + whole-matrix == per-group-concat identity of the rowwise blocked scales.""" + if rows % ROW_GROUP_ALIGNMENT != 0: + raise ValueError( + f"{name} must be a multiple of {ROW_GROUP_ALIGNMENT}, got {rows}" + ) + + +# Small per-(groups, dtype, device) caches for the kernels' alpha/beta and +# norm-const tensors, each stored with the event recorded after its fill: +# the fill runs on the FIRST caller's stream, so a cache hit on any other +# stream must order after it (the buffer is immutable once filled, so one +# event covers every later consumer). Never cached: the CUDA stream itself +# (looked up per call). +_ones_cache: dict = {} + + +def _cached_ones(numel: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + key = (numel, dtype, device) + hit = _ones_cache.get(key) + if hit is None: + out = torch.ones(numel, dtype=dtype, device=device) + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(device)) + _ones_cache[key] = (out, event) + return out + out, event = hit + event.wait(torch.cuda.current_stream(device)) + return out + + +# The always-on validation tier is metadata-only, so its verdict is a pure +# function of the operands' metadata (the pointer-alignment gate is covered +# by storage_offset: torch's CUDA caching allocator hands out aligned storage +# bases). A training step calls each op hundreds of times with identical +# metadata; the full battery runs once per distinct signature and repeats +# skip straight to the derived dims. Signatures are recorded only AFTER a +# REAL-tensor pass (a rejected call never poisons the cache; a fake pass has +# no data pointer to prove alignment). The opt-in offsets-VALUES check +# (TORCHAO_MXFP8_VALIDATE_OFFSETS) reads data, not metadata, so it runs on +# every call while enabled. +_validated_sigs: set = set() +_VALIDATED_SIGS_CAP = 4096 + + +# SymInt ships with every torch new enough to compile these ops; the empty +# tuple keeps the isinstance gate a no-op elsewhere. +_SYMBOLIC_TYPES = (torch.SymInt,) if hasattr(torch, "SymInt") else () + + +def _meta_sig(tag: str, *tensors: torch.Tensor) -> Optional[tuple]: + # torch.Size and stride() are hashable tuples; device/dtype hash directly. + # Symbolic metadata (SymInt dims/strides/offsets under dynamic-shape + # compile) is unhashable, so those calls get no signature and never touch + # the memo; the full battery still runs. + for t in tensors: + for d in (*t.shape, *t.stride(), t.storage_offset()): + if isinstance(d, _SYMBOLIC_TYPES): + return None + return (tag,) + tuple( + (t.shape, t.stride(), t.dtype, t.device, t.storage_offset()) for t in tensors + ) + + +def _remember_sig(sig: Optional[tuple], *tensors: torch.Tensor) -> None: + # Fake passes skip the data_ptr alignment gates, so a fake-recorded + # signature would exempt the first REAL call from them: record only + # real-tensor passes (fakes revalidate every time; metadata is cheap). + if sig is None or any(_is_fake(t) for t in tensors): + return + if len(_validated_sigs) < _VALIDATED_SIGS_CAP: + _validated_sigs.add(sig) + + +def _require_cuda_device(device: torch.device, name: str) -> None: + if device.type != "cuda": + raise ValueError( + f"{name} must be a CUDA tensor, got device {device}; these kernels " + "run only on CUDA SM100 devices" + ) + + +def _as_e8m0(scales: torch.Tensor) -> torch.Tensor: + return scales if scales.dtype == _E8M0 else scales.view(_E8M0) + + +def _act_scale_view(sf_flat: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Flat blocked scales of a logical [rows, cols/32] matrix -> the wrapper's + 6-D activation view (32, 4, rows/128, 4, cols/128, 1).""" + return ( + _as_e8m0(sf_flat) + .view(1, rows // 128, cols // 128, 32, 4, 4) + .permute(3, 4, 1, 5, 2, 0) + ) + + +def _weight_scale_view( + sf_flat: torch.Tensor, groups: int, n: int, k: int +) -> torch.Tensor: + """Per-group-concat flat blocked scales of logical [n, k/32] per expert -> + the wrapper's 6-D weight view (32, 4, n/128, 4, k/128, G).""" + return ( + _as_e8m0(sf_flat) + .view(groups, n // 128, k // 128, 32, 4, 4) + .permute(3, 4, 1, 5, 2, 0) + ) + + +def _flat_scales(sf_6d: torch.Tensor) -> torch.Tensor: + """Kernel-returned 6-D scale view -> the flat blocked buffer (a free view: + the inverse permute restores the allocation's contiguous order).""" + return sf_6d.permute(5, 2, 4, 0, 1, 3).reshape(-1) + + +def _check_normalized( + tensor: torch.Tensor, *, name: str, shape: tuple, dtype: torch.dtype +) -> torch.Tensor: + """Guard against wrapper-output metadata drifting from the fake spec.""" + if tuple(tensor.shape) != tuple(shape) or tensor.dtype != dtype: + raise RuntimeError( + f"cudnn wrapper output {name} has shape {tuple(tensor.shape)} dtype " + f"{tensor.dtype}; expected {tuple(shape)} {dtype}. The installed " + "cudnn-frontend's output contract changed; the registered fake no " + "longer matches eager." + ) + if not tensor.is_contiguous(): + return tensor.contiguous() + return tensor + + +def _stream() -> int: + return torch.cuda.current_stream().cuda_stream + + +# -------------------------------------------------------------------------- +# Op 1: FC1 grouped GEMM + SwiGLU + dual quantization (glu wrapper) +# -------------------------------------------------------------------------- + + +def _fwd_output_specs(rows: int, hidden: int): + two_hidden = 2 * hidden + return ( + ("z_bf16", (rows, two_hidden), torch.bfloat16), + ("h_row_q", (rows, hidden), _E4M3), + ("h_row_sf", (rows * hidden // _BLOCK,), _E8M0), + ("h_col_q", (rows, hidden), _E4M3), + ("h_col_sf", (hidden * rows // _BLOCK,), _E8M0), + ) + + +def _allocate_from_specs(specs, device) -> Tuple[torch.Tensor, ...]: + return tuple( + torch.empty(shape, dtype=dtype, device=device) for _, shape, dtype in specs + ) + + +def _validate_fwd_inputs(x_q, x_sf, w13_q, w13_sf, offsets): + sig = _meta_sig("fwd", x_q, x_sf, w13_q, w13_sf, offsets) + if sig is not None and sig in _validated_sigs: + rows, model_dim = x_q.shape + groups, two_hidden, _ = w13_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=x_q.device + ) + return rows, model_dim, two_hidden // 2, groups + if x_q.ndim != 2: + raise ValueError(f"x_q must be 2D [R, D], got shape {tuple(x_q.shape)}") + if w13_q.ndim != 3: + raise ValueError(f"w13_q must be 3D [G, 2F, D], got shape {tuple(w13_q.shape)}") + rows, model_dim = x_q.shape + groups, two_hidden, w_k = w13_q.shape + if w_k != model_dim: + raise ValueError(f"w13_q contraction dim {w_k} must match x_q's D {model_dim}") + if two_hidden % 2 != 0: + raise ValueError( + f"w13_q's row dim must be 2F (32-block interleaved gate/up), " + f"got {two_hidden}" + ) + hidden = two_hidden // 2 + device = x_q.device + + _require_cuda_device(device, "x_q") + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + if rows * max(model_dim, two_hidden) >= 2**31: + raise ValueError( + f"R * max(D, 2F) = {rows * max(model_dim, two_hidden)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + x_q, + name="x_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + # The rowwise weight cast delivers a contiguous [G, 2F, D] stack; the + # kernel-facing (2F, D, G) view is built from exactly that layout. + validate_operand( + w13_q, + name="w13_q", + shape=(groups, two_hidden, model_dim), + stride=(two_hidden * model_dim, model_dim, 1), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + x_sf, + name="x_sf", + logical_rows=rows, + logical_cols=model_dim // _BLOCK, + device=device, + ) + validate_blocked_scales( + w13_sf, + name="w13_sf", + logical_rows=two_hidden, + logical_cols=model_dim // _BLOCK, + device=device, + groups=groups, + ) + _remember_sig(sig, x_q, x_sf, w13_q, w13_sf, offsets) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_swiglu_fwd", mutates_args=()) +def _mxfp8_grouped_gemm_swiglu_fwd( + x_q: torch.Tensor, + x_sf: torch.Tensor, + w13_q: torch.Tensor, + w13_sf: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC1 grouped GEMM + SwiGLU + dual MXFP8 RCEIL quantization (one cuDNN launch). + + Inputs (all CUDA, one device; prequantized outside): + x_q E4M3 ``[R, D]`` stride ``(D, 1)``, rowwise 1x32 quantized. + x_sf flat blocked E8M0 scales for logical ``[R, D/32]`` (whole-matrix + blocked == per-group concat because R and every group are %256). + w13_q E4M3 ``[G, 2F, D]`` contiguous, rowwise quantized, rows in the + cuDNN 32-BLOCK GLU order ``[gate0(32) | up0(32) | gate1 | ...]``. + w13_sf per-group flat blocked E8M0, logical ``[2F, D/32]`` per expert. + offsets int32 CUDA ``[G]`` exclusive end rows; per-expert counts %256 + (caller invariant; see the validation module). + + Returns ``(z_bf16, h_row_q, h_row_sf, h_col_q, h_col_sf)``: + z_bf16 BF16 ``[R, 2F]`` contiguous pre-activation in the same 32-block + order; consumed unchanged by the bwd op. + h_row_q E4M3 ``[R, F]`` contiguous; h_row_sf flat blocked for ``[R, F/32]``. + h_col_q E4M3 ``[R, F]`` contiguous columnwise-quantized bytes + (un-transposed kernel layout); h_col_sf PER-GROUP flat blocked + for ``[F, rows_g/32]`` per expert. + + Rows past ``offsets[-1]`` of every output are GARBAGE (kernel-computed from + the quantized input tail) and read-forbidden. ``R == 0`` returns empty + outputs without touching cudnn; ``G == 0`` raises ValueError. + """ + rows, model_dim, hidden, groups = _validate_fwd_inputs( + x_q, x_sf, w13_q, w13_sf, offsets + ) + specs = _fwd_output_specs(rows, hidden) + if rows == 0: + return _allocate_from_specs(specs, x_q.device) + + import cudnn + + out = cudnn.grouped_gemm_glu_wrapper_sm100( + a_tensor=x_q.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(x_sf, rows, model_dim), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, x_q.device), + b_tensor=w13_q.permute(1, 2, 0), + sfb_tensor=_weight_scale_view(w13_sf, groups, 2 * hidden, model_dim), + norm_const_tensor=_cached_ones(1, torch.float32, x_q.device), + prob_tensor=None, + acc_dtype=torch.float32, + c_dtype=torch.bfloat16, + d_dtype=_E4M3, + cd_major="n", + sf_vec_size=_BLOCK, + act_func="swiglu", + discrete_col_sfd=True, + use_dynamic_sched=True, + current_stream=_stream(), + ) + results = ( + out["c_tensor"].view(rows, 2 * hidden), + out["d_tensor"].view(rows, hidden), + _flat_scales(out["sfd_row_tensor"]), + out["d_col_tensor"].view(rows, hidden), + _flat_scales(out["sfd_col_tensor"]), + ) + return tuple( + _check_normalized(t, name=spec[0], shape=spec[1], dtype=spec[2]) + for t, spec in zip(results, specs) + ) + + +@_mxfp8_grouped_gemm_swiglu_fwd.register_fake +def _(x_q, x_sf, w13_q, w13_sf, offsets): + rows, _model_dim, hidden, _groups = _validate_fwd_inputs( + x_q, x_sf, w13_q, w13_sf, offsets + ) + return _allocate_from_specs(_fwd_output_specs(rows, hidden), x_q.device) + + +# -------------------------------------------------------------------------- +# Op 2: grouped GEMM on prequantized operands -> BF16 (quant wrapper) +# -------------------------------------------------------------------------- + + +def _validate_mm_inputs(a_q, a_sf, b_q, b_sf, offsets): + sig = _meta_sig("mm", a_q, a_sf, b_q, b_sf, offsets) + if sig is not None and sig in _validated_sigs: + rows, contraction = a_q.shape + groups, out_features, _ = b_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=a_q.device + ) + return rows, out_features, contraction, groups + if a_q.ndim != 2: + raise ValueError(f"a_q must be 2D [R, K], got shape {tuple(a_q.shape)}") + if b_q.ndim != 3: + raise ValueError(f"b_q must be 3D [G, N, K], got shape {tuple(b_q.shape)}") + rows, contraction = a_q.shape + groups, out_features, b_k = b_q.shape + if b_k != contraction: + raise ValueError(f"b_q contraction dim {b_k} must match a_q's K {contraction}") + device = a_q.device + + _require_cuda_device(device, "a_q") + # N and K are both feature dims here (D/F/2F at the two call sites). + validate_feature_dims( + model_dim=out_features, + hidden_dim=contraction, + model_dim_name="b_q's output feature dim N", + hidden_dim_name="the contraction dim K", + ) + validate_allocated_rows(rows) + if rows * max(out_features, contraction) >= 2**31: + raise ValueError( + f"R * max(N, K) = {rows * max(out_features, contraction)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + a_q, + name="a_q", + shape=(rows, contraction), + stride=(contraction, 1), + dtype=_E4M3, + device=device, + ) + # b_q strides are free: rowwise weight casts arrive [G, N, K] contiguous + # and dim1-colwise casts arrive transposed to [G, N, K] (also row-major in + # this orientation); the wrapper reads the strides (both probe-proven). + validate_operand( + b_q, + name="b_q", + shape=(groups, out_features, contraction), + dtype=_E4M3, + device=device, + ) + validate_blocked_scales( + a_sf, + name="a_sf", + logical_rows=rows, + logical_cols=contraction // _BLOCK, + device=device, + ) + validate_blocked_scales( + b_sf, + name="b_sf", + logical_rows=out_features, + logical_cols=contraction // _BLOCK, + device=device, + groups=groups, + ) + _remember_sig(sig, a_q, a_sf, b_q, b_sf, offsets) + return rows, out_features, contraction, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm", mutates_args=()) +def _mxfp8_grouped_gemm( + a_q: torch.Tensor, + a_sf: torch.Tensor, + b_q: torch.Tensor, + b_sf: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + """Ragged grouped GEMM ``out[r] = dequant(a[r]) @ dequant(b[g(r)]).T`` -> BF16. + + Inputs: + a_q E4M3 ``[R, K]`` stride ``(K, 1)``, rowwise 1x32 quantized; a_sf flat + blocked for logical ``[R, K/32]``. + b_q E4M3 ``[G, N, K]``-logical, quantized ALONG K, any strides (rowwise + weight casts pass as-is; dim1-colwise casts pass transposed into + this orientation). + b_sf per-group flat blocked for the ``[N, K/32]``-oriented scale matrix + (uniform for both quantization axes). + offsets int32 CUDA ``[G]`` exclusive end rows. + + Covers FC2 forward (b = w2 rowwise: N=D, K=F) and FC1 dgrad (b = w13 + colwise: N=D, K=2F). Returns contiguous BF16 ``[R, N]``; rows past + ``offsets[-1]`` are left uninitialized (probe-verified untouched). + ``R == 0`` returns an empty output without touching cudnn. + """ + rows, out_features, contraction, groups = _validate_mm_inputs( + a_q, a_sf, b_q, b_sf, offsets + ) + out = torch.empty(rows, out_features, dtype=torch.bfloat16, device=a_q.device) + if rows == 0: + return out + + import cudnn + + cudnn.grouped_gemm_quant_wrapper_sm100( + a_tensor=a_q.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(a_sf, rows, contraction), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, a_q.device), + b_tensor=b_q.permute(1, 2, 0), + sfb_tensor=_weight_scale_view(b_sf, groups, out_features, contraction), + norm_const_tensor=None, + prob_tensor=None, + acc_dtype=torch.float32, + d_dtype=torch.bfloat16, + d_tensor=out.as_strided( + (rows, out_features, 1), (out_features, 1, rows * out_features) + ), + cd_major="n", + sf_vec_size=_BLOCK, + use_dynamic_sched=True, + current_stream=_stream(), + ) + return out + + +@_mxfp8_grouped_gemm.register_fake +def _(a_q, a_sf, b_q, b_sf, offsets): + rows, out_features, _contraction, _groups = _validate_mm_inputs( + a_q, a_sf, b_q, b_sf, offsets + ) + return torch.empty(rows, out_features, dtype=torch.bfloat16, device=a_q.device) + + +# -------------------------------------------------------------------------- +# Op 3: FC2 dgrad + dSwiGLU + dual quantization (dglu wrapper) +# -------------------------------------------------------------------------- + + +def _bwd_output_specs(rows: int, hidden: int): + two_hidden = 2 * hidden + return ( + ("dz_row_q", (rows, two_hidden), _E4M3), + ("dz_row_sf", (rows * two_hidden // _BLOCK,), _E8M0), + ("dz_col_q", (rows, two_hidden), _E4M3), + ("dz_col_sf", (two_hidden * rows // _BLOCK,), _E8M0), + ) + + +def _validate_bwd_inputs(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + sig = _meta_sig("bwd", dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets) + if sig is not None and sig in _validated_sigs: + rows, model_dim = dy_q.shape + groups, _, hidden = w2_col_q.shape + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=dy_q.device + ) + return rows, model_dim, hidden, groups + if dy_q.ndim != 2: + raise ValueError(f"dy_q must be 2D [R, D], got shape {tuple(dy_q.shape)}") + if w2_col_q.ndim != 3: + raise ValueError( + f"w2_col_q must be 3D [G, D, F], got shape {tuple(w2_col_q.shape)}" + ) + rows, model_dim = dy_q.shape + groups, w_d, hidden = w2_col_q.shape + if w_d != model_dim: + raise ValueError(f"w2_col_q's D dim {w_d} must match dy_q's D {model_dim}") + if z_bf16.ndim != 2 or tuple(z_bf16.shape) != (rows, 2 * hidden): + raise ValueError( + f"z_bf16 must be [{rows}, {2 * hidden}] (32-block interleaved, the " + f"exact fwd-op output), got shape {tuple(z_bf16.shape)}" + ) + device = dy_q.device + + _require_cuda_device(device, "dy_q") + validate_feature_dims(model_dim=model_dim, hidden_dim=hidden) + validate_allocated_rows(rows) + if rows * max(model_dim, 2 * hidden) >= 2**31: + raise ValueError( + f"R * max(D, 2F) = {rows * max(model_dim, 2 * hidden)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + validate_operand( + dy_q, + name="dy_q", + shape=(rows, model_dim), + stride=(model_dim, 1), + dtype=_E4M3, + device=device, + ) + # Colwise-quantized w2; strides free (dim1-native layout probe-proven). + validate_operand( + w2_col_q, + name="w2_col_q", + shape=(groups, model_dim, hidden), + dtype=_E4M3, + device=device, + ) + validate_operand( + z_bf16, + name="z_bf16", + shape=(rows, 2 * hidden), + stride=(2 * hidden, 1), + dtype=torch.bfloat16, + device=device, + ) + validate_blocked_scales( + dy_sf, + name="dy_sf", + logical_rows=rows, + logical_cols=model_dim // _BLOCK, + device=device, + ) + # Colwise weight scales: logical [F, D/32] per expert. + validate_blocked_scales( + w2_col_sf, + name="w2_col_sf", + logical_rows=hidden, + logical_cols=model_dim // _BLOCK, + device=device, + groups=groups, + ) + _remember_sig(sig, dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets) + return rows, model_dim, hidden, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_dswiglu_bwd", mutates_args=()) +def _mxfp8_grouped_gemm_dswiglu_bwd( + dy_q: torch.Tensor, + dy_sf: torch.Tensor, + w2_col_q: torch.Tensor, + w2_col_sf: torch.Tensor, + z_bf16: torch.Tensor, + offsets: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FC2 dgrad grouped GEMM + dSwiGLU + dual MXFP8 quantization (one launch). + + Inputs: + dy_q / dy_sf rowwise-quantized FC2 output gradient ``[R, D]``. + w2_col_q E4M3 ``[G, D, F]``-logical, quantized along D, any + strides (dim1-native accepted). + w2_col_sf per-group flat blocked for logical ``[F, D/32]``. + z_bf16 the EXACT ``[R, 2F]`` output of the fwd op (32-block + interleaved). Rows past ``offsets[-1]`` never read. + offsets int32 CUDA ``[G]``. + + Returns ``(dz_row_q, dz_row_sf, dz_col_q, dz_col_sf)``: the FC1 gradient + ``[R, 2F]`` in the same 32-block order, rowwise + columnwise quantized + (columnwise: un-transposed kernel bytes, PER-GROUP flat blocked scales). + Tails garbage/read-forbidden as in the fwd op. + """ + rows, model_dim, hidden, groups = _validate_bwd_inputs( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + specs = _bwd_output_specs(rows, hidden) + if rows == 0: + return _allocate_from_specs(specs, dy_q.device) + + import cudnn + + out = cudnn.grouped_gemm_dglu_wrapper_sm100( + a_tensor=dy_q.unsqueeze(0).permute(1, 2, 0), + c_tensor=z_bf16.unsqueeze(0).permute(1, 2, 0), + sfa_tensor=_act_scale_view(dy_sf, rows, model_dim), + padded_offsets=offsets, + alpha_tensor=_cached_ones(groups, torch.bfloat16, dy_q.device), + beta_tensor=_cached_ones(groups, torch.bfloat16, dy_q.device), + prob_tensor=None, + dprob_tensor=None, + b_tensor=w2_col_q.permute(2, 1, 0), + sfb_tensor=_weight_scale_view(w2_col_sf, groups, hidden, model_dim), + norm_const_tensor=_cached_ones(1, torch.float32, dy_q.device), + acc_dtype=torch.float32, + d_dtype=_E4M3, + cd_major="n", + sf_vec_size=_BLOCK, + act_func="dswiglu", + discrete_col_sfd=True, + use_dynamic_sched=True, + current_stream=_stream(), + ) + results = ( + out["d_row_tensor"].view(rows, 2 * hidden), + _flat_scales(out["sfd_row_tensor"]), + out["d_col_tensor"].view(rows, 2 * hidden), + _flat_scales(out["sfd_col_tensor"]), + ) + return tuple( + _check_normalized(t, name=spec[0], shape=spec[1], dtype=spec[2]) + for t, spec in zip(results, specs) + ) + + +@_mxfp8_grouped_gemm_dswiglu_bwd.register_fake +def _(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + rows, _model_dim, hidden, _groups = _validate_bwd_inputs( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + return _allocate_from_specs(_bwd_output_specs(rows, hidden), dy_q.device) + + +# -------------------------------------------------------------------------- +# Op 4: grouped weight gradient (wgrad wrapper, dense output mode) +# -------------------------------------------------------------------------- + + +def _validate_wgrad_inputs(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + sig = _meta_sig("wgrad", dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets) + if sig is not None and sig in _validated_sigs: + rows, out_features = dy_col_q.shape + in_features = x_col_q.shape[1] + groups = offsets.numel() + if host_offsets_validation_enabled(): + validate_group_offsets( + offsets, + num_groups=groups, + allocated_rows=rows, + device=dy_col_q.device, + ) + return rows, out_features, in_features, groups + if dy_col_q.ndim != 2 or x_col_q.ndim != 2: + raise ValueError( + "dy_col_q and x_col_q must both be 2D logical [R, N] / [R, K], got " + f"{tuple(dy_col_q.shape)} and {tuple(x_col_q.shape)}" + ) + rows, out_features = dy_col_q.shape + x_rows, in_features = x_col_q.shape + if x_rows != rows: + raise ValueError( + f"dy_col_q and x_col_q must share the row dim: {rows} vs {x_rows}" + ) + groups = offsets.numel() if isinstance(offsets, torch.Tensor) else 0 + device = dy_col_q.device + + _require_cuda_device(device, "dy_col_q") + validate_allocated_rows(rows) + validate_feature_dims( + model_dim=out_features, + hidden_dim=in_features, + model_dim_name="dy_col_q's feature dim N", + hidden_dim_name="x_col_q's feature dim K", + ) + if rows * max(out_features, in_features) >= 2**31: + raise ValueError( + f"R * max(N, K) = {rows * max(out_features, in_features)} does not " + "fit an int32 element index" + ) + validate_group_offsets( + offsets, num_groups=groups, allocated_rows=rows, device=device + ) + # Both operands accept ANY major: dim1-native transposed memory, the fwd/ + # bwd ops' un-transposed kernel bytes, and mixes -- all four combinations + # probe-proven. + validate_operand( + dy_col_q, + name="dy_col_q", + shape=(rows, out_features), + dtype=_E4M3, + device=device, + ) + validate_operand( + x_col_q, + name="x_col_q", + shape=(rows, in_features), + dtype=_E4M3, + device=device, + ) + # Columnwise scale buffers are sized by offsets[-1] (a device value), not + # by R: only dtype/device/divisibility are host-checkable. + validate_ragged_colwise_scales( + dy_col_sf, + name="dy_col_sf", + features=out_features, + allocated_rows=rows, + device=device, + ) + validate_ragged_colwise_scales( + x_col_sf, + name="x_col_sf", + features=in_features, + allocated_rows=rows, + device=device, + ) + # No cross-buffer size check: a kernel-produced operand's scales are sized + # by the ALLOCATED rows while a composite-produced operand's are sized by + # the ROUTED total offsets[-1] -- mixing the two is legitimate and + # probe-proven (tail case); the kernel reads only within offsets. + _remember_sig(sig, dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets) + return rows, out_features, in_features, groups + + +@torch.library.custom_op("torchao::mxfp8_grouped_gemm_wgrad", mutates_args=()) +def _mxfp8_grouped_gemm_wgrad( + dy_col_q: torch.Tensor, + dy_col_sf: torch.Tensor, + x_col_q: torch.Tensor, + x_col_sf: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Inputs: + dy_col_q E4M3 logical ``[R, N]``, columnwise (32x1) quantized, ANY major. + dy_col_sf PER-GROUP flat blocked scales (each expert's ``[N, rows_g/32]`` + block concatenated; the K-groups layout). Sized by the routed + total ``offsets[-1]``, which may be < R. + x_col_q / x_col_sf likewise for logical ``[R, K]``. + offsets int32 CUDA ``[G]`` exclusive ends over the shared row axis. + + Do NOT feed whole-matrix ``to_blocked`` scales here: the per-group and + whole-matrix orders coincide in byte count but not content whenever G > 1, + and the mismatch is silent (probe: 2-5 dB instead of 100+). + + Returns contiguous BF16 ``dw [G, N, K]`` (FP32 accumulation). Zero-token + experts ARE written (all-zero slices, probe-verified); ``R == 0`` returns + zeros without launching. FC1 wgrad: N=2F, K=D. FC2 wgrad: N=D, K=F. + """ + rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + dw = torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) + if rows == 0: + return dw.zero_() + + import cudnn + + cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=dy_col_q.t(), + b_tensor=x_col_q, + sfa_tensor=_as_e8m0(dy_col_sf).view(out_features, -1), + sfb_tensor=_as_e8m0(x_col_sf).view(in_features, -1), + offsets_tensor=offsets, + acc_dtype=torch.float32, + sf_vec_size=_BLOCK, + accumulate_on_output=False, + output_mode="dense", + wgrad_tensor=dw, + wgrad_dtype=torch.bfloat16, + current_stream=_stream(), + ) + return dw + + +@_mxfp8_grouped_gemm_wgrad.register_fake +def _(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + _rows, out_features, in_features, groups = _validate_wgrad_inputs( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + ) + return torch.empty( + (groups, out_features, in_features), + dtype=torch.bfloat16, + device=dy_col_q.device, + ) + + +# -------------------------------------------------------------------------- +# Public wrappers: availability-gated entry points over the four custom ops. +# -------------------------------------------------------------------------- + + +def mxfp8_grouped_gemm_swiglu_fwd(x_q, x_sf, w13_q, w13_sf, offsets): + """FC1 grouped GEMM + SwiGLU + rowwise/columnwise MXFP8 quantization. + + See ``torchao::mxfp8_grouped_gemm_swiglu_fwd`` for the full ABI. ``w13_q`` + is E4M3 ``[G, 2F, D]`` contiguous with rows in 32-block GLU order; returns + ``(z_bf16 [R, 2F], h_row_q [R, F], h_row_sf, h_col_q [R, F], h_col_sf)`` + where the columnwise scales are PER-GROUP blocked. Rows past + ``offsets[-1]`` of every output are garbage and read-forbidden. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_q, x_sf, w13_q, w13_sf, offsets + ) + + +def mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets): + """Ragged grouped GEMM on prequantized MXFP8 operands, BF16 output. + + ``b_q`` is ``[G, N, K]``-logical quantized along K with free strides + (rowwise casts as-is; dim1-colwise casts transposed into this + orientation); ``b_sf`` is always the per-group blocked ``[N, K/32]`` + orientation. Returns BF16 ``[R, N]`` with rows past ``offsets[-1]`` + uninitialized. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm(a_q, a_sf, b_q, b_sf, offsets) + + +def mxfp8_grouped_gemm_dswiglu_bwd(dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets): + """FC2 dgrad + dSwiGLU + dual MXFP8 quantization of the FC1 gradient. + + ``z_bf16`` must be the exact fwd-op output. Returns + ``(dz_row_q [R, 2F], dz_row_sf, dz_col_q [R, 2F], dz_col_sf)`` in the same + 32-block order. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets + ) + + +def mxfp8_grouped_gemm_wgrad(dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets): + """Grouped MXFP8 weight gradient ``dw[g] = dequant(dy_g).T @ dequant(x_g)``. + + Both operands columnwise (32x1) quantized with PER-GROUP blocked scales + (never whole-matrix ``to_blocked`` -- same byte count, silently wrong + block order). Returns contiguous BF16 ``[G, N, K]``. + """ + _require_available() + return torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets + )