diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 41771a8389..c036c62ba9 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -2617,15 +2617,28 @@ def get_block_m() -> int: _tile_m, _s1_sfx, _s2_sfx = 32, "_w2", "_bnt2" elif token < 4096: _tile_m, _s1_sfx, _s2_sfx = 64, "_w3_bnt0", "" - elif token < 16384: + elif token < 8192: _tile_m, _s1_sfx, _s2_sfx = 128, "_w2_bnt0", "" else: _tile_m, _s1_sfx, _s2_sfx = 64, "_w4_bnt0", "" + # Stage2 tiling is picked independently of the sorting block. + if token >= 4096: + _s2_tm, _s2_tn, _s2_mode = 64, 256, "reduce" + else: + _s2_tm, _s2_tn, _s2_mode = _tile_m, 128, "atomic" _base_kn1 = flydsl_kernel_name( 1, _a_type, _w_type, _out_type, _tile_m, 128, 256 ) _base_kn2 = flydsl_kernel_name( - 2, _a_type, _w_type, _out_type, _tile_m, 128, _s2_tk, "atomic" + 2, + _a_type, + _w_type, + _out_type, + _s2_tm, + _s2_tn, + _s2_tk, + _s2_mode, + sort_block_m=_tile_m, ) kn1 = f"{_base_kn1}{_s1_sfx}" kn2 = f"{_base_kn2}{_s2_sfx}" diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py index 40e7da86e4..f066e656ac 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py @@ -3559,7 +3559,17 @@ def scale_elem_type(): ) bytes_per_thread_x = bytes_x_per_tile // total_threads - lds_stride = tile_k + # swizzle_xor16 only stays within a row when the LDS X row stride is a power + # of two, so pad it up and leave the tail of the row unused. + _lds_stride_raw = int(tile_k) + lds_stride = 1 << (_lds_stride_raw - 1).bit_length() + lds_stride_bytes = int(lds_stride) * int(a_elem_bytes) + if const_expr(lds_stride != _lds_stride_raw and bool(use_async_copy)): + # buffer_load_lds fills the destination linearly, so it cannot leave a + # gap at the end of each padded LDS row. + raise ValueError( + f"use_async_copy requires a power-of-2 tile_k, got tile_k={tile_k}" + ) if const_expr(out_is_f32): _use_cshuffle_epilog = bool(use_cshuffle_epilog) @@ -3612,7 +3622,10 @@ def load_bias_scalar(bias_rsrc, offset): f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" f"_t{tile_m}x{tile_n}x{tile_k}{variant_tags}" ).replace("-", "_") - lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(a_elem_bytes) + # A single K tile per batch never touches the ping buffer, so skip + # double-buffering X and save the LDS. + _x_buf_count = 1 if num_k_tiles_per_batch == 1 else 2 + lds_x_bytes = _x_buf_count * int(tile_m) * int(lds_stride) * int(a_elem_bytes) lds_out_bytes = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 lds_tid_bytes = int(tile_m) * 4 lds_tw_bytes = (int(tile_m) * 4) if bool(doweight_stage2) else 0 @@ -3695,10 +3708,10 @@ def check_c_k_valid_gate(base_k): if const_expr(use_async_copy and a_elem_vec_pack > 1): eff_lds_stride = lds_stride // a_elem_vec_pack - eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + eff_tile_k_bytes = lds_stride_bytes // a_elem_vec_pack else: eff_lds_stride = lds_stride - eff_tile_k_bytes = tile_k_bytes + eff_tile_k_bytes = lds_stride_bytes shape_lds = fx.make_shape(tile_m, eff_lds_stride) stride_lds = fx.make_stride(eff_lds_stride, 1) @@ -3765,7 +3778,7 @@ def check_c_k_valid_gate(base_k): else None ) - lds_x_b = 2 * int(tile_m) * int(lds_stride) * int(a_elem_bytes) + lds_x_b = _x_buf_count * int(tile_m) * int(lds_stride) * int(a_elem_bytes) lds_out_b = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 lds_tid_off = max(lds_x_b, lds_out_b) lds_tid = SmemPtr( @@ -4191,7 +4204,10 @@ def load_x_tile(base_k): m_repeat = tile_m // 16 k_unroll = tile_k_bytes // 128 - k_unroll_packed = k_unroll // pack_K + # Round up: an odd k_unroll (e.g. tile_k=384 -> 3) still needs the + # packed-scale i32 holding its last 128-K block, and the host-side + # e8m0 shuffle pads K so the extra index stays in bounds. + k_unroll_packed = (k_unroll + pack_K - 1) // pack_K m_repeat_packed = m_repeat // pack_M num_acc_n_packed = body_num_acc_n // pack_N @@ -5344,6 +5360,15 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): ) e_vec = 2 if accumulate else min(body_tile_n // 32, 8) + if const_expr(need_fp8_out and e_vec != 8): + # With e_vec < 8 two fragments share one e8m0 scale byte and + # race on it, so half the row dequantizes with the wrong + # exponent. + raise ValueError( + "fp8 stage2 route-out (AITER_FLYDSL_STAGE2_FP8=1) needs " + f"tile_n >= 256 to form whole 8-column MXFP8 groups, got " + f"tile_n={body_tile_n} (e_vec={e_vec})" + ) rocdl.s_setprio(3) c_shuffle_epilog( arith=arith, diff --git a/aiter/ops/flydsl/kernels/mxmoe_dispatcher.py b/aiter/ops/flydsl/kernels/mxmoe_dispatcher.py index 14cbf1f8ca..32ac43ed85 100644 --- a/aiter/ops/flydsl/kernels/mxmoe_dispatcher.py +++ b/aiter/ops/flydsl/kernels/mxmoe_dispatcher.py @@ -2,10 +2,12 @@ # Copyright (C) 2025-2026 FlyDSL Project Contributors """Compile + launch dispatch for the layout-API MXFP4 MoE gemm (BM32, opus-sort); a4w4/a8w4 entry point.""" +import contextlib import os import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import Int8, T @@ -78,6 +80,19 @@ def _spart_output_tile_index(block_1d_id, M0, N0, group_num, m01): return m_block_idx, n_block_idx +# Tuned scheduling/staging defaults for the gfx950 a8w4 stage2 gemm; all are +# bit-exact against the stock kernel. +_G2_KUNROLL_DEFAULT = 1 +# bit2 = 4-way C-slab split, bit3 = wide route-out store, bit5 = LDS bank swizzle. +_G2_EPI_DEFAULT = 44 +# GroupNum=32, M01=1 block->tile swizzle, for L2 reuse of the B weights. +_G2_SPART_DEFAULT = 3201 +_G2_WPE_DEFAULT = 0 # 0 = let the compiler pick +_G2_WCPL_DEFAULT = 0 +# 0/1 = off. G>1 = each block owns G consecutive n-tiles of one m-tile. +_G2_NLOOP_DEFAULT = 2 + + def compile_gemm2_a4w4_port( BM=32, BN=256, @@ -97,6 +112,12 @@ def compile_gemm2_a4w4_port( g2_ascale_pf=None, g2_spart=None, g2_bf16_lds=None, + g2_diag=None, + g2_kunroll=None, + g2_epi=None, + g2_wpe=None, + g2_wcpl=None, + g2_nloop=None, out_dtype="bf16", ): """Compile gemm2 a4w4 down-proj; epilog 'atomic' (weighted atomic-fadd) or 'reduce' (store into out[token_id*topk+slot]). inter_dim runtime; SBM None -> SBM==BM byte-identical.""" @@ -131,7 +152,7 @@ def compile_gemm2_a4w4_port( g2_ascale_pf = os.environ.get("MXFP4_G2_ASCALE_PF", "1") == "1" g2_ascale_pf = bool(g2_ascale_pf) if g2_spart is None: - g2_spart = int(os.environ.get("MXFP4_G2_SPART", "402")) + g2_spart = int(os.environ.get("MXFP4_G2_SPART", str(_G2_SPART_DEFAULT))) g2_spart = int(g2_spart) g2_group_num = g2_spart // 100 if g2_spart > 0 else 0 g2_m01 = g2_spart % 100 if g2_spart > 0 else 0 @@ -146,11 +167,67 @@ def compile_gemm2_a4w4_port( if g2_bf16_lds is None: g2_bf16_lds = os.environ.get("MXFP4_G2_BF16_LDS", "0") == "1" g2_bf16_lds = bool(g2_bf16_lds) + # MXFP4_G2_DIAG: PERF-ATTRIBUTION ONLY, produces WRONG results. Each bit drops + # one phase of the kernel to size how much of the runtime it owns. + if g2_diag is None: + g2_diag = int(os.environ.get("MXFP4_G2_DIAG", "0")) + g2_diag = int(g2_diag) + # g2_kunroll: fully unroll the K loop, one LDS slot per A K-tile. Legal only + # when the runtime inter_dim equals INTER_MAX (the host entry asserts this). + if g2_kunroll is None: + g2_kunroll = int(os.environ.get("MXFP4_G2_KUNROLL", _G2_KUNROLL_DEFAULT)) + # 0 = rolled K loop, 1 = unroll with per-tile registers, 2 = unroll with shared ones. + g2_kunroll = int(g2_kunroll) + # g2_epi bitmask, all bit-exact with g2_epi=0: + # 1 = 32-byte XOR-swizzled C slab + vectorised cshuffle readback + fused + # 8-byte fp8 value store. + # 2 = stage the cshuffle in 2 row slices, 4 = in 4 slices, which keeps the C + # slab off the LDS high-water mark so more workgroups fit per CU. + if g2_epi is None: + g2_epi = int(os.environ.get("MXFP4_G2_EPI", str(_G2_EPI_DEFAULT))) + g2_epi = int(g2_epi) + # g2_wpe: --amdgpu-waves-per-eu. Occupancy here is capped by the register file + # (64 f32 accumulators per wave), not LDS, and every non-zero value measures + # worse than the compiler's own heuristic. + if g2_wpe is None: + g2_wpe = int(os.environ.get("MXFP4_G2_WPE", str(_G2_WPE_DEFAULT))) + g2_wpe = int(g2_wpe) + if g2_wcpl is None: + g2_wcpl = int(os.environ.get("MXFP4_G2_WCPL", str(_G2_WCPL_DEFAULT))) + if g2_nloop is None: + g2_nloop = int(os.environ.get("MXFP4_G2_NLOOP", str(_G2_NLOOP_DEFAULT))) + # The group must tile the n range exactly, or the tail n-tiles are never computed. + g2_nloop = int(g2_nloop) + if g2_nloop > 1 and (HIDDEN_MAX // BN) % g2_nloop: + g2_nloop = 1 + g2_wcpl = int(g2_wcpl) KH_TILE_A = BK // (1 if is_f8 else 2) # A LDS K-tile bytes (fp8 256, fp4 128) slot_bytes = BM * KH_TILE_A - aStages = 2 if g2_bf16_lds else 3 - c_lds_bytes = BM * BN * (2 if g2_bf16_lds else 4) - lds_bytes = max(c_lds_bytes, aStages * slot_bytes) + # A-slot count must exceed kStages (=2): with aStages==2 the DMA for tile + # kt+kStages targets the slot tile kt is still being ds-read from, and since a + # wave DMAs only its own BM/4 rows but reads all BM, that is a cross-wave WAR + # race on the A operands. aStages=3 is free -- the C slab dominates the union. + aStages = 3 + a_prologue_tiles = kStages + if g2_kunroll: + # One slot per K-tile: the prologue DMAs the whole contraction up front. + a_prologue_tiles = aStages = INTER_MAX // BK + # g2_epi>=2 stages the cshuffle in 2 (or 4) row slices, so only part of the + # C slab is live at a time. + c_split = 4 if (g2_epi & 4) else (2 if (g2_epi & 2) else 1) + while c_split > 1 and ((BM // 16) % c_split or (BM // 8) % c_split): + c_split //= 2 + c_lds_bytes = BM * BN * (2 if g2_bf16_lds else 4) // c_split + # N-loop keeps A live across every n-tile, so the C slab cannot union with it. + c_lds_off = (aStages * slot_bytes) if g2_nloop > 1 else 0 + lds_bytes = ( + c_lds_off + c_lds_bytes + if g2_nloop > 1 + else max(c_lds_bytes, aStages * slot_bytes) + ) + # MXFP4_G2_LDSPAD: DIAGNOSTIC. Inflate the LDS request without using it, to find + # the workgroups/CU cliff. Correct results; perf only. + lds_bytes = max(lds_bytes, int(os.environ.get("MXFP4_G2_LDSPAD", "0"))) # N_OUT = model_dim/hidden is runtime; HIDDEN_MAX is a compile/cache bucket # so different runtime hidden sizes can reuse one compiled launcher. assert ( @@ -178,9 +255,14 @@ def compile_gemm2_a4w4_port( apf_tag = "_apf" if g2_ascale_pf else "" spart_tag = f"_spart{g2_group_num}x{g2_m01}" if g2_spart > 0 else "" bf16lds_tag = "_bf16lds" if g2_bf16_lds else "" + diag_tag = f"_diag{g2_diag}" if g2_diag else "" + kunroll_tag = f"_kunroll{g2_kunroll}" if g2_kunroll else "" + epi_tag = f"_epi{g2_epi}" if g2_epi else "" + epi_tag += f"_wcpl{g2_wcpl}" if (g2_epi & 8) and g2_wcpl else "" + nloop_tag = f"_nloop{g2_nloop}" if g2_nloop > 1 else "" out_tag = "_fp8out" if route_out_fp8 else "" tile_tag = "" if (BN, BK) == (256, 256) else f"_bn{BN}_bk{BK}" - tag = f"hmax{HIDDEN_MAX}_imax{INTER_MAX}_bm{BM}{tile_tag}{'_nt' if use_nt else ''}_{etag}{atag}{sbm_tag}{persist_tag}{pad_tag}{ks_tag}{bh_tag}{apf_tag}{spart_tag}{bf16lds_tag}{out_tag}_v2" + tag = f"hmax{HIDDEN_MAX}_imax{INTER_MAX}_bm{BM}{tile_tag}{'_nt' if use_nt else ''}_{etag}{atag}{sbm_tag}{persist_tag}{pad_tag}{ks_tag}{bh_tag}{apf_tag}{spart_tag}{bf16lds_tag}{diag_tag}{kunroll_tag}{epi_tag}{nloop_tag}{out_tag}_v2" name = f"gemm2_a4w4_port_{tag}" @fx.struct @@ -221,7 +303,9 @@ def _gemm2_kernel_body( # Preload the first kStages K-tiles (the streaming prologue). def issue_all_a_loads(m_row0): - for slot in range_constexpr(kStages): + if const_expr(bool(g2_diag & 32)): # attribution: drop the A->LDS DMA + return + for slot in range_constexpr(a_prologue_tiles): issue_a_load_lds_dt( arg_aq, aq_num, @@ -273,10 +357,37 @@ def run_unit(unit_bx): g2_bhoist=g2_bhoist, g2_ascale_pf=g2_ascale_pf, g2_bf16_lds=g2_bf16_lds, + g2_diag=g2_diag, + g2_kunroll=g2_kunroll, + g2_epi=g2_epi, + g2_wcpl=g2_wcpl, route_out_fp8=route_out_fp8, + c_lds_off=c_lds_off, ) - if const_expr(not persist and g2_spart <= 0): + if const_expr(not persist and g2_nloop > 1): + # N-group schedule: a block owns g2_nloop CONSECUTIVE n-tiles of one + # m-tile, so the A->LDS prologue and the sorted_token_ids/sorted_weights + # row metadata are paid once per group instead of once per n-tile. The + # group is unrolled at compile time, and the grid stays large enough to + # keep independent workgroups per CU hiding each other's epilogues. + num_n_groups = num_n_blocks // fx.Int32(g2_nloop) + cumsum0 = global_typed_ptr(arg_cumsum, T.i32)[0] + total_m_blocks = cumsum0 // BM + bound = total_m_blocks * num_n_groups + + if fx.Int32(bx_i32) < bound: + m_block_idx, n_group_idx = _spart_output_tile_index( + bx_i32, total_m_blocks, num_n_groups, g2_group_num, g2_m01 + ) + issue_all_a_loads(m_block_idx * fx.Int32(BM)) + rocdl.sched_barrier(0) + n_block0 = n_group_idx * fx.Int32(g2_nloop) + for _j in range_constexpr(g2_nloop): + run_unit( + m_block_idx * fx.Int32(num_n_blocks) + n_block0 + fx.Int32(_j) + ) + elif const_expr(not persist and g2_spart <= 0): # One-shot naive linear block->(m,n): issue A->LDS before the cumsum load (latency overlap). issue_all_a_loads((bx_i32 // num_n_blocks) * fx.Int32(BM)) rocdl.sched_barrier(0) @@ -395,7 +506,12 @@ def launch_gemm2( ): # i32_max_m_blocks sizes buffer resources; i32_grid_blocks bounds the launch to real m-blocks. num_n_blocks = fx.Int32(i32_hidden) // fx.Int32(BN) - grid_x = i32_grid_blocks * num_n_blocks + # N-loop folds the n dimension into the kernel, so the grid is m-tiles only. + grid_x = ( + i32_grid_blocks * (num_n_blocks // fx.Int32(g2_nloop)) + if (g2_nloop > 1 and not persist) + else i32_grid_blocks * num_n_blocks + ) gemm2_kernel( arg_aq, arg_ascale, @@ -448,8 +564,19 @@ def get_g2( g2_kstages = int(os.environ.get("MXFP4_G2_KSTAGES", "2")) g2_bhoist = os.environ.get("MXFP4_G2_BHOIST", "1") == "1" g2_ascale_pf = os.environ.get("MXFP4_G2_ASCALE_PF", "1") == "1" - g2_spart = int(os.environ.get("MXFP4_G2_SPART", "402")) + # These MUST use the same fallbacks as compile_gemm2_a4w4_port: the key is built + # here but the kernel is compiled there, so any disagreement silently hands back + # a launcher that does not match the key. + g2_spart = int(os.environ.get("MXFP4_G2_SPART", str(_G2_SPART_DEFAULT))) g2_bf16_lds = os.environ.get("MXFP4_G2_BF16_LDS", "0") == "1" + g2_kunroll = int(os.environ.get("MXFP4_G2_KUNROLL", _G2_KUNROLL_DEFAULT)) + g2_diag = int(os.environ.get("MXFP4_G2_DIAG", "0")) + g2_epi = int(os.environ.get("MXFP4_G2_EPI", str(_G2_EPI_DEFAULT))) + g2_wpe = int(os.environ.get("MXFP4_G2_WPE", str(_G2_WPE_DEFAULT))) + g2_wcpl = int(os.environ.get("MXFP4_G2_WCPL", str(_G2_WCPL_DEFAULT))) + g2_nloop = int(os.environ.get("MXFP4_G2_NLOOP", str(_G2_NLOOP_DEFAULT))) + if g2_nloop > 1 and (HIDDEN_MAX // BN) % g2_nloop: + g2_nloop = 1 key = ( BM, BN, @@ -469,6 +596,12 @@ def get_g2( g2_ascale_pf, g2_spart, g2_bf16_lds, + g2_kunroll, + g2_diag, + g2_epi, + g2_wpe, + g2_wcpl, + g2_nloop, out_dtype, ) launch = G2_CACHE.get(key) @@ -492,6 +625,11 @@ def get_g2( g2_ascale_pf=g2_ascale_pf, g2_spart=g2_spart, g2_bf16_lds=g2_bf16_lds, + g2_kunroll=g2_kunroll, + g2_diag=g2_diag, + g2_epi=g2_epi, + g2_wpe=g2_wpe, + g2_wcpl=g2_wcpl, out_dtype=out_dtype, ) G2_CACHE[key] = launch @@ -560,6 +698,15 @@ def mxfp4_moe_gemm2( raise AssertionError( f"D_INTER ({D_INTER}) exceeds compile cap INTER_MAX ({INTER_MAX})" ) + if int(os.environ.get("MXFP4_G2_KUNROLL", _G2_KUNROLL_DEFAULT)): + # The unrolled K path folds the tile count at compile time, so the compile + # bucket must be the exact runtime K (one launcher per inter_dim). + INTER_MAX = D_INTER + if int(os.environ.get("MXFP4_G2_NLOOP", _G2_NLOOP_DEFAULT)) > 1: + # Same for the n-group schedule: the group count must divide the n range + # exactly, so a bucket shared with a different model_dim could drop a tail + # n-tile. One launcher per model_dim. + HIDDEN_MAX = D_HIDDEN launch = get_g2( BM, BN, @@ -590,25 +737,34 @@ def mxfp4_moe_gemm2( out_scale = out # unused by the atomic epilog; any valid device ptr is fine # i32_kpad (inter_dim_pad) + i32_npad (model_dim_pad) are always threaded after # i32_hidden; when has_pad is False they are 0 and the kernel folds pad math away. - run_compiled( - launch, - inter_sorted_quant.data_ptr(), - inter_sorted_shuffled_scale.data_ptr(), - w2_u8.data_ptr(), - w2_scale_u8.data_ptr(), - sorted_expert_ids.data_ptr(), - cumsum_tensor.data_ptr(), - sorted_token_ids.data_ptr(), - sorted_weights.data_ptr(), - M_logical, - max_m_blocks, - grid_blocks, - D_INTER, - D_HIDDEN, - int(inter_dim_pad), - int(model_dim_pad), - out.data_ptr(), - out_scale.data_ptr(), - torch.cuda.current_stream() if stream is None else stream, + # Codegen happens on the first run_compiled for a given launcher, so the + # waves-per-EU hint has to be live here rather than at definition time. + _wpe = int(os.environ.get("MXFP4_G2_WPE", str(_G2_WPE_DEFAULT))) + _hint = ( + CompilationContext.compile_hints({"waves_per_eu": _wpe}) + if _wpe + else contextlib.nullcontext() ) + with _hint: + run_compiled( + launch, + inter_sorted_quant.data_ptr(), + inter_sorted_shuffled_scale.data_ptr(), + w2_u8.data_ptr(), + w2_scale_u8.data_ptr(), + sorted_expert_ids.data_ptr(), + cumsum_tensor.data_ptr(), + sorted_token_ids.data_ptr(), + sorted_weights.data_ptr(), + M_logical, + max_m_blocks, + grid_blocks, + D_INTER, + D_HIDDEN, + int(inter_dim_pad), + int(model_dim_pad), + out.data_ptr(), + out_scale.data_ptr(), + torch.cuda.current_stream() if stream is None else stream, + ) return out diff --git a/aiter/ops/flydsl/kernels/mxmoe_gemm_v2.py b/aiter/ops/flydsl/kernels/mxmoe_gemm_v2.py index a5b2ac493c..cad9915523 100644 --- a/aiter/ops/flydsl/kernels/mxmoe_gemm_v2.py +++ b/aiter/ops/flydsl/kernels/mxmoe_gemm_v2.py @@ -12,11 +12,13 @@ Float8E4M3FN, Float32, Int8, + Int16, Int32, T, ) from flydsl.expr.typing import Vector as Vec +from . import dpp_utils from .mxfp4_gemm_common import _fabs_f32 as fabs_f32 from .mxfp4_gemm_common import _lds_swizzle_mask as lds_swizzle_mask from .mxfp4_gemm_common import ( @@ -218,8 +220,19 @@ def gemm2_body_v2( g2_bhoist=True, g2_ascale_pf=True, g2_bf16_lds=False, + g2_diag=0, + g2_kunroll=False, + g2_epi=0, + g2_wcpl=0, route_out_fp8=False, + c_lds_off=0, ): + # g2_diag: PERF-ATTRIBUTION ONLY (wrong results); each bit drops one phase. + diag_no_barrier = bool(g2_diag & 1) + diag_no_epilog = bool(g2_diag & 2) + diag_no_ads = bool(g2_diag & 4) # kunroll path only: skip the A ds-reads + diag_no_bld = bool(g2_diag & 8) # kunroll path only: skip the B global loads + diag_no_mfma = bool(g2_diag & 16) # kunroll path only: skip the MFMA cluster # gemm2 K-loop perf knobs (default ON, no-op unless g2_kstages==2): kstages=2 double-buffers B weight+scale one tile ahead; bhoist issues that prefetch above the LDS barrier; ascale_pf prefetches A-scale one tile ahead. if g2_kstages not in (1, 2): raise AssertionError(f"g2_kstages must be 1 or 2, got {g2_kstages}") @@ -277,7 +290,9 @@ def gemm2_body_v2( lane_mod_16 = lane % 16 s_aq_base = lds_base_i32 - lds_acc_base = lds_base_i32 # f32 acc unions the A-tile LDS region (shared union) + # The f32 C slab normally unions the A-tile LDS region. Under the N-loop schedule + # A must survive every n-tile's epilogue, so the caller passes a disjoint offset. + lds_acc_base = lds_base_i32 + fx.Int32(c_lds_off) if c_lds_off else lds_base_i32 mma_atoms = scale_mma_atoms(a_dtype) aq_num_records = fx.Int64(i32_max_m_blocks) * fx.Int64(BM * K_BYTES) @@ -303,8 +318,10 @@ def issue_a_load_lds(slot, kt): BM=BM, ) - def issue_a_ds_read(slot): + def issue_a_ds_read(slot, frags=None): # A ds-read for one slot into a_frags: fp8 -> i32<8:1> (two 128-K halves), fp4 -> i32<4:1>. + if frags is None: + frags = a_frags for k in range_constexpr(kHalves): for i in range_constexpr(kMChunks): lds_row = lane_mod_16 + i * 16 @@ -333,7 +350,7 @@ def issue_a_ds_read(slot): ) ) a64 = Vec.from_elements([lo[0], lo[1], hi[0], hi[1]], fx.Int64) - a_frags[i][k].store(a64.bitcast(fx.Int32)) + frags[i][k].store(a64.bitcast(fx.Int32)) else: mask = lds_swizzle_mask(lane_mod_16, KH_TILE_A) lds_col = (lane_div_16 * 16 + k * 64) ^ mask @@ -344,7 +361,7 @@ def issue_a_ds_read(slot): fx.Int32, align=16, ) - a_frags[i][k].store(Vec(vec)) + frags[i][k].store(Vec(vec)) # Scale words (e8m0): shared scale_view / copy atom for both A and B. A-scale is one # word per 32-row chunk, each view bounded to bytes remaining after its baked base. @@ -460,7 +477,9 @@ def shift_scale_word(scale, kt_rt): scale_shift = (kt_rt % fx.Int32(tilesPerScaleChunk)) * fx.Int32(16) return scale.shrui(scale_shift) - def mfma_cluster(bqf, bsf, sa, kt_rt): + def mfma_cluster(bqf, bsf, sa, kt_rt, afrg=None): + if afrg is None: + afrg = a_frags # opsel (no gate/up split): mni=J//2, in_b=J%2; sa is a per-32-row-chunk list. sa = [ shift_scale_word(sa[sub], kt_rt) for sub in range_constexpr(kScaleSubBlocks) @@ -479,7 +498,7 @@ def mfma_cluster(bqf, bsf, sa, kt_rt): sa[0], sb, bqf, - a_frags, + afrg, c_frags, mma_atoms, i0=0, @@ -494,7 +513,7 @@ def mfma_cluster(bqf, bsf, sa, kt_rt): sa[sub], sb, bqf, - a_frags, + afrg, c_frags, mma_atoms, i0=2 * sub, @@ -522,7 +541,64 @@ def store_c_carry(state): n += 1 return n - if const_expr(g2_kstages == 1): + if const_expr(g2_kunroll): + # Fully-unrolled K (inter_dim == INTER_MAX, K_TILES compile-time). All + # K_TILES A-tiles were DMA'd to their own LDS slot by the prologue, so the + # whole contraction needs ONE barrier and carries no scf.for state. + KT = INTER_MAX // BK + # g2_kunroll==2: reuse ONE register set across all K-tiles (A frags and the + # B double-buffer) instead of KT private sets, trading some ds_read/MFMA + # overlap for ~128 VGPRs of occupancy headroom. + nA = 1 if g2_kunroll == 2 else KT + nB = 2 if g2_kunroll == 2 else KT + a_sets = [ + [ + [fx.make_rmem_tensor(A_NDW, Int32) for _ in range_constexpr(kHalves)] + for _ in range_constexpr(kMChunks) + ] + for _ in range_constexpr(nA) + ] + bqfs = [make_bq_fragments() for _ in range_constexpr(nB)] + bsfs = [make_scale_fragments(nPairs) for _ in range_constexpr(nB)] + safs = [make_scale_fragments(kScaleSubBlocks) for _ in range_constexpr(nB)] + + def load_a_scale_into(saf, kt_ct): + sa = load_a_scale_tile(fx.Int32(kt_ct)) + for sub in range_constexpr(kScaleSubBlocks): + saf[sub].store(Vec.from_elements([sa[sub]], Int32)) + + # Issue tile 0's B before the barrier so its vmem latency hides behind the + # A-DMA wait that the barrier resolves. + if const_expr(not diag_no_bld): + issue_b_load_into(bqfs[0], bsfs[0], fx.Int32(0)) + load_a_scale_into(safs[0], 0) + if const_expr(not diag_no_barrier): + gpu.barrier() + for kt_ct in range_constexpr(KT): + cur_b, nxt_b = kt_ct % nB, (kt_ct + 1) % nB + if const_expr(kt_ct + 1 < KT): + if const_expr(not diag_no_bld): + issue_b_load_into(bqfs[nxt_b], bsfs[nxt_b], fx.Int32(kt_ct + 1)) + load_a_scale_into(safs[nxt_b], kt_ct + 1) + if const_expr(not diag_no_ads): + issue_a_ds_read(kt_ct, frags=a_sets[kt_ct % nA]) + sa = [ + Vec(safs[cur_b][sub].load())[0] + for sub in range_constexpr(kScaleSubBlocks) + ] + rocdl.sched_barrier(0) + rocdl.s_setprio(1) + if const_expr(not diag_no_mfma): + mfma_cluster( + bqfs[cur_b], + bsfs[cur_b], + sa, + fx.Int32(kt_ct), + afrg=a_sets[kt_ct % nA], + ) + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + elif const_expr(g2_kstages == 1): # 1-deep pipe: synchronous B load per K-tile. for kt_iv, state in range( fx.Int32(0), @@ -532,7 +608,8 @@ def store_c_carry(state): ): store_c_carry(state) kt_rt = fx.Int32(kt_iv) - gpu.barrier() + if const_expr(not diag_no_barrier): + gpu.barrier() issue_a_ds_read(kt_rt % fx.Int32(aStages)) nxt = kt_rt + fx.Int32(kStages) if nxt < K_TILES_RT: @@ -629,7 +706,8 @@ def prefetch_next_b(kt_rt): kt_rt = fx.Int32(kt_iv) if const_expr(g2_bhoist): prefetch_next_b(kt_rt) - gpu.barrier() + if const_expr(not diag_no_barrier): + gpu.barrier() issue_a_ds_read(kt_rt % fx.Int32(aStages)) nxt_a = kt_rt + fx.Int32(kStages) if nxt_a < K_TILES_RT: @@ -655,6 +733,18 @@ def prefetch_next_b(kt_rt): accm_vecs = [ [c_frags[i][J].load() for J in range(numAccN)] for i in range(kMChunks) ] + if const_expr(diag_no_epilog): + # Attribution build: one dependent store keeps the whole MFMA chain live. + acc = fx.Float32(0.0) + for i in range_constexpr(kMChunks): + for J in range_constexpr(numAccN): + v = Vec(accm_vecs[i][J]) + for q in range_constexpr(4): + acc = acc + fx.Float32(v[q]) + out_i8 = global_typed_ptr(arg_out, T.i8) + if acc > fx.Float32(1.0e30): + out_i8[lane] = fx.Int8(1) + return atomic_bf16_epilog( lds_acc_base, accm_vecs, @@ -673,6 +763,9 @@ def prefetch_next_b(kt_rt): topk=topk, SBM=SBM, g2_bf16_lds=g2_bf16_lds, + g2_epi=g2_epi, + g2_wcpl=g2_wcpl, + g2_diag=g2_diag, route_out_fp8=route_out_fp8, ) @@ -697,8 +790,18 @@ def atomic_bf16_epilog( topk=1, SBM=None, g2_bf16_lds=False, + g2_epi=0, + g2_wcpl=0, + g2_diag=0, route_out_fp8=False, ): + # Epilogue attribution bits (WRONG results): 64 = skip the C-slab cshuffle write, + # 128 = skip the readback/quantise/store loop, 256 = skip the stids/sweights + # loads, 512 = skip the e8m0 scale stores. + diag_no_cwrite = bool(g2_diag & 64) + diag_no_cread = bool(g2_diag & 128) + diag_no_meta = bool(g2_diag & 256) + diag_no_scale = bool(g2_diag & 512) if SBM is None: SBM = BM kMChunks = BM // 16 @@ -712,6 +815,48 @@ def atomic_bf16_epilog( if const_expr(g2_bf16_lds) else None ) + # C-slab XOR swizzle (g2_epi>=1). Row stride BN is a multiple of 32 banks, so the + # four lane_div_16 write groups (4 rows apart) all land on the same banks -> 4-way + # conflict on every cshuffle ds_write. XOR the column by 32 bytes per 4-row group + # to spread them over banks 0/8/16/24. 32 bytes is >= every readback run, so + # contiguous runs survive the XOR intact and stay addressable. + C_ELEM_BYTES = 2 if g2_bf16_lds else 4 + SWZ = (32 // C_ELEM_BYTES) if g2_epi >= 1 else 0 + # g2_epi>=2: stage the cshuffle in CSPLIT row-slices instead of the whole BM x BN + # slab at once. The slab is the LDS high-water mark, so splitting it buys + # workgroups/CU. Purely a staging change: the same values are written and read + # back in the same order, so it is bit-exact (unlike g2_bf16_lds). + _want_split = 4 if (g2_epi & 4) else (2 if (g2_epi & 2) else 1) + CSPLIT = _want_split + while CSPLIT > 1 and (kMChunks % CSPLIT or M_REPS % CSPLIT): + CSPLIT //= 2 + C_CHUNKS = kMChunks // CSPLIT # cshuffle-write chunks per slice + C_MREPS = M_REPS // CSPLIT # readback rows per thread per slice + C_SLICE_ROWS = BM // CSPLIT + + # g2_epi bit 5: bank-conflict swizzle for the slab READBACK. cswz's row_grp4 XOR + # only permutes within a 16-element block, so the 16 lanes of a readback row-group + # (columns 16 apart) start in just 2 of the 8 bank quads -> 8-way conflict. XOR + # bits 2-4 of the column with bits 5-7 of the same column: bits 5-7 survive, so it + # is an involution for every fixed row_grp4, and it moves the 16 lanes to 2 per + # quad -- the 32-bank minimum for a 64-lane b128 read. Granularity drops from 8 + # elements to 4, so readbacks swizzle each 4-float chunk instead of adding 4 to a + # swizzled base. + CSWZ4 = bool(g2_epi & 32) and SWZ != 0 and C_ELEM_BYTES == 4 + # ...and the WRITE side needs a wider row XOR. LDS retires a b32 wave access in + # two 32-lane halves, so lane groups 0/1 must cover 32 distinct banks; a 32-byte + # XOR leaves both on the same 16. A 64-byte XOR on the row-group parity puts them + # on opposite halves (groups 2/3 reuse those banks, but are the other half-wave). + SWZ_R = (64 // C_ELEM_BYTES) if CSWZ4 else SWZ + + def cswz(col, row_grp4): + # row_grp4 = (row >> 2) & 3, the write-group / readback-row selector. + if const_expr(SWZ == 0): + return col + if const_expr(CSWZ4): + col = col ^ ((row_grp4 & fx.Int32(1)) * fx.Int32(SWZ_R)) + return col ^ (((col >> fx.Int32(5)) & fx.Int32(7)) << fx.Int32(2)) + return col ^ (row_grp4 * fx.Int32(SWZ)) tx_i32 = fx.Int32(gpu.thread_id("x")) m_lane = tx_i32 // 32 @@ -721,6 +866,52 @@ def atomic_bf16_epilog( col_start = n_lane * store_vec wave_n = BN // 4 + # g2_epi bit 3: "wide" route-out store. The default readback gives each lane 8 + # output columns of 8 different rows, so per thread the epilogue issues 8 dwordx2 + # value stores + 8 single-byte e8m0 scale stores and redoes the i64 row-address + # math and the token_id bit-identical output. + WIDE = bool(g2_epi & 8) and use_reduce and route_out_fp8 and not g2_bf16_lds + # Columns per lane. The natural choice keeps all 256 threads busy + # (C_SLICE_ROWS * BN / 256), but that can leave W_CPL/8 == 1 e8m0 byte per lane, + # and single-byte global stores dominate the epilogue. g2_wcpl forces a wider + # slice so the scales go out as i16/i32, at the price of idling some threads. + W_CPL = (BN * C_SLICE_ROWS) // 256 if WIDE else 0 + if WIDE and g2_wcpl: + W_CPL = max(W_CPL, g2_wcpl) + W_CPL = min(W_CPL, BN) if WIDE else 0 + if WIDE and ( + W_CPL < 8 + or W_CPL % 8 + or BN % W_CPL + or (W_CPL // 8) not in (1, 2, 4) + or C_SLICE_ROWS % 4 + or (C_SLICE_ROWS * (BN // W_CPL)) > 256 + ): + WIDE = False + W_NL = (BN // W_CPL) if WIDE else 0 + W_NS = (W_CPL // 8) if WIDE else 0 + W_ACTIVE = (C_SLICE_ROWS * W_NL) if WIDE else 0 + # g2_epi bit 4: DPP-combine the e8m0 bytes of W_NPACK adjacent lanes into one + # dword before storing. Those lanes hold the same output row and consecutive + # scale indices, so the merge is a pure quad_perm shuffle. Unlike widening the + # per-lane column slice, this drops the sub-dword stores with no idle threads. + W_NPACK = (4 // W_NS) if (WIDE and (g2_epi & 16)) else 1 + if WIDE and (W_NPACK > 1) and (W_NL % W_NPACK or W_NL < W_NPACK): + W_NPACK = 1 + if const_expr(WIDE): + w_mlane = tx_i32 // fx.Int32(W_NL) + w_nlane = tx_i32 - w_mlane * fx.Int32(W_NL) + # Rows within a wave must stride by 4 so that the 4-row cswz groups differ and + # the readback keeps the bank spread the 8-col mapping had (a 4 x C_SLICE_ROWS/4 + # transpose; bijective on [0, C_SLICE_ROWS)). + w_m4 = w_mlane * fx.Int32(4) + w_row_local = (w_m4 % fx.Int32(C_SLICE_ROWS)) + (w_m4 // fx.Int32(C_SLICE_ROWS)) + w_col_base = w_nlane * fx.Int32(W_CPL) + def flat_buffer(arg, elem_ty, align): ptr = global_typed_ptr(arg, elem_ty, align=align) view = fx.Tensor(fx.make_view(ptr, fx.make_layout((1, 1), (1, 1)))) @@ -736,51 +927,113 @@ def flat_buffer(arg, elem_ty, align): store_bf16x2 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(2), BFloat16) atomic_bf16x2 = fx.make_copy_atom(fx.rocdl.BufferAtomicPkAdd(BFloat16), BFloat16) store_i32 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(2), Int32) + store_i32x2 = fx.make_copy_atom(fx.rocdl.BufferCopy64b(2), Int32) store_i8 = fx.make_copy_atom(fx.rocdl.BufferCopy8b(2), Int8) + store_i32x4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(2), Int32) + store_i16 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(2), Int16) - def load_scalar(atom, src, index, elem_ty): + def issue_scalar(atom, src, index, elem_ty): + # Issue only: the register read is deferred to collect_scalar so that ONE + # s_waitcnt can cover the whole batch instead of one stall per load. frag = fx.make_rmem_tensor(1, elem_ty) fx.copy(atom, src[None, index], frag) + return frag + + def collect_scalar(frag): return Vec(frag.load())[0] + def load_scalar(atom, src, index, elem_ty): + return collect_scalar(issue_scalar(atom, src, index, elem_ty)) + # Prefetch sorted_token_ids / sorted_weights (invariant); latency overlaps stores+barriers. + # Issue every load back to back and only then read the destination registers, + # otherwise each `frag.load()` forces its own waitcnt and the batch costs one + # serialised L2 round-trip per row right before the barrier. packed = [] weight = [] - for mr in range_constexpr(M_REPS): - sorted_pos = m_row + mr * 8 + m_lane - packed.append(load_scalar(load_i32, stids, sorted_pos, Int32)) - weight.append(load_scalar(load_f32, sweights, sorted_pos, Float32)) - - # pre-store fence+barrier (HIP run_one __syncthreads() before the epilog). - gpu.barrier() + _pf = [] + # WIDE needs one (token_id, weight) pair per C slice, not one per M_REPS row. + _meta_reps = ( + range_constexpr(CSPLIT) if const_expr(WIDE) else range_constexpr(M_REPS) + ) + # NOTE: these loads are invariant across the n-tiles of an MXFP4_G2_NLOOP group, + # but caching them across tiles measures slower -- the extra values live across + # the next n-tile's whole K loop, and the K loop is more sensitive to register + # pressure than to a handful of L2-resident scalar loads. + for mr in _meta_reps: + if const_expr(WIDE): + sorted_pos = m_row + fx.Int32(mr * C_SLICE_ROWS) + w_row_local + else: + sorted_pos = m_row + mr * 8 + m_lane + if const_expr(diag_no_meta): + _pf.append(None) + else: + _pf.append( + ( + issue_scalar(load_i32, stids, sorted_pos, Int32), + issue_scalar(load_f32, sweights, sorted_pos, Float32), + ) + ) - # write accm -> lds_acc cshuffle. f32 path: scalar f32 stores (weight applied on readback). - if const_expr(g2_bf16_lds): - for i in range_constexpr(kMChunks): - row_base = fx.Int32(i * 16) + lane_div_16 * 4 - w_row = [ - load_scalar(load_f32, sweights, m_row + row_base + v, Float32) + # bf16-LDS path bakes the routing weight in at cshuffle-write time; prefetch all + # kMChunks*4 of those sweights rows ABOVE the barrier so the vmem latency overlaps + # the barrier + the tail of the K loop instead of stalling each chunk in turn. + w_rows = None + _wpf = None + if const_expr(g2_bf16_lds and diag_no_meta): + w_rows = [ + [fx.Float32(1.0) for _ in range_constexpr(4)] + for _ in range_constexpr(kMChunks) + ] + elif const_expr(g2_bf16_lds): + _wpf = [ + [ + issue_scalar( + load_f32, + sweights, + m_row + fx.Int32(i * 16) + lane_div_16 * 4 + v, + Float32, + ) for v in range_constexpr(4) ] + for i in range_constexpr(kMChunks) + ] + + # All meta loads are in flight; now read the destination registers. + for mr in _meta_reps: + if const_expr(diag_no_meta): + packed.append(fx.Int32(0)) + weight.append(fx.Float32(1.0)) + else: + packed.append(collect_scalar(_pf[mr][0])) + weight.append(collect_scalar(_pf[mr][1])) + + if _wpf is not None: + w_rows = [ + [collect_scalar(_wpf[i][v]) for v in range_constexpr(4)] + for i in range_constexpr(kMChunks) + ] + + def cshuffle_write(sp): + # accm -> lds_acc for row slice `sp`. f32 path stores raw accumulators (the + # routing weight is applied on readback); bf16 path bakes the weight in. + if const_expr(diag_no_cwrite): + return + for i in range_constexpr(sp * C_CHUNKS, (sp + 1) * C_CHUNKS): + # LDS row is slice-relative; the global row is i*16 + ... as before. + row_base = fx.Int32((i - sp * C_CHUNKS) * 16) + lane_div_16 * 4 for J in range_constexpr(numAccN): col = wave * wave_n + J * 16 + lane_mod_16 vec = Vec(accm[i][J]) + col_s = cswz(col, lane_div_16) for v in range_constexpr(4): - idx = (row_base + v) * BN + col - lds_base_bf16[idx] = fx.BFloat16( - fx.Float32(vec[v]) * fx.Float32(w_row[v]) - ) - else: - for i in range_constexpr(kMChunks): - row_base = fx.Int32(i * 16) + lane_div_16 * 4 - for J in range_constexpr(numAccN): - col = wave * wave_n + J * 16 + lane_mod_16 - vec = Vec(accm[i][J]) - for v in range_constexpr(4): - idx = (row_base + v) * BN + col - lds_base_fptr[idx] = fx.Float32(vec[v]) - - gpu.barrier() + idx = (row_base + v) * BN + col_s + if const_expr(g2_bf16_lds): + lds_base_bf16[idx] = fx.BFloat16( + fx.Float32(vec[v]) * fx.Float32(w_rows[i][v]) + ) + else: + lds_base_fptr[idx] = fx.Float32(vec[v]) # read back + weighted store (atomic: fadd out[token_id]; reduce: store out[token_id*topk+slot]); # token_id95K, 932us->107us). Always gate; reduce already gated via # use_reduce. (fp4-atomic families are gated too -> strictly correct OOB-skip, kernel IR changes.) - def store_one_mr(mr): - row_in_block = fx.Int32(mr * 8) + m_lane + def _quant8(vals): + """8 weighted f32 -> (lo_i32, hi_i32, e8m0), as in store_one_mr's route path.""" + local_max = fabs_f32(vals[0]) + for q in range_constexpr(1, 8): + local_max = local_max.maximumf(fabs_f32(vals[q])) + amax_bits = fx.Int32(_raw(local_max).bitcast(T.i32)) + ax_e = (amax_bits >> fx.Int32(23)) & fx.Int32(0xFF) + e8m0 = ax_e - fx.Int32(7) + e8m0 = (e8m0 < fx.Int32(1)).select(fx.Int32(1), e8m0) + e8m0 = (amax_bits == fx.Int32(0)).select(fx.Int32(0), e8m0) + block_scale = fx.Float32(_raw(e8m0 << fx.Int32(23)).bitcast(T.f32)) + bs_raw = _raw(block_scale) + pk_ty = T.vec(2, T.i16) + lo = _raw(Vec.filled([2], 0, fx.Int16)) + lo = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, lo, _raw(vals[0]), _raw(vals[1]), bs_raw, 0 + ) + lo = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, lo, _raw(vals[2]), _raw(vals[3]), bs_raw, 1 + ) + hi = _raw(Vec.filled([2], 0, fx.Int16)) + hi = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, hi, _raw(vals[4]), _raw(vals[5]), bs_raw, 0 + ) + hi = rocdl.cvt_scalef32_pk_fp8_f32( + pk_ty, hi, _raw(vals[6]), _raw(vals[7]), bs_raw, 1 + ) + return ( + Vec(Vec(lo).bitcast(Int32))[0], + Vec(Vec(hi).bitcast(Int32))[0], + e8m0, + ) + + def store_wide(sp): + # One C-slice row per lane, W_CPL contiguous output columns. + grow = fx.Int32(sp * C_SLICE_ROWS) + w_row_local + row_grp4 = (grow >> fx.Int32(2)) & fx.Int32(3) + pk = packed[sp] + wt = weight[sp] + out_row = fx.Int64( + (pk & fx.Int32(0x00FFFFFF)) * fx.Int32(topk) + (pk >> fx.Int32(24)) + ) + row_base_addr = out_row * fx.Int64(N_OUT + (N_OUT // fx.Int32(8))) + col_g0_base = n_block_idx * BN + w_col_base + words = [] + scales = [] + for g in range_constexpr(W_NS): + row_e = w_row_local * BN + vals = [] + for h in range_constexpr(2): + base_e = row_e + cswz(w_col_base + fx.Int32(g * 8 + h * 4), row_grp4) + v4 = Vec( + lds_vec_load( + lds_acc_base, + base_e * 4, + Vec.make_type(4, Float32), + Float32, + align=16, + ) + ) + for q in range_constexpr(4): + vals.append(fx.Float32(v4[q]) * wt) + lo, hi, e8m0 = _quant8(vals) + words.append(lo) + words.append(hi) + scales.append(e8m0) + # W_CPL values = 2*W_NS dwords: dwordx4 while four remain, then a dwordx2 for + # the odd pair (W_NS==1, i.e. BN/W_NL == 8). row_base_addr (4032 B/row) and + # col_g0_base are both 16 B aligned, so every store is naturally aligned. + _nw = len(words) + for w in range_constexpr(_nw // 4): + f4 = fx.make_rmem_tensor(4, Int32) + f4.store(Vec.from_elements(words[w * 4 : w * 4 + 4], Int32)) + fx.copy( + store_i32x4, + f4, + out_i8[ + None, + row_base_addr + fx.Int64(col_g0_base + fx.Int32(w * 16)), + ], + ) + if const_expr(_nw % 4 == 2): + f2 = fx.make_rmem_tensor(2, Int32) + f2.store(Vec.from_elements(words[_nw - 2 :], Int32)) + fx.copy( + store_i32x2, + f2, + out_i8[ + None, + row_base_addr + fx.Int64(col_g0_base + fx.Int32((_nw // 4) * 16)), + ], + ) + scale_off = ( + row_base_addr + fx.Int64(N_OUT) + fx.Int64(col_g0_base // fx.Int32(8)) + ) + acc = scales[0] & fx.Int32(0xFF) + for q in range_constexpr(1, W_NS): + acc = acc | ((scales[q] & fx.Int32(0xFF)) << fx.Int32(8 * q)) + if const_expr(diag_no_scale): + pass + elif const_expr(W_NPACK > 1): + # quad_perm broadcasts: [i,i,i,i] for a 4-lane merge, [0,0,2,2]/[1,1,3,3] + # for a 2-lane one. Every lane ends up with the full dword; only the + # group leader stores it, at its own (lowest) scale offset. + ctrls = (0x00, 0x55, 0xAA, 0xFF) if W_NPACK == 4 else (0xA0, 0xF5) + merged = None + for j in range_constexpr(W_NPACK): + part = fx.Int32( + dpp_utils.update_dpp_i32( + _raw(acc), _raw(acc), ctrls[j], 0xF, 0xF, True + ) + ) + part = part << fx.Int32(8 * W_NS * j) + merged = part if merged is None else (merged | part) + sf = fx.make_rmem_tensor(1, Int32) + sf.store(Vec.from_elements([merged], Int32)) + + @flyc.jit + def store_scale_leader(sf, scale_off, tx_i32): + if (tx_i32 & fx.Int32(W_NPACK - 1)) == fx.Int32(0): + fx.copy(store_i32, sf, out_i8[None, scale_off]) + + store_scale_leader(sf, scale_off, tx_i32) + elif const_expr(W_NS == 4): + sf = fx.make_rmem_tensor(1, Int32) + sf.store(Vec.from_elements([acc], Int32)) + fx.copy(store_i32, sf, out_i8[None, scale_off]) + elif const_expr(W_NS == 2): + sf = fx.make_rmem_tensor(1, Int16) + sf.store(Vec.from_elements([acc.to(Int16)], Int16)) + fx.copy(store_i16, sf, out_i8[None, scale_off]) + else: + sf = fx.make_rmem_tensor(1, Int8) + sf.store(Vec.from_elements([acc.to(Int8)], Int8)) + fx.copy(store_i8, sf, out_i8[None, scale_off]) + + def store_one_mr(mr, sp=0): + # Global row within the BM tile drives the swizzle selector (it must match the + # write side, which indexes by the global chunk); the LDS row is slice-relative. + row_in_block = fx.Int32((mr - sp * C_MREPS) * 8) + m_lane + row_grp4 = ((fx.Int32(mr * 8) + m_lane) >> fx.Int32(2)) & fx.Int32(3) token_id = packed[mr] & fx.Int32(0x00FFFFFF) if const_expr(use_reduce): # reduce out_row can reach tokens*topk (large-M) so compute the element base in i64 (atomic i32 path byte-identical). @@ -813,13 +1205,50 @@ def store_one_mr(mr): def store_route_group(col_lane8): col_g0 = n_block_idx * BN + col_lane8 vals = [] - for q in range_constexpr(route_vec): - idx_q = row_in_block * BN + col_lane8 + fx.Int32(q) + if const_expr(g2_epi >= 1): + # route_vec cshuffle slots are contiguous in the C slab (the + # 32-byte XOR swizzle preserves 8-element runs), so pull them + # in as b128 vector ds_reads instead of route_vec scalar ones. + base_e = row_in_block * BN + cswz(col_lane8, row_grp4) + row_e = row_in_block * BN if const_expr(g2_bf16_lds): - # bf16 LDS already has routing weight baked in at write time. - vals.append(fx.Float32(lds_base_bf16[idx_q])) + v8 = Vec( + lds_vec_load( + lds_acc_base, + base_e * 2, + Vec.make_type(route_vec, BFloat16), + BFloat16, + align=16, + ) + ) + for q in range_constexpr(route_vec): + vals.append(fx.Float32(v8[q])) else: - vals.append(fx.Float32(lds_base_fptr[idx_q]) * weight[mr]) + for h in range_constexpr(route_vec // 4): + e_h = row_e + cswz( + col_lane8 + fx.Int32(h * 4), row_grp4 + ) + v4 = Vec( + lds_vec_load( + lds_acc_base, + e_h * 4, + Vec.make_type(4, Float32), + Float32, + align=16, + ) + ) + for q in range_constexpr(4): + vals.append(fx.Float32(v4[q]) * weight[mr]) + else: + for q in range_constexpr(route_vec): + idx_q = row_in_block * BN + col_lane8 + fx.Int32(q) + if const_expr(g2_bf16_lds): + # bf16 LDS already has the routing weight baked in. + vals.append(fx.Float32(lds_base_bf16[idx_q])) + else: + vals.append( + fx.Float32(lds_base_fptr[idx_q]) * weight[mr] + ) local_max = fabs_f32(vals[0]) for q in range_constexpr(1, route_vec): local_max = local_max.maximumf(fabs_f32(vals[q])) @@ -846,21 +1275,39 @@ def store_route_group(col_lane8): pk_ty, packed_hi, _raw(vals[6]), _raw(vals[7]), bs_raw, 1 ) row_val_off = row_base_addr + fx.Int64(col_g0) - packed_frag = fx.make_rmem_tensor(1, Int32) - packed_frag.store(Vec(packed_lo).bitcast(Int32)) - fx.copy(store_i32, packed_frag, out_i8[None, row_val_off]) - packed_frag.store(Vec(packed_hi).bitcast(Int32)) - fx.copy( - store_i32, packed_frag, out_i8[None, row_val_off + fx.Int64(4)] - ) + if const_expr(g2_epi >= 1): + # The two fp8 dwords are adjacent (col_g0 is 8-aligned) -> one + # dwordx2 store instead of two dwordx1. + pf2 = fx.make_rmem_tensor(2, Int32) + pf2.store( + Vec.from_elements( + [ + Vec(Vec(packed_lo).bitcast(Int32))[0], + Vec(Vec(packed_hi).bitcast(Int32))[0], + ], + Int32, + ) + ) + fx.copy(store_i32x2, pf2, out_i8[None, row_val_off]) + else: + packed_frag = fx.make_rmem_tensor(1, Int32) + packed_frag.store(Vec(packed_lo).bitcast(Int32)) + fx.copy(store_i32, packed_frag, out_i8[None, row_val_off]) + packed_frag.store(Vec(packed_hi).bitcast(Int32)) + fx.copy( + store_i32, + packed_frag, + out_i8[None, row_val_off + fx.Int64(4)], + ) scale_off = ( row_base_addr + fx.Int64(N_OUT) + fx.Int64(col_g0 // fx.Int32(route_vec)) ) - scale_frag = fx.make_rmem_tensor(1, Int8) - scale_frag.store(Vec.from_elements([e8m0.to(Int8)], Int8)) - fx.copy(store_i8, scale_frag, out_i8[None, scale_off]) + if const_expr(not diag_no_scale): + scale_frag = fx.make_rmem_tensor(1, Int8) + scale_frag.store(Vec.from_elements([e8m0.to(Int8)], Int8)) + fx.copy(store_i8, scale_frag, out_i8[None, scale_off]) @flyc.jit def store_route_group_if_valid(col_lane8): @@ -871,7 +1318,9 @@ def store_route_group_if_valid(col_lane8): else: for s in range_constexpr(BN // store_group_n): # adjacent ee=0,1 contiguous -> one 2-wide load. - idx0 = row_in_block * BN + col_start + s * store_group_n + idx0 = row_in_block * BN + cswz( + col_start + fx.Int32(s * store_group_n), row_grp4 + ) if const_expr(g2_bf16_lds): pk = Vec( lds_vec_load( @@ -903,12 +1352,35 @@ def store_route_group_if_valid(col_lane8): else: fx.copy(atomic_bf16x2, out_frag, out_bf16[None, out_off]) - for mr in range_constexpr(M_REPS): - token_id = packed[mr] & fx.Int32(0x00FFFFFF) - - @flyc.jit - def store_if_valid(token_id, mr): - if token_id < i32_M: - store_one_mr(mr) - - store_if_valid(token_id, mr) + for sp in range_constexpr(CSPLIT): + # Barrier before each slice: for sp>0 it also fences the previous slice's + # readback against this slice's overwrite of the same LDS bytes. The sp==0 + # barrier is needed too: when the C slab unions the A LDS region + # (c_lds_off == 0), the first slice's write otherwise races other waves' + # last-K-tile A ds_reads and the route-out differs run to run. + gpu.barrier() + cshuffle_write(sp) + gpu.barrier() + if const_expr(diag_no_cread): + continue + if const_expr(WIDE): + token_id = packed[sp] & fx.Int32(0x00FFFFFF) + + @flyc.jit + def store_wide_if_valid(token_id, sp, tx_i32): + # tx >= W_ACTIVE only when g2_wcpl widened the lane slice past the + # thread count; those lanes have no row and must not store. + if token_id < i32_M and tx_i32 < fx.Int32(W_ACTIVE): + store_wide(sp) + + store_wide_if_valid(token_id, sp, tx_i32) + continue + for mr in range_constexpr(sp * C_MREPS, (sp + 1) * C_MREPS): + token_id = packed[mr] & fx.Int32(0x00FFFFFF) + + @flyc.jit + def store_if_valid(token_id, mr, sp): + if token_id < i32_M: + store_one_mr(mr, sp) + + store_if_valid(token_id, mr, sp) diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index a9e7a4639c..bac27626ec 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -82,6 +82,9 @@ def pick_flydsl_stage2_tile_k(inter_dim: int) -> int: ``inter_dim % 256 != 0`` (e.g. DSV4 TP8 ``inter=640``) must use ``tile_k=128``; ``tile_k=256`` only tiles cleanly when K is 256-aligned. Matches ``fused_moe.get_2stage_cfgs`` FlyDSL fallback (``_s2_tk``). + + A single-pass ``tile_k == inter_dim`` (e.g. 384) is registered but not + auto-picked; the tuner has to name it explicitly. """ inter_dim = int(inter_dim) return 256 if (inter_dim % 256 == 0) else 128 @@ -291,7 +294,8 @@ def get_flydsl_stage2_kernels( # fp4 stage2 supports tile_k=128 (pack_K=1 scale sub-group shift path) as # well as 256. tile_k=128 cleanly tiles K=inter_dim for TP-sharded shapes # whose inter_dim is a multiple of 128 but not 256 (e.g. MiniMax TP4=384). - tile_ks = [128, 256] if (is_fp4 or is_fp8) else [128] + # 384 covers the single-pass K case for inter_dim=384 (MiniMax/KimiK3 TP). + tile_ks = [128, 256, 384] if (is_fp4 or is_fp8) else [128] tile_ms = [16, 32, 64, 128] if is_fp4 else [32, 64, 128] modes = ["atomic", "reduce"]