[gfx1250] Fuse A4W4 stage1 FP4 quantization

aiter/ops/flydsl/batched_gemm_mxfp4.py

@@ -4,7 +4,7 @@
 """Thin strided-batched MXFP4/MXFP6/MXFP8 preshuffle GEMM launcher: out[b] =
 dequant(x[b]) @ dequant(w[b]).T, per-1x32 e8m0 scales folded into a scaled 16x16x128
 matrix op. gfx950 uses the wave64 MFMA path (a4w4/a8w4, fp4/fp6/fp8 A); gfx1250 uses the
-wave32 WMMA path (a8w4 only: MXFP8 E4M3 A x MXFP4 B). Operands are preshuffled + laid out
+wave32 WMMA path (MXFP4/MXFP8 A x MXFP4 B). Operands are preshuffled + laid out
 by the caller (once, off the launch path) -- see the arch-specific preshuffle in the tests.
 layout 'bmn' = contiguous [B,M,N], 'mbn' = the deepseek-v4 grouped-output [M,B,N] (returned
 as a non-contiguous [B,M,N] view)."""
@@ -59,8 +59,9 @@ def flydsl_grouped_gemm_a8w4_masked(
     Mirrors the MoE grouped-gemm contiguous-M scheduling: a compact grid over the
     (1, contiguous_m, *) buffers, with a per-M-tile expert id (``tile_expert``)
     selecting the per-expert B / B-scale slab. Only valid M tiles launch.
-      out          (1, contiguous_m, N)  bf16/f16  (or fp8 payload when stage1_quant_out)
-      a            (1, contiguous_m, K)  uint8 (fp8 payload)
+      out          (1, contiguous_m, N) bf16/f16; quantized stage1 output is
+                   uint8 with N//2 bytes for MXFP8 or N//4 bytes for MXFP4
+      a            (1, contiguous_m, K) logical elements in MXFP4/MXFP8 payload
       w            (E, N, K//2) uint8 (moe_shuffle_weight == cat_e shuffle_weight_gfx1250)
       a_scales     grouped A-scale (1, contiguous_m//wmma_rep, (K//32)*wmma_rep) viewed int32
       w_scales     n32k4 B-scale (E, N//32, (K//32)*32) viewed int32
@@ -72,11 +73,11 @@ def flydsl_grouped_gemm_a8w4_masked(
     The betas are runtime kernel arguments, so all SiTUv2 shapes share one
     compiled kernel.
 
-    When ``stage1_quant_out=1`` (fp8), the epilogue fuses the activation + MX
-    fp8 quantization + e8m0 scale preshuffle into the kernel.  ``out`` receives
-    the fp8 payload (uint8, 1 byte/elem) and ``quant_scale`` receives the
-    preshuffled e8m0 scale (uint8).  ``quant_wmma_rep`` is gemm2's
-    ``warp_tile_m // 16``, controlling the scale preshuffle tile geometry.
+    ``stage1_quant_out`` selects the fused activation-quant epilogue: 0 emits
+    bf16/f16, 1 emits MXFP8 (one byte/element), and 2 emits packed MXFP4 (one
+    byte/two elements). ``quant_scale`` receives the preshuffled E8M0 scales.
+    ``quant_wmma_rep`` is gemm2's ``warp_tile_m // 16``, controlling the scale
+    preshuffle tile geometry.
     """
     from .kernels.mxfp4_preshuffle_gfx1250_tdm import launch_gemm_a8w4_tdm
 
@@ -89,6 +90,15 @@ def flydsl_grouped_gemm_a8w4_masked(
             raise ValueError(f"situ_beta must be > 0, got {situ_beta!r}")
         if float(situ_linear_beta) <= 0.0:
             raise ValueError(f"situ_linear_beta must be > 0, got {situ_linear_beta!r}")
+    if stage1_quant_out not in (0, 1, 2):
+        raise ValueError(
+            f"stage1_quant_out must be 0 (bf16), 1 (fp8), or 2 (fp4), "
+            f"got {stage1_quant_out!r}"
+        )
+    if stage1_quant_out and (not stage1_act or quant_scale is None):
+        raise ValueError(
+            "stage1_quant_out requires an activation epilogue and quant_scale"
+        )
     nb = min(num_buffers, max(1, K // tile_k))
     has_bias = 1 if bias is not None else 0
     bias_ptr = ptr_arg(bias) if bias is not None else ptr_arg(a)

aiter/ops/flydsl/grouped_moe_gfx1250.py

@@ -526,17 +526,20 @@ def _grouped_a8w4_tdm_moe(
         num_valid_routes=_ep_nvr,
     )
 
-    # Fuse gemm1 silu/swiglu + fp8 quantization + scale preshuffle into the
-    # kernel epilogue (a8w4 only), eliminating the standalone
-    # flydsl_moe_fused_quant_preshuffle call between gemm1 and gemm2.
-    _fuse_quant = (not _is_fp4) and (_b1 is None)
+    # Fuse gemm1 activation + MX quantization + scale preshuffle into the
+    # kernel epilogue, eliminating the standalone quant pass between gemm1 and
+    # gemm2. The epilogue currently requires a bias-free stage1 (Kimi-K3) and
+    # four WN subtiles per MX block within each wave.
+    _wmma_n_rep = tile_n // (4 * 16)  # TDM launcher uses n_warp=4.
+    _fuse_quant = _b1 is None and _wmma_n_rep >= 4 and _wmma_n_rep % 4 == 0
+    _stage1_quant_out = 2 if _is_fp4 else 1
     w1_u8 = _grouped_weight_uint8(w1)
     w1s_i32 = w1_scale.reshape(-1).view(torch.int32)
 
     if _fuse_quant:
-        # Pre-allocate fp8 payload + preshuffled e8m0 scale for gemm1 output.
+        # Pre-allocate MX payload + preshuffled E8M0 scale for gemm1 output.
         # These are written directly by the kernel's fused quant epilogue.
-        payload_bytes = inter_dim  # fp8: 1 byte per element
+        payload_bytes = inter_dim // 2 if _is_fp4 else inter_dim
         scale_bytes = inter_dim // 32  # one e8m0 byte per 32-element MX block
         a2_payload = torch.empty(
             (1, contiguous_m, payload_bytes), dtype=torch.uint8, device=device
@@ -546,9 +549,8 @@ def _grouped_a8w4_tdm_moe(
             dtype=torch.uint8,
             device=device,
         )
-        # The gemm1 kernel writes fp8 payload to `a2_payload` (passed as
-        # `out` / arg_c) and preshuffled e8m0 scale to `a2_scale` (passed via
-        # quant_scale / arg_quant_scale).
+        # The gemm1 kernel writes packed MX payload to `a2_payload` (passed as
+        # `out` / arg_c) and preshuffled E8M0 scale to `a2_scale`.
         flydsl_grouped_gemm_a8w4_masked(
             a2_payload.view(torch.uint8),
             a1_payload,
@@ -569,7 +571,7 @@ def _grouped_a8w4_tdm_moe(
             bias=_b1,
             swiglu_limit=sl,
             num_buffers=num_buffers,
-            stage1_quant_out=1,
+            stage1_quant_out=_stage1_quant_out,
             quant_scale=a2_scale,
             quant_wmma_rep=wmma_rep2,
             **_situ_kw,
@@ -659,7 +661,7 @@ def _grouped_a8w4_tdm_moe(
                         bias=_b1,
                         swiglu_limit=sl,
                         num_buffers=num_buffers,
-                        stage1_quant_out=1,
+                        stage1_quant_out=_stage1_quant_out,
                         quant_scale=a2_scale,
                         quant_wmma_rep=wmma_rep2,
                         **_situ_kw,

aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2025 FlyDSL Project Contributors
 
-"""Grouped contiguous-M A8W4 preshuffle MoE GEMM for gfx1250 (TDM pipeline)."""
+"""Grouped contiguous-M A4W4/A8W4 preshuffle MoE GEMM for gfx1250."""
 
 import math
 from collections import namedtuple
@@ -26,6 +26,7 @@
 )
 from .quant_utils import (
     emit_amax_e8m0_native_scale,
+    emit_cvt_scalef32_pk8_fp4_bf16,
     emit_cvt_scalef32_pk8_fp8_f32,
 )
 from .tensor_shim import AITER_FLYDSL_MOE_EXPERT_SCHEDULING_MODE
@@ -551,7 +552,12 @@ def compute_ktile(buf, prefetch_kt):
 
             accs = [c_frags[idx].load().ir_value() for idx in range_constexpr(n_acc)]
             pipeline_fence(outstanding=0)
-            STORE_N = (tile_n // 2) if stage1_act else tile_n
+            quant_fp4 = stage1_quant_out == 2
+            STORE_N = (
+                ((tile_n // 4) if quant_fp4 else (tile_n // 2))
+                if stage1_act
+                else tile_n
+            )
             neg_limit = fx.Float32(0.0) - f32_swiglu_limit
             is_swiglu = stage1_act == 2
             is_situv2 = stage1_act == 3
@@ -566,7 +572,9 @@ def compute_ktile(buf, prefetch_kt):
 
             # -- Activate + stage to LDS --
             if const_expr(stage1_quant_out and stage1_act):
-                # Fused silu/swiglu -> fp8 quant; stage fp8 payload to LDS, scatter scale to global.
+                # Fused activation -> MX quant. FP8 stores one byte/element;
+                # FP4 packs two elements/byte. Both scatter gemm2's preshuffled
+                # E8M0 scale layout directly from the epilogue.
                 i32_ptr_g = fx.PointerType.get(
                     elem_ty=fx.Int8.ir_type,
                     address_space=fx.AddressSpace.Global,
@@ -616,37 +624,89 @@ def compute_ktile(buf, prefetch_kt):
                                 range_constexpr=range_constexpr,
                             )
 
+                        if const_expr(quant_fp4):
+                            # Match the previous bf16-intermediate + standalone
+                            # FP4 quant path: round the activated result to bf16
+                            # before computing its MX scale and packed payload.
+                            rounded_vec = (
+                                Vec.from_elements(all_vals, fx.Float32)
+                                .to(fx.BFloat16)
+                                .to(fx.Float32)
+                            )
+                            quant_vals = [
+                                rounded_vec[i] for i in range_constexpr(len(all_vals))
+                            ]
+                            quant_dtype = MxDtype.FP4_E2M1
+                        else:
+                            quant_vals = all_vals
+                            quant_dtype = MxDtype.FP8_E4M3
+
                         scale_f32, e8m0_byte = emit_amax_e8m0_native_scale(
-                            all_vals, wave_size=WAVE, dtype=MxDtype.FP8_E4M3
+                            quant_vals, wave_size=WAVE, dtype=quant_dtype
                         )
                         mx_blk_i = (
                             fx.Int32(blk_n + wnb + mx_blk * WN_PER_MX_BLOCK * 16) >> 6
                         )
                         e8m0_bytes.append(e8m0_byte)
                         mx_blk_is.append(mx_blk_i)
 
-                        for half in range_constexpr(WN_PER_MX_BLOCK // 2):
-                            src_f32 = Vec.from_elements(
-                                all_vals[half * 8 : half * 8 + 8],
-                                fx.Float32,
-                            ).ir_value()
-                            packed_v2i32 = emit_cvt_scalef32_pk8_fp8_f32(
-                                src_f32, scale_f32, v2i32_ty=v2i32_ty, rocdl=rocdl
-                            )
-                            for sub in range_constexpr(2):
-                                sub_wn = half * 2 + sub
-                                wn = mx_blk * WN_PER_MX_BLOCK + sub_wn
-                                packed_i32 = vector.extract(
-                                    packed_v2i32,
-                                    static_position=[sub],
-                                    dynamic_position=[],
-                                )
-                                col_fp8 = (wnb + wn * 16 + kgrp * 8) // 2
-                                lds_store_b32(
-                                    stC_idx,
-                                    row_rel * STORE_N + col_fp8,
-                                    Vec.from_elements([packed_i32], fx.Int32),
+                        if const_expr(quant_fp4):
+                            # kgrp0 owns activated columns [0:4] of each wn and
+                            # kgrp1 owns [4:8]. Exchange the peer values so
+                            # kgrp0 can issue one native pk8 conversion and one
+                            # aligned b32 LDS store for all eight FP4 values.
+                            peer_vals = [
+                                v.shuffle_xor(fx.Int32(16), fx.Int32(WAVE))
+                                for v in quant_vals
+                            ]
+                            if row_rel < mn_oob and is_kgrp0:
+                                for sub_wn in range_constexpr(WN_PER_MX_BLOCK):
+                                    begin = sub_wn * 4
+                                    src_bf16 = (
+                                        Vec.from_elements(
+                                            quant_vals[begin : begin + 4]
+                                            + peer_vals[begin : begin + 4],
+                                            fx.Float32,
+                                        )
+                                        .to(fx.BFloat16)
+                                        .ir_value()
+                                    )
+                                    packed_i32 = emit_cvt_scalef32_pk8_fp4_bf16(
+                                        src_bf16, scale_f32, i32_ty=T.i32
+                                    )
+                                    wn = mx_blk * WN_PER_MX_BLOCK + sub_wn
+                                    col_fp4 = (wnb + wn * 16) // 4
+                                    lds_store_b32(
+                                        stC_idx,
+                                        row_rel * STORE_N + col_fp4,
+                                        Vec.from_elements([packed_i32], fx.Int32),
+                                    )
+                        else:
+                            for half in range_constexpr(WN_PER_MX_BLOCK // 2):
+                                src_f32 = Vec.from_elements(
+                                    quant_vals[half * 8 : half * 8 + 8],
+                                    fx.Float32,
+                                ).ir_value()
+                                packed_v2i32 = emit_cvt_scalef32_pk8_fp8_f32(
+                                    src_f32,
+                                    scale_f32,
+                                    v2i32_ty=v2i32_ty,
+                                    rocdl=rocdl,
                                 )
+                                for sub in range_constexpr(2):
+                                    sub_wn = half * 2 + sub
+                                    wn = mx_blk * WN_PER_MX_BLOCK + sub_wn
+                                    packed_i32 = vector.extract(
+                                        packed_v2i32,
+                                        static_position=[sub],
+                                        dynamic_position=[],
+                                    )
+                                    col_fp8 = (wnb + wn * 16 + kgrp * 8) // 2
+                                    lds_store_b32(
+                                        stC_idx,
+                                        row_rel * STORE_N + col_fp8,
+                                        Vec.from_elements([packed_i32], fx.Int32),
+                                    )
 
                     # Preshuffled e8m0 scale: one branch per wm (not per mx_blk).
                     if row_rel < mn_oob and is_kgrp0:
@@ -718,8 +778,12 @@ def compute_ktile(buf, prefetch_kt):
             # -- Shared LDS -> TDM store to global --
             workgroup_barrier()
             if const_expr(stage1_act):
-                out_stride = i32_n // 2
-                out_col_off = blk_n64 // 2
+                if const_expr(quant_fp4):
+                    out_stride = i32_n // 4
+                    out_col_off = blk_n64 // 4
+                else:
+                    out_stride = i32_n // 2
+                    out_col_off = blk_n64 // 2
             else:
                 out_stride = c_stride
                 out_col_off = c_inner_off

aiter/ops/flydsl/kernels/quant_utils.py

@@ -297,3 +297,19 @@ def emit_cvt_scalef32_pk8_fp8_f32(src_v8f32, scale_f32, *, v2i32_ty, rocdl):
         _raw(src_v8f32),
         _raw(scale_f32),
     )
+
+
+def emit_cvt_scalef32_pk8_fp4_bf16(src_v8bf16, scale_f32, *, i32_ty):
+    """Native gfx1250 ``v_cvt_scalef32_pk8_fp4_bf16`` scaled pack.
+
+    Converts eight bf16 values to eight FP4 E2M1 nibbles packed into one i32.
+    ``scale_f32`` carries the forward E8M0 block scale; the instruction divides
+    by it before round-to-nearest-even conversion.
+    """
+    return llvm.inline_asm(
+        i32_ty,
+        [_raw(src_v8bf16), _raw(scale_f32)],
+        "v_cvt_scalef32_pk8_fp4_bf16 $0, $1, $2",
+        "=v,v,v",
+        has_side_effects=False,
+    )

op_tests/test_flydsl_grouped_gemm_gfx1250.py

@@ -685,9 +685,44 @@ def test_grouped_a4w4_situv2_matches_torch_ref():
     )
 
 
+@pytest.mark.parametrize(
+    "activation",
+    [ActivationType.Silu, ActivationType.Swiglu, ActivationType.Situv2],
+)
+def test_grouped_a4w4_fused_quant_activations_match_torch_ref(activation):
+    # A bias-free GEMM1 writes the packed MXFP4 payload and preshuffled E8M0
+    # scales for GEMM2 directly from its activation epilogue.
+    run_moe(
+        "a4w4",
+        activation=activation,
+        model_dim=512,
+        inter_dim=512,
+        use_bias=False,
+        check_aot_cache=False,
+    )
+
+
+@pytest.mark.parametrize("tile_m", [16, 128])
+def test_grouped_a4w4_situv2_fused_quant_scale_layouts_match_torch_ref(
+    monkeypatch, tile_m
+):
+    # Cover quant_wmma_rep=1 and 8 in addition to the default tile_m=64/rep=4.
+    # Keep both GEMMs on the same M tile so every expert boundary is legal.
+    monkeypatch.setenv("AITER_TDM_TILE_M", str(tile_m))
+    monkeypatch.setenv("AITER_TDM_TILE_M2", str(tile_m))
+    run_moe(
+        "a4w4",
+        activation=ActivationType.Situv2,
+        model_dim=512,
+        inter_dim=512,
+        use_bias=False,
+        check_aot_cache=False,
+    )
+
+
 def test_grouped_a8w4_situv2_matches_torch_ref():
-    # a8w4 takes the fused stage1 quant epilogue (batched activation), which is
-    # a separate code path from a4w4's bf16 intermediate (element-wise).
+    # A8W4 uses the FP8 form of the fused stage1 quant epilogue; A4W4 uses its
+    # packed FP4 form.
     run_moe(
         "a8w4",
         activation=ActivationType.Situv2,