From a1c3a29f5e4f2f8fd95c7a519afef2e83dc514b3 Mon Sep 17 00:00:00 2001 From: Surya Jasper <45545431+suryajasper@users.noreply.github.com> Date: Tue, 21 Apr 2026 07:38:43 +0000 Subject: [PATCH] Add dynamic MXFP4 quantization kernel and tests Implement a Wave kernel for dynamic MXFP4 (E2M1) quantization of f32 activations with E8M0 block scales, following the algorithm from ROCm/aiter's dynamic_mxfp4_quant. The pipeline is split into three stages: - Scale computation (compute_mxfp4_scales): per-32-element block amax via IEEE 754 bit manipulation, producing E8M0 scale bytes and a broadcasted quant_scale tensor. Runs as pure PyTorch on GPU since the 32-element reduction is incompatible with AMD wave64 (which requires >=64 elements for wave-level reduce). - FP4 encoding (Wave kernel): comparison-based E2M1 encoding using seven threshold comparisons and select, entirely elementwise. - Nibble packing (pack_mxfp4_codes): pairs of i8 codes into uint8. Also adds f32_to_mxfp4 PyTorch reference in mxfp_utils.py for validation. Tests: - Lit test: FileCheck verifying key MLIR ops (absf, cmpf, select, fptosi) in the compiled kernel. - Unit pytest: 4 parametrized shapes comparing Wave kernel output against f32_to_mxfp4 reference with dequantize round-trip. - E2E pytest: Wave quant -> MXFP4 GEMM validated against the full PyTorch reference pipeline (f32_to_mxfp4 + torchScaledGemmMXFP4). Signed-off-by: Surya Jasper <45545431+suryajasper@users.noreply.github.com> Made-with: Cursor --- lit_tests/kernel/wave/dynamic_mxfp4_quant.py | 87 +++++++++ tests/kernel/wave_dynamic_mxfp4_quant_test.py | 159 ++++++++++++++++ wave_lang/kernel/wave/templates/__init__.py | 8 + .../wave/templates/dynamic_mxfp4_quant.py | 171 ++++++++++++++++++ wave_lang/kernel/wave/utils/mxfp_utils.py | 68 +++++++ 5 files changed, 493 insertions(+) create mode 100644 lit_tests/kernel/wave/dynamic_mxfp4_quant.py create mode 100644 tests/kernel/wave_dynamic_mxfp4_quant_test.py create mode 100644 wave_lang/kernel/wave/templates/dynamic_mxfp4_quant.py diff --git a/lit_tests/kernel/wave/dynamic_mxfp4_quant.py b/lit_tests/kernel/wave/dynamic_mxfp4_quant.py new file mode 100644 index 0000000000..f63ad6c78a --- /dev/null +++ b/lit_tests/kernel/wave/dynamic_mxfp4_quant.py @@ -0,0 +1,87 @@ +# RUN: python %s | FileCheck %s + +import wave_lang.kernel.lang as tkl +import wave_lang.kernel.wave as tkw +from wave_lang.kernel.lang.global_symbols import * +from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile +from wave_lang.kernel.wave.utils.general_utils import run_test + +M = tkl.sym.M +K = tkl.sym.K +BLOCK_M = tkl.sym.BLOCK_M +BLOCK_K = tkl.sym.BLOCK_K +ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + + +@run_test +def test_dynamic_mxfp4_quant(): + constraints: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(K, BLOCK_K, 1), + tkw.WaveConstraint(M, BLOCK_M), + tkw.WaveConstraint(K, BLOCK_K), + tkw.HardwareConstraint( + threads_per_wave=64, + vector_shapes={M: BLOCK_M, K: BLOCK_K}, + ), + ] + + @tkw.wave(constraints) + def quant( + x: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f32], + quant_scale: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f32], + x_codes: tkl.Memory[M, K, ADDRESS_SPACE, tkl.i8], + ): + x_reg = tkw.read(x) + qs = tkw.read(quant_scale) + qx = x_reg * qs + abs_qx = tkw.abs(qx) + + T025 = tkl.Register[M, K, tkl.f32](0.25) + T075 = tkl.Register[M, K, tkl.f32](0.75) + T125 = tkl.Register[M, K, tkl.f32](1.25) + T175 = tkl.Register[M, K, tkl.f32](1.75) + T250 = tkl.Register[M, K, tkl.f32](2.5) + T350 = tkl.Register[M, K, tkl.f32](3.5) + T500 = tkl.Register[M, K, tkl.f32](5.0) + + mag = ( + tkw.cast(abs_qx >= T025, tkl.f32) + + tkw.cast(abs_qx >= T075, tkl.f32) + + tkw.cast(abs_qx >= T125, tkl.f32) + + tkw.cast(abs_qx >= T175, tkl.f32) + + tkw.cast(abs_qx >= T250, tkl.f32) + + tkw.cast(abs_qx >= T350, tkl.f32) + + tkw.cast(abs_qx >= T500, tkl.f32) + ) + + ZERO_F32 = tkl.Register[M, K, tkl.f32](0.0) + EIGHT_F32 = tkl.Register[M, K, tkl.f32](8.0) + sign = tkw.select(qx < ZERO_F32, EIGHT_F32, ZERO_F32) + + codes = tkw.cast(mag + sign, tkl.i8) + tkw.write(codes, x_codes) + + options = WaveCompileOptions( + subs={ + M: 4, + K: 64, + BLOCK_M: 2, + BLOCK_K: 32, + ADDRESS_SPACE: GLOBAL_ADDRESS_SPACE, + }, + canonicalize=True, + compile_to_mlir=True, + ) + + quant = wave_compile(options, quant) + print(quant.asm) + + # CHECK-LABEL: test_dynamic_mxfp4_quant + # CHECK: func.func @quant + # CHECK: math.absf {{.*}} : vector<{{.*}}xf32> + # CHECK: arith.cmpf oge, {{.*}} : vector<{{.*}}xf32> + # CHECK: arith.select {{.*}} : vector<{{.*}}xi1>, vector<{{.*}}xf32> + # CHECK: arith.cmpf olt, {{.*}} : vector<{{.*}}xf32> + # CHECK: arith.select {{.*}} : vector<{{.*}}xi1>, vector<{{.*}}xf32> + # CHECK: arith.fptosi {{.*}} : vector<{{.*}}xf32> to vector<{{.*}}xi8> diff --git a/tests/kernel/wave_dynamic_mxfp4_quant_test.py b/tests/kernel/wave_dynamic_mxfp4_quant_test.py new file mode 100644 index 0000000000..f098394d87 --- /dev/null +++ b/tests/kernel/wave_dynamic_mxfp4_quant_test.py @@ -0,0 +1,159 @@ +# Copyright 2025 Advanced Micro Devices, Inc. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +import torch +import pytest + +from wave_lang.kernel.wave.compile import wave_compile +from wave_lang.kernel.wave.constraints import ScaledMMAType +from wave_lang.kernel.wave.utils.run_utils import set_default_run_config +from wave_lang.kernel.wave.utils.torch_utils import ( + device_randn, + device_zeros, +) +from wave_lang.kernel.wave.utils.mxfp_utils import ( + SCALE_GROUP_SIZE, + f32_to_mxfp4, + mxfp4_to_f32, + e8m0_to_f32, + generate_gemm_afp4wfp4_inputs, + torchScaledGemmMXFP4, +) +from wave_lang.kernel.wave.templates.dynamic_mxfp4_quant import ( + compute_mxfp4_scales, + get_dynamic_mxfp4_quant_kernel, + pack_mxfp4_codes, +) +from wave_lang.kernel.wave.templates import ( + get_tagged_mxfp4_gemm, +) +from wave_lang.kernel.wave.schedules import ( + get_mxfp4_dbuf_schedule, +) + +from .common.utils import ( + require_cdna4, + require_e2e, +) + + +# --------------------------------------------------------------------------- +# Unit test: Wave quant kernel vs PyTorch reference +# --------------------------------------------------------------------------- + + +@require_e2e +@pytest.mark.parametrize( + "shape", + [ + (4, 64), + (128, 256), + (1, 1024), + (64, 8192), + ], +) +def test_dynamic_mxfp4_quant(shape): + """Compile and run the Wave dynamic MXFP4 quant kernel, then compare + against the pure-PyTorch ``f32_to_mxfp4`` reference.""" + M_val, K_val = shape + + quant, options = get_dynamic_mxfp4_quant_kernel(shape, block_m=2) + options = set_default_run_config(options) + quant = wave_compile(options, quant) + + torch.manual_seed(42) + x = device_randn(shape, dtype=torch.float32) + + # Compute scales on GPU via PyTorch + qs, bs_e8m0, _ = compute_mxfp4_scales(x) + + # Snapshot values to CPU *before* the IREE kernel call, which may + # invalidate GPU tensors that share the PyTorch CUDA caching allocator. + x_cpu = x.cpu().clone() + bs_cpu = bs_e8m0.cpu().clone() + + # Allocate output and run Wave kernel + codes = device_zeros(shape, dtype=torch.int8) + quant(x, qs, codes) + + # PyTorch reference (runs on CPU) + ref_fp4, ref_scales = f32_to_mxfp4(x_cpu.float()) + + # Pack the Wave codes for comparison + wave_fp4 = pack_mxfp4_codes(codes.cpu()) + + torch.testing.assert_close( + wave_fp4.to(torch.uint8), ref_fp4, atol=0, rtol=0 + ) + torch.testing.assert_close(bs_cpu, ref_scales, atol=0, rtol=0) + + # Cross-check via dequantize round-trip + wave_deq = mxfp4_to_f32(wave_fp4.to(torch.uint8)) + wave_s = e8m0_to_f32( + bs_cpu.repeat_interleave(SCALE_GROUP_SIZE, dim=1).float() + ) + ref_deq = mxfp4_to_f32(ref_fp4) + ref_s = e8m0_to_f32( + ref_scales.repeat_interleave(SCALE_GROUP_SIZE, dim=1).float() + ) + torch.testing.assert_close(wave_deq * wave_s, ref_deq * ref_s) + + +# --------------------------------------------------------------------------- +# End-to-end test: Wave quant -> MXFP4 GEMM +# --------------------------------------------------------------------------- + + +@require_e2e +@require_cdna4 +@pytest.mark.parametrize("shape", [(256, 256, 1024)]) +def test_dynamic_mxfp4_quant_gemm_e2e(shape): + """Quantise f32 activations with the Wave kernel, feed packed results into + the MXFP4 scaled GEMM, and compare against the reference pipeline.""" + M_val, N_val, K_val = shape + + # -- Step 1: compile the quant kernel -- + quant, q_options = get_dynamic_mxfp4_quant_kernel( + (M_val, K_val), block_m=2 + ) + q_options = set_default_run_config(q_options) + quant = wave_compile(q_options, quant) + + # -- Step 2: compile the GEMM kernel -- + block_shape = (32, 32, 256) + gemm, g_options = get_tagged_mxfp4_gemm( + shape, + block_shape, + wave_shape=(2, 2), + mfma_variant=ScaledMMAType.F32_16x16x128_F8F6F4, + ) + schedule = get_mxfp4_dbuf_schedule() + g_options = set_default_run_config(g_options) + gemm = wave_compile(g_options, gemm, schedule) + + # -- Step 3: generate test data -- + torch.manual_seed(7) + a = device_randn((M_val, K_val), dtype=torch.float32) + _, w, _, w_scales = generate_gemm_afp4wfp4_inputs(shape) + + # -- Step 4: run Wave quant -- + qs, a_scales, _ = compute_mxfp4_scales(a) + codes = device_zeros((M_val, K_val), dtype=torch.int8) + quant(a, qs, codes) + a_fp4 = pack_mxfp4_codes(codes) + + # -- Step 5: run GEMM -- + out = device_zeros(M_val, N_val, dtype=torch.float32) + w_t = w.T.contiguous() + gemm(a_fp4, a_scales, w_t, w_scales, out) + + # -- Step 6: reference pipeline -- + ref_fp4, ref_scales = f32_to_mxfp4(a.cpu().float()) + ref_fp4 = ref_fp4.to(a.device) + ref_scales = ref_scales.to(a.device) + ref_out = torchScaledGemmMXFP4(ref_fp4, w, ref_scales, w_scales) + + torch.testing.assert_close(ref_out, out, check_dtype=False) diff --git a/wave_lang/kernel/wave/templates/__init__.py b/wave_lang/kernel/wave/templates/__init__.py index a74a263d22..08c748013b 100644 --- a/wave_lang/kernel/wave/templates/__init__.py +++ b/wave_lang/kernel/wave/templates/__init__.py @@ -6,6 +6,11 @@ from .attention_common import AttentionShape from .tagged_attention import get_tagged_bshd_attention_kernel +from .dynamic_mxfp4_quant import ( + compute_mxfp4_scales, + get_dynamic_mxfp4_quant_kernel, + pack_mxfp4_codes, +) from .tagged_mxfp4_gemm import ( get_tagged_mxfp4_gemm, get_tagged_mxfp4_gemm_preshuffle_b, @@ -15,6 +20,9 @@ ) __all__ = [ + "compute_mxfp4_scales", + "get_dynamic_mxfp4_quant_kernel", + "pack_mxfp4_codes", "AttentionShape", "get_tagged_bshd_attention_kernel", "get_tagged_mxfp4_gemm", diff --git a/wave_lang/kernel/wave/templates/dynamic_mxfp4_quant.py b/wave_lang/kernel/wave/templates/dynamic_mxfp4_quant.py new file mode 100644 index 0000000000..7d8cd95db6 --- /dev/null +++ b/wave_lang/kernel/wave/templates/dynamic_mxfp4_quant.py @@ -0,0 +1,171 @@ +# Copyright 2025 Advanced Micro Devices, Inc. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +""" +Dynamic MXFP4 quantization kernel template. + +Quantizes f32 activations to per-element E2M1 codes (unpacked, 0-15 in i8) +given precomputed per-element quantization scales. + +The full dynamic quantization pipeline is: + +1. **Scale computation** (``compute_mxfp4_scales``, pure PyTorch on GPU): + per-32-element block amax → E8M0 scale bytes + broadcasted quant_scale. +2. **FP4 encoding** (Wave kernel): comparison-based E2M1 encoding of + ``x * quant_scale`` into i8 codes [0..15]. +3. **Nibble packing** (``pack_mxfp4_codes``): pairs of i8 codes → uint8 bytes. + +The scale computation uses a wave-incompatible 32-element reduction (AMD wave64 +requires ≥ 64 elements for a wave-level reduce), so it stays in PyTorch. +""" + +import torch +from torch import Tensor + +import wave_lang.kernel.lang as tkl +import wave_lang.kernel.wave as tkw +from wave_lang.kernel.lang.global_symbols import GLOBAL_ADDRESS_SPACE +from wave_lang.kernel.wave.compile import WaveCompileOptions + + +SCALE_GROUP_SIZE = 32 + + +def compute_mxfp4_scales( + x: Tensor, +) -> tuple[Tensor, Tensor, Tensor]: + """Compute per-block E8M0 scales and broadcasted quant_scale for MXFP4. + + Runs entirely on the same device as *x* using standard PyTorch ops. + + Args: + x: float32 tensor ``[M, K]``. *K* must be divisible by 32. + + Returns: + quant_scale: ``f32[M, K]`` -- per-element scale factor (broadcasted) + bs_e8m0: ``uint8[M, K // 32]`` -- biased E8M0 scale bytes + scale_exp: ``f32[M, K // 32, 1]`` -- unbiased exponents (for debug) + """ + M, K = x.shape + x_blocked = x.reshape(M, K // SCALE_GROUP_SIZE, SCALE_GROUP_SIZE) + + amax = x_blocked.abs().amax(dim=-1, keepdim=True) + + amax_bits = amax.contiguous().view(torch.int32) + amax_bits = amax_bits + 0x200000 + mask = torch.tensor(-8388608, dtype=torch.int32, device=x.device) + amax_bits = amax_bits & mask + amax_pow2 = amax_bits.contiguous().view(torch.float32) + + scale_exp = torch.log2(amax_pow2).floor() - 2 + scale_exp = scale_exp.clamp(-127, 127) + + quant_scale_group = torch.exp2(-scale_exp) + quant_scale = quant_scale_group.expand_as(x_blocked).reshape(M, K) + + bs_e8m0 = (scale_exp + 127).to(torch.uint8).squeeze(-1) + + return quant_scale, bs_e8m0, scale_exp + + +def get_dynamic_mxfp4_quant_kernel( + shape: tuple[int, int], + block_m: int = 2, +): + """Return a Wave kernel + compile options for MXFP4 FP4 encoding. + + The kernel reads ``x[M, K]`` (f32) and precomputed ``quant_scale[M, K]`` + (f32), multiplies each element by its scale, and comparison-encodes the + result to a 4-bit E2M1 code (stored as i8 in [0, 15]). + + Args: + shape: ``(M, K)`` problem size. *K* must be divisible by 32. + block_m: Rows per workgroup tile. ``block_m * 32`` must equal 64. + + Returns: + ``(kernel_fn, WaveCompileOptions)`` + """ + M_val, K_val = shape + assert K_val % 32 == 0 + block_k = 32 + + M = tkl.sym.M + K = tkl.sym.K + BLOCK_M = tkl.sym.BLOCK_M + BLOCK_K = tkl.sym.BLOCK_K + ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE + + constraints: list[tkw.Constraint] = [ + tkw.WorkgroupConstraint(M, BLOCK_M, 0), + tkw.WorkgroupConstraint(K, BLOCK_K, 1), + tkw.WaveConstraint(M, BLOCK_M), + tkw.WaveConstraint(K, BLOCK_K), + tkw.HardwareConstraint( + threads_per_wave=64, + vector_shapes={M: BLOCK_M, K: BLOCK_K}, + ), + ] + + @tkw.wave(constraints) + def quant( + x: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f32], + quant_scale: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f32], + x_codes: tkl.Memory[M, K, ADDRESS_SPACE, tkl.i8], + ): + x_reg = tkw.read(x) + qs = tkw.read(quant_scale) + qx = x_reg * qs + abs_qx = tkw.abs(qx) + + T025 = tkl.Register[M, K, tkl.f32](0.25) + T075 = tkl.Register[M, K, tkl.f32](0.75) + T125 = tkl.Register[M, K, tkl.f32](1.25) + T175 = tkl.Register[M, K, tkl.f32](1.75) + T250 = tkl.Register[M, K, tkl.f32](2.5) + T350 = tkl.Register[M, K, tkl.f32](3.5) + T500 = tkl.Register[M, K, tkl.f32](5.0) + + mag = ( + tkw.cast(abs_qx >= T025, tkl.f32) + + tkw.cast(abs_qx >= T075, tkl.f32) + + tkw.cast(abs_qx >= T125, tkl.f32) + + tkw.cast(abs_qx >= T175, tkl.f32) + + tkw.cast(abs_qx >= T250, tkl.f32) + + tkw.cast(abs_qx >= T350, tkl.f32) + + tkw.cast(abs_qx >= T500, tkl.f32) + ) + + ZERO_F32 = tkl.Register[M, K, tkl.f32](0.0) + EIGHT_F32 = tkl.Register[M, K, tkl.f32](8.0) + sign = tkw.select(qx < ZERO_F32, EIGHT_F32, ZERO_F32) + + codes = tkw.cast(mag + sign, tkl.i8) + tkw.write(codes, x_codes) + + hyperparams = { + ADDRESS_SPACE: GLOBAL_ADDRESS_SPACE, + BLOCK_M: block_m, + BLOCK_K: block_k, + M: M_val, + K: K_val, + } + + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + ) + + return quant, options + + +def pack_mxfp4_codes(codes): + """Pack unpacked i8 E2M1 codes [M, K] into nibble-packed uint8 [M, K/2]. + + Even-indexed codes go into the low nibble, odd-indexed into the high nibble, + matching the layout expected by ``mxfp4_to_f32`` and the MXFP4 GEMM kernel. + """ + c = codes.to(torch.uint8) + return c[:, 0::2] + c[:, 1::2] * 16 diff --git a/wave_lang/kernel/wave/utils/mxfp_utils.py b/wave_lang/kernel/wave/utils/mxfp_utils.py index 901828c412..5fcfff775c 100644 --- a/wave_lang/kernel/wave/utils/mxfp_utils.py +++ b/wave_lang/kernel/wave/utils/mxfp_utils.py @@ -95,6 +95,74 @@ def e8m0_to_f32(x: Tensor) -> Tensor: return x_f32 +def f32_to_mxfp4( + x: Tensor, block_size: int = SCALE_GROUP_SIZE +) -> tuple[Tensor, Tensor]: + """Quantize an f32 tensor to packed MXFP4 (E2M1) with E8M0 block scales. + + Reference implementation matching the dynamic_mxfp4_quant algorithm from + ROCm/aiter: per-block power-of-two scale via IEEE 754 bit manipulation, + then comparison-based FP4 encoding with round-half-up at midpoints. + + Args: + x: float32 tensor of shape ``[M, K]``. *K* must be divisible by + *block_size*. + block_size: Number of elements per scale group (default 32 per the + OCP MX specification). + + Returns: + x_fp4: ``uint8[M, K // 2]`` -- packed FP4 nibble pairs + scales: ``uint8[M, K // block_size]`` -- biased E8M0 block scales + """ + assert x.dtype == torch.float32 + M, K = x.shape + assert K % block_size == 0 + + x_blocked = x.reshape(M, K // block_size, block_size) + + # ---- Phase 1: block-scale computation ---- + amax = x_blocked.abs().amax(dim=-1, keepdim=True) + + # Round amax up to a power-of-two boundary via IEEE 754 bit trick: + # add half-ULP at bit 21, then zero the mantissa. + amax_bits = amax.contiguous().view(torch.int32) + amax_bits = amax_bits + 0x200000 + mask = torch.tensor(-8388608, dtype=torch.int32, device=x.device) # 0xFF800000 + amax_bits = amax_bits & mask + amax_pow2 = amax_bits.contiguous().view(torch.float32) + + # Unbiased E8M0 exponent: floor(log2(amax_pow2)) - 2, clamped. + # -2 because FP4 E2M1 emax = 2: max representable = 6.0 = 1.5 * 2^2. + scale_exp = torch.log2(amax_pow2).floor() - 2 + scale_exp = scale_exp.clamp(-127, 127) + + quant_scale = torch.exp2(-scale_exp) + bs_e8m0 = (scale_exp + 127).to(torch.uint8).squeeze(-1) + + # ---- Phase 2: per-element FP4 encoding ---- + qx = x_blocked * quant_scale + abs_qx = qx.abs() + + # Comparison-based E2M1 magnitude code (round-half-up at midpoints). + # Midpoints between adjacent representable values: + # 0↔0.5 → 0.25, 0.5↔1 → 0.75, 1↔1.5 → 1.25, 1.5↔2 → 1.75, + # 2↔3 → 2.5, 3↔4 → 3.5, 4↔6 → 5.0 + thresholds = torch.tensor( + [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], + dtype=torch.float32, + device=x.device, + ) + mag_code = (abs_qx.unsqueeze(-1) >= thresholds).sum(dim=-1).to(torch.uint8) + sign_code = (qx < 0).to(torch.uint8) * 8 + fp4_code = mag_code + sign_code + + # ---- Phase 3: nibble packing ---- + fp4_flat = fp4_code.reshape(M, K) + x_fp4 = fp4_flat[:, 0::2] + fp4_flat[:, 1::2] * 16 + + return x_fp4, bs_e8m0 + + def b_preshuffle(b: Tensor) -> Tensor: """Preshuffle packed MXFP4 B weights for direct global reads.