From 0556570b4099016959b03d5a02923b208553dffb Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 28 Aug 2026 13:37:08 -0500 Subject: [PATCH 01/11] Triton mxfp4 grouped GEMM kernel --- .../triton_kernels/grouped_gemm_mxfp4.py | 601 ++++++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4_impl.py | 270 ++++++++ 2 files changed, 871 insertions(+) create mode 100644 transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py create mode 100644 transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py 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..0f9764c0ca --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -0,0 +1,601 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Grouped MXFP4 (MX_BLOCKWISE) GEMM Triton persistent kernels (gfx950). + +Adapted for TransformerEngine from Primus-Turbo +(https://github.com/AMD-AGI/Primus-Turbo, MIT) -- the block-scaled grouped FP4 +kernels in ``primus_turbo/triton/grouped_gemm/grouped_gemm_fp4_kernel.py``. The +kernel bodies are kept faithful to the source; the shared scaffolding +(``NUM_XCDS``, ``_chiplet_transform_chunked``, the AMD Triton-knob scope, the +output padding-tail pass) is inlined here so the module has no Primus-Turbo +dependency. + +The operands are E2M1 FP4 packed two-values-per-byte along the contraction (K) +axis, fed to the hardware block-scaled MMA via ``tl.dot_scaled(..., "e2m1", ...)`` +with E8M0 (VEC_SIZE=32) scales. + + - _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 + +FP4 packing notes: + * data tensors store K/2 (resp. M/2) bytes along the contraction axis; + * scale tensors store K/32 (resp. M/32) E8M0 bytes (one per 1x32 block); + * scales are stored free-major (free, K/32) / (G, free, K/32); each K-iter + loads a (K/32, free) tile (coalesced) and transposes it back in-reg for + tl.dot_scaled. + +Only EVEN contraction (a multiple of BLOCK_SIZE_K = 128 logical elements) is +supported; the FP4 quantizer pads K to 128 and the MX wrapper pads per-group M +to 128, so this always holds in practice. +""" + +from __future__ import annotations + +import contextlib +import functools +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + + +# =============================================================================== +# Arch helper +# =============================================================================== + + +@functools.lru_cache +def _is_gfx950() -> bool: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + return (props.major, props.minor) == (9, 5) + + +# =============================================================================== +# 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. +# =============================================================================== + +_AMD_KNOB_ATTRS = ("use_async_copy", "scalarize_packed_fops", "use_block_pingpong") +_AMD_KNOB_ENVS = ( + "TRITON_HIP_USE_ASYNC_COPY", + "AMDGCN_SCALARIZE_PACKED_FOPS", + "TRITON_HIP_USE_BLOCK_PINGPONG", +) + + +def _amd_knobs_available() -> bool: + return hasattr(triton, "knobs") and hasattr(triton.knobs, "amd") + + +@contextlib.contextmanager +def _scoped_amd_triton_knobs(): + """Snapshot AMD Triton compiler knobs on entry, restore them on exit. + + On gfx950 the full gfx950 knob set is enabled for the duration of the scope; + all changes are restored on exit. + """ + saved_attrs = {} + if _amd_knobs_available(): + amd = triton.knobs.amd + for name in _AMD_KNOB_ATTRS: + if hasattr(amd, name): + saved_attrs[name] = getattr(amd, name) + saved_env = {name: os.environ.get(name) for name in _AMD_KNOB_ENVS} + try: + if _is_gfx950(): + if _amd_knobs_available(): + amd = triton.knobs.amd + for name in _AMD_KNOB_ATTRS: + if hasattr(amd, name): + setattr(amd, name, True) + else: + for name in _AMD_KNOB_ENVS: + os.environ[name] = "1" + yield + finally: + if saved_attrs: + amd = triton.knobs.amd + for name, val in saved_attrs.items(): + setattr(amd, name, val) + for name, val in saved_env.items(): + if val is None: + os.environ.pop(name, None) + else: + os.environ[name] = val + + +def _scoped_amd_knobs(func): + """Decorator: run ``func`` under :func:`_scoped_amd_triton_knobs`.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + with _scoped_amd_triton_knobs(): + return func(*args, **kwargs) + + return wrapper + + +# =============================================================================== +# 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 # MI300 / MI350 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 + + +# ########################################################################### +# 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.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, +): + 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_PACK: tl.constexpr = BLOCK_SIZE_K // 2 # packed 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 = tl.arange(0, BK_PACK) # packed K bytes + 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[None, :] * stride_ak + B_BASE = B + group_idx.to(tl.int64) * stride_bg + rk[:, 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/2) + b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) # (BK/2, 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, "e2m1", b, b_s, "e2m1", acc) + A_BASE += BK_PACK * stride_ak + B_BASE += BK_PACK * 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 +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 BLOCK_SIZE_K=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, BK = 256, 256, 128 + 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 + grid = (num_sms,) + _grouped_mxfp4_persistent_gemm_kernel[grid]( + 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, + BLOCK_SIZE_K=BK, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + num_warps=8, + num_stages=2, + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + 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.jit +def _grouped_mxfp4_variable_k_gemm_kernel( + LHS, + RHS, + C, + LHS_scale, + RHS_scale, + go_pad_ptr, + G, + OUT_M, + OUT_N, + 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, +): + 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 + + c = acc.to(C.type.element_ty) + 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 + tl.store(C_, c, cmask) + + +@_scoped_amd_knobs +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 +): + """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). + """ + 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, BK = 256, 128, 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 + _grouped_mxfp4_variable_k_gemm_kernel[(num_sms,)]( + l_u8, + r_u8, + c, + ls, + rs, + go_pad, + G, + OUT_M, + OUT_N, + 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, + BLOCK_SIZE_K=BK, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + num_warps=8, + num_stages=3, + waves_per_eu=2, + matrix_instr_nonkdim=16, + kpack=1, + ) + return c + + +# ########################################################################### +# Output padding tail +# ########################################################################### + + +@triton.jit +def _grouped_gemm_output_tail_kernel( + C, + group_offs_ptr, + G, + M_total, + N, + stride_cm, + stride_cn, + NUM_SMS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, +): + covered = tl.load(group_offs_ptr + G).to(tl.int64) + m_total = tl.cast(M_total, tl.int64) + num_pad = m_total - covered + if num_pad <= 0: + return + + num_m_blocks = tl.cdiv(num_pad, BLOCK_SIZE_M) + num_n_blocks = tl.cdiv(N, BLOCK_SIZE_N) + total_tiles = num_m_blocks * num_n_blocks + + zeros = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=C.dtype.element_ty) + + pid = tl.program_id(0) + for tile in range(pid, total_tiles, NUM_SMS): + bm = tile // num_n_blocks + bn = tile % num_n_blocks + rows = covered + bm * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + cols = bn * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + mask = (rows[:, None] < m_total) & (cols[None, :] < N) + C_ = C + rows[:, None] * stride_cm + cols[None, :] * stride_cn + tl.store(C_, zeros, mask) + + +def grouped_gemm_output_tail_kernel( + out: torch.Tensor, + group_offs: torch.Tensor, +) -> torch.Tensor: + """Zero the uncovered padding tail of a grouped GEMM output, in place. + + Clears rows ``[group_offs[-1], out.shape[0])`` -- the rows the persistent + forward/dgrad kernel never writes when the output is over-allocated to + padded token counts. CPU-sync-free (the covered boundary is read on device) + and a no-op when there is no padding. + """ + assert out.ndim == 2, f"expected 2D grouped GEMM output, got {tuple(out.shape)}" + M_total, N = out.shape + G = group_offs.shape[0] - 1 + if G <= 0 or M_total == 0 or N == 0: + return out + + num_sms = torch.cuda.get_device_properties(out.device).multi_processor_count + _grouped_gemm_output_tail_kernel[(num_sms,)]( + out, + group_offs, + G, + M_total, + N, + out.stride(0), + out.stride(1), + NUM_SMS=num_sms, + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=256, + num_warps=4, + ) + return out 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..8aa667ba51 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -0,0 +1,270 @@ +# 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). + +This is the glue between TransformerEngine's MXFP4 quantizer and the ported +Triton block-scaled grouped kernels in ``grouped_gemm_mxfp4.py``. It reuses +TE's :class:`MXFP4Quantizer` (rather than porting Primus-Turbo's fused grouped +dual-quant) to produce the packed FP4 operands the kernels consume, then builds +the per-group offsets and calls the kernels. + +Three grouped ops, paired exactly like the dense MXFP4 recipe (all NT): + + fprop : C = A_row @ W_row^T (contract K, fwd kernel) + dgrad : dA = gradO_row @ W_col^T (contract N, fwd kernel) + wgrad : dW = gradO_col @ A_col^T (contract M, variable-K kernel) + +Operand layout the kernels expect (plain, un-swizzled OCP MXFP4): + * data : E2M1 packed two-per-byte along the contraction axis (``uint8``); + * scale : one E8M0 (``uint8``) per 1x32 logical-element block, contiguous + along the block axis. +TE's ``MXFP4Quantizer`` emits exactly this in its row-wise ``_rowwise_data`` / +``_rowwise_scale_inv`` and col-wise ``_columnwise_data`` / ``_columnwise_scale_inv`` +members when the shuffle / GEMM-swizzle flags are all off. The scale buffers are +over-allocated (rows to 256, block axis to 8); we slice the live ``[:, :F/32]`` +region and make it contiguous so the kernel sees tight strides. + +RHT (Hadamard) on the wgrad operands and the ``main_grad`` beta=1 accumulate +fusion are deferred to the autograd layer; this module runs plain +MX (``use_hadamard=False``) and returns a fresh weight-gradient tensor. +""" + +from __future__ import annotations + +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 ( + 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 +# 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 _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, Feat = x.shape + q = _quantizer(rowwise=True, columnwise=False).quantize(x.contiguous()) + data = q._rowwise_data.view(torch.uint8) + 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. Columnwise-only quant is a less-exercised + config, so we take the safe route here. + """ + M, Feat = x.shape + q = _quantizer(rowwise=True, columnwise=True).quantize(x.contiguous()) + data = q._columnwise_data.view(torch.uint8) + scale = q._columnwise_scale_inv[:Feat, : M // MXFP4_BLOCK].contiguous() + return data, scale + + +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 BLOCK_SIZE_K. + + ``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, BLOCK_SIZE_K) + 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, +) -> 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). + + Returns: + [total_M, N] output in ``out_dtype``. + """ + K = a.shape[1] + N = weights[0].shape[0] + _check_contract(K, "K") + + a_data, a_scale = _row_operand(a) # (total_M, K/2), (total_M, K/32) + 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) + + 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_mxfp4_dgrad( + grad_out: torch.Tensor, + weights: List[torch.Tensor], + m_splits: Sequence[int], + *, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = 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]. + + Returns: + [total_M, K] input gradient in ``out_dtype``. + """ + N = grad_out.shape[1] + K = weights[0].shape[1] + _check_contract(N, "N") + + go_data, go_scale = _row_operand(grad_out) # (total_M, N/2), (total_M, N/32) + 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) + + 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, +) -> 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``. + + Returns: + [G, N, K] weight gradient in ``out_dtype``. + """ + N = grad_out.shape[1] + K = a.shape[1] + G = len(m_splits) + + 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) + + 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, + ) From e586a623811e14a04956b02940067b58f8474d72 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 31 Aug 2026 13:57:37 -0500 Subject: [PATCH 02/11] add tests --- tests/pytorch/test_grouped_gemm_mxfp4.py | 188 ++++++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4.py | 4 +- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 63 ++++++ 3 files changed, 253 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/test_grouped_gemm_mxfp4.py diff --git a/tests/pytorch/test_grouped_gemm_mxfp4.py b/tests/pytorch/test_grouped_gemm_mxfp4.py new file mode 100644 index 0000000000..18984c3c19 --- /dev/null +++ b/tests/pytorch/test_grouped_gemm_mxfp4.py @@ -0,0 +1,188 @@ +# 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). + +These validate the two layout assumptions of the port: that TE's plain +E8M0/E2M1 quantizer output is interpreted the same way by +``tl.dot_scaled(..., "e2m1", ...)``, and that zero-padded rows contribute +nothing to the variable-K (wgrad) contraction. + +Each op is checked against a *precise* reference: the same packed MXFP4 operands +the kernel consumes are dequantized here by an independent OCP E2M1/E8M0 decoder +(``_dequant_mxfp4``) and matmul'd in fp32. Kernel and reference start from +identical fp4 values, so a correct kernel matches to ~bf16-rounding while a +layout / nibble-order / scale-bias / transpose bug makes them disagree grossly. +wgrad additionally checks against the true bf16 grouped matmul, which exercises +the per-group zero-padding (padded rows must contribute nothing). +""" + +import pytest +import torch + +triton = pytest.importorskip("triton") + +try: + from transformer_engine.pytorch.quantization import check_mxfp4_support + + _MXFP4_OK, _MXFP4_REASON = check_mxfp4_support() +except Exception as exc: # pragma: no cover - import/support probe + _MXFP4_OK, _MXFP4_REASON = False, str(exc) + +from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( + MXFP4_BLOCK, + _col_operand, + _col_operand_grouped_padded, + _row_operand, + grouped_gemm_mxfp4_dgrad, + grouped_gemm_mxfp4_fprop, + grouped_gemm_mxfp4_wgrad, + grouped_linear_mxfp4, +) + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm device required"), + 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] +N, K = 256, 128 +# 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) + + +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_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_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) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 0f9764c0ca..f05652cb23 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -125,11 +125,11 @@ def wrapper(*args, **kwargs): # =============================================================================== -# Chiplet transform -- map the AMD round-robin program_id onto XCD-contiguous +# 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 # MI300 / MI350 chiplet count +NUM_XCDS = 8 # gfx942/gfx950 chiplet count @triton.jit diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py index 8aa667ba51..31c228bb85 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -40,6 +40,7 @@ from ..tensor.mxfp4_tensor import MXFP4Quantizer from ..utils import round_up_to_nearest_multiple from .grouped_gemm_mxfp4 import ( + _is_gfx950, grouped_gemm_mxfp4_triton_kernel, grouped_gemm_mxfp4_variable_k_triton_kernel, ) @@ -78,6 +79,15 @@ def _check_contract(dim: int, name: str) -> None: ) +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, Feat = x.shape @@ -153,6 +163,7 @@ def grouped_gemm_mxfp4_fprop( """ K = a.shape[1] N = weights[0].shape[0] + _require_gfx950() _check_contract(K, "K") a_data, a_scale = _row_operand(a) # (total_M, K/2), (total_M, K/32) @@ -202,6 +213,7 @@ def grouped_gemm_mxfp4_dgrad( """ N = grad_out.shape[1] K = weights[0].shape[1] + _require_gfx950() _check_contract(N, "N") go_data, go_scale = _row_operand(grad_out) # (total_M, N/2), (total_M, N/32) @@ -252,6 +264,7 @@ def grouped_gemm_mxfp4_wgrad( N = grad_out.shape[1] K = a.shape[1] G = len(m_splits) + _require_gfx950() 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) @@ -268,3 +281,53 @@ def grouped_gemm_mxfp4_wgrad( out_dtype=out_dtype, num_cu=num_cu, ) + + +class _GroupedGemmMXFP4Func(torch.autograd.Function): + """Autograd for grouped MXFP4 linear: ``out[g] = a[g] @ weight[g]^T``. + + Forward/backward pair the MXFP4 recipe exactly (fprop / dgrad / wgrad). Each + op quantizes its own operands independently; sharing the quantization across + the three passes is a future optimization. + """ + + @staticmethod + def forward(ctx, a, weight, m_splits): + out = grouped_gemm_mxfp4_fprop( + a, list(torch.unbind(weight, dim=0)), m_splits, out_dtype=a.dtype + ) + ctx.save_for_backward(a, weight) + ctx.m_splits = m_splits + return out + + @staticmethod + def backward(ctx, grad_out): + a, weight = 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, list(torch.unbind(weight, dim=0)), m_splits, out_dtype=a.dtype + ) + 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)) From 4cd50f0702ea6134c4594bb25f86e7525b4d362e Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 10 Sep 2026 17:24:49 -0500 Subject: [PATCH 03/11] update copyright years --- transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index f05652cb23..9b14e1a554 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. # License for AMD contributions = MIT. See LICENSE for more information """Grouped MXFP4 (MX_BLOCKWISE) GEMM Triton persistent kernels (gfx950). From 24ae165560dcd3c972e2f7a783cf1c64fd899bb9 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 10 Sep 2026 17:49:27 -0500 Subject: [PATCH 04/11] add accumulate support --- tests/pytorch/test_grouped_gemm_mxfp4.py | 15 +++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4.py | 17 ++++++++++++++--- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 6 ++++++ 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_grouped_gemm_mxfp4.py b/tests/pytorch/test_grouped_gemm_mxfp4.py index 18984c3c19..6c5bf29aff 100644 --- a/tests/pytorch/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/test_grouped_gemm_mxfp4.py @@ -168,6 +168,21 @@ def test_wgrad_precise_and_padding(): 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) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 9b14e1a554..a3fc8d8f75 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -397,6 +397,7 @@ def _grouped_mxfp4_variable_k_gemm_kernel( CHUNK_SIZE: tl.constexpr, CACHE_MODIFIER: tl.constexpr, VEC: tl.constexpr, + ACCUMULATE: tl.constexpr, ): pid = tl.program_id(0) if NUM_XCDS != 1: @@ -457,24 +458,33 @@ def _grouped_mxfp4_variable_k_gemm_kernel( LS_BASE += BK_SCALE * stride_lsk RS_BASE += BK_SCALE * stride_rsk - c = acc.to(C.type.element_ty) 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 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 + lhs, lhs_scale, rhs, rhs_scale, go_pad, OUT_M, OUT_N, G, + out_dtype=torch.bfloat16, num_cu=None, out=None, accumulate=False, ): """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). """ - c = torch.empty((G, OUT_M, OUT_N), dtype=out_dtype, device=lhs.device) + if out is not None: + 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) @@ -517,6 +527,7 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, + ACCUMULATE=accumulate, num_warps=8, num_stages=3, waves_per_eu=2, diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py index 31c228bb85..e48f02f377 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -248,6 +248,8 @@ def grouped_gemm_mxfp4_wgrad( *, 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). @@ -257,6 +259,8 @@ def grouped_gemm_mxfp4_wgrad( 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``. @@ -280,6 +284,8 @@ def grouped_gemm_mxfp4_wgrad( G, out_dtype=out_dtype, num_cu=num_cu, + out=out, + accumulate=accumulate, ) From ccff4b0c79902449213f4f6a39deb46acc22e3a3 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Sep 2026 10:51:31 -0500 Subject: [PATCH 05/11] address review comments --- tests/pytorch/test_grouped_gemm_mxfp4.py | 18 ++++ .../triton_kernels/grouped_gemm_mxfp4.py | 94 +++---------------- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 46 ++++++++- 3 files changed, 73 insertions(+), 85 deletions(-) diff --git a/tests/pytorch/test_grouped_gemm_mxfp4.py b/tests/pytorch/test_grouped_gemm_mxfp4.py index 6c5bf29aff..7b446ac268 100644 --- a/tests/pytorch/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/test_grouped_gemm_mxfp4.py @@ -116,6 +116,24 @@ def test_fprop_precise(): 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_dgrad_precise(): total_m = sum(M_SPLITS) grad_out = _rand(total_m, N) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index a3fc8d8f75..149dafda67 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -7,8 +7,8 @@ (https://github.com/AMD-AGI/Primus-Turbo, MIT) -- the block-scaled grouped FP4 kernels in ``primus_turbo/triton/grouped_gemm/grouped_gemm_fp4_kernel.py``. The kernel bodies are kept faithful to the source; the shared scaffolding -(``NUM_XCDS``, ``_chiplet_transform_chunked``, the AMD Triton-knob scope, the -output padding-tail pass) is inlined here so the module has no Primus-Turbo +(``NUM_XCDS``, ``_chiplet_transform_chunked``, the AMD Triton-knob scope) is +inlined here so the module has no Primus-Turbo dependency. The operands are E2M1 FP4 packed two-values-per-byte along the contraction (K) @@ -49,12 +49,17 @@ # =============================================================================== -@functools.lru_cache -def _is_gfx950() -> bool: - props = torch.cuda.get_device_properties(torch.cuda.current_device()) +@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 # @@ -481,6 +486,10 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( ``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" @@ -535,78 +544,3 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( kpack=1, ) return c - - -# ########################################################################### -# Output padding tail -# ########################################################################### - - -@triton.jit -def _grouped_gemm_output_tail_kernel( - C, - group_offs_ptr, - G, - M_total, - N, - stride_cm, - stride_cn, - NUM_SMS: tl.constexpr, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, -): - covered = tl.load(group_offs_ptr + G).to(tl.int64) - m_total = tl.cast(M_total, tl.int64) - num_pad = m_total - covered - if num_pad <= 0: - return - - num_m_blocks = tl.cdiv(num_pad, BLOCK_SIZE_M) - num_n_blocks = tl.cdiv(N, BLOCK_SIZE_N) - total_tiles = num_m_blocks * num_n_blocks - - zeros = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=C.dtype.element_ty) - - pid = tl.program_id(0) - for tile in range(pid, total_tiles, NUM_SMS): - bm = tile // num_n_blocks - bn = tile % num_n_blocks - rows = covered + bm * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) - cols = bn * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - mask = (rows[:, None] < m_total) & (cols[None, :] < N) - C_ = C + rows[:, None] * stride_cm + cols[None, :] * stride_cn - tl.store(C_, zeros, mask) - - -def grouped_gemm_output_tail_kernel( - out: torch.Tensor, - group_offs: torch.Tensor, -) -> torch.Tensor: - """Zero the uncovered padding tail of a grouped GEMM output, in place. - - Clears rows ``[group_offs[-1], out.shape[0])`` -- the rows the persistent - forward/dgrad kernel never writes when the output is over-allocated to - padded token counts. CPU-sync-free (the covered boundary is read on device) - and a no-op when there is no padding. - """ - assert out.ndim == 2, f"expected 2D grouped GEMM output, got {tuple(out.shape)}" - M_total, N = out.shape - G = group_offs.shape[0] - 1 - if G <= 0 or M_total == 0 or N == 0: - return out - - num_sms = torch.cuda.get_device_properties(out.device).multi_processor_count - _grouped_gemm_output_tail_kernel[(num_sms,)]( - out, - group_offs, - G, - M_total, - N, - out.stride(0), - out.stride(1), - NUM_SMS=num_sms, - BLOCK_SIZE_M=64, - BLOCK_SIZE_N=256, - num_warps=4, - ) - return out diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py index e48f02f377..8dfae398dc 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -79,6 +79,26 @@ def _check_contract(dim: int, name: str) -> None: ) +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(): @@ -89,10 +109,17 @@ def _require_gfx950() -> None: 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).""" + """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) + data = q._rowwise_data.view(torch.uint8)[:M] scale = q._rowwise_scale_inv[:M, : Feat // MXFP4_BLOCK].contiguous() return data, scale @@ -101,12 +128,16 @@ 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. Columnwise-only quant is a less-exercised - config, so we take the safe route here. + 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) + data = q._columnwise_data.view(torch.uint8)[:Feat, : M // 2] scale = q._columnwise_scale_inv[:Feat, : M // MXFP4_BLOCK].contiguous() return data, scale @@ -165,6 +196,7 @@ def grouped_gemm_mxfp4_fprop( N = weights[0].shape[0] _require_gfx950() _check_contract(K, "K") + _check_splits(m_splits, a.shape[0], len(weights)) a_data, a_scale = _row_operand(a) # (total_M, K/2), (total_M, K/32) b_datas, b_scales = [], [] @@ -215,6 +247,7 @@ def grouped_gemm_mxfp4_dgrad( K = weights[0].shape[1] _require_gfx950() _check_contract(N, "N") + _check_splits(m_splits, grad_out.shape[0], len(weights)) go_data, go_scale = _row_operand(grad_out) # (total_M, N/2), (total_M, N/32) b_datas, b_scales = [], [] @@ -269,6 +302,9 @@ def grouped_gemm_mxfp4_wgrad( 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) From 89bc588dbf56db80960e463e2ffcdc0c1a127ce5 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Sep 2026 14:58:02 -0500 Subject: [PATCH 06/11] add autotune --- .../triton_kernels/grouped_gemm_mxfp4.py | 273 +++++++++++++----- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 137 +++++++-- 2 files changed, 302 insertions(+), 108 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 149dafda67..cfc0a23a18 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -37,6 +37,7 @@ import contextlib import functools import os +import sys from typing import Optional import torch @@ -284,6 +285,81 @@ def _grouped_mxfp4_persistent_gemm_kernel( tl.store(C_, c, c_mask) +# ########################################################################### +# Optional launch-param autotuning (env-gated; default off = fixed configs). +# NVTE_MXFP4_AUTOTUNE=1 sweeps a few (BLOCK_K, num_warps, num_stages, +# waves_per_eu) configs once per shape and caches the fastest (measured with +# triton.testing.do_bench). BLOCK_K is a reduction-loop tile, so varying it does +# not change the grid. (kpack is deprecated on gfx950 and coerced to 1.) +# ########################################################################### + +_MXFP4_AUTOTUNE = bool(int(os.environ.get("NVTE_MXFP4_AUTOTUNE", "0"))) +# When set, print the winning (BLOCK_K, num_warps, num_stages, waves_per_eu) per shape. +_MXFP4_AUTOTUNE_VERBOSE = bool(int(os.environ.get("NVTE_MXFP4_AUTOTUNE_VERBOSE", "0"))) + +# (BLOCK_K, num_warps, num_stages, waves_per_eu); index 0 = current default. +_FWD_CONFIGS = [ + (128, 8, 2, 0), + (128, 8, 3, 2), + (128, 8, 4, 2), + (128, 4, 2, 0), + (256, 8, 2, 0), + (256, 8, 3, 2), +] +_VARK_CONFIGS = [ + (128, 8, 3, 2), + (128, 8, 2, 0), + (128, 8, 4, 2), + (128, 4, 3, 2), + (128, 8, 3, 0), + (256, 8, 2, 0), + (256, 8, 3, 2), +] +_fwd_tune_cache: dict = {} +_vark_tune_cache: dict = {} + + +def _resolve_config(cache, configs, key, launch, valid, label=""): + """Return the fastest valid config for ``key``, tuning ``launch`` once if needed. + + ``launch(cfg)`` runs the kernel into a benchmarking buffer; ``valid(cfg)`` + filters configs that don't fit the shape. Autotune off (or a cache hit) returns + without benchmarking; the default is ``configs[0]`` (unchanged behavior). + """ + if not _MXFP4_AUTOTUNE: + return configs[0] + best = cache.get(key) + if best is None: + import triton.testing # only needed when tuning + + best_t = float("inf") + best = configs[0] + for cfg_try in configs: + if not valid(cfg_try): + continue + try: + t = float(triton.testing.do_bench(lambda cfg=cfg_try: launch(cfg))) + except Exception: # a config the compiler rejects -> skip it + continue + if t < best_t: + best_t, best = t, cfg_try + cache[key] = best + if _MXFP4_AUTOTUNE_VERBOSE: + print( + f"[mxfp4-autotune] {label} shape={key} -> " + f"(BLOCK_K, num_warps, num_stages, waves_per_eu)={best} " + f"({best_t:.4f} ms)", + file=sys.stderr, + flush=True, + ) + return best + + +def _autotune_launch(cache, configs, key, launch, valid, label=""): + """Tune (if enabled) then launch the winning config into the caller's output.""" + launch(_resolve_config(cache, configs, key, launch, valid, label)) + + @_scoped_amd_knobs def grouped_gemm_mxfp4_triton_kernel( a, @@ -312,7 +388,7 @@ def grouped_gemm_mxfp4_triton_kernel( 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, BK = 256, 256, 128 + BM, BN = 256, 256 m_alloc = a.shape[0] avg_m = max(m_alloc // max(G, 1), 1) tiles_n = (N + BN - 1) // BN @@ -320,44 +396,55 @@ def grouped_gemm_mxfp4_triton_kernel( 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 - grid = (num_sms,) - _grouped_mxfp4_persistent_gemm_kernel[grid]( - 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, - BLOCK_SIZE_K=BK, - GROUP_SIZE_M=GM, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=chunk, - CACHE_MODIFIER=".ca", - VEC=VEC_SIZE, - num_warps=8, - num_stages=2, - waves_per_eu=0, - matrix_instr_nonkdim=16, - kpack=1, + + def _launch(cfg): + bk, nw, ns, wpe = cfg + _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, + BLOCK_SIZE_K=bk, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + num_warps=nw, + num_stages=ns, + waves_per_eu=wpe, + matrix_instr_nonkdim=16, + kpack=1, + ) + + _autotune_launch( + _fwd_tune_cache, + _FWD_CONFIGS, + (m_alloc, N, K, G), + _launch, + lambda cfg: K % cfg[0] == 0, + label="fwd", ) return c @@ -499,48 +586,80 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( 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, BK = 256, 128, 128 + 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 - _grouped_mxfp4_variable_k_gemm_kernel[(num_sms,)]( - l_u8, - r_u8, - c, - ls, - rs, - go_pad, - G, - OUT_M, - OUT_N, - 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, - BLOCK_SIZE_K=BK, - GROUP_SIZE_M=GM, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=chunk, - CACHE_MODIFIER=".ca", - VEC=VEC_SIZE, - ACCUMULATE=accumulate, - num_warps=8, - num_stages=3, - waves_per_eu=2, - matrix_instr_nonkdim=16, - kpack=1, - ) + + def _launch(cfg, c_out=c): + bk, nw, ns, wpe = cfg + _grouped_mxfp4_variable_k_gemm_kernel[(num_sms,)]( + l_u8, + r_u8, + c_out, + ls, + rs, + go_pad, + G, + OUT_M, + OUT_N, + l_u8.stride(0), + l_u8.stride(1), + r_u8.stride(0), + r_u8.stride(1), + c_out.stride(0), + c_out.stride(1), + c_out.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, + BLOCK_SIZE_K=bk, + GROUP_SIZE_M=GM, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CHUNK_SIZE=chunk, + CACHE_MODIFIER=".ca", + VEC=VEC_SIZE, + ACCUMULATE=accumulate, + num_warps=nw, + num_stages=ns, + waves_per_eu=wpe, + matrix_instr_nonkdim=16, + kpack=1, + ) + + # A vark config is valid only if its BLOCK_K divides every padded group length + # (the reduction loop is unmasked); guard 256 against 128-only padding. + if _MXFP4_AUTOTUNE: + _lens = [b - a for a, b in zip(go_pad.tolist(), go_pad[1:].tolist())] + valid = lambda cfg: all(n % cfg[0] == 0 for n in _lens) + else: + valid = lambda cfg: True + key = (OUT_M, OUT_N, G, lhs.shape[1]) + + if out is None: + # Fresh output: safe for do_bench to overwrite it directly. + _autotune_launch(_vark_tune_cache, _VARK_CONFIGS, key, _launch, valid, label="wgrad") + else: + # Caller buffer (accumulate into e.g. main_grad): do_bench would re-run the + # beta=1 kernel into `out` every warmup/timing rep, so tune on a scratch + # buffer (never during CUDA-graph capture), then do the real accumulate into + # `out` (primus-turbo grouped_gemm_fp4_variable_k_accum_impl pattern). + cfg = _VARK_CONFIGS[0] + if _MXFP4_AUTOTUNE: + if key in _vark_tune_cache: + cfg = _vark_tune_cache[key] + elif not torch.cuda.is_current_stream_capturing(): + scratch = torch.zeros_like(out) + cfg = _resolve_config( + _vark_tune_cache, _VARK_CONFIGS, key, + lambda cfg_try: _launch(cfg_try, scratch), valid, label="wgrad", + ) + _launch(cfg, out) 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 index 8dfae398dc..a891e8f89f 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -32,6 +32,7 @@ from __future__ import annotations +import os from typing import List, Optional, Sequence, Tuple import torch @@ -47,6 +48,10 @@ # 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 @@ -142,11 +147,56 @@ def _col_operand(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 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 _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 BLOCK_SIZE_K. + """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 @@ -161,7 +211,7 @@ def _col_operand_grouped_padded( m = int(m) xg = x[start : start + m] start += m - m_pad = round_up_to_nearest_multiple(m, BLOCK_SIZE_K) + 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) @@ -181,6 +231,7 @@ def grouped_gemm_mxfp4_fprop( *, 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). @@ -188,25 +239,33 @@ def grouped_gemm_mxfp4_fprop( 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] - N = weights[0].shape[0] _require_gfx950() _check_contract(K, "K") - _check_splits(m_splits, a.shape[0], len(weights)) - a_data, a_scale = _row_operand(a) # (total_M, K/2), (total_M, K/32) - 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) + 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, @@ -224,11 +283,12 @@ def grouped_gemm_mxfp4_fprop( def grouped_gemm_mxfp4_dgrad( grad_out: torch.Tensor, - weights: List[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). @@ -238,26 +298,36 @@ def grouped_gemm_mxfp4_dgrad( 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]. + 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] - K = weights[0].shape[1] _require_gfx950() _check_contract(N, "N") - _check_splits(m_splits, grad_out.shape[0], len(weights)) - go_data, go_scale = _row_operand(grad_out) # (total_M, N/2), (total_M, N/32) - 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) + 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( @@ -328,29 +398,34 @@ def grouped_gemm_mxfp4_wgrad( class _GroupedGemmMXFP4Func(torch.autograd.Function): """Autograd for grouped MXFP4 linear: ``out[g] = a[g] @ weight[g]^T``. - Forward/backward pair the MXFP4 recipe exactly (fprop / dgrad / wgrad). Each - op quantizes its own operands independently; sharing the quantization across - the three passes is a future optimization. + 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, list(torch.unbind(weight, dim=0)), m_splits, out_dtype=a.dtype + a, weights, m_splits, out_dtype=a.dtype, weight_row=w_row ) - ctx.save_for_backward(a, weight) + 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 = ctx.saved_tensors + 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, list(torch.unbind(weight, dim=0)), m_splits, out_dtype=a.dtype + 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) From 8c89493ee4190099e024ee4faddc110bf3c16f36 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Sep 2026 15:57:42 -0500 Subject: [PATCH 07/11] use triton.autotune --- .../triton_kernels/grouped_gemm_mxfp4.py | 335 ++++++++---------- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 8 + 2 files changed, 151 insertions(+), 192 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index cfc0a23a18..4a1a776520 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -37,7 +37,6 @@ import contextlib import functools import os -import sys from typing import Optional import torch @@ -160,6 +159,62 @@ def _chiplet_transform_chunked( VEC_SIZE = 32 +# ########################################################################### +# Launch-param autotuning. Each shape is swept over a few (BLOCK_K, num_warps, +# num_stages, waves_per_eu) tiles via @triton.autotune and the fastest is cached. +# Set TRITON_PRINT_AUTOTUNING=1 to log the winner. BLOCK_K is a reduction-loop +# tile, so sweeping it doesn't change the grid. +# ########################################################################### + +# (BLOCK_K, num_warps, num_stages, waves_per_eu). +_FWD_CONFIGS = [ + (128, 8, 2, 0), + (128, 8, 3, 2), + (128, 8, 4, 2), + (128, 4, 2, 0), + (256, 8, 2, 0), + (256, 8, 3, 2), +] +_VARK_CONFIGS = [ + (128, 8, 3, 2), + (128, 8, 2, 0), + (128, 8, 4, 2), + (128, 4, 3, 2), + (128, 8, 3, 0), + (256, 8, 2, 0), + (256, 8, 3, 2), +] + + +def _make_configs(specs): + """Build a ``triton.Config`` list from (BLOCK_K, warps, stages, waves) tuples. + + ``kpack=1`` is deprecated on gfx950; ``matrix_instr_nonkdim=16`` pins the + 16x16 scaled-MFMA. + """ + return [ + triton.Config( + {"BLOCK_SIZE_K": bk, "waves_per_eu": wpe, "matrix_instr_nonkdim": 16, "kpack": 1}, + num_warps=nw, + num_stages=ns, + ) + for (bk, nw, ns, 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 @@ -167,6 +222,11 @@ def _chiplet_transform_chunked( # ########################################################################### +@triton.autotune( + configs=_make_configs(_FWD_CONFIGS), + key=["N", "K", "G"], + prune_configs_by={"early_config_prune": _prune_fwd_by_k}, +) @triton.jit def _grouped_mxfp4_persistent_gemm_kernel( A, @@ -285,81 +345,6 @@ def _grouped_mxfp4_persistent_gemm_kernel( tl.store(C_, c, c_mask) -# ########################################################################### -# Optional launch-param autotuning (env-gated; default off = fixed configs). -# NVTE_MXFP4_AUTOTUNE=1 sweeps a few (BLOCK_K, num_warps, num_stages, -# waves_per_eu) configs once per shape and caches the fastest (measured with -# triton.testing.do_bench). BLOCK_K is a reduction-loop tile, so varying it does -# not change the grid. (kpack is deprecated on gfx950 and coerced to 1.) -# ########################################################################### - -_MXFP4_AUTOTUNE = bool(int(os.environ.get("NVTE_MXFP4_AUTOTUNE", "0"))) -# When set, print the winning (BLOCK_K, num_warps, num_stages, waves_per_eu) per shape. -_MXFP4_AUTOTUNE_VERBOSE = bool(int(os.environ.get("NVTE_MXFP4_AUTOTUNE_VERBOSE", "0"))) - -# (BLOCK_K, num_warps, num_stages, waves_per_eu); index 0 = current default. -_FWD_CONFIGS = [ - (128, 8, 2, 0), - (128, 8, 3, 2), - (128, 8, 4, 2), - (128, 4, 2, 0), - (256, 8, 2, 0), - (256, 8, 3, 2), -] -_VARK_CONFIGS = [ - (128, 8, 3, 2), - (128, 8, 2, 0), - (128, 8, 4, 2), - (128, 4, 3, 2), - (128, 8, 3, 0), - (256, 8, 2, 0), - (256, 8, 3, 2), -] -_fwd_tune_cache: dict = {} -_vark_tune_cache: dict = {} - - -def _resolve_config(cache, configs, key, launch, valid, label=""): - """Return the fastest valid config for ``key``, tuning ``launch`` once if needed. - - ``launch(cfg)`` runs the kernel into a benchmarking buffer; ``valid(cfg)`` - filters configs that don't fit the shape. Autotune off (or a cache hit) returns - without benchmarking; the default is ``configs[0]`` (unchanged behavior). - """ - if not _MXFP4_AUTOTUNE: - return configs[0] - best = cache.get(key) - if best is None: - import triton.testing # only needed when tuning - - best_t = float("inf") - best = configs[0] - for cfg_try in configs: - if not valid(cfg_try): - continue - try: - t = float(triton.testing.do_bench(lambda cfg=cfg_try: launch(cfg))) - except Exception: # a config the compiler rejects -> skip it - continue - if t < best_t: - best_t, best = t, cfg_try - cache[key] = best - if _MXFP4_AUTOTUNE_VERBOSE: - print( - f"[mxfp4-autotune] {label} shape={key} -> " - f"(BLOCK_K, num_warps, num_stages, waves_per_eu)={best} " - f"({best_t:.4f} ms)", - file=sys.stderr, - flush=True, - ) - return best - - -def _autotune_launch(cache, configs, key, launch, valid, label=""): - """Tune (if enabled) then launch the winning config into the caller's output.""" - launch(_resolve_config(cache, configs, key, launch, valid, label)) - - @_scoped_amd_knobs def grouped_gemm_mxfp4_triton_kernel( a, @@ -397,54 +382,40 @@ def grouped_gemm_mxfp4_triton_kernel( num_sms = min(total_tiles, cu) chunk = 64 if num_sms >= NUM_XCDS * 64 else 32 - def _launch(cfg): - bk, nw, ns, wpe = cfg - _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, - BLOCK_SIZE_K=bk, - GROUP_SIZE_M=GM, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=chunk, - CACHE_MODIFIER=".ca", - VEC=VEC_SIZE, - num_warps=nw, - num_stages=ns, - waves_per_eu=wpe, - matrix_instr_nonkdim=16, - kpack=1, - ) - - _autotune_launch( - _fwd_tune_cache, - _FWD_CONFIGS, - (m_alloc, N, K, G), - _launch, - lambda cfg: K % cfg[0] == 0, - label="fwd", + # 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, ) return c @@ -458,6 +429,12 @@ def _launch(cfg): # ########################################################################### +@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, @@ -469,6 +446,7 @@ def _grouped_mxfp4_variable_k_gemm_kernel( 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, @@ -563,7 +541,7 @@ def _grouped_mxfp4_variable_k_gemm_kernel( @_scoped_amd_knobs 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, + out_dtype=torch.bfloat16, num_cu=None, out=None, accumulate=False, can256=None, ): """C[g] (OUT_M, OUT_N) = lhs[:,g] @ rhs[:,g]^T. @@ -594,72 +572,45 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( num_sms = min(total_tiles, cu) chunk = 64 if num_sms >= NUM_XCDS * 64 else 32 - def _launch(cfg, c_out=c): - bk, nw, ns, wpe = cfg - _grouped_mxfp4_variable_k_gemm_kernel[(num_sms,)]( - l_u8, - r_u8, - c_out, - ls, - rs, - go_pad, - G, - OUT_M, - OUT_N, - l_u8.stride(0), - l_u8.stride(1), - r_u8.stride(0), - r_u8.stride(1), - c_out.stride(0), - c_out.stride(1), - c_out.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, - BLOCK_SIZE_K=bk, - GROUP_SIZE_M=GM, - NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, - CHUNK_SIZE=chunk, - CACHE_MODIFIER=".ca", - VEC=VEC_SIZE, - ACCUMULATE=accumulate, - num_warps=nw, - num_stages=ns, - waves_per_eu=wpe, - matrix_instr_nonkdim=16, - kpack=1, - ) - - # A vark config is valid only if its BLOCK_K divides every padded group length - # (the reduction loop is unmasked); guard 256 against 128-only padding. - if _MXFP4_AUTOTUNE: + # 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())] - valid = lambda cfg: all(n % cfg[0] == 0 for n in _lens) - else: - valid = lambda cfg: True - key = (OUT_M, OUT_N, G, lhs.shape[1]) - - if out is None: - # Fresh output: safe for do_bench to overwrite it directly. - _autotune_launch(_vark_tune_cache, _VARK_CONFIGS, key, _launch, valid, label="wgrad") - else: - # Caller buffer (accumulate into e.g. main_grad): do_bench would re-run the - # beta=1 kernel into `out` every warmup/timing rep, so tune on a scratch - # buffer (never during CUDA-graph capture), then do the real accumulate into - # `out` (primus-turbo grouped_gemm_fp4_variable_k_accum_impl pattern). - cfg = _VARK_CONFIGS[0] - if _MXFP4_AUTOTUNE: - if key in _vark_tune_cache: - cfg = _vark_tune_cache[key] - elif not torch.cuda.is_current_stream_capturing(): - scratch = torch.zeros_like(out) - cfg = _resolve_config( - _vark_tune_cache, _VARK_CONFIGS, key, - lambda cfg_try: _launch(cfg_try, scratch), valid, label="wgrad", - ) - _launch(cfg, out) + 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 index a891e8f89f..f978e9eddd 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -379,6 +379,13 @@ def grouped_gemm_mxfp4_wgrad( 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, @@ -392,6 +399,7 @@ def grouped_gemm_mxfp4_wgrad( num_cu=num_cu, out=out, accumulate=accumulate, + can256=can256, ) From d710e0ce927793d45b41f4dff4c1a0c923f37198 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Sep 2026 16:02:23 -0500 Subject: [PATCH 08/11] fix comments --- .../pytorch/triton_kernels/grouped_gemm_mxfp4.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 4a1a776520..b8fd3ab010 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -27,9 +27,11 @@ loads a (K/32, free) tile (coalesced) and transposes it back in-reg for tl.dot_scaled. -Only EVEN contraction (a multiple of BLOCK_SIZE_K = 128 logical elements) is -supported; the FP4 quantizer pads K to 128 and the MX wrapper pads per-group M -to 128, so this always holds in practice. +The reduction loop is unmasked, so the contraction must be a multiple of the +autotuned BLOCK_K (128 or 256). This always holds: fprop/dgrad require the +contraction K to be a multiple of 128 (enforced in the impl), the wgrad per-group +M is padded to a 128-multiple, and config pruning only selects BLOCK_K=256 when +the contraction is also 256-divisible. """ from __future__ import annotations @@ -362,7 +364,7 @@ def grouped_gemm_mxfp4_triton_kernel( 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 BLOCK_SIZE_K=128). + K is the logical contraction size (must be a multiple of 128). """ if group_offs_out is None: group_offs_out = group_offs From 431b2c7d0ccd83ffc244eed29dfdb3a140105682 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 17 Sep 2026 12:04:23 -0500 Subject: [PATCH 09/11] add a8w4 variant --- tests/pytorch/test_grouped_gemm_mxfp4.py | 89 +++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4.py | 107 ++++++++++++++++-- 2 files changed, 186 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_grouped_gemm_mxfp4.py b/tests/pytorch/test_grouped_gemm_mxfp4.py index 7b446ac268..850ce25dad 100644 --- a/tests/pytorch/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/test_grouped_gemm_mxfp4.py @@ -33,12 +33,16 @@ MXFP4_BLOCK, _col_operand, _col_operand_grouped_padded, + _prefix_offsets, _row_operand, grouped_gemm_mxfp4_dgrad, grouped_gemm_mxfp4_fprop, grouped_gemm_mxfp4_wgrad, grouped_linear_mxfp4, ) +from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4 import ( + grouped_gemm_a8w4_triton_kernel, +) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm device required"), @@ -96,6 +100,26 @@ def _dequant_mxfp4(data_u8, scale_u8, feat): return vals * scale.repeat_interleave(MXFP4_BLOCK, dim=1) +# MXFP8 (e4m3) activation side, for the a8w4 forward +def _quant_mxfp8_e4m3(x): + """OCP MXFP8 row-wise: ``x`` [R, K] -> (e4m3 [R, K], e8m0 u8 [R, K/32]).""" + fmax = torch.finfo(torch.float8_e4m3fn).max + rows, feat = x.shape + xb = x.reshape(rows, feat // 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) + return q.reshape(rows, feat).contiguous(), (exp.squeeze(-1) + 127).to(torch.uint8) + + +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) @@ -219,3 +243,68 @@ def test_autograd_matches_ops(): 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. +def _run_a8w4(a, weights, m_splits): + """Quantize (a->e4m3, weights->e2m1) and run the a8w4 grouped forward kernel.""" + feat_k = a.shape[1] + n = weights[0].shape[0] + a_q, a_s = _quant_mxfp8_e4m3(a) + b_ops = [_row_operand(w) for w in weights] # each -> (e2m1 [N,K/2], e8m0 [N,K/32]) + b_data = torch.stack([d for d, _s in b_ops], dim=0) + b_scale = torch.stack([s for _d, s in b_ops], dim=0) + offs = _prefix_offsets(m_splits, a.device) + out = grouped_gemm_a8w4_triton_kernel(a_q, a_s, b_data, b_scale, offs, n, feat_k, out_dtype=DTYPE) + return out, a_q, a_s + + +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, a_q, a_s = _run_a8w4(a, weights, M_SPLITS) + 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, a_q, a_s = _run_a8w4(a, weights, splits) + 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, a_q, a_s = _run_a8w4(a, weights, splits) + 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 index b8fd3ab010..56a1b38284 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -226,7 +226,7 @@ def _prune_vark_by_pad(configs, named_args, **kwargs): @triton.autotune( configs=_make_configs(_FWD_CONFIGS), - key=["N", "K", "G"], + key=["N", "K", "G", "A_PACK", "B_PACK"], prune_configs_by={"early_config_prune": _prune_fwd_by_k}, ) @triton.jit @@ -262,6 +262,10 @@ def _grouped_mxfp4_persistent_gemm_kernel( 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: @@ -276,7 +280,8 @@ def _grouped_mxfp4_persistent_gemm_kernel( tl.assume(stride_am > 0) tl.assume(stride_cm > 0) - BK_PACK: tl.constexpr = BLOCK_SIZE_K // 2 # packed bytes per K-iter + 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): @@ -311,12 +316,13 @@ def _grouped_mxfp4_persistent_gemm_kernel( 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 = tl.arange(0, BK_PACK) # packed K bytes + 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[None, :] * stride_ak - B_BASE = B + group_idx.to(tl.int64) * stride_bg + rk[:, None] * stride_bk + rn[None, :] * stride_bn + 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 @@ -328,13 +334,13 @@ def _grouped_mxfp4_persistent_gemm_kernel( 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/2) - b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) # (BK/2, BN) + 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, "e2m1", b, b_s, "e2m1", acc) - A_BASE += BK_PACK * stride_ak - B_BASE += BK_PACK * stride_bk + 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 @@ -418,6 +424,87 @@ def grouped_gemm_mxfp4_triton_kernel( 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 +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 From 136a6f541aacfe3496156ebf8402c4f82768340c Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 18 Sep 2026 14:40:13 -0500 Subject: [PATCH 10/11] a8w4 fprop op, more autotune config --- tests/pytorch/test_grouped_gemm_mxfp4.py | 80 ++++++++++-------- .../triton_kernels/grouped_gemm_mxfp4.py | 49 +++++------ .../triton_kernels/grouped_gemm_mxfp4_impl.py | 82 +++++++++++++++++++ 3 files changed, 156 insertions(+), 55 deletions(-) diff --git a/tests/pytorch/test_grouped_gemm_mxfp4.py b/tests/pytorch/test_grouped_gemm_mxfp4.py index 850ce25dad..1b4125c241 100644 --- a/tests/pytorch/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/test_grouped_gemm_mxfp4.py @@ -33,16 +33,14 @@ MXFP4_BLOCK, _col_operand, _col_operand_grouped_padded, - _prefix_offsets, _row_operand, + _row_operand_mxfp8, + grouped_gemm_a8w4_fprop, grouped_gemm_mxfp4_dgrad, grouped_gemm_mxfp4_fprop, grouped_gemm_mxfp4_wgrad, grouped_linear_mxfp4, ) -from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4 import ( - grouped_gemm_a8w4_triton_kernel, -) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm device required"), @@ -100,18 +98,8 @@ def _dequant_mxfp4(data_u8, scale_u8, feat): return vals * scale.repeat_interleave(MXFP4_BLOCK, dim=1) -# MXFP8 (e4m3) activation side, for the a8w4 forward -def _quant_mxfp8_e4m3(x): - """OCP MXFP8 row-wise: ``x`` [R, K] -> (e4m3 [R, K], e8m0 u8 [R, K/32]).""" - fmax = torch.finfo(torch.float8_e4m3fn).max - rows, feat = x.shape - xb = x.reshape(rows, feat // 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) - return q.reshape(rows, feat).contiguous(), (exp.squeeze(-1) + 127).to(torch.uint8) - - +# 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 @@ -246,19 +234,8 @@ def test_autograd_matches_ops(): # a8w4 forward: MXFP8 (e4m3) activation x MXFP4 (e2m1) weight, grouped along M. -def _run_a8w4(a, weights, m_splits): - """Quantize (a->e4m3, weights->e2m1) and run the a8w4 grouped forward kernel.""" - feat_k = a.shape[1] - n = weights[0].shape[0] - a_q, a_s = _quant_mxfp8_e4m3(a) - b_ops = [_row_operand(w) for w in weights] # each -> (e2m1 [N,K/2], e8m0 [N,K/32]) - b_data = torch.stack([d for d, _s in b_ops], dim=0) - b_scale = torch.stack([s for _d, s in b_ops], dim=0) - offs = _prefix_offsets(m_splits, a.device) - out = grouped_gemm_a8w4_triton_kernel(a_q, a_s, b_data, b_scale, offs, n, feat_k, out_dtype=DTYPE) - return out, a_q, a_s - - +# 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) @@ -277,7 +254,8 @@ def test_a8w4_fprop_precise(): a = _rand(total_m, K) weights = [_rand(N, K) for _ in M_SPLITS] - out, a_q, a_s = _run_a8w4(a, weights, 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) @@ -290,7 +268,8 @@ def test_a8w4_fprop_single_group(): a = _rand(512, K) weights = [_rand(N, K)] - out, a_q, a_s = _run_a8w4(a, weights, 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 _rel_err(out, ref) < _TIGHT_TOL @@ -303,8 +282,45 @@ def test_a8w4_fprop_unaligned_total_m(): a = _rand(total_m, K) weights = [_rand(N, K) for _ in splits] - out, a_q, a_s = _run_a8w4(a, weights, 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 index 56a1b38284..178beab677 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -163,44 +163,47 @@ def _chiplet_transform_chunked( # ########################################################################### # Launch-param autotuning. Each shape is swept over a few (BLOCK_K, num_warps, -# num_stages, waves_per_eu) tiles via @triton.autotune and the fastest is cached. -# Set TRITON_PRINT_AUTOTUNING=1 to log the winner. BLOCK_K is a reduction-loop -# tile, so sweeping it doesn't change the grid. +# 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, waves_per_eu). +# (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, 0), - (128, 8, 3, 2), - (128, 8, 4, 2), - (128, 4, 2, 0), - (256, 8, 2, 0), - (256, 8, 3, 2), + (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, 2), - (128, 8, 2, 0), - (128, 8, 4, 2), - (128, 4, 3, 2), - (128, 8, 3, 0), - (256, 8, 2, 0), - (256, 8, 3, 2), + (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, waves) tuples. - - ``kpack=1`` is deprecated on gfx950; ``matrix_instr_nonkdim=16`` pins the - 16x16 scaled-MFMA. + """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": 16, "kpack": 1}, + {"BLOCK_SIZE_K": bk, "waves_per_eu": wpe, "matrix_instr_nonkdim": nonk, "kpack": 1}, num_warps=nw, num_stages=ns, ) - for (bk, nw, ns, wpe) in specs + for (bk, nw, ns, nonk, wpe) in specs ] diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py index f978e9eddd..0b76b8cc3d 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -42,6 +42,7 @@ 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, ) @@ -171,6 +172,27 @@ def _row_col_operand( 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]]: @@ -281,6 +303,66 @@ def grouped_gemm_mxfp4_fprop( ) +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]], From 929ec97f1cd3fac3ec676df9ce329a3c1f9ec830 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 18 Sep 2026 15:19:37 -0500 Subject: [PATCH 11/11] address review comments, add to CI --- ci/pytorch.sh | 1 + .../test_grouped_gemm_mxfp4.py | 53 ++++---- .../triton_kernels/grouped_gemm_mxfp4.py | 117 ++++-------------- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 29 +---- 4 files changed, 56 insertions(+), 144 deletions(-) rename tests/pytorch/{ => triton_kernels}/test_grouped_gemm_mxfp4.py (90%) diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 1d6aef2922..2abd0b410b 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/test_grouped_gemm_mxfp4.py b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py similarity index 90% rename from tests/pytorch/test_grouped_gemm_mxfp4.py rename to tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py index 1b4125c241..f255b6ad0b 100644 --- a/tests/pytorch/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py @@ -1,33 +1,14 @@ # 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). - -These validate the two layout assumptions of the port: that TE's plain -E8M0/E2M1 quantizer output is interpreted the same way by -``tl.dot_scaled(..., "e2m1", ...)``, and that zero-padded rows contribute -nothing to the variable-K (wgrad) contraction. - -Each op is checked against a *precise* reference: the same packed MXFP4 operands -the kernel consumes are dequantized here by an independent OCP E2M1/E8M0 decoder -(``_dequant_mxfp4``) and matmul'd in fp32. Kernel and reference start from -identical fp4 values, so a correct kernel matches to ~bf16-rounding while a -layout / nibble-order / scale-bias / transpose bug makes them disagree grossly. -wgrad additionally checks against the true bf16 grouped matmul, which exercises -the per-group zero-padding (padded rows must contribute nothing). -""" +"""Numeric correctness tests for the grouped MXFP4 Triton GEMM (gfx950).""" import pytest import torch -triton = pytest.importorskip("triton") +from transformer_engine.pytorch.quantization import check_mxfp4_support -try: - from transformer_engine.pytorch.quantization import check_mxfp4_support - - _MXFP4_OK, _MXFP4_REASON = check_mxfp4_support() -except Exception as exc: # pragma: no cover - import/support probe - _MXFP4_OK, _MXFP4_REASON = False, str(exc) +_MXFP4_OK, _MXFP4_REASON = check_mxfp4_support() from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( MXFP4_BLOCK, @@ -43,7 +24,6 @@ ) pytestmark = [ - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm device required"), pytest.mark.skipif(not _MXFP4_OK, reason=f"MXFP4 unsupported: {_MXFP4_REASON}"), ] @@ -51,7 +31,9 @@ # 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] -N, K = 256, 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 @@ -146,6 +128,29 @@ def test_fprop_unaligned_total_m(): 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) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 178beab677..1093813361 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -1,50 +1,26 @@ # Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. # License for AMD contributions = MIT. See LICENSE for more information -"""Grouped MXFP4 (MX_BLOCKWISE) GEMM Triton persistent kernels (gfx950). - -Adapted for TransformerEngine from Primus-Turbo -(https://github.com/AMD-AGI/Primus-Turbo, MIT) -- the block-scaled grouped FP4 -kernels in ``primus_turbo/triton/grouped_gemm/grouped_gemm_fp4_kernel.py``. The -kernel bodies are kept faithful to the source; the shared scaffolding -(``NUM_XCDS``, ``_chiplet_transform_chunked``, the AMD Triton-knob scope) is -inlined here so the module has no Primus-Turbo -dependency. - -The operands are E2M1 FP4 packed two-values-per-byte along the contraction (K) -axis, fed to the hardware block-scaled MMA via ``tl.dot_scaled(..., "e2m1", ...)`` -with E8M0 (VEC_SIZE=32) scales. - - - _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 - -FP4 packing notes: - * data tensors store K/2 (resp. M/2) bytes along the contraction axis; - * scale tensors store K/32 (resp. M/32) E8M0 bytes (one per 1x32 block); - * scales are stored free-major (free, K/32) / (G, free, K/32); each K-iter - loads a (K/32, free) tile (coalesced) and transposes it back in-reg for - tl.dot_scaled. - -The reduction loop is unmasked, so the contraction must be a multiple of the -autotuned BLOCK_K (128 or 256). This always holds: fprop/dgrad require the -contraction K to be a multiple of 128 (enforced in the impl), the wgrad per-group -M is padded to a 128-multiple, and config pruning only selects BLOCK_K=256 when -the contraction is also 256-divisible. +"""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 contextlib import functools -import os -from typing import Optional import torch import triton import triton.language as tl +from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs + # =============================================================================== # Arch helper @@ -71,64 +47,21 @@ def _is_gfx950() -> bool: # read at compile time and are not part of Triton's compile cache key. # =============================================================================== -_AMD_KNOB_ATTRS = ("use_async_copy", "scalarize_packed_fops", "use_block_pingpong") -_AMD_KNOB_ENVS = ( - "TRITON_HIP_USE_ASYNC_COPY", - "AMDGCN_SCALARIZE_PACKED_FOPS", - "TRITON_HIP_USE_BLOCK_PINGPONG", -) - - -def _amd_knobs_available() -> bool: - return hasattr(triton, "knobs") and hasattr(triton.knobs, "amd") +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) -@contextlib.contextmanager -def _scoped_amd_triton_knobs(): - """Snapshot AMD Triton compiler knobs on entry, restore them on exit. + return wrapper - On gfx950 the full gfx950 knob set is enabled for the duration of the scope; - all changes are restored on exit. - """ - saved_attrs = {} - if _amd_knobs_available(): - amd = triton.knobs.amd - for name in _AMD_KNOB_ATTRS: - if hasattr(amd, name): - saved_attrs[name] = getattr(amd, name) - saved_env = {name: os.environ.get(name) for name in _AMD_KNOB_ENVS} - try: - if _is_gfx950(): - if _amd_knobs_available(): - amd = triton.knobs.amd - for name in _AMD_KNOB_ATTRS: - if hasattr(amd, name): - setattr(amd, name, True) - else: - for name in _AMD_KNOB_ENVS: - os.environ[name] = "1" - yield - finally: - if saved_attrs: - amd = triton.knobs.amd - for name, val in saved_attrs.items(): - setattr(amd, name, val) - for name, val in saved_env.items(): - if val is None: - os.environ.pop(name, None) - else: - os.environ[name] = val - - -def _scoped_amd_knobs(func): - """Decorator: run ``func`` under :func:`_scoped_amd_triton_knobs`.""" - - @functools.wraps(func) - def wrapper(*args, **kwargs): - with _scoped_amd_triton_knobs(): - return func(*args, **kwargs) - - return wrapper + return deco # =============================================================================== @@ -356,7 +289,7 @@ def _grouped_mxfp4_persistent_gemm_kernel( tl.store(C_, c, c_mask) -@_scoped_amd_knobs +@_scoped_amd_knobs(is_tn=False) def grouped_gemm_mxfp4_triton_kernel( a, a_scale, @@ -433,7 +366,7 @@ def grouped_gemm_mxfp4_triton_kernel( return c -@_scoped_amd_knobs +@_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 @@ -630,7 +563,7 @@ def _grouped_mxfp4_variable_k_gemm_kernel( tl.store(C_, c, cmask) -@_scoped_amd_knobs +@_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, diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py index 0b76b8cc3d..81702da6e7 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -1,34 +1,7 @@ # 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). - -This is the glue between TransformerEngine's MXFP4 quantizer and the ported -Triton block-scaled grouped kernels in ``grouped_gemm_mxfp4.py``. It reuses -TE's :class:`MXFP4Quantizer` (rather than porting Primus-Turbo's fused grouped -dual-quant) to produce the packed FP4 operands the kernels consume, then builds -the per-group offsets and calls the kernels. - -Three grouped ops, paired exactly like the dense MXFP4 recipe (all NT): - - fprop : C = A_row @ W_row^T (contract K, fwd kernel) - dgrad : dA = gradO_row @ W_col^T (contract N, fwd kernel) - wgrad : dW = gradO_col @ A_col^T (contract M, variable-K kernel) - -Operand layout the kernels expect (plain, un-swizzled OCP MXFP4): - * data : E2M1 packed two-per-byte along the contraction axis (``uint8``); - * scale : one E8M0 (``uint8``) per 1x32 logical-element block, contiguous - along the block axis. -TE's ``MXFP4Quantizer`` emits exactly this in its row-wise ``_rowwise_data`` / -``_rowwise_scale_inv`` and col-wise ``_columnwise_data`` / ``_columnwise_scale_inv`` -members when the shuffle / GEMM-swizzle flags are all off. The scale buffers are -over-allocated (rows to 256, block axis to 8); we slice the live ``[:, :F/32]`` -region and make it contiguous so the kernel sees tight strides. - -RHT (Hadamard) on the wgrad operands and the ``main_grad`` beta=1 accumulate -fusion are deferred to the autograd layer; this module runs plain -MX (``use_hadamard=False``) and returns a fresh weight-gradient tensor. -""" +"""TE-side quantization bridge + dispatch for grouped MXFP4 GEMM (gfx950).""" from __future__ import annotations