extend MXFP4 recipe to GroupedLinear and a8w4 (Triton grouped GEMM kernels)

tests/pytorch/mxfp4/test_a8w4_recipe.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# License for AMD contributions = MIT. See LICENSE for more 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
+from transformer_engine.pytorch.triton_kernels.grouped_gemm_mxfp4_impl import MXFP4_GEMM_ARCHS
+from transformer_engine.pytorch.utils import get_device_compute_capability
+
+pytestmark = pytest.mark.skipif(
+    get_device_compute_capability() not in MXFP4_GEMM_ARCHS,
+    reason="a8w4 recipe requires gfx950 or gfx1250",
+)
+
+
+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 grouped Triton 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()

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)

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`` -- a4w4: 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 gemm
+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: 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: gemm.GEMMType = gemm.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"`` -> a4w4 (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
+
+
+# a4w4: 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")