Skip to content

Fix/mxfp4 a4w4 bm16 streaming accuracy #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
benenzhu wants to merge 2 commits into
base: dev/randomflow_pr
Choose a base branch
from
Open

Fix/mxfp4 a4w4 bm16 streaming accuracy #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3fb973e
mxfp4 moe gemm2: fix a4w4 BM16 streaming K-loop accuracy (K_TILES==8)
benenzhu Jun 23, 2026
7e06ff8
mxfp4 moe tuner: fix a4w4 run_config layout + gate_mode
benenzhu Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 36 additions & 15 deletions aiter/ops/flydsl/kernels/mxfp4_gemm2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down
17 changes: 15 additions & 2 deletions csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)
Expand Down