diff --git a/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py new file mode 100644 index 0000000000..a3adf33aa2 --- /dev/null +++ b/tests/pytorch/mxfp4/test_a8w4_grouped_recipe.py @@ -0,0 +1,499 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""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 +``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 ( + MXFP4_GEMM_ARCHS, + _col_operand_grouped_padded, + _row_operand, + _row_operand_mxfp8, + _row_operand_mxfp8_torch, + grouped_gemm_mxfp4_dgrad, + 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, + 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, +) + +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 + + +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.""" + _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 grouped Triton 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_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: routing + exact output/dgrad/wgrad vs the reference.""" + _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) + 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(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 _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 _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() + 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)) + + +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" + + +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/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) + ), 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)}" + + +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) + + +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" + + +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, + single_grouped_weight=False, + single_grouped_bias=False, + 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/tests/pytorch/mxfp4/test_a8w4_recipe.py b/tests/pytorch/mxfp4/test_a8w4_recipe.py new file mode 100644 index 0000000000..5588b8e350 --- /dev/null +++ b/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() 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/custom_recipes/quantization_mxfp4_grouped.py b/transformer_engine/pytorch/custom_recipes/quantization_mxfp4_grouped.py new file mode 100644 index 0000000000..a3feabc245 --- /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`` -- 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") diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1e039816d7..7ab649a909 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -749,6 +749,109 @@ def _is_blockwise_fp8_triton_grouped_gemm_supported( ) ) + @staticmethod + def _is_grouped_mxfp4_triton_supported( + *, + fp8, + recipe, + input_quantizers, + weight_quantizers, + grad_output_quantizers, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + dgrad_out, + fuse_wgrad_accumulation, + wgrad_store, + single_grouped_weight, + single_grouped_bias, + in_features, + out_features, + ) -> bool: + """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 + dgrad/wgrad (see :meth:`_forward_grouped_mxfp4_triton`). Gated off for options this + 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. (A forward ``out`` buffer is honored.) + """ + if not (IS_HIP_EXTENSION and fp8 and recipe is not None and recipe.custom()): + return False + # 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 + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits 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()) + # The fast path works on per-expert weights and returns per-expert grads. + or single_grouped_weight + or single_grouped_bias + ): + return False + 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 + + 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 + if not isinstance(q, MXFP4QuantizerRef): + return False + 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). + # 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 + 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 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 def _forward_blockwise_fp8_triton( ctx, @@ -1052,6 +1155,208 @@ 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, + is_first_microbatch, + weight_workspaces, + cache_weight, + out=None, + ): + """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`). 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, + ) + from ..custom_recipes.quantization_mxfp4_grouped import MXFP8E4M3QuantizerRef + + 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] + ) + + # Activation: quantize the whole grouped [total_M, K] once (row-wise scales + # 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] + a_is_mxfp8 = isinstance(input_quantizer, MXFP8E4M3QuantizerRef) + if a_is_mxfp8: + # 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: + 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. + # 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. + # 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 + # 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: + _, b_data, b_scale, w_col_data, w_col_scale = cached + else: + 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: + (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() + if cache_weight: + 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 + ) + result = grouped_gemm_mxfp4_fprop_prequantized( + a_data, + a_scale, + b_data, + b_scale, + m_splits_list, + a_is_mxfp8=a_is_mxfp8, + out_dtype=activation_dtype, + scale_transposed=a_is_mxfp8, + out=result_out, + ) + + 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 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 + ctx.save_for_backward(a, w_col_data, w_col_scale) + + return result.view(-1, *inp.shape[1:-1], result.shape[-1]), new_workspaces + + @staticmethod + def _backward_grouped_mxfp4_triton(ctx, grad_output): + """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 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 Triton 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( @@ -1161,6 +1466,45 @@ def forward( pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, ) + # 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( + fp8=fp8, + 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, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + dgrad_out=dgrad_out, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + 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, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + out=out, + ) + backward_needs_input = is_grad_enabled and weight_requires_grad if backward_override is None and save_original_input and backward_needs_input: if recipe is not None and recipe.delayed(): @@ -1806,6 +2150,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.grouped_tensor_supported: @@ -2958,11 +3304,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 and IS_HIP_EXTENSION: + from ..custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef + + if isinstance(weight_quantizers[0], MXFP4QuantizerRef): + ws_prefix = "mxfp4_grouped_weight" if self.single_grouped_weight: weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None] else: 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 ) @@ -3012,7 +3368,7 @@ def forward( if ws is not None: if isinstance(ws, torch.Tensor): ws = ws.detach() - key = "weight" if self.single_grouped_weight else f"weight{i}" + key = "weight" if self.single_grouped_weight else f"{ws_prefix}{i}" self._fp8_workspaces[key] = ws finally: 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..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,71 +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): - """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. - """ - amd = getattr(getattr(triton, "knobs", None), "amd", None) - if amd is None: - yield - return - 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/cast_transpose.py b/transformer_engine/pytorch/triton_kernels/cast_transpose.py index 377e6a6331..c2fcd097fa 100644 --- a/transformer_engine/pytorch/triton_kernels/cast_transpose.py +++ b/transformer_engine/pytorch/triton_kernels/cast_transpose.py @@ -868,6 +868,45 @@ 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, 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))``. + """ + 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) + 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) + 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 + 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_m, scale_stride_k, + 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/common.py b/transformer_engine/pytorch/triton_kernels/common.py index 212c8cd695..ebdeb24a8f 100644 --- a/transformer_engine/pytorch/triton_kernels/common.py +++ b/transformer_engine/pytorch/triton_kernels/common.py @@ -1,11 +1,14 @@ # 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 import transformer_engine_torch as tex from functools import cache +from ..utils import get_device_compute_capability @cache def get_arch(): @@ -14,6 +17,77 @@ 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). + 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_device_compute_capability() == (12, 5): + return {"use_async_copy": False} + 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 1093813361..8b5630e5e1 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. @@ -19,38 +19,24 @@ import triton import triton.language as tl -from .blockwise_fp8_grouped_gemm import _amd_compiler_knobs - - -# =============================================================================== -# Arch helper -# =============================================================================== - - -@functools.lru_cache(maxsize=None) -def _is_gfx950_for(device_index: int) -> bool: - props = torch.cuda.get_device_properties(device_index) - return (props.major, props.minor) == (9, 5) - - -def _is_gfx950() -> bool: - # Key the cache by device index so a set_device change isn't masked. - return _is_gfx950_for(torch.cuda.current_device()) +from ..utils import get_device_compute_capability +from .common import _amd_compiler_knobs # =============================================================================== # 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. # =============================================================================== + 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. + """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): @@ -72,6 +58,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_device_compute_capability() == (12, 5) else NUM_XCDS + + @triton.jit def _chiplet_transform_chunked( pid, @@ -123,6 +115,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): @@ -140,19 +153,38 @@ 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"] + 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): - """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).""" + 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_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) + + # ########################################################################### # Forward (NT): C[g] = A[g] @ B[g]^T # A (total_M, K/2) fp4-packed, A_scale (total_M, K/32) e8m0 @@ -161,9 +193,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( @@ -202,6 +235,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 +294,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 +339,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 +360,8 @@ def grouped_gemm_mxfp4_triton_kernel( group_offs_out=None, 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. @@ -311,11 +372,14 @@ 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) 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] @@ -324,7 +388,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 @@ -347,21 +412,22 @@ 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, NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, + NUM_XCDS=num_xcds, CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, A_PACK=2, # fp4 A; passed explicitly so A_PACK/B_PACK bind into the autotune key B_PACK=2, + SCALE_TRANSPOSED=scale_transposed, ) return c @@ -378,6 +444,8 @@ def grouped_gemm_a8w4_triton_kernel( group_offs_out=None, 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). @@ -391,11 +459,14 @@ 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) 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] @@ -404,7 +475,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, @@ -424,16 +496,16 @@ 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, NUM_SMS=num_sms, - NUM_XCDS=NUM_XCDS, + NUM_XCDS=num_xcds, CHUNK_SIZE=chunk, CACHE_MODIFIER=".ca", VEC=VEC_SIZE, @@ -441,6 +513,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 @@ -455,7 +528,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"], @@ -589,13 +662,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_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 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 @@ -632,7 +706,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 81702da6e7..022f56c3f2 100644 --- a/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py +++ b/transformer_engine/pytorch/triton_kernels/grouped_gemm_mxfp4_impl.py @@ -1,33 +1,37 @@ # 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 -import os from typing import List, Optional, Sequence, Tuple import torch import torch.nn.functional as F from ..tensor.mxfp4_tensor import MXFP4Quantizer -from ..utils import round_up_to_nearest_multiple +from ..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, ) +from .cast_transpose import mxfp8_e4m3_rowwise_downcast # 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: @@ -78,12 +82,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 not _is_gfx950(): +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 +109,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 @@ -145,15 +147,14 @@ def _row_col_operand( return row_data, row_scale, col_data, col_scale -def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """Row-wise MXFP8 (e4m3): ``x`` [M, K] -> (data [M, K] e4m3, scale [M, K/32] u8). +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. 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/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]``. """ fmax = torch.finfo(torch.float8_e4m3fn).max M, K = x.shape @@ -163,9 +164,22 @@ def _row_operand_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 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, transpose_scale: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """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() in MXFP4_GEMM_ARCHS: + 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( weights: Sequence[torch.Tensor], ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: @@ -191,30 +205,29 @@ 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 - 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 = _wgrad_pad_multiple() + pad_list = [round_up_to_nearest_multiple(m, pad) 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 @@ -242,7 +255,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: @@ -284,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). @@ -297,18 +311,22 @@ 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``. """ K = a.shape[1] - _require_gfx950() + _require_mxfp4_gemm_arch() _check_contract(K, "K") if weight_row is not None: b_data, b_scale = weight_row N = b_data.shape[1] _check_splits(m_splits, a.shape[0], b_data.shape[0]) + 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)) @@ -319,8 +337,14 @@ 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) - - a_data, a_scale = _row_operand_mxfp8(a) # (total_M, K) e4m3, (total_M, K/32) e8m0 + 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 ([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( a_data, @@ -333,6 +357,153 @@ def grouped_gemm_a8w4_fprop( group_offs_out=group_offs, out_dtype=out_dtype, num_cu=num_cu, + scale_transposed=True, + ) + + +def _check_prequantized_operands( + a_data: torch.Tensor, + a_scale: torch.Tensor, + b_data: torch.Tensor, + b_scale: torch.Tensor, + *, + a_is_mxfp8: bool, + scale_transposed: bool = False, +) -> 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]. 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: + 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, + 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, + scale_transposed: bool = False, + out: Optional[torch.Tensor] = 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. + 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``. + """ + _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] + 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, scale_transposed=scale_transposed + ) + 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 + 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, + scale_transposed=scale_transposed, + out=out, ) @@ -363,7 +534,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: @@ -426,7 +597,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]) @@ -437,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,