From 3fb973e0265e843ac22070670cf9b364d1c25fb9 Mon Sep 17 00:00:00 2001 From: zhutaoyu Date: Tue, 23 Jun 2026 13:10:04 +0000 Subject: [PATCH 1/2] mxfp4 moe gemm2: fix a4w4 BM16 streaming K-loop accuracy (K_TILES==8) The FlyDSL gemm2 a4w4 port lost accuracy on the BM16 (M<=128) streaming K-loop for K_TILES_TOTAL==8 (D_INTER=2048, e.g. dsv3_a): the triple-buffered (3-slot) A->LDS pipeline did not reliably drain the async raw.ptr.buffer.load.lds of a streamed slot before that slot's ds_read 3 iterations later, producing non-deterministic garbage (e2e logits_diff ~0.02-0.05 vs the ~0.009 fp4 floor; cosine ~0.90). Other K_TILES (4/6/10/12) happened to schedule cleanly at 3 slots. - Widen the streaming A-slot buffer from 3 to 4 (quad buffer): adds slack between a slot's read and its next overwrite, hiding the load latency and making all K_TILES_TOTAL deterministic. LDS cost is free -- the s_Aq slots stay <= the lds_acc cshuffle union they overlap, so per-WG LDS is unchanged. - Gate the streaming issue_a_load_lds on wave < n_load_waves (matching the prologue): at BM16 only waves 0,1 hold A rows; ungated, waves 2,3 streamed into lds_row 16/24 (past the BM16 slot) from the next m-block's rows. Verified: dsv3_a / dsv4_ep8 / kimi / dsv3_b / minimax_a / dsv4_tp2 all at the ~0.009 fp4 floor and deterministic across runs; large-M (BM32/BM128) unchanged; 33 standalone gemm1/gemm2 port tests pass. Co-Authored-By: Claude Opus 4 (1M context) --- aiter/ops/flydsl/kernels/mxfp4_gemm2.py | 51 +++++++++++++++++-------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py index 1a174b9ead0..2c8b7f27400 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py @@ -336,11 +336,20 @@ def compile_gemm2_a4w4_port( _slot_bytes = saq_slot_bytes(BM) # Number of LDS A-slots (s_Aq stages). The K_TILES_TOTAL==2 fast path preloads # all 2 tiles into 2 slots (double buffer). The streaming K_TILES_TOTAL>2 path - # uses 3 slots (triple buffer) so the per-iteration read-slot (tile kt) and - # write-slot (tile kt+kStages, streamed in) never alias -- the read tile maps to - # kt%3 and the streamed tile to (kt+2)%3, which differ for all kt. (gemm1 uses - # the same kAStages=3 triple-buffer for its streaming non-128 path.) - _aStages = kStages if _K_TILES_TOTAL <= kStages else 3 + # uses a quad buffer (4 slots): read-slot (tile kt) and write-slot (tile + # kt+kStages, streamed in) never alias, with two slots of slack between a + # slot's read and its next overwrite. + # + # A triple buffer (3 slots) is enough to avoid aliasing on paper, and works for + # most K_TILES_TOTAL, but K_TILES_TOTAL==8 (e.g. dsv3_a INTER=2048) miscompiled + # at BM16: the in-flight async raw.ptr.buffer.load.lds for a streamed slot was + # not reliably drained before that slot's ds_read 3 iterations later, producing + # non-deterministic garbage (e2e logits_diff ~0.02-0.05 vs the ~0.009 fp4 + # floor; cosine ~0.90). The extra slot widens the read->overwrite window enough + # to hide the load latency and makes all K_TILES_TOTAL deterministic. LDS cost + # is free: the s_Aq slots (_aStages * BM*KH_TILE) stay <= the lds_acc cshuffle + # union (BM*BN*4) it overlaps, so the per-WG LDS footprint is unchanged. + _aStages = kStages if _K_TILES_TOTAL <= kStages else 4 # atomic / mxfp4 epilog reuses LDS for the cshuffle (BM*BN f32); nonatomic # bf16 writes direct, so only s_Aq (_aStages slots) is needed. nonatomic_cshuffle # cshuffles in 64-row passes -> only min(BM,64)*BN f32 (BM128 -> 64KB not 128KB). @@ -775,12 +784,25 @@ def load_b_tile(kt): # only; the K_TILES_TOTAL==2 prologue is preloaded by the kernel). Mirrors # gemm1.issue_a_load_lds: m_row-derived cached rows + lds_swizzle. ------- def issue_a_load_lds(slot, kt): - for sub in range_constexpr(_kSubBlocks): - lds_row = wave * fx.Int32(_rows_per_wave) + fx.Int32(sub * 8) - car = m_row + lds_row + (lane // fx.Int32(8)) - _issue_a_load_lds( - aq_rsrc, saq, slot, kt, car, lane, _slot_bytes, lds_row, k_half=_K_HALF - ) + # Only the first n_load_waves hold A rows (BM16: waves 0,1; rows_per_wave=8 + # -> 2 waves cover the 16-row block). Ungated, waves 2,3 would stream into + # lds_row 16/24 (past the BM16 slot) from car = m_row+16/24 (the next + # m-block's rows / wrong expert), corrupting the tile -- the prologue + # _issue_all_a_loads gates the same way (see wave < _n_load_waves below). + def _do(): + for sub in range_constexpr(_kSubBlocks): + lds_row = wave * fx.Int32(_rows_per_wave) + fx.Int32(sub * 8) + car = m_row + lds_row + (lane // fx.Int32(8)) + _issue_a_load_lds( + aq_rsrc, saq, slot, kt, car, lane, _slot_bytes, lds_row, + k_half=_K_HALF, + ) + + if const_expr(_n_load_waves < 4): + if wave < fx.Int32(_n_load_waves): + _do() + else: + _do() # -- ds_read(slot) -> a[i][k] (i32x4) ; i in [0,kMChunks) ----------------- def issue_a_ds_read(slot): @@ -920,10 +942,9 @@ def _kloop_fence(vmcnt_atomic): # main loop: OFFSET in [0, kUnroll). Process tile kt=OFFSET (read from LDS # slot kt%_aStages), and stream the next tile next_kt=kStages+OFFSET into - # slot next_kt%_aStages. With _aStages=3 (triple buffer) read-slot and - # write-slot never alias, so the streamed load cannot clobber the tile this - # iteration is reading. A loop-top barrier (inside _kloop_fence) guards the - # cross-iteration reuse of each slot. + # slot next_kt%_aStages. The quad buffer (_aStages=4) keeps read-slot and + # write-slot distinct with slack (see the _aStages note above). A loop-top + # barrier (inside _kloop_fence) guards the cross-iteration reuse of a slot. for OFFSET in range_constexpr(_kUnroll): kt = OFFSET slot = kt % _aStages From 7e06ff88c723bcaad7b6f9f2e5ab37bb0af0772b Mon Sep 17 00:00:00 2001 From: zhutaoyu Date: Tue, 23 Jun 2026 13:10:16 +0000 Subject: [PATCH 2/2] mxfp4 moe tuner: fix a4w4 run_config layout + gate_mode gemm_moe_tune.py's run_config benchmarked a4w4 (fp4 act + fp4 weight) with the wrong weight layout and gate_mode, so its accuracy check reported garbage (err_ratio ~1.0, 99.7% elements off) for every a4w4 shape -- the same class of bug already fixed in op_tests/test_moe_2stage.py (commit 6a580ab6) but missed in the tuner. - include fp4x2 activations in the a16w4 shuffle branch so a4w4 gets shuffle_weight_a16w4 / shuffle_scale_a16w4 like a16w4/a8w4 (the FlyDSL port the a4w4 MoE runs on expects the gate-up-interleaved a16w4 layout). - pass gate_mode=interleave to fused_moe for mxfp4 weights (a4w4/a8w4/a16w4) and mxfp8, mirroring op_tests _effective_gate_mode (the a4w4 tuned rows are interleave-keyed; the default SEPARATED reaches the wrong dispatch). - add the top-level GateMode import the new code needs. Co-Authored-By: Claude Opus 4 (1M context) --- .../gemm_moe_tune.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index fae1e77fb4e..4db4de30259 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -44,6 +44,7 @@ convert_int8_to_uint32_int4, ) from aiter.ops.quant import per_1x32_i4_quant, per_1x32_f8_scale_f8_quant +from aiter.ops.flydsl.moe_common import GateMode from aiter import dtypes from aiter import ActivationType as ActivationType from aiter.jit.utils.chip_info import get_gfx, get_gfx_runtime, gfx_from_cu_num @@ -3532,9 +3533,9 @@ def run_config(self, args): ) elif ( q_type == QuantType.per_1x32 - and q_dtype_a in [dtypes.bf16, dtypes.fp16, dtypes.fp8] + and q_dtype_a in [dtypes.bf16, dtypes.fp16, dtypes.fp8, dtypes.fp4x2] and q_dtype_w == dtypes.fp4x2 - ): + ): # a16w4 / a8w4 / a4w4 — all use the gate-up-interleaved a16w4 layout w1_qt_fmoe = shuffle_weight_a16w4(w1_qt_fmoe, 16, True) w1_scale_fmoe = shuffle_scale_a16w4(w1_scale, expert, True) w2_qt_fmoe = shuffle_weight_a16w4(w2_qt_fmoe, 16, False) @@ -3589,6 +3590,17 @@ def run_config(self, args): torch_quant = aiter.get_torch_quant(q_type) a1_qt, a1_scale = torch_quant(hidden, quant_dtype=q_dtype_a) + # mxfp4 weights (a4w4/a8w4/a16w4) and mxfp8 (a8w8) run the + # gate-up-interleaved layout the tuned rows are keyed on + # (matches serving ATOM_MOE_GU_ITLV=1); everything else stays + # separated. Mirror op_tests/test_moe_2stage._effective_gate_mode. + if ( + q_dtype_a in [dtypes.fp8, dtypes.bf16, dtypes.fp4x2] + and q_dtype_w == dtypes.fp4x2 + ) or (q_dtype_a == dtypes.fp8 and q_dtype_w == dtypes.fp8): + gate_mode = GateMode.INTERLEAVE.value + else: + gate_mode = GateMode.SEPARATED.value out, us = run_perftest( fused_moe, hidden, @@ -3602,6 +3614,7 @@ def run_config(self, args): w1_scale=w1_scale_fmoe, w2_scale=w2_scale_fmoe, dtype=dtype, + gate_mode=gate_mode, num_warmup=args.warmup, num_iters=args.iters, )