From 349e95b825502d857ab5b91433055e50a2db2b59 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 12:11:17 -0500 Subject: [PATCH 01/26] add a8w4/mxfp4 reference recipe --- tests/pytorch/mxfp4/test_a8w4_recipe.py | 85 ++++++++ .../quantization_mxfp4_grouped.py | 191 ++++++++++++++++++ 2 files changed, 276 insertions(+) create mode 100644 tests/pytorch/mxfp4/test_a8w4_recipe.py create mode 100644 transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py diff --git a/tests/pytorch/mxfp4/test_a8w4_recipe.py b/tests/pytorch/mxfp4/test_a8w4_recipe.py new file mode 100644 index 0000000000..b95a640460 --- /dev/null +++ b/tests/pytorch/mxfp4/test_a8w4_recipe.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""End-to-end test: te.Linear under the a8w4 CustomRecipe.""" + +import gc + +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.common import recipe + +from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_mxfp4_grouped import ( + MXFP8E4M3QuantizerRef, + a8w4_quantizer_factory, +) +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + +_recipe_ok, _reason = te.is_mxfp4_available(return_reason=True) +pytestmark = pytest.mark.skipif(not _recipe_ok, reason=_reason) + + +def _isolate(): + FP8GlobalStateManager.reset() + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _e2m1_ref(): + return MXFP4QuantizerRef( + rowwise=True, + columnwise=True, + shuffle_rowwise_data=False, + shuffle_columnwise_data=False, + with_gemm_swizzled_scales=False, + use_hadamard=False, + ) + + +def test_a8w4_custom_recipe_linear_forward(): + """te.Linear forward under CustomRecipe(a8w4) == the recipe's qgemm reference.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + B, S, K, N = 2, 128, 256, 256 # K=256 -> loop_k>1 in the eventual fast path + torch.manual_seed(0) + + module = te.Linear(K, N, bias=False, device=device, params_dtype=dtype) + x = torch.randn(B, S, K, device=device, dtype=dtype) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with te.autocast(enabled=True, recipe=rec): + y = module(x) + + # Reference + xq = MXFP8E4M3QuantizerRef(rowwise=True, columnwise=False).quantize(x.reshape(-1, K)) + wq = _e2m1_ref().quantize(module.weight.detach()) + y_ref = MXFP8E4M3QuantizerRef().qgemm( + xq.data, wq.data, None, dtype, xq.scale, wq.scale, qresult_x=xq, qresult_w=wq, + ).reshape(B, S, N) + + rel = (y.float() - y_ref.float()).norm() / y_ref.float().norm().clamp_min(1e-12) + assert y.shape == (B, S, N) + assert rel < 5e-2, f"a8w4 recipe forward disagrees with its qgemm reference: rel={rel:.4f}" + + +def test_a8w4_custom_recipe_linear_runs_backward(): + """Smoke: fwd+bwd under the a8w4 recipe complete without error and grads are finite.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + B, S, K, N = 2, 128, 256, 256 + torch.manual_seed(0) + + module = te.Linear(K, N, bias=False, device=device, params_dtype=dtype) + x = torch.randn(B, S, K, device=device, dtype=dtype, requires_grad=True) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with te.autocast(enabled=True, recipe=rec): + y = module(x) + y.backward(torch.randn_like(y)) + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert module.weight.grad is not None and torch.isfinite(module.weight.grad).all() diff --git a/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py new file mode 100644 index 0000000000..9d2b0a6258 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Reference MXFP4-family recipes. + +Extensible ``CustomRecipe`` quantizer factories built on the existing +``MXFP4QuantizerRef``. + + * ``mxfp4_grouped_quantizer_factory`` -- w4a4: E2M1 weights and activations. + * ``a8w4_quantizer_factory`` -- Kimi-K3 a8w4: MXFP8 E4M3 activation x + MXFP4 E2M1 weight. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from transformer_engine.pytorch.custom_recipes import quantization +from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer + +from .quantization_mxfp4 import ( + MXFP4_BLOCK_SIZE, + MXFP4QuantizerRef, + MXFP4TensorRef, + e8m0_to_f32, + mxfp4_to_f32, +) + + +def _deq_ref_operand(data, scale, storage, *, default_e4m3: bool) -> torch.Tensor: + """Dequantize one reference operand to f32 by its own stored format.""" + quantizer = getattr(storage, "_quantizer", None) + if quantizer is not None: + is_e4m3 = isinstance(quantizer, MXFP8E4M3QuantizerRef) + else: + is_e4m3 = default_e4m3 + hp = data.float() if is_e4m3 else mxfp4_to_f32(data) + rows, k = hp.shape + grid_k = k // MXFP4_BLOCK_SIZE + s = e8m0_to_f32(scale[:rows, :grid_k]).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + return hp * s + + +class MXFP8E4M3QuantizerRef(Quantizer): + """Reference MXFP8 (E4M3) block-scaled quantizer for the a8w4 activation.""" + + def __init__(self, rowwise: bool = True, columnwise: bool = True): + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.internal = True + + @property + def custom(self) -> bool: + return True + + def _quantize(self, tensor: torch.Tensor): + """``tensor`` [M, K] -> (E4M3 data [M, K], E8M0 scale u8 [M, K/32]).""" + fmax = torch.finfo(torch.float8_e4m3fn).max # 448 + m, k = tensor.shape + assert k % MXFP4_BLOCK_SIZE == 0, "K must be a multiple of the 32-elem block" + xb = tensor.float().reshape(m, k // MXFP4_BLOCK_SIZE, MXFP4_BLOCK_SIZE) + 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) # E8M0, biased by 127 + return data, scale + + def quantize(self, tensor: torch.Tensor, **kwargs) -> MXFP4TensorRef: # noqa: ARG002 + """Quantize a high-precision tensor, returning an MXFP4TensorRef holder.""" + original_shape = tensor.shape + if tensor.ndim > 2: + tensor = tensor.view(-1, tensor.shape[-1]) + qx, sx = self._quantize(tensor) if self.rowwise_usage else (None, None) + if self.columnwise_usage: + qx_t, sx_t = self._quantize(tensor.t().contiguous()) + else: + qx_t = sx_t = None + return MXFP4TensorRef( + data=qx, + scale=sx, + data_t=qx_t, + scale_t=sx_t, + dtype=tensor.dtype, + device=tensor.device, + original_shape=original_shape, + _quantizer=self, + ) + + # ------------------------------------------------------------------ + # Reference GEMM + # ------------------------------------------------------------------ + def qgemm( + self, + qx: torch.Tensor, + qw: torch.Tensor, + m_params: quantization.MMParams, + out_dtype: torch.dtype, + sx: torch.Tensor, + sw: torch.Tensor, + bias: torch.Tensor | None = None, + out: torch.Tensor | None = None, + accumulate: bool = False, + gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, + qresult_x: QuantizedTensorStorage | None = None, + qresult_w: QuantizedTensorStorage | None = None, + layout: str = "TN", + ) -> torch.Tensor: + """``Y = dequant(qx) @ dequant(qw)^T``; each operand by its own format.""" + assert bias is None, "Bias not yet supported in the a8w4 reference GEMM." + hp_x = _deq_ref_operand(qx, sx, qresult_x, default_e4m3=True) # A: E4M3 here + hp_w = _deq_ref_operand(qw, sw, qresult_w, default_e4m3=False) # B: E2M1 weight + assert hp_x.shape[1] == hp_w.shape[1], "K mismatch between operands" + y = torch.mm(hp_x, hp_w.t()).to(out_dtype) + if accumulate: + assert out is not None + y = y + out.to(out_dtype) + return y + + @staticmethod + def _dequantize(data_e4m3, scale_u8, dtype=torch.float32) -> torch.Tensor: + m, k = data_e4m3.shape + vals = data_e4m3.float().reshape(m, k // MXFP4_BLOCK_SIZE, MXFP4_BLOCK_SIZE) + scale = e8m0_to_f32(scale_u8[:m, : k // MXFP4_BLOCK_SIZE]).unsqueeze(-1) + return (vals * scale).reshape(m, k).to(dtype) + + def dequantize_rowwise(self, ref: MXFP4TensorRef, dtype=torch.float32) -> torch.Tensor: + assert ref.data is not None and ref.scale is not None + return self._dequantize(ref.data, ref.scale, dtype=dtype) + + def dequantize_columnwise(self, ref: MXFP4TensorRef, dtype=torch.float32) -> torch.Tensor: + assert ref.data_t is not None and ref.scale_t is not None + return self._dequantize(ref.data_t, ref.scale_t, dtype=dtype).t().contiguous() + + @property + def supports_allgather_fp8(self) -> bool: + return False + + @property + def supports_dequantize(self) -> bool: + return True + + def transpose_qresult(self, qresult): # noqa: ARG002 + raise NotImplementedError + + @property + def is_data_t_transposed_in_memory(self) -> bool: + raise NotImplementedError + + +def _e2m1_ref(*, rowwise: bool, columnwise: bool) -> MXFP4QuantizerRef: + """Plain (un-swizzled, no-Hadamard) E2M1 reference quantizer.""" + return MXFP4QuantizerRef( + rowwise=rowwise, + columnwise=columnwise, + shuffle_rowwise_data=False, + shuffle_columnwise_data=False, + with_gemm_swizzled_scales=False, + use_hadamard=False, + ) + + +def _make_mxfp4_family_factory(*, activation: str): + """Build a role-based MXFP4-family ``qfactory``. + + ``activation`` selects the non-weight format: ``"e2m1"`` -> w4a4 (E2M1 + everywhere); ``"e4m3"`` -> a8w4 (E4M3 activation and grad, E2M1 weight). + Weights are always E2M1 and need both layouts (fprop row + dgrad col). + """ + + def factory(role: Optional[QuantizerRole]) -> Quantizer: + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + tensor_type = role.tensor_type if role is not None else "" + if is_linear and tensor_type == "weight": + return _e2m1_ref(rowwise=True, columnwise=True) + if activation == "e2m1": + return _e2m1_ref(rowwise=True, columnwise=True) + # a8w4: E4M3 for activation and grad_output (uniform E4M3 keeps every + # backward GEMM E4M3xE2M1 / E4M3xE4M3, all handled by the ref qgemm). + return MXFP8E4M3QuantizerRef(rowwise=True, columnwise=True) + + return factory + + +# w4a4: pure E2M1 +mxfp4_grouped_quantizer_factory = _make_mxfp4_family_factory(activation="e2m1") + +# Kimi-K3 a8w4: E4M3 activation x E2M1 weight +a8w4_quantizer_factory = _make_mxfp4_family_factory(activation="e4m3") From f4c9b082a0a84e77c1091b8abedade83cdafbe9a Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 13:27:21 -0500 Subject: [PATCH 02/26] enable fast path for groupedlinear --- tests/pytorch/mxfp4/test_a8w4_grouped_fast.py | 204 ++++++++++++++++ .../quantization_mxfp4_grouped.py | 6 +- .../pytorch/module/grouped_linear.py | 221 ++++++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4_impl.py | 52 +++++ 4 files changed, 480 insertions(+), 3 deletions(-) create mode 100644 tests/pytorch/mxfp4/test_a8w4_grouped_fast.py diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_fast.py b/tests/pytorch/mxfp4/test_a8w4_grouped_fast.py new file mode 100644 index 0000000000..a15d115448 --- /dev/null +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_fast.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Forward fast path: grouped a8w4 reference bytes -> persistent Triton kernel. + +Covers both the standalone pre-quantized entry +(:func:`grouped_gemm_mxfp4_fprop_prequantized`) and its end-to-end use inside +``te.GroupedLinear`` under ``CustomRecipe(a8w4_quantizer_factory)``. +""" + +import gc + +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.common import recipe + +from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( + _row_operand, + grouped_gemm_mxfp4_dgrad, + grouped_gemm_mxfp4_fprop_prequantized, +) +from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import ( + MXFP4_BLOCK_SIZE, + e8m0_to_f32, + mxfp4_to_f32, +) +from transformer_engine.pytorch.custom_recipes.quantization_mxfp4_grouped import ( + MXFP8E4M3QuantizerRef, + _deq_ref_operand, + _e2m1_ref, + a8w4_quantizer_factory, + mxfp4_grouped_quantizer_factory, +) + +_recipe_ok, _reason = te.is_mxfp4_available(return_reason=True) +pytestmark = pytest.mark.skipif(not _recipe_ok, reason=_reason) + +# fp32 matmul vs MFMA on identical quantized bytes: only accumulation differs. +_REL_TOL = 5e-2 + + +def _isolate(): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _reference_grouped_a8w4(aq, weight_refs, m_splits, N): + """Independent reference: dequant the same reference bytes, per-group fp32 matmul.""" + total_M = aq.data.shape[0] + a_hp = _deq_ref_operand(aq.data, aq.scale, aq, default_e4m3=True) + ref = torch.zeros(total_M, N, device=aq.data.device, dtype=torch.float32) + start = 0 + for g, m in enumerate(m_splits): + if m > 0: + w_hp = _deq_ref_operand( + weight_refs[g].data, weight_refs[g].scale, weight_refs[g], default_e4m3=False + ) + ref[start : start + m] = a_hp[start : start + m] @ w_hp.t() + start += m + return ref + + +def test_grouped_prequantized_a8w4_matches_reference(): + """The kernel on reference bytes agrees with the reference qgemm (Option B).""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + K, N = 256, 256 # K % 128 == 0 -> loop_k > 1 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + + a = torch.randn(total_M, K, device=device, dtype=dtype) + weights = [torch.randn(N, K, device=device, dtype=dtype) for _ in m_splits] + + aq = MXFP8E4M3QuantizerRef(rowwise=True, columnwise=False).quantize(a) + weight_refs = [_e2m1_ref(rowwise=True, columnwise=False).quantize(w) for w in weights] + b_data = torch.stack([w.data for w in weight_refs], dim=0) + b_scale = torch.stack([w.scale for w in weight_refs], dim=0) + + out = grouped_gemm_mxfp4_fprop_prequantized( + aq.data, aq.scale, b_data, b_scale, m_splits, a_is_mxfp8=True, out_dtype=dtype + ) + ref = _reference_grouped_a8w4(aq, weight_refs, m_splits, N) + + rel = (out.float() - ref).norm() / ref.norm().clamp_min(1e-12) + assert out.shape == (total_M, N) + assert rel < _REL_TOL, f"prequantized grouped a8w4 disagrees with reference: rel={rel:.4f}" + + +def test_grouped_linear_a8w4_recipe_forward(): + """te.GroupedLinear forward under CustomRecipe(a8w4) routes to the fast kernel.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with torch.no_grad(), te.autocast(enabled=True, recipe=rec): + out = model(inp, m_splits) + + # Independent reference on the module's own weights, quantized by the same + # reference quantizers the recipe factory produces. + aq = MXFP8E4M3QuantizerRef(rowwise=True, columnwise=False).quantize(inp) + weight_refs = [ + _e2m1_ref(rowwise=True, columnwise=False).quantize(getattr(model, f"weight{g}").detach()) + for g in range(num_gemms) + ] + ref = _reference_grouped_a8w4(aq, weight_refs, m_splits, N) + + rel = (out.float() - ref).norm() / ref.norm().clamp_min(1e-12) + assert out.shape == (total_M, N) + assert rel < _REL_TOL, f"GroupedLinear a8w4 forward disagrees with reference: rel={rel:.4f}" + + +def test_grouped_prequantized_a4w4_dgrad_matches_reference(): + """dgrad kernel on reference col-weight bytes agrees with the dequant reference.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + K, N = 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + grad_out = torch.randn(total_M, N, device=device, dtype=dtype) + weights = [torch.randn(N, K, device=device, dtype=dtype) for _ in m_splits] + + # Reference col weight (data_t/scale_t) + wrefs = [_e2m1_ref(rowwise=False, columnwise=True).quantize(w) for w in weights] + w_col_data = torch.stack([w.data_t for w in wrefs], dim=0) # [G, K, N/2] + w_col_scale = torch.stack([w.scale_t for w in wrefs], dim=0) # [G, K, N/32] + + grad_a = grouped_gemm_mxfp4_dgrad( + grad_out, None, m_splits, out_dtype=dtype, weight_col=(w_col_data, w_col_scale) + ) + + # Reference: dequant the same operands (the kernel's native mxfp4 grad + the + # reference col weight) and matmul in fp32. dX[g] = dY[g] @ W[g]. + gd, gs = _row_operand(grad_out) # native mxfp4 grad bytes the kernel casts to + g_hp = mxfp4_to_f32(gd) * e8m0_to_f32(gs).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + ref = torch.zeros(total_M, K, device=device, dtype=torch.float32) + start = 0 + for g, m in enumerate(m_splits): + if m > 0: + w_t_hp = _deq_ref_operand( + wrefs[g].data_t, wrefs[g].scale_t, wrefs[g], default_e4m3=False + ) # [K, N] == W^T + ref[start : start + m] = g_hp[start : start + m] @ w_t_hp.t() + start += m + + rel = (grad_a.float() - ref).norm() / ref.norm().clamp_min(1e-12) + assert grad_a.shape == (total_M, K) + assert rel < _REL_TOL, f"a4w4 dgrad on reference col weight disagrees: rel={rel:.4f}" + + +def test_grouped_linear_a4w4_recipe_backward(): + """te.GroupedLinear a4w4 fwd+bwd routes to the low-precision dgrad/wgrad kernels.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype, requires_grad=True) + + rec = recipe.CustomRecipe(qfactory=mxfp4_grouped_quantizer_factory) + with te.autocast(enabled=True, recipe=rec): + out = model(inp, m_splits) + out.backward(torch.randn_like(out)) + + assert out.shape == (total_M, N) + assert inp.grad is not None and inp.grad.shape == inp.shape + assert torch.isfinite(inp.grad).all() + for g in range(num_gemms): + w = getattr(model, f"weight{g}") + assert w.grad is not None and w.grad.shape == w.shape + assert torch.isfinite(w.grad).all() + + +def test_grouped_linear_a8w4_backward_raises(): + """a8w4 is forward-only QAT: the grouped fast-path backward must raise clearly.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 2, 256, 256 + m_splits = [128, 128] + total_M = sum(m_splits) + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype, requires_grad=True) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with te.autocast(enabled=True, recipe=rec): + out = model(inp, m_splits) + with pytest.raises(NotImplementedError, match="forward-only QAT"): + out.backward(torch.randn_like(out)) + diff --git a/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py index 9d2b0a6258..85ec536610 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py @@ -6,7 +6,7 @@ Extensible ``CustomRecipe`` quantizer factories built on the existing ``MXFP4QuantizerRef``. - * ``mxfp4_grouped_quantizer_factory`` -- w4a4: E2M1 weights and activations. + * ``mxfp4_grouped_quantizer_factory`` -- a4w4: E2M1 weights and activations. * ``a8w4_quantizer_factory`` -- Kimi-K3 a8w4: MXFP8 E4M3 activation x MXFP4 E2M1 weight. """ @@ -165,7 +165,7 @@ def _e2m1_ref(*, rowwise: bool, columnwise: bool) -> MXFP4QuantizerRef: def _make_mxfp4_family_factory(*, activation: str): """Build a role-based MXFP4-family ``qfactory``. - ``activation`` selects the non-weight format: ``"e2m1"`` -> w4a4 (E2M1 + ``activation`` selects the non-weight format: ``"e2m1"`` -> a4w4 (E2M1 everywhere); ``"e4m3"`` -> a8w4 (E4M3 activation and grad, E2M1 weight). Weights are always E2M1 and need both layouts (fprop row + dgrad col). """ @@ -184,7 +184,7 @@ def factory(role: Optional[QuantizerRole]) -> Quantizer: return factory -# w4a4: pure E2M1 +# a4w4: pure E2M1 mxfp4_grouped_quantizer_factory = _make_mxfp4_family_factory(activation="e2m1") # Kimi-K3 a8w4: E4M3 activation x E2M1 weight diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 059c3a4ce6..42804a388a 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -577,6 +577,59 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) ) + @staticmethod + def _is_grouped_mxfp4_triton_supported( + *, + fp8, + recipe, + input_quantizers, + weight_quantizers, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ) -> bool: + """Triton grouped MXFP4 / a8w4 path: gfx950 + a ``CustomRecipe`` whose + factory yields the plain-layout MXFP4 reference quantizers. + + a8w4 is forward-only (backward raises); a4w4 also wires the low-precision + dgrad/wgrad (see :meth:`_forward_grouped_mxfp4_triton`). + """ + if not (IS_HIP_EXTENSION and fp8 and recipe is not None and recipe.custom()): + return False + # Contractions the kernels step in 128s: fprop/wgrad over K=in_features, + # dgrad over N=out_features. Require both so the a4w4 backward is safe. + if in_features % 128 != 0 or out_features % 128 != 0: + return False + if ( + use_bias + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits is not None + ): + return False + try: + from ..triton_kernels.grouped_gemm_mxfp4 import _is_gfx950 + from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef + from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef + except ImportError: + return False + if not _is_gfx950(): + return False + wq, iq = weight_quantizers[0], input_quantizers[0] + # Weight is always plain-layout E2M1; activation is E2M1 (a4w4) or E4M3 (a8w4). + return isinstance(wq, MXFP4QuantizerRef) and isinstance( + iq, (MXFP4QuantizerRef, MXFP8E4M3QuantizerRef) + ) + @staticmethod def _forward_blockwise_fp8_triton( ctx, @@ -880,6 +933,143 @@ def grouped_gemm_wgrad(*_unused): *grad_biases, ) + @staticmethod + def _forward_grouped_mxfp4_triton( + ctx, + *, + inp, + m_splits, + weights, + input_quantizers, + weight_quantizers, + activation_dtype, + is_grad_enabled, + ): + """Grouped MXFP4 / a8w4 forward (ROCm Triton, gfx950). + + Selected from :meth:`forward` under a ``CustomRecipe`` whose factory yields + the plain-layout MXFP4 reference quantizers (see + :meth:`_is_grouped_mxfp4_triton_supported`). The recipe's own reference + quantizers own the quantization; their packed E2M1/E4M3 data + E8M0 scales + feed straight into the persistent grouped kernel + (:func:`grouped_gemm_mxfp4_fprop_prequantized`), so the fast-path numerics + match the reference ``qgemm``. Forward-only: the low-precision backward is + not yet wired (:meth:`_backward_grouped_mxfp4_triton`). + """ + from ..triton_kernels.grouped_gemm_mxfp4_impl import ( + grouped_gemm_mxfp4_fprop_prequantized, + ) + from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef + + num_gemms = len(weights) + in_features = weights[0].size(-1) + m_splits_list = ( + m_splits.tolist() if isinstance(m_splits, torch.Tensor) else [int(m) for m in m_splits] + ) + + # Activation: quantize the whole grouped [total_M, K] once (row-wise scales + # do not cross groups). Detach: the reference quantizer casts via NumPy and + # the low-precision forward is not differentiated through (custom backward). + a = inp.reshape(-1, in_features).to(activation_dtype).contiguous().detach() + input_quantizer = input_quantizers[0] + input_quantizer.set_usage(rowwise=True, columnwise=False) + a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) + aq = input_quantizer.quantize(a) + + # Weights: per-expert row-wise E2M1, stacked to the kernel's [G, N, *] layout. + # The a4w4 backward additionally needs the col-wise weight for the dgrad kernel. + want_wcol = is_grad_enabled and not a_is_mxfp8 + b_datas, b_scales = [], [] + wcol_datas, wcol_scales = [], [] + for i in range(num_gemms): + wq = weight_quantizers[i] + wq.set_usage(rowwise=True, columnwise=want_wcol) + wref = wq.quantize(weights[i].detach()) + b_datas.append(wref.data) + b_scales.append(wref.scale) + if want_wcol: + wcol_datas.append(wref.data_t) + wcol_scales.append(wref.scale_t) + b_data = torch.stack(b_datas, dim=0) + b_scale = torch.stack(b_scales, dim=0) + + out = grouped_gemm_mxfp4_fprop_prequantized( + aq.data, + aq.scale, + b_data, + b_scale, + m_splits_list, + a_is_mxfp8=a_is_mxfp8, + out_dtype=activation_dtype, + ) + + if is_grad_enabled: + ctx.use_grouped_mxfp4_triton = True + ctx.grouped_mxfp4_is_a8w4 = a_is_mxfp8 + ctx.num_gemms = num_gemms + ctx.grouped_mxfp4_inp_shape = inp.shape + ctx.grouped_mxfp4_m_splits = m_splits_list + if not a_is_mxfp8: + # a4w4 low-precision backward: dgrad consumes the reference col weight; + # wgrad re-quantizes the saved hp activation + grad_out in the kernel. + ctx.grouped_mxfp4_requires_dgrad = inp.requires_grad + ctx.grouped_mxfp4_weight_requires_grad = weights[0].requires_grad + ctx.save_for_backward( + a, torch.stack(wcol_datas, dim=0), torch.stack(wcol_scales, dim=0) + ) + + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), [None] * num_gemms + + @staticmethod + def _backward_grouped_mxfp4_triton(ctx, grad_output): + """Backward for the grouped MXFP4 / a8w4 fast path. + + a4w4 uses the low-precision ``grouped_gemm_mxfp4_dgrad`` / ``_wgrad`` kernels. + a8w4 (Kimi-K3) is a forward-only QAT recipe: its high-precision (STE) backward + is plain bf16, not a fast kernel, and is intentionally not wired -- run a8w4 + under inference (``torch.no_grad``). + """ + if getattr(ctx, "grouped_mxfp4_is_a8w4", False): + raise NotImplementedError( + "a8w4 (Kimi-K3) is a forward-only QAT recipe on the grouped fast path; its " + "high-precision backward is not wired. Run a8w4 under inference (torch.no_grad)." + ) + from ..triton_kernels.grouped_gemm_mxfp4_impl import ( + grouped_gemm_mxfp4_dgrad, + grouped_gemm_mxfp4_wgrad, + ) + + a, w_col_data, w_col_scale = ctx.saved_tensors + m_splits = ctx.grouped_mxfp4_m_splits + go = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() + + grad_a = grad_weight = None + if ctx.grouped_mxfp4_requires_dgrad: + grad_a = grouped_gemm_mxfp4_dgrad( + go, None, m_splits, out_dtype=a.dtype, weight_col=(w_col_data, w_col_scale) + ) + if ctx.grouped_mxfp4_weight_requires_grad: + grad_weight = grouped_gemm_mxfp4_wgrad(a, go, m_splits, out_dtype=a.dtype) + + wgrad_list = ( + list(torch.unbind(grad_weight, dim=0)) + if grad_weight is not None + else [None] * ctx.num_gemms + ) + grad_biases = [None] * ctx.num_gemms + dgrad = grad_a.view(ctx.grouped_mxfp4_inp_shape) if grad_a is not None else None + # Grads match forward inputs: (inp, m_splits, non_tensor_args, out, dgrad_out, + # *weights, *biases). + return ( + dgrad, + None, # m_splits + None, # non_tensor_args + None, # out + None, # dgrad_out + *wgrad_list, + *grad_biases, + ) + # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( @@ -981,6 +1171,35 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) + # Grouped MXFP4 / a8w4 (Triton, gfx950) under a CustomRecipe. Runs its + # own reference-quantizer quantization + persistent grouped GEMM and returns + # early. Forward-only. + if _GroupedLinear._is_grouped_mxfp4_triton_supported( + fp8=fp8, + recipe=blockwise_recipe, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + use_bias=use_bias, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=weights[0].size(-1), + out_features=weights[0].size(0), + ): + return _GroupedLinear._forward_grouped_mxfp4_triton( + ctx, + inp=inp, + m_splits=m_splits, + weights=weights, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + activation_dtype=activation_dtype, + is_grad_enabled=is_grad_enabled, + ) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -1499,6 +1718,8 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): + if getattr(ctx, "use_grouped_mxfp4_triton", False): + return _GroupedLinear._backward_grouped_mxfp4_triton(ctx, grad_output) if getattr(ctx, "use_blockwise_fp8_triton", False): return _GroupedLinear._backward_blockwise_fp8_triton(ctx, grad_output) if ctx.use_grouped_tensor_path: 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 81702da6e7..a199972035 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -336,6 +336,58 @@ def grouped_gemm_a8w4_fprop( ) +def grouped_gemm_mxfp4_fprop_prequantized( + a_data: torch.Tensor, + a_scale: torch.Tensor, + b_data: torch.Tensor, + b_scale: torch.Tensor, + m_splits: Sequence[int], + *, + a_is_mxfp8: bool, + out_dtype: torch.dtype = torch.bfloat16, + num_cu: Optional[int] = None, +) -> torch.Tensor: + """Grouped forward from already-quantized operands: ``C[g] = A[g] @ W[g]^T``. + + Unlike :func:`grouped_gemm_mxfp4_fprop` / :func:`grouped_gemm_a8w4_fprop`, this + takes the packed data + E8M0 scales directly (no internal cast), so a recipe's + own reference quantizer owns the quantization and this only drives the GEMM. + + Args: + a_data: activation data grouped along M by ``m_splits`` -- ``[total_M, K]`` + e4m3 when ``a_is_mxfp8`` else ``[total_M, K/2]`` packed e2m1. + a_scale: ``[total_M, K/32]`` uint8 E8M0 activation scales. + b_data: ``[G, N, K/2]`` packed e2m1 weights. + b_scale: ``[G, N, K/32]`` uint8 E8M0 weight scales. + m_splits: per-group token counts (len G). + a_is_mxfp8: ``True`` routes the MXFP8(e4m3) x MXFP4(e2m1) a8w4 kernel; + ``False`` the MXFP4 x MXFP4 kernel. + + Returns: + ``[total_M, N]`` output in ``out_dtype``. + """ + N = b_data.shape[1] + K = b_data.shape[2] * 2 # weights are packed 2 e2m1 elems/byte along K + _require_gfx950() + _check_contract(K, "K") + _check_splits(m_splits, a_data.shape[0], b_data.shape[0]) + + group_offs = _prefix_offsets(m_splits, a_data.device) + kernel = grouped_gemm_a8w4_triton_kernel if a_is_mxfp8 else grouped_gemm_mxfp4_triton_kernel + return 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 a42c6c1cd82900d68157e1846c39e74713b472d2 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 14:38:16 -0500 Subject: [PATCH 03/26] address copilot review --- ...ed_fast.py => test_a8w4_grouped_recipe.py} | 6 +- tests/pytorch/mxfp4/test_a8w4_recipe.py | 5 +- .../pytorch/module/grouped_linear.py | 40 ++++++++++---- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 55 ++++++++++++++++++- 4 files changed, 89 insertions(+), 17 deletions(-) rename tests/pytorch/mxfp4/{test_a8w4_grouped_fast.py => test_a8w4_grouped_recipe.py} (97%) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_fast.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py similarity index 97% rename from tests/pytorch/mxfp4/test_a8w4_grouped_fast.py rename to tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index a15d115448..7338051f92 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_fast.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -1,7 +1,7 @@ # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # License for AMD contributions = MIT. See LICENSE for more information -"""Forward fast path: grouped a8w4 reference bytes -> persistent Triton kernel. +"""Grouped MXFP4/a8w4 recipe path: reference bytes -> persistent Triton kernel. Covers both the standalone pre-quantized entry (:func:`grouped_gemm_mxfp4_fprop_prequantized`) and its end-to-end use inside @@ -90,7 +90,7 @@ def test_grouped_prequantized_a8w4_matches_reference(): def test_grouped_linear_a8w4_recipe_forward(): - """te.GroupedLinear forward under CustomRecipe(a8w4) routes to the fast kernel.""" + """te.GroupedLinear forward under CustomRecipe(a8w4) routes to the grouped Triton kernel.""" _isolate() device, dtype = "cuda", torch.bfloat16 num_gemms, K, N = 4, 256, 256 @@ -185,7 +185,7 @@ def test_grouped_linear_a4w4_recipe_backward(): def test_grouped_linear_a8w4_backward_raises(): - """a8w4 is forward-only QAT: the grouped fast-path backward must raise clearly.""" + """a8w4 is forward-only QAT: the grouped-path backward must raise clearly.""" _isolate() device, dtype = "cuda", torch.bfloat16 num_gemms, K, N = 2, 256, 256 diff --git a/tests/pytorch/mxfp4/test_a8w4_recipe.py b/tests/pytorch/mxfp4/test_a8w4_recipe.py index b95a640460..334b1454c8 100644 --- a/tests/pytorch/mxfp4/test_a8w4_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_recipe.py @@ -1,6 +1,5 @@ # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. +# License for AMD contributions = MIT. See LICENSE for more information """End-to-end test: te.Linear under the a8w4 CustomRecipe.""" @@ -44,7 +43,7 @@ def test_a8w4_custom_recipe_linear_forward(): """te.Linear forward under CustomRecipe(a8w4) == the recipe's qgemm reference.""" _isolate() device, dtype = "cuda", torch.bfloat16 - B, S, K, N = 2, 128, 256, 256 # K=256 -> loop_k>1 in the eventual fast path + B, S, K, N = 2, 128, 256, 256 # K=256 -> loop_k>1 in the grouped Triton path torch.manual_seed(0) module = te.Linear(K, N, bias=False, device=device, params_dtype=dtype) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 42804a388a..8932fce0e0 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -624,10 +624,30 @@ def _is_grouped_mxfp4_triton_supported( return False if not _is_gfx950(): return False - wq, iq = weight_quantizers[0], input_quantizers[0] + + def _plain_layout(q) -> bool: + # The grouped Triton kernels need plain (un-shuffled, un-swizzled, + # no-Hadamard) operands. MXFP8E4M3QuantizerRef is always plain; an + # MXFP4QuantizerRef must have every layout-transform flag off. + if isinstance(q, MXFP8E4M3QuantizerRef): + return True + return not ( + q.shuffle_rowwise_data + or q.shuffle_columnwise_data + or q.with_gemm_swizzled_scales + or q.use_hadamard + ) + # Weight is always plain-layout E2M1; activation is E2M1 (a4w4) or E4M3 (a8w4). - return isinstance(wq, MXFP4QuantizerRef) and isinstance( - iq, (MXFP4QuantizerRef, MXFP8E4M3QuantizerRef) + # Check every expert quantizer -- a shuffled/swizzled/Hadamard variant of the + # same class would silently feed the wrong layout to the kernel. + if not all( + isinstance(q, MXFP4QuantizerRef) and _plain_layout(q) for q in weight_quantizers + ): + return False + return all( + isinstance(q, (MXFP4QuantizerRef, MXFP8E4M3QuantizerRef)) and _plain_layout(q) + for q in input_quantizers ) @staticmethod @@ -952,9 +972,9 @@ def _forward_grouped_mxfp4_triton( :meth:`_is_grouped_mxfp4_triton_supported`). The recipe's own reference quantizers own the quantization; their packed E2M1/E4M3 data + E8M0 scales feed straight into the persistent grouped kernel - (:func:`grouped_gemm_mxfp4_fprop_prequantized`), so the fast-path numerics - match the reference ``qgemm``. Forward-only: the low-precision backward is - not yet wired (:meth:`_backward_grouped_mxfp4_triton`). + (:func:`grouped_gemm_mxfp4_fprop_prequantized`), so the grouped-path numerics + match the reference ``qgemm``. a4w4 wires the low-precision backward; a8w4 + is forward-only (see :meth:`_backward_grouped_mxfp4_triton`). """ from ..triton_kernels.grouped_gemm_mxfp4_impl import ( grouped_gemm_mxfp4_fprop_prequantized, @@ -1022,16 +1042,16 @@ def _forward_grouped_mxfp4_triton( @staticmethod def _backward_grouped_mxfp4_triton(ctx, grad_output): - """Backward for the grouped MXFP4 / a8w4 fast path. + """Backward for the grouped MXFP4 / a8w4 Triton path. a4w4 uses the low-precision ``grouped_gemm_mxfp4_dgrad`` / ``_wgrad`` kernels. a8w4 (Kimi-K3) is a forward-only QAT recipe: its high-precision (STE) backward - is plain bf16, not a fast kernel, and is intentionally not wired -- run a8w4 + is plain bf16, not a dedicated grouped kernel, and is intentionally not wired -- run a8w4 under inference (``torch.no_grad``). """ if getattr(ctx, "grouped_mxfp4_is_a8w4", False): raise NotImplementedError( - "a8w4 (Kimi-K3) is a forward-only QAT recipe on the grouped fast path; its " + "a8w4 (Kimi-K3) is a forward-only QAT recipe on the grouped Triton path; its " "high-precision backward is not wired. Run a8w4 under inference (torch.no_grad)." ) from ..triton_kernels.grouped_gemm_mxfp4_impl import ( @@ -1173,7 +1193,7 @@ def forward( # Grouped MXFP4 / a8w4 (Triton, gfx950) under a CustomRecipe. Runs its # own reference-quantizer quantization + persistent grouped GEMM and returns - # early. Forward-only. + # early. a4w4 also runs the backward; a8w4 is forward-only. if _GroupedLinear._is_grouped_mxfp4_triton_supported( fp8=fp8, recipe=blockwise_recipe, 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 a199972035..84fa5a0dc8 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -336,6 +336,56 @@ def grouped_gemm_a8w4_fprop( ) +def _check_prequantized_operands( + a_data: torch.Tensor, + a_scale: torch.Tensor, + b_data: torch.Tensor, + b_scale: torch.Tensor, + *, + a_is_mxfp8: bool, +) -> None: + """Validate pre-quantized operand ranks/shapes/dtypes/devices before launch. + + The raw kernel indexes straight into these buffers, so a malformed operand would + read out of bounds rather than fail; reject it up front. Assumes ``b_data`` is + 3-D ``[G, N, K/2]`` with ``K`` a 128-multiple (checked by the caller). + """ + G, N = b_data.shape[0], b_data.shape[1] + K = b_data.shape[2] * 2 + grid_k = K // MXFP4_BLOCK # one E8M0 scale per 32-elem block along K + total_M = a_data.shape[0] + a_cols = K if a_is_mxfp8 else K // 2 # e4m3 unpacked vs e2m1 packed + + if a_data.ndim != 2 or a_data.shape[1] != a_cols: + raise ValueError( + f"a_data must be [total_M, {a_cols}]" + f" ({'e4m3' if a_is_mxfp8 else 'packed e2m1'}), got {tuple(a_data.shape)}" + ) + # Scales may be over-allocated (the reference quantizer pads rows to 256 and + # cols to 8); the kernel reads only the valid region via group offsets + strides, + # so require they merely *cover* [rows, K/32]. + if a_scale.ndim != 2 or a_scale.shape[0] < total_M or a_scale.shape[1] < grid_k: + raise ValueError( + f"a_scale must cover at least [{total_M}, {grid_k}], got {tuple(a_scale.shape)}" + ) + if ( + b_scale.ndim != 3 + or b_scale.shape[0] != G + or b_scale.shape[1] < N + or b_scale.shape[2] < grid_k + ): + raise ValueError( + f"b_scale must cover at least [{G}, {N}, {grid_k}], got {tuple(b_scale.shape)}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("E8M0 scales (a_scale, b_scale) must be uint8") + if a_data.element_size() != 1 or b_data.element_size() != 1: + raise TypeError("packed data operands (a_data, b_data) must be 1-byte") + devices = {a_data.device, a_scale.device, b_data.device, b_scale.device} + if len(devices) != 1: + raise ValueError(f"pre-quantized operands must share one device, got {devices}") + + def grouped_gemm_mxfp4_fprop_prequantized( a_data: torch.Tensor, a_scale: torch.Tensor, @@ -366,11 +416,14 @@ def grouped_gemm_mxfp4_fprop_prequantized( Returns: ``[total_M, N]`` output in ``out_dtype``. """ + _require_gfx950() + if b_data.ndim != 3: + raise ValueError(f"b_data (weights) must be 3-D [G, N, K/2], got rank {b_data.ndim}") N = b_data.shape[1] K = b_data.shape[2] * 2 # weights are packed 2 e2m1 elems/byte along K - _require_gfx950() _check_contract(K, "K") _check_splits(m_splits, a_data.shape[0], b_data.shape[0]) + _check_prequantized_operands(a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8) group_offs = _prefix_offsets(m_splits, a_data.device) kernel = grouped_gemm_a8w4_triton_kernel if a_is_mxfp8 else grouped_gemm_mxfp4_triton_kernel From 7ff27ca0a3fe0fa6aa449b838525faf2be256dc3 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 15:42:27 -0500 Subject: [PATCH 04/26] weight caching --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 21 ++++++++ .../pytorch/module/grouped_linear.py | 50 ++++++++++++------- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 7338051f92..b253f4881b 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -202,3 +202,24 @@ def test_grouped_linear_a8w4_backward_raises(): with pytest.raises(NotImplementedError, match="forward-only QAT"): out.backward(torch.randn_like(out)) + +def test_grouped_linear_a8w4_weight_cache_is_consistent(): + """is_first_microbatch weight caching must not change the a8w4 forward output.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with torch.no_grad(), te.autocast(enabled=True, recipe=rec): + out_nocache = model(inp, m_splits) # is_first_microbatch=None -> no caching + out_first = model(inp, m_splits, is_first_microbatch=True) # quantize + cache + out_reuse = model(inp, m_splits, is_first_microbatch=False) # reuse the cache + + assert torch.equal(out_first, out_reuse), "cached weights diverged from the first microbatch" + assert torch.equal(out_nocache, out_first), "weight caching changed the forward output" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8932fce0e0..eb29ae4708 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -964,6 +964,9 @@ def _forward_grouped_mxfp4_triton( weight_quantizers, activation_dtype, is_grad_enabled, + is_first_microbatch, + weight_workspaces, + cache_weight, ): """Grouped MXFP4 / a8w4 forward (ROCm Triton, gfx950). @@ -998,20 +1001,32 @@ def _forward_grouped_mxfp4_triton( # Weights: per-expert row-wise E2M1, stacked to the kernel's [G, N, *] layout. # The a4w4 backward additionally needs the col-wise weight for the dgrad kernel. + # Cache the quantized weights across microbatches (is_first_microbatch): the + # weight quant is a fixed per-call cost (~G*N*K) that otherwise dominates the + # grouped GEMM, so reuse the cache whenever this is not the first microbatch. want_wcol = is_grad_enabled and not a_is_mxfp8 - b_datas, b_scales = [], [] - wcol_datas, wcol_scales = [], [] - for i in range(num_gemms): - wq = weight_quantizers[i] - wq.set_usage(rowwise=True, columnwise=want_wcol) - wref = wq.quantize(weights[i].detach()) - b_datas.append(wref.data) - b_scales.append(wref.scale) - if want_wcol: - wcol_datas.append(wref.data_t) - wcol_scales.append(wref.scale_t) - b_data = torch.stack(b_datas, dim=0) - b_scale = torch.stack(b_scales, dim=0) + update_ws = is_first_microbatch is None or is_first_microbatch + cached = weight_workspaces[0] if weight_workspaces else None + new_workspaces = [None] * num_gemms + if not update_ws and cached is not None and (not want_wcol or cached[2] is not None): + b_data, b_scale, w_col_data, w_col_scale = cached + else: + b_datas, b_scales, wcol_datas, wcol_scales = [], [], [], [] + for i in range(num_gemms): + wq = weight_quantizers[i] + wq.set_usage(rowwise=True, columnwise=want_wcol) + wref = wq.quantize(weights[i].detach()) + b_datas.append(wref.data) + b_scales.append(wref.scale) + if want_wcol: + wcol_datas.append(wref.data_t) + wcol_scales.append(wref.scale_t) + b_data = torch.stack(b_datas, dim=0) + b_scale = torch.stack(b_scales, dim=0) + w_col_data = torch.stack(wcol_datas, dim=0) if want_wcol else None + w_col_scale = torch.stack(wcol_scales, dim=0) if want_wcol else None + if cache_weight: + new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) out = grouped_gemm_mxfp4_fprop_prequantized( aq.data, @@ -1034,11 +1049,9 @@ def _forward_grouped_mxfp4_triton( # wgrad re-quantizes the saved hp activation + grad_out in the kernel. ctx.grouped_mxfp4_requires_dgrad = inp.requires_grad ctx.grouped_mxfp4_weight_requires_grad = weights[0].requires_grad - ctx.save_for_backward( - a, torch.stack(wcol_datas, dim=0), torch.stack(wcol_scales, dim=0) - ) + ctx.save_for_backward(a, w_col_data, w_col_scale) - return out.view(-1, *inp.shape[1:-1], out.shape[-1]), [None] * num_gemms + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod def _backward_grouped_mxfp4_triton(ctx, grad_output): @@ -1218,6 +1231,9 @@ def forward( weight_quantizers=weight_quantizers, activation_dtype=activation_dtype, is_grad_enabled=is_grad_enabled, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, ) # Configure quantizers From 94be09de5e3319753cdf3470facef7ee4b1d931c Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 16:11:57 -0500 Subject: [PATCH 05/26] fused downcast --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 17 ++++++++++ .../pytorch/module/grouped_linear.py | 12 +++++-- .../pytorch/triton_kernels/cast_transpose.py | 34 +++++++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4_impl.py | 23 +++++++++---- 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index b253f4881b..857290da43 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -17,6 +17,8 @@ from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( _row_operand, + _row_operand_mxfp8, + _row_operand_mxfp8_torch, grouped_gemm_mxfp4_dgrad, grouped_gemm_mxfp4_fprop_prequantized, ) @@ -223,3 +225,18 @@ def test_grouped_linear_a8w4_weight_cache_is_consistent(): assert torch.equal(out_first, out_reuse), "cached weights diverged from the first microbatch" assert torch.equal(out_nocache, out_first), "weight caching changed the forward output" + + +def test_row_operand_mxfp8_fused_matches_torch(): + """The fused MXFP8 e4m3 downcast is bit-identical to the torch reference.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + torch.manual_seed(0) + for M, K in [(384, 256), (130, 256), (293, 3584)]: + x = torch.randn(M, K, device=device, dtype=dtype) + d_ref, s_ref = _row_operand_mxfp8_torch(x) + d_fused, s_fused = _row_operand_mxfp8(x) # fused on gfx950 + assert torch.equal(s_fused, s_ref), f"E8M0 scales differ at {(M, K)}" + assert torch.equal( + d_fused.view(torch.uint8), d_ref.view(torch.uint8) + ), f"e4m3 data bytes differ at {(M, K)}" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index eb29ae4708..08fc4cb135 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -980,6 +980,7 @@ def _forward_grouped_mxfp4_triton( is forward-only (see :meth:`_backward_grouped_mxfp4_triton`). """ from ..triton_kernels.grouped_gemm_mxfp4_impl import ( + _row_operand_mxfp8, grouped_gemm_mxfp4_fprop_prequantized, ) from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef @@ -997,7 +998,12 @@ def _forward_grouped_mxfp4_triton( input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=False) a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) - aq = input_quantizer.quantize(a) + if a_is_mxfp8: + # a8w4 activation: fused single-pass MXFP8 e4m3 downcast on gfx950 + a_data, a_scale = _row_operand_mxfp8(a) + else: + aq = input_quantizer.quantize(a) + a_data, a_scale = aq.data, aq.scale # Weights: per-expert row-wise E2M1, stacked to the kernel's [G, N, *] layout. # The a4w4 backward additionally needs the col-wise weight for the dgrad kernel. @@ -1029,8 +1035,8 @@ def _forward_grouped_mxfp4_triton( new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) out = grouped_gemm_mxfp4_fprop_prequantized( - aq.data, - aq.scale, + a_data, + a_scale, b_data, b_scale, m_splits_list, diff --git a/transformer_engine/pytorch/triton_kernels/cast_transpose.py b/transformer_engine/pytorch/triton_kernels/cast_transpose.py index 377e6a6331..bcd432fe5a 100644 --- a/transformer_engine/pytorch/triton_kernels/cast_transpose.py +++ b/transformer_engine/pytorch/triton_kernels/cast_transpose.py @@ -868,6 +868,40 @@ def te_cast_transpose_mxfp8_triton(input, out, noop_flag=None): colwise_scale_M, colwise_scale_N, max_fp8, BLOCK_X, BLOCK_Y, GROUP_Y, MXFP8_BLOCK_SCALING_SIZE, USE_ROWWISE_SCALING, USE_COLWISE_SCALING) + +def mxfp8_e4m3_rowwise_downcast(x): + """Fused single-pass MXFP8 e4m3 row-wise downcast (Kimi-K3 a8w4 activation). + + ``x`` [M, K] (K a multiple of 32) -> (data [M, K] e4m3, scale [M, K/32] uint8 + E8M0), using the ceil scale rule ``ceil(log2(amax_block / 448))``. Reuses + :func:`_cast_transpose_triton_mxfp8` row-wise only, so it is bit-identical to the + torch reference ``_row_operand_mxfp8_torch`` but single-pass. + """ + assert x.shape[-1] % MXFP8_BLOCK_SCALING_SIZE == 0, "K must be a multiple of 32" + x = x.contiguous() + M, K = x.shape + grid_k = K // MXFP8_BLOCK_SCALING_SIZE + fp8_dtype = tex.DType.kFloat8E4M3 + data = torch.empty((M, K), dtype=te_dtype_to_torch_dtype(fp8_dtype), device=x.device) + scale = torch.empty((M, grid_k), dtype=torch.uint8, device=x.device) + rowwise_y = triton.reinterpret(data, te_dtype_to_triton_dtype(fp8_dtype)) + BLOCK_X = BLOCK_Y = 64 + GROUP_Y = MXFP8_BLOCK_SCALING_SIZE + grid = lambda META: (triton.cdiv(M, META["BLOCK_Y"]) * triton.cdiv(K, META["BLOCK_X"]),) + _cast_transpose_triton_mxfp8[grid]( + x, rowwise_y, None, + x.stride(0), x.stride(1), + M, K, + scale, scale.stride(0), scale.stride(1), + M, grid_k, + None, 1, 1, + 1, 1, + get_fp8_max(fp8_dtype), BLOCK_X, BLOCK_Y, GROUP_Y, MXFP8_BLOCK_SCALING_SIZE, + True, False, + ) + return data, scale + + def te_dequantize_mxfp8_triton(input, dtype): input_metadata = input.get_metadata() use_rowwise_scaling = input_metadata["rowwise_data"] is not 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 84fa5a0dc8..6a7c67d55a 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -19,6 +19,7 @@ grouped_gemm_mxfp4_triton_kernel, grouped_gemm_mxfp4_variable_k_triton_kernel, ) +from .cast_transpose import mxfp8_e4m3_rowwise_downcast # Logical contraction tile the kernels step by (see grouped_gemm_mxfp4.py). BLOCK_SIZE_K = 128 @@ -145,15 +146,13 @@ 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]: +def _row_operand_mxfp8_torch(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. + Plain OCP MXFP8: unpacked e4m3 element bytes + one E8M0 scale per 1x32 block, with + the Kimi-K3 ``ceil(log2(amax / 448))`` block-scale rule. Torch multi-pass reference; + :func:`_row_operand_mxfp8` uses the fused single-pass kernel on gfx950. ``K`` (the + contraction) is already a 32-multiple; ``M`` needs no padding (scale is per-row). """ fmax = torch.finfo(torch.float8_e4m3fn).max M, K = x.shape @@ -166,6 +165,16 @@ def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: return data, scale +def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Row-wise MXFP8 (e4m3) downcast (ceil rule): fused single-pass Triton on gfx950 + (:func:`mxfp8_e4m3_rowwise_downcast`), else the bit-identical torch reference + :func:`_row_operand_mxfp8_torch`. Used by the a8w4 forward (GATE-2 fused quant). + """ + if _is_gfx950(): + return mxfp8_e4m3_rowwise_downcast(x) + return _row_operand_mxfp8_torch(x) + + def _quantize_weights_row_col( weights: Sequence[torch.Tensor], ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: From 898ef53a329315868e3a6ccf6d0e95a2e7c54068 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 22:22:28 -0500 Subject: [PATCH 06/26] adjust comments Updated test description for clarity. --- tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 857290da43..945bc1ea1f 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -38,7 +38,6 @@ _recipe_ok, _reason = te.is_mxfp4_available(return_reason=True) pytestmark = pytest.mark.skipif(not _recipe_ok, reason=_reason) -# fp32 matmul vs MFMA on identical quantized bytes: only accumulation differs. _REL_TOL = 5e-2 @@ -65,7 +64,7 @@ def _reference_grouped_a8w4(aq, weight_refs, m_splits, N): def test_grouped_prequantized_a8w4_matches_reference(): - """The kernel on reference bytes agrees with the reference qgemm (Option B).""" + """The kernel on reference bytes agrees with the reference qgemm.""" _isolate() device, dtype = "cuda", torch.bfloat16 K, N = 256, 256 # K % 128 == 0 -> loop_k > 1 From 41e94da7634ae4d301b6fa2adab39c5caf3b9bf8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 28 Sep 2026 16:58:30 -0500 Subject: [PATCH 07/26] swizzle a8w4 --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 7 ++ .../pytorch/module/grouped_linear.py | 9 ++- .../pytorch/triton_kernels/cast_transpose.py | 19 +++-- .../triton_kernels/grouped_gemm_mxfp4.py | 74 +++++++++++++----- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 77 ++++++++++++++----- 5 files changed, 137 insertions(+), 49 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 945bc1ea1f..9c5454abbe 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -239,3 +239,10 @@ def test_row_operand_mxfp8_fused_matches_torch(): assert torch.equal( d_fused.view(torch.uint8), d_ref.view(torch.uint8) ), f"e4m3 data bytes differ at {(M, K)}" + # transposed-scale (a8w4 swizzle) path: [K/32, M], must equal the plain scale.T + _, s_fused_t = _row_operand_mxfp8(x, transpose_scale=True) + _, s_ref_t = _row_operand_mxfp8_torch(x, transpose_scale=True) + assert torch.equal(s_fused_t, s_ref_t), f"transposed E8M0 scales differ at {(M, K)}" + assert torch.equal( + s_fused_t, s_ref.t().contiguous() + ), f"transposed scale != plain scale.T at {(M, K)}" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 08fc4cb135..c999d5a730 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -999,8 +999,9 @@ def _forward_grouped_mxfp4_triton( input_quantizer.set_usage(rowwise=True, columnwise=False) a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) if a_is_mxfp8: - # a8w4 activation: fused single-pass MXFP8 e4m3 downcast on gfx950 - a_data, a_scale = _row_operand_mxfp8(a) + # a8w4 activation: fused single-pass MXFP8 e4m3 downcast on gfx950; emit the + # scale transposed ([K/32, M]) for the kernel's swizzle path. + a_data, a_scale = _row_operand_mxfp8(a, transpose_scale=True) else: aq = input_quantizer.quantize(a) a_data, a_scale = aq.data, aq.scale @@ -1029,6 +1030,9 @@ def _forward_grouped_mxfp4_triton( wcol_scales.append(wref.scale_t) b_data = torch.stack(b_datas, dim=0) b_scale = torch.stack(b_scales, dim=0) + if a_is_mxfp8: + # a8w4 swizzle: cache the weight scale transposed ([G, K/32, N]) + b_scale = b_scale.transpose(1, 2).contiguous() w_col_data = torch.stack(wcol_datas, dim=0) if want_wcol else None w_col_scale = torch.stack(wcol_scales, dim=0) if want_wcol else None if cache_weight: @@ -1042,6 +1046,7 @@ def _forward_grouped_mxfp4_triton( m_splits_list, a_is_mxfp8=a_is_mxfp8, out_dtype=activation_dtype, + scale_transposed=a_is_mxfp8, ) if is_grad_enabled: diff --git a/transformer_engine/pytorch/triton_kernels/cast_transpose.py b/transformer_engine/pytorch/triton_kernels/cast_transpose.py index bcd432fe5a..5bd0d30315 100644 --- a/transformer_engine/pytorch/triton_kernels/cast_transpose.py +++ b/transformer_engine/pytorch/triton_kernels/cast_transpose.py @@ -869,13 +869,13 @@ def te_cast_transpose_mxfp8_triton(input, out, noop_flag=None): max_fp8, BLOCK_X, BLOCK_Y, GROUP_Y, MXFP8_BLOCK_SCALING_SIZE, USE_ROWWISE_SCALING, USE_COLWISE_SCALING) -def mxfp8_e4m3_rowwise_downcast(x): +def mxfp8_e4m3_rowwise_downcast(x, transpose_scale=False): """Fused single-pass MXFP8 e4m3 row-wise downcast (Kimi-K3 a8w4 activation). - ``x`` [M, K] (K a multiple of 32) -> (data [M, K] e4m3, scale [M, K/32] uint8 - E8M0), using the ceil scale rule ``ceil(log2(amax_block / 448))``. Reuses - :func:`_cast_transpose_triton_mxfp8` row-wise only, so it is bit-identical to the - torch reference ``_row_operand_mxfp8_torch`` but single-pass. + ``x`` [M, K] (K a multiple of 32) -> (data [M, K] e4m3, scale uint8 E8M0), using the + ceil scale rule ``ceil(log2(amax_block / 448))``. Reuses :func:`_cast_transpose_triton_mxfp8` + row-wise only, so it is bit-identical to the torch reference ``_row_operand_mxfp8_torch`` + but single-pass. """ assert x.shape[-1] % MXFP8_BLOCK_SCALING_SIZE == 0, "K must be a multiple of 32" x = x.contiguous() @@ -883,7 +883,12 @@ def mxfp8_e4m3_rowwise_downcast(x): grid_k = K // MXFP8_BLOCK_SCALING_SIZE fp8_dtype = tex.DType.kFloat8E4M3 data = torch.empty((M, K), dtype=te_dtype_to_torch_dtype(fp8_dtype), device=x.device) - scale = torch.empty((M, grid_k), dtype=torch.uint8, device=x.device) + if transpose_scale: + scale = torch.empty((grid_k, M), dtype=torch.uint8, device=x.device) + scale_stride_m, scale_stride_k = scale.stride(1), scale.stride(0) + else: + scale = torch.empty((M, grid_k), dtype=torch.uint8, device=x.device) + scale_stride_m, scale_stride_k = scale.stride(0), scale.stride(1) rowwise_y = triton.reinterpret(data, te_dtype_to_triton_dtype(fp8_dtype)) BLOCK_X = BLOCK_Y = 64 GROUP_Y = MXFP8_BLOCK_SCALING_SIZE @@ -892,7 +897,7 @@ def mxfp8_e4m3_rowwise_downcast(x): x, rowwise_y, None, x.stride(0), x.stride(1), M, K, - scale, scale.stride(0), scale.stride(1), + scale, scale_stride_m, scale_stride_k, M, grid_k, None, 1, 1, 1, 1, diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 1093813361..e11dcb46a2 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -202,6 +202,8 @@ def _grouped_mxfp4_persistent_gemm_kernel( B_PACK: tl.constexpr = 2, A_FMT: tl.constexpr = "e2m1", # per-operand tl.dot_scaled format ("e2m1"/"e4m3"/"e5m2") B_FMT: tl.constexpr = "e2m1", + SCALE_TRANSPOSED: tl.constexpr = False, # scales pre-transposed ([K/32,M]/[G,K/32,N]): + # load (BM,BK/32)/(BN,BK/32) directly ): pid = tl.program_id(0) if NUM_XCDS != 1: @@ -259,21 +261,36 @@ def _grouped_mxfp4_persistent_gemm_kernel( A_BASE = A + (m_rd + rm[:, None]) * stride_am + rk_a[None, :] * stride_ak B_BASE = B + group_idx.to(tl.int64) * stride_bg + rk_b[:, None] * stride_bk + rn[None, :] * stride_bn - AS_BASE = A_scale + rks[:, None] * stride_ask + (m_rd + rm[None, :]) * stride_asm - BS_BASE = ( - B_scale - + group_idx.to(tl.int64) * stride_bsg - + rks[:, None] * stride_bsk - + rn[None, :] * stride_bsn - ) + if SCALE_TRANSPOSED: + # scales stored [K/32, M] / [G, K/32, N] -> a (BM, BK/32) / (BN, BK/32) tile loads + # directly (M/N contiguous) + AS_BASE = A_scale + (m_rd + rm[:, None]) * stride_asm + rks[None, :] * stride_ask + BS_BASE = ( + B_scale + + group_idx.to(tl.int64) * stride_bsg + + rn[:, None] * stride_bsn + + rks[None, :] * stride_bsk + ) + else: + AS_BASE = A_scale + rks[:, None] * stride_ask + (m_rd + rm[None, :]) * stride_asm + BS_BASE = ( + B_scale + + group_idx.to(tl.int64) * stride_bsg + + rks[:, None] * stride_bsk + + rn[None, :] * stride_bsn + ) acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) loop_k = K // BLOCK_SIZE_K for ki in range(0, loop_k): a = tl.load(tl.multiple_of(A_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) # (BM, BK_A) b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) # (BK_B, BN) - a_s = tl.trans(tl.load(AS_BASE)) # (BK/32, BM) -> (BM, BK/32) - b_s = tl.trans(tl.load(BS_BASE)) # (BK/32, BN) -> (BN, BK/32) + if SCALE_TRANSPOSED: + a_s = tl.load(AS_BASE) # (BM, BK/32) directly + b_s = tl.load(BS_BASE) # (BN, BK/32) directly + else: + a_s = tl.trans(tl.load(AS_BASE)) # (BK/32, BM) -> (BM, BK/32) + b_s = tl.trans(tl.load(BS_BASE)) # (BK/32, BN) -> (BN, BK/32) acc = tl.dot_scaled(a, a_s, A_FMT, b, b_s, B_FMT, acc) A_BASE += BK_A * stride_ak B_BASE += BK_B * stride_bk @@ -289,6 +306,15 @@ def _grouped_mxfp4_persistent_gemm_kernel( tl.store(C_, c, c_mask) +def _scale_strides(a_s, b_s, transposed): + """(stride_asm, stride_ask, stride_bsg, stride_bsn, stride_bsk) for the persistent + kernel. Plain scales are [M,K/32] / [G,N,K/32]; when ``transposed`` they are + [K/32,M] / [G,K/32,N] so the kernel loads (BM/BN, BK/32) directly.""" + if transposed: + return a_s.stride(1), a_s.stride(0), b_s.stride(0), b_s.stride(2), b_s.stride(1) + return a_s.stride(0), a_s.stride(1), b_s.stride(0), b_s.stride(1), b_s.stride(2) + + @_scoped_amd_knobs(is_tn=False) def grouped_gemm_mxfp4_triton_kernel( a, @@ -301,6 +327,7 @@ def grouped_gemm_mxfp4_triton_kernel( group_offs_out=None, out_dtype=torch.bfloat16, num_cu=None, + scale_transposed=False, ): """A(total_M, K/2) @ B(G, N, K/2)^T -> C. FP4-packed data, e8m0 uint8 scales. @@ -316,6 +343,9 @@ def grouped_gemm_mxfp4_triton_kernel( b_s = b_scale.view(torch.uint8) a_u8 = a.view(torch.uint8) b_u8 = b.view(torch.uint8) + stride_asm, stride_ask, stride_bsg, stride_bsn, stride_bsk = _scale_strides( + a_s, b_s, scale_transposed + ) 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] @@ -347,11 +377,11 @@ def grouped_gemm_mxfp4_triton_kernel( 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) + stride_asm, # a_s is (total_M, K/32) plain, or (K/32, total_M) transposed + stride_ask, + stride_bsg, + stride_bsn, # b_s is (G, N, K/32) plain, or (G, K/32, N) transposed + stride_bsk, BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, GROUP_SIZE_M=GM, @@ -362,6 +392,7 @@ def grouped_gemm_mxfp4_triton_kernel( VEC=VEC_SIZE, A_PACK=2, # fp4 A; passed explicitly so A_PACK/B_PACK bind into the autotune key B_PACK=2, + SCALE_TRANSPOSED=scale_transposed, ) return c @@ -378,6 +409,7 @@ def grouped_gemm_a8w4_triton_kernel( group_offs_out=None, out_dtype=torch.bfloat16, num_cu=None, + scale_transposed=False, ): """a8w4 forward: C = A(e4m3) @ B(e2m1)^T, grouped along M (routed experts). @@ -396,6 +428,9 @@ def grouped_gemm_a8w4_triton_kernel( 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) + stride_asm, stride_ask, stride_bsg, stride_bsn, stride_bsk = _scale_strides( + a_s, b_s, scale_transposed + ) 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] @@ -424,11 +459,11 @@ def grouped_gemm_a8w4_triton_kernel( 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), + stride_asm, + stride_ask, + stride_bsg, + stride_bsn, + stride_bsk, BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, GROUP_SIZE_M=GM, @@ -441,6 +476,7 @@ def grouped_gemm_a8w4_triton_kernel( A_FMT="e4m3", B_PACK=2, # e2m1: 2 elems/byte -> K/2 bytes B_FMT="e2m1", + SCALE_TRANSPOSED=scale_transposed, ) 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 6a7c67d55a..dee157bf99 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -146,13 +146,14 @@ def _row_col_operand( return row_data, row_scale, col_data, col_scale -def _row_operand_mxfp8_torch(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """Row-wise MXFP8 (e4m3): ``x`` [M, K] -> (data [M, K] e4m3, scale [M, K/32] u8). +def _row_operand_mxfp8_torch(x: torch.Tensor, transpose_scale: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """Row-wise MXFP8 (e4m3): ``x`` [M, K] -> (data [M, K] e4m3, scale 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. Torch multi-pass reference; :func:`_row_operand_mxfp8` uses the fused single-pass kernel on gfx950. ``K`` (the contraction) is already a 32-multiple; ``M`` needs no padding (scale is per-row). + ``transpose_scale`` emits the scale ``[K/32, M]`` instead of ``[M, K/32]``. """ fmax = torch.finfo(torch.float8_e4m3fn).max M, K = x.shape @@ -162,17 +163,20 @@ def _row_operand_mxfp8_torch(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso 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) + if transpose_scale: + scale = scale.t().contiguous() return data, scale -def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: +def _row_operand_mxfp8(x: torch.Tensor, transpose_scale: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: """Row-wise MXFP8 (e4m3) downcast (ceil rule): fused single-pass Triton on gfx950 (:func:`mxfp8_e4m3_rowwise_downcast`), else the bit-identical torch reference :func:`_row_operand_mxfp8_torch`. Used by the a8w4 forward (GATE-2 fused quant). + ``transpose_scale`` emits the scale ``[K/32, M]``. """ if _is_gfx950(): - return mxfp8_e4m3_rowwise_downcast(x) - return _row_operand_mxfp8_torch(x) + return mxfp8_e4m3_rowwise_downcast(x, transpose_scale=transpose_scale) + return _row_operand_mxfp8_torch(x, transpose_scale=transpose_scale) def _quantize_weights_row_col( @@ -329,7 +333,12 @@ def grouped_gemm_a8w4_fprop( 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 + a_data, a_scale = _row_operand_mxfp8(a, transpose_scale=True) # a_scale [K/32, total_M] + # a8w4 swizzle: the kernel consumes both scales transposed. + # Normalise the weight scale to [G, K/32, N] (plain weight quant / weight_row is + # [G, N, K/32]; a caller may also pass it already transposed to keep the cache free). + if b_scale.shape[-1] != N: + b_scale = b_scale.transpose(1, 2).contiguous() group_offs = _prefix_offsets(m_splits, a.device) return grouped_gemm_a8w4_triton_kernel( a_data, @@ -342,6 +351,7 @@ def grouped_gemm_a8w4_fprop( group_offs_out=group_offs, out_dtype=out_dtype, num_cu=num_cu, + scale_transposed=True, ) @@ -352,6 +362,7 @@ def _check_prequantized_operands( b_scale: torch.Tensor, *, a_is_mxfp8: bool, + scale_transposed: bool = False, ) -> None: """Validate pre-quantized operand ranks/shapes/dtypes/devices before launch. @@ -372,20 +383,38 @@ def _check_prequantized_operands( ) # Scales may be over-allocated (the reference quantizer pads rows to 256 and # cols to 8); the kernel reads only the valid region via group offsets + strides, - # so require they merely *cover* [rows, K/32]. - if a_scale.ndim != 2 or a_scale.shape[0] < total_M or a_scale.shape[1] < grid_k: - raise ValueError( - f"a_scale must cover at least [{total_M}, {grid_k}], got {tuple(a_scale.shape)}" - ) - if ( - b_scale.ndim != 3 - or b_scale.shape[0] != G - or b_scale.shape[1] < N - or b_scale.shape[2] < grid_k - ): - raise ValueError( - f"b_scale must cover at least [{G}, {N}, {grid_k}], got {tuple(b_scale.shape)}" - ) + # so require they merely *cover* [rows, K/32]. When ``scale_transposed`` the + # scales are stored [K/32, total_M] / [G, K/32, N]. + if scale_transposed: + if a_scale.ndim != 2 or a_scale.shape[0] < grid_k or a_scale.shape[1] < total_M: + raise ValueError( + f"transposed a_scale must cover at least [{grid_k}, {total_M}]," + f" got {tuple(a_scale.shape)}" + ) + if ( + b_scale.ndim != 3 + or b_scale.shape[0] != G + or b_scale.shape[1] < grid_k + or b_scale.shape[2] < N + ): + raise ValueError( + f"transposed b_scale must cover at least [{G}, {grid_k}, {N}]," + f" got {tuple(b_scale.shape)}" + ) + else: + if a_scale.ndim != 2 or a_scale.shape[0] < total_M or a_scale.shape[1] < grid_k: + raise ValueError( + f"a_scale must cover at least [{total_M}, {grid_k}], got {tuple(a_scale.shape)}" + ) + if ( + b_scale.ndim != 3 + or b_scale.shape[0] != G + or b_scale.shape[1] < N + or b_scale.shape[2] < grid_k + ): + raise ValueError( + f"b_scale must cover at least [{G}, {N}, {grid_k}], got {tuple(b_scale.shape)}" + ) if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: raise TypeError("E8M0 scales (a_scale, b_scale) must be uint8") if a_data.element_size() != 1 or b_data.element_size() != 1: @@ -405,6 +434,7 @@ def grouped_gemm_mxfp4_fprop_prequantized( a_is_mxfp8: bool, out_dtype: torch.dtype = torch.bfloat16, num_cu: Optional[int] = None, + scale_transposed: bool = False, ) -> torch.Tensor: """Grouped forward from already-quantized operands: ``C[g] = A[g] @ W[g]^T``. @@ -421,6 +451,8 @@ def grouped_gemm_mxfp4_fprop_prequantized( m_splits: per-group token counts (len G). a_is_mxfp8: ``True`` routes the MXFP8(e4m3) x MXFP4(e2m1) a8w4 kernel; ``False`` the MXFP4 x MXFP4 kernel. + scale_transposed: when ``True`` the E8M0 scales are laid out ``[K/32, total_M]`` + / ``[G, K/32, N]`` so the kernel loads them without the per-iter ``tl.trans``. Returns: ``[total_M, N]`` output in ``out_dtype``. @@ -432,7 +464,9 @@ def grouped_gemm_mxfp4_fprop_prequantized( K = b_data.shape[2] * 2 # weights are packed 2 e2m1 elems/byte along K _check_contract(K, "K") _check_splits(m_splits, a_data.shape[0], b_data.shape[0]) - _check_prequantized_operands(a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8) + _check_prequantized_operands( + a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8, scale_transposed=scale_transposed + ) group_offs = _prefix_offsets(m_splits, a_data.device) kernel = grouped_gemm_a8w4_triton_kernel if a_is_mxfp8 else grouped_gemm_mxfp4_triton_kernel @@ -447,6 +481,7 @@ def grouped_gemm_mxfp4_fprop_prequantized( group_offs_out=group_offs, out_dtype=out_dtype, num_cu=num_cu, + scale_transposed=scale_transposed, ) From 92a19b8990087137c1c30f27e7d6c7d7af95e68c Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 10:22:23 -0500 Subject: [PATCH 08/26] use native quantizer for a4w4 --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 20 +++++++++++++++++++ .../pytorch/module/grouped_linear.py | 9 ++++----- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 9c5454abbe..4b7c502793 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -159,6 +159,26 @@ def test_grouped_prequantized_a4w4_dgrad_matches_reference(): assert rel < _REL_TOL, f"a4w4 dgrad on reference col weight disagrees: rel={rel:.4f}" +def test_a4w4_activation_native_matches_reference(): + """ + The a4w4 module activation uses the native GPU MXFP4 quantizer; its dequantized values must match the reference CPU quantizer. + """ + _isolate() + device, dtype = "cuda", torch.bfloat16 + M, K = 384, 256 + torch.manual_seed(0) + x = torch.randn(M, K, device=device, dtype=dtype) + + d_nat, s_nat = _row_operand(x) # native GPU path + hp_nat = mxfp4_to_f32(d_nat) * e8m0_to_f32(s_nat).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + + ref = _e2m1_ref(rowwise=True, columnwise=False).quantize(x) # pure-NumPy reference + hp_ref = _deq_ref_operand(ref.data, ref.scale, ref, default_e4m3=False) + + rel = (hp_nat - hp_ref).norm() / hp_ref.norm().clamp_min(1e-12) + assert rel < _REL_TOL, f"native a4w4 activation disagrees with reference: rel={rel:.4f}" + + def test_grouped_linear_a4w4_recipe_backward(): """te.GroupedLinear a4w4 fwd+bwd routes to the low-precision dgrad/wgrad kernels.""" _isolate() diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c999d5a730..851cbcd512 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -980,6 +980,7 @@ def _forward_grouped_mxfp4_triton( is forward-only (see :meth:`_backward_grouped_mxfp4_triton`). """ from ..triton_kernels.grouped_gemm_mxfp4_impl import ( + _row_operand, _row_operand_mxfp8, grouped_gemm_mxfp4_fprop_prequantized, ) @@ -992,19 +993,17 @@ def _forward_grouped_mxfp4_triton( ) # Activation: quantize the whole grouped [total_M, K] once (row-wise scales - # do not cross groups). Detach: the reference quantizer casts via NumPy and - # the low-precision forward is not differentiated through (custom backward). + # do not cross groups). Detach: the low-precision forward is not differentiated + # through (a custom backward handles grads). a = inp.reshape(-1, in_features).to(activation_dtype).contiguous().detach() input_quantizer = input_quantizers[0] - input_quantizer.set_usage(rowwise=True, columnwise=False) a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) if a_is_mxfp8: # a8w4 activation: fused single-pass MXFP8 e4m3 downcast on gfx950; emit the # scale transposed ([K/32, M]) for the kernel's swizzle path. a_data, a_scale = _row_operand_mxfp8(a, transpose_scale=True) else: - aq = input_quantizer.quantize(a) - a_data, a_scale = aq.data, aq.scale + a_data, a_scale = _row_operand(a) # Weights: per-expert row-wise E2M1, stacked to the kernel's [G, N, *] layout. # The a4w4 backward additionally needs the col-wise weight for the dgrad kernel. From c893901108d61c7c8b5da1e55a4d7e47cecbda22 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 12:05:15 -0500 Subject: [PATCH 09/26] address copilot reviews --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 22 +++++++++++++++++++ .../pytorch/module/grouped_linear.py | 16 +++++++++++++- .../pytorch/triton_kernels/cast_transpose.py | 7 ++++-- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 2 ++ 4 files changed, 44 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 4b7c502793..f0472bb64b 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -266,3 +266,25 @@ def test_row_operand_mxfp8_fused_matches_torch(): assert torch.equal( s_fused_t, s_ref.t().contiguous() ), f"transposed scale != plain scale.T at {(M, K)}" + + +def test_grouped_prequantized_empty_batch(): + """Empty MoE routing batch (total_M=0): the grouped path returns [0, N] with no + zero-sized grid launch (the fused downcast and the GEMM both short-circuit).""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + G, K, N = 4, 256, 256 + m_splits = [0] * G + weights = [torch.randn(N, K, device=device, dtype=dtype) for _ in range(G)] + b_ops = [_row_operand(w) for w in weights] + b_data = torch.stack([d for d, _ in b_ops]) + b_scale = torch.stack([s for _, s in b_ops]) + + a = torch.empty(0, K, device=device, dtype=dtype) + a_data, a_scale = _row_operand_mxfp8(a) # fused downcast must skip the zero-sized grid + assert a_data.shape == (0, K) and a_scale.numel() == 0 + + out = grouped_gemm_mxfp4_fprop_prequantized( + a_data, a_scale, b_data, b_scale, m_splits, a_is_mxfp8=True, out_dtype=dtype + ) + assert out.shape == (0, N) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 851cbcd512..edcf13b7bf 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1014,7 +1014,21 @@ def _forward_grouped_mxfp4_triton( update_ws = is_first_microbatch is None or is_first_microbatch cached = weight_workspaces[0] if weight_workspaces else None new_workspaces = [None] * num_gemms - if not update_ws and cached is not None and (not want_wcol or cached[2] is not None): + reuse = not update_ws and cached is not None and (not want_wcol or cached[2] is not None) + if reuse: + # Validate the cached workspace before trusting it (mirrors the blockwise path): + # our 4-tuple whose stacked b_data matches the current weights. + expected = (num_gemms, weights[0].size(0), in_features // 2) + if not ( + isinstance(cached, tuple) + and len(cached) == 4 + and isinstance(cached[0], torch.Tensor) + and tuple(cached[0].shape) == expected + ): + raise RuntimeError( + "Cached grouped MXFP4 weight workspace is incompatible with the current " + f"weights; expected a 4-tuple whose b_data has shape {expected}." + ) b_data, b_scale, w_col_data, w_col_scale = cached else: b_datas, b_scales, wcol_datas, wcol_scales = [], [], [], [] diff --git a/transformer_engine/pytorch/triton_kernels/cast_transpose.py b/transformer_engine/pytorch/triton_kernels/cast_transpose.py index 5bd0d30315..9848d72363 100644 --- a/transformer_engine/pytorch/triton_kernels/cast_transpose.py +++ b/transformer_engine/pytorch/triton_kernels/cast_transpose.py @@ -874,8 +874,9 @@ def mxfp8_e4m3_rowwise_downcast(x, transpose_scale=False): ``x`` [M, K] (K a multiple of 32) -> (data [M, K] e4m3, scale uint8 E8M0), using the ceil scale rule ``ceil(log2(amax_block / 448))``. Reuses :func:`_cast_transpose_triton_mxfp8` - row-wise only, so it is bit-identical to the torch reference ``_row_operand_mxfp8_torch`` - but single-pass. + row-wise only, so it matches the torch reference ``_row_operand_mxfp8_torch`` bit-for-bit for + non-zero blocks (an all-zero 32-block may get a different E8M0 scale byte, but its data is zero + so it dequantizes to zero either way) but is single-pass. """ assert x.shape[-1] % MXFP8_BLOCK_SCALING_SIZE == 0, "K must be a multiple of 32" x = x.contiguous() @@ -889,6 +890,8 @@ def mxfp8_e4m3_rowwise_downcast(x, transpose_scale=False): else: scale = torch.empty((M, grid_k), dtype=torch.uint8, device=x.device) scale_stride_m, scale_stride_k = scale.stride(0), scale.stride(1) + if M == 0: # empty MoE routing batch: nothing to quantize, skip the zero-sized grid + return data, scale rowwise_y = triton.reinterpret(data, te_dtype_to_triton_dtype(fp8_dtype)) BLOCK_X = BLOCK_Y = 64 GROUP_Y = MXFP8_BLOCK_SCALING_SIZE 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 dee157bf99..19c6768831 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -463,6 +463,8 @@ def grouped_gemm_mxfp4_fprop_prequantized( N = b_data.shape[1] K = b_data.shape[2] * 2 # weights are packed 2 e2m1 elems/byte along K _check_contract(K, "K") + if a_data.shape[0] == 0: # empty MoE routing batch: no tokens for this GPU's experts + return torch.empty((0, N), dtype=out_dtype, device=a_data.device) _check_splits(m_splits, a_data.shape[0], b_data.shape[0]) _check_prequantized_operands( a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8, scale_transposed=scale_transposed From fcd002d3d1d74cfe49ed09015faec27e90988851 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 13:48:42 -0500 Subject: [PATCH 10/26] address comments --- .../pytorch/module/grouped_linear.py | 17 ++++++++++++++++- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 18 ++++++++++++------ 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index edcf13b7bf..6c7002c5b4 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -591,6 +591,10 @@ def _is_grouped_mxfp4_triton_supported( debug, unpad_output, actual_m_splits, + out, + dgrad_out, + fuse_wgrad_accumulation, + wgrad_store, in_features, out_features, ) -> bool: @@ -598,7 +602,10 @@ def _is_grouped_mxfp4_triton_supported( factory yields the plain-layout MXFP4 reference quantizers. a8w4 is forward-only (backward raises); a4w4 also wires the low-precision - dgrad/wgrad (see :meth:`_forward_grouped_mxfp4_triton`). + dgrad/wgrad (see :meth:`_forward_grouped_mxfp4_triton`). Gated off for options this + self-contained path does not honor -- caller ``out``/``dgrad_out`` buffers, + ``fuse_wgrad_accumulation``, and delayed wgrad (``wgrad_store``) -- so those fall + back to the standard path. """ if not (IS_HIP_EXTENSION and fp8 and recipe is not None and recipe.custom()): return False @@ -614,6 +621,10 @@ def _is_grouped_mxfp4_triton_supported( or debug or unpad_output or actual_m_splits is not None + or out is not None + or dgrad_out is not None + or fuse_wgrad_accumulation + or (wgrad_store is not None and wgrad_store.delay_wgrad_compute()) ): return False try: @@ -1243,6 +1254,10 @@ def forward( debug=debug, unpad_output=unpad_output, actual_m_splits=actual_m_splits, + out=out, + dgrad_out=dgrad_out, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + wgrad_store=wgrad_store, in_features=weights[0].size(-1), out_features=weights[0].size(0), ): 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 19c6768831..06aca167a5 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -297,6 +297,7 @@ def grouped_gemm_a8w4_fprop( out_dtype: torch.dtype = torch.bfloat16, num_cu: Optional[int] = None, weight_row: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + weight_scale_transposed: bool = False, ) -> torch.Tensor: """Grouped a8w4 forward: ``C[g] = A[g] @ W[g]^T`` (contract K). @@ -310,6 +311,9 @@ def grouped_gemm_a8w4_fprop( 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. + weight_scale_transposed: set ``True`` when ``weight_row``'s scale is already the + transposed ``[G, K/32, N]`` swizzle layout (a cached transpose); otherwise the + plain ``[G, N, K/32]`` scale is transposed here. Returns: [total_M, N] output in ``out_dtype``. @@ -322,6 +326,7 @@ def grouped_gemm_a8w4_fprop( b_data, b_scale = weight_row N = b_data.shape[1] _check_splits(m_splits, a.shape[0], b_data.shape[0]) + b_scale_is_transposed = weight_scale_transposed # caller declares the cached layout else: N = weights[0].shape[0] _check_splits(m_splits, a.shape[0], len(weights)) @@ -332,12 +337,13 @@ def grouped_gemm_a8w4_fprop( 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) + b_scale_is_transposed = False # fresh weight quant is always plain [G, N, K/32] a_data, a_scale = _row_operand_mxfp8(a, transpose_scale=True) # a_scale [K/32, total_M] - # a8w4 swizzle: the kernel consumes both scales transposed. - # Normalise the weight scale to [G, K/32, N] (plain weight quant / weight_row is - # [G, N, K/32]; a caller may also pass it already transposed to keep the cache free). - if b_scale.shape[-1] != N: + # a8w4 swizzle: the kernel consumes both scales transposed ([G, K/32, N]). The plain + # weight quant / documented weight_row is [G, N, K/32]; transpose unless the caller + # already cached it transposed (weight_scale_transposed). + if not b_scale_is_transposed: b_scale = b_scale.transpose(1, 2).contiguous() group_offs = _prefix_offsets(m_splits, a.device) return grouped_gemm_a8w4_triton_kernel( @@ -463,12 +469,12 @@ def grouped_gemm_mxfp4_fprop_prequantized( N = b_data.shape[1] K = b_data.shape[2] * 2 # weights are packed 2 e2m1 elems/byte along K _check_contract(K, "K") - if a_data.shape[0] == 0: # empty MoE routing batch: no tokens for this GPU's experts - return torch.empty((0, N), dtype=out_dtype, device=a_data.device) _check_splits(m_splits, a_data.shape[0], b_data.shape[0]) _check_prequantized_operands( a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8, scale_transposed=scale_transposed ) + if a_data.shape[0] == 0: # empty MoE routing batch: no tokens for this GPU's experts + return torch.empty((0, N), dtype=out_dtype, device=a_data.device) group_offs = _prefix_offsets(m_splits, a_data.device) kernel = grouped_gemm_a8w4_triton_kernel if a_is_mxfp8 else grouped_gemm_mxfp4_triton_kernel From 9aaa49b251925c9b24cb68961b5b83912963dd88 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 14:06:55 -0500 Subject: [PATCH 11/26] wire in out --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 48 +++++++++++++++++++ .../pytorch/module/grouped_linear.py | 14 +++--- .../triton_kernels/grouped_gemm_mxfp4.py | 6 ++- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 18 ++++++- 4 files changed, 75 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index f0472bb64b..47e317e197 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -288,3 +288,51 @@ def test_grouped_prequantized_empty_batch(): a_data, a_scale, b_data, b_scale, m_splits, a_is_mxfp8=True, out_dtype=dtype ) assert out.shape == (0, N) + + +def test_grouped_prequantized_out_buffer(): + """grouped_gemm_mxfp4_fprop_prequantized writes into a caller-provided out buffer.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + K, N = 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + a = torch.randn(total_M, K, device=device, dtype=dtype) + weights = [torch.randn(N, K, device=device, dtype=dtype) for _ in m_splits] + + aq = MXFP8E4M3QuantizerRef(rowwise=True, columnwise=False).quantize(a) + weight_refs = [_e2m1_ref(rowwise=True, columnwise=False).quantize(w) for w in weights] + b_data = torch.stack([w.data for w in weight_refs], dim=0) + b_scale = torch.stack([w.scale for w in weight_refs], dim=0) + + ref = grouped_gemm_mxfp4_fprop_prequantized( + aq.data, aq.scale, b_data, b_scale, m_splits, a_is_mxfp8=True, out_dtype=dtype + ) + buf = torch.empty(total_M, N, device=device, dtype=dtype) + got = grouped_gemm_mxfp4_fprop_prequantized( + aq.data, aq.scale, b_data, b_scale, m_splits, a_is_mxfp8=True, out_dtype=dtype, out=buf + ) + assert got.data_ptr() == buf.data_ptr(), "the caller's out buffer was not used" + assert torch.equal(buf, ref), "out-buffer result differs from the allocating path" + + +def test_grouped_linear_a8w4_out_buffer(): + """te.GroupedLinear a8w4 honors a caller-provided out buffer in the fast path.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with torch.no_grad(), te.autocast(enabled=True, recipe=rec): + ref = model(inp, m_splits) + buf = torch.empty(total_M, N, device=device, dtype=dtype) + got = model(inp, m_splits, out=buf) + + assert torch.equal(buf, ref), "out buffer does not hold the forward result" + assert torch.equal(got, ref), "returned output differs when an out buffer is given" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 6c7002c5b4..c2e6f0ce3e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -591,7 +591,6 @@ def _is_grouped_mxfp4_triton_supported( debug, unpad_output, actual_m_splits, - out, dgrad_out, fuse_wgrad_accumulation, wgrad_store, @@ -603,9 +602,9 @@ def _is_grouped_mxfp4_triton_supported( a8w4 is forward-only (backward raises); a4w4 also wires the low-precision dgrad/wgrad (see :meth:`_forward_grouped_mxfp4_triton`). Gated off for options this - self-contained path does not honor -- caller ``out``/``dgrad_out`` buffers, + self-contained path does not honor -- a caller ``dgrad_out`` buffer, ``fuse_wgrad_accumulation``, and delayed wgrad (``wgrad_store``) -- so those fall - back to the standard path. + back to the standard path. (A forward ``out`` buffer is honored.) """ if not (IS_HIP_EXTENSION and fp8 and recipe is not None and recipe.custom()): return False @@ -621,7 +620,6 @@ def _is_grouped_mxfp4_triton_supported( or debug or unpad_output or actual_m_splits is not None - or out is not None or dgrad_out is not None or fuse_wgrad_accumulation or (wgrad_store is not None and wgrad_store.delay_wgrad_compute()) @@ -978,6 +976,7 @@ def _forward_grouped_mxfp4_triton( is_first_microbatch, weight_workspaces, cache_weight, + out=None, ): """Grouped MXFP4 / a8w4 forward (ROCm Triton, gfx950). @@ -1062,7 +1061,7 @@ def _forward_grouped_mxfp4_triton( if cache_weight: new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) - out = grouped_gemm_mxfp4_fprop_prequantized( + result = grouped_gemm_mxfp4_fprop_prequantized( a_data, a_scale, b_data, @@ -1071,6 +1070,7 @@ def _forward_grouped_mxfp4_triton( a_is_mxfp8=a_is_mxfp8, out_dtype=activation_dtype, scale_transposed=a_is_mxfp8, + out=None if out is None else out.reshape(-1, out.shape[-1]), ) if is_grad_enabled: @@ -1086,7 +1086,7 @@ def _forward_grouped_mxfp4_triton( ctx.grouped_mxfp4_weight_requires_grad = weights[0].requires_grad ctx.save_for_backward(a, w_col_data, w_col_scale) - return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + return result.view(-1, *inp.shape[1:-1], result.shape[-1]), new_workspaces @staticmethod def _backward_grouped_mxfp4_triton(ctx, grad_output): @@ -1254,7 +1254,6 @@ def forward( debug=debug, unpad_output=unpad_output, actual_m_splits=actual_m_splits, - out=out, dgrad_out=dgrad_out, fuse_wgrad_accumulation=fuse_wgrad_accumulation, wgrad_store=wgrad_store, @@ -1273,6 +1272,7 @@ def forward( is_first_microbatch=is_first_microbatch, weight_workspaces=weight_workspaces, cache_weight=cache_weight, + out=out, ) # Configure quantizers diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index e11dcb46a2..bd73fba520 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -328,6 +328,7 @@ def grouped_gemm_mxfp4_triton_kernel( out_dtype=torch.bfloat16, num_cu=None, scale_transposed=False, + out=None, ): """A(total_M, K/2) @ B(G, N, K/2)^T -> C. FP4-packed data, e8m0 uint8 scales. @@ -338,7 +339,7 @@ def grouped_gemm_mxfp4_triton_kernel( 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) + c = out if out is not None else 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) @@ -410,6 +411,7 @@ def grouped_gemm_a8w4_triton_kernel( out_dtype=torch.bfloat16, num_cu=None, scale_transposed=False, + out=None, ): """a8w4 forward: C = A(e4m3) @ B(e2m1)^T, grouped along M (routed experts). @@ -423,7 +425,7 @@ def grouped_gemm_a8w4_triton_kernel( 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) + c = out if out is not None else 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) 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 06aca167a5..9ff45e9eec 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -441,6 +441,7 @@ def grouped_gemm_mxfp4_fprop_prequantized( out_dtype: torch.dtype = torch.bfloat16, num_cu: Optional[int] = None, scale_transposed: bool = False, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Grouped forward from already-quantized operands: ``C[g] = A[g] @ W[g]^T``. @@ -459,6 +460,8 @@ def grouped_gemm_mxfp4_fprop_prequantized( ``False`` the MXFP4 x MXFP4 kernel. scale_transposed: when ``True`` the E8M0 scales are laid out ``[K/32, total_M]`` / ``[G, K/32, N]`` so the kernel loads them without the per-iter ``tl.trans``. + out: optional ``[total_M, N]`` ``out_dtype`` buffer the kernel writes into instead + of allocating (caller-provided output; must match device). Returns: ``[total_M, N]`` output in ``out_dtype``. @@ -473,8 +476,18 @@ def grouped_gemm_mxfp4_fprop_prequantized( _check_prequantized_operands( a_data, a_scale, b_data, b_scale, a_is_mxfp8=a_is_mxfp8, scale_transposed=scale_transposed ) - if a_data.shape[0] == 0: # empty MoE routing batch: no tokens for this GPU's experts - return torch.empty((0, N), dtype=out_dtype, device=a_data.device) + total_M = a_data.shape[0] + if out is not None and ( + tuple(out.shape) != (total_M, N) + or out.dtype != out_dtype + or out.device != a_data.device + ): + raise ValueError( + f"out must be [{total_M}, {N}] {out_dtype} on {a_data.device}, got" + f" {tuple(out.shape)} {out.dtype} on {out.device}" + ) + if total_M == 0: # empty MoE routing batch: no tokens for this GPU's experts + return out if out is not None else torch.empty((0, N), dtype=out_dtype, device=a_data.device) group_offs = _prefix_offsets(m_splits, a_data.device) kernel = grouped_gemm_a8w4_triton_kernel if a_is_mxfp8 else grouped_gemm_mxfp4_triton_kernel @@ -490,6 +503,7 @@ def grouped_gemm_mxfp4_fprop_prequantized( out_dtype=out_dtype, num_cu=num_cu, scale_transposed=scale_transposed, + out=out, ) From aef244aca7bc8e9b2aa7f38631b643112f23d97b Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 14:27:31 -0500 Subject: [PATCH 12/26] address review --- transformer_engine/pytorch/module/grouped_linear.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c2e6f0ce3e..9f5e4fcee5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -998,6 +998,11 @@ def _forward_grouped_mxfp4_triton( num_gemms = len(weights) in_features = weights[0].size(-1) + if inp.size(-1) != in_features: + raise ValueError( + f"Input tensor (shape={tuple(inp.size())}) is not compatible with " + f"weight tensor (shape={tuple(weights[0].size())})" + ) m_splits_list = ( m_splits.tolist() if isinstance(m_splits, torch.Tensor) else [int(m) for m in m_splits] ) From 4f44d8a8d7541cae7785a52126e8cf20a605a29b Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 14:36:50 -0500 Subject: [PATCH 13/26] validate buffer --- transformer_engine/pytorch/module/grouped_linear.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9f5e4fcee5..d984f5af32 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1066,6 +1066,11 @@ def _forward_grouped_mxfp4_triton( if cache_weight: new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) + # Validate the caller's out buffer (2D/contiguous/dtype/device, and reject a + # requires-grad buffer -- the kernel writes in place) or allocate one. + result_out = _GroupedLinear._validate_or_alloc_output( + out, a_data.shape[0], b_data.shape[1], activation_dtype, a_data.device + ) result = grouped_gemm_mxfp4_fprop_prequantized( a_data, a_scale, @@ -1075,7 +1080,7 @@ def _forward_grouped_mxfp4_triton( a_is_mxfp8=a_is_mxfp8, out_dtype=activation_dtype, scale_transposed=a_is_mxfp8, - out=None if out is None else out.reshape(-1, out.shape[-1]), + out=result_out, ) if is_grad_enabled: From 8da10664feefe9fff59b6be92f9055dbf3cad3ae Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 14:59:56 -0500 Subject: [PATCH 14/26] address comments --- .../pytorch/module/grouped_linear.py | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d984f5af32..1a3c9bc6be 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -608,9 +608,10 @@ def _is_grouped_mxfp4_triton_supported( """ if not (IS_HIP_EXTENSION and fp8 and recipe is not None and recipe.custom()): return False - # Contractions the kernels step in 128s: fprop/wgrad over K=in_features, - # dgrad over N=out_features. Require both so the a4w4 backward is safe. - if in_features % 128 != 0 or out_features % 128 != 0: + # The persistent forward kernel steps the K=in_features contraction in 128s (both + # a8w4 and a4w4). out_features (the dgrad N contraction) is required only for a4w4 + # below, since a8w4 is forward-only and the forward masks arbitrary N. + if in_features % 128 != 0: return False if ( use_bias @@ -654,10 +655,17 @@ def _plain_layout(q) -> bool: isinstance(q, MXFP4QuantizerRef) and _plain_layout(q) for q in weight_quantizers ): return False - return all( - isinstance(q, (MXFP4QuantizerRef, MXFP8E4M3QuantizerRef)) and _plain_layout(q) - for q in input_quantizers - ) + if not all(_plain_layout(q) for q in input_quantizers): + return False + # The forward quantizes the whole grouped activation in one format chosen from + # input_quantizers[0], so every expert's input quantizer must be the same format; + # a mixed list would silently drop the later formats. + if all(isinstance(q, MXFP8E4M3QuantizerRef) for q in input_quantizers): + return True # a8w4: forward-only, only the K=in_features contraction matters + if not all(isinstance(q, MXFP4QuantizerRef) for q in input_quantizers): + return False # mixed activation formats -> fall back to the standard path + # a4w4 runs the low-precision dgrad, which contracts over N=out_features in 128s. + return out_features % 128 == 0 @staticmethod def _forward_blockwise_fp8_triton( @@ -1031,8 +1039,7 @@ def _forward_grouped_mxfp4_triton( new_workspaces = [None] * num_gemms reuse = not update_ws and cached is not None and (not want_wcol or cached[2] is not None) if reuse: - # Validate the cached workspace before trusting it (mirrors the blockwise path): - # our 4-tuple whose stacked b_data matches the current weights. + # Validate the cached workspace before trusting it expected = (num_gemms, weights[0].size(0), in_features // 2) if not ( isinstance(cached, tuple) @@ -1066,8 +1073,6 @@ def _forward_grouped_mxfp4_triton( if cache_weight: new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) - # Validate the caller's out buffer (2D/contiguous/dtype/device, and reject a - # requires-grad buffer -- the kernel writes in place) or allocate one. result_out = _GroupedLinear._validate_or_alloc_output( out, a_data.shape[0], b_data.shape[1], activation_dtype, a_data.device ) From 69b9ba6b17b7c271b634b23a9355631b0ec16e0f Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 15:09:56 -0500 Subject: [PATCH 15/26] check quantizer type before accessing MXFP4-only attributes Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- transformer_engine/pytorch/module/grouped_linear.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1a3c9bc6be..25b59b355d 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -641,6 +641,8 @@ def _plain_layout(q) -> bool: # MXFP4QuantizerRef must have every layout-transform flag off. if isinstance(q, MXFP8E4M3QuantizerRef): return True + if not isinstance(q, MXFP4QuantizerRef): + return False return not ( q.shuffle_rowwise_data or q.shuffle_columnwise_data From 1580d77b8c1f83be12661d08ef6734833b524ed6 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 15:34:59 -0500 Subject: [PATCH 16/26] address review --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 55 +++++++++++++++++-- .../pytorch/module/grouped_linear.py | 33 ++++++----- 2 files changed, 69 insertions(+), 19 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 47e317e197..c263ea962e 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -16,6 +16,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( + _col_operand_grouped_padded, _row_operand, _row_operand_mxfp8, _row_operand_mxfp8_torch, @@ -180,7 +181,7 @@ def test_a4w4_activation_native_matches_reference(): def test_grouped_linear_a4w4_recipe_backward(): - """te.GroupedLinear a4w4 fwd+bwd routes to the low-precision dgrad/wgrad kernels.""" + """te.GroupedLinear a4w4 fwd+bwd: routing + exact output/dgrad/wgrad vs the reference.""" _isolate() device, dtype = "cuda", torch.bfloat16 num_gemms, K, N = 4, 256, 256 @@ -190,19 +191,65 @@ def test_grouped_linear_a4w4_recipe_backward(): model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() inp = torch.randn(total_M, K, device=device, dtype=dtype, requires_grad=True) + grad_out = torch.randn(total_M, N, device=device, dtype=dtype) rec = recipe.CustomRecipe(qfactory=mxfp4_grouped_quantizer_factory) with te.autocast(enabled=True, recipe=rec): out = model(inp, m_splits) - out.backward(torch.randn_like(out)) + out.backward(grad_out) + + # Reference: dequant the SAME native-e2m1 activation/grad and reference-e2m1 weights + # the kernels consume, per expert, in fp32. Forward Y = Aq @ Wq^T (row weight); dgrad + # dX = dYq @ Wq (reference col weight, matching the dgrad kernel). + def _deq_rows(x): + d, s = _row_operand(x) + return mxfp4_to_f32(d) * e8m0_to_f32(s).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + + a_hp = _deq_rows(inp.detach()) + go_hp = _deq_rows(grad_out) + weight_refs = [ + _e2m1_ref(rowwise=True, columnwise=True).quantize(getattr(model, f"weight{g}").detach()) + for g in range(num_gemms) + ] + y_ref = torch.zeros(total_M, N, device=device, dtype=torch.float32) + dx_ref = torch.zeros(total_M, K, device=device, dtype=torch.float32) + start = 0 + for g, m in enumerate(m_splits): + if m > 0: + seg = slice(start, start + m) + w_hp = _deq_ref_operand( + weight_refs[g].data, weight_refs[g].scale, weight_refs[g], default_e4m3=False + ) # [N, K] + w_t_hp = _deq_ref_operand( + weight_refs[g].data_t, weight_refs[g].scale_t, weight_refs[g], default_e4m3=False + ) # [K, N] == W^T + y_ref[seg] = a_hp[seg] @ w_hp.t() + dx_ref[seg] = go_hp[seg] @ w_t_hp.t() + start += m + + # wgrad dW[g] = dYq^T @ Aq contracts over M with the columnwise, per-group-padded + # operands the kernel consumes -- build them with the exact prep fn so the quantization + # matches; the zero padding rows contribute nothing to the reduction. + go_ld, go_ls, go_off = _col_operand_grouped_padded(grad_out, m_splits) + a_ld, a_ls, _ = _col_operand_grouped_padded(inp.detach(), m_splits) + go_col = mxfp4_to_f32(go_ld) * e8m0_to_f32(go_ls).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + a_col = mxfp4_to_f32(a_ld) * e8m0_to_f32(a_ls).repeat_interleave(MXFP4_BLOCK_SIZE, dim=1) + off = go_off.tolist() + dw_ref = [ + go_col[:, off[g] : off[g + 1]] @ a_col[:, off[g] : off[g + 1]].t() for g in range(num_gemms) + ] + + def _rel(x, ref): + return (x.float() - ref).norm() / ref.norm().clamp_min(1e-12) assert out.shape == (total_M, N) + assert _rel(out, y_ref) < _REL_TOL, f"a4w4 forward disagrees with reference: {_rel(out, y_ref):.4f}" assert inp.grad is not None and inp.grad.shape == inp.shape - assert torch.isfinite(inp.grad).all() + assert _rel(inp.grad, dx_ref) < _REL_TOL, f"a4w4 dgrad disagrees: {_rel(inp.grad, dx_ref):.4f}" for g in range(num_gemms): w = getattr(model, f"weight{g}") assert w.grad is not None and w.grad.shape == w.shape - assert torch.isfinite(w.grad).all() + assert _rel(w.grad, dw_ref[g]) < _REL_TOL, f"a4w4 wgrad[{g}] disagrees: {_rel(w.grad, dw_ref[g]):.4f}" def test_grouped_linear_a8w4_backward_raises(): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 25b59b355d..11d37d4f53 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1039,21 +1039,24 @@ def _forward_grouped_mxfp4_triton( update_ws = is_first_microbatch is None or is_first_microbatch cached = weight_workspaces[0] if weight_workspaces else None new_workspaces = [None] * num_gemms - reuse = not update_ws and cached is not None and (not want_wcol or cached[2] is not None) + # Reuse the cached quantized weights only when they came from a compatible call. + # _fp8_workspaces is shared and only cleared when the recipe class changes, so a + # workspace from another CustomRecipe (or a prior fallback) may sit here. The cache + # tuple tags the operand format (a_is_mxfp8) -- a8w4 stores the weight scale + # transposed [G, K/32, N], a4w4 plain [G, N, K/32] -- so the two are never mixed. + # Validate the structure before indexing, and re-quantize on any mismatch. + expected_b = (num_gemms, weights[0].size(0), in_features // 2) + reuse = ( + not update_ws + and isinstance(cached, tuple) + and len(cached) == 5 + and cached[0] == a_is_mxfp8 + and isinstance(cached[1], torch.Tensor) + and tuple(cached[1].shape) == expected_b + and (not want_wcol or cached[3] is not None) + ) if reuse: - # Validate the cached workspace before trusting it - expected = (num_gemms, weights[0].size(0), in_features // 2) - if not ( - isinstance(cached, tuple) - and len(cached) == 4 - and isinstance(cached[0], torch.Tensor) - and tuple(cached[0].shape) == expected - ): - raise RuntimeError( - "Cached grouped MXFP4 weight workspace is incompatible with the current " - f"weights; expected a 4-tuple whose b_data has shape {expected}." - ) - b_data, b_scale, w_col_data, w_col_scale = cached + _, b_data, b_scale, w_col_data, w_col_scale = cached else: b_datas, b_scales, wcol_datas, wcol_scales = [], [], [], [] for i in range(num_gemms): @@ -1073,7 +1076,7 @@ def _forward_grouped_mxfp4_triton( w_col_data = torch.stack(wcol_datas, dim=0) if want_wcol else None w_col_scale = torch.stack(wcol_scales, dim=0) if want_wcol else None if cache_weight: - new_workspaces[0] = (b_data, b_scale, w_col_data, w_col_scale) + new_workspaces[0] = (a_is_mxfp8, b_data, b_scale, w_col_data, w_col_scale) result_out = _GroupedLinear._validate_or_alloc_output( out, a_data.shape[0], b_data.shape[1], activation_dtype, a_data.device From d026ae2246df7db5faabc1bd7c0a0ae775697e59 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 15:54:07 -0500 Subject: [PATCH 17/26] only quantize a4w4 col weight when input grad is needed --- transformer_engine/pytorch/module/grouped_linear.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 11d37d4f53..f575c8e379 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1035,7 +1035,10 @@ def _forward_grouped_mxfp4_triton( # Cache the quantized weights across microbatches (is_first_microbatch): the # weight quant is a fixed per-call cost (~G*N*K) that otherwise dominates the # grouped GEMM, so reuse the cache whenever this is not the first microbatch. - want_wcol = is_grad_enabled and not a_is_mxfp8 + # a4w4 dgrad consumes the col-wise weight, so only request/quantize it when an + # input gradient is actually needed (matches the standard grouped path); a8w4 is + # forward-only and never needs it. + want_wcol = is_grad_enabled and not a_is_mxfp8 and inp.requires_grad update_ws = is_first_microbatch is None or is_first_microbatch cached = weight_workspaces[0] if weight_workspaces else None new_workspaces = [None] * num_gemms From b08bd19001ef4614d3dd91289d5fe906423dceee Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 16:39:52 -0500 Subject: [PATCH 18/26] isolate MXFP4 weight-cache namespace; gate a4w4 on grad_output quantize --- .../pytorch/module/grouped_linear.py | 28 +++++++++++++++++-- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f575c8e379..93d996ae35 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -584,6 +584,7 @@ def _is_grouped_mxfp4_triton_supported( recipe, input_quantizers, weight_quantizers, + grad_output_quantizers, use_bias, backward_override, cpu_offloading, @@ -666,7 +667,14 @@ def _plain_layout(q) -> bool: return True # a8w4: forward-only, only the K=in_features contraction matters if not all(isinstance(q, MXFP4QuantizerRef) for q in input_quantizers): return False # mixed activation formats -> fall back to the standard path - # a4w4 runs the low-precision dgrad, which contracts over N=out_features in 128s. + # a4w4 runs the low-precision backward, which quantizes grad_output as plain E2M1; + # the recipe's grad_output quantizer must match (plain MXFP4) or the fast path would + # silently use a different grad quantization than the recipe specifies. + if not all( + isinstance(q, MXFP4QuantizerRef) and _plain_layout(q) for q in grad_output_quantizers + ): + return False + # dgrad contracts over N=out_features in 128s. return out_features % 128 == 0 @staticmethod @@ -1270,6 +1278,7 @@ def forward( recipe=blockwise_recipe, input_quantizers=input_quantizers, weight_quantizers=weight_quantizers, + grad_output_quantizers=grad_output_quantizers, use_bias=use_bias, backward_override=backward_override, cpu_offloading=cpu_offloading, @@ -2772,8 +2781,21 @@ def forward( autograd_ctx = [None] cache_weight = is_first_microbatch is not None + # The MXFP4 grouped fast path caches a raw tuple; isolate it in its own workspace + # namespace so it never lands in the shared weight{i} slot that the standard / + # blockwise paths consume as a quantized weight (_fp8_workspaces is shared and only + # clears on a recipe *class* change, not between two CustomRecipe factories). + ws_prefix = "weight" + if cache_weight and self.fp8: + try: + from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef + + if isinstance(weight_quantizers[0], MXFP4QuantizerRef): + ws_prefix = "mxfp4_grouped_weight" + except ImportError: + pass weight_workspaces = ( - [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] + [self._fp8_workspaces.get(f"{ws_prefix}{i}") for i in range(num_gemms)] if cache_weight else [None] * num_gemms ) @@ -2820,7 +2842,7 @@ def forward( if ws is not None: if isinstance(ws, torch.Tensor): ws = ws.detach() - self._fp8_workspaces[f"weight{i}"] = ws + self._fp8_workspaces[f"{ws_prefix}{i}"] = ws finally: self.end_forward() From c29e03e2bf12c80b6599397fa4c3caef118d891d Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 29 Sep 2026 17:15:06 -0500 Subject: [PATCH 19/26] add tests --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 84 +++++++++++++++++++ .../pytorch/triton_kernels/cast_transpose.py | 5 +- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index c263ea962e..67a9856810 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -23,8 +23,10 @@ grouped_gemm_mxfp4_dgrad, grouped_gemm_mxfp4_fprop_prequantized, ) +from transformer_engine.pytorch.module.grouped_linear import _GroupedLinear from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import ( MXFP4_BLOCK_SIZE, + MXFP4QuantizerRef, e8m0_to_f32, mxfp4_to_f32, ) @@ -383,3 +385,85 @@ def test_grouped_linear_a8w4_out_buffer(): assert torch.equal(buf, ref), "out buffer does not hold the forward result" assert torch.equal(got, ref), "returned output differs when an out buffer is given" + + +def test_a4w4_gate_rejects_nonplain_grad_output_quantizer(): + """a4w4 fast path must fall back unless the recipe's grad_output quantizer is plain MXFP4. + + The low-precision backward quantizes grad_output as native E2M1, so a recipe that keeps a + plain-MXFP4 forward but a different grad format (E4M3) or a swizzled/Hadamard MXFP4 grad + quantizer would silently disagree with what the recipe specifies -- gate must reject it. + """ + _isolate() + G, K, N = 4, 256, 256 + rec = recipe.CustomRecipe(qfactory=mxfp4_grouped_quantizer_factory) + base = dict( + fp8=True, + recipe=rec, + input_quantizers=[_e2m1_ref(rowwise=True, columnwise=False) for _ in range(G)], + weight_quantizers=[_e2m1_ref(rowwise=True, columnwise=True) for _ in range(G)], + use_bias=False, + backward_override=False, + cpu_offloading=False, + save_original_input=False, + debug=False, + unpad_output=False, + actual_m_splits=None, + dgrad_out=None, + fuse_wgrad_accumulation=False, + wgrad_store=None, + in_features=K, + out_features=N, + ) + + def _supported(grad_quantizers): + return _GroupedLinear._is_grouped_mxfp4_triton_supported( + grad_output_quantizers=grad_quantizers, **base + ) + + # Plain E2M1 grad_output -> a4w4 fast path is supported. + assert _supported([_e2m1_ref(rowwise=True, columnwise=False) for _ in range(G)]) + # E4M3 grad_output (not MXFP4) -> fall back to the standard path. + assert not _supported( + [MXFP8E4M3QuantizerRef(rowwise=True, columnwise=True) for _ in range(G)] + ) + # Swizzled MXFP4 grad_output (MXFP4 but not plain) -> fall back. + swz = [ + MXFP4QuantizerRef( + rowwise=True, + columnwise=False, + shuffle_rowwise_data=False, + shuffle_columnwise_data=False, + with_gemm_swizzled_scales=True, + use_hadamard=False, + ) + for _ in range(G) + ] + assert not _supported(swz) + + +def test_a8w4_weight_cache_uses_isolated_namespace(): + """The fast-path weight cache is a raw tuple; it must live in an MXFP4-only workspace key, + never the shared 'weight{i}' slot the standard/blockwise weight-prep paths read (a stale + tuple there would break a later call under a different CustomRecipe, since _fp8_workspaces + only clears on a recipe *class* change).""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + total_M = sum(m_splits) + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(total_M, K, device=device, dtype=dtype) + + rec = recipe.CustomRecipe(qfactory=a8w4_quantizer_factory) + with torch.no_grad(), te.autocast(enabled=True, recipe=rec): + model(inp, m_splits, is_first_microbatch=True) # quantize + cache the weights + + ws = model._fp8_workspaces + # the raw fast-path tuple lives under the isolated MXFP4 namespace... + assert isinstance(ws.get("mxfp4_grouped_weight0"), tuple), list(ws.keys()) + # ...and never in the shared 'weight{i}' slots the other weight-prep paths read. + for i in range(num_gemms): + assert not isinstance(ws.get(f"weight{i}"), tuple) diff --git a/transformer_engine/pytorch/triton_kernels/cast_transpose.py b/transformer_engine/pytorch/triton_kernels/cast_transpose.py index 9848d72363..c2fcd097fa 100644 --- a/transformer_engine/pytorch/triton_kernels/cast_transpose.py +++ b/transformer_engine/pytorch/triton_kernels/cast_transpose.py @@ -873,10 +873,7 @@ def mxfp8_e4m3_rowwise_downcast(x, transpose_scale=False): """Fused single-pass MXFP8 e4m3 row-wise downcast (Kimi-K3 a8w4 activation). ``x`` [M, K] (K a multiple of 32) -> (data [M, K] e4m3, scale uint8 E8M0), using the - ceil scale rule ``ceil(log2(amax_block / 448))``. Reuses :func:`_cast_transpose_triton_mxfp8` - row-wise only, so it matches the torch reference ``_row_operand_mxfp8_torch`` bit-for-bit for - non-zero blocks (an all-zero 32-block may get a different E8M0 scale byte, but its data is zero - so it dequantizes to zero either way) but is single-pass. + ceil scale rule ``ceil(log2(amax_block / 448))``. """ assert x.shape[-1] % MXFP8_BLOCK_SCALING_SIZE == 0, "K must be a multiple of 32" x = x.contiguous() From f7f98a5795c7687c2d1c5a8477561418282cf491 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 9 Oct 2026 10:49:36 -0500 Subject: [PATCH 20/26] use get_device_compute_capability for the gfx950 gate; drop ImportError guards --- .../pytorch/module/grouped_linear.py | 21 +++++++------------ .../triton_kernels/grouped_gemm_mxfp4.py | 16 -------------- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 7 +++---- 3 files changed, 10 insertions(+), 34 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 93d996ae35..1bfe2fc2bc 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -627,14 +627,10 @@ def _is_grouped_mxfp4_triton_supported( or (wgrad_store is not None and wgrad_store.delay_wgrad_compute()) ): return False - try: - from ..triton_kernels.grouped_gemm_mxfp4 import _is_gfx950 - from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef - from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef - except ImportError: - return False - if not _is_gfx950(): + if get_device_compute_capability() != (9, 5): return False + from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef + from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef def _plain_layout(q) -> bool: # The grouped Triton kernels need plain (un-shuffled, un-swizzled, @@ -2786,14 +2782,11 @@ def forward( # blockwise paths consume as a quantized weight (_fp8_workspaces is shared and only # clears on a recipe *class* change, not between two CustomRecipe factories). ws_prefix = "weight" - if cache_weight and self.fp8: - try: - from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef + if cache_weight and self.fp8 and IS_HIP_EXTENSION: + from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef - if isinstance(weight_quantizers[0], MXFP4QuantizerRef): - ws_prefix = "mxfp4_grouped_weight" - except ImportError: - pass + if isinstance(weight_quantizers[0], MXFP4QuantizerRef): + ws_prefix = "mxfp4_grouped_weight" weight_workspaces = ( [self._fp8_workspaces.get(f"{ws_prefix}{i}") for i in range(num_gemms)] if cache_weight diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index bd73fba520..9206bfc85b 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -22,22 +22,6 @@ from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs -# =============================================================================== -# Arch helper -# =============================================================================== - - -@functools.lru_cache(maxsize=None) -def _is_gfx950_for(device_index: int) -> bool: - props = torch.cuda.get_device_properties(device_index) - return (props.major, props.minor) == (9, 5) - - -def _is_gfx950() -> bool: - # Key the cache by device index so a set_device change isn't masked. - return _is_gfx950_for(torch.cuda.current_device()) - - # =============================================================================== # AMD Triton compiler knobs # 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 9ff45e9eec..50d4b3ef76 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -12,9 +12,8 @@ import torch.nn.functional as F from ..tensor.mxfp4_tensor import MXFP4Quantizer -from ..utils import round_up_to_nearest_multiple +from ..utils import get_device_compute_capability, 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, @@ -81,7 +80,7 @@ def _check_splits( 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(): + if get_device_compute_capability() != (9, 5): raise RuntimeError( "grouped MXFP4 GEMM requires gfx950 (CDNA4); the current device lacks the" " scaled-FP4 MFMA used by tl.dot_scaled(..., \"e2m1\", ...)." @@ -174,7 +173,7 @@ def _row_operand_mxfp8(x: torch.Tensor, transpose_scale: bool = False) -> Tuple[ :func:`_row_operand_mxfp8_torch`. Used by the a8w4 forward (GATE-2 fused quant). ``transpose_scale`` emits the scale ``[K/32, M]``. """ - if _is_gfx950(): + if get_device_compute_capability() == (9, 5): return mxfp8_e4m3_rowwise_downcast(x, transpose_scale=transpose_scale) return _row_operand_mxfp8_torch(x, transpose_scale=transpose_scale) From c4c0e64f75812b7a8ae5ed20d5ca2dc8e24e0663 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 13 Oct 2026 17:24:12 -0500 Subject: [PATCH 21/26] fix post-2.19 tests --- tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py | 2 ++ .../pytorch/custom_recipes/quantization_mxfp4_grouped.py | 6 +++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index 67a9856810..fed53c3b0f 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -412,6 +412,8 @@ def test_a4w4_gate_rejects_nonplain_grad_output_quantizer(): dgrad_out=None, fuse_wgrad_accumulation=False, wgrad_store=None, + single_grouped_weight=False, + single_grouped_bias=False, in_features=K, out_features=N, ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py index 85ec536610..a3feabc245 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py @@ -17,7 +17,7 @@ import torch -from transformer_engine.pytorch.custom_recipes import quantization +from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer @@ -96,14 +96,14 @@ def qgemm( self, qx: torch.Tensor, qw: torch.Tensor, - m_params: quantization.MMParams, + m_params: gemm.MMParams, out_dtype: torch.dtype, sx: torch.Tensor, sw: torch.Tensor, bias: torch.Tensor | None = None, out: torch.Tensor | None = None, accumulate: bool = False, - gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, + gemm_type: gemm.GEMMType = gemm.GEMMType.FPROP, qresult_x: QuantizedTensorStorage | None = None, qresult_w: QuantizedTensorStorage | None = None, layout: str = "TN", From 395f42d7e83a3b50998851f123b38c6a23107975 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 19 Oct 2026 20:39:10 +0000 Subject: [PATCH 22/26] gfx1250 support, GPU weight cast, batched wgrad cast --- .../pytorch/mxfp4/test_a8w4_grouped_recipe.py | 34 +++++++++- tests/pytorch/mxfp4/test_a8w4_recipe.py | 8 ++- .../triton_kernels/test_grouped_gemm_mxfp4.py | 30 ++++++-- .../pytorch/module/grouped_linear.py | 59 ++++++++++------ .../blockwise_fp8_grouped_gemm.py | 6 +- .../triton_kernels/grouped_gemm_mxfp4.py | 21 +++++- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 68 +++++++++---------- 7 files changed, 153 insertions(+), 73 deletions(-) diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py index fed53c3b0f..a3adf33aa2 100644 --- a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -16,6 +16,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( + MXFP4_GEMM_ARCHS, _col_operand_grouped_padded, _row_operand, _row_operand_mxfp8, @@ -24,6 +25,7 @@ grouped_gemm_mxfp4_fprop_prequantized, ) from transformer_engine.pytorch.module.grouped_linear import _GroupedLinear +from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import ( MXFP4_BLOCK_SIZE, MXFP4QuantizerRef, @@ -38,8 +40,10 @@ mxfp4_grouped_quantizer_factory, ) -_recipe_ok, _reason = te.is_mxfp4_available(return_reason=True) -pytestmark = pytest.mark.skipif(not _recipe_ok, reason=_reason) +pytestmark = pytest.mark.skipif( + get_device_compute_capability() not in MXFP4_GEMM_ARCHS, + reason="grouped MXFP4 Triton GEMM requires gfx950 or gfx1250", +) _REL_TOL = 5e-2 @@ -254,6 +258,30 @@ def _rel(x, ref): assert _rel(w.grad, dw_ref[g]) < _REL_TOL, f"a4w4 wgrad[{g}] disagrees: {_rel(w.grad, dw_ref[g]):.4f}" +def test_grouped_linear_weight_cast_matches_reference(): + """The fast path's GPU weight cast is bit-identical to the recipe's reference quantizer.""" + _isolate() + device, dtype = "cuda", torch.bfloat16 + num_gemms, K, N = 4, 256, 256 + m_splits = [128, 64, 0, 192] + torch.manual_seed(0) + + model = te.GroupedLinear(num_gemms, K, N, bias=False, params_dtype=dtype).cuda() + inp = torch.randn(sum(m_splits), K, device=device, dtype=dtype, requires_grad=True) + + rec = recipe.CustomRecipe(qfactory=mxfp4_grouped_quantizer_factory) + with te.autocast(enabled=True, recipe=rec): + model(inp, m_splits, is_first_microbatch=True) + + _, b_data, b_scale, w_col_data, w_col_scale = model._fp8_workspaces["mxfp4_grouped_weight0"] + for g in range(num_gemms): + ref = _e2m1_ref(rowwise=True, columnwise=True).quantize(getattr(model, f"weight{g}").detach()) + assert torch.equal(b_data[g], ref.data) + assert torch.equal(b_scale[g], ref.scale[:N, : K // MXFP4_BLOCK_SIZE]) + assert torch.equal(w_col_data[g], ref.data_t) + assert torch.equal(w_col_scale[g], ref.scale_t[:K, : N // MXFP4_BLOCK_SIZE]) + + def test_grouped_linear_a8w4_backward_raises(): """a8w4 is forward-only QAT: the grouped-path backward must raise clearly.""" _isolate() @@ -303,7 +331,7 @@ def test_row_operand_mxfp8_fused_matches_torch(): for M, K in [(384, 256), (130, 256), (293, 3584)]: x = torch.randn(M, K, device=device, dtype=dtype) d_ref, s_ref = _row_operand_mxfp8_torch(x) - d_fused, s_fused = _row_operand_mxfp8(x) # fused on gfx950 + d_fused, s_fused = _row_operand_mxfp8(x) # fused on gfx950/gfx1250 assert torch.equal(s_fused, s_ref), f"E8M0 scales differ at {(M, K)}" assert torch.equal( d_fused.view(torch.uint8), d_ref.view(torch.uint8) diff --git a/tests/pytorch/mxfp4/test_a8w4_recipe.py b/tests/pytorch/mxfp4/test_a8w4_recipe.py index 334b1454c8..5588b8e350 100644 --- a/tests/pytorch/mxfp4/test_a8w4_recipe.py +++ b/tests/pytorch/mxfp4/test_a8w4_recipe.py @@ -16,9 +16,13 @@ a8w4_quantizer_factory, ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import MXFP4_GEMM_ARCHS +from transformer_engine.pytorch.utils import get_device_compute_capability -_recipe_ok, _reason = te.is_mxfp4_available(return_reason=True) -pytestmark = pytest.mark.skipif(not _recipe_ok, reason=_reason) +pytestmark = pytest.mark.skipif( + get_device_compute_capability() not in MXFP4_GEMM_ARCHS, + reason="a8w4 recipe requires gfx950 or gfx1250", +) def _isolate(): diff --git a/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py index f255b6ad0b..21e6fc6552 100644 --- a/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py +++ b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py @@ -1,17 +1,16 @@ # 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).""" +"""Numeric correctness tests for the grouped MXFP4 Triton GEMM (gfx950, gfx1250).""" import pytest import torch +import torch.nn.functional as F -from transformer_engine.pytorch.quantization import check_mxfp4_support - -_MXFP4_OK, _MXFP4_REASON = check_mxfp4_support() - +from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import ( MXFP4_BLOCK, + MXFP4_GEMM_ARCHS, _col_operand, _col_operand_grouped_padded, _row_operand, @@ -24,7 +23,10 @@ ) pytestmark = [ - pytest.mark.skipif(not _MXFP4_OK, reason=f"MXFP4 unsupported: {_MXFP4_REASON}"), + pytest.mark.skipif( + get_device_compute_capability() not in MXFP4_GEMM_ARCHS, + reason="grouped MXFP4 Triton GEMM requires gfx950 or gfx1250", + ), ] DTYPE = torch.bfloat16 @@ -203,6 +205,22 @@ def test_wgrad_precise_and_padding(): assert _rel_err(wgrad, ref_true) < _REL_TOL +def test_col_operand_grouped_padded_matches_per_group(): + """The single batched col-wise cast equals casting each zero-padded group alone.""" + splits = [*M_SPLITS, 0, 1] + x = _rand(sum(splits), K) + data, scale, go_pad = _col_operand_grouped_padded(x, splits) + offs = go_pad.tolist() + start = 0 + for g, m in enumerate(splits): + s, e = offs[g], offs[g + 1] + if e > s: + d_g, s_g = _col_operand(F.pad(x[start : start + m], (0, 0, 0, e - s - m))) + assert torch.equal(data[:, s // 2 : e // 2], d_g) + assert torch.equal(scale[:, s // MXFP4_BLOCK : e // MXFP4_BLOCK], s_g) + start += m + + def test_wgrad_accumulate(): total_m = sum(M_SPLITS) a = _rand(total_m, K) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9ca2756acd..7ab649a909 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -772,7 +772,7 @@ def _is_grouped_mxfp4_triton_supported( in_features, out_features, ) -> bool: - """Triton grouped MXFP4 / a8w4 path: gfx950 + a ``CustomRecipe`` whose + """Triton grouped MXFP4 / a8w4 path: gfx950/gfx1250 + a ``CustomRecipe`` whose factory yields the plain-layout MXFP4 reference quantizers. a8w4 is forward-only (backward raises); a4w4 also wires the low-precision @@ -804,7 +804,9 @@ def _is_grouped_mxfp4_triton_supported( or single_grouped_bias ): return False - if get_device_compute_capability() != (9, 5): + from ..triton_kernels.grouped_gemm_mxfp4_impl import MXFP4_GEMM_ARCHS + + if get_device_compute_capability() not in MXFP4_GEMM_ARCHS: return False from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef @@ -1169,18 +1171,20 @@ def _forward_grouped_mxfp4_triton( cache_weight, out=None, ): - """Grouped MXFP4 / a8w4 forward (ROCm Triton, gfx950). + """Grouped MXFP4 / a8w4 forward (ROCm Triton, gfx950/gfx1250). Selected from :meth:`forward` under a ``CustomRecipe`` whose factory yields the plain-layout MXFP4 reference quantizers (see - :meth:`_is_grouped_mxfp4_triton_supported`). The recipe's own reference - quantizers own the quantization; their packed E2M1/E4M3 data + E8M0 scales - feed straight into the persistent grouped kernel + :meth:`_is_grouped_mxfp4_triton_supported`). Operands are cast on the GPU with + the native MXFP4 / fused MXFP8 kernels, which are bit-identical to those + reference quantizers (non-bf16 weights fall back to the reference quantizer + itself), and feed the persistent grouped kernel (:func:`grouped_gemm_mxfp4_fprop_prequantized`), so the grouped-path numerics match the reference ``qgemm``. a4w4 wires the low-precision backward; a8w4 is forward-only (see :meth:`_backward_grouped_mxfp4_triton`). """ from ..triton_kernels.grouped_gemm_mxfp4_impl import ( + _quantize_weights_row_col, _row_operand, _row_operand_mxfp8, grouped_gemm_mxfp4_fprop_prequantized, @@ -1205,7 +1209,7 @@ def _forward_grouped_mxfp4_triton( input_quantizer = input_quantizers[0] a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) if a_is_mxfp8: - # a8w4 activation: fused single-pass MXFP8 e4m3 downcast on gfx950; emit the + # a8w4 activation: fused single-pass MXFP8 e4m3 downcast; emit the # scale transposed ([K/32, M]) for the kernel's swizzle path. a_data, a_scale = _row_operand_mxfp8(a, transpose_scale=True) else: @@ -1242,23 +1246,34 @@ def _forward_grouped_mxfp4_triton( if reuse: _, b_data, b_scale, w_col_data, w_col_scale = cached else: - b_datas, b_scales, wcol_datas, wcol_scales = [], [], [], [] - for i in range(num_gemms): - wq = weight_quantizers[i] - wq.set_usage(rowwise=True, columnwise=want_wcol) - wref = wq.quantize(weights[i].detach()) - b_datas.append(wref.data) - b_scales.append(wref.scale) + if weights[0].dtype == torch.bfloat16: + # Native GPU cast (bf16-only), bit-identical to the plain E2M1 reference. + ws = [w.detach() for w in weights] + w_col_data = w_col_scale = None if want_wcol: - wcol_datas.append(wref.data_t) - wcol_scales.append(wref.scale_t) - b_data = torch.stack(b_datas, dim=0) - b_scale = torch.stack(b_scales, dim=0) + (b_data, b_scale), (w_col_data, w_col_scale) = _quantize_weights_row_col(ws) + else: + rows = [_row_operand(w) for w in ws] + b_data = torch.stack([d for d, _ in rows], dim=0) + b_scale = torch.stack([s for _, s in rows], dim=0) + else: + b_datas, b_scales, wcol_datas, wcol_scales = [], [], [], [] + for i in range(num_gemms): + wq = weight_quantizers[i] + wq.set_usage(rowwise=True, columnwise=want_wcol) + wref = wq.quantize(weights[i].detach()) + b_datas.append(wref.data) + b_scales.append(wref.scale) + if want_wcol: + wcol_datas.append(wref.data_t) + wcol_scales.append(wref.scale_t) + b_data = torch.stack(b_datas, dim=0) + b_scale = torch.stack(b_scales, dim=0) + w_col_data = torch.stack(wcol_datas, dim=0) if want_wcol else None + w_col_scale = torch.stack(wcol_scales, dim=0) if want_wcol else None if a_is_mxfp8: # a8w4 swizzle: cache the weight scale transposed ([G, K/32, N]) b_scale = b_scale.transpose(1, 2).contiguous() - w_col_data = torch.stack(wcol_datas, dim=0) if want_wcol else None - w_col_scale = torch.stack(wcol_scales, dim=0) if want_wcol else None if cache_weight: new_workspaces[0] = (a_is_mxfp8, b_data, b_scale, w_col_data, w_col_scale) @@ -1284,7 +1299,7 @@ def _forward_grouped_mxfp4_triton( ctx.grouped_mxfp4_inp_shape = inp.shape ctx.grouped_mxfp4_m_splits = m_splits_list if not a_is_mxfp8: - # a4w4 low-precision backward: dgrad consumes the reference col weight; + # a4w4 low-precision backward: dgrad consumes the cached col weight; # wgrad re-quantizes the saved hp activation + grad_out in the kernel. ctx.grouped_mxfp4_requires_dgrad = inp.requires_grad ctx.grouped_mxfp4_weight_requires_grad = weights[0].requires_grad @@ -1451,7 +1466,7 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) - # Grouped MXFP4 / a8w4 (Triton, gfx950) under a CustomRecipe. Runs its + # Grouped MXFP4 / a8w4 (Triton, gfx950/gfx1250) under a CustomRecipe. Runs its # own reference-quantizer quantization + persistent grouped GEMM and returns # early. a4w4 also runs the backward; a8w4 is forward-only. if _GroupedLinear._is_grouped_mxfp4_triton_supported( diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index fd8b98ae09..b37be0071f 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -52,17 +52,19 @@ def _amd_knob_overrides(*, is_tn: bool) -> dict: @contextlib.contextmanager -def _amd_compiler_knobs(*, is_tn: bool): +def _amd_compiler_knobs(*, is_tn: bool, overrides: dict | None = None): """Temporarily apply AMD Triton compiler knobs, restoring prior state on exit. Scoped rather than set once because the knobs are process-global and read at kernel compile time; leaving them set would leak into other kernels' codegen. + ``overrides`` replaces the default :func:`_amd_knob_overrides` set. """ amd = getattr(getattr(triton, "knobs", None), "amd", None) if amd is None: yield return - overrides = _amd_knob_overrides(is_tn=is_tn) + if overrides is None: + overrides = _amd_knob_overrides(is_tn=is_tn) saved = { name: (name in amd.__dict__, amd.__dict__.get(name), os.environ.get(env_key)) for name, env_key in _AMD_KNOB_ENV.items() diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 9206bfc85b..6bd9074da9 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. # License for AMD contributions = MIT. See LICENSE for more information -"""Grouped MXFP4 GEMM Triton persistent kernels (gfx950). +"""Grouped MXFP4 GEMM Triton persistent kernels (gfx950, gfx1250). Adapted for TransformerEngine from Primus-Turbo. @@ -14,12 +14,14 @@ from __future__ import annotations import functools +import warnings import torch import triton import triton.language as tl from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs +from .common import get_arch # =============================================================================== @@ -31,16 +33,29 @@ # read at compile time and are not part of Triton's compile cache key. # =============================================================================== +# gfx1250: Triton's async copy (on by default there) corrupts the pipelined +# prologue K-iteration of the scaled-WMMA dot, so it is turned off. +_GFX1250_KNOBS = {"use_async_copy": False} + + 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. + async_copy / scalarize / block_pingpong set is enabled either way; on gfx1250 + only async_copy is disabled. """ def deco(func): @functools.wraps(func) def wrapper(*args, **kwargs): - with _amd_compiler_knobs(is_tn=is_tn): + overrides = None + if get_arch() == "gfx1250": + overrides = _GFX1250_KNOBS + warnings.warn( + "Grouped MXFP4 Triton GEMM: disabling Triton async_copy on gfx1250 (it" + " corrupts the pipelined scaled-WMMA dot); performance may be reduced." + ) + with _amd_compiler_knobs(is_tn=is_tn, overrides=overrides): return func(*args, **kwargs) return wrapper 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 50d4b3ef76..f3fb2a404f 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -1,7 +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).""" +"""TE-side quantization bridge + dispatch for grouped MXFP4 GEMM (gfx950, gfx1250).""" from __future__ import annotations @@ -28,6 +28,8 @@ 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 +# Archs with the scaled-FP4 matrix op used by tl.dot_scaled: gfx950 (MFMA), gfx1250 (WMMA). +MXFP4_GEMM_ARCHS = ((9, 5), (12, 5)) def _quantizer(*, rowwise: bool, columnwise: bool) -> MXFP4Quantizer: @@ -78,12 +80,11 @@ def _check_splits( ) -def _require_gfx950() -> None: - # The kernels use the CDNA4 scaled-FP4 MFMA (tl.dot_scaled e2m1); no other arch has it. - if get_device_compute_capability() != (9, 5): +def _require_mxfp4_gemm_arch() -> None: + if get_device_compute_capability() not in MXFP4_GEMM_ARCHS: raise RuntimeError( - "grouped MXFP4 GEMM requires gfx950 (CDNA4); the current device lacks the" - " scaled-FP4 MFMA used by tl.dot_scaled(..., \"e2m1\", ...)." + "grouped MXFP4 GEMM requires gfx950 or gfx1250; the current device lacks the" + " scaled-FP4 matrix op used by tl.dot_scaled(..., \"e2m1\", ...)." ) @@ -106,16 +107,15 @@ def _row_operand(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: def _col_operand(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Col-wise MXFP4: ``x`` [M, F] -> (data [F, M/2] u8, scale [F, M/32] u8). - Requests both directions (matching the production transpose-quant path); the - row-wise result is discarded. Both dims are zero-padded to a 32-multiple for - the quantizer and sliced back (both are free/masked in the GEMM). + 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()) + q = _quantizer(rowwise=False, columnwise=True).quantize(x.contiguous()) data = q._columnwise_data.view(torch.uint8)[:Feat, : M // 2] scale = q._columnwise_scale_inv[:Feat, : M // MXFP4_BLOCK].contiguous() return data, scale @@ -150,7 +150,7 @@ def _row_operand_mxfp8_torch(x: torch.Tensor, transpose_scale: bool = False) -> Plain OCP MXFP8: unpacked e4m3 element bytes + one E8M0 scale per 1x32 block, with the Kimi-K3 ``ceil(log2(amax / 448))`` block-scale rule. Torch multi-pass reference; - :func:`_row_operand_mxfp8` uses the fused single-pass kernel on gfx950. ``K`` (the + :func:`_row_operand_mxfp8` uses the fused single-pass kernel on gfx950/gfx1250. ``K`` (the contraction) is already a 32-multiple; ``M`` needs no padding (scale is per-row). ``transpose_scale`` emits the scale ``[K/32, M]`` instead of ``[M, K/32]``. """ @@ -168,12 +168,12 @@ def _row_operand_mxfp8_torch(x: torch.Tensor, transpose_scale: bool = False) -> def _row_operand_mxfp8(x: torch.Tensor, transpose_scale: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: - """Row-wise MXFP8 (e4m3) downcast (ceil rule): fused single-pass Triton on gfx950 + """Row-wise MXFP8 (e4m3) downcast (ceil rule): fused single-pass Triton on gfx950/gfx1250 (:func:`mxfp8_e4m3_rowwise_downcast`), else the bit-identical torch reference :func:`_row_operand_mxfp8_torch`. Used by the a8w4 forward (GATE-2 fused quant). ``transpose_scale`` emits the scale ``[K/32, M]``. """ - if get_device_compute_capability() == (9, 5): + if get_device_compute_capability() in MXFP4_GEMM_ARCHS: return mxfp8_e4m3_rowwise_downcast(x, transpose_scale=transpose_scale) return _row_operand_mxfp8_torch(x, transpose_scale=transpose_scale) @@ -207,26 +207,24 @@ def _col_operand_grouped_padded( ``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. + they contribute nothing to the wgrad contraction over M. All groups are cast in + one quantize: the padded group boundaries are 32-aligned, so no scale block spans + two groups and the result equals casting each group separately. """ - Feat = x.shape[1] - datas: List[torch.Tensor] = [] - scales: List[torch.Tensor] = [] + m_list = [int(m) for m in m_splits] + pad_list = [round_up_to_nearest_multiple(m, WGRAD_PAD_MULTIPLE) for m in m_list] offs = [0] - start = 0 - for m in m_splits: - m = int(m) - xg = x[start : start + m] - start += m - m_pad = round_up_to_nearest_multiple(m, WGRAD_PAD_MULTIPLE) - if m_pad != m: - xg = F.pad(xg, (0, 0, 0, m_pad - m)) - data, scale = _col_operand(xg) # (F, m_pad/2), (F, m_pad/32) - datas.append(data) - scales.append(scale) + for m_pad in pad_list: offs.append(offs[-1] + m_pad) - data = torch.cat(datas, dim=1).contiguous() - scale = torch.cat(scales, dim=1).contiguous() + if pad_list == m_list: + xp = x + else: + xp = x.new_zeros((offs[-1], x.shape[1])) + start = 0 + for m, off in zip(m_list, offs): + xp[off : off + m] = x[start : start + m] + start += m + data, scale = _col_operand(xp) # (F, total_pad/2), (F, total_pad/32) go_pad = torch.tensor(offs, dtype=torch.int64, device=x.device) return data, scale, go_pad @@ -254,7 +252,7 @@ def grouped_gemm_mxfp4_fprop( [total_M, N] output in ``out_dtype``. """ K = a.shape[1] - _require_gfx950() + _require_mxfp4_gemm_arch() _check_contract(K, "K") if weight_row is not None: @@ -318,7 +316,7 @@ def grouped_gemm_a8w4_fprop( [total_M, N] output in ``out_dtype``. """ K = a.shape[1] - _require_gfx950() + _require_mxfp4_gemm_arch() _check_contract(K, "K") if weight_row is not None: @@ -465,7 +463,7 @@ def grouped_gemm_mxfp4_fprop_prequantized( Returns: ``[total_M, N]`` output in ``out_dtype``. """ - _require_gfx950() + _require_mxfp4_gemm_arch() if b_data.ndim != 3: raise ValueError(f"b_data (weights) must be 3-D [G, N, K/2], got rank {b_data.ndim}") N = b_data.shape[1] @@ -533,7 +531,7 @@ def grouped_gemm_mxfp4_dgrad( [total_M, K] input gradient in ``out_dtype``. """ N = grad_out.shape[1] - _require_gfx950() + _require_mxfp4_gemm_arch() _check_contract(N, "N") if weight_col is not None: @@ -596,7 +594,7 @@ def grouped_gemm_mxfp4_wgrad( N = grad_out.shape[1] K = a.shape[1] G = len(m_splits) - _require_gfx950() + _require_mxfp4_gemm_arch() 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]) From 92f50263e1dabe50e4737b43e1ab7eca645051f3 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 19 Oct 2026 23:38:01 +0000 Subject: [PATCH 23/26] gfx1250 support --- .../triton_kernels/grouped_gemm_mxfp4.py | 70 ++++++++++++++++--- .../triton_kernels/grouped_gemm_mxfp4_impl.py | 22 +++--- 2 files changed, 71 insertions(+), 21 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 6bd9074da9..89392721a8 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -71,6 +71,12 @@ def wrapper(*args, **kwargs): NUM_XCDS = 8 # gfx942/gfx950 chiplet count +def _num_xcds(): + """Chiplet-remap width. gfx1250 measured faster with the plain round-robin pid + order (up to ~15% on MoE grouped GEMMs), so the remap is skipped there.""" + return 1 if get_arch() == "gfx1250" else NUM_XCDS + + @triton.jit def _chiplet_transform_chunked( pid, @@ -122,6 +128,27 @@ def _chiplet_transform_chunked( (256, 8, 2, 16, 0), (256, 8, 3, 16, 2), ] +# gfx1250 forward configs. num_warps=16 spreads the 256x256 fp32 accumulator over +# twice the waves, so fp4 x fp4 fits 256 VGPRs at 4 waves/SIMD (num_warps=8 needs +# ~360 VGPRs, 2 waves/SIMD); e4m3 x e2m1 spills at num_warps=16 and is fastest at +# num_warps=8 with BLOCK_K=256. The BLOCK_K=128 rows cover K % 256 != 0. +_FWD_CONFIGS_GFX1250 = [ + (256, 16, 3, 16, 0), + (256, 16, 2, 16, 0), + (256, 8, 3, 16, 0), + (256, 8, 2, 16, 0), + (512, 8, 3, 16, 0), + (128, 16, 2, 16, 0), + (128, 8, 2, 16, 0), +] +# gfx1250 wgrad configs (host tile 128x256 there). BLOCK_K=256 needs every padded +# group length 256-divisible (see _wgrad_pad_multiple in grouped_gemm_mxfp4_impl.py). +_VARK_CONFIGS_GFX1250 = [ + (256, 16, 3, 16, 0), + (256, 16, 2, 16, 0), + (128, 16, 3, 16, 0), + (128, 8, 2, 16, 0), +] def _make_configs(specs): @@ -139,19 +166,36 @@ def _make_configs(specs): ] +_FWD_TRITON_CONFIGS = _make_configs(_FWD_CONFIGS) +_FWD_TRITON_CONFIGS_GFX1250 = _make_configs(_FWD_CONFIGS_GFX1250) +_VARK_TRITON_CONFIGS = _make_configs(_VARK_CONFIGS) +_VARK_TRITON_CONFIGS_GFX1250 = _make_configs(_VARK_CONFIGS_GFX1250) + + def _prune_fwd_by_k(configs, named_args, **kwargs): - """Keep only configs whose BLOCK_K divides K (the fwd reduction loop is unmasked).""" + """Keep only the current arch's configs whose BLOCK_K divides K (the fwd + reduction loop is unmasked).""" K = named_args["K"] + configs = _FWD_TRITON_CONFIGS_GFX1250 if get_arch() == "gfx1250" else _FWD_TRITON_CONFIGS 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).""" + """Keep the current arch's configs; drop BLOCK_K=256 unless every padded group + length is 256-divisible (can256).""" + configs = _VARK_TRITON_CONFIGS_GFX1250 if get_arch() == "gfx1250" else _VARK_TRITON_CONFIGS if named_args["can256"]: return list(configs) return [c for c in configs if c.kwargs["BLOCK_SIZE_K"] != 256] +def _fwd_do_bench(kernel_call, quantiles): + if get_arch() == "gfx1250": + # gfx1250 downclocks under back-to-back load, which do_bench's per-launch L2 flush hides. + return triton.testing.do_bench_cudagraph(kernel_call, rep=100, quantiles=quantiles) + return triton.runtime.driver.active.get_benchmarker()(kernel_call, quantiles=quantiles) + + # ########################################################################### # Forward (NT): C[g] = A[g] @ B[g]^T # A (total_M, K/2) fp4-packed, A_scale (total_M, K/32) e8m0 @@ -160,9 +204,10 @@ def _prune_vark_by_pad(configs, named_args, **kwargs): @triton.autotune( - configs=_make_configs(_FWD_CONFIGS), + configs=_FWD_TRITON_CONFIGS + _FWD_TRITON_CONFIGS_GFX1250, key=["N", "K", "G", "A_PACK", "B_PACK"], prune_configs_by={"early_config_prune": _prune_fwd_by_k}, + do_bench=_fwd_do_bench, ) @triton.jit def _grouped_mxfp4_persistent_gemm_kernel( @@ -354,7 +399,8 @@ def grouped_gemm_mxfp4_triton_kernel( 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 + num_xcds = _num_xcds() + chunk = 64 if num_sms >= num_xcds * 64 else 32 # BLOCK_SIZE_K / num_warps / num_stages / waves_per_eu / matrix_instr_nonkdim / # kpack come from the @triton.autotune config; the grid and the host-chosen @@ -386,7 +432,7 @@ def grouped_gemm_mxfp4_triton_kernel( BLOCK_SIZE_N=BN, GROUP_SIZE_M=GM, NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, + NUM_XCDS=num_xcds, CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, @@ -440,7 +486,8 @@ def grouped_gemm_a8w4_triton_kernel( 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 + num_xcds = _num_xcds() + chunk = 64 if num_sms >= num_xcds * 64 else 32 _grouped_mxfp4_persistent_gemm_kernel[(num_sms,)]( a_e4m3, @@ -469,7 +516,7 @@ def grouped_gemm_a8w4_triton_kernel( BLOCK_SIZE_N=BN, GROUP_SIZE_M=GM, NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, + NUM_XCDS=num_xcds, CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, @@ -492,7 +539,7 @@ def grouped_gemm_a8w4_triton_kernel( @triton.autotune( - configs=_make_configs(_VARK_CONFIGS), + configs=_VARK_TRITON_CONFIGS + _VARK_TRITON_CONFIGS_GFX1250, key=["OUT_M", "OUT_N", "G", "can256"], prune_configs_by={"early_config_prune": _prune_vark_by_pad}, restore_value=["C"], @@ -626,13 +673,14 @@ 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 = 256, 128 + BM, BN = (128, 256) if get_arch() == "gfx1250" else (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 + num_xcds = _num_xcds() + chunk = 64 if num_sms >= num_xcds * 64 else 32 # can256 (256-divisibility of every padded group length) is part of the # autotune key so a cached BLOCK_K=256 config is never reused on a shape that @@ -669,7 +717,7 @@ def grouped_gemm_mxfp4_variable_k_triton_kernel( BLOCK_SIZE_N=BN, GROUP_SIZE_M=GM, NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, + NUM_XCDS=num_xcds, CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, 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 f3fb2a404f..022f56c3f2 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -5,7 +5,6 @@ from __future__ import annotations -import os from typing import List, Optional, Sequence, Tuple import torch @@ -22,16 +21,19 @@ # 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 # Archs with the scaled-FP4 matrix op used by tl.dot_scaled: gfx950 (MFMA), gfx1250 (WMMA). MXFP4_GEMM_ARCHS = ((9, 5), (12, 5)) +def _wgrad_pad_multiple() -> int: + """Per-group wgrad M padding multiple; must be >= the largest wgrad BLOCK_K the + kernel may pick. 256 on gfx1250, where the BLOCK_K=256 wgrad is 1.2-1.6x faster + (padding included).""" + return 256 if get_device_compute_capability() == (12, 5) else BLOCK_SIZE_K + + def _quantizer(*, rowwise: bool, columnwise: bool) -> MXFP4Quantizer: """MXFP4 quantizer configured for the plain (un-swizzled) Triton layout.""" return MXFP4Quantizer( @@ -203,7 +205,7 @@ def _col_operand_grouped_padded( x: torch.Tensor, m_splits: Sequence[int], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Per-group col-wise MXFP4 with each group's M padded to WGRAD_PAD_MULTIPLE. + """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 @@ -212,7 +214,8 @@ def _col_operand_grouped_padded( two groups and the result equals casting each group separately. """ m_list = [int(m) for m in m_splits] - pad_list = [round_up_to_nearest_multiple(m, WGRAD_PAD_MULTIPLE) for m in m_list] + pad = _wgrad_pad_multiple() + pad_list = [round_up_to_nearest_multiple(m, pad) for m in m_list] offs = [0] for m_pad in pad_list: offs.append(offs[-1] + m_pad) @@ -605,9 +608,8 @@ def grouped_gemm_mxfp4_wgrad( # 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) - ) + pad = _wgrad_pad_multiple() + can256 = int(all(round_up_to_nearest_multiple(int(m), pad) % 256 == 0 for m in m_splits)) return grouped_gemm_mxfp4_variable_k_triton_kernel( lhs_data, From b4e9e1956bc571bd3440512e1e488819c90ad81c Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 20 Oct 2026 20:59:59 +0000 Subject: [PATCH 24/26] factor out knob helper --- .../blockwise_fp8_grouped_gemm.py | 72 +------------------ .../pytorch/triton_kernels/common.py | 68 ++++++++++++++++++ .../triton_kernels/grouped_gemm_mxfp4.py | 5 +- 3 files changed, 71 insertions(+), 74 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py index b37be0071f..7bf6d1d3a3 100644 --- a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -5,14 +5,11 @@ # Blockwise FP8 grouped GEMM Triton kernels (MoE), adapted from AMD Primus-Turbo # (primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py). -import contextlib -import os - import torch import triton import triton.language as tl -from .common import is_cdna3 +from .common import _amd_compiler_knobs from .gmm.pid_preprocessing import remap_xcd_chunked @@ -20,73 +17,6 @@ def get_num_cus() -> int: return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count -# ── AMD compiler knobs (scoped) ── -# The AMD codegen toggles live on ``triton.knobs.amd`` and are process-global: they -# are read when a kernel is compiled, so leaving them set would silently change the -# codegen of every unrelated Triton kernel compiled afterwards (cast, gmm, MXFP8, -# ...). We therefore apply them only around our own launches and restore the prior -# state on exit. Triton's setter keeps the backing env var in sync with the -# attribute, so we snapshot/restore both. -_AMD_KNOB_ENV = { - "use_async_copy": "TRITON_HIP_USE_ASYNC_COPY", - "scalarize_packed_fops": "AMDGCN_SCALARIZE_PACKED_FOPS", - "use_block_pingpong": "TRITON_HIP_USE_BLOCK_PINGPONG", -} - - -def _amd_knob_overrides(*, is_tn: bool) -> dict: - """Knob values to apply for the enclosed GEMM launch. - - gfx950 (CDNA4): force async_copy / scalarize / block_pingpong on. - gfx942 (CDNA3): async_copy / scalarize help NT/NN but regress TN/wgrad - ~5-8%, so gate them on the GEMM layout (block_pingpong left at its default). - """ - if is_cdna3(): - enable = not is_tn - return {"use_async_copy": enable, "scalarize_packed_fops": enable} - return { - "use_async_copy": True, - "scalarize_packed_fops": True, - "use_block_pingpong": True, - } - - -@contextlib.contextmanager -def _amd_compiler_knobs(*, is_tn: bool, overrides: dict | None = None): - """Temporarily apply AMD Triton compiler knobs, restoring prior state on exit. - - Scoped rather than set once because the knobs are process-global and read at - kernel compile time; leaving them set would leak into other kernels' codegen. - ``overrides`` replaces the default :func:`_amd_knob_overrides` set. - """ - amd = getattr(getattr(triton, "knobs", None), "amd", None) - if amd is None: - yield - return - if overrides is None: - overrides = _amd_knob_overrides(is_tn=is_tn) - saved = { - name: (name in amd.__dict__, amd.__dict__.get(name), os.environ.get(env_key)) - for name, env_key in _AMD_KNOB_ENV.items() - if name in overrides - } - try: - for name, value in overrides.items(): - setattr(amd, name, value) - yield - finally: - for name, (had_override, override_val, env_val) in saved.items(): - if had_override: - amd.__dict__[name] = override_val - else: - amd.__dict__.pop(name, None) - env_key = _AMD_KNOB_ENV[name] - if env_val is None: - os.environ.pop(env_key, None) - else: - os.environ[env_key] = env_val - - # Number of XCDs used for round-robin PID swizzling (see ``remap_xcd_chunked``) # to improve L2 cache locality. This cannot be queried programmatically, so it is # hardcoded. NOTE: 8 is only correct for a full-config SPX-mode device (e.g. MI300X/ diff --git a/transformer_engine/pytorch/triton_kernels/common.py b/transformer_engine/pytorch/triton_kernels/common.py index 212c8cd695..fe7d09cfcf 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -1,6 +1,8 @@ # Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. # License for AMD contributions = MIT. See LICENSE for more information +import contextlib +import os import torch import triton import triton.language as tl @@ -14,6 +16,72 @@ def get_arch(): def is_cdna3(): return get_arch() == "gfx942" +# ── AMD compiler knobs (scoped) ── +# The AMD codegen toggles live on ``triton.knobs.amd`` and are process-global: they +# are read when a kernel is compiled, so leaving them set would silently change the +# codegen of every unrelated Triton kernel compiled afterwards (cast, gmm, MXFP8, +# ...). We therefore apply them only around our own launches and restore the prior +# state on exit. Triton's setter keeps the backing env var in sync with the +# attribute, so we snapshot/restore both. +_AMD_KNOB_ENV = { + "use_async_copy": "TRITON_HIP_USE_ASYNC_COPY", + "scalarize_packed_fops": "AMDGCN_SCALARIZE_PACKED_FOPS", + "use_block_pingpong": "TRITON_HIP_USE_BLOCK_PINGPONG", +} + + +def _amd_knob_overrides(*, is_tn: bool) -> dict: + """Knob values to apply for the enclosed GEMM launch. + + gfx950 (CDNA4): force async_copy / scalarize / block_pingpong on. + gfx942 (CDNA3): async_copy / scalarize help NT/NN but regress TN/wgrad + ~5-8%, so gate them on the GEMM layout (block_pingpong left at its default). + """ + if is_cdna3(): + enable = not is_tn + return {"use_async_copy": enable, "scalarize_packed_fops": enable} + return { + "use_async_copy": True, + "scalarize_packed_fops": True, + "use_block_pingpong": True, + } + + +@contextlib.contextmanager +def _amd_compiler_knobs(*, is_tn: bool, overrides: dict | None = None): + """Temporarily apply AMD Triton compiler knobs, restoring prior state on exit. + + Scoped rather than set once because the knobs are process-global and read at + kernel compile time; leaving them set would leak into other kernels' codegen. + ``overrides`` replaces the default :func:`_amd_knob_overrides` set. + """ + amd = getattr(getattr(triton, "knobs", None), "amd", None) + if amd is None: + yield + return + if overrides is None: + overrides = _amd_knob_overrides(is_tn=is_tn) + saved = { + name: (name in amd.__dict__, amd.__dict__.get(name), os.environ.get(env_key)) + for name, env_key in _AMD_KNOB_ENV.items() + if name in overrides + } + try: + for name, value in overrides.items(): + setattr(amd, name, value) + yield + finally: + for name, (had_override, override_val, env_val) in saved.items(): + if had_override: + amd.__dict__[name] = override_val + else: + amd.__dict__.pop(name, None) + env_key = _AMD_KNOB_ENV[name] + if env_val is None: + os.environ.pop(env_key, None) + else: + os.environ[env_key] = env_val + get_torch_e4m3_type = lambda: torch.float8_e4m3fn if not is_cdna3() else torch.float8_e4m3fnuz get_torch_e5m2_type = lambda: torch.float8_e5m2 if not is_cdna3() else torch.float8_e5m2fnuz diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index 89392721a8..f11a3fa8ff 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -20,8 +20,7 @@ import triton import triton.language as tl -from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs -from .common import get_arch +from .common import _amd_compiler_knobs, get_arch # =============================================================================== @@ -39,7 +38,7 @@ def _scoped_amd_knobs(*, is_tn): - """Decorator applying blockwise_fp8's scoped AMD compiler knobs around a + """Decorator applying the scoped AMD compiler knobs (common.py) 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; on gfx1250 only async_copy is disabled. From 8780523f80f5eb2af3d82f7ba727971a106de528 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 27 Oct 2026 11:01:47 -0500 Subject: [PATCH 25/26] centralize gfx1250 Triton knob override in _amd_knob_overrides --- .../pytorch/triton_kernels/common.py | 5 ++++ .../triton_kernels/grouped_gemm_mxfp4.py | 29 +++++-------------- 2 files changed, 13 insertions(+), 21 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/common.py b/transformer_engine/pytorch/triton_kernels/common.py index fe7d09cfcf..ddde544a79 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -36,10 +36,15 @@ def _amd_knob_overrides(*, is_tn: bool) -> dict: gfx950 (CDNA4): force async_copy / scalarize / block_pingpong on. gfx942 (CDNA3): async_copy / scalarize help NT/NN but regress TN/wgrad ~5-8%, so gate them on the GEMM layout (block_pingpong left at its default). + gfx1250 (CDNA5): Triton's async copy (default-on) corrupts the pipelined + prologue K-iteration of the scaled-WMMA dot, so force it off; the other knobs + are left at their defaults. """ if is_cdna3(): enable = not is_tn return {"use_async_copy": enable, "scalarize_packed_fops": enable} + if get_arch() == "gfx1250": + return {"use_async_copy": False} return { "use_async_copy": True, "scalarize_packed_fops": True, diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index f11a3fa8ff..a332308d8c 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -14,7 +14,6 @@ from __future__ import annotations import functools -import warnings import torch import triton @@ -26,35 +25,23 @@ # =============================================================================== # 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. +# The grouped FP4 kernels want the arch-specific knob set (see +# common._amd_knob_overrides) 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. # =============================================================================== -# gfx1250: Triton's async copy (on by default there) corrupts the pipelined -# prologue K-iteration of the scaled-WMMA dot, so it is turned off. -_GFX1250_KNOBS = {"use_async_copy": False} - def _scoped_amd_knobs(*, is_tn): - """Decorator applying the scoped AMD compiler knobs (common.py) 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; on gfx1250 - only async_copy is disabled. + """Decorator applying the scoped, arch-specific AMD compiler knobs (common.py) + around a launch. ``is_tn`` selects the gfx942 layout-gated set; the gfx1250 + async_copy override lives in :func:`common._amd_knob_overrides`. """ def deco(func): @functools.wraps(func) def wrapper(*args, **kwargs): - overrides = None - if get_arch() == "gfx1250": - overrides = _GFX1250_KNOBS - warnings.warn( - "Grouped MXFP4 Triton GEMM: disabling Triton async_copy on gfx1250 (it" - " corrupts the pipelined scaled-WMMA dot); performance may be reduced." - ) - with _amd_compiler_knobs(is_tn=is_tn, overrides=overrides): + with _amd_compiler_knobs(is_tn=is_tn): return func(*args, **kwargs) return wrapper From 09dda8ad097f364fb0c89413442848d634842a0d Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 28 Oct 2026 11:20:29 -0500 Subject: [PATCH 26/26] change to get_device_compute_capability --- .../pytorch/triton_kernels/common.py | 3 ++- .../pytorch/triton_kernels/grouped_gemm_mxfp4.py | 15 +++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/triton_kernels/common.py b/transformer_engine/pytorch/triton_kernels/common.py index ddde544a79..ebdeb24a8f 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -8,6 +8,7 @@ import triton.language as tl import transformer_engine_torch as tex from functools import cache +from ..utils import get_device_compute_capability @cache def get_arch(): @@ -43,7 +44,7 @@ def _amd_knob_overrides(*, is_tn: bool) -> dict: if is_cdna3(): enable = not is_tn return {"use_async_copy": enable, "scalarize_packed_fops": enable} - if get_arch() == "gfx1250": + if get_device_compute_capability() == (12, 5): return {"use_async_copy": False} return { "use_async_copy": True, diff --git a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py index a332308d8c..8b5630e5e1 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4.py @@ -19,7 +19,8 @@ import triton import triton.language as tl -from .common import _amd_compiler_knobs, get_arch +from ..utils import get_device_compute_capability +from .common import _amd_compiler_knobs # =============================================================================== @@ -60,7 +61,7 @@ def wrapper(*args, **kwargs): def _num_xcds(): """Chiplet-remap width. gfx1250 measured faster with the plain round-robin pid order (up to ~15% on MoE grouped GEMMs), so the remap is skipped there.""" - return 1 if get_arch() == "gfx1250" else NUM_XCDS + return 1 if get_device_compute_capability() == (12, 5) else NUM_XCDS @triton.jit @@ -162,21 +163,23 @@ def _prune_fwd_by_k(configs, named_args, **kwargs): """Keep only the current arch's configs whose BLOCK_K divides K (the fwd reduction loop is unmasked).""" K = named_args["K"] - configs = _FWD_TRITON_CONFIGS_GFX1250 if get_arch() == "gfx1250" else _FWD_TRITON_CONFIGS + is_gfx1250 = get_device_compute_capability() == (12, 5) + configs = _FWD_TRITON_CONFIGS_GFX1250 if is_gfx1250 else _FWD_TRITON_CONFIGS return [c for c in configs if K % c.kwargs["BLOCK_SIZE_K"] == 0] def _prune_vark_by_pad(configs, named_args, **kwargs): """Keep the current arch's configs; drop BLOCK_K=256 unless every padded group length is 256-divisible (can256).""" - configs = _VARK_TRITON_CONFIGS_GFX1250 if get_arch() == "gfx1250" else _VARK_TRITON_CONFIGS + is_gfx1250 = get_device_compute_capability() == (12, 5) + configs = _VARK_TRITON_CONFIGS_GFX1250 if is_gfx1250 else _VARK_TRITON_CONFIGS if named_args["can256"]: return list(configs) return [c for c in configs if c.kwargs["BLOCK_SIZE_K"] != 256] def _fwd_do_bench(kernel_call, quantiles): - if get_arch() == "gfx1250": + if get_device_compute_capability() == (12, 5): # gfx1250 downclocks under back-to-back load, which do_bench's per-launch L2 flush hides. return triton.testing.do_bench_cudagraph(kernel_call, rep=100, quantiles=quantiles) return triton.runtime.driver.active.get_benchmarker()(kernel_call, quantiles=quantiles) @@ -659,7 +662,7 @@ 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 = (128, 256) if get_arch() == "gfx1250" else (256, 128) + BM, BN = (128, 256) if get_device_compute_capability() == (12, 5) else (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