diff --git a/kernels/moe/moe_gemm_2stage/__init__.py b/kernels/moe/moe_gemm_2stage/__init__.py new file mode 100644 index 000000000..4784b58e3 --- /dev/null +++ b/kernels/moe/moe_gemm_2stage/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""MoE 2-stage MFMA kernels (stage1 / stage2 / reduction). + +Split from the former monolithic ``moe_gemm_2stage.py``; public API unchanged. + +.. deprecated:: + These kernels use the legacy FlyDSL authoring API (``SmemAllocator`` / + ``SmemPtr`` + raw ``buffer_ops``) and will be deprecated soon. New MoE work + should use the current ``fx.*`` surface (``make_buffer_tensor`` + + ``SharedAllocator`` + ``fx.copy`` / ``fx.gemm``); see ``kernels/moe/mxfp_moe`` + for the fused a4w4/a8w4 pipeline and the ``kernel-code-cleanup`` skill for the + migration map. +""" + +from kernels.moe.moe_gemm_2stage.gemm1 import compile_moe_gemm1 +from kernels.moe.moe_gemm_2stage.gemm2 import ( + MoeGemm2Mode, + _MoeGemm2ReduceWrapper, + compile_moe_gemm2, + compile_moe_gemm2_ex, +) +from kernels.moe.moe_gemm_2stage.moe_reduce import compile_moe_reduction + +__all__ = [ + "MoeGemm2Mode", + "_MoeGemm2ReduceWrapper", + "compile_moe_gemm1", + "compile_moe_gemm2", + "compile_moe_gemm2_ex", + "compile_moe_reduction", +] diff --git a/kernels/moe/moe_gemm_2stage/gemm1.py b/kernels/moe/moe_gemm_2stage/gemm1.py new file mode 100644 index 000000000..44a1dac70 --- /dev/null +++ b/kernels/moe/moe_gemm_2stage/gemm1.py @@ -0,0 +1,2218 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""MoE GEMM stage1 (MFMA) kernel builder. + +Legacy authoring API (SmemAllocator/SmemPtr + raw buffer_ops); slated for +deprecation -- refactor to the current fx.* surface (make_buffer_tensor + +SharedAllocator + fx.copy/fx.gemm). See kernels/moe/mxfp_moe and the +kernel-code-cleanup skill. +""" + +import functools +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, vector +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.common import buffer_ops +from kernels.common.kernels_common import _if_then, default_f8_type +from kernels.common.mem_ops import buffer_atomic_add +from kernels.common.mma.mfma_epilogues import c_shuffle_epilog, mfma_epilog +from kernels.common.mma.mfma_preshuffle_pipeline import ( + buffer_copy_gmem16_dwordx4, + extract_bf16_scale, + lds_store_4b_xor16, + lds_store_8b_xor16, + lds_store_16b_xor16, + load_b_pack_k32, + load_b_raw_w4a16, + load_b_raw_w4a16_groupwise, + make_preshuffle_b_layout, + preshuffle_crd2idx, + swizzle_xor16, + tile_chunk_coord_i32, + unpack_b_w4a16, +) +from kernels.moe.moe_common import ( + i64_to_v4f16 as _i64_to_v4f16, +) +from kernels.moe.moe_common import ( + i64_to_v4i16 as _i64_to_v4i16, +) +from kernels.moe.moe_common import ( + i64x2_to_v8bf16 as _i64x2_to_v8bf16, +) +from kernels.moe.moe_common import ( + i64x2_to_v8f16 as _i64x2_to_v8f16, +) +from kernels.moe.moe_gemm_2stage import layout_helpers as fxh + + +def _build_moe_gemm1_fp8_gateup( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + out_dtype: str, + in_dtype: str = "fp8", +): + """Native gate-up GEMM (B-first MFMA, fp8/bf16): out[t,slot,inter] = + silu(gate*sx*sw_g)*(up*sx*sw_u)[*routed], scattered by sorted token ids. + fp8 byte-identical to the original; bf16 shares the pipeline (unscaled). + """ + _is_bf16 = in_dtype == "bf16" + if _is_bf16: + elem_t = fx.BFloat16 # gfx950 bf16 uses native MFMA(16,16,32) + else: + elem_t = fx.Float8E4M3FNUZ + MFMA_K = 32 + elem_bytes = elem_t.width // 8 # fp8=1, bf16=2 + + K = int(model_dim) + N_e = int(2 * inter_dim) # per-expert output cols (gate+up) + BM = int(tile_m) + BN = int(tile_n) + TILE_K = int(tile_k) + TOPK = int(topk) + out_bf16 = out_dtype == "bf16" + + assert TILE_K in (128, 256), f"native gate-up needs tile_k in (128,256), got {TILE_K}" + assert K % TILE_K == 0 and (K // TILE_K) % 2 == 0, f"K={K} must be an even multiple of TILE_K={TILE_K}" + assert 64 <= BN <= 256 and BN % 64 == 0, f"tile_n must be in [64,256] multiple of 64, got {BN}" + assert 16 <= BM <= 256 and BM % 16 == 0, f"tile_m must be a 16-multiple in [16,256], got {BM}" + + # 4 waves tile the channel dim at 16 ch/wave, so a block needs >=64 real channels; + # round tile_n=64 (contiguous_n 32) up to 64 (covers two tile_n=64 groups at once). + contiguous_n = max(BN // 2, 64) + assert inter_dim % contiguous_n == 0, ( + f"inter_dim={inter_dim} must be divisible by the effective channel tile " + f"contiguous_n={contiguous_n} (from tile_n={BN})" + ) + + fp8_t = elem_t + a_lds_size = BM * TILE_K + + @fx.struct + class GemmBuffers: + a_ping: fx.Array[fp8_t, a_lds_size, 16] + a_pong: fx.Array[fp8_t, a_lds_size, 16] + + @fx.union + class SharedStorage: + sorted_lds: fx.Array[fx.Int32, 256, 16] + gemm: GemmBuffers + + _val_per_thr = 16 // elem_bytes # elements per 128b buffer_load (fp8=16, bf16=8) + # A-LDS swizzle (matches preshuffle_gemm): 8-bit (fp8) uses (3,4,3) at + # tile_k=128 (128b atom = 16 elems); 16-bit (bf16) uses (3,3,3) (128b = 8 elems). + _swz_params = (3, 3, 3) if _is_bf16 else (3, 4, 3) + _thrs_k = TILE_K // _val_per_thr + _thrs_m = 256 // _thrs_k + _m_per_wave = _thrs_m // 4 + + def _gemm_1x4(blk_n, arg_p_input, arg_p_weight, lds, M): + """B-first native-fp8 gate/up GEMM with A-gather + LDS ping-pong.""" + tid = gpu.thread_idx.x + mma_atom, tiled_mma = fxh.make_1x4_tiled_mma(fp8_t) + + # Explicit record bound so sentinel-row gathers (token_id == M padding) read + # 0 via hardware OOB instead of garbage/NaN. + a_tensor = fx.rocdl.make_buffer_tensor( + arg_p_input, max_size=False, num_records_bytes=fx.Int64(M) * fx.Int64(K) * fx.Int64(elem_bytes) + ) + b_tensor = fx.rocdl.make_buffer_tensor(arg_p_weight, max_size=False) + + # A (activation): static (BM,K) fake keeps flat_divide static; rows gathered via index. + a_size_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.get_iter(arg_p_input), fx.make_layout((BM, K), (K, 1))), max_size=False + ) + a_tile = fx.flat_divide(a_size_buf, fx.make_tile(BM, TILE_K))[None, None, 0, None] + buf_cp_atom_r = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fp8_t) + g2r_tv_layout = fx.make_layout( + ((_thrs_k, _thrs_m), (1, _val_per_thr)), + ((_thrs_m * _val_per_thr, 1), (1, _thrs_m)), + ) + a_mem_cp_g2r = fx.make_tiled_copy(buf_cp_atom_r, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)) + cp_atom_sortid_a = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32) + tiled_copy_sortid_a = fx.make_tiled_copy( + cp_atom_sortid_a, + fx.make_layout(((_thrs_k, _m_per_wave, 4), 1), ((0, 1, _m_per_wave), 0)), + fx.make_tile(_thrs_m), + ) + a_index_frag = fxh.read_sorted_index(tiled_copy_sortid_a, tid, lds.sorted_lds, BM) + a_idx = fxh.make_tensor_with_index(a_tensor, BM, TILE_K, a_index_frag, a_mem_cp_g2r, tid, TOPK) + a_mem_thr = a_mem_cp_g2r.get_slice(tid).partition_S(a_tile) + a_cp_frag = fx.make_fragment_like(a_mem_thr[None, None, None, 0]) + gpu.barrier() # sorted_lds reads done before overwriting with A tile + + # 2-stage A LDS ping-pong: overlap the next K-tile's global load + LDS + # write with the MFMA that consumes the current tile. a_ping/a_pong are the + # two staging buffers; the loop below alternates between them each K-tile. + swz = fx.SwizzleType.get(*_swz_params) + uni_cp_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fp8_t) + a_lds_bufs = [] + a_lds_w_bufs = [] + a_lds_r_bufs = [] + a_frag_bufs = [] + a_frag_retile_bufs = [] + for _buf_ptr in (lds.gemm.a_ping.ptr, lds.gemm.a_pong.ptr): + a_lds = fx.make_view( + _buf_ptr, + fx.make_composed_layout(fx.static(swz), fx.make_ordered_layout((BM, TILE_K), order=(1, 0))), + ) + a_r2s = fx.make_tiled_copy(uni_cp_atom, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)).get_slice(tid) + a_lds_bufs.append(a_lds) + a_lds_w_bufs.append(a_r2s.partition_D(a_lds)) + a_lds_r_bufs.append(fx.make_tiled_copy_B(uni_cp_atom, tiled_mma).get_slice(tid).partition_S(a_lds)) + a_frag = tiled_mma.make_fragment_B(a_lds) + a_frag_bufs.append(a_frag) + a_frag_retile_bufs.append(fx.make_tiled_copy_B(uni_cp_atom, tiled_mma).get_slice(tid).retile(a_frag)) + # Per-stage A gather staging fragments (one per in-flight buffer). + a_cp_frag_bufs = [a_cp_frag, fx.make_fragment_like(a_mem_thr[None, None, None, 0])] + a_cp_frag_retile_bufs = [ + fx.make_tiled_copy(uni_cp_atom, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)).get_slice(tid).retile(f) + for f in a_cp_frag_bufs + ] + + # B (weight gate/up): direct global->register, prefetched one K-tile ahead + # into a ping-pong pair of fragments so weight VMEM overlaps the MFMA. + bl_tile = fx.flat_divide(b_tensor, fx.make_tile(contiguous_n, TILE_K))[None, None, blk_n * 2 + 0, None] + br_tile = fx.flat_divide(b_tensor, fx.make_tile(contiguous_n, TILE_K))[None, None, blk_n * 2 + 1, None] + b_g2r = fx.make_tiled_copy_A(buf_cp_atom_r, tiled_mma).get_slice(tid) + bl_g2r = b_g2r.partition_S(bl_tile) + br_g2r = b_g2r.partition_S(br_tile) + bl_frag_bufs = [tiled_mma.make_fragment_A(bl_tile[None, None, 0]) for _ in range(2)] + br_frag_bufs = [tiled_mma.make_fragment_A(br_tile[None, None, 0]) for _ in range(2)] + bl_ret_bufs = [b_g2r.retile(f) for f in bl_frag_bufs] + br_ret_bufs = [b_g2r.retile(f) for f in br_frag_bufs] + + # Number of 128b buffer_loads issued per thread for the B (gate+up) tile. + # Used to build a targeted s_waitcnt that awaits only the A-gather (issued + # first) while leaving the prefetched B loads in flight to overlap the MFMA. + # Each BufferCopy128b load moves _val_per_thr elements; gate+up => 2 fragments. + _b_loads_per_tile = 2 * (fx.size(fx.get_shape(bl_frag_bufs[0])).to_py_value() // _val_per_thr) + + c_fake_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.get_iter(arg_p_input), fx.make_layout((contiguous_n, BM), (BM, 1))), max_size=False + ) + c_fake = fx.flat_divide(c_fake_buf, fx.make_tile(contiguous_n, BM))[None, None, 0, 0] + c_gate = tiled_mma.make_fragment_C(c_fake) + c_up = tiled_mma.make_fragment_C(c_fake) + c_gate.fill(0) + c_up.fill(0) + + _m_reps = fxh.reps(c_gate, 1) + _n_reps = fxh.reps(c_gate, 2) + k_iters = TILE_K // (2 * MFMA_K) + num_tiles = K // TILE_K + + def _load_gmem(kt, s): + """Issue global A-gather + B(gate/up) loads for K-tile kt into stage s.""" + kb = fx.Int32(kt) + a_idx.copy(buf_cp_atom_r, kb, a_cp_frag_bufs[s]) + fx.copy(buf_cp_atom_r, bl_g2r[None, None, None, kb], bl_ret_bufs[s]) + fx.copy(buf_cp_atom_r, br_g2r[None, None, None, kb], br_ret_bufs[s]) + + def _write_a_lds(s): + fx.copy(uni_cp_atom, a_cp_frag_retile_bufs[s], a_lds_w_bufs[s]) + + def _read_a_lds(s): + """Issue buffer-s LDS A-reads separately from the MFMA so ds_read latency + hides: read right after the validating barrier, ahead of next iter's MFMA.""" + for ki in range_constexpr(k_iters): + fx.copy(uni_cp_atom, a_lds_r_bufs[s][None, None, ki], a_frag_retile_bufs[s][None, None, ki]) + + def _mfma(s): + for ki in range_constexpr(k_iters): + for n in range_constexpr(_n_reps): + for m in range_constexpr(_m_reps): + for k in range_constexpr(2): + fx.mma_atom_call( + mma_atom, + c_gate[None, m, n], + bl_frag_bufs[s][None, m, (k, ki)], + a_frag_bufs[s][None, n, (k, ki)], + c_gate[None, m, n], + ) + fx.mma_atom_call( + mma_atom, + c_up[None, m, n], + br_frag_bufs[s][None, m, (k, ki)], + a_frag_bufs[s][None, n, (k, ki)], + c_up[None, m, n], + ) + + # Prologue: stage 0 global loads + LDS write for K-tile 0. Await only the + # A-gather (issued before B) so the following ds_write sees valid A; the B + # register loads for tile 0 stay in flight to overlap the first MFMA. Then + # pre-read tile 0's A fragment from LDS so its latency overlaps the next + # tile's global loads / ds_write below. + _load_gmem(0, 0) + rocdl.s_waitcnt(fxh._encode_waitcnt(vmcnt=_b_loads_per_tile)) + _write_a_lds(0) + gpu.barrier() + _read_a_lds(0) + + # Unrolled ping-pong: compute tile kt on `cur` while prefetching kt+1 (A+B) + # into `nxt`. gemm1 is low-VGPR (134 -> 3 blocks/CU) so it keeps this cross-tile + # overlap; rolling to a single buffer (like stage2) regresses ~17%. + for kt in range_constexpr(num_tiles): + cur = kt % 2 + if kt + 1 < num_tiles: + nxt = (kt + 1) % 2 + _load_gmem(kt + 1, nxt) + rocdl.s_waitcnt(fxh._encode_waitcnt(vmcnt=_b_loads_per_tile)) + _write_a_lds(nxt) + _mfma(cur) + gpu.barrier() + _read_a_lds(nxt) + else: + _mfma(cur) + return c_gate, c_up + + _gemm_1x4 = ASTRewriter.transform(_gemm_1x4) + + def _apply_fp8_dequant(c_gate_frag, c_up_frag, tid, expert_id, blk_n, asc_idx, M, arg_scale_w, arg_scale_x): + # ptpc: per-channel weight scale (gate [0,inter), up [inter,2inter)), per-token act scale. + m_reps = fxh.reps(c_gate_frag, 1) + n_reps = fxh.reps(c_gate_frag, 2) + sw_ptr = fx.recast_iter(fx.Float32, fx.get_iter(arg_scale_w)) + scale_gate = fx.make_view(sw_ptr + expert_id * N_e + blk_n * contiguous_n, fx.make_layout(contiguous_n, 1)) + scale_up = fx.make_view( + sw_ptr + expert_id * N_e + fx.Int32(inter_dim) + blk_n * contiguous_n, fx.make_layout(contiguous_n, 1) + ) + cp_atom_scale = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32) + scale_copy = fx.make_tiled_copy( + cp_atom_scale, fx.make_layout(((16, 4, 4), 4), ((0, 4, 16), 1)), fx.make_tile(64) + ) + sg_thr = scale_copy.get_slice(tid).partition_S(scale_gate) + su_thr = scale_copy.get_slice(tid).partition_S(scale_up) + gate_scale = fx.make_fragment_like(sg_thr) + up_scale = fx.make_fragment_like(su_thr) + fx.copy(cp_atom_scale, sg_thr, gate_scale) + fx.copy(cp_atom_scale, su_thr, up_scale) + + a_scale_tensor = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.recast_iter(fx.Float32, fx.get_iter(arg_scale_x)), fx.make_layout(M, 1)), + max_size=False, + num_records_bytes=fx.Int64(M) * fx.Int64(4), + ) + a_sc_n = [a_scale_tensor[asc_idx[0, n] & 0xFFFFFF] for n in range_constexpr(n_reps)] + for m in range_constexpr(m_reps): + sg_v = gate_scale[None, m].load() + su_v = up_scale[None, m].load() + for n in range_constexpr(n_reps): + a_sc = a_sc_n[n] + cg = c_gate_frag[None, m, n].load() + cu = c_up_frag[None, m, n].load() + cg_items = [] + cu_items = [] + for v in range_constexpr(4): + cg_items.append(cg[v] * sg_v[v] * a_sc) + cu_items.append(cu[v] * su_v[v] * a_sc) + c_gate_frag[None, m, n].store(fxh.Vec.from_elements(cg_items, fx.Float32)) + c_up_frag[None, m, n].store(fxh.Vec.from_elements(cu_items, fx.Float32)) + + _apply_fp8_dequant = ASTRewriter.transform(_apply_fp8_dequant) + + def _apply_doweight(c_gate_frag, c_up_frag, tid, e_idx, arg_sorted_weights): + # Per-sorted-row routed weight (one per token_rep n), folded into gate. + m_reps = fxh.reps(c_gate_frag, 1) + n_reps = fxh.reps(c_gate_frag, 2) + sw_ptr = fx.recast_iter(fx.Float32, fx.get_iter(arg_sorted_weights) + e_idx * fx.Int32(BM)) + tw_view = fx.make_view(sw_ptr, fx.make_layout(BM, 1)) + tw_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), + fx.make_layout(((16, 4, 4), 1), ((1, 0, 0), 0)), + fx.make_tile(16), + ) + tw_thr = tw_copy.get_slice(tid).partition_S(tw_view) + tw_frag = fx.make_fragment_like(tw_thr) + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), tw_thr, tw_frag) + for n in range_constexpr(n_reps): + tw = tw_frag[0, n] + for m in range_constexpr(m_reps): + c_gate_frag[None, m, n].store(c_gate_frag[None, m, n].load() * tw) + + _apply_doweight = ASTRewriter.transform(_apply_doweight) + + @flyc.kernel + def moe_gemm1_fp8_gateup( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_max_token_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_inter_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + ): + tid = gpu.thread_idx.x + blk_n = gpu.block_idx.x # tile along inter (channel/N) + e_idx = gpu.block_idx.y # expert-block id (sorted M-block) + + M = i32_tokens_in + + # Pointers / views. + in_ptr = fx.recast_iter(fp8_t, fx.get_iter(arg_x)) + arg_p_input = fx.make_view(in_ptr, fx.make_layout((M, fx.Int32(K)), (fx.Int32(K), 1))) + + max_valid_id = fxh.view_as_torch_tensor(fx.get_iter(arg_max_token_ids), (1,), fx.Int32)[0] + + if e_idx * fx.Int32(BM) < max_valid_id: + lds = fx.SharedAllocator().allocate(SharedStorage) + lds.sorted_lds = lds.sorted_lds.peek() + lds.gemm = lds.gemm.peek() + + arg_p_sorted_ids = fx.make_view( + fx.recast_iter(fx.Int32, fx.get_iter(arg_sorted_token_ids) + e_idx * fx.Int32(BM)), + fx.make_layout(BM, 1), + ) + expert_id = fxh.view_as_torch_tensor(fx.get_iter(arg_expert_ids), (1,), fx.Int32)[e_idx] + + w_ptr = fx.recast_iter(fp8_t, fx.get_iter(arg_w)) + arg_p_weight = fxh.make_gateup_weight_view(w_ptr, expert_id, contiguous_n, N_e, K) + + # Seed sorted ids into LDS (used by A-gather + output scatter index). + sorted_ids_buf = fx.rocdl.make_buffer_tensor(arg_p_sorted_ids, max_size=False) + if tid < fx.Int32(BM): + lds_view = fx.make_view(lds.sorted_lds.ptr, fx.make_layout(BM, 1)) + lds_view[tid] = sorted_ids_buf[tid] + gpu.barrier() + + # Output [M, TOPK, inter] fp16/bf16; scatter index seeded from sorted_lds now. + out_elem = fx.BFloat16 if out_bf16 else fx.Float16 + arg_p_output = fx.make_view( + fx.recast_iter(out_elem, fx.get_iter(arg_out)), + fx.make_layout( + (M, fx.Int32(TOPK), fx.Int32(inter_dim)), (fx.Int32(TOPK * inter_dim), fx.Int32(inter_dim), 1) + ), + ) + out_tensor = fx.rocdl.make_buffer_tensor(arg_p_output, max_size=False) + buf_atom_w128 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), out_elem) + # CShuffle read/scatter: 4-wave 2x2 thread grid over (BM x contiguous_n). + c_rw_copy = fx.make_tiled_copy( + buf_atom_w128, + fx.make_layout(((4, 16, 2, 2), 8), ((256, 1, 16, 1024), 32)), + fx.make_tile(32, 64), + ) + c_index_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), + fx.make_layout(((4, 16, 2, 2), 1), ((0, 1, 16, 0), 0)), + fx.make_tile(32), + ) + c_out_index_frag = fxh.read_sorted_index(c_index_copy, tid, lds.sorted_lds, BM) + c_out = fxh.make_tensor_with_index( + out_tensor, BM, contiguous_n, c_out_index_frag, c_rw_copy, tid, TOPK, is_read_from_mem=False + ) + + # Per-token activation scale index (ptpc): one id per token_rep. + asc_index_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), + fx.make_layout(((16, 4, 4), 1), ((1, 0, 0), 0)), + fx.make_tile(16), + ) + asc_lds = fx.make_view(lds.sorted_lds.ptr, fx.make_layout(BM, 1)) + asc_thr = asc_index_copy.get_slice(tid).partition_S(asc_lds) + asc_idx = fx.make_fragment_like(asc_thr) + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), asc_thr, asc_idx) + + c_gate_frag, c_up_frag = _gemm_1x4(blk_n, arg_p_input, arg_p_weight, lds, M) + + # fp8 dequant: sx (per token) * sw (per channel), folded into gate/up. + # bf16 inputs are unscaled, so there is no dequant step. + if const_expr(not _is_bf16): + _apply_fp8_dequant(c_gate_frag, c_up_frag, tid, expert_id, blk_n, asc_idx, M, arg_scale_w, arg_scale_x) + + # Optional routed-weight scale (per sorted row). + if const_expr(doweight_stage1): + _apply_doweight(c_gate_frag, c_up_frag, tid, e_idx, arg_sorted_weights) + + # silu output dtype MUST match the CShuffle LDS staging / output store + # dtype (out_elem); otherwise the raw fragment bits are reinterpreted + # (bf16 0x4480 == 1024.0 read back as f16 == 4.5). + c_out_bf16 = fxh.silu_pair_bf16(c_gate_frag, c_up_frag, out_dtype=out_elem) + + # CShuffle epilogue: stage silu output to LDS (transpose, swz 3,3,3), read back + # channel-contiguous, scatter to out[t, slot, inter] via the sorted-id index. + _, _tiled_mma = fxh.make_1x4_tiled_mma(fp8_t) + cshuf_atom_w = fx.make_copy_atom(fx.UniversalCopy64b(), out_elem) + cshuf_atom_r = fx.make_copy_atom(fx.UniversalCopy128b(), out_elem) + cshuf_ptr = fx.recast_iter(out_elem, lds.gemm.a_ping.ptr) + swz_c = fx.SwizzleType.get(3, 3, 3) + lds_c_store = fx.make_view( + cshuf_ptr, + fx.make_composed_layout(fx.static(swz_c), fx.make_ordered_layout((contiguous_n, BM), order=(0, 1))), + ) + lds_c = fx.make_view( + cshuf_ptr, + fx.make_composed_layout(fx.static(swz_c), fx.make_ordered_layout((BM, contiguous_n), order=(1, 0))), + ) + gpu.barrier() + store_c = fx.make_tiled_copy_C(cshuf_atom_w, _tiled_mma).get_slice(tid) + fx.copy(cshuf_atom_w, store_c.retile(c_out_bf16), store_c.partition_D(lds_c_store)) + gpu.barrier() + rd = fx.make_fragment_like(c_rw_copy.get_slice(tid).partition_S(lds_c)) + fx.copy(cshuf_atom_r, c_rw_copy.get_slice(tid).partition_S(lds_c), rd) + c_out.copy(buf_atom_w128, blk_n, rd) + + @flyc.jit + def launch_moe_gemm1( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_max_token_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_inter_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + stream: fx.Stream, + ): + inter_in = arith.index_cast(T.index, i32_inter_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + # Each block produces `contiguous_n` output channels (gate/up combine via + # silu into one output per channel pair), so gx = inter / contiguous_n. + # contiguous_n = max(tile_n//2, 64) mirrors the 4-wave MFMA channel minimum + # computed in the builder (see comment there). + gx = inter_in // fx.Index(contiguous_n) + gy = size_expert_ids_in + moe_gemm1_fp8_gateup( + arg_out, + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_max_token_ids, + i32_tokens_in, + i32_inter_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch(grid=(gx, gy, 1), block=(256, 1, 1), stream=stream) + + return launch_moe_gemm1 + + +@functools.lru_cache(maxsize=1024) +def compile_moe_gemm1( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + # NOTE: aiter swap passes these for API symmetry; stage1 uses dynamic memrefs so they are ignored. + doweight_stage1: bool, + in_dtype: str = "fp8", + group_size: int = -1, + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + scale_is_bf16: bool = False, + k_batch: int = 1, +): + """Compile stage1 kernel (`moe_gemm1`) and return the compiled executable. + + in_dtype: + - "fp8": X/W are fp8 + - "fp16": X/W are fp16 + - "bf16": X/W are bf16 + - "int8": X/W are int8 (X is [tokens, K]) + - "int8smooth": X/W are int8, but X is pre-expanded to [tokens*topk, K] with per-(token,slot) + quant scales (used to emulate MoE smoothquant behavior where each (token,slot)->expert route can + have a distinct input scaling before quantization). + - "int4": W4A8 path: X is int8, W is packed int4 (2 values per byte) unpacked to int8 in-kernel + - "int4_bf16": W4A16 path: X is bf16, W is packed int4 unpacked to bf16 in-kernel + scale_is_bf16: When True, groupwise scales are bf16 (halves scale bandwidth). + k_batch: Split-K factor. When >1, K is partitioned across k_batch CTAs that + atomically accumulate gate/up partials. Caller must pre-zero output. + """ + + # Native fp8 gate-up (new pipeline), hoisted above the legacy preamble for a clean + # MLIR context. fp8 non-split-K on CDNA3/CDNA4 routes here (stage1 has no atomics, so + # no CDNA3 concern); bf16 not routed yet. Everything else -> legacy body. + if in_dtype == "fp8" and k_batch == 1 and ("gfx95" in get_rocm_arch() or "gfx94" in get_rocm_arch()): + return _build_moe_gemm1_fp8_gateup( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=doweight_stage1, + out_dtype=out_dtype, + in_dtype=in_dtype, + ) + + gpu_arch = get_rocm_arch() + allocator = SmemAllocator(None, arch=gpu_arch) + _state = {} # legacy; kept until stage2/reduction are migrated + + _valid_dtypes = ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16") + if in_dtype not in _valid_dtypes: + raise ValueError(f"in_dtype must be one of {_valid_dtypes}, got {in_dtype!r}") + is_int4_bf16 = in_dtype == "int4_bf16" # W4A16: bf16 activations, packed int4 weights + is_f16 = in_dtype == "fp16" + is_bf16 = is_int4_bf16 or in_dtype == "bf16" + is_f16_or_bf16 = is_f16 or is_bf16 + needs_scale_w = (not is_f16_or_bf16) or is_int4_bf16 + elem_bytes = 2 if is_f16_or_bf16 else 1 + if out_dtype not in ("f16", "bf16"): + raise ValueError(f"out_dtype must be 'f16' or 'bf16', got {out_dtype!r}") + + # NOTE: don't materialize MLIR types outside an active MLIR Context. + def out_mlir(): + return (lambda ty: ty() if callable(ty) else ty)(T.f16 if out_dtype == "f16" else T.bf16) + + tile_k_bytes = int(tile_k) * int(elem_bytes) + # K64-byte micro-step: always 64 bytes per `ku`. For fp16 this is 32 elements. + if (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={elem_bytes})" + ) + is_int4 = in_dtype == "int4" + # INT4 here means W4A8: X is int8, W is packed int4 and unpacked to int8 in-kernel. + is_int8 = (in_dtype == "int8") or is_int4 + x_is_token_slot = in_dtype == "int8smooth" + # "int8smooth" still uses int8 MFMA, but X/scale_x are provided per (token,slot). + is_int8 = is_int8 or x_is_token_slot + + # w_is_int4: True for any variant where weights are packed int4. + w_is_int4 = is_int4 or is_int4_bf16 + + # Group-wise scale support for W4A16 + # NOTE: Only group_size=32 is supported due to int4 preshuffle layout constraints. + use_groupwise_scale = w_is_int4 and group_size > 0 + if use_groupwise_scale and group_size != 32: + raise ValueError( + f"FlyDSL groupwise scale only supports group_size=32, got {group_size}. " + f"This is due to int4 preshuffle layout constraints. " + f"Please use Triton kernel for other group sizes." + ) + is_int4_bf16_groupwise = is_int4_bf16 and use_groupwise_scale + num_groups = model_dim // group_size if use_groupwise_scale else 1 + _scale_is_bf16 = scale_is_bf16 and use_groupwise_scale + experts * (2 * inter_dim) * num_groups + # For groupwise scale, weight scale is applied per-group in the K loop, + # so epilogue can skip weight scale multiplication (uses 1.0 for sw). + + _is_gfx950 = "gfx95" in get_rocm_arch() + _has_cvt_off_f32_i4 = hasattr(rocdl, "cvt_off_f32_i4") + use_gfx950_cvt = is_int4_bf16 and _is_gfx950 and _has_cvt_off_f32_i4 + + # Split-K validation + _is_splitk = k_batch > 1 + if _is_splitk: + _k_per_batch = model_dim // k_batch + assert model_dim % k_batch == 0, f"model_dim={model_dim} not divisible by k_batch={k_batch}" + assert _k_per_batch % tile_k == 0, f"K_per_batch={_k_per_batch} not divisible by tile_k={tile_k}" + # The ping-pong K-loop requires an even number of K tiles (>=4). + _k_tiles = _k_per_batch // tile_k + assert _k_tiles >= 4 and _k_tiles % 2 == 0, ( + f"K_per_batch/tile_k={_k_tiles} must be even and >=4 for the ping-pong pipeline. " + f"Try a different k_batch (model_dim={model_dim}, tile_k={tile_k})." + ) + else: + _k_per_batch = model_dim + + mfma_i32_k32 = None + if 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: + raise AttributeError( + "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` " "(or `rocdl.mfma_i32_16x16x32_i8`)." + ) + + mfma_f32_bf16_k16 = None + if is_bf16: + mfma_f32_bf16_k16 = getattr(rocdl, "mfma_f32_16x16x16bf16_1k", None) or getattr( + rocdl, "mfma_f32_16x16x16_bf16_1k", None + ) + if mfma_f32_bf16_k16 is None: + raise AttributeError( + "BF16 K16 MFMA op not found: expected `rocdl.mfma_f32_16x16x16bf16_1k` " + "(or `rocdl.mfma_f32_16x16x16_bf16_1k`)." + ) + + # gfx950: use 16x16x32 MFMA for f16/bf16 (K=32 per MFMA, vs K=16 on gfx942). + # Check if K=32 MFMA supports the (result_type, operands_list) calling convention. + _has_k32_mfma_compat = False + if _is_gfx950 and (is_f16 or is_bf16): + import inspect + + _k32_fn = rocdl.mfma_f32_16x16x32_bf16 if is_bf16 else rocdl.mfma_f32_16x16x32_f16 + try: + _k32_sig = inspect.signature(_k32_fn) + _k32_params = list(_k32_sig.parameters.keys()) + # Compatible if second param is "operands" (list-based API) + _has_k32_mfma_compat = len(_k32_params) >= 2 and _k32_params[1] == "operands" + except (ValueError, TypeError): + _has_k32_mfma_compat = False + _use_mfma_k32 = _is_gfx950 and (is_f16 or is_bf16) and _has_k32_mfma_compat + + ir.ShapedType.get_dynamic_size() + # W is packed int4 for W4A8/W4A16/W4A_FP8: 2 values per byte. + ((experts * (2 * inter_dim) * model_dim) // 2 if w_is_int4 else (experts * (2 * inter_dim) * model_dim)) + + total_threads = 256 + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(elem_bytes) + if 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={elem_bytes}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + # Keep MoE stage1 X gmem->LDS pipeline consistent with the optimized GEMM kernel: + # split into <=16B pieces and use direct buffer_load for smaller widths. + # (Compute the split lens inside the kernel so the code matches GEMM structure.) + + # LDS128 mode (same idea as test_preshuffle_gemm.py): + # - LDS stride == tile_k (no extra padding) + XOR16 swizzle + # - Use ds_{read,write}_b128 (16B) and extract 8B halves for MFMA steps + _ck_lds128 = os.environ.get("FLYDSL_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _ck_lds128 else 8 + lds_stride = tile_k + pad_k + if use_cshuffle_epilog is None: + use_cshuffle_epilog = os.environ.get("FLYDSL_MOE_STAGE1_CSHUFFLE", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + use_cshuffle_epilog = bool(use_cshuffle_epilog) + # Split-K uses f32 atomic CShuffle regardless of out_dtype, so skip this check. + if out_dtype != "f16" and use_cshuffle_epilog and not _is_splitk: + raise ValueError("stage1 cshuffle epilog currently supports only f16 output (out_dtype='f16')") + + epilog_tag = "cshuffle" if use_cshuffle_epilog else "direct" + # IMPORTANT: module name participates in FlyDSL's compile cache key. + # Keep an explicit ABI tag so signature changes can't accidentally reuse an old binary. + _gs_tag = f"_g{group_size}" if use_groupwise_scale else "" + scale_tag = "_sbf16" if _scale_is_bf16 else "" + _split_k_tag = f"_splitk{k_batch}" if _is_splitk else "" + ( + f"mfma_moe1_{in_dtype}_{out_dtype}_{epilog_tag}" + f"_t{tile_m}x{tile_n}x{tile_k}" + f"{_gs_tag}{scale_tag}{_split_k_tag}" + f"_abi3" # also mask sentinel token ids on loads (X/scale_x) to avoid illegal address faults + ).replace("-", "_") + + # ── LDS sizing (pure Python; no MLIR Context needed) ───────────────────── + # Reuse the same LDS bytes for both: + # - ping-pong X tiles (2 * tile_m * lds_stride bytes) + # - optional epilogue CShuffle tile (tile_m * tile_n f16 -> 2 * tile_m * tile_n bytes) + _use_cshuffle_epilog = bool(use_cshuffle_epilog) + # Split-K requires CShuffle epilogue (atomic adds via store_pair callback) + if _is_splitk: + _use_cshuffle_epilog = True + # bf16 split-K: use bf16 atomics (halves bandwidth, gfx950 has buffer_atomic_pk_add_bf16). + # Other dtypes keep f32 for precision. + _splitk_use_bf16 = _is_splitk and is_bf16 + _cshuffle_elem_bytes = 2 if (not _is_splitk or _splitk_use_bf16) else 4 + lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(elem_bytes) + lds_out_bytes = _cshuffle_elem_bytes * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + lds_total_bytes = max(lds_x_bytes, lds_out_bytes) + lds_total_elems = lds_total_bytes if elem_bytes == 1 else (lds_total_bytes // 2) + + lds_alloc_bytes = int(lds_total_elems) * int(elem_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_alloc_bytes + + if True: + + @flyc.kernel + def moe_gemm1( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_max_token_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_inter_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + ): + tokens_in = arith.index_cast(T.index, i32_tokens_in) + inter_in = arith.index_cast(T.index, i32_inter_in) + k_in = arith.index_cast(T.index, i32_k_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + # i32 versions for layout construction (fly.make_shape requires i32/i64) + tokens_i32_v = i32_tokens_in + k_i32_v = i32_k_in + x_elem = T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type())) + # For int4/int4_bf16, weights are stored as packed bytes (i8) and unpacked in-kernel. + w_elem = ( + T.i8 + if w_is_int4 + else (T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type()))) + ) + scale_dtype = T.bf16 if _scale_is_bf16 else T.f32 + vec16_elems = 16 if elem_bytes == 1 else 8 + vec8_elems = 8 if elem_bytes == 1 else 4 + vec8_x = T.vec(vec8_elems, x_elem) + vec16_x = T.vec(vec16_elems, x_elem) + + def silu(x): + # device fast path: + # emu = exp(-x) ~= exp2(log2e * (-x)) -> v_exp_f32 + # sig = rcp(1 + emu) -> v_rcp_f32 + # y = x * sig + # + # Using llvm.amdgcn intrinsics prevents lowering to the div_scale/div_fixup + # sequences that introduce extra compares/cndmasks. + t = x * (-1.4426950408889634) # -log2(e) + emu = rocdl.exp2(T.f32, t) + den = 1.0 + emu + sig = rocdl.rcp(T.f32, den) + return x * sig + + acc_init = arith.constant_vector(0, T.i32x4) if is_int8 else arith.constant_vector(0.0, T.f32x4) + zero_f32_acc = arith.constant_vector(0.0, T.f32x4) if is_int4_bf16_groupwise else None + + # Layouts (use i32 values; fly.make_shape requires i32/i64, not index) + fx.make_layout((tokens_i32_v, k_i32_v), stride=(k_i32_v, 1)) + + # B preshuffle layout: match GEMM test helper exactly. + c_n_total = arith.index(experts * (2 * inter_dim)) + # For packed int4 (W4A8/W4A16/W4A_FP8), kpack_bytes=8. + kpack_bytes = 8 if w_is_int4 else 16 + w_elem_bytes = 1 if w_is_int4 else elem_bytes + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in, + kpack_bytes=kpack_bytes, + elem_bytes=w_elem_bytes, + ) + layout_b = b_layout.layout_b + (k_in * arith.index(int(elem_bytes))) // fx.Index(64) + + shape_lds = fx.make_shape(tile_m, tile_k) + stride_lds = fx.make_stride(lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + # Align with Aiter launch mapping (NSwizzle==false): + # - blockIdx.x -> N dimension (tile along inter_dim) + # - blockIdx.y -> expert-block id / M dimension (tile along sorted M) + by = gpu.block_id("x") # tile along inter_dim + bx = gpu.block_id("y") # tile along sorted M + + if const_expr(_is_splitk): + bz = gpu.block_id("z") # K-batch id + k_base_idx = bz * arith.index(_k_per_batch) + else: + k_base_idx = arith.index(0) + + # Block validity: compute as early as possible so invalid blocks skip all buffer-resource + # setup, LDS pointer math, and gmem prefetch work. + bx_m = bx * fx.Index(tile_m) + maxids_rsrc = buffer_ops.create_buffer_resource( + arg_max_token_ids, + max_size=False, + num_records_bytes=fx.Index(4), + ) + max_token_id_i32 = buffer_ops.buffer_load(maxids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32) + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, max_token_id_i32) + # Common constants/atoms (hoisted): keep IR small like GEMM. + # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). + k_blocks16 = arith.index(tile_k_bytes // 16) + layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + # Everything below is gated by `blk_valid` to avoid doing buffer-resource setup and + # gmem work for padding blocks. + _if_blk = scf.IfOp(blk_valid) + with _if_then(_if_blk): + base_ptr = allocator.get_base() + lds_x_ptr = SmemPtr( + base_ptr, + lds_alloc_offset, + (T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type()))), + shape=(lds_total_elems,), + ) + lds_x = lds_x_ptr.get() + # Alias LDS bytes for optional CShuffle epilogue. + # bf16 split-K uses bf16 (2B); other split-K uses f32 (4B); normal uses f16/bf16 (2B). + _lds_out_elem_type = T.f32 if (_is_splitk and not _splitk_use_bf16) else (T.bf16 if is_bf16 else T.f16) + lds_out = ( + SmemPtr( + base_ptr, + lds_x_ptr.byte_offset, + _lds_out_elem_type, + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + + # Buffer resources: for dynamic memrefs, provide `num_records_bytes` explicitly so + # hardware OOB behavior is stable (otherwise it falls back to a large max size). + c_topk = fx.Index(topk) + + # X: [tokens, k] bytes = tokens*k*elem_bytes + x_rows = tokens_in * (c_topk if x_is_token_slot else fx.Index(1)) + x_nbytes_idx = x_rows * k_in * arith.index(int(elem_bytes)) + x_rsrc = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_idx) + + w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False) + + # OUT: normal=[tokens, topk, inter] f16/bf16, + # split-K=[tokens*topk, 2*inter] f32 (or bf16 for bf16 split-K) + out_elem_bytes = 4 if (_is_splitk and not _splitk_use_bf16) else 2 + if const_expr(_is_splitk): + out_nbytes_idx = tokens_in * c_topk * inter_in * fx.Index(2 * out_elem_bytes) + else: + out_nbytes_idx = tokens_in * c_topk * inter_in * fx.Index(out_elem_bytes) + out_rsrc = buffer_ops.create_buffer_resource(arg_out, max_size=False, num_records_bytes=out_nbytes_idx) + + # scale_x: fp16/bf16 path ignores (implicit scale=1.0); int4_bf16 also uses 1.0. + if const_expr(is_f16_or_bf16): + sx_rsrc = None + else: + sx_rows = tokens_in * (c_topk if x_is_token_slot else fx.Index(1)) + sx_nbytes_idx = sx_rows * fx.Index(4) + sx_rsrc = buffer_ops.create_buffer_resource( + arg_scale_x, max_size=False, num_records_bytes=sx_nbytes_idx + ) + # scale_w: fp16/bf16 (non-int4) path ignores; int4_bf16 needs dequant scale. + if const_expr(not needs_scale_w): + sw_rsrc = None + else: + sw_rsrc = buffer_ops.create_buffer_resource(arg_scale_w, max_size=False) + + sorted_rsrc = buffer_ops.create_buffer_resource(arg_sorted_token_ids, max_size=False) + sorted_w_rsrc = buffer_ops.create_buffer_resource(arg_sorted_weights, max_size=False) + + # expert ids: [blocks] i32 -> bytes = size_expert_ids_in*4 + expert_rsrc = buffer_ops.create_buffer_resource( + arg_expert_ids, + max_size=False, + num_records_bytes=(size_expert_ids_in * fx.Index(4)), + ) + + # Expert id for this M tile (keep address math in `index`) + expert_i32 = buffer_ops.buffer_load(expert_rsrc, bx, vec_width=1, dtype=T.i32) + expert_idx = arith.index_cast(T.index, expert_i32) + inter2_idx = arith.index(2 * inter_dim) + expert_off_idx = expert_idx * inter2_idx # index + + # ---- 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/bf16 we require 16B. + if const_expr(is_f16_or_bf16): + 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 const_expr(bytes_per_thread_x % 16 == 0): + x_load_bytes = 16 + elif const_expr(bytes_per_thread_x % 8 == 0): + x_load_bytes = 8 + elif const_expr(bytes_per_thread_x % 4 == 0): + x_load_bytes = 4 + else: + raise ValueError( + f"bytes_per_thread_x ({bytes_per_thread_x}) must be divisible " + "by 4 to use the dword-indexed load mapping." + ) + num_x_loads = bytes_per_thread_x // x_load_bytes + chunk_i32 = x_load_bytes // 4 # dwords per chunk (1/2/4) + + c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4) + c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4) + fx.make_layout((tokens_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1)) + tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4 + layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1)) + c_chunk_i32 = fx.Index(chunk_i32) + tx_i32_base = tx * c_chunk_i32 + mask24 = fx.Int32(0xFFFFFF) + tokens_i32 = arith.index_cast(T.i32, tokens_in) + topk_i32 = fx.Int32(topk) + + def x_tile_chunk_coord_i32(i: int): + return tile_chunk_coord_i32( + arith, + tx_i32_base=tx_i32_base, + i=i, + total_threads=total_threads, + layout_tile_div4=layout_x_tile_div4, + chunk_i32=chunk_i32, + ) + + # decode token once (per thread's M-slice) and build a base row offset. + x_row_base_div4 = [] + x_col_local_i32 = [] + x_row_local = [] + for i in range_constexpr(num_x_loads): + row_local, col_local_i32 = x_tile_chunk_coord_i32(i) + x_row_local.append(row_local) + x_col_local_i32.append(col_local_i32) + + sorted_row_i = bx_m + row_local + # NOTE: rows beyond `num_valid_ids` can contain garbage (within the allocated + # buffer). That's OK as long as we never use an out-of-range token id to index X. + fused_i = buffer_ops.buffer_load(sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32) + t_raw = fused_i & mask24 + # NOTE: aiter moe_sorting uses sentinel token_id == tokens for padding. + # Do NOT rely on buffer OOB semantics for X loads; explicitly mask to a safe row. + t_valid_i32 = arith.cmpi(arith.CmpIPredicate.ult, t_raw, tokens_i32) + if const_expr(x_is_token_slot): + s_raw = fused_i >> 24 + # X is indexed by token-slot in **slot-major** order: + # row_ts = slot * tokens + token + # This matches CK's moe_smoothquant output layout. + row_ts_i32 = s_raw * tokens_i32 + t_raw + row_ts_idx = arith.index_cast(T.index, row_ts_i32) + # Apply bounds check to token-slot index + row_ts_safe = t_valid_i32.select(row_ts_idx, fx.Index(0)) + x_row_base_div4.append(row_ts_safe * c_k_div4) + else: + t_idx = arith.index_cast(T.index, t_raw) + t_safe = t_valid_i32.select(t_idx, fx.Index(0)) + x_row_base_div4.append(t_safe * c_k_div4) + + vec4_x = T.vec(4, x_elem) + + def load_x(idx_i32): + """Load `x_load_bytes` bytes from X (gmem) into regs. + + For 16B, keep the fast dwordx4 path. For 8B/4B, use byte offsets. + idx_i32 is in dword units; convert to element index for _buffer_load_vec. + """ + if const_expr(x_load_bytes == 16): + idx_elem = idx_i32 if elem_bytes == 1 else (idx_i32 * fx.Index(2)) + return buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + elem_type=x_elem, + idx_i32=idx_elem, + rsrc=x_rsrc, + vec_elems=vec16_elems, + elem_bytes=elem_bytes, + ) + # For 8B/4B, load raw i32 dwords directly. + if const_expr(x_load_bytes == 8): + return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=2, dtype=T.i32) + return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=1, dtype=T.i32) + + def load_x_tile(base_k): + """Prefetch the per-thread X tile portion (gmem -> regs) for a given K base (in elements).""" + base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4) + parts = [] + for i in range_constexpr(num_x_loads): + 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): + parts.append(vector.bitcast(T.i32x4, as_ir_value(x_vec))) + elif const_expr(x_load_bytes == 8): + parts.append(x_vec) + else: + parts.append(x_vec) + return parts + + # tx -> wave/lane (GEMM-style decomposition). + coord_wl = fx.idx2crd(fx.Int32(tx), layout_tx_wave_lane) + wave_id = fx.get(coord_wl, 0) + lane_id = fx.get(coord_wl, 1) + coord_l16 = fx.idx2crd(fx.Int32(lane_id), layout_lane16) + lane_div_16 = fx.get(coord_l16, 0) + lane_mod_16 = fx.get(coord_l16, 1) + + # Match GEMM naming/pattern: row in LDS is lane_mod_16, and col base is lane_div_16 * a_kpack_elems. + # A-side kpack is always 16 bytes (activation elements); B-side kpack_bytes + # may differ (e.g. 8 for int4 weights), but that only affects B preshuffle. + row_a_lds = lane_mod_16 + a_kpack_elems = 16 // elem_bytes + col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems)) + col_offset_base_bytes = ( + col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes))) + ) + + # Dynamic N tiling within block (same as existing kernels) + by_n = by * fx.Index(tile_n) + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + c_n_per_wave = fx.Index(n_per_wave) + wave_mod_4 = wave_id % fx.Index(4) + n_tile_base = wave_mod_4 * c_n_per_wave + + # Precompute n_blk/n_intra for gate and up rows (GEMM-style: idx2crd/get) + n_intra_gate = [] + n_blk_gate = [] + n_intra_up = [] + n_blk_up = [] + col_g_list = [] + inter_idx = fx.Index(inter_dim) + c_n_total // fx.Index(16) + c_n0_static = experts * (2 * inter_dim) // 16 + layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1)) + for ni in range_constexpr(num_acc_n): + offset = arith.index(ni * 16) + col_g = by_n + n_tile_base + col_g = col_g + offset + col_g = col_g + lane_mod_16 + col_g_list.append(col_g) + + row_gate = expert_off_idx + col_g + row_up = row_gate + inter_idx + + coord_gate = fx.idx2crd(fx.Int32(row_gate), layout_n_blk_intra) + n_blk_gate.append(fx.get(coord_gate, 0)) + n_intra_gate.append(fx.get(coord_gate, 1)) + + coord_up = fx.idx2crd(fx.Int32(row_up), layout_n_blk_intra) + n_blk_up.append(fx.get(coord_up, 0)) + n_intra_up.append(fx.get(coord_up, 1)) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 64 # K64-byte micro-step (2x MFMA) + + # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- + def load_b_pack(base_k, ki_step, ni, blk_list, intra_list): + return load_b_pack_k32( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ki_step=ki_step, + n_blk=blk_list[ni], + n_intra=intra_list[ni], + lane_div_16=lane_div_16, # 0..3 + elem_type=w_elem, + kpack_bytes=kpack_bytes, + elem_bytes=w_elem_bytes, + unpack_int4=is_int4, + ) + + def load_b_tile(base_k, blk_list, intra_list): + """Prefetch the entire per-thread B tile (gmem -> regs) for a given K base. + + Returns a list of length `k_unroll`, where each entry is a tuple: + (packs_half0[ni], packs_half1[ni]) for the K64 micro-step. + For groupwise variants, each entry also includes per-group scales: + (packs0[ni], packs1[ni], scales0[ni], scales1[ni]) + """ + if const_expr(is_int4_bf16_groupwise): + # W4A16 groupwise: load raw packed32 + scale; defer dequant to compute_tile. + raw_data = [] + for ku in range_constexpr(k_unroll): + raw_ku = [] + for ni in range_constexpr(num_acc_n): + packed32, scale_val = load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=blk_list[ni], + n_intra=intra_list[ni], + lane_div_16=lane_div_16, + elem_type=w_elem, + scale_rsrc=sw_rsrc, + expert_offset=expert_off_idx, + num_groups=num_groups, + group_size=group_size, + n_per_expert=2 * inter_dim, + kpack_bytes=kpack_bytes, + scale_dtype=scale_dtype, + ) + raw_ku.append((packed32, scale_val)) + raw_data.append(raw_ku) + return raw_data + elif const_expr(is_int4_bf16): + # W4A16 per-row: load raw packed32; defer dequant to compute_tile. + raw_data = [] + for ku in range_constexpr(k_unroll): + raw_ku = [] + for ni in range_constexpr(num_acc_n): + raw = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=blk_list[ni], + n_intra=intra_list[ni], + lane_div_16=lane_div_16, + elem_type=w_elem, + kpack_bytes=kpack_bytes, + ) + raw_ku.append(raw) + raw_data.append(raw_ku) + return raw_data + else: + # fp8/int8/bf16/fp16: original code path + b_tile = [] + for ku in range_constexpr(k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + ki0 = (ku * 2) + 0 + ki1 = (ku * 2) + 1 + b0 = load_b_pack(base_k, ki0, ni, blk_list, intra_list) + b1 = load_b_pack(base_k, ki1, ni, blk_list, intra_list) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + acc_gate = [acc_init] * (num_acc_n * m_repeat) + acc_up = [acc_init] * (num_acc_n * m_repeat) + + # ---- 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): + lds_store_16b_xor16( + arith, + vector, + lds_memref=lds_x, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x4=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + elif const_expr(x_load_bytes == 8): + lds_store_8b_xor16( + arith, + vector, + lds_memref=lds_x, + vec8_ty=vec8_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x2=vec_x_in_parts[i], + ) + else: + lds_store_4b_xor16( + arith, + vector, + lds_memref=lds_x, + vec4_ty=vec4_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x1=vec_x_in_parts[i], + ) + + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- + def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base): + col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16) + col_base_swz = ( + col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes))) + ) + idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds) + idx_a16 = idx_a16 + lds_base + loaded_a16 = vector.load(vec16_x, as_ir_value(lds_x), [as_ir_value(idx_a16)]) + a_i64x2 = vector.bitcast(T.i64x2, as_ir_value(loaded_a16)) + a0 = vector.extract(as_ir_value(a_i64x2), dynamic_position=[], static_position=[0]) + a1 = vector.extract(as_ir_value(a_i64x2), dynamic_position=[], static_position=[1]) + return a0, a1 + + def compute_tile( + acc_gate_in, + acc_up_in, + b_gate_tile_in, + b_up_tile_in, + lds_base, + *, + prefetch_epilogue: bool = False, + a0_prefetch=None, + ): + gate_list = list(acc_gate_in) + up_list = list(acc_up_in) + mfma_res_ty = T.i32x4 if is_int8 else T.f32x4 + if const_expr(_use_mfma_k32): + mfma_fn = rocdl.mfma_f32_16x16x32_f16 if is_f16 else rocdl.mfma_f32_16x16x32_bf16 + else: + mfma_fn = ( + mfma_i32_k32 + if is_int8 + else ( + mfma_f32_bf16_k16 + if is_bf16 + else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) + ) + ) + + # Optional: prefetch epilogue scales while we are about to run the last MFMA tile, + # matching the preshuffle GEMM pattern of overlapping scale loads with MFMA. + epilogue_pf = None + if const_expr(prefetch_epilogue and not use_groupwise_scale): + expert_off_pf = expert_off_idx + sw_gate_pf = [] + sw_up_pf = [] + for ni in range_constexpr(num_acc_n): + col_g = col_g_list[ni] + row_gate_idx = expert_off_pf + col_g + row_up_idx = row_gate_idx + inter_idx + sw_gate_pf.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_gate_idx, vec_width=1, dtype=T.f32) + ) + sw_up_pf.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_up_idx, vec_width=1, dtype=T.f32) + ) + epilogue_pf = (sw_gate_pf, sw_up_pf) + + def mfma_k64(acc_in, a0, a1, b0, b1): + if const_expr(_use_mfma_k32): + # gfx950: single 16x16x32 MFMA consuming all 128 bits (K=32 f16/bf16) + if const_expr(is_f16): + av = _i64x2_to_v8f16(a0, a1) + bv = _i64x2_to_v8f16(b0, b1) + else: + av = _i64x2_to_v8bf16(a0, a1) + bv = _i64x2_to_v8bf16(b0, b1) + return mfma_fn(mfma_res_ty, [av, bv, acc_in, 0, 0, 0]) + if const_expr(is_f16): + a0v = _i64_to_v4f16(a0) + a1v = _i64_to_v4f16(a1) + b0v = _i64_to_v4f16(b0) + b1v = _i64_to_v4f16(b1) + acc_mid = mfma_fn(mfma_res_ty, [a0v, b0v, acc_in, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1v, b1v, acc_mid, 0, 0, 0]) + if const_expr(is_bf16): + a0v = _i64_to_v4i16(a0) + a1v = _i64_to_v4i16(a1) + b0v = _i64_to_v4i16(b0) + b1v = _i64_to_v4i16(b1) + acc_mid = mfma_fn(mfma_res_ty, [a0v, b0v, acc_in, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1v, b1v, acc_mid, 0, 0, 0]) + acc_mid = mfma_fn(mfma_res_ty, [a0, b0, acc_in, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1, b1, acc_mid, 0, 0, 0]) + + def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val): + """MFMA f32 partial -> scale -> add to f32 accumulator via math.fma on vector.""" + from flydsl._mlir.dialects._math_ops_gen import fma as _math_fma + + _uw = arith._to_raw + scale_vec = _uw(vector.broadcast(T.f32x4, scale_val)) + return arith.ArithValue(_math_fma(scale_vec, _uw(f32_partial_vec), _uw(f32_acc_vec))) + + if const_expr(is_int4_bf16 or is_int4_bf16_groupwise): + # W4A16: deferred dequant — unpack int4->bf16 right before MFMA + # to minimize VGPR lifetime of dequantized bf16 values. + _pending_gate_up = None + for ku in range_constexpr(k_unroll): + b_gate_raw = b_gate_tile_in[ku] + b_up_raw = b_up_tile_in[ku] + ki64 = arith.index(ku * 64) + col_base = col_offset_base_bytes + ki64 + + for mi in range_constexpr(m_repeat): + mi_val = arith.index(mi * 16) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) + + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + if const_expr(is_int4_bf16_groupwise): + packed_g, sc_g = b_gate_raw[ni] + packed_u, sc_u = b_up_raw[ni] + if const_expr(_scale_is_bf16): + sc_g = extract_bf16_scale(arith, sc_g, ku) + sc_u = extract_bf16_scale(arith, sc_u, ku) + else: + packed_g, sc_g = b_gate_raw[ni], None + packed_u, sc_u = b_up_raw[ni], None + if const_expr(is_int4_bf16_groupwise and use_gfx950_cvt): + # Defer group scale to post-MFMA FMA with pipeline: + # Issue current MFMA, then apply FMA for previous iteration's result. + bg0, bg1 = unpack_b_w4a16( + packed_g, + arith, + vector, + scale_val=None, + use_gfx950_cvt=True, + defer_scale16=True, + ) + tmp_g = mfma_k64(zero_f32_acc, a0, a1, bg0, bg1) + bu0, bu1 = unpack_b_w4a16( + packed_u, + arith, + vector, + scale_val=None, + use_gfx950_cvt=True, + defer_scale16=True, + ) + tmp_u = mfma_k64(zero_f32_acc, a0, a1, bu0, bu1) + # Apply FMA for previous pending result (MFMA already completed). + if _pending_gate_up is not None: + p_idx, p_g, p_u, p_sc_g, p_sc_u = _pending_gate_up + gate_list[p_idx] = _acc_scaled_f32(gate_list[p_idx], p_g, p_sc_g) + up_list[p_idx] = _acc_scaled_f32(up_list[p_idx], p_u, p_sc_u) + _pending_gate_up = ( + acc_idx, + tmp_g, + tmp_u, + sc_g, + sc_u, + ) + else: + bg0, bg1 = unpack_b_w4a16( + packed_g, + arith, + vector, + scale_val=sc_g, + use_gfx950_cvt=use_gfx950_cvt, + defer_scale16=use_gfx950_cvt, + ) + gate_list[acc_idx] = mfma_k64(gate_list[acc_idx], a0, a1, bg0, bg1) + bu0, bu1 = unpack_b_w4a16( + packed_u, + arith, + vector, + scale_val=sc_u, + use_gfx950_cvt=use_gfx950_cvt, + defer_scale16=use_gfx950_cvt, + ) + up_list[acc_idx] = mfma_k64(up_list[acc_idx], a0, a1, bu0, bu1) + # Drain last pending FMA. + if _pending_gate_up is not None: + p_idx, p_g, p_u, p_sc_g, p_sc_u = _pending_gate_up + gate_list[p_idx] = _acc_scaled_f32(gate_list[p_idx], p_g, p_sc_g) + up_list[p_idx] = _acc_scaled_f32(up_list[p_idx], p_u, p_sc_u) + else: + for ku in range_constexpr(k_unroll): + b_gate_packs0, b_gate_packs1 = b_gate_tile_in[ku] + b_up_packs0, b_up_packs1 = b_up_tile_in[ku] + ki64 = arith.index(ku * 64) + col_base = col_offset_base_bytes + ki64 + + for mi in range_constexpr(m_repeat): + mi_val = arith.index(mi * 16) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) + + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + gate_list[acc_idx] = mfma_k64( + gate_list[acc_idx], + a0, + a1, + b_gate_packs0[ni], + b_gate_packs1[ni], + ) + up_list[acc_idx] = mfma_k64( + up_list[acc_idx], + a0, + a1, + b_up_packs0[ni], + b_up_packs1[ni], + ) + return gate_list, up_list, epilogue_pf + + # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- + lds_tile_elems = arith.index(tile_m * lds_stride) + lds_base_cur = fx.Index(0) + lds_base_nxt = lds_tile_elems + + # Optional scheduler hints (copied from tuned GEMM); can be disabled via env. + rocdl.sched_barrier(0) + + def hot_loop_scheduler(): + rocdl.sched_barrier(0) + return + mfma_group = num_acc_n * 2 + # K64 micro-step: 2x K32 MFMA per gemm. + 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(2) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + + # DS-write hints near the end: match total X LDS-store micro-ops per thread. + dswr_tail = num_x_loads + if const_expr(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 const_expr(sche_i >= dswr_start - 1): + rocdl.sched_dswr(1) + rocdl.sched_barrier(0) + + # Prologue: prefetch tile0, store to LDS(cur), sync. + k0 = k_base_idx + x_regs0 = load_x_tile(k0) + b_gate_cur = load_b_tile(k0, n_blk_gate, n_intra_gate) + b_up_cur = load_b_tile(k0, n_blk_up, n_intra_up) + store_x_tile_to_lds(x_regs0, lds_base_cur) + gpu.barrier() + + # Loop-carried ping/pong state. + lds_base_pong = lds_base_cur # current/compute + lds_base_ping = lds_base_nxt # next/load+store + + # Cross-tile A0 LDS prefetch (default-on): prefetch the first A-pack (K64) for the + # tile we are about to compute from LDS, to overlap with upcoming VMEM. + a0_prefetch_pong = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_pong) + + # Ping-pong main loop (2 tiles per iteration), leaving 2 tail tiles. + # Uses scf.for with loop-carried accumulators, B-tile prefetch, and A0 LDS prefetch. + arith.index(tile_k * 2) + c_tile_k = arith.index(tile_k) + total_tiles = int(_k_per_batch) // int(tile_k) + pair_iters = max((total_tiles - 2) // 2, 0) + + # B-tile data layout per k_unroll entry (3 variants): + # + # 1) int4 + groupwise scale (is_int4_bf16_groupwise): + # [(packed_w4, scale), (packed_w4, scale), ...] per ni + # Each ni has a (packed_weights, groupwise_scale) pair. + # Flattened as: [packed_0..N, scale_0..N] → 2 * num_acc_n values + # + # 2) int4_bf16 without groupwise scale (int4_bf16_single_field): + # [raw_i64, raw_i64, ...] per ni + # Single packed i64 per ni, already contains both weight halves. + # Flattened as: [raw_0..N] → 1 * num_acc_n values + # + # 3) fp8/int8/bf16/fp16 (default — two register packs per ku): + # (packs_even_list, packs_odd_list) + # Two lists of num_acc_n regs for even/odd MFMA operands. + # Flattened as: [even_0..N, odd_0..N] → 2 * num_acc_n values + # + int4_bf16_single_field = is_int4_bf16 and not is_int4_bf16_groupwise + _fields_per_ku = 1 if int4_bf16_single_field else 2 + _vals_per_b_tile = k_unroll * _fields_per_ku * num_acc_n + + def _flatten_b_tile(b_tile): + """Flatten B tile to a 1-D list for scf.for loop-carried state.""" + flat = [] + for ku_entry in b_tile: + if is_int4_bf16_groupwise: + # [(packed, scale), ...] → [packed_0..N, scale_0..N] + flat.extend(t[0] for t in ku_entry) + flat.extend(t[1] for t in ku_entry) + elif int4_bf16_single_field: + # [raw_i64, ...] → [raw_0..N] + flat.extend(ku_entry) + else: + # (packs_even, packs_odd) → [even_0..N, odd_0..N] + flat.extend(ku_entry[0]) + flat.extend(ku_entry[1]) + return flat + + def _unflatten_b_tile(vals): + """Reconstruct B tile from flattened scf.for loop-carried state.""" + b_tile, idx = [], 0 + for _ in range_constexpr(k_unroll): + if is_int4_bf16_groupwise: + packed = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + scales = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + b_tile.append([(packed[ni], scales[ni]) for ni in range_constexpr(num_acc_n)]) + elif int4_bf16_single_field: + b_tile.append(list(vals[idx : idx + num_acc_n])) + idx += num_acc_n + else: + packs_even = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + packs_odd = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + b_tile.append((packs_even, packs_odd)) + return b_tile + + init_state = ( + list(acc_gate) + + list(acc_up) + + _flatten_b_tile(b_gate_cur) + + _flatten_b_tile(b_up_cur) + + list(a0_prefetch_pong) + ) + + _n_acc = m_repeat * num_acc_n + _p_bg = 2 * _n_acc + _p_bu = _p_bg + _vals_per_b_tile + _p_a0 = _p_bu + _vals_per_b_tile + + for pair_iv, state in range(0, pair_iters, 1, init=init_state): + _ag = list(state[:_n_acc]) + _au = list(state[_n_acc:_p_bg]) + _bg = _unflatten_b_tile(list(state[_p_bg:_p_bu])) + _bu = _unflatten_b_tile(list(state[_p_bu:_p_a0])) + _a0pf = (state[_p_a0], state[_p_a0 + 1]) + + k_iv = k_base_idx + pair_iv * (c_tile_k + c_tile_k) + + # ---- stage 0: prefetch+store ping, compute pong ---- + next_k1 = k_iv + c_tile_k + x_regs_ping = load_x_tile(next_k1) + _bg_ping = load_b_tile(next_k1, n_blk_gate, n_intra_gate) + _bu_ping = load_b_tile(next_k1, n_blk_up, n_intra_up) + + _ag, _au, _ = compute_tile(_ag, _au, _bg, _bu, lds_base_pong, a0_prefetch=_a0pf) + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + hot_loop_scheduler() + gpu.barrier() + + _a0pf_ping = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) + + # ---- stage 1: prefetch+store pong, compute ping ---- + next_k2 = k_iv + c_tile_k + c_tile_k + x_regs_pong = load_x_tile(next_k2) + _bg_next = load_b_tile(next_k2, n_blk_gate, n_intra_gate) + _bu_next = load_b_tile(next_k2, n_blk_up, n_intra_up) + + _ag, _au, _ = compute_tile( + _ag, + _au, + _bg_ping, + _bu_ping, + lds_base_ping, + a0_prefetch=_a0pf_ping, + ) + store_x_tile_to_lds(x_regs_pong, lds_base_pong) + hot_loop_scheduler() + gpu.barrier() + + _a0pf_new = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_pong) + + loop_results = yield ( + list(_ag) + list(_au) + _flatten_b_tile(_bg_next) + _flatten_b_tile(_bu_next) + list(_a0pf_new) + ) + + # After scf.for: extract final state from yielded results. + SmemPtr._view_cache = None + if pair_iters > 0: + acc_gate = list(loop_results[:_n_acc]) + acc_up = list(loop_results[_n_acc:_p_bg]) + b_gate_cur = _unflatten_b_tile(list(loop_results[_p_bg:_p_bu])) + b_up_cur = _unflatten_b_tile(list(loop_results[_p_bu:_p_a0])) + a0_prefetch_pong = (loop_results[_p_a0], loop_results[_p_a0 + 1]) + k_tail1 = k_base_idx + arith.index(_k_per_batch - tile_k) + x_regs_ping = load_x_tile(k_tail1) + b_gate_ping = load_b_tile(k_tail1, n_blk_gate, n_intra_gate) + b_up_ping = load_b_tile(k_tail1, n_blk_up, n_intra_up) + + acc_gate, acc_up, _ = compute_tile( + acc_gate, + acc_up, + b_gate_cur, + b_up_cur, + lds_base_pong, + a0_prefetch=a0_prefetch_pong, + ) + a0_prefetch_pong = None + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + hot_loop_scheduler() + gpu.barrier() + + # Cross-tile prefetch for the final ping tile. + a0_prefetch_ping = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) + + # Epilogue: compute last tile with epilogue scale prefetch to overlap loads with MFMA. + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + b_gate_ping, + b_up_ping, + lds_base_ping, + prefetch_epilogue=True, + a0_prefetch=a0_prefetch_ping, + ) + + # Store epilogue to out[t, slot, inter] + expert_off = expert_off_idx + tokens_i32_v = tokens_i32 + topk_i32_v = topk_i32 + inter_i32_v = fx.Int32(inter_dim) + mask24_i32 = fx.Int32(0xFFFFFF) + + if const_expr(use_groupwise_scale): + sw_gate_vals = [arith.constant(1.0, type=T.f32)] * num_acc_n + sw_up_vals = [arith.constant(1.0, type=T.f32)] * num_acc_n + elif const_expr(epilogue_pf is not None): + sw_gate_vals, sw_up_vals = epilogue_pf + else: + sw_gate_vals = [] + sw_up_vals = [] + for ni in range_constexpr(num_acc_n): + col_g = col_g_list[ni] + row_gate_idx = expert_off + col_g + row_up_idx = row_gate_idx + inter_idx + sw_gate_vals.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_gate_idx, vec_width=1, dtype=T.f32) + ) + sw_up_vals.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_up_idx, vec_width=1, dtype=T.f32) + ) + + # When defer_scale16 was used, the x16 correction for v_cvt_off_f32_i4 + # was omitted from the hot loop. Fold it into the epilogue scale. + if const_expr(use_gfx950_cvt): + _c16 = fx.Float32(16.0) + sw_gate_vals = [v * _c16 for v in sw_gate_vals] + sw_up_vals = [v * _c16 for v in sw_up_vals] + + # Epilogue hoists to keep IR + Python build time small: + col_i32_list = [] + for ni in range_constexpr(num_acc_n): + col_i32_list.append(arith.index_cast(T.i32, col_g_list[ni])) + + lane_div_16 * fx.Index(4) + inter_i32_local = inter_i32_v + + # Uses EVec=4 (buffer store "x4" of fp16 elements). + use_cshuffle_epilog_flag = _use_cshuffle_epilog + + # ─── Split-K epilogue: two-pass gate/up with atomic fadd ─── + # bf16 split-K uses bf16 atomics; other dtypes use f32 atomics. + if const_expr(_is_splitk): + if const_expr(lds_out is None): + raise RuntimeError("Split-K epilogue requires lds_out (CShuffle)") + + _has_buffer_atomic_bf16_s1 = str(gpu_arch).startswith(("gfx95", "gfx12")) + _needs_global_atomic_bf16_s1 = _splitk_use_bf16 and not _has_buffer_atomic_bf16_s1 + + out_base_idx = buffer_ops.extract_base_index(arg_out) + _split_k_out_row_stride = inter_dim * 2 * out_elem_bytes # bytes per row + _split_k_e_vec = 2 # vec2 for atomic fadd (f32 or bf16) + + # Mutable slot: 0 for gate pass, inter_dim for up pass + _split_k_n_offset = [0] + + # Mutable slots for two-pass gate/up selection + _split_k_acc = [acc_gate] + _split_k_sw_vals = [sw_gate_vals] + + _splitk_lds_elem = T.bf16 if _splitk_use_bf16 else T.f32 + _splitk_lds_align = 2 if _splitk_use_bf16 else 4 + + def write_row_to_lds_splitk( + *, + mi: int, + ii: int, + row_in_tile, + row, + row_base_lds, + col_base_local, + num_acc_n: int, + lds_out, + ): + """Write scaled partial sums to LDS (no silu, no doweight).""" + _acc = _split_k_acc[0] + _sw = _split_k_sw_vals[0] + # Load per-row scale_x (sx) — same logic as normal epilogue. + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) + if const_expr(x_is_token_slot): + s2 = fused2 >> 24 + ts2 = s2 * tokens_i32_v + t2 + sx = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, ts2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + else: + sx = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, t2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + for ni in range_constexpr(num_acc_n): + col_local = col_base_local + (ni * 16) + acc_idx = mi * num_acc_n + ni + v = vector.extract(as_ir_value(_acc[acc_idx]), dynamic_position=[], static_position=[ii]) + if is_int8: + v = arith.sitofp(T.f32, v) + v = v * sx * _sw[ni] + if _splitk_use_bf16: + v = arith.trunc_f(T.bf16, v) + lds_idx = row_base_lds + col_local + v1 = vector.from_elements(T.vec(1, _splitk_lds_elem), [as_ir_value(v)]) + vector.store( + as_ir_value(v1), + as_ir_value(lds_out), + [as_ir_value(lds_idx)], + alignment=_splitk_lds_align, + ) + + def precompute_row_splitk(*, row_local, row): + 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(arith.CmpIPredicate.ult, t2, tokens_i32_v) + t_idx = arith.index_cast(T.index, t2) + s_idx = arith.index_cast(T.index, s2) + ts_idx = t_idx * arith.index(topk) + s_idx + if const_expr(_splitk_use_bf16 and not _needs_global_atomic_bf16_s1): + # For buffer atomics: compute relative byte offset from buffer base + row_byte_off = ts_idx * arith.index(_split_k_out_row_stride) + return (row_byte_off, t_ok) + else: + # For global atomics: compute absolute address + row_byte_base = out_base_idx + ts_idx * arith.index(_split_k_out_row_stride) + return (row_byte_base, t_ok) + + _splitk_zero_i32 = [fx.Int32(0) if _splitk_use_bf16 else None] + + def store_pair_splitk(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + row_byte_ctx = row_ctx + col_idx = col_g0 + arith.index(_split_k_n_offset[0]) + byte_off_col = col_idx * arith.index(out_elem_bytes) + if const_expr(_splitk_use_bf16): + _z = _splitk_zero_i32[0] + if const_expr(_needs_global_atomic_bf16_s1): + # gfx942: global atomicrmw fadd for bf16 + ptr_addr_idx = row_byte_ctx + byte_off_col + out_ptr = buffer_ops.create_llvm_ptr(ptr_addr_idx, address_space=1) + out_ptr_v = out_ptr._value if hasattr(out_ptr, "_value") else out_ptr + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.AtomicRMWOp( + llvm.AtomicBinOp.fadd, + out_ptr_v, + frag_v, + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=_split_k_e_vec * out_elem_bytes, + ) + else: + # gfx950+: buffer_atomic_pk_add_bf16 + byte_off_i32 = arith.index_cast(T.i32, row_byte_ctx + byte_off_col) + buffer_atomic_add(frag, out_rsrc, byte_off_i32, _z, _z) + else: + # f32 atomic: global atomicrmw fadd + ptr_addr_idx = row_byte_ctx + byte_off_col + out_ptr = buffer_ops.create_llvm_ptr(ptr_addr_idx, address_space=1) + out_ptr_v = out_ptr._value if hasattr(out_ptr, "_value") else out_ptr + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.AtomicRMWOp( + llvm.AtomicBinOp.fadd, + out_ptr_v, + frag_v, + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=_split_k_e_vec * out_elem_bytes, + ) + + _cshuffle_nlane_splitk = min(32, tile_n // _split_k_e_vec) + _splitk_frag_elem = ir.BF16Type.get() if _splitk_use_bf16 else ir.F32Type.get() + + # Pass 1: gate (offset=0) + _split_k_acc[0] = acc_gate + _split_k_sw_vals[0] = sw_gate_vals + _split_k_n_offset[0] = 0 + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=tile_n, + e_vec=_split_k_e_vec, + cshuffle_nlane=_cshuffle_nlane_splitk, + block_size=total_threads, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + frag_elem_type=_splitk_frag_elem, + write_row_to_lds=write_row_to_lds_splitk, + precompute_row=precompute_row_splitk, + store_pair=store_pair_splitk, + ) + + gpu.barrier() + + # Pass 2: up (offset=inter_dim) + _split_k_acc[0] = acc_up + _split_k_sw_vals[0] = sw_up_vals + _split_k_n_offset[0] = inter_dim + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=tile_n, + e_vec=_split_k_e_vec, + cshuffle_nlane=_cshuffle_nlane_splitk, + block_size=total_threads, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + frag_elem_type=_splitk_frag_elem, + write_row_to_lds=write_row_to_lds_splitk, + precompute_row=precompute_row_splitk, + store_pair=store_pair_splitk, + ) + return + + if const_expr(use_cshuffle_epilog_flag): + if const_expr(lds_out is None): + raise RuntimeError("CShuffle epilogue enabled but lds_out is not allocated/aliased.") + + def write_row_to_lds( + *, + mi: int, + ii: int, + row_in_tile, + row, + row_base_lds, + col_base_local, + num_acc_n: int, + lds_out, + ): + # `row` is the sorted-row index (bx_m + row_in_tile). + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + # aiter moe_sorting uses sentinel token_id == tokens for padding. + # Do NOT rely on buffer OOB semantics for scale loads; explicitly mask. + t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) + if const_expr(x_is_token_slot): + # slot-major: slot*tokens + token + ts2 = s2 * tokens_i32_v + t2 + sx = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, ts2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + else: + sx = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, t2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + + # Sorted weight aligned with `row` (matches aiter moe_sorting output). + if const_expr(doweight_stage1): + tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) + + for ni in range_constexpr(num_acc_n): + col_local = col_base_local + (ni * 16) + sw_gate = sw_gate_vals[ni] + sw_up = sw_up_vals[ni] + + acc_idx = mi * num_acc_n + ni + vg = vector.extract( + as_ir_value(acc_gate[acc_idx]), + dynamic_position=[], + static_position=[ii], + ) + vu = vector.extract( + as_ir_value(acc_up[acc_idx]), + dynamic_position=[], + static_position=[ii], + ) + + if const_expr(is_int8): + vg = arith.sitofp(T.f32, vg) + vu = arith.sitofp(T.f32, vu) + vg = vg * sx * sw_gate + vu = vu * sx * sw_up + + y = silu(vg) * vu + if const_expr(doweight_stage1): + y = y * tw + y16 = arith.trunc_f(T.f16, y) + + lds_idx = row_base_lds + col_local + v1 = vector.from_elements(T.vec(1, T.f16), [as_ir_value(y16)]) + vector.store( + as_ir_value(v1), + as_ir_value(lds_out), + [as_ir_value(lds_idx)], + alignment=2, + ) + + def precompute_row(*, row_local, row): + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + return (t2 * topk_i32_v + s2) * inter_i32_local + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + # Guard against sentinel token ids (t == tokens) produced by aiter moe_sorting padding. + # OOB buffer stores are not guaranteed to be safe on all paths, so predicate explicitly. + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) + _if_valid = scf.IfOp(t_valid) + with _if_then(_if_valid): + idx0 = row_ctx + col_i32 = arith.index_cast(T.i32, col_g0) + idx_out = idx0 + col_i32 + # Vectorized fp16 store (EVec=4). + buffer_ops.buffer_store(frag, out_rsrc, idx_out) + + mfma_epilog( + use_cshuffle=True, + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=tile_n, + e_vec=4, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + return + + def _stage1_store_row(*, mi: int, ii: int, row_in_tile, row): + # `row` is the sorted-row index (bx_m + row_in_tile). + # Block-level early-exit already guards `bx_m` range. + # Here we rely on buffer OOB semantics for any tail rows. + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2_raw = fused2 & mask24_i32 + s2_raw = fused2 >> 24 + t2 = t2_raw + s2 = s2_raw + t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) + + # Do NOT rely on buffer OOB semantics for scale loads; explicitly mask. + if const_expr(x_is_token_slot): + # slot-major: slot*tokens + token + ts2 = s2 * tokens_i32_v + t2 + sx0 = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, ts2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + else: + sx0 = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + t_valid, + buffer_ops.buffer_load(sx_rsrc, t2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + sx = sx0 + arith.constant(0.0, type=out_mlir()) + + # out linear index base = ((t*topk + s)*inter_dim) (invariant across ni) + idx0 = (t2 * topk_i32_v + s2) * inter_i32_local + + # Sorted weight aligned with `row` (matches aiter moe_sorting output). + if const_expr(doweight_stage1): + tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) + + _if_valid = scf.IfOp(t_valid) + with _if_then(_if_valid): + for ni in range_constexpr(num_acc_n): + col_i32 = col_i32_list[ni] + sw_gate = sw_gate_vals[ni] + sw_up = sw_up_vals[ni] + + acc_idx = mi * num_acc_n + ni + vg = vector.extract( + as_ir_value(acc_gate[acc_idx]), + dynamic_position=[], + static_position=[ii], + ) + vu = vector.extract( + as_ir_value(acc_up[acc_idx]), + dynamic_position=[], + static_position=[ii], + ) + + if const_expr(is_int8): + vg = arith.sitofp(T.f32, vg) + vu = arith.sitofp(T.f32, vu) + vg = vg * sx * sw_gate + vu = vu * sx * sw_up + + y = silu(vg) * vu + if const_expr(doweight_stage1): + y = y * tw + y = arith.trunc_f(out_mlir(), y) + idx_out0 = idx0 + col_i32 + buffer_ops.buffer_store(y, out_rsrc, idx_out0) + + mfma_epilog( + use_cshuffle=False, + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_stage1_store_row, + ) + + # ── Host launcher (flyc.jit + .launch) ──────────────────────────────── + @flyc.jit + def launch_moe_gemm1( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_max_token_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_inter_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + stream: fx.Stream, + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + inter_in = arith.index_cast(T.index, i32_inter_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + gx = inter_in // fx.Index(tile_n) + gy = size_expert_ids_in + + moe_gemm1( + arg_out, + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_max_token_ids, + i32_tokens_in, + i32_inter_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch( + grid=(gx, gy, k_batch), + block=(256, 1, 1), + stream=stream, + ) + + return launch_moe_gemm1 diff --git a/kernels/moe/moe_gemm_2stage/gemm2.py b/kernels/moe/moe_gemm_2stage/gemm2.py new file mode 100644 index 000000000..76f81e527 --- /dev/null +++ b/kernels/moe/moe_gemm_2stage/gemm2.py @@ -0,0 +1,2222 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""MoE GEMM stage2 (MFMA) kernel builder + reduce-mode dispatch. + +Legacy authoring API (SmemAllocator/SmemPtr + raw buffer_ops); slated for +deprecation -- refactor to the current fx.* surface (make_buffer_tensor + +SharedAllocator + fx.copy/fx.gemm). See kernels/moe/mxfp_moe and the +kernel-code-cleanup skill. +""" + +import functools +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +try: + from flydsl.runtime.device import ( + bf16_global_atomics_arch_description, + supports_bf16_global_atomics, + ) +except ImportError: + # Backward compatibility for runtime.device versions that only expose get_rocm_arch. + def supports_bf16_global_atomics(arch: str) -> bool: + return str(arch).startswith(("gfx94", "gfx95", "gfx12")) + + def bf16_global_atomics_arch_description() -> str: + return "gfx94+/gfx95+/gfx12+" + + +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, vector +from flydsl.expr.typing import T +from kernels.common import buffer_ops +from kernels.common.kernels_common import _if_then, default_f8_type +from kernels.common.mem_ops import buffer_atomic_add +from kernels.common.mma.mfma_epilogues import c_shuffle_epilog, default_epilog +from kernels.common.mma.mfma_preshuffle_pipeline import ( + buffer_copy_gmem16_dwordx4, + extract_bf16_scale, + lds_store_4b_xor16, + lds_store_8b_xor16, + lds_store_16b_xor16, + load_b_pack_k32, + load_b_raw_w4a16, + load_b_raw_w4a16_groupwise, + make_preshuffle_b_layout, + preshuffle_crd2idx, + swizzle_xor16, + tile_chunk_coord_i32, + unpack_b_w4a16, +) +from kernels.common.tensor_shim import _run_compiled +from kernels.moe.moe_common import ( + i64_to_v4f16 as _i64_to_v4f16, +) +from kernels.moe.moe_common import ( + i64_to_v4i16 as _i64_to_v4i16, +) +from kernels.moe.moe_common import ( + i64x2_to_v8bf16 as _i64x2_to_v8bf16, +) +from kernels.moe.moe_common import ( + i64x2_to_v8f16 as _i64x2_to_v8f16, +) +from kernels.moe.moe_gemm_2stage import layout_helpers as fxh +from kernels.moe.moe_gemm_2stage.moe_reduce import compile_moe_reduction + + +def _build_moe_gemm2_fp8( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + out_dtype: str, + accumulate: bool, + in_dtype: str = "fp8", +): + """Native stage2 down-projection (B-first MFMA, fp8/bf16); out=(A2@W2^T)*a_scale*w_scale. + Epilogue: atomic (out f16/bf16/f32) or reduce; prefer reduce (atomic is occupancy-bound + on compute-bound -- mitigated by the rolled k-loop below -- and memory-bound on decode). + """ + _is_bf16 = in_dtype == "bf16" + if _is_bf16: + elem_t = fx.BFloat16 # gfx950 bf16 uses native MFMA(16,16,32) + else: + elem_t = fx.Float8E4M3FNUZ + MFMA_K = 32 + elem_bytes = elem_t.width // 8 # fp8=1, bf16=2 + + K = int(inter_dim) # stage2 K dimension + N = int(model_dim) # stage2 output/N dimension + BM = int(tile_m) + BN = int(tile_n) + TILE_K = int(tile_k) + TOPK = int(topk) + + out_s = str(out_dtype).strip().lower() + out_is_f32 = out_s in ("f32", "fp32", "float") + out_is_bf16 = out_s in ("bf16", "bfloat16") + out_elem = fx.Float32 if out_is_f32 else (fx.BFloat16 if out_is_bf16 else fx.Float16) + out_bytes = 4 if out_is_f32 else 2 + + assert TILE_K in (128, 256), f"native stage2 needs tile_k in (128,256), got {TILE_K}" + assert K % TILE_K == 0, f"K(inter_dim)={K} must be a multiple of TILE_K={TILE_K}" + assert 64 <= BN <= 256 and BN % 64 == 0, f"tile_n must be in [64,256] multiple of 64, got {BN}" + assert 16 <= BM <= 256 and BM % 16 == 0, f"tile_m must be a 16-multiple in [16,256], got {BM}" + + # The B-first tiled_mma puts all 4 waves on the output (N) dim at 16 channels/wave, + # so each block must cover >=64 output channels. tile_n>=64 already satisfies this. + contiguous_n = BN + assert contiguous_n % 64 == 0, f"tile_n={BN} must be a 64-multiple for the 4-wave B-first MMA" + assert model_dim % contiguous_n == 0, f"model_dim={model_dim} must be divisible by tile_n={BN}" + + fp8_t = elem_t + # A LDS holds BM*TILE_K activation elements; byte size scales with elem width. + a_lds_bytes = BM * TILE_K * elem_bytes + # CShuffle staging reuses the A LDS bytes; f32 output needs BM*BN*4 bytes there. + cshuf_bytes = BM * BN * out_bytes + # Single LDS region: the A ping/pong pair (2*a_lds) during the loop, reused by the + # CShuffle epilogue (cshuf) after. max() of the two fits CDNA3's 64KB (vs CDNA4 160KB) + # for large f32 tiles. ping = region[0:], pong = +a_lds. + region_bytes = max(2 * a_lds_bytes, cshuf_bytes) + + @fx.struct + class GemmBuffers: + region: fx.Array[fx.Int8, region_bytes, 16] + + @fx.union + class SharedStorage: + sorted_lds: fx.Array[fx.Int32, 256, 16] + gemm: GemmBuffers + + _val_per_thr = 16 // elem_bytes # elements per 128b buffer_load (fp8=16, bf16=8) + # A-LDS swizzle (matches preshuffle_gemm): 8-bit (fp8) uses (3,4,3); + # 16-bit (bf16) uses (3,3,3). + _swz_params = (3, 3, 3) if _is_bf16 else (3, 4, 3) + _thrs_k = TILE_K // _val_per_thr + _thrs_m = 256 // _thrs_k + _m_per_wave = _thrs_m // 4 + + def _gemm_1x4(blk_n, arg_p_input, arg_p_weight, lds, M): + """B-first native-fp8 down-projection GEMM with A-gather + LDS ping-pong.""" + tid = gpu.thread_idx.x + mma_atom, tiled_mma = fxh.make_1x4_tiled_mma(fp8_t) + + a_tensor = fx.rocdl.make_buffer_tensor( + arg_p_input, max_size=False, num_records_bytes=fx.Int64(M) * fx.Int64(K) * fx.Int64(elem_bytes) + ) + b_tensor = fx.rocdl.make_buffer_tensor(arg_p_weight, max_size=False) + + a_size_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.get_iter(arg_p_input), fx.make_layout((BM, K), (K, 1))), max_size=False + ) + a_tile = fx.flat_divide(a_size_buf, fx.make_tile(BM, TILE_K))[None, None, 0, None] + buf_cp_atom_r = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fp8_t) + g2r_tv_layout = fx.make_layout( + ((_thrs_k, _thrs_m), (1, _val_per_thr)), + ((_thrs_m * _val_per_thr, 1), (1, _thrs_m)), + ) + a_mem_cp_g2r = fx.make_tiled_copy(buf_cp_atom_r, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)) + cp_atom_sortid_a = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32) + tiled_copy_sortid_a = fx.make_tiled_copy( + cp_atom_sortid_a, + fx.make_layout(((_thrs_k, _m_per_wave, 4), 1), ((0, 1, _m_per_wave), 0)), + fx.make_tile(_thrs_m), + ) + a_index_frag = fxh.read_sorted_index(tiled_copy_sortid_a, tid, lds.sorted_lds, BM) + a_idx = fxh.make_tensor_with_index(a_tensor, BM, TILE_K, a_index_frag, a_mem_cp_g2r, tid, TOPK) + a_mem_thr = a_mem_cp_g2r.get_slice(tid).partition_S(a_tile) + a_cp_frag = fx.make_fragment_like(a_mem_thr[None, None, None, 0]) + gpu.barrier() # sorted_lds reads done before overwriting with A tile + + swz = fx.SwizzleType.get(*_swz_params) + uni_cp_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fp8_t) + a_lds_bufs = [] + a_lds_w_bufs = [] + a_lds_r_bufs = [] + a_frag_bufs = [] + a_frag_retile_bufs = [] + for _buf_ptr in (lds.gemm.region.ptr, lds.gemm.region.ptr + a_lds_bytes): + a_lds = fx.make_view( + fx.recast_iter(fp8_t, _buf_ptr), + fx.make_composed_layout(fx.static(swz), fx.make_ordered_layout((BM, TILE_K), order=(1, 0))), + ) + a_r2s = fx.make_tiled_copy(uni_cp_atom, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)).get_slice(tid) + a_lds_bufs.append(a_lds) + a_lds_w_bufs.append(a_r2s.partition_D(a_lds)) + a_lds_r_bufs.append(fx.make_tiled_copy_B(uni_cp_atom, tiled_mma).get_slice(tid).partition_S(a_lds)) + a_frag = tiled_mma.make_fragment_B(a_lds) + a_frag_bufs.append(a_frag) + a_frag_retile_bufs.append(fx.make_tiled_copy_B(uni_cp_atom, tiled_mma).get_slice(tid).retile(a_frag)) + a_cp_frag_bufs = [a_cp_frag, fx.make_fragment_like(a_mem_thr[None, None, None, 0])] + a_cp_frag_retile_bufs = [ + fx.make_tiled_copy(uni_cp_atom, g2r_tv_layout, fx.make_tile(_thrs_m, TILE_K)).get_slice(tid).retile(f) + for f in a_cp_frag_bufs + ] + + # B (weight): single tile per block (no gate/up split); direct global->register, + # prefetched one K-tile ahead into a ping-pong pair of fragments. + b_tile = fx.flat_divide(b_tensor, fx.make_tile(contiguous_n, TILE_K))[None, None, blk_n, None] + b_g2r = fx.make_tiled_copy_A(buf_cp_atom_r, tiled_mma).get_slice(tid) + b_g2r_s = b_g2r.partition_S(b_tile) + b_frag_bufs = [tiled_mma.make_fragment_A(b_tile[None, None, 0]) for _ in range(2)] + b_ret_bufs = [b_g2r.retile(f) for f in b_frag_bufs] + + _b_loads_per_tile = fx.size(fx.get_shape(b_frag_bufs[0])).to_py_value() // _val_per_thr + + c_fake_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.get_iter(arg_p_input), fx.make_layout((contiguous_n, BM), (BM, 1))), max_size=False + ) + c_fake = fx.flat_divide(c_fake_buf, fx.make_tile(contiguous_n, BM))[None, None, 0, 0] + c_frag = tiled_mma.make_fragment_C(c_fake) + c_frag.fill(0) + + k_iters = TILE_K // (2 * MFMA_K) + num_tiles = K // TILE_K + + def _load_gmem(kb, s): + # kb is an fx.Int32 K-tile index (may be a runtime scf.for value). + a_idx.copy(buf_cp_atom_r, kb, a_cp_frag_bufs[s]) + fx.copy(buf_cp_atom_r, b_g2r_s[None, None, None, kb], b_ret_bufs[s]) + + def _write_a_lds(s): + fx.copy(uni_cp_atom, a_cp_frag_retile_bufs[s], a_lds_w_bufs[s]) + + def _read_a_lds(s): + for ki in range_constexpr(k_iters): + fx.copy(uni_cp_atom, a_lds_r_bufs[s][None, None, ki], a_frag_retile_bufs[s][None, None, ki]) + + _m_reps = fxh.reps(c_frag, 1) + _n_reps = fxh.reps(c_frag, 2) + + def _mfma(s): + for ki in range_constexpr(k_iters): + for n in range_constexpr(_n_reps): + for m in range_constexpr(_m_reps): + for k in range_constexpr(2): + fx.mma_atom_call( + mma_atom, + c_frag[None, m, n], + b_frag_bufs[s][None, m, (k, ki)], + a_frag_bufs[s][None, n, (k, ki)], + c_frag[None, m, n], + ) + + # Rolled single-buffer scf.for: the unrolled ping-pong kept both buffers live + # over the whole unroll (196 VGPR, 2 blocks/CU) and couldn't hide the atomic + # drain tail; rolling drops VGPR to 130 (3 blocks/CU), keeping intra-tile overlap. + for iv in range(0, num_tiles, 1): + kb = arith.index_cast(T.i32, iv) + _load_gmem(kb, 0) + rocdl.s_waitcnt(fxh._encode_waitcnt(vmcnt=_b_loads_per_tile)) + gpu.barrier() # WAR: all waves finished reading the prior tile's A-LDS + _write_a_lds(0) + gpu.barrier() # RAW: A-LDS write visible before the read below + _read_a_lds(0) + _mfma(0) + return c_frag + + _gemm_1x4 = ASTRewriter.transform(_gemm_1x4) + + def _apply_fp8_dequant(c_frag, tid, expert_id, blk_n, asc_idx, M, arg_scale_w, arg_scale_x): + # ptpc: per-channel (model_dim) weight scale, per-row act scale. + # Sentinel rows (token id decode invalid) get a_scale=0 so their atomic contribution is 0. + m_reps = fxh.reps(c_frag, 1) + n_reps = fxh.reps(c_frag, 2) + sw_ptr = fx.recast_iter(fx.Float32, fx.get_iter(arg_scale_w)) + scale_w = fx.make_view(sw_ptr + expert_id * N + blk_n * contiguous_n, fx.make_layout(contiguous_n, 1)) + cp_atom_scale = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32) + scale_copy = fx.make_tiled_copy( + cp_atom_scale, fx.make_layout(((16, 4, 4), 4), ((0, 4, 16), 1)), fx.make_tile(64) + ) + sw_thr = scale_copy.get_slice(tid).partition_S(scale_w) + w_scale = fx.make_fragment_like(sw_thr) + fx.copy(cp_atom_scale, sw_thr, w_scale) + + tokens = M // TOPK + a_scale_tensor = fx.rocdl.make_buffer_tensor( + fx.make_view(fx.recast_iter(fx.Float32, fx.get_iter(arg_scale_x)), fx.make_layout(M, 1)), + max_size=False, + num_records_bytes=fx.Int64(M) * fx.Int64(4), + ) + # A2 scale row = token*TOPK + slot; sentinel (token>=tokens) -> scale 0. + a_sc_n = [] + for n in range_constexpr(n_reps): + packed = asc_idx[0, n] + tok = packed & 0xFFFFFF + slot = packed >> 24 + valid = tok < fx.Int32(tokens) + row = valid.select(tok * fx.Int32(TOPK) + slot, fx.Int32(0)) + sc = valid.select(a_scale_tensor[row], fx.Float32(0.0)) + a_sc_n.append(sc) + + for m in range_constexpr(m_reps): + sw_v = w_scale[None, m].load() + for n in range_constexpr(n_reps): + a_sc = a_sc_n[n] + c = c_frag[None, m, n].load() + items = [] + for v in range_constexpr(4): + items.append(c[v] * sw_v[v] * a_sc) + c_frag[None, m, n].store(fxh.Vec.from_elements(items, fx.Float32)) + + _apply_fp8_dequant = ASTRewriter.transform(_apply_fp8_dequant) + + def _apply_doweight(c_frag, tid, e_idx, arg_sorted_weights): + # Per-sorted-row routed weight (one per token_rep n). + m_reps = fxh.reps(c_frag, 1) + n_reps = fxh.reps(c_frag, 2) + sw_ptr = fx.recast_iter(fx.Float32, fx.get_iter(arg_sorted_weights) + e_idx * fx.Int32(BM)) + tw_view = fx.make_view(sw_ptr, fx.make_layout(BM, 1)) + tw_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), + fx.make_layout(((16, 4, 4), 1), ((1, 0, 0), 0)), + fx.make_tile(16), + ) + tw_thr = tw_copy.get_slice(tid).partition_S(tw_view) + tw_frag = fx.make_fragment_like(tw_thr) + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32), tw_thr, tw_frag) + for n in range_constexpr(n_reps): + tw = tw_frag[0, n] + for m in range_constexpr(m_reps): + c_frag[None, m, n].store(c_frag[None, m, n].load() * tw) + + _apply_doweight = ASTRewriter.transform(_apply_doweight) + + def _c_to_out_frag(c_frag): + """Convert an f32 C fragment to an out_elem fragment (bf16 rounds via round_bit).""" + round_bit = fx.Uint32(0x8000) + out_frag = fx.make_fragment_like(c_frag, dtype=out_elem) + m_reps = fxh.reps(c_frag, 1) + n_reps = fxh.reps(c_frag, 2) + for m in range_constexpr(m_reps): + for n in range_constexpr(n_reps): + acc = c_frag[None, m, n].load() + if const_expr(out_is_f32): + pass + elif const_expr(out_is_bf16): + acc = ((acc.bitcast(fx.Uint32) + round_bit) >> 16).to(fx.Uint16).bitcast(fx.BFloat16) + else: + acc = acc.to(out_elem) + out_frag[None, m, n].store(acc) + return out_frag + + _c_to_out_frag = ASTRewriter.transform(_c_to_out_frag) + + @flyc.kernel + def moe_gemm2( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_num_valid_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_n_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + ): + tid = gpu.thread_idx.x + blk_n = gpu.block_idx.x # tile along model_dim (output/N) + e_idx = gpu.block_idx.y # expert-block id (sorted M-block) + + tokens = i32_tokens_in + M = tokens * fx.Int32(TOPK) + + in_ptr = fx.recast_iter(fp8_t, fx.get_iter(arg_x)) + # A2 is [tokens, topk, inter(K)] flattened; the A-gather decodes the sorted + # id into (token, slot), so the gather view MUST be rank-3 (a rank-2 view + # would drop the slot and read slot-0 of every token -- correct only when + # all topk slots of a token are near-identical, which masks the bug on tame + # data). Row = token*topk + slot. + arg_p_input = fx.make_view( + in_ptr, + fx.make_layout((tokens, fx.Int32(TOPK), fx.Int32(K)), (fx.Int32(TOPK * K), fx.Int32(K), 1)), + ) + + num_valid_id = fxh.view_as_torch_tensor(fx.get_iter(arg_num_valid_ids), (1,), fx.Int32)[0] + + if e_idx * fx.Int32(BM) < num_valid_id: + lds = fx.SharedAllocator().allocate(SharedStorage) + lds.sorted_lds = lds.sorted_lds.peek() + lds.gemm = lds.gemm.peek() + + arg_p_sorted_ids = fx.make_view( + fx.recast_iter(fx.Int32, fx.get_iter(arg_sorted_token_ids) + e_idx * fx.Int32(BM)), + fx.make_layout(BM, 1), + ) + expert_id = fxh.view_as_torch_tensor(fx.get_iter(arg_expert_ids), (1,), fx.Int32)[e_idx] + + w_ptr = fx.recast_iter(fp8_t, fx.get_iter(arg_w)) + arg_p_weight = fxh.make_weight_view(w_ptr, expert_id, N, K) + + # Seed sorted ids into LDS (A-gather + output scatter index). + sorted_ids_buf = fx.rocdl.make_buffer_tensor(arg_p_sorted_ids, max_size=False) + if tid < fx.Int32(BM): + lds_view = fx.make_view(lds.sorted_lds.ptr, fx.make_layout(BM, 1)) + lds_view[tid] = sorted_ids_buf[tid] + gpu.barrier() + + # Output tensor: reduce -> [tokens*topk, model_dim] buffer tensor for + # BufferCopy128b stores. Atomic -> a rank-2 [tokens, model_dim] buffer + # resource; scatter issues buffer_atomic_add at explicit element indices + # (out_tensor here only supplies the tile_m/tile_k shape to the index + # helper; its blocks are unused on the atomic path). + out_atomic_rsrc = None + if const_expr(accumulate): + out_atomic_rsrc = buffer_ops.create_buffer_resource( + arg_out, + max_size=False, + num_records_bytes=(fx.Index(tokens) * fx.Index(N * out_bytes)), + ) + arg_p_output = fx.make_view( + fx.recast_iter(out_elem, fx.get_iter(arg_out)), + fx.make_layout((tokens, fx.Int32(N)), (fx.Int32(N), 1)), + ) + else: + arg_p_output = fx.make_view( + fx.recast_iter(out_elem, fx.get_iter(arg_out)), + fx.make_layout((tokens, fx.Int32(TOPK), fx.Int32(N)), (fx.Int32(TOPK * N), fx.Int32(N), 1)), + ) + out_tensor = fx.rocdl.make_buffer_tensor(arg_p_output, max_size=False) + buf_atom_w128 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), out_elem) + _c_vec = 128 // out_elem.width # values per 128b atom (f16/bf16=8, f32=4) + # CShuffle read/scatter: 4-wave 2x2 thread grid over (BM x contiguous_n). + c_rw_copy = fx.make_tiled_copy( + buf_atom_w128, + fx.make_layout(((4, 16, 2, 2), _c_vec), ((256, 1, 16, 1024), 32)), + fx.make_tile(32, 64), + ) + c_index_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), + fx.make_layout(((4, 16, 2, 2), 1), ((0, 1, 16, 0), 0)), + fx.make_tile(32), + ) + c_out_index_frag = fxh.read_sorted_index(c_index_copy, tid, lds.sorted_lds, BM) + c_out = fxh.make_tensor_with_index( + out_tensor, BM, contiguous_n, c_out_index_frag, c_rw_copy, tid, TOPK, is_read_from_mem=False + ) + + # Per-row activation scale index (ptpc): one packed id per token_rep. + asc_index_copy = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), + fx.make_layout(((16, 4, 4), 1), ((1, 0, 0), 0)), + fx.make_tile(16), + ) + asc_lds = fx.make_view(lds.sorted_lds.ptr, fx.make_layout(BM, 1)) + asc_thr = asc_index_copy.get_slice(tid).partition_S(asc_lds) + asc_idx = fx.make_fragment_like(asc_thr) + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32), asc_thr, asc_idx) + + c_frag = _gemm_1x4(blk_n, arg_p_input, arg_p_weight, lds, M) + + # fp8 dequant: a_scale (per row) * w_scale (per channel), folded into C. + # bf16 inputs are unscaled, so there is no dequant step. + if const_expr(not _is_bf16): + _apply_fp8_dequant(c_frag, tid, expert_id, blk_n, asc_idx, M, arg_scale_w, arg_scale_x) + if const_expr(doweight_stage2): + _apply_doweight(c_frag, tid, e_idx, arg_sorted_weights) + + c_out_frag = _c_to_out_frag(c_frag) + + # CShuffle epilogue: stage output to LDS (transpose, swz 3,3,3 for 2B / 3,2,3 for 4B), + # read back channel-contiguous, scatter to out via the sorted-id index. + _, _tiled_mma = fxh.make_1x4_tiled_mma(fp8_t) + _log2_vec = 3 if out_bytes == 2 else 2 # 8 (2B) or 4 (4B) elems per 128b + cshuf_atom_w = fx.make_copy_atom(fx.UniversalCopy64b(), out_elem) + cshuf_atom_r = fx.make_copy_atom(fx.UniversalCopy128b(), out_elem) + cshuf_ptr = fx.recast_iter(out_elem, lds.gemm.region.ptr) + swz_c = fx.SwizzleType.get(3, _log2_vec, 3) + lds_c_store = fx.make_view( + cshuf_ptr, + fx.make_composed_layout(fx.static(swz_c), fx.make_ordered_layout((contiguous_n, BM), order=(0, 1))), + ) + lds_c = fx.make_view( + cshuf_ptr, + fx.make_composed_layout(fx.static(swz_c), fx.make_ordered_layout((BM, contiguous_n), order=(1, 0))), + ) + gpu.barrier() + store_c = fx.make_tiled_copy_C(cshuf_atom_w, _tiled_mma).get_slice(tid) + fx.copy(cshuf_atom_w, store_c.retile(c_out_frag), store_c.partition_D(lds_c_store)) + gpu.barrier() + rd = fx.make_fragment_like(c_rw_copy.get_slice(tid).partition_S(lds_c)) + fx.copy(cshuf_atom_r, c_rw_copy.get_slice(tid).partition_S(lds_c), rd) + if const_expr(not accumulate): + c_out.copy(buf_atom_w128, blk_n, rd) + else: + _atomic_mode = "f32" if out_is_f32 else "pk" + c_out.copy( + buf_atom_w128, + blk_n, + rd, + atomic=_atomic_mode, + atomic_rsrc=out_atomic_rsrc, + out_bytes=out_bytes, + row_stride=N, + row_limit=tokens, + ) + + @flyc.jit + def launch_moe_gemm2( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_num_valid_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_n_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + stream: fx.Stream, + ): + n_in = arith.index_cast(T.index, i32_n_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + gx = n_in // fx.Index(contiguous_n) + gy = size_expert_ids_in + moe_gemm2( + arg_out, + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_num_valid_ids, + i32_tokens_in, + i32_n_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch(grid=(gx, gy, 1), block=(256, 1, 1), stream=stream) + + return launch_moe_gemm2 + + +@functools.lru_cache(maxsize=1024) +def compile_moe_gemm2( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + in_dtype: str = "fp8", + group_size: int = -1, + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + accumulate: bool = True, + scale_is_bf16: bool = False, +): + """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + + in_dtype: + - "fp8": A2/W are fp8 + - "fp16": A2/W are fp16 + - "bf16": A2/W are bf16 + - "int8": A2/W are int8 + - "int4": W4A8 path: A2 is int8, W is packed int4 unpacked to int8 in-kernel + - "int4_bf16": W4A16 path: A2 is bf16, W is packed int4 unpacked to bf16 in-kernel + scale_is_bf16: When True, groupwise scales are bf16 (halves scale bandwidth). + + Stage2 output supports: + - out_dtype="f16": fp16 half2 atomics (fast, can overflow to +/-inf for bf16 workloads) + - out_dtype="f32": fp32 scalar atomics (slower, but avoids fp16 atomic overflow) + + `use_cshuffle_epilog` controls whether we use the LDS CShuffle epilogue before + global atomics (recommended for performance). + """ + # Native fp8 (new pipeline), non-groupwise, CDNA3/CDNA4. gfx942 carve-out: CDNA3 + # lacks buffer_atomic_pk_add_bf16, so atomic+bf16-output stays on legacy (global + # bf16 atomics) there; every other fp8 combo uses the new path. bf16 not routed yet. + _arch = get_rocm_arch() + _cdna = "gfx94" in _arch or "gfx95" in _arch + _g942_bf16_atomic = "gfx94" in _arch and accumulate and str(out_dtype).strip().lower() in ("bf16", "bfloat16") + if in_dtype == "fp8" and group_size <= 0 and _cdna and not _g942_bf16_atomic: + _out_s = str(out_dtype).strip().lower() + if _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}") + if (not bool(accumulate)) and _out_s in ("f32", "fp32", "float"): + raise ValueError("compile_moe_gemm2(accumulate=False) only supports out_dtype in {'f16','bf16'}") + return _build_moe_gemm2_fp8( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=doweight_stage2, + out_dtype=out_dtype, + accumulate=accumulate, + in_dtype=in_dtype, + ) + + gpu_arch = get_rocm_arch() + allocator = SmemAllocator(None, arch=gpu_arch) + _state = {} + + _valid_dtypes = ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16") + if in_dtype not in _valid_dtypes: + raise ValueError(f"in_dtype must be one of {_valid_dtypes}, got {in_dtype!r}") + is_int4_bf16 = in_dtype == "int4_bf16" # W4A16: bf16 activations, packed int4 weights + is_f16 = in_dtype == "fp16" + is_bf16 = is_int4_bf16 or in_dtype == "bf16" + is_f16_or_bf16 = is_f16 or is_bf16 + needs_scale_w = (not is_f16_or_bf16) or is_int4_bf16 + elem_bytes = 2 if is_f16_or_bf16 else 1 + out_s = str(out_dtype).strip().lower() + if 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: + raise ValueError("compile_moe_gemm2(accumulate=False) only supports out_dtype in {'f16','bf16'}") + is_int4 = in_dtype == "int4" + # w_is_int4: True for any variant where weights are packed int4. + w_is_int4 = is_int4 or is_int4_bf16 + # INT4 here means W4A8: A2 is int8, W is packed int4 and unpacked to int8 in-kernel. + is_int8 = (in_dtype in ("int8", "int8smooth")) or is_int4 + + # Group-wise scale support for W4A16 + use_groupwise_scale = w_is_int4 and group_size > 0 + if use_groupwise_scale and group_size != 32: + raise ValueError( + f"FlyDSL groupwise scale only supports group_size=32, got {group_size}. " + f"This is due to int4 preshuffle layout constraints. " + f"Please use Triton kernel for other group sizes." + ) + is_int4_bf16_groupwise = is_int4_bf16 and use_groupwise_scale + # Stage2 K dimension is inter_dim (weight shape: [E, model_dim, inter_dim]) + num_groups = inter_dim // group_size if use_groupwise_scale else 1 + _scale_is_bf16 = scale_is_bf16 and use_groupwise_scale + experts * model_dim * num_groups + + _is_gfx950 = "gfx95" in get_rocm_arch() + _has_cvt_off_f32_i4 = hasattr(rocdl, "cvt_off_f32_i4") + use_gfx950_cvt = is_int4_bf16 and _is_gfx950 and _has_cvt_off_f32_i4 + + mfma_i32_k32 = None + if 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: + raise AttributeError( + "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` " "(or `rocdl.mfma_i32_16x16x32_i8`)." + ) + + mfma_f32_bf16_k16 = None + if is_bf16: + mfma_f32_bf16_k16 = getattr(rocdl, "mfma_f32_16x16x16bf16_1k", None) or getattr( + rocdl, "mfma_f32_16x16x16_bf16_1k", None + ) + if mfma_f32_bf16_k16 is None: + raise AttributeError( + "BF16 K16 MFMA op not found: expected `rocdl.mfma_f32_16x16x16bf16_1k` " + "(or `rocdl.mfma_f32_16x16x16_bf16_1k`)." + ) + + # gfx950: use 16x16x32 MFMA for f16/bf16 (K=32 per MFMA, vs K=16 on gfx942). + # Check if K=32 MFMA supports the (result_type, operands_list) calling convention. + _has_k32_mfma_compat = False + if _is_gfx950 and (is_f16 or is_bf16): + import inspect + + _k32_fn = rocdl.mfma_f32_16x16x32_bf16 if is_bf16 else rocdl.mfma_f32_16x16x32_f16 + try: + _k32_sig = inspect.signature(_k32_fn) + _k32_params = list(_k32_sig.parameters.keys()) + # Compatible if second param is "operands" (list-based API) + _has_k32_mfma_compat = len(_k32_params) >= 2 and _k32_params[1] == "operands" + except (ValueError, TypeError): + _has_k32_mfma_compat = False + _use_mfma_k32 = _is_gfx950 and (is_f16 or is_bf16) and _has_k32_mfma_compat + + ir.ShapedType.get_dynamic_size() + # W is packed int4 for W4A8/W4A16/W4A_FP8: 2 values per byte. + ((experts * model_dim * inter_dim) // 2 if w_is_int4 else (experts * model_dim * inter_dim)) + + total_threads = 256 + tile_k_bytes = int(tile_k) * int(elem_bytes) + if (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={elem_bytes})" + ) + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(elem_bytes) + if 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={elem_bytes}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _ck_lds128 = os.environ.get("FLYDSL_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _ck_lds128 else 8 + lds_stride = tile_k + pad_k + # gfx950+ has buffer_atomic_pk_add_bf16 → bf16 can use buffer atomics (same as f16). + # gfx942 only has global_atomic_pk_add_bf16 → must use global atomics with raw pointer. + _has_buffer_atomic_bf16 = str(gpu_arch).startswith(("gfx95", "gfx12")) + _needs_global_atomic_bf16 = out_is_bf16 and not _has_buffer_atomic_bf16 + if out_is_bf16: + if not supports_bf16_global_atomics(gpu_arch): + raise ValueError( + f"out_dtype='bf16' requires bf16 global atomics " + f"({bf16_global_atomics_arch_description()}), got arch={gpu_arch!r}" + ) + + if 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: + raise ValueError("out_dtype='f32' does not support CShuffle epilogue (set use_cshuffle_epilog=False).") + else: + if use_cshuffle_epilog is None: + _use_cshuffle_epilog = os.environ.get("FLYDSL_MOE_STAGE2_CSHUFFLE", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + else: + _use_cshuffle_epilog = bool(use_cshuffle_epilog) + if not _use_cshuffle_epilog: + raise ValueError("stage2 f16 output currently requires CShuffle epilogue (FLYDSL_MOE_STAGE2_CSHUFFLE=1).") + + # NOTE: Keep this as a callable so we don't require an MLIR Context at Python-time. + def out_elem(): + ty = T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + return ty() if callable(ty) else ty + + epilog_tag = "cshuffle" + # IMPORTANT: include tiling in the module name to avoid accidentally reusing a compiled + # binary for a different (tile_m, tile_n, tile_k) configuration. + # See stage1 note: include ABI tag to prevent binary reuse across signature changes. + # IMPORTANT: module name participates in FlyDSL's compile cache key. + # 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. + _gs_tag = f"_g{group_size}" if use_groupwise_scale else "" + scale_tag = "_sbf16" if _scale_is_bf16 else "" + ( + f"mfma_moe2_{in_dtype}_{out_s}_{epilog_tag}" + f"_t{tile_m}x{tile_n}x{tile_k}" + f"{_gs_tag}{scale_tag}" + f"_abi2" # mask sentinel token ids on loads/stores to avoid illegal address faults + ).replace("-", "_") + + # ── CShuffle epilogue e_vec (pure Python; must be computed before @flyc.kernel + # because the AST rewriter intercepts `if` statements inside kernel bodies and + # turns them into closure dispatches, which breaks variable reassignment) ──── + _cshuffle_nlane = 32 + if bool(accumulate): + _e_vec = 2 + else: + _e_vec = 8 if int(tile_n) % (_cshuffle_nlane * 8) == 0 else 2 + _cshuffle_stride = _cshuffle_nlane * _e_vec + if int(tile_n) % _cshuffle_stride != 0: + raise ValueError(f"tile_n={tile_n} must be divisible by {_cshuffle_stride} when accumulate=False") + + # ── LDS sizing (pure Python; no MLIR Context needed) ───────────────────── + lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(elem_bytes) + lds_out_bytes = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 # f16 bytes + lds_total_bytes = max(lds_x_bytes, lds_out_bytes) + lds_total_elems = lds_total_bytes if elem_bytes == 1 else (lds_total_bytes // 2) + + lds_alloc_bytes = int(lds_total_elems) * int(elem_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_alloc_bytes + + if True: + + @flyc.kernel + def moe_gemm2( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_num_valid_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_n_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + ): + tokens_in = arith.index_cast(T.index, i32_tokens_in) + n_in = arith.index_cast(T.index, i32_n_in) + k_in = arith.index_cast(T.index, i32_k_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + # i32 versions for layout construction (fly.make_shape requires i32/i64) + k_i32_v = i32_k_in + x_elem = T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type())) + # For int4/int4_bf16, weights are stored as packed bytes (i8) and unpacked in-kernel. + w_elem = ( + T.i8 + if w_is_int4 + else (T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type()))) + ) + scale_dtype = T.bf16 if _scale_is_bf16 else T.f32 + vec16_elems = 16 if elem_bytes == 1 else 8 + vec8_elems = 8 if elem_bytes == 1 else 4 + vec8_x = T.vec(vec8_elems, x_elem) + vec16_x = T.vec(vec16_elems, x_elem) + + acc_init = arith.constant_vector(0, T.i32x4) if is_int8 else arith.constant_vector(0.0, T.f32x4) + zero_f32_acc = arith.constant_vector(0.0, T.f32x4) if is_int4_bf16_groupwise else None + + # A2 layout (flatten token-slot -> M; use i32 for fly.make_shape). + topk_idx = fx.Index(topk) + m_in = tokens_in * topk_idx + m_i32_v = arith.index_cast(T.i32, m_in) + fx.make_layout((m_i32_v, k_i32_v), stride=(k_i32_v, 1)) + + # B preshuffle layout: [experts*model_dim, inter_dim] + c_n_total = arith.index(experts * model_dim) + # For packed int4 (W4A8/W4A16/W4A_FP8), kpack_bytes=8. + kpack_bytes = 8 if w_is_int4 else 16 + w_elem_bytes = 1 if w_is_int4 else elem_bytes + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in, + kpack_bytes=kpack_bytes, + elem_bytes=w_elem_bytes, + ) + layout_b = b_layout.layout_b + (k_in * arith.index(int(elem_bytes))) // fx.Index(64) + + shape_lds = fx.make_shape(tile_m, tile_k) + stride_lds = fx.make_stride(lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + # Align with Aiter launch mapping: + # - blockIdx.x -> N dimension (tile along model_dim) + # - blockIdx.y -> expert-block id / M dimension (tile along sorted M) + by = gpu.block_id("x") # tile along model_dim + bx = gpu.block_id("y") # tile along sorted M + + # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). + k_blocks16 = arith.index(tile_k_bytes // 16) + layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + fx.make_layout((tile_m, tile_k), stride=(tile_k, 1)) + + base_ptr = allocator.get_base() + lds_x_ptr = SmemPtr( + base_ptr, + lds_alloc_offset, + (T.bf16 if is_bf16 else (T.f16 if is_f16 else (T.i8 if is_int8 else default_f8_type()))), + 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, + 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 + ) + + # Buffer resources. + # For dynamic memrefs, `max_size=False` cannot infer the logical size from the memref *type*, + # so we should pass `num_records_bytes` explicitly for stable hardware OOB behavior. + c_topk = fx.Index(topk) + + # X(A2): [tokens*topk, inter_dim] bytes = tokens*topk*k*elem_bytes + x_nbytes_idx = (tokens_in * c_topk) * k_in * arith.index(int(elem_bytes)) + x_rsrc = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_idx) + + w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False) + + # OUT: [tokens, model_dim] -> clamp to descriptor max (i32 bytes) to avoid overflow on huge tokens. + out_elem_bytes = 4 if out_is_f32 else 2 + out_nbytes_idx = tokens_in * n_in * fx.Index(out_elem_bytes) + if const_expr(not bool(accumulate)): + out_nbytes_idx = tokens_in * fx.Index(topk) * n_in * fx.Index(out_elem_bytes) + out_rsrc = buffer_ops.create_buffer_resource(arg_out, max_size=False, num_records_bytes=out_nbytes_idx) + # scale_x: fp16/bf16 path ignores (implicit scale=1.0); int4_bf16 also uses 1.0. + if const_expr(is_f16_or_bf16): + sx_rsrc = None + else: + # scale_x (A2 scale): [tokens*topk] f32 -> bytes = tokens*topk*4 + sx_nbytes_idx = (tokens_in * c_topk) * fx.Index(4) + sx_rsrc = buffer_ops.create_buffer_resource( + arg_scale_x, max_size=False, num_records_bytes=sx_nbytes_idx + ) + # scale_w: fp16/bf16 (non-int4) path ignores; int4_bf16 needs dequant scale. + if const_expr(not needs_scale_w): + sw_rsrc = None + else: + # scale_w: [experts*model_dim] f32 (static shape in practice) + sw_rsrc = buffer_ops.create_buffer_resource(arg_scale_w, max_size=False) + + # sorted_token_ids / sorted_weights: [blocks*tile_m] (CK-style padded length) + sorted_nbytes_idx = size_expert_ids_in * fx.Index(tile_m) * fx.Index(4) + sorted_rsrc = buffer_ops.create_buffer_resource( + arg_sorted_token_ids, + max_size=False, + num_records_bytes=sorted_nbytes_idx, + ) + sorted_w_rsrc = buffer_ops.create_buffer_resource( + arg_sorted_weights, max_size=False, num_records_bytes=sorted_nbytes_idx + ) + + # expert ids: [blocks] i32 -> bytes = size_expert_ids_in*4 + eid_nbytes_idx = size_expert_ids_in * fx.Index(4) + expert_rsrc = buffer_ops.create_buffer_resource( + arg_expert_ids, max_size=False, num_records_bytes=eid_nbytes_idx + ) + bx_m = bx * fx.Index(tile_m) + + # Early-exit guard (as in 2ce65fb): some routing paths can produce extra/garbage + # expert blocks beyond `num_valid_ids`. Skip those blocks entirely to avoid OOB. + numids_rsrc = buffer_ops.create_buffer_resource( + arg_num_valid_ids, + max_size=False, + num_records_bytes=fx.Index(4), + ) + num_valid_i32 = buffer_ops.buffer_load(numids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32) + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, num_valid_i32) + + def _moe_gemm2_then_body(): + # Expert id for this M tile. + expert_i32 = buffer_ops.buffer_load(expert_rsrc, bx, vec_width=1, dtype=T.i32) + expert_idx = arith.index_cast(T.index, expert_i32) + n_idx = fx.Index(model_dim) + expert_off_idx = expert_idx * n_idx # index + + # ---- 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/bf16 we require 16B. + if const_expr(is_f16_or_bf16): + 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 const_expr(bytes_per_thread_x % 16 == 0): + x_load_bytes = 16 + elif const_expr(bytes_per_thread_x % 8 == 0): + x_load_bytes = 8 + elif const_expr(bytes_per_thread_x % 4 == 0): + x_load_bytes = 4 + else: + raise ValueError( + f"bytes_per_thread_x ({bytes_per_thread_x}) must be divisible " + "by 4 to use the dword-indexed load mapping." + ) + num_x_loads = bytes_per_thread_x // x_load_bytes + chunk_i32 = x_load_bytes // 4 # dwords per chunk (1/2/4) + + c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4) + c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4) + fx.make_layout((m_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1)) + tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4 + layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1)) + c_chunk_i32 = fx.Index(chunk_i32) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = fx.Int32(topk) + mask24 = fx.Int32(0xFFFFFF) + # Sentinel clamp uses `tokens` as the upper bound: t_valid = (t < tokens). + tokens_i32 = arith.index_cast(T.i32, tokens_in) + + def x_tile_chunk_coord_i32(i: int): + return tile_chunk_coord_i32( + arith, + tx_i32_base=tx_i32_base, + i=i, + total_threads=total_threads, + layout_tile_div4=layout_x_tile_div4, + chunk_i32=chunk_i32, + ) + + vec4_x = T.vec(4, x_elem) + + def load_x(idx_i32): + if const_expr(x_load_bytes == 16): + idx_elem = idx_i32 if elem_bytes == 1 else (idx_i32 * fx.Index(2)) + return buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + elem_type=x_elem, + idx_i32=idx_elem, + rsrc=x_rsrc, + vec_elems=vec16_elems, + elem_bytes=elem_bytes, + ) + if const_expr(x_load_bytes == 8): + return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=2, dtype=T.i32) + return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=1, dtype=T.i32) + + # decode routed token once (per thread's M-slice) and build a base offset. + x_row_base_div4 = [] + x_col_local_i32 = [] + x_row_local = [] + for i in range_constexpr(num_x_loads): + row_local, col_local_i32 = x_tile_chunk_coord_i32(i) + 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 = fused_i & mask24 + s_i32 = fused_i >> 24 + # aiter moe_sorting uses sentinel token_id == tokens for padding. + # Do NOT rely on buffer OOB semantics for A2/scale loads; explicitly mask. + t_valid = arith.cmpi(arith.CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(arith.CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = t_valid & s_valid + t_safe = ts_valid.select(t_i32, fx.Int32(0)) + s_safe = ts_valid.select(s_i32, fx.Int32(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(T.index, row_ts_i32) + # Base row offset in dword units: row_ts_idx * (k_in/4) + x_row_base_div4.append(row_ts_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4) + parts = [] + for i in range_constexpr(num_x_loads): + 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): + parts.append(vector.bitcast(T.i32x4, as_ir_value(x_vec))) + elif const_expr(x_load_bytes == 8): + parts.append(vector.bitcast(T.vec(2, T.i32), as_ir_value(x_vec))) + else: + parts.append(vector.bitcast(T.vec(1, T.i32), as_ir_value(x_vec))) + return parts + + # tx -> wave/lane (GEMM-style decomposition). + coord_wl = fx.idx2crd(fx.Int32(tx), layout_tx_wave_lane) + wave_id = fx.get(coord_wl, 0) + lane_id = fx.get(coord_wl, 1) + coord_l16 = fx.idx2crd(fx.Int32(lane_id), layout_lane16) + lane_div_16 = fx.get(coord_l16, 0) + lane_mod_16 = fx.get(coord_l16, 1) + + row_a_lds = lane_mod_16 + # A-side kpack is always 16 bytes; kpack_bytes is B-side (may be 8 for int4). + a_kpack_elems = 16 // elem_bytes + col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems)) + col_offset_base_bytes = ( + col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes))) + ) + + # Dynamic N tiling within block. + by_n = by * fx.Index(tile_n) + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + c_n_per_wave = fx.Index(n_per_wave) + wave_mod_4 = wave_id % fx.Index(4) + n_tile_base = wave_mod_4 * c_n_per_wave + + # Precompute (n_blk, n_intra) for B, and col indices for output. + n_intra_list = [] + n_blk_list = [] + col_g_list = [] + c_n_total // fx.Index(16) + c_n0_static = experts * model_dim // 16 + layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1)) + for ni in range_constexpr(num_acc_n): + offset = arith.index(ni * 16) + col_g = by_n + n_tile_base + offset + lane_mod_16 + col_g_list.append(col_g) + + row_w = expert_off_idx + col_g + coord_w = fx.idx2crd(fx.Int32(row_w), layout_n_blk_intra) + n_blk_list.append(fx.get(coord_w, 0)) + n_intra_list.append(fx.get(coord_w, 1)) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 64 # K64-byte micro-step (2x MFMA) + + # --- B Load Logic (K64) --- + def load_b_pack(base_k, ki_step, ni): + return load_b_pack_k32( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ki_step=ki_step, + n_blk=n_blk_list[ni], + n_intra=n_intra_list[ni], + lane_div_16=lane_div_16, # 0..3 + elem_type=w_elem, + kpack_bytes=kpack_bytes, + elem_bytes=w_elem_bytes, + unpack_int4=is_int4, + ) + + def load_b_tile(base_k): + """Prefetch the entire per-thread B tile (gmem -> regs) for a given K base. + + Returns a list of length `k_unroll`, where each entry is a tuple: + (packs_half0[ni], packs_half1[ni]) for the K64 micro-step. + For groupwise variants, each entry also includes per-group scales: + (packs0[ni], packs1[ni], scales0[ni], scales1[ni]) + """ + if const_expr(is_int4_bf16_groupwise): + # W4A16 groupwise: load raw packed32 + scale; defer dequant to compute_tile. + raw_data = [] + for ku in range_constexpr(k_unroll): + raw_ku = [] + for ni in range_constexpr(num_acc_n): + packed32, scale_val = load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk_list[ni], + n_intra=n_intra_list[ni], + lane_div_16=lane_div_16, + elem_type=w_elem, + scale_rsrc=sw_rsrc, + expert_offset=expert_off_idx, + num_groups=num_groups, + group_size=group_size, + n_per_expert=model_dim, + kpack_bytes=kpack_bytes, + scale_dtype=scale_dtype, + ) + raw_ku.append((packed32, scale_val)) + raw_data.append(raw_ku) + return raw_data + elif const_expr(is_int4_bf16): + # W4A16 per-row: load raw packed32; defer dequant to compute_tile. + raw_data = [] + for ku in range_constexpr(k_unroll): + raw_ku = [] + for ni in range_constexpr(num_acc_n): + raw = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_w, + b_rsrc=w_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk_list[ni], + n_intra=n_intra_list[ni], + lane_div_16=lane_div_16, + elem_type=w_elem, + kpack_bytes=kpack_bytes, + ) + raw_ku.append(raw) + raw_data.append(raw_ku) + return raw_data + else: + # fp8/int8/bf16/fp16: original code path + b_tile = [] + for ku in range_constexpr(k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + ki0 = (ku * 2) + 0 + ki1 = (ku * 2) + 1 + b0 = load_b_pack(base_k, ki0, ni) + b1 = load_b_pack(base_k, ki1, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + # ---- 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): + lds_store_16b_xor16( + arith, + vector, + lds_memref=lds_x, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x4=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + elif const_expr(x_load_bytes == 8): + lds_store_8b_xor16( + arith, + vector, + lds_memref=lds_x, + vec8_ty=vec8_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x2=vec_x_in_parts[i], + ) + else: + lds_store_4b_xor16( + arith, + vector, + lds_memref=lds_x, + vec4_ty=vec4_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=fx.Index(4), + k_blocks16=k_blocks16, + lds_base=lds_base, + vec_part_i32x1=vec_x_in_parts[i], + ) + + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- + def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base): + col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16) + col_base_swz = ( + col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes))) + ) + idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds) + idx_a16 = idx_a16 + lds_base + loaded_a16 = vector.load(vec16_x, as_ir_value(lds_x), [as_ir_value(idx_a16)]) + a_i64x2 = vector.bitcast(T.i64x2, as_ir_value(loaded_a16)) + a0 = vector.extract(as_ir_value(a_i64x2), dynamic_position=[], static_position=[0]) + a1 = vector.extract(as_ir_value(a_i64x2), dynamic_position=[], static_position=[1]) + return a0, a1 + + def compute_tile( + acc_in, + b_tile_in, + lds_base, + *, + prefetch_epilogue: bool = False, + a0_prefetch=None, + ): + acc_list = list(acc_in) + mfma_res_ty = T.i32x4 if is_int8 else T.f32x4 + if const_expr(_use_mfma_k32): + mfma_fn = rocdl.mfma_f32_16x16x32_f16 if is_f16 else rocdl.mfma_f32_16x16x32_bf16 + else: + mfma_fn = ( + mfma_i32_k32 + if is_int8 + else ( + mfma_f32_bf16_k16 + if is_bf16 + else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) + ) + ) + + epilogue_pf = None + if const_expr(prefetch_epilogue and not use_groupwise_scale): + expert_off_pf = expert_off_idx + sw_pf = [] + for ni in range_constexpr(num_acc_n): + col_g = col_g_list[ni] + row_w_idx = expert_off_pf + col_g + sw_pf.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_w_idx, vec_width=1, dtype=T.f32) + ) + # Also prefetch per-row routed/topk weights (sorted_weights) when enabled. + tw_pf = None + if const_expr(doweight_stage2): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * fx.Index(4) + ii_idx_list_pf = [fx.Index(ii) for ii in range(4)] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.index(mi * 16) + for ii in range_constexpr(4): + row_off_pf = lane_div_16_mul4_pf + ii_idx_list_pf[ii] + 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=T.f32, + ) + ) + epilogue_pf = (sw_pf, tw_pf) + + def mfma_k64(acc0, a0, a1, b0, b1): + if const_expr(_use_mfma_k32): + # gfx950: single 16x16x32 MFMA consuming all 128 bits (K=32 f16/bf16) + if const_expr(is_f16): + av = _i64x2_to_v8f16(a0, a1) + bv = _i64x2_to_v8f16(b0, b1) + else: + av = _i64x2_to_v8bf16(a0, a1) + bv = _i64x2_to_v8bf16(b0, b1) + return mfma_fn(mfma_res_ty, [av, bv, acc0, 0, 0, 0]) + if const_expr(is_f16): + a0v = _i64_to_v4f16(a0) + a1v = _i64_to_v4f16(a1) + b0v = _i64_to_v4f16(b0) + b1v = _i64_to_v4f16(b1) + acc1 = mfma_fn(mfma_res_ty, [a0v, b0v, acc0, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1v, b1v, acc1, 0, 0, 0]) + if const_expr(is_bf16): + a0v = _i64_to_v4i16(a0) + a1v = _i64_to_v4i16(a1) + b0v = _i64_to_v4i16(b0) + b1v = _i64_to_v4i16(b1) + acc1 = mfma_fn(mfma_res_ty, [a0v, b0v, acc0, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1v, b1v, acc1, 0, 0, 0]) + acc1 = mfma_fn(mfma_res_ty, [a0, b0, acc0, 0, 0, 0]) + return mfma_fn(mfma_res_ty, [a1, b1, acc1, 0, 0, 0]) + + def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val): + """MFMA f32 partial -> scale -> add to f32 accumulator via math.fma on vector.""" + from flydsl._mlir.dialects._math_ops_gen import fma as _math_fma + + _uw = arith._to_raw + scale_vec = _uw(vector.broadcast(T.f32x4, scale_val)) + return arith.ArithValue(_math_fma(scale_vec, _uw(f32_partial_vec), _uw(f32_acc_vec))) + + if const_expr(is_int4_bf16 or is_int4_bf16_groupwise): + # W4A16: deferred dequant -- unpack int4->bf16 right before MFMA + # to minimize VGPR lifetime of dequantized bf16 values. + _pending_acc = None + for ku in range_constexpr(k_unroll): + b_raw = b_tile_in[ku] + ki64 = arith.index(ku * 64) + col_base = col_offset_base_bytes + ki64 + + for mi in range_constexpr(m_repeat): + mi_val = arith.index(mi * 16) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) + + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + if const_expr(is_int4_bf16_groupwise): + packed, sc = b_raw[ni] + if const_expr(_scale_is_bf16): + sc = extract_bf16_scale(arith, sc, ku) + else: + packed, sc = b_raw[ni], None + if const_expr(is_int4_bf16_groupwise and use_gfx950_cvt): + b0, b1 = unpack_b_w4a16( + packed, + arith, + vector, + scale_val=None, + use_gfx950_cvt=True, + defer_scale16=True, + ) + tmp = mfma_k64(zero_f32_acc, a0, a1, b0, b1) + if _pending_acc is not None: + p_idx, p_tmp, p_sc = _pending_acc + acc_list[p_idx] = _acc_scaled_f32(acc_list[p_idx], p_tmp, p_sc) + _pending_acc = (acc_idx, tmp, sc) + else: + b0, b1 = unpack_b_w4a16( + packed, + arith, + vector, + scale_val=sc, + use_gfx950_cvt=use_gfx950_cvt, + defer_scale16=use_gfx950_cvt, + ) + acc_list[acc_idx] = mfma_k64(acc_list[acc_idx], a0, a1, b0, b1) + # Drain last pending FMA. + if _pending_acc is not None: + p_idx, p_tmp, p_sc = _pending_acc + acc_list[p_idx] = _acc_scaled_f32(acc_list[p_idx], p_tmp, p_sc) + else: + for ku in range_constexpr(k_unroll): + b_packs0, b_packs1 = b_tile_in[ku] + ki64 = arith.index(ku * 64) + col_base = col_offset_base_bytes + ki64 + + for mi in range_constexpr(m_repeat): + mi_val = arith.index(mi * 16) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): + a0, a1 = a0_prefetch + else: + a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) + + for ni in range_constexpr(num_acc_n): + acc_idx = mi * num_acc_n + ni + acc_list[acc_idx] = mfma_k64( + acc_list[acc_idx], + a0, + a1, + b_packs0[ni], + b_packs1[ni], + ) + return acc_list, epilogue_pf + + # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- + lds_tile_elems = arith.index(tile_m * lds_stride) + lds_base_cur = fx.Index(0) + lds_base_nxt = lds_tile_elems + + rocdl.sched_barrier(0) + + # def hot_loop_scheduler(): + # mfma_group = num_acc_n + # # K64 micro-step: 2x K32 MFMA per accumulator update. + # 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) + # rocdl.sched_mfma(1) + # if num_acc_n < 4: + # rocdl.sched_dsrd(1) + # rocdl.sched_mfma(1) + # rocdl.sched_dsrd(1) + # rocdl.sched_mfma(1) + # rocdl.sched_vmem(1) + # rocdl.sched_mfma(1) + # rocdl.sched_vmem(1) + # rocdl.sched_mfma(2) + # rocdl.sched_dsrd(1) + # rocdl.sched_mfma(2) + # rocdl.sched_vmem(1) + + # 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_mfma(mfma_group // 2) + # rocdl.sched_dsrd(1) + # rocdl.sched_mfma(mfma_group // 2) + # rocdl.sched_vmem(1) + # rocdl.sched_mfma(mfma_group) + # if sche_i >= dswr_start - 1: + # rocdl.sched_dswr(1) + # rocdl.sched_barrier(0) + + def hot_loop_scheduler(): + rocdl.sched_barrier(0) + return + # - 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 const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + if const_expr(num_acc_n < 4): + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + + # DS-write hints near the end: match total A LDS-store micro-ops per thread. + dswr_tail = num_x_loads + if const_expr(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 const_expr(sche_i >= dswr_start - 1): + rocdl.sched_dswr(1) + + rocdl.sched_barrier(0) + + # Prologue. + k0 = fx.Index(0) + x_regs0 = load_x_tile(k0) + b_cur = load_b_tile(k0) + store_x_tile_to_lds(x_regs0, lds_base_cur) + 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 LDS prefetch (default-on): prefetch the first A-pack (K64) for the + # tile we are about to compute from LDS, to overlap with upcoming VMEM. + a0_prefetch_pong = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_pong) + + # Main loop: process K tiles in 2-tile ping-pong steps. + # + # 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) + 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): + k_main2_py = 0 + + arith.index(tile_k * 2) + c_tile_k_s2 = arith.index(tile_k) + pair_iters = k_main2_py // (int(tile_k) * 2) + + # B-tile data layout per k_unroll entry (3 variants): + # See gemm1 _flatten_b_tile for full layout documentation. + int4_bf16_single_field = is_int4_bf16 and not is_int4_bf16_groupwise + _fields_per_ku = 1 if int4_bf16_single_field else 2 + _vals_per_b_tile = k_unroll * _fields_per_ku * num_acc_n + _n_acc = m_repeat * num_acc_n + _p_b = _n_acc + _p_a0 = _p_b + _vals_per_b_tile + + def _flatten_b_tile(b_tile): + """Flatten B tile to a 1-D list for scf.for loop-carried state.""" + flat = [] + for ku_entry in b_tile: + if is_int4_bf16_groupwise: + flat.extend(t[0] for t in ku_entry) + flat.extend(t[1] for t in ku_entry) + elif int4_bf16_single_field: + flat.extend(ku_entry) + else: + flat.extend(ku_entry[0]) + flat.extend(ku_entry[1]) + return flat + + def _unflatten_b_tile(vals): + """Reconstruct B tile from flattened scf.for loop-carried state.""" + b_tile, idx = [], 0 + for _ in range_constexpr(k_unroll): + if is_int4_bf16_groupwise: + packed = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + scales = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + b_tile.append([(packed[ni], scales[ni]) for ni in range_constexpr(num_acc_n)]) + elif int4_bf16_single_field: + b_tile.append(list(vals[idx : idx + num_acc_n])) + idx += num_acc_n + else: + packs_even = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + packs_odd = list(vals[idx : idx + num_acc_n]) + idx += num_acc_n + b_tile.append((packs_even, packs_odd)) + return b_tile + + init_state = list(acc) + _flatten_b_tile(b_cur) + list(a0_prefetch_pong) + + for pair_iv, state in range(0, pair_iters, 1, init=init_state): + _ac = list(state[:_n_acc]) + _bc = _unflatten_b_tile(list(state[_p_b:_p_a0])) + _a0 = (state[_p_a0], state[_p_a0 + 1]) + + k_iv = pair_iv * (c_tile_k_s2 + c_tile_k_s2) + + next_k1 = k_iv + c_tile_k_s2 + x_regs_ping = load_x_tile(next_k1) + _bp = load_b_tile(next_k1) + + _ac, _ = compute_tile(_ac, _bc, lds_base_pong, a0_prefetch=_a0) + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + hot_loop_scheduler() + gpu.barrier() + + _a0p = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) + + next_k2 = k_iv + c_tile_k_s2 + c_tile_k_s2 + x_regs_pong = load_x_tile(next_k2) + _bn = load_b_tile(next_k2) + + _ac, _ = compute_tile(_ac, _bp, lds_base_ping, a0_prefetch=_a0p) + store_x_tile_to_lds(x_regs_pong, lds_base_pong) + hot_loop_scheduler() + gpu.barrier() + + _a0n = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_pong) + + loop_results = yield list(_ac) + _flatten_b_tile(_bn) + list(_a0n) + + SmemPtr._view_cache = None + if pair_iters > 0: + acc = list(loop_results[:_n_acc]) + b_cur = _unflatten_b_tile(list(loop_results[_p_b:_p_a0])) + a0_prefetch_pong = (loop_results[_p_a0], loop_results[_p_a0 + 1]) + + if const_expr(odd_k_tiles): + # Tail: single remaining tile (already in `b_cur` / `lds_base_pong`). + acc, epilogue_pf = compute_tile( + acc, + b_cur, + lds_base_pong, + prefetch_epilogue=True, + a0_prefetch=a0_prefetch_pong, + ) + else: + k_tail1 = k_in - tile_k + x_regs_ping = load_x_tile(k_tail1) + b_ping = load_b_tile(k_tail1) + + acc, _ = compute_tile(acc, b_cur, lds_base_pong, a0_prefetch=a0_prefetch_pong) + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + hot_loop_scheduler() + gpu.barrier() + + a0_prefetch_ping = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) + acc, epilogue_pf = compute_tile( + acc, + b_ping, + lds_base_ping, + prefetch_epilogue=True, + a0_prefetch=a0_prefetch_ping, + ) + + # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- + # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. + expert_off = expert_off_idx + mask24_i32 = fx.Int32(0xFFFFFF) + model_i32 = fx.Int32(model_dim) + topk_i32_v = topk_i32 + + zero_i32 = fx.Int32(0) + c2_i32 = fx.Int32(2) # 2B element size for f16/bf16 + mask_even_i32 = fx.Int32(0xFFFFFFFE) # align element index to even for half2 atomics + + e_vec = _e_vec + + def atomic_add_f16x2(val_f16x2, byte_off_i32): + buffer_atomic_add(val_f16x2, out_rsrc, byte_off_i32, zero_i32, zero_i32) + + sw_pf = None + tw_pf = None + if const_expr(epilogue_pf is not None): + sw_pf, tw_pf = epilogue_pf + + # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). + if const_expr(use_groupwise_scale): + # Groupwise: weight scale already applied per-group in K-loop. + sw_vals = [arith.constant(1.0, type=T.f32)] * num_acc_n + elif const_expr(sw_pf is not None): + sw_vals = sw_pf + else: + sw_vals = [] + for ni in range_constexpr(num_acc_n): + col_g = col_g_list[ni] + row_w_idx = expert_off + col_g + sw_vals.append( + fx.Float32(1.0) + if not needs_scale_w + else buffer_ops.buffer_load(sw_rsrc, row_w_idx, vec_width=1, dtype=T.f32) + ) + + # When defer_scale16 was used, the x16 correction for v_cvt_off_f32_i4 + # was omitted from the hot loop. Fold it into the epilogue scale. + if const_expr(use_gfx950_cvt): + _c16 = fx.Float32(16.0) + sw_vals = [v * _c16 for v in sw_vals] + + if const_expr(out_is_f32): + # origin/dev_a16w4: f32 output uses scalar f32 atomics and skips CShuffle/LDS. + c4_i32 = fx.Int32(4) + + def atomic_add_f32(val_f32, byte_off_i32): + buffer_atomic_add(val_f32, out_rsrc, byte_off_i32, zero_i32, zero_i32) + + def _stage2_row_atomic(*, mi: int, ii: int, row_in_tile, row): + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + + # Mask sentinel (token_id==tokens, slot==topk) to avoid OOB scale_x loads. + # For invalid rows, force sx=0 so they contribute exactly 0 to output. + t_ok = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32) + s_ok = arith.cmpi(arith.CmpIPredicate.ult, s2, topk_i32_v) + ts_ok = t_ok & s_ok + t2_safe = ts_ok.select(t2, fx.Int32(0)) + s2_safe = ts_ok.select(s2, fx.Int32(0)) + ts2 = t2_safe * topk_i32_v + s2_safe + sx = ( + arith.select(ts_ok, fx.Float32(1.0), fx.Float32(0.0)) + if is_f16_or_bf16 + else arith.select( + ts_ok, + buffer_ops.buffer_load(sx_rsrc, ts2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + + if const_expr(doweight_stage2): + tw_idx = (mi * 4) + ii + if const_expr(tw_pf is not None): + tw = ts_ok.select(tw_pf[tw_idx], fx.Float32(0.0)) + else: + tw = arith.select( + ts_ok, + buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + + idx0 = t2_safe * model_i32 # i32 element index base (safe for sentinel rows) + + for ni in range_constexpr(num_acc_n): + col_g = col_g_list[ni] + sw = sw_vals[ni] + acc_idx = mi * num_acc_n + ni + v = vector.extract(as_ir_value(acc[acc_idx]), dynamic_position=[], static_position=[ii]) + if const_expr(is_int8): + v = arith.sitofp(T.f32, v) + v = v * sx * sw + if const_expr(doweight_stage2): + v = v * tw + col_i32 = arith.index_cast(T.i32, col_g) + idx_elem = idx0 + col_i32 + byte_off = idx_elem * c4_i32 + atomic_add_f32(v, byte_off) + + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_stage2_row_atomic, + ) + else: + if const_expr(lds_out is None): + raise RuntimeError("FLYDSL_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased.") + + # For bf16 global atomics (gfx942 only), precompute the output base address. + # gfx950+ has buffer_atomic_pk_add_bf16, so bf16 uses buffer atomics there. + out_base_idx = None + if const_expr(_needs_global_atomic_bf16): + out_base_idx = buffer_ops.extract_base_index(arg_out) + + def write_row_to_lds( + *, + mi: int, + ii: int, + row_in_tile, + row, + row_base_lds, + col_base_local, + num_acc_n: int, + lds_out, + ): + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + # Explicitly mask sentinel token/slot to avoid OOB scale_x loads. + t_ok = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32) + s_ok = arith.cmpi(arith.CmpIPredicate.ult, s2, topk_i32_v) + ts_ok = t_ok & s_ok + t2_safe = ts_ok.select(t2, fx.Int32(0)) + s2_safe = ts_ok.select(s2, fx.Int32(0)) + ts2 = t2_safe * topk_i32_v + s2_safe + sx = ( + fx.Float32(1.0) + if is_f16_or_bf16 + else arith.select( + ts_ok, + buffer_ops.buffer_load(sx_rsrc, ts2, vec_width=1, dtype=T.f32), + fx.Float32(0.0), + ) + ) + + if const_expr(doweight_stage2): + tw_idx = (mi * 4) + ii + if const_expr(tw_pf is not None): + tw = tw_pf[tw_idx] + else: + tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) + + for ni in range_constexpr(num_acc_n): + col_local = col_base_local + (ni * 16) + sw = sw_vals[ni] + acc_idx = mi * num_acc_n + ni + v = vector.extract(as_ir_value(acc[acc_idx]), dynamic_position=[], static_position=[ii]) + if const_expr(is_int8): + v = arith.sitofp(T.f32, v) + v = v * sx * sw + if const_expr(doweight_stage2): + v = v * tw + v_out = arith.trunc_f(out_elem(), v) + + lds_idx = row_base_lds + col_local + vec1_out = T.vec(1, out_elem()) + v1 = vector.from_elements(vec1_out, [as_ir_value(v_out)]) + vector.store( + as_ir_value(v1), + as_ir_value(lds_out), + [as_ir_value(lds_idx)], + alignment=2, + ) + + def precompute_row(*, row_local, row): + # Precompute row context for cshuffle stores. + # Return (fused_i32, row_valid_i1) so the epilogue can skip the entire row + # for invalid tail rows (CK-style), avoiding per-store branching. + fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(arith.CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(arith.CmpIPredicate.ult, t, tokens_i32) + s_ok = arith.cmpi(arith.CmpIPredicate.ult, s, topk_i32_v) + row_valid = row_valid0 & t_ok & s_ok + return (fused2, row_valid) + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused = row_ctx + t = fused & mask24_i32 + s = fused >> 24 + idx0 = t * model_i32 + if const_expr(not bool(accumulate)): + ts = t * topk_i32_v + s + idx0 = ts * model_i32 + col_i32 = arith.index_cast(T.i32, col_g0) + idx_elem = idx0 + col_i32 + idx_elem_even = idx_elem & mask_even_i32 + if const_expr(_needs_global_atomic_bf16): + # gfx942: no buffer_atomic_pk_add_bf16, use global atomicrmw fadd + if const_expr(bool(accumulate)): + byte_off = idx_elem_even * c2_i32 + byte_off_idx = arith.index_cast(T.index, byte_off) + ptr_addr_idx = out_base_idx + byte_off_idx + out_ptr = buffer_ops.create_llvm_ptr(ptr_addr_idx, address_space=1) + out_ptr_v = out_ptr._value if const_expr(hasattr(out_ptr, "_value")) else out_ptr + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.AtomicRMWOp( + llvm.AtomicBinOp.fadd, + out_ptr_v, + frag_v, + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ) + else: + buffer_ops.buffer_store(frag, out_rsrc, idx_elem_even) + else: + # f16, or bf16 on gfx950+ (has buffer_atomic_pk_add_bf16) + byte_off = idx_elem_even * c2_i32 + if const_expr(bool(accumulate)): + atomic_add_f16x2(frag, byte_off) + else: + buffer_ops.buffer_store(frag, out_rsrc, idx_elem_even) + + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=tile_n, + e_vec=e_vec, + m_repeat=m_repeat, + num_acc_n=num_acc_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=by_n, + n_tile_base=n_tile_base, + lds_out=lds_out, + frag_elem_type=(T.bf16 if out_is_bf16 else T.f16), + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + _if_blk = scf.IfOp(blk_valid) + with _if_then(_if_blk): + _moe_gemm2_then_body() + + # ── Host launcher (flyc.jit + .launch) ──────────────────────────────── + @flyc.jit + def launch_moe_gemm2( + arg_out: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_scale_x: fx.Tensor, + arg_scale_w: fx.Tensor, + arg_sorted_token_ids: fx.Tensor, + arg_expert_ids: fx.Tensor, + arg_sorted_weights: fx.Tensor, + arg_num_valid_ids: fx.Tensor, + i32_tokens_in: fx.Int32, + i32_n_in: fx.Int32, + i32_k_in: fx.Int32, + i32_size_expert_ids_in: fx.Int32, + stream: fx.Stream, + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + n_in = arith.index_cast(T.index, i32_n_in) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) + gx = n_in // fx.Index(tile_n) + gy = size_expert_ids_in + + moe_gemm2( + arg_out, + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_num_valid_ids, + i32_tokens_in, + i32_n_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch( + grid=(gx, gy, 1), + block=(256, 1, 1), + stream=stream, + ) + + return launch_moe_gemm2 + + +class MoeGemm2Mode: + """Execution mode for MoE GEMM2.""" + + ATOMIC = "atomic" # Use atomic accumulation (default) + REDUCE = "reduce" # Use non-atomic write + reduce kernel + + +class _MoeGemm2ReduceWrapper: + """Wrapper combining GEMM2 (no atomics) with reduction kernel. + + This wrapper handles the intermediate buffer allocation and orchestrates + the two-phase computation: + 1. GEMM2 outputs to [tokens*topk, model_dim] without atomics + 2. Reduce sums over topk to produce [tokens, model_dim] + """ + + def __init__( + self, + gemm2_exe, + reduce_exe, + topk: int, + model_dim: int, + out_dtype_str: str = "f16", + use_mask: bool = False, + zero_intermediate: bool = True, + ): + self._gemm2_exe = gemm2_exe + self._reduce_exe = reduce_exe + self._topk = topk + self._model_dim = model_dim + self._out_dtype_str = out_dtype_str + self._use_mask = use_mask + self._zero_intermediate = zero_intermediate + + def _get_torch_dtype(self): + """Convert dtype string to torch dtype.""" + import torch + + dtype_map = { + "f16": torch.float16, + "fp16": torch.float16, + "bf16": torch.bfloat16, + "f32": torch.float32, + } + return dtype_map.get(self._out_dtype_str, torch.float16) + + def __call__( + self, + arg_out, + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_num_valid_ids, + tokens_in, + n_in, + k_in, + size_expert_ids_in, + expert_mask=None, + topk_ids=None, + stream=None, + ): + """Execute GEMM2 + reduce. + + Args match moe_gemm2 kernel signature (see compile_moe_gemm2). + When ``self._use_mask`` is True, expert_mask + topk_ids are required and + the reduction fuses ``valid = expert_mask[topk_ids[t, k]] != 0``. + """ + import torch + + if stream is None: + stream = torch.cuda.current_stream() + intermediate = torch.empty( + tokens_in * self._topk, + self._model_dim, + device=arg_out.device, + dtype=self._get_torch_dtype(), + ) + if self._zero_intermediate and not self._use_mask: + intermediate.zero_() + # Phase 1: GEMM2 (no atomics) -> [tokens*topk, model_dim] + self._gemm2_exe( + intermediate.view(-1), + arg_x, + arg_w, + arg_scale_x, + arg_scale_w, + arg_sorted_token_ids, + arg_expert_ids, + arg_sorted_weights, + arg_num_valid_ids, + tokens_in, + n_in, + k_in, + size_expert_ids_in, + stream, + ) + # Phase 2: Reduce over topk -> [tokens, model_dim]. The reduce launcher + # takes fx.Pointer args (base-ptr folding keeps voffsets i32-safe for + # X > 4 GiB), so dispatch it through the _run_compiled pointer shim. + X = intermediate.view(tokens_in, self._topk, self._model_dim) + Y = arg_out.view(tokens_in, self._model_dim) + if self._use_mask: + if expert_mask is None or topk_ids is None: + raise ValueError("expert_mask and topk_ids are required when use_mask=True") + em = expert_mask.to(torch.int32).contiguous() + tk = topk_ids.to(torch.int32).contiguous() + else: + # Placeholders; kernel ignores them when use_mask=False (compile-time). + em = torch.empty(0, device=arg_out.device, dtype=torch.int32) + tk = torch.empty(0, device=arg_out.device, dtype=torch.int32) + + def _ptr_arg(t): + type_name = type(t).__name__ + module_name = type(t).__module__ + if type_name == "FakeTensor" or "fake_tensor" in module_name: + return flyc.from_c_void_p(fx.Uint8, 0) + return flyc.from_c_void_p(fx.Uint8, t.data_ptr()) + + _run_compiled( + self._reduce_exe, + _ptr_arg(X), + _ptr_arg(Y), + _ptr_arg(em), + _ptr_arg(tk), + tokens_in, + stream, + ) + + @property + def mode(self) -> str: + """Return the execution mode.""" + return MoeGemm2Mode.REDUCE + + +def compile_moe_gemm2_ex( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + in_dtype: str = "fp8", + group_size: int = -1, + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + # Extended parameters for mode control + mode: str = MoeGemm2Mode.ATOMIC, + use_mask: bool = False, + zero_intermediate: bool = True, + scale_is_bf16: bool = False, +): + """Compile MoE GEMM2 kernel with optional reduction. + + This is the extended interface that supports explicit mode control. + + Args: + mode: Execution mode selection: + - "atomic": Use atomic accumulation (original behavior) + - "reduce": Use non-atomic write + reduce kernel + + use_mask: If True, the reduction kernel fuses the EP gather + ``valid = expert_mask[topk_ids[t, k]] != 0`` and only sums + valid slots. Caller must pass expert_mask + topk_ids at call time. + + zero_intermediate: If all output slots are valid, + set False to increase performance + + Returns: + Compiled executable (either wrapped or raw depending on mode). + """ + # Compile based on mode + if mode == MoeGemm2Mode.REDUCE: + # Compile GEMM2 with accumulate=False + gemm2_exe = compile_moe_gemm2( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=doweight_stage2, + in_dtype=in_dtype, + group_size=group_size, + out_dtype=out_dtype, + use_cshuffle_epilog=use_cshuffle_epilog, + accumulate=False, + scale_is_bf16=scale_is_bf16, + ) + # Compile reduction kernel with masking support + out_s = str(out_dtype).strip().lower() + if out_s in ("f16", "fp16", "half"): + dtype_str = "f16" + elif out_s in ("bf16", "bfloat16"): + dtype_str = "bf16" + else: + dtype_str = "f32" + reduce_exe = compile_moe_reduction( + topk=topk, + model_dim=model_dim, + dtype_str=dtype_str, + use_mask=use_mask, + # expert_mask is sized by global expert count (≠ w2.shape[0] under EP). + num_experts=experts, + ) + return _MoeGemm2ReduceWrapper( + gemm2_exe=gemm2_exe, + reduce_exe=reduce_exe, + topk=topk, + model_dim=model_dim, + out_dtype_str=dtype_str, + use_mask=use_mask, + zero_intermediate=zero_intermediate, + ) + else: + # Compile GEMM2 with accumulate=True (atomic mode) + return compile_moe_gemm2( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=doweight_stage2, + in_dtype=in_dtype, + group_size=group_size, + out_dtype=out_dtype, + use_cshuffle_epilog=use_cshuffle_epilog, + accumulate=True, + ) diff --git a/kernels/moe/moe_gemm_2stage/layout_helpers.py b/kernels/moe/moe_gemm_2stage/layout_helpers.py new file mode 100644 index 000000000..6f5cf6d11 --- /dev/null +++ b/kernels/moe/moe_gemm_2stage/layout_helpers.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +# Portions Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Layout-API helper layer for the MoE 2-stage MFMA kernels (gemm1.py / gemm2.py), +ported from the aiter reference kernel to this repo's ``fx.*`` surface.""" + +import flydsl.expr as fx +from flydsl._mlir.dialects import rocdl +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import _to_raw as _raw + + +def reps(tensor, mode): + """Static repeat count of ``tensor``'s ``mode`` (shape size, as a Python int).""" + return fx.size(fx.get_shape(tensor)[mode]).to_py_value() + + +def _encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=63): + """Encode s_waitcnt bitfield for CDNA3 (gfx94x).""" + vm_lo = vmcnt & 0xF + vm_hi = (vmcnt >> 4) & 0x3 + return vm_lo | (expcnt << 4) | (lgkmcnt << 8) | (vm_hi << 14) + + +def _as_ptr(p, dtype=None): + """Iterator for ``fx.make_view`` from a raw pointer or a runtime memref (opt. recast).""" + try: + p = fx.get_iter(p) + finally: + if dtype is not None and p.dtype != dtype: + p = fx.recast_iter(dtype, p) + return p # noqa: B012 + + +def torch_layout(*shape): + if len(shape) == 1: + return fx.make_layout(shape[0], 1) + order = [i for i in range(len(shape) - 1, -1, -1)] + return fx.make_ordered_layout(shape, order) + + +def view_as_torch_tensor(ptr, shape, dtype=None): + ptr = _as_ptr(ptr, dtype) + return fx.make_view(ptr, torch_layout(*shape)) + + +# ── Native-fp8 (MFMA 16x16x32) gate-up building blocks ─────────────────────── +# Ported from the aiter reference kernel's native-fp8 prefill_1x4 gate-up path, +# with the compile-time closures (N, K, TOPK, BLOCK_M, weight_dtype) made +# explicit args so the helpers are reusable across tile configs. + + +def _buffer_atomic_pk(rsrc, elem_idx, reg_vec, elem_bytes): + """Pairwise buffer atomic-add of an f16/bf16 vector into out[elem_idx..] + (buffer rsrc + byte offset; OOB lanes dropped by hardware clamp).""" + from kernels.common.mem_ops import buffer_atomic_add + + _z = fx.Int32(0) + for i in range_constexpr(reg_vec.numel // 2): + pair = Vec.from_elements([reg_vec[i * 2], reg_vec[i * 2 + 1]], reg_vec.dtype) + byte_off = (elem_idx + fx.Int32(i * 2)) * fx.Int32(elem_bytes) + buffer_atomic_add(pair, rsrc, byte_off, _z, _z) + + +def _buffer_atomic_f32(rsrc, elem_idx, reg_vec): + """Scalar buffer atomic-add of an f32 vector into out[elem_idx..].""" + from kernels.common.mem_ops import buffer_atomic_add + + _z = fx.Int32(0) + for i in range_constexpr(reg_vec.numel): + byte_off = (elem_idx + fx.Int32(i)) * fx.Int32(4) + buffer_atomic_add(reg_vec[i], rsrc, byte_off, _z, _z) + + +def make_1x4_tiled_mma(weight_dtype): + """B-first 1x4 tiled_mma (weight=A, activation=B; 4 waves tile the M/channel dim). + fp8 and gfx950 bf16 both use native MFMA(16,16,32), same k_perm, differing only in dtype.""" + mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 32, weight_dtype)) + k_perm = fx.make_layout((8, 4, 2), (1, 16, 8)) + tiled_mma = fx.make_tiled_mma( + mma_atom, + fx.make_layout((4, 1, 1), (1, 0, 0)), + fx.make_tile(None, None, k_perm), + ) + return mma_atom, tiled_mma + + +def make_gateup_weight_view(p_weight, expert_id, contiguous_n, N, K): + """Per-expert logical (N,K) view over the shuffle_weight-ordered weight, composed + with the gate/up silu grouping (N = 2*inter_dim).""" + group_layout_silu = fx.make_layout( + ((contiguous_n, 2, N // (contiguous_n * 2)), K), + ((1, N // 2, contiguous_n), N), + ) + element_num = 16 // (p_weight.dtype.width // 8) + return fx.make_view( + p_weight + fx.Int64(expert_id * N * K), + fx.composition( + fx.make_layout( + ((16, N // 16), (element_num, K // element_num)), + ((element_num, 16 * K), (1, 16 * element_num)), + ), + group_layout_silu, + ), + ) + + +def make_weight_view(p_weight, expert_id, N, K): + """Per-expert logical (N,K) view over shuffle_weight-ordered weight, no gate/up + grouping (stage2 analog of make_gateup_weight_view; N=model_dim, K=inter_dim).""" + element_num = 16 // (p_weight.dtype.width // 8) + return fx.make_view( + p_weight + fx.Int64(expert_id * N * K), + fx.make_layout( + ((16, N // 16), (element_num, K // element_num)), + ((element_num, 16 * K), (1, 16 * element_num)), + ), + ) + + +def read_sorted_index(tiled_copy_index, tid, lds_index, index_size, index_offset=0): + """Read the sorted M-row index from LDS into a per-thread fragment (explicit so it + happens before the CShuffle epilogue overwrites sorted_lds).""" + lds = fx.make_view(lds_index.ptr + index_offset, fx.make_layout(index_size, 1)) + cp_atom_lds = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32) + lds_thr = tiled_copy_index.get_slice(tid).partition_S(lds) + index_frag = fx.make_fragment_like(lds_thr) + fx.copy(cp_atom_lds, lds_thr, index_frag) + return index_frag + + +def silu_pair_bf16(gate_frag, up_frag, gate_scale=None, up_scale=None, a_scale=None, out_dtype=fx.BFloat16): + """silu(gate)*up -> out_dtype (optional fp8 weight/act scales folded in pre-silu). + out_dtype MUST match the caller's CShuffle staging/store dtype: the fragment holds + raw bits, so a mismatch silently reinterprets them (1024.0 bf16 0x4480 -> f16 4.5).""" + log2_exp1 = -1.4426950408889634 + round_bit = fx.Uint32(0x8000) + out_frag = fx.make_fragment_like(gate_frag, dtype=out_dtype) + m_reps = reps(gate_frag, 1) + n_reps = reps(gate_frag, 2) + for m in range_constexpr(m_reps): + if const_expr(a_scale is not None): + a_sc = a_scale[m] + for n in range_constexpr(n_reps): + gate = gate_frag[None, m, n].load() + up = up_frag[None, m, n].load() + if const_expr(gate_scale is not None): + sc_g = gate_scale[None, n].load() + sc_u = up_scale[None, n].load() + acc = [] + for j in range_constexpr(gate.numel): + g = gate[j] + u = up[j] + if const_expr(gate_scale is not None): + g = g * sc_g[j] + u = u * sc_u[j] + if const_expr(a_scale is not None): + g = g * a_sc + u = u * a_sc + tmp = rocdl.exp2(T.f32, _raw(g * log2_exp1)) + acc.append((g * rocdl.rcp(T.f32, 1.0 + tmp)) * u) + acc = Vec.from_elements(acc, fx.Float32) + if const_expr(out_dtype == fx.BFloat16): + acc = ((acc.bitcast(fx.Uint32) + round_bit) >> 16).to(fx.Uint16).bitcast(fx.BFloat16) + else: + acc = acc.to(out_dtype) + out_frag[None, m, n].store(acc) + return out_frag + + +def make_tensor_with_index(view, tile_m, tile_k, index_frag, tiled_copy, tid, topk, is_read_from_mem=True): + """MoE gather/scatter helper: returns an object whose ``.copy(copy_atom, k_idx, frag)`` + gathers/scatters per-thread tiles by ``index_frag`` (packed token|slot ids).""" + return _TensorWithIndex(view, tile_m, tile_k, index_frag, tiled_copy, tid, topk, is_read_from_mem) + + +class _TensorWithIndex: + def __init__(self, view, tile_m, tile_k, index_frag, tiled_copy, tid, topk, is_read_from_mem=True): + self.view = view + self.tile_m = tile_m + self.tile_k = tile_k + self.is_read_from_mem = is_read_from_mem + self.TOPK = topk + self.index_frag = index_frag + + rank = fx.get_shape(self.view).rank + dims = [1] * (rank - 1) + self.tensor_blocks_in_k = fx.zipped_divide(view, fx.make_tile(*dims, tile_k)) + + dtype = fx.PointerType.get(fx.Int8.ir_type, 1, 512) + ptr = fx.inttoptr(dtype, fx.Int32(0)) + self.fake_tensor = fx.make_view(ptr, fx.make_layout((tile_m, tile_k), (1, tile_m))) + self.fake_tensor_thr = ( + tiled_copy.get_slice(tid).partition_S(self.fake_tensor) + if is_read_from_mem + else tiled_copy.get_slice(tid).partition_D(self.fake_tensor) + ) + offset_thread = fx.Int32(fx.ptrtoint(fx.get_iter(self.fake_tensor_thr))) + self.offset_thread = offset_thread + self.offset_thread_k = offset_thread // tile_m + # Row-guard fake: a tall column-major tile whose row count exceeds any + # tiled_copy grid, so partitioning does NOT wrap OOB grid rows into the + # column dim. Lets the atomic epilogue detect grid slots whose row is + # outside [0, tile_m) (the plain-store path ignores this via buffer OOB). + self._guard_rows = 256 + guard_fake = fx.make_view(ptr, fx.make_layout((self._guard_rows, tile_k), (1, self._guard_rows))) + guard_thr = ( + tiled_copy.get_slice(tid).partition_S(guard_fake) + if is_read_from_mem + else tiled_copy.get_slice(tid).partition_D(guard_fake) + ) + self.guard_offset = fx.Int32(fx.ptrtoint(fx.get_iter(guard_thr))) + self.guard_layout = fx.get_layout(guard_thr) + + def copy( + self, + copy_atom, + k_idx, + frag, + atomic=None, + atomic_rsrc=None, + out_bytes=None, + row_stride=None, + row_limit=None, + ): + """Gather/scatter per-thread tiles: plain buffer-view store, or (atomic in + {"pk","f32"}) buffer atomic-add into atomic_rsrc at tok*row_stride+k_idx*tile_k + +chan; sentinel/out-of-tile lanes go OOB (dropped by the buffer clamp).""" + layout = fx.get_layout(self.fake_tensor_thr) + rep_m = reps(self.fake_tensor_thr, 1) + rep_k = reps(self.fake_tensor_thr, 2) + value_size = fx.get_shape(frag)[0].to_py_value() + stride_size = fx.get_stride(frag)[0].to_py_value() + + rank = fx.get_shape(self.view).rank + block_cord = [None] * (rank - 1) + [k_idx] + tensor_block = self.tensor_blocks_in_k[None, (*block_cord,)] + for m in range_constexpr(rep_m): + if const_expr(atomic is not None): + tok = self.index_frag[0, m] & 0xFFFFFF + if const_expr(row_limit is not None): + tok = (tok < row_limit).select(tok, fx.Int32(0)) + row_base_i32 = tok * fx.Int32(row_stride) + fx.Int32(k_idx) * fx.Int32(self.tile_k) + for k in range_constexpr(rep_k): + offset_block = fx.crd2idx((0, m, k), layout).to_py_value() + offset_block_k = offset_block // self.tile_m + chan_off = offset_block_k + self.offset_thread_k + # `valid` = this grid slot maps inside the real (tile_m, tile_k) block. + guard_full = fx.crd2idx((0, m, k), self.guard_layout).to_py_value() + self.guard_offset + g_row = guard_full % fx.Int32(self._guard_rows) + g_col = guard_full // fx.Int32(self._guard_rows) + valid = (g_row < fx.Int32(self.tile_m)) & (g_col < fx.Int32(self.tile_k)) + reg_vec = frag[None, m, k].load() + # Out-of-tile lanes -> OOB element index so the buffer bounds-check DROPS + # the atomic; redirecting them to out[0] serializes padding lanes (~600x). + _va = reg_vec.numel + aligned = (row_base_i32 + chan_off) & fx.Int32(~(_va - 1)) + elem_idx = valid.select(aligned, fx.Int32(row_limit) * fx.Int32(row_stride)) + if const_expr(atomic == "f32"): + _buffer_atomic_f32(atomic_rsrc, elem_idx, reg_vec) + else: + _buffer_atomic_pk(atomic_rsrc, elem_idx, reg_vec, out_bytes) + continue + if const_expr(rank == 2): + tensor_sub_block = tensor_block[None, self.index_frag[0, m] & 0xFFFFFF] + else: + tensor_sub_block = tensor_block[ + None, + self.index_frag[0, m] & 0xFFFFFF, + (self.index_frag[0, m] >> 24), + ] + for k in range_constexpr(rep_k): + offset_block = fx.crd2idx((0, m, k), layout).to_py_value() + offset_block_k = offset_block // self.tile_m + offset_k_in_tile = offset_block_k + self.offset_thread_k + reg = frag[None, m, k] + mem = fx.make_view( + fx.get_iter(tensor_sub_block) + offset_k_in_tile, + fx.make_layout(value_size, stride_size), + ) + if const_expr(self.is_read_from_mem): + fx.copy(copy_atom, mem, reg) + else: + fx.copy(copy_atom, reg, mem) + + +_TensorWithIndex.copy = ASTRewriter.transform(_TensorWithIndex.copy) diff --git a/kernels/moe/moe_gemm_2stage/moe_reduce.py b/kernels/moe/moe_gemm_2stage/moe_reduce.py new file mode 100644 index 000000000..94b3785b4 --- /dev/null +++ b/kernels/moe/moe_gemm_2stage/moe_reduce.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""MoE topk-reduction kernel (FlyDSL, layout API). + +``Y[t, d] = sum_k X[t, k, d]``, optionally gated by the EP validity mask +(``valid[t,k] = expert_mask[topk_ids[t,k]] != 0``). Epilogue of stage2 +``mode="reduce"``, shared by every dtype's reduce path. Extracted from the +legacy ``reduction.py``. Build a per-shape launcher with ``compile_moe_reduction`` +(cached); the kernel's compile-time params are ``Constexpr`` so flyc specializes +per shape/dtype. + +``dtype_str="fp8"`` reduces MXFP8 route-out rows (a flat uint8 buffer of +``[model_dim fp8 bytes | model_dim/8 e8m0 scale bytes]`` per row): each fp8 +value is scaled by its e8m0 microscale, accumulated in f32 and written to +``out_dtype_str`` (bf16/f16). The dense (f32/f16/bf16) path reduces a +contiguous ``X[tokens, topk, model_dim]`` tensor. + +The launcher takes ``fx.Pointer`` args; dispatch it through +``kernels.common.tensor_shim._run_compiled`` (each shape is a distinct launcher +object, so the shim's per-exe ``_cf`` cache stays correct). +""" + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr, gpu, ptrtoint, range_constexpr +from flydsl.expr.typing import T + +BLOCK = 256 +FP8_VEC = 8 # fp8 values per 64b buffer load (also the store granularity) + + +@flyc.kernel +def moe_reduction_kernel( + X: fx.Pointer, + Y: fx.Pointer, + expert_mask: fx.Pointer, + topk_ids: fx.Pointer, + i32_m_tokens: fx.Int32, + topk: fx.Constexpr[int], + model_dim: fx.Constexpr[int], + dtype_str: fx.Constexpr[str], + use_mask: fx.Constexpr[bool], + num_experts: fx.Constexpr[int], + out_dtype_str: fx.Constexpr[str], +): + # One tiled-copy reduce for every dtype. Dense (f16/bf16/f32) loads V elems + # and extends to f32; fp8 route-out loads 8 fp8 bytes + their e8m0 microscale + # and decodes to f32. Both then run the same masked f32 topk-accumulate + # (uniform soffset = k*row_stride) and truncating store. row_stride differs: + # an fp8 row is padded with its N/8 scale bytes ([N fp8 | N/8 e8m0]). + is_fp8 = dtype_str == "fp8" + if const_expr(is_fp8): + in_elem, in_bytes, V = fx.Int8, 1, FP8_VEC + row_stride = model_dim + model_dim // 8 + out_numeric = fx.Float16 if (out_dtype_str or "bf16") == "f16" else fx.BFloat16 + load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), fx.Int8) + else: + in_elem = fx.Float32 if dtype_str == "f32" else (fx.Float16 if dtype_str == "f16" else fx.BFloat16) + in_bytes = 4 if dtype_str == "f32" else 2 + row_stride, out_numeric = model_dim, in_elem + V = 128 // (8 * in_bytes) # 4 (f32), 8 (16b) + load_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), in_elem) + out_bytes = out_numeric.width // 8 + is_16b = out_numeric.width < 32 + TILE = BLOCK * V + store_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), out_numeric) + + token, tile, tid = gpu.block_id("x"), gpu.block_id("y"), gpu.thread_id("x") + tok64 = fx.Int64(token) + vec_f32, vec_out = T.vec(V, T.f32), T.vec(V, out_numeric.ir_type) + + def _view(elem, ptr_i64, ncols, nbytes): # 2D [1, ncols] V# buffer descriptor + pt = fx.PointerType.get( + elem.ir_type, + address_space=fx.AddressSpace.Global, + alignment=elem.width // 8, + ) + view = fx.make_view(fx.inttoptr(pt, ptr_i64), fx.make_layout((1, ncols), (ncols, 1))) + return fx.rocdl.make_buffer_tensor(view, num_records_bytes=fx.Int64(nbytes)) + + # Fold the per-token byte offset into the base ptr: keeps voffsets i32-safe + # for X > 4 GiB. + x_row_bytes = topk * row_stride * in_bytes + xbase = fx.Int64(ptrtoint(X)) + tok64 * fx.Int64(x_row_bytes) + xbuf = _view(in_elem, xbase, model_dim, x_row_bytes) + ybuf = _view( + out_numeric, + fx.Int64(ptrtoint(Y)) + tok64 * fx.Int64(model_dim * out_bytes), + model_dim, + model_dim * out_bytes, + ) + if const_expr(is_fp8): + # e8m0 scales trail the values in each row (one byte per FP8_VEC elems). + scbuf = _view( + fx.Int8, + xbase + fx.Int64(model_dim), + model_dim // 8, + x_row_bytes - model_dim, + ) + if const_expr(use_mask): + i32pt = fx.PointerType.get(T.i32, address_space=fx.AddressSpace.Global, alignment=4) + tk_ptr = fx.inttoptr(i32pt, fx.Int64(ptrtoint(topk_ids)) + tok64 * fx.Int64(topk * 4)) + em_ptr = fx.inttoptr(i32pt, fx.Int64(ptrtoint(expert_mask))) + + # Tiled copy: BLOCK threads across the tile, V contiguous elems per thread. + tile_mn, tv_layout = fx.make_layout_tv(fx.make_layout((1, BLOCK), (1, 1)), fx.make_layout((1, V), (1, 1))) + thr_load = fx.make_tiled_copy(load_atom, tv_layout, tile_mn).get_slice(tid) + thr_store = fx.make_tiled_copy(store_atom, tv_layout, tile_mn).get_slice(tid) + if const_expr(is_fp8): + sc_atom = fx.make_copy_atom(fx.rocdl.BufferCopy8b(), fx.Int8) + sc_tile, sc_tv = fx.make_layout_tv(fx.make_layout((1, BLOCK), (1, 1)), fx.make_layout((1, 1), (1, 1))) + thr_sc = fx.make_tiled_copy(sc_atom, sc_tv, sc_tile).get_slice(tid) + + def _decode_fp8(vfrag, sfrag): # 8 fp8 bytes + 1 e8m0 -> Vector(8, f32) + w = fx.Vector(fx.memref_load_vec(vfrag)).bitcast(fx.Int32) + e8m0 = fx.Uint32(fx.Uint8(fx.Vector(fx.memref_load_vec(sfrag))[0])) + scale = (e8m0 << fx.Uint32(23)).bitcast(fx.Float32) + words = (w[0], w[0], w[1], w[1]) + lanes = [] + for pi in range_constexpr(4): + pair = fx.Vector(fx.rocdl.cvt_pk_f32_fp8(T.f32x2, words[pi], bool(pi & 1))) + lanes.append(pair[0] * scale) + lanes.append(pair[1] * scale) + return fx.Vector.from_elements(lanes, fx.Float32) + + def _reduce_tile(): + p_src = thr_load.partition_S(fx.slice(fx.zipped_divide(xbuf, tile_mn), (None, (0, tile)))) + p_dst = thr_store.partition_D(fx.slice(fx.zipped_divide(ybuf, tile_mn), (None, (0, tile)))) + # topk rows share one per-thread voffset via a uniform scalar + # soffset = k*row_stride, so the loads issue back-to-back. + frags = [fx.make_fragment_like(p_src) for _ in range_constexpr(topk)] + for k in range_constexpr(topk): + fx.copy(load_atom, p_src, frags[k], soffset=fx.Int32(k * row_stride)) + if const_expr(is_fp8): + p_sc = thr_sc.partition_S(fx.slice(fx.zipped_divide(scbuf, sc_tile), (None, (0, tile)))) + sfrags = [fx.make_fragment_like(p_sc) for _ in range_constexpr(topk)] + for k in range_constexpr(topk): + fx.copy(sc_atom, p_sc, sfrags[k], soffset=fx.Int32(k * row_stride)) + + acc = fx.Vector.filled(V, 0.0, fx.Float32) + for k in range_constexpr(topk): + if const_expr(is_fp8): + vk = _decode_fp8(frags[k], sfrags[k]) + else: + vk = fx.Vector(fx.memref_load_vec(frags[k])) + vk = vk.extf(vec_f32) if is_16b else vk + if const_expr(use_mask): + vk = (em_ptr[tk_ptr[k]] != fx.Int32(0)).select(vk, fx.Vector.filled(V, 0.0, fx.Float32)) + acc = acc + vk + ofrag = fx.make_fragment_like(p_dst) + fx.memref_store_vec(acc.truncf(vec_out) if is_16b else acc, ofrag) + fx.copy(store_atom, ofrag, p_dst) + + # Skip threads whose column group starts past model_dim (their loads would + # read the next row -- in-descriptor, wasted BW); only needed when TILE ∤ md. + if const_expr(model_dim % TILE != 0): + if fx.Int32(tile) * fx.Int32(TILE) + fx.Int32(tid) * fx.Int32(V) < fx.Int32(model_dim): + _reduce_tile() + else: + _reduce_tile() + + +@functools.lru_cache(maxsize=1024) +def compile_moe_reduction( + *, + topk: int, + model_dim: int, + dtype_str: str = "f16", + use_mask: bool = False, + num_experts: int = 0, + out_dtype_str: str | None = None, +): + """Compile the topk-reduce launcher for one Constexpr set (cached per shape). + + Returns a ``@flyc.jit`` taking ``(X, Y, expert_mask, topk_ids, i32_m_tokens, + stream)``; dispatch it through ``kernels.common.tensor_shim._run_compiled``. + The launcher is a distinct object per shape, so the shim's per-exe ``_cf`` + cache stays correct. + """ + V = FP8_VEC if dtype_str == "fp8" else 128 // (32 if dtype_str == "f32" else 16) + gy = (model_dim + BLOCK * V - 1) // (BLOCK * V) + out_tag = out_dtype_str or dtype_str + + @flyc.jit + def launch( + X: fx.Pointer, + Y: fx.Pointer, + expert_mask: fx.Pointer, + topk_ids: fx.Pointer, + i32_m_tokens: fx.Int32, + stream: fx.Stream, + ): + moe_reduction_kernel( + X, + Y, + expert_mask, + topk_ids, + i32_m_tokens, + topk, + model_dim, + dtype_str, + use_mask, + num_experts, + out_tag, + ).launch(grid=(fx.Int64(i32_m_tokens), gy, 1), block=(BLOCK, 1, 1), stream=stream) + + return launch diff --git a/tests/kernels/test_moe_gemm_2stage.py b/tests/kernels/test_moe_gemm_2stage.py new file mode 100644 index 000000000..8162a516d --- /dev/null +++ b/tests/kernels/test_moe_gemm_2stage.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +""" +Basic smoke tests for the ported ``moe_gemm_2stage`` fp8 2-stage kernels. + +These verify that the stage1 (``compile_moe_gemm1``) and stage2 +(``compile_moe_gemm2``) builders trace and lower without error for small, +tile-valid shapes on both the fp8 and bf16 paths. Full numerical / e2e +coverage lives in the routing-based MoE tests; this file just guards the +port's builders against API drift after the layout-API migration. +""" + +import os +import sys + +import pytest +import torch + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +for _p in (os.path.join(_REPO_ROOT, "build", "python_packages"), _REPO_ROOT): + if os.path.isdir(_p) and _p not in sys.path: + sys.path.insert(0, _p) + +from flydsl.runtime.device import get_rocm_arch # noqa: E402 +from kernels.moe.moe_gemm_2stage import compile_moe_gemm1, compile_moe_gemm2 # noqa: E402 + +if not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + + +def _fp8_supported() -> bool: + arch = get_rocm_arch() + return ("gfx95" in arch) or ("gfx94" in arch) + + +# Small, tile-valid shape (mirrors the "S" case of the former 2-stage suite): +# model_dim=256, inter_dim=128, experts=4, topk=2 +# stage1 tile = (tile_m=16, tile_n1=64, tile_k1=128) +# stage2 tile = (tile_m=16, tile_n2=64, tile_k2=128) +_SHAPE = dict(model_dim=256, inter_dim=128, experts=4, topk=2) + + +@pytest.mark.parametrize("in_dtype", ["fp8", "bf16"]) +def test_moe_gemm1_builds(in_dtype: str): + """Stage1 builder traces/lowers for a small tile-valid shape.""" + if in_dtype == "fp8" and not _fp8_supported(): + pytest.skip("fp8 stage1 requires gfx94*/gfx95*") + + exe = compile_moe_gemm1( + **_SHAPE, + tile_m=16, + tile_n=64, + tile_k=128, + doweight_stage1=False, + in_dtype=in_dtype, + out_dtype="f16", + ) + assert callable(exe) + + +@pytest.mark.parametrize("in_dtype", ["fp8", "bf16"]) +def test_moe_gemm2_builds(in_dtype: str): + """Stage2 builder traces/lowers for a small tile-valid shape.""" + if in_dtype == "fp8" and not _fp8_supported(): + pytest.skip("fp8 stage2 requires gfx94*/gfx95*") + + exe = compile_moe_gemm2( + **_SHAPE, + tile_m=16, + tile_n=64, + tile_k=128, + doweight_stage2=False, + in_dtype=in_dtype, + out_dtype="f16", + ) + assert callable(exe) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/kernels/test_moe_reduce.py b/tests/kernels/test_moe_reduce.py new file mode 100644 index 000000000..33e1568bc --- /dev/null +++ b/tests/kernels/test_moe_reduce.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +""" +MoE Reduction Kernel Test + +Reduces [tokens, topk, model_dim] along the topk dimension. +Designed for MoE stage-2 shapes where topk is small +and model_dim is large and aligned (e.g. 5120, 7168). + +MoeReduce(x) = sum(x, dim=1) +""" + +import argparse +import logging +import os +import sys + +import pytest +import torch + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +_PYTHON_CANDIDATES = [ + os.path.join(_REPO_ROOT, "build", "python_packages"), + _REPO_ROOT, +] +for _p in reversed(_PYTHON_CANDIDATES): + if os.path.isdir(_p) and _p not in sys.path: + sys.path.insert(0, _p) + +import flydsl.compiler as flyc # noqa: E402 +import flydsl.expr as fx # noqa: E402 +from kernels.common.tensor_shim import _run_compiled # noqa: E402 +from kernels.moe.moe_gemm_2stage import compile_moe_reduction # noqa: E402 +from tests.test_common import run_perftest, verify_output # noqa: E402 + +logging.basicConfig(level=logging.INFO) + +if not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + + +def _ptr(t): + """torch tensor -> fx pointer arg for the fx.Pointer reduce launcher.""" + return flyc.from_c_void_p(fx.Uint8, t.data_ptr()) + + +def run_reduce_test( + tokens: int, + topk: int, + model_dim: int, + dtype_str: str = "f16", + use_mask: bool = False, + num_experts: int = 8, + num_iters: int = 20, + num_warmup: int = 5, + compare_torch: bool = True, +): + """Run reduce kernel test: correctness + performance. + + Masking uses the fused EP gather ``valid[t, k] = expert_mask[topk_ids[t, k]] + != 0``; the kernel takes ``expert_mask``/``topk_ids`` (not a precomputed + mask). The launcher takes ``fx.Pointer`` args, dispatched via ``_run_compiled``. + """ + dtype_map = {"f16": torch.float16, "bf16": torch.bfloat16, "f32": torch.float32} + dtype = dtype_map[dtype_str] + device = torch.device("cuda") + + print( + f"=== MoE Reduce Kernel: tokens={tokens}, topk={topk}, model_dim={model_dim}, " + f"use_mask={use_mask}, dtype={dtype_str} ===" + ) + + reduce_exe = compile_moe_reduction( + topk=topk, + model_dim=model_dim, + dtype_str=dtype_str, + use_mask=use_mask, + num_experts=num_experts if use_mask else 0, + ) + + X = torch.randn(tokens, topk, model_dim, device=device, dtype=dtype) + Y = torch.empty(tokens, model_dim, device=device, dtype=dtype) + + if use_mask: + topk_ids = torch.randint(0, num_experts, (tokens, topk), device=device, dtype=torch.int32) + expert_mask = torch.randint(0, 2, (num_experts,), device=device, dtype=torch.int32) + else: + topk_ids = torch.empty(0, device=device, dtype=torch.int32) + expert_mask = torch.empty(0, device=device, dtype=torch.int32) + + stream = torch.cuda.current_stream() + + def launch(): + _run_compiled(reduce_exe, _ptr(X), _ptr(Y), _ptr(expert_mask), _ptr(topk_ids), tokens, stream) + + _, us = run_perftest( + launch, + num_iters=num_iters, + num_warmup=num_warmup, + ) + torch.cuda.synchronize() + + # Correctness + if use_mask: + valid = (expert_mask[topk_ids.long()] != 0).to(torch.bool) # [tokens, topk] + X_ref = X * valid.unsqueeze(-1) + else: + X_ref = X + Y_ref = torch.sum(X_ref, dim=1) + assert verify_output(Y.float(), Y_ref.float(), rtol=1e-2, atol=1e-2, msg="[reduce kernel]") + + # Bandwidth + elem_bytes = X.element_size() + bytes_moved = (tokens * topk * model_dim + tokens * model_dim) * elem_bytes + bw_gb_s = bytes_moved / 1e9 / (us / 1e6) + print(f"[FlyDSL reduce] {us:.1f} us, Bandwidth: {bw_gb_s:.2f} GB/s") + + if compare_torch: + + def launch_torch(y, x_ref): + torch.sum(x_ref, dim=1, out=y) + + Y_torch = torch.empty_like(Y) + _, us_torch = run_perftest( + launch_torch, + Y_torch, + X_ref, + num_iters=num_iters, + num_warmup=num_warmup, + ) + torch.cuda.synchronize() + bw_torch = bytes_moved / 1e9 / (us_torch / 1e6) + speedup = us_torch / us if us > 0 else 0 + print(f"[torch.sum] {us_torch:.1f} us, Bandwidth: {bw_torch:.2f} GB/s") + print(f"[speedup] {speedup:.2f}x") + + +@pytest.mark.parametrize( + "tokens, topk, model_dim, use_mask", + [ + pytest.param(32769, 8, 7168, False, id="DS-TP8-prefill-L", marks=pytest.mark.large_shape), + pytest.param(1, 8, 7168, False, id="DS-TP8-decode-S"), + pytest.param(5, 8, 7168, False, id="DS-TP8-decode-M"), + pytest.param(65, 8, 7168, False, id="DS-TP8-decode-L"), + pytest.param(16384, 6, 5120, False, id="EP-K6-prefill", marks=pytest.mark.large_shape), + pytest.param(1, 6, 5120, False, id="EP-K6-decode-S"), + pytest.param(5, 6, 5120, False, id="EP-K6-decode-M"), + pytest.param(65, 6, 5120, False, id="EP-K6-decode-L"), + pytest.param(129, 8, 7168, True, id="DS-TP8-masked"), + pytest.param(129, 6, 5120, True, id="EP-K6-masked"), + ], +) +def test_moe_reduce_kernel(tokens: int, topk: int, model_dim: int, use_mask: bool): + """Test reduce kernel correctness and performance vs torch.sum.""" + run_reduce_test(tokens=tokens, topk=topk, model_dim=model_dim, use_mask=use_mask) + + +if __name__ == "__main__": + torch.set_default_device("cuda") + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description=( + "MoE Reduction Kernel — correctness & performance test.\n" + "\n" + "Reduces [tokens, topk, model_dim] along the topk dimension.\n" + "Designed for MoE stage-2 shapes where topk is small \n" + "and model_dim is large and aligned (e.g. 5120, 7168)." + ), + ) + parser.add_argument("--tokens", "-t", type=int, default=16384) + parser.add_argument("--topk", "-k", type=int, default=8) + parser.add_argument("--model_dim", "-d", type=int, default=7168) + parser.add_argument("--dtype", type=str, default="f16", choices=["f16", "bf16", "f32"]) + parser.add_argument("--use_mask", action="store_true", default=False) + parser.add_argument("--num_iters", type=int, default=20) + parser.add_argument("--num_warmup", type=int, default=5) + parser.add_argument("--no_compare_torch", action="store_true", default=False) + + args = parser.parse_args() + run_reduce_test( + tokens=args.tokens, + topk=args.topk, + model_dim=args.model_dim, + dtype_str=args.dtype, + use_mask=args.use_mask, + num_iters=args.num_iters, + num_warmup=args.num_warmup, + compare_torch=not args.no_compare_torch, + )