diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 0fddf63b33..7e99ae256d 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -103,6 +103,7 @@ run_test_config(){ run_default_fa 1 triton_kernels/test_cast_mxfp8.py run_default_fa 1 triton_kernels/test_cast_mxfp4.py run_default_fa 1 triton_kernels/test_grouped_gemm.py + run_default_fa 1 triton_kernels/test_grouped_gemm_mxfp4.py run_default_fa 1 triton_kernels/test_utils.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 triton_kernels/test_norms.py NVTE_ROCM_ENABLE_MXFP8=1 NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py diff --git a/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py new file mode 100644 index 0000000000..f255b6ad0b --- /dev/null +++ b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py @@ -0,0 +1,331 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Numeric correctness tests for the grouped MXFP4 Triton GEMM (gfx950).""" + +import pytest +import torch + +from transformer_engine.pytorch.quantization import check_mxfp4_support + +_MXFP4_OK, _MXFP4_REASON = check_mxfp4_support() + +from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( + MXFP4_BLOCK, + _col_operand, + _col_operand_grouped_padded, + _row_operand, + _row_operand_mxfp8, + grouped_gemm_a8w4_fprop, + grouped_gemm_mxfp4_dgrad, + grouped_gemm_mxfp4_fprop, + grouped_gemm_mxfp4_wgrad, + grouped_linear_mxfp4, +) + +pytestmark = [ + pytest.mark.skipif(not _MXFP4_OK, reason=f"MXFP4 unsupported: {_MXFP4_REASON}"), +] + +DTYPE = torch.bfloat16 +# Uneven, non-128-multiple group sizes exercise fprop/dgrad masking and the +# wgrad per-group zero-padding to 128. +M_SPLITS = [96, 128, 160, 128] +# K = 256 (>128) so the reduction loop runs multiple iterations (loop_k = +# K/BLOCK_K), exercising the in-loop operand pointer advance. +N, K = 256, 256 +# Precise-reference bar: kernel vs dequant-of-the-same-operands differ only by +# ~bf16 output rounding. +_TIGHT_TOL = 3.0e-2 +# Loose bar vs the true (unquantized) bf16 matmul: ~0.16 MXFP4 noise floor for +# these shapes, well below the ~1.4 an uncorrelated (layout-bug) output gives. +_REL_TOL = 2.0e-1 + +# OCP E2M1 magnitude indexed by the 3 low bits (exp2 | mantissa1); bit 3 = sign. +_E2M1_MAG = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +@pytest.fixture(autouse=True) +def _mxfp4_env(monkeypatch): + # Triton MXFP4 quantizer (no aiter dependency); seed for a deterministic + # quantization error. + monkeypatch.setenv("NVTE_USE_CAST_TRANSPOSE_TRITON", "1") + torch.manual_seed(0) + + +def _rand(*shape): + return torch.randn(*shape, dtype=DTYPE, device="cuda") + + +def _rel_err(out, ref): + out = out.float() + ref = ref.float() + return (out - ref).norm() / ref.norm().clamp_min(1e-12) + + +def _dequant_mxfp4(data_u8, scale_u8, feat): + """Independent OCP MXFP4 dequant of packed operands -> fp32 ``[R, feat]``. + + ``data_u8`` ``[R, feat/2]`` packs two E2M1 codes per byte (low nibble = even + index along ``feat``); ``scale_u8`` ``[R, feat/32]`` is one E8M0 scale + (value ``2**(x-127)``) per 1x32 block. + """ + lut = torch.tensor(_E2M1_MAG, dtype=torch.float32, device=data_u8.device) + lo = (data_u8 & 0xF).to(torch.long) + hi = ((data_u8 >> 4) & 0xF).to(torch.long) + codes = torch.stack((lo, hi), dim=-1).reshape(data_u8.shape[0], feat) + mag = lut[codes & 0x7] + vals = torch.where((codes & 0x8).bool(), -mag, mag) + scale = torch.exp2(scale_u8.to(torch.float32) - 127.0) + return vals * scale.repeat_interleave(MXFP4_BLOCK, dim=1) + + +# MXFP8 (e4m3) activation side, for the a8w4 forward: the quantizer lives in the +# impl module (_row_operand_mxfp8); here we only need the independent decoder. +def _dequant_mxfp8(data_e4m3, scale_u8): + """Independent MXFP8 dequant: e4m3 [R, K] + e8m0 [R, K/32] -> fp32 [R, K].""" + rows, feat = data_e4m3.shape + vals = data_e4m3.float().reshape(rows, feat // MXFP4_BLOCK, MXFP4_BLOCK) + scale = torch.exp2(scale_u8.to(torch.float32) - 127.0) + return (vals * scale.unsqueeze(-1)).reshape(rows, feat) + + +def test_fprop_precise(): + total_m = sum(M_SPLITS) + a = _rand(total_m, K) + weights = [_rand(N, K) for _ in M_SPLITS] + + out = grouped_gemm_mxfp4_fprop(a, weights, M_SPLITS, out_dtype=DTYPE) + + # Dequantize the same row-wise operands the kernel used, then grouped matmul. + a_deq = _dequant_mxfp4(*_row_operand(a), K) + ref = torch.empty((total_m, N), dtype=torch.float32, device="cuda") + start = 0 + for w, m in zip(weights, M_SPLITS): + w_deq = _dequant_mxfp4(*_row_operand(w), K) # [N, K] + ref[start : start + m] = a_deq[start : start + m] @ w_deq.t() + start += m + + assert out.shape == (total_m, N) + assert _rel_err(out, ref) < _TIGHT_TOL + + +def test_fprop_unaligned_total_m(): + # total_M not divisible by 32 must still quantize (leading dim is padded). + splits = [100, 130] # sum = 230, not a multiple of 32 + total_m = sum(splits) + a = _rand(total_m, K) + weights = [_rand(N, K) for _ in splits] + + out = grouped_gemm_mxfp4_fprop(a, weights, splits, out_dtype=DTYPE) + + ref = torch.empty((total_m, N), dtype=torch.float32, device="cuda") + start = 0 + for w, m in zip(weights, splits): + ref[start : start + m] = a[start : start + m].float() @ w.float().t() + start += m + assert out.shape == (total_m, N) + assert _rel_err(out, ref) < _REL_TOL + + +def test_fprop_zero_group(): + # An expert may receive no tokens (m_split == 0): it must contribute no tiles + # and leave the other groups' outputs correct. + splits = [96, 0, 160, 0, 128] # two empty experts + total_m = sum(splits) + a = _rand(total_m, K) + weights = [_rand(N, K) for _ in splits] + + out = grouped_gemm_mxfp4_fprop(a, weights, splits, out_dtype=DTYPE) + + a_deq = _dequant_mxfp4(*_row_operand(a), K) + ref = torch.empty((total_m, N), dtype=torch.float32, device="cuda") + start = 0 + for w, m in zip(weights, splits): + if m == 0: + continue + w_deq = _dequant_mxfp4(*_row_operand(w), K) + ref[start : start + m] = a_deq[start : start + m] @ w_deq.t() + start += m + assert out.shape == (total_m, N) + assert _rel_err(out, ref) < _TIGHT_TOL + + +def test_dgrad_precise(): + total_m = sum(M_SPLITS) + grad_out = _rand(total_m, N) + weights = [_rand(N, K) for _ in M_SPLITS] + + dgrad = grouped_gemm_mxfp4_dgrad(grad_out, weights, M_SPLITS, out_dtype=DTYPE) + + # gradO row-wise, weight col-wise (the transposed operand): dA = gradO @ W. + go_deq = _dequant_mxfp4(*_row_operand(grad_out), N) # [total_M, N] + ref = torch.empty((total_m, K), dtype=torch.float32, device="cuda") + start = 0 + for w, m in zip(weights, M_SPLITS): + w_col_deq = _dequant_mxfp4(*_col_operand(w), N) # [K, N] ~ W^T + ref[start : start + m] = go_deq[start : start + m] @ w_col_deq.t() + start += m + + assert dgrad.shape == (total_m, K) + assert _rel_err(dgrad, ref) < _TIGHT_TOL + + +def test_wgrad_precise_and_padding(): + total_m = sum(M_SPLITS) + a = _rand(total_m, K) + grad_out = _rand(total_m, N) + + wgrad = grouped_gemm_mxfp4_wgrad(a, grad_out, M_SPLITS, out_dtype=DTYPE) + assert wgrad.shape == (len(M_SPLITS), N, K) + + # Precise: dequant the exact per-group padded col operands the kernel reduced + # over, sliced by the padded offsets. C[g] = lhs[:, g] @ rhs[:, g]^T. + lhs_data, lhs_scale, go_pad = _col_operand_grouped_padded(grad_out, M_SPLITS) # [N, Mpad/2] + rhs_data, rhs_scale, _ = _col_operand_grouped_padded(a, M_SPLITS) # [K, Mpad/2] + m_pad_total = lhs_data.shape[1] * 2 + lhs_deq = _dequant_mxfp4(lhs_data, lhs_scale, m_pad_total) # [N, Mpad] + rhs_deq = _dequant_mxfp4(rhs_data, rhs_scale, m_pad_total) # [K, Mpad] + go = go_pad.tolist() + ref = torch.empty((len(M_SPLITS), N, K), dtype=torch.float32, device="cuda") + for g in range(len(M_SPLITS)): + s, e = go[g], go[g + 1] + ref[g] = lhs_deq[:, s:e] @ rhs_deq[:, s:e].t() + assert _rel_err(wgrad, ref) < _TIGHT_TOL + + # End-to-end vs the true (unquantized) matmul: exercises the per-group + # zero-padding -- the padded rows must contribute nothing. + ref_true = torch.empty((len(M_SPLITS), N, K), dtype=torch.float32, device="cuda") + start = 0 + for g, m in enumerate(M_SPLITS): + ref_true[g] = grad_out[start : start + m].float().t() @ a[start : start + m].float() + start += m + assert _rel_err(wgrad, ref_true) < _REL_TOL + + +def test_wgrad_accumulate(): + total_m = sum(M_SPLITS) + a = _rand(total_m, K) + grad_out = _rand(total_m, N) + G = len(M_SPLITS) + + # beta=1: dW accumulates into an existing (main_grad-style) fp32 buffer. + base = torch.randn(G, N, K, dtype=torch.float32, device="cuda") + out = base.clone() + grouped_gemm_mxfp4_wgrad(a, grad_out, M_SPLITS, out_dtype=torch.float32, out=out, accumulate=True) + + fresh = grouped_gemm_mxfp4_wgrad(a, grad_out, M_SPLITS, out_dtype=torch.float32) + torch.testing.assert_close(out, base + fresh) + + +def test_autograd_matches_ops(): + total_m = sum(M_SPLITS) + a = _rand(total_m, K).requires_grad_(True) + weight = torch.stack([_rand(N, K) for _ in M_SPLITS], dim=0).requires_grad_(True) + + out = grouped_linear_mxfp4(a, weight, M_SPLITS) + assert out.shape == (total_m, N) + + grad_out = _rand(total_m, N) + out.backward(grad_out) + + # The autograd grads must equal the direct op calls (same code path). + ref_da = grouped_gemm_mxfp4_dgrad( + grad_out, list(weight.detach().unbind(0)), M_SPLITS, out_dtype=a.dtype + ) + ref_dw = grouped_gemm_mxfp4_wgrad(a.detach(), grad_out, M_SPLITS, out_dtype=weight.dtype) + torch.testing.assert_close(a.grad, ref_da) + torch.testing.assert_close(weight.grad, ref_dw) + + +# a8w4 forward: MXFP8 (e4m3) activation x MXFP4 (e2m1) weight, grouped along M. +# The op (grouped_gemm_a8w4_fprop) does the quantize + call; each test re-derives +# the same e4m3 operands via _row_operand_mxfp8 to build the precise reference. +def _a8w4_reference(a_q, a_s, weights, m_splits): + """Per-group fp32 ref on the same operands: dequant(e4m3 A) @ dequant(e2m1 W)^T.""" + a_deq = _dequant_mxfp8(a_q, a_s) + n, feat_k = weights[0].shape + ref = torch.empty((a_q.shape[0], n), dtype=torch.float32, device="cuda") + start = 0 + for w, m in zip(weights, m_splits): + w_deq = _dequant_mxfp4(*_row_operand(w), feat_k) # [N, K] + ref[start : start + m] = a_deq[start : start + m] @ w_deq.t() + start += m + return ref + + +def test_a8w4_fprop_precise(): + total_m = sum(M_SPLITS) + a = _rand(total_m, K) + weights = [_rand(N, K) for _ in M_SPLITS] + + out = grouped_gemm_a8w4_fprop(a, weights, M_SPLITS, out_dtype=DTYPE) + a_q, a_s = _row_operand_mxfp8(a) + ref = _a8w4_reference(a_q, a_s, weights, M_SPLITS) + + assert out.shape == (total_m, N) + assert _rel_err(out, ref) < _TIGHT_TOL + + +def test_a8w4_fprop_single_group(): + # G=1 dense case (Gluon-style): the group scan degenerates to one expert. + splits = [512] + a = _rand(512, K) + weights = [_rand(N, K)] + + out = grouped_gemm_a8w4_fprop(a, weights, splits, out_dtype=DTYPE) + a_q, a_s = _row_operand_mxfp8(a) + ref = _a8w4_reference(a_q, a_s, weights, splits) + + assert _rel_err(out, ref) < _TIGHT_TOL + + +def test_a8w4_fprop_unaligned_total_m(): + # total_M not a multiple of 32 must still quantize + run (leading dim padded). + splits = [100, 130] # sum = 230 + total_m = sum(splits) + a = _rand(total_m, K) + weights = [_rand(N, K) for _ in splits] + + out = grouped_gemm_a8w4_fprop(a, weights, splits, out_dtype=DTYPE) + a_q, a_s = _row_operand_mxfp8(a) + ref = _a8w4_reference(a_q, a_s, weights, splits) + + assert out.shape == (total_m, N) + assert _rel_err(out, ref) < _TIGHT_TOL + + +# Kimi-K3 MoE expert-GEMM shapes: fc1 gate-up and fc2 down, G=28 experts. +# Small per-group M keeps the check light while still exercising the K3 N/K +# tiling (deep K, multi-tile N) and the 28-group scan under balanced and +# imbalanced (~4:1 routing skew) token distributions. +_K3_SHAPES = {"gateup": (6144, 3584), "down": (3584, 3072)} # (N, K) + + +def _k3_splits(kind, G=28, per=8): + total = G * per + if kind == "balanced": + return [per] * G + # deterministic ~4:1 max:min skew, integer, exact sum + w = [1.0 + 3.0 * i / (G - 1) for i in range(G)] + s = sum(w) + sp = [max(1, round(total * wi / s)) for wi in w] + sp[-1] += total - sum(sp) + return sp + + +@pytest.mark.parametrize("shape", list(_K3_SHAPES), ids=list(_K3_SHAPES)) +@pytest.mark.parametrize("split", ["balanced", "imbalanced"]) +def test_a8w4_fprop_k3(shape, split): + n, k = _K3_SHAPES[shape] + splits = _k3_splits(split) + total_m = sum(splits) + a = _rand(total_m, k) + weights = [_rand(n, k) for _ in splits] + + out = grouped_gemm_a8w4_fprop(a, weights, splits, out_dtype=DTYPE) + a_q, a_s = _row_operand_mxfp8(a) + ref = _a8w4_reference(a_q, a_s, weights, splits) + + assert out.shape == (total_m, n) + assert _rel_err(out, ref) < _TIGHT_TOL diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py new file mode 100644 index 0000000000..1093813361 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -0,0 +1,641 @@ +# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Grouped MXFP4 GEMM Triton persistent kernels (gfx950). + +Adapted for TransformerEngine from Primus-Turbo. + +- _grouped_mxfp4_persistent_gemm_kernel: Forward (NT) C[g] = A[g] @ B[g]^T +- grouped_gemm_mxfp4_triton_kernel: Forward public API +- _grouped_mxfp4_variable_k_gemm_kernel: Backward wgrad C[g] = LHS[g] @ RHS[g]^T +- grouped_gemm_mxfp4_variable_k_triton_kernel: Backward (variable-K) public API +""" + +from __future__ import annotations + +import functools + +import torch +import triton +import triton.language as tl + +from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs + + +# =============================================================================== +# Arch helper +# =============================================================================== + + +@functools.lru_cache(maxsize=None) +def _is_gfx950_for(device_index: int) -> bool: + props = torch.cuda.get_device_properties(device_index) + return (props.major, props.minor) == (9, 5) + + +def _is_gfx950() -> bool: + # Key the cache by device index so a set_device change isn't masked. + return _is_gfx950_for(torch.cuda.current_device()) + + +# =============================================================================== +# AMD Triton compiler knobs +# +# The grouped FP4 kernels want the gfx950 knob set (async_copy, +# scalarize_packed_fops, block_pingpong) active only while they compile/launch, +# then restored so they never leak into other Triton kernels. These knobs are +# read at compile time and are not part of Triton's compile cache key. +# =============================================================================== + +def _scoped_amd_knobs(*, is_tn): + """Decorator applying blockwise_fp8's scoped AMD compiler knobs around a + launch. ``is_tn`` selects the layout-gated set (gfx942); on gfx950 the full + async_copy / scalarize / block_pingpong set is enabled either way. + """ + + def deco(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + with _amd_compiler_knobs(is_tn=is_tn): + return func(*args, **kwargs) + + return wrapper + + return deco + + +# =============================================================================== +# Chiplet transform: map the AMD round-robin program_id onto XCD-contiguous +# chunks so a band of programs shares B columns in L2. +# =============================================================================== + +NUM_XCDS = 8 # gfx942/gfx950 chiplet count + + +@triton.jit +def _chiplet_transform_chunked( + pid, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, +): + if pid > (NUM_SMS // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE): + return pid + local_pid = pid // NUM_XCDS + chunk_idx = local_pid // CHUNK_SIZE + pos_in_chunk = local_pid % CHUNK_SIZE + xcd = pid % NUM_XCDS + return chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk + + +# E8M0 block size: one scale per 1x32 logical-element block (same as MXFP8). +# tl.dot_scaled consumes the "e2m1" FP4 format string (inlined in the kernels; +# @triton.jit cannot read module globals). +VEC_SIZE = 32 + + +# ########################################################################### +# Launch-param autotuning. Each shape is swept over a few (BLOCK_K, num_warps, +# num_stages, nonkdim, waves_per_eu) configs via @triton.autotune and the +# fastest is cached. Set TRITON_PRINT_AUTOTUNING=1 to log the winner. Neither +# BLOCK_K nor nonkdim (the scaled-MFMA size) changes the grid. +# ########################################################################### + +# (BLOCK_K, num_warps, num_stages, matrix_instr_nonkdim, waves_per_eu). +# The nonkdim=32 rows use the 32x32 scaled-MFMA, which is faster for a8w4 +# (e4m3 x e2m1); fp4 x fp4 stays on the 16x16 op the tuner already prefers. +_FWD_CONFIGS = [ + (128, 8, 2, 16, 0), + (128, 8, 3, 16, 2), + (128, 8, 4, 16, 2), + (128, 4, 2, 16, 0), + (256, 8, 2, 16, 0), + (256, 8, 3, 16, 2), + (128, 8, 2, 32, 0), + (128, 8, 3, 32, 2), +] +_VARK_CONFIGS = [ + (128, 8, 3, 16, 2), + (128, 8, 2, 16, 0), + (128, 8, 4, 16, 2), + (128, 4, 3, 16, 2), + (128, 8, 3, 16, 0), + (256, 8, 2, 16, 0), + (256, 8, 3, 16, 2), +] + + +def _make_configs(specs): + """Build a ``triton.Config`` list from (BLOCK_K, warps, stages, nonkdim, + waves) tuples. ``kpack=1`` is deprecated on gfx950; ``matrix_instr_nonkdim`` + selects the scaled-MFMA size (16 or 32). + """ + return [ + triton.Config( + {"BLOCK_SIZE_K": bk, "waves_per_eu": wpe, "matrix_instr_nonkdim": nonk, "kpack": 1}, + num_warps=nw, + num_stages=ns, + ) + for (bk, nw, ns, nonk, wpe) in specs + ] + + +def _prune_fwd_by_k(configs, named_args, **kwargs): + """Keep only configs whose BLOCK_K divides K (the fwd reduction loop is unmasked).""" + K = named_args["K"] + return [c for c in configs if K % c.kwargs["BLOCK_SIZE_K"] == 0] + + +def _prune_vark_by_pad(configs, named_args, **kwargs): + """Drop BLOCK_K=256 unless every padded group length is 256-divisible (can256).""" + if named_args["can256"]: + return list(configs) + return [c for c in configs if c.kwargs["BLOCK_SIZE_K"] != 256] + + +# ########################################################################### +# Forward (NT): C[g] = A[g] @ B[g]^T +# A (total_M, K/2) fp4-packed, A_scale (total_M, K/32) e8m0 +# B (G, N, K/2) fp4-packed, B_scale (G, N, K/32) e8m0 +# ########################################################################### + + +@triton.autotune( + configs=_make_configs(_FWD_CONFIGS), + key=["N", "K", "G", "A_PACK", "B_PACK"], + prune_configs_by={"early_config_prune": _prune_fwd_by_k}, +) +@triton.jit +def _grouped_mxfp4_persistent_gemm_kernel( + A, + B, + C, + A_scale, # (total_M, K//32) uint8 e8m0 + B_scale, # (G, N, K//32) uint8 e8m0 + rd_offs_ptr, # (G+1) int64 — read base along M (A / A_scale) + out_offs_ptr, # (G+1) int64 — write base along M (C, tile count) + G, + N, + K, # logical contraction size (multiple of BLOCK_SIZE_K) + stride_am, + stride_ak, # along packed K bytes + stride_bg, + stride_bn, + stride_bk, # along packed K bytes + stride_cm, + stride_cn, + stride_asm, + stride_ask, + stride_bsg, + stride_bsn, + stride_bsk, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, # logical K per iter + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, + VEC: tl.constexpr, + A_PACK: tl.constexpr = 2, # elems/byte along K: 2 => fp4 (K/2 bytes/iter), 1 => fp8 (K bytes) + B_PACK: tl.constexpr = 2, + A_FMT: tl.constexpr = "e2m1", # per-operand tl.dot_scaled format ("e2m1"/"e4m3"/"e5m2") + B_FMT: tl.constexpr = "e2m1", +): + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = _chiplet_transform_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + total_tiles: tl.int32 = 0 + for _g in range(G): + m_g = (tl.load(out_offs_ptr + _g + 1) - tl.load(out_offs_ptr + _g)).to(tl.int32) + total_tiles += tl.cdiv(m_g, BLOCK_SIZE_M) * num_pid_n + + tl.assume(stride_am > 0) + tl.assume(stride_cm > 0) + + BK_A: tl.constexpr = BLOCK_SIZE_K // A_PACK # A bytes per K-iter (fp4=K/2, fp8=K) + BK_B: tl.constexpr = BLOCK_SIZE_K // B_PACK # B bytes per K-iter + BK_SCALE: tl.constexpr = BLOCK_SIZE_K // VEC # scale entries per K-iter + + for global_tile_id in range(pid, total_tiles, NUM_SMS): + # ── locate group (linear scan) ── + group_idx: tl.int32 = 0 + tile_start: tl.int32 = 0 + cumsum: tl.int32 = 0 + for _g in range(G): + m_g_i = (tl.load(out_offs_ptr + _g + 1) - tl.load(out_offs_ptr + _g)).to(tl.int32) + tiles_g = tl.cdiv(m_g_i, BLOCK_SIZE_M) * num_pid_n + new_cumsum = cumsum + tiles_g + if global_tile_id >= new_cumsum: + group_idx = _g + 1 + tile_start = new_cumsum + cumsum = new_cumsum + + local_tile = global_tile_id - tile_start + m_rd = tl.load(rd_offs_ptr + group_idx) # int64 read base (A / A_scale) + m_out = tl.load(out_offs_ptr + group_idx) # int64 write base (C) + M_g = (tl.load(out_offs_ptr + group_idx + 1) - m_out).to(tl.int32) + tiles_m_g = tl.cdiv(M_g, BLOCK_SIZE_M) + + # ── swizzle ── + num_pid_in_group = GROUP_SIZE_M * num_pid_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m_g - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rk_a = tl.arange(0, BK_A) # A bytes along K + rk_b = tl.arange(0, BK_B) # B bytes along K + rks = tl.arange(0, BK_SCALE) # scale entries + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + A_BASE = A + (m_rd + rm[:, None]) * stride_am + rk_a[None, :] * stride_ak + B_BASE = B + group_idx.to(tl.int64) * stride_bg + rk_b[:, None] * stride_bk + rn[None, :] * stride_bn + AS_BASE = A_scale + rks[:, None] * stride_ask + (m_rd + rm[None, :]) * stride_asm + BS_BASE = ( + B_scale + + group_idx.to(tl.int64) * stride_bsg + + rks[:, None] * stride_bsk + + rn[None, :] * stride_bsn + ) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + loop_k = K // BLOCK_SIZE_K + for ki in range(0, loop_k): + a = tl.load(tl.multiple_of(A_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) # (BM, BK_A) + b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) # (BK_B, BN) + a_s = tl.trans(tl.load(AS_BASE)) # (BK/32, BM) -> (BM, BK/32) + b_s = tl.trans(tl.load(BS_BASE)) # (BK/32, BN) -> (BN, BK/32) + acc = tl.dot_scaled(a, a_s, A_FMT, b, b_s, B_FMT, acc) + A_BASE += BK_A * stride_ak + B_BASE += BK_B * stride_bk + AS_BASE += BK_SCALE * stride_ask + BS_BASE += BK_SCALE * stride_bsk + + c = acc.to(C.type.element_ty) + rm_s = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn_s = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) + c_mask = (rm_s[:, None] < M_g) & (rn_s[None, :] < N) + C_ = C + (m_out + rm_s[:, None]) * stride_cm + rn_s[None, :] * stride_cn + tl.store(C_, c, c_mask) + + +@_scoped_amd_knobs(is_tn=False) +def grouped_gemm_mxfp4_triton_kernel( + a, + a_scale, + b, + b_scale, + group_offs, + N, + K, + group_offs_out=None, + out_dtype=torch.bfloat16, + num_cu=None, +): + """A(total_M, K/2) @ B(G, N, K/2)^T -> C. FP4-packed data, e8m0 uint8 scales. + + group_offs: read offsets along M for A / A_scale. + group_offs_out: write offsets along M for C (defaults to group_offs). + K is the logical contraction size (must be a multiple of 128). + """ + if group_offs_out is None: + group_offs_out = group_offs + G = b.shape[0] + c = torch.empty((a.shape[0], N), dtype=out_dtype, device=a.device) + a_s = a_scale.view(torch.uint8) + b_s = b_scale.view(torch.uint8) + a_u8 = a.view(torch.uint8) + b_u8 = b.view(torch.uint8) + cu = num_cu if num_cu is not None else torch.cuda.get_device_properties(a.device).multi_processor_count + BM, BN = 256, 256 + m_alloc = a.shape[0] + avg_m = max(m_alloc // max(G, 1), 1) + tiles_n = (N + BN - 1) // BN + GM = 8 if min((avg_m + BM - 1) // BM, tiles_n) < 16 else 4 + total_tiles = ((m_alloc + BM - 1) // BM + G) * tiles_n + num_sms = min(total_tiles, cu) + chunk = 64 if num_sms >= NUM_XCDS * 64 else 32 + + # BLOCK_SIZE_K / num_warps / num_stages / waves_per_eu / matrix_instr_nonkdim / + # kpack come from the @triton.autotune config; the grid and the host-chosen + # tiles (BLOCK_M/N, GROUP_M, CHUNK, NUM_SMS) are passed here. + _grouped_mxfp4_persistent_gemm_kernel[(num_sms,)]( + a_u8, + b_u8, + c, + a_s, + b_s, + group_offs, + group_offs_out, + G, + N, + K, + a_u8.stride(0), + a_u8.stride(1), + b_u8.stride(0), + b_u8.stride(1), + b_u8.stride(2), + c.stride(0), + c.stride(1), + a_s.stride(0), # stride_asm (M); a_s is (total_M, K/32) + a_s.stride(1), # stride_ask (K/32 contiguous) + b_s.stride(0), # stride_bsg + b_s.stride(1), # stride_bsn (N); b_s is (G, N, K/32) + b_s.stride(2), # stride_bsk (K/32 contiguous) + BLOCK_SIZE_M=BM, + BLOCK_SIZE_N=BN, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + A_PACK=2, # fp4 A; passed explicitly so A_PACK/B_PACK bind into the autotune key + B_PACK=2, + ) + return c + + +@_scoped_amd_knobs(is_tn=False) +def grouped_gemm_a8w4_triton_kernel( + a, # activation MXFP8 e4m3, (total_M, K) float8_e4m3fn + a_scale, # (total_M, K/32) uint8 e8m0 + b, # weight MXFP4 e2m1, (G, N, K/2) uint8-packed + b_scale, # (G, N, K/32) uint8 e8m0 + group_offs, + N, + K, + group_offs_out=None, + out_dtype=torch.bfloat16, + num_cu=None, +): + """a8w4 forward: C = A(e4m3) @ B(e2m1)^T, grouped along M (routed experts). + + Same persistent NT kernel as the FP4xFP4 path, but the A operand is MXFP8 + (1 byte/elem, K bytes/row) instead of packed MXFP4. Both operands carry E8M0 + (VEC_SIZE=32) scales along K; K is the logical contraction (multiple of 128). + + group_offs: read offsets along M for A / A_scale. + group_offs_out: write offsets along M for C (defaults to group_offs). + """ + if group_offs_out is None: + group_offs_out = group_offs + G = b.shape[0] + c = torch.empty((a.shape[0], N), dtype=out_dtype, device=a.device) + a_s = a_scale.view(torch.uint8) + b_s = b_scale.view(torch.uint8) + a_e4m3 = a.view(torch.float8_e4m3fn) # e4m3 operand, unpacked (K bytes/row) + b_u8 = b.view(torch.uint8) # e2m1 packed (K/2 bytes/row) + cu = num_cu if num_cu is not None else torch.cuda.get_device_properties(a.device).multi_processor_count + BM, BN = 256, 256 + m_alloc = a.shape[0] + avg_m = max(m_alloc // max(G, 1), 1) + tiles_n = (N + BN - 1) // BN + GM = 8 if min((avg_m + BM - 1) // BM, tiles_n) < 16 else 4 + total_tiles = ((m_alloc + BM - 1) // BM + G) * tiles_n + num_sms = min(total_tiles, cu) + chunk = 64 if num_sms >= NUM_XCDS * 64 else 32 + + _grouped_mxfp4_persistent_gemm_kernel[(num_sms,)]( + a_e4m3, + b_u8, + c, + a_s, + b_s, + group_offs, + group_offs_out, + G, + N, + K, + a_e4m3.stride(0), + a_e4m3.stride(1), # stride_ak: 1 e4m3 byte per K element + b_u8.stride(0), + b_u8.stride(1), + b_u8.stride(2), # stride_bk: packed K/2 bytes + c.stride(0), + c.stride(1), + a_s.stride(0), + a_s.stride(1), + b_s.stride(0), + b_s.stride(1), + b_s.stride(2), + BLOCK_SIZE_M=BM, + BLOCK_SIZE_N=BN, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + A_PACK=1, # e4m3: 1 elem/byte -> K bytes per K-iter + A_FMT="e4m3", + B_PACK=2, # e2m1: 2 elems/byte -> K/2 bytes + B_FMT="e2m1", + ) + return c + + +# ########################################################################### +# Variable-K backward (wgrad): C[g] = LHS[g] @ RHS[g]^T, reduction over M_g +# LHS (OUT_M, M_total/2) fp4-packed, LHS_scale (OUT_M, M_total/32) e8m0 +# RHS (OUT_N, M_total/2) fp4-packed, RHS_scale (OUT_N, M_total/32) e8m0 +# C (G, OUT_M, OUT_N) +# go_pad: padded per-group offsets along M (each M_g a multiple of 128) +# ########################################################################### + + +@triton.autotune( + configs=_make_configs(_VARK_CONFIGS), + key=["OUT_M", "OUT_N", "G", "can256"], + prune_configs_by={"early_config_prune": _prune_vark_by_pad}, + restore_value=["C"], +) +@triton.jit +def _grouped_mxfp4_variable_k_gemm_kernel( + LHS, + RHS, + C, + LHS_scale, + RHS_scale, + go_pad_ptr, + G, + OUT_M, + OUT_N, + can256, # 1 if every padded group length is 256-divisible (autotune key only) + stride_lm, + stride_lk, # along packed M bytes + stride_rm, + stride_rk, # along packed M bytes + stride_cg, + stride_cm, + stride_cn, + stride_lsm, + stride_lsk, + stride_rsm, + stride_rsk, + BLOCK_SIZE_M: tl.constexpr, # over OUT_M + BLOCK_SIZE_N: tl.constexpr, # over OUT_N + BLOCK_SIZE_K: tl.constexpr, # logical reduction over M_g + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, + VEC: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = _chiplet_transform_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + tiles_m = tl.cdiv(OUT_M, BLOCK_SIZE_M) + tiles_n = tl.cdiv(OUT_N, BLOCK_SIZE_N) + tiles_per_group = tiles_m * tiles_n + total_tiles = G * tiles_per_group + + tl.assume(stride_lm > 0) + tl.assume(stride_rm > 0) + tl.assume(stride_cm > 0) + + BK_PACK: tl.constexpr = BLOCK_SIZE_K // 2 + BK_SCALE: tl.constexpr = BLOCK_SIZE_K // VEC + + for global_tile in range(pid, total_tiles, NUM_SMS): + group_idx = global_tile // tiles_per_group + local_tile = global_tile - group_idx * tiles_per_group + + num_pid_in_group = GROUP_SIZE_M * tiles_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + m_start = tl.load(go_pad_ptr + group_idx) # int64 (logical M, multiple of 128) + M_g = (tl.load(go_pad_ptr + group_idx + 1) - m_start).to(tl.int32) + + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % OUT_M + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % OUT_N + rk = tl.arange(0, BK_PACK) + rks = tl.arange(0, BK_SCALE) + rm = tl.max_contiguous(tl.multiple_of(rm, BLOCK_SIZE_M), BLOCK_SIZE_M) + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + # packed M offset = m_start // 2 (M_g multiple of 128 -> even) + mp0 = (m_start // 2).to(tl.int64) + L_BASE = LHS + rm[:, None] * stride_lm + (mp0 + rk[None, :]) * stride_lk + R_BASE = RHS + (mp0 + rk[:, None]) * stride_rk + rn[None, :] * stride_rm + sk0 = (m_start // VEC).to(tl.int32) + LS_BASE = LHS_scale + (sk0 + rks[:, None]) * stride_lsk + rm[None, :] * stride_lsm + RS_BASE = RHS_scale + (sk0 + rks[:, None]) * stride_rsk + rn[None, :] * stride_rsm + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + loop_k = M_g // BLOCK_SIZE_K # padded -> no mask + for _ in range(loop_k): + l = tl.load(tl.multiple_of(L_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) # (BM, BK/2) + r = tl.load(tl.multiple_of(R_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) # (BK/2, BN) + ls = tl.trans(tl.load(LS_BASE)) # (BK/32, BM) -> (BM, BK/32) + rs = tl.trans(tl.load(RS_BASE)) + acc = tl.dot_scaled(l, ls, "e2m1", r, rs, "e2m1", acc) + L_BASE += BK_PACK * stride_lk + R_BASE += BK_PACK * stride_rk + LS_BASE += BK_SCALE * stride_lsk + RS_BASE += BK_SCALE * stride_rsk + + rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + cmask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) + C_ = C + group_idx.to(tl.int64) * stride_cg + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + if ACCUMULATE: # beta=1: fold into the existing buffer (e.g. main_grad) + acc += tl.load(C_, mask=cmask, other=0.0).to(tl.float32) + c = acc.to(C.type.element_ty) + tl.store(C_, c, cmask) + + +@_scoped_amd_knobs(is_tn=True) +def grouped_gemm_mxfp4_variable_k_triton_kernel( + lhs, lhs_scale, rhs, rhs_scale, go_pad, OUT_M, OUT_N, G, + out_dtype=torch.bfloat16, num_cu=None, out=None, accumulate=False, can256=None, +): + """C[g] (OUT_M, OUT_N) = lhs[:,g] @ rhs[:,g]^T. + + lhs (OUT_M, M_total/2) fp4-packed, rhs (OUT_N, M_total/2) fp4-packed. + go_pad: padded per-group offsets along M (each M_g a multiple of 128). + ``out`` (G, OUT_M, OUT_N): reuse this caller buffer instead of allocating; with + ``accumulate=True`` the result is added into it (beta=1, e.g. main_grad). + """ + if out is not None: + assert tuple(out.shape) == (G, OUT_M, OUT_N), ( + f"out shape {tuple(out.shape)} != {(G, OUT_M, OUT_N)}" + ) + assert out.device == lhs.device, f"out device {out.device} != operand device {lhs.device}" + c = out + else: + assert not accumulate, "accumulate=True requires an existing out buffer" + c = torch.empty((G, OUT_M, OUT_N), dtype=out_dtype, device=lhs.device) + ls = lhs_scale.view(torch.uint8) + rs = rhs_scale.view(torch.uint8) + l_u8 = lhs.view(torch.uint8) + r_u8 = rhs.view(torch.uint8) + cu = num_cu if num_cu is not None else torch.cuda.get_device_properties(lhs.device).multi_processor_count + BM, BN = 256, 128 + tiles_m = (OUT_M + BM - 1) // BM + tiles_n = (OUT_N + BN - 1) // BN + GM = 8 if min(tiles_m, tiles_n) < 16 else 4 + total_tiles = G * tiles_m * tiles_n + num_sms = min(total_tiles, cu) + chunk = 64 if num_sms >= NUM_XCDS * 64 else 32 + + # can256 (256-divisibility of every padded group length) is part of the + # autotune key so a cached BLOCK_K=256 config is never reused on a shape that + # can't take it. The op supplies it (no sync); fall back to a go_pad read here. + if can256 is None: + _lens = [b - a for a, b in zip(go_pad.tolist(), go_pad[1:].tolist())] + can256 = 1 if all(n % 256 == 0 for n in _lens) else 0 + + # BLOCK_SIZE_K / warps / stages / waves come from @triton.autotune; restore_value + # on the kernel keeps the beta=1 accumulate buffer intact across benchmark reps. + _grouped_mxfp4_variable_k_gemm_kernel[(num_sms,)]( + l_u8, + r_u8, + c, + ls, + rs, + go_pad, + G, + OUT_M, + OUT_N, + can256, + l_u8.stride(0), + l_u8.stride(1), + r_u8.stride(0), + r_u8.stride(1), + c.stride(0), + c.stride(1), + c.stride(2), + ls.stride(0), # stride_lsm (OUT_M); ls is (OUT_M, M/32) + ls.stride(1), # stride_lsk (M/32 contiguous) + rs.stride(0), # stride_rsm (OUT_N); rs is (OUT_N, M/32) + rs.stride(1), # stride_rsk (M/32 contiguous) + BLOCK_SIZE_M=BM, + BLOCK_SIZE_N=BN, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + ACCUMULATE=accumulate, + ) + return c diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py new file mode 100644 index 0000000000..81702da6e7 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -0,0 +1,513 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""TE-side quantization bridge + dispatch for grouped MXFP4 GEMM (gfx950).""" + +from __future__ import annotations + +import os +from typing import List, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F + +from ..tensor.mxfp4_tensor import MXFP4Quantizer +from ..utils import round_up_to_nearest_multiple +from .grouped_gemm_mxfp4 import ( + _is_gfx950, + grouped_gemm_a8w4_triton_kernel, + grouped_gemm_mxfp4_triton_kernel, + grouped_gemm_mxfp4_variable_k_triton_kernel, +) + +# Logical contraction tile the kernels step by (see grouped_gemm_mxfp4.py). +BLOCK_SIZE_K = 128 +# Per-group wgrad M padding multiple; must be >= the largest wgrad BLOCK_K the +# kernel may pick. 256 lets the variable-K autotune try BLOCK_K=256 (at the cost +# of more padding for small groups). Default 128 = unchanged behavior. +WGRAD_PAD_MULTIPLE = int(os.environ.get("NVTE_MXFP4_WGRAD_PAD", str(BLOCK_SIZE_K))) +# One E8M0 scale per 32-element block. +MXFP4_BLOCK = 32 + + +def _quantizer(*, rowwise: bool, columnwise: bool) -> MXFP4Quantizer: + """MXFP4 quantizer configured for the plain (un-swizzled) Triton layout.""" + return MXFP4Quantizer( + rowwise=rowwise, + columnwise=columnwise, + shuffle_rowwise_data=False, + shuffle_columnwise_data=False, + with_gemm_swizzled_scales=False, + use_hadamard=False, + ) + + +def _prefix_offsets(m_splits: Sequence[int], device: torch.device) -> torch.Tensor: + """Tight prefix-sum offsets ``[0, m0, m0+m1, ...]`` as int64 on device.""" + offs = [0] + for m in m_splits: + offs.append(offs[-1] + int(m)) + return torch.tensor(offs, dtype=torch.int64, device=device) + + +def _check_contract(dim: int, name: str) -> None: + if dim % BLOCK_SIZE_K != 0: + raise ValueError( + f"grouped MXFP4 GEMM requires the contraction dim ({name}={dim}) to be a" + f" multiple of {BLOCK_SIZE_K}." + ) + + +def _check_splits( + m_splits: Sequence[int], total_rows: int, num_experts: Optional[int] = None +) -> None: + """Validate ``m_splits`` before it reaches the raw kernel (memory safety). + + Bad splits would otherwise make the Triton kernel read past ``group_offs`` or + store past the output. Checks: non-negative splits, ``sum == total_rows``, and + (when given) ``len == num_experts``. + """ + if any(int(m) < 0 for m in m_splits): + raise ValueError(f"m_splits must be non-negative, got {list(m_splits)}") + got = sum(int(m) for m in m_splits) + if got != total_rows: + raise ValueError(f"sum(m_splits)={got} must equal the grouped row count {total_rows}") + if num_experts is not None and len(m_splits) != num_experts: + raise ValueError( + f"len(m_splits)={len(m_splits)} must equal the number of experts {num_experts}" + ) + + +def _require_gfx950() -> None: + # The kernels use the CDNA4 scaled-FP4 MFMA (tl.dot_scaled e2m1); no other arch has it. + if not _is_gfx950(): + raise RuntimeError( + "grouped MXFP4 GEMM requires gfx950 (CDNA4); the current device lacks the" + " scaled-FP4 MFMA used by tl.dot_scaled(..., \"e2m1\", ...)." + ) + + +def _row_operand(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Row-wise MXFP4: ``x`` [M, F] -> (data [M, F/2] u8, scale [M, F/32] u8). + + ``M`` is a free (masked) GEMM dim of any size; it is zero-padded to a + 32-multiple for the quantizer (which requires it) and sliced back. + """ + M, Feat = x.shape + M_pad = round_up_to_nearest_multiple(M, MXFP4_BLOCK) + if M_pad != M: + x = F.pad(x.contiguous(), (0, 0, 0, M_pad - M)) + q = _quantizer(rowwise=True, columnwise=False).quantize(x.contiguous()) + data = q._rowwise_data.view(torch.uint8)[:M] + scale = q._rowwise_scale_inv[:M, : Feat // MXFP4_BLOCK].contiguous() + return data, scale + + +def _col_operand(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Col-wise MXFP4: ``x`` [M, F] -> (data [F, M/2] u8, scale [F, M/32] u8). + + Requests both directions (matching the production transpose-quant path); the + row-wise result is discarded. Both dims are zero-padded to a 32-multiple for + the quantizer and sliced back (both are free/masked in the GEMM). + """ + M, Feat = x.shape + M_pad = round_up_to_nearest_multiple(M, MXFP4_BLOCK) + Feat_pad = round_up_to_nearest_multiple(Feat, MXFP4_BLOCK) + if M_pad != M or Feat_pad != Feat: + x = F.pad(x.contiguous(), (0, Feat_pad - Feat, 0, M_pad - M)) + q = _quantizer(rowwise=True, columnwise=True).quantize(x.contiguous()) + data = q._columnwise_data.view(torch.uint8)[:Feat, : M // 2] + scale = q._columnwise_scale_inv[:Feat, : M // MXFP4_BLOCK].contiguous() + return data, scale + + +def _row_col_operand( + x: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """One cast_transpose: MXFP4 row-wise *and* col-wise from a single quantize. + + ``x`` [M, F] -> (row_data [M, F/2], row_scale [M, F/32], col_data [F, M/2], + col_scale [F, M/32]), all ``uint8``. Lets a caller that needs both layouts of + the same tensor (e.g. a weight consumed row-wise by fprop and col-wise by + dgrad) pay the transpose-quant once instead of twice. Both dims are padded to + a 32-multiple for the quantizer and sliced back. + """ + M, Feat = x.shape + M_pad = round_up_to_nearest_multiple(M, MXFP4_BLOCK) + Feat_pad = round_up_to_nearest_multiple(Feat, MXFP4_BLOCK) + if M_pad != M or Feat_pad != Feat: + x = F.pad(x.contiguous(), (0, Feat_pad - Feat, 0, M_pad - M)) + q = _quantizer(rowwise=True, columnwise=True).quantize(x.contiguous()) + row_data = q._rowwise_data.view(torch.uint8)[:M, : Feat // 2] + row_scale = q._rowwise_scale_inv[:M, : Feat // MXFP4_BLOCK].contiguous() + col_data = q._columnwise_data.view(torch.uint8)[:Feat, : M // 2] + col_scale = q._columnwise_scale_inv[:Feat, : M // MXFP4_BLOCK].contiguous() + return row_data, row_scale, col_data, col_scale + + +def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Row-wise MXFP8 (e4m3): ``x`` [M, K] -> (data [M, K] e4m3, scale [M, K/32] u8). + + Plain OCP MXFP8: unpacked e4m3 element bytes + one E8M0 scale per 1x32 block, + with the Kimi-K3 ``ceil(log2(amax / 448))`` block-scale rule. This is a + reference cast matching the a8w4 kernel's tested operand layout; the GATE-2 + quantizer-parity pass swaps in the recipe quantizer for exact QAT parity. + ``K`` (the contraction) is already a 32-multiple; ``M`` needs no padding since + the scale is per-row. + """ + fmax = torch.finfo(torch.float8_e4m3fn).max + M, K = x.shape + xb = x.reshape(M, K // MXFP4_BLOCK, MXFP4_BLOCK).float() + amax = xb.abs().amax(-1, keepdim=True).clamp_min(1e-30) + exp = torch.ceil(torch.log2(amax / fmax)).clamp(-127, 127) + q = (xb / torch.exp2(exp)).clamp(-fmax, fmax).to(torch.float8_e4m3fn) + data = q.reshape(M, K).contiguous() + scale = (exp.squeeze(-1) + 127).to(torch.uint8) + return data, scale + + +def _quantize_weights_row_col( + weights: Sequence[torch.Tensor], +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Per-expert weight cast_transpose once -> stacked row and col MXFP4 operands. + + Returns ``((row_data [G,N,K/2], row_scale [G,N,K/32]), + (col_data [G,K,N/2], col_scale [G,K,N/32]))`` -- row for fprop, col for dgrad, + both from a single quantize per expert. + """ + rows_d, rows_s, cols_d, cols_s = [], [], [], [] + for w in weights: + rd, rs, cd, cs = _row_col_operand(w) + rows_d.append(rd) + rows_s.append(rs) + cols_d.append(cd) + cols_s.append(cs) + row = (torch.stack(rows_d, dim=0), torch.stack(rows_s, dim=0)) + col = (torch.stack(cols_d, dim=0), torch.stack(cols_s, dim=0)) + return row, col + + +def _col_operand_grouped_padded( + x: torch.Tensor, + m_splits: Sequence[int], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-group col-wise MXFP4 with each group's M padded to WGRAD_PAD_MULTIPLE. + + ``x`` [total_M, F] grouped along M -> (data [F, total_pad/2] u8, + scale [F, total_pad/32] u8, go_pad [G+1] int64). Padding rows are zeros, so + they contribute nothing to the wgrad contraction over M. + """ + Feat = x.shape[1] + datas: List[torch.Tensor] = [] + scales: List[torch.Tensor] = [] + offs = [0] + start = 0 + for m in m_splits: + m = int(m) + xg = x[start : start + m] + start += m + m_pad = round_up_to_nearest_multiple(m, WGRAD_PAD_MULTIPLE) + if m_pad != m: + xg = F.pad(xg, (0, 0, 0, m_pad - m)) + data, scale = _col_operand(xg) # (F, m_pad/2), (F, m_pad/32) + datas.append(data) + scales.append(scale) + offs.append(offs[-1] + m_pad) + data = torch.cat(datas, dim=1).contiguous() + scale = torch.cat(scales, dim=1).contiguous() + go_pad = torch.tensor(offs, dtype=torch.int64, device=x.device) + return data, scale, go_pad + + +def grouped_gemm_mxfp4_fprop( + a: torch.Tensor, + weights: List[torch.Tensor], + m_splits: Sequence[int], + *, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = None, + weight_row: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> torch.Tensor: + """Grouped MXFP4 forward: ``C[g] = A[g] @ W[g]^T`` (contract K). + + Args: + a: [total_M, K] activations, grouped along M by ``m_splits``. + weights: list of G per-expert weight tensors, each [N, K]. + m_splits: per-group token counts (len G). + weight_row: optional pre-quantized row-wise weights + ``(data [G, N, K/2], scale [G, N, K/32])`` -- skips the weight cast + (e.g. shared with dgrad's col-wise cast, see ``_quantize_weights_row_col``). + + Returns: + [total_M, N] output in ``out_dtype``. + """ + K = a.shape[1] + _require_gfx950() + _check_contract(K, "K") + + if weight_row is not None: + b_data, b_scale = weight_row + N = b_data.shape[1] + _check_splits(m_splits, a.shape[0], b_data.shape[0]) + else: + N = weights[0].shape[0] + _check_splits(m_splits, a.shape[0], len(weights)) + b_datas, b_scales = [], [] + for w in weights: + d, s = _row_operand(w) # (N, K/2), (N, K/32) + b_datas.append(d) + b_scales.append(s) + b_data = torch.stack(b_datas, dim=0) # (G, N, K/2) + b_scale = torch.stack(b_scales, dim=0) # (G, N, K/32) + + a_data, a_scale = _row_operand(a) # (total_M, K/2), (total_M, K/32) + group_offs = _prefix_offsets(m_splits, a.device) + return grouped_gemm_mxfp4_triton_kernel( + a_data, + a_scale, + b_data, + b_scale, + group_offs, + N, + K, + group_offs_out=group_offs, + out_dtype=out_dtype, + num_cu=num_cu, + ) + + +def grouped_gemm_a8w4_fprop( + a: torch.Tensor, + weights: List[torch.Tensor], + m_splits: Sequence[int], + *, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = None, + weight_row: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> torch.Tensor: + """Grouped a8w4 forward: ``C[g] = A[g] @ W[g]^T`` (contract K). + + Same as :func:`grouped_gemm_mxfp4_fprop`, but the activation is cast to MXFP8 + (e4m3) instead of MXFP4; the weight stays MXFP4 (e2m1). Forward-only, per the + Kimi-K3 QAT recipe (no a8w4 dgrad/wgrad). + + Args: + a: [total_M, K] activations, grouped along M by ``m_splits``. + weights: list of G per-expert weight tensors, each [N, K]. + m_splits: per-group token counts (len G). + weight_row: optional pre-quantized row-wise MXFP4 weights + ``(data [G, N, K/2], scale [G, N, K/32])`` -- skips the weight cast. + + Returns: + [total_M, N] output in ``out_dtype``. + """ + K = a.shape[1] + _require_gfx950() + _check_contract(K, "K") + + if weight_row is not None: + b_data, b_scale = weight_row + N = b_data.shape[1] + _check_splits(m_splits, a.shape[0], b_data.shape[0]) + else: + N = weights[0].shape[0] + _check_splits(m_splits, a.shape[0], len(weights)) + b_datas, b_scales = [], [] + for w in weights: + d, s = _row_operand(w) # (N, K/2), (N, K/32) + b_datas.append(d) + b_scales.append(s) + b_data = torch.stack(b_datas, dim=0) # (G, N, K/2) + b_scale = torch.stack(b_scales, dim=0) # (G, N, K/32) + + a_data, a_scale = _row_operand_mxfp8(a) # (total_M, K) e4m3, (total_M, K/32) e8m0 + group_offs = _prefix_offsets(m_splits, a.device) + return grouped_gemm_a8w4_triton_kernel( + a_data, + a_scale, + b_data, + b_scale, + group_offs, + N, + K, + group_offs_out=group_offs, + out_dtype=out_dtype, + num_cu=num_cu, + ) + + +def grouped_gemm_mxfp4_dgrad( + grad_out: torch.Tensor, + weights: Optional[List[torch.Tensor]], + m_splits: Sequence[int], + *, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = None, + weight_col: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, +) -> torch.Tensor: + """Grouped MXFP4 dgrad: ``dA[g] = gradO[g] @ W[g]`` (contract N). + + Uses the forward NT kernel with the col-wise weight as the transposed + operand: ``B[g] = W_col`` has shape [K, N/2], so the kernel's free dim is K + (in_features) and its contraction is N (out_features). + + Args: + grad_out: [total_M, N] output gradient, grouped along M by ``m_splits``. + weights: list of G per-expert weight tensors, each [N, K]. May be ``None`` + when ``weight_col`` supplies the pre-quantized weights. + weight_col: optional pre-quantized col-wise weights + ``(data [G, K, N/2], scale [G, K, N/32])`` -- skips the weight cast + (e.g. shared with fprop's row-wise cast, see ``_quantize_weights_row_col``). + + Returns: + [total_M, K] input gradient in ``out_dtype``. + """ + N = grad_out.shape[1] + _require_gfx950() + _check_contract(N, "N") + + if weight_col is not None: + b_data, b_scale = weight_col + K = b_data.shape[1] + _check_splits(m_splits, grad_out.shape[0], b_data.shape[0]) + else: + assert weights is not None, "dgrad needs weights or weight_col" + K = weights[0].shape[1] + _check_splits(m_splits, grad_out.shape[0], len(weights)) + b_datas, b_scales = [], [] + for w in weights: + d, s = _col_operand(w) # (K, N/2), (K, N/32) + b_datas.append(d) + b_scales.append(s) + b_data = torch.stack(b_datas, dim=0) # (G, K, N/2) + b_scale = torch.stack(b_scales, dim=0) # (G, K, N/32) + + go_data, go_scale = _row_operand(grad_out) # (total_M, N/2), (total_M, N/32) + group_offs = _prefix_offsets(m_splits, grad_out.device) + # Kernel free dim = b.shape[-2] = K; kernel contraction = b.shape[-1]*2 = N. + return grouped_gemm_mxfp4_triton_kernel( + go_data, + go_scale, + b_data, + b_scale, + group_offs, + K, + N, + group_offs_out=group_offs, + out_dtype=out_dtype, + num_cu=num_cu, + ) + + +def grouped_gemm_mxfp4_wgrad( + a: torch.Tensor, + grad_out: torch.Tensor, + m_splits: Sequence[int], + *, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = None, + out: Optional[torch.Tensor] = None, + accumulate: bool = False, +) -> torch.Tensor: + """Grouped MXFP4 wgrad: ``dW[g] = gradO[g]^T @ A[g]`` (contract M, variable-K). + + Both operands are col-wise, per-group padded to a BLOCK_SIZE_K-multiple M so + the reduction over M needs no masking. + + Args: + a: [total_M, K] activations, grouped along M by ``m_splits``. + grad_out: [total_M, N] output gradient, grouped along M by ``m_splits``. + out: [G, N, K] buffer to write into instead of allocating; with + ``accumulate=True`` the wgrad is added into it (beta=1, e.g. main_grad). + + Returns: + [G, N, K] weight gradient in ``out_dtype``. + """ + N = grad_out.shape[1] + K = a.shape[1] + G = len(m_splits) + _require_gfx950() + if a.shape[0] != grad_out.shape[0]: + raise ValueError(f"a rows ({a.shape[0]}) must equal grad_out rows ({grad_out.shape[0]})") + _check_splits(m_splits, a.shape[0]) + + lhs_data, lhs_scale, go_pad = _col_operand_grouped_padded(grad_out, m_splits) # (N, Mpad/2) + rhs_data, rhs_scale, _ = _col_operand_grouped_padded(a, m_splits) # (K, Mpad/2) + + # Whether every padded group length is 256-divisible; part of the wgrad autotune + # key so a cached BLOCK_K=256 config is never reused on a shape that can't take + # it. Computed from the Python splits, so the hot path needs no device sync. + can256 = int( + all(round_up_to_nearest_multiple(int(m), WGRAD_PAD_MULTIPLE) % 256 == 0 for m in m_splits) + ) + + return grouped_gemm_mxfp4_variable_k_triton_kernel( + lhs_data, + lhs_scale, + rhs_data, + rhs_scale, + go_pad, + N, + K, + G, + out_dtype=out_dtype, + num_cu=num_cu, + out=out, + accumulate=accumulate, + can256=can256, + ) + + +class _GroupedGemmMXFP4Func(torch.autograd.Function): + """Autograd for grouped MXFP4 linear: ``out[g] = a[g] @ weight[g]^T``. + + The weight is cast_transposed once in forward (``_quantize_weights_row_col``): + the row-wise layout feeds fprop and the col-wise layout is stashed for dgrad, + so the weight quant is paid once per step instead of twice. Activation and + grad_out still cast per pass -- their col layouts are per-group padded for the + variable-K wgrad reduction, so they don't share with the un-padded row cast. + """ + + @staticmethod + def forward(ctx, a, weight, m_splits): + weights = list(torch.unbind(weight, dim=0)) + w_row, w_col = _quantize_weights_row_col(weights) + out = grouped_gemm_mxfp4_fprop( + a, weights, m_splits, out_dtype=a.dtype, weight_row=w_row + ) + ctx.save_for_backward(a, weight, w_col[0], w_col[1]) + ctx.m_splits = m_splits + return out + + @staticmethod + def backward(ctx, grad_out): + a, weight, w_col_data, w_col_scale = ctx.saved_tensors + m_splits = ctx.m_splits + grad_out = grad_out.contiguous() + grad_a = grad_weight = None + if ctx.needs_input_grad[0]: + grad_a = grouped_gemm_mxfp4_dgrad( + grad_out, None, m_splits, out_dtype=a.dtype, + weight_col=(w_col_data, w_col_scale), + ) + if ctx.needs_input_grad[1]: + grad_weight = grouped_gemm_mxfp4_wgrad(a, grad_out, m_splits, out_dtype=weight.dtype) + return grad_a, grad_weight, None + + +def grouped_linear_mxfp4( + a: torch.Tensor, + weight: torch.Tensor, + m_splits: Sequence[int], +) -> torch.Tensor: + """Autograd-enabled grouped MXFP4 linear. + + Args: + a: [total_M, K] activations, grouped along M by ``m_splits``. + weight: [G, N, K] stacked per-expert weights. + m_splits: per-group token counts (len G). + + Returns: + [total_M, N] output. + """ + return _GroupedGemmMXFP4Func.apply(a, weight, tuple(int(m) for m in m_splits))