From ff9609fe4b57725960c6455ddd354049c0893e0e Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 6 Aug 2026 11:08:17 +0000 Subject: [PATCH 1/5] [gfx1250] Fuse A4W4 stage1 quantization Write packed MXFP4 payloads and preshuffled E8M0 scales directly from the grouped GEMM1 epilogue, removing the standalone quant launch while preserving BF16 rounding semantics. Co-authored-by: Cursor --- aiter/ops/flydsl/batched_gemm_mxfp4.py | 26 ++-- aiter/ops/flydsl/grouped_moe_gfx1250.py | 24 ++-- .../kernels/mxfp4_preshuffle_gfx1250_tdm.py | 118 ++++++++++++++---- aiter/ops/flydsl/kernels/quant_utils.py | 16 +++ op_tests/test_flydsl_grouped_gemm_gfx1250.py | 39 +++++- 5 files changed, 175 insertions(+), 48 deletions(-) diff --git a/aiter/ops/flydsl/batched_gemm_mxfp4.py b/aiter/ops/flydsl/batched_gemm_mxfp4.py index 233d6a711b..221549a588 100644 --- a/aiter/ops/flydsl/batched_gemm_mxfp4.py +++ b/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) diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index eda290d693..028692c86c 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/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, diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index 9cdbd5184e..a4c45dfbe7 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py +++ b/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,8 +624,25 @@ 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 @@ -625,28 +650,63 @@ def compute_ktile(buf, prefetch_kt): 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 diff --git a/aiter/ops/flydsl/kernels/quant_utils.py b/aiter/ops/flydsl/kernels/quant_utils.py index 4e70f28170..a202d7f995 100644 --- a/aiter/ops/flydsl/kernels/quant_utils.py +++ b/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, + ) diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/test_flydsl_grouped_gemm_gfx1250.py index dc9ef7bb7b..efb3f879f4 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/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, From 8855cdcf1c35adacbcfd0d97dbc99a4fc790ec60 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 6 Aug 2026 11:36:04 +0000 Subject: [PATCH 2/5] [gfx1250] Reduce fused FP4 epilogue conversions Compute MX scale from the FP32 activation and exchange packed BF16 words so the native FP4 pack avoids a redundant round trip and half of the cross-lane shuffles. Co-authored-by: Cursor --- .../kernels/mxfp4_preshuffle_gfx1250_tdm.py | 54 +++++++++++-------- 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index a4c45dfbe7..91ba5ddb1e 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py +++ b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py @@ -625,17 +625,14 @@ def compute_ktile(buf, prefetch_kt): ) 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) + # Compute the MX scale directly from the f32 activation. + # gfx1250's native FP4 pack consumes bf16, so keep one + # rounded view for that instruction without widening it + # back to f32. + pack_bf16 = Vec.from_elements(all_vals, fx.Float32).to( + fx.BFloat16 ) - quant_vals = [ - rounded_vec[i] for i in range_constexpr(len(all_vals)) - ] + quant_vals = all_vals quant_dtype = MxDtype.FP4_E2M1 else: quant_vals = all_vals @@ -652,23 +649,36 @@ def compute_ktile(buf, prefetch_kt): 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 - ] + # kgrp1 owns [4:8]. Pack each lane's four bf16 values + # into two dwords before exchanging them, halving the + # cross-kgrp shuffles. + packed_bf16 = pack_bf16.bitcast(fx.Int32) + local_bf16_words = [] + peer_bf16_words = [] + for sub_wn in range_constexpr(WN_PER_MX_BLOCK): + begin = sub_wn * 2 + local_words = [ + packed_bf16[begin + i] + for i in range_constexpr(2) + ] + local_bf16_words.append(local_words) + peer_bf16_words.append( + [ + word.shuffle_xor( + fx.Int32(16), fx.Int32(WAVE) + ) + for word in local_words + ] + ) 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, + local_bf16_words[sub_wn] + + peer_bf16_words[sub_wn], + fx.Int32, ) - .to(fx.BFloat16) + .bitcast(fx.BFloat16) .ir_value() ) packed_i32 = emit_cvt_scalef32_pk8_fp4_bf16( From 669545b95fcc79b2bb97dac2fb1b7f959ba0e65a Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 6 Aug 2026 11:48:26 +0000 Subject: [PATCH 3/5] [gfx1250] Balance fused MX amax reduction Shorten the local amax dependency chain from 16 max operations to four levels before the cross-wave exchange. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/quant_utils.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/aiter/ops/flydsl/kernels/quant_utils.py b/aiter/ops/flydsl/kernels/quant_utils.py index a202d7f995..c1844c5eba 100644 --- a/aiter/ops/flydsl/kernels/quant_utils.py +++ b/aiter/ops/flydsl/kernels/quant_utils.py @@ -265,10 +265,26 @@ def emit_amax_e8m0_native_scale(all_vals, *, wave_size, dtype=_D.FP8_E4M3): c23 = arith.constant(23, type=T.i32) c_wave = arith.constant(wave_size, type=T.i32) - block_amax = arith.constant(0.0, type=T.f32) + abs_vals = [] for v in all_vals: - abs_v = llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], []) - block_amax = arith.maxnumf(block_amax, abs_v) + abs_vals.append( + llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], []) + ) + + # Build a balanced max tree instead of a serial dependency chain. The + # epilogue contributes 16 values per lane, so this reduces local amax depth + # from 16 maxnum operations to 4 before the cross-kgrp shuffle. + while len(abs_vals) > 1: + next_vals = [] + for i in range(0, len(abs_vals), 2): + if i + 1 < len(abs_vals): + next_vals.append(arith.maxnumf(abs_vals[i], abs_vals[i + 1])) + else: + next_vals.append(abs_vals[i]) + abs_vals = next_vals + + c0_f32 = arith.constant(0.0, type=T.f32) + block_amax = arith.maxnumf(c0_f32, abs_vals[0]) if abs_vals else c0_f32 block_amax = arith.minnumf(block_amax, c_flt_max) peer = block_amax.shuffle_xor(c16, c_wave) block_amax = arith.maxnumf(block_amax, peer) From 82673d8ba9097ce0d88f7be37d3bb4a9a7000325 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 6 Aug 2026 11:51:22 +0000 Subject: [PATCH 4/5] Apply Black formatting to FP4 epilogue changes Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py | 7 ++----- aiter/ops/flydsl/kernels/quant_utils.py | 4 +--- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py index 91ba5ddb1e..37522f7c73 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py +++ b/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py @@ -658,15 +658,12 @@ def compute_ktile(buf, prefetch_kt): for sub_wn in range_constexpr(WN_PER_MX_BLOCK): begin = sub_wn * 2 local_words = [ - packed_bf16[begin + i] - for i in range_constexpr(2) + packed_bf16[begin + i] for i in range_constexpr(2) ] local_bf16_words.append(local_words) peer_bf16_words.append( [ - word.shuffle_xor( - fx.Int32(16), fx.Int32(WAVE) - ) + word.shuffle_xor(fx.Int32(16), fx.Int32(WAVE)) for word in local_words ] ) diff --git a/aiter/ops/flydsl/kernels/quant_utils.py b/aiter/ops/flydsl/kernels/quant_utils.py index c1844c5eba..430c77e18b 100644 --- a/aiter/ops/flydsl/kernels/quant_utils.py +++ b/aiter/ops/flydsl/kernels/quant_utils.py @@ -267,9 +267,7 @@ def emit_amax_e8m0_native_scale(all_vals, *, wave_size, dtype=_D.FP8_E4M3): abs_vals = [] for v in all_vals: - abs_vals.append( - llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], []) - ) + abs_vals.append(llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], [])) # Build a balanced max tree instead of a serial dependency chain. The # epilogue contributes 16 values per lane, so this reduces local amax depth From 190d8d8ac607d15bd09dbf5665305b06092f27b6 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Thu, 6 Aug 2026 11:58:39 +0000 Subject: [PATCH 5/5] [gfx1250] Drop inconclusive balanced amax tree Keep the original serial reduction because paired end-to-end measurements showed no reliable benefit from the balanced form. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/quant_utils.py | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/aiter/ops/flydsl/kernels/quant_utils.py b/aiter/ops/flydsl/kernels/quant_utils.py index 430c77e18b..a202d7f995 100644 --- a/aiter/ops/flydsl/kernels/quant_utils.py +++ b/aiter/ops/flydsl/kernels/quant_utils.py @@ -265,24 +265,10 @@ def emit_amax_e8m0_native_scale(all_vals, *, wave_size, dtype=_D.FP8_E4M3): c23 = arith.constant(23, type=T.i32) c_wave = arith.constant(wave_size, type=T.i32) - abs_vals = [] + block_amax = arith.constant(0.0, type=T.f32) for v in all_vals: - abs_vals.append(llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], [])) - - # Build a balanced max tree instead of a serial dependency chain. The - # epilogue contributes 16 values per lane, so this reduces local amax depth - # from 16 maxnum operations to 4 before the cross-kgrp shuffle. - while len(abs_vals) > 1: - next_vals = [] - for i in range(0, len(abs_vals), 2): - if i + 1 < len(abs_vals): - next_vals.append(arith.maxnumf(abs_vals[i], abs_vals[i + 1])) - else: - next_vals.append(abs_vals[i]) - abs_vals = next_vals - - c0_f32 = arith.constant(0.0, type=T.f32) - block_amax = arith.maxnumf(c0_f32, abs_vals[0]) if abs_vals else c0_f32 + abs_v = llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(v)], [], []) + block_amax = arith.maxnumf(block_amax, abs_v) block_amax = arith.minnumf(block_amax, c_flt_max) peer = block_amax.shuffle_xor(c16, c_wave) block_amax = arith.maxnumf(block_amax, peer)