From 2f0cdab2edd2b823e873f5c2032b94d1f10a104c Mon Sep 17 00:00:00 2001 From: inkcherry Date: Mon, 11 May 2026 05:34:31 +0000 Subject: [PATCH 1/8] perf(flydsl): MXFP4 fused-MoE stage2 optimization for EP4 prefill Optimize the FlyDSL stage2 (down-projection) kernel for the production fp4xfp4 EP4 DeepSeek prefill shape (tile_m=64, tile_n=128, tile_k=256, M=49152) on MI355X. ## Summary of changes aiter/ops/flydsl/moe_kernels.py - Add new wrapper params (waves_per_eu, use_async_copy, cu_num_mul) and plumb them through compile_flydsl_moe_stage2 -> compile_mixed_moe_gemm2. - Register a production fp4xfp4 stage2 variant `flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic_persist_async_w4_cumul3` with waves_per_eu=4, use_async_copy=True, cu_num_mul=3. - Drop the defensive `out.fill_(0)` added in PR #2863: when accumulate=True the standard fused_moe path already zeros moe_buf via `moe_buf_set_zero_kernel_2d` inside `moe_sorting_*_fwd`. The extra fill is a ~token_num*model_dim HBM write (~130 us per call at MI355X HBM bw on this shape). Callers using flydsl_moe_stage2 directly with accumulate=True remain responsible for zeroing `out` (documented). aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py - compile_mixed_moe_gemm2 signature: add waves_per_eu, use_async_copy, cu_num_mul. Drop unused legacy b_nt, xcd_swizzle on this path. - cu_num_mul: multiplies persistent-mode grid_y (CU count) by an integer factor; cu_num_mul>1 launches more persistent CTAs that each cover fewer M tiles, increasing in-flight parallelism. cu_num_mul=3 is the EP4-prefill sweet spot; cu_num_mul=4 regresses ~2.4% on the same shape. - use_async_copy: enable raw_ptr_buffer_load_lds for the X (activation) tile in the prologue so the X DMA overlaps with B/scale VMEM. Gated to the production tile/dtype combo so other configs are unaffected. - Asymmetric b_lo/b_hi split (b_lo loads 3/4 of K, b_hi loads 1/4): the smaller b_hi burst frees VMEM issue slots right before s_setprio(1) and the K-loop; the extra ku of b_lo is issued in the prologue where it overlaps with the X DMA. - Stage sorted topk weights into LDS once per M-tile and prefetch via vec4 LDS reads in the cshuffle epilogue, eliminating per-mi VMEM dwordx4 loads from the MFMA-heavy hot loop. - Emit AMDGPU `disable_xdl_arb_stall` + `s_setprio(1)` around the K-loop on gfx950 so the high-prio K-loop drains MFMAs back-to-back. - Issue scales BEFORE B-VMEM in the steady-state K-loop so the extract-time stalls in the first MFMA chunk are reduced. - Defer the lds_tid prologue (sorted_idx/sorted_w VMEM dwords + LDS stores) to AFTER the long-latency X DMA + B/scale loads, hiding the lds_tid scalar load latency behind the dominant VMEM phase. - Merge the lds_tid + lds_tw prologue store guards into a single scf.IfOp so MLIR can co-schedule the two VMEM buffer_loads. - For the buffer-atomic path: fast-path the cshuffle row token-index load (move the AND + multiply + validity-mask into the lds_tid preload site where it overlaps with B/scale/X VMEM loads). - Balanced persistent tile distribution (`tD4balPersist`) and balanced LDS-write hint accounting on the async path (`dswr_async`). aiter/fused_moe.py - `_flydsl_stage2_wrapper`: forward waves_per_eu, use_async_copy, cu_num_mul from the parsed kernel-name config to flydsl_moe_stage2. ## Performance Benchmark harness: /home/mingzliu/sgl_opt/ep4_mxfp4_opt/coopt_loop/scripts/bench.sh Workload: /home/mingzliu/sgl_opt/test_fused_moe_ep4_mxfp4.py (M=49152, model_dim=7168, inter_dim=2048, 64 local experts, topk=8, QuantType.per_1x32, fp4x2 weights & activations, bf16 out) Metric: 5-run median latency, 3 reps, MI355X (smci355-ccs-aus-n06-21). baseline (kernel=_atomic_persist): 3438.8 us (3-rep avg) this PR (kernel=_atomic_persist_async_w4_cumul3): 3173.6 us (3-rep avg) delta: -265 us (-7.71%) --- aiter/fused_moe.py | 3 + .../flydsl/kernels/mixed_moe_gemm_2stage.py | 1179 +++++++++++------ aiter/ops/flydsl/moe_kernels.py | 40 +- 3 files changed, 798 insertions(+), 424 deletions(-) diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 1b3ed4067b1..3514f0c0613 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -795,6 +795,9 @@ def _flydsl_stage2_wrapper( a2_scale=a2_scale, sorted_weights=sorted_weights, sort_block_m=parsed.get("sort_block_m", 0), + waves_per_eu=parsed.get("waves_per_eu", None), + use_async_copy=parsed.get("use_async_copy", False), + cu_num_mul=parsed.get("cu_num_mul", 1), b_nt=parsed.get("b_nt", 0), persist=parsed.get("persist", None), bias=bias2, diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 83371d35b0a..2831ae9d2cd 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -11,6 +11,7 @@ import os from contextlib import contextmanager +from typing import Optional import flydsl.compiler as flyc import flydsl.expr as fx @@ -2750,11 +2751,18 @@ def compile_mixed_moe_gemm2( inter_dim_pad: int = 0, persist_m: int = 4, sort_block_m: int = 0, - b_nt: int = 2, - xcd_swizzle: int = 0, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, ): """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + waves_per_eu: + - None (default): let the backend pick. + - 1..4: hint the ROCm backend with ``rocdl.waves_per_eu`` to cap or + request a specific occupancy. Useful to trade VGPR pressure for + better latency hiding on the CU. + persist_m: - > 0: legacy mode -- each CTA processes exactly persist_m consecutive M tiles. - <= 0: **persistent mode** -- grid_y = cu_num (auto-detected), each CTA @@ -2790,10 +2798,15 @@ def compile_mixed_moe_gemm2( f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" ) + # issue async X DMA early in prologue for max overlap with B/scale VMEM. + _r139_xdma_first = bool( + use_async_copy and tile_m == 64 and tile_n == 128 and tile_k == 256 + and a_dtype == "fp4" and b_dtype == "fp4" and bool(accumulate) + ) + + gpu_arch = get_hip_arch() - allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") - allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") - _state = {} + allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") if a_dtype not in ("fp8", "fp16", "int8", "fp4"): raise ValueError( @@ -2910,13 +2923,6 @@ def _scale_elem_type(): pad_k = 0 if _use_lds128 else 8 lds_stride = tile_k + pad_k - if 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 - else: - _eff_lds_stride = lds_stride - _eff_tile_k_bytes = tile_k_bytes - if out_is_f32: # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. _use_cshuffle_epilog = ( @@ -2957,45 +2963,56 @@ def _load_bias_scalar(bias_rsrc, offset): # Dynamic-shape variant: safe to reuse across (tokens/sorted_size/size_expert_ids) at runtime. # Keep a distinct ABI tag so the compile cache never mixes with historical signatures. _persistent = persist_m <= 0 + # cu_num_mul: multiplies persistent-mode grid_y (CU count). =1 byte- + # identical to legacy; >1 launches more persistent CTAs each covering + # fewer M tiles. cu_num_mul=3 is the EP4 prefill sweet spot on MI355X + # (cu_num_mul=4 regresses ~2.4%). + if not isinstance(cu_num_mul, int) or cu_num_mul < 1: + raise ValueError(f"cu_num_mul must be int >= 1, got {cu_num_mul}") if _persistent: from aiter.jit.utils.chip_info import get_cu_num - _cu_num = get_cu_num() + _cu_num = get_cu_num() * int(cu_num_mul) else: _cu_num = 0 _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" - _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + _wpe_tag = f"_w{waves_per_eu}" if waves_per_eu is not None else "" + if waves_per_eu is not None and not (1 <= int(waves_per_eu) <= 10): + raise ValueError( + f"waves_per_eu must be in [1, 10] or None, got {waves_per_eu}" + ) + _num_k_tiles_per_batch = int(inter_dim) // int(tile_k) + _async_tag = "_async" if use_async_copy else "" + _cumul_tag = f"_cumul{int(cu_num_mul)}" if int(cu_num_mul) != 1 else "" module_name = ( f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" f"_t{tile_m}x{tile_n}x{tile_k}" - f"_vscale_fix3{_pm_tag}{_sbm_tag}{_xcd_tag}" + f"_vscale_fix3_fp4opt_v1{_pm_tag}{_sbm_tag}{_wpe_tag}{_async_tag}{_cumul_tag}" ).replace("-", "_") # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- - # Ping-pong A2 tiles via separate allocators (like stage1). - _single_x_bytes = int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) - _cshuffle_elem_bytes_s2 = 2 # f16/bf16 = 2 bytes + # Reuse a single allocation for both: + # - ping-pong A2 tiles (2 * tile_m * lds_stride * elem_bytes bytes) + # - epilogue CShuffle tile (tile_m * tile_n f16 -> 2 * tile_m * tile_n bytes) + lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(a_elem_bytes) lds_out_bytes = ( - _cshuffle_elem_bytes_s2 * int(tile_m) * int(tile_n) - if _use_cshuffle_epilog - else 0 - ) + 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + ) # f16 bytes lds_tid_bytes = int(tile_m) * 4 - _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) - - _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) - _ping_buffer_bytes = _single_x_bytes + # also stage sorted topk weights into LDS once per + # M-tile so the cshuffle epilogue path can prefetch them via fast LDS + # vec4 reads instead of doing a per-mi VMEM buffer_load_dwordx4 inside + # compute_tile. Only ~tile_m*4 bytes of additional LDS budget. + lds_tw_bytes = (int(tile_m) * 4) if bool(doweight_stage2) else 0 + lds_total_bytes = max(lds_x_bytes, lds_out_bytes) + lds_tid_bytes + lds_tw_bytes + lds_total_elems = lds_total_bytes if a_elem_bytes == 1 else (lds_total_bytes // 2) def x_lds_elem(): return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) - lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) - allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes - _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) - allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes - - lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) - allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + lds_alloc_bytes = int(lds_total_elems) * int(a_elem_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_alloc_bytes if True: @@ -3020,10 +3037,9 @@ def moe_gemm2( tokens_in = arith.index_cast(ir.IndexType.get(), i32_tokens_in.ir_value()) n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) - size_expert_ids_in = arith.index_cast( - ir.IndexType.get(), i32_size_expert_ids_in.ir_value() - ) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + # For int4, weights are stored as packed bytes (i8) and unpacked to i8 packs. f32 = T.f32 i32 = T.i32 i64 = T.i64 @@ -3066,72 +3082,73 @@ def check_c_k_valid_gate(base_k): arith, c_mn=c_n_total, c_k=c_k_orig ) + # When use_async_copy is enabled with packed-elem (fp4) A, the + # DMA (buffer_load_lds) writes packed bytes into LDS, so the + # effective stride is halved. Keep legacy (non-async) path + # bit-identical by using lds_stride directly. + if 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 + else: + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) stride_lds = fx.make_stride(_eff_lds_stride, 1) layout_lds = fx.make_layout(shape_lds, stride_lds) tx = gpu.thread_id("x") - by = gpu.block_id("x") # tile along model_dim (N-dim) + by_outer = gpu.block_id("x") bx_persist = gpu.block_id("y") # persistent WG index (M-dim) - - if const_expr(xcd_swizzle > 0): - _NUM_XCDS_S = 8 - _c1_sw = arith.constant(1, index=True) - _c_tn_sw = arith.constant(tile_n, index=True) - _c_mdp_sw = arith.constant(model_dim_pad, index=True) - _gx = (n_in - _c_mdp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw - if const_expr(_persistent): - _gy = arith.constant(_cu_num, index=True) - else: - _c_pm_sw = arith.constant(persist_m, index=True) - _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw - - _linear_id = bx_persist * _gx + by - _num_wgs = _gx * _gy - - _c_xcds = arith.constant(_NUM_XCDS_S, index=True) - _wgs_per_xcd = _num_wgs / _c_xcds - _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) - - _WGM_S = xcd_swizzle - _c_wgm = arith.constant(_WGM_S, index=True) - _num_wgid_in_group = _c_wgm * _gx - _group_id = _wgid / _num_wgid_in_group - _first_pid_m = _group_id * _c_wgm - _remaining_m = _gy - _first_pid_m - _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) - _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) - - _wgid_in_group = _wgid % _num_wgid_in_group - bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) - by = _wgid_in_group / _group_size_m + by = by_outer # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) - base_ptr_pong = allocator_pong.get_base() - base_ptr_ping = allocator_ping.get_base() - lds_x_pong = SmemPtr( - base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) - ).get() - lds_x_ping = SmemPtr( - base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) - ).get() + base_ptr = allocator.get_base() + lds_x_ptr = SmemPtr( + base_ptr, + lds_alloc_offset, + x_lds_elem(), + shape=(lds_total_elems,), + ) + lds_x = lds_x_ptr.get() + # Alias the same underlying LDS bytes as f16/bf16 for epilogue shuffle. lds_out = ( SmemPtr( - base_ptr_pong, - lds_pong_offset, + base_ptr, + lds_x_ptr.byte_offset, (T.bf16 if out_is_bf16 else T.f16), shape=(tile_m * tile_n,), ).get() if _use_cshuffle_epilog else None ) + + # lds_tid: alias LDS after max(x, out) for sorted_idx preload + _lds_x_b = 2 * 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( - base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + base_ptr, lds_x_ptr.byte_offset + _lds_tid_off, T.i32, shape=(tile_m,) ).get() + # lds_tw aliases the LDS slot immediately after + # lds_tid for staging the sorted topk weights (f32) once per M-tile. + # Only allocated when doweight_stage2; lives in a disjoint LDS region + # from lds_x / lds_out / lds_tid so it is safe to alias here. + _lds_tw_off = _lds_tid_off + int(tile_m) * 4 + lds_tw = ( + SmemPtr( + base_ptr, + lds_x_ptr.byte_offset + _lds_tw_off, + T.f32, + shape=(tile_m,), + ).get() + if doweight_stage2 + else None + ) # Buffer resources. # For dynamic memrefs, `max_size=False` cannot infer the logical size from the memref *type*, @@ -3142,6 +3159,7 @@ def check_c_k_valid_gate(base_k): # fp8/int8: 1 byte per element -> bytes = tokens*topk * K # fp4: 2 elements per byte -> bytes = tokens*topk * K / 2 c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + c_a_pack = arith.constant(int(a_elem_vec_pack), index=True) x_nbytes_idx = _div_pow2( (tokens_in * c_topk) * k_in * c_elem_bytes, int(a_elem_vec_pack) ) @@ -3157,7 +3175,7 @@ def check_c_k_valid_gate(base_k): out_nbytes_idx = ( tokens_in * n_in * arith.constant(out_elem_bytes, index=True) ) - if const_expr(not bool(accumulate)): + if not bool(accumulate): out_nbytes_idx = ( tokens_in * arith.index(topk) @@ -3186,10 +3204,10 @@ def check_c_k_valid_gate(base_k): num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) # fp16 path ignores scales completely (implicit scale=1.0). - sx_rsrc = 1 - sw_rsrc = 1 - if const_expr(not is_f16_a): - if const_expr(is_f4_a or is_f8_a): + if is_f16_a: + sx_rsrc = None + else: + if is_f4_a: # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. # Caller must pre-scatter a2_scale via moe_mxfp4_sort. kblk = _div_pow2(k_in, 32) @@ -3206,7 +3224,9 @@ def check_c_k_valid_gate(base_k): arg_scale_x, max_size=False, num_records_bytes=sx_nbytes_i32 ) - if const_expr(not is_f16_b): + if is_f16_b: + sw_rsrc = None + else: # Weight microscale buffer (packed i32 holding e8m0 bytes). # Use an exact descriptor size so hardware OOB checking works. kblk_w = _div_pow2(k_in, 32) # K/32 @@ -3255,16 +3275,26 @@ def check_c_k_valid_gate(base_k): _c0_p = arith.constant(0, index=True) _c1_p = arith.constant(1, index=True) - if const_expr(_persistent): - # Expert-phase scheduling: contiguous M-tile dispatch. - # grid_y = cu_num, each CTA handles a contiguous chunk of M-tiles: - # [bx_persist * tiles_per_block, ..., (bx_persist+1) * tiles_per_block - 1] - # Adjacent blocks process adjacent M-tiles -> same expert -> B weight L2 reuse. + if _persistent: + # Balanced contiguous M-tile dispatch: each CTA gets + # floor(T/CU) or ceil(T/CU) tiles, distributing the remainder + # across the first (T%CU) CTAs (avoids the no-work tail band + # of the old ceil mapping). Expert-contiguous ordering kept. _c_cu = arith.constant(_cu_num, index=True) _c_tm_p = arith.constant(tile_m, index=True) _num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) _total_m_tiles = (_num_valid_idx + _c_tm_p - _c1_p) / _c_tm_p - _tiles_per_block = (_total_m_tiles + _c_cu - _c1_p) / _c_cu + _tiles_per_block_base = _total_m_tiles / _c_cu + _tiles_remainder = _total_m_tiles - (_tiles_per_block_base * _c_cu) + _has_extra_tile = arith.cmpi( + CmpIPredicate.ult, bx_persist, _tiles_remainder + ) + _extra_tile = arith.select(_has_extra_tile, _c1_p, _c0_p) + _tiles_per_block = _tiles_per_block_base + _extra_tile + _start_tail = arith.select( + _has_extra_tile, bx_persist, _tiles_remainder + ) + _persist_start_tile = bx_persist * _tiles_per_block_base + _start_tail _i1 = ir.IntegerType.get_signless(1) _init_active = arith.constant(1, type=_i1) _for_persist = scf.ForOp(_c0_p, _tiles_per_block, _c1_p, [_init_active]) @@ -3284,9 +3314,9 @@ def check_c_k_valid_gate(base_k): _for_ip.__enter__() _mi_p = _for_persist.induction_variable - if const_expr(_persistent): + if _persistent: _still_active = _for_persist.inner_iter_args[0] - bx = bx_persist * _tiles_per_block + _mi_p + bx = _persist_start_tile + _mi_p else: _prev_expert_i32 = _for_persist.inner_iter_args[0] _prev_expert_b_base = _for_persist.inner_iter_args[1] @@ -3302,12 +3332,12 @@ def check_c_k_valid_gate(base_k): expert_i32 = buffer_ops.buffer_load( expert_rsrc, sort_blk, vec_width=1, dtype=T.i32 ) - expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + expert_idx = arith.index_cast(T.index, expert_i32) exp_valid = arith.cmpi( CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) ) - if const_expr(_persistent): + if _persistent: # Absolute B-base: no cross-iteration state needed. _expert_b_base = expert_idx * arith.constant( _expert_b_stride, index=True @@ -3334,7 +3364,7 @@ def check_c_k_valid_gate(base_k): # For tile_m < 32 (pack_M < _scale_pack_m): shift a_scale i32 so the # correct bytes land at the op_sel positions we use. - if const_expr(pack_M < _scale_pack_m): + if pack_M < _scale_pack_m: _m_off = _mod_pow2(_div_pow2(bx_m, 16), _scale_pack_m) _m_scale_shift_i32 = arith.index_cast( T.i32, _m_off * arith.constant(8, index=True) @@ -3350,18 +3380,18 @@ def _moe_gemm2_then_body(): # ---- X gmem->reg prefetch (match preshuffle GEMM mapping) ---- # Prefer 16B buffer-load (dwordx4). If the per-thread byte count isn't divisible by # 16, fall back to 8B (dwordx2) or 4B (dword) loads. For fp16 we require 16B. - if const_expr(is_f16_a): - if const_expr(bytes_per_thread_x % 16 != 0): + if is_f16_a: + if bytes_per_thread_x % 16 != 0: raise ValueError( f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16" ) x_load_bytes = 16 else: - if const_expr(bytes_per_thread_x % 16 == 0): + if bytes_per_thread_x % 16 == 0: x_load_bytes = 16 - elif const_expr(bytes_per_thread_x % 8 == 0): + elif bytes_per_thread_x % 8 == 0: x_load_bytes = 8 - elif const_expr(bytes_per_thread_x % 4 == 0): + elif bytes_per_thread_x % 4 == 0: x_load_bytes = 4 else: raise ValueError( @@ -3411,7 +3441,7 @@ def load_x(idx_i32): For 16B, keep the fast dwordx4 path. For 8B/4B, use byte offsets. """ - if const_expr(x_load_bytes == 16): + if x_load_bytes == 16: idx_elem = ( idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.index(2)) ) @@ -3437,6 +3467,20 @@ def load_x(idx_i32): ) # decode routed token once (per thread's M-slice) and build a base offset. + # Async X-copy + packed-elem (fp4): the DMA uses only + # _num_dma_loads address slots, not the full num_x_loads. + # Decode only the physically needed rows to skip 2 redundant + # VMEM+decode chains per thread. + if use_async_copy and a_elem_vec_pack > 1: + _dma_bytes_pre = 16 + _eff_bytes_pre = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_x_addr_loads = max( + 1, _eff_bytes_pre // (total_threads * _dma_bytes_pre) + ) + else: + _num_x_addr_loads = num_x_loads x_row_base_div4 = [] x_col_local_i32 = [] x_row_local = [] @@ -3445,22 +3489,26 @@ def load_x(idx_i32): x_row_local.append(row_local) x_col_local_i32.append(col_local_i32) - sorted_row_i = bx_m + row_local - fused_i = buffer_ops.buffer_load( - sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32 - ) - t_i32 = arith.andi(fused_i, mask24) - s_i32 = arith.shrui(fused_i, arith.constant(24)) - - t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) - s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) - ts_valid = arith.andi(t_valid, s_valid) - t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) - s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) - row_ts_i32 = t_safe * topk_i32 + s_safe - row_ts_idx = arith.index_cast(ir.IndexType.get(), row_ts_i32) - - x_row_base_div4.append(row_ts_idx * c_k_div4) + if i < _num_x_addr_loads: + sorted_row_i = bx_m + row_local + fused_i = buffer_ops.buffer_load( + sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32 + ) + t_i32 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(T.index, row_ts_i32) + + x_row_base_div4.append(row_ts_idx * c_k_div4) + else: + # Placeholder: not used in async DMA path. + x_row_base_div4.append(arith.index(0)) def load_x_tile(base_k): base_k_div4 = _div_pow2( @@ -3473,9 +3521,9 @@ def load_x_tile(base_k): idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i] x_vec = load_x(idx_i32) - if const_expr(x_load_bytes == 16): + if x_load_bytes == 16: parts.append(vector.bitcast(vec4_i32, x_vec)) - elif const_expr(x_load_bytes == 8): + elif x_load_bytes == 8: parts.append(vector.bitcast(vec2_i32, x_vec)) else: parts.append(vector.bitcast(vec1_i32, x_vec)) @@ -3503,7 +3551,7 @@ def load_x_tile(base_k): by_n = by * arith.constant(tile_n, index=True) - if const_expr(pack_N < _scale_pack_n): + if pack_N < _scale_pack_n: _global_n_base = expert_off_idx + by_n + n_tile_base _n_off = _mod_pow2(_div_pow2(_global_n_base, 16), _scale_pack_n) _n_scale_shift_i32 = arith.index_cast( @@ -3513,13 +3561,8 @@ def load_x_tile(base_k): _n_scale_shift_i32 = None n_intra_list = [None] * num_acc_n n_blk_list = [None] * num_acc_n - col_g_list = [None] * num_acc_n for i in range_constexpr(num_acc_n): offset = i * 16 - col_g = by_n + n_tile_base - col_g = _div_pow2(col_g, 2) + offset - col_g = col_g + lane_mod_16 - col_g_list[i] = col_g c_offset = arith.constant(offset, index=True) global_n = by_n + n_tile_base + c_offset + lane_mod_16 n_blk_list[i] = _div_pow2(global_n, 16) @@ -3533,19 +3576,18 @@ def load_x_tile(base_k): m_repeat_packed = m_repeat // pack_M num_acc_n_packed = num_acc_n // pack_N - _K_per_ku_s2 = tile_k // k_unroll - _pad_k_elems_s2 = (inter_dim_pad % tile_k) if inter_dim_pad > 0 else 0 - _pad_ku_skip_s2 = _pad_k_elems_s2 // _K_per_ku_s2 - _tail_ku_s2 = k_unroll - _pad_ku_skip_s2 - _tail_ku_packed_s2 = ( - (_tail_ku_s2 + pack_K - 1) // pack_K - if _pad_ku_skip_s2 > 0 - else None - ) - # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- - def load_b_packs_k64(base_k, ku: int, ni: int): - """Load one K64-byte B micro-step: single 16B load, split into 2x i64.""" + def load_b_packs_k64( + base_k, ku: int, ni: int, *, n_blk_p=None, n_intra_p=None + ): + """Load one K64-byte B micro-step: single 16B load, split into 2x i64. + + `n_blk_p` / `n_intra_p` allow callers to override the per-N + list. When ``None`` we fall back to the body-level defaults + (``n_blk_list`` / ``n_intra_list``). + """ + _blk = n_blk_p if n_blk_p is not None else n_blk_list + _intra = n_intra_p if n_intra_p is not None else n_intra_list base_k_bytes = base_k * arith.constant( int(b_elem_bytes), index=True ) @@ -3559,10 +3601,10 @@ def load_b_packs_k64(base_k, ku: int, ni: int): # compile-time constants -> shift/mul, no Barrett. idx_pack = ( _expert_b_base - + n_blk_list[ni] * arith.constant(_b_stride_n0, index=True) + + _blk[ni] * arith.constant(_b_stride_n0, index=True) + k0 * arith.constant(_b_stride_k0, index=True) + k1 * arith.constant(_b_stride_klane, index=True) - + n_intra_list[ni] * arith.constant(_b_stride_nlane, index=True) + + _intra[ni] * arith.constant(_b_stride_nlane, index=True) ) vec_elems = kpack_bytes // int(b_elem_bytes) @@ -3575,7 +3617,6 @@ def load_b_packs_k64(base_k, ku: int, ni: int): vec_elems=vec_elems, elem_bytes=b_elem_bytes, offset_in_bytes=(b_elem_bytes == 1), - cache_modifier=b_nt, ) b_i64x2 = vector.bitcast(vec2_i64, b16) b0 = vector.extract( @@ -3586,42 +3627,65 @@ def load_b_packs_k64(base_k, ku: int, ni: int): ) return b0, b1 - def load_b_tile(base_k, ku_limit=k_unroll): + def load_b_tile(base_k, *, n_blk_p=None, n_intra_p=None): b_tile = [] - for ku in range_constexpr(ku_limit): + for ku in range_constexpr(k_unroll): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, ku, ni, + n_blk_p=n_blk_p, n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) return b_tile _b_split_enabled = k_unroll >= 2 - _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll + # Asymmetric b_lo/b_hi split for prod prefill: b_lo loads + # 3/4 of K (extra ku in prologue, overlaps X DMA), b_hi loads + # 1/4 (smaller burst before K-loop frees VMEM slots so the + # first MFMAs don't stall on VMEM issue). + _is_prefill_shape = ( + int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + and bool(use_async_copy) + and a_elem_vec_pack > 1 + ) + if _b_split_enabled and _is_prefill_shape and k_unroll >= 4: + _b_split_ku = k_unroll - 1 + else: + _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll - def load_b_tile_lo(base_k): + def load_b_tile_lo(base_k, *, n_blk_p=None, n_intra_p=None): """Load first half of B tile (ku < _b_split_ku).""" b_tile = [] for ku in range_constexpr(_b_split_ku): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, ku, ni, + n_blk_p=n_blk_p, n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) return b_tile - def load_b_tile_hi(base_k): + def load_b_tile_hi(base_k, *, n_blk_p=None, n_intra_p=None): """Load second half of B tile (ku >= _b_split_ku).""" b_tile = [] for ku in range_constexpr(_b_split_ku, k_unroll): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, ku, ni, + n_blk_p=n_blk_p, n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) @@ -3641,7 +3705,7 @@ def load_scale(arg_scale, rsrc, scale_info, ku, mni): return vector.from_elements(T.vec(1, T.i32), [s]) def _apply_k_shift(scale_vec, k_shift_bits): - if const_expr(k_shift_bits > 0): + if k_shift_bits > 0: val = vector.extract( scale_vec, static_position=[0], dynamic_position=[] ) @@ -3649,11 +3713,14 @@ def _apply_k_shift(scale_vec, k_shift_bits): return vector.from_elements(T.vec(1, T.i32), [val]) return scale_vec - def load_b_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def load_b_scale_tile(base_k, k_shift_bits=0, *, by_n_p=None): + """Load B-scale tile. + + `by_n_p` overrides the body-default `by_n`. + """ + _byn = by_n_p if by_n_p is not None else by_n b_scale_tile = [] - for ku in range_constexpr(ku_packed_limit): + for ku in range_constexpr(k_unroll_packed): for ni in range_constexpr(num_acc_n_packed): scale = load_scale( arg_scale_w, @@ -3663,7 +3730,7 @@ def load_b_scale_tile( ni + _div_pow2( _div_pow2( - expert_off_idx + by_n + n_tile_base, + expert_off_idx + _byn + n_tile_base, _scale_pack_n, ), 16, @@ -3673,11 +3740,9 @@ def load_b_scale_tile( b_scale_tile.append(scale) return b_scale_tile - def load_a_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def load_a_scale_tile(base_k, k_shift_bits=0): a_scale_tile = [] - for ku in range_constexpr(ku_packed_limit): + for ku in range_constexpr(k_unroll_packed): for mi in range_constexpr(m_repeat_packed): scale = load_scale( arg_scale_x, @@ -3690,55 +3755,47 @@ def load_a_scale_tile( a_scale_tile.append(scale) return a_scale_tile - def prefetch_ab_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def prefetch_ab_scale_tile(base_k, k_shift_bits=0, *, by_n_p=None): return [ - load_a_scale_tile( - base_k, k_shift_bits, ku_packed_limit=ku_packed_limit - ), - load_b_scale_tile( - base_k, k_shift_bits, ku_packed_limit=ku_packed_limit - ), + load_a_scale_tile(base_k, k_shift_bits), + load_b_scale_tile(base_k, k_shift_bits, by_n_p=by_n_p), ] vec8_x = T.vec(vec8_elems, x_elem) vec4_x_lds = T.vec(vec4_elems, x_elem) - # ---- Pipeline helpers: store X tile to LDS (unused in DMA path) ---- - _lds_base_zero = arith.index(0) - - def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + # ---- Pipeline helpers: store X tile to LDS with ping-pong base ---- + def store_x_tile_to_lds(vec_x_in_parts, lds_base): for i in range_constexpr(num_x_loads): row_local = x_row_local[i] col_local_i32 = x_col_local_i32[i] - if const_expr(x_load_bytes == 16): + if x_load_bytes == 16: lds_store_16b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec16_ty=vec16_x, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x4=vec_x_in_parts[i], elem_bytes=elem_bytes, ) - elif const_expr(x_load_bytes == 8): + elif x_load_bytes == 8: lds_store_8b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec8_ty=vec8_x, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x2=vec_x_in_parts[i], elem_bytes=elem_bytes, ) @@ -3746,20 +3803,88 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): lds_store_4b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec4_ty=vec4_x_lds, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x1=vec_x_in_parts[i], elem_bytes=elem_bytes, ) + # Async X-copy (buffer_load_lds): DMA X tile gmem -> LDS + # bypassing per-thread reg staging. Mirrors stage1 but uses + # stage2's single lds_x memref with ping/pong offsets. + if use_async_copy: + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + _c_a_elem_bytes_dma = arith.constant( + int(a_elem_bytes), index=True + ) + _c_wave_dma_bytes = arith.constant( + _wave_size * _dma_bytes, index=True + ) + + def dma_x_tile_to_lds(base_k, lds_base): + c4_idx = arith.index(4) + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(elem_bytes), index=True) + ) / arith.index(4) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if i == 0: + lds_addr = ( + memref.extract_aligned_pointer_as_index(lds_x) + + lds_base * _c_a_elem_bytes_dma + + wave_id * _c_wave_dma_bytes + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_base): + dma_x_tile_to_lds(base_k, lds_base) + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- - def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_base): + # Swizzle in bytes, then convert to element offset for memref indexing. col_base_swz_bytes = swizzle_xor16( curr_row_a_lds, col_base, k_blocks16 ) @@ -3768,8 +3893,10 @@ def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): if elem_bytes == 1 else (col_base_swz_bytes / arith.index(2)) ) + # Pass as list so layout_utils.crd2idx uses static arith path idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) - loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + idx_a16 = idx_a16 + lds_base + loaded_a16 = vector.load_op(vec16_x, lds_x, [idx_a16]) a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) a0 = vector.extract( a_i64x2, static_position=[0], dynamic_position=[] @@ -3782,7 +3909,7 @@ def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): def compute_tile( acc_in, b_tile_in, - lds_buffer, + lds_base, a_scale=None, b_scale=None, *, @@ -3790,9 +3917,17 @@ def compute_tile( a0_prefetch=None, a1_prefetch=None, b_hi_loader=None, - ku_count=k_unroll, + n_scale_shift_p=None, ): - if const_expr(b_hi_loader is not None): + # n_scale_shift_p overrides body-default `_n_scale_shift_i32` + # so the K-outer/N-inner pipeline can apply the right shift + # for each ni slot. + _nss = ( + n_scale_shift_p + if n_scale_shift_p is not None + else _n_scale_shift_i32 + ) + if b_hi_loader is not None: b_tile_full = [None] * k_unroll for i in range_constexpr(_b_split_ku): b_tile_full[i] = b_tile_in[i] @@ -3803,36 +3938,64 @@ def compute_tile( epilogue_pf = None bias = None - if const_expr(prefetch_epilogue): - if const_expr(enable_bias): + if prefetch_epilogue: + if enable_bias: bias = [] for ni in range_constexpr(num_acc_n): global_n = by_n + n_tile_base + ni * 16 + lane_mod_16 bias_offset = expert_off_idx + global_n bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) tw_pf = None - if const_expr(doweight_stage2): + if doweight_stage2: + # When lds_tw is staged in the prologue (prod + # prefill default), prefetch topk-weight tile via + # vec4 LDS reads (LDS ~10cyc vs VMEM ~500cyc), + # saving m_repeat VMEM ops from compute_tile. + # Falls back to vec4 VMEM if lds_tw unavailable. tw_pf = [] lane_div_16_mul4_pf = lane_div_16 * arith.index(4) - ii_idx_list_pf = [ - arith.constant(ii, index=True) for ii in range(4) - ] - for mi in range_constexpr(m_repeat): - mi_base_pf = arith.constant(mi * 16, index=True) - for ii in range_constexpr(4): - row_off_pf = ( - lane_div_16_mul4_pf + ii_idx_list_pf[ii] + _vec4_f32_pf = T.vec(4, f32) + if lds_tw is not None: + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant( + mi * 16, index=True ) - row_in_tile_pf = mi_base_pf + row_off_pf - sorted_row_pf = bx_m + row_in_tile_pf - tw_pf.append( - buffer_ops.buffer_load( - sorted_w_rsrc, - sorted_row_pf, - vec_width=1, - dtype=f32, + lds_row_pf = ( + mi_base_pf + lane_div_16_mul4_pf + ) + tw_v4 = vector.load_op( + _vec4_f32_pf, lds_tw, [lds_row_pf] + ) + for ii in range_constexpr(4): + tw_pf.append( + vector.extract( + tw_v4, + static_position=[ii], + dynamic_position=[], + ) ) + else: + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant( + mi * 16, index=True + ) + base_row_pf = ( + bx_m + mi_base_pf + lane_div_16_mul4_pf + ) + tw_v4 = buffer_ops.buffer_load( + sorted_w_rsrc, + base_row_pf, + vec_width=4, + dtype=f32, ) + for ii in range_constexpr(4): + tw_pf.append( + vector.extract( + tw_v4, + static_position=[ii], + dynamic_position=[], + ) + ) epilogue_pf = (None, tw_pf, bias) c0_i64 = arith.constant(0, type=T.i64) @@ -3849,12 +4012,31 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): _pack_K_shift = (pack_K - 1).bit_length() _pack_K_mask = pack_K - 1 - if const_expr(b_hi_loader is not None): + # Emit AMDGPU `disable_xdl_arb_stall` before b_hi prefetch + # + s_setprio(1) so the high-prio K-loop drains MFMAs back- + # to-back on gfx950 (XDL arbiter normally throttles MFMA + # issue against sibling waves on the same SIMD). + _xdl_arb_hint = ( + int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + and bool(use_async_copy) + and a_elem_vec_pack > 1 + ) + if _xdl_arb_hint: + rocdl.disable_xdl_arb_stall() + + if b_hi_loader is not None: _b_hi = b_hi_loader() for _bhi_i in range_constexpr(len(_b_hi)): b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] - for k_idx in range_constexpr(ku_count): + # Raise wave priority around MFMA-heavy compute so this wave + # is favored on the EU vs other waves doing async copies. + # Mirrors stage1's pattern (lines ~1410/1438). + rocdl.s_setprio(1) + + for k_idx in range_constexpr(k_unroll): ku128 = k_idx >> _pack_K_shift ikxdl = k_idx & _pack_K_mask @@ -3867,10 +4049,52 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): a_scale_val = vector.extract( a_scale_i32, static_position=[0], dynamic_position=[] ) - if const_expr(_m_scale_shift_i32 is not None): + if _m_scale_shift_i32 is not None: a_scale_val = arith.shrui( a_scale_val, _m_scale_shift_i32 ) + # Hoist A-side LDS loads + a128 pack OUT of the + # ni loop (addr depends only on k_idx,mi,imxdl). + # Previously emitted num_acc_n_packed times per + # (mi,imxdl); hoisting removes redundant IR and + # groups LDS reads at the head of the mi block. + a128_list = [None] * pack_M + for imxdl in range_constexpr(pack_M): + col_base0 = col_base + mi_idx = mi * pack_M + imxdl + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row_a_lds = row_a_lds + mi_val + + if ( + (a0_prefetch is not None) + and (k_idx == 0) + and (mi_idx == 0) + ): + a0, a1 = a0_prefetch + elif ( + (a1_prefetch is not None) + and (k_idx == 1) + and (mi_idx == 0) + ): + a0, a1 = a1_prefetch + else: + a0, a1 = lds_load_packs_k64( + curr_row_a_lds, col_base0, lds_base + ) + + if is_f8_a: + col_base1 = col_base + 64 + a2, a3 = lds_load_packs_k64( + curr_row_a_lds, col_base1, lds_base + ) + a128_list[imxdl] = pack_i64x4_to_i32x8( + a0, a1, a2, a3 + ) + else: + a128_list[imxdl] = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + for ni in range_constexpr(num_acc_n_packed): b_scale_i32 = b_scale[ku128 * num_acc_n_packed + ni] b_scale_val = vector.extract( @@ -3878,55 +4102,37 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): static_position=[0], dynamic_position=[], ) - if const_expr(_n_scale_shift_i32 is not None): + if _nss is not None: b_scale_val = arith.shrui( - b_scale_val, _n_scale_shift_i32 + b_scale_val, _nss + ) + + # Hoist B-pack i64x4_to_i32x8 outside the + # imxdl loop — b128 is a function only of + # (ni, inxdl), not (mi, imxdl). Build once per + # (ni, inxdl) and reuse across all mi/imxdl. + b128_list = [None] * pack_N + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + b0 = b_packs0[ni_idx] + b1 = b_packs1[ni_idx] + b128_list[inxdl] = pack_i64x4_to_i32x8( + b0, b1, c0_i64, c0_i64 ) for imxdl in range_constexpr(pack_M): - col_base0 = col_base mi_idx = mi * pack_M + imxdl - mi_val = arith.constant(mi_idx * 16, index=True) - curr_row_a_lds = row_a_lds + mi_val - - if const_expr( - (a0_prefetch is not None) - and (k_idx == 0) - and (mi_idx == 0) - ): - a0, a1 = a0_prefetch - elif const_expr( - (a1_prefetch is not None) - and (k_idx == 1) - and (mi_idx == 0) - ): - a0, a1 = a1_prefetch - else: - a0, a1 = lds_load_packs_k64( - curr_row_a_lds, col_base0, lds_buffer - ) - - if const_expr(is_f8_a): - col_base1 = col_base + 64 - a2, a3 = lds_load_packs_k64( - curr_row_a_lds, col_base1, lds_buffer - ) - a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) - else: - a128 = pack_i64x4_to_i32x8( - a0, a1, c0_i64, c0_i64 - ) + a128 = a128_list[imxdl] for inxdl in range_constexpr(pack_N): ni_idx = ni * pack_N + inxdl - b0 = b_packs0[ni_idx] - b1 = b_packs1[ni_idx] - b128 = pack_i64x4_to_i32x8( - b0, b1, c0_i64, c0_i64 - ) + b128 = b128_list[inxdl] acc_idx = mi_idx * num_acc_n + ni_idx + # Removed per-MFMA sched_barrier(0): with + # s_setprio(1) around the MFMA loop, LLVM can + # interleave MFMAs with next b_hi VMEM+DSREADs. acc_list[acc_idx] = ( rocdl.mfma_scale_f32_16x16x128_f8f6f4( mfma_res_ty, @@ -3944,75 +4150,73 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): ) ) + # Keep s_setprio(1) across the post-compute barrier and + # the VMEM-heavy prefetch_x/scale phase so this wave stays + # favored vs sibling WGs in MFMA loops. Next compute_tile's + # setprio(1) becomes a no-op; the cshuffle epilog handles + # its own prio (bumps to 3, drops to 0 at end). return acc_list, epilogue_pf # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- - # ---- Async DMA: GMEM -> LDS (bypasses VGPR, like stage1) ---- - _dma_bytes = 16 - _wave_size = 64 - _eff_bytes_per_buffer = ( - int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) - ) - _num_dma_loads = max( - 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) - ) - - def dma_x_tile_to_lds(base_k, lds_buffer): - c4_idx = arith.index(4) - base_k_div4 = _div_pow2( - _div_pow2(base_k, int(a_elem_vec_pack)) - * arith.constant(int(a_elem_bytes), index=True), - 4, - ) - - lds_ptr_i64 = None - for i in range_constexpr(_num_dma_loads): - row_local_i = x_row_local[i] - col_local_i32_i = x_col_local_i32[i] - col_local_sw = swizzle_xor16( - row_local_i, col_local_i32_i * c4_idx, k_blocks16 - ) - row_k_dw = x_row_base_div4[i] + base_k_div4 - global_byte_idx = row_k_dw * c4_idx + col_local_sw - global_offset = arith.index_cast(T.i32, global_byte_idx) - - if const_expr(i == 0): - lds_addr = memref.extract_aligned_pointer_as_index( - lds_buffer - ) + wave_id * arith.constant( - _wave_size * _dma_bytes, index=True - ) - lds_ptr_i64 = rocdl.readfirstlane( - T.i64, arith.index_cast(T.i64, lds_addr) - ) - else: - lds_ptr_i64 = lds_ptr_i64 + arith.constant( - total_threads * _dma_bytes, type=T.i64 - ) - - lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") - lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) - - rocdl.raw_ptr_buffer_load_lds( - x_rsrc, - lds_ptr, - arith.constant(_dma_bytes, type=T.i32), - global_offset, - arith.constant(0, type=T.i32), - arith.constant(0, type=T.i32), - arith.constant(0, type=T.i32), - ) - - def prefetch_x_to_lds(base_k, lds_buffer): - dma_x_tile_to_lds(base_k, lds_buffer) + lds_tile_elems = arith.constant(tile_m * lds_stride, index=True) + lds_base_cur = arith.index(0) + lds_base_nxt = lds_tile_elems rocdl.sched_barrier(0) def hot_loop_scheduler(): + # - MFMA group size per "slot": num_acc_n + # - Total MFMA per tile: (2*K32 per K64) * k_unroll * m_repeat * num_acc_n + # - We emit (mfma_group + dsrd + mfma_group) per scheduler iteration. + mfma_group = num_acc_n + mfma_total = (k_unroll * 2) * m_repeat * mfma_group + mfma_per_iter = 2 * mfma_group + sche_iters = ( + 0 if mfma_per_iter == 0 else (mfma_total // mfma_per_iter) + ) + + rocdl.sched_dsrd(2) + rocdl.sched_mfma(1) + if tile_m == 16: + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + if tile_m == 16: + rocdl.sched_vmem(1) + if num_acc_n < 4: + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if tile_m == 16: + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if tile_m == 16: + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + + # DS-write hints (match total A LDS-store micro-ops/thread): + # async path's raw_ptr_buffer_load_lds writes LDS inside the + # buffer-load, so there is no register-mediated ds_write. + # Legacy dswr_tail=num_x_loads overcounts; zero it for async. + if use_async_copy: + dswr_tail = 0 + else: + dswr_tail = num_x_loads + if dswr_tail > sche_iters: + dswr_tail = sche_iters + dswr_start = sche_iters - dswr_tail + + for sche_i in range_constexpr(sche_iters): + rocdl.sched_vmem(1) + rocdl.sched_mfma(mfma_group) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(mfma_group) + if dswr_tail > 0 and sche_i >= dswr_start - 1: + rocdl.sched_dswr(1) + rocdl.sched_barrier(0) def _k_shift_bits(k_py): - if const_expr(pack_K >= _scale_pack_k): + if pack_K >= _scale_pack_k: return 0 return ((k_py // 128) % _scale_pack_k) * _scale_pack_m * 8 @@ -4021,43 +4225,125 @@ def _k_base(k_py): # Preload sorted_idx into lds_tid for epilogue precompute_row # (N-independent; placed before N-tile loop so it's done once per M-tile.) + # + # For the buffer-atomic prefill path (_use_buf_atomic_pre), + # pre-compute `(t & 0xFFFFFF) * row_stride_bytes` and OOB-mask + # invalid rows here in the prologue (overlapped with B/scale/X + # VMEM); the cshuffle precompute_row then needs only an LDS load. + _row_stride_bytes_pre = int(model_dim) * int(out_elem_bytes) + _use_buf_atomic_pre = bool(accumulate) and ( + _row_stride_bytes_pre <= 16384 + ) _c_tile_m_idx = arith.constant(tile_m, index=True) _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) - _if_tid = scf.IfOp(_tid_in_range) - with ir.InsertionPoint(_if_tid.then_block): - _tid_row = bx_m + tx - _tid_val = buffer_ops.buffer_load( - sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 - ) - _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) - vector.store(_tid_vec1, lds_tid, [tx]) - scf.YieldOp([]) + # Defer lds_tid prologue (sorted_idx/sorted_w VMEM + LDS + # stores) to AFTER the long-latency X-DMA + B/scale loads, + # hiding lds_tid's scalar VMEM behind the dominant VMEM phase. + # lds_tid is consumed only by the cshuffle epilogue; the + # post-prologue barrier provides the needed LDS visibility. + _r216_defer_tid = bool( + _r139_xdma_first + and use_async_copy + and int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + ) - gpu.barrier() + def _emit_tid_lds_prologue(): + # Merge lds_tid + lds_tw prologue store guards into a + # single scf.IfOp (same `tx pong. + # Prologue -- B-first. k0 = arith.index(0) - if const_expr(_b_split_enabled): - b_cur = load_b_tile_lo(k0) + k0_bk = k0 + if _r139_xdma_first: + prefetch_x_to_lds(k0, lds_base_cur) + rocdl.sched_barrier(0) + if _b_split_enabled: + b_cur = load_b_tile_lo(k0_bk) else: - b_cur = load_b_tile(k0) + b_cur = load_b_tile(k0_bk) a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( _k_base(0), _k_shift_bits(0) ) rocdl.sched_barrier(0) - prefetch_x_to_lds(k0, lds_x_pong) - rocdl.s_waitcnt(0) + if not _r139_xdma_first: + if use_async_copy: + prefetch_x_to_lds(k0, lds_base_cur) + else: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_base_cur) + elif not use_async_copy: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_base_cur) + if _r216_defer_tid: + # Deferred emit: see _r216_defer_tid above for rationale. + _emit_tid_lds_prologue() + rocdl.sched_barrier(0) gpu.barrier() acc = [acc_init] * num_acc_n * m_repeat + lds_base_pong = lds_base_cur + lds_base_ping = lds_base_nxt - # Cross-tile A0+A1 LDS prefetch from pong buffer. + # Cross-tile A0+A1 LDS prefetch: issue both ds_reads back-to-back + # (??2) so LDS bandwidth is fully utilized and the second read + # completes during MFMA #1/#2 execution. a0_prefetch_pong = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_pong + row_a_lds, col_offset_base, lds_base_pong ) _a1_col_base = col_offset_base + 128 // a_elem_vec_pack a1_prefetch_pong = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_pong) if pack_K >= 2 else None ) @@ -4067,16 +4353,17 @@ def _k_base(k_py): # IMPORTANT: for odd number of K tiles, leave **1** tail tile; for even, leave **2**. # Otherwise the 2-tile tail below would double-count the last tile when num_tiles is odd # (e.g. inter_dim=192, tile_k=64 -> 3 tiles). - num_k_tiles_py = int(inter_dim) // int(tile_k) + num_k_tiles_py = _num_k_tiles_per_batch odd_k_tiles = (num_k_tiles_py % 2) == 1 tail_tiles = 1 if odd_k_tiles else 2 k_main2_py = (num_k_tiles_py - tail_tiles) * int(tile_k) - if const_expr(k_main2_py < 0): + if k_main2_py < 0: k_main2_py = 0 c2_tile_k = arith.constant(tile_k * 2, index=True) b_pong = b_cur - k0_pong_bk = k0 + # b_pong/b_hi loader uses packed-byte K units (see note above). + k0_pong_bk = k0_bk # Only emit the scf.for when there are actually iterations to run. # When k_main2_py == 0 the loop body is empty; emitting an scf.for @@ -4086,27 +4373,32 @@ def _make_b_hi_loader(base_k): """Create a b_hi_loader callable for a given base_k.""" return lambda _bk=base_k: load_b_tile_hi(_bk) - if const_expr(k_main2_py > 0): + if k_main2_py > 0: for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): - rocdl.sched_barrier(0) k_iv = arith.index(k_iv_py) next_k1 = k_iv + tile_k next_k1_bk = next_k1 // 2 - # DMA X(next_k1) -> ping (non-blocking, overlaps with compute) - prefetch_x_to_lds(next_k1, lds_x_ping) + # Issue scales BEFORE B-VMEM in the steady-state K-loop: + # scales are 1-dword (low VMEM latency) and feed + # `vector.extract` at the head of compute_tile's MFMA + # loop, so landing them first hides B's longer loads. + if use_async_copy: + prefetch_x_to_lds(next_k1, lds_base_ping) + else: + x_regs_ping = load_x_tile(next_k1) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(next_k1), _k_shift_bits(next_k1) + ) b_ping_lo = ( load_b_tile_lo(next_k1_bk) if _b_split_enabled else load_b_tile(next_k1_bk) ) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(next_k1), _k_shift_bits(next_k1) - ) acc, _ = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4117,16 +4409,17 @@ def _make_b_hi_loader(base_k): else None ), ) - hot_loop_scheduler() - rocdl.s_waitcnt(0) + if not use_async_copy: + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + # hot_loop_scheduler() gpu.barrier() # Cross-tile prefetch for the ping tile we are about to compute. a0_prefetch_ping = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_ping + row_a_lds, col_offset_base, lds_base_ping ) a1_prefetch_ping = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_ping) if pack_K >= 2 else None ) @@ -4134,21 +4427,24 @@ def _make_b_hi_loader(base_k): next_k2 = k_iv + c2_tile_k next_k2_py = k_iv_py + tile_k * 2 next_k2_bk = next_k2 // 2 - # DMA X(next_k2) -> pong (non-blocking, overlaps with compute) - prefetch_x_to_lds(next_k2, lds_x_pong) + # scales-before-B (see above for rationale). + if use_async_copy: + prefetch_x_to_lds(next_k2, lds_base_pong) + else: + x_regs_pong = load_x_tile(next_k2) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(next_k2_py), _k_shift_bits(next_k2_py) + ) b_pong = ( load_b_tile_lo(next_k2_bk) if _b_split_enabled else load_b_tile(next_k2_bk) ) - a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( - _k_base(next_k2_py), _k_shift_bits(next_k2_py) - ) acc, _ = compute_tile( acc, b_ping_lo, - lds_x_ping, + lds_base_ping, a_scale_ping, b_scale_ping, a0_prefetch=a0_prefetch_ping, @@ -4160,25 +4456,27 @@ def _make_b_hi_loader(base_k): ), ) k0_pong_bk = next_k2_bk - hot_loop_scheduler() + if not use_async_copy: + store_x_tile_to_lds(x_regs_pong, lds_base_pong) + # hot_loop_scheduler() gpu.barrier() # Cross-tile prefetch for the next pong tile. a0_prefetch_pong = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_pong + row_a_lds, col_offset_base, lds_base_pong ) a1_prefetch_pong = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_pong) if pack_K >= 2 else None ) - if const_expr(odd_k_tiles): - # Tail: single remaining tile (already in pong buffer). + if odd_k_tiles: + # Tail: single remaining tile (already in `b_cur` / `lds_base_pong`). acc, epilogue_pf = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4187,7 +4485,6 @@ def _make_b_hi_loader(base_k): b_hi_loader=( _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None ), - ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, ) else: @@ -4197,29 +4494,23 @@ def _make_b_hi_loader(base_k): int(inter_dim) + tile_k - 1 ) // tile_k * tile_k - tile_k k_tail1_bk = k_tail1 // 2 - # DMA tail X -> ping - prefetch_x_to_lds(k_tail1, lds_x_ping) - if const_expr(_pad_ku_skip_s2 > 0): - b_ping_lo = load_b_tile(k_tail1_bk, ku_limit=_tail_ku_s2) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(k_tail1_py), - _k_shift_bits(k_tail1_py), - ku_packed_limit=_tail_ku_packed_s2, - ) + if use_async_copy: + prefetch_x_to_lds(k_tail1, lds_base_ping) else: - b_ping_lo = ( - load_b_tile_lo(k_tail1_bk) - if _b_split_enabled - else load_b_tile(k_tail1_bk) - ) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) - ) + x_regs_ping = load_x_tile(k_tail1) + b_ping_lo = ( + load_b_tile_lo(k_tail1_bk) + if _b_split_enabled + else load_b_tile(k_tail1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) + ) acc, _ = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4229,48 +4520,41 @@ def _make_b_hi_loader(base_k): ), ) + if not use_async_copy: + store_x_tile_to_lds(x_regs_ping, lds_base_ping) # hot_loop_scheduler() - rocdl.s_waitcnt(0) gpu.barrier() # Epilogue tile with sw prefetch. a0_prefetch_ping = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_ping + row_a_lds, col_offset_base, lds_base_ping ) a1_prefetch_ping = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) - if pack_K >= 2 and (_pad_ku_skip_s2 == 0 or _tail_ku_s2 >= 2) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_ping) + if pack_K >= 2 else None ) acc, epilogue_pf = compute_tile( acc, b_ping_lo, - lds_x_ping, + lds_base_ping, a_scale_ping, b_scale_ping, a0_prefetch=a0_prefetch_ping, a1_prefetch=a1_prefetch_ping, prefetch_epilogue=True, b_hi_loader=( - None - if _pad_ku_skip_s2 > 0 - else ( - _make_b_hi_loader(k_tail1_bk) - if _b_split_enabled - else None - ) + _make_b_hi_loader(k_tail1_bk) if _b_split_enabled else None ), - ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, ) # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. - sw_pf = None tw_pf = None bias_pf = None - if const_expr(epilogue_pf is not None): - sw_pf, tw_pf, bias_pf = epilogue_pf + if epilogue_pf is not None: + _, tw_pf, bias_pf = epilogue_pf mask24_i32 = arith.constant(0xFFFFFF) topk_i32_v = topk_i32 @@ -4287,7 +4571,7 @@ def atomic_add_f16x2(val_f16x2, byte_off_i32): ) # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). - if const_expr(lds_out is None): + if lds_out is None: raise RuntimeError( "FLIR_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased." ) @@ -4303,7 +4587,7 @@ def atomic_add_f16x2(val_f16x2, byte_off_i32): _llvm_ptr_ty, arg_out ) out_base_i64 = llvm.ptrtoint(T.i64, out_base_ptr) - out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + out_base_idx = arith.index_cast(T.index, out_base_i64) def write_row_to_lds( *, @@ -4316,23 +4600,9 @@ def write_row_to_lds( num_acc_n: int, lds_out, ): - # Match origin/dev_a16w4: rely on sentinel padded rows + hardware OOB behavior. - fused2 = buffer_ops.buffer_load( - sorted_rsrc, row, vec_width=1, dtype=T.i32 - ) - t2 = fused2 & mask24_i32 - s2 = fused2 >> 24 - - t_ok = arith.cmpi(CmpIPredicate.ult, t2, tokens_i32) - s_ok = arith.cmpi(CmpIPredicate.ult, s2, topk_i32_v) - ts_ok = arith.andi(t_ok, s_ok) - t2_safe = arith.select(ts_ok, t2, arith.constant(0)) - s2_safe = arith.select(ts_ok, s2, arith.constant(0)) - t2_safe * topk_i32_v + s2_safe - - if const_expr(doweight_stage2): + if doweight_stage2: tw_idx = (mi * 4) + ii - if const_expr(tw_pf is not None): + if tw_pf is not None: tw = tw_pf[tw_idx] else: tw = buffer_ops.buffer_load( @@ -4345,12 +4615,12 @@ def write_row_to_lds( v = vector.extract( acc[acc_idx], static_position=[ii], dynamic_position=[] ) - if const_expr(is_int8): + if is_int8: v = arith.sitofp(f32, v) - if const_expr(enable_bias): + if enable_bias: v = v + bias_pf[ni] - if const_expr(doweight_stage2): + if doweight_stage2: v = v * tw v_out = arith.trunc_f(out_elem(), v) @@ -4360,7 +4630,45 @@ def write_row_to_lds( vector.store(v1, lds_out, [lds_idx], alignment=2) + # Buffer-atomic eligibility for accumulate=True path: + # the byte-offset must fit in i32. The relevant max is + # max_t * model_dim * out_elem_bytes; we conservatively + # cap row_stride <= 16384 bytes (covers model_dim<=8192 bf16). + _row_stride_bytes_py = int(model_dim) * int(out_elem_bytes) + _use_buf_atomic = bool(accumulate) and ( + _row_stride_bytes_py <= 16384 + ) + _row_stride_bytes_i32 = ( + arith.constant(_row_stride_bytes_py, type=T.i32) + if _use_buf_atomic + else None + ) + _out_elem_bytes_i32 = ( + arith.constant(int(out_elem_bytes), type=T.i32) + if _use_buf_atomic + else None + ) + # 0x7FFF0000 is safely beyond out_rsrc.num_records_bytes + # (~470MB max for prefill DSV3 accumulate=True) and below + # INT32_MAX with room for column offset (≤tile_n*out_elem_bytes). + # HW drops OOB buffer atomics, so routing invalid rows here + # lets us drop the per-mr scf.IfOp in the cshuffle epilogue. + _OOB_safe_i32 = ( + arith.constant(0x7FFF0000, type=T.i32) + if _use_buf_atomic + else None + ) + def precompute_row(*, row_local, row): + # Fast path: lds_tid was prefilled with the per-row byte + # offset (see _use_buf_atomic_pre at the prologue), so this + # is just one ds_read on the per-row critical path. + if _use_buf_atomic: + row_byte_off_i32 = memref.load(lds_tid, [row_local]) + return ( + (None, None, row_byte_off_i32), + None, + ) # Use lds_tid (sorted_idx preloaded to LDS) instead of buffer_load # to avoid extra VMEM round-trips in the epilogue. fused2 = memref.load(lds_tid, [row_local]) @@ -4374,7 +4682,7 @@ def precompute_row(*, row_local, row): t_idx = arith.index_cast(ir.IndexType.get(), t) s_idx = arith.index_cast(ir.IndexType.get(), s) ts_idx = t_idx * arith.constant(topk, index=True) + s_idx - if const_expr(accumulate): + if accumulate: row_byte_base = out_base_idx + t_idx * arith.constant( model_dim * out_elem_bytes, index=True ) @@ -4382,7 +4690,8 @@ def precompute_row(*, row_local, row): row_byte_base = out_base_idx + ts_idx * arith.constant( model_dim * out_elem_bytes, index=True ) - return ((fused2, row_byte_base), row_valid) + row_byte_off_i32 = None + return ((fused2, row_byte_base, row_byte_off_i32), row_valid) def _idx_to_llvm_ptr(idx_val, addr_space=1): """Convert an index-typed byte address to !llvm.ptr.""" @@ -4393,8 +4702,8 @@ def _idx_to_llvm_ptr(idx_val, addr_space=1): return llvm.inttoptr(ptr_ty, i64_raw) def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): - fused, row_byte_base = row_ctx - if const_expr(not bool(accumulate)): + fused, row_byte_base, row_byte_off_i32 = row_ctx + if not bool(accumulate): # ---- 64-bit global store path (avoids i32 offset overflow) ---- col_idx = col_g0 byte_off_col = col_idx * arith.constant( @@ -4409,8 +4718,23 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): alignment=_e_vec * out_elem_bytes, nontemporal=True, ) + elif _use_buf_atomic: + # Buffer-atomic path -> BUFFER_ATOMIC_PK_ADD_BF16. + # out_rsrc keeps the 64-bit base in SGPRs so only the + # 32-bit voffset lives in VGPRs, vs llvm.AtomicRMWOp + # which materializes a full 64-bit pointer per op. + col_i32 = arith.index_cast(T.i32, col_g0) + col_byte_off_i32 = col_i32 * _out_elem_bytes_i32 + byte_off_i32 = row_byte_off_i32 + col_byte_off_i32 + rocdl.raw_ptr_buffer_atomic_fadd( + frag, + out_rsrc, + byte_off_i32, + zero_i32, + zero_i32, + ) else: - # ---- accumulate=True: 64-bit global atomic path ---- + # ---- accumulate=True (large-output): 64-bit global atomic path ---- col_idx = col_g0 byte_off_col = col_idx * arith.constant( out_elem_bytes, index=True @@ -4428,6 +4752,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): ) _e_vec = 2 if accumulate else min(tile_n // 32, 8) + rocdl.s_setprio(3) c_shuffle_epilog( arith=arith, vector=vector, @@ -4453,10 +4778,11 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): precompute_row=precompute_row, store_pair=store_pair, ) + rocdl.s_setprio(0) _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) - if const_expr(_persistent): + if _persistent: # Short-circuit: contiguous tiles are monotonically increasing, # so once bx_m >= num_valid_ids all remaining tiles are invalid. _cur_active = arith.andi(_still_active, blk_valid) @@ -4498,7 +4824,8 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): persist_m, _sort_block_m, _cu_num if _persistent else 0, - xcd_swizzle, + waves_per_eu, + use_async_copy, ) @flyc.jit @@ -4520,12 +4847,10 @@ def launch_mixed_moe_gemm2( stream: fx.Stream, ): _ = _cache_tag - allocator_pong.finalized = False - allocator_ping.finalized = False + allocator.finalized = False ctx = CompilationContext.get_current() with ir.InsertionPoint(ctx.gpu_module_body): - allocator_pong.finalize() - allocator_ping.finalize() + allocator.finalize() n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) _tile_n_idx = arith.constant(tile_n, index=True) @@ -4543,7 +4868,7 @@ def launch_mixed_moe_gemm2( - arith.constant(1, index=True) ) / _c_pm_l - moe_gemm2( + _launcher = moe_gemm2( arg_out, arg_x, arg_w, @@ -4558,7 +4883,21 @@ def launch_mixed_moe_gemm2( i32_n_in, i32_k_in, i32_size_expert_ids_in, - ).launch( + ) + # Stamp rocdl.waves_per_eu on the emitted gpu.func so the backend + # respects the occupancy request. Done here (not at kernel-body time) + # because the gpu.func only exists after the kernel has been emitted. + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + for _op in ctx.gpu_module_body.operations: + if ( + hasattr(_op, "attributes") + and _op.OPERATION_NAME == "gpu.func" + ): + _op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, _wpe + ) + _launcher.launch( grid=(gx, gy, 1), block=(256, 1, 1), stream=stream, diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 719bdcae225..989075ffd2a 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -186,6 +186,24 @@ def get_flydsl_stage2_kernels( **base_params, "persist": True, } + if base_name == "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic": + # Production fp4xfp4 stage2 variant for the EP4 + # DeepSeek prefill shape on MI355X. Adds: + # - use_async_copy=True (async X DMA in prologue + # overlaps with B/scale VMEM) + # - cu_num_mul=3 (persistent grid 3x CU count + # to fill in-flight slack from small per-WG + # M tile counts; cu_num_mul=4 regresses ~2.4% + # on the same shape) + # - waves_per_eu=4 (best on EP4 prefill at + # cu_num_mul=3; wpe=5/6 underperform here) + kernels[f"{base_name}_persist_async_w4_cumul3"] = { + **base_params, + "persist": True, + "use_async_copy": True, + "waves_per_eu": 4, + "cu_num_mul": 3, + } return kernels @@ -378,6 +396,9 @@ def compile_flydsl_moe_stage2( accumulate: bool = True, persist_m: int = 1, sort_block_m: int = 0, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, b_nt: int = 0, model_dim_pad: int = 0, inter_dim_pad: int = 0, @@ -403,10 +424,11 @@ def compile_flydsl_moe_stage2( accumulate=accumulate, persist_m=persist_m, sort_block_m=sort_block_m, - b_nt=b_nt, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad, - xcd_swizzle=xcd_swizzle, enable_bias=enable_bias, ) elif a_dtype == "bf16" and b_dtype == "int4": @@ -1035,6 +1057,9 @@ def flydsl_moe_stage2( sorted_weights: Optional[torch.Tensor] = None, sort_block_m: int = 0, persist: Optional[bool] = None, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, b_nt: int = 0, model_dim_pad: int = 0, inter_dim_pad: int = 0, @@ -1070,8 +1095,12 @@ def flydsl_moe_stage2( out = alloc_fn( (token_num, model_dim), dtype=torch_out_dtype, device=inter_states.device ) - elif accumulate: - out.fill_(0) + # NOTE: when ``accumulate=True`` (atomic mode), the caller is responsible + # for ensuring ``out`` is zero-initialized. In the standard ``fused_moe`` + # dispatch path this is handled by ``moe_sorting_*_fwd`` which already + # zeros ``moe_buf`` via ``moe_buf_set_zero_kernel_2d``, so an extra + # ``out.fill_(0)`` here would be a redundant ~``token_num * model_dim`` + # HBM write (~130us per call at MI355X HBM bw on EP4 prefill shape). dev = inter_states.device flat_a_scale = ( @@ -1166,6 +1195,9 @@ def flydsl_moe_stage2( accumulate=accumulate, persist_m=_persist_m, sort_block_m=sort_block_m, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, b_nt=b_nt, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad, From 83ae8ad77c8c1fff2584c820ac75ff4238045ee2 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Mon, 11 May 2026 06:25:31 +0000 Subject: [PATCH 2/8] compat(flydsl): wrap Python-compile-time conditions with const_expr The strict flydsl in current aiter-main rejects bare `if X:` / `elif X:` over Python-level booleans inside a kernel emit body -- it lowers them into `scf.if` so any Python variable assigned in either branch is trapped inside the scf.if scope and the outer reference raises `NameError: name '' is not defined`. `compile_mixed_moe_gemm1` already wraps every such condition with `const_expr(...)`. The stage2 emit body in the parent commit had ~98 bare-if conditions over compile-time identifiers (`accumulate`, `_persistent`, `doweight_stage2`, `_b_split_enabled`, `is_f4_a`, `use_async_copy`, `_r139_xdma_first`, `_r216_defer_tid`, `_use_buf_atomic_pre`, ...). Wrap them with `const_expr(...)` so the condition is evaluated at Python trace time and the variable binding escapes to the enclosing scope, matching the stage1 convention. No semantic change: only adds `const_expr(...)` around if/elif heads inside compile_mixed_moe_gemm2; DSL/runtime conditions (those whose sub-expressions reference `arith.`, `buffer_ops.`, `rocdl.`, ...) are left as bare `if` so they remain runtime `scf.if`s. Repro (without this commit) on aiter-main runtime: File ".../mixed_moe_gemm_2stage.py", line 3096, in moe_gemm2 shape_lds = fx.make_shape(tile_m, _eff_lds_stride) NameError: name '_eff_lds_stride' is not defined After this commit the kernel compiles, the end-to-end fused_moe perf is unchanged from the parent commit on the legacy flydsl, and is restored on aiter-main flydsl. --- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 196 +++++++++--------- 1 file changed, 98 insertions(+), 98 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 2831ae9d2cd..6abdf0d9040 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -2793,7 +2793,7 @@ def compile_mixed_moe_gemm2( sorting/stage1. Requires sort_block_m % tile_m == 0. """ _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m - if _sort_block_m != tile_m and _sort_block_m % tile_m != 0: + if const_expr(_sort_block_m != tile_m and _sort_block_m % tile_m != 0): raise ValueError( f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" ) @@ -2808,11 +2808,11 @@ def compile_mixed_moe_gemm2( gpu_arch = get_hip_arch() allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") - if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + if const_expr(a_dtype not in ("fp8", "fp16", "int8", "fp4")): raise ValueError( f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" ) - if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + if const_expr(b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4")): raise ValueError( f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" ) @@ -2859,20 +2859,20 @@ def compile_mixed_moe_gemm2( _expert_b_stride = (model_dim // 16) * _b_stride_n0 # K64-byte micro-step: always 64 bytes per `ku`. For fp16, this is 32 elements (2xK16 MFMA). - if (tile_k_bytes % 64) != 0: + if const_expr((tile_k_bytes % 64) != 0): raise ValueError( f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " f"(tile_k={tile_k}, elem_bytes={a_elem_bytes})" ) out_s = str(out_dtype).strip().lower() - if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): + if const_expr(out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float")): raise ValueError( f"out_dtype must be 'f16', 'bf16', or 'f32', got {out_dtype!r}" ) out_is_f32 = out_s in ("f32", "fp32", "float") out_is_bf16 = out_s in ("bf16", "bfloat16") - if (not bool(accumulate)) and out_is_f32: + if const_expr((not bool(accumulate)) and out_is_f32): raise ValueError( "compile_moe_gemm2(accumulate=False) only supports out_dtype in {'f16','bf16'}" ) @@ -2881,23 +2881,23 @@ def compile_mixed_moe_gemm2( is_int8 = False mfma_i32_k32 = None - if is_int8: + if const_expr(is_int8): mfma_i32_k32 = getattr(rocdl, "mfma_i32_16x16x32i8", None) or getattr( rocdl, "mfma_i32_16x16x32_i8", None ) - if mfma_i32_k32 is None: + if const_expr(mfma_i32_k32 is None): raise AttributeError( "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` " "(or `rocdl.mfma_i32_16x16x32_i8`)." ) def _x_elem_type(): - if is_f4_b: + if const_expr(is_f4_b): return T.f8 if is_f8_a else T.i8 return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) def _w_elem_type(): - if is_f4_b: + if const_expr(is_f4_b): return T.i8 return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) @@ -2906,7 +2906,7 @@ def _scale_elem_type(): total_threads = 256 bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) - if bytes_x_per_tile % total_threads != 0: + if const_expr(bytes_x_per_tile % total_threads != 0): raise ValueError( "tile_m*tile_k*elem_bytes must be divisible by " f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={a_elem_bytes}" @@ -2923,17 +2923,17 @@ def _scale_elem_type(): pad_k = 0 if _use_lds128 else 8 lds_stride = tile_k + pad_k - if out_is_f32: + if const_expr(out_is_f32): # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. _use_cshuffle_epilog = ( False if use_cshuffle_epilog is None else bool(use_cshuffle_epilog) ) - if _use_cshuffle_epilog: + if const_expr(_use_cshuffle_epilog): raise ValueError( "out_dtype='f32' does not support CShuffle epilogue (set use_cshuffle_epilog=False)." ) else: - if use_cshuffle_epilog is None: + if const_expr(use_cshuffle_epilog is None): _use_cshuffle_epilog = os.environ.get("FLIR_MOE_STAGE2_CSHUFFLE", "1") in ( "1", "true", @@ -2943,7 +2943,7 @@ def _scale_elem_type(): ) else: _use_cshuffle_epilog = bool(use_cshuffle_epilog) - if not _use_cshuffle_epilog: + if const_expr(not _use_cshuffle_epilog): raise ValueError( "stage2 f16 output currently requires CShuffle epilogue (FLIR_MOE_STAGE2_CSHUFFLE=1)." ) @@ -2967,9 +2967,9 @@ def _load_bias_scalar(bias_rsrc, offset): # identical to legacy; >1 launches more persistent CTAs each covering # fewer M tiles. cu_num_mul=3 is the EP4 prefill sweet spot on MI355X # (cu_num_mul=4 regresses ~2.4%). - if not isinstance(cu_num_mul, int) or cu_num_mul < 1: + if const_expr(not isinstance(cu_num_mul, int) or cu_num_mul < 1): raise ValueError(f"cu_num_mul must be int >= 1, got {cu_num_mul}") - if _persistent: + if const_expr(_persistent): from aiter.jit.utils.chip_info import get_cu_num _cu_num = get_cu_num() * int(cu_num_mul) @@ -2978,7 +2978,7 @@ def _load_bias_scalar(bias_rsrc, offset): _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" _wpe_tag = f"_w{waves_per_eu}" if waves_per_eu is not None else "" - if waves_per_eu is not None and not (1 <= int(waves_per_eu) <= 10): + if const_expr(waves_per_eu is not None and not (1 <= int(waves_per_eu) <= 10)): raise ValueError( f"waves_per_eu must be in [1, 10] or None, got {waves_per_eu}" ) @@ -3014,7 +3014,7 @@ def x_lds_elem(): lds_alloc_offset = allocator._align(allocator.ptr, 16) allocator.ptr = lds_alloc_offset + lds_alloc_bytes - if True: + if const_expr(True): @flyc.kernel def moe_gemm2( @@ -3086,7 +3086,7 @@ def check_c_k_valid_gate(base_k): # DMA (buffer_load_lds) writes packed bytes into LDS, so the # effective stride is halved. Keep legacy (non-async) path # bit-identical by using lds_stride directly. - if use_async_copy and a_elem_vec_pack > 1: + 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 else: @@ -3175,7 +3175,7 @@ def check_c_k_valid_gate(base_k): out_nbytes_idx = ( tokens_in * n_in * arith.constant(out_elem_bytes, index=True) ) - if not bool(accumulate): + if const_expr(not bool(accumulate)): out_nbytes_idx = ( tokens_in * arith.index(topk) @@ -3204,10 +3204,10 @@ def check_c_k_valid_gate(base_k): num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) # fp16 path ignores scales completely (implicit scale=1.0). - if is_f16_a: + if const_expr(is_f16_a): sx_rsrc = None else: - if is_f4_a: + if const_expr(is_f4_a): # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. # Caller must pre-scatter a2_scale via moe_mxfp4_sort. kblk = _div_pow2(k_in, 32) @@ -3224,7 +3224,7 @@ def check_c_k_valid_gate(base_k): arg_scale_x, max_size=False, num_records_bytes=sx_nbytes_i32 ) - if is_f16_b: + if const_expr(is_f16_b): sw_rsrc = None else: # Weight microscale buffer (packed i32 holding e8m0 bytes). @@ -3275,7 +3275,7 @@ def check_c_k_valid_gate(base_k): _c0_p = arith.constant(0, index=True) _c1_p = arith.constant(1, index=True) - if _persistent: + if const_expr(_persistent): # Balanced contiguous M-tile dispatch: each CTA gets # floor(T/CU) or ceil(T/CU) tiles, distributing the remainder # across the first (T%CU) CTAs (avoids the no-work tail band @@ -3314,7 +3314,7 @@ def check_c_k_valid_gate(base_k): _for_ip.__enter__() _mi_p = _for_persist.induction_variable - if _persistent: + if const_expr(_persistent): _still_active = _for_persist.inner_iter_args[0] bx = _persist_start_tile + _mi_p else: @@ -3337,7 +3337,7 @@ def check_c_k_valid_gate(base_k): CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) ) - if _persistent: + if const_expr(_persistent): # Absolute B-base: no cross-iteration state needed. _expert_b_base = expert_idx * arith.constant( _expert_b_stride, index=True @@ -3364,7 +3364,7 @@ def check_c_k_valid_gate(base_k): # For tile_m < 32 (pack_M < _scale_pack_m): shift a_scale i32 so the # correct bytes land at the op_sel positions we use. - if pack_M < _scale_pack_m: + if const_expr(pack_M < _scale_pack_m): _m_off = _mod_pow2(_div_pow2(bx_m, 16), _scale_pack_m) _m_scale_shift_i32 = arith.index_cast( T.i32, _m_off * arith.constant(8, index=True) @@ -3380,18 +3380,18 @@ def _moe_gemm2_then_body(): # ---- X gmem->reg prefetch (match preshuffle GEMM mapping) ---- # Prefer 16B buffer-load (dwordx4). If the per-thread byte count isn't divisible by # 16, fall back to 8B (dwordx2) or 4B (dword) loads. For fp16 we require 16B. - if is_f16_a: - if bytes_per_thread_x % 16 != 0: + if const_expr(is_f16_a): + if const_expr(bytes_per_thread_x % 16 != 0): raise ValueError( f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16" ) x_load_bytes = 16 else: - if bytes_per_thread_x % 16 == 0: + if const_expr(bytes_per_thread_x % 16 == 0): x_load_bytes = 16 - elif bytes_per_thread_x % 8 == 0: + elif const_expr(bytes_per_thread_x % 8 == 0): x_load_bytes = 8 - elif bytes_per_thread_x % 4 == 0: + elif const_expr(bytes_per_thread_x % 4 == 0): x_load_bytes = 4 else: raise ValueError( @@ -3441,7 +3441,7 @@ def load_x(idx_i32): For 16B, keep the fast dwordx4 path. For 8B/4B, use byte offsets. """ - if x_load_bytes == 16: + if const_expr(x_load_bytes == 16): idx_elem = ( idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.index(2)) ) @@ -3471,7 +3471,7 @@ def load_x(idx_i32): # _num_dma_loads address slots, not the full num_x_loads. # Decode only the physically needed rows to skip 2 redundant # VMEM+decode chains per thread. - if use_async_copy and a_elem_vec_pack > 1: + if const_expr(use_async_copy and a_elem_vec_pack > 1): _dma_bytes_pre = 16 _eff_bytes_pre = ( int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) @@ -3489,7 +3489,7 @@ def load_x(idx_i32): x_row_local.append(row_local) x_col_local_i32.append(col_local_i32) - if i < _num_x_addr_loads: + if const_expr(i < _num_x_addr_loads): sorted_row_i = bx_m + row_local fused_i = buffer_ops.buffer_load( sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32 @@ -3521,9 +3521,9 @@ def load_x_tile(base_k): idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i] x_vec = load_x(idx_i32) - if x_load_bytes == 16: + if const_expr(x_load_bytes == 16): parts.append(vector.bitcast(vec4_i32, x_vec)) - elif x_load_bytes == 8: + elif const_expr(x_load_bytes == 8): parts.append(vector.bitcast(vec2_i32, x_vec)) else: parts.append(vector.bitcast(vec1_i32, x_vec)) @@ -3551,7 +3551,7 @@ def load_x_tile(base_k): by_n = by * arith.constant(tile_n, index=True) - if pack_N < _scale_pack_n: + if const_expr(pack_N < _scale_pack_n): _global_n_base = expert_off_idx + by_n + n_tile_base _n_off = _mod_pow2(_div_pow2(_global_n_base, 16), _scale_pack_n) _n_scale_shift_i32 = arith.index_cast( @@ -3654,7 +3654,7 @@ def load_b_tile(base_k, *, n_blk_p=None, n_intra_p=None): and bool(use_async_copy) and a_elem_vec_pack > 1 ) - if _b_split_enabled and _is_prefill_shape and k_unroll >= 4: + if const_expr(_b_split_enabled and _is_prefill_shape and k_unroll >= 4): _b_split_ku = k_unroll - 1 else: _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll @@ -3705,7 +3705,7 @@ def load_scale(arg_scale, rsrc, scale_info, ku, mni): return vector.from_elements(T.vec(1, T.i32), [s]) def _apply_k_shift(scale_vec, k_shift_bits): - if k_shift_bits > 0: + if const_expr(k_shift_bits > 0): val = vector.extract( scale_vec, static_position=[0], dynamic_position=[] ) @@ -3769,7 +3769,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base): for i in range_constexpr(num_x_loads): row_local = x_row_local[i] col_local_i32 = x_col_local_i32[i] - if x_load_bytes == 16: + if const_expr(x_load_bytes == 16): lds_store_16b_xor16( arith, vector, @@ -3784,7 +3784,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base): vec_part_i32x4=vec_x_in_parts[i], elem_bytes=elem_bytes, ) - elif x_load_bytes == 8: + elif const_expr(x_load_bytes == 8): lds_store_8b_xor16( arith, vector, @@ -3818,7 +3818,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base): # Async X-copy (buffer_load_lds): DMA X tile gmem -> LDS # bypassing per-thread reg staging. Mirrors stage1 but uses # stage2's single lds_x memref with ping/pong offsets. - if use_async_copy: + if const_expr(use_async_copy): _dma_bytes = 16 _wave_size = 64 _eff_bytes_per_buffer = ( @@ -3852,7 +3852,7 @@ def dma_x_tile_to_lds(base_k, lds_base): global_byte_idx = row_k_dw * c4_idx + col_local_sw global_offset = arith.index_cast(T.i32, global_byte_idx) - if i == 0: + if const_expr(i == 0): lds_addr = ( memref.extract_aligned_pointer_as_index(lds_x) + lds_base * _c_a_elem_bytes_dma @@ -3927,7 +3927,7 @@ def compute_tile( if n_scale_shift_p is not None else _n_scale_shift_i32 ) - if b_hi_loader is not None: + if const_expr(b_hi_loader is not None): b_tile_full = [None] * k_unroll for i in range_constexpr(_b_split_ku): b_tile_full[i] = b_tile_in[i] @@ -3938,15 +3938,15 @@ def compute_tile( epilogue_pf = None bias = None - if prefetch_epilogue: - if enable_bias: + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): bias = [] for ni in range_constexpr(num_acc_n): global_n = by_n + n_tile_base + ni * 16 + lane_mod_16 bias_offset = expert_off_idx + global_n bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) tw_pf = None - if doweight_stage2: + if const_expr(doweight_stage2): # When lds_tw is staged in the prologue (prod # prefill default), prefetch topk-weight tile via # vec4 LDS reads (LDS ~10cyc vs VMEM ~500cyc), @@ -3955,7 +3955,7 @@ def compute_tile( tw_pf = [] lane_div_16_mul4_pf = lane_div_16 * arith.index(4) _vec4_f32_pf = T.vec(4, f32) - if lds_tw is not None: + if const_expr(lds_tw is not None): for mi in range_constexpr(m_repeat): mi_base_pf = arith.constant( mi * 16, index=True @@ -4023,10 +4023,10 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): and bool(use_async_copy) and a_elem_vec_pack > 1 ) - if _xdl_arb_hint: + if const_expr(_xdl_arb_hint): rocdl.disable_xdl_arb_stall() - if b_hi_loader is not None: + if const_expr(b_hi_loader is not None): _b_hi = b_hi_loader() for _bhi_i in range_constexpr(len(_b_hi)): b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] @@ -4049,7 +4049,7 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): a_scale_val = vector.extract( a_scale_i32, static_position=[0], dynamic_position=[] ) - if _m_scale_shift_i32 is not None: + if const_expr(_m_scale_shift_i32 is not None): a_scale_val = arith.shrui( a_scale_val, _m_scale_shift_i32 ) @@ -4065,13 +4065,13 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): mi_val = arith.constant(mi_idx * 16, index=True) curr_row_a_lds = row_a_lds + mi_val - if ( + if const_expr( (a0_prefetch is not None) and (k_idx == 0) and (mi_idx == 0) ): a0, a1 = a0_prefetch - elif ( + elif const_expr( (a1_prefetch is not None) and (k_idx == 1) and (mi_idx == 0) @@ -4082,7 +4082,7 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): curr_row_a_lds, col_base0, lds_base ) - if is_f8_a: + if const_expr(is_f8_a): col_base1 = col_base + 64 a2, a3 = lds_load_packs_k64( curr_row_a_lds, col_base1, lds_base @@ -4102,7 +4102,7 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): static_position=[0], dynamic_position=[], ) - if _nss is not None: + if const_expr(_nss is not None): b_scale_val = arith.shrui( b_scale_val, _nss ) @@ -4177,19 +4177,19 @@ def hot_loop_scheduler(): rocdl.sched_dsrd(2) rocdl.sched_mfma(1) - if tile_m == 16: + if const_expr(tile_m == 16): rocdl.sched_vmem(1) rocdl.sched_mfma(1) - if tile_m == 16: + if const_expr(tile_m == 16): rocdl.sched_vmem(1) - if num_acc_n < 4: + if const_expr(num_acc_n < 4): rocdl.sched_dsrd(1) rocdl.sched_mfma(1) - if tile_m == 16: + if const_expr(tile_m == 16): rocdl.sched_vmem(1) rocdl.sched_dsrd(1) rocdl.sched_mfma(1) - if tile_m == 16: + if const_expr(tile_m == 16): rocdl.sched_vmem(1) rocdl.sched_mfma(1) @@ -4197,11 +4197,11 @@ def hot_loop_scheduler(): # async path's raw_ptr_buffer_load_lds writes LDS inside the # buffer-load, so there is no register-mediated ds_write. # Legacy dswr_tail=num_x_loads overcounts; zero it for async. - if use_async_copy: + if const_expr(use_async_copy): dswr_tail = 0 else: dswr_tail = num_x_loads - if dswr_tail > sche_iters: + if const_expr(dswr_tail > sche_iters): dswr_tail = sche_iters dswr_start = sche_iters - dswr_tail @@ -4210,13 +4210,13 @@ def hot_loop_scheduler(): rocdl.sched_mfma(mfma_group) rocdl.sched_dsrd(1) rocdl.sched_mfma(mfma_group) - if dswr_tail > 0 and sche_i >= dswr_start - 1: + if const_expr(dswr_tail > 0 and sche_i >= dswr_start - 1): rocdl.sched_dswr(1) rocdl.sched_barrier(0) def _k_shift_bits(k_py): - if pack_K >= _scale_pack_k: + if const_expr(pack_K >= _scale_pack_k): return 0 return ((k_py // 128) % _scale_pack_k) * _scale_pack_m * 8 @@ -4260,11 +4260,11 @@ def _emit_tid_lds_prologue(): _tid_val = buffer_ops.buffer_load( sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 ) - if doweight_stage2: + if const_expr(doweight_stage2): _tw_val_m = buffer_ops.buffer_load( sorted_w_rsrc, _tid_row, vec_width=1, dtype=f32 ) - if _use_buf_atomic_pre: + if const_expr(_use_buf_atomic_pre): # 0x7FFF0000 = OOB constant; HW drops OOB buffer # atomics (see comment at the cshuffle-epilogue use). _t_pre = _tid_val & arith.constant(0xFFFFFF, type=T.i32) @@ -4287,14 +4287,14 @@ def _emit_tid_lds_prologue(): _stored_val = _tid_val _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_stored_val]) vector.store(_tid_vec1, lds_tid, [tx]) - if doweight_stage2: + if const_expr(doweight_stage2): _tw_vec1 = vector.from_elements( T.vec(1, T.f32), [_tw_val_m] ) vector.store(_tw_vec1, lds_tw, [tx], alignment=4) scf.YieldOp([]) - if not _r216_defer_tid: + if const_expr(not _r216_defer_tid): _emit_tid_lds_prologue() # Removed the post-lds_tid gpu.barrier(): lds_tid is consumed @@ -4305,10 +4305,10 @@ def _emit_tid_lds_prologue(): # Prologue -- B-first. k0 = arith.index(0) k0_bk = k0 - if _r139_xdma_first: + if const_expr(_r139_xdma_first): prefetch_x_to_lds(k0, lds_base_cur) rocdl.sched_barrier(0) - if _b_split_enabled: + if const_expr(_b_split_enabled): b_cur = load_b_tile_lo(k0_bk) else: b_cur = load_b_tile(k0_bk) @@ -4316,16 +4316,16 @@ def _emit_tid_lds_prologue(): _k_base(0), _k_shift_bits(0) ) rocdl.sched_barrier(0) - if not _r139_xdma_first: - if use_async_copy: + if const_expr(not _r139_xdma_first): + if const_expr(use_async_copy): prefetch_x_to_lds(k0, lds_base_cur) else: x_regs0 = load_x_tile(k0) store_x_tile_to_lds(x_regs0, lds_base_cur) - elif not use_async_copy: + elif const_expr(not use_async_copy): x_regs0 = load_x_tile(k0) store_x_tile_to_lds(x_regs0, lds_base_cur) - if _r216_defer_tid: + if const_expr(_r216_defer_tid): # Deferred emit: see _r216_defer_tid above for rationale. _emit_tid_lds_prologue() rocdl.sched_barrier(0) @@ -4357,7 +4357,7 @@ def _emit_tid_lds_prologue(): odd_k_tiles = (num_k_tiles_py % 2) == 1 tail_tiles = 1 if odd_k_tiles else 2 k_main2_py = (num_k_tiles_py - tail_tiles) * int(tile_k) - if k_main2_py < 0: + if const_expr(k_main2_py < 0): k_main2_py = 0 c2_tile_k = arith.constant(tile_k * 2, index=True) @@ -4373,7 +4373,7 @@ def _make_b_hi_loader(base_k): """Create a b_hi_loader callable for a given base_k.""" return lambda _bk=base_k: load_b_tile_hi(_bk) - if k_main2_py > 0: + if const_expr(k_main2_py > 0): for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): k_iv = arith.index(k_iv_py) next_k1 = k_iv + tile_k @@ -4382,7 +4382,7 @@ def _make_b_hi_loader(base_k): # scales are 1-dword (low VMEM latency) and feed # `vector.extract` at the head of compute_tile's MFMA # loop, so landing them first hides B's longer loads. - if use_async_copy: + if const_expr(use_async_copy): prefetch_x_to_lds(next_k1, lds_base_ping) else: x_regs_ping = load_x_tile(next_k1) @@ -4409,7 +4409,7 @@ def _make_b_hi_loader(base_k): else None ), ) - if not use_async_copy: + if const_expr(not use_async_copy): store_x_tile_to_lds(x_regs_ping, lds_base_ping) # hot_loop_scheduler() gpu.barrier() @@ -4428,7 +4428,7 @@ def _make_b_hi_loader(base_k): next_k2_py = k_iv_py + tile_k * 2 next_k2_bk = next_k2 // 2 # scales-before-B (see above for rationale). - if use_async_copy: + if const_expr(use_async_copy): prefetch_x_to_lds(next_k2, lds_base_pong) else: x_regs_pong = load_x_tile(next_k2) @@ -4456,7 +4456,7 @@ def _make_b_hi_loader(base_k): ), ) k0_pong_bk = next_k2_bk - if not use_async_copy: + if const_expr(not use_async_copy): store_x_tile_to_lds(x_regs_pong, lds_base_pong) # hot_loop_scheduler() gpu.barrier() @@ -4471,7 +4471,7 @@ def _make_b_hi_loader(base_k): else None ) - if odd_k_tiles: + if const_expr(odd_k_tiles): # Tail: single remaining tile (already in `b_cur` / `lds_base_pong`). acc, epilogue_pf = compute_tile( acc, @@ -4494,7 +4494,7 @@ def _make_b_hi_loader(base_k): int(inter_dim) + tile_k - 1 ) // tile_k * tile_k - tile_k k_tail1_bk = k_tail1 // 2 - if use_async_copy: + if const_expr(use_async_copy): prefetch_x_to_lds(k_tail1, lds_base_ping) else: x_regs_ping = load_x_tile(k_tail1) @@ -4520,7 +4520,7 @@ def _make_b_hi_loader(base_k): ), ) - if not use_async_copy: + if const_expr(not use_async_copy): store_x_tile_to_lds(x_regs_ping, lds_base_ping) # hot_loop_scheduler() gpu.barrier() @@ -4553,7 +4553,7 @@ def _make_b_hi_loader(base_k): tw_pf = None bias_pf = None - if epilogue_pf is not None: + if const_expr(epilogue_pf is not None): _, tw_pf, bias_pf = epilogue_pf mask24_i32 = arith.constant(0xFFFFFF) @@ -4571,7 +4571,7 @@ def atomic_add_f16x2(val_f16x2, byte_off_i32): ) # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). - if lds_out is None: + if const_expr(lds_out is None): raise RuntimeError( "FLIR_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased." ) @@ -4600,9 +4600,9 @@ def write_row_to_lds( num_acc_n: int, lds_out, ): - if doweight_stage2: + if const_expr(doweight_stage2): tw_idx = (mi * 4) + ii - if tw_pf is not None: + if const_expr(tw_pf is not None): tw = tw_pf[tw_idx] else: tw = buffer_ops.buffer_load( @@ -4615,12 +4615,12 @@ def write_row_to_lds( v = vector.extract( acc[acc_idx], static_position=[ii], dynamic_position=[] ) - if is_int8: + if const_expr(is_int8): v = arith.sitofp(f32, v) - if enable_bias: + if const_expr(enable_bias): v = v + bias_pf[ni] - if doweight_stage2: + if const_expr(doweight_stage2): v = v * tw v_out = arith.trunc_f(out_elem(), v) @@ -4663,7 +4663,7 @@ def precompute_row(*, row_local, row): # Fast path: lds_tid was prefilled with the per-row byte # offset (see _use_buf_atomic_pre at the prologue), so this # is just one ds_read on the per-row critical path. - if _use_buf_atomic: + if const_expr(_use_buf_atomic): row_byte_off_i32 = memref.load(lds_tid, [row_local]) return ( (None, None, row_byte_off_i32), @@ -4682,7 +4682,7 @@ def precompute_row(*, row_local, row): t_idx = arith.index_cast(ir.IndexType.get(), t) s_idx = arith.index_cast(ir.IndexType.get(), s) ts_idx = t_idx * arith.constant(topk, index=True) + s_idx - if accumulate: + if const_expr(accumulate): row_byte_base = out_base_idx + t_idx * arith.constant( model_dim * out_elem_bytes, index=True ) @@ -4703,7 +4703,7 @@ def _idx_to_llvm_ptr(idx_val, addr_space=1): def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): fused, row_byte_base, row_byte_off_i32 = row_ctx - if not bool(accumulate): + if const_expr(not bool(accumulate)): # ---- 64-bit global store path (avoids i32 offset overflow) ---- col_idx = col_g0 byte_off_col = col_idx * arith.constant( @@ -4718,7 +4718,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): alignment=_e_vec * out_elem_bytes, nontemporal=True, ) - elif _use_buf_atomic: + elif const_expr(_use_buf_atomic): # Buffer-atomic path -> BUFFER_ATOMIC_PK_ADD_BF16. # out_rsrc keeps the 64-bit base in SGPRs so only the # 32-bit voffset lives in VGPRs, vs llvm.AtomicRMWOp @@ -4782,7 +4782,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) - if _persistent: + if const_expr(_persistent): # Short-circuit: contiguous tiles are monotonically increasing, # so once bx_m >= num_valid_ids all remaining tiles are invalid. _cur_active = arith.andi(_still_active, blk_valid) @@ -4887,7 +4887,7 @@ def launch_mixed_moe_gemm2( # Stamp rocdl.waves_per_eu on the emitted gpu.func so the backend # respects the occupancy request. Done here (not at kernel-body time) # because the gpu.func only exists after the kernel has been emitted. - if waves_per_eu is not None: + if const_expr(waves_per_eu is not None): _wpe = int(waves_per_eu) for _op in ctx.gpu_module_body.operations: if ( From ffb62273b594353f48102b36fedba19bc95b8009 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Mon, 11 May 2026 09:04:50 +0000 Subject: [PATCH 3/8] update --- aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv b/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv index 21af02626e5..3e7047c709f 100644 --- a/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv +++ b/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv @@ -48,3 +48,5 @@ cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w, 256,16384,7168,512,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,1082.6229,moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,1997.4238,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0284,3080.0467,0,0.0,0.0,flydsl_fallback 256,32768,7168,512,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,1633.0245,moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,4180.9635,moe_ck2stages_gemm2_64x64x128x128_1x1_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0291,5813.988,0,0.0,0.0,flydsl_fallback 256,2,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,25.0445,moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,9.0041,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0144,34.0486,0,0.0,0.0,flydsl_fallback +256,32768,7168,2048,64,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,4655.916200000001,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,4252.6419,flydsl_moe2_afp4_wfp4_bf16_t128x128x256_atomic_persist,6.8%,8908.5581,0,2591.86,395.49, +256,32768,7168,2048,64,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,4655.916200000001,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,6110.1531,moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0681,10766.0693,0,0.0,0.0,flydsl_fallback From 9c255d5e9463609f2cbec18232c51909e89708b0 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Mon, 11 May 2026 09:38:33 +0000 Subject: [PATCH 4/8] style(flydsl): apply Black formatting Re-flow a handful of long if-conditions, error message strings, and multi-arg call sites in `aiter/ops/flydsl/moe_kernels.py` and `aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py` so they pass the `psf/black@stable` check in `.github/workflows/pre-checks.yaml`. No semantic change -- pure whitespace/line-wrap. `ruff check` on the three changed files is also clean. --- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 96 ++++++++++--------- aiter/ops/flydsl/moe_kernels.py | 5 +- 2 files changed, 57 insertions(+), 44 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 6abdf0d9040..f18d81f4ffa 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -2800,11 +2800,15 @@ def compile_mixed_moe_gemm2( # issue async X DMA early in prologue for max overlap with B/scale VMEM. _r139_xdma_first = bool( - use_async_copy and tile_m == 64 and tile_n == 128 and tile_k == 256 - and a_dtype == "fp4" and b_dtype == "fp4" and bool(accumulate) + use_async_copy + and tile_m == 64 + and tile_n == 128 + and tile_k == 256 + and a_dtype == "fp4" + and b_dtype == "fp4" + and bool(accumulate) ) - gpu_arch = get_hip_arch() allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") @@ -2866,7 +2870,9 @@ def compile_mixed_moe_gemm2( ) out_s = str(out_dtype).strip().lower() - if const_expr(out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float")): + if const_expr( + out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float") + ): raise ValueError( f"out_dtype must be 'f16', 'bf16', or 'f32', got {out_dtype!r}" ) @@ -2979,9 +2985,7 @@ def _load_bias_scalar(bias_rsrc, offset): _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" _wpe_tag = f"_w{waves_per_eu}" if waves_per_eu is not None else "" if const_expr(waves_per_eu is not None and not (1 <= int(waves_per_eu) <= 10)): - raise ValueError( - f"waves_per_eu must be in [1, 10] or None, got {waves_per_eu}" - ) + raise ValueError(f"waves_per_eu must be in [1, 10] or None, got {waves_per_eu}") _num_k_tiles_per_batch = int(inter_dim) // int(tile_k) _async_tag = "_async" if use_async_copy else "" _cumul_tag = f"_cumul{int(cu_num_mul)}" if int(cu_num_mul) != 1 else "" @@ -3634,8 +3638,11 @@ def load_b_tile(base_k, *, n_blk_p=None, n_intra_p=None): packs1 = [] for ni in range_constexpr(num_acc_n): b0, b1 = load_b_packs_k64( - base_k, ku, ni, - n_blk_p=n_blk_p, n_intra_p=n_intra_p, + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, ) packs0.append(b0) packs1.append(b1) @@ -3667,8 +3674,11 @@ def load_b_tile_lo(base_k, *, n_blk_p=None, n_intra_p=None): packs1 = [] for ni in range_constexpr(num_acc_n): b0, b1 = load_b_packs_k64( - base_k, ku, ni, - n_blk_p=n_blk_p, n_intra_p=n_intra_p, + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, ) packs0.append(b0) packs1.append(b1) @@ -3683,8 +3693,11 @@ def load_b_tile_hi(base_k, *, n_blk_p=None, n_intra_p=None): packs1 = [] for ni in range_constexpr(num_acc_n): b0, b1 = load_b_packs_k64( - base_k, ku, ni, - n_blk_p=n_blk_p, n_intra_p=n_intra_p, + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, ) packs0.append(b0) packs1.append(b1) @@ -3827,9 +3840,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base): _num_dma_loads = max( 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) ) - _c_a_elem_bytes_dma = arith.constant( - int(a_elem_bytes), index=True - ) + _c_a_elem_bytes_dma = arith.constant(int(a_elem_bytes), index=True) _c_wave_dma_bytes = arith.constant( _wave_size * _dma_bytes, index=True ) @@ -3957,12 +3968,8 @@ def compute_tile( _vec4_f32_pf = T.vec(4, f32) if const_expr(lds_tw is not None): for mi in range_constexpr(m_repeat): - mi_base_pf = arith.constant( - mi * 16, index=True - ) - lds_row_pf = ( - mi_base_pf + lane_div_16_mul4_pf - ) + mi_base_pf = arith.constant(mi * 16, index=True) + lds_row_pf = mi_base_pf + lane_div_16_mul4_pf tw_v4 = vector.load_op( _vec4_f32_pf, lds_tw, [lds_row_pf] ) @@ -3976,9 +3983,7 @@ def compute_tile( ) else: for mi in range_constexpr(m_repeat): - mi_base_pf = arith.constant( - mi * 16, index=True - ) + mi_base_pf = arith.constant(mi * 16, index=True) base_row_pf = ( bx_m + mi_base_pf + lane_div_16_mul4_pf ) @@ -4103,9 +4108,7 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): dynamic_position=[], ) if const_expr(_nss is not None): - b_scale_val = arith.shrui( - b_scale_val, _nss - ) + b_scale_val = arith.shrui(b_scale_val, _nss) # Hoist B-pack i64x4_to_i32x8 outside the # imxdl loop — b128 is a function only of @@ -4268,20 +4271,34 @@ def _emit_tid_lds_prologue(): # 0x7FFF0000 = OOB constant; HW drops OOB buffer # atomics (see comment at the cshuffle-epilogue use). _t_pre = _tid_val & arith.constant(0xFFFFFF, type=T.i32) - _s_pre = arith.shrui(_tid_val, arith.constant(24, type=T.i32)) + _s_pre = arith.shrui( + _tid_val, arith.constant(24, type=T.i32) + ) _row_byte_off = _t_pre * arith.constant( _row_stride_bytes_pre, type=T.i32 ) _global_row_i32 = arith.index_cast(T.i32, _tid_row) _valid = arith.andi( arith.andi( - arith.cmpi(CmpIPredicate.ult, _global_row_i32, num_valid_i32), - arith.cmpi(CmpIPredicate.ult, _t_pre, _tokens_i32_guard), + arith.cmpi( + CmpIPredicate.ult, + _global_row_i32, + num_valid_i32, + ), + arith.cmpi( + CmpIPredicate.ult, _t_pre, _tokens_i32_guard + ), + ), + arith.cmpi( + CmpIPredicate.ult, + _s_pre, + arith.constant(topk, type=T.i32), ), - arith.cmpi(CmpIPredicate.ult, _s_pre, arith.constant(topk, type=T.i32)), ) _stored_val = arith.select( - _valid, _row_byte_off, arith.constant(0x7FFF0000, type=T.i32) + _valid, + _row_byte_off, + arith.constant(0x7FFF0000, type=T.i32), ) else: _stored_val = _tid_val @@ -4635,9 +4652,7 @@ def write_row_to_lds( # max_t * model_dim * out_elem_bytes; we conservatively # cap row_stride <= 16384 bytes (covers model_dim<=8192 bf16). _row_stride_bytes_py = int(model_dim) * int(out_elem_bytes) - _use_buf_atomic = bool(accumulate) and ( - _row_stride_bytes_py <= 16384 - ) + _use_buf_atomic = bool(accumulate) and (_row_stride_bytes_py <= 16384) _row_stride_bytes_i32 = ( arith.constant(_row_stride_bytes_py, type=T.i32) if _use_buf_atomic @@ -4654,9 +4669,7 @@ def write_row_to_lds( # HW drops OOB buffer atomics, so routing invalid rows here # lets us drop the per-mr scf.IfOp in the cshuffle epilogue. _OOB_safe_i32 = ( - arith.constant(0x7FFF0000, type=T.i32) - if _use_buf_atomic - else None + arith.constant(0x7FFF0000, type=T.i32) if _use_buf_atomic else None ) def precompute_row(*, row_local, row): @@ -4890,10 +4903,7 @@ def launch_mixed_moe_gemm2( if const_expr(waves_per_eu is not None): _wpe = int(waves_per_eu) for _op in ctx.gpu_module_body.operations: - if ( - hasattr(_op, "attributes") - and _op.OPERATION_NAME == "gpu.func" - ): + if hasattr(_op, "attributes") and _op.OPERATION_NAME == "gpu.func": _op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( T.i32, _wpe ) diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 989075ffd2a..1431c4c0430 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -186,7 +186,10 @@ def get_flydsl_stage2_kernels( **base_params, "persist": True, } - if base_name == "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic": + if ( + base_name + == "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic" + ): # Production fp4xfp4 stage2 variant for the EP4 # DeepSeek prefill shape on MI355X. Adds: # - use_async_copy=True (async X DMA in prologue From f9984b8cddac437b64c82cad71c06be9b6dfc0cb Mon Sep 17 00:00:00 2001 From: inkcherry Date: Thu, 14 May 2026 08:11:41 +0000 Subject: [PATCH 5/8] review(flydsl): fix review comments --- aiter/aot/flydsl/moe.py | 8 +++ aiter/fused_moe.py | 11 ++- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 9 +++ aiter/ops/flydsl/moe_kernels.py | 68 +++++++++++++------ 4 files changed, 74 insertions(+), 22 deletions(-) diff --git a/aiter/aot/flydsl/moe.py b/aiter/aot/flydsl/moe.py index e2789aea9ff..0d932192b84 100644 --- a/aiter/aot/flydsl/moe.py +++ b/aiter/aot/flydsl/moe.py @@ -165,6 +165,11 @@ def _precompile_to_cache( a_scale_one: bool = False, xcd_swizzle: int = 0, enable_bias: bool = False, + # Stage2-only kernel tuning knobs (registered by the production-variant + # entries in `get_flydsl_stage2_kernels`). Forwarded into + # `compile_flydsl_moe_stage2` for stage 2 AOT compilation. + use_async_copy: bool = False, + cu_num_mul: int = 1, **kwargs, ): """Trigger MLIR compilation with dummy tensors and COMPILE_ONLY=1. @@ -505,6 +510,9 @@ def _precompile_silu_fused(): accumulate=accumulate, persist_m=_persist_m, sort_block_m=sort_block_m, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, b_nt=b_nt, xcd_swizzle=xcd_swizzle, enable_bias=enable_bias, diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index fe20fdd9965..700e31e4b18 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -771,7 +771,16 @@ def _flydsl_stage2_wrapper( bias2=None, **_kwargs, ): - + # `parsed` is the static per-kernel params dict registered at + # import time by `get_flydsl_stage2_kernels` (see + # aiter/ops/flydsl/moe_kernels.py) and pre-populated into + # `_KERNEL_PARAMS` by `_register_all_configs()`. Variant-specific + # knobs that live on the kernel name (e.g. the + # `..._persist_async_w4_cumul3` production variant adds + # `use_async_copy=True / waves_per_eu=4 / cu_num_mul=3`) are + # already baked into this dict, so the `parsed.get(..., default)` + # calls below pick up the registered values for that kernel name + # rather than always falling back to defaults. parsed = aiter.ops.flydsl.moe_kernels.get_flydsl_kernel_params(kernelName) if parsed is None: raise ValueError(f"Invalid FlyDSL kernel name: {kernelName}") diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 916c3f3a102..08054b159f1 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -2764,6 +2764,14 @@ def compile_mixed_moe_gemm2( waves_per_eu: Optional[int] = None, use_async_copy: bool = False, cu_num_mul: int = 1, + # API parity (reviewer #3): the fp4xfp4 stage2 fold doesn't consume + # `b_nt` or `xcd_swizzle` directly -- they're only meaningful for the + # stage1 kernel and the legacy stage2 implementations. Accept them as + # ignored kwargs so callers parsing `_bnt{N}` / `_xcd{N}` registry + # suffixes (unifying their stage1/stage2 dispatch) keep working without + # having to special-case the fp4xfp4 path. + b_nt: int = 0, + xcd_swizzle: int = 0, ): """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. @@ -2802,6 +2810,7 @@ def compile_mixed_moe_gemm2( assumed equal to `tile_m`. When set, stage2 can use a different tile_m from sorting/stage1. Requires sort_block_m % tile_m == 0. """ + del b_nt, xcd_swizzle # no-op on the fp4xfp4 stage2 path; see signature note _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m if const_expr(_sort_block_m != tile_m and _sort_block_m % tile_m != 0): raise ValueError( diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 1431c4c0430..e0657026cd6 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -186,30 +186,49 @@ def get_flydsl_stage2_kernels( **base_params, "persist": True, } - if ( - base_name - == "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_atomic" - ): - # Production fp4xfp4 stage2 variant for the EP4 - # DeepSeek prefill shape on MI355X. Adds: - # - use_async_copy=True (async X DMA in prologue - # overlaps with B/scale VMEM) - # - cu_num_mul=3 (persistent grid 3x CU count - # to fill in-flight slack from small per-WG - # M tile counts; cu_num_mul=4 regresses ~2.4% - # on the same shape) - # - waves_per_eu=4 (best on EP4 prefill at - # cu_num_mul=3; wpe=5/6 underperform here) - kernels[f"{base_name}_persist_async_w4_cumul3"] = { - **base_params, - "persist": True, - "use_async_copy": True, - "waves_per_eu": 4, - "cu_num_mul": 3, - } + _register_production_variants_stage2(kernels, a_dtype, b_dtype, out_dtype) return kernels +def _register_production_variants_stage2( + kernels: Dict[str, Dict], a_dtype: str, b_dtype: str, out_dtype: str +) -> None: + """Append hand-tuned stage2 variants to ``kernels`` in-place. + + Pulled out of the 6-deep tile/mode/bnt/xcd cartesian product in + ``get_flydsl_stage2_kernels`` so we don't pay a ``base_name == "..."`` + string special-case on every iteration. Each entry pins a specific + shape (tile/mode/dtype) and applies a hand-tuned override dict. + """ + # (a, b, out, tile_m, tile_n, tile_k, mode, suffix, overrides) + PRODUCTION_VARIANTS = ( + # EP4 DeepSeek prefill on MI355X (M=49152, model_dim=7168, inter_dim=2048): + # use_async_copy=True -- async X DMA in prologue overlaps with B/scale VMEM + # cu_num_mul=3 -- persistent grid 3x CU count fills in-flight + # slack from small per-WG M tile counts; + # cu_num_mul=4 regresses ~2.4% on the same shape + # waves_per_eu=4 -- best on EP4 prefill at cu_num_mul=3; + # wpe=5/6 underperform here + ( + "fp4", "fp4", "bf16", 64, 128, 256, "atomic", + "_persist_async_w4_cumul3", + { + "persist": True, + "use_async_copy": True, + "waves_per_eu": 4, + "cu_num_mul": 3, + }, + ), + ) + for pa, pb, pout, ptm, ptn, ptk, pmode, psuffix, povr in PRODUCTION_VARIANTS: + if (pa, pb, pout) != (a_dtype, b_dtype, out_dtype): + continue + _base = flydsl_kernel_name(2, pa, pb, pout, ptm, ptn, ptk, pmode) + if _base not in kernels: + continue + kernels[_base + psuffix] = {**kernels[_base], **povr} + + def get_flydsl_stage1_kernels_int4_bf16(out_dtype: str) -> Dict[str, Dict]: """Return {kernelName: params} for all supported int4_bf16 stage1 configs.""" kernels = {} @@ -430,6 +449,13 @@ def compile_flydsl_moe_stage2( waves_per_eu=waves_per_eu, use_async_copy=use_async_copy, cu_num_mul=cu_num_mul, + # API parity (reviewer #3): forward `b_nt` and `xcd_swizzle` + # from the kernel-name parser. They are accepted as ignored + # kwargs on the fp4xfp4 path so callers parsing the + # `_bnt{N}` / `_xcd{N}` registry suffixes don't need + # per-dtype special cases. + b_nt=b_nt, + xcd_swizzle=xcd_swizzle, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad, enable_bias=enable_bias, From 04ee79b50e026f17876c58df68a3d5be7c0d52dc Mon Sep 17 00:00:00 2001 From: inkcherry Date: Thu, 14 May 2026 09:41:36 +0000 Subject: [PATCH 6/8] review(flydsl): restore stage2 xcd_swizzle + format Wire xcd_swizzle end-to-end through compile_mixed_moe_gemm2 so the fp4xfp4 stage2 path actually consumes it (mirrors stage1): - signature: drop the no-op `del xcd_swizzle`, keep b_nt as no-op - module_name: append `_xcd{N}` tag when xcd_swizzle > 0 - kernel body: remap (bx_persist, by_outer) by NUM_XCDS=8 with the same group-of-M scheme stage1 uses, so consecutive WGs hit the same XCD and share L2 on MI355X - _cache_tag: include xcd_swizzle so distinct values JIT separately Also pick up the black reflow on the production-variants tuple in moe_kernels.py. --- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 53 ++++++++++++++++--- aiter/ops/flydsl/moe_kernels.py | 8 ++- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 08054b159f1..672ecf252cf 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -2765,12 +2765,14 @@ def compile_mixed_moe_gemm2( use_async_copy: bool = False, cu_num_mul: int = 1, # API parity (reviewer #3): the fp4xfp4 stage2 fold doesn't consume - # `b_nt` or `xcd_swizzle` directly -- they're only meaningful for the - # stage1 kernel and the legacy stage2 implementations. Accept them as - # ignored kwargs so callers parsing `_bnt{N}` / `_xcd{N}` registry - # suffixes (unifying their stage1/stage2 dispatch) keep working without - # having to special-case the fp4xfp4 path. + # `b_nt` directly -- it's only meaningful for the stage1 kernel's B-side + # cache modifier. Accept it as an ignored kwarg so callers parsing + # `_bnt{N}` registry suffixes (unifying their stage1/stage2 dispatch) + # keep working without having to special-case the fp4xfp4 path. b_nt: int = 0, + # XCD WG-ID swizzle (lalala-sh): remap (bx_persist, by_outer) so that + # consecutive WGs land on the same XCD and share L2, which lifts the + # stage2 hit-rate on MI355X. `0` disables the remap. xcd_swizzle: int = 0, ): """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. @@ -2810,7 +2812,7 @@ def compile_mixed_moe_gemm2( assumed equal to `tile_m`. When set, stage2 can use a different tile_m from sorting/stage1. Requires sort_block_m % tile_m == 0. """ - del b_nt, xcd_swizzle # no-op on the fp4xfp4 stage2 path; see signature note + del b_nt # no-op on the fp4xfp4 stage2 path; see signature note above _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m if const_expr(_sort_block_m != tile_m and _sort_block_m % tile_m != 0): raise ValueError( @@ -3008,10 +3010,11 @@ def _load_bias_scalar(bias_rsrc, offset): _num_k_tiles_per_batch = int(inter_dim) // int(tile_k) _async_tag = "_async" if use_async_copy else "" _cumul_tag = f"_cumul{int(cu_num_mul)}" if int(cu_num_mul) != 1 else "" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" module_name = ( f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" f"_t{tile_m}x{tile_n}x{tile_k}" - f"_vscale_fix3_fp4opt_v1{_pm_tag}{_sbm_tag}{_wpe_tag}{_async_tag}{_cumul_tag}" + f"_vscale_fix3_fp4opt_v1{_pm_tag}{_sbm_tag}{_wpe_tag}{_async_tag}{_cumul_tag}{_xcd_tag}" ).replace("-", "_") # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- # Reuse a single allocation for both: @@ -3123,6 +3126,41 @@ def check_c_k_valid_gate(base_k): tx = gpu.thread_id("x") by_outer = gpu.block_id("x") bx_persist = gpu.block_id("y") # persistent WG index (M-dim) + + if const_expr(xcd_swizzle > 0): + # Remap (bx_persist, by_outer) so consecutive WGs hit the + # same XCD and share L2 -- mirrors the stage1 XCD remap. + _NUM_XCDS_S = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_mdp_sw = arith.constant(model_dim_pad, index=True) + _gx = (n_in - _c_mdp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + if const_expr(_persistent): + _gy = arith.constant(_cu_num, index=True) + else: + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by_outer + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S = xcd_swizzle + _c_wgm = arith.constant(_WGM_S, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by_outer = _wgid_in_group / _group_size_m + by = by_outer # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). @@ -4858,6 +4896,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): _cu_num if _persistent else 0, waves_per_eu, use_async_copy, + xcd_swizzle, ) @flyc.jit diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index e0657026cd6..6a79ec9b235 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -210,7 +210,13 @@ def _register_production_variants_stage2( # waves_per_eu=4 -- best on EP4 prefill at cu_num_mul=3; # wpe=5/6 underperform here ( - "fp4", "fp4", "bf16", 64, 128, 256, "atomic", + "fp4", + "fp4", + "bf16", + 64, + 128, + 256, + "atomic", "_persist_async_w4_cumul3", { "persist": True, From 970228acc71cb17f5246c16f040dfc4de4c62f45 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Fri, 15 May 2026 05:20:45 +0000 Subject: [PATCH 7/8] review(flydsl): fix a8w4 (fp8a + fp4w) microscale gating The 2stage compile path was mis-classifying fp8 activations: only is_f4_a triggered the e8m0 microscale buffer-resource setup, so fp8xfp4 per_1x32 quantized inputs fell through to a wrong/empty scale-buffer path. This produced large output mismatches in op_tests/test_moe_2stage.py -q 7 (logits_diff ~0.98). Widen the gate to (is_f4_a or is_f8_a) so fp8a + fp4w (a8w4) correctly takes the e8m0 microscale buffer-descriptor branch, matching the pre-perf-commit behaviour. Verified locally on MI355X: - q=7 (a8w4) all 100 configs now PASS, max logits_diff = 1.19e-3 (same magnitude as stock baseline, vs ~0.98 before) - 42 xcd_swizzle=4 routed variants also pass - no change to fp4xfp4 / bf16xint4 / fp16 paths (gate only widens the scale-buffer branch for fp8a) --- aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index 672ecf252cf..6d7a7f16be7 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -3268,7 +3268,7 @@ def check_c_k_valid_gate(base_k): if const_expr(is_f16_a): sx_rsrc = None else: - if const_expr(is_f4_a): + if const_expr(is_f4_a or is_f8_a): # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. # Caller must pre-scatter a2_scale via moe_mxfp4_sort. kblk = _div_pow2(k_in, 32) From 91a047b18aec1c442e11b43031514853175618a6 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Thu, 21 May 2026 02:53:18 +0000 Subject: [PATCH 8/8] fix(aot/flydsl): default waves_per_eu to None to match runtime cache_tag This PR added `waves_per_eu` to the stage2 `_cache_tag` tuple, but aot/flydsl/moe.py::_precompile_to_cache still defaulted it to `3` while compile_mixed_moe_gemm2 defaults to `None`. AOT wrote the .pkl keyed on (..., 3, ...), runtime looked up (..., None, ...) -> AOT cache miss on every legacy stage2 kernel (test_moe_2stage.py first fails at t32x128x256_atomic, token=4, fp4xfp4 per_1x32). Align AOT default with runtime default. Stage1 and the new `_persist_async_w4_cumul3` variant are unaffected because their base_params always pin waves_per_eu explicitly. Signed-off-by: inkcherry --- aiter/aot/flydsl/moe.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/aiter/aot/flydsl/moe.py b/aiter/aot/flydsl/moe.py index f070e70de7c..9b50043ef2c 100644 --- a/aiter/aot/flydsl/moe.py +++ b/aiter/aot/flydsl/moe.py @@ -27,6 +27,7 @@ import os import sys import time +from typing import Optional from aiter.aot.flydsl.common import ( collect_aot_jobs, @@ -186,7 +187,14 @@ def _precompile_to_cache( out_dtype: str = "bf16", act: str = "silu", doweight_stage1: bool = False, - waves_per_eu: int = 3, + # Must match the runtime default of ``compile_mixed_moe_gemm2`` / + # ``compile_mixed_moe_gemm1`` (``Optional[int] = None``). A scalar default + # (e.g. ``3``) would make the AOT-side ``_cache_tag`` tuple disagree with + # the runtime-side tuple for any legacy kernel that does not explicitly + # pin ``waves_per_eu`` in ``get_flydsl_stage{1,2}_kernels`` (only the + # production-variant ``_persist_async_w4_cumul3`` does), causing + # ``AOT cache miss`` at runtime even though the .pkl is present on disk. + waves_per_eu: Optional[int] = None, k_batch: int = 1, b_nt: int = 2, gate_mode: str = "separated",