From eaddb05aaa27ae2c578ec8b1229cc7869e64099a Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Tue, 14 Jul 2026 06:03:23 +0000 Subject: [PATCH] Add deferred torch2flydsl MoE kernel tasks --- .../torch2flydsl/moe_a8w4_kernel/config.yaml | 23 + tasks/torch2flydsl/moe_a8w4_kernel/kernel.py | 6674 +++++++++++++++++ tasks/torch2flydsl/moe_a8w4_kernel/model.py | 227 + .../moe_a8w4_kernel/test_kernel_harness.py | 240 + tasks/torch2flydsl/moe_kernel/config.yaml | 23 + tasks/torch2flydsl/moe_kernel/kernel.py | 6640 ++++++++++++++++ tasks/torch2flydsl/moe_kernel/model.py | 211 + .../moe_kernel/test_kernel_harness.py | 238 + .../moe_swiglu_kernel/config.yaml | 23 + .../torch2flydsl/moe_swiglu_kernel/kernel.py | 6651 ++++++++++++++++ tasks/torch2flydsl/moe_swiglu_kernel/model.py | 223 + .../moe_swiglu_kernel/test_kernel_harness.py | 238 + 12 files changed, 21411 insertions(+) create mode 100644 tasks/torch2flydsl/moe_a8w4_kernel/config.yaml create mode 100644 tasks/torch2flydsl/moe_a8w4_kernel/kernel.py create mode 100644 tasks/torch2flydsl/moe_a8w4_kernel/model.py create mode 100644 tasks/torch2flydsl/moe_a8w4_kernel/test_kernel_harness.py create mode 100644 tasks/torch2flydsl/moe_kernel/config.yaml create mode 100644 tasks/torch2flydsl/moe_kernel/kernel.py create mode 100644 tasks/torch2flydsl/moe_kernel/model.py create mode 100644 tasks/torch2flydsl/moe_kernel/test_kernel_harness.py create mode 100644 tasks/torch2flydsl/moe_swiglu_kernel/config.yaml create mode 100644 tasks/torch2flydsl/moe_swiglu_kernel/kernel.py create mode 100644 tasks/torch2flydsl/moe_swiglu_kernel/model.py create mode 100644 tasks/torch2flydsl/moe_swiglu_kernel/test_kernel_harness.py diff --git a/tasks/torch2flydsl/moe_a8w4_kernel/config.yaml b/tasks/torch2flydsl/moe_a8w4_kernel/config.yaml new file mode 100644 index 00000000..8a9ee4b8 --- /dev/null +++ b/tasks/torch2flydsl/moe_a8w4_kernel/config.yaml @@ -0,0 +1,23 @@ +source_file_path: +- kernel.py +target_kernel_functions: +- flydsl_moe_a8w4 +- build_moe_stage1_module +- build_moe_stage2_module +- compile_mixed_moe_gemm1 +- compile_mixed_moe_gemm2 +compile_command: +- python3 -c "import torch; from kernel import build_moe_stage1_module, build_moe_stage2_module; + build_moe_stage1_module(model_dim=7168, inter_dim=256, experts=384, topk=8, tile_m=32, + tile_n=256, tile_k=256, doweight_stage1=False, a_dtype='fp8', b_dtype='fp4', out_dtype='bf16'); + build_moe_stage2_module(model_dim=7168, inter_dim=256, experts=384, topk=8, tile_m=32, + tile_n=256, tile_k=256, doweight_stage2=True, a_dtype='fp8', b_dtype='fp4', out_dtype='bf16'); + print('compile ok')" +correctness_command: +- python3 test_kernel_harness.py --correctness +performance_command: +- python3 test_kernel_harness.py --full-benchmark +task_type: torch2flydsl +supported_archs: +- gfx950 +task_result_template: null diff --git a/tasks/torch2flydsl/moe_a8w4_kernel/kernel.py b/tasks/torch2flydsl/moe_a8w4_kernel/kernel.py new file mode 100644 index 00000000..410df148 --- /dev/null +++ b/tasks/torch2flydsl/moe_a8w4_kernel/kernel.py @@ -0,0 +1,6674 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""FlyDSL a8w4 two-stage fused MoE kernel (MXFP8 activations, MXFP4 weights). + +Defines the stage-1 (gate/up GEMM with fused silu-gated activation) and stage-2 +(down GEMM with weighted top-k combine) device kernels in FlyDSL. The device +builders ``compile_mixed_moe_gemm1`` / ``compile_mixed_moe_gemm2`` -- exposed via +``build_moe_stage1_module`` / ``build_moe_stage2_module`` and built with +``a_dtype='fp8'`` / ``b_dtype='fp4'`` -- are adapted from AITER's +mixed_moe_gemm_2stage path together with its preshuffle pipeline, CShuffle MFMA +epilogue, layout helpers and GateMode enum. + +``flydsl_moe_a8w4`` is the launcher: activations are quantized to MXFP8 (e4m3) +and expert weights to MXFP4, both with e8m0 per_1x32 block scales. It takes bf16 +weights and a precomputed routing (so the reference and kernel share identical +top-k selection) and returns a bf16 ``[T, model_dim]`` tensor. Host-side data +prep -- quantization, weight/scale pre-shuffle and the sorted token/expert +dispatch (``moe_sorting``) -- uses AITER utilities to shape inputs into the +layout the device kernels consume. +""" +from __future__ import annotations + +import functools +import math as _math +import os +import re +import builtins as _builtins +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Dict, Optional + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, memref +from flydsl._mlir.dialects.arith import CmpIPredicate +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr import arith as _arith +from flydsl.expr.arith import ArithValue +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + + +# =========================================================================== +# Inlined from aiter/ops/flydsl/moe_common.py :: GateMode +# =========================================================================== +class GateMode(str, Enum): + """Gate/Up computation strategy for stage1 GEMM (see AITER moe_common).""" + + SEPARATED = "separated" + MOCK_GATE_ONLY = "mock_gate_only" + GATE_ONLY = "gate_only" + INTERLEAVE = "interleave" + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/layout_utils.py +# ========================================================================= +def _wrap(v): + """Wrap raw ir.Value in ArithValue for operator overloading compatibility.""" + if isinstance(v, ArithValue): + return v + if isinstance(v, ir.Value): + return ArithValue(v) + return v + + +def _is_pow2(n): + """Return True when *n* is a positive power of two.""" + return n > 0 and (n & (n - 1)) == 0 + + +def _div_pow2(val, divisor): + """Unsigned divide index *val* by a **compile-time** power-of-2 *divisor*. + + Emits ``arith.shrui`` (1 VALU cycle) instead of ``arith.divui`` + (10-15 VALU cycles on CDNA). + """ + shift = _math.log2(divisor) + assert shift == int(shift), f"{divisor} is not a power of 2" + return arith.shrui(val, arith.index(int(shift))) + + +def _mod_pow2(val, modulus): + """Unsigned remainder of index *val* by a **compile-time** power-of-2 *modulus*. + + Emits ``arith.andi`` (1 VALU cycle) instead of ``arith.remui``. + """ + return arith.andi(val, arith.index(modulus - 1)) + + +def _parse_dim(tok): + """Parse a single dimension token: '?' -> None, otherwise int.""" + tok = tok.strip() + return None if tok == "?" else int(tok) + + +def _parse_layout(ly): + """Parse '(s0,s1,...):(d0,d1,...)' -> (shapes, strides) as lists (None for '?').""" + ly_str = str(ly.type) if hasattr(ly, "type") else str(ly) + m = re.search(r"\(([^)]+)\):\(([^)]+)\)", ly_str) + if not m: + return None + shapes = [_parse_dim(s) for s in m.group(1).split(",")] + strides = [_parse_dim(s) for s in m.group(2).split(",")] + return shapes, strides + + +def _has_dynamic_strides(strides): + """Check if any stride is dynamic (None).""" + return any(s is None for s in strides) + + +def idx2crd(idx, layout): + """Decompose flat index into a list of coordinate values. + + For static layouts, computes coordinates with plain arith ops. + Power-of-2 strides/shapes use shift/mask instead of div/rem. + For dynamic layouts, falls back to fx.idx2crd + fx.get. + """ + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + result = fx.idx2crd(idx, layout) + ndims = len(parsed[1]) if parsed else 1 + return [_wrap(fx.get(result, i)) for i in range(ndims)] + + if hasattr(idx, "type") and str(idx.type) != "index": + idx = arith.index_cast(T.index, idx) + shapes, strides = parsed + ndims = len(strides) + + ordered = sorted( + [ + (i, s, sz) + for i, s, sz in _builtins.zip(range(ndims), strides, shapes) + if s != 0 + ], + key=lambda x: x[1], + reverse=True, + ) + coords = [None] * ndims + remaining = idx + for i, stride_val, size_val in ordered: + if stride_val == 1: + c = remaining + elif _is_pow2(stride_val): + c = _div_pow2(remaining, stride_val) + else: + c = remaining / arith.index(stride_val) + if size_val is not None: + if _is_pow2(size_val): + c = _mod_pow2(c, size_val) + else: + c = c % arith.index(size_val) + coords[i] = c + for i in range(ndims): + if coords[i] is None: + coords[i] = remaining + return coords + + +def crd2idx(crd, layout): + """Compute flat index from a coordinate tuple/list. + + For static layouts, computes with plain arith ops. + For dynamic layouts, falls back to fx.crd2idx with fx.make_coord. + """ + if not isinstance(crd, (list, tuple)): + crd = [crd] + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + # fly.make_coord requires i32/i64, not index + crd_i32 = [] + for c in crd: + cv = c + if isinstance(cv, ArithValue): + cv = cv.ir_value() if hasattr(cv, "ir_value") else cv + if isinstance(cv, ir.Value) and isinstance(cv.type, ir.IndexType): + cv = arith.index_cast(T.i32, cv) + crd_i32.append(cv) + coord_val = fx.make_coord(*crd_i32) + result = fx.crd2idx(coord_val, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = arith.index_cast(T.index, scalar) + return _wrap(scalar) + + _, strides = parsed + result = None + for coord_v, stride_v in _builtins.zip(crd, strides): + if stride_v == 0: + continue + term = coord_v if stride_v == 1 else coord_v * arith.index(stride_v) + result = term if result is None else result + term + return result if result is not None else arith.index(0) + + +def get(int_tuple, mode): + """Extract element at `mode` from a Python list/tuple.""" + return int_tuple[mode] + + +layout_get = get + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_preshuffle_pipeline.py (crd2idx -> _pre_crd2idx) +# ========================================================================= +def _pre_crd2idx(crd, layout): + """crd2idx returning an index-type scalar (unwraps fly.int_tuple).""" + result = fx.crd2idx(crd, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = _arith.IndexCastOp(T.index, scalar).result + return scalar + + +def swizzle_xor16(row, col, k_blocks16): + """XOR-with-row swizzle on the K dimension at 16B granularity. + + Computes: col XOR ((row & (k_blocks16 - 1)) * 16) + + k_blocks16 is always a power of 2 (tile_k_bytes / 16), so use + bitwise AND instead of remui to save ~10 VALU cycles on CDNA. + """ + from flydsl.expr import arith as _swz_arith + + mask = k_blocks16 - _swz_arith.index(1) + rem = _swz_arith.andi(row, mask) + return col ^ (rem * 16) + + +def lds_row_major_idx(row, col, row_stride, base=None): + """Linearize a 2D LDS coordinate with explicit index arithmetic.""" + idx = row * row_stride + col + return idx if base is None else idx + base + + +def split_row_major_2d(index, minor_extent): + """Split a linear row-major index into (major, minor).""" + return index // minor_extent, index % minor_extent + + +def _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx, + *, + elem_type, + vec_elems, + elem_bytes, + offset_in_bytes, + cache_modifier=0, +): + """Load vec_elems elements via buffer_load dwordx[1,2,4] + bitcast.""" + from flydsl.expr import arith as _ld_arith + + elem_size = int(elem_bytes) + load_bytes = int(vec_elems) * elem_size + vec_width = load_bytes // 4 + + if offset_in_bytes: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(2)) + elif elem_bytes == 2: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(1)) + else: + idx_i32 = idx + + i32_val = buffer_ops.buffer_load( + rsrc, + idx_i32, + vec_width=vec_width, + dtype=T.i32, + cache_modifier=cache_modifier, + ) + if vec_width == 1: + i32_vec = vector.from_elements(T.vec(1, T.i32), [i32_val]) + else: + i32_vec = i32_val + return vector.bitcast(T.vec(int(vec_elems), elem_type), i32_vec) + + +@dataclass(frozen=True) +class PreshuffleScaleLayout: + """Container returned by `make_preshuffle_scale_layout`. + + The scale layout is ``(c_mn1, c_k1, 4, 16) : (stride_n0, stride_k0, stride_klane, 1)``. + Callers compute flat index directly with plain arith:: + + idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + """ + + layout_scale: object + stride_n0: object + stride_k0: object + stride_klane: object + + +def make_preshuffle_scale_layout( + arith, + *, + c_mn: ir.Value, + c_k: ir.Value, + mn_pack: int = 2, + k_pack: int = 2, + elem_bytes: int = 4, + scale_block_size: int = 32, +) -> PreshuffleScaleLayout: + """Build scale layout matching aiter/CK preshuffle for FP4/FP8 microscale. + + Layout shape: ``(c_mn1, c_k1, 4, 16)`` where + ``c_mn1 = c_mn / 16 / mn_pack`` and ``c_k1 = (c_k / scale_block_size) / 4 / k_pack``. + """ + c16 = fx.Index(16) + c4 = fx.Index(4) + c_k_scale = c_k // fx.Index(scale_block_size) + + c_mn1 = (c_mn // c16) // fx.Index(mn_pack) + c_k1 = (c_k_scale // c4) // fx.Index(k_pack) + if elem_bytes != mn_pack * k_pack: + raise ValueError( + f"elem_bytes of scale must be {mn_pack} * {k_pack}, got {elem_bytes!r}" + ) + + stride_klane = c16 + stride_k0 = c4 * stride_klane + stride_n0 = c_k1 * stride_k0 + + c_mn1_i32 = arith.index_cast(T.i32, c_mn1) + c_k1_i32 = arith.index_cast(T.i32, c_k1) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + + layout_scale = fx.make_layout( + (c_mn1_i32, c_k1_i32, 4, 16), + stride=(stride_n0_i32, stride_k0_i32, stride_klane_i32, 1), + ) + + return PreshuffleScaleLayout( + layout_scale=layout_scale, + stride_n0=stride_n0, + stride_k0=stride_k0, + stride_klane=stride_klane, + ) + + +@dataclass(frozen=True) +class PreshuffleBLayout: + """Container returned by `make_preshuffle_b_layout`.""" + + layout_b: object + kpack_bytes: int + + +def make_preshuffle_b_layout( + arith, + *, + c_n: ir.Value, + c_k: ir.Value, + kpack_bytes: int = 16, + elem_bytes: int = 1, + k_major: bool = False, +) -> PreshuffleBLayout: + """Build B layout matching aiter/CK preshuffle for A8 MFMA kernels. + + When *k_major* is True the block-level order is K-major (``k_blk`` outermost), + matching the ``(0,3,1,4,2,5)`` shuffle permutation. The default N-major + order (``k_major=False``) matches the legacy ``(0,1,3,4,2,5)`` permutation. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + + c16 = fx.Index(16) + c_kpack = fx.Index(kpack_bytes) + + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + c_k_bytes = c_k * arith.constant(int(elem_bytes), index=True) + n0 = c_n // c16 + + c_kpack_elems = ( + c_kpack + if elem_bytes == 1 + else (c_kpack // arith.constant(int(elem_bytes), index=True)) + ) + + stride_nlane = c_kpack_elems + + if k_major: + c32 = fx.Index(32) + c2 = fx.Index(2) + c_k0 = c_k_bytes // c32 + klane_dim = 2 + stride_klane = c16 * stride_nlane + stride_n0 = c2 * stride_klane + stride_k0 = n0 * stride_n0 + else: + c64 = fx.Index(64) + c4 = fx.Index(4) + c_k0 = c_k_bytes // c64 + klane_dim = 4 + stride_klane = c16 * stride_nlane + stride_k0 = c4 * stride_klane + stride_n0 = c_k0 * stride_k0 + + kpack_elems_static = kpack_bytes if elem_bytes == 1 else kpack_bytes // elem_bytes + n0_i32 = arith.index_cast(T.i32, n0) + c_k0_i32 = arith.index_cast(T.i32, c_k0) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + stride_nlane_i32 = arith.index_cast(T.i32, stride_nlane) + + stride_b = (stride_n0_i32, stride_k0_i32, stride_klane_i32, stride_nlane_i32, 1) + layout_b = fx.make_layout( + (n0_i32, c_k0_i32, klane_dim, 16, kpack_elems_static), stride_b + ) + return PreshuffleBLayout(layout_b=layout_b, kpack_bytes=kpack_bytes) + + +def _unpack_int4_to_int8_pair(packed32): + """Split packed int4 dword into two int8 dwords (even/odd nibbles). + + 7-op bit manipulation shared by all int4 unpack paths (W4A8, W4A16, W4A_FP8). + """ + c_08 = fx.Int32(0x08080808) + c_0f = fx.Int32(0x0F0F0F0F) + c_1e = fx.Int32(0x1E) + c_4 = fx.Int32(4) + s0 = (packed32 & c_08) * c_1e + even = (packed32 & c_0f) | s0 + t = packed32 >> c_4 + s1 = (t & c_08) * c_1e + odd = (t & c_0f) | s1 + return even, odd + + +def _pack_i32_pair_to_i64(lo, hi, vector): + """Pack two i32 values into one i64 via vector bitcast.""" + v2 = vector.from_elements(T.vec(2, T.i32), [lo, hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def _i8x4_in_i32_to_bf16x4_i64(val_i32, arith, vector, scale_val=None): + """Convert one i32 (4 signed int8 bytes) to 4 bf16 packed as i64. + + Uses shift-based f32->bf16 truncation (lshr 16) instead of arith.truncf + which on gfx942 expands to ~5 VALU per element. The shift is exact for + unscaled int8 values and introduces <0.5 ULP error for scaled values. + """ + vec1_i32_t = T.vec(1, T.i32) + vec2_i32 = T.i32x2 + vec4_i8 = T.i8x4 + vec1_i64 = T.vec(1, T.i64) + + v1 = vector.from_elements(vec1_i32_t, [val_i32]) + i8x4 = vector.bitcast(vec4_i8, v1) + + f32_vals = [] + for i in range(4): + val_i8 = vector.extract(i8x4, static_position=[i], dynamic_position=[]) + v = arith.sitofp(T.f32, val_i8) + if scale_val is not None: + v = v * scale_val + f32_vals.append(v) + + c16 = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bits0 = arith.bitcast(T.i32, f32_vals[0]) + bits1 = arith.bitcast(T.i32, f32_vals[1]) + bits2 = arith.bitcast(T.i32, f32_vals[2]) + bits3 = arith.bitcast(T.i32, f32_vals[3]) + i32_lo = (bits0 >> c16) | (bits1 & c_ffff0000) + i32_hi = (bits2 >> c16) | (bits3 & c_ffff0000) + + v2 = vector.from_elements(vec2_i32, [i32_lo, i32_hi]) + v64 = vector.bitcast(vec1_i64, v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def load_b_raw_w4a16( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ku: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 8, +): + """Phase 1 of W4A16 B load: issue buffer_load_dword, return raw packed i32. + + Same address calculation as the int4 unpack path in load_b_pack_k32 + but using ku-based indexing for 2-phase latency hiding. + """ + if kpack_bytes != 8: + raise ValueError(f"W4A16 requires kpack_bytes=8, got {kpack_bytes!r}") + + c64 = fx.Index(64) + half_bytes = kpack_bytes // 2 + c2_idx = fx.Index(2) + c4_idx = fx.Index(4) + + k0_base = base_k // c64 + + k1_layout_offset = ku * 2 + lane_div_32 = lane_div_16 // c2_idx + total_k1 = fx.Index(k1_layout_offset) + lane_div_32 + k0 = k0_base + (total_k1 // c4_idx) + k1_local = total_k1 % c4_idx + lane_odd = lane_div_16 % c2_idx + k2_base = lane_odd * fx.Index(half_bytes) + + coord_pack = (n_blk, k0, k1_local, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + idx_bytes = idx_pack + k2_base + + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + return packed32 + + +def _int4_to_bf16x4_i64_gfx950( + packed32, nibble_offsets, arith, vector, scale_val=None, defer_scale16=False +): + """Convert 4 int4 nibbles to 4 bf16 packed as i64 using gfx950 instructions. + + Uses v_cvt_off_f32_i4_sdwa with byte_sel to avoid per-nibble shifts. + Even nibbles (0,2,4,6) → SDWA BYTE_0/1/2/3 on original src. + Odd nibbles (1,3,5,7) → SDWA BYTE_0/1/2/3 on (src >> 4). + Only 1 shift total instead of 7. + + When defer_scale16=True, the ×16 correction factor for v_cvt_off_f32_i4 is + omitted and must be applied later (e.g. in the epilogue). This saves VALU + in the hot loop and uses v_cvt_pk_bf16_f32 for proper f32→bf16 conversion. + """ + from flydsl.expr import rocdl + from flydsl._mlir.dialects._arith_ops_gen import MulFOp as _MulFOp + + _uw = _arith._to_raw + _av = _arith.ArithValue + + src_even = packed32 + src_odd = packed32 >> fx.Int32(4) + + f32_vals = [] + for nib in nibble_offsets: + byte_idx = nib // 2 + src = src_odd if (nib % 2) else src_even + v = rocdl.cvt_off_f32_i4(src, byte_sel=byte_idx) + f32_vals.append(v) + + if defer_scale16: + # Skip ×16; multiply by scale_val only if groupwise. + if scale_val is not None: + raw_scale = _uw(scale_val) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Use v_cvt_pk_bf16_f32 for proper f32→bf16 (no bit-shift trick needed). + i32_lo = rocdl.cvt_pk_bf16_f32(f32_vals[0], f32_vals[1]) + i32_hi = rocdl.cvt_pk_bf16_f32(f32_vals[2], f32_vals[3]) + else: + c16 = fx.Float32(16.0) + if scale_val is not None: + effective_scale = scale_val * c16 + else: + effective_scale = c16 + raw_scale = _uw(effective_scale) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Truncate f32→bf16 via bit-shift (exact for scaled int values). + c16_shift = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bf16_vals = [arith.bitcast(T.i32, _av(v)) for v in f32_vals] + i32_lo = (bf16_vals[0] >> c16_shift) | (bf16_vals[1] & c_ffff0000) + i32_hi = (bf16_vals[2] >> c16_shift) | (bf16_vals[3] & c_ffff0000) + + v2 = vector.from_elements(T.vec(2, T.i32), [i32_lo, i32_hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def unpack_b_w4a16( + packed32, arith, vector, scale_val=None, use_gfx950_cvt=False, defer_scale16=False +): + """Phase 2 of W4A16 B load: unpack int4->int8 + convert int8->bf16. + + Takes raw packed32 from load_b_raw_w4a16 and produces (b0, b1) -- + two i64 values each containing 4 bf16 for one MFMA. + + When use_gfx950_cvt=True, uses v_cvt_off_f32_i4 + v_cvt_pk_bf16_f32 + for ~2x fewer VALU instructions. + + When defer_scale16=True (requires use_gfx950_cvt=True), the ×16 + correction for v_cvt_off_f32_i4 is omitted; caller must apply it + in the epilogue. + """ + if use_gfx950_cvt: + b0 = _int4_to_bf16x4_i64_gfx950( + packed32, + [0, 2, 4, 6], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + b1 = _int4_to_bf16x4_i64_gfx950( + packed32, + [1, 3, 5, 7], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + return (b0, b1) + even, odd = _unpack_int4_to_int8_pair(packed32) + b0 = _i8x4_in_i32_to_bf16x4_i64(even, arith, vector, scale_val=scale_val) + b1 = _i8x4_in_i32_to_bf16x4_i64(odd, arith, vector, scale_val=scale_val) + return (b0, b1) + + +def load_b_pack_k32( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ki_step: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 16, + elem_bytes: int = 1, + unpack_int4: bool = False, +) -> ir.Value: + """Load one B pack for one MFMA(x32) micro-step. + + Returns an i64 Value containing 8 bytes consumed by MFMA. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + if unpack_int4 and kpack_bytes != 8: + raise ValueError("unpack_int4 requires kpack_bytes=8 (packed int4 layout)") + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + + c64 = fx.Index(64) + base_k_bytes = base_k * arith.constant(int(elem_bytes), index=True) + k0_base = base_k_bytes // c64 + k0 = k0_base + arith.constant(ki_step // 2, index=True) + k1 = lane_div_16 + half_bytes = kpack_bytes // 2 + k2_base = arith.constant((ki_step % 2) * half_bytes, index=True) + + coord_pack = (n_blk, k0, k1, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + + if unpack_int4: + idx_bytes = idx_pack + k2_base + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + even, odd = _unpack_int4_to_int8_pair(packed32) + return _pack_i32_pair_to_i64(even, odd, vector) + + vec_elems = kpack_bytes // int(elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_pack, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=(elem_bytes == 1), + ) + + b_i32x4 = vector.bitcast(T.i32x4, b16) + + half = ki_step % 2 + if half == 0: + d0 = vector.extract(b_i32x4, static_position=[0], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[1], dynamic_position=[]) + else: + d0 = vector.extract(b_i32x4, static_position=[2], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[3], dynamic_position=[]) + + v2 = vector.from_elements(T.vec(2, T.i32), [d0, d1]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def tile_chunk_coord_i32( + arith, + *, + tx_i32_base: ir.Value, + i: int, + total_threads: int, + layout_tile_div4, + chunk_i32: int = 4, +): + """Map (thread, chunk_id) -> (row_local, col_local_i32) for X/A loads.""" + if chunk_i32 not in (1, 2, 4): + raise ValueError(f"chunk_i32 must be one of (1,2,4), got {chunk_i32!r}") + chunk_off_i32 = arith.constant(i * total_threads * chunk_i32, index=True) + tile_idx_i32 = tx_i32_base + chunk_off_i32 + coord_local = fx.idx2crd(tile_idx_i32, layout_tile_div4) + row_local = fx.get(coord_local, 0) + col_local_i32 = fx.get(coord_local, 1) + return row_local, col_local_i32 + + +def buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + *, + elem_type, + idx_i32: ir.Value, + rsrc, + vec_elems: int = 16, + elem_bytes: int = 1, +): + """Copy 16 bytes from global memory into regs via buffer-load dwordx4 lowering.""" + if int(vec_elems) <= 0: + raise ValueError(f"vec_elems must be > 0, got {vec_elems!r}") + return _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx_i32, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=False, + ) + + +def lds_store_16b_xor16( + arith, + vector, + *, + lds_memref, + vec16_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x4: ir.Value, + elem_bytes: int = 1, +): + """Store one 16B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v16 = vector.bitcast(vec16_ty, vec_part_i32x4) + vector.store(v16, lds_memref, [idx0]) + + +def lds_store_8b_xor16( + arith, + vector, + *, + lds_memref, + vec8_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x2: ir.Value, + elem_bytes: int = 1, +): + """Store one 8B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v8 = vector.bitcast(vec8_ty, vec_part_i32x2) + vector.store(v8, lds_memref, [idx0]) + + +def lds_store_4b_xor16( + arith, + vector, + *, + lds_memref, + vec4_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x1: ir.Value, + elem_bytes: int = 1, +): + """Store one 4B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v4 = vector.bitcast(vec4_ty, vec_part_i32x1) + vector.store(v4, lds_memref, [idx0]) + + +def lds_load_pack_k32( + arith, + vector, + *, + lds_memref, + layout_lds, + k_blocks16: ir.Value, + curr_row_a_lds: ir.Value, + col_base: ir.Value, + half: int, + lds_base: ir.Value, + ck_lds128: bool, + vec16_ty, + vec8_ty, + vec2_i64_ty, + vec1_i64_ty, +): + """Load one i64 A-pack for an MFMA K32 micro-step from LDS.""" + col_base_swz = swizzle_xor16(curr_row_a_lds, col_base, k_blocks16) + if ck_lds128: + coord_a16 = (curr_row_a_lds, col_base_swz) + idx_a16 = _pre_crd2idx(coord_a16, layout_lds) + lds_base + loaded_a16 = vector.load_op(vec16_ty, lds_memref, [idx_a16]) + a_vec128 = vector.bitcast(vec2_i64_ty, loaded_a16) + return vector.extract(a_vec128, static_position=[half], dynamic_position=[]) + else: + col_swizzled = col_base_swz + (half * 8) + coord_a = (curr_row_a_lds, col_swizzled) + idx_a = _pre_crd2idx(coord_a, layout_lds) + lds_base + loaded_a8 = vector.load_op(vec8_ty, lds_memref, [idx_a]) + a_vec64 = vector.bitcast(vec1_i64_ty, loaded_a8) + return vector.extract(a_vec64, static_position=[0], dynamic_position=[]) + + +def xcd_remap_bx_by( + bx, + by, + c_m, + *, + tile_m: int, + tile_n: int, + N: int, + xcd_swizzle: int, + num_xcds: int = 8, +): + """Remap (bx, by) for L2-cache reuse via XCD swizzle. + + No-op when ``xcd_swizzle <= 0``. Otherwise: + 1. Linearize the original (bx, by) grid round-robin across ``num_xcds`` + XCDs so that contiguous workgroup ids stay on the same XCD. + 2. Re-tile that 1-D order with an M-major group of size ``xcd_swizzle``, + folding the tail group when ``gy`` does not divide evenly. + + Designed to be called inside a ``@flyc.kernel`` immediately after:: + + bx = gpu.block_id("x") + by = gpu.block_id("y") + bx, by = xcd_remap_bx_by(bx, by, c_m, tile_m=..., tile_n=..., N=..., + xcd_swizzle=xcd_swizzle) + + ``c_m`` is the dynamic ``fx.Index`` for runtime ``M``; ``tile_m``, + ``tile_n``, ``N`` and ``xcd_swizzle`` are compile-time Python ints. + """ + if xcd_swizzle <= 0: + return bx, by + + _c1 = fx.arith.constant(1, index=True) + _c_tm = fx.arith.constant(tile_m, index=True) + _gx = fx.arith.constant(N // tile_n, index=True) + _gy = (c_m + _c_tm - _c1) / _c_tm + + _linear_id = bx * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = fx.arith.constant(num_xcds, index=True) + _q = _num_wgs / _c_xcds + _r = _num_wgs % _c_xcds + _xcd = _linear_id % _c_xcds + _in_xcd = _linear_id / _c_xcds + _xcd_lt_r = fx.arith.cmpi(CmpIPredicate.ult, _xcd, _r) + _clip = fx.arith.select(_xcd_lt_r, _xcd, _r) + _wgid = _xcd * _q + _clip + _in_xcd + + _c_wgm = fx.arith.constant(xcd_swizzle, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = fx.arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = fx.arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + new_bx = _first_pid_m + (_wgid_in_group % _group_size_m) + new_by = _wgid_in_group / _group_size_m + return new_bx, new_by + + +__all__ = [ + "PreshuffleBLayout", + "PreshuffleScaleLayout", + "buffer_copy_gmem16_dwordx4", + "lds_load_pack_k32", + "lds_row_major_idx", + "lds_store_4b_xor16", + "lds_store_8b_xor16", + "lds_store_16b_xor16", + "make_preshuffle_b_layout", + "make_preshuffle_scale_layout", + "load_b_pack_k32", + "load_b_raw_w4a16", + "unpack_b_w4a16", + "load_b_raw_w4a16_groupwise", + "unpack_b_w4a16_groupwise", + "extract_bf16_scale", + "split_row_major_2d", + "swizzle_xor16", + "tile_chunk_coord_i32", + "xcd_remap_bx_by", +] + + +# --------------------------------------------------------------------------- +# Groupwise scale load helper (shared by W4A16 and W4A8 groupwise paths) +# --------------------------------------------------------------------------- + + +def _load_groupwise_scale( + buffer_ops, + arith, + *, + scale_rsrc, + expert_offset, + n_blk, + n_intra, + k_pos, + num_groups: int, + group_size: int, + n_per_expert: int, + scale_dtype=None, +): + """Load one per-group scale value from the scale buffer. + + Computes the linear index into the scale tensor from expert offset, + N position, and group index derived from ``k_pos``. + + For bf16 scales the tensor uses ``(E, G//2, N, 2)`` layout — two + adjacent groups for the same N position are packed into one dword. + We load the raw i32 dword (no extraction) so it can be carried as + loop state without register copies. Use :func:`extract_bf16_scale` + in the compute phase to obtain the f32 value. + """ + c16 = fx.Index(16) + n_global = n_blk * c16 + n_intra + c_group_size = fx.Index(group_size) + c_npe = fx.Index(n_per_expert) + group_idx = k_pos // c_group_size + if scale_dtype is None: + scale_dtype = T.f32 + + if scale_dtype == T.bf16: + # (E, G//2, N, 2) layout: dword at [e, pair, n] holds bf16 scales + # for groups 2*pair and 2*pair+1. + pair_idx = group_idx >> fx.Index(1) # group_idx // 2 + # Dword index: same flat formula but with G//2 groups + num_pairs = num_groups // 2 + c_npm1 = fx.Index(num_pairs - 1) + dword_base = expert_offset * c_npm1 + n_global + dword_elem = dword_base + pair_idx * c_npe + dword_idx = arith.index_cast(T.i32, dword_elem) + # Return raw i32 dword — extraction deferred to compute phase. + scale_val = buffer_ops.buffer_load( + scale_rsrc, dword_idx, vec_width=1, dtype=T.i32 + ) + else: + # (E, G, N) layout with f32 dtype + c_gm1 = fx.Index(num_groups - 1) + base_scale = expert_offset * c_gm1 + n_global + elem_idx = base_scale + group_idx * c_npe + scale_idx_i32 = arith.index_cast(T.i32, elem_idx) + scale_val = buffer_ops.buffer_load( + scale_rsrc, scale_idx_i32, vec_width=1, dtype=T.f32 + ) + return scale_val + + +def extract_bf16_scale(arith, scale_raw_i32, ku: int): + """Extract f32 scale from raw i32 dword loaded by bf16 groupwise path. + + In the ``(E, G//2, N, 2)`` layout two adjacent groups share one dword. + ``ku`` determines which half: even ku → low bf16, odd ku → high bf16. + """ + if ku % 2 == 0: + # Low bf16: shift left by 16 to place in upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 << fx.Int32(16)) + else: + # High bf16: mask upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 & fx.Int32(0xFFFF0000)) + + +# --------------------------------------------------------------------------- +# W4A16 groupwise load / unpack helpers +# --------------------------------------------------------------------------- + + +def load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k, + ku: int, + n_blk, + n_intra, + lane_div_16, + elem_type, + scale_rsrc, + expert_offset, + num_groups: int, + group_size: int, + n_per_expert: int, + kpack_bytes: int = 8, + scale_dtype=None, +): + """Phase 1 of W4A16 groupwise B load: buffer_loads for weight + scale. + + Reuses :func:`load_b_raw_w4a16` for the weight load, then issues an + additional ``buffer_load_dword`` for the per-group scale. + + Returns ``(packed32, scale_val)``. + """ + packed32 = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk, + n_intra=n_intra, + lane_div_16=lane_div_16, + elem_type=elem_type, + kpack_bytes=kpack_bytes, + ) + k_pos = base_k + fx.Index(ku * 32) + scale_val = _load_groupwise_scale( + buffer_ops, + arith, + scale_rsrc=scale_rsrc, + expert_offset=expert_offset, + n_blk=n_blk, + n_intra=n_intra, + k_pos=k_pos, + num_groups=num_groups, + group_size=group_size, + n_per_expert=n_per_expert, + scale_dtype=scale_dtype, + ) + return (packed32, scale_val) + + +def unpack_b_w4a16_groupwise(packed32, scale_val, arith, vector, use_gfx950_cvt=False): + """Phase 2 of W4A16 groupwise: unpack + scale + convert to bf16.""" + return unpack_b_w4a16( + packed32, arith, vector, scale_val=scale_val, use_gfx950_cvt=use_gfx950_cvt + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_epilogues.py (_if_then -> _epi_if_then) +# ========================================================================= +@contextmanager +def _epi_if_then(if_op, scf): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def default_epilog( + *, + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + body_row: Callable, +): + """Iterate the standard MFMA 16x16 row mapping and call `body_row(...)`. + + The mapping matches the common MFMA fragment layout used across kernels in this repo. + + Args: + arith: flydsl arith ext module. + range_constexpr: compile-time unrolled range helper. + m_repeat: tile_m // 16 (python int). + lane_div_16: index Value (0..3). + bx_m: base row (index Value). For MoE, this is the base sorted-row for the tile. + body_row: callback invoked as: + body_row(mi=, ii=, row_in_tile=, row=) + """ + bx_m_v = bx_m + lane_div_16_mul4 = lane_div_16 * 4 + ii_idx_list = [fx.Index(ii) for ii in range(4)] + + for mi in range_constexpr(m_repeat): + mi_base = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off = lane_div_16_mul4 + ii_idx_list[ii] + row_in_tile = mi_base + row_off + row = bx_m_v + row_in_tile + body_row(mi=mi, ii=ii, row_in_tile=row_in_tile, row=row) + + +def c_shuffle_epilog( + *, + arith, + vector, + gpu, + scf=None, + range_constexpr, + # Tile params + tile_m: int, + tile_n: int, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + m_repeat: int, + num_acc_n: int, + # Thread mapping inputs + tx, + lane_div_16, + lane_mod_16, + bx_m, + by_n, + n_tile_base, + # LDS buffer (f16 view, row-major [tile_m, tile_n] flattened) + lds_out, + # Element type for LDS loads (defaults to f16). Pass bf16 to support bf16 epilogues. + frag_elem_type: ir.Type | None = None, + # Callbacks + write_row_to_lds: Callable, + precompute_row: Callable | None = None, + store_pair: Callable, + # When LDS overflows, split lds_out across two buffers by wave-group. + # Pass the second buffer here; first buffer is `lds_out`. + lds_out_split=None, + # Row offset in lds_out for 8-wave mode (MLIR index value). + # Shifts both write and read LDS indices by lds_row_offset * tile_n elements. + lds_row_offset=None, +): + """LDS CShuffle epilogue skeleton. + + Call pattern: + - `write_row_to_lds(...)` is called once per MFMA row produced by this thread. + It is responsible for writing all ni columns for that row into `lds_out`. + - `store_pair(...)` is called for each (row_local, col_pair0) half2 after shuffle. + + `store_pair` can implement either global stores or atomics. + """ + if int(block_size) <= 0 or (int(block_size) % int(cshuffle_nlane)) != 0: + raise ValueError( + f"block_size ({block_size}) must be divisible by cshuffle_nlane ({cshuffle_nlane})" + ) + cshuffle_mlane = int(block_size) // int(cshuffle_nlane) + if (int(tile_m) % cshuffle_mlane) != 0: + raise ValueError( + f"tile_m must be divisible by CShuffleMLane ({cshuffle_mlane}), got tile_m={tile_m}" + ) + if int(e_vec) <= 0: + raise ValueError(f"e_vec must be positive, got {e_vec}") + if (int(tile_n) % (int(cshuffle_nlane) * int(e_vec))) != 0: + raise ValueError( + f"tile_n must be divisible by (CShuffleNLane*EVec) = {cshuffle_nlane*e_vec}, got tile_n={tile_n}" + ) + + # ===================== Split-LDS mode (early return) ===================== + # When lds_out_split is provided, waves are divided into two groups: + # Group A (waves 0..N/2-1) uses lds_out, columns [0, tile_n/2) + # Group B (waves N/2..N-1) uses lds_out_split, columns [tile_n/2, tile_n) + # Each group writes/reads independently; same barriers synchronise all waves. + if lds_out_split is not None: + if scf is None: + raise ValueError("scf module is required for split-LDS cshuffle") + + _half_n = int(tile_n) // 2 + _half_threads = int(block_size) // 2 + EVec = int(e_vec) + + CShuffleNLane_s = min(int(cshuffle_nlane), _half_n // EVec) + if _half_threads % CShuffleNLane_s != 0: + raise ValueError( + f"half_threads={_half_threads} not divisible by CShuffleNLane_split={CShuffleNLane_s}" + ) + CShuffleMLane_s = _half_threads // CShuffleNLane_s + if int(tile_m) % CShuffleMLane_s != 0: + raise ValueError( + f"tile_m={tile_m} not divisible by CShuffleMLane_split={CShuffleMLane_s}" + ) + m_reps_s = int(tile_m) // CShuffleMLane_s + n_reps_s = _half_n // (CShuffleNLane_s * EVec) + + _half_n_idx = arith.constant(_half_n, index=True) + _half_thr_idx = arith.constant(_half_threads, index=True) + _zero_idx = arith.constant(0, index=True) + + _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + + # -- write phase (all waves, each to its group's LDS buffer) -- + n_tile_base_v = n_tile_base + col_base_local_a = n_tile_base_v + lane_mod_16 + col_base_local_b = col_base_local_a - _half_n_idx + + def _write_row_split(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * _half_n_idx + _if_g = scf.IfOp(_is_group_b, has_else=True) + with ir.InsertionPoint(_if_g.then_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_b, + num_acc_n=num_acc_n, + lds_out=lds_out_split, + ) + scf.YieldOp([]) + with ir.InsertionPoint(_if_g.else_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_a, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + scf.YieldOp([]) + + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row_split, + ) + gpu.barrier() + + # -- read phase (each group reads from its own LDS buffer) -- + tx_local = tx - arith.select(_is_group_b, _half_thr_idx, _zero_idx) + c_nlane_s = arith.constant(CShuffleNLane_s, index=True) + m_lane_s = tx_local / c_nlane_s + n_lane_s = tx_local % c_nlane_s + c_evec = arith.constant(EVec, index=True) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + _precomputed_rows_s = [] + for mr in range_constexpr(m_reps_s): + row_base_m = arith.constant(mr * CShuffleMLane_s, index=True) + row_local = row_base_m + m_lane_s + row = bx_m_v + row_local + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + _precomputed_rows_s.append((row_local, row, row_ctx, row_pred)) + + for mr in range_constexpr(m_reps_s): + row_local, row, row_ctx, row_pred = _precomputed_rows_s[mr] + + def _do_store_row_split(): + row_base_lds = row_local * _half_n_idx + for nr in range_constexpr(n_reps_s): + col_base_nr = arith.constant( + nr * (CShuffleNLane_s * EVec), index=True + ) + col_pair0_local = col_base_nr + (n_lane_s * c_evec) + lds_idx = row_base_lds + col_pair0_local + + _if_ld = scf.IfOp(_is_group_b, [vec_frag], has_else=True) + with ir.InsertionPoint(_if_ld.then_block): + fb = vector.load_op(vec_frag, lds_out_split, [lds_idx]) + scf.YieldOp([fb]) + with ir.InsertionPoint(_if_ld.else_block): + fa = vector.load_op(vec_frag, lds_out, [lds_idx]) + scf.YieldOp([fa]) + frag = _if_ld.results[0] + + col_pair0 = col_pair0_local + arith.select( + _is_group_b, _half_n_idx, _zero_idx + ) + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row_split() + else: + _do_store_row_split() + + return # split path complete + + # ===================== Standard (non-split) path below ===================== + + # ---------------- Step 1: write C tile to LDS (row-major, fp16) ---------------- + tile_n_idx = arith.constant(int(tile_n), index=True) + n_tile_base_v = n_tile_base + col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) + + _lds_row_base_offset = ( + lds_row_offset * tile_n_idx if lds_row_offset is not None else None + ) + + def _write_row(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + + # Ensure all LDS reads finished before the lds write. + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row, + ) + + # Ensure all LDS writes are visible before the shuffle-read. + gpu.barrier() + + # ---------------- Step 2: shuffle mapping + half2 store/atomic ---------------- + CShuffleNLane = int(cshuffle_nlane) + CShuffleMLane = int(cshuffle_mlane) + EVec = int(e_vec) + + m_reps_shuffle = int(tile_m) // CShuffleMLane + n_reps_shuffle = int(tile_n) // (CShuffleNLane * EVec) + + c_nlane = fx.Index(CShuffleNLane) + m_lane = tx // c_nlane + n_lane = tx % c_nlane + c_evec = fx.Index(EVec) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + # Batch-precompute all row contexts (sorted_idx loads) before the store loop. + # This issues all buffer_load instructions upfront so the compiler can pipeline + # them instead of serializing each load with s_waitcnt vmcnt(0). + _precomputed_rows = [] + for mr in range_constexpr(m_reps_shuffle): + row_base_m = arith.constant(mr * CShuffleMLane, index=True) + row_local = row_base_m + m_lane + row = bx_m_v + row_local + + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + + # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)` and `scf` + # is provided, we can skip the entire N-loop for invalid rows (cheaper than per-store checks). + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + + _precomputed_rows.append((row_local, row, row_ctx, row_pred)) + + # Now perform LDS reads and stores using the pre-fetched row contexts. + for mr in range_constexpr(m_reps_shuffle): + row_local, row, row_ctx, row_pred = _precomputed_rows[mr] + + def _do_store_row(): + row_base_lds = row_local * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + for nr in range_constexpr(n_reps_shuffle): + col_base_nr = arith.constant(nr * (CShuffleNLane * EVec), index=True) + col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile + + lds_idx_pair = row_base_lds + col_pair0 + frag = vector.load_op(vec_frag, lds_out, [lds_idx_pair]) + + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row() + else: + _do_store_row() + + +def mfma_epilog( + *, + use_cshuffle: bool, + # Common (always required) + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + # Default epilog (required when use_cshuffle=False) + body_row: Callable | None = None, + # CShuffle epilog (required when use_cshuffle=True) + vector=None, + gpu=None, + scf=None, + tile_m: int | None = None, + tile_n: int | None = None, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + num_acc_n: int | None = None, + tx=None, + lane_mod_16=None, + by_n=None, + n_tile_base=None, + lds_out=None, + write_row_to_lds: Callable | None = None, + precompute_row: Callable | None = None, + store_pair: Callable | None = None, + frag_elem_type: ir.Type | None = None, +): + if not use_cshuffle: + if body_row is None: + raise ValueError("mfma_epilog(use_cshuffle=False) requires `body_row`.") + return default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=body_row, + ) + + return c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=int(tile_m), + tile_n=int(tile_n), + e_vec=int(e_vec), + cshuffle_nlane=int(cshuffle_nlane), + block_size=int(block_size), + m_repeat=m_repeat, + num_acc_n=int(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=frag_elem_type, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +# ========================================================================= +@contextmanager +def _if_then(if_op): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def _barrier(vmcnt=63, lgkmcnt=63): + """Emit s_waitcnt + s_barrier via inline asm. + + Bypasses LLVM SIInsertWaitcnts which would insert a conservative + s_waitcnt vmcnt(0) lgkmcnt(0) before every S_BARRIER MI. + """ + parts = [] + needs_waitcnt = vmcnt < 63 or lgkmcnt < 63 + if needs_waitcnt: + wc = [] + if vmcnt < 63: + wc.append(f"vmcnt({vmcnt})") + if lgkmcnt < 63: + wc.append(f"lgkmcnt({lgkmcnt})") + parts.append("s_waitcnt " + " ".join(wc)) + parts.append("s_barrier") + llvm.InlineAsmOp( + res=None, + operands_=[], + asm_string="\n".join(parts), + constraints="", + has_side_effects=True, + is_align_stack=False, + ) + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm1( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + act: str = "silu", + use_cshuffle_epilog: bool | None = None, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 1, + use_async_copy: bool = False, + waves_per_eu: int = 4, + k_batch: int = 1, + b_nt: int = 0, + gate_mode: GateMode = GateMode.SEPARATED, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Compile stage1 kernel (gate+up with silu/swiglu). + + GEMM: act(X @ W_gate.T, X @ W_up.T) -> [tokens*topk, inter_dim] + Direct store (no atomic). When k_batch>1 (split-K), each CTA + computes a K-slice and atomically adds gate/up partials. + Note: persist_m=1 (no persistence) is optimal for stage1 because K=model_dim + is large, so each CTA is already compute-heavy. persist_m>1 serializes M blocks + that the GPU can process in parallel. + + gate_mode controls the gate/up computation strategy — see GateMode enum. + """ + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + sort_block_m = max(32, tile_m) + num_waves = min(4, tile_n // 32) + total_threads = num_waves * 64 + pack_M = 1 if tile_m < 32 else 2 + n_per_wave = tile_n // num_waves + pack_N = min(2, n_per_wave // 16) + pack_K = 2 + scale_mn_pack = 2 + elem_bytes = 1 + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + if (tile_k_bytes % 64) != 0: + raise ValueError(f"tile_k_bytes must be divisible by 64, got {tile_k_bytes}") + + out_s = str(out_dtype).strip().lower() + out_is_f32 = out_s in ("f32", "fp32", "float") + out_is_bf16 = out_s in ("bf16", "bfloat16") + is_int4 = b_dtype == "int4" + is_int8 = False + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + mock_gate_only = gate_mode is GateMode.MOCK_GATE_ONLY + gate_up_interleave = gate_mode is GateMode.INTERLEAVE + gate_only = gate_mode is GateMode.GATE_ONLY + + # Padding semantics: model_dim and inter_dim INCLUDE padding. + # model_dim = model_dim_true + model_dim_pad (K direction) + # inter_dim = inter_dim_true + inter_dim_pad (N direction) + # Tensor sizes use the padded dimensions (inter_dim, model_dim). + # Padding only affects kernel internal logic and grid computation. + _inter_dim_valid = inter_dim - inter_dim_pad + + # Split-K validation + _is_splitk = k_batch > 1 + if mock_gate_only and not _is_splitk: + raise ValueError("mock_gate_only requires k_batch > 1 (split-K)") + 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}" + + out_dtype = "bf16" + else: + _k_per_batch = model_dim + _k_dim = _k_per_batch + + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + f"tile_m*tile_k*elem_bytes must be divisible by {total_threads}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if use_cshuffle_epilog is None: + _use_cshuffle_epilog = os.environ.get("FLIR_MOE_STAGE1_CSHUFFLE", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + else: + _use_cshuffle_epilog = bool(use_cshuffle_epilog) + + _need_fp4 = out_dtype == "fp4" + _need_fp8 = out_dtype == "fp8" + _need_quant = _need_fp4 or _need_fp8 + _need_sort = _need_quant + + if _need_quant: + _use_cshuffle_epilog = True + + _fp4q_tag = "_fp4q" if _need_fp4 else "" + _fp8q_tag = "_fp8q" if _need_fp8 else "" + _sort_tag = "_sort" if _need_sort else "" + _async_tag = "_async" if use_async_copy else "" + _sk_tag = f"_sk{k_batch}" if _is_splitk else "" + _go_tag = "_go" if mock_gate_only else "" + _gui_tag = "_gui" if gate_up_interleave else "" + _as1_tag = "_as1" if a_scale_one else "" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe1_silu_mul_a{a_dtype}_w{b_dtype}_{out_s}" + f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{_fp4q_tag}{_fp8q_tag}{_sort_tag}{_async_tag}{_sk_tag}{_go_tag}{_gui_tag}{_as1_tag}{_xcd_tag}_v32" + ).replace("-", "_") + + # -- LDS sizing -- + _cshuffle_elem_bytes = 4 if _need_quant else (4 if out_is_f32 else 2) + _single_x_bytes = int(tile_m) * int(lds_stride) * int(a_elem_bytes) + lds_out_bytes = ( + _cshuffle_elem_bytes * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + # Determine whether we need wave-group split for lds_out. + # Standard layout: pong = max(input, lds_out) + tid, ping = input. + # When this overflows, split lds_out into two halves across pong & ping. + _GLOBAL_ALIGN = 1024 + _std_pong = max(_single_x_bytes, lds_out_bytes) + lds_tid_bytes + _std_ping = _single_x_bytes + _std_pong_aligned = allocator_pong._align(_std_pong, 128) + _std_total = allocator_pong._align( + _std_pong_aligned, _GLOBAL_ALIGN + ) + allocator_pong._align(_std_ping, 128) + _lds_limit = {"gfx950": 163840, "gfx942": 65536}.get(gpu_arch, 0) + + _split_lds_out = ( + _lds_limit > 0 + and lds_out_bytes > 0 + and _std_total > _lds_limit + and num_waves >= 2 + ) + + if _split_lds_out: + _half_out_bytes = _cshuffle_elem_bytes * int(tile_m) * (int(tile_n) // 2) + _pong_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + _ping_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + else: + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if waves_per_eu is not None and waves_per_eu >= 1: + _total_cu_lds = 160 * 1024 + _min_lds = _total_cu_lds // (waves_per_eu + 1) + 1 + _pong_sz = allocator_pong._align(allocator_pong.ptr, 128) + _ping_sz = allocator_ping._align(allocator_ping.ptr, 128) + _cur_lds = _pong_sz + _ping_sz + if _cur_lds < _min_lds: + allocator_ping.ptr += _min_lds - _cur_lds + + kpack_bytes = 8 if is_int4 else 16 + out_elem_bytes = 4 if out_is_f32 else 2 + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * (2 * inter_dim) * model_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * (2 * inter_dim) * 4 + + _e_vec_s1 = min(tile_n // 32, 8) + if _need_quant: + _e_vec_s1 = max(2, _e_vec_s1) + _num_threads_per_quant_blk_s1 = 32 // _e_vec_s1 + _shuffle_dists_s1 = [] + _sh_val = 1 + while _sh_val < _num_threads_per_quant_blk_s1: + _shuffle_dists_s1.append(_sh_val) + _sh_val *= 2 + _num_shuffle_steps_s1 = len(_shuffle_dists_s1) + + # ---- Unified pipeline schedule (outside @flyc.kernel) ---- + # Each scheduling phase is a dict: + # mfma: [(k_idx, mi_idx, ikxdl, imxdl, asv_idx), ...] + # a_reads: [(k, mi), ...] # A ds_read subtiles + # b_loads: [('gate'/'up', ku, ni), ...] # B VMEM loads + # has_scale: bool # A/B scale VMEM loads + _pipe_m_repeat = tile_m // 16 + _pipe_k_unroll = tile_k_bytes // 128 + _pipe_k_unroll_packed = _pipe_k_unroll // pack_K + _pipe_m_repeat_packed = _pipe_m_repeat // pack_M + _pipe_num_acc_n = n_per_wave // 16 + + # A ds_read groups: group by mi (same mi, all k values together) + _pipe_a_groups = [] + for _mi in range(_pipe_m_repeat): + _grp = [] + for _k in range(_pipe_k_unroll): + _grp.append((_k, _mi)) + if len(_grp) == 2: + _pipe_a_groups.append(_grp) + _grp = [] + if _grp: + _pipe_a_groups.append(_grp) + + # B VMEM loads: individual gate/up loads + _pipe_b_loads = [] + for ku in range(_pipe_k_unroll): + for ni in range(_pipe_num_acc_n): + _pipe_b_loads.append(("gate", ku, ni)) + if not mock_gate_only and not gate_up_interleave: + _pipe_b_loads.append(("up", ku, ni)) + + # MFMA order: B-major (fix B, cycle all A tiles before next B) + # Each entry: one (k, ni) pair; the compute function loops over all mi. + # This keeps B operands (from VMEM) fixed while cycling A (from LDS, no wait). + _pipe_num_acc_n_packed = _pipe_num_acc_n // pack_N + _pipe_all_mfma = [] + for _ku128 in range(_pipe_k_unroll_packed): + for _ni_packed in range(_pipe_num_acc_n_packed): + for _ikxdl in range(pack_K): + for _inxdl in range(pack_N): + _k_idx = _ku128 * pack_K + _ikxdl + _ni_idx = _ni_packed * pack_N + _inxdl + _pipe_all_mfma.append((_k_idx, _ni_idx, _ikxdl, _inxdl, _ku128)) + + # Group MFMAs per scheduling phase (wider M -> more MFMAs per phase) + _pipe_mfma_per_phase = max(1, len(_pipe_all_mfma) // 4) + _pipe_n_phases = len(_pipe_all_mfma) // _pipe_mfma_per_phase + + # Build unified phase descriptors + _a_groups_per_phase = (len(_pipe_a_groups) + _pipe_n_phases - 1) // _pipe_n_phases + _pipe_phases = [] + _mfma_i = 0 + _a_i = 0 + for _p in range(_pipe_n_phases): + _a_reads = [] + for _ in range(_a_groups_per_phase): + if _a_i < len(_pipe_a_groups): + _a_reads.extend(_pipe_a_groups[_a_i]) + _a_i += 1 + _phase = { + "mfma": _pipe_all_mfma[_mfma_i : _mfma_i + _pipe_mfma_per_phase], + "a_reads": _a_reads, + "b_loads": [], + "has_scale": (_p == 0), + } + _mfma_i += _pipe_mfma_per_phase + _pipe_phases.append(_phase) + + # Distribute B loads evenly across phases 1..n-1 (phase 0 has scales) + _bi = 0 + for _p in range(1, _pipe_n_phases): + _rem_b = len(_pipe_b_loads) - _bi + _rem_p = _pipe_n_phases - _p + _n_b = (_rem_b + _rem_p - 1) // _rem_p if _rem_p > 0 else 0 + for _ in range(_n_b): + if _bi < len(_pipe_b_loads): + _pipe_phases[_p]["b_loads"].append(_pipe_b_loads[_bi]) + _bi += 1 + + # Extract flat lists for kernel access (avoids dict access in AST rewriter) + _pp_mfma = [p["mfma"] for p in _pipe_phases] + _pp_a_reads = [p["a_reads"] for p in _pipe_phases] + _pp_b_loads = [p["b_loads"] for p in _pipe_phases] + _pp_has_scale = [p["has_scale"] for p in _pipe_phases] + + fp4_ratio = 2 if a_dtype == "fp4" else 1 + gui_ratio = 1 if gate_up_interleave else 2 + _vmcnt_before_barrier = tile_m // 32 // fp4_ratio + tile_n // 32 * gui_ratio + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = arith.constant_vector(0.0, vec4_f32) + + # --- Stage1 dimension mapping --- + # X: [tokens, model_dim] -- M = sorted tokens, K = model_dim + # W: [E*2*inter_dim, model_dim] gate portion -- N = inter_dim + # Out: [tokens*topk, inter_dim] + + # B preshuffle layout: [E*2*inter_dim, model_dim] + # Gate rows for expert e: [e*2*inter_dim, e*2*inter_dim + inter_dim) + c_n_total = arith.constant(experts * (2 * inter_dim), index=True) + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in // pack_K, + kpack_bytes=kpack_bytes, + elem_bytes=b_elem_bytes, + # k_major=True, + ) + layout_b = b_layout.layout_b + + # A-scale: [sorted_size, K/32] -- pre-scattered by caller into sorted layout + # Same as stage2: indexed by sorted_row position, not by token_id. + sorted_m = size_expert_ids_in * arith.constant(sort_block_m, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=sorted_m, c_k=arith.constant(model_dim, index=True) + ) + # B-scale: [E*2*inter_dim, K/32] + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=arith.constant(model_dim, index=True) + ) + + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + if const_expr(use_async_copy and a_elem_vec_pack > 1): + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along inter_dim (N) + bx_persist = gpu.block_id("y") # persistent WG index + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S1 = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_idp_sw = arith.constant(2 * inter_dim_pad, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + _gx = (n_in - _c_idp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + else: + _c2_sw = arith.constant(2, index=True) + _gx = ( + (n_in - _c_idp_sw + _c2_sw * _c_tn_sw - _c1_sw) + / _c_tn_sw + / _c2_sw + ) + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S1, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S1 = xcd_swizzle + _c_wgm = arith.constant(_WGM_S1, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + by_n = by * arith.constant(tile_n, index=True) + + k_base_idx = arith.index(0) + if const_expr(_is_splitk): + bz = gpu.block_id("z") # K-batch id + k_base_idx = bz * arith.constant(_k_dim, index=True) + + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((num_waves, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + _lds_out_elem_type = ( + T.f32 if _need_quant else (T.bf16 if out_is_bf16 else T.f16) + ) + if const_expr(_split_lds_out and _use_cshuffle_epilog): + _half_out_elems = int(tile_m) * (int(tile_n) // 2) + lds_out = SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + lds_out_B = SmemPtr( + base_ptr_ping, + lds_ping_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + else: + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_out_B = None + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # Buffer resources + c_a_pack = arith.constant(int(a_elem_vec_pack), index=True) + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + + # X: [tokens, model_dim] + x_nbytes_idx = (tokens_in * k_in * c_elem_bytes) / c_a_pack + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # Out: [tokens*topk, inter_dim] + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not (is_f16_a or a_scale_one)): + # A scale: [sorted_size, model_dim/32] pre-scattered by caller + c32 = arith.constant(32, index=True) + kblk = k_in / c32 + sx_nbytes_idx = sorted_m * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + c32 = arith.constant(32, index=True) + kblk_w = k_in / c32 + mn_w = arith.constant(experts * (2 * inter_dim), index=True) + sw_nbytes_idx = mn_w * kblk_w + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + sorted_nbytes_idx = size_expert_ids_in * arith.constant( + sort_block_m * 4, index=True + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + eid_nbytes_idx = size_expert_ids_in * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # Sorted-scale buffer resource for fused mxfp4 quantization + _sorted_scale_cols = inter_dim // 32 + _sorted_scale_cols_i32 = arith.constant(_sorted_scale_cols, type=T.i32) + sorted_scale_rsrc = None + if const_expr(_need_sort): + _sort_rows_idx = size_expert_ids_in * arith.constant( + sort_block_m, index=True + ) + _sort_padded_rows = ( + (_sort_rows_idx + arith.constant(255, index=True)) + / arith.constant(256, index=True) + * arith.constant(256, index=True) + ) + _sort_padded_cols = arith.constant( + ((_sorted_scale_cols + 7) // 8) * 8, index=True + ) + _sort_scale_nbytes = arith.index_cast( + T.i32, _sort_padded_rows * _sort_padded_cols + ) + sorted_scale_rsrc = _ptr_buffer_resource( + arg_out_scale_sorted, _sort_scale_nbytes + ) + + # ---- persist_m loop (same pattern as stage2) ---- + _PERSIST_M = persist_m + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + _c_pm = arith.constant(_PERSIST_M, index=True) + _for_persist = scf.ForOp(_c0_p, _c_pm, _c1_p) + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + bx = bx_persist * _c_pm + _mi_p + bx_m = bx * arith.constant(sort_block_m, index=True) + + # Block validity + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, bx, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + def _moe_gemm1_body(): + # Gate expert offset: first inter_dim rows of each expert's 2*inter_dim block + expert_off_idx = expert_idx * arith.constant(2 * inter_dim, index=True) + + # X loading -- KEY DIFFERENCE from stage2: X row = token_id only + x_load_bytes = 16 + num_x_loads = bytes_per_thread_x // x_load_bytes + chunk_i32 = x_load_bytes // 4 + + c_k_div4 = ( + (k_in / c_a_pack) * arith.constant(int(a_elem_bytes), index=True) + ) / arith.index(4) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(0xFFFFFF) + 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, + ) + + def load_x(idx_i32): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + + # Decode sorted token ids -- stage1: X row = token_id (not t*topk+s) + x_row_base_div4 = [] + x_col_local_i32 = [] + x_row_local = [] + # Also store token_id and slot_id for output indexing + + 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + + # KEY: X row base uses token_id only (not t*topk+s) + t_idx = arith.index_cast(ir.IndexType.get(), t_safe) + x_row_base_div4.append(t_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(a_elem_bytes), index=True) + ) / arith.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) + parts.append(vector.bitcast(T.vec(4, i32), x_vec)) + return parts + + # Wave/lane decomposition (identical to stage2) + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + row_a_lds = lane_mod_16 + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_n_id = wave_id % arith.constant(num_waves, index=True) + n_tile_base = wave_n_id * c_n_per_wave + + # N-tile precompute for gate AND up weights + gate_n_intra_list = [] + gate_n_blk_list = [] + up_n_intra_list = [] + up_n_blk_list = [] + col_g_list = [] + c_n0_static = experts * (2 * inter_dim) // 16 + layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1)) + inter_idx = arith.constant(inter_dim, index=True) + + for i in range_constexpr(num_acc_n): + offset = i * 16 + c_offset = arith.constant(offset, index=True) + if const_expr(not gate_up_interleave): + col_g = by_n + n_tile_base + c_offset + lane_mod_16 + col_g_list.append(col_g) + + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + # Gate/interleave: rows [expert_off, expert_off + 2*inter_dim) + gate_row_w = expert_off_idx + global_n + gate_coord = idx2crd(gate_row_w, layout_n_blk_intra) + gate_n_blk_list.append(layout_get(gate_coord, 0)) + gate_n_intra_list.append(layout_get(gate_coord, 1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_row_w = gate_row_w + inter_idx + up_coord = idx2crd(up_row_w, layout_n_blk_intra) + up_n_blk_list.append(layout_get(up_coord, 0)) + up_n_intra_list.append(layout_get(up_coord, 1)) + + if const_expr(gate_up_interleave): + _gui_num_acc_n_out = num_acc_n // pack_N + for _gui_i in range_constexpr(_gui_num_acc_n_out): + _gui_offset = _gui_i * 16 + _gui_c_offset = arith.constant(_gui_offset, index=True) + _gui_col_g = ( + (by_n + n_tile_base) // arith.constant(2, index=True) + + _gui_c_offset + + lane_mod_16 + ) + col_g_list.append(_gui_col_g) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku = tile_k // k_unroll + _pad_k_elems = ( + (model_dim_pad % tile_k) + if (not _is_splitk and model_dim_pad > 0) + else 0 + ) + _pad_ku_skip = _pad_k_elems // _K_per_ku + _tail_ku = k_unroll - _pad_ku_skip + _tail_ku_packed = ( + (_tail_ku + pack_K - 1) // pack_K if _pad_ku_skip > 0 else None + ) + + # B load for gate and up separately + def load_b_packs_k64(base_k, ku: int, n_blk, n_intra): + c64 = arith.constant(64, index=True) + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0 = base_k_bytes // c64 + arith.constant(ku, index=True) + k1 = lane_div_16 + coord_pack = (n_blk, k0, k1, n_intra, arith.constant(0, index=True)) + idx_pack = crd2idx(coord_pack, layout_b) + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + """Load B tiles. Returns (gate_b_tile, up_b_tile). + When mock_gate_only or gate_up_interleave, up_b_tile is None.""" + gate_b_tile = [] + up_b_tile = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_limit): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + gb0, gb1 = load_b_packs_k64( + base_k, ku, gate_n_blk_list[ni], gate_n_intra_list[ni] + ) + g_packs0.append(gb0) + g_packs1.append(gb1) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + ub0, ub1 = load_b_packs_k64( + base_k, ku, up_n_blk_list[ni], up_n_intra_list[ni] + ) + u_packs0.append(ub0) + u_packs1.append(ub1) + gate_b_tile.append((g_packs0, g_packs1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_b_tile.append((u_packs0, u_packs1)) + return gate_b_tile, up_b_tile + + # Pre-compute scale base element indices (K-loop invariant). + # idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + # Split into: base_elem = mni * stride_n0 + lane_elem (invariant) + # k_elem = ku * stride_k0 (per-iteration) + _scale_lane_elem = ( + lane_div_16 * layout_b_scale.stride_klane + lane_mod_16 + ) + + _gate_scale_bases = [] + _up_scale_bases = [] + for _ni in range_constexpr(num_acc_n_packed): + _col_base = ( + by_n + + n_tile_base + + arith.constant(_ni * 16 * pack_N, index=True) + ) + _gate_mni = (expert_off_idx + _col_base) // arith.constant( + 32, index=True + ) + _gate_scale_bases.append( + _gate_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + if const_expr(not mock_gate_only and not gate_up_interleave): + _up_mni = ( + expert_off_idx + inter_idx + _col_base + ) // arith.constant(32, index=True) + _up_scale_bases.append( + _up_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + + if const_expr(not a_scale_one): + _a_scale_bases = [] + for _mi in range_constexpr(m_repeat_packed): + _a_mni = _mi + bx_m // scale_mn_pack // 16 + _a_scale_bases.append( + _a_mni * layout_a_scale.stride_n0 + _scale_lane_elem + ) + + _c16_idx = arith.constant(16, index=True) + _c2_idx = arith.constant(2, index=True) + _scale_mask_lo = arith.constant(0xFF, type=T.i32) + + _m_half_idx = arith.constant(0, type=T.i32) + _m_half_i32 = arith.constant(0, type=T.i32) + _scale_shift = arith.constant(0, type=T.i32) + _scale_shift_hi = arith.constant(0, type=T.i32) + _n_half_idx = arith.constant(0, type=T.i32) + _n_half_i32 = arith.constant(0, type=T.i32) + _bscale_shift = arith.constant(0, type=T.i32) + _bscale_shift_hi = arith.constant(0, type=T.i32) + if const_expr(pack_M < scale_mn_pack): + _m_half_idx = (bx_m // _c16_idx) % _c2_idx + _m_half_i32 = arith.index_cast(T.i32, _m_half_idx) + _scale_shift = _m_half_i32 * arith.constant(8, type=T.i32) + _scale_shift_hi = _scale_shift + arith.constant(16, type=T.i32) + + if const_expr(pack_N < scale_mn_pack): + _n_half_idx = (n_tile_base // _c16_idx) % _c2_idx + _n_half_i32 = arith.index_cast(T.i32, _n_half_idx) + _bscale_shift = _n_half_i32 * arith.constant(8, type=T.i32) + _bscale_shift_hi = _bscale_shift + arith.constant(16, type=T.i32) + + def _rearrange_a_scale(raw_i32): + """Rearrange scale bytes for pack_M=1: extract m_half's k0,k1 bytes.""" + if const_expr(pack_M >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _scale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _scale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + def _rearrange_b_scale(raw_i32): + """Rearrange scale bytes for pack_N=1: extract n_half's k0,k1 bytes.""" + if const_expr(pack_N >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _bscale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _bscale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + if const_expr(a_scale_one): + _as1_const = arith.constant(0x7F7F7F7F, type=T.i32) + _as1_vec = vector.from_elements(T.vec(1, T.i32), [_as1_const]) + + def prefetch_ab_scale_tile(base_k, ku_packed_limit=k_unroll_packed): + a_scale_tile = [] + gate_b_scale = [] + up_b_scale = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_packed_limit): + k_off = (ku + base_k) * layout_b_scale.stride_k0 + for mi in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + a_scale_tile.append(_as1_vec) + else: + s = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[mi] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + s = _rearrange_a_scale(s) + a_scale_tile.append( + vector.from_elements(T.vec(1, T.i32), [s]) + ) + for ni in range_constexpr(num_acc_n_packed): + gs = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + gs = _rearrange_b_scale(gs) + gate_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [gs]) + ) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + us = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + us = _rearrange_b_scale(us) + up_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [us]) + ) + return [a_scale_tile, gate_b_scale, up_b_scale] + + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x4=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + if const_expr(use_async_copy): + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(elem_bytes), index=True) + ) / arith.index(4) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def prefetch_full_a_from_lds(lds_buffer, ku_limit=k_unroll): + """Load entire A tile from LDS into registers before compute.""" + a_regs = [] + for k_idx in range_constexpr(ku_limit): + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + for mi_idx in range_constexpr(m_repeat): + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64( + curr_row, col_base + 64, lds_buffer + ) + a_regs.append((a0, a1, a2, a3)) + else: + a_regs.append((a0, a1)) + return a_regs + + # Compute tile: gate + up MFMA interleaved, same A data, different B data. + # Two accumulator sets; after all K tiles, acc = acc_gate + acc_up (f32 add). + def compute_tile( + acc_gate_in, + acc_up_in, + gate_b_tile_in, + up_b_tile_in, + a_tile_regs, + a_scale=None, + gate_b_scale=None, + up_b_scale=None, + *, + prefetch_epilogue=False, + ku_count=k_unroll, + ): + gate_list = list(acc_gate_in) + _single_b = mock_gate_only or gate_up_interleave + up_list = None if _single_b else list(acc_up_in) + mfma_res_ty = vec4_f32 + epilogue_pf = None + bias_pf = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + if const_expr(gate_up_interleave): + bias_pf = [] + for ni in range_constexpr(num_acc_n): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (ni % 2 == 1) + else arith.constant(0, index=True) + ) + bias_offset = ( + expert_off_idx + _up_off + _logical_col + ) + bias_pf.append( + _load_bias_scalar(bias_rsrc, bias_offset) + ) + else: + gate_bias_pf = [] + up_bias_pf = ( + [] if const_expr(not mock_gate_only) else None + ) + for ni in range_constexpr(num_acc_n): + global_n = ( + by_n + + n_tile_base + + arith.constant(ni * 16, index=True) + + lane_mod_16 + ) + gate_bias_pf.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + global_n + ) + ) + if const_expr(not mock_gate_only): + up_bias_pf.append( + _load_bias_scalar( + bias_rsrc, + expert_off_idx + inter_idx + global_n, + ) + ) + bias_pf = (gate_bias_pf, up_bias_pf) + tw_pf = None + if const_expr(doweight_stage1): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + sorted_row_pf = bx_m + mi_base_pf + row_off_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias_pf) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + _eff_packed = (ku_count + pack_K - 1) // pack_K + # B-major: fix B (ni), cycle A (mi) -- B from VMEM stays + # in registers while A from LDS is repacked per mi. + for ku128 in range_constexpr(_eff_packed): + for ni in range_constexpr(num_acc_n_packed): + gate_bs_i32 = gate_b_scale[ku128 * num_acc_n_packed + ni] + gate_bs_val = vector.extract( + gate_bs_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(not _single_b): + up_bs_i32 = up_b_scale[ku128 * num_acc_n_packed + ni] + up_bs_val = vector.extract( + up_bs_i32, static_position=[0], dynamic_position=[] + ) + for ikxdl in range_constexpr(pack_K): + k_idx = ku128 * pack_K + ikxdl + if const_expr(k_idx < ku_count): + gate_bp0, gate_bp1 = gate_b_tile_in[k_idx] + if const_expr(not _single_b): + up_bp0, up_bp1 = up_b_tile_in[k_idx] + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + gb0 = gate_bp0[ni_idx] + gb1 = gate_bp1[ni_idx] + gb128 = pack_i64x4_to_i32x8( + gb0, gb1, c0_i64, c0_i64 + ) + if const_expr(not _single_b): + ub0 = up_bp0[ni_idx] + ub1 = up_bp1[ni_idx] + ub128 = pack_i64x4_to_i32x8( + ub0, ub1, c0_i64, c0_i64 + ) + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ + ku128 * m_repeat_packed + mi + ] + a_scale_val = vector.extract( + a_scale_i32, + static_position=[0], + dynamic_position=[], + ) + for imxdl in range_constexpr(pack_M): + mi_idx = mi * pack_M + imxdl + _a_reg_idx = k_idx * m_repeat + mi_idx + if const_expr(is_f8_a): + a0, a1, a2, a3 = a_tile_regs[ + _a_reg_idx + ] + a128 = pack_i64x4_to_i32x8( + a0, a1, a2, a3 + ) + else: + a0, a1 = a_tile_regs[_a_reg_idx] + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + ) + if const_expr(not _single_b): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + return gate_list, up_list, epilogue_pf + + def load_a_subtile(k_idx, mi_idx, lds_buffer): + """Load a single A sub-tile from LDS (one ds_read).""" + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64(curr_row, col_base + 64, lds_buffer) + return (a0, a1, a2, a3) + else: + return (a0, a1) + + _single_b_pipe = mock_gate_only or gate_up_interleave + + def compute_bmajor_mfma_phase( + all_a_tiles, + gate_b_single, + up_b_single, + a_scale_vals, + gate_bs_val, + up_bs_val, + gate_list, + up_list, + k_idx, + ni_idx, + ikxdl, + inxdl, + ): + """B-major MFMA: fix one B (ni), cycle all A tiles (mi). + + Packs B once and reuses across all mi iterations. + A tiles come from LDS (already available, no VMEM wait). + + all_a_tiles: flat list indexed by [k*m_repeat + mi]. + gate_b_single/up_b_single: (b0, b1) for one specific ni. + When _single_b_pipe (mock_gate_only or interleave), up_b_single is None. + a_scale_vals: list of A scale scalars indexed by mi_packed. + """ + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def _pack(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + mfma_res_ty = vec4_f32 + gb128 = _pack(gate_b_single[0], gate_b_single[1], c0_i64, c0_i64) + if const_expr(not _single_b_pipe): + ub128 = _pack(up_b_single[0], up_b_single[1], c0_i64, c0_i64) + + for mi_p in range_constexpr(m_repeat_packed): + a_scale_val = a_scale_vals[mi_p] + for imxdl in range_constexpr(pack_M): + mi_idx = mi_p * pack_M + imxdl + a_reg = all_a_tiles[k_idx * m_repeat + mi_idx] + + if const_expr(is_f8_a): + a128 = _pack(a_reg[0], a_reg[1], a_reg[2], a_reg[3]) + else: + a128 = _pack(a_reg[0], a_reg[1], c0_i64, c0_i64) + + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + if const_expr(not _single_b_pipe): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + + def _interleaved_half( + lds_read, + lds_write, + next_k_dma_py, + next_k_load, + prev_a_tile, + prev_gate_w, + prev_up_w, + prev_a_scale, + prev_gate_bs, + prev_up_bs, + acc_gate, + acc_up, + ): + """One flatmm-style interleaved half-iteration (deep pipeline). + + Generalized for arbitrary m_repeat (block_m=32, 64, ...). + DMA targets lds_write (OTHER buffer) while ds_read uses + lds_read (already DMA'd in previous half). + + Interleaving schedule (per half): + Phase 0: scale VMEM + 2 ds_read(A) -> 4 MFMA(prev) + Phase 1..N: B VMEM(distributed) + 2 ds_read(A, if avail) -> 4 MFMA(prev) + Phase N+1..: remaining B VMEM -> 4 MFMA(prev) + """ + _abs_k = k_base_idx + arith.constant(next_k_load, index=True) + _bk = _abs_k // arith.constant(2, index=True) + _sk = _abs_k // arith.constant(pack_K * 128, index=True) + _k_off = _sk * layout_b_scale.stride_k0 + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(_vmcnt_before_barrier) + _barrier() + rocdl.sched_barrier(0) + + # DMA A to OTHER buffer (for next half), non-blocking + _abs_k_dma = k_base_idx + arith.constant(next_k_dma_py, index=True) + if const_expr(use_async_copy and next_k_dma_py < int(_k_dim)): + prefetch_x_to_lds(_abs_k_dma, lds_write) + if const_expr(not use_async_copy): + _x_regs = load_x_tile(_abs_k_dma) + + # ---- Extract previous scale values ---- + _prev_asvs = [] + for _mi_p in range_constexpr(m_repeat_packed): + _prev_asvs.append( + vector.extract( + prev_a_scale[_mi_p], + static_position=[0], + dynamic_position=[], + ) + ) + _prev_gsv_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _prev_gsv_list.append( + vector.extract( + prev_gate_bs[_gs_ni], + static_position=[0], + dynamic_position=[], + ) + ) + if const_expr(not _single_b_pipe): + _prev_usv_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _prev_usv_list.append( + vector.extract( + prev_up_bs[_us_ni], + static_position=[0], + dynamic_position=[], + ) + ) + + # ---- Execute phases from unified schedule ---- + _a_all = {} + _b_gate_all = {} + _b_up_all = {} + + for _p in range_constexpr(_pipe_n_phases): + # Scale VMEM loads (phase 0 only) + if const_expr(_pp_has_scale[_p]): + _new_as_list = [] + for _mi_p in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + _new_as_list.append(_as1_const) + else: + _raw_as = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[_mi_p] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_as_list.append(_rearrange_a_scale(_raw_as)) + _new_gs_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _gs_raw = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[_gs_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_gs_list.append(_rearrange_b_scale(_gs_raw)) + if const_expr(not _single_b_pipe): + _new_us_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _us_raw = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[_us_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_us_list.append(_rearrange_b_scale(_us_raw)) + + # B VMEM loads + for _b_j in range_constexpr(len(_pp_b_loads[_p])): + _b_type, _b_ku, _b_ni = _pp_b_loads[_p][_b_j] + if const_expr(_b_type == "gate"): + _b_gate_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + gate_n_blk_list[_b_ni], + gate_n_intra_list[_b_ni], + ) + else: + _b_up_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + up_n_blk_list[_b_ni], + up_n_intra_list[_b_ni], + ) + + # A ds_reads + rocdl.sched_barrier(0) + for _a_j in range_constexpr(len(_pp_a_reads[_p])): + _ak, _ami = _pp_a_reads[_p][_a_j] + _a_all[(_ak, _ami)] = load_a_subtile( + _ak, + _ami, + lds_read, + ) + rocdl.sched_barrier(0) + + # MFMAs on prev data + rocdl.s_setprio(1) + for _m_j in range_constexpr(len(_pp_mfma[_p])): + _k_idx, _ni_idx, _ikxdl, _inxdl, _ku128 = _pp_mfma[_p][_m_j] + _ni_packed_idx = _ni_idx // pack_N + _up_b_single = ( + ( + prev_up_w[_k_idx][0][_ni_idx], + prev_up_w[_k_idx][1][_ni_idx], + ) + if not _single_b_pipe + else None + ) + compute_bmajor_mfma_phase( + prev_a_tile, + ( + prev_gate_w[_k_idx][0][_ni_idx], + prev_gate_w[_k_idx][1][_ni_idx], + ), + _up_b_single, + _prev_asvs, + _prev_gsv_list[_ni_packed_idx], + ( + _prev_usv_list[_ni_packed_idx] + if not _single_b_pipe + else None + ), + acc_gate, + acc_up, + _k_idx, + _ni_idx, + _ikxdl, + _inxdl, + ) + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + + # ---- Assemble loaded data for next half-iteration ---- + cur_a_tile = [] + for _k in range_constexpr(k_unroll): + for _mi in range_constexpr(m_repeat): + cur_a_tile.append(_a_all[(_k, _mi)]) + + cur_gate_w = [] + cur_up_w = None if _single_b_pipe else [] + for ku in range_constexpr(k_unroll): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + g = _b_gate_all[(ku, ni)] + g_packs0.append(g[0]) + g_packs1.append(g[1]) + if const_expr(not _single_b_pipe): + u = _b_up_all[(ku, ni)] + u_packs0.append(u[0]) + u_packs1.append(u[1]) + cur_gate_w.append((g_packs0, g_packs1)) + if const_expr(not _single_b_pipe): + cur_up_w.append((u_packs0, u_packs1)) + + cur_a_scale = [] + for _mi_p in range_constexpr(m_repeat_packed): + cur_a_scale.append( + vector.from_elements( + T.vec(1, T.i32), + [_new_as_list[_mi_p]], + ) + ) + cur_gate_bs = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + cur_gate_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_gs_list[_gs_ni]] + ) + ) + if const_expr(not _single_b_pipe): + cur_up_bs = [] + for _us_ni in range_constexpr(num_acc_n_packed): + cur_up_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_us_list[_us_ni]] + ) + ) + else: + cur_up_bs = None + + if const_expr(not use_async_copy): + store_x_tile_to_lds(_x_regs, lds_write) + + return ( + cur_a_tile, + cur_gate_w, + cur_up_w, + cur_a_scale, + cur_gate_bs, + cur_up_bs, + acc_gate, + acc_up, + ) + + # Pipeline (split ping/pong allocators) + rocdl.sched_barrier(0) + + k0 = k_base_idx + if const_expr(use_async_copy): + prefetch_x_to_lds(k0, lds_x_pong) + else: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_x_pong) + rocdl.sched_barrier(0) + _k0_scale = k_base_idx // arith.constant(pack_K * 128, index=True) + a_scale_pong, gate_bs_pong, up_bs_pong = prefetch_ab_scale_tile( + _k0_scale + ) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + acc_gate = [acc_init] * num_acc_n * m_repeat + acc_up = ( + [acc_init] * num_acc_n * m_repeat if not _single_b_pipe else None + ) + + _k1 = k_base_idx + arith.constant(tile_k, index=True) + rocdl.sched_barrier(0) + if const_expr(use_async_copy): + prefetch_x_to_lds(_k1, lds_x_ping) + else: + _x_regs_prime = load_x_tile(_k1) + store_x_tile_to_lds(_x_regs_prime, lds_x_ping) + + _k0_b = k_base_idx // arith.constant(2, index=True) + gate_w0, up_w0 = load_b_tile(_k0_b) + # Prime the deep pipeline: DMA K=tile_k -> ping (1 tile ahead) + if const_expr(use_async_copy): + rocdl.s_waitcnt(0) + gpu.barrier() + rocdl.sched_barrier(0) + a_tile_pong = prefetch_full_a_from_lds(lds_x_pong) + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(6) + + num_k_tiles_py = int(_k_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 + + gate_w_pong = gate_w0 + up_w_pong = up_w0 + + rocdl.sched_barrier(0) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + next_k_load_1 = k_iv_py + tile_k + next_k_load_2 = k_iv_py + tile_k * 2 + next_k_dma_1 = k_iv_py + tile_k * 2 + next_k_dma_2 = k_iv_py + tile_k * 3 + + # Half 1: read ping (DMA'd prev half), DMA->pong, MFMA(pong) + ( + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_ping, + lds_x_pong, + next_k_dma_1, + next_k_load_1, + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) + + # Half 2: read pong (DMA'd Half 1), DMA->ping, MFMA(ping) + ( + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_pong, + lds_x_ping, + next_k_dma_2, + next_k_load_2, + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) + + # _wave_mod2_b = wave_id % arith.constant(2, index=True) + # _wave_odd = arith.cmpi( + # CmpIPredicate.eq, _wave_mod2_b, arith.constant(1, index=True) + # ) + # _if_wave_odd = scf.IfOp(_wave_odd) + # with ir.InsertionPoint(_if_wave_odd.then_block): + # # gpu.barrier() + # _barrier() + # scf.YieldOp([]) + + if const_expr(odd_k_tiles): + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + else: + _k_tail_rel = arith.constant(_k_dim - tile_k, index=True) + k_tail1 = k_base_idx + _k_tail_rel + x_regs_ping = [] + if const_expr(use_async_copy): + prefetch_x_to_lds(k_tail1, lds_x_ping) + else: + x_regs_ping = load_x_tile(k_tail1) + if const_expr(_pad_ku_skip > 0): + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True), + ku_limit=_tail_ku, + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True), + ku_packed_limit=_tail_ku_packed, + ) + else: + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True) + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True) + ) + acc_gate, acc_up, _ = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + ) + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_ping, lds_x_ping) + rocdl.s_waitcnt(0) + _barrier() + if const_expr(_pad_ku_skip > 0): + a_tile_ping = prefetch_full_a_from_lds( + lds_x_ping, ku_limit=_tail_ku + ) + else: + a_tile_ping = prefetch_full_a_from_lds(lds_x_ping) + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_ping, + up_w_ping, + a_tile_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + + bias_pf = None + if const_expr(epilogue_pf is not None): + _, _, bias_pf = epilogue_pf + + # Activation helpers (f32 element-wise on vec4_f32) + def _silu_elem(g): + """silu(x) = x * sigmoid(x); HW fast path: exp2, rcp""" + neg_log2e = arith.constant(-1.4426950408889634, type=f32) + t = g * neg_log2e + emu = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [t], [], []) + one = arith.constant(1.0, type=f32) + den = one + emu + sig = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den], [], []) + return g * sig + + def _silu_mul_vec4(gate_v4, up_v4): + """Element-wise silu(gate) * up on vec4_f32. + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before applying silu(gate) * up. + """ + result_elems = [] + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + if const_expr(swiglu_limit != 0): + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + result_elems.append(_silu_elem(g) * u) + return vector.from_elements(vec4_f32, result_elems) + + def _swiglu_mul_vec4(gate_v4, up_v4): + """Element-wise swiglu(gate, up) on vec4_f32. + swiglu(g, u) = g * sigmoid(alpha * g) * (u + 1) + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before the activation. + """ + result_elems = [] + _alpha = arith.constant(1.702, type=f32) + _one = arith.constant(1.0, type=f32) + _neg_log2e = arith.constant(-1.4426950408889634, type=f32) + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + else: + _limit = arith.constant(float(7.0), type=f32) + _neg_limit = arith.constant(-float(7.0), type=f32) + + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + t = g * _alpha * _neg_log2e + emu = llvm.call_intrinsic( + f32, "llvm.amdgcn.exp2.f32", [t], [], [] + ) + den = _one + emu + sig = llvm.call_intrinsic( + f32, "llvm.amdgcn.rcp.f32", [den], [], [] + ) + result_elems.append(g * sig * (u + _one)) + return vector.from_elements(vec4_f32, result_elems) + + def _act_vec4(gate_v4, up_v4): + """Dispatch activation based on `act` parameter.""" + if const_expr(act == "swiglu"): + return _swiglu_mul_vec4(gate_v4, up_v4) + else: + return _silu_mul_vec4(gate_v4, up_v4) + + # Add bias to raw GEMM accumulators before activation. + # bias layout: [E, 2*inter_dim] flat f32 (non-interleaved: gate then up). + # For gate_up_interleave, map physical column to logical bias offset. + if const_expr(enable_bias and not _is_splitk): + _bias_up_vals = None + if const_expr(bias_pf is not None): + if const_expr(gate_up_interleave): + _bias_gate_vals = bias_pf + else: + _bias_gate_vals, _bias_up_vals = bias_pf + else: + _bias_gate_vals = [] + for _ni in range_constexpr(num_acc_n): + if const_expr(gate_up_interleave): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((_ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (_ni % 2 == 1) + else arith.constant(0, index=True) + ) + _bias_off = expert_off_idx + _up_off + _logical_col + else: + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_off = expert_off_idx + _bn + _bias_gate_vals.append( + _load_bias_scalar(bias_rsrc, _bias_off) + ) + if const_expr(not (mock_gate_only or gate_up_interleave)): + _bias_up_vals = [] + for _ni in range_constexpr(num_acc_n): + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_up_vals.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + inter_idx + _bn + ) + ) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_gate_vals[_ni]] * 4 + ) + acc_gate[_aidx] = arith.addf(acc_gate[_aidx], _bsplat) + + if const_expr(not (mock_gate_only or gate_up_interleave)): + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_up_vals[_ni]] * 4 + ) + acc_up[_aidx] = arith.addf(acc_up[_aidx], _bsplat) + + if const_expr(gate_up_interleave and not _is_splitk): + _gui_out_n = num_acc_n // pack_N + acc = [None] * (_gui_out_n * m_repeat) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(_gui_out_n): + _g_idx = _mi * num_acc_n + _ni * pack_N + _u_idx = _g_idx + 1 + _out_idx = _mi * _gui_out_n + _ni + acc[_out_idx] = _act_vec4( + acc_gate[_g_idx], acc_gate[_u_idx] + ) + elif const_expr(not _is_splitk): + acc = [None] * (int(num_acc_n) * int(m_repeat)) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + acc[_aidx] = _act_vec4(acc_gate[_aidx], acc_up[_aidx]) + + # ---- Epilogue: CShuffle + direct store (accumulate=False) ---- + # Output: out[(t*topk+s) * inter_dim + col] = silu(gate) * up + # For split-K: skip silu, output gate/up separately with atomic add + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + _, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + tokens_i32_v = tokens_i32 + + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + if const_expr(lds_out is None): + raise RuntimeError("CShuffle epilogue requires lds_out") + + _apply_weight = doweight_stage1 and not _is_splitk + + 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, + ): + if const_expr(_apply_weight): + 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=f32 + ) + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(_apply_weight): + v = v * tw + if const_expr(_need_quant): + lds_idx = row_base_lds + col_local + vec1_f32 = T.vec(1, f32) + v1 = vector.from_elements(vec1_f32, [v]) + vector.store(v1, lds_out, [lds_idx], alignment=4) + else: + 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, [v_out]) + vector.store(v1, lds_out, [lds_idx], alignment=2) + + _out_row_stride = ( + inter_dim * 2 * out_elem_bytes + if _is_splitk + else ( + inter_dim // 2 + if _need_fp4 + else (inter_dim if _need_fp8 else inter_dim * out_elem_bytes) + ) + ) + + def precompute_row(*, row_local, row): + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32_v) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + row_byte_base = out_base_idx + ts_idx * arith.constant( + _out_row_stride, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + _e_vec = _e_vec_s1 + _e_vec_sk = 2 + _cshuffle_nlane = min(32, tile_n // _e_vec) + _cshuffle_nlane_sk = min(32, tile_n // _e_vec_sk) + _num_threads_per_quant_blk = _num_threads_per_quant_blk_s1 + + _c0_i32 = arith.constant(0, type=T.i32) + _c1_i32 = arith.constant(1, type=T.i32) + _c2_i32 = arith.constant(2, type=T.i32) + _c3_i32 = arith.constant(3, type=T.i32) + _c4_i32 = arith.constant(4, type=T.i32) + _c5_i32 = arith.constant(5, type=T.i32) + _c15_i32 = arith.constant(15, type=T.i32) + _c22_i32 = arith.constant(22, type=T.i32) + _c23_i32 = arith.constant(23, type=T.i32) + _c28_i32 = arith.constant(28, type=T.i32) + _c31_i32 = arith.constant(31, type=T.i32) + _c32_i32 = arith.constant(32, type=T.i32) + _c64_i32 = arith.constant(64, type=T.i32) + _c254_i32 = arith.constant(254, type=T.i32) + _c256_i32 = arith.constant(256, type=T.i32) + _c0xFF800000_i32 = arith.constant(0xFF800000, type=T.i32) + _c0x400000_i32 = arith.constant(0x400000, type=T.i32) + _c0x7FFFFFFF_i32 = arith.constant(0x7FFFFFFF, type=T.i32) + _c0x80000000_i32 = arith.constant(0x80000000, type=T.i32) + _c0x3F800000_i32 = arith.constant(0x3F800000, type=T.i32) # 1.0f + _c0x40C00000_i32 = arith.constant(0x40C00000, type=T.i32) # 6.0f + _c0x4A800000_i32 = arith.constant(0x4A800000, type=T.i32) + _c0xC11FFFFF_i32 = arith.constant(0xC11FFFFF, type=T.i32) + _c0x7_i32 = arith.constant(0x7, type=T.i32) + _c0_f32 = arith.constant(0.0, type=T.f32) + + _c8_i32 = arith.constant(8, type=T.i32) + _fp_headroom = 2 if _need_fp4 else (8 if _need_fp8 else 0) + _c_headroom_i32 = arith.constant(_fp_headroom, type=T.i32) + + def _f32_to_e2m1(qx_f32): + """Convert a scaled f32 value to fp4 (e2m1) 4-bit integer.""" + # Match fp4_utils.f32_to_mxfp4 / HIP quant: saturate, denorm, + # and normal round-to-nearest-even paths. + qx = qx_f32.bitcast(T.i32) + s = qx & _c0x80000000_i32 + qx_abs = qx & _c0x7FFFFFFF_i32 + denormal_mask = arith.cmpi( + CmpIPredicate.ult, qx_abs, _c0x3F800000_i32 + ) + normal_mask = arith.andi( + arith.cmpi(CmpIPredicate.ult, qx_abs, _c0x40C00000_i32), + arith.cmpi(CmpIPredicate.uge, qx_abs, _c0x3F800000_i32), + ) + + denorm_f32 = qx_abs.bitcast(T.f32) + _c0x4A800000_i32.bitcast(T.f32) + denormal_x = denorm_f32.bitcast(T.i32) - _c0x4A800000_i32 + + mant_odd = (qx_abs >> _c22_i32) & _c1_i32 + normal_x = qx_abs + _c0xC11FFFFF_i32 + mant_odd + normal_x = normal_x >> _c22_i32 + + e2m1 = arith.select(normal_mask, normal_x, _c0x7_i32) + e2m1 = arith.select(denormal_mask, denormal_x, e2m1) + return (s >> _c28_i32) | e2m1 + + if const_expr(_need_sort): + _n32_sort = _sorted_scale_cols_i32 * _c32_i32 + + # Mutable slot for split-K N-offset (gate=0, up=inter_dim) + _sk_n_offset = [0] + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(_need_quant and not _is_splitk): + frag_vals = [] + for i in range_constexpr(_e_vec): + frag_vals.append( + vector.extract( + frag, static_position=[i], dynamic_position=[] + ) + ) + + local_max = _c0_f32 + for i in range_constexpr(_e_vec): + abs_v = llvm.call_intrinsic( + f32, "llvm.fabs.f32", [frag_vals[i]], [], [] + ) + local_max = arith.maximumf(local_max, abs_v) + + for _si in range_constexpr(_num_shuffle_steps_s1): + off = arith.constant(_shuffle_dists_s1[_si], type=T.i32) + peer = local_max.shuffle_xor(off, _c64_i32) + local_max = arith.maximumf(local_max, peer) + + max_i32 = local_max.bitcast(T.i32) + # Match fp4_utils.f32_to_e8m0(max_abs / 4): round the + # exponent at the 1.5x threshold before dropping mantissa. + max_rounded = (max_i32 + _c0x400000_i32) & _c0xFF800000_i32 + exp_field = max_rounded >> _c23_i32 + e8m0_biased = arith.maxsi(exp_field - _c_headroom_i32, _c0_i32) + + quant_exp = _c254_i32 - e8m0_biased + quant_scale = (quant_exp << _c23_i32).bitcast(T.f32) + + if const_expr(_need_fp4): + fp4_vals = [] + for i in range_constexpr(_e_vec): + scaled_v = frag_vals[i] * quant_scale + fp4_vals.append(_f32_to_e2m1(scaled_v)) + + packed_i32 = fp4_vals[0] | (fp4_vals[1] << _c4_i32) + for k in range_constexpr(1, _e_vec // 2): + byte_k = fp4_vals[2 * k] | ( + fp4_vals[2 * k + 1] << _c4_i32 + ) + packed_i32 = packed_i32 | ( + byte_k << arith.constant(k * 8, type=T.i32) + ) + + ptr_addr_idx = row_byte_base + col_g0 / arith.constant( + 2, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + _pack_bytes = _e_vec // 2 + if const_expr(_pack_bytes == 1): + store_val = arith.TruncIOp(T.i8, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=1, nontemporal=True + ) + elif const_expr(_pack_bytes == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=2, nontemporal=True + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, out_ptr_v, alignment=4, nontemporal=True + ) + + elif const_expr(_need_fp8): + scaled_vals = [] + for i in range_constexpr(_e_vec): + scaled_vals.append(frag_vals[i] * quant_scale) + + ptr_addr_idx = row_byte_base + col_g0 + if const_expr(_e_vec <= 4): + packed_i32 = _c0_i32 + for _w in range_constexpr(_e_vec // 2): + packed_i32 = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[2 * _w], + scaled_vals[2 * _w + 1], + packed_i32, + _w, + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + if const_expr(_e_vec == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, + out_ptr_v, + alignment=2, + nontemporal=True, + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + else: + for _wg in range_constexpr(_e_vec // 4): + _b = _wg * 4 + packed_w = _c0_i32 + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b], + scaled_vals[_b + 1], + packed_w, + 0, + ) + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b + 2], + scaled_vals[_b + 3], + packed_w, + 1, + ) + word_ptr = ptr_addr_idx + arith.constant( + _wg * 4, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(word_ptr) + packed_raw = ( + packed_w._value + if hasattr(packed_w, "_value") + else packed_w + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + + if const_expr(_need_sort): + col_g0_i32 = arith.index_cast(T.i32, col_g0) + is_scale_writer = arith.cmpi( + CmpIPredicate.eq, col_g0_i32 & _c31_i32, _c0_i32 + ) + _if_scale = scf.IfOp(is_scale_writer) + with ir.InsertionPoint(_if_scale.then_block): + row_i32_s = arith.index_cast(T.i32, row) + col_s_i32 = col_g0_i32 >> _c5_i32 + d0 = row_i32_s >> _c5_i32 + d1 = (row_i32_s >> _c4_i32) & _c1_i32 + d2 = row_i32_s & _c15_i32 + d3 = col_s_i32 >> _c3_i32 + d4 = (col_s_i32 >> _c2_i32) & _c1_i32 + d5 = col_s_i32 & _c3_i32 + byte_off = ( + d0 * _n32_sort + + d3 * _c256_i32 + + d5 * _c64_i32 + + d2 * _c4_i32 + + d4 * _c2_i32 + + d1 + ) + e8m0_i8 = arith.TruncIOp(T.i8, e8m0_biased) + buffer_ops.buffer_store( + e8m0_i8, + sorted_scale_rsrc, + byte_off, + offset_is_bytes=True, + ) + scf.YieldOp([]) + elif const_expr(_is_splitk): + col_idx = col_g0 + arith.constant(_sk_n_offset[0], index=True) + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec_sk * out_elem_bytes, + ) + else: + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + + _frag_elem = ( + ir.F32Type.get() + if _need_quant + else (ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get()) + ) + + if const_expr(gate_up_interleave and not _is_splitk): + # gui without splitk: acc has activation applied, halved N + _gui_eff_n = _gui_out_n + _gui_tile_n = tile_n // 2 + _gui_cshuffle_nlane = min(32, _gui_tile_n // _e_vec) + _gui_by_n = by_n / arith.constant(2, index=True) + _gui_n_tile_base = n_tile_base / arith.constant(2, index=True) + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=_gui_tile_n, + e_vec=_e_vec, + cshuffle_nlane=_gui_cshuffle_nlane, + block_size=total_threads, + m_repeat=m_repeat, + num_acc_n=_gui_eff_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=_gui_by_n, + n_tile_base=_gui_n_tile_base, + lds_out=lds_out, + frag_elem_type=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + elif const_expr(mock_gate_only or (gate_up_interleave and _is_splitk)): + # mock_gate_only: single pass, by_n covers full [0, 2*inter_dim) + _eff_e_vec = _e_vec_sk + acc = acc_gate + 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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + elif const_expr(_is_splitk): + # Two-pass epilogue: gate then up, each with atomic add + _eff_e_vec = _e_vec_sk + + # Pass 1: gate + acc = acc_gate + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + gpu.barrier() + + # Pass 2: up + acc = acc_up + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + else: + 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, + cshuffle_nlane=_cshuffle_nlane, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + _if_blk = scf.IfOp(blk_valid) + with ir.InsertionPoint(_if_blk.then_block): + _ifexpert_of = scf.IfOp(exp_valid) + with ir.InsertionPoint(_ifexpert_of.then_block): + _moe_gemm1_body() + scf.YieldOp([]) + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher -- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage1, + act, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + use_async_copy, + waves_per_eu, + k_batch, + gate_mode, + a_scale_one, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_max_token_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + inter_dim_pad_total = arith.constant(2 * inter_dim_pad, index=True) + tile2_pad = 0 + if const_expr(not gate_only): + tile_k_stage2 = tile_k // 2 + tile2_pad = ( + tile_k_stage2 - (inter_dim - inter_dim_pad) % tile_k_stage2 + ) % tile_k_stage2 + + inter_in = arith.index_cast(ir.IndexType.get(), i32_inter_in.ir_value()) + tile_n_index = arith.constant(tile_n, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + gx = ( + inter_in - inter_dim_pad_total + tile2_pad + tile_n_index - 1 + ) / tile_n_index + else: + gx = ( + (inter_in - inter_dim_pad_total + tile2_pad + 2 * tile_n_index - 1) + / tile_n_index + / arith.constant(2, index=True) + ) + + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + arg_out_scale_sorted, + i32_tokens_in, + i32_inter_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch(grid=(gx, gy, k_batch), block=(total_threads, 1, 1), stream=stream) + + return launch_mixed_moe_gemm1 + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm2( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + # Optional experiment: write per-(token,slot) output (no atomics) into an output shaped + # [tokens*topk, model_dim] (or [tokens, topk, model_dim] flattened), then reduce over topk outside. + # This can reduce atomic contention for small tokens at the cost of extra bandwidth / reduction. + accumulate: bool = True, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 4, + sort_block_m: int = 0, + b_nt: int = 2, + xcd_swizzle: int = 0, +): + """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + + persist_m: + - > 0: legacy mode -- each CTA processes exactly persist_m consecutive M tiles. + - <= 0: **persistent mode** -- grid_y = cu_num (auto-detected), each CTA + round-robins over M tiles with stride cu_num. + + a_dtype: + - "fp8": A2 is fp8 + - "fp16": A2 is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": A2 is int8 + - "fp4": A2 is fp4 + + b_dtype: + - "fp8": W is fp8 + - "fp16": W is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": W is int8 + - "int4": W4A8 path: A2 is int8, W is packed int4 (2 values per byte) unpacked to int8 in-kernel + - "fp4": W is fp4 + + 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). + + `sort_block_m` is the block_size used by moe_sorting / stage1. When 0 (default), + assumed equal to `tile_m`. When set, stage2 can use a different tile_m from + sorting/stage1. Requires sort_block_m % tile_m == 0. + """ + _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m + if _sort_block_m != tile_m and _sort_block_m % tile_m != 0: + raise ValueError( + f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" + ) + + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + _scale_pack_m = 2 # physical mn_pack in preshuffle microscale layout + _scale_pack_n = 2 + _scale_pack_k = 2 # physical k_pack in preshuffle scale layout + pack_M = min(_scale_pack_m, tile_m // 16) + pack_N = min(_scale_pack_n, tile_n // 64) + _k_unroll_raw = (int(tile_k) * (2 if a_dtype == "fp16" else 1)) // 128 + pack_K = min(_scale_pack_k, _k_unroll_raw) + + elem_bytes = 1 + + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + # ---- Static B preshuffle strides (compile-time) ---- + # All values below are Python ints computable at kernel-compile time. + # Using them in an explicit multiply-add replaces the fly dialect's + # dynamic ``crd2idx`` path which emits Barrett reduction for the + # non-power-of-2 ``n0 = experts*model_dim//16`` shape. + _b_kpack_bytes_s = 8 if (b_dtype == "int4") else 16 + _b_kpack_elems_s = _b_kpack_bytes_s // b_elem_bytes + _b_c_k_s = inter_dim // _scale_pack_k + _b_c_k0_s = (_b_c_k_s * b_elem_bytes) // 64 + _b_stride_nlane = _b_kpack_elems_s # 16 + _b_stride_klane = 16 * _b_stride_nlane # 256 + _b_stride_k0 = 4 * _b_stride_klane # 1024 + _b_stride_n0 = _b_c_k0_s * _b_stride_k0 # c_k0 * 1024 + assert model_dim % 16 == 0, "model_dim must be divisible by 16" + _expert_b_stride = (model_dim // 16) * _b_stride_n0 + + # K64-byte micro-step: always 64 bytes per `ku`. For fp16, this is 32 elements (2xK16 MFMA). + if (tile_k_bytes % 64) != 0: + raise ValueError( + f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " + f"(tile_k={tile_k}, elem_bytes={a_elem_bytes})" + ) + + out_s = str(out_dtype).strip().lower() + if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): + 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 = b_dtype == "int4" + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * model_dim * inter_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * model_dim * 4 + # INT4 here means W4A8: A2 is int8, W is packed int4 and unpacked to int8 in-kernel. + is_int8 = False + + 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`)." + ) + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def _scale_elem_type(): + return T.i32 + + total_threads = 256 + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + "tile_m*tile_k*elem_bytes must be divisible by " + f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={a_elem_bytes}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if a_elem_vec_pack > 1: + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + else: + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + + if out_is_f32: + # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. + _use_cshuffle_epilog = ( + 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("FLIR_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 (FLIR_MOE_STAGE2_CSHUFFLE=1)." + ) + + # NOTE: Keep this as a callable so we don't require an MLIR Context at Python-time. + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + 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 the compiler 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. + _persistent = persist_m <= 0 + if _persistent: + from aiter.jit.utils.chip_info import get_cu_num + + _cu_num = get_cu_num() + else: + _cu_num = 0 + _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" + _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" + f"_t{tile_m}x{tile_n}x{tile_k}" + f"_vscale_fix3{_pm_tag}{_sbm_tag}{_xcd_tag}" + ).replace("-", "_") + # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- + # Ping-pong A2 tiles via separate allocators (like stage1). + _single_x_bytes = int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + _cshuffle_elem_bytes_s2 = 2 # f16/bf16 = 2 bytes + lds_out_bytes = ( + _cshuffle_elem_bytes_s2 * int(tile_m) * int(tile_n) + if _use_cshuffle_epilog + else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec4_i32 = T.vec(4, i32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec8_elems = 8 if a_elem_bytes == 1 else 4 + vec4_elems = 4 if a_elem_bytes == 1 else 2 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = ( + arith.constant_vector(0, vec4_i32) + if is_int8 + else arith.constant_vector(0.0, vec4_f32) + ) + + # A2 layout (flatten token-slot -> M; use i32 for fly.make_shape). + topk_idx = arith.constant(topk, index=True) + m_in = tokens_in * topk_idx + + # B preshuffle layout: [experts*model_dim, inter_dim] + c_n_total = arith.constant(experts * model_dim, index=True) + kpack_bytes = 8 if is_int4 else 16 + # (inlined: _div_pow2, _mod_pow2 are module-global) + + def check_c_n_valid_gate(base_n): + return arith.cmpi(CmpIPredicate.ult, base_n, model_dim - model_dim_pad) + + def check_c_k_valid_gate(base_k): + return arith.cmpi(CmpIPredicate.ult, base_k, inter_dim - inter_dim_pad) + + # A&B's scale preshuffle layout + # For fp4, k_in is already packed (inter_dim // a_elem_vec_pack), so we need original inter_dim + c_k_orig = arith.constant(inter_dim, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=m_in, c_k=c_k_orig + ) + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=c_k_orig + ) + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along model_dim (N-dim) + bx_persist = gpu.block_id("y") # persistent WG index (M-dim) + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_mdp_sw = arith.constant(model_dim_pad, index=True) + _gx = (n_in - _c_mdp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + if const_expr(_persistent): + _gy = arith.constant(_cu_num, index=True) + else: + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S = xcd_swizzle + _c_wgm = arith.constant(_WGM_S, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + (T.bf16 if out_is_bf16 else T.f16), + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # 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 = arith.constant(topk, index=True) + + # X(A2): buffer size in bytes, accounting for FP4 packing (2 elements per byte). + # fp8/int8: 1 byte per element -> bytes = tokens*topk * K + # fp4: 2 elements per byte -> bytes = tokens*topk * K / 2 + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + x_nbytes_idx = _div_pow2( + (tokens_in * c_topk) * k_in * c_elem_bytes, int(a_elem_vec_pack) + ) + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # 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 * arith.constant(out_elem_bytes, index=True) + ) + if const_expr(not bool(accumulate)): + out_nbytes_idx = ( + tokens_in + * arith.index(topk) + * n_in + * arith.constant(out_elem_bytes, index=True) + ) + out_nbytes_i32 = arith.index_cast(T.i32, out_nbytes_idx) + out_rsrc = _ptr_buffer_resource(arg_out, out_nbytes_i32) + + # num_valid_ids (sorted padded MN) for scale sizing / guards. + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + # num_valid_ids is a scalar (same value for all lanes) loaded into + # VGPR. Promote to SGPR so downstream buffer resource descriptors + # that use it for num_records stay in SGPRs, eliminating the + # expensive waterfall loop the compiler would otherwise emit. + num_valid_i32 = rocdl.ReadfirstlaneOp(T.i32, num_valid_i32).res + num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + + # fp16 path ignores scales completely (implicit scale=1.0). + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not is_f16_a): + if const_expr(is_f4_a or is_f8_a): + # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. + # Caller must pre-scatter a2_scale via moe_mxfp4_sort. + kblk = _div_pow2(k_in, 32) + sx_nbytes_idx = num_valid_idx * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + else: + # scale_x (A2 scale): [tokens*topk] f32 -> bytes = tokens*topk*4 + sx_nbytes_idx = (tokens_in * c_topk) * arith.constant(4, index=True) + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + # Weight microscale buffer (packed i32 holding e8m0 bytes). + # Use an exact descriptor size so hardware OOB checking works. + kblk_w = _div_pow2(k_in, 32) # K/32 + mn_w = arith.constant(experts * model_dim, index=True) + sw_nbytes_idx = mn_w * kblk_w # bytes (e8m0) + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + # sorted_token_ids / sorted_weights: [blocks*tile_m] (padded length) + sorted_nbytes_idx = ( + size_expert_ids_in + * arith.constant(tile_m, index=True) + * arith.constant(4, index=True) + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + # expert ids: [sort_blocks] i32. + _c_sbm = arith.constant(_sort_block_m, index=True) + _c_tm = arith.constant(tile_m, index=True) + _c1 = arith.constant(1, index=True) + _sort_blocks_ub = _div_pow2( + size_expert_ids_in * _c_tm + _c_sbm - _c1, _sort_block_m + ) + eid_nbytes_idx = _sort_blocks_ub * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # ---- persist loop ---- + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + + if const_expr(_persistent): + # Expert-phase scheduling: contiguous M-tile dispatch. + # grid_y = cu_num, each CTA handles a contiguous chunk of M-tiles: + # [bx_persist * tiles_per_block, ..., (bx_persist+1) * tiles_per_block - 1] + # Adjacent blocks process adjacent M-tiles -> same expert -> B weight L2 reuse. + _c_cu = arith.constant(_cu_num, index=True) + _c_tm_p = arith.constant(tile_m, index=True) + _num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + _total_m_tiles = (_num_valid_idx + _c_tm_p - _c1_p) / _c_tm_p + _tiles_per_block = (_total_m_tiles + _c_cu - _c1_p) / _c_cu + _i1 = ir.IntegerType.get_signless(1) + _init_active = arith.constant(1, type=_i1) + _for_persist = scf.ForOp(_c0_p, _tiles_per_block, _c1_p, [_init_active]) + else: + # Legacy mode: fixed persist_m consecutive tiles. + _c_pm = arith.constant(persist_m, index=True) + _init_prev_expert = arith.constant(0, type=T.i32) + _init_prev_b_base = arith.constant(0, index=True) + _for_persist = scf.ForOp( + _c0_p, + _c_pm, + _c1_p, + [_init_prev_expert, _init_prev_b_base], + ) + + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + + if const_expr(_persistent): + _still_active = _for_persist.inner_iter_args[0] + bx = bx_persist * _tiles_per_block + _mi_p + else: + _prev_expert_i32 = _for_persist.inner_iter_args[0] + _prev_expert_b_base = _for_persist.inner_iter_args[1] + bx = bx_persist * arith.constant(persist_m, index=True) + _mi_p + + bx_m = bx * arith.constant(tile_m, index=True) + + # Early-exit guard: skip garbage expert blocks beyond `num_valid_ids`. + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + + sort_blk = _div_pow2(bx_m, _sort_block_m) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, sort_blk, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + if const_expr(_persistent): + # Absolute B-base: no cross-iteration state needed. + _expert_b_base = expert_idx * arith.constant( + _expert_b_stride, index=True + ) + else: + # Legacy incremental B-base: delta = (cur - prev) * stride + _delta_expert = arith.subi(expert_i32, _prev_expert_i32) + _delta_expert_idx = arith.index_cast(ir.IndexType.get(), _delta_expert) + _delta_b = _delta_expert_idx * arith.constant( + _expert_b_stride, index=True + ) + _expert_b_base = _prev_expert_b_base + _delta_b + + # Early-exit: if the first row of this tile is a sentinel (all-padding tile), + # skip the entire GEMM. + _first_tok = buffer_ops.buffer_load( + sorted_rsrc, bx_m, vec_width=1, dtype=T.i32 + ) + _first_tid = arith.andi(_first_tok, arith.constant(0xFFFFFF, type=T.i32)) + _tokens_i32_guard = arith.index_cast(T.i32, tokens_in) + tile_has_tokens = arith.cmpi( + CmpIPredicate.ult, _first_tid, _tokens_i32_guard + ) + + # For tile_m < 32 (pack_M < _scale_pack_m): shift a_scale i32 so the + # correct bytes land at the op_sel positions we use. + if const_expr(pack_M < _scale_pack_m): + _m_off = _mod_pow2(_div_pow2(bx_m, 16), _scale_pack_m) + _m_scale_shift_i32 = arith.index_cast( + T.i32, _m_off * arith.constant(8, index=True) + ) + else: + _m_scale_shift_i32 = None + + def _moe_gemm2_then_body(): + # Expert id for this M tile. + n_idx = arith.constant(model_dim, index=True) + 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 we require 16B. + if const_expr(is_f16_a): + if const_expr(bytes_per_thread_x % 16 != 0): + raise ValueError( + f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16" + ) + x_load_bytes = 16 + else: + if 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) + vec4_i32 = T.vec(4, i32) + + c_k_div4 = _div_pow2( + _div_pow2(k_in, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(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, + ) + + vec1_i32 = T.vec(1, i32) + vec2_i32 = T.vec(2, i32) + x_load_vec_elems = ( + x_load_bytes if a_elem_bytes == 1 else x_load_bytes // a_elem_bytes + ) + + 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. + """ + if const_expr(x_load_bytes == 16): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + # 8B/4B: convert dword index to byte offset and use offset_in_bytes path. + idx_bytes = idx_i32 * arith.index(4) + return _buffer_load_vec( + buffer_ops, + vector, + x_rsrc, + idx_bytes, + elem_type=x_elem, + vec_elems=x_load_vec_elems, + elem_bytes=a_elem_bytes, + offset_in_bytes=True, + ) + + # 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(ir.IndexType.get(), row_ts_i32) + + x_row_base_div4.append(row_ts_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 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(vec4_i32, x_vec)) + elif const_expr(x_load_bytes == 8): + parts.append(vector.bitcast(vec2_i32, x_vec)) + else: + parts.append(vector.bitcast(vec1_i32, x_vec)) + return parts + + # tx -> wave/lane (GEMM-style decomposition). + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + + row_a_lds = lane_mod_16 + + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + # Dynamic N tiling within block. + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_mod_4 = _mod_pow2(wave_id, 4) + n_tile_base = wave_mod_4 * c_n_per_wave + + by_n = by * arith.constant(tile_n, index=True) + + if const_expr(pack_N < _scale_pack_n): + _global_n_base = expert_off_idx + by_n + n_tile_base + _n_off = _mod_pow2(_div_pow2(_global_n_base, 16), _scale_pack_n) + _n_scale_shift_i32 = arith.index_cast( + T.i32, _n_off * arith.constant(8, index=True) + ) + else: + _n_scale_shift_i32 = None + n_intra_list = [None] * num_acc_n + n_blk_list = [None] * num_acc_n + col_g_list = [None] * num_acc_n + for i in range_constexpr(num_acc_n): + offset = i * 16 + col_g = by_n + n_tile_base + col_g = _div_pow2(col_g, 2) + offset + col_g = col_g + lane_mod_16 + col_g_list[i] = col_g + c_offset = arith.constant(offset, index=True) + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + n_blk_list[i] = _div_pow2(global_n, 16) + n_intra_list[i] = _mod_pow2(global_n, 16) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 # K64-byte micro-step (2x MFMA) + + # fp4 pack + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku_s2 = tile_k // k_unroll + _pad_k_elems_s2 = (inter_dim_pad % tile_k) if inter_dim_pad > 0 else 0 + _pad_ku_skip_s2 = _pad_k_elems_s2 // _K_per_ku_s2 + _tail_ku_s2 = k_unroll - _pad_ku_skip_s2 + _tail_ku_packed_s2 = ( + (_tail_ku_s2 + pack_K - 1) // pack_K + if _pad_ku_skip_s2 > 0 + else None + ) + + # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- + def load_b_packs_k64(base_k, ku: int, ni: int): + """Load one K64-byte B micro-step: single 16B load, split into 2x i64.""" + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0_base = _div_pow2(base_k_bytes, 64) + k0 = k0_base + arith.constant(ku, index=True) + k1 = lane_div_16 + # Incremental B addressing: _expert_b_base carries the + # expert's preshuffle offset (updated via delta each + # persist_m iteration); local n_blk/n_intra contribute + # the per-lane within-tile offset. All strides are + # compile-time constants -> shift/mul, no Barrett. + idx_pack = ( + _expert_b_base + + n_blk_list[ni] * arith.constant(_b_stride_n0, index=True) + + k0 * arith.constant(_b_stride_k0, index=True) + + k1 * arith.constant(_b_stride_klane, index=True) + + n_intra_list[ni] * arith.constant(_b_stride_nlane, index=True) + ) + + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + b_tile = [] + for ku in range_constexpr(ku_limit): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + _b_split_enabled = k_unroll >= 2 + _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll + + def load_b_tile_lo(base_k): + """Load first half of B tile (ku < _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_b_tile_hi(base_k): + """Load second half of B tile (ku >= _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku, k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_scale(arg_scale, rsrc, scale_info, ku, mni): + k_lane = lane_div_16 + n_lane = lane_mod_16 + # Direct arith crd2idx: idx = mni*stride_n0 + ku*stride_k0 + k_lane*stride_klane + n_lane + idx_pack = ( + mni * scale_info.stride_n0 + + ku * scale_info.stride_k0 + + k_lane * scale_info.stride_klane + + n_lane + ) + s = buffer_ops.buffer_load(rsrc, idx_pack, vec_width=1, dtype=T.i32) + return vector.from_elements(T.vec(1, T.i32), [s]) + + def _apply_k_shift(scale_vec, k_shift_bits): + if const_expr(k_shift_bits > 0): + val = vector.extract( + scale_vec, static_position=[0], dynamic_position=[] + ) + val = arith.shrui(val, arith.constant(k_shift_bits, type=T.i32)) + return vector.from_elements(T.vec(1, T.i32), [val]) + return scale_vec + + def load_b_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + b_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for ni in range_constexpr(num_acc_n_packed): + scale = load_scale( + arg_scale_w, + sw_rsrc, + layout_b_scale, + ku + base_k, + ni + + _div_pow2( + _div_pow2( + expert_off_idx + by_n + n_tile_base, + _scale_pack_n, + ), + 16, + ), + ) + scale = _apply_k_shift(scale, k_shift_bits) + b_scale_tile.append(scale) + return b_scale_tile + + def load_a_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + a_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for mi in range_constexpr(m_repeat_packed): + scale = load_scale( + arg_scale_x, + sx_rsrc, + layout_a_scale, + ku + base_k, + mi + _div_pow2(_div_pow2(bx_m, _scale_pack_m), 16), + ) + scale = _apply_k_shift(scale, k_shift_bits) + a_scale_tile.append(scale) + return a_scale_tile + + def prefetch_ab_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + return [ + load_a_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + load_b_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + ] + + vec8_x = T.vec(vec8_elems, x_elem) + vec4_x_lds = T.vec(vec4_elems, x_elem) + + # ---- Pipeline helpers: store X tile to LDS (unused in DMA path) ---- + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + 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_buffer, + vec8_ty=vec8_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x2=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + else: # x_load_bytes == 4 + lds_store_4b_xor16( + arith, + vector, + lds_memref=lds_buffer, + vec4_ty=vec4_x_lds, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x1=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def compute_tile( + acc_in, + b_tile_in, + lds_buffer, + a_scale=None, + b_scale=None, + *, + prefetch_epilogue: bool = False, + a0_prefetch=None, + a1_prefetch=None, + b_hi_loader=None, + ku_count=k_unroll, + ): + if const_expr(b_hi_loader is not None): + b_tile_full = [None] * k_unroll + for i in range_constexpr(_b_split_ku): + b_tile_full[i] = b_tile_in[i] + else: + b_tile_full = b_tile_in + acc_list = list(acc_in) + mfma_res_ty = vec4_i32 if is_int8 else vec4_f32 + + epilogue_pf = None + bias = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + bias = [] + for ni in range_constexpr(num_acc_n): + global_n = by_n + n_tile_base + ni * 16 + lane_mod_16 + bias_offset = expert_off_idx + global_n + bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) + tw_pf = None + if const_expr(doweight_stage2): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + row_in_tile_pf = mi_base_pf + row_off_pf + sorted_row_pf = bx_m + row_in_tile_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + # fp4 path -- single k_idx loop [0, k_unroll). + # b_hi load is issued at the very start so all k_unroll + # MFMAs can overlap the VMEM latency. + _pack_K_shift = (pack_K - 1).bit_length() + _pack_K_mask = pack_K - 1 + + if const_expr(b_hi_loader is not None): + _b_hi = b_hi_loader() + for _bhi_i in range_constexpr(len(_b_hi)): + b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] + + for k_idx in range_constexpr(ku_count): + ku128 = k_idx >> _pack_K_shift + ikxdl = k_idx & _pack_K_mask + + b_packs0, b_packs1 = b_tile_full[k_idx] + + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ku128 * m_repeat_packed + mi] + a_scale_val = vector.extract( + a_scale_i32, static_position=[0], dynamic_position=[] + ) + if const_expr(_m_scale_shift_i32 is not None): + a_scale_val = arith.shrui( + a_scale_val, _m_scale_shift_i32 + ) + for ni in range_constexpr(num_acc_n_packed): + b_scale_i32 = b_scale[ku128 * num_acc_n_packed + ni] + b_scale_val = vector.extract( + b_scale_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(_n_scale_shift_i32 is not None): + b_scale_val = arith.shrui( + b_scale_val, _n_scale_shift_i32 + ) + + for imxdl in range_constexpr(pack_M): + col_base0 = col_base + mi_idx = mi * pack_M + imxdl + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr( + (a0_prefetch is not None) + and (k_idx == 0) + and (mi_idx == 0) + ): + a0, a1 = a0_prefetch + elif const_expr( + (a1_prefetch is not None) + and (k_idx == 1) + and (mi_idx == 0) + ): + a0, a1 = a1_prefetch + else: + a0, a1 = lds_load_packs_k64( + curr_row_a_lds, col_base0, lds_buffer + ) + + if const_expr(is_f8_a): + col_base1 = col_base + 64 + a2, a3 = lds_load_packs_k64( + curr_row_a_lds, col_base1, lds_buffer + ) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + else: + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + + b0 = b_packs0[ni_idx] + b1 = b_packs1[ni_idx] + b128 = pack_i64x4_to_i32x8( + b0, b1, c0_i64, c0_i64 + ) + + acc_idx = mi_idx * num_acc_n + ni_idx + acc_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + b128, + acc_list[acc_idx], + cbsz, + blgp, + ikxdl * _scale_pack_m + imxdl, + a_scale_val, + ikxdl * _scale_pack_n + inxdl, + b_scale_val, + ], + ) + ) + + return acc_list, epilogue_pf + + # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- + # ---- Async DMA: GMEM -> LDS (bypasses VGPR, like stage1) ---- + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler(): + rocdl.sched_barrier(0) + + def _k_shift_bits(k_py): + if const_expr(pack_K >= _scale_pack_k): + return 0 + return ((k_py // 128) % _scale_pack_k) * _scale_pack_m * 8 + + def _k_base(k_py): + return k_py // _scale_pack_k // 128 + + # Preload sorted_idx into lds_tid for epilogue precompute_row + # (N-independent; placed before N-tile loop so it's done once per M-tile.) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + gpu.barrier() + + # Prologue -- B-first + async DMA X(0) -> pong. + k0 = arith.index(0) + if const_expr(_b_split_enabled): + b_cur = load_b_tile_lo(k0) + else: + b_cur = load_b_tile(k0) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(0), _k_shift_bits(0) + ) + rocdl.sched_barrier(0) + prefetch_x_to_lds(k0, lds_x_pong) + rocdl.s_waitcnt(0) + gpu.barrier() + + acc = [acc_init] * num_acc_n * m_repeat + + # Cross-tile A0+A1 LDS prefetch from pong buffer. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + _a1_col_base = col_offset_base + 128 // a_elem_vec_pack + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + # 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 + + c2_tile_k = arith.constant(tile_k * 2, index=True) + b_pong = b_cur + k0_pong_bk = k0 + + # Only emit the scf.for when there are actually iterations to run. + # When k_main2_py == 0 the loop body is empty; emitting an scf.for + # would create a region whose internal SSA values cannot be used + # by the post-loop tail code. + def _make_b_hi_loader(base_k): + """Create a b_hi_loader callable for a given base_k.""" + return lambda _bk=base_k: load_b_tile_hi(_bk) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + rocdl.sched_barrier(0) + k_iv = arith.index(k_iv_py) + next_k1 = k_iv + tile_k + next_k1_bk = next_k1 // 2 + # DMA X(next_k1) -> ping (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k1, lds_x_ping) + b_ping_lo = ( + load_b_tile_lo(next_k1_bk) + if _b_split_enabled + else load_b_tile(next_k1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(next_k1), _k_shift_bits(next_k1) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) + if _b_split_enabled + else None + ), + ) + hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Cross-tile prefetch for the ping tile we are about to compute. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 + else None + ) + + next_k2 = k_iv + c2_tile_k + next_k2_py = k_iv_py + tile_k * 2 + next_k2_bk = next_k2 // 2 + # DMA X(next_k2) -> pong (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k2, lds_x_pong) + b_pong = ( + load_b_tile_lo(next_k2_bk) + if _b_split_enabled + else load_b_tile(next_k2_bk) + ) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(next_k2_py), _k_shift_bits(next_k2_py) + ) + + acc, _ = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + b_hi_loader=( + _make_b_hi_loader(next_k1_bk) + if _b_split_enabled + else None + ), + ) + k0_pong_bk = next_k2_bk + hot_loop_scheduler() + gpu.barrier() + + # Cross-tile prefetch for the next pong tile. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + if const_expr(odd_k_tiles): + # Tail: single remaining tile (already in pong buffer). + acc, epilogue_pf = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + prefetch_epilogue=True, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + else: + # Tail: 2 remaining tiles. + k_tail1 = (k_in + tile_k - 1) // tile_k * tile_k - tile_k + k_tail1_py = ( + int(inter_dim) + tile_k - 1 + ) // tile_k * tile_k - tile_k + k_tail1_bk = k_tail1 // 2 + # DMA tail X -> ping + prefetch_x_to_lds(k_tail1, lds_x_ping) + if const_expr(_pad_ku_skip_s2 > 0): + b_ping_lo = load_b_tile(k_tail1_bk, ku_limit=_tail_ku_s2) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), + _k_shift_bits(k_tail1_py), + ku_packed_limit=_tail_ku_packed_s2, + ) + else: + b_ping_lo = ( + load_b_tile_lo(k_tail1_bk) + if _b_split_enabled + else load_b_tile(k_tail1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ) + + # hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Epilogue tile with sw prefetch. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 and (_pad_ku_skip_s2 == 0 or _tail_ku_s2 >= 2) + else None + ) + acc, epilogue_pf = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + prefetch_epilogue=True, + b_hi_loader=( + None + if _pad_ku_skip_s2 > 0 + else ( + _make_b_hi_loader(k_tail1_bk) + if _b_split_enabled + else None + ) + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- + # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. + + sw_pf = None + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + sw_pf, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + + zero_i32 = arith.constant(0) + + def atomic_add_f16x2(val_f16x2, byte_off_i32): + rocdl.raw_ptr_buffer_atomic_fadd( + val_f16x2, + out_rsrc, + byte_off_i32, + zero_i32, + zero_i32, + ) + + # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). + if const_expr(lds_out is None): + raise RuntimeError( + "FLIR_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased." + ) + + # Precompute the output base address (i64 index) for ALL paths. + # Both accumulate=True (global atomic) and accumulate=False (global store) + # need 64-bit addressing to avoid i32 offset overflow when + # tokens * model_dim * elem_bytes > INT32_MAX (~150K tokens for model_dim=7168). + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + 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, + ): + # Match origin/dev_a16w4: rely on sentinel padded rows + hardware OOB behavior. + fused2 = buffer_ops.buffer_load( + sorted_rsrc, row, vec_width=1, dtype=T.i32 + ) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + + t_ok = arith.cmpi(CmpIPredicate.ult, t2, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s2, topk_i32_v) + ts_ok = arith.andi(t_ok, s_ok) + t2_safe = arith.select(ts_ok, t2, arith.constant(0)) + s2_safe = arith.select(ts_ok, s2, arith.constant(0)) + t2_safe * topk_i32_v + s2_safe + + if const_expr(doweight_stage2): + 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=f32 + ) + + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(is_int8): + v = arith.sitofp(f32, v) + if const_expr(enable_bias): + v = v + bias_pf[ni] + + 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, [v_out]) + + vector.store(v1, lds_out, [lds_idx], alignment=2) + + def precompute_row(*, row_local, row): + # Use lds_tid (sorted_idx preloaded to LDS) instead of buffer_load + # to avoid extra VMEM round-trips in the epilogue. + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + if const_expr(accumulate): + row_byte_base = out_base_idx + t_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + else: + row_byte_base = out_base_idx + ts_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + """Convert an index-typed byte address to !llvm.ptr.""" + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(not bool(accumulate)): + # ---- 64-bit global store path (avoids i32 offset overflow) ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + else: + # ---- accumulate=True: 64-bit global atomic path ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec * out_elem_bytes, + ) + + _e_vec = 2 if accumulate else min(tile_n // 32, 8) + 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=( + ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get() + ), + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) + + if const_expr(_persistent): + # Short-circuit: contiguous tiles are monotonically increasing, + # so once bx_m >= num_valid_ids all remaining tiles are invalid. + _cur_active = arith.andi(_still_active, blk_valid) + _do_gemm = arith.andi( + _cur_active, arith.andi(exp_valid, tile_has_tokens) + ) + _if_valid = scf.IfOp(_do_gemm) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([_cur_active]) + else: + _if_valid = scf.IfOp(_all_valid) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([expert_i32, _expert_b_base]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher (flyc.jit + .launch) -------------------------------- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage2, + accumulate, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + _sort_block_m, + _cu_num if _persistent else 0, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + _tile_n_idx = arith.constant(tile_n, index=True) + _model_dim_pad_idx = arith.constant(model_dim_pad, index=True) + gx = ( + n_in - _model_dim_pad_idx + _tile_n_idx - arith.constant(1, index=True) + ) / _tile_n_idx + if const_expr(_persistent): + gy = arith.constant(_cu_num, index=True) + else: + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + 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_mixed_moe_gemm2 + +# =========================================================================== +# Host-side launchers (adapted from aiter/ops/flydsl/moe_kernels.py). +# These pack pointer args and drive the inline compile_mixed_moe_gemm1/2 +# builders above. +# =========================================================================== +_DLPACK_SAFE = (torch.uint8, torch.float16, torch.bfloat16, torch.float32) + + +def _view_safe(t: torch.Tensor) -> torch.Tensor: + """View as uint8 if dtype is not dlpack-safe, otherwise return as-is.""" + return ( + t.view(torch.uint8) + if t is not None and t.numel() > 0 and t.dtype not in _DLPACK_SAFE + else t + ) + + +def _ptr_view_safe(t: torch.Tensor): + """Pass only the device data pointer; shape is carried by explicit args.""" + view = _view_safe(t) + type_name = type(view).__name__ + module_name = type(view).__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, view.data_ptr()) + + +def _s1_args_fp4( + out, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + out_scale_sorted, + token_num, + n_in, + k_in, + size_expert_ids_in, + dev, + bias=None, + stream=None, +): + empty_f32 = torch.empty(0, device=dev, dtype=torch.float32) + _bias = bias if bias is not None else empty_f32 + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(out), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + _ptr_view_safe(out_scale_sorted), + token_num, + n_in, + k_in, + size_expert_ids_in, + stream, + ) + + +def _s2_args_fp4( + target, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + token_num, + n_in, + k_in, + blocks, + dev, + bias=None, + stream=None, +): + _bias = ( + bias.view(-1) + if bias is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(target), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + token_num, + n_in, + k_in, + blocks, + stream, + ) + + +def _run_compiled(exe, args): + """Call the JitFunction with the given args (handles compile caching).""" + try: + exe(*args) + except Exception: + # JitFunction.__call__ leaks ir.Context on compilation failure; clean up + # leaked contexts so subsequent calls do not take a wrong code path. + try: + while ir.Context.current is not None: + ir.Context.current.__exit__(None, None, None) + except Exception: + pass + raise + + +def build_moe_stage1_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + act: str = "silu", + persist_m: int = 1, + use_async_copy: bool = False, + k_batch: int = 1, + waves_per_eu: int = 3, + b_nt: int = 0, + gate_mode: str = "separated", + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + enable_bias: bool = False, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Build (and cache) the inline FlyDSL a4w4 stage1 device kernel.""" + return compile_mixed_moe_gemm1( + 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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + act=act, + persist_m=persist_m, + use_async_copy=use_async_copy, + k_batch=k_batch, + waves_per_eu=waves_per_eu, + b_nt=b_nt, + gate_mode=GateMode(gate_mode), + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + enable_bias=enable_bias, + a_scale_one=a_scale_one, + xcd_swizzle=xcd_swizzle, + swiglu_limit=swiglu_limit, + ) + + +def build_moe_stage2_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + accumulate: bool = True, + persist_m: int = 1, + sort_block_m: int = 0, + b_nt: int = 0, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + xcd_swizzle: int = 0, + enable_bias: bool = False, +): + """Build (and cache) the inline FlyDSL a4w4 stage2 device kernel.""" + return compile_mixed_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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=persist_m, + sort_block_m=sort_block_m, + b_nt=b_nt, + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + xcd_swizzle=xcd_swizzle, + enable_bias=enable_bias, + ) + + +def _moe_stage1( + a, + w1, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + w1_scale, + a1_scale, + sorted_weights=None, + a_dtype="fp8", + act="silu", + swiglu_limit=0.0, +): + """Host runner for the inline a8w4 stage1 (fp8 act / fp4 weight -> bf16, k_batch=1).""" + token_num = a.shape[0] + E = w1.shape[0] + inter_dim = w1.shape[1] // 2 + # fp8 activations store one element per byte; fp4 packs two per byte. + model_dim = a.shape[1] * 2 if a_dtype == "fp4" else a.shape[1] + dev = a.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + out = torch.empty((token_num, topk, inter_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a1_scale.view(-1) if a1_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w1_scale.view(-1) if w1_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + + _sort_block_m = tile_m + _all_blks = sorted_expert_ids.shape[0] + _dense_blks = ( + min(token_num * topk * _sort_block_m, sorted_token_ids.shape[0]) + // _sort_block_m + ) + _grid_y = min(_dense_blks, _all_blks) + + out_scale_sorted_flat = torch.empty(0, dtype=torch.uint8, device=dev) + _n_in = inter_dim * 2 + _k_in = model_dim + + args = _s1_args_fp4( + out.view(-1), + a.view(-1), + w1.view(-1), + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + out_scale_sorted_flat.view(-1), + token_num, + _n_in, + _k_in, + _grid_y, + dev, + bias=torch.empty(0, device=dev), + ) + + exe = build_moe_stage1_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=(sorted_weights is not None), + a_dtype=a_dtype, + b_dtype="fp4", + out_dtype=out_dtype, + act=act, + persist_m=1, + k_batch=1, + waves_per_eu=3, + b_nt=0, + gate_mode="separated", + swiglu_limit=swiglu_limit, + ) + _run_compiled(exe, args) + return out + + +def _moe_stage2( + inter_states, + w2, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + mode, + w2_scale, + a2_scale, + sorted_weights, + a_dtype="fp8", +): + """Host runner for the inline a8w4 stage2 (fp8 act / fp4 weight -> bf16, atomic).""" + token_num = inter_states.shape[0] + E = w2.shape[0] + model_dim = w2.shape[1] + # fp8 activations store one element per byte; fp4 packs two per byte. + inter_dim = inter_states.shape[2] * 2 if a_dtype == "fp4" else inter_states.shape[2] + accumulate = mode != "reduce" + dev = inter_states.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + alloc_fn = torch.zeros if accumulate else torch.empty + out = alloc_fn((token_num, model_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a2_scale.view(-1) if a2_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w2_scale.view(-1) if w2_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(sorted_token_ids.shape, dtype=torch.float32, device=dev) + ) + + m_blocks = min(sorted_expert_ids.shape[0], token_num * topk) + _persist_m = -1 if m_blocks > 256 else 1 + + _n_in = model_dim + _k_in = inter_dim + target = out + if not accumulate: + target = torch.empty( + (token_num * topk * model_dim,), device=out.device, dtype=out.dtype + ) + + args = _s2_args_fp4( + target, + inter_states, + w2, + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + token_num, + _n_in, + _k_in, + m_blocks, + dev, + bias=None, + ) + + exe = build_moe_stage2_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=(sorted_weights is not None), + a_dtype=a_dtype, + b_dtype="fp4", + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=_persist_m, + sort_block_m=0, + b_nt=0, + ) + _run_compiled(exe, args) + + if not accumulate: + torch.sum(target.view(token_num, topk, model_dim), dim=1, out=out) + return out + + +def _quant_act_fp8(x, block_size=32): + """MXFP8 (e4m3) per-1x32 block quant with e8m0 scales (matches the a8w4 kernel). + + Returns (fp8 values [m, n], e8m0 scale [m, n//block_size]). Mirrors + aiter.ops.quant.per_1x32_f8_scale_f8_quant(scale_type=fp8_e8m0). + """ + from aiter import dtypes + from aiter.ops.quant import per_1x32_f8_scale_f8_quant + + y, scale = per_1x32_f8_scale_f8_quant( + x.contiguous(), quant_dtype=dtypes.fp8, scale_type=dtypes.fp8_e8m0 + ) + m = x.shape[0] + return y.view(m, -1), scale.view(m, -1) + + +def flydsl_moe_a8w4( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + block_m: int = 32, + tile_n: int = 256, + tile_k: int = 256, + mode: str = "atomic", +) -> torch.Tensor: + """Run the inline FlyDSL a8w4 MoE stage1+stage2 path (fp8 act / fp4 weight). + + Activations are MXFP8 (e4m3, per-1x32 e8m0 block scales); expert weights are + MXFP4 (e2m1, per-1x32 e8m0). Returns [T, model_dim] bf16. + """ + import aiter + from aiter import QuantType, dtypes + from aiter.fused_moe import moe_sorting + from aiter.ops.shuffle import ( + shuffle_scale_a16w4, + shuffle_weight, + shuffle_weight_a16w4, + ) + from aiter.utility.fp4_utils import e8m0_shuffle, moe_mxfp4_sort + + experts = w1.shape[0] + inter_dim = w1.shape[1] // 2 + model_dim = w1.shape[2] + token = hidden_states.shape[0] + topk = topk_ids.shape[1] + torch_dtype = hidden_states.dtype + out_dtype = "bf16" if torch_dtype == torch.bfloat16 else "f16" + + topk_ids = topk_ids.to(torch.int32).contiguous() + topk_weights = topk_weights.to(torch.float32).contiguous() + + # --- harness-side prep (NOT the kernel): fp4 weight quant + e8m0 scales, + # fp8 ACTIVATION quant + e8m0 scales, weight/scale preshuffle, + # sorted token dispatch --- + q_dtype = dtypes.fp4x2 + torch_quant = aiter.get_torch_quant(QuantType.per_1x32) + + w1_qt, w1_scale = torch_quant(w1.contiguous(), quant_dtype=q_dtype) + w2_qt, w2_scale = torch_quant(w2.contiguous(), quant_dtype=q_dtype) + w1_qt = w1_qt.view(experts, inter_dim * 2, model_dim // 2) + w2_qt = w2_qt.view(experts, model_dim, inter_dim // 2) + # fp8 activations (1 byte/elem) + e8m0 per-1x32 scales. + a1_qt, a1_scale = _quant_act_fp8(hidden_states) + + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, _ = moe_sorting( + topk_ids, topk_weights, experts, model_dim, torch_dtype, block_m + ) + + w1_qt_shuf = shuffle_weight(w1_qt, (16, 16)) + w2_qt_shuf = shuffle_weight_a16w4(w2_qt, 16, False) + w1_scale_shuf = e8m0_shuffle(w1_scale) + w2_scale_shuf = shuffle_scale_a16w4(w2_scale, experts, False) + a1_scale_sort = moe_mxfp4_sort( + a1_scale[:token, :].view(token, 1, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage1 gate/up GEMM + fused gated activation === + stage1_out = _moe_stage1( + a1_qt, + w1_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + w1_scale=w1_scale_shuf, + a1_scale=a1_scale_sort, + sorted_weights=None, + a_dtype="fp8", + act="silu", + ) + torch.cuda.synchronize() + + # re-quantize the stage-1 output to fp8 (the stage2 activation is also fp8). + a2_qt, a2_scale = _quant_act_fp8(stage1_out.view(-1, inter_dim)) + a2_qt = a2_qt.view(token, topk, -1) + a2_scale_sort = moe_mxfp4_sort( + a2_scale[: token * topk, :].view(token, topk, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage2 down GEMM + weighted top-k combine === + out = _moe_stage2( + a2_qt, + w2_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + mode=mode, + w2_scale=w2_scale_shuf, + a2_scale=a2_scale_sort, + sorted_weights=sorted_weights, + a_dtype="fp8", + ) + torch.cuda.synchronize() + return out diff --git a/tasks/torch2flydsl/moe_a8w4_kernel/model.py b/tasks/torch2flydsl/moe_a8w4_kernel/model.py new file mode 100644 index 00000000..4c57baea --- /dev/null +++ b/tasks/torch2flydsl/moe_a8w4_kernel/model.py @@ -0,0 +1,227 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Pure-PyTorch reference for the quantized a8w4 fused MoE. + +The op is a top-k Mixture-of-Experts feed-forward block with MXFP8 +(``float8_e4m3fn``) activations and MXFP4 (``float4_e2m1fn_x2``) expert weights, +both using e8m0 per-1x32 block scales. A softmax router selects ``topk`` experts +per token; stage 1 runs a grouped gate/up GEMM followed by ``silu(gate) * up``; +stage 2 runs the down GEMM and combines the experts with the renormalized router +weights. GEMMs accumulate in fp32 over dequantized operands. + +Activations are quantized to MXFP8 (e4m3, e8m0 block scale) and dequantized so +the matmul sees the fp8-rounded values; weights are quantized to MXFP4. The +stage-1 result is re-quantized to MXFP8 before the down GEMM, and the output is +returned in bf16. The MXFP4/MXFP8 rounding and e8m0 block-scale numerics +implemented here match AMD's reference quantizer bit-for-bit. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +# MXFP4 (e2m1) decode table indexed by the 4-bit code (sign in bit 3). +_MXFP4_VALUES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) +_BLOCK = 32 +# log2(F4E2M1_MAX=6) floored -> dtypeMax = 2**2 used as the e8m0 fp4 divisor. +_FP4_DTYPE_MAX = 4.0 +# log2(F8E4M3_MAX=448) floored -> dtypeMax = 2**8 used as the e8m0 fp8 divisor. +_FP8_DTYPE_MAX = 256.0 + + +def _f32_to_e8m0(x): + """Round positive fp32 magnitudes to biased e8m0 exponents (uint8).""" + u32 = x.contiguous().view(torch.int32) + exponent = ((u32 >> 23) & 0xFF).view(torch.uint32).to(torch.uint8) + nan_case = exponent == 0xFF + round_case = ((u32 & 0x400000) > 0) & ( + ((u32 & 0x200000) > 0) | ((u32 & 0x1FFFFF) > 0) | (exponent > 0) + ) + exponent[round_case] += 1 + exponent[nan_case] = 0xFF + return exponent + + +def _e8m0_to_f32(scale_e8m0_biased): + """Decode biased e8m0 exponents (uint8) back to fp32 power-of-two scales.""" + scale_e8m0_biased = scale_e8m0_biased.view(torch.uint8) + zero_case = scale_e8m0_biased == 0 + nan_case = scale_e8m0_biased == 0xFF + scale_f32 = scale_e8m0_biased.to(torch.int32) << 23 + scale_f32[zero_case] = 0x00400000 + scale_f32[nan_case] = 0x7F800001 + return scale_f32.view(torch.float32) + + +def _f32_to_e2m1_codes(x): + """Round fp32 values to MXFP4 (e2m1) 4-bit codes, saturating out-of-range + magnitudes and handling denormals (adapted from the torchao FP utilities).""" + EBITS, MBITS = 2, 1 + EBITS_F32, MBITS_F32 = 8, 23 + F32_EXP_BIAS = (1 << (EBITS_F32 - 1)) - 1 + exp_bias = (1 << (EBITS - 1)) - 1 + max_int = (1 << (EBITS + MBITS)) - 1 + sign_mask = 1 << (EBITS + MBITS) + magic_adder = (1 << (MBITS_F32 - MBITS - 1)) - 1 + max_normal = 2 ** ((1 << EBITS) - 1 - exp_bias) * ( + ((1 << (MBITS + 1)) - 1) / (2**MBITS) + ) + min_normal = 2 ** (1 - exp_bias) + denorm_exp = (F32_EXP_BIAS - exp_bias) + (MBITS_F32 - MBITS) + 1 + denorm_mask_int = denorm_exp << MBITS_F32 + denorm_mask_float = torch.tensor( + denorm_mask_int, dtype=torch.int32 + ).view(torch.float32) + + x = x.float().view(torch.int32) + sign = x & 0x80000000 + x = x ^ sign + x = x.view(torch.float) + + saturate_mask = x >= max_normal + denormal_mask = torch.logical_and( + torch.logical_not(saturate_mask), x < min_normal + ) + normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask)) + + denormal_x = x + denorm_mask_float + denormal_x = denormal_x.view(torch.int32) + denormal_x -= denorm_mask_int + denormal_x = denormal_x.to(torch.uint8) + + normal_x = x.view(torch.int32) + mant_odd = (normal_x >> (MBITS_F32 - MBITS)) & 1 + val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder + normal_x += val_to_add + normal_x += mant_odd + normal_x = normal_x >> (MBITS_F32 - MBITS) + normal_x = normal_x.to(torch.uint8) + + codes = torch.full_like(x, max_int, dtype=torch.uint8) + codes = torch.where(denormal_mask, denormal_x, codes) + codes = torch.where(normal_mask, normal_x, codes) + + sign_lp = sign >> (MBITS_F32 + EBITS_F32 - MBITS - EBITS) + sign_lp = sign_lp.to(torch.uint8) & sign_mask + return (codes | sign_lp).to(torch.uint8) + + +def _mxfp4_dequant(x): + """MXFP4 per-1x32 e8m0 quantize+dequantize over the last dim, returning the + fp32 values the hardware GEMM sees.""" + shape = x.shape + xb = x.float().reshape(-1, _BLOCK) + max_abs = torch.amax(torch.abs(xb), dim=1) + scale_e8m0 = _f32_to_e8m0(max_abs / _FP4_DTYPE_MAX) + scale_f32 = _e8m0_to_f32(scale_e8m0).view(-1, 1) + codes = _f32_to_e2m1_codes(xb / scale_f32) + table = _MXFP4_VALUES.to(x.device) + deq = table[codes.long()] * scale_f32 + return deq.reshape(shape) + + +def _mxfp8_dequant(x): + """MXFP8 (e4m3) per-1x32 e8m0 quantize+dequantize over the last dim. The + kernel reads the fp8-rounded activations and multiplies by the e8m0 block + scale in the GEMM, so the reference reproduces that rounding here.""" + shape = x.shape + xb = x.float().reshape(-1, _BLOCK) + max_abs = torch.amax(torch.abs(xb), dim=1) + scale_e8m0 = _f32_to_e8m0(max_abs / _FP8_DTYPE_MAX) + scale_f32 = _e8m0_to_f32(scale_e8m0).view(-1, 1) + y_fp8 = (xb / scale_f32).to(torch.float8_e4m3fn).float() + deq = (y_fp8 * scale_f32).reshape(shape) + return deq.to(torch.bfloat16) + + +def _grouped_gemm_stage1(acts, weights, topk_ids): + """Per-expert grouped GEMM: out[b, k] = acts[b] @ weights[topk_ids[b, k]].T.""" + acts = acts.float() + B, D = acts.shape + topk = topk_ids.shape[1] + N = weights.shape[1] + h = acts.view(B, 1, D).repeat(1, topk, 1) + out = torch.zeros(B, topk, N, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = h[mask] @ weights[e].transpose(0, 1) + return out + + +def _grouped_gemm_stage2(acts, weights, topk_ids, topk_weights): + """Per-expert down GEMM with weighted top-k combine to a single output row.""" + acts = acts.float() + B, topk = topk_ids.shape + model_dim = weights.shape[1] + out = torch.zeros(B, topk, model_dim, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = acts[mask] @ weights[e].transpose(0, 1) + out = out * topk_weights.view(B, topk, 1) + return out.sum(1) + + +def route_topk(logits, topk): + """Softmax router + top-k with renormalized weights. Shared by the harness. + + Ties are broken by ascending expert index via a stable descending sort. The + bf16 gate produces many duplicate logits across the large expert count, and a + nondeterministic top-k tie-break would let the reference and the runtime op + select different experts; the stable order keeps both routings identical. + """ + gate = torch.softmax(logits.float(), dim=-1) + order = torch.sort(gate, dim=-1, descending=True, stable=True).indices + ids = order[..., :topk] + weights = torch.gather(gate, -1, ids) + weights = weights / weights.sum(dim=-1, keepdim=True) + return weights.float(), ids.to(torch.int32) + + +class Model(nn.Module): + def __init__(self, model_dim, inter_dim, experts, topk, activation="silu"): + super().__init__() + self.model_dim = model_dim + self.inter_dim = inter_dim + self.experts = experts + self.topk = topk + self.activation = activation + self.gate = nn.Linear(model_dim, experts, bias=False).to(torch.bfloat16) + self.w1 = nn.Parameter( + (torch.randn(experts, 2 * inter_dim, model_dim) / 10).to(torch.bfloat16) + ) + self.w2 = nn.Parameter( + (torch.randn(experts, model_dim, inter_dim) / 10).to(torch.bfloat16) + ) + + def forward(self, hidden_states): + I = self.inter_dim + + logits = self.gate(hidden_states) + topk_weights, topk_ids = route_topk(logits, self.topk) + + a1 = _mxfp8_dequant(hidden_states) + w1 = _mxfp4_dequant(self.w1) + w2 = _mxfp4_dequant(self.w2) + + stage1 = _grouped_gemm_stage1(a1, w1, topk_ids) + gate, up = stage1.split([I, I], dim=-1) + stage1 = (F.silu(gate) * up).to(torch.bfloat16) + + a2 = _mxfp8_dequant(stage1.reshape(-1, I)).reshape( + hidden_states.shape[0], self.topk, I + ) + + out = _grouped_gemm_stage2(a2, w2, topk_ids, topk_weights) + return out.to(torch.float16).to(torch.bfloat16) + + +def get_inputs(): + return [torch.randn(16, 7168, dtype=torch.bfloat16)] + + +def get_init_inputs(): + return [7168, 256, 384, 8] diff --git a/tasks/torch2flydsl/moe_a8w4_kernel/test_kernel_harness.py b/tasks/torch2flydsl/moe_a8w4_kernel/test_kernel_harness.py new file mode 100644 index 00000000..0f747f64 --- /dev/null +++ b/tasks/torch2flydsl/moe_a8w4_kernel/test_kernel_harness.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Correctness and performance harness for the a8w4 (fp8 act / fp4 weight) fused +MoE task. + +The pure-torch reference in ``model.py`` and the FlyDSL kernel share the same +top-k routing (``model.route_topk``) so expert selection is identical. The +correctness gate is the normalized max error ``max|ref - out| / max|ref|``, which +must stay <= ``REL_TOL``; element-wise close% at 1e-2 and 1e-1 is also reported. +The check asserts and exits non-zero on failure. + +Modes: + --correctness compare the kernel against the reference + --full-benchmark time the kernel vs the reference and write a perf report +""" +import argparse +import importlib.util +import json +import math +import os +import sys +from pathlib import Path + +KERNEL_FILE = "kernel.py" +MODEL_FILE = "model.py" + + +def _resolve_kernel_dir(): + here = os.path.dirname(os.path.abspath(__file__)) + if os.path.isfile(os.path.join(here, KERNEL_FILE)): + return here + cwd = os.getcwd() + if os.path.isfile(os.path.join(cwd, KERNEL_FILE)): + return cwd + return here + + +def _load_module(kernel_dir, filename, alias): + entry = os.path.join(kernel_dir, filename) + if not os.path.isfile(entry): + return None + if kernel_dir not in sys.path: + sys.path.insert(0, kernel_dir) + spec = importlib.util.spec_from_file_location(alias, entry) + if spec is None or spec.loader is None: + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +_KERNEL_DIR = _resolve_kernel_dir() + +# Real a8w4 fp8fp4 fused-MoE shapes (q_dtype_a = float8_e4m3fn, +# q_dtype_w = float4_e2m1fn_x2, per_1x32, Silu): +# kimik2_fp8fp4_untuned_fmoe.csv (Kimi-K2): D=7168, I=256, E=384, topk=8 +# dsv4_fp8fp4_untuned_fmoe.csv (DeepSeek-V4/V3.x): D=7168, I=512, E=385, topk=7 +SHAPES = [ + {"name": "kimik2_t16_e384_k8", "tokens": 16, "model_dim": 7168, "inter_dim": 256, "experts": 384, "topk": 8}, + {"name": "dsv4_t32_e385_k7", "tokens": 32, "model_dim": 7168, "inter_dim": 512, "experts": 385, "topk": 7}, +] + +# Tight element-wise gate: normalized max error <= REL_TOL. +REL_TOL = 1e-2 +SEED = 20260401 +BLOCK_M, TILE_N, TILE_K, MODE = 32, 256, 256, "atomic" +# Correctness uses the deterministic "reduce" combine. The "atomic" stage-2 +# combine sums per-expert partials with order-dependent fp32 atomic-adds, so its +# result is nondeterministic run-to-run; "reduce" computes the identical math +# with a deterministic reduction (same numerics, reproducible comparison). +CORRECTNESS_MODE = "reduce" + + +def _build_model(mmod, shape, device="cuda"): + import torch + + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + model = mmod.Model( + model_dim=shape["model_dim"], inter_dim=shape["inter_dim"], + experts=shape["experts"], topk=shape["topk"], + ).to(device).eval() + hidden = torch.randn( + shape["tokens"], shape["model_dim"], dtype=torch.bfloat16, device=device + ) + return model, hidden + + +def _kernel_out(kmod, mmod, model, hidden, topk): + # Recompute the SAME routing the reference used and run the FlyDSL kernel. + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + return kmod.flydsl_moe_a8w4( + hidden, model.w1.detach(), model.w2.detach(), topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=CORRECTNESS_MODE, + ) + + +def run_correctness(verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + failures = [] + for shape in SHAPES: + model, hidden = _build_model(mmod, shape) + with torch.no_grad(): + ref = model(hidden).float() + out = _kernel_out(kmod, mmod, model, hidden, shape["topk"]).float() + torch.cuda.synchronize() + + max_abs = (ref - out).abs().max().item() + ref_scale = ref.abs().max().item() + 1e-9 + rel_err = max_abs / ref_scale + max_rel = ((ref - out).abs() / (ref.abs() + 1e-9)).max().item() + pct1e2 = torch.isclose(ref, out, atol=1e-2, rtol=1e-2).float().mean().item() * 100 + pct1e1 = torch.isclose(ref, out, atol=1e-1, rtol=1e-1).float().mean().item() * 100 + ok = rel_err <= REL_TOL + if verbose: + print( + f" {'PASS' if ok else 'FAIL'}: {shape['name']} " + f"(D{shape['model_dim']}/I{shape['inter_dim']}/E{shape['experts']}/k{shape['topk']}) " + f"norm_max_err={rel_err:.5f} (tol={REL_TOL}) " + f"max_abs={max_abs:.4f} max_rel={max_rel:.3f} " + f"close%@1e-2={pct1e2:.2f} @1e-1={pct1e1:.2f}" + ) + if not ok: + failures.append(shape["name"]) + + status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(SHAPES)})" + print(f"Status: {status}") + print(f"correctness: {'pass' if not failures else 'fail'}") + assert not failures, f"correctness FAILED for: {failures}" + return True + + +def run_benchmark(warmup=10, iters=100, verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + latencies, speedups, report = [], [], [] + print(f"{'Config':<24} {'Ref':>10} {'FlyDSL':>10} {'Speedup':>10}") + print("-" * 60) + for idx, shape in enumerate(SHAPES): + model, hidden = _build_model(mmod, shape) + topk = shape["topk"] + with torch.no_grad(): + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + w1, w2 = model.w1.detach(), model.w2.detach() + + def run_kernel(): + return kmod.flydsl_moe_a8w4( + hidden, w1, w2, topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=MODE, + ) + + run_kernel() + torch.cuda.synchronize() + for _ in range(warmup): + run_kernel() + torch.cuda.synchronize() + ktimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); run_kernel(); e.record(); torch.cuda.synchronize() + ktimes.append(s.elapsed_time(e)) + kernel_ms = sum(ktimes) / len(ktimes) + + rtimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); model(hidden); e.record(); torch.cuda.synchronize() + rtimes.append(s.elapsed_time(e)) + ref_ms = sum(rtimes) / len(rtimes) + + speedup = ref_ms / kernel_ms if kernel_ms > 0 else 1.0 + latencies.append(kernel_ms); speedups.append(speedup) + report.append({ + "test_case_id": f"test_case_{idx}", + "execution_time_ms": kernel_ms, + "shape": [shape["tokens"], shape["model_dim"], shape["inter_dim"]], + "params": {k: shape[k] for k in ("tokens", "model_dim", "inter_dim", "experts", "topk")}, + }) + if verbose: + print(f"{shape['name']:<24} {ref_ms:>8.4f}ms {kernel_ms:>8.4f}ms {speedup:>8.2f}x") + del model, hidden + torch.cuda.empty_cache() + + geomean_latency = math.exp(sum(math.log(x) for x in latencies) / len(latencies)) + geomean_speedup = math.exp(sum(math.log(x) for x in speedups) / len(speedups)) + + build_dir = Path(_KERNEL_DIR) / "build" + build_dir.mkdir(exist_ok=True) + with open(build_dir / "performance_report.json", "w") as f: + json.dump(report, f, indent=2) + + print("-" * 60) + print(f"Geometric mean latency: {geomean_latency:.4f} ms") + print(f"Geometric mean speedup: {geomean_speedup:.2f}x") + return {"geomean_latency_ms": geomean_latency, "geomean_speedup": geomean_speedup} + + +if __name__ == "__main__": + try: + import torch as _t + _arch = _t.cuda.get_device_properties(0).gcnArchName.split(":")[0] + except Exception: + _arch = "" + if _arch != "gfx950": + print(f"SKIPPED: gfx950-only task on arch={_arch or 'unknown'} (FP4/MX scaled-MFMA requires CDNA4/gfx950)") + print("correctness: skip") + sys.exit(0) + parser = argparse.ArgumentParser(description="torch2flydsl moe harness") + parser.add_argument("--correctness", action="store_true") + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--full-benchmark", action="store_true") + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + print("=" * 60) + print("torch2flydsl MoE (a8w4: fp8 act / fp4 weight, quantized reference)") + print("=" * 60) + + if args.correctness: + try: + run_correctness() + except AssertionError as exc: + print(f"ASSERTION: {exc}") + sys.exit(1) + sys.exit(0) + else: + run_benchmark(warmup=args.warmup, iters=args.iterations) diff --git a/tasks/torch2flydsl/moe_kernel/config.yaml b/tasks/torch2flydsl/moe_kernel/config.yaml new file mode 100644 index 00000000..e328badd --- /dev/null +++ b/tasks/torch2flydsl/moe_kernel/config.yaml @@ -0,0 +1,23 @@ +source_file_path: +- kernel.py +target_kernel_functions: +- flydsl_moe_a4w4 +- build_moe_stage1_module +- build_moe_stage2_module +- compile_mixed_moe_gemm1 +- compile_mixed_moe_gemm2 +compile_command: +- python3 -c "import torch; from kernel import build_moe_stage1_module, build_moe_stage2_module; + build_moe_stage1_module(model_dim=7168, inter_dim=256, experts=257, topk=9, tile_m=32, + tile_n=256, tile_k=256, doweight_stage1=False, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16'); + build_moe_stage2_module(model_dim=7168, inter_dim=256, experts=257, topk=9, tile_m=32, + tile_n=256, tile_k=256, doweight_stage2=True, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16'); + print('compile ok')" +correctness_command: +- python3 test_kernel_harness.py --correctness +performance_command: +- python3 test_kernel_harness.py --full-benchmark +task_type: torch2flydsl +supported_archs: +- gfx950 +task_result_template: null diff --git a/tasks/torch2flydsl/moe_kernel/kernel.py b/tasks/torch2flydsl/moe_kernel/kernel.py new file mode 100644 index 00000000..7838e2b8 --- /dev/null +++ b/tasks/torch2flydsl/moe_kernel/kernel.py @@ -0,0 +1,6640 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""FlyDSL a4w4 (MXFP4) two-stage fused MoE kernel. + +Defines the stage-1 (gate/up GEMM with fused silu-gated activation) and stage-2 +(down GEMM with weighted top-k combine) device kernels in FlyDSL. The device +builders ``compile_mixed_moe_gemm1`` / ``compile_mixed_moe_gemm2`` -- exposed via +``build_moe_stage1_module`` / ``build_moe_stage2_module`` and driven by the +``_moe_stage1`` / ``_moe_stage2`` host runners -- are adapted from AITER's +mixed_moe_gemm_2stage path together with its preshuffle pipeline, CShuffle MFMA +epilogue, layout helpers and GateMode enum. + +``flydsl_moe_a4w4`` is the launcher: it takes bf16 weights and a precomputed +routing (so the reference and kernel share identical top-k selection) and returns +a bf16 ``[T, model_dim]`` tensor. Host-side data prep -- MXFP4/e8m0 per_1x32 +quantization, weight/scale pre-shuffle and the sorted token/expert dispatch +(``moe_sorting``) -- uses AITER utilities to shape inputs into the layout the +device kernels consume. +""" +from __future__ import annotations + +import functools +import math as _math +import os +import re +import builtins as _builtins +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Dict, Optional + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, memref +from flydsl._mlir.dialects.arith import CmpIPredicate +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr import arith as _arith +from flydsl.expr.arith import ArithValue +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + + +# =========================================================================== +# Inlined from aiter/ops/flydsl/moe_common.py :: GateMode +# =========================================================================== +class GateMode(str, Enum): + """Gate/Up computation strategy for stage1 GEMM (see AITER moe_common).""" + + SEPARATED = "separated" + MOCK_GATE_ONLY = "mock_gate_only" + GATE_ONLY = "gate_only" + INTERLEAVE = "interleave" + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/layout_utils.py +# ========================================================================= +def _wrap(v): + """Wrap raw ir.Value in ArithValue for operator overloading compatibility.""" + if isinstance(v, ArithValue): + return v + if isinstance(v, ir.Value): + return ArithValue(v) + return v + + +def _is_pow2(n): + """Return True when *n* is a positive power of two.""" + return n > 0 and (n & (n - 1)) == 0 + + +def _div_pow2(val, divisor): + """Unsigned divide index *val* by a **compile-time** power-of-2 *divisor*. + + Emits ``arith.shrui`` (1 VALU cycle) instead of ``arith.divui`` + (10-15 VALU cycles on CDNA). + """ + shift = _math.log2(divisor) + assert shift == int(shift), f"{divisor} is not a power of 2" + return arith.shrui(val, arith.index(int(shift))) + + +def _mod_pow2(val, modulus): + """Unsigned remainder of index *val* by a **compile-time** power-of-2 *modulus*. + + Emits ``arith.andi`` (1 VALU cycle) instead of ``arith.remui``. + """ + return arith.andi(val, arith.index(modulus - 1)) + + +def _parse_dim(tok): + """Parse a single dimension token: '?' -> None, otherwise int.""" + tok = tok.strip() + return None if tok == "?" else int(tok) + + +def _parse_layout(ly): + """Parse '(s0,s1,...):(d0,d1,...)' -> (shapes, strides) as lists (None for '?').""" + ly_str = str(ly.type) if hasattr(ly, "type") else str(ly) + m = re.search(r"\(([^)]+)\):\(([^)]+)\)", ly_str) + if not m: + return None + shapes = [_parse_dim(s) for s in m.group(1).split(",")] + strides = [_parse_dim(s) for s in m.group(2).split(",")] + return shapes, strides + + +def _has_dynamic_strides(strides): + """Check if any stride is dynamic (None).""" + return any(s is None for s in strides) + + +def idx2crd(idx, layout): + """Decompose flat index into a list of coordinate values. + + For static layouts, computes coordinates with plain arith ops. + Power-of-2 strides/shapes use shift/mask instead of div/rem. + For dynamic layouts, falls back to fx.idx2crd + fx.get. + """ + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + result = fx.idx2crd(idx, layout) + ndims = len(parsed[1]) if parsed else 1 + return [_wrap(fx.get(result, i)) for i in range(ndims)] + + if hasattr(idx, "type") and str(idx.type) != "index": + idx = arith.index_cast(T.index, idx) + shapes, strides = parsed + ndims = len(strides) + + ordered = sorted( + [ + (i, s, sz) + for i, s, sz in _builtins.zip(range(ndims), strides, shapes) + if s != 0 + ], + key=lambda x: x[1], + reverse=True, + ) + coords = [None] * ndims + remaining = idx + for i, stride_val, size_val in ordered: + if stride_val == 1: + c = remaining + elif _is_pow2(stride_val): + c = _div_pow2(remaining, stride_val) + else: + c = remaining / arith.index(stride_val) + if size_val is not None: + if _is_pow2(size_val): + c = _mod_pow2(c, size_val) + else: + c = c % arith.index(size_val) + coords[i] = c + for i in range(ndims): + if coords[i] is None: + coords[i] = remaining + return coords + + +def crd2idx(crd, layout): + """Compute flat index from a coordinate tuple/list. + + For static layouts, computes with plain arith ops. + For dynamic layouts, falls back to fx.crd2idx with fx.make_coord. + """ + if not isinstance(crd, (list, tuple)): + crd = [crd] + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + # fly.make_coord requires i32/i64, not index + crd_i32 = [] + for c in crd: + cv = c + if isinstance(cv, ArithValue): + cv = cv.ir_value() if hasattr(cv, "ir_value") else cv + if isinstance(cv, ir.Value) and isinstance(cv.type, ir.IndexType): + cv = arith.index_cast(T.i32, cv) + crd_i32.append(cv) + coord_val = fx.make_coord(*crd_i32) + result = fx.crd2idx(coord_val, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = arith.index_cast(T.index, scalar) + return _wrap(scalar) + + _, strides = parsed + result = None + for coord_v, stride_v in _builtins.zip(crd, strides): + if stride_v == 0: + continue + term = coord_v if stride_v == 1 else coord_v * arith.index(stride_v) + result = term if result is None else result + term + return result if result is not None else arith.index(0) + + +def get(int_tuple, mode): + """Extract element at `mode` from a Python list/tuple.""" + return int_tuple[mode] + + +layout_get = get + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_preshuffle_pipeline.py (crd2idx -> _pre_crd2idx) +# ========================================================================= +def _pre_crd2idx(crd, layout): + """crd2idx returning an index-type scalar (unwraps fly.int_tuple).""" + result = fx.crd2idx(crd, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = _arith.IndexCastOp(T.index, scalar).result + return scalar + + +def swizzle_xor16(row, col, k_blocks16): + """XOR-with-row swizzle on the K dimension at 16B granularity. + + Computes: col XOR ((row & (k_blocks16 - 1)) * 16) + + k_blocks16 is always a power of 2 (tile_k_bytes / 16), so use + bitwise AND instead of remui to save ~10 VALU cycles on CDNA. + """ + from flydsl.expr import arith as _swz_arith + + mask = k_blocks16 - _swz_arith.index(1) + rem = _swz_arith.andi(row, mask) + return col ^ (rem * 16) + + +def lds_row_major_idx(row, col, row_stride, base=None): + """Linearize a 2D LDS coordinate with explicit index arithmetic.""" + idx = row * row_stride + col + return idx if base is None else idx + base + + +def split_row_major_2d(index, minor_extent): + """Split a linear row-major index into (major, minor).""" + return index // minor_extent, index % minor_extent + + +def _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx, + *, + elem_type, + vec_elems, + elem_bytes, + offset_in_bytes, + cache_modifier=0, +): + """Load vec_elems elements via buffer_load dwordx[1,2,4] + bitcast.""" + from flydsl.expr import arith as _ld_arith + + elem_size = int(elem_bytes) + load_bytes = int(vec_elems) * elem_size + vec_width = load_bytes // 4 + + if offset_in_bytes: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(2)) + elif elem_bytes == 2: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(1)) + else: + idx_i32 = idx + + i32_val = buffer_ops.buffer_load( + rsrc, + idx_i32, + vec_width=vec_width, + dtype=T.i32, + cache_modifier=cache_modifier, + ) + if vec_width == 1: + i32_vec = vector.from_elements(T.vec(1, T.i32), [i32_val]) + else: + i32_vec = i32_val + return vector.bitcast(T.vec(int(vec_elems), elem_type), i32_vec) + + +@dataclass(frozen=True) +class PreshuffleScaleLayout: + """Container returned by `make_preshuffle_scale_layout`. + + The scale layout is ``(c_mn1, c_k1, 4, 16) : (stride_n0, stride_k0, stride_klane, 1)``. + Callers compute flat index directly with plain arith:: + + idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + """ + + layout_scale: object + stride_n0: object + stride_k0: object + stride_klane: object + + +def make_preshuffle_scale_layout( + arith, + *, + c_mn: ir.Value, + c_k: ir.Value, + mn_pack: int = 2, + k_pack: int = 2, + elem_bytes: int = 4, + scale_block_size: int = 32, +) -> PreshuffleScaleLayout: + """Build scale layout matching aiter/CK preshuffle for FP4/FP8 microscale. + + Layout shape: ``(c_mn1, c_k1, 4, 16)`` where + ``c_mn1 = c_mn / 16 / mn_pack`` and ``c_k1 = (c_k / scale_block_size) / 4 / k_pack``. + """ + c16 = fx.Index(16) + c4 = fx.Index(4) + c_k_scale = c_k // fx.Index(scale_block_size) + + c_mn1 = (c_mn // c16) // fx.Index(mn_pack) + c_k1 = (c_k_scale // c4) // fx.Index(k_pack) + if elem_bytes != mn_pack * k_pack: + raise ValueError( + f"elem_bytes of scale must be {mn_pack} * {k_pack}, got {elem_bytes!r}" + ) + + stride_klane = c16 + stride_k0 = c4 * stride_klane + stride_n0 = c_k1 * stride_k0 + + c_mn1_i32 = arith.index_cast(T.i32, c_mn1) + c_k1_i32 = arith.index_cast(T.i32, c_k1) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + + layout_scale = fx.make_layout( + (c_mn1_i32, c_k1_i32, 4, 16), + stride=(stride_n0_i32, stride_k0_i32, stride_klane_i32, 1), + ) + + return PreshuffleScaleLayout( + layout_scale=layout_scale, + stride_n0=stride_n0, + stride_k0=stride_k0, + stride_klane=stride_klane, + ) + + +@dataclass(frozen=True) +class PreshuffleBLayout: + """Container returned by `make_preshuffle_b_layout`.""" + + layout_b: object + kpack_bytes: int + + +def make_preshuffle_b_layout( + arith, + *, + c_n: ir.Value, + c_k: ir.Value, + kpack_bytes: int = 16, + elem_bytes: int = 1, + k_major: bool = False, +) -> PreshuffleBLayout: + """Build B layout matching aiter/CK preshuffle for A8 MFMA kernels. + + When *k_major* is True the block-level order is K-major (``k_blk`` outermost), + matching the ``(0,3,1,4,2,5)`` shuffle permutation. The default N-major + order (``k_major=False``) matches the legacy ``(0,1,3,4,2,5)`` permutation. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + + c16 = fx.Index(16) + c_kpack = fx.Index(kpack_bytes) + + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + c_k_bytes = c_k * arith.constant(int(elem_bytes), index=True) + n0 = c_n // c16 + + c_kpack_elems = ( + c_kpack + if elem_bytes == 1 + else (c_kpack // arith.constant(int(elem_bytes), index=True)) + ) + + stride_nlane = c_kpack_elems + + if k_major: + c32 = fx.Index(32) + c2 = fx.Index(2) + c_k0 = c_k_bytes // c32 + klane_dim = 2 + stride_klane = c16 * stride_nlane + stride_n0 = c2 * stride_klane + stride_k0 = n0 * stride_n0 + else: + c64 = fx.Index(64) + c4 = fx.Index(4) + c_k0 = c_k_bytes // c64 + klane_dim = 4 + stride_klane = c16 * stride_nlane + stride_k0 = c4 * stride_klane + stride_n0 = c_k0 * stride_k0 + + kpack_elems_static = kpack_bytes if elem_bytes == 1 else kpack_bytes // elem_bytes + n0_i32 = arith.index_cast(T.i32, n0) + c_k0_i32 = arith.index_cast(T.i32, c_k0) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + stride_nlane_i32 = arith.index_cast(T.i32, stride_nlane) + + stride_b = (stride_n0_i32, stride_k0_i32, stride_klane_i32, stride_nlane_i32, 1) + layout_b = fx.make_layout( + (n0_i32, c_k0_i32, klane_dim, 16, kpack_elems_static), stride_b + ) + return PreshuffleBLayout(layout_b=layout_b, kpack_bytes=kpack_bytes) + + +def _unpack_int4_to_int8_pair(packed32): + """Split packed int4 dword into two int8 dwords (even/odd nibbles). + + 7-op bit manipulation shared by all int4 unpack paths (W4A8, W4A16, W4A_FP8). + """ + c_08 = fx.Int32(0x08080808) + c_0f = fx.Int32(0x0F0F0F0F) + c_1e = fx.Int32(0x1E) + c_4 = fx.Int32(4) + s0 = (packed32 & c_08) * c_1e + even = (packed32 & c_0f) | s0 + t = packed32 >> c_4 + s1 = (t & c_08) * c_1e + odd = (t & c_0f) | s1 + return even, odd + + +def _pack_i32_pair_to_i64(lo, hi, vector): + """Pack two i32 values into one i64 via vector bitcast.""" + v2 = vector.from_elements(T.vec(2, T.i32), [lo, hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def _i8x4_in_i32_to_bf16x4_i64(val_i32, arith, vector, scale_val=None): + """Convert one i32 (4 signed int8 bytes) to 4 bf16 packed as i64. + + Uses shift-based f32->bf16 truncation (lshr 16) instead of arith.truncf + which on gfx942 expands to ~5 VALU per element. The shift is exact for + unscaled int8 values and introduces <0.5 ULP error for scaled values. + """ + vec1_i32_t = T.vec(1, T.i32) + vec2_i32 = T.i32x2 + vec4_i8 = T.i8x4 + vec1_i64 = T.vec(1, T.i64) + + v1 = vector.from_elements(vec1_i32_t, [val_i32]) + i8x4 = vector.bitcast(vec4_i8, v1) + + f32_vals = [] + for i in range(4): + val_i8 = vector.extract(i8x4, static_position=[i], dynamic_position=[]) + v = arith.sitofp(T.f32, val_i8) + if scale_val is not None: + v = v * scale_val + f32_vals.append(v) + + c16 = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bits0 = arith.bitcast(T.i32, f32_vals[0]) + bits1 = arith.bitcast(T.i32, f32_vals[1]) + bits2 = arith.bitcast(T.i32, f32_vals[2]) + bits3 = arith.bitcast(T.i32, f32_vals[3]) + i32_lo = (bits0 >> c16) | (bits1 & c_ffff0000) + i32_hi = (bits2 >> c16) | (bits3 & c_ffff0000) + + v2 = vector.from_elements(vec2_i32, [i32_lo, i32_hi]) + v64 = vector.bitcast(vec1_i64, v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def load_b_raw_w4a16( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ku: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 8, +): + """Phase 1 of W4A16 B load: issue buffer_load_dword, return raw packed i32. + + Same address calculation as the int4 unpack path in load_b_pack_k32 + but using ku-based indexing for 2-phase latency hiding. + """ + if kpack_bytes != 8: + raise ValueError(f"W4A16 requires kpack_bytes=8, got {kpack_bytes!r}") + + c64 = fx.Index(64) + half_bytes = kpack_bytes // 2 + c2_idx = fx.Index(2) + c4_idx = fx.Index(4) + + k0_base = base_k // c64 + + k1_layout_offset = ku * 2 + lane_div_32 = lane_div_16 // c2_idx + total_k1 = fx.Index(k1_layout_offset) + lane_div_32 + k0 = k0_base + (total_k1 // c4_idx) + k1_local = total_k1 % c4_idx + lane_odd = lane_div_16 % c2_idx + k2_base = lane_odd * fx.Index(half_bytes) + + coord_pack = (n_blk, k0, k1_local, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + idx_bytes = idx_pack + k2_base + + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + return packed32 + + +def _int4_to_bf16x4_i64_gfx950( + packed32, nibble_offsets, arith, vector, scale_val=None, defer_scale16=False +): + """Convert 4 int4 nibbles to 4 bf16 packed as i64 using gfx950 instructions. + + Uses v_cvt_off_f32_i4_sdwa with byte_sel to avoid per-nibble shifts. + Even nibbles (0,2,4,6) → SDWA BYTE_0/1/2/3 on original src. + Odd nibbles (1,3,5,7) → SDWA BYTE_0/1/2/3 on (src >> 4). + Only 1 shift total instead of 7. + + When defer_scale16=True, the ×16 correction factor for v_cvt_off_f32_i4 is + omitted and must be applied later (e.g. in the epilogue). This saves VALU + in the hot loop and uses v_cvt_pk_bf16_f32 for proper f32→bf16 conversion. + """ + from flydsl.expr import rocdl + from flydsl._mlir.dialects._arith_ops_gen import MulFOp as _MulFOp + + _uw = _arith._to_raw + _av = _arith.ArithValue + + src_even = packed32 + src_odd = packed32 >> fx.Int32(4) + + f32_vals = [] + for nib in nibble_offsets: + byte_idx = nib // 2 + src = src_odd if (nib % 2) else src_even + v = rocdl.cvt_off_f32_i4(src, byte_sel=byte_idx) + f32_vals.append(v) + + if defer_scale16: + # Skip ×16; multiply by scale_val only if groupwise. + if scale_val is not None: + raw_scale = _uw(scale_val) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Use v_cvt_pk_bf16_f32 for proper f32→bf16 (no bit-shift trick needed). + i32_lo = rocdl.cvt_pk_bf16_f32(f32_vals[0], f32_vals[1]) + i32_hi = rocdl.cvt_pk_bf16_f32(f32_vals[2], f32_vals[3]) + else: + c16 = fx.Float32(16.0) + if scale_val is not None: + effective_scale = scale_val * c16 + else: + effective_scale = c16 + raw_scale = _uw(effective_scale) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Truncate f32→bf16 via bit-shift (exact for scaled int values). + c16_shift = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bf16_vals = [arith.bitcast(T.i32, _av(v)) for v in f32_vals] + i32_lo = (bf16_vals[0] >> c16_shift) | (bf16_vals[1] & c_ffff0000) + i32_hi = (bf16_vals[2] >> c16_shift) | (bf16_vals[3] & c_ffff0000) + + v2 = vector.from_elements(T.vec(2, T.i32), [i32_lo, i32_hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def unpack_b_w4a16( + packed32, arith, vector, scale_val=None, use_gfx950_cvt=False, defer_scale16=False +): + """Phase 2 of W4A16 B load: unpack int4->int8 + convert int8->bf16. + + Takes raw packed32 from load_b_raw_w4a16 and produces (b0, b1) -- + two i64 values each containing 4 bf16 for one MFMA. + + When use_gfx950_cvt=True, uses v_cvt_off_f32_i4 + v_cvt_pk_bf16_f32 + for ~2x fewer VALU instructions. + + When defer_scale16=True (requires use_gfx950_cvt=True), the ×16 + correction for v_cvt_off_f32_i4 is omitted; caller must apply it + in the epilogue. + """ + if use_gfx950_cvt: + b0 = _int4_to_bf16x4_i64_gfx950( + packed32, + [0, 2, 4, 6], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + b1 = _int4_to_bf16x4_i64_gfx950( + packed32, + [1, 3, 5, 7], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + return (b0, b1) + even, odd = _unpack_int4_to_int8_pair(packed32) + b0 = _i8x4_in_i32_to_bf16x4_i64(even, arith, vector, scale_val=scale_val) + b1 = _i8x4_in_i32_to_bf16x4_i64(odd, arith, vector, scale_val=scale_val) + return (b0, b1) + + +def load_b_pack_k32( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ki_step: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 16, + elem_bytes: int = 1, + unpack_int4: bool = False, +) -> ir.Value: + """Load one B pack for one MFMA(x32) micro-step. + + Returns an i64 Value containing 8 bytes consumed by MFMA. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + if unpack_int4 and kpack_bytes != 8: + raise ValueError("unpack_int4 requires kpack_bytes=8 (packed int4 layout)") + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + + c64 = fx.Index(64) + base_k_bytes = base_k * arith.constant(int(elem_bytes), index=True) + k0_base = base_k_bytes // c64 + k0 = k0_base + arith.constant(ki_step // 2, index=True) + k1 = lane_div_16 + half_bytes = kpack_bytes // 2 + k2_base = arith.constant((ki_step % 2) * half_bytes, index=True) + + coord_pack = (n_blk, k0, k1, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + + if unpack_int4: + idx_bytes = idx_pack + k2_base + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + even, odd = _unpack_int4_to_int8_pair(packed32) + return _pack_i32_pair_to_i64(even, odd, vector) + + vec_elems = kpack_bytes // int(elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_pack, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=(elem_bytes == 1), + ) + + b_i32x4 = vector.bitcast(T.i32x4, b16) + + half = ki_step % 2 + if half == 0: + d0 = vector.extract(b_i32x4, static_position=[0], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[1], dynamic_position=[]) + else: + d0 = vector.extract(b_i32x4, static_position=[2], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[3], dynamic_position=[]) + + v2 = vector.from_elements(T.vec(2, T.i32), [d0, d1]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def tile_chunk_coord_i32( + arith, + *, + tx_i32_base: ir.Value, + i: int, + total_threads: int, + layout_tile_div4, + chunk_i32: int = 4, +): + """Map (thread, chunk_id) -> (row_local, col_local_i32) for X/A loads.""" + if chunk_i32 not in (1, 2, 4): + raise ValueError(f"chunk_i32 must be one of (1,2,4), got {chunk_i32!r}") + chunk_off_i32 = arith.constant(i * total_threads * chunk_i32, index=True) + tile_idx_i32 = tx_i32_base + chunk_off_i32 + coord_local = fx.idx2crd(tile_idx_i32, layout_tile_div4) + row_local = fx.get(coord_local, 0) + col_local_i32 = fx.get(coord_local, 1) + return row_local, col_local_i32 + + +def buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + *, + elem_type, + idx_i32: ir.Value, + rsrc, + vec_elems: int = 16, + elem_bytes: int = 1, +): + """Copy 16 bytes from global memory into regs via buffer-load dwordx4 lowering.""" + if int(vec_elems) <= 0: + raise ValueError(f"vec_elems must be > 0, got {vec_elems!r}") + return _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx_i32, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=False, + ) + + +def lds_store_16b_xor16( + arith, + vector, + *, + lds_memref, + vec16_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x4: ir.Value, + elem_bytes: int = 1, +): + """Store one 16B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v16 = vector.bitcast(vec16_ty, vec_part_i32x4) + vector.store(v16, lds_memref, [idx0]) + + +def lds_store_8b_xor16( + arith, + vector, + *, + lds_memref, + vec8_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x2: ir.Value, + elem_bytes: int = 1, +): + """Store one 8B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v8 = vector.bitcast(vec8_ty, vec_part_i32x2) + vector.store(v8, lds_memref, [idx0]) + + +def lds_store_4b_xor16( + arith, + vector, + *, + lds_memref, + vec4_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x1: ir.Value, + elem_bytes: int = 1, +): + """Store one 4B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v4 = vector.bitcast(vec4_ty, vec_part_i32x1) + vector.store(v4, lds_memref, [idx0]) + + +def lds_load_pack_k32( + arith, + vector, + *, + lds_memref, + layout_lds, + k_blocks16: ir.Value, + curr_row_a_lds: ir.Value, + col_base: ir.Value, + half: int, + lds_base: ir.Value, + ck_lds128: bool, + vec16_ty, + vec8_ty, + vec2_i64_ty, + vec1_i64_ty, +): + """Load one i64 A-pack for an MFMA K32 micro-step from LDS.""" + col_base_swz = swizzle_xor16(curr_row_a_lds, col_base, k_blocks16) + if ck_lds128: + coord_a16 = (curr_row_a_lds, col_base_swz) + idx_a16 = _pre_crd2idx(coord_a16, layout_lds) + lds_base + loaded_a16 = vector.load_op(vec16_ty, lds_memref, [idx_a16]) + a_vec128 = vector.bitcast(vec2_i64_ty, loaded_a16) + return vector.extract(a_vec128, static_position=[half], dynamic_position=[]) + else: + col_swizzled = col_base_swz + (half * 8) + coord_a = (curr_row_a_lds, col_swizzled) + idx_a = _pre_crd2idx(coord_a, layout_lds) + lds_base + loaded_a8 = vector.load_op(vec8_ty, lds_memref, [idx_a]) + a_vec64 = vector.bitcast(vec1_i64_ty, loaded_a8) + return vector.extract(a_vec64, static_position=[0], dynamic_position=[]) + + +def xcd_remap_bx_by( + bx, + by, + c_m, + *, + tile_m: int, + tile_n: int, + N: int, + xcd_swizzle: int, + num_xcds: int = 8, +): + """Remap (bx, by) for L2-cache reuse via XCD swizzle. + + No-op when ``xcd_swizzle <= 0``. Otherwise: + 1. Linearize the original (bx, by) grid round-robin across ``num_xcds`` + XCDs so that contiguous workgroup ids stay on the same XCD. + 2. Re-tile that 1-D order with an M-major group of size ``xcd_swizzle``, + folding the tail group when ``gy`` does not divide evenly. + + Designed to be called inside a ``@flyc.kernel`` immediately after:: + + bx = gpu.block_id("x") + by = gpu.block_id("y") + bx, by = xcd_remap_bx_by(bx, by, c_m, tile_m=..., tile_n=..., N=..., + xcd_swizzle=xcd_swizzle) + + ``c_m`` is the dynamic ``fx.Index`` for runtime ``M``; ``tile_m``, + ``tile_n``, ``N`` and ``xcd_swizzle`` are compile-time Python ints. + """ + if xcd_swizzle <= 0: + return bx, by + + _c1 = fx.arith.constant(1, index=True) + _c_tm = fx.arith.constant(tile_m, index=True) + _gx = fx.arith.constant(N // tile_n, index=True) + _gy = (c_m + _c_tm - _c1) / _c_tm + + _linear_id = bx * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = fx.arith.constant(num_xcds, index=True) + _q = _num_wgs / _c_xcds + _r = _num_wgs % _c_xcds + _xcd = _linear_id % _c_xcds + _in_xcd = _linear_id / _c_xcds + _xcd_lt_r = fx.arith.cmpi(CmpIPredicate.ult, _xcd, _r) + _clip = fx.arith.select(_xcd_lt_r, _xcd, _r) + _wgid = _xcd * _q + _clip + _in_xcd + + _c_wgm = fx.arith.constant(xcd_swizzle, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = fx.arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = fx.arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + new_bx = _first_pid_m + (_wgid_in_group % _group_size_m) + new_by = _wgid_in_group / _group_size_m + return new_bx, new_by + + +__all__ = [ + "PreshuffleBLayout", + "PreshuffleScaleLayout", + "buffer_copy_gmem16_dwordx4", + "lds_load_pack_k32", + "lds_row_major_idx", + "lds_store_4b_xor16", + "lds_store_8b_xor16", + "lds_store_16b_xor16", + "make_preshuffle_b_layout", + "make_preshuffle_scale_layout", + "load_b_pack_k32", + "load_b_raw_w4a16", + "unpack_b_w4a16", + "load_b_raw_w4a16_groupwise", + "unpack_b_w4a16_groupwise", + "extract_bf16_scale", + "split_row_major_2d", + "swizzle_xor16", + "tile_chunk_coord_i32", + "xcd_remap_bx_by", +] + + +# --------------------------------------------------------------------------- +# Groupwise scale load helper (shared by W4A16 and W4A8 groupwise paths) +# --------------------------------------------------------------------------- + + +def _load_groupwise_scale( + buffer_ops, + arith, + *, + scale_rsrc, + expert_offset, + n_blk, + n_intra, + k_pos, + num_groups: int, + group_size: int, + n_per_expert: int, + scale_dtype=None, +): + """Load one per-group scale value from the scale buffer. + + Computes the linear index into the scale tensor from expert offset, + N position, and group index derived from ``k_pos``. + + For bf16 scales the tensor uses ``(E, G//2, N, 2)`` layout — two + adjacent groups for the same N position are packed into one dword. + We load the raw i32 dword (no extraction) so it can be carried as + loop state without register copies. Use :func:`extract_bf16_scale` + in the compute phase to obtain the f32 value. + """ + c16 = fx.Index(16) + n_global = n_blk * c16 + n_intra + c_group_size = fx.Index(group_size) + c_npe = fx.Index(n_per_expert) + group_idx = k_pos // c_group_size + if scale_dtype is None: + scale_dtype = T.f32 + + if scale_dtype == T.bf16: + # (E, G//2, N, 2) layout: dword at [e, pair, n] holds bf16 scales + # for groups 2*pair and 2*pair+1. + pair_idx = group_idx >> fx.Index(1) # group_idx // 2 + # Dword index: same flat formula but with G//2 groups + num_pairs = num_groups // 2 + c_npm1 = fx.Index(num_pairs - 1) + dword_base = expert_offset * c_npm1 + n_global + dword_elem = dword_base + pair_idx * c_npe + dword_idx = arith.index_cast(T.i32, dword_elem) + # Return raw i32 dword — extraction deferred to compute phase. + scale_val = buffer_ops.buffer_load( + scale_rsrc, dword_idx, vec_width=1, dtype=T.i32 + ) + else: + # (E, G, N) layout with f32 dtype + c_gm1 = fx.Index(num_groups - 1) + base_scale = expert_offset * c_gm1 + n_global + elem_idx = base_scale + group_idx * c_npe + scale_idx_i32 = arith.index_cast(T.i32, elem_idx) + scale_val = buffer_ops.buffer_load( + scale_rsrc, scale_idx_i32, vec_width=1, dtype=T.f32 + ) + return scale_val + + +def extract_bf16_scale(arith, scale_raw_i32, ku: int): + """Extract f32 scale from raw i32 dword loaded by bf16 groupwise path. + + In the ``(E, G//2, N, 2)`` layout two adjacent groups share one dword. + ``ku`` determines which half: even ku → low bf16, odd ku → high bf16. + """ + if ku % 2 == 0: + # Low bf16: shift left by 16 to place in upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 << fx.Int32(16)) + else: + # High bf16: mask upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 & fx.Int32(0xFFFF0000)) + + +# --------------------------------------------------------------------------- +# W4A16 groupwise load / unpack helpers +# --------------------------------------------------------------------------- + + +def load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k, + ku: int, + n_blk, + n_intra, + lane_div_16, + elem_type, + scale_rsrc, + expert_offset, + num_groups: int, + group_size: int, + n_per_expert: int, + kpack_bytes: int = 8, + scale_dtype=None, +): + """Phase 1 of W4A16 groupwise B load: buffer_loads for weight + scale. + + Reuses :func:`load_b_raw_w4a16` for the weight load, then issues an + additional ``buffer_load_dword`` for the per-group scale. + + Returns ``(packed32, scale_val)``. + """ + packed32 = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk, + n_intra=n_intra, + lane_div_16=lane_div_16, + elem_type=elem_type, + kpack_bytes=kpack_bytes, + ) + k_pos = base_k + fx.Index(ku * 32) + scale_val = _load_groupwise_scale( + buffer_ops, + arith, + scale_rsrc=scale_rsrc, + expert_offset=expert_offset, + n_blk=n_blk, + n_intra=n_intra, + k_pos=k_pos, + num_groups=num_groups, + group_size=group_size, + n_per_expert=n_per_expert, + scale_dtype=scale_dtype, + ) + return (packed32, scale_val) + + +def unpack_b_w4a16_groupwise(packed32, scale_val, arith, vector, use_gfx950_cvt=False): + """Phase 2 of W4A16 groupwise: unpack + scale + convert to bf16.""" + return unpack_b_w4a16( + packed32, arith, vector, scale_val=scale_val, use_gfx950_cvt=use_gfx950_cvt + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_epilogues.py (_if_then -> _epi_if_then) +# ========================================================================= +@contextmanager +def _epi_if_then(if_op, scf): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def default_epilog( + *, + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + body_row: Callable, +): + """Iterate the standard MFMA 16x16 row mapping and call `body_row(...)`. + + The mapping matches the common MFMA fragment layout used across kernels in this repo. + + Args: + arith: flydsl arith ext module. + range_constexpr: compile-time unrolled range helper. + m_repeat: tile_m // 16 (python int). + lane_div_16: index Value (0..3). + bx_m: base row (index Value). For MoE, this is the base sorted-row for the tile. + body_row: callback invoked as: + body_row(mi=, ii=, row_in_tile=, row=) + """ + bx_m_v = bx_m + lane_div_16_mul4 = lane_div_16 * 4 + ii_idx_list = [fx.Index(ii) for ii in range(4)] + + for mi in range_constexpr(m_repeat): + mi_base = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off = lane_div_16_mul4 + ii_idx_list[ii] + row_in_tile = mi_base + row_off + row = bx_m_v + row_in_tile + body_row(mi=mi, ii=ii, row_in_tile=row_in_tile, row=row) + + +def c_shuffle_epilog( + *, + arith, + vector, + gpu, + scf=None, + range_constexpr, + # Tile params + tile_m: int, + tile_n: int, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + m_repeat: int, + num_acc_n: int, + # Thread mapping inputs + tx, + lane_div_16, + lane_mod_16, + bx_m, + by_n, + n_tile_base, + # LDS buffer (f16 view, row-major [tile_m, tile_n] flattened) + lds_out, + # Element type for LDS loads (defaults to f16). Pass bf16 to support bf16 epilogues. + frag_elem_type: ir.Type | None = None, + # Callbacks + write_row_to_lds: Callable, + precompute_row: Callable | None = None, + store_pair: Callable, + # When LDS overflows, split lds_out across two buffers by wave-group. + # Pass the second buffer here; first buffer is `lds_out`. + lds_out_split=None, + # Row offset in lds_out for 8-wave mode (MLIR index value). + # Shifts both write and read LDS indices by lds_row_offset * tile_n elements. + lds_row_offset=None, +): + """LDS CShuffle epilogue skeleton. + + Call pattern: + - `write_row_to_lds(...)` is called once per MFMA row produced by this thread. + It is responsible for writing all ni columns for that row into `lds_out`. + - `store_pair(...)` is called for each (row_local, col_pair0) half2 after shuffle. + + `store_pair` can implement either global stores or atomics. + """ + if int(block_size) <= 0 or (int(block_size) % int(cshuffle_nlane)) != 0: + raise ValueError( + f"block_size ({block_size}) must be divisible by cshuffle_nlane ({cshuffle_nlane})" + ) + cshuffle_mlane = int(block_size) // int(cshuffle_nlane) + if (int(tile_m) % cshuffle_mlane) != 0: + raise ValueError( + f"tile_m must be divisible by CShuffleMLane ({cshuffle_mlane}), got tile_m={tile_m}" + ) + if int(e_vec) <= 0: + raise ValueError(f"e_vec must be positive, got {e_vec}") + if (int(tile_n) % (int(cshuffle_nlane) * int(e_vec))) != 0: + raise ValueError( + f"tile_n must be divisible by (CShuffleNLane*EVec) = {cshuffle_nlane*e_vec}, got tile_n={tile_n}" + ) + + # ===================== Split-LDS mode (early return) ===================== + # When lds_out_split is provided, waves are divided into two groups: + # Group A (waves 0..N/2-1) uses lds_out, columns [0, tile_n/2) + # Group B (waves N/2..N-1) uses lds_out_split, columns [tile_n/2, tile_n) + # Each group writes/reads independently; same barriers synchronise all waves. + if lds_out_split is not None: + if scf is None: + raise ValueError("scf module is required for split-LDS cshuffle") + + _half_n = int(tile_n) // 2 + _half_threads = int(block_size) // 2 + EVec = int(e_vec) + + CShuffleNLane_s = min(int(cshuffle_nlane), _half_n // EVec) + if _half_threads % CShuffleNLane_s != 0: + raise ValueError( + f"half_threads={_half_threads} not divisible by CShuffleNLane_split={CShuffleNLane_s}" + ) + CShuffleMLane_s = _half_threads // CShuffleNLane_s + if int(tile_m) % CShuffleMLane_s != 0: + raise ValueError( + f"tile_m={tile_m} not divisible by CShuffleMLane_split={CShuffleMLane_s}" + ) + m_reps_s = int(tile_m) // CShuffleMLane_s + n_reps_s = _half_n // (CShuffleNLane_s * EVec) + + _half_n_idx = arith.constant(_half_n, index=True) + _half_thr_idx = arith.constant(_half_threads, index=True) + _zero_idx = arith.constant(0, index=True) + + _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + + # -- write phase (all waves, each to its group's LDS buffer) -- + n_tile_base_v = n_tile_base + col_base_local_a = n_tile_base_v + lane_mod_16 + col_base_local_b = col_base_local_a - _half_n_idx + + def _write_row_split(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * _half_n_idx + _if_g = scf.IfOp(_is_group_b, has_else=True) + with ir.InsertionPoint(_if_g.then_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_b, + num_acc_n=num_acc_n, + lds_out=lds_out_split, + ) + scf.YieldOp([]) + with ir.InsertionPoint(_if_g.else_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_a, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + scf.YieldOp([]) + + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row_split, + ) + gpu.barrier() + + # -- read phase (each group reads from its own LDS buffer) -- + tx_local = tx - arith.select(_is_group_b, _half_thr_idx, _zero_idx) + c_nlane_s = arith.constant(CShuffleNLane_s, index=True) + m_lane_s = tx_local / c_nlane_s + n_lane_s = tx_local % c_nlane_s + c_evec = arith.constant(EVec, index=True) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + _precomputed_rows_s = [] + for mr in range_constexpr(m_reps_s): + row_base_m = arith.constant(mr * CShuffleMLane_s, index=True) + row_local = row_base_m + m_lane_s + row = bx_m_v + row_local + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + _precomputed_rows_s.append((row_local, row, row_ctx, row_pred)) + + for mr in range_constexpr(m_reps_s): + row_local, row, row_ctx, row_pred = _precomputed_rows_s[mr] + + def _do_store_row_split(): + row_base_lds = row_local * _half_n_idx + for nr in range_constexpr(n_reps_s): + col_base_nr = arith.constant( + nr * (CShuffleNLane_s * EVec), index=True + ) + col_pair0_local = col_base_nr + (n_lane_s * c_evec) + lds_idx = row_base_lds + col_pair0_local + + _if_ld = scf.IfOp(_is_group_b, [vec_frag], has_else=True) + with ir.InsertionPoint(_if_ld.then_block): + fb = vector.load_op(vec_frag, lds_out_split, [lds_idx]) + scf.YieldOp([fb]) + with ir.InsertionPoint(_if_ld.else_block): + fa = vector.load_op(vec_frag, lds_out, [lds_idx]) + scf.YieldOp([fa]) + frag = _if_ld.results[0] + + col_pair0 = col_pair0_local + arith.select( + _is_group_b, _half_n_idx, _zero_idx + ) + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row_split() + else: + _do_store_row_split() + + return # split path complete + + # ===================== Standard (non-split) path below ===================== + + # ---------------- Step 1: write C tile to LDS (row-major, fp16) ---------------- + tile_n_idx = arith.constant(int(tile_n), index=True) + n_tile_base_v = n_tile_base + col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) + + _lds_row_base_offset = ( + lds_row_offset * tile_n_idx if lds_row_offset is not None else None + ) + + def _write_row(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + + # Ensure all LDS reads finished before the lds write. + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row, + ) + + # Ensure all LDS writes are visible before the shuffle-read. + gpu.barrier() + + # ---------------- Step 2: shuffle mapping + half2 store/atomic ---------------- + CShuffleNLane = int(cshuffle_nlane) + CShuffleMLane = int(cshuffle_mlane) + EVec = int(e_vec) + + m_reps_shuffle = int(tile_m) // CShuffleMLane + n_reps_shuffle = int(tile_n) // (CShuffleNLane * EVec) + + c_nlane = fx.Index(CShuffleNLane) + m_lane = tx // c_nlane + n_lane = tx % c_nlane + c_evec = fx.Index(EVec) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + # Batch-precompute all row contexts (sorted_idx loads) before the store loop. + # This issues all buffer_load instructions upfront so the compiler can pipeline + # them instead of serializing each load with s_waitcnt vmcnt(0). + _precomputed_rows = [] + for mr in range_constexpr(m_reps_shuffle): + row_base_m = arith.constant(mr * CShuffleMLane, index=True) + row_local = row_base_m + m_lane + row = bx_m_v + row_local + + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + + # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)` and `scf` + # is provided, we can skip the entire N-loop for invalid rows (cheaper than per-store checks). + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + + _precomputed_rows.append((row_local, row, row_ctx, row_pred)) + + # Now perform LDS reads and stores using the pre-fetched row contexts. + for mr in range_constexpr(m_reps_shuffle): + row_local, row, row_ctx, row_pred = _precomputed_rows[mr] + + def _do_store_row(): + row_base_lds = row_local * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + for nr in range_constexpr(n_reps_shuffle): + col_base_nr = arith.constant(nr * (CShuffleNLane * EVec), index=True) + col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile + + lds_idx_pair = row_base_lds + col_pair0 + frag = vector.load_op(vec_frag, lds_out, [lds_idx_pair]) + + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row() + else: + _do_store_row() + + +def mfma_epilog( + *, + use_cshuffle: bool, + # Common (always required) + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + # Default epilog (required when use_cshuffle=False) + body_row: Callable | None = None, + # CShuffle epilog (required when use_cshuffle=True) + vector=None, + gpu=None, + scf=None, + tile_m: int | None = None, + tile_n: int | None = None, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + num_acc_n: int | None = None, + tx=None, + lane_mod_16=None, + by_n=None, + n_tile_base=None, + lds_out=None, + write_row_to_lds: Callable | None = None, + precompute_row: Callable | None = None, + store_pair: Callable | None = None, + frag_elem_type: ir.Type | None = None, +): + if not use_cshuffle: + if body_row is None: + raise ValueError("mfma_epilog(use_cshuffle=False) requires `body_row`.") + return default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=body_row, + ) + + return c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=int(tile_m), + tile_n=int(tile_n), + e_vec=int(e_vec), + cshuffle_nlane=int(cshuffle_nlane), + block_size=int(block_size), + m_repeat=m_repeat, + num_acc_n=int(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=frag_elem_type, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +# ========================================================================= +@contextmanager +def _if_then(if_op): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def _barrier(vmcnt=63, lgkmcnt=63): + """Emit s_waitcnt + s_barrier via inline asm. + + Bypasses LLVM SIInsertWaitcnts which would insert a conservative + s_waitcnt vmcnt(0) lgkmcnt(0) before every S_BARRIER MI. + """ + parts = [] + needs_waitcnt = vmcnt < 63 or lgkmcnt < 63 + if needs_waitcnt: + wc = [] + if vmcnt < 63: + wc.append(f"vmcnt({vmcnt})") + if lgkmcnt < 63: + wc.append(f"lgkmcnt({lgkmcnt})") + parts.append("s_waitcnt " + " ".join(wc)) + parts.append("s_barrier") + llvm.InlineAsmOp( + res=None, + operands_=[], + asm_string="\n".join(parts), + constraints="", + has_side_effects=True, + is_align_stack=False, + ) + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm1( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + act: str = "silu", + use_cshuffle_epilog: bool | None = None, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 1, + use_async_copy: bool = False, + waves_per_eu: int = 4, + k_batch: int = 1, + b_nt: int = 0, + gate_mode: GateMode = GateMode.SEPARATED, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Compile stage1 kernel (gate+up with silu/swiglu). + + GEMM: act(X @ W_gate.T, X @ W_up.T) -> [tokens*topk, inter_dim] + Direct store (no atomic). When k_batch>1 (split-K), each CTA + computes a K-slice and atomically adds gate/up partials. + Note: persist_m=1 (no persistence) is optimal for stage1 because K=model_dim + is large, so each CTA is already compute-heavy. persist_m>1 serializes M blocks + that the GPU can process in parallel. + + gate_mode controls the gate/up computation strategy — see GateMode enum. + """ + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + sort_block_m = max(32, tile_m) + num_waves = min(4, tile_n // 32) + total_threads = num_waves * 64 + pack_M = 1 if tile_m < 32 else 2 + n_per_wave = tile_n // num_waves + pack_N = min(2, n_per_wave // 16) + pack_K = 2 + scale_mn_pack = 2 + elem_bytes = 1 + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + if (tile_k_bytes % 64) != 0: + raise ValueError(f"tile_k_bytes must be divisible by 64, got {tile_k_bytes}") + + out_s = str(out_dtype).strip().lower() + out_is_f32 = out_s in ("f32", "fp32", "float") + out_is_bf16 = out_s in ("bf16", "bfloat16") + is_int4 = b_dtype == "int4" + is_int8 = False + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + mock_gate_only = gate_mode is GateMode.MOCK_GATE_ONLY + gate_up_interleave = gate_mode is GateMode.INTERLEAVE + gate_only = gate_mode is GateMode.GATE_ONLY + + # Padding semantics: model_dim and inter_dim INCLUDE padding. + # model_dim = model_dim_true + model_dim_pad (K direction) + # inter_dim = inter_dim_true + inter_dim_pad (N direction) + # Tensor sizes use the padded dimensions (inter_dim, model_dim). + # Padding only affects kernel internal logic and grid computation. + _inter_dim_valid = inter_dim - inter_dim_pad + + # Split-K validation + _is_splitk = k_batch > 1 + if mock_gate_only and not _is_splitk: + raise ValueError("mock_gate_only requires k_batch > 1 (split-K)") + 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}" + + out_dtype = "bf16" + else: + _k_per_batch = model_dim + _k_dim = _k_per_batch + + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + f"tile_m*tile_k*elem_bytes must be divisible by {total_threads}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if use_cshuffle_epilog is None: + _use_cshuffle_epilog = os.environ.get("FLIR_MOE_STAGE1_CSHUFFLE", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + else: + _use_cshuffle_epilog = bool(use_cshuffle_epilog) + + _need_fp4 = out_dtype == "fp4" + _need_fp8 = out_dtype == "fp8" + _need_quant = _need_fp4 or _need_fp8 + _need_sort = _need_quant + + if _need_quant: + _use_cshuffle_epilog = True + + _fp4q_tag = "_fp4q" if _need_fp4 else "" + _fp8q_tag = "_fp8q" if _need_fp8 else "" + _sort_tag = "_sort" if _need_sort else "" + _async_tag = "_async" if use_async_copy else "" + _sk_tag = f"_sk{k_batch}" if _is_splitk else "" + _go_tag = "_go" if mock_gate_only else "" + _gui_tag = "_gui" if gate_up_interleave else "" + _as1_tag = "_as1" if a_scale_one else "" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe1_silu_mul_a{a_dtype}_w{b_dtype}_{out_s}" + f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{_fp4q_tag}{_fp8q_tag}{_sort_tag}{_async_tag}{_sk_tag}{_go_tag}{_gui_tag}{_as1_tag}{_xcd_tag}_v32" + ).replace("-", "_") + + # -- LDS sizing -- + _cshuffle_elem_bytes = 4 if _need_quant else (4 if out_is_f32 else 2) + _single_x_bytes = int(tile_m) * int(lds_stride) * int(a_elem_bytes) + lds_out_bytes = ( + _cshuffle_elem_bytes * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + # Determine whether we need wave-group split for lds_out. + # Standard layout: pong = max(input, lds_out) + tid, ping = input. + # When this overflows, split lds_out into two halves across pong & ping. + _GLOBAL_ALIGN = 1024 + _std_pong = max(_single_x_bytes, lds_out_bytes) + lds_tid_bytes + _std_ping = _single_x_bytes + _std_pong_aligned = allocator_pong._align(_std_pong, 128) + _std_total = allocator_pong._align( + _std_pong_aligned, _GLOBAL_ALIGN + ) + allocator_pong._align(_std_ping, 128) + _lds_limit = {"gfx950": 163840, "gfx942": 65536}.get(gpu_arch, 0) + + _split_lds_out = ( + _lds_limit > 0 + and lds_out_bytes > 0 + and _std_total > _lds_limit + and num_waves >= 2 + ) + + if _split_lds_out: + _half_out_bytes = _cshuffle_elem_bytes * int(tile_m) * (int(tile_n) // 2) + _pong_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + _ping_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + else: + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if waves_per_eu is not None and waves_per_eu >= 1: + _total_cu_lds = 160 * 1024 + _min_lds = _total_cu_lds // (waves_per_eu + 1) + 1 + _pong_sz = allocator_pong._align(allocator_pong.ptr, 128) + _ping_sz = allocator_ping._align(allocator_ping.ptr, 128) + _cur_lds = _pong_sz + _ping_sz + if _cur_lds < _min_lds: + allocator_ping.ptr += _min_lds - _cur_lds + + kpack_bytes = 8 if is_int4 else 16 + out_elem_bytes = 4 if out_is_f32 else 2 + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * (2 * inter_dim) * model_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * (2 * inter_dim) * 4 + + _e_vec_s1 = min(tile_n // 32, 8) + if _need_quant: + _e_vec_s1 = max(2, _e_vec_s1) + _num_threads_per_quant_blk_s1 = 32 // _e_vec_s1 + _shuffle_dists_s1 = [] + _sh_val = 1 + while _sh_val < _num_threads_per_quant_blk_s1: + _shuffle_dists_s1.append(_sh_val) + _sh_val *= 2 + _num_shuffle_steps_s1 = len(_shuffle_dists_s1) + + # ---- Unified pipeline schedule (outside @flyc.kernel) ---- + # Each scheduling phase is a dict: + # mfma: [(k_idx, mi_idx, ikxdl, imxdl, asv_idx), ...] + # a_reads: [(k, mi), ...] # A ds_read subtiles + # b_loads: [('gate'/'up', ku, ni), ...] # B VMEM loads + # has_scale: bool # A/B scale VMEM loads + _pipe_m_repeat = tile_m // 16 + _pipe_k_unroll = tile_k_bytes // 128 + _pipe_k_unroll_packed = _pipe_k_unroll // pack_K + _pipe_m_repeat_packed = _pipe_m_repeat // pack_M + _pipe_num_acc_n = n_per_wave // 16 + + # A ds_read groups: group by mi (same mi, all k values together) + _pipe_a_groups = [] + for _mi in range(_pipe_m_repeat): + _grp = [] + for _k in range(_pipe_k_unroll): + _grp.append((_k, _mi)) + if len(_grp) == 2: + _pipe_a_groups.append(_grp) + _grp = [] + if _grp: + _pipe_a_groups.append(_grp) + + # B VMEM loads: individual gate/up loads + _pipe_b_loads = [] + for ku in range(_pipe_k_unroll): + for ni in range(_pipe_num_acc_n): + _pipe_b_loads.append(("gate", ku, ni)) + if not mock_gate_only and not gate_up_interleave: + _pipe_b_loads.append(("up", ku, ni)) + + # MFMA order: B-major (fix B, cycle all A tiles before next B) + # Each entry: one (k, ni) pair; the compute function loops over all mi. + # This keeps B operands (from VMEM) fixed while cycling A (from LDS, no wait). + _pipe_num_acc_n_packed = _pipe_num_acc_n // pack_N + _pipe_all_mfma = [] + for _ku128 in range(_pipe_k_unroll_packed): + for _ni_packed in range(_pipe_num_acc_n_packed): + for _ikxdl in range(pack_K): + for _inxdl in range(pack_N): + _k_idx = _ku128 * pack_K + _ikxdl + _ni_idx = _ni_packed * pack_N + _inxdl + _pipe_all_mfma.append((_k_idx, _ni_idx, _ikxdl, _inxdl, _ku128)) + + # Group MFMAs per scheduling phase (wider M -> more MFMAs per phase) + _pipe_mfma_per_phase = max(1, len(_pipe_all_mfma) // 4) + _pipe_n_phases = len(_pipe_all_mfma) // _pipe_mfma_per_phase + + # Build unified phase descriptors + _a_groups_per_phase = (len(_pipe_a_groups) + _pipe_n_phases - 1) // _pipe_n_phases + _pipe_phases = [] + _mfma_i = 0 + _a_i = 0 + for _p in range(_pipe_n_phases): + _a_reads = [] + for _ in range(_a_groups_per_phase): + if _a_i < len(_pipe_a_groups): + _a_reads.extend(_pipe_a_groups[_a_i]) + _a_i += 1 + _phase = { + "mfma": _pipe_all_mfma[_mfma_i : _mfma_i + _pipe_mfma_per_phase], + "a_reads": _a_reads, + "b_loads": [], + "has_scale": (_p == 0), + } + _mfma_i += _pipe_mfma_per_phase + _pipe_phases.append(_phase) + + # Distribute B loads evenly across phases 1..n-1 (phase 0 has scales) + _bi = 0 + for _p in range(1, _pipe_n_phases): + _rem_b = len(_pipe_b_loads) - _bi + _rem_p = _pipe_n_phases - _p + _n_b = (_rem_b + _rem_p - 1) // _rem_p if _rem_p > 0 else 0 + for _ in range(_n_b): + if _bi < len(_pipe_b_loads): + _pipe_phases[_p]["b_loads"].append(_pipe_b_loads[_bi]) + _bi += 1 + + # Extract flat lists for kernel access (avoids dict access in AST rewriter) + _pp_mfma = [p["mfma"] for p in _pipe_phases] + _pp_a_reads = [p["a_reads"] for p in _pipe_phases] + _pp_b_loads = [p["b_loads"] for p in _pipe_phases] + _pp_has_scale = [p["has_scale"] for p in _pipe_phases] + + fp4_ratio = 2 if a_dtype == "fp4" else 1 + gui_ratio = 1 if gate_up_interleave else 2 + _vmcnt_before_barrier = tile_m // 32 // fp4_ratio + tile_n // 32 * gui_ratio + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = arith.constant_vector(0.0, vec4_f32) + + # --- Stage1 dimension mapping --- + # X: [tokens, model_dim] -- M = sorted tokens, K = model_dim + # W: [E*2*inter_dim, model_dim] gate portion -- N = inter_dim + # Out: [tokens*topk, inter_dim] + + # B preshuffle layout: [E*2*inter_dim, model_dim] + # Gate rows for expert e: [e*2*inter_dim, e*2*inter_dim + inter_dim) + c_n_total = arith.constant(experts * (2 * inter_dim), index=True) + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in // pack_K, + kpack_bytes=kpack_bytes, + elem_bytes=b_elem_bytes, + # k_major=True, + ) + layout_b = b_layout.layout_b + + # A-scale: [sorted_size, K/32] -- pre-scattered by caller into sorted layout + # Same as stage2: indexed by sorted_row position, not by token_id. + sorted_m = size_expert_ids_in * arith.constant(sort_block_m, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=sorted_m, c_k=arith.constant(model_dim, index=True) + ) + # B-scale: [E*2*inter_dim, K/32] + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=arith.constant(model_dim, index=True) + ) + + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + if const_expr(use_async_copy and a_elem_vec_pack > 1): + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along inter_dim (N) + bx_persist = gpu.block_id("y") # persistent WG index + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S1 = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_idp_sw = arith.constant(2 * inter_dim_pad, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + _gx = (n_in - _c_idp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + else: + _c2_sw = arith.constant(2, index=True) + _gx = ( + (n_in - _c_idp_sw + _c2_sw * _c_tn_sw - _c1_sw) + / _c_tn_sw + / _c2_sw + ) + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S1, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S1 = xcd_swizzle + _c_wgm = arith.constant(_WGM_S1, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + by_n = by * arith.constant(tile_n, index=True) + + k_base_idx = arith.index(0) + if const_expr(_is_splitk): + bz = gpu.block_id("z") # K-batch id + k_base_idx = bz * arith.constant(_k_dim, index=True) + + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((num_waves, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + _lds_out_elem_type = ( + T.f32 if _need_quant else (T.bf16 if out_is_bf16 else T.f16) + ) + if const_expr(_split_lds_out and _use_cshuffle_epilog): + _half_out_elems = int(tile_m) * (int(tile_n) // 2) + lds_out = SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + lds_out_B = SmemPtr( + base_ptr_ping, + lds_ping_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + else: + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_out_B = None + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # Buffer resources + c_a_pack = arith.constant(int(a_elem_vec_pack), index=True) + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + + # X: [tokens, model_dim] + x_nbytes_idx = (tokens_in * k_in * c_elem_bytes) / c_a_pack + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # Out: [tokens*topk, inter_dim] + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not (is_f16_a or a_scale_one)): + # A scale: [sorted_size, model_dim/32] pre-scattered by caller + c32 = arith.constant(32, index=True) + kblk = k_in / c32 + sx_nbytes_idx = sorted_m * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + c32 = arith.constant(32, index=True) + kblk_w = k_in / c32 + mn_w = arith.constant(experts * (2 * inter_dim), index=True) + sw_nbytes_idx = mn_w * kblk_w + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + sorted_nbytes_idx = size_expert_ids_in * arith.constant( + sort_block_m * 4, index=True + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + eid_nbytes_idx = size_expert_ids_in * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # Sorted-scale buffer resource for fused mxfp4 quantization + _sorted_scale_cols = inter_dim // 32 + _sorted_scale_cols_i32 = arith.constant(_sorted_scale_cols, type=T.i32) + sorted_scale_rsrc = None + if const_expr(_need_sort): + _sort_rows_idx = size_expert_ids_in * arith.constant( + sort_block_m, index=True + ) + _sort_padded_rows = ( + (_sort_rows_idx + arith.constant(255, index=True)) + / arith.constant(256, index=True) + * arith.constant(256, index=True) + ) + _sort_padded_cols = arith.constant( + ((_sorted_scale_cols + 7) // 8) * 8, index=True + ) + _sort_scale_nbytes = arith.index_cast( + T.i32, _sort_padded_rows * _sort_padded_cols + ) + sorted_scale_rsrc = _ptr_buffer_resource( + arg_out_scale_sorted, _sort_scale_nbytes + ) + + # ---- persist_m loop (same pattern as stage2) ---- + _PERSIST_M = persist_m + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + _c_pm = arith.constant(_PERSIST_M, index=True) + _for_persist = scf.ForOp(_c0_p, _c_pm, _c1_p) + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + bx = bx_persist * _c_pm + _mi_p + bx_m = bx * arith.constant(sort_block_m, index=True) + + # Block validity + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, bx, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + def _moe_gemm1_body(): + # Gate expert offset: first inter_dim rows of each expert's 2*inter_dim block + expert_off_idx = expert_idx * arith.constant(2 * inter_dim, index=True) + + # X loading -- KEY DIFFERENCE from stage2: X row = token_id only + x_load_bytes = 16 + num_x_loads = bytes_per_thread_x // x_load_bytes + chunk_i32 = x_load_bytes // 4 + + c_k_div4 = ( + (k_in / c_a_pack) * arith.constant(int(a_elem_bytes), index=True) + ) / arith.index(4) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(0xFFFFFF) + 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, + ) + + def load_x(idx_i32): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + + # Decode sorted token ids -- stage1: X row = token_id (not t*topk+s) + x_row_base_div4 = [] + x_col_local_i32 = [] + x_row_local = [] + # Also store token_id and slot_id for output indexing + + 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + + # KEY: X row base uses token_id only (not t*topk+s) + t_idx = arith.index_cast(ir.IndexType.get(), t_safe) + x_row_base_div4.append(t_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(a_elem_bytes), index=True) + ) / arith.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) + parts.append(vector.bitcast(T.vec(4, i32), x_vec)) + return parts + + # Wave/lane decomposition (identical to stage2) + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + row_a_lds = lane_mod_16 + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_n_id = wave_id % arith.constant(num_waves, index=True) + n_tile_base = wave_n_id * c_n_per_wave + + # N-tile precompute for gate AND up weights + gate_n_intra_list = [] + gate_n_blk_list = [] + up_n_intra_list = [] + up_n_blk_list = [] + col_g_list = [] + c_n0_static = experts * (2 * inter_dim) // 16 + layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1)) + inter_idx = arith.constant(inter_dim, index=True) + + for i in range_constexpr(num_acc_n): + offset = i * 16 + c_offset = arith.constant(offset, index=True) + if const_expr(not gate_up_interleave): + col_g = by_n + n_tile_base + c_offset + lane_mod_16 + col_g_list.append(col_g) + + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + # Gate/interleave: rows [expert_off, expert_off + 2*inter_dim) + gate_row_w = expert_off_idx + global_n + gate_coord = idx2crd(gate_row_w, layout_n_blk_intra) + gate_n_blk_list.append(layout_get(gate_coord, 0)) + gate_n_intra_list.append(layout_get(gate_coord, 1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_row_w = gate_row_w + inter_idx + up_coord = idx2crd(up_row_w, layout_n_blk_intra) + up_n_blk_list.append(layout_get(up_coord, 0)) + up_n_intra_list.append(layout_get(up_coord, 1)) + + if const_expr(gate_up_interleave): + _gui_num_acc_n_out = num_acc_n // pack_N + for _gui_i in range_constexpr(_gui_num_acc_n_out): + _gui_offset = _gui_i * 16 + _gui_c_offset = arith.constant(_gui_offset, index=True) + _gui_col_g = ( + (by_n + n_tile_base) // arith.constant(2, index=True) + + _gui_c_offset + + lane_mod_16 + ) + col_g_list.append(_gui_col_g) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku = tile_k // k_unroll + _pad_k_elems = ( + (model_dim_pad % tile_k) + if (not _is_splitk and model_dim_pad > 0) + else 0 + ) + _pad_ku_skip = _pad_k_elems // _K_per_ku + _tail_ku = k_unroll - _pad_ku_skip + _tail_ku_packed = ( + (_tail_ku + pack_K - 1) // pack_K if _pad_ku_skip > 0 else None + ) + + # B load for gate and up separately + def load_b_packs_k64(base_k, ku: int, n_blk, n_intra): + c64 = arith.constant(64, index=True) + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0 = base_k_bytes // c64 + arith.constant(ku, index=True) + k1 = lane_div_16 + coord_pack = (n_blk, k0, k1, n_intra, arith.constant(0, index=True)) + idx_pack = crd2idx(coord_pack, layout_b) + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + """Load B tiles. Returns (gate_b_tile, up_b_tile). + When mock_gate_only or gate_up_interleave, up_b_tile is None.""" + gate_b_tile = [] + up_b_tile = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_limit): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + gb0, gb1 = load_b_packs_k64( + base_k, ku, gate_n_blk_list[ni], gate_n_intra_list[ni] + ) + g_packs0.append(gb0) + g_packs1.append(gb1) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + ub0, ub1 = load_b_packs_k64( + base_k, ku, up_n_blk_list[ni], up_n_intra_list[ni] + ) + u_packs0.append(ub0) + u_packs1.append(ub1) + gate_b_tile.append((g_packs0, g_packs1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_b_tile.append((u_packs0, u_packs1)) + return gate_b_tile, up_b_tile + + # Pre-compute scale base element indices (K-loop invariant). + # idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + # Split into: base_elem = mni * stride_n0 + lane_elem (invariant) + # k_elem = ku * stride_k0 (per-iteration) + _scale_lane_elem = ( + lane_div_16 * layout_b_scale.stride_klane + lane_mod_16 + ) + + _gate_scale_bases = [] + _up_scale_bases = [] + for _ni in range_constexpr(num_acc_n_packed): + _col_base = ( + by_n + + n_tile_base + + arith.constant(_ni * 16 * pack_N, index=True) + ) + _gate_mni = (expert_off_idx + _col_base) // arith.constant( + 32, index=True + ) + _gate_scale_bases.append( + _gate_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + if const_expr(not mock_gate_only and not gate_up_interleave): + _up_mni = ( + expert_off_idx + inter_idx + _col_base + ) // arith.constant(32, index=True) + _up_scale_bases.append( + _up_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + + if const_expr(not a_scale_one): + _a_scale_bases = [] + for _mi in range_constexpr(m_repeat_packed): + _a_mni = _mi + bx_m // scale_mn_pack // 16 + _a_scale_bases.append( + _a_mni * layout_a_scale.stride_n0 + _scale_lane_elem + ) + + _c16_idx = arith.constant(16, index=True) + _c2_idx = arith.constant(2, index=True) + _scale_mask_lo = arith.constant(0xFF, type=T.i32) + + _m_half_idx = arith.constant(0, type=T.i32) + _m_half_i32 = arith.constant(0, type=T.i32) + _scale_shift = arith.constant(0, type=T.i32) + _scale_shift_hi = arith.constant(0, type=T.i32) + _n_half_idx = arith.constant(0, type=T.i32) + _n_half_i32 = arith.constant(0, type=T.i32) + _bscale_shift = arith.constant(0, type=T.i32) + _bscale_shift_hi = arith.constant(0, type=T.i32) + if const_expr(pack_M < scale_mn_pack): + _m_half_idx = (bx_m // _c16_idx) % _c2_idx + _m_half_i32 = arith.index_cast(T.i32, _m_half_idx) + _scale_shift = _m_half_i32 * arith.constant(8, type=T.i32) + _scale_shift_hi = _scale_shift + arith.constant(16, type=T.i32) + + if const_expr(pack_N < scale_mn_pack): + _n_half_idx = (n_tile_base // _c16_idx) % _c2_idx + _n_half_i32 = arith.index_cast(T.i32, _n_half_idx) + _bscale_shift = _n_half_i32 * arith.constant(8, type=T.i32) + _bscale_shift_hi = _bscale_shift + arith.constant(16, type=T.i32) + + def _rearrange_a_scale(raw_i32): + """Rearrange scale bytes for pack_M=1: extract m_half's k0,k1 bytes.""" + if const_expr(pack_M >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _scale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _scale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + def _rearrange_b_scale(raw_i32): + """Rearrange scale bytes for pack_N=1: extract n_half's k0,k1 bytes.""" + if const_expr(pack_N >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _bscale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _bscale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + if const_expr(a_scale_one): + _as1_const = arith.constant(0x7F7F7F7F, type=T.i32) + _as1_vec = vector.from_elements(T.vec(1, T.i32), [_as1_const]) + + def prefetch_ab_scale_tile(base_k, ku_packed_limit=k_unroll_packed): + a_scale_tile = [] + gate_b_scale = [] + up_b_scale = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_packed_limit): + k_off = (ku + base_k) * layout_b_scale.stride_k0 + for mi in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + a_scale_tile.append(_as1_vec) + else: + s = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[mi] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + s = _rearrange_a_scale(s) + a_scale_tile.append( + vector.from_elements(T.vec(1, T.i32), [s]) + ) + for ni in range_constexpr(num_acc_n_packed): + gs = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + gs = _rearrange_b_scale(gs) + gate_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [gs]) + ) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + us = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + us = _rearrange_b_scale(us) + up_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [us]) + ) + return [a_scale_tile, gate_b_scale, up_b_scale] + + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x4=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + if const_expr(use_async_copy): + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(elem_bytes), index=True) + ) / arith.index(4) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def prefetch_full_a_from_lds(lds_buffer, ku_limit=k_unroll): + """Load entire A tile from LDS into registers before compute.""" + a_regs = [] + for k_idx in range_constexpr(ku_limit): + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + for mi_idx in range_constexpr(m_repeat): + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64( + curr_row, col_base + 64, lds_buffer + ) + a_regs.append((a0, a1, a2, a3)) + else: + a_regs.append((a0, a1)) + return a_regs + + # Compute tile: gate + up MFMA interleaved, same A data, different B data. + # Two accumulator sets; after all K tiles, acc = acc_gate + acc_up (f32 add). + def compute_tile( + acc_gate_in, + acc_up_in, + gate_b_tile_in, + up_b_tile_in, + a_tile_regs, + a_scale=None, + gate_b_scale=None, + up_b_scale=None, + *, + prefetch_epilogue=False, + ku_count=k_unroll, + ): + gate_list = list(acc_gate_in) + _single_b = mock_gate_only or gate_up_interleave + up_list = None if _single_b else list(acc_up_in) + mfma_res_ty = vec4_f32 + epilogue_pf = None + bias_pf = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + if const_expr(gate_up_interleave): + bias_pf = [] + for ni in range_constexpr(num_acc_n): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (ni % 2 == 1) + else arith.constant(0, index=True) + ) + bias_offset = ( + expert_off_idx + _up_off + _logical_col + ) + bias_pf.append( + _load_bias_scalar(bias_rsrc, bias_offset) + ) + else: + gate_bias_pf = [] + up_bias_pf = ( + [] if const_expr(not mock_gate_only) else None + ) + for ni in range_constexpr(num_acc_n): + global_n = ( + by_n + + n_tile_base + + arith.constant(ni * 16, index=True) + + lane_mod_16 + ) + gate_bias_pf.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + global_n + ) + ) + if const_expr(not mock_gate_only): + up_bias_pf.append( + _load_bias_scalar( + bias_rsrc, + expert_off_idx + inter_idx + global_n, + ) + ) + bias_pf = (gate_bias_pf, up_bias_pf) + tw_pf = None + if const_expr(doweight_stage1): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + sorted_row_pf = bx_m + mi_base_pf + row_off_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias_pf) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + _eff_packed = (ku_count + pack_K - 1) // pack_K + # B-major: fix B (ni), cycle A (mi) -- B from VMEM stays + # in registers while A from LDS is repacked per mi. + for ku128 in range_constexpr(_eff_packed): + for ni in range_constexpr(num_acc_n_packed): + gate_bs_i32 = gate_b_scale[ku128 * num_acc_n_packed + ni] + gate_bs_val = vector.extract( + gate_bs_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(not _single_b): + up_bs_i32 = up_b_scale[ku128 * num_acc_n_packed + ni] + up_bs_val = vector.extract( + up_bs_i32, static_position=[0], dynamic_position=[] + ) + for ikxdl in range_constexpr(pack_K): + k_idx = ku128 * pack_K + ikxdl + if const_expr(k_idx < ku_count): + gate_bp0, gate_bp1 = gate_b_tile_in[k_idx] + if const_expr(not _single_b): + up_bp0, up_bp1 = up_b_tile_in[k_idx] + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + gb0 = gate_bp0[ni_idx] + gb1 = gate_bp1[ni_idx] + gb128 = pack_i64x4_to_i32x8( + gb0, gb1, c0_i64, c0_i64 + ) + if const_expr(not _single_b): + ub0 = up_bp0[ni_idx] + ub1 = up_bp1[ni_idx] + ub128 = pack_i64x4_to_i32x8( + ub0, ub1, c0_i64, c0_i64 + ) + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ + ku128 * m_repeat_packed + mi + ] + a_scale_val = vector.extract( + a_scale_i32, + static_position=[0], + dynamic_position=[], + ) + for imxdl in range_constexpr(pack_M): + mi_idx = mi * pack_M + imxdl + _a_reg_idx = k_idx * m_repeat + mi_idx + if const_expr(is_f8_a): + a0, a1, a2, a3 = a_tile_regs[ + _a_reg_idx + ] + a128 = pack_i64x4_to_i32x8( + a0, a1, a2, a3 + ) + else: + a0, a1 = a_tile_regs[_a_reg_idx] + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + ) + if const_expr(not _single_b): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + return gate_list, up_list, epilogue_pf + + def load_a_subtile(k_idx, mi_idx, lds_buffer): + """Load a single A sub-tile from LDS (one ds_read).""" + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64(curr_row, col_base + 64, lds_buffer) + return (a0, a1, a2, a3) + else: + return (a0, a1) + + _single_b_pipe = mock_gate_only or gate_up_interleave + + def compute_bmajor_mfma_phase( + all_a_tiles, + gate_b_single, + up_b_single, + a_scale_vals, + gate_bs_val, + up_bs_val, + gate_list, + up_list, + k_idx, + ni_idx, + ikxdl, + inxdl, + ): + """B-major MFMA: fix one B (ni), cycle all A tiles (mi). + + Packs B once and reuses across all mi iterations. + A tiles come from LDS (already available, no VMEM wait). + + all_a_tiles: flat list indexed by [k*m_repeat + mi]. + gate_b_single/up_b_single: (b0, b1) for one specific ni. + When _single_b_pipe (mock_gate_only or interleave), up_b_single is None. + a_scale_vals: list of A scale scalars indexed by mi_packed. + """ + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def _pack(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + mfma_res_ty = vec4_f32 + gb128 = _pack(gate_b_single[0], gate_b_single[1], c0_i64, c0_i64) + if const_expr(not _single_b_pipe): + ub128 = _pack(up_b_single[0], up_b_single[1], c0_i64, c0_i64) + + for mi_p in range_constexpr(m_repeat_packed): + a_scale_val = a_scale_vals[mi_p] + for imxdl in range_constexpr(pack_M): + mi_idx = mi_p * pack_M + imxdl + a_reg = all_a_tiles[k_idx * m_repeat + mi_idx] + + if const_expr(is_f8_a): + a128 = _pack(a_reg[0], a_reg[1], a_reg[2], a_reg[3]) + else: + a128 = _pack(a_reg[0], a_reg[1], c0_i64, c0_i64) + + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + if const_expr(not _single_b_pipe): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + + def _interleaved_half( + lds_read, + lds_write, + next_k_dma_py, + next_k_load, + prev_a_tile, + prev_gate_w, + prev_up_w, + prev_a_scale, + prev_gate_bs, + prev_up_bs, + acc_gate, + acc_up, + ): + """One flatmm-style interleaved half-iteration (deep pipeline). + + Generalized for arbitrary m_repeat (block_m=32, 64, ...). + DMA targets lds_write (OTHER buffer) while ds_read uses + lds_read (already DMA'd in previous half). + + Interleaving schedule (per half): + Phase 0: scale VMEM + 2 ds_read(A) -> 4 MFMA(prev) + Phase 1..N: B VMEM(distributed) + 2 ds_read(A, if avail) -> 4 MFMA(prev) + Phase N+1..: remaining B VMEM -> 4 MFMA(prev) + """ + _abs_k = k_base_idx + arith.constant(next_k_load, index=True) + _bk = _abs_k // arith.constant(2, index=True) + _sk = _abs_k // arith.constant(pack_K * 128, index=True) + _k_off = _sk * layout_b_scale.stride_k0 + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(_vmcnt_before_barrier) + _barrier() + rocdl.sched_barrier(0) + + # DMA A to OTHER buffer (for next half), non-blocking + _abs_k_dma = k_base_idx + arith.constant(next_k_dma_py, index=True) + if const_expr(use_async_copy and next_k_dma_py < int(_k_dim)): + prefetch_x_to_lds(_abs_k_dma, lds_write) + if const_expr(not use_async_copy): + _x_regs = load_x_tile(_abs_k_dma) + + # ---- Extract previous scale values ---- + _prev_asvs = [] + for _mi_p in range_constexpr(m_repeat_packed): + _prev_asvs.append( + vector.extract( + prev_a_scale[_mi_p], + static_position=[0], + dynamic_position=[], + ) + ) + _prev_gsv_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _prev_gsv_list.append( + vector.extract( + prev_gate_bs[_gs_ni], + static_position=[0], + dynamic_position=[], + ) + ) + if const_expr(not _single_b_pipe): + _prev_usv_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _prev_usv_list.append( + vector.extract( + prev_up_bs[_us_ni], + static_position=[0], + dynamic_position=[], + ) + ) + + # ---- Execute phases from unified schedule ---- + _a_all = {} + _b_gate_all = {} + _b_up_all = {} + + for _p in range_constexpr(_pipe_n_phases): + # Scale VMEM loads (phase 0 only) + if const_expr(_pp_has_scale[_p]): + _new_as_list = [] + for _mi_p in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + _new_as_list.append(_as1_const) + else: + _raw_as = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[_mi_p] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_as_list.append(_rearrange_a_scale(_raw_as)) + _new_gs_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _gs_raw = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[_gs_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_gs_list.append(_rearrange_b_scale(_gs_raw)) + if const_expr(not _single_b_pipe): + _new_us_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _us_raw = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[_us_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_us_list.append(_rearrange_b_scale(_us_raw)) + + # B VMEM loads + for _b_j in range_constexpr(len(_pp_b_loads[_p])): + _b_type, _b_ku, _b_ni = _pp_b_loads[_p][_b_j] + if const_expr(_b_type == "gate"): + _b_gate_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + gate_n_blk_list[_b_ni], + gate_n_intra_list[_b_ni], + ) + else: + _b_up_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + up_n_blk_list[_b_ni], + up_n_intra_list[_b_ni], + ) + + # A ds_reads + rocdl.sched_barrier(0) + for _a_j in range_constexpr(len(_pp_a_reads[_p])): + _ak, _ami = _pp_a_reads[_p][_a_j] + _a_all[(_ak, _ami)] = load_a_subtile( + _ak, + _ami, + lds_read, + ) + rocdl.sched_barrier(0) + + # MFMAs on prev data + rocdl.s_setprio(1) + for _m_j in range_constexpr(len(_pp_mfma[_p])): + _k_idx, _ni_idx, _ikxdl, _inxdl, _ku128 = _pp_mfma[_p][_m_j] + _ni_packed_idx = _ni_idx // pack_N + _up_b_single = ( + ( + prev_up_w[_k_idx][0][_ni_idx], + prev_up_w[_k_idx][1][_ni_idx], + ) + if not _single_b_pipe + else None + ) + compute_bmajor_mfma_phase( + prev_a_tile, + ( + prev_gate_w[_k_idx][0][_ni_idx], + prev_gate_w[_k_idx][1][_ni_idx], + ), + _up_b_single, + _prev_asvs, + _prev_gsv_list[_ni_packed_idx], + ( + _prev_usv_list[_ni_packed_idx] + if not _single_b_pipe + else None + ), + acc_gate, + acc_up, + _k_idx, + _ni_idx, + _ikxdl, + _inxdl, + ) + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + + # ---- Assemble loaded data for next half-iteration ---- + cur_a_tile = [] + for _k in range_constexpr(k_unroll): + for _mi in range_constexpr(m_repeat): + cur_a_tile.append(_a_all[(_k, _mi)]) + + cur_gate_w = [] + cur_up_w = None if _single_b_pipe else [] + for ku in range_constexpr(k_unroll): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + g = _b_gate_all[(ku, ni)] + g_packs0.append(g[0]) + g_packs1.append(g[1]) + if const_expr(not _single_b_pipe): + u = _b_up_all[(ku, ni)] + u_packs0.append(u[0]) + u_packs1.append(u[1]) + cur_gate_w.append((g_packs0, g_packs1)) + if const_expr(not _single_b_pipe): + cur_up_w.append((u_packs0, u_packs1)) + + cur_a_scale = [] + for _mi_p in range_constexpr(m_repeat_packed): + cur_a_scale.append( + vector.from_elements( + T.vec(1, T.i32), + [_new_as_list[_mi_p]], + ) + ) + cur_gate_bs = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + cur_gate_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_gs_list[_gs_ni]] + ) + ) + if const_expr(not _single_b_pipe): + cur_up_bs = [] + for _us_ni in range_constexpr(num_acc_n_packed): + cur_up_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_us_list[_us_ni]] + ) + ) + else: + cur_up_bs = None + + if const_expr(not use_async_copy): + store_x_tile_to_lds(_x_regs, lds_write) + + return ( + cur_a_tile, + cur_gate_w, + cur_up_w, + cur_a_scale, + cur_gate_bs, + cur_up_bs, + acc_gate, + acc_up, + ) + + # Pipeline (split ping/pong allocators) + rocdl.sched_barrier(0) + + k0 = k_base_idx + if const_expr(use_async_copy): + prefetch_x_to_lds(k0, lds_x_pong) + else: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_x_pong) + rocdl.sched_barrier(0) + _k0_scale = k_base_idx // arith.constant(pack_K * 128, index=True) + a_scale_pong, gate_bs_pong, up_bs_pong = prefetch_ab_scale_tile( + _k0_scale + ) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + acc_gate = [acc_init] * num_acc_n * m_repeat + acc_up = ( + [acc_init] * num_acc_n * m_repeat if not _single_b_pipe else None + ) + + _k1 = k_base_idx + arith.constant(tile_k, index=True) + rocdl.sched_barrier(0) + if const_expr(use_async_copy): + prefetch_x_to_lds(_k1, lds_x_ping) + else: + _x_regs_prime = load_x_tile(_k1) + store_x_tile_to_lds(_x_regs_prime, lds_x_ping) + + _k0_b = k_base_idx // arith.constant(2, index=True) + gate_w0, up_w0 = load_b_tile(_k0_b) + # Prime the deep pipeline: DMA K=tile_k -> ping (1 tile ahead) + if const_expr(use_async_copy): + rocdl.s_waitcnt(0) + gpu.barrier() + rocdl.sched_barrier(0) + a_tile_pong = prefetch_full_a_from_lds(lds_x_pong) + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(6) + + num_k_tiles_py = int(_k_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 + + gate_w_pong = gate_w0 + up_w_pong = up_w0 + + rocdl.sched_barrier(0) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + next_k_load_1 = k_iv_py + tile_k + next_k_load_2 = k_iv_py + tile_k * 2 + next_k_dma_1 = k_iv_py + tile_k * 2 + next_k_dma_2 = k_iv_py + tile_k * 3 + + # Half 1: read ping (DMA'd prev half), DMA->pong, MFMA(pong) + ( + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_ping, + lds_x_pong, + next_k_dma_1, + next_k_load_1, + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) + + # Half 2: read pong (DMA'd Half 1), DMA->ping, MFMA(ping) + ( + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_pong, + lds_x_ping, + next_k_dma_2, + next_k_load_2, + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) + + # _wave_mod2_b = wave_id % arith.constant(2, index=True) + # _wave_odd = arith.cmpi( + # CmpIPredicate.eq, _wave_mod2_b, arith.constant(1, index=True) + # ) + # _if_wave_odd = scf.IfOp(_wave_odd) + # with ir.InsertionPoint(_if_wave_odd.then_block): + # # gpu.barrier() + # _barrier() + # scf.YieldOp([]) + + if const_expr(odd_k_tiles): + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + else: + _k_tail_rel = arith.constant(_k_dim - tile_k, index=True) + k_tail1 = k_base_idx + _k_tail_rel + x_regs_ping = [] + if const_expr(use_async_copy): + prefetch_x_to_lds(k_tail1, lds_x_ping) + else: + x_regs_ping = load_x_tile(k_tail1) + if const_expr(_pad_ku_skip > 0): + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True), + ku_limit=_tail_ku, + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True), + ku_packed_limit=_tail_ku_packed, + ) + else: + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True) + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True) + ) + acc_gate, acc_up, _ = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + ) + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_ping, lds_x_ping) + rocdl.s_waitcnt(0) + _barrier() + if const_expr(_pad_ku_skip > 0): + a_tile_ping = prefetch_full_a_from_lds( + lds_x_ping, ku_limit=_tail_ku + ) + else: + a_tile_ping = prefetch_full_a_from_lds(lds_x_ping) + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_ping, + up_w_ping, + a_tile_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + + bias_pf = None + if const_expr(epilogue_pf is not None): + _, _, bias_pf = epilogue_pf + + # Activation helpers (f32 element-wise on vec4_f32) + def _silu_elem(g): + """silu(x) = x * sigmoid(x); HW fast path: exp2, rcp""" + neg_log2e = arith.constant(-1.4426950408889634, type=f32) + t = g * neg_log2e + emu = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [t], [], []) + one = arith.constant(1.0, type=f32) + den = one + emu + sig = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den], [], []) + return g * sig + + def _silu_mul_vec4(gate_v4, up_v4): + """Element-wise silu(gate) * up on vec4_f32. + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before applying silu(gate) * up. + """ + result_elems = [] + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + if const_expr(swiglu_limit != 0): + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + result_elems.append(_silu_elem(g) * u) + return vector.from_elements(vec4_f32, result_elems) + + def _swiglu_mul_vec4(gate_v4, up_v4): + """Element-wise swiglu(gate, up) on vec4_f32. + swiglu(g, u) = g * sigmoid(alpha * g) * (u + 1) + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before the activation. + """ + result_elems = [] + _alpha = arith.constant(1.702, type=f32) + _one = arith.constant(1.0, type=f32) + _neg_log2e = arith.constant(-1.4426950408889634, type=f32) + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + else: + _limit = arith.constant(float(7.0), type=f32) + _neg_limit = arith.constant(-float(7.0), type=f32) + + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + t = g * _alpha * _neg_log2e + emu = llvm.call_intrinsic( + f32, "llvm.amdgcn.exp2.f32", [t], [], [] + ) + den = _one + emu + sig = llvm.call_intrinsic( + f32, "llvm.amdgcn.rcp.f32", [den], [], [] + ) + result_elems.append(g * sig * (u + _one)) + return vector.from_elements(vec4_f32, result_elems) + + def _act_vec4(gate_v4, up_v4): + """Dispatch activation based on `act` parameter.""" + if const_expr(act == "swiglu"): + return _swiglu_mul_vec4(gate_v4, up_v4) + else: + return _silu_mul_vec4(gate_v4, up_v4) + + # Add bias to raw GEMM accumulators before activation. + # bias layout: [E, 2*inter_dim] flat f32 (non-interleaved: gate then up). + # For gate_up_interleave, map physical column to logical bias offset. + if const_expr(enable_bias and not _is_splitk): + _bias_up_vals = None + if const_expr(bias_pf is not None): + if const_expr(gate_up_interleave): + _bias_gate_vals = bias_pf + else: + _bias_gate_vals, _bias_up_vals = bias_pf + else: + _bias_gate_vals = [] + for _ni in range_constexpr(num_acc_n): + if const_expr(gate_up_interleave): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((_ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (_ni % 2 == 1) + else arith.constant(0, index=True) + ) + _bias_off = expert_off_idx + _up_off + _logical_col + else: + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_off = expert_off_idx + _bn + _bias_gate_vals.append( + _load_bias_scalar(bias_rsrc, _bias_off) + ) + if const_expr(not (mock_gate_only or gate_up_interleave)): + _bias_up_vals = [] + for _ni in range_constexpr(num_acc_n): + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_up_vals.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + inter_idx + _bn + ) + ) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_gate_vals[_ni]] * 4 + ) + acc_gate[_aidx] = arith.addf(acc_gate[_aidx], _bsplat) + + if const_expr(not (mock_gate_only or gate_up_interleave)): + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_up_vals[_ni]] * 4 + ) + acc_up[_aidx] = arith.addf(acc_up[_aidx], _bsplat) + + if const_expr(gate_up_interleave and not _is_splitk): + _gui_out_n = num_acc_n // pack_N + acc = [None] * (_gui_out_n * m_repeat) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(_gui_out_n): + _g_idx = _mi * num_acc_n + _ni * pack_N + _u_idx = _g_idx + 1 + _out_idx = _mi * _gui_out_n + _ni + acc[_out_idx] = _act_vec4( + acc_gate[_g_idx], acc_gate[_u_idx] + ) + elif const_expr(not _is_splitk): + acc = [None] * (int(num_acc_n) * int(m_repeat)) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + acc[_aidx] = _act_vec4(acc_gate[_aidx], acc_up[_aidx]) + + # ---- Epilogue: CShuffle + direct store (accumulate=False) ---- + # Output: out[(t*topk+s) * inter_dim + col] = silu(gate) * up + # For split-K: skip silu, output gate/up separately with atomic add + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + _, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + tokens_i32_v = tokens_i32 + + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + if const_expr(lds_out is None): + raise RuntimeError("CShuffle epilogue requires lds_out") + + _apply_weight = doweight_stage1 and not _is_splitk + + 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, + ): + if const_expr(_apply_weight): + 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=f32 + ) + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(_apply_weight): + v = v * tw + if const_expr(_need_quant): + lds_idx = row_base_lds + col_local + vec1_f32 = T.vec(1, f32) + v1 = vector.from_elements(vec1_f32, [v]) + vector.store(v1, lds_out, [lds_idx], alignment=4) + else: + 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, [v_out]) + vector.store(v1, lds_out, [lds_idx], alignment=2) + + _out_row_stride = ( + inter_dim * 2 * out_elem_bytes + if _is_splitk + else ( + inter_dim // 2 + if _need_fp4 + else (inter_dim if _need_fp8 else inter_dim * out_elem_bytes) + ) + ) + + def precompute_row(*, row_local, row): + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32_v) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + row_byte_base = out_base_idx + ts_idx * arith.constant( + _out_row_stride, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + _e_vec = _e_vec_s1 + _e_vec_sk = 2 + _cshuffle_nlane = min(32, tile_n // _e_vec) + _cshuffle_nlane_sk = min(32, tile_n // _e_vec_sk) + _num_threads_per_quant_blk = _num_threads_per_quant_blk_s1 + + _c0_i32 = arith.constant(0, type=T.i32) + _c1_i32 = arith.constant(1, type=T.i32) + _c2_i32 = arith.constant(2, type=T.i32) + _c3_i32 = arith.constant(3, type=T.i32) + _c4_i32 = arith.constant(4, type=T.i32) + _c5_i32 = arith.constant(5, type=T.i32) + _c15_i32 = arith.constant(15, type=T.i32) + _c22_i32 = arith.constant(22, type=T.i32) + _c23_i32 = arith.constant(23, type=T.i32) + _c28_i32 = arith.constant(28, type=T.i32) + _c31_i32 = arith.constant(31, type=T.i32) + _c32_i32 = arith.constant(32, type=T.i32) + _c64_i32 = arith.constant(64, type=T.i32) + _c254_i32 = arith.constant(254, type=T.i32) + _c256_i32 = arith.constant(256, type=T.i32) + _c0xFF800000_i32 = arith.constant(0xFF800000, type=T.i32) + _c0x400000_i32 = arith.constant(0x400000, type=T.i32) + _c0x7FFFFFFF_i32 = arith.constant(0x7FFFFFFF, type=T.i32) + _c0x80000000_i32 = arith.constant(0x80000000, type=T.i32) + _c0x3F800000_i32 = arith.constant(0x3F800000, type=T.i32) # 1.0f + _c0x40C00000_i32 = arith.constant(0x40C00000, type=T.i32) # 6.0f + _c0x4A800000_i32 = arith.constant(0x4A800000, type=T.i32) + _c0xC11FFFFF_i32 = arith.constant(0xC11FFFFF, type=T.i32) + _c0x7_i32 = arith.constant(0x7, type=T.i32) + _c0_f32 = arith.constant(0.0, type=T.f32) + + _c8_i32 = arith.constant(8, type=T.i32) + _fp_headroom = 2 if _need_fp4 else (8 if _need_fp8 else 0) + _c_headroom_i32 = arith.constant(_fp_headroom, type=T.i32) + + def _f32_to_e2m1(qx_f32): + """Convert a scaled f32 value to fp4 (e2m1) 4-bit integer.""" + # Match fp4_utils.f32_to_mxfp4 / HIP quant: saturate, denorm, + # and normal round-to-nearest-even paths. + qx = qx_f32.bitcast(T.i32) + s = qx & _c0x80000000_i32 + qx_abs = qx & _c0x7FFFFFFF_i32 + denormal_mask = arith.cmpi( + CmpIPredicate.ult, qx_abs, _c0x3F800000_i32 + ) + normal_mask = arith.andi( + arith.cmpi(CmpIPredicate.ult, qx_abs, _c0x40C00000_i32), + arith.cmpi(CmpIPredicate.uge, qx_abs, _c0x3F800000_i32), + ) + + denorm_f32 = qx_abs.bitcast(T.f32) + _c0x4A800000_i32.bitcast(T.f32) + denormal_x = denorm_f32.bitcast(T.i32) - _c0x4A800000_i32 + + mant_odd = (qx_abs >> _c22_i32) & _c1_i32 + normal_x = qx_abs + _c0xC11FFFFF_i32 + mant_odd + normal_x = normal_x >> _c22_i32 + + e2m1 = arith.select(normal_mask, normal_x, _c0x7_i32) + e2m1 = arith.select(denormal_mask, denormal_x, e2m1) + return (s >> _c28_i32) | e2m1 + + if const_expr(_need_sort): + _n32_sort = _sorted_scale_cols_i32 * _c32_i32 + + # Mutable slot for split-K N-offset (gate=0, up=inter_dim) + _sk_n_offset = [0] + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(_need_quant and not _is_splitk): + frag_vals = [] + for i in range_constexpr(_e_vec): + frag_vals.append( + vector.extract( + frag, static_position=[i], dynamic_position=[] + ) + ) + + local_max = _c0_f32 + for i in range_constexpr(_e_vec): + abs_v = llvm.call_intrinsic( + f32, "llvm.fabs.f32", [frag_vals[i]], [], [] + ) + local_max = arith.maximumf(local_max, abs_v) + + for _si in range_constexpr(_num_shuffle_steps_s1): + off = arith.constant(_shuffle_dists_s1[_si], type=T.i32) + peer = local_max.shuffle_xor(off, _c64_i32) + local_max = arith.maximumf(local_max, peer) + + max_i32 = local_max.bitcast(T.i32) + # Match fp4_utils.f32_to_e8m0(max_abs / 4): round the + # exponent at the 1.5x threshold before dropping mantissa. + max_rounded = (max_i32 + _c0x400000_i32) & _c0xFF800000_i32 + exp_field = max_rounded >> _c23_i32 + e8m0_biased = arith.maxsi(exp_field - _c_headroom_i32, _c0_i32) + + quant_exp = _c254_i32 - e8m0_biased + quant_scale = (quant_exp << _c23_i32).bitcast(T.f32) + + if const_expr(_need_fp4): + fp4_vals = [] + for i in range_constexpr(_e_vec): + scaled_v = frag_vals[i] * quant_scale + fp4_vals.append(_f32_to_e2m1(scaled_v)) + + packed_i32 = fp4_vals[0] | (fp4_vals[1] << _c4_i32) + for k in range_constexpr(1, _e_vec // 2): + byte_k = fp4_vals[2 * k] | ( + fp4_vals[2 * k + 1] << _c4_i32 + ) + packed_i32 = packed_i32 | ( + byte_k << arith.constant(k * 8, type=T.i32) + ) + + ptr_addr_idx = row_byte_base + col_g0 / arith.constant( + 2, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + _pack_bytes = _e_vec // 2 + if const_expr(_pack_bytes == 1): + store_val = arith.TruncIOp(T.i8, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=1, nontemporal=True + ) + elif const_expr(_pack_bytes == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=2, nontemporal=True + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, out_ptr_v, alignment=4, nontemporal=True + ) + + elif const_expr(_need_fp8): + scaled_vals = [] + for i in range_constexpr(_e_vec): + scaled_vals.append(frag_vals[i] * quant_scale) + + ptr_addr_idx = row_byte_base + col_g0 + if const_expr(_e_vec <= 4): + packed_i32 = _c0_i32 + for _w in range_constexpr(_e_vec // 2): + packed_i32 = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[2 * _w], + scaled_vals[2 * _w + 1], + packed_i32, + _w, + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + if const_expr(_e_vec == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, + out_ptr_v, + alignment=2, + nontemporal=True, + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + else: + for _wg in range_constexpr(_e_vec // 4): + _b = _wg * 4 + packed_w = _c0_i32 + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b], + scaled_vals[_b + 1], + packed_w, + 0, + ) + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b + 2], + scaled_vals[_b + 3], + packed_w, + 1, + ) + word_ptr = ptr_addr_idx + arith.constant( + _wg * 4, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(word_ptr) + packed_raw = ( + packed_w._value + if hasattr(packed_w, "_value") + else packed_w + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + + if const_expr(_need_sort): + col_g0_i32 = arith.index_cast(T.i32, col_g0) + is_scale_writer = arith.cmpi( + CmpIPredicate.eq, col_g0_i32 & _c31_i32, _c0_i32 + ) + _if_scale = scf.IfOp(is_scale_writer) + with ir.InsertionPoint(_if_scale.then_block): + row_i32_s = arith.index_cast(T.i32, row) + col_s_i32 = col_g0_i32 >> _c5_i32 + d0 = row_i32_s >> _c5_i32 + d1 = (row_i32_s >> _c4_i32) & _c1_i32 + d2 = row_i32_s & _c15_i32 + d3 = col_s_i32 >> _c3_i32 + d4 = (col_s_i32 >> _c2_i32) & _c1_i32 + d5 = col_s_i32 & _c3_i32 + byte_off = ( + d0 * _n32_sort + + d3 * _c256_i32 + + d5 * _c64_i32 + + d2 * _c4_i32 + + d4 * _c2_i32 + + d1 + ) + e8m0_i8 = arith.TruncIOp(T.i8, e8m0_biased) + buffer_ops.buffer_store( + e8m0_i8, + sorted_scale_rsrc, + byte_off, + offset_is_bytes=True, + ) + scf.YieldOp([]) + elif const_expr(_is_splitk): + col_idx = col_g0 + arith.constant(_sk_n_offset[0], index=True) + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec_sk * out_elem_bytes, + ) + else: + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + + _frag_elem = ( + ir.F32Type.get() + if _need_quant + else (ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get()) + ) + + if const_expr(gate_up_interleave and not _is_splitk): + # gui without splitk: acc has activation applied, halved N + _gui_eff_n = _gui_out_n + _gui_tile_n = tile_n // 2 + _gui_cshuffle_nlane = min(32, _gui_tile_n // _e_vec) + _gui_by_n = by_n / arith.constant(2, index=True) + _gui_n_tile_base = n_tile_base / arith.constant(2, index=True) + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=_gui_tile_n, + e_vec=_e_vec, + cshuffle_nlane=_gui_cshuffle_nlane, + block_size=total_threads, + m_repeat=m_repeat, + num_acc_n=_gui_eff_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=_gui_by_n, + n_tile_base=_gui_n_tile_base, + lds_out=lds_out, + frag_elem_type=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + elif const_expr(mock_gate_only or (gate_up_interleave and _is_splitk)): + # mock_gate_only: single pass, by_n covers full [0, 2*inter_dim) + _eff_e_vec = _e_vec_sk + acc = acc_gate + 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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + elif const_expr(_is_splitk): + # Two-pass epilogue: gate then up, each with atomic add + _eff_e_vec = _e_vec_sk + + # Pass 1: gate + acc = acc_gate + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + gpu.barrier() + + # Pass 2: up + acc = acc_up + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + else: + 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, + cshuffle_nlane=_cshuffle_nlane, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + _if_blk = scf.IfOp(blk_valid) + with ir.InsertionPoint(_if_blk.then_block): + _ifexpert_of = scf.IfOp(exp_valid) + with ir.InsertionPoint(_ifexpert_of.then_block): + _moe_gemm1_body() + scf.YieldOp([]) + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher -- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage1, + act, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + use_async_copy, + waves_per_eu, + k_batch, + gate_mode, + a_scale_one, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_max_token_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + inter_dim_pad_total = arith.constant(2 * inter_dim_pad, index=True) + tile2_pad = 0 + if const_expr(not gate_only): + tile_k_stage2 = tile_k // 2 + tile2_pad = ( + tile_k_stage2 - (inter_dim - inter_dim_pad) % tile_k_stage2 + ) % tile_k_stage2 + + inter_in = arith.index_cast(ir.IndexType.get(), i32_inter_in.ir_value()) + tile_n_index = arith.constant(tile_n, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + gx = ( + inter_in - inter_dim_pad_total + tile2_pad + tile_n_index - 1 + ) / tile_n_index + else: + gx = ( + (inter_in - inter_dim_pad_total + tile2_pad + 2 * tile_n_index - 1) + / tile_n_index + / arith.constant(2, index=True) + ) + + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + arg_out_scale_sorted, + i32_tokens_in, + i32_inter_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch(grid=(gx, gy, k_batch), block=(total_threads, 1, 1), stream=stream) + + return launch_mixed_moe_gemm1 + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm2( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + # Optional experiment: write per-(token,slot) output (no atomics) into an output shaped + # [tokens*topk, model_dim] (or [tokens, topk, model_dim] flattened), then reduce over topk outside. + # This can reduce atomic contention for small tokens at the cost of extra bandwidth / reduction. + accumulate: bool = True, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 4, + sort_block_m: int = 0, + b_nt: int = 2, + xcd_swizzle: int = 0, +): + """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + + persist_m: + - > 0: legacy mode -- each CTA processes exactly persist_m consecutive M tiles. + - <= 0: **persistent mode** -- grid_y = cu_num (auto-detected), each CTA + round-robins over M tiles with stride cu_num. + + a_dtype: + - "fp8": A2 is fp8 + - "fp16": A2 is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": A2 is int8 + - "fp4": A2 is fp4 + + b_dtype: + - "fp8": W is fp8 + - "fp16": W is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": W is int8 + - "int4": W4A8 path: A2 is int8, W is packed int4 (2 values per byte) unpacked to int8 in-kernel + - "fp4": W is fp4 + + 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). + + `sort_block_m` is the block_size used by moe_sorting / stage1. When 0 (default), + assumed equal to `tile_m`. When set, stage2 can use a different tile_m from + sorting/stage1. Requires sort_block_m % tile_m == 0. + """ + _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m + if _sort_block_m != tile_m and _sort_block_m % tile_m != 0: + raise ValueError( + f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" + ) + + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + _scale_pack_m = 2 # physical mn_pack in preshuffle microscale layout + _scale_pack_n = 2 + _scale_pack_k = 2 # physical k_pack in preshuffle scale layout + pack_M = min(_scale_pack_m, tile_m // 16) + pack_N = min(_scale_pack_n, tile_n // 64) + _k_unroll_raw = (int(tile_k) * (2 if a_dtype == "fp16" else 1)) // 128 + pack_K = min(_scale_pack_k, _k_unroll_raw) + + elem_bytes = 1 + + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + # ---- Static B preshuffle strides (compile-time) ---- + # All values below are Python ints computable at kernel-compile time. + # Using them in an explicit multiply-add replaces the fly dialect's + # dynamic ``crd2idx`` path which emits Barrett reduction for the + # non-power-of-2 ``n0 = experts*model_dim//16`` shape. + _b_kpack_bytes_s = 8 if (b_dtype == "int4") else 16 + _b_kpack_elems_s = _b_kpack_bytes_s // b_elem_bytes + _b_c_k_s = inter_dim // _scale_pack_k + _b_c_k0_s = (_b_c_k_s * b_elem_bytes) // 64 + _b_stride_nlane = _b_kpack_elems_s # 16 + _b_stride_klane = 16 * _b_stride_nlane # 256 + _b_stride_k0 = 4 * _b_stride_klane # 1024 + _b_stride_n0 = _b_c_k0_s * _b_stride_k0 # c_k0 * 1024 + assert model_dim % 16 == 0, "model_dim must be divisible by 16" + _expert_b_stride = (model_dim // 16) * _b_stride_n0 + + # K64-byte micro-step: always 64 bytes per `ku`. For fp16, this is 32 elements (2xK16 MFMA). + if (tile_k_bytes % 64) != 0: + raise ValueError( + f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " + f"(tile_k={tile_k}, elem_bytes={a_elem_bytes})" + ) + + out_s = str(out_dtype).strip().lower() + if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): + 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 = b_dtype == "int4" + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * model_dim * inter_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * model_dim * 4 + # INT4 here means W4A8: A2 is int8, W is packed int4 and unpacked to int8 in-kernel. + is_int8 = False + + 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`)." + ) + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def _scale_elem_type(): + return T.i32 + + total_threads = 256 + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + "tile_m*tile_k*elem_bytes must be divisible by " + f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={a_elem_bytes}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if a_elem_vec_pack > 1: + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + else: + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + + if out_is_f32: + # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. + _use_cshuffle_epilog = ( + 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("FLIR_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 (FLIR_MOE_STAGE2_CSHUFFLE=1)." + ) + + # NOTE: Keep this as a callable so we don't require an MLIR Context at Python-time. + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + 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 the compiler 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. + _persistent = persist_m <= 0 + if _persistent: + from aiter.jit.utils.chip_info import get_cu_num + + _cu_num = get_cu_num() + else: + _cu_num = 0 + _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" + _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" + f"_t{tile_m}x{tile_n}x{tile_k}" + f"_vscale_fix3{_pm_tag}{_sbm_tag}{_xcd_tag}" + ).replace("-", "_") + # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- + # Ping-pong A2 tiles via separate allocators (like stage1). + _single_x_bytes = int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + _cshuffle_elem_bytes_s2 = 2 # f16/bf16 = 2 bytes + lds_out_bytes = ( + _cshuffle_elem_bytes_s2 * int(tile_m) * int(tile_n) + if _use_cshuffle_epilog + else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec4_i32 = T.vec(4, i32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec8_elems = 8 if a_elem_bytes == 1 else 4 + vec4_elems = 4 if a_elem_bytes == 1 else 2 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = ( + arith.constant_vector(0, vec4_i32) + if is_int8 + else arith.constant_vector(0.0, vec4_f32) + ) + + # A2 layout (flatten token-slot -> M; use i32 for fly.make_shape). + topk_idx = arith.constant(topk, index=True) + m_in = tokens_in * topk_idx + + # B preshuffle layout: [experts*model_dim, inter_dim] + c_n_total = arith.constant(experts * model_dim, index=True) + kpack_bytes = 8 if is_int4 else 16 + # (inlined: _div_pow2, _mod_pow2 are module-global) + + def check_c_n_valid_gate(base_n): + return arith.cmpi(CmpIPredicate.ult, base_n, model_dim - model_dim_pad) + + def check_c_k_valid_gate(base_k): + return arith.cmpi(CmpIPredicate.ult, base_k, inter_dim - inter_dim_pad) + + # A&B's scale preshuffle layout + # For fp4, k_in is already packed (inter_dim // a_elem_vec_pack), so we need original inter_dim + c_k_orig = arith.constant(inter_dim, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=m_in, c_k=c_k_orig + ) + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=c_k_orig + ) + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along model_dim (N-dim) + bx_persist = gpu.block_id("y") # persistent WG index (M-dim) + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_mdp_sw = arith.constant(model_dim_pad, index=True) + _gx = (n_in - _c_mdp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + if const_expr(_persistent): + _gy = arith.constant(_cu_num, index=True) + else: + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S = xcd_swizzle + _c_wgm = arith.constant(_WGM_S, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + (T.bf16 if out_is_bf16 else T.f16), + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # 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 = arith.constant(topk, index=True) + + # X(A2): buffer size in bytes, accounting for FP4 packing (2 elements per byte). + # fp8/int8: 1 byte per element -> bytes = tokens*topk * K + # fp4: 2 elements per byte -> bytes = tokens*topk * K / 2 + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + x_nbytes_idx = _div_pow2( + (tokens_in * c_topk) * k_in * c_elem_bytes, int(a_elem_vec_pack) + ) + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # 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 * arith.constant(out_elem_bytes, index=True) + ) + if const_expr(not bool(accumulate)): + out_nbytes_idx = ( + tokens_in + * arith.index(topk) + * n_in + * arith.constant(out_elem_bytes, index=True) + ) + out_nbytes_i32 = arith.index_cast(T.i32, out_nbytes_idx) + out_rsrc = _ptr_buffer_resource(arg_out, out_nbytes_i32) + + # num_valid_ids (sorted padded MN) for scale sizing / guards. + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + # num_valid_ids is a scalar (same value for all lanes) loaded into + # VGPR. Promote to SGPR so downstream buffer resource descriptors + # that use it for num_records stay in SGPRs, eliminating the + # expensive waterfall loop the compiler would otherwise emit. + num_valid_i32 = rocdl.ReadfirstlaneOp(T.i32, num_valid_i32).res + num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + + # fp16 path ignores scales completely (implicit scale=1.0). + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not is_f16_a): + if const_expr(is_f4_a or is_f8_a): + # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. + # Caller must pre-scatter a2_scale via moe_mxfp4_sort. + kblk = _div_pow2(k_in, 32) + sx_nbytes_idx = num_valid_idx * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + else: + # scale_x (A2 scale): [tokens*topk] f32 -> bytes = tokens*topk*4 + sx_nbytes_idx = (tokens_in * c_topk) * arith.constant(4, index=True) + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + # Weight microscale buffer (packed i32 holding e8m0 bytes). + # Use an exact descriptor size so hardware OOB checking works. + kblk_w = _div_pow2(k_in, 32) # K/32 + mn_w = arith.constant(experts * model_dim, index=True) + sw_nbytes_idx = mn_w * kblk_w # bytes (e8m0) + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + # sorted_token_ids / sorted_weights: [blocks*tile_m] (padded length) + sorted_nbytes_idx = ( + size_expert_ids_in + * arith.constant(tile_m, index=True) + * arith.constant(4, index=True) + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + # expert ids: [sort_blocks] i32. + _c_sbm = arith.constant(_sort_block_m, index=True) + _c_tm = arith.constant(tile_m, index=True) + _c1 = arith.constant(1, index=True) + _sort_blocks_ub = _div_pow2( + size_expert_ids_in * _c_tm + _c_sbm - _c1, _sort_block_m + ) + eid_nbytes_idx = _sort_blocks_ub * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # ---- persist loop ---- + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + + if const_expr(_persistent): + # Expert-phase scheduling: contiguous M-tile dispatch. + # grid_y = cu_num, each CTA handles a contiguous chunk of M-tiles: + # [bx_persist * tiles_per_block, ..., (bx_persist+1) * tiles_per_block - 1] + # Adjacent blocks process adjacent M-tiles -> same expert -> B weight L2 reuse. + _c_cu = arith.constant(_cu_num, index=True) + _c_tm_p = arith.constant(tile_m, index=True) + _num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + _total_m_tiles = (_num_valid_idx + _c_tm_p - _c1_p) / _c_tm_p + _tiles_per_block = (_total_m_tiles + _c_cu - _c1_p) / _c_cu + _i1 = ir.IntegerType.get_signless(1) + _init_active = arith.constant(1, type=_i1) + _for_persist = scf.ForOp(_c0_p, _tiles_per_block, _c1_p, [_init_active]) + else: + # Legacy mode: fixed persist_m consecutive tiles. + _c_pm = arith.constant(persist_m, index=True) + _init_prev_expert = arith.constant(0, type=T.i32) + _init_prev_b_base = arith.constant(0, index=True) + _for_persist = scf.ForOp( + _c0_p, + _c_pm, + _c1_p, + [_init_prev_expert, _init_prev_b_base], + ) + + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + + if const_expr(_persistent): + _still_active = _for_persist.inner_iter_args[0] + bx = bx_persist * _tiles_per_block + _mi_p + else: + _prev_expert_i32 = _for_persist.inner_iter_args[0] + _prev_expert_b_base = _for_persist.inner_iter_args[1] + bx = bx_persist * arith.constant(persist_m, index=True) + _mi_p + + bx_m = bx * arith.constant(tile_m, index=True) + + # Early-exit guard: skip garbage expert blocks beyond `num_valid_ids`. + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + + sort_blk = _div_pow2(bx_m, _sort_block_m) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, sort_blk, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + if const_expr(_persistent): + # Absolute B-base: no cross-iteration state needed. + _expert_b_base = expert_idx * arith.constant( + _expert_b_stride, index=True + ) + else: + # Legacy incremental B-base: delta = (cur - prev) * stride + _delta_expert = arith.subi(expert_i32, _prev_expert_i32) + _delta_expert_idx = arith.index_cast(ir.IndexType.get(), _delta_expert) + _delta_b = _delta_expert_idx * arith.constant( + _expert_b_stride, index=True + ) + _expert_b_base = _prev_expert_b_base + _delta_b + + # Early-exit: if the first row of this tile is a sentinel (all-padding tile), + # skip the entire GEMM. + _first_tok = buffer_ops.buffer_load( + sorted_rsrc, bx_m, vec_width=1, dtype=T.i32 + ) + _first_tid = arith.andi(_first_tok, arith.constant(0xFFFFFF, type=T.i32)) + _tokens_i32_guard = arith.index_cast(T.i32, tokens_in) + tile_has_tokens = arith.cmpi( + CmpIPredicate.ult, _first_tid, _tokens_i32_guard + ) + + # For tile_m < 32 (pack_M < _scale_pack_m): shift a_scale i32 so the + # correct bytes land at the op_sel positions we use. + if const_expr(pack_M < _scale_pack_m): + _m_off = _mod_pow2(_div_pow2(bx_m, 16), _scale_pack_m) + _m_scale_shift_i32 = arith.index_cast( + T.i32, _m_off * arith.constant(8, index=True) + ) + else: + _m_scale_shift_i32 = None + + def _moe_gemm2_then_body(): + # Expert id for this M tile. + n_idx = arith.constant(model_dim, index=True) + 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 we require 16B. + if const_expr(is_f16_a): + if const_expr(bytes_per_thread_x % 16 != 0): + raise ValueError( + f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16" + ) + x_load_bytes = 16 + else: + if 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) + vec4_i32 = T.vec(4, i32) + + c_k_div4 = _div_pow2( + _div_pow2(k_in, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(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, + ) + + vec1_i32 = T.vec(1, i32) + vec2_i32 = T.vec(2, i32) + x_load_vec_elems = ( + x_load_bytes if a_elem_bytes == 1 else x_load_bytes // a_elem_bytes + ) + + 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. + """ + if const_expr(x_load_bytes == 16): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + # 8B/4B: convert dword index to byte offset and use offset_in_bytes path. + idx_bytes = idx_i32 * arith.index(4) + return _buffer_load_vec( + buffer_ops, + vector, + x_rsrc, + idx_bytes, + elem_type=x_elem, + vec_elems=x_load_vec_elems, + elem_bytes=a_elem_bytes, + offset_in_bytes=True, + ) + + # 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(ir.IndexType.get(), row_ts_i32) + + x_row_base_div4.append(row_ts_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 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(vec4_i32, x_vec)) + elif const_expr(x_load_bytes == 8): + parts.append(vector.bitcast(vec2_i32, x_vec)) + else: + parts.append(vector.bitcast(vec1_i32, x_vec)) + return parts + + # tx -> wave/lane (GEMM-style decomposition). + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + + row_a_lds = lane_mod_16 + + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + # Dynamic N tiling within block. + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_mod_4 = _mod_pow2(wave_id, 4) + n_tile_base = wave_mod_4 * c_n_per_wave + + by_n = by * arith.constant(tile_n, index=True) + + if const_expr(pack_N < _scale_pack_n): + _global_n_base = expert_off_idx + by_n + n_tile_base + _n_off = _mod_pow2(_div_pow2(_global_n_base, 16), _scale_pack_n) + _n_scale_shift_i32 = arith.index_cast( + T.i32, _n_off * arith.constant(8, index=True) + ) + else: + _n_scale_shift_i32 = None + n_intra_list = [None] * num_acc_n + n_blk_list = [None] * num_acc_n + col_g_list = [None] * num_acc_n + for i in range_constexpr(num_acc_n): + offset = i * 16 + col_g = by_n + n_tile_base + col_g = _div_pow2(col_g, 2) + offset + col_g = col_g + lane_mod_16 + col_g_list[i] = col_g + c_offset = arith.constant(offset, index=True) + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + n_blk_list[i] = _div_pow2(global_n, 16) + n_intra_list[i] = _mod_pow2(global_n, 16) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 # K64-byte micro-step (2x MFMA) + + # fp4 pack + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku_s2 = tile_k // k_unroll + _pad_k_elems_s2 = (inter_dim_pad % tile_k) if inter_dim_pad > 0 else 0 + _pad_ku_skip_s2 = _pad_k_elems_s2 // _K_per_ku_s2 + _tail_ku_s2 = k_unroll - _pad_ku_skip_s2 + _tail_ku_packed_s2 = ( + (_tail_ku_s2 + pack_K - 1) // pack_K + if _pad_ku_skip_s2 > 0 + else None + ) + + # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- + def load_b_packs_k64(base_k, ku: int, ni: int): + """Load one K64-byte B micro-step: single 16B load, split into 2x i64.""" + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0_base = _div_pow2(base_k_bytes, 64) + k0 = k0_base + arith.constant(ku, index=True) + k1 = lane_div_16 + # Incremental B addressing: _expert_b_base carries the + # expert's preshuffle offset (updated via delta each + # persist_m iteration); local n_blk/n_intra contribute + # the per-lane within-tile offset. All strides are + # compile-time constants -> shift/mul, no Barrett. + idx_pack = ( + _expert_b_base + + n_blk_list[ni] * arith.constant(_b_stride_n0, index=True) + + k0 * arith.constant(_b_stride_k0, index=True) + + k1 * arith.constant(_b_stride_klane, index=True) + + n_intra_list[ni] * arith.constant(_b_stride_nlane, index=True) + ) + + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + b_tile = [] + for ku in range_constexpr(ku_limit): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + _b_split_enabled = k_unroll >= 2 + _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll + + def load_b_tile_lo(base_k): + """Load first half of B tile (ku < _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_b_tile_hi(base_k): + """Load second half of B tile (ku >= _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku, k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_scale(arg_scale, rsrc, scale_info, ku, mni): + k_lane = lane_div_16 + n_lane = lane_mod_16 + # Direct arith crd2idx: idx = mni*stride_n0 + ku*stride_k0 + k_lane*stride_klane + n_lane + idx_pack = ( + mni * scale_info.stride_n0 + + ku * scale_info.stride_k0 + + k_lane * scale_info.stride_klane + + n_lane + ) + s = buffer_ops.buffer_load(rsrc, idx_pack, vec_width=1, dtype=T.i32) + return vector.from_elements(T.vec(1, T.i32), [s]) + + def _apply_k_shift(scale_vec, k_shift_bits): + if const_expr(k_shift_bits > 0): + val = vector.extract( + scale_vec, static_position=[0], dynamic_position=[] + ) + val = arith.shrui(val, arith.constant(k_shift_bits, type=T.i32)) + return vector.from_elements(T.vec(1, T.i32), [val]) + return scale_vec + + def load_b_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + b_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for ni in range_constexpr(num_acc_n_packed): + scale = load_scale( + arg_scale_w, + sw_rsrc, + layout_b_scale, + ku + base_k, + ni + + _div_pow2( + _div_pow2( + expert_off_idx + by_n + n_tile_base, + _scale_pack_n, + ), + 16, + ), + ) + scale = _apply_k_shift(scale, k_shift_bits) + b_scale_tile.append(scale) + return b_scale_tile + + def load_a_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + a_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for mi in range_constexpr(m_repeat_packed): + scale = load_scale( + arg_scale_x, + sx_rsrc, + layout_a_scale, + ku + base_k, + mi + _div_pow2(_div_pow2(bx_m, _scale_pack_m), 16), + ) + scale = _apply_k_shift(scale, k_shift_bits) + a_scale_tile.append(scale) + return a_scale_tile + + def prefetch_ab_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + return [ + load_a_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + load_b_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + ] + + vec8_x = T.vec(vec8_elems, x_elem) + vec4_x_lds = T.vec(vec4_elems, x_elem) + + # ---- Pipeline helpers: store X tile to LDS (unused in DMA path) ---- + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + 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_buffer, + vec8_ty=vec8_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x2=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + else: # x_load_bytes == 4 + lds_store_4b_xor16( + arith, + vector, + lds_memref=lds_buffer, + vec4_ty=vec4_x_lds, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x1=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def compute_tile( + acc_in, + b_tile_in, + lds_buffer, + a_scale=None, + b_scale=None, + *, + prefetch_epilogue: bool = False, + a0_prefetch=None, + a1_prefetch=None, + b_hi_loader=None, + ku_count=k_unroll, + ): + if const_expr(b_hi_loader is not None): + b_tile_full = [None] * k_unroll + for i in range_constexpr(_b_split_ku): + b_tile_full[i] = b_tile_in[i] + else: + b_tile_full = b_tile_in + acc_list = list(acc_in) + mfma_res_ty = vec4_i32 if is_int8 else vec4_f32 + + epilogue_pf = None + bias = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + bias = [] + for ni in range_constexpr(num_acc_n): + global_n = by_n + n_tile_base + ni * 16 + lane_mod_16 + bias_offset = expert_off_idx + global_n + bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) + tw_pf = None + if const_expr(doweight_stage2): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + row_in_tile_pf = mi_base_pf + row_off_pf + sorted_row_pf = bx_m + row_in_tile_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + # fp4 path -- single k_idx loop [0, k_unroll). + # b_hi load is issued at the very start so all k_unroll + # MFMAs can overlap the VMEM latency. + _pack_K_shift = (pack_K - 1).bit_length() + _pack_K_mask = pack_K - 1 + + if const_expr(b_hi_loader is not None): + _b_hi = b_hi_loader() + for _bhi_i in range_constexpr(len(_b_hi)): + b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] + + for k_idx in range_constexpr(ku_count): + ku128 = k_idx >> _pack_K_shift + ikxdl = k_idx & _pack_K_mask + + b_packs0, b_packs1 = b_tile_full[k_idx] + + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ku128 * m_repeat_packed + mi] + a_scale_val = vector.extract( + a_scale_i32, static_position=[0], dynamic_position=[] + ) + if const_expr(_m_scale_shift_i32 is not None): + a_scale_val = arith.shrui( + a_scale_val, _m_scale_shift_i32 + ) + for ni in range_constexpr(num_acc_n_packed): + b_scale_i32 = b_scale[ku128 * num_acc_n_packed + ni] + b_scale_val = vector.extract( + b_scale_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(_n_scale_shift_i32 is not None): + b_scale_val = arith.shrui( + b_scale_val, _n_scale_shift_i32 + ) + + for imxdl in range_constexpr(pack_M): + col_base0 = col_base + mi_idx = mi * pack_M + imxdl + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr( + (a0_prefetch is not None) + and (k_idx == 0) + and (mi_idx == 0) + ): + a0, a1 = a0_prefetch + elif const_expr( + (a1_prefetch is not None) + and (k_idx == 1) + and (mi_idx == 0) + ): + a0, a1 = a1_prefetch + else: + a0, a1 = lds_load_packs_k64( + curr_row_a_lds, col_base0, lds_buffer + ) + + if const_expr(is_f8_a): + col_base1 = col_base + 64 + a2, a3 = lds_load_packs_k64( + curr_row_a_lds, col_base1, lds_buffer + ) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + else: + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + + b0 = b_packs0[ni_idx] + b1 = b_packs1[ni_idx] + b128 = pack_i64x4_to_i32x8( + b0, b1, c0_i64, c0_i64 + ) + + acc_idx = mi_idx * num_acc_n + ni_idx + acc_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + b128, + acc_list[acc_idx], + cbsz, + blgp, + ikxdl * _scale_pack_m + imxdl, + a_scale_val, + ikxdl * _scale_pack_n + inxdl, + b_scale_val, + ], + ) + ) + + return acc_list, epilogue_pf + + # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- + # ---- Async DMA: GMEM -> LDS (bypasses VGPR, like stage1) ---- + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler(): + rocdl.sched_barrier(0) + + def _k_shift_bits(k_py): + if const_expr(pack_K >= _scale_pack_k): + return 0 + return ((k_py // 128) % _scale_pack_k) * _scale_pack_m * 8 + + def _k_base(k_py): + return k_py // _scale_pack_k // 128 + + # Preload sorted_idx into lds_tid for epilogue precompute_row + # (N-independent; placed before N-tile loop so it's done once per M-tile.) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + gpu.barrier() + + # Prologue -- B-first + async DMA X(0) -> pong. + k0 = arith.index(0) + if const_expr(_b_split_enabled): + b_cur = load_b_tile_lo(k0) + else: + b_cur = load_b_tile(k0) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(0), _k_shift_bits(0) + ) + rocdl.sched_barrier(0) + prefetch_x_to_lds(k0, lds_x_pong) + rocdl.s_waitcnt(0) + gpu.barrier() + + acc = [acc_init] * num_acc_n * m_repeat + + # Cross-tile A0+A1 LDS prefetch from pong buffer. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + _a1_col_base = col_offset_base + 128 // a_elem_vec_pack + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + # 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 + + c2_tile_k = arith.constant(tile_k * 2, index=True) + b_pong = b_cur + k0_pong_bk = k0 + + # Only emit the scf.for when there are actually iterations to run. + # When k_main2_py == 0 the loop body is empty; emitting an scf.for + # would create a region whose internal SSA values cannot be used + # by the post-loop tail code. + def _make_b_hi_loader(base_k): + """Create a b_hi_loader callable for a given base_k.""" + return lambda _bk=base_k: load_b_tile_hi(_bk) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + rocdl.sched_barrier(0) + k_iv = arith.index(k_iv_py) + next_k1 = k_iv + tile_k + next_k1_bk = next_k1 // 2 + # DMA X(next_k1) -> ping (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k1, lds_x_ping) + b_ping_lo = ( + load_b_tile_lo(next_k1_bk) + if _b_split_enabled + else load_b_tile(next_k1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(next_k1), _k_shift_bits(next_k1) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) + if _b_split_enabled + else None + ), + ) + hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Cross-tile prefetch for the ping tile we are about to compute. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 + else None + ) + + next_k2 = k_iv + c2_tile_k + next_k2_py = k_iv_py + tile_k * 2 + next_k2_bk = next_k2 // 2 + # DMA X(next_k2) -> pong (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k2, lds_x_pong) + b_pong = ( + load_b_tile_lo(next_k2_bk) + if _b_split_enabled + else load_b_tile(next_k2_bk) + ) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(next_k2_py), _k_shift_bits(next_k2_py) + ) + + acc, _ = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + b_hi_loader=( + _make_b_hi_loader(next_k1_bk) + if _b_split_enabled + else None + ), + ) + k0_pong_bk = next_k2_bk + hot_loop_scheduler() + gpu.barrier() + + # Cross-tile prefetch for the next pong tile. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + if const_expr(odd_k_tiles): + # Tail: single remaining tile (already in pong buffer). + acc, epilogue_pf = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + prefetch_epilogue=True, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + else: + # Tail: 2 remaining tiles. + k_tail1 = (k_in + tile_k - 1) // tile_k * tile_k - tile_k + k_tail1_py = ( + int(inter_dim) + tile_k - 1 + ) // tile_k * tile_k - tile_k + k_tail1_bk = k_tail1 // 2 + # DMA tail X -> ping + prefetch_x_to_lds(k_tail1, lds_x_ping) + if const_expr(_pad_ku_skip_s2 > 0): + b_ping_lo = load_b_tile(k_tail1_bk, ku_limit=_tail_ku_s2) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), + _k_shift_bits(k_tail1_py), + ku_packed_limit=_tail_ku_packed_s2, + ) + else: + b_ping_lo = ( + load_b_tile_lo(k_tail1_bk) + if _b_split_enabled + else load_b_tile(k_tail1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ) + + # hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Epilogue tile with sw prefetch. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 and (_pad_ku_skip_s2 == 0 or _tail_ku_s2 >= 2) + else None + ) + acc, epilogue_pf = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + prefetch_epilogue=True, + b_hi_loader=( + None + if _pad_ku_skip_s2 > 0 + else ( + _make_b_hi_loader(k_tail1_bk) + if _b_split_enabled + else None + ) + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- + # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. + + sw_pf = None + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + sw_pf, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + + zero_i32 = arith.constant(0) + + def atomic_add_f16x2(val_f16x2, byte_off_i32): + rocdl.raw_ptr_buffer_atomic_fadd( + val_f16x2, + out_rsrc, + byte_off_i32, + zero_i32, + zero_i32, + ) + + # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). + if const_expr(lds_out is None): + raise RuntimeError( + "FLIR_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased." + ) + + # Precompute the output base address (i64 index) for ALL paths. + # Both accumulate=True (global atomic) and accumulate=False (global store) + # need 64-bit addressing to avoid i32 offset overflow when + # tokens * model_dim * elem_bytes > INT32_MAX (~150K tokens for model_dim=7168). + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + 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, + ): + # Match origin/dev_a16w4: rely on sentinel padded rows + hardware OOB behavior. + fused2 = buffer_ops.buffer_load( + sorted_rsrc, row, vec_width=1, dtype=T.i32 + ) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + + t_ok = arith.cmpi(CmpIPredicate.ult, t2, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s2, topk_i32_v) + ts_ok = arith.andi(t_ok, s_ok) + t2_safe = arith.select(ts_ok, t2, arith.constant(0)) + s2_safe = arith.select(ts_ok, s2, arith.constant(0)) + t2_safe * topk_i32_v + s2_safe + + if const_expr(doweight_stage2): + 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=f32 + ) + + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(is_int8): + v = arith.sitofp(f32, v) + if const_expr(enable_bias): + v = v + bias_pf[ni] + + 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, [v_out]) + + vector.store(v1, lds_out, [lds_idx], alignment=2) + + def precompute_row(*, row_local, row): + # Use lds_tid (sorted_idx preloaded to LDS) instead of buffer_load + # to avoid extra VMEM round-trips in the epilogue. + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + if const_expr(accumulate): + row_byte_base = out_base_idx + t_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + else: + row_byte_base = out_base_idx + ts_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + """Convert an index-typed byte address to !llvm.ptr.""" + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(not bool(accumulate)): + # ---- 64-bit global store path (avoids i32 offset overflow) ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + else: + # ---- accumulate=True: 64-bit global atomic path ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec * out_elem_bytes, + ) + + _e_vec = 2 if accumulate else min(tile_n // 32, 8) + 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=( + ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get() + ), + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) + + if const_expr(_persistent): + # Short-circuit: contiguous tiles are monotonically increasing, + # so once bx_m >= num_valid_ids all remaining tiles are invalid. + _cur_active = arith.andi(_still_active, blk_valid) + _do_gemm = arith.andi( + _cur_active, arith.andi(exp_valid, tile_has_tokens) + ) + _if_valid = scf.IfOp(_do_gemm) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([_cur_active]) + else: + _if_valid = scf.IfOp(_all_valid) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([expert_i32, _expert_b_base]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher (flyc.jit + .launch) -------------------------------- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage2, + accumulate, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + _sort_block_m, + _cu_num if _persistent else 0, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + _tile_n_idx = arith.constant(tile_n, index=True) + _model_dim_pad_idx = arith.constant(model_dim_pad, index=True) + gx = ( + n_in - _model_dim_pad_idx + _tile_n_idx - arith.constant(1, index=True) + ) / _tile_n_idx + if const_expr(_persistent): + gy = arith.constant(_cu_num, index=True) + else: + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + 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_mixed_moe_gemm2 + +# =========================================================================== +# Host-side launchers (adapted from aiter/ops/flydsl/moe_kernels.py). +# These pack pointer args and drive the inline compile_mixed_moe_gemm1/2 +# builders above. +# =========================================================================== +_DLPACK_SAFE = (torch.uint8, torch.float16, torch.bfloat16, torch.float32) + + +def _view_safe(t: torch.Tensor) -> torch.Tensor: + """View as uint8 if dtype is not dlpack-safe, otherwise return as-is.""" + return ( + t.view(torch.uint8) + if t is not None and t.numel() > 0 and t.dtype not in _DLPACK_SAFE + else t + ) + + +def _ptr_view_safe(t: torch.Tensor): + """Pass only the device data pointer; shape is carried by explicit args.""" + view = _view_safe(t) + type_name = type(view).__name__ + module_name = type(view).__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, view.data_ptr()) + + +def _s1_args_fp4( + out, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + out_scale_sorted, + token_num, + n_in, + k_in, + size_expert_ids_in, + dev, + bias=None, + stream=None, +): + empty_f32 = torch.empty(0, device=dev, dtype=torch.float32) + _bias = bias if bias is not None else empty_f32 + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(out), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + _ptr_view_safe(out_scale_sorted), + token_num, + n_in, + k_in, + size_expert_ids_in, + stream, + ) + + +def _s2_args_fp4( + target, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + token_num, + n_in, + k_in, + blocks, + dev, + bias=None, + stream=None, +): + _bias = ( + bias.view(-1) + if bias is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(target), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + token_num, + n_in, + k_in, + blocks, + stream, + ) + + +def _run_compiled(exe, args): + """Call the JitFunction with the given args (handles compile caching).""" + try: + exe(*args) + except Exception: + # JitFunction.__call__ leaks ir.Context on compilation failure; clean up + # leaked contexts so subsequent calls do not take a wrong code path. + try: + while ir.Context.current is not None: + ir.Context.current.__exit__(None, None, None) + except Exception: + pass + raise + + +def build_moe_stage1_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + act: str = "silu", + persist_m: int = 1, + use_async_copy: bool = False, + k_batch: int = 1, + waves_per_eu: int = 3, + b_nt: int = 0, + gate_mode: str = "separated", + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + enable_bias: bool = False, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Build (and cache) the inline FlyDSL a4w4 stage1 device kernel.""" + return compile_mixed_moe_gemm1( + 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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + act=act, + persist_m=persist_m, + use_async_copy=use_async_copy, + k_batch=k_batch, + waves_per_eu=waves_per_eu, + b_nt=b_nt, + gate_mode=GateMode(gate_mode), + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + enable_bias=enable_bias, + a_scale_one=a_scale_one, + xcd_swizzle=xcd_swizzle, + swiglu_limit=swiglu_limit, + ) + + +def build_moe_stage2_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + accumulate: bool = True, + persist_m: int = 1, + sort_block_m: int = 0, + b_nt: int = 0, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + xcd_swizzle: int = 0, + enable_bias: bool = False, +): + """Build (and cache) the inline FlyDSL a4w4 stage2 device kernel.""" + return compile_mixed_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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=persist_m, + sort_block_m=sort_block_m, + b_nt=b_nt, + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + xcd_swizzle=xcd_swizzle, + enable_bias=enable_bias, + ) + + +def _moe_stage1( + a, + w1, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + w1_scale, + a1_scale, + sorted_weights=None, +): + """Host runner for the inline a4w4 stage1 (fp4/fp4 -> bf16, k_batch=1).""" + token_num = a.shape[0] + E = w1.shape[0] + inter_dim = w1.shape[1] // 2 + model_dim = a.shape[1] * 2 # a_dtype == "fp4": packed 2 values per byte + dev = a.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + out = torch.empty((token_num, topk, inter_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a1_scale.view(-1) if a1_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w1_scale.view(-1) if w1_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + + _sort_block_m = tile_m + _all_blks = sorted_expert_ids.shape[0] + _dense_blks = ( + min(token_num * topk * _sort_block_m, sorted_token_ids.shape[0]) + // _sort_block_m + ) + _grid_y = min(_dense_blks, _all_blks) + + out_scale_sorted_flat = torch.empty(0, dtype=torch.uint8, device=dev) + _n_in = inter_dim * 2 + _k_in = model_dim + + args = _s1_args_fp4( + out.view(-1), + a.view(-1), + w1.view(-1), + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + out_scale_sorted_flat.view(-1), + token_num, + _n_in, + _k_in, + _grid_y, + dev, + bias=torch.empty(0, device=dev), + ) + + exe = build_moe_stage1_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=(sorted_weights is not None), + a_dtype="fp4", + b_dtype="fp4", + out_dtype=out_dtype, + act="silu", + persist_m=1, + k_batch=1, + waves_per_eu=3, + b_nt=0, + gate_mode="separated", + ) + _run_compiled(exe, args) + return out + + +def _moe_stage2( + inter_states, + w2, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + mode, + w2_scale, + a2_scale, + sorted_weights, +): + """Host runner for the inline a4w4 stage2 (fp4/fp4 -> bf16, atomic).""" + token_num = inter_states.shape[0] + E = w2.shape[0] + model_dim = w2.shape[1] + inter_dim = inter_states.shape[2] * 2 # a_dtype == "fp4" + accumulate = mode != "reduce" + dev = inter_states.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + alloc_fn = torch.zeros if accumulate else torch.empty + out = alloc_fn((token_num, model_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a2_scale.view(-1) if a2_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w2_scale.view(-1) if w2_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(sorted_token_ids.shape, dtype=torch.float32, device=dev) + ) + + m_blocks = min(sorted_expert_ids.shape[0], token_num * topk) + _persist_m = -1 if m_blocks > 256 else 1 + + _n_in = model_dim + _k_in = inter_dim + target = out + if not accumulate: + target = torch.empty( + (token_num * topk * model_dim,), device=out.device, dtype=out.dtype + ) + + args = _s2_args_fp4( + target, + inter_states, + w2, + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + token_num, + _n_in, + _k_in, + m_blocks, + dev, + bias=None, + ) + + exe = build_moe_stage2_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=(sorted_weights is not None), + a_dtype="fp4", + b_dtype="fp4", + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=_persist_m, + sort_block_m=0, + b_nt=0, + ) + _run_compiled(exe, args) + + if not accumulate: + torch.sum(target.view(token_num, topk, model_dim), dim=1, out=out) + return out + + +def flydsl_moe_a4w4( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + block_m: int = 32, + tile_n: int = 256, + tile_k: int = 256, + mode: str = "atomic", +) -> torch.Tensor: + """Run the inline FlyDSL a4w4 MoE stage1+stage2 path. Returns [T, model_dim] bf16.""" + import aiter + from aiter import QuantType, dtypes + from aiter.fused_moe import moe_sorting + from aiter.ops.shuffle import ( + shuffle_scale_a16w4, + shuffle_weight, + shuffle_weight_a16w4, + ) + from aiter.utility.fp4_utils import e8m0_shuffle, moe_mxfp4_sort + + experts = w1.shape[0] + inter_dim = w1.shape[1] // 2 + model_dim = w1.shape[2] + token = hidden_states.shape[0] + topk = topk_ids.shape[1] + torch_dtype = hidden_states.dtype + out_dtype = "bf16" if torch_dtype == torch.bfloat16 else "f16" + + topk_ids = topk_ids.to(torch.int32).contiguous() + topk_weights = topk_weights.to(torch.float32).contiguous() + + # --- harness-side prep (NOT the kernel): mxfp4 quant, e8m0 scales, + # weight/scale preshuffle, sorted token dispatch --- + q_dtype = dtypes.fp4x2 + torch_quant = aiter.get_torch_quant(QuantType.per_1x32) + + w1_qt, w1_scale = torch_quant(w1.contiguous(), quant_dtype=q_dtype) + w2_qt, w2_scale = torch_quant(w2.contiguous(), quant_dtype=q_dtype) + w1_qt = w1_qt.view(experts, inter_dim * 2, model_dim // 2) + w2_qt = w2_qt.view(experts, model_dim, inter_dim // 2) + a1_qt, a1_scale = torch_quant(hidden_states.contiguous(), quant_dtype=q_dtype) + + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, _ = moe_sorting( + topk_ids, topk_weights, experts, model_dim, torch_dtype, block_m + ) + + w1_qt_shuf = shuffle_weight(w1_qt, (16, 16)) + w2_qt_shuf = shuffle_weight_a16w4(w2_qt, 16, False) + w1_scale_shuf = e8m0_shuffle(w1_scale) + w2_scale_shuf = shuffle_scale_a16w4(w2_scale, experts, False) + a1_scale_sort = moe_mxfp4_sort( + a1_scale[:token, :].view(token, 1, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage1 gate/up GEMM + fused gated activation === + stage1_out = _moe_stage1( + a1_qt, + w1_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + w1_scale=w1_scale_shuf, + a1_scale=a1_scale_sort, + sorted_weights=None, + ) + torch.cuda.synchronize() + + a2_qt, a2_scale = torch_quant(stage1_out.view(-1, inter_dim), quant_dtype=q_dtype) + a2_qt = a2_qt.view(token, topk, -1) + a2_scale_sort = moe_mxfp4_sort( + a2_scale[: token * topk, :].view(token, topk, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage2 down GEMM + weighted top-k combine === + out = _moe_stage2( + a2_qt, + w2_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + mode=mode, + w2_scale=w2_scale_shuf, + a2_scale=a2_scale_sort, + sorted_weights=sorted_weights, + ) + torch.cuda.synchronize() + return out diff --git a/tasks/torch2flydsl/moe_kernel/model.py b/tasks/torch2flydsl/moe_kernel/model.py new file mode 100644 index 00000000..47f37a1a --- /dev/null +++ b/tasks/torch2flydsl/moe_kernel/model.py @@ -0,0 +1,211 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Pure-PyTorch reference for the quantized a4w4 fused MoE. + +The op is a top-k Mixture-of-Experts feed-forward block evaluated in MXFP4 +(``float4_e2m1fn_x2`` values with e8m0 per-1x32 block scales). A softmax router +selects ``topk`` experts per token; stage 1 runs a grouped gate/up GEMM followed +by ``silu(gate) * up``; stage 2 runs the down GEMM and combines the experts with +the renormalized router weights. All quantized GEMMs accumulate in fp32 over +dequantized operands. + +Activations and weights are quantized to MXFP4 with e8m0 block scales, the +stage-1 result is re-quantized to MXFP4 before the down GEMM, and the output is +returned in bf16. The MXFP4 (f32->e2m1 rounding, saturation, denormals) and +e8m0 block-scale numerics implemented here match AMD's reference quantizer +bit-for-bit so the dequantized values are identical to the hardware path. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +# MXFP4 (e2m1) decode table indexed by the 4-bit code (sign in bit 3). +_MXFP4_VALUES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) +_BLOCK = 32 +# log2(F4E2M1_MAX=6) floored -> dtypeMax = 2**2 used as the e8m0 scale divisor. +_FP4_DTYPE_MAX = 4.0 + + +def _f32_to_e8m0(x): + """Round positive fp32 magnitudes to biased e8m0 exponents (uint8).""" + u32 = x.contiguous().view(torch.int32) + exponent = ((u32 >> 23) & 0xFF).view(torch.uint32).to(torch.uint8) + nan_case = exponent == 0xFF + round_case = ((u32 & 0x400000) > 0) & ( + ((u32 & 0x200000) > 0) | ((u32 & 0x1FFFFF) > 0) | (exponent > 0) + ) + exponent[round_case] += 1 + exponent[nan_case] = 0xFF + return exponent + + +def _e8m0_to_f32(scale_e8m0_biased): + """Decode biased e8m0 exponents (uint8) back to fp32 power-of-two scales.""" + scale_e8m0_biased = scale_e8m0_biased.view(torch.uint8) + zero_case = scale_e8m0_biased == 0 + nan_case = scale_e8m0_biased == 0xFF + scale_f32 = scale_e8m0_biased.to(torch.int32) << 23 + scale_f32[zero_case] = 0x00400000 + scale_f32[nan_case] = 0x7F800001 + return scale_f32.view(torch.float32) + + +def _f32_to_e2m1_codes(x): + """Round fp32 values to MXFP4 (e2m1) 4-bit codes, saturating out-of-range + magnitudes and handling denormals (adapted from the torchao FP utilities).""" + EBITS, MBITS = 2, 1 + EBITS_F32, MBITS_F32 = 8, 23 + F32_EXP_BIAS = (1 << (EBITS_F32 - 1)) - 1 + exp_bias = (1 << (EBITS - 1)) - 1 + max_int = (1 << (EBITS + MBITS)) - 1 + sign_mask = 1 << (EBITS + MBITS) + magic_adder = (1 << (MBITS_F32 - MBITS - 1)) - 1 + max_normal = 2 ** ((1 << EBITS) - 1 - exp_bias) * ( + ((1 << (MBITS + 1)) - 1) / (2**MBITS) + ) + min_normal = 2 ** (1 - exp_bias) + denorm_exp = (F32_EXP_BIAS - exp_bias) + (MBITS_F32 - MBITS) + 1 + denorm_mask_int = denorm_exp << MBITS_F32 + denorm_mask_float = torch.tensor( + denorm_mask_int, dtype=torch.int32 + ).view(torch.float32) + + x = x.float().view(torch.int32) + sign = x & 0x80000000 + x = x ^ sign + x = x.view(torch.float) + + saturate_mask = x >= max_normal + denormal_mask = torch.logical_and( + torch.logical_not(saturate_mask), x < min_normal + ) + normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask)) + + denormal_x = x + denorm_mask_float + denormal_x = denormal_x.view(torch.int32) + denormal_x -= denorm_mask_int + denormal_x = denormal_x.to(torch.uint8) + + normal_x = x.view(torch.int32) + mant_odd = (normal_x >> (MBITS_F32 - MBITS)) & 1 + val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder + normal_x += val_to_add + normal_x += mant_odd + normal_x = normal_x >> (MBITS_F32 - MBITS) + normal_x = normal_x.to(torch.uint8) + + codes = torch.full_like(x, max_int, dtype=torch.uint8) + codes = torch.where(denormal_mask, denormal_x, codes) + codes = torch.where(normal_mask, normal_x, codes) + + sign_lp = sign >> (MBITS_F32 + EBITS_F32 - MBITS - EBITS) + sign_lp = sign_lp.to(torch.uint8) & sign_mask + return (codes | sign_lp).to(torch.uint8) + + +def _mxfp4_dequant(x): + """MXFP4 per-1x32 e8m0 quantize+dequantize over the last dim, returning the + fp32 values the hardware GEMM sees.""" + shape = x.shape + xb = x.float().reshape(-1, _BLOCK) + max_abs = torch.amax(torch.abs(xb), dim=1) + scale_e8m0 = _f32_to_e8m0(max_abs / _FP4_DTYPE_MAX) + scale_f32 = _e8m0_to_f32(scale_e8m0).view(-1, 1) + codes = _f32_to_e2m1_codes(xb / scale_f32) + table = _MXFP4_VALUES.to(x.device) + deq = table[codes.long()] * scale_f32 + return deq.reshape(shape) + + +def _grouped_gemm_stage1(acts, weights, topk_ids): + """Per-expert grouped GEMM: out[b, k] = acts[b] @ weights[topk_ids[b, k]].T.""" + acts = acts.float() + B, D = acts.shape + topk = topk_ids.shape[1] + N = weights.shape[1] + h = acts.view(B, 1, D).repeat(1, topk, 1) + out = torch.zeros(B, topk, N, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = h[mask] @ weights[e].transpose(0, 1) + return out + + +def _grouped_gemm_stage2(acts, weights, topk_ids, topk_weights): + """Per-expert down GEMM with weighted top-k combine to a single output row.""" + acts = acts.float() + B, topk = topk_ids.shape + model_dim = weights.shape[1] + out = torch.zeros(B, topk, model_dim, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = acts[mask] @ weights[e].transpose(0, 1) + out = out * topk_weights.view(B, topk, 1) + return out.sum(1) + + +def route_topk(logits, topk): + """Softmax router + top-k with renormalized weights. Shared by the harness. + + Ties are broken by ascending expert index via a stable descending sort. The + bf16 gate produces many duplicate logits across the large expert count, and a + nondeterministic top-k tie-break would let the reference and the runtime op + select different experts; the stable order keeps both routings identical. + """ + gate = torch.softmax(logits.float(), dim=-1) + order = torch.sort(gate, dim=-1, descending=True, stable=True).indices + ids = order[..., :topk] + weights = torch.gather(gate, -1, ids) + weights = weights / weights.sum(dim=-1, keepdim=True) + return weights.float(), ids.to(torch.int32) + + +class Model(nn.Module): + def __init__(self, model_dim, inter_dim, experts, topk, activation="silu"): + super().__init__() + self.model_dim = model_dim + self.inter_dim = inter_dim + self.experts = experts + self.topk = topk + self.activation = activation + self.gate = nn.Linear(model_dim, experts, bias=False).to(torch.bfloat16) + self.w1 = nn.Parameter( + (torch.randn(experts, 2 * inter_dim, model_dim) / 10).to(torch.bfloat16) + ) + self.w2 = nn.Parameter( + (torch.randn(experts, model_dim, inter_dim) / 10).to(torch.bfloat16) + ) + + def forward(self, hidden_states): + I = self.inter_dim + + logits = self.gate(hidden_states) + topk_weights, topk_ids = route_topk(logits, self.topk) + + a1 = _mxfp4_dequant(hidden_states) + w1 = _mxfp4_dequant(self.w1) + w2 = _mxfp4_dequant(self.w2) + + stage1 = _grouped_gemm_stage1(a1, w1, topk_ids) + gate, up = stage1.split([I, I], dim=-1) + stage1 = (F.silu(gate) * up).to(torch.bfloat16) + + a2 = _mxfp4_dequant(stage1.reshape(-1, I)).reshape( + hidden_states.shape[0], self.topk, I + ) + + out = _grouped_gemm_stage2(a2, w2, topk_ids, topk_weights) + return out.to(torch.float16).to(torch.bfloat16) + + +def get_inputs(): + return [torch.randn(16, 7168, dtype=torch.bfloat16)] + + +def get_init_inputs(): + return [7168, 256, 257, 9] diff --git a/tasks/torch2flydsl/moe_kernel/test_kernel_harness.py b/tasks/torch2flydsl/moe_kernel/test_kernel_harness.py new file mode 100644 index 00000000..aaf2cf7a --- /dev/null +++ b/tasks/torch2flydsl/moe_kernel/test_kernel_harness.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Correctness and performance harness for the a4w4 fused MoE task. + +The pure-torch reference in ``model.py`` and the FlyDSL kernel share the same +top-k routing (``model.route_topk``) so expert selection is identical. The +correctness gate is the normalized max error ``max|ref - out| / max|ref|``, which +must stay <= ``REL_TOL``; element-wise close% at 1e-2 and 1e-1 is also reported. +The check asserts and exits non-zero on failure. + +Modes: + --correctness compare the kernel against the reference + --full-benchmark time the kernel vs the reference and write a perf report +""" +import argparse +import importlib.util +import json +import math +import os +import sys +from pathlib import Path + +KERNEL_FILE = "kernel.py" +MODEL_FILE = "model.py" + + +def _resolve_kernel_dir(): + here = os.path.dirname(os.path.abspath(__file__)) + if os.path.isfile(os.path.join(here, KERNEL_FILE)): + return here + cwd = os.getcwd() + if os.path.isfile(os.path.join(cwd, KERNEL_FILE)): + return cwd + return here + + +def _load_module(kernel_dir, filename, alias): + entry = os.path.join(kernel_dir, filename) + if not os.path.isfile(entry): + return None + if kernel_dir not in sys.path: + sys.path.insert(0, kernel_dir) + spec = importlib.util.spec_from_file_location(alias, entry) + if spec is None or spec.loader is None: + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +_KERNEL_DIR = _resolve_kernel_dir() + +# Real a4w4 fp4 fused-MoE shapes (q_dtype_a/w = float4_e2m1fn_x2, per_1x32): +# dsv3_fp4_untuned_fmoe.csv (DeepSeek-V3): D=7168, I=256, E=257, topk=9 +# kimik2_fp4_untuned_fmoe.csv (Kimi-K2): D=7168, I=256, E=384, topk=8 +SHAPES = [ + {"name": "dsv3_t16_e257_k9", "tokens": 16, "model_dim": 7168, "inter_dim": 256, "experts": 257, "topk": 9}, + {"name": "kimik2_t32_e384_k8", "tokens": 32, "model_dim": 7168, "inter_dim": 256, "experts": 384, "topk": 8}, +] + +# Tight element-wise gate: normalized max error <= REL_TOL. +REL_TOL = 1e-2 +SEED = 20260401 +BLOCK_M, TILE_N, TILE_K, MODE = 32, 256, 256, "atomic" +# Correctness uses the deterministic "reduce" combine. The "atomic" stage-2 +# combine sums per-expert partials with order-dependent fp32 atomic-adds, so its +# result is nondeterministic run-to-run; "reduce" computes the identical math +# with a deterministic reduction (same numerics, reproducible comparison). +CORRECTNESS_MODE = "reduce" + + +def _build_model(mmod, shape, device="cuda"): + import torch + + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + model = mmod.Model( + model_dim=shape["model_dim"], inter_dim=shape["inter_dim"], + experts=shape["experts"], topk=shape["topk"], + ).to(device).eval() + hidden = torch.randn( + shape["tokens"], shape["model_dim"], dtype=torch.bfloat16, device=device + ) + return model, hidden + + +def _kernel_out(kmod, mmod, model, hidden, topk): + # Recompute the SAME routing the reference used and run the FlyDSL kernel. + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + return kmod.flydsl_moe_a4w4( + hidden, model.w1.detach(), model.w2.detach(), topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=CORRECTNESS_MODE, + ) + + +def run_correctness(verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + failures = [] + for shape in SHAPES: + model, hidden = _build_model(mmod, shape) + with torch.no_grad(): + ref = model(hidden).float() + out = _kernel_out(kmod, mmod, model, hidden, shape["topk"]).float() + torch.cuda.synchronize() + + max_abs = (ref - out).abs().max().item() + ref_scale = ref.abs().max().item() + 1e-9 + rel_err = max_abs / ref_scale + max_rel = ((ref - out).abs() / (ref.abs() + 1e-9)).max().item() + pct1e2 = torch.isclose(ref, out, atol=1e-2, rtol=1e-2).float().mean().item() * 100 + pct1e1 = torch.isclose(ref, out, atol=1e-1, rtol=1e-1).float().mean().item() * 100 + ok = rel_err <= REL_TOL + if verbose: + print( + f" {'PASS' if ok else 'FAIL'}: {shape['name']} " + f"(D{shape['model_dim']}/I{shape['inter_dim']}/E{shape['experts']}/k{shape['topk']}) " + f"norm_max_err={rel_err:.5f} (tol={REL_TOL}) " + f"max_abs={max_abs:.4f} max_rel={max_rel:.3f} " + f"close%@1e-2={pct1e2:.2f} @1e-1={pct1e1:.2f}" + ) + if not ok: + failures.append(shape["name"]) + + status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(SHAPES)})" + print(f"Status: {status}") + print(f"correctness: {'pass' if not failures else 'fail'}") + assert not failures, f"correctness FAILED for: {failures}" + return True + + +def run_benchmark(warmup=10, iters=100, verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + latencies, speedups, report = [], [], [] + print(f"{'Config':<24} {'Ref':>10} {'FlyDSL':>10} {'Speedup':>10}") + print("-" * 60) + for idx, shape in enumerate(SHAPES): + model, hidden = _build_model(mmod, shape) + topk = shape["topk"] + with torch.no_grad(): + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + w1, w2 = model.w1.detach(), model.w2.detach() + + def run_kernel(): + return kmod.flydsl_moe_a4w4( + hidden, w1, w2, topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=MODE, + ) + + run_kernel() + torch.cuda.synchronize() + for _ in range(warmup): + run_kernel() + torch.cuda.synchronize() + ktimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); run_kernel(); e.record(); torch.cuda.synchronize() + ktimes.append(s.elapsed_time(e)) + kernel_ms = sum(ktimes) / len(ktimes) + + rtimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); model(hidden); e.record(); torch.cuda.synchronize() + rtimes.append(s.elapsed_time(e)) + ref_ms = sum(rtimes) / len(rtimes) + + speedup = ref_ms / kernel_ms if kernel_ms > 0 else 1.0 + latencies.append(kernel_ms); speedups.append(speedup) + report.append({ + "test_case_id": f"test_case_{idx}", + "execution_time_ms": kernel_ms, + "shape": [shape["tokens"], shape["model_dim"], shape["inter_dim"]], + "params": {k: shape[k] for k in ("tokens", "model_dim", "inter_dim", "experts", "topk")}, + }) + if verbose: + print(f"{shape['name']:<24} {ref_ms:>8.4f}ms {kernel_ms:>8.4f}ms {speedup:>8.2f}x") + del model, hidden + torch.cuda.empty_cache() + + geomean_latency = math.exp(sum(math.log(x) for x in latencies) / len(latencies)) + geomean_speedup = math.exp(sum(math.log(x) for x in speedups) / len(speedups)) + + build_dir = Path(_KERNEL_DIR) / "build" + build_dir.mkdir(exist_ok=True) + with open(build_dir / "performance_report.json", "w") as f: + json.dump(report, f, indent=2) + + print("-" * 60) + print(f"Geometric mean latency: {geomean_latency:.4f} ms") + print(f"Geometric mean speedup: {geomean_speedup:.2f}x") + return {"geomean_latency_ms": geomean_latency, "geomean_speedup": geomean_speedup} + + +if __name__ == "__main__": + try: + import torch as _t + _arch = _t.cuda.get_device_properties(0).gcnArchName.split(":")[0] + except Exception: + _arch = "" + if _arch != "gfx950": + print(f"SKIPPED: gfx950-only task on arch={_arch or 'unknown'} (FP4/MX scaled-MFMA requires CDNA4/gfx950)") + print("correctness: skip") + sys.exit(0) + parser = argparse.ArgumentParser(description="torch2flydsl moe harness") + parser.add_argument("--correctness", action="store_true") + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--full-benchmark", action="store_true") + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + print("=" * 60) + print("torch2flydsl MoE (a4w4, quantized reference)") + print("=" * 60) + + if args.correctness: + try: + run_correctness() + except AssertionError as exc: + print(f"ASSERTION: {exc}") + sys.exit(1) + sys.exit(0) + else: + run_benchmark(warmup=args.warmup, iters=args.iterations) diff --git a/tasks/torch2flydsl/moe_swiglu_kernel/config.yaml b/tasks/torch2flydsl/moe_swiglu_kernel/config.yaml new file mode 100644 index 00000000..90f5bdd5 --- /dev/null +++ b/tasks/torch2flydsl/moe_swiglu_kernel/config.yaml @@ -0,0 +1,23 @@ +source_file_path: +- kernel.py +target_kernel_functions: +- flydsl_moe_swiglu +- build_moe_stage1_module +- build_moe_stage2_module +- compile_mixed_moe_gemm1 +- compile_mixed_moe_gemm2 +compile_command: +- python3 -c "import torch; from kernel import build_moe_stage1_module, build_moe_stage2_module; + build_moe_stage1_module(model_dim=3072, inter_dim=512, experts=128, topk=4, tile_m=32, + tile_n=256, tile_k=256, doweight_stage1=False, a_dtype='fp4', b_dtype='fp4', out_dtype='bf16', + act='swiglu'); build_moe_stage2_module(model_dim=3072, inter_dim=512, experts=128, + topk=4, tile_m=32, tile_n=256, tile_k=256, doweight_stage2=True, a_dtype='fp4', + b_dtype='fp4', out_dtype='bf16'); print('compile ok')" +correctness_command: +- python3 test_kernel_harness.py --correctness +performance_command: +- python3 test_kernel_harness.py --full-benchmark +task_type: torch2flydsl +supported_archs: +- gfx950 +task_result_template: null diff --git a/tasks/torch2flydsl/moe_swiglu_kernel/kernel.py b/tasks/torch2flydsl/moe_swiglu_kernel/kernel.py new file mode 100644 index 00000000..1f1a8dbf --- /dev/null +++ b/tasks/torch2flydsl/moe_swiglu_kernel/kernel.py @@ -0,0 +1,6651 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""FlyDSL a4w4 (MXFP4) two-stage fused MoE kernel with SWIGLU activation. + +Defines the stage-1 (gate/up GEMM with fused clamped SWIGLU activation) and +stage-2 (down GEMM with weighted top-k combine) device kernels in FlyDSL. The +device builders ``compile_mixed_moe_gemm1`` / ``compile_mixed_moe_gemm2`` -- +exposed via ``build_moe_stage1_module`` / ``build_moe_stage2_module`` and driven +by the ``_moe_stage1`` / ``_moe_stage2`` host runners -- are adapted from AITER's +mixed_moe_gemm_2stage path together with its preshuffle pipeline, CShuffle MFMA +epilogue, layout helpers and GateMode enum. + +``flydsl_moe_swiglu`` is the launcher: it takes bf16 weights and a precomputed +routing (so the reference and kernel share identical top-k selection) and returns +a bf16 ``[T, model_dim]`` tensor. Host-side data prep -- MXFP4/e8m0 per_1x32 +quantization, weight/scale pre-shuffle and the sorted token/expert dispatch +(``moe_sorting``) -- uses AITER utilities to shape inputs into the layout the +device kernels consume. +""" +from __future__ import annotations + +import functools +import math as _math +import os +import re +import builtins as _builtins +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Dict, Optional + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf, memref +from flydsl._mlir.dialects.arith import CmpIPredicate +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr import arith as _arith +from flydsl.expr.arith import ArithValue +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + + +# =========================================================================== +# Inlined from aiter/ops/flydsl/moe_common.py :: GateMode +# =========================================================================== +class GateMode(str, Enum): + """Gate/Up computation strategy for stage1 GEMM (see AITER moe_common).""" + + SEPARATED = "separated" + MOCK_GATE_ONLY = "mock_gate_only" + GATE_ONLY = "gate_only" + INTERLEAVE = "interleave" + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/layout_utils.py +# ========================================================================= +def _wrap(v): + """Wrap raw ir.Value in ArithValue for operator overloading compatibility.""" + if isinstance(v, ArithValue): + return v + if isinstance(v, ir.Value): + return ArithValue(v) + return v + + +def _is_pow2(n): + """Return True when *n* is a positive power of two.""" + return n > 0 and (n & (n - 1)) == 0 + + +def _div_pow2(val, divisor): + """Unsigned divide index *val* by a **compile-time** power-of-2 *divisor*. + + Emits ``arith.shrui`` (1 VALU cycle) instead of ``arith.divui`` + (10-15 VALU cycles on CDNA). + """ + shift = _math.log2(divisor) + assert shift == int(shift), f"{divisor} is not a power of 2" + return arith.shrui(val, arith.index(int(shift))) + + +def _mod_pow2(val, modulus): + """Unsigned remainder of index *val* by a **compile-time** power-of-2 *modulus*. + + Emits ``arith.andi`` (1 VALU cycle) instead of ``arith.remui``. + """ + return arith.andi(val, arith.index(modulus - 1)) + + +def _parse_dim(tok): + """Parse a single dimension token: '?' -> None, otherwise int.""" + tok = tok.strip() + return None if tok == "?" else int(tok) + + +def _parse_layout(ly): + """Parse '(s0,s1,...):(d0,d1,...)' -> (shapes, strides) as lists (None for '?').""" + ly_str = str(ly.type) if hasattr(ly, "type") else str(ly) + m = re.search(r"\(([^)]+)\):\(([^)]+)\)", ly_str) + if not m: + return None + shapes = [_parse_dim(s) for s in m.group(1).split(",")] + strides = [_parse_dim(s) for s in m.group(2).split(",")] + return shapes, strides + + +def _has_dynamic_strides(strides): + """Check if any stride is dynamic (None).""" + return any(s is None for s in strides) + + +def idx2crd(idx, layout): + """Decompose flat index into a list of coordinate values. + + For static layouts, computes coordinates with plain arith ops. + Power-of-2 strides/shapes use shift/mask instead of div/rem. + For dynamic layouts, falls back to fx.idx2crd + fx.get. + """ + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + result = fx.idx2crd(idx, layout) + ndims = len(parsed[1]) if parsed else 1 + return [_wrap(fx.get(result, i)) for i in range(ndims)] + + if hasattr(idx, "type") and str(idx.type) != "index": + idx = arith.index_cast(T.index, idx) + shapes, strides = parsed + ndims = len(strides) + + ordered = sorted( + [ + (i, s, sz) + for i, s, sz in _builtins.zip(range(ndims), strides, shapes) + if s != 0 + ], + key=lambda x: x[1], + reverse=True, + ) + coords = [None] * ndims + remaining = idx + for i, stride_val, size_val in ordered: + if stride_val == 1: + c = remaining + elif _is_pow2(stride_val): + c = _div_pow2(remaining, stride_val) + else: + c = remaining / arith.index(stride_val) + if size_val is not None: + if _is_pow2(size_val): + c = _mod_pow2(c, size_val) + else: + c = c % arith.index(size_val) + coords[i] = c + for i in range(ndims): + if coords[i] is None: + coords[i] = remaining + return coords + + +def crd2idx(crd, layout): + """Compute flat index from a coordinate tuple/list. + + For static layouts, computes with plain arith ops. + For dynamic layouts, falls back to fx.crd2idx with fx.make_coord. + """ + if not isinstance(crd, (list, tuple)): + crd = [crd] + parsed = _parse_layout(layout) + + if parsed is None or _has_dynamic_strides(parsed[1]): + # fly.make_coord requires i32/i64, not index + crd_i32 = [] + for c in crd: + cv = c + if isinstance(cv, ArithValue): + cv = cv.ir_value() if hasattr(cv, "ir_value") else cv + if isinstance(cv, ir.Value) and isinstance(cv.type, ir.IndexType): + cv = arith.index_cast(T.i32, cv) + crd_i32.append(cv) + coord_val = fx.make_coord(*crd_i32) + result = fx.crd2idx(coord_val, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = arith.index_cast(T.index, scalar) + return _wrap(scalar) + + _, strides = parsed + result = None + for coord_v, stride_v in _builtins.zip(crd, strides): + if stride_v == 0: + continue + term = coord_v if stride_v == 1 else coord_v * arith.index(stride_v) + result = term if result is None else result + term + return result if result is not None else arith.index(0) + + +def get(int_tuple, mode): + """Extract element at `mode` from a Python list/tuple.""" + return int_tuple[mode] + + +layout_get = get + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_preshuffle_pipeline.py (crd2idx -> _pre_crd2idx) +# ========================================================================= +def _pre_crd2idx(crd, layout): + """crd2idx returning an index-type scalar (unwraps fly.int_tuple).""" + result = fx.crd2idx(crd, layout) + scalar = fx.get_scalar(result) + if isinstance(scalar, ir.Value) and not isinstance(scalar.type, ir.IndexType): + scalar = _arith.IndexCastOp(T.index, scalar).result + return scalar + + +def swizzle_xor16(row, col, k_blocks16): + """XOR-with-row swizzle on the K dimension at 16B granularity. + + Computes: col XOR ((row & (k_blocks16 - 1)) * 16) + + k_blocks16 is always a power of 2 (tile_k_bytes / 16), so use + bitwise AND instead of remui to save ~10 VALU cycles on CDNA. + """ + from flydsl.expr import arith as _swz_arith + + mask = k_blocks16 - _swz_arith.index(1) + rem = _swz_arith.andi(row, mask) + return col ^ (rem * 16) + + +def lds_row_major_idx(row, col, row_stride, base=None): + """Linearize a 2D LDS coordinate with explicit index arithmetic.""" + idx = row * row_stride + col + return idx if base is None else idx + base + + +def split_row_major_2d(index, minor_extent): + """Split a linear row-major index into (major, minor).""" + return index // minor_extent, index % minor_extent + + +def _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx, + *, + elem_type, + vec_elems, + elem_bytes, + offset_in_bytes, + cache_modifier=0, +): + """Load vec_elems elements via buffer_load dwordx[1,2,4] + bitcast.""" + from flydsl.expr import arith as _ld_arith + + elem_size = int(elem_bytes) + load_bytes = int(vec_elems) * elem_size + vec_width = load_bytes // 4 + + if offset_in_bytes: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(2)) + elif elem_bytes == 2: + idx_i32 = _ld_arith.shrui(idx, _ld_arith.index(1)) + else: + idx_i32 = idx + + i32_val = buffer_ops.buffer_load( + rsrc, + idx_i32, + vec_width=vec_width, + dtype=T.i32, + cache_modifier=cache_modifier, + ) + if vec_width == 1: + i32_vec = vector.from_elements(T.vec(1, T.i32), [i32_val]) + else: + i32_vec = i32_val + return vector.bitcast(T.vec(int(vec_elems), elem_type), i32_vec) + + +@dataclass(frozen=True) +class PreshuffleScaleLayout: + """Container returned by `make_preshuffle_scale_layout`. + + The scale layout is ``(c_mn1, c_k1, 4, 16) : (stride_n0, stride_k0, stride_klane, 1)``. + Callers compute flat index directly with plain arith:: + + idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + """ + + layout_scale: object + stride_n0: object + stride_k0: object + stride_klane: object + + +def make_preshuffle_scale_layout( + arith, + *, + c_mn: ir.Value, + c_k: ir.Value, + mn_pack: int = 2, + k_pack: int = 2, + elem_bytes: int = 4, + scale_block_size: int = 32, +) -> PreshuffleScaleLayout: + """Build scale layout matching aiter/CK preshuffle for FP4/FP8 microscale. + + Layout shape: ``(c_mn1, c_k1, 4, 16)`` where + ``c_mn1 = c_mn / 16 / mn_pack`` and ``c_k1 = (c_k / scale_block_size) / 4 / k_pack``. + """ + c16 = fx.Index(16) + c4 = fx.Index(4) + c_k_scale = c_k // fx.Index(scale_block_size) + + c_mn1 = (c_mn // c16) // fx.Index(mn_pack) + c_k1 = (c_k_scale // c4) // fx.Index(k_pack) + if elem_bytes != mn_pack * k_pack: + raise ValueError( + f"elem_bytes of scale must be {mn_pack} * {k_pack}, got {elem_bytes!r}" + ) + + stride_klane = c16 + stride_k0 = c4 * stride_klane + stride_n0 = c_k1 * stride_k0 + + c_mn1_i32 = arith.index_cast(T.i32, c_mn1) + c_k1_i32 = arith.index_cast(T.i32, c_k1) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + + layout_scale = fx.make_layout( + (c_mn1_i32, c_k1_i32, 4, 16), + stride=(stride_n0_i32, stride_k0_i32, stride_klane_i32, 1), + ) + + return PreshuffleScaleLayout( + layout_scale=layout_scale, + stride_n0=stride_n0, + stride_k0=stride_k0, + stride_klane=stride_klane, + ) + + +@dataclass(frozen=True) +class PreshuffleBLayout: + """Container returned by `make_preshuffle_b_layout`.""" + + layout_b: object + kpack_bytes: int + + +def make_preshuffle_b_layout( + arith, + *, + c_n: ir.Value, + c_k: ir.Value, + kpack_bytes: int = 16, + elem_bytes: int = 1, + k_major: bool = False, +) -> PreshuffleBLayout: + """Build B layout matching aiter/CK preshuffle for A8 MFMA kernels. + + When *k_major* is True the block-level order is K-major (``k_blk`` outermost), + matching the ``(0,3,1,4,2,5)`` shuffle permutation. The default N-major + order (``k_major=False``) matches the legacy ``(0,1,3,4,2,5)`` permutation. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + + c16 = fx.Index(16) + c_kpack = fx.Index(kpack_bytes) + + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + c_k_bytes = c_k * arith.constant(int(elem_bytes), index=True) + n0 = c_n // c16 + + c_kpack_elems = ( + c_kpack + if elem_bytes == 1 + else (c_kpack // arith.constant(int(elem_bytes), index=True)) + ) + + stride_nlane = c_kpack_elems + + if k_major: + c32 = fx.Index(32) + c2 = fx.Index(2) + c_k0 = c_k_bytes // c32 + klane_dim = 2 + stride_klane = c16 * stride_nlane + stride_n0 = c2 * stride_klane + stride_k0 = n0 * stride_n0 + else: + c64 = fx.Index(64) + c4 = fx.Index(4) + c_k0 = c_k_bytes // c64 + klane_dim = 4 + stride_klane = c16 * stride_nlane + stride_k0 = c4 * stride_klane + stride_n0 = c_k0 * stride_k0 + + kpack_elems_static = kpack_bytes if elem_bytes == 1 else kpack_bytes // elem_bytes + n0_i32 = arith.index_cast(T.i32, n0) + c_k0_i32 = arith.index_cast(T.i32, c_k0) + stride_n0_i32 = arith.index_cast(T.i32, stride_n0) + stride_k0_i32 = arith.index_cast(T.i32, stride_k0) + stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + stride_nlane_i32 = arith.index_cast(T.i32, stride_nlane) + + stride_b = (stride_n0_i32, stride_k0_i32, stride_klane_i32, stride_nlane_i32, 1) + layout_b = fx.make_layout( + (n0_i32, c_k0_i32, klane_dim, 16, kpack_elems_static), stride_b + ) + return PreshuffleBLayout(layout_b=layout_b, kpack_bytes=kpack_bytes) + + +def _unpack_int4_to_int8_pair(packed32): + """Split packed int4 dword into two int8 dwords (even/odd nibbles). + + 7-op bit manipulation shared by all int4 unpack paths (W4A8, W4A16, W4A_FP8). + """ + c_08 = fx.Int32(0x08080808) + c_0f = fx.Int32(0x0F0F0F0F) + c_1e = fx.Int32(0x1E) + c_4 = fx.Int32(4) + s0 = (packed32 & c_08) * c_1e + even = (packed32 & c_0f) | s0 + t = packed32 >> c_4 + s1 = (t & c_08) * c_1e + odd = (t & c_0f) | s1 + return even, odd + + +def _pack_i32_pair_to_i64(lo, hi, vector): + """Pack two i32 values into one i64 via vector bitcast.""" + v2 = vector.from_elements(T.vec(2, T.i32), [lo, hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def _i8x4_in_i32_to_bf16x4_i64(val_i32, arith, vector, scale_val=None): + """Convert one i32 (4 signed int8 bytes) to 4 bf16 packed as i64. + + Uses shift-based f32->bf16 truncation (lshr 16) instead of arith.truncf + which on gfx942 expands to ~5 VALU per element. The shift is exact for + unscaled int8 values and introduces <0.5 ULP error for scaled values. + """ + vec1_i32_t = T.vec(1, T.i32) + vec2_i32 = T.i32x2 + vec4_i8 = T.i8x4 + vec1_i64 = T.vec(1, T.i64) + + v1 = vector.from_elements(vec1_i32_t, [val_i32]) + i8x4 = vector.bitcast(vec4_i8, v1) + + f32_vals = [] + for i in range(4): + val_i8 = vector.extract(i8x4, static_position=[i], dynamic_position=[]) + v = arith.sitofp(T.f32, val_i8) + if scale_val is not None: + v = v * scale_val + f32_vals.append(v) + + c16 = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bits0 = arith.bitcast(T.i32, f32_vals[0]) + bits1 = arith.bitcast(T.i32, f32_vals[1]) + bits2 = arith.bitcast(T.i32, f32_vals[2]) + bits3 = arith.bitcast(T.i32, f32_vals[3]) + i32_lo = (bits0 >> c16) | (bits1 & c_ffff0000) + i32_hi = (bits2 >> c16) | (bits3 & c_ffff0000) + + v2 = vector.from_elements(vec2_i32, [i32_lo, i32_hi]) + v64 = vector.bitcast(vec1_i64, v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def load_b_raw_w4a16( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ku: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 8, +): + """Phase 1 of W4A16 B load: issue buffer_load_dword, return raw packed i32. + + Same address calculation as the int4 unpack path in load_b_pack_k32 + but using ku-based indexing for 2-phase latency hiding. + """ + if kpack_bytes != 8: + raise ValueError(f"W4A16 requires kpack_bytes=8, got {kpack_bytes!r}") + + c64 = fx.Index(64) + half_bytes = kpack_bytes // 2 + c2_idx = fx.Index(2) + c4_idx = fx.Index(4) + + k0_base = base_k // c64 + + k1_layout_offset = ku * 2 + lane_div_32 = lane_div_16 // c2_idx + total_k1 = fx.Index(k1_layout_offset) + lane_div_32 + k0 = k0_base + (total_k1 // c4_idx) + k1_local = total_k1 % c4_idx + lane_odd = lane_div_16 % c2_idx + k2_base = lane_odd * fx.Index(half_bytes) + + coord_pack = (n_blk, k0, k1_local, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + idx_bytes = idx_pack + k2_base + + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + return packed32 + + +def _int4_to_bf16x4_i64_gfx950( + packed32, nibble_offsets, arith, vector, scale_val=None, defer_scale16=False +): + """Convert 4 int4 nibbles to 4 bf16 packed as i64 using gfx950 instructions. + + Uses v_cvt_off_f32_i4_sdwa with byte_sel to avoid per-nibble shifts. + Even nibbles (0,2,4,6) → SDWA BYTE_0/1/2/3 on original src. + Odd nibbles (1,3,5,7) → SDWA BYTE_0/1/2/3 on (src >> 4). + Only 1 shift total instead of 7. + + When defer_scale16=True, the ×16 correction factor for v_cvt_off_f32_i4 is + omitted and must be applied later (e.g. in the epilogue). This saves VALU + in the hot loop and uses v_cvt_pk_bf16_f32 for proper f32→bf16 conversion. + """ + from flydsl.expr import rocdl + from flydsl._mlir.dialects._arith_ops_gen import MulFOp as _MulFOp + + _uw = _arith._to_raw + _av = _arith.ArithValue + + src_even = packed32 + src_odd = packed32 >> fx.Int32(4) + + f32_vals = [] + for nib in nibble_offsets: + byte_idx = nib // 2 + src = src_odd if (nib % 2) else src_even + v = rocdl.cvt_off_f32_i4(src, byte_sel=byte_idx) + f32_vals.append(v) + + if defer_scale16: + # Skip ×16; multiply by scale_val only if groupwise. + if scale_val is not None: + raw_scale = _uw(scale_val) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Use v_cvt_pk_bf16_f32 for proper f32→bf16 (no bit-shift trick needed). + i32_lo = rocdl.cvt_pk_bf16_f32(f32_vals[0], f32_vals[1]) + i32_hi = rocdl.cvt_pk_bf16_f32(f32_vals[2], f32_vals[3]) + else: + c16 = fx.Float32(16.0) + if scale_val is not None: + effective_scale = scale_val * c16 + else: + effective_scale = c16 + raw_scale = _uw(effective_scale) + f32_vals = [_MulFOp(v, raw_scale).result for v in f32_vals] + # Truncate f32→bf16 via bit-shift (exact for scaled int values). + c16_shift = fx.Int32(16) + c_ffff0000 = fx.Int32(0xFFFF0000) + bf16_vals = [arith.bitcast(T.i32, _av(v)) for v in f32_vals] + i32_lo = (bf16_vals[0] >> c16_shift) | (bf16_vals[1] & c_ffff0000) + i32_hi = (bf16_vals[2] >> c16_shift) | (bf16_vals[3] & c_ffff0000) + + v2 = vector.from_elements(T.vec(2, T.i32), [i32_lo, i32_hi]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def unpack_b_w4a16( + packed32, arith, vector, scale_val=None, use_gfx950_cvt=False, defer_scale16=False +): + """Phase 2 of W4A16 B load: unpack int4->int8 + convert int8->bf16. + + Takes raw packed32 from load_b_raw_w4a16 and produces (b0, b1) -- + two i64 values each containing 4 bf16 for one MFMA. + + When use_gfx950_cvt=True, uses v_cvt_off_f32_i4 + v_cvt_pk_bf16_f32 + for ~2x fewer VALU instructions. + + When defer_scale16=True (requires use_gfx950_cvt=True), the ×16 + correction for v_cvt_off_f32_i4 is omitted; caller must apply it + in the epilogue. + """ + if use_gfx950_cvt: + b0 = _int4_to_bf16x4_i64_gfx950( + packed32, + [0, 2, 4, 6], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + b1 = _int4_to_bf16x4_i64_gfx950( + packed32, + [1, 3, 5, 7], + arith, + vector, + scale_val, + defer_scale16=defer_scale16, + ) + return (b0, b1) + even, odd = _unpack_int4_to_int8_pair(packed32) + b0 = _i8x4_in_i32_to_bf16x4_i64(even, arith, vector, scale_val=scale_val) + b1 = _i8x4_in_i32_to_bf16x4_i64(odd, arith, vector, scale_val=scale_val) + return (b0, b1) + + +def load_b_pack_k32( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k: ir.Value, + ki_step: int, + n_blk: ir.Value, + n_intra: ir.Value, + lane_div_16: ir.Value, + elem_type: ir.Type, + kpack_bytes: int = 16, + elem_bytes: int = 1, + unpack_int4: bool = False, +) -> ir.Value: + """Load one B pack for one MFMA(x32) micro-step. + + Returns an i64 Value containing 8 bytes consumed by MFMA. + """ + if kpack_bytes not in (8, 16): + raise ValueError(f"kpack_bytes must be 8 or 16, got {kpack_bytes!r}") + if unpack_int4 and kpack_bytes != 8: + raise ValueError("unpack_int4 requires kpack_bytes=8 (packed int4 layout)") + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + + c64 = fx.Index(64) + base_k_bytes = base_k * arith.constant(int(elem_bytes), index=True) + k0_base = base_k_bytes // c64 + k0 = k0_base + arith.constant(ki_step // 2, index=True) + k1 = lane_div_16 + half_bytes = kpack_bytes // 2 + k2_base = arith.constant((ki_step % 2) * half_bytes, index=True) + + coord_pack = (n_blk, k0, k1, n_intra, fx.Index(0)) + idx_pack = _pre_crd2idx(coord_pack, layout_b) + + if unpack_int4: + idx_bytes = idx_pack + k2_base + b4 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_bytes, + elem_type=elem_type, + vec_elems=4, + elem_bytes=1, + offset_in_bytes=True, + ) + packed32 = vector.extract( + vector.bitcast(T.vec(1, T.i32), b4), + static_position=[0], + dynamic_position=[], + ) + even, odd = _unpack_int4_to_int8_pair(packed32) + return _pack_i32_pair_to_i64(even, odd, vector) + + vec_elems = kpack_bytes // int(elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + b_rsrc, + idx_pack, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=(elem_bytes == 1), + ) + + b_i32x4 = vector.bitcast(T.i32x4, b16) + + half = ki_step % 2 + if half == 0: + d0 = vector.extract(b_i32x4, static_position=[0], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[1], dynamic_position=[]) + else: + d0 = vector.extract(b_i32x4, static_position=[2], dynamic_position=[]) + d1 = vector.extract(b_i32x4, static_position=[3], dynamic_position=[]) + + v2 = vector.from_elements(T.vec(2, T.i32), [d0, d1]) + v64 = vector.bitcast(T.vec(1, T.i64), v2) + return vector.extract(v64, static_position=[0], dynamic_position=[]) + + +def tile_chunk_coord_i32( + arith, + *, + tx_i32_base: ir.Value, + i: int, + total_threads: int, + layout_tile_div4, + chunk_i32: int = 4, +): + """Map (thread, chunk_id) -> (row_local, col_local_i32) for X/A loads.""" + if chunk_i32 not in (1, 2, 4): + raise ValueError(f"chunk_i32 must be one of (1,2,4), got {chunk_i32!r}") + chunk_off_i32 = arith.constant(i * total_threads * chunk_i32, index=True) + tile_idx_i32 = tx_i32_base + chunk_off_i32 + coord_local = fx.idx2crd(tile_idx_i32, layout_tile_div4) + row_local = fx.get(coord_local, 0) + col_local_i32 = fx.get(coord_local, 1) + return row_local, col_local_i32 + + +def buffer_copy_gmem16_dwordx4( + buffer_ops, + vector, + *, + elem_type, + idx_i32: ir.Value, + rsrc, + vec_elems: int = 16, + elem_bytes: int = 1, +): + """Copy 16 bytes from global memory into regs via buffer-load dwordx4 lowering.""" + if int(vec_elems) <= 0: + raise ValueError(f"vec_elems must be > 0, got {vec_elems!r}") + return _buffer_load_vec( + buffer_ops, + vector, + rsrc, + idx_i32, + elem_type=elem_type, + vec_elems=vec_elems, + elem_bytes=elem_bytes, + offset_in_bytes=False, + ) + + +def lds_store_16b_xor16( + arith, + vector, + *, + lds_memref, + vec16_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x4: ir.Value, + elem_bytes: int = 1, +): + """Store one 16B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v16 = vector.bitcast(vec16_ty, vec_part_i32x4) + vector.store(v16, lds_memref, [idx0]) + + +def lds_store_8b_xor16( + arith, + vector, + *, + lds_memref, + vec8_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x2: ir.Value, + elem_bytes: int = 1, +): + """Store one 8B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v8 = vector.bitcast(vec8_ty, vec_part_i32x2) + vector.store(v8, lds_memref, [idx0]) + + +def lds_store_4b_xor16( + arith, + vector, + *, + lds_memref, + vec4_ty, + layout_lds, + row_local: ir.Value, + col_local_i32: ir.Value, + tx_c4: ir.Value, + k_blocks16: ir.Value, + lds_base: ir.Value, + vec_part_i32x1: ir.Value, + elem_bytes: int = 1, +): + """Store one 4B chunk into LDS with CK-style XOR16 swizzle on the K dimension.""" + if elem_bytes not in (1, 2): + raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") + col_local_bytes = col_local_i32 * tx_c4 + col_swz_bytes = swizzle_xor16(row_local, col_local_bytes, k_blocks16) + col_swz = col_swz_bytes if elem_bytes == 1 else col_swz_bytes // 2 + coord_store = (row_local, col_swz) + idx0 = _pre_crd2idx(coord_store, layout_lds) + lds_base + v4 = vector.bitcast(vec4_ty, vec_part_i32x1) + vector.store(v4, lds_memref, [idx0]) + + +def lds_load_pack_k32( + arith, + vector, + *, + lds_memref, + layout_lds, + k_blocks16: ir.Value, + curr_row_a_lds: ir.Value, + col_base: ir.Value, + half: int, + lds_base: ir.Value, + ck_lds128: bool, + vec16_ty, + vec8_ty, + vec2_i64_ty, + vec1_i64_ty, +): + """Load one i64 A-pack for an MFMA K32 micro-step from LDS.""" + col_base_swz = swizzle_xor16(curr_row_a_lds, col_base, k_blocks16) + if ck_lds128: + coord_a16 = (curr_row_a_lds, col_base_swz) + idx_a16 = _pre_crd2idx(coord_a16, layout_lds) + lds_base + loaded_a16 = vector.load_op(vec16_ty, lds_memref, [idx_a16]) + a_vec128 = vector.bitcast(vec2_i64_ty, loaded_a16) + return vector.extract(a_vec128, static_position=[half], dynamic_position=[]) + else: + col_swizzled = col_base_swz + (half * 8) + coord_a = (curr_row_a_lds, col_swizzled) + idx_a = _pre_crd2idx(coord_a, layout_lds) + lds_base + loaded_a8 = vector.load_op(vec8_ty, lds_memref, [idx_a]) + a_vec64 = vector.bitcast(vec1_i64_ty, loaded_a8) + return vector.extract(a_vec64, static_position=[0], dynamic_position=[]) + + +def xcd_remap_bx_by( + bx, + by, + c_m, + *, + tile_m: int, + tile_n: int, + N: int, + xcd_swizzle: int, + num_xcds: int = 8, +): + """Remap (bx, by) for L2-cache reuse via XCD swizzle. + + No-op when ``xcd_swizzle <= 0``. Otherwise: + 1. Linearize the original (bx, by) grid round-robin across ``num_xcds`` + XCDs so that contiguous workgroup ids stay on the same XCD. + 2. Re-tile that 1-D order with an M-major group of size ``xcd_swizzle``, + folding the tail group when ``gy`` does not divide evenly. + + Designed to be called inside a ``@flyc.kernel`` immediately after:: + + bx = gpu.block_id("x") + by = gpu.block_id("y") + bx, by = xcd_remap_bx_by(bx, by, c_m, tile_m=..., tile_n=..., N=..., + xcd_swizzle=xcd_swizzle) + + ``c_m`` is the dynamic ``fx.Index`` for runtime ``M``; ``tile_m``, + ``tile_n``, ``N`` and ``xcd_swizzle`` are compile-time Python ints. + """ + if xcd_swizzle <= 0: + return bx, by + + _c1 = fx.arith.constant(1, index=True) + _c_tm = fx.arith.constant(tile_m, index=True) + _gx = fx.arith.constant(N // tile_n, index=True) + _gy = (c_m + _c_tm - _c1) / _c_tm + + _linear_id = bx * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = fx.arith.constant(num_xcds, index=True) + _q = _num_wgs / _c_xcds + _r = _num_wgs % _c_xcds + _xcd = _linear_id % _c_xcds + _in_xcd = _linear_id / _c_xcds + _xcd_lt_r = fx.arith.cmpi(CmpIPredicate.ult, _xcd, _r) + _clip = fx.arith.select(_xcd_lt_r, _xcd, _r) + _wgid = _xcd * _q + _clip + _in_xcd + + _c_wgm = fx.arith.constant(xcd_swizzle, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = fx.arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = fx.arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + new_bx = _first_pid_m + (_wgid_in_group % _group_size_m) + new_by = _wgid_in_group / _group_size_m + return new_bx, new_by + + +__all__ = [ + "PreshuffleBLayout", + "PreshuffleScaleLayout", + "buffer_copy_gmem16_dwordx4", + "lds_load_pack_k32", + "lds_row_major_idx", + "lds_store_4b_xor16", + "lds_store_8b_xor16", + "lds_store_16b_xor16", + "make_preshuffle_b_layout", + "make_preshuffle_scale_layout", + "load_b_pack_k32", + "load_b_raw_w4a16", + "unpack_b_w4a16", + "load_b_raw_w4a16_groupwise", + "unpack_b_w4a16_groupwise", + "extract_bf16_scale", + "split_row_major_2d", + "swizzle_xor16", + "tile_chunk_coord_i32", + "xcd_remap_bx_by", +] + + +# --------------------------------------------------------------------------- +# Groupwise scale load helper (shared by W4A16 and W4A8 groupwise paths) +# --------------------------------------------------------------------------- + + +def _load_groupwise_scale( + buffer_ops, + arith, + *, + scale_rsrc, + expert_offset, + n_blk, + n_intra, + k_pos, + num_groups: int, + group_size: int, + n_per_expert: int, + scale_dtype=None, +): + """Load one per-group scale value from the scale buffer. + + Computes the linear index into the scale tensor from expert offset, + N position, and group index derived from ``k_pos``. + + For bf16 scales the tensor uses ``(E, G//2, N, 2)`` layout — two + adjacent groups for the same N position are packed into one dword. + We load the raw i32 dword (no extraction) so it can be carried as + loop state without register copies. Use :func:`extract_bf16_scale` + in the compute phase to obtain the f32 value. + """ + c16 = fx.Index(16) + n_global = n_blk * c16 + n_intra + c_group_size = fx.Index(group_size) + c_npe = fx.Index(n_per_expert) + group_idx = k_pos // c_group_size + if scale_dtype is None: + scale_dtype = T.f32 + + if scale_dtype == T.bf16: + # (E, G//2, N, 2) layout: dword at [e, pair, n] holds bf16 scales + # for groups 2*pair and 2*pair+1. + pair_idx = group_idx >> fx.Index(1) # group_idx // 2 + # Dword index: same flat formula but with G//2 groups + num_pairs = num_groups // 2 + c_npm1 = fx.Index(num_pairs - 1) + dword_base = expert_offset * c_npm1 + n_global + dword_elem = dword_base + pair_idx * c_npe + dword_idx = arith.index_cast(T.i32, dword_elem) + # Return raw i32 dword — extraction deferred to compute phase. + scale_val = buffer_ops.buffer_load( + scale_rsrc, dword_idx, vec_width=1, dtype=T.i32 + ) + else: + # (E, G, N) layout with f32 dtype + c_gm1 = fx.Index(num_groups - 1) + base_scale = expert_offset * c_gm1 + n_global + elem_idx = base_scale + group_idx * c_npe + scale_idx_i32 = arith.index_cast(T.i32, elem_idx) + scale_val = buffer_ops.buffer_load( + scale_rsrc, scale_idx_i32, vec_width=1, dtype=T.f32 + ) + return scale_val + + +def extract_bf16_scale(arith, scale_raw_i32, ku: int): + """Extract f32 scale from raw i32 dword loaded by bf16 groupwise path. + + In the ``(E, G//2, N, 2)`` layout two adjacent groups share one dword. + ``ku`` determines which half: even ku → low bf16, odd ku → high bf16. + """ + if ku % 2 == 0: + # Low bf16: shift left by 16 to place in upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 << fx.Int32(16)) + else: + # High bf16: mask upper 16 bits → f32 + return arith.bitcast(T.f32, scale_raw_i32 & fx.Int32(0xFFFF0000)) + + +# --------------------------------------------------------------------------- +# W4A16 groupwise load / unpack helpers +# --------------------------------------------------------------------------- + + +def load_b_raw_w4a16_groupwise( + buffer_ops, + arith, + vector, + *, + arg_b, + b_rsrc, + layout_b, + base_k, + ku: int, + n_blk, + n_intra, + lane_div_16, + elem_type, + scale_rsrc, + expert_offset, + num_groups: int, + group_size: int, + n_per_expert: int, + kpack_bytes: int = 8, + scale_dtype=None, +): + """Phase 1 of W4A16 groupwise B load: buffer_loads for weight + scale. + + Reuses :func:`load_b_raw_w4a16` for the weight load, then issues an + additional ``buffer_load_dword`` for the per-group scale. + + Returns ``(packed32, scale_val)``. + """ + packed32 = load_b_raw_w4a16( + buffer_ops, + arith, + vector, + arg_b=arg_b, + b_rsrc=b_rsrc, + layout_b=layout_b, + base_k=base_k, + ku=ku, + n_blk=n_blk, + n_intra=n_intra, + lane_div_16=lane_div_16, + elem_type=elem_type, + kpack_bytes=kpack_bytes, + ) + k_pos = base_k + fx.Index(ku * 32) + scale_val = _load_groupwise_scale( + buffer_ops, + arith, + scale_rsrc=scale_rsrc, + expert_offset=expert_offset, + n_blk=n_blk, + n_intra=n_intra, + k_pos=k_pos, + num_groups=num_groups, + group_size=group_size, + n_per_expert=n_per_expert, + scale_dtype=scale_dtype, + ) + return (packed32, scale_val) + + +def unpack_b_w4a16_groupwise(packed32, scale_val, arith, vector, use_gfx950_cvt=False): + """Phase 2 of W4A16 groupwise: unpack + scale + convert to bf16.""" + return unpack_b_w4a16( + packed32, arith, vector, scale_val=scale_val, use_gfx950_cvt=use_gfx950_cvt + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mfma_epilogues.py (_if_then -> _epi_if_then) +# ========================================================================= +@contextmanager +def _epi_if_then(if_op, scf): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def default_epilog( + *, + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + body_row: Callable, +): + """Iterate the standard MFMA 16x16 row mapping and call `body_row(...)`. + + The mapping matches the common MFMA fragment layout used across kernels in this repo. + + Args: + arith: flydsl arith ext module. + range_constexpr: compile-time unrolled range helper. + m_repeat: tile_m // 16 (python int). + lane_div_16: index Value (0..3). + bx_m: base row (index Value). For MoE, this is the base sorted-row for the tile. + body_row: callback invoked as: + body_row(mi=, ii=, row_in_tile=, row=) + """ + bx_m_v = bx_m + lane_div_16_mul4 = lane_div_16 * 4 + ii_idx_list = [fx.Index(ii) for ii in range(4)] + + for mi in range_constexpr(m_repeat): + mi_base = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off = lane_div_16_mul4 + ii_idx_list[ii] + row_in_tile = mi_base + row_off + row = bx_m_v + row_in_tile + body_row(mi=mi, ii=ii, row_in_tile=row_in_tile, row=row) + + +def c_shuffle_epilog( + *, + arith, + vector, + gpu, + scf=None, + range_constexpr, + # Tile params + tile_m: int, + tile_n: int, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + m_repeat: int, + num_acc_n: int, + # Thread mapping inputs + tx, + lane_div_16, + lane_mod_16, + bx_m, + by_n, + n_tile_base, + # LDS buffer (f16 view, row-major [tile_m, tile_n] flattened) + lds_out, + # Element type for LDS loads (defaults to f16). Pass bf16 to support bf16 epilogues. + frag_elem_type: ir.Type | None = None, + # Callbacks + write_row_to_lds: Callable, + precompute_row: Callable | None = None, + store_pair: Callable, + # When LDS overflows, split lds_out across two buffers by wave-group. + # Pass the second buffer here; first buffer is `lds_out`. + lds_out_split=None, + # Row offset in lds_out for 8-wave mode (MLIR index value). + # Shifts both write and read LDS indices by lds_row_offset * tile_n elements. + lds_row_offset=None, +): + """LDS CShuffle epilogue skeleton. + + Call pattern: + - `write_row_to_lds(...)` is called once per MFMA row produced by this thread. + It is responsible for writing all ni columns for that row into `lds_out`. + - `store_pair(...)` is called for each (row_local, col_pair0) half2 after shuffle. + + `store_pair` can implement either global stores or atomics. + """ + if int(block_size) <= 0 or (int(block_size) % int(cshuffle_nlane)) != 0: + raise ValueError( + f"block_size ({block_size}) must be divisible by cshuffle_nlane ({cshuffle_nlane})" + ) + cshuffle_mlane = int(block_size) // int(cshuffle_nlane) + if (int(tile_m) % cshuffle_mlane) != 0: + raise ValueError( + f"tile_m must be divisible by CShuffleMLane ({cshuffle_mlane}), got tile_m={tile_m}" + ) + if int(e_vec) <= 0: + raise ValueError(f"e_vec must be positive, got {e_vec}") + if (int(tile_n) % (int(cshuffle_nlane) * int(e_vec))) != 0: + raise ValueError( + f"tile_n must be divisible by (CShuffleNLane*EVec) = {cshuffle_nlane*e_vec}, got tile_n={tile_n}" + ) + + # ===================== Split-LDS mode (early return) ===================== + # When lds_out_split is provided, waves are divided into two groups: + # Group A (waves 0..N/2-1) uses lds_out, columns [0, tile_n/2) + # Group B (waves N/2..N-1) uses lds_out_split, columns [tile_n/2, tile_n) + # Each group writes/reads independently; same barriers synchronise all waves. + if lds_out_split is not None: + if scf is None: + raise ValueError("scf module is required for split-LDS cshuffle") + + _half_n = int(tile_n) // 2 + _half_threads = int(block_size) // 2 + EVec = int(e_vec) + + CShuffleNLane_s = min(int(cshuffle_nlane), _half_n // EVec) + if _half_threads % CShuffleNLane_s != 0: + raise ValueError( + f"half_threads={_half_threads} not divisible by CShuffleNLane_split={CShuffleNLane_s}" + ) + CShuffleMLane_s = _half_threads // CShuffleNLane_s + if int(tile_m) % CShuffleMLane_s != 0: + raise ValueError( + f"tile_m={tile_m} not divisible by CShuffleMLane_split={CShuffleMLane_s}" + ) + m_reps_s = int(tile_m) // CShuffleMLane_s + n_reps_s = _half_n // (CShuffleNLane_s * EVec) + + _half_n_idx = arith.constant(_half_n, index=True) + _half_thr_idx = arith.constant(_half_threads, index=True) + _zero_idx = arith.constant(0, index=True) + + _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + + # -- write phase (all waves, each to its group's LDS buffer) -- + n_tile_base_v = n_tile_base + col_base_local_a = n_tile_base_v + lane_mod_16 + col_base_local_b = col_base_local_a - _half_n_idx + + def _write_row_split(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * _half_n_idx + _if_g = scf.IfOp(_is_group_b, has_else=True) + with ir.InsertionPoint(_if_g.then_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_b, + num_acc_n=num_acc_n, + lds_out=lds_out_split, + ) + scf.YieldOp([]) + with ir.InsertionPoint(_if_g.else_block): + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_a, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + scf.YieldOp([]) + + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row_split, + ) + gpu.barrier() + + # -- read phase (each group reads from its own LDS buffer) -- + tx_local = tx - arith.select(_is_group_b, _half_thr_idx, _zero_idx) + c_nlane_s = arith.constant(CShuffleNLane_s, index=True) + m_lane_s = tx_local / c_nlane_s + n_lane_s = tx_local % c_nlane_s + c_evec = arith.constant(EVec, index=True) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + _precomputed_rows_s = [] + for mr in range_constexpr(m_reps_s): + row_base_m = arith.constant(mr * CShuffleMLane_s, index=True) + row_local = row_base_m + m_lane_s + row = bx_m_v + row_local + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + _precomputed_rows_s.append((row_local, row, row_ctx, row_pred)) + + for mr in range_constexpr(m_reps_s): + row_local, row, row_ctx, row_pred = _precomputed_rows_s[mr] + + def _do_store_row_split(): + row_base_lds = row_local * _half_n_idx + for nr in range_constexpr(n_reps_s): + col_base_nr = arith.constant( + nr * (CShuffleNLane_s * EVec), index=True + ) + col_pair0_local = col_base_nr + (n_lane_s * c_evec) + lds_idx = row_base_lds + col_pair0_local + + _if_ld = scf.IfOp(_is_group_b, [vec_frag], has_else=True) + with ir.InsertionPoint(_if_ld.then_block): + fb = vector.load_op(vec_frag, lds_out_split, [lds_idx]) + scf.YieldOp([fb]) + with ir.InsertionPoint(_if_ld.else_block): + fa = vector.load_op(vec_frag, lds_out, [lds_idx]) + scf.YieldOp([fa]) + frag = _if_ld.results[0] + + col_pair0 = col_pair0_local + arith.select( + _is_group_b, _half_n_idx, _zero_idx + ) + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row_split() + else: + _do_store_row_split() + + return # split path complete + + # ===================== Standard (non-split) path below ===================== + + # ---------------- Step 1: write C tile to LDS (row-major, fp16) ---------------- + tile_n_idx = arith.constant(int(tile_n), index=True) + n_tile_base_v = n_tile_base + col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) + + _lds_row_base_offset = ( + lds_row_offset * tile_n_idx if lds_row_offset is not None else None + ) + + def _write_row(mi: int, ii: int, row_in_tile, row): + row_base_lds = row_in_tile * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + + # Ensure all LDS reads finished before the lds write. + gpu.barrier() + default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=_write_row, + ) + + # Ensure all LDS writes are visible before the shuffle-read. + gpu.barrier() + + # ---------------- Step 2: shuffle mapping + half2 store/atomic ---------------- + CShuffleNLane = int(cshuffle_nlane) + CShuffleMLane = int(cshuffle_mlane) + EVec = int(e_vec) + + m_reps_shuffle = int(tile_m) // CShuffleMLane + n_reps_shuffle = int(tile_n) // (CShuffleNLane * EVec) + + c_nlane = fx.Index(CShuffleNLane) + m_lane = tx // c_nlane + n_lane = tx % c_nlane + c_evec = fx.Index(EVec) + + if frag_elem_type is None: + frag_elem_type = T.f16 + vec_frag = T.vec(EVec, frag_elem_type) + bx_m_v = bx_m + by_n_v = by_n + + # Batch-precompute all row contexts (sorted_idx loads) before the store loop. + # This issues all buffer_load instructions upfront so the compiler can pipeline + # them instead of serializing each load with s_waitcnt vmcnt(0). + _precomputed_rows = [] + for mr in range_constexpr(m_reps_shuffle): + row_base_m = arith.constant(mr * CShuffleMLane, index=True) + row_local = row_base_m + m_lane + row = bx_m_v + row_local + + row_ctx_raw = ( + precompute_row(row_local=row_local, row=row) + if precompute_row is not None + else None + ) + + # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)` and `scf` + # is provided, we can skip the entire N-loop for invalid rows (cheaper than per-store checks). + row_ctx = row_ctx_raw + row_pred = None + if ( + scf is not None + and row_ctx_raw is not None + and isinstance(row_ctx_raw, tuple) + and len(row_ctx_raw) == 2 + ): + row_ctx, row_pred = row_ctx_raw + + _precomputed_rows.append((row_local, row, row_ctx, row_pred)) + + # Now perform LDS reads and stores using the pre-fetched row contexts. + for mr in range_constexpr(m_reps_shuffle): + row_local, row, row_ctx, row_pred = _precomputed_rows[mr] + + def _do_store_row(): + row_base_lds = row_local * tile_n_idx + if _lds_row_base_offset is not None: + row_base_lds = row_base_lds + _lds_row_base_offset + for nr in range_constexpr(n_reps_shuffle): + col_base_nr = arith.constant(nr * (CShuffleNLane * EVec), index=True) + col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile + + lds_idx_pair = row_base_lds + col_pair0 + frag = vector.load_op(vec_frag, lds_out, [lds_idx_pair]) + + store_pair( + row_local=row_local, + row=row, + row_ctx=row_ctx, + col_pair0=col_pair0, + col_g0=by_n_v + col_pair0, + frag=frag, + ) + + if row_pred is not None: + _if_row = scf.IfOp(row_pred) + with _epi_if_then(_if_row, scf): + _do_store_row() + else: + _do_store_row() + + +def mfma_epilog( + *, + use_cshuffle: bool, + # Common (always required) + arith, + range_constexpr, + m_repeat: int, + lane_div_16, + bx_m, + # Default epilog (required when use_cshuffle=False) + body_row: Callable | None = None, + # CShuffle epilog (required when use_cshuffle=True) + vector=None, + gpu=None, + scf=None, + tile_m: int | None = None, + tile_n: int | None = None, + e_vec: int = 2, + cshuffle_nlane: int = 32, + block_size: int = 256, + num_acc_n: int | None = None, + tx=None, + lane_mod_16=None, + by_n=None, + n_tile_base=None, + lds_out=None, + write_row_to_lds: Callable | None = None, + precompute_row: Callable | None = None, + store_pair: Callable | None = None, + frag_elem_type: ir.Type | None = None, +): + if not use_cshuffle: + if body_row is None: + raise ValueError("mfma_epilog(use_cshuffle=False) requires `body_row`.") + return default_epilog( + arith=arith, + range_constexpr=range_constexpr, + m_repeat=m_repeat, + lane_div_16=lane_div_16, + bx_m=bx_m, + body_row=body_row, + ) + + return c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=int(tile_m), + tile_n=int(tile_n), + e_vec=int(e_vec), + cshuffle_nlane=int(cshuffle_nlane), + block_size=int(block_size), + m_repeat=m_repeat, + num_acc_n=int(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=frag_elem_type, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + +# ========================================================================= +# Inlined from aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +# ========================================================================= +@contextmanager +def _if_then(if_op): + """Compat helper for SCF IfOp then-region across old/new Python APIs.""" + with ir.InsertionPoint(if_op.then_block): + try: + yield if_op.then_block + finally: + blk = if_op.then_block + if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): + scf.YieldOp([]) + + +def _barrier(vmcnt=63, lgkmcnt=63): + """Emit s_waitcnt + s_barrier via inline asm. + + Bypasses LLVM SIInsertWaitcnts which would insert a conservative + s_waitcnt vmcnt(0) lgkmcnt(0) before every S_BARRIER MI. + """ + parts = [] + needs_waitcnt = vmcnt < 63 or lgkmcnt < 63 + if needs_waitcnt: + wc = [] + if vmcnt < 63: + wc.append(f"vmcnt({vmcnt})") + if lgkmcnt < 63: + wc.append(f"lgkmcnt({lgkmcnt})") + parts.append("s_waitcnt " + " ".join(wc)) + parts.append("s_barrier") + llvm.InlineAsmOp( + res=None, + operands_=[], + asm_string="\n".join(parts), + constraints="", + has_side_effects=True, + is_align_stack=False, + ) + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm1( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + act: str = "silu", + use_cshuffle_epilog: bool | None = None, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 1, + use_async_copy: bool = False, + waves_per_eu: int = 4, + k_batch: int = 1, + b_nt: int = 0, + gate_mode: GateMode = GateMode.SEPARATED, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Compile stage1 kernel (gate+up with silu/swiglu). + + GEMM: act(X @ W_gate.T, X @ W_up.T) -> [tokens*topk, inter_dim] + Direct store (no atomic). When k_batch>1 (split-K), each CTA + computes a K-slice and atomically adds gate/up partials. + Note: persist_m=1 (no persistence) is optimal for stage1 because K=model_dim + is large, so each CTA is already compute-heavy. persist_m>1 serializes M blocks + that the GPU can process in parallel. + + gate_mode controls the gate/up computation strategy — see GateMode enum. + """ + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + sort_block_m = max(32, tile_m) + num_waves = min(4, tile_n // 32) + total_threads = num_waves * 64 + pack_M = 1 if tile_m < 32 else 2 + n_per_wave = tile_n // num_waves + pack_N = min(2, n_per_wave // 16) + pack_K = 2 + scale_mn_pack = 2 + elem_bytes = 1 + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + if (tile_k_bytes % 64) != 0: + raise ValueError(f"tile_k_bytes must be divisible by 64, got {tile_k_bytes}") + + out_s = str(out_dtype).strip().lower() + out_is_f32 = out_s in ("f32", "fp32", "float") + out_is_bf16 = out_s in ("bf16", "bfloat16") + is_int4 = b_dtype == "int4" + is_int8 = False + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + mock_gate_only = gate_mode is GateMode.MOCK_GATE_ONLY + gate_up_interleave = gate_mode is GateMode.INTERLEAVE + gate_only = gate_mode is GateMode.GATE_ONLY + + # Padding semantics: model_dim and inter_dim INCLUDE padding. + # model_dim = model_dim_true + model_dim_pad (K direction) + # inter_dim = inter_dim_true + inter_dim_pad (N direction) + # Tensor sizes use the padded dimensions (inter_dim, model_dim). + # Padding only affects kernel internal logic and grid computation. + _inter_dim_valid = inter_dim - inter_dim_pad + + # Split-K validation + _is_splitk = k_batch > 1 + if mock_gate_only and not _is_splitk: + raise ValueError("mock_gate_only requires k_batch > 1 (split-K)") + 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}" + + out_dtype = "bf16" + else: + _k_per_batch = model_dim + _k_dim = _k_per_batch + + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + f"tile_m*tile_k*elem_bytes must be divisible by {total_threads}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if use_cshuffle_epilog is None: + _use_cshuffle_epilog = os.environ.get("FLIR_MOE_STAGE1_CSHUFFLE", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + else: + _use_cshuffle_epilog = bool(use_cshuffle_epilog) + + _need_fp4 = out_dtype == "fp4" + _need_fp8 = out_dtype == "fp8" + _need_quant = _need_fp4 or _need_fp8 + _need_sort = _need_quant + + if _need_quant: + _use_cshuffle_epilog = True + + _fp4q_tag = "_fp4q" if _need_fp4 else "" + _fp8q_tag = "_fp8q" if _need_fp8 else "" + _sort_tag = "_sort" if _need_sort else "" + _async_tag = "_async" if use_async_copy else "" + _sk_tag = f"_sk{k_batch}" if _is_splitk else "" + _go_tag = "_go" if mock_gate_only else "" + _gui_tag = "_gui" if gate_up_interleave else "" + _as1_tag = "_as1" if a_scale_one else "" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe1_silu_mul_a{a_dtype}_w{b_dtype}_{out_s}" + f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{_fp4q_tag}{_fp8q_tag}{_sort_tag}{_async_tag}{_sk_tag}{_go_tag}{_gui_tag}{_as1_tag}{_xcd_tag}_v32" + ).replace("-", "_") + + # -- LDS sizing -- + _cshuffle_elem_bytes = 4 if _need_quant else (4 if out_is_f32 else 2) + _single_x_bytes = int(tile_m) * int(lds_stride) * int(a_elem_bytes) + lds_out_bytes = ( + _cshuffle_elem_bytes * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + # Determine whether we need wave-group split for lds_out. + # Standard layout: pong = max(input, lds_out) + tid, ping = input. + # When this overflows, split lds_out into two halves across pong & ping. + _GLOBAL_ALIGN = 1024 + _std_pong = max(_single_x_bytes, lds_out_bytes) + lds_tid_bytes + _std_ping = _single_x_bytes + _std_pong_aligned = allocator_pong._align(_std_pong, 128) + _std_total = allocator_pong._align( + _std_pong_aligned, _GLOBAL_ALIGN + ) + allocator_pong._align(_std_ping, 128) + _lds_limit = {"gfx950": 163840, "gfx942": 65536}.get(gpu_arch, 0) + + _split_lds_out = ( + _lds_limit > 0 + and lds_out_bytes > 0 + and _std_total > _lds_limit + and num_waves >= 2 + ) + + if _split_lds_out: + _half_out_bytes = _cshuffle_elem_bytes * int(tile_m) * (int(tile_n) // 2) + _pong_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + _ping_buffer_bytes = max(_single_x_bytes, _half_out_bytes) + else: + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if waves_per_eu is not None and waves_per_eu >= 1: + _total_cu_lds = 160 * 1024 + _min_lds = _total_cu_lds // (waves_per_eu + 1) + 1 + _pong_sz = allocator_pong._align(allocator_pong.ptr, 128) + _ping_sz = allocator_ping._align(allocator_ping.ptr, 128) + _cur_lds = _pong_sz + _ping_sz + if _cur_lds < _min_lds: + allocator_ping.ptr += _min_lds - _cur_lds + + kpack_bytes = 8 if is_int4 else 16 + out_elem_bytes = 4 if out_is_f32 else 2 + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * (2 * inter_dim) * model_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * (2 * inter_dim) * 4 + + _e_vec_s1 = min(tile_n // 32, 8) + if _need_quant: + _e_vec_s1 = max(2, _e_vec_s1) + _num_threads_per_quant_blk_s1 = 32 // _e_vec_s1 + _shuffle_dists_s1 = [] + _sh_val = 1 + while _sh_val < _num_threads_per_quant_blk_s1: + _shuffle_dists_s1.append(_sh_val) + _sh_val *= 2 + _num_shuffle_steps_s1 = len(_shuffle_dists_s1) + + # ---- Unified pipeline schedule (outside @flyc.kernel) ---- + # Each scheduling phase is a dict: + # mfma: [(k_idx, mi_idx, ikxdl, imxdl, asv_idx), ...] + # a_reads: [(k, mi), ...] # A ds_read subtiles + # b_loads: [('gate'/'up', ku, ni), ...] # B VMEM loads + # has_scale: bool # A/B scale VMEM loads + _pipe_m_repeat = tile_m // 16 + _pipe_k_unroll = tile_k_bytes // 128 + _pipe_k_unroll_packed = _pipe_k_unroll // pack_K + _pipe_m_repeat_packed = _pipe_m_repeat // pack_M + _pipe_num_acc_n = n_per_wave // 16 + + # A ds_read groups: group by mi (same mi, all k values together) + _pipe_a_groups = [] + for _mi in range(_pipe_m_repeat): + _grp = [] + for _k in range(_pipe_k_unroll): + _grp.append((_k, _mi)) + if len(_grp) == 2: + _pipe_a_groups.append(_grp) + _grp = [] + if _grp: + _pipe_a_groups.append(_grp) + + # B VMEM loads: individual gate/up loads + _pipe_b_loads = [] + for ku in range(_pipe_k_unroll): + for ni in range(_pipe_num_acc_n): + _pipe_b_loads.append(("gate", ku, ni)) + if not mock_gate_only and not gate_up_interleave: + _pipe_b_loads.append(("up", ku, ni)) + + # MFMA order: B-major (fix B, cycle all A tiles before next B) + # Each entry: one (k, ni) pair; the compute function loops over all mi. + # This keeps B operands (from VMEM) fixed while cycling A (from LDS, no wait). + _pipe_num_acc_n_packed = _pipe_num_acc_n // pack_N + _pipe_all_mfma = [] + for _ku128 in range(_pipe_k_unroll_packed): + for _ni_packed in range(_pipe_num_acc_n_packed): + for _ikxdl in range(pack_K): + for _inxdl in range(pack_N): + _k_idx = _ku128 * pack_K + _ikxdl + _ni_idx = _ni_packed * pack_N + _inxdl + _pipe_all_mfma.append((_k_idx, _ni_idx, _ikxdl, _inxdl, _ku128)) + + # Group MFMAs per scheduling phase (wider M -> more MFMAs per phase) + _pipe_mfma_per_phase = max(1, len(_pipe_all_mfma) // 4) + _pipe_n_phases = len(_pipe_all_mfma) // _pipe_mfma_per_phase + + # Build unified phase descriptors + _a_groups_per_phase = (len(_pipe_a_groups) + _pipe_n_phases - 1) // _pipe_n_phases + _pipe_phases = [] + _mfma_i = 0 + _a_i = 0 + for _p in range(_pipe_n_phases): + _a_reads = [] + for _ in range(_a_groups_per_phase): + if _a_i < len(_pipe_a_groups): + _a_reads.extend(_pipe_a_groups[_a_i]) + _a_i += 1 + _phase = { + "mfma": _pipe_all_mfma[_mfma_i : _mfma_i + _pipe_mfma_per_phase], + "a_reads": _a_reads, + "b_loads": [], + "has_scale": (_p == 0), + } + _mfma_i += _pipe_mfma_per_phase + _pipe_phases.append(_phase) + + # Distribute B loads evenly across phases 1..n-1 (phase 0 has scales) + _bi = 0 + for _p in range(1, _pipe_n_phases): + _rem_b = len(_pipe_b_loads) - _bi + _rem_p = _pipe_n_phases - _p + _n_b = (_rem_b + _rem_p - 1) // _rem_p if _rem_p > 0 else 0 + for _ in range(_n_b): + if _bi < len(_pipe_b_loads): + _pipe_phases[_p]["b_loads"].append(_pipe_b_loads[_bi]) + _bi += 1 + + # Extract flat lists for kernel access (avoids dict access in AST rewriter) + _pp_mfma = [p["mfma"] for p in _pipe_phases] + _pp_a_reads = [p["a_reads"] for p in _pipe_phases] + _pp_b_loads = [p["b_loads"] for p in _pipe_phases] + _pp_has_scale = [p["has_scale"] for p in _pipe_phases] + + fp4_ratio = 2 if a_dtype == "fp4" else 1 + gui_ratio = 1 if gate_up_interleave else 2 + _vmcnt_before_barrier = tile_m // 32 // fp4_ratio + tile_n // 32 * gui_ratio + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = arith.constant_vector(0.0, vec4_f32) + + # --- Stage1 dimension mapping --- + # X: [tokens, model_dim] -- M = sorted tokens, K = model_dim + # W: [E*2*inter_dim, model_dim] gate portion -- N = inter_dim + # Out: [tokens*topk, inter_dim] + + # B preshuffle layout: [E*2*inter_dim, model_dim] + # Gate rows for expert e: [e*2*inter_dim, e*2*inter_dim + inter_dim) + c_n_total = arith.constant(experts * (2 * inter_dim), index=True) + b_layout = make_preshuffle_b_layout( + arith, + c_n=c_n_total, + c_k=k_in // pack_K, + kpack_bytes=kpack_bytes, + elem_bytes=b_elem_bytes, + # k_major=True, + ) + layout_b = b_layout.layout_b + + # A-scale: [sorted_size, K/32] -- pre-scattered by caller into sorted layout + # Same as stage2: indexed by sorted_row position, not by token_id. + sorted_m = size_expert_ids_in * arith.constant(sort_block_m, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=sorted_m, c_k=arith.constant(model_dim, index=True) + ) + # B-scale: [E*2*inter_dim, K/32] + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=arith.constant(model_dim, index=True) + ) + + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + if const_expr(use_async_copy and a_elem_vec_pack > 1): + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along inter_dim (N) + bx_persist = gpu.block_id("y") # persistent WG index + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S1 = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_idp_sw = arith.constant(2 * inter_dim_pad, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + _gx = (n_in - _c_idp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + else: + _c2_sw = arith.constant(2, index=True) + _gx = ( + (n_in - _c_idp_sw + _c2_sw * _c_tn_sw - _c1_sw) + / _c_tn_sw + / _c2_sw + ) + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S1, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S1 = xcd_swizzle + _c_wgm = arith.constant(_WGM_S1, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + by_n = by * arith.constant(tile_n, index=True) + + k_base_idx = arith.index(0) + if const_expr(_is_splitk): + bz = gpu.block_id("z") # K-batch id + k_base_idx = bz * arith.constant(_k_dim, index=True) + + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((num_waves, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + _lds_out_elem_type = ( + T.f32 if _need_quant else (T.bf16 if out_is_bf16 else T.f16) + ) + if const_expr(_split_lds_out and _use_cshuffle_epilog): + _half_out_elems = int(tile_m) * (int(tile_n) // 2) + lds_out = SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + lds_out_B = SmemPtr( + base_ptr_ping, + lds_ping_offset, + _lds_out_elem_type, + shape=(_half_out_elems,), + ).get() + else: + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + _lds_out_elem_type, + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_out_B = None + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # Buffer resources + c_a_pack = arith.constant(int(a_elem_vec_pack), index=True) + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + + # X: [tokens, model_dim] + x_nbytes_idx = (tokens_in * k_in * c_elem_bytes) / c_a_pack + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # Out: [tokens*topk, inter_dim] + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not (is_f16_a or a_scale_one)): + # A scale: [sorted_size, model_dim/32] pre-scattered by caller + c32 = arith.constant(32, index=True) + kblk = k_in / c32 + sx_nbytes_idx = sorted_m * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + c32 = arith.constant(32, index=True) + kblk_w = k_in / c32 + mn_w = arith.constant(experts * (2 * inter_dim), index=True) + sw_nbytes_idx = mn_w * kblk_w + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + sorted_nbytes_idx = size_expert_ids_in * arith.constant( + sort_block_m * 4, index=True + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + eid_nbytes_idx = size_expert_ids_in * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # Sorted-scale buffer resource for fused mxfp4 quantization + _sorted_scale_cols = inter_dim // 32 + _sorted_scale_cols_i32 = arith.constant(_sorted_scale_cols, type=T.i32) + sorted_scale_rsrc = None + if const_expr(_need_sort): + _sort_rows_idx = size_expert_ids_in * arith.constant( + sort_block_m, index=True + ) + _sort_padded_rows = ( + (_sort_rows_idx + arith.constant(255, index=True)) + / arith.constant(256, index=True) + * arith.constant(256, index=True) + ) + _sort_padded_cols = arith.constant( + ((_sorted_scale_cols + 7) // 8) * 8, index=True + ) + _sort_scale_nbytes = arith.index_cast( + T.i32, _sort_padded_rows * _sort_padded_cols + ) + sorted_scale_rsrc = _ptr_buffer_resource( + arg_out_scale_sorted, _sort_scale_nbytes + ) + + # ---- persist_m loop (same pattern as stage2) ---- + _PERSIST_M = persist_m + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + _c_pm = arith.constant(_PERSIST_M, index=True) + _for_persist = scf.ForOp(_c0_p, _c_pm, _c1_p) + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + bx = bx_persist * _c_pm + _mi_p + bx_m = bx * arith.constant(sort_block_m, index=True) + + # Block validity + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, bx, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + def _moe_gemm1_body(): + # Gate expert offset: first inter_dim rows of each expert's 2*inter_dim block + expert_off_idx = expert_idx * arith.constant(2 * inter_dim, index=True) + + # X loading -- KEY DIFFERENCE from stage2: X row = token_id only + x_load_bytes = 16 + num_x_loads = bytes_per_thread_x // x_load_bytes + chunk_i32 = x_load_bytes // 4 + + c_k_div4 = ( + (k_in / c_a_pack) * arith.constant(int(a_elem_bytes), index=True) + ) / arith.index(4) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(0xFFFFFF) + 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, + ) + + def load_x(idx_i32): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + + # Decode sorted token ids -- stage1: X row = token_id (not t*topk+s) + x_row_base_div4 = [] + x_col_local_i32 = [] + x_row_local = [] + # Also store token_id and slot_id for output indexing + + 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + + # KEY: X row base uses token_id only (not t*topk+s) + t_idx = arith.index_cast(ir.IndexType.get(), t_safe) + x_row_base_div4.append(t_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(a_elem_bytes), index=True) + ) / arith.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) + parts.append(vector.bitcast(T.vec(4, i32), x_vec)) + return parts + + # Wave/lane decomposition (identical to stage2) + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + row_a_lds = lane_mod_16 + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_n_id = wave_id % arith.constant(num_waves, index=True) + n_tile_base = wave_n_id * c_n_per_wave + + # N-tile precompute for gate AND up weights + gate_n_intra_list = [] + gate_n_blk_list = [] + up_n_intra_list = [] + up_n_blk_list = [] + col_g_list = [] + c_n0_static = experts * (2 * inter_dim) // 16 + layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1)) + inter_idx = arith.constant(inter_dim, index=True) + + for i in range_constexpr(num_acc_n): + offset = i * 16 + c_offset = arith.constant(offset, index=True) + if const_expr(not gate_up_interleave): + col_g = by_n + n_tile_base + c_offset + lane_mod_16 + col_g_list.append(col_g) + + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + # Gate/interleave: rows [expert_off, expert_off + 2*inter_dim) + gate_row_w = expert_off_idx + global_n + gate_coord = idx2crd(gate_row_w, layout_n_blk_intra) + gate_n_blk_list.append(layout_get(gate_coord, 0)) + gate_n_intra_list.append(layout_get(gate_coord, 1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_row_w = gate_row_w + inter_idx + up_coord = idx2crd(up_row_w, layout_n_blk_intra) + up_n_blk_list.append(layout_get(up_coord, 0)) + up_n_intra_list.append(layout_get(up_coord, 1)) + + if const_expr(gate_up_interleave): + _gui_num_acc_n_out = num_acc_n // pack_N + for _gui_i in range_constexpr(_gui_num_acc_n_out): + _gui_offset = _gui_i * 16 + _gui_c_offset = arith.constant(_gui_offset, index=True) + _gui_col_g = ( + (by_n + n_tile_base) // arith.constant(2, index=True) + + _gui_c_offset + + lane_mod_16 + ) + col_g_list.append(_gui_col_g) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku = tile_k // k_unroll + _pad_k_elems = ( + (model_dim_pad % tile_k) + if (not _is_splitk and model_dim_pad > 0) + else 0 + ) + _pad_ku_skip = _pad_k_elems // _K_per_ku + _tail_ku = k_unroll - _pad_ku_skip + _tail_ku_packed = ( + (_tail_ku + pack_K - 1) // pack_K if _pad_ku_skip > 0 else None + ) + + # B load for gate and up separately + def load_b_packs_k64(base_k, ku: int, n_blk, n_intra): + c64 = arith.constant(64, index=True) + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0 = base_k_bytes // c64 + arith.constant(ku, index=True) + k1 = lane_div_16 + coord_pack = (n_blk, k0, k1, n_intra, arith.constant(0, index=True)) + idx_pack = crd2idx(coord_pack, layout_b) + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + """Load B tiles. Returns (gate_b_tile, up_b_tile). + When mock_gate_only or gate_up_interleave, up_b_tile is None.""" + gate_b_tile = [] + up_b_tile = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_limit): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + gb0, gb1 = load_b_packs_k64( + base_k, ku, gate_n_blk_list[ni], gate_n_intra_list[ni] + ) + g_packs0.append(gb0) + g_packs1.append(gb1) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + ub0, ub1 = load_b_packs_k64( + base_k, ku, up_n_blk_list[ni], up_n_intra_list[ni] + ) + u_packs0.append(ub0) + u_packs1.append(ub1) + gate_b_tile.append((g_packs0, g_packs1)) + if const_expr(not mock_gate_only and not gate_up_interleave): + up_b_tile.append((u_packs0, u_packs1)) + return gate_b_tile, up_b_tile + + # Pre-compute scale base element indices (K-loop invariant). + # idx = mni * stride_n0 + ku * stride_k0 + k_lane * stride_klane + n_lane + # Split into: base_elem = mni * stride_n0 + lane_elem (invariant) + # k_elem = ku * stride_k0 (per-iteration) + _scale_lane_elem = ( + lane_div_16 * layout_b_scale.stride_klane + lane_mod_16 + ) + + _gate_scale_bases = [] + _up_scale_bases = [] + for _ni in range_constexpr(num_acc_n_packed): + _col_base = ( + by_n + + n_tile_base + + arith.constant(_ni * 16 * pack_N, index=True) + ) + _gate_mni = (expert_off_idx + _col_base) // arith.constant( + 32, index=True + ) + _gate_scale_bases.append( + _gate_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + if const_expr(not mock_gate_only and not gate_up_interleave): + _up_mni = ( + expert_off_idx + inter_idx + _col_base + ) // arith.constant(32, index=True) + _up_scale_bases.append( + _up_mni * layout_b_scale.stride_n0 + _scale_lane_elem + ) + + if const_expr(not a_scale_one): + _a_scale_bases = [] + for _mi in range_constexpr(m_repeat_packed): + _a_mni = _mi + bx_m // scale_mn_pack // 16 + _a_scale_bases.append( + _a_mni * layout_a_scale.stride_n0 + _scale_lane_elem + ) + + _c16_idx = arith.constant(16, index=True) + _c2_idx = arith.constant(2, index=True) + _scale_mask_lo = arith.constant(0xFF, type=T.i32) + + _m_half_idx = arith.constant(0, type=T.i32) + _m_half_i32 = arith.constant(0, type=T.i32) + _scale_shift = arith.constant(0, type=T.i32) + _scale_shift_hi = arith.constant(0, type=T.i32) + _n_half_idx = arith.constant(0, type=T.i32) + _n_half_i32 = arith.constant(0, type=T.i32) + _bscale_shift = arith.constant(0, type=T.i32) + _bscale_shift_hi = arith.constant(0, type=T.i32) + if const_expr(pack_M < scale_mn_pack): + _m_half_idx = (bx_m // _c16_idx) % _c2_idx + _m_half_i32 = arith.index_cast(T.i32, _m_half_idx) + _scale_shift = _m_half_i32 * arith.constant(8, type=T.i32) + _scale_shift_hi = _scale_shift + arith.constant(16, type=T.i32) + + if const_expr(pack_N < scale_mn_pack): + _n_half_idx = (n_tile_base // _c16_idx) % _c2_idx + _n_half_i32 = arith.index_cast(T.i32, _n_half_idx) + _bscale_shift = _n_half_i32 * arith.constant(8, type=T.i32) + _bscale_shift_hi = _bscale_shift + arith.constant(16, type=T.i32) + + def _rearrange_a_scale(raw_i32): + """Rearrange scale bytes for pack_M=1: extract m_half's k0,k1 bytes.""" + if const_expr(pack_M >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _scale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _scale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + def _rearrange_b_scale(raw_i32): + """Rearrange scale bytes for pack_N=1: extract n_half's k0,k1 bytes.""" + if const_expr(pack_N >= scale_mn_pack): + return raw_i32 + b_k0 = arith.andi( + arith.shrui(raw_i32, _bscale_shift), _scale_mask_lo + ) + b_k1 = arith.andi( + arith.shrui(raw_i32, _bscale_shift_hi), _scale_mask_lo + ) + return arith.ori( + b_k0, arith.shli(b_k1, arith.constant(8, type=T.i32)) + ) + + if const_expr(a_scale_one): + _as1_const = arith.constant(0x7F7F7F7F, type=T.i32) + _as1_vec = vector.from_elements(T.vec(1, T.i32), [_as1_const]) + + def prefetch_ab_scale_tile(base_k, ku_packed_limit=k_unroll_packed): + a_scale_tile = [] + gate_b_scale = [] + up_b_scale = ( + [] if (not mock_gate_only and not gate_up_interleave) else None + ) + for ku in range_constexpr(ku_packed_limit): + k_off = (ku + base_k) * layout_b_scale.stride_k0 + for mi in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + a_scale_tile.append(_as1_vec) + else: + s = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[mi] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + s = _rearrange_a_scale(s) + a_scale_tile.append( + vector.from_elements(T.vec(1, T.i32), [s]) + ) + for ni in range_constexpr(num_acc_n_packed): + gs = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + gs = _rearrange_b_scale(gs) + gate_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [gs]) + ) + if const_expr( + not mock_gate_only and not gate_up_interleave + ): + us = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[ni] + k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + us = _rearrange_b_scale(us) + up_b_scale.append( + vector.from_elements(T.vec(1, T.i32), [us]) + ) + return [a_scale_tile, gate_b_scale, up_b_scale] + + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x4=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + if const_expr(use_async_copy): + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(elem_bytes), index=True) + ) / arith.index(4) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def prefetch_full_a_from_lds(lds_buffer, ku_limit=k_unroll): + """Load entire A tile from LDS into registers before compute.""" + a_regs = [] + for k_idx in range_constexpr(ku_limit): + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + for mi_idx in range_constexpr(m_repeat): + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64( + curr_row, col_base + 64, lds_buffer + ) + a_regs.append((a0, a1, a2, a3)) + else: + a_regs.append((a0, a1)) + return a_regs + + # Compute tile: gate + up MFMA interleaved, same A data, different B data. + # Two accumulator sets; after all K tiles, acc = acc_gate + acc_up (f32 add). + def compute_tile( + acc_gate_in, + acc_up_in, + gate_b_tile_in, + up_b_tile_in, + a_tile_regs, + a_scale=None, + gate_b_scale=None, + up_b_scale=None, + *, + prefetch_epilogue=False, + ku_count=k_unroll, + ): + gate_list = list(acc_gate_in) + _single_b = mock_gate_only or gate_up_interleave + up_list = None if _single_b else list(acc_up_in) + mfma_res_ty = vec4_f32 + epilogue_pf = None + bias_pf = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + if const_expr(gate_up_interleave): + bias_pf = [] + for ni in range_constexpr(num_acc_n): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (ni % 2 == 1) + else arith.constant(0, index=True) + ) + bias_offset = ( + expert_off_idx + _up_off + _logical_col + ) + bias_pf.append( + _load_bias_scalar(bias_rsrc, bias_offset) + ) + else: + gate_bias_pf = [] + up_bias_pf = ( + [] if const_expr(not mock_gate_only) else None + ) + for ni in range_constexpr(num_acc_n): + global_n = ( + by_n + + n_tile_base + + arith.constant(ni * 16, index=True) + + lane_mod_16 + ) + gate_bias_pf.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + global_n + ) + ) + if const_expr(not mock_gate_only): + up_bias_pf.append( + _load_bias_scalar( + bias_rsrc, + expert_off_idx + inter_idx + global_n, + ) + ) + bias_pf = (gate_bias_pf, up_bias_pf) + tw_pf = None + if const_expr(doweight_stage1): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + sorted_row_pf = bx_m + mi_base_pf + row_off_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias_pf) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + _eff_packed = (ku_count + pack_K - 1) // pack_K + # B-major: fix B (ni), cycle A (mi) -- B from VMEM stays + # in registers while A from LDS is repacked per mi. + for ku128 in range_constexpr(_eff_packed): + for ni in range_constexpr(num_acc_n_packed): + gate_bs_i32 = gate_b_scale[ku128 * num_acc_n_packed + ni] + gate_bs_val = vector.extract( + gate_bs_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(not _single_b): + up_bs_i32 = up_b_scale[ku128 * num_acc_n_packed + ni] + up_bs_val = vector.extract( + up_bs_i32, static_position=[0], dynamic_position=[] + ) + for ikxdl in range_constexpr(pack_K): + k_idx = ku128 * pack_K + ikxdl + if const_expr(k_idx < ku_count): + gate_bp0, gate_bp1 = gate_b_tile_in[k_idx] + if const_expr(not _single_b): + up_bp0, up_bp1 = up_b_tile_in[k_idx] + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + gb0 = gate_bp0[ni_idx] + gb1 = gate_bp1[ni_idx] + gb128 = pack_i64x4_to_i32x8( + gb0, gb1, c0_i64, c0_i64 + ) + if const_expr(not _single_b): + ub0 = up_bp0[ni_idx] + ub1 = up_bp1[ni_idx] + ub128 = pack_i64x4_to_i32x8( + ub0, ub1, c0_i64, c0_i64 + ) + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ + ku128 * m_repeat_packed + mi + ] + a_scale_val = vector.extract( + a_scale_i32, + static_position=[0], + dynamic_position=[], + ) + for imxdl in range_constexpr(pack_M): + mi_idx = mi * pack_M + imxdl + _a_reg_idx = k_idx * m_repeat + mi_idx + if const_expr(is_f8_a): + a0, a1, a2, a3 = a_tile_regs[ + _a_reg_idx + ] + a128 = pack_i64x4_to_i32x8( + a0, a1, a2, a3 + ) + else: + a0, a1 = a_tile_regs[_a_reg_idx] + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + ) + if const_expr(not _single_b): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + return gate_list, up_list, epilogue_pf + + def load_a_subtile(k_idx, mi_idx, lds_buffer): + """Load a single A sub-tile from LDS (one ds_read).""" + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row = row_a_lds + mi_val + a0, a1 = lds_load_packs_k64(curr_row, col_base, lds_buffer) + if const_expr(is_f8_a): + a2, a3 = lds_load_packs_k64(curr_row, col_base + 64, lds_buffer) + return (a0, a1, a2, a3) + else: + return (a0, a1) + + _single_b_pipe = mock_gate_only or gate_up_interleave + + def compute_bmajor_mfma_phase( + all_a_tiles, + gate_b_single, + up_b_single, + a_scale_vals, + gate_bs_val, + up_bs_val, + gate_list, + up_list, + k_idx, + ni_idx, + ikxdl, + inxdl, + ): + """B-major MFMA: fix one B (ni), cycle all A tiles (mi). + + Packs B once and reuses across all mi iterations. + A tiles come from LDS (already available, no VMEM wait). + + all_a_tiles: flat list indexed by [k*m_repeat + mi]. + gate_b_single/up_b_single: (b0, b1) for one specific ni. + When _single_b_pipe (mock_gate_only or interleave), up_b_single is None. + a_scale_vals: list of A scale scalars indexed by mi_packed. + """ + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def _pack(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + mfma_res_ty = vec4_f32 + gb128 = _pack(gate_b_single[0], gate_b_single[1], c0_i64, c0_i64) + if const_expr(not _single_b_pipe): + ub128 = _pack(up_b_single[0], up_b_single[1], c0_i64, c0_i64) + + for mi_p in range_constexpr(m_repeat_packed): + a_scale_val = a_scale_vals[mi_p] + for imxdl in range_constexpr(pack_M): + mi_idx = mi_p * pack_M + imxdl + a_reg = all_a_tiles[k_idx * m_repeat + mi_idx] + + if const_expr(is_f8_a): + a128 = _pack(a_reg[0], a_reg[1], a_reg[2], a_reg[3]) + else: + a128 = _pack(a_reg[0], a_reg[1], c0_i64, c0_i64) + + acc_idx = mi_idx * num_acc_n + ni_idx + gate_list[acc_idx] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + gb128, + gate_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + gate_bs_val, + ], + ) + if const_expr(not _single_b_pipe): + up_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + ub128, + up_list[acc_idx], + cbsz, + blgp, + ikxdl * pack_M + imxdl, + a_scale_val, + ikxdl * pack_N + inxdl, + up_bs_val, + ], + ) + ) + + def _interleaved_half( + lds_read, + lds_write, + next_k_dma_py, + next_k_load, + prev_a_tile, + prev_gate_w, + prev_up_w, + prev_a_scale, + prev_gate_bs, + prev_up_bs, + acc_gate, + acc_up, + ): + """One flatmm-style interleaved half-iteration (deep pipeline). + + Generalized for arbitrary m_repeat (block_m=32, 64, ...). + DMA targets lds_write (OTHER buffer) while ds_read uses + lds_read (already DMA'd in previous half). + + Interleaving schedule (per half): + Phase 0: scale VMEM + 2 ds_read(A) -> 4 MFMA(prev) + Phase 1..N: B VMEM(distributed) + 2 ds_read(A, if avail) -> 4 MFMA(prev) + Phase N+1..: remaining B VMEM -> 4 MFMA(prev) + """ + _abs_k = k_base_idx + arith.constant(next_k_load, index=True) + _bk = _abs_k // arith.constant(2, index=True) + _sk = _abs_k // arith.constant(pack_K * 128, index=True) + _k_off = _sk * layout_b_scale.stride_k0 + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(_vmcnt_before_barrier) + _barrier() + rocdl.sched_barrier(0) + + # DMA A to OTHER buffer (for next half), non-blocking + _abs_k_dma = k_base_idx + arith.constant(next_k_dma_py, index=True) + if const_expr(use_async_copy and next_k_dma_py < int(_k_dim)): + prefetch_x_to_lds(_abs_k_dma, lds_write) + if const_expr(not use_async_copy): + _x_regs = load_x_tile(_abs_k_dma) + + # ---- Extract previous scale values ---- + _prev_asvs = [] + for _mi_p in range_constexpr(m_repeat_packed): + _prev_asvs.append( + vector.extract( + prev_a_scale[_mi_p], + static_position=[0], + dynamic_position=[], + ) + ) + _prev_gsv_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _prev_gsv_list.append( + vector.extract( + prev_gate_bs[_gs_ni], + static_position=[0], + dynamic_position=[], + ) + ) + if const_expr(not _single_b_pipe): + _prev_usv_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _prev_usv_list.append( + vector.extract( + prev_up_bs[_us_ni], + static_position=[0], + dynamic_position=[], + ) + ) + + # ---- Execute phases from unified schedule ---- + _a_all = {} + _b_gate_all = {} + _b_up_all = {} + + for _p in range_constexpr(_pipe_n_phases): + # Scale VMEM loads (phase 0 only) + if const_expr(_pp_has_scale[_p]): + _new_as_list = [] + for _mi_p in range_constexpr(m_repeat_packed): + if const_expr(a_scale_one): + _new_as_list.append(_as1_const) + else: + _raw_as = buffer_ops.buffer_load( + sx_rsrc, + _a_scale_bases[_mi_p] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_as_list.append(_rearrange_a_scale(_raw_as)) + _new_gs_list = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + _gs_raw = buffer_ops.buffer_load( + sw_rsrc, + _gate_scale_bases[_gs_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_gs_list.append(_rearrange_b_scale(_gs_raw)) + if const_expr(not _single_b_pipe): + _new_us_list = [] + for _us_ni in range_constexpr(num_acc_n_packed): + _us_raw = buffer_ops.buffer_load( + sw_rsrc, + _up_scale_bases[_us_ni] + _k_off, + vec_width=1, + dtype=T.i32, + cache_modifier=0, + ) + _new_us_list.append(_rearrange_b_scale(_us_raw)) + + # B VMEM loads + for _b_j in range_constexpr(len(_pp_b_loads[_p])): + _b_type, _b_ku, _b_ni = _pp_b_loads[_p][_b_j] + if const_expr(_b_type == "gate"): + _b_gate_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + gate_n_blk_list[_b_ni], + gate_n_intra_list[_b_ni], + ) + else: + _b_up_all[(_b_ku, _b_ni)] = load_b_packs_k64( + _bk, + _b_ku, + up_n_blk_list[_b_ni], + up_n_intra_list[_b_ni], + ) + + # A ds_reads + rocdl.sched_barrier(0) + for _a_j in range_constexpr(len(_pp_a_reads[_p])): + _ak, _ami = _pp_a_reads[_p][_a_j] + _a_all[(_ak, _ami)] = load_a_subtile( + _ak, + _ami, + lds_read, + ) + rocdl.sched_barrier(0) + + # MFMAs on prev data + rocdl.s_setprio(1) + for _m_j in range_constexpr(len(_pp_mfma[_p])): + _k_idx, _ni_idx, _ikxdl, _inxdl, _ku128 = _pp_mfma[_p][_m_j] + _ni_packed_idx = _ni_idx // pack_N + _up_b_single = ( + ( + prev_up_w[_k_idx][0][_ni_idx], + prev_up_w[_k_idx][1][_ni_idx], + ) + if not _single_b_pipe + else None + ) + compute_bmajor_mfma_phase( + prev_a_tile, + ( + prev_gate_w[_k_idx][0][_ni_idx], + prev_gate_w[_k_idx][1][_ni_idx], + ), + _up_b_single, + _prev_asvs, + _prev_gsv_list[_ni_packed_idx], + ( + _prev_usv_list[_ni_packed_idx] + if not _single_b_pipe + else None + ), + acc_gate, + acc_up, + _k_idx, + _ni_idx, + _ikxdl, + _inxdl, + ) + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + + # ---- Assemble loaded data for next half-iteration ---- + cur_a_tile = [] + for _k in range_constexpr(k_unroll): + for _mi in range_constexpr(m_repeat): + cur_a_tile.append(_a_all[(_k, _mi)]) + + cur_gate_w = [] + cur_up_w = None if _single_b_pipe else [] + for ku in range_constexpr(k_unroll): + g_packs0, g_packs1 = [], [] + u_packs0, u_packs1 = [], [] + for ni in range_constexpr(num_acc_n): + g = _b_gate_all[(ku, ni)] + g_packs0.append(g[0]) + g_packs1.append(g[1]) + if const_expr(not _single_b_pipe): + u = _b_up_all[(ku, ni)] + u_packs0.append(u[0]) + u_packs1.append(u[1]) + cur_gate_w.append((g_packs0, g_packs1)) + if const_expr(not _single_b_pipe): + cur_up_w.append((u_packs0, u_packs1)) + + cur_a_scale = [] + for _mi_p in range_constexpr(m_repeat_packed): + cur_a_scale.append( + vector.from_elements( + T.vec(1, T.i32), + [_new_as_list[_mi_p]], + ) + ) + cur_gate_bs = [] + for _gs_ni in range_constexpr(num_acc_n_packed): + cur_gate_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_gs_list[_gs_ni]] + ) + ) + if const_expr(not _single_b_pipe): + cur_up_bs = [] + for _us_ni in range_constexpr(num_acc_n_packed): + cur_up_bs.append( + vector.from_elements( + T.vec(1, T.i32), [_new_us_list[_us_ni]] + ) + ) + else: + cur_up_bs = None + + if const_expr(not use_async_copy): + store_x_tile_to_lds(_x_regs, lds_write) + + return ( + cur_a_tile, + cur_gate_w, + cur_up_w, + cur_a_scale, + cur_gate_bs, + cur_up_bs, + acc_gate, + acc_up, + ) + + # Pipeline (split ping/pong allocators) + rocdl.sched_barrier(0) + + k0 = k_base_idx + if const_expr(use_async_copy): + prefetch_x_to_lds(k0, lds_x_pong) + else: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_x_pong) + rocdl.sched_barrier(0) + _k0_scale = k_base_idx // arith.constant(pack_K * 128, index=True) + a_scale_pong, gate_bs_pong, up_bs_pong = prefetch_ab_scale_tile( + _k0_scale + ) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + acc_gate = [acc_init] * num_acc_n * m_repeat + acc_up = ( + [acc_init] * num_acc_n * m_repeat if not _single_b_pipe else None + ) + + _k1 = k_base_idx + arith.constant(tile_k, index=True) + rocdl.sched_barrier(0) + if const_expr(use_async_copy): + prefetch_x_to_lds(_k1, lds_x_ping) + else: + _x_regs_prime = load_x_tile(_k1) + store_x_tile_to_lds(_x_regs_prime, lds_x_ping) + + _k0_b = k_base_idx // arith.constant(2, index=True) + gate_w0, up_w0 = load_b_tile(_k0_b) + # Prime the deep pipeline: DMA K=tile_k -> ping (1 tile ahead) + if const_expr(use_async_copy): + rocdl.s_waitcnt(0) + gpu.barrier() + rocdl.sched_barrier(0) + a_tile_pong = prefetch_full_a_from_lds(lds_x_pong) + + rocdl.sched_barrier(0) + rocdl.s_waitcnt(6) + + num_k_tiles_py = int(_k_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 + + gate_w_pong = gate_w0 + up_w_pong = up_w0 + + rocdl.sched_barrier(0) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + next_k_load_1 = k_iv_py + tile_k + next_k_load_2 = k_iv_py + tile_k * 2 + next_k_dma_1 = k_iv_py + tile_k * 2 + next_k_dma_2 = k_iv_py + tile_k * 3 + + # Half 1: read ping (DMA'd prev half), DMA->pong, MFMA(pong) + ( + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_ping, + lds_x_pong, + next_k_dma_1, + next_k_load_1, + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) + + # Half 2: read pong (DMA'd Half 1), DMA->ping, MFMA(ping) + ( + a_tile_pong, + gate_w_pong, + up_w_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + acc_gate, + acc_up, + ) = _interleaved_half( + lds_x_pong, + lds_x_ping, + next_k_dma_2, + next_k_load_2, + a_tile_ping, + gate_w_ping, + up_w_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + acc_gate, + acc_up, + ) + + # _wave_mod2_b = wave_id % arith.constant(2, index=True) + # _wave_odd = arith.cmpi( + # CmpIPredicate.eq, _wave_mod2_b, arith.constant(1, index=True) + # ) + # _if_wave_odd = scf.IfOp(_wave_odd) + # with ir.InsertionPoint(_if_wave_odd.then_block): + # # gpu.barrier() + # _barrier() + # scf.YieldOp([]) + + if const_expr(odd_k_tiles): + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + else: + _k_tail_rel = arith.constant(_k_dim - tile_k, index=True) + k_tail1 = k_base_idx + _k_tail_rel + x_regs_ping = [] + if const_expr(use_async_copy): + prefetch_x_to_lds(k_tail1, lds_x_ping) + else: + x_regs_ping = load_x_tile(k_tail1) + if const_expr(_pad_ku_skip > 0): + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True), + ku_limit=_tail_ku, + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True), + ku_packed_limit=_tail_ku_packed, + ) + else: + gate_w_ping, up_w_ping = load_b_tile( + k_tail1 // arith.constant(2, index=True) + ) + a_scale_ping, gate_bs_ping, up_bs_ping = prefetch_ab_scale_tile( + k_tail1 // arith.constant(pack_K * 128, index=True) + ) + acc_gate, acc_up, _ = compute_tile( + acc_gate, + acc_up, + gate_w_pong, + up_w_pong, + a_tile_pong, + a_scale_pong, + gate_bs_pong, + up_bs_pong, + ) + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_ping, lds_x_ping) + rocdl.s_waitcnt(0) + _barrier() + if const_expr(_pad_ku_skip > 0): + a_tile_ping = prefetch_full_a_from_lds( + lds_x_ping, ku_limit=_tail_ku + ) + else: + a_tile_ping = prefetch_full_a_from_lds(lds_x_ping) + acc_gate, acc_up, epilogue_pf = compute_tile( + acc_gate, + acc_up, + gate_w_ping, + up_w_ping, + a_tile_ping, + a_scale_ping, + gate_bs_ping, + up_bs_ping, + prefetch_epilogue=True, + ku_count=_tail_ku if _pad_ku_skip > 0 else k_unroll, + ) + + bias_pf = None + if const_expr(epilogue_pf is not None): + _, _, bias_pf = epilogue_pf + + # Activation helpers (f32 element-wise on vec4_f32) + def _silu_elem(g): + """silu(x) = x * sigmoid(x); HW fast path: exp2, rcp""" + neg_log2e = arith.constant(-1.4426950408889634, type=f32) + t = g * neg_log2e + emu = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [t], [], []) + one = arith.constant(1.0, type=f32) + den = one + emu + sig = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den], [], []) + return g * sig + + def _silu_mul_vec4(gate_v4, up_v4): + """Element-wise silu(gate) * up on vec4_f32. + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before applying silu(gate) * up. + """ + result_elems = [] + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + if const_expr(swiglu_limit != 0): + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + result_elems.append(_silu_elem(g) * u) + return vector.from_elements(vec4_f32, result_elems) + + def _swiglu_mul_vec4(gate_v4, up_v4): + """Element-wise swiglu(gate, up) on vec4_f32. + swiglu(g, u) = g * sigmoid(alpha * g) * (u + 1) + When swiglu_limit != 0, clamp gate <= limit and + -limit <= up <= limit before the activation. + """ + result_elems = [] + _alpha = arith.constant(1.702, type=f32) + _one = arith.constant(1.0, type=f32) + _neg_log2e = arith.constant(-1.4426950408889634, type=f32) + if const_expr(swiglu_limit != 0): + _limit = arith.constant(float(swiglu_limit), type=f32) + _neg_limit = arith.constant(-float(swiglu_limit), type=f32) + else: + _limit = arith.constant(float(7.0), type=f32) + _neg_limit = arith.constant(-float(7.0), type=f32) + + for ei in range_constexpr(4): + g = vector.extract( + gate_v4, static_position=[ei], dynamic_position=[] + ) + u = vector.extract( + up_v4, static_position=[ei], dynamic_position=[] + ) + g = arith.minimumf(g, _limit) + u = arith.minimumf(u, _limit) + u = arith.maximumf(u, _neg_limit) + t = g * _alpha * _neg_log2e + emu = llvm.call_intrinsic( + f32, "llvm.amdgcn.exp2.f32", [t], [], [] + ) + den = _one + emu + sig = llvm.call_intrinsic( + f32, "llvm.amdgcn.rcp.f32", [den], [], [] + ) + result_elems.append(g * sig * (u + _one)) + return vector.from_elements(vec4_f32, result_elems) + + def _act_vec4(gate_v4, up_v4): + """Dispatch activation based on `act` parameter.""" + if const_expr(act == "swiglu"): + return _swiglu_mul_vec4(gate_v4, up_v4) + else: + return _silu_mul_vec4(gate_v4, up_v4) + + # Add bias to raw GEMM accumulators before activation. + # bias layout: [E, 2*inter_dim] flat f32 (non-interleaved: gate then up). + # For gate_up_interleave, map physical column to logical bias offset. + if const_expr(enable_bias and not _is_splitk): + _bias_up_vals = None + if const_expr(bias_pf is not None): + if const_expr(gate_up_interleave): + _bias_gate_vals = bias_pf + else: + _bias_gate_vals, _bias_up_vals = bias_pf + else: + _bias_gate_vals = [] + for _ni in range_constexpr(num_acc_n): + if const_expr(gate_up_interleave): + _logical_col = ( + (by_n + n_tile_base) + // arith.constant(2, index=True) + + arith.constant((_ni // 2) * 16, index=True) + + lane_mod_16 + ) + _up_off = ( + inter_idx + if (_ni % 2 == 1) + else arith.constant(0, index=True) + ) + _bias_off = expert_off_idx + _up_off + _logical_col + else: + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_off = expert_off_idx + _bn + _bias_gate_vals.append( + _load_bias_scalar(bias_rsrc, _bias_off) + ) + if const_expr(not (mock_gate_only or gate_up_interleave)): + _bias_up_vals = [] + for _ni in range_constexpr(num_acc_n): + _bn = ( + by_n + + n_tile_base + + arith.constant(_ni * 16, index=True) + + lane_mod_16 + ) + _bias_up_vals.append( + _load_bias_scalar( + bias_rsrc, expert_off_idx + inter_idx + _bn + ) + ) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_gate_vals[_ni]] * 4 + ) + acc_gate[_aidx] = arith.addf(acc_gate[_aidx], _bsplat) + + if const_expr(not (mock_gate_only or gate_up_interleave)): + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + _bsplat = vector.from_elements( + vec4_f32, [_bias_up_vals[_ni]] * 4 + ) + acc_up[_aidx] = arith.addf(acc_up[_aidx], _bsplat) + + if const_expr(gate_up_interleave and not _is_splitk): + _gui_out_n = num_acc_n // pack_N + acc = [None] * (_gui_out_n * m_repeat) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(_gui_out_n): + _g_idx = _mi * num_acc_n + _ni * pack_N + _u_idx = _g_idx + 1 + _out_idx = _mi * _gui_out_n + _ni + acc[_out_idx] = _act_vec4( + acc_gate[_g_idx], acc_gate[_u_idx] + ) + elif const_expr(not _is_splitk): + acc = [None] * (int(num_acc_n) * int(m_repeat)) + for _mi in range_constexpr(m_repeat): + for _ni in range_constexpr(num_acc_n): + _aidx = _mi * num_acc_n + _ni + acc[_aidx] = _act_vec4(acc_gate[_aidx], acc_up[_aidx]) + + # ---- Epilogue: CShuffle + direct store (accumulate=False) ---- + # Output: out[(t*topk+s) * inter_dim + col] = silu(gate) * up + # For split-K: skip silu, output gate/up separately with atomic add + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + _, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + tokens_i32_v = tokens_i32 + + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + if const_expr(lds_out is None): + raise RuntimeError("CShuffle epilogue requires lds_out") + + _apply_weight = doweight_stage1 and not _is_splitk + + 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, + ): + if const_expr(_apply_weight): + 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=f32 + ) + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(_apply_weight): + v = v * tw + if const_expr(_need_quant): + lds_idx = row_base_lds + col_local + vec1_f32 = T.vec(1, f32) + v1 = vector.from_elements(vec1_f32, [v]) + vector.store(v1, lds_out, [lds_idx], alignment=4) + else: + 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, [v_out]) + vector.store(v1, lds_out, [lds_idx], alignment=2) + + _out_row_stride = ( + inter_dim * 2 * out_elem_bytes + if _is_splitk + else ( + inter_dim // 2 + if _need_fp4 + else (inter_dim if _need_fp8 else inter_dim * out_elem_bytes) + ) + ) + + def precompute_row(*, row_local, row): + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32_v) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + row_byte_base = out_base_idx + ts_idx * arith.constant( + _out_row_stride, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + _e_vec = _e_vec_s1 + _e_vec_sk = 2 + _cshuffle_nlane = min(32, tile_n // _e_vec) + _cshuffle_nlane_sk = min(32, tile_n // _e_vec_sk) + _num_threads_per_quant_blk = _num_threads_per_quant_blk_s1 + + _c0_i32 = arith.constant(0, type=T.i32) + _c1_i32 = arith.constant(1, type=T.i32) + _c2_i32 = arith.constant(2, type=T.i32) + _c3_i32 = arith.constant(3, type=T.i32) + _c4_i32 = arith.constant(4, type=T.i32) + _c5_i32 = arith.constant(5, type=T.i32) + _c15_i32 = arith.constant(15, type=T.i32) + _c22_i32 = arith.constant(22, type=T.i32) + _c23_i32 = arith.constant(23, type=T.i32) + _c28_i32 = arith.constant(28, type=T.i32) + _c31_i32 = arith.constant(31, type=T.i32) + _c32_i32 = arith.constant(32, type=T.i32) + _c64_i32 = arith.constant(64, type=T.i32) + _c254_i32 = arith.constant(254, type=T.i32) + _c256_i32 = arith.constant(256, type=T.i32) + _c0xFF800000_i32 = arith.constant(0xFF800000, type=T.i32) + _c0x400000_i32 = arith.constant(0x400000, type=T.i32) + _c0x7FFFFFFF_i32 = arith.constant(0x7FFFFFFF, type=T.i32) + _c0x80000000_i32 = arith.constant(0x80000000, type=T.i32) + _c0x3F800000_i32 = arith.constant(0x3F800000, type=T.i32) # 1.0f + _c0x40C00000_i32 = arith.constant(0x40C00000, type=T.i32) # 6.0f + _c0x4A800000_i32 = arith.constant(0x4A800000, type=T.i32) + _c0xC11FFFFF_i32 = arith.constant(0xC11FFFFF, type=T.i32) + _c0x7_i32 = arith.constant(0x7, type=T.i32) + _c0_f32 = arith.constant(0.0, type=T.f32) + + _c8_i32 = arith.constant(8, type=T.i32) + _fp_headroom = 2 if _need_fp4 else (8 if _need_fp8 else 0) + _c_headroom_i32 = arith.constant(_fp_headroom, type=T.i32) + + def _f32_to_e2m1(qx_f32): + """Convert a scaled f32 value to fp4 (e2m1) 4-bit integer.""" + # Match fp4_utils.f32_to_mxfp4 / HIP quant: saturate, denorm, + # and normal round-to-nearest-even paths. + qx = qx_f32.bitcast(T.i32) + s = qx & _c0x80000000_i32 + qx_abs = qx & _c0x7FFFFFFF_i32 + denormal_mask = arith.cmpi( + CmpIPredicate.ult, qx_abs, _c0x3F800000_i32 + ) + normal_mask = arith.andi( + arith.cmpi(CmpIPredicate.ult, qx_abs, _c0x40C00000_i32), + arith.cmpi(CmpIPredicate.uge, qx_abs, _c0x3F800000_i32), + ) + + denorm_f32 = qx_abs.bitcast(T.f32) + _c0x4A800000_i32.bitcast(T.f32) + denormal_x = denorm_f32.bitcast(T.i32) - _c0x4A800000_i32 + + mant_odd = (qx_abs >> _c22_i32) & _c1_i32 + normal_x = qx_abs + _c0xC11FFFFF_i32 + mant_odd + normal_x = normal_x >> _c22_i32 + + e2m1 = arith.select(normal_mask, normal_x, _c0x7_i32) + e2m1 = arith.select(denormal_mask, denormal_x, e2m1) + return (s >> _c28_i32) | e2m1 + + if const_expr(_need_sort): + _n32_sort = _sorted_scale_cols_i32 * _c32_i32 + + # Mutable slot for split-K N-offset (gate=0, up=inter_dim) + _sk_n_offset = [0] + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(_need_quant and not _is_splitk): + frag_vals = [] + for i in range_constexpr(_e_vec): + frag_vals.append( + vector.extract( + frag, static_position=[i], dynamic_position=[] + ) + ) + + local_max = _c0_f32 + for i in range_constexpr(_e_vec): + abs_v = llvm.call_intrinsic( + f32, "llvm.fabs.f32", [frag_vals[i]], [], [] + ) + local_max = arith.maximumf(local_max, abs_v) + + for _si in range_constexpr(_num_shuffle_steps_s1): + off = arith.constant(_shuffle_dists_s1[_si], type=T.i32) + peer = local_max.shuffle_xor(off, _c64_i32) + local_max = arith.maximumf(local_max, peer) + + max_i32 = local_max.bitcast(T.i32) + # Match fp4_utils.f32_to_e8m0(max_abs / 4): round the + # exponent at the 1.5x threshold before dropping mantissa. + max_rounded = (max_i32 + _c0x400000_i32) & _c0xFF800000_i32 + exp_field = max_rounded >> _c23_i32 + e8m0_biased = arith.maxsi(exp_field - _c_headroom_i32, _c0_i32) + + quant_exp = _c254_i32 - e8m0_biased + quant_scale = (quant_exp << _c23_i32).bitcast(T.f32) + + if const_expr(_need_fp4): + fp4_vals = [] + for i in range_constexpr(_e_vec): + scaled_v = frag_vals[i] * quant_scale + fp4_vals.append(_f32_to_e2m1(scaled_v)) + + packed_i32 = fp4_vals[0] | (fp4_vals[1] << _c4_i32) + for k in range_constexpr(1, _e_vec // 2): + byte_k = fp4_vals[2 * k] | ( + fp4_vals[2 * k + 1] << _c4_i32 + ) + packed_i32 = packed_i32 | ( + byte_k << arith.constant(k * 8, type=T.i32) + ) + + ptr_addr_idx = row_byte_base + col_g0 / arith.constant( + 2, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + _pack_bytes = _e_vec // 2 + if const_expr(_pack_bytes == 1): + store_val = arith.TruncIOp(T.i8, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=1, nontemporal=True + ) + elif const_expr(_pack_bytes == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, out_ptr_v, alignment=2, nontemporal=True + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, out_ptr_v, alignment=4, nontemporal=True + ) + + elif const_expr(_need_fp8): + scaled_vals = [] + for i in range_constexpr(_e_vec): + scaled_vals.append(frag_vals[i] * quant_scale) + + ptr_addr_idx = row_byte_base + col_g0 + if const_expr(_e_vec <= 4): + packed_i32 = _c0_i32 + for _w in range_constexpr(_e_vec // 2): + packed_i32 = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[2 * _w], + scaled_vals[2 * _w + 1], + packed_i32, + _w, + ) + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + if const_expr(_e_vec == 2): + store_val = arith.TruncIOp(T.i16, packed_i32) + store_raw = ( + store_val._value + if hasattr(store_val, "_value") + else store_val + ) + llvm.StoreOp( + store_raw, + out_ptr_v, + alignment=2, + nontemporal=True, + ) + else: + packed_raw = ( + packed_i32._value + if hasattr(packed_i32, "_value") + else packed_i32 + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + else: + for _wg in range_constexpr(_e_vec // 4): + _b = _wg * 4 + packed_w = _c0_i32 + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b], + scaled_vals[_b + 1], + packed_w, + 0, + ) + packed_w = rocdl.cvt_pk_fp8_f32( + T.i32, + scaled_vals[_b + 2], + scaled_vals[_b + 3], + packed_w, + 1, + ) + word_ptr = ptr_addr_idx + arith.constant( + _wg * 4, index=True + ) + out_ptr_v = _idx_to_llvm_ptr(word_ptr) + packed_raw = ( + packed_w._value + if hasattr(packed_w, "_value") + else packed_w + ) + llvm.StoreOp( + packed_raw, + out_ptr_v, + alignment=4, + nontemporal=True, + ) + + if const_expr(_need_sort): + col_g0_i32 = arith.index_cast(T.i32, col_g0) + is_scale_writer = arith.cmpi( + CmpIPredicate.eq, col_g0_i32 & _c31_i32, _c0_i32 + ) + _if_scale = scf.IfOp(is_scale_writer) + with ir.InsertionPoint(_if_scale.then_block): + row_i32_s = arith.index_cast(T.i32, row) + col_s_i32 = col_g0_i32 >> _c5_i32 + d0 = row_i32_s >> _c5_i32 + d1 = (row_i32_s >> _c4_i32) & _c1_i32 + d2 = row_i32_s & _c15_i32 + d3 = col_s_i32 >> _c3_i32 + d4 = (col_s_i32 >> _c2_i32) & _c1_i32 + d5 = col_s_i32 & _c3_i32 + byte_off = ( + d0 * _n32_sort + + d3 * _c256_i32 + + d5 * _c64_i32 + + d2 * _c4_i32 + + d4 * _c2_i32 + + d1 + ) + e8m0_i8 = arith.TruncIOp(T.i8, e8m0_biased) + buffer_ops.buffer_store( + e8m0_i8, + sorted_scale_rsrc, + byte_off, + offset_is_bytes=True, + ) + scf.YieldOp([]) + elif const_expr(_is_splitk): + col_idx = col_g0 + arith.constant(_sk_n_offset[0], index=True) + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec_sk * out_elem_bytes, + ) + else: + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + + _frag_elem = ( + ir.F32Type.get() + if _need_quant + else (ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get()) + ) + + if const_expr(gate_up_interleave and not _is_splitk): + # gui without splitk: acc has activation applied, halved N + _gui_eff_n = _gui_out_n + _gui_tile_n = tile_n // 2 + _gui_cshuffle_nlane = min(32, _gui_tile_n // _e_vec) + _gui_by_n = by_n / arith.constant(2, index=True) + _gui_n_tile_base = n_tile_base / arith.constant(2, index=True) + c_shuffle_epilog( + arith=arith, + vector=vector, + gpu=gpu, + scf=scf, + range_constexpr=range_constexpr, + tile_m=tile_m, + tile_n=_gui_tile_n, + e_vec=_e_vec, + cshuffle_nlane=_gui_cshuffle_nlane, + block_size=total_threads, + m_repeat=m_repeat, + num_acc_n=_gui_eff_n, + tx=tx, + lane_div_16=lane_div_16, + lane_mod_16=lane_mod_16, + bx_m=bx_m, + by_n=_gui_by_n, + n_tile_base=_gui_n_tile_base, + lds_out=lds_out, + frag_elem_type=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + elif const_expr(mock_gate_only or (gate_up_interleave and _is_splitk)): + # mock_gate_only: single pass, by_n covers full [0, 2*inter_dim) + _eff_e_vec = _e_vec_sk + acc = acc_gate + 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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + elif const_expr(_is_splitk): + # Two-pass epilogue: gate then up, each with atomic add + _eff_e_vec = _e_vec_sk + + # Pass 1: gate + acc = acc_gate + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + gpu.barrier() + + # Pass 2: up + acc = acc_up + _sk_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=_eff_e_vec, + cshuffle_nlane=_cshuffle_nlane_sk, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + else: + 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, + cshuffle_nlane=_cshuffle_nlane, + 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=_frag_elem, + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + lds_out_split=lds_out_B, + ) + + _if_blk = scf.IfOp(blk_valid) + with ir.InsertionPoint(_if_blk.then_block): + _ifexpert_of = scf.IfOp(exp_valid) + with ir.InsertionPoint(_ifexpert_of.then_block): + _moe_gemm1_body() + scf.YieldOp([]) + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher -- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage1, + act, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + use_async_copy, + waves_per_eu, + k_batch, + gate_mode, + a_scale_one, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm1( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_max_token_ids: fx.Pointer, + arg_bias: fx.Pointer, + arg_out_scale_sorted: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + inter_dim_pad_total = arith.constant(2 * inter_dim_pad, index=True) + tile2_pad = 0 + if const_expr(not gate_only): + tile_k_stage2 = tile_k // 2 + tile2_pad = ( + tile_k_stage2 - (inter_dim - inter_dim_pad) % tile_k_stage2 + ) % tile_k_stage2 + + inter_in = arith.index_cast(ir.IndexType.get(), i32_inter_in.ir_value()) + tile_n_index = arith.constant(tile_n, index=True) + if const_expr(mock_gate_only or gate_up_interleave): + gx = ( + inter_in - inter_dim_pad_total + tile2_pad + tile_n_index - 1 + ) / tile_n_index + else: + gx = ( + (inter_in - inter_dim_pad_total + tile2_pad + 2 * tile_n_index - 1) + / tile_n_index + / arith.constant(2, index=True) + ) + + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + arg_out_scale_sorted, + i32_tokens_in, + i32_inter_in, + i32_k_in, + i32_size_expert_ids_in, + ).launch(grid=(gx, gy, k_batch), block=(total_threads, 1, 1), stream=stream) + + return launch_mixed_moe_gemm1 + + +@functools.lru_cache(maxsize=None) +def compile_mixed_moe_gemm2( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp8", + b_dtype: str = "fp4", + out_dtype: str = "f16", + use_cshuffle_epilog: bool | None = None, + # Optional experiment: write per-(token,slot) output (no atomics) into an output shaped + # [tokens*topk, model_dim] (or [tokens, topk, model_dim] flattened), then reduce over topk outside. + # This can reduce atomic contention for small tokens at the cost of extra bandwidth / reduction. + accumulate: bool = True, + enable_bias: bool = False, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + persist_m: int = 4, + sort_block_m: int = 0, + b_nt: int = 2, + xcd_swizzle: int = 0, +): + """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + + persist_m: + - > 0: legacy mode -- each CTA processes exactly persist_m consecutive M tiles. + - <= 0: **persistent mode** -- grid_y = cu_num (auto-detected), each CTA + round-robins over M tiles with stride cu_num. + + a_dtype: + - "fp8": A2 is fp8 + - "fp16": A2 is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": A2 is int8 + - "fp4": A2 is fp4 + + b_dtype: + - "fp8": W is fp8 + - "fp16": W is fp16 (caller uses tile_k halved vs fp8 to match MFMA K halving) + - "int8": W is int8 + - "int4": W4A8 path: A2 is int8, W is packed int4 (2 values per byte) unpacked to int8 in-kernel + - "fp4": W is fp4 + + 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). + + `sort_block_m` is the block_size used by moe_sorting / stage1. When 0 (default), + assumed equal to `tile_m`. When set, stage2 can use a different tile_m from + sorting/stage1. Requires sort_block_m % tile_m == 0. + """ + _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m + if _sort_block_m != tile_m and _sort_block_m % tile_m != 0: + raise ValueError( + f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" + ) + + gpu_arch = get_hip_arch() + allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") + allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") + _state = {} + + if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + raise ValueError( + f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" + ) + if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + raise ValueError( + f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" + ) + + is_f16_a = a_dtype == "fp16" + is_f16_b = b_dtype == "fp16" + + is_f8_a = a_dtype == "fp8" + is_f4_a = a_dtype == "fp4" + is_f4_b = b_dtype == "fp4" + + _scale_pack_m = 2 # physical mn_pack in preshuffle microscale layout + _scale_pack_n = 2 + _scale_pack_k = 2 # physical k_pack in preshuffle scale layout + pack_M = min(_scale_pack_m, tile_m // 16) + pack_N = min(_scale_pack_n, tile_n // 64) + _k_unroll_raw = (int(tile_k) * (2 if a_dtype == "fp16" else 1)) // 128 + pack_K = min(_scale_pack_k, _k_unroll_raw) + + elem_bytes = 1 + + a_elem_bytes = 2 if is_f16_a else 1 + b_elem_bytes = 1 + tile_k_bytes = int(tile_k) * int(a_elem_bytes) + + a_elem_vec_pack = 2 if is_f4_a else 1 + cbsz = 0 if is_f8_a else 4 + blgp = 4 + + # ---- Static B preshuffle strides (compile-time) ---- + # All values below are Python ints computable at kernel-compile time. + # Using them in an explicit multiply-add replaces the fly dialect's + # dynamic ``crd2idx`` path which emits Barrett reduction for the + # non-power-of-2 ``n0 = experts*model_dim//16`` shape. + _b_kpack_bytes_s = 8 if (b_dtype == "int4") else 16 + _b_kpack_elems_s = _b_kpack_bytes_s // b_elem_bytes + _b_c_k_s = inter_dim // _scale_pack_k + _b_c_k0_s = (_b_c_k_s * b_elem_bytes) // 64 + _b_stride_nlane = _b_kpack_elems_s # 16 + _b_stride_klane = 16 * _b_stride_nlane # 256 + _b_stride_k0 = 4 * _b_stride_klane # 1024 + _b_stride_n0 = _b_c_k0_s * _b_stride_k0 # c_k0 * 1024 + assert model_dim % 16 == 0, "model_dim must be divisible by 16" + _expert_b_stride = (model_dim // 16) * _b_stride_n0 + + # K64-byte micro-step: always 64 bytes per `ku`. For fp16, this is 32 elements (2xK16 MFMA). + if (tile_k_bytes % 64) != 0: + raise ValueError( + f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " + f"(tile_k={tile_k}, elem_bytes={a_elem_bytes})" + ) + + out_s = str(out_dtype).strip().lower() + if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): + 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 = b_dtype == "int4" + w_elem_bytes = 2 if is_f16_b else 1 + w_elem_pack = 2 if (is_f4_b or is_int4) else 1 + w_nbytes = (experts * model_dim * inter_dim * w_elem_bytes) // w_elem_pack + bias_nbytes = experts * model_dim * 4 + # INT4 here means W4A8: A2 is int8, W is packed int4 and unpacked to int8 in-kernel. + is_int8 = False + + 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`)." + ) + + def _x_elem_type(): + if is_f4_b: + return T.f8 if is_f8_a else T.i8 + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + def _w_elem_type(): + if is_f4_b: + return T.i8 + return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) + + def _scale_elem_type(): + return T.i32 + + total_threads = 256 + bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) + if bytes_x_per_tile % total_threads != 0: + raise ValueError( + "tile_m*tile_k*elem_bytes must be divisible by " + f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={a_elem_bytes}" + ) + bytes_per_thread_x = bytes_x_per_tile // total_threads + + _use_lds128 = os.environ.get("FLIR_CK_LDS128", "1") in ( + "1", + "true", + "True", + "YES", + "yes", + ) + pad_k = 0 if _use_lds128 else 8 + lds_stride = tile_k + pad_k + + if a_elem_vec_pack > 1: + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + else: + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + + if out_is_f32: + # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. + _use_cshuffle_epilog = ( + 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("FLIR_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 (FLIR_MOE_STAGE2_CSHUFFLE=1)." + ) + + # NOTE: Keep this as a callable so we don't require an MLIR Context at Python-time. + def out_elem(): + return T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) + + def _load_bias_scalar(bias_rsrc, offset): + return buffer_ops.buffer_load(bias_rsrc, offset, vec_width=1, dtype=T.f32) + + 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 the compiler 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. + _persistent = persist_m <= 0 + if _persistent: + from aiter.jit.utils.chip_info import get_cu_num + + _cu_num = get_cu_num() + else: + _cu_num = 0 + _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" + _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" + _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" + module_name = ( + f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" + f"_t{tile_m}x{tile_n}x{tile_k}" + f"_vscale_fix3{_pm_tag}{_sbm_tag}{_xcd_tag}" + ).replace("-", "_") + # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- + # Ping-pong A2 tiles via separate allocators (like stage1). + _single_x_bytes = int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + _cshuffle_elem_bytes_s2 = 2 # f16/bf16 = 2 bytes + lds_out_bytes = ( + _cshuffle_elem_bytes_s2 * int(tile_m) * int(tile_n) + if _use_cshuffle_epilog + else 0 + ) + lds_tid_bytes = int(tile_m) * 4 + _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) + + _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) + _ping_buffer_bytes = _single_x_bytes + + def x_lds_elem(): + return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + + lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) + allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes + _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) + allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes + + lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) + allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + + if True: + + @flyc.kernel(name=module_name) + def moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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(ir.IndexType.get(), i32_tokens_in.ir_value()) + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) + size_expert_ids_in = arith.index_cast( + ir.IndexType.get(), i32_size_expert_ids_in.ir_value() + ) + x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + f32 = T.f32 + i32 = T.i32 + i64 = T.i64 + vec4_f32 = T.vec(4, f32) + vec4_i32 = T.vec(4, i32) + vec16_elems = 16 if a_elem_bytes == 1 else 8 + vec8_elems = 8 if a_elem_bytes == 1 else 4 + vec4_elems = 4 if a_elem_bytes == 1 else 2 + vec16_x = T.vec(vec16_elems, x_elem) + vec2_i64 = T.vec(2, i64) + + def _ptr_buffer_resource(ptr, num_records_bytes): + addr = fx.ptrtoint(ptr) + addr_i64 = arith.index_cast(T.i64, addr) + return buffer_ops.create_buffer_resource_from_addr( + addr_i64, num_records_bytes=num_records_bytes + ) + + acc_init = ( + arith.constant_vector(0, vec4_i32) + if is_int8 + else arith.constant_vector(0.0, vec4_f32) + ) + + # A2 layout (flatten token-slot -> M; use i32 for fly.make_shape). + topk_idx = arith.constant(topk, index=True) + m_in = tokens_in * topk_idx + + # B preshuffle layout: [experts*model_dim, inter_dim] + c_n_total = arith.constant(experts * model_dim, index=True) + kpack_bytes = 8 if is_int4 else 16 + # (inlined: _div_pow2, _mod_pow2 are module-global) + + def check_c_n_valid_gate(base_n): + return arith.cmpi(CmpIPredicate.ult, base_n, model_dim - model_dim_pad) + + def check_c_k_valid_gate(base_k): + return arith.cmpi(CmpIPredicate.ult, base_k, inter_dim - inter_dim_pad) + + # A&B's scale preshuffle layout + # For fp4, k_in is already packed (inter_dim // a_elem_vec_pack), so we need original inter_dim + c_k_orig = arith.constant(inter_dim, index=True) + layout_a_scale = make_preshuffle_scale_layout( + arith, c_mn=m_in, c_k=c_k_orig + ) + layout_b_scale = make_preshuffle_scale_layout( + arith, c_mn=c_n_total, c_k=c_k_orig + ) + + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) + stride_lds = fx.make_stride(_eff_lds_stride, 1) + layout_lds = fx.make_layout(shape_lds, stride_lds) + + tx = gpu.thread_id("x") + by = gpu.block_id("x") # tile along model_dim (N-dim) + bx_persist = gpu.block_id("y") # persistent WG index (M-dim) + + if const_expr(xcd_swizzle > 0): + _NUM_XCDS_S = 8 + _c1_sw = arith.constant(1, index=True) + _c_tn_sw = arith.constant(tile_n, index=True) + _c_mdp_sw = arith.constant(model_dim_pad, index=True) + _gx = (n_in - _c_mdp_sw + _c_tn_sw - _c1_sw) / _c_tn_sw + if const_expr(_persistent): + _gy = arith.constant(_cu_num, index=True) + else: + _c_pm_sw = arith.constant(persist_m, index=True) + _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw + + _linear_id = bx_persist * _gx + by + _num_wgs = _gx * _gy + + _c_xcds = arith.constant(_NUM_XCDS_S, index=True) + _wgs_per_xcd = _num_wgs / _c_xcds + _wgid = (_linear_id % _c_xcds) * _wgs_per_xcd + (_linear_id / _c_xcds) + + _WGM_S = xcd_swizzle + _c_wgm = arith.constant(_WGM_S, index=True) + _num_wgid_in_group = _c_wgm * _gx + _group_id = _wgid / _num_wgid_in_group + _first_pid_m = _group_id * _c_wgm + _remaining_m = _gy - _first_pid_m + _cmp_m = arith.cmpi(CmpIPredicate.ult, _remaining_m, _c_wgm) + _group_size_m = arith.select(_cmp_m, _remaining_m, _c_wgm) + + _wgid_in_group = _wgid % _num_wgid_in_group + bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) + by = _wgid_in_group / _group_size_m + + # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). + k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) + layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) + layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) + + base_ptr_pong = allocator_pong.get_base() + base_ptr_ping = allocator_ping.get_base() + lds_x_pong = SmemPtr( + base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_x_ping = SmemPtr( + base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) + ).get() + lds_out = ( + SmemPtr( + base_ptr_pong, + lds_pong_offset, + (T.bf16 if out_is_bf16 else T.f16), + shape=(tile_m * tile_n,), + ).get() + if _use_cshuffle_epilog + else None + ) + lds_tid = SmemPtr( + base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + ).get() + + # 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 = arith.constant(topk, index=True) + + # X(A2): buffer size in bytes, accounting for FP4 packing (2 elements per byte). + # fp8/int8: 1 byte per element -> bytes = tokens*topk * K + # fp4: 2 elements per byte -> bytes = tokens*topk * K / 2 + c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + x_nbytes_idx = _div_pow2( + (tokens_in * c_topk) * k_in * c_elem_bytes, int(a_elem_vec_pack) + ) + x_nbytes_i32 = arith.index_cast(T.i32, x_nbytes_idx) + x_rsrc = _ptr_buffer_resource(arg_x, x_nbytes_i32) + + w_rsrc = _ptr_buffer_resource(arg_w, w_nbytes) + + # 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 * arith.constant(out_elem_bytes, index=True) + ) + if const_expr(not bool(accumulate)): + out_nbytes_idx = ( + tokens_in + * arith.index(topk) + * n_in + * arith.constant(out_elem_bytes, index=True) + ) + out_nbytes_i32 = arith.index_cast(T.i32, out_nbytes_idx) + out_rsrc = _ptr_buffer_resource(arg_out, out_nbytes_i32) + + # num_valid_ids (sorted padded MN) for scale sizing / guards. + numids_rsrc = _ptr_buffer_resource( + arg_num_valid_ids, arith.constant(4, type=T.i32) + ) + num_valid_i32 = buffer_ops.buffer_load( + numids_rsrc, arith.constant(0, index=True), vec_width=1, dtype=T.i32 + ) + # num_valid_ids is a scalar (same value for all lanes) loaded into + # VGPR. Promote to SGPR so downstream buffer resource descriptors + # that use it for num_records stay in SGPRs, eliminating the + # expensive waterfall loop the compiler would otherwise emit. + num_valid_i32 = rocdl.ReadfirstlaneOp(T.i32, num_valid_i32).res + num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + + # fp16 path ignores scales completely (implicit scale=1.0). + sx_rsrc = 1 + sw_rsrc = 1 + if const_expr(not is_f16_a): + if const_expr(is_f4_a or is_f8_a): + # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. + # Caller must pre-scatter a2_scale via moe_mxfp4_sort. + kblk = _div_pow2(k_in, 32) + sx_nbytes_idx = num_valid_idx * kblk + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + else: + # scale_x (A2 scale): [tokens*topk] f32 -> bytes = tokens*topk*4 + sx_nbytes_idx = (tokens_in * c_topk) * arith.constant(4, index=True) + sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx) + sx_rsrc = _ptr_buffer_resource(arg_scale_x, sx_nbytes_i32) + + if const_expr(not is_f16_b): + # Weight microscale buffer (packed i32 holding e8m0 bytes). + # Use an exact descriptor size so hardware OOB checking works. + kblk_w = _div_pow2(k_in, 32) # K/32 + mn_w = arith.constant(experts * model_dim, index=True) + sw_nbytes_idx = mn_w * kblk_w # bytes (e8m0) + sw_nbytes_i32 = arith.index_cast(T.i32, sw_nbytes_idx) + sw_rsrc = _ptr_buffer_resource(arg_scale_w, sw_nbytes_i32) + + # sorted_token_ids / sorted_weights: [blocks*tile_m] (padded length) + sorted_nbytes_idx = ( + size_expert_ids_in + * arith.constant(tile_m, index=True) + * arith.constant(4, index=True) + ) + sorted_nbytes_i32 = arith.index_cast(T.i32, sorted_nbytes_idx) + sorted_rsrc = _ptr_buffer_resource(arg_sorted_token_ids, sorted_nbytes_i32) + sorted_w_rsrc = _ptr_buffer_resource(arg_sorted_weights, sorted_nbytes_i32) + + # expert ids: [sort_blocks] i32. + _c_sbm = arith.constant(_sort_block_m, index=True) + _c_tm = arith.constant(tile_m, index=True) + _c1 = arith.constant(1, index=True) + _sort_blocks_ub = _div_pow2( + size_expert_ids_in * _c_tm + _c_sbm - _c1, _sort_block_m + ) + eid_nbytes_idx = _sort_blocks_ub * arith.constant(4, index=True) + eid_nbytes_i32 = arith.index_cast(T.i32, eid_nbytes_idx) + expert_rsrc = _ptr_buffer_resource(arg_expert_ids, eid_nbytes_i32) + bias_rsrc = ( + _ptr_buffer_resource(arg_bias, bias_nbytes) if enable_bias else None + ) + + # ---- persist loop ---- + _c0_p = arith.constant(0, index=True) + _c1_p = arith.constant(1, index=True) + + if const_expr(_persistent): + # Expert-phase scheduling: contiguous M-tile dispatch. + # grid_y = cu_num, each CTA handles a contiguous chunk of M-tiles: + # [bx_persist * tiles_per_block, ..., (bx_persist+1) * tiles_per_block - 1] + # Adjacent blocks process adjacent M-tiles -> same expert -> B weight L2 reuse. + _c_cu = arith.constant(_cu_num, index=True) + _c_tm_p = arith.constant(tile_m, index=True) + _num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) + _total_m_tiles = (_num_valid_idx + _c_tm_p - _c1_p) / _c_tm_p + _tiles_per_block = (_total_m_tiles + _c_cu - _c1_p) / _c_cu + _i1 = ir.IntegerType.get_signless(1) + _init_active = arith.constant(1, type=_i1) + _for_persist = scf.ForOp(_c0_p, _tiles_per_block, _c1_p, [_init_active]) + else: + # Legacy mode: fixed persist_m consecutive tiles. + _c_pm = arith.constant(persist_m, index=True) + _init_prev_expert = arith.constant(0, type=T.i32) + _init_prev_b_base = arith.constant(0, index=True) + _for_persist = scf.ForOp( + _c0_p, + _c_pm, + _c1_p, + [_init_prev_expert, _init_prev_b_base], + ) + + _for_ip = ir.InsertionPoint(_for_persist.body) + _for_ip.__enter__() + _mi_p = _for_persist.induction_variable + + if const_expr(_persistent): + _still_active = _for_persist.inner_iter_args[0] + bx = bx_persist * _tiles_per_block + _mi_p + else: + _prev_expert_i32 = _for_persist.inner_iter_args[0] + _prev_expert_b_base = _for_persist.inner_iter_args[1] + bx = bx_persist * arith.constant(persist_m, index=True) + _mi_p + + bx_m = bx * arith.constant(tile_m, index=True) + + # Early-exit guard: skip garbage expert blocks beyond `num_valid_ids`. + bx_m_i32 = arith.index_cast(T.i32, bx_m) + blk_valid = arith.cmpi(CmpIPredicate.ult, bx_m_i32, num_valid_i32) + + sort_blk = _div_pow2(bx_m, _sort_block_m) + expert_i32 = buffer_ops.buffer_load( + expert_rsrc, sort_blk, vec_width=1, dtype=T.i32 + ) + expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + exp_valid = arith.cmpi( + CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) + ) + + if const_expr(_persistent): + # Absolute B-base: no cross-iteration state needed. + _expert_b_base = expert_idx * arith.constant( + _expert_b_stride, index=True + ) + else: + # Legacy incremental B-base: delta = (cur - prev) * stride + _delta_expert = arith.subi(expert_i32, _prev_expert_i32) + _delta_expert_idx = arith.index_cast(ir.IndexType.get(), _delta_expert) + _delta_b = _delta_expert_idx * arith.constant( + _expert_b_stride, index=True + ) + _expert_b_base = _prev_expert_b_base + _delta_b + + # Early-exit: if the first row of this tile is a sentinel (all-padding tile), + # skip the entire GEMM. + _first_tok = buffer_ops.buffer_load( + sorted_rsrc, bx_m, vec_width=1, dtype=T.i32 + ) + _first_tid = arith.andi(_first_tok, arith.constant(0xFFFFFF, type=T.i32)) + _tokens_i32_guard = arith.index_cast(T.i32, tokens_in) + tile_has_tokens = arith.cmpi( + CmpIPredicate.ult, _first_tid, _tokens_i32_guard + ) + + # For tile_m < 32 (pack_M < _scale_pack_m): shift a_scale i32 so the + # correct bytes land at the op_sel positions we use. + if const_expr(pack_M < _scale_pack_m): + _m_off = _mod_pow2(_div_pow2(bx_m, 16), _scale_pack_m) + _m_scale_shift_i32 = arith.index_cast( + T.i32, _m_off * arith.constant(8, index=True) + ) + else: + _m_scale_shift_i32 = None + + def _moe_gemm2_then_body(): + # Expert id for this M tile. + n_idx = arith.constant(model_dim, index=True) + 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 we require 16B. + if const_expr(is_f16_a): + if const_expr(bytes_per_thread_x % 16 != 0): + raise ValueError( + f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16" + ) + x_load_bytes = 16 + else: + if 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) + vec4_i32 = T.vec(4, i32) + + c_k_div4 = _div_pow2( + _div_pow2(k_in, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + tile_k_dwords = (int(tile_k) * int(a_elem_bytes)) // ( + 4 * int(a_elem_vec_pack) + ) + layout_x_tile_div4 = fx.make_layout( + (tile_m, tile_k_dwords), stride=(tile_k_dwords, 1) + ) + c_chunk_i32 = arith.constant(chunk_i32, index=True) + tx_i32_base = tx * c_chunk_i32 + + topk_i32 = arith.constant(topk) + mask24 = arith.constant(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, + ) + + vec1_i32 = T.vec(1, i32) + vec2_i32 = T.vec(2, i32) + x_load_vec_elems = ( + x_load_bytes if a_elem_bytes == 1 else x_load_bytes // a_elem_bytes + ) + + 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. + """ + if const_expr(x_load_bytes == 16): + idx_elem = ( + idx_i32 if a_elem_bytes == 1 else (idx_i32 * arith.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, + ) + # 8B/4B: convert dword index to byte offset and use offset_in_bytes path. + idx_bytes = idx_i32 * arith.index(4) + return _buffer_load_vec( + buffer_ops, + vector, + x_rsrc, + idx_bytes, + elem_type=x_elem, + vec_elems=x_load_vec_elems, + elem_bytes=a_elem_bytes, + offset_in_bytes=True, + ) + + # 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 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(ir.IndexType.get(), row_ts_i32) + + x_row_base_div4.append(row_ts_idx * c_k_div4) + + def load_x_tile(base_k): + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 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(vec4_i32, x_vec)) + elif const_expr(x_load_bytes == 8): + parts.append(vector.bitcast(vec2_i32, x_vec)) + else: + parts.append(vector.bitcast(vec1_i32, x_vec)) + return parts + + # tx -> wave/lane (GEMM-style decomposition). + coord_wl = idx2crd(tx, layout_tx_wave_lane) + wave_id = layout_get(coord_wl, 0) + lane_id = layout_get(coord_wl, 1) + coord_l16 = idx2crd(lane_id, layout_lane16) + lane_div_16 = layout_get(coord_l16, 0) + lane_mod_16 = layout_get(coord_l16, 1) + + row_a_lds = lane_mod_16 + + col_offset_base = lane_div_16 * arith.constant(16, index=True) + + # Dynamic N tiling within block. + num_waves = 4 + n_per_wave = tile_n // num_waves + num_acc_n = n_per_wave // 16 + c_n_per_wave = arith.constant(n_per_wave, index=True) + wave_mod_4 = _mod_pow2(wave_id, 4) + n_tile_base = wave_mod_4 * c_n_per_wave + + by_n = by * arith.constant(tile_n, index=True) + + if const_expr(pack_N < _scale_pack_n): + _global_n_base = expert_off_idx + by_n + n_tile_base + _n_off = _mod_pow2(_div_pow2(_global_n_base, 16), _scale_pack_n) + _n_scale_shift_i32 = arith.index_cast( + T.i32, _n_off * arith.constant(8, index=True) + ) + else: + _n_scale_shift_i32 = None + n_intra_list = [None] * num_acc_n + n_blk_list = [None] * num_acc_n + col_g_list = [None] * num_acc_n + for i in range_constexpr(num_acc_n): + offset = i * 16 + col_g = by_n + n_tile_base + col_g = _div_pow2(col_g, 2) + offset + col_g = col_g + lane_mod_16 + col_g_list[i] = col_g + c_offset = arith.constant(offset, index=True) + global_n = by_n + n_tile_base + c_offset + lane_mod_16 + n_blk_list[i] = _div_pow2(global_n, 16) + n_intra_list[i] = _mod_pow2(global_n, 16) + + m_repeat = tile_m // 16 + k_unroll = tile_k_bytes // 128 # K64-byte micro-step (2x MFMA) + + # fp4 pack + k_unroll_packed = k_unroll // pack_K + m_repeat_packed = m_repeat // pack_M + num_acc_n_packed = num_acc_n // pack_N + + _K_per_ku_s2 = tile_k // k_unroll + _pad_k_elems_s2 = (inter_dim_pad % tile_k) if inter_dim_pad > 0 else 0 + _pad_ku_skip_s2 = _pad_k_elems_s2 // _K_per_ku_s2 + _tail_ku_s2 = k_unroll - _pad_ku_skip_s2 + _tail_ku_packed_s2 = ( + (_tail_ku_s2 + pack_K - 1) // pack_K + if _pad_ku_skip_s2 > 0 + else None + ) + + # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- + def load_b_packs_k64(base_k, ku: int, ni: int): + """Load one K64-byte B micro-step: single 16B load, split into 2x i64.""" + base_k_bytes = base_k * arith.constant( + int(b_elem_bytes), index=True + ) + k0_base = _div_pow2(base_k_bytes, 64) + k0 = k0_base + arith.constant(ku, index=True) + k1 = lane_div_16 + # Incremental B addressing: _expert_b_base carries the + # expert's preshuffle offset (updated via delta each + # persist_m iteration); local n_blk/n_intra contribute + # the per-lane within-tile offset. All strides are + # compile-time constants -> shift/mul, no Barrett. + idx_pack = ( + _expert_b_base + + n_blk_list[ni] * arith.constant(_b_stride_n0, index=True) + + k0 * arith.constant(_b_stride_k0, index=True) + + k1 * arith.constant(_b_stride_klane, index=True) + + n_intra_list[ni] * arith.constant(_b_stride_nlane, index=True) + ) + + vec_elems = kpack_bytes // int(b_elem_bytes) + b16 = _buffer_load_vec( + buffer_ops, + vector, + w_rsrc, + idx_pack, + elem_type=_w_elem_type(), + vec_elems=vec_elems, + elem_bytes=b_elem_bytes, + offset_in_bytes=(b_elem_bytes == 1), + cache_modifier=b_nt, + ) + b_i64x2 = vector.bitcast(vec2_i64, b16) + b0 = vector.extract( + b_i64x2, static_position=[0], dynamic_position=[] + ) + b1 = vector.extract( + b_i64x2, static_position=[1], dynamic_position=[] + ) + return b0, b1 + + def load_b_tile(base_k, ku_limit=k_unroll): + b_tile = [] + for ku in range_constexpr(ku_limit): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + _b_split_enabled = k_unroll >= 2 + _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll + + def load_b_tile_lo(base_k): + """Load first half of B tile (ku < _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_b_tile_hi(base_k): + """Load second half of B tile (ku >= _b_split_ku).""" + b_tile = [] + for ku in range_constexpr(_b_split_ku, k_unroll): + packs0 = [] + packs1 = [] + for ni in range_constexpr(num_acc_n): + b0, b1 = load_b_packs_k64(base_k, ku, ni) + packs0.append(b0) + packs1.append(b1) + b_tile.append((packs0, packs1)) + return b_tile + + def load_scale(arg_scale, rsrc, scale_info, ku, mni): + k_lane = lane_div_16 + n_lane = lane_mod_16 + # Direct arith crd2idx: idx = mni*stride_n0 + ku*stride_k0 + k_lane*stride_klane + n_lane + idx_pack = ( + mni * scale_info.stride_n0 + + ku * scale_info.stride_k0 + + k_lane * scale_info.stride_klane + + n_lane + ) + s = buffer_ops.buffer_load(rsrc, idx_pack, vec_width=1, dtype=T.i32) + return vector.from_elements(T.vec(1, T.i32), [s]) + + def _apply_k_shift(scale_vec, k_shift_bits): + if const_expr(k_shift_bits > 0): + val = vector.extract( + scale_vec, static_position=[0], dynamic_position=[] + ) + val = arith.shrui(val, arith.constant(k_shift_bits, type=T.i32)) + return vector.from_elements(T.vec(1, T.i32), [val]) + return scale_vec + + def load_b_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + b_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for ni in range_constexpr(num_acc_n_packed): + scale = load_scale( + arg_scale_w, + sw_rsrc, + layout_b_scale, + ku + base_k, + ni + + _div_pow2( + _div_pow2( + expert_off_idx + by_n + n_tile_base, + _scale_pack_n, + ), + 16, + ), + ) + scale = _apply_k_shift(scale, k_shift_bits) + b_scale_tile.append(scale) + return b_scale_tile + + def load_a_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + a_scale_tile = [] + for ku in range_constexpr(ku_packed_limit): + for mi in range_constexpr(m_repeat_packed): + scale = load_scale( + arg_scale_x, + sx_rsrc, + layout_a_scale, + ku + base_k, + mi + _div_pow2(_div_pow2(bx_m, _scale_pack_m), 16), + ) + scale = _apply_k_shift(scale, k_shift_bits) + a_scale_tile.append(scale) + return a_scale_tile + + def prefetch_ab_scale_tile( + base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed + ): + return [ + load_a_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + load_b_scale_tile( + base_k, k_shift_bits, ku_packed_limit=ku_packed_limit + ), + ] + + vec8_x = T.vec(vec8_elems, x_elem) + vec4_x_lds = T.vec(vec4_elems, x_elem) + + # ---- Pipeline helpers: store X tile to LDS (unused in DMA path) ---- + _lds_base_zero = arith.index(0) + + def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + 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_buffer, + vec16_ty=vec16_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + 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_buffer, + vec8_ty=vec8_x, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x2=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + else: # x_load_bytes == 4 + lds_store_4b_xor16( + arith, + vector, + lds_memref=lds_buffer, + vec4_ty=vec4_x_lds, + layout_lds=layout_lds, + row_local=row_local, + col_local_i32=col_local_i32, + tx_c4=arith.index(4), + k_blocks16=k_blocks16, + lds_base=_lds_base_zero, + vec_part_i32x1=vec_x_in_parts[i], + elem_bytes=elem_bytes, + ) + + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + col_base_swz_bytes = swizzle_xor16( + curr_row_a_lds, col_base, k_blocks16 + ) + col_base_swz = ( + col_base_swz_bytes + if elem_bytes == 1 + else (col_base_swz_bytes / arith.index(2)) + ) + idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) + loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) + a0 = vector.extract( + a_i64x2, static_position=[0], dynamic_position=[] + ) + a1 = vector.extract( + a_i64x2, static_position=[1], dynamic_position=[] + ) + return a0, a1 + + def compute_tile( + acc_in, + b_tile_in, + lds_buffer, + a_scale=None, + b_scale=None, + *, + prefetch_epilogue: bool = False, + a0_prefetch=None, + a1_prefetch=None, + b_hi_loader=None, + ku_count=k_unroll, + ): + if const_expr(b_hi_loader is not None): + b_tile_full = [None] * k_unroll + for i in range_constexpr(_b_split_ku): + b_tile_full[i] = b_tile_in[i] + else: + b_tile_full = b_tile_in + acc_list = list(acc_in) + mfma_res_ty = vec4_i32 if is_int8 else vec4_f32 + + epilogue_pf = None + bias = None + if const_expr(prefetch_epilogue): + if const_expr(enable_bias): + bias = [] + for ni in range_constexpr(num_acc_n): + global_n = by_n + n_tile_base + ni * 16 + lane_mod_16 + bias_offset = expert_off_idx + global_n + bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) + tw_pf = None + if const_expr(doweight_stage2): + tw_pf = [] + lane_div_16_mul4_pf = lane_div_16 * arith.index(4) + ii_idx_list_pf = [ + arith.constant(ii, index=True) for ii in range(4) + ] + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + for ii in range_constexpr(4): + row_off_pf = ( + lane_div_16_mul4_pf + ii_idx_list_pf[ii] + ) + row_in_tile_pf = mi_base_pf + row_off_pf + sorted_row_pf = bx_m + row_in_tile_pf + tw_pf.append( + buffer_ops.buffer_load( + sorted_w_rsrc, + sorted_row_pf, + vec_width=1, + dtype=f32, + ) + ) + epilogue_pf = (None, tw_pf, bias) + + c0_i64 = arith.constant(0, type=T.i64) + vec4_i64 = T.vec(4, T.i64) + vec8_i32 = T.vec(8, T.i32) + + def pack_i64x4_to_i32x8(x0, x1, x2, x3): + v4 = vector.from_elements(vec4_i64, [x0, x1, x2, x3]) + return vector.bitcast(vec8_i32, v4) + + # fp4 path -- single k_idx loop [0, k_unroll). + # b_hi load is issued at the very start so all k_unroll + # MFMAs can overlap the VMEM latency. + _pack_K_shift = (pack_K - 1).bit_length() + _pack_K_mask = pack_K - 1 + + if const_expr(b_hi_loader is not None): + _b_hi = b_hi_loader() + for _bhi_i in range_constexpr(len(_b_hi)): + b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] + + for k_idx in range_constexpr(ku_count): + ku128 = k_idx >> _pack_K_shift + ikxdl = k_idx & _pack_K_mask + + b_packs0, b_packs1 = b_tile_full[k_idx] + + col_base = col_offset_base + (k_idx * 128) // a_elem_vec_pack + + for mi in range_constexpr(m_repeat_packed): + a_scale_i32 = a_scale[ku128 * m_repeat_packed + mi] + a_scale_val = vector.extract( + a_scale_i32, static_position=[0], dynamic_position=[] + ) + if const_expr(_m_scale_shift_i32 is not None): + a_scale_val = arith.shrui( + a_scale_val, _m_scale_shift_i32 + ) + for ni in range_constexpr(num_acc_n_packed): + b_scale_i32 = b_scale[ku128 * num_acc_n_packed + ni] + b_scale_val = vector.extract( + b_scale_i32, + static_position=[0], + dynamic_position=[], + ) + if const_expr(_n_scale_shift_i32 is not None): + b_scale_val = arith.shrui( + b_scale_val, _n_scale_shift_i32 + ) + + for imxdl in range_constexpr(pack_M): + col_base0 = col_base + mi_idx = mi * pack_M + imxdl + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr( + (a0_prefetch is not None) + and (k_idx == 0) + and (mi_idx == 0) + ): + a0, a1 = a0_prefetch + elif const_expr( + (a1_prefetch is not None) + and (k_idx == 1) + and (mi_idx == 0) + ): + a0, a1 = a1_prefetch + else: + a0, a1 = lds_load_packs_k64( + curr_row_a_lds, col_base0, lds_buffer + ) + + if const_expr(is_f8_a): + col_base1 = col_base + 64 + a2, a3 = lds_load_packs_k64( + curr_row_a_lds, col_base1, lds_buffer + ) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + else: + a128 = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + + b0 = b_packs0[ni_idx] + b1 = b_packs1[ni_idx] + b128 = pack_i64x4_to_i32x8( + b0, b1, c0_i64, c0_i64 + ) + + acc_idx = mi_idx * num_acc_n + ni_idx + acc_list[acc_idx] = ( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + mfma_res_ty, + [ + a128, + b128, + acc_list[acc_idx], + cbsz, + blgp, + ikxdl * _scale_pack_m + imxdl, + a_scale_val, + ikxdl * _scale_pack_n + inxdl, + b_scale_val, + ], + ) + ) + + return acc_list, epilogue_pf + + # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- + # ---- Async DMA: GMEM -> LDS (bypasses VGPR, like stage1) ---- + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + + def dma_x_tile_to_lds(base_k, lds_buffer): + c4_idx = arith.index(4) + base_k_div4 = _div_pow2( + _div_pow2(base_k, int(a_elem_vec_pack)) + * arith.constant(int(a_elem_bytes), index=True), + 4, + ) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = memref.extract_aligned_pointer_as_index( + lds_buffer + ) + wave_id * arith.constant( + _wave_size * _dma_bytes, index=True + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_buffer): + dma_x_tile_to_lds(base_k, lds_buffer) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler(): + rocdl.sched_barrier(0) + + def _k_shift_bits(k_py): + if const_expr(pack_K >= _scale_pack_k): + return 0 + return ((k_py // 128) % _scale_pack_k) * _scale_pack_m * 8 + + def _k_base(k_py): + return k_py // _scale_pack_k // 128 + + # Preload sorted_idx into lds_tid for epilogue precompute_row + # (N-independent; placed before N-tile loop so it's done once per M-tile.) + _c_tile_m_idx = arith.constant(tile_m, index=True) + _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) + _if_tid = scf.IfOp(_tid_in_range) + with ir.InsertionPoint(_if_tid.then_block): + _tid_row = bx_m + tx + _tid_val = buffer_ops.buffer_load( + sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 + ) + _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) + vector.store(_tid_vec1, lds_tid, [tx]) + scf.YieldOp([]) + + gpu.barrier() + + # Prologue -- B-first + async DMA X(0) -> pong. + k0 = arith.index(0) + if const_expr(_b_split_enabled): + b_cur = load_b_tile_lo(k0) + else: + b_cur = load_b_tile(k0) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(0), _k_shift_bits(0) + ) + rocdl.sched_barrier(0) + prefetch_x_to_lds(k0, lds_x_pong) + rocdl.s_waitcnt(0) + gpu.barrier() + + acc = [acc_init] * num_acc_n * m_repeat + + # Cross-tile A0+A1 LDS prefetch from pong buffer. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + _a1_col_base = col_offset_base + 128 // a_elem_vec_pack + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + # 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 + + c2_tile_k = arith.constant(tile_k * 2, index=True) + b_pong = b_cur + k0_pong_bk = k0 + + # Only emit the scf.for when there are actually iterations to run. + # When k_main2_py == 0 the loop body is empty; emitting an scf.for + # would create a region whose internal SSA values cannot be used + # by the post-loop tail code. + def _make_b_hi_loader(base_k): + """Create a b_hi_loader callable for a given base_k.""" + return lambda _bk=base_k: load_b_tile_hi(_bk) + + if const_expr(k_main2_py > 0): + for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): + rocdl.sched_barrier(0) + k_iv = arith.index(k_iv_py) + next_k1 = k_iv + tile_k + next_k1_bk = next_k1 // 2 + # DMA X(next_k1) -> ping (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k1, lds_x_ping) + b_ping_lo = ( + load_b_tile_lo(next_k1_bk) + if _b_split_enabled + else load_b_tile(next_k1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(next_k1), _k_shift_bits(next_k1) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) + if _b_split_enabled + else None + ), + ) + hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Cross-tile prefetch for the ping tile we are about to compute. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 + else None + ) + + next_k2 = k_iv + c2_tile_k + next_k2_py = k_iv_py + tile_k * 2 + next_k2_bk = next_k2 // 2 + # DMA X(next_k2) -> pong (non-blocking, overlaps with compute) + prefetch_x_to_lds(next_k2, lds_x_pong) + b_pong = ( + load_b_tile_lo(next_k2_bk) + if _b_split_enabled + else load_b_tile(next_k2_bk) + ) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(next_k2_py), _k_shift_bits(next_k2_py) + ) + + acc, _ = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + b_hi_loader=( + _make_b_hi_loader(next_k1_bk) + if _b_split_enabled + else None + ), + ) + k0_pong_bk = next_k2_bk + hot_loop_scheduler() + gpu.barrier() + + # Cross-tile prefetch for the next pong tile. + a0_prefetch_pong = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_pong + ) + a1_prefetch_pong = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + if pack_K >= 2 + else None + ) + + if const_expr(odd_k_tiles): + # Tail: single remaining tile (already in pong buffer). + acc, epilogue_pf = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + prefetch_epilogue=True, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + else: + # Tail: 2 remaining tiles. + k_tail1 = (k_in + tile_k - 1) // tile_k * tile_k - tile_k + k_tail1_py = ( + int(inter_dim) + tile_k - 1 + ) // tile_k * tile_k - tile_k + k_tail1_bk = k_tail1 // 2 + # DMA tail X -> ping + prefetch_x_to_lds(k_tail1, lds_x_ping) + if const_expr(_pad_ku_skip_s2 > 0): + b_ping_lo = load_b_tile(k_tail1_bk, ku_limit=_tail_ku_s2) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), + _k_shift_bits(k_tail1_py), + ku_packed_limit=_tail_ku_packed_s2, + ) + else: + b_ping_lo = ( + load_b_tile_lo(k_tail1_bk) + if _b_split_enabled + else load_b_tile(k_tail1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) + ) + + acc, _ = compute_tile( + acc, + b_pong, + lds_x_pong, + a_scale_pong, + b_scale_pong, + a0_prefetch=a0_prefetch_pong, + a1_prefetch=a1_prefetch_pong, + b_hi_loader=( + _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None + ), + ) + + # hot_loop_scheduler() + rocdl.s_waitcnt(0) + gpu.barrier() + + # Epilogue tile with sw prefetch. + a0_prefetch_ping = lds_load_packs_k64( + row_a_lds, col_offset_base, lds_x_ping + ) + a1_prefetch_ping = ( + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + if pack_K >= 2 and (_pad_ku_skip_s2 == 0 or _tail_ku_s2 >= 2) + else None + ) + acc, epilogue_pf = compute_tile( + acc, + b_ping_lo, + lds_x_ping, + a_scale_ping, + b_scale_ping, + a0_prefetch=a0_prefetch_ping, + a1_prefetch=a1_prefetch_ping, + prefetch_epilogue=True, + b_hi_loader=( + None + if _pad_ku_skip_s2 > 0 + else ( + _make_b_hi_loader(k_tail1_bk) + if _b_split_enabled + else None + ) + ), + ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, + ) + + # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- + # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. + + sw_pf = None + tw_pf = None + bias_pf = None + if const_expr(epilogue_pf is not None): + sw_pf, tw_pf, bias_pf = epilogue_pf + + mask24_i32 = arith.constant(0xFFFFFF) + topk_i32_v = topk_i32 + + zero_i32 = arith.constant(0) + + def atomic_add_f16x2(val_f16x2, byte_off_i32): + rocdl.raw_ptr_buffer_atomic_fadd( + val_f16x2, + out_rsrc, + byte_off_i32, + zero_i32, + zero_i32, + ) + + # Weight scales for the N tile (col_g depends on lane/wave/by but not on (t,s)). + if const_expr(lds_out is None): + raise RuntimeError( + "FLIR_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased." + ) + + # Precompute the output base address (i64 index) for ALL paths. + # Both accumulate=True (global atomic) and accumulate=False (global store) + # need 64-bit addressing to avoid i32 offset overflow when + # tokens * model_dim * elem_bytes > INT32_MAX (~150K tokens for model_dim=7168). + out_base_i64 = arith.index_cast(T.i64, fx.ptrtoint(arg_out)) + out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + + 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, + ): + # Match origin/dev_a16w4: rely on sentinel padded rows + hardware OOB behavior. + fused2 = buffer_ops.buffer_load( + sorted_rsrc, row, vec_width=1, dtype=T.i32 + ) + t2 = fused2 & mask24_i32 + s2 = fused2 >> 24 + + t_ok = arith.cmpi(CmpIPredicate.ult, t2, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s2, topk_i32_v) + ts_ok = arith.andi(t_ok, s_ok) + t2_safe = arith.select(ts_ok, t2, arith.constant(0)) + s2_safe = arith.select(ts_ok, s2, arith.constant(0)) + t2_safe * topk_i32_v + s2_safe + + if const_expr(doweight_stage2): + 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=f32 + ) + + 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( + acc[acc_idx], static_position=[ii], dynamic_position=[] + ) + if const_expr(is_int8): + v = arith.sitofp(f32, v) + if const_expr(enable_bias): + v = v + bias_pf[ni] + + 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, [v_out]) + + vector.store(v1, lds_out, [lds_idx], alignment=2) + + def precompute_row(*, row_local, row): + # Use lds_tid (sorted_idx preloaded to LDS) instead of buffer_load + # to avoid extra VMEM round-trips in the epilogue. + fused2 = memref.load(lds_tid, [row_local]) + row_i32 = arith.index_cast(T.i32, row) + row_valid0 = arith.cmpi(CmpIPredicate.ult, row_i32, num_valid_i32) + t = fused2 & mask24_i32 + s = fused2 >> 24 + t_ok = arith.cmpi(CmpIPredicate.ult, t, tokens_i32) + s_ok = arith.cmpi(CmpIPredicate.ult, s, topk_i32_v) + row_valid = arith.andi(row_valid0, arith.andi(t_ok, s_ok)) + t_idx = arith.index_cast(ir.IndexType.get(), t) + s_idx = arith.index_cast(ir.IndexType.get(), s) + ts_idx = t_idx * arith.constant(topk, index=True) + s_idx + if const_expr(accumulate): + row_byte_base = out_base_idx + t_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + else: + row_byte_base = out_base_idx + ts_idx * arith.constant( + model_dim * out_elem_bytes, index=True + ) + return ((fused2, row_byte_base), row_valid) + + def _idx_to_llvm_ptr(idx_val, addr_space=1): + """Convert an index-typed byte address to !llvm.ptr.""" + idx_v = idx_val._value if hasattr(idx_val, "_value") else idx_val + i64_v = arith.index_cast(T.i64, idx_v) + i64_raw = i64_v._value if hasattr(i64_v, "_value") else i64_v + ptr_ty = ir.Type.parse(f"!llvm.ptr<{addr_space}>") + return llvm.inttoptr(ptr_ty, i64_raw) + + def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): + fused, row_byte_base = row_ctx + if const_expr(not bool(accumulate)): + # ---- 64-bit global store path (avoids i32 offset overflow) ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + frag_v = frag._value if hasattr(frag, "_value") else frag + llvm.StoreOp( + frag_v, + out_ptr_v, + alignment=_e_vec * out_elem_bytes, + nontemporal=True, + ) + else: + # ---- accumulate=True: 64-bit global atomic path ---- + col_idx = col_g0 + byte_off_col = col_idx * arith.constant( + out_elem_bytes, index=True + ) + ptr_addr_idx = row_byte_base + byte_off_col + out_ptr_v = _idx_to_llvm_ptr(ptr_addr_idx) + 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=_e_vec * out_elem_bytes, + ) + + _e_vec = 2 if accumulate else min(tile_n // 32, 8) + 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=( + ir.BF16Type.get() if out_is_bf16 else ir.F16Type.get() + ), + write_row_to_lds=write_row_to_lds, + precompute_row=precompute_row, + store_pair=store_pair, + ) + + _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) + + if const_expr(_persistent): + # Short-circuit: contiguous tiles are monotonically increasing, + # so once bx_m >= num_valid_ids all remaining tiles are invalid. + _cur_active = arith.andi(_still_active, blk_valid) + _do_gemm = arith.andi( + _cur_active, arith.andi(exp_valid, tile_has_tokens) + ) + _if_valid = scf.IfOp(_do_gemm) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([_cur_active]) + else: + _if_valid = scf.IfOp(_all_valid) + with ir.InsertionPoint(_if_valid.then_block): + _moe_gemm2_then_body() + scf.YieldOp([]) + + gpu.barrier() + scf.YieldOp([expert_i32, _expert_b_base]) + _for_ip.__exit__(None, None, None) + + # -- Host launcher (flyc.jit + .launch) -------------------------------- + _cache_tag = ( + module_name, + a_dtype, + b_dtype, + out_dtype, + tile_m, + tile_n, + tile_k, + doweight_stage2, + accumulate, + enable_bias, + model_dim_pad, + inter_dim_pad, + use_cshuffle_epilog, + persist_m, + _sort_block_m, + _cu_num if _persistent else 0, + xcd_swizzle, + ) + + @flyc.jit + def launch_mixed_moe_gemm2( + arg_out: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_scale_x: fx.Pointer, + arg_scale_w: fx.Pointer, + arg_sorted_token_ids: fx.Pointer, + arg_expert_ids: fx.Pointer, + arg_sorted_weights: fx.Pointer, + arg_num_valid_ids: fx.Pointer, + arg_bias: fx.Pointer, + 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, + ): + _ = _cache_tag + allocator_pong.finalized = False + allocator_ping.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator_pong.finalize() + allocator_ping.finalize() + + n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) + _tile_n_idx = arith.constant(tile_n, index=True) + _model_dim_pad_idx = arith.constant(model_dim_pad, index=True) + gx = ( + n_in - _model_dim_pad_idx + _tile_n_idx - arith.constant(1, index=True) + ) / _tile_n_idx + if const_expr(_persistent): + gy = arith.constant(_cu_num, index=True) + else: + _c_pm_l = arith.constant(persist_m, index=True) + gy = ( + arith.index_cast(ir.IndexType.get(), i32_size_expert_ids_in.ir_value()) + + _c_pm_l + - arith.constant(1, index=True) + ) / _c_pm_l + + 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, + arg_bias, + 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_mixed_moe_gemm2 + +# =========================================================================== +# Host-side launchers (adapted from aiter/ops/flydsl/moe_kernels.py). +# These pack pointer args and drive the inline compile_mixed_moe_gemm1/2 +# builders above. +# =========================================================================== +_DLPACK_SAFE = (torch.uint8, torch.float16, torch.bfloat16, torch.float32) + + +def _view_safe(t: torch.Tensor) -> torch.Tensor: + """View as uint8 if dtype is not dlpack-safe, otherwise return as-is.""" + return ( + t.view(torch.uint8) + if t is not None and t.numel() > 0 and t.dtype not in _DLPACK_SAFE + else t + ) + + +def _ptr_view_safe(t: torch.Tensor): + """Pass only the device data pointer; shape is carried by explicit args.""" + view = _view_safe(t) + type_name = type(view).__name__ + module_name = type(view).__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, view.data_ptr()) + + +def _s1_args_fp4( + out, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + out_scale_sorted, + token_num, + n_in, + k_in, + size_expert_ids_in, + dev, + bias=None, + stream=None, +): + empty_f32 = torch.empty(0, device=dev, dtype=torch.float32) + _bias = bias if bias is not None else empty_f32 + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(out), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + _ptr_view_safe(out_scale_sorted), + token_num, + n_in, + k_in, + size_expert_ids_in, + stream, + ) + + +def _s2_args_fp4( + target, + a, + w, + a_scale, + w_scale, + sorted_ids, + sorted_expert_ids, + sorted_weights, + num_valid_ids, + token_num, + n_in, + k_in, + blocks, + dev, + bias=None, + stream=None, +): + _bias = ( + bias.view(-1) + if bias is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + if stream is None: + stream = torch.cuda.current_stream() + return ( + _ptr_view_safe(target), + _ptr_view_safe(a), + _ptr_view_safe(w), + _ptr_view_safe(a_scale), + _ptr_view_safe(w_scale), + _ptr_view_safe(sorted_ids), + _ptr_view_safe(sorted_expert_ids), + _ptr_view_safe(sorted_weights), + _ptr_view_safe(num_valid_ids), + _ptr_view_safe(_bias), + token_num, + n_in, + k_in, + blocks, + stream, + ) + + +def _run_compiled(exe, args): + """Call the JitFunction with the given args (handles compile caching).""" + try: + exe(*args) + except Exception: + # JitFunction.__call__ leaks ir.Context on compilation failure; clean up + # leaked contexts so subsequent calls do not take a wrong code path. + try: + while ir.Context.current is not None: + ir.Context.current.__exit__(None, None, None) + except Exception: + pass + raise + + +def build_moe_stage1_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage1: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + act: str = "silu", + persist_m: int = 1, + use_async_copy: bool = False, + k_batch: int = 1, + waves_per_eu: int = 3, + b_nt: int = 0, + gate_mode: str = "separated", + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + enable_bias: bool = False, + a_scale_one: bool = False, + xcd_swizzle: int = 0, + swiglu_limit: float = 0.0, +): + """Build (and cache) the inline FlyDSL a4w4 stage1 device kernel.""" + return compile_mixed_moe_gemm1( + 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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + act=act, + persist_m=persist_m, + use_async_copy=use_async_copy, + k_batch=k_batch, + waves_per_eu=waves_per_eu, + b_nt=b_nt, + gate_mode=GateMode(gate_mode), + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + enable_bias=enable_bias, + a_scale_one=a_scale_one, + xcd_swizzle=xcd_swizzle, + swiglu_limit=swiglu_limit, + ) + + +def build_moe_stage2_module( + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, + doweight_stage2: bool, + a_dtype: str = "fp4", + b_dtype: str = "fp4", + out_dtype: str = "bf16", + accumulate: bool = True, + persist_m: int = 1, + sort_block_m: int = 0, + b_nt: int = 0, + model_dim_pad: int = 0, + inter_dim_pad: int = 0, + xcd_swizzle: int = 0, + enable_bias: bool = False, +): + """Build (and cache) the inline FlyDSL a4w4 stage2 device kernel.""" + return compile_mixed_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, + a_dtype=a_dtype, + b_dtype=b_dtype, + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=persist_m, + sort_block_m=sort_block_m, + b_nt=b_nt, + model_dim_pad=model_dim_pad, + inter_dim_pad=inter_dim_pad, + xcd_swizzle=xcd_swizzle, + enable_bias=enable_bias, + ) + + +def _moe_stage1( + a, + w1, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + w1_scale, + a1_scale, + sorted_weights=None, + act="silu", + swiglu_limit=0.0, +): + """Host runner for the inline a4w4 stage1 (fp4/fp4 -> bf16, k_batch=1).""" + token_num = a.shape[0] + E = w1.shape[0] + inter_dim = w1.shape[1] // 2 + model_dim = a.shape[1] * 2 # a_dtype == "fp4": packed 2 values per byte + dev = a.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + out = torch.empty((token_num, topk, inter_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a1_scale.view(-1) if a1_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w1_scale.view(-1) if w1_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(0, device=dev, dtype=torch.float32) + ) + + _sort_block_m = tile_m + _all_blks = sorted_expert_ids.shape[0] + _dense_blks = ( + min(token_num * topk * _sort_block_m, sorted_token_ids.shape[0]) + // _sort_block_m + ) + _grid_y = min(_dense_blks, _all_blks) + + out_scale_sorted_flat = torch.empty(0, dtype=torch.uint8, device=dev) + _n_in = inter_dim * 2 + _k_in = model_dim + + args = _s1_args_fp4( + out.view(-1), + a.view(-1), + w1.view(-1), + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + out_scale_sorted_flat.view(-1), + token_num, + _n_in, + _k_in, + _grid_y, + dev, + bias=torch.empty(0, device=dev), + ) + + exe = build_moe_stage1_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=(sorted_weights is not None), + a_dtype="fp4", + b_dtype="fp4", + out_dtype=out_dtype, + act=act, + persist_m=1, + k_batch=1, + waves_per_eu=3, + b_nt=0, + gate_mode="separated", + swiglu_limit=swiglu_limit, + ) + _run_compiled(exe, args) + return out + + +def _moe_stage2( + inter_states, + w2, + sorted_token_ids, + sorted_expert_ids, + num_valid_ids, + topk, + *, + tile_m, + tile_n, + tile_k, + out_dtype, + mode, + w2_scale, + a2_scale, + sorted_weights, +): + """Host runner for the inline a4w4 stage2 (fp4/fp4 -> bf16, atomic).""" + token_num = inter_states.shape[0] + E = w2.shape[0] + model_dim = w2.shape[1] + inter_dim = inter_states.shape[2] * 2 # a_dtype == "fp4" + accumulate = mode != "reduce" + dev = inter_states.device + torch_out_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + alloc_fn = torch.zeros if accumulate else torch.empty + out = alloc_fn((token_num, model_dim), dtype=torch_out_dtype, device=dev) + + flat_a_scale = ( + a2_scale.view(-1) if a2_scale is not None else torch.empty(0, device=dev) + ) + flat_w_scale = ( + w2_scale.view(-1) if w2_scale is not None else torch.empty(0, device=dev) + ) + sw = ( + sorted_weights + if sorted_weights is not None + else torch.empty(sorted_token_ids.shape, dtype=torch.float32, device=dev) + ) + + m_blocks = min(sorted_expert_ids.shape[0], token_num * topk) + _persist_m = -1 if m_blocks > 256 else 1 + + _n_in = model_dim + _k_in = inter_dim + target = out + if not accumulate: + target = torch.empty( + (token_num * topk * model_dim,), device=out.device, dtype=out.dtype + ) + + args = _s2_args_fp4( + target, + inter_states, + w2, + flat_a_scale, + flat_w_scale, + sorted_token_ids, + sorted_expert_ids, + sw, + num_valid_ids, + token_num, + _n_in, + _k_in, + m_blocks, + dev, + bias=None, + ) + + exe = build_moe_stage2_module( + model_dim=model_dim, + inter_dim=inter_dim, + experts=E, + topk=topk, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=(sorted_weights is not None), + a_dtype="fp4", + b_dtype="fp4", + out_dtype=out_dtype, + accumulate=accumulate, + persist_m=_persist_m, + sort_block_m=0, + b_nt=0, + ) + _run_compiled(exe, args) + + if not accumulate: + torch.sum(target.view(token_num, topk, model_dim), dim=1, out=out) + return out + + +def flydsl_moe_swiglu( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + block_m: int = 32, + tile_n: int = 256, + tile_k: int = 256, + mode: str = "atomic", + swiglu_limit: float = 0.0, +) -> torch.Tensor: + """Run the inline FlyDSL a4w4 MoE stage1+stage2 path with SWIGLU activation. + + fp4/fp4 (MXFP4 acts + weights), but stage1 fuses swiglu (gate * sigmoid(1.702 + * gate) * (up + 1)) instead of silu. swiglu_limit=0.0 -> default clamp 7.0 on + both the kernel and the reference. Returns [T, model_dim] bf16. + """ + import aiter + from aiter import QuantType, dtypes + from aiter.fused_moe import moe_sorting + from aiter.ops.shuffle import ( + shuffle_scale_a16w4, + shuffle_weight, + shuffle_weight_a16w4, + ) + from aiter.utility.fp4_utils import e8m0_shuffle, moe_mxfp4_sort + + experts = w1.shape[0] + inter_dim = w1.shape[1] // 2 + model_dim = w1.shape[2] + token = hidden_states.shape[0] + topk = topk_ids.shape[1] + torch_dtype = hidden_states.dtype + out_dtype = "bf16" if torch_dtype == torch.bfloat16 else "f16" + + topk_ids = topk_ids.to(torch.int32).contiguous() + topk_weights = topk_weights.to(torch.float32).contiguous() + + # --- harness-side prep (NOT the kernel): mxfp4 quant, e8m0 scales, + # weight/scale preshuffle, sorted token dispatch --- + q_dtype = dtypes.fp4x2 + torch_quant = aiter.get_torch_quant(QuantType.per_1x32) + + w1_qt, w1_scale = torch_quant(w1.contiguous(), quant_dtype=q_dtype) + w2_qt, w2_scale = torch_quant(w2.contiguous(), quant_dtype=q_dtype) + w1_qt = w1_qt.view(experts, inter_dim * 2, model_dim // 2) + w2_qt = w2_qt.view(experts, model_dim, inter_dim // 2) + a1_qt, a1_scale = torch_quant(hidden_states.contiguous(), quant_dtype=q_dtype) + + sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, _ = moe_sorting( + topk_ids, topk_weights, experts, model_dim, torch_dtype, block_m + ) + + w1_qt_shuf = shuffle_weight(w1_qt, (16, 16)) + w2_qt_shuf = shuffle_weight_a16w4(w2_qt, 16, False) + w1_scale_shuf = e8m0_shuffle(w1_scale) + w2_scale_shuf = shuffle_scale_a16w4(w2_scale, experts, False) + a1_scale_sort = moe_mxfp4_sort( + a1_scale[:token, :].view(token, 1, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage1 gate/up GEMM + fused gated activation === + stage1_out = _moe_stage1( + a1_qt, + w1_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + w1_scale=w1_scale_shuf, + a1_scale=a1_scale_sort, + sorted_weights=None, + act="swiglu", + swiglu_limit=swiglu_limit, + ) + torch.cuda.synchronize() + + a2_qt, a2_scale = torch_quant(stage1_out.view(-1, inter_dim), quant_dtype=q_dtype) + a2_qt = a2_qt.view(token, topk, -1) + a2_scale_sort = moe_mxfp4_sort( + a2_scale[: token * topk, :].view(token, topk, -1), + sorted_ids=sorted_ids, + num_valid_ids=num_valid_ids, + token_num=token, + block_size=block_m, + ) + + # === FlyDSL device kernel: stage2 down GEMM + weighted top-k combine === + out = _moe_stage2( + a2_qt, + w2_qt_shuf, + sorted_ids, + sorted_expert_ids, + num_valid_ids, + topk, + tile_m=block_m, + tile_n=tile_n, + tile_k=tile_k, + out_dtype=out_dtype, + mode=mode, + w2_scale=w2_scale_shuf, + a2_scale=a2_scale_sort, + sorted_weights=sorted_weights, + ) + torch.cuda.synchronize() + return out diff --git a/tasks/torch2flydsl/moe_swiglu_kernel/model.py b/tasks/torch2flydsl/moe_swiglu_kernel/model.py new file mode 100644 index 00000000..b149f57d --- /dev/null +++ b/tasks/torch2flydsl/moe_swiglu_kernel/model.py @@ -0,0 +1,223 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Pure-PyTorch reference for the quantized a4w4 SWIGLU fused MoE. + +The op is a top-k Mixture-of-Experts feed-forward block evaluated in MXFP4 +(``float4_e2m1fn_x2`` values with e8m0 per-1x32 block scales), using the GPT-OSS +clamped SWIGLU gate instead of silu. A softmax router selects ``topk`` experts +per token; stage 1 runs a grouped gate/up GEMM followed by the SWIGLU activation; +stage 2 runs the down GEMM and combines the experts with the renormalized router +weights. GEMMs accumulate in fp32 over dequantized operands. + +The SWIGLU activation clamps the gate to ``limit`` and the linear branch to +``[-limit, limit]`` (limit = 7.0), then returns +``gate * sigmoid(1.702 * gate) * (linear + 1)``. Activations and weights are +quantized to MXFP4, the stage-1 result is re-quantized to MXFP4 before the down +GEMM, and the output is returned in bf16. The MXFP4 rounding and e8m0 +block-scale numerics implemented here match AMD's reference quantizer +bit-for-bit. +""" +import torch +import torch.nn as nn + +# MXFP4 (e2m1) decode table indexed by the 4-bit code (sign in bit 3). +_MXFP4_VALUES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) +_BLOCK = 32 +# log2(F4E2M1_MAX=6) floored -> dtypeMax = 2**2 used as the e8m0 scale divisor. +_FP4_DTYPE_MAX = 4.0 +_SWIGLU_ALPHA = 1.702 +_SWIGLU_LIMIT = 7.0 + + +def _f32_to_e8m0(x): + """Round positive fp32 magnitudes to biased e8m0 exponents (uint8).""" + u32 = x.contiguous().view(torch.int32) + exponent = ((u32 >> 23) & 0xFF).view(torch.uint32).to(torch.uint8) + nan_case = exponent == 0xFF + round_case = ((u32 & 0x400000) > 0) & ( + ((u32 & 0x200000) > 0) | ((u32 & 0x1FFFFF) > 0) | (exponent > 0) + ) + exponent[round_case] += 1 + exponent[nan_case] = 0xFF + return exponent + + +def _e8m0_to_f32(scale_e8m0_biased): + """Decode biased e8m0 exponents (uint8) back to fp32 power-of-two scales.""" + scale_e8m0_biased = scale_e8m0_biased.view(torch.uint8) + zero_case = scale_e8m0_biased == 0 + nan_case = scale_e8m0_biased == 0xFF + scale_f32 = scale_e8m0_biased.to(torch.int32) << 23 + scale_f32[zero_case] = 0x00400000 + scale_f32[nan_case] = 0x7F800001 + return scale_f32.view(torch.float32) + + +def _f32_to_e2m1_codes(x): + """Round fp32 values to MXFP4 (e2m1) 4-bit codes, saturating out-of-range + magnitudes and handling denormals (adapted from the torchao FP utilities).""" + EBITS, MBITS = 2, 1 + EBITS_F32, MBITS_F32 = 8, 23 + F32_EXP_BIAS = (1 << (EBITS_F32 - 1)) - 1 + exp_bias = (1 << (EBITS - 1)) - 1 + max_int = (1 << (EBITS + MBITS)) - 1 + sign_mask = 1 << (EBITS + MBITS) + magic_adder = (1 << (MBITS_F32 - MBITS - 1)) - 1 + max_normal = 2 ** ((1 << EBITS) - 1 - exp_bias) * ( + ((1 << (MBITS + 1)) - 1) / (2**MBITS) + ) + min_normal = 2 ** (1 - exp_bias) + denorm_exp = (F32_EXP_BIAS - exp_bias) + (MBITS_F32 - MBITS) + 1 + denorm_mask_int = denorm_exp << MBITS_F32 + denorm_mask_float = torch.tensor( + denorm_mask_int, dtype=torch.int32 + ).view(torch.float32) + + x = x.float().view(torch.int32) + sign = x & 0x80000000 + x = x ^ sign + x = x.view(torch.float) + + saturate_mask = x >= max_normal + denormal_mask = torch.logical_and( + torch.logical_not(saturate_mask), x < min_normal + ) + normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask)) + + denormal_x = x + denorm_mask_float + denormal_x = denormal_x.view(torch.int32) + denormal_x -= denorm_mask_int + denormal_x = denormal_x.to(torch.uint8) + + normal_x = x.view(torch.int32) + mant_odd = (normal_x >> (MBITS_F32 - MBITS)) & 1 + val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder + normal_x += val_to_add + normal_x += mant_odd + normal_x = normal_x >> (MBITS_F32 - MBITS) + normal_x = normal_x.to(torch.uint8) + + codes = torch.full_like(x, max_int, dtype=torch.uint8) + codes = torch.where(denormal_mask, denormal_x, codes) + codes = torch.where(normal_mask, normal_x, codes) + + sign_lp = sign >> (MBITS_F32 + EBITS_F32 - MBITS - EBITS) + sign_lp = sign_lp.to(torch.uint8) & sign_mask + return (codes | sign_lp).to(torch.uint8) + + +def _mxfp4_dequant(x): + """MXFP4 per-1x32 e8m0 quantize+dequantize over the last dim, returning the + fp32 values the hardware GEMM sees.""" + shape = x.shape + xb = x.float().reshape(-1, _BLOCK) + max_abs = torch.amax(torch.abs(xb), dim=1) + scale_e8m0 = _f32_to_e8m0(max_abs / _FP4_DTYPE_MAX) + scale_f32 = _e8m0_to_f32(scale_e8m0).view(-1, 1) + codes = _f32_to_e2m1_codes(xb / scale_f32) + table = _MXFP4_VALUES.to(x.device) + deq = table[codes.long()] * scale_f32 + return deq.reshape(shape) + + +def _swiglu(gate, linear, alpha=_SWIGLU_ALPHA, limit=_SWIGLU_LIMIT): + """GPT-OSS clamped SWIGLU: clamp(gate)<=limit, clamp(linear) in [-limit, limit], + then gate*sigmoid(alpha*gate)*(linear + 1).""" + gate = gate.clamp(min=None, max=limit) + linear = linear.clamp(min=-limit, max=limit) + return gate * torch.sigmoid(alpha * gate) * (linear + 1) + + +def _grouped_gemm_stage1(acts, weights, topk_ids): + """Per-expert grouped GEMM: out[b, k] = acts[b] @ weights[topk_ids[b, k]].T.""" + acts = acts.float() + B, D = acts.shape + topk = topk_ids.shape[1] + N = weights.shape[1] + h = acts.view(B, 1, D).repeat(1, topk, 1) + out = torch.zeros(B, topk, N, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = h[mask] @ weights[e].transpose(0, 1) + return out + + +def _grouped_gemm_stage2(acts, weights, topk_ids, topk_weights): + """Per-expert down GEMM with weighted top-k combine to a single output row.""" + acts = acts.float() + B, topk = topk_ids.shape + model_dim = weights.shape[1] + out = torch.zeros(B, topk, model_dim, dtype=torch.float32, device=acts.device) + for e in range(weights.shape[0]): + mask = topk_ids == e + if mask.any(): + out[mask] = acts[mask] @ weights[e].transpose(0, 1) + out = out * topk_weights.view(B, topk, 1) + return out.sum(1) + + +def route_topk(logits, topk): + """Softmax router + top-k with renormalized weights. Shared by the harness. + + Ties are broken by ascending expert index via a stable descending sort. The + bf16 gate produces many duplicate logits across the large expert count, and a + nondeterministic top-k tie-break would let the reference and the runtime op + select different experts; the stable order keeps both routings identical. + """ + gate = torch.softmax(logits.float(), dim=-1) + order = torch.sort(gate, dim=-1, descending=True, stable=True).indices + ids = order[..., :topk] + weights = torch.gather(gate, -1, ids) + weights = weights / weights.sum(dim=-1, keepdim=True) + return weights.float(), ids.to(torch.int32) + + +class Model(nn.Module): + def __init__(self, model_dim, inter_dim, experts, topk, activation="swiglu"): + super().__init__() + self.model_dim = model_dim + self.inter_dim = inter_dim + self.experts = experts + self.topk = topk + self.activation = activation + self.swiglu_limit = _SWIGLU_LIMIT + self.gate = nn.Linear(model_dim, experts, bias=False).to(torch.bfloat16) + self.w1 = nn.Parameter( + (torch.randn(experts, 2 * inter_dim, model_dim) / 10).to(torch.bfloat16) + ) + self.w2 = nn.Parameter( + (torch.randn(experts, model_dim, inter_dim) / 10).to(torch.bfloat16) + ) + + def forward(self, hidden_states): + I = self.inter_dim + + logits = self.gate(hidden_states) + topk_weights, topk_ids = route_topk(logits, self.topk) + + a1 = _mxfp4_dequant(hidden_states) + w1 = _mxfp4_dequant(self.w1) + w2 = _mxfp4_dequant(self.w2) + + stage1 = _grouped_gemm_stage1(a1, w1, topk_ids) + gate, up = stage1.split([I, I], dim=-1) + stage1 = _swiglu(gate, up, limit=self.swiglu_limit).to(torch.bfloat16) + + a2 = _mxfp4_dequant(stage1.reshape(-1, I)).reshape( + hidden_states.shape[0], self.topk, I + ) + + out = _grouped_gemm_stage2(a2, w2, topk_ids, topk_weights) + return out.to(torch.float16).to(torch.bfloat16) + + +def get_inputs(): + return [torch.randn(16, 3072, dtype=torch.bfloat16)] + + +def get_init_inputs(): + return [3072, 512, 128, 4] diff --git a/tasks/torch2flydsl/moe_swiglu_kernel/test_kernel_harness.py b/tasks/torch2flydsl/moe_swiglu_kernel/test_kernel_harness.py new file mode 100644 index 00000000..ed0168a7 --- /dev/null +++ b/tasks/torch2flydsl/moe_swiglu_kernel/test_kernel_harness.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Correctness and performance harness for the a4w4 SWIGLU fused MoE task. + +The pure-torch reference in ``model.py`` and the FlyDSL kernel share the same +top-k routing (``model.route_topk``) so expert selection is identical. The +correctness gate is the normalized max error ``max|ref - out| / max|ref|``, which +must stay <= ``REL_TOL``; element-wise close% at 1e-2 and 1e-1 is also reported. +The check asserts and exits non-zero on failure. + +Modes: + --correctness compare the kernel against the reference + --full-benchmark time the kernel vs the reference and write a perf report +""" +import argparse +import importlib.util +import json +import math +import os +import sys +from pathlib import Path + +KERNEL_FILE = "kernel.py" +MODEL_FILE = "model.py" + + +def _resolve_kernel_dir(): + here = os.path.dirname(os.path.abspath(__file__)) + if os.path.isfile(os.path.join(here, KERNEL_FILE)): + return here + cwd = os.getcwd() + if os.path.isfile(os.path.join(cwd, KERNEL_FILE)): + return cwd + return here + + +def _load_module(kernel_dir, filename, alias): + entry = os.path.join(kernel_dir, filename) + if not os.path.isfile(entry): + return None + if kernel_dir not in sys.path: + sys.path.insert(0, kernel_dir) + spec = importlib.util.spec_from_file_location(alias, entry) + if spec is None or spec.loader is None: + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +_KERNEL_DIR = _resolve_kernel_dir() + +# Real a4w4 fp4 SWIGLU fused-MoE shapes (q_dtype_a/w = float4_e2m1fn_x2, +# per_1x32, act_type = ActivationType.Swiglu) from gptoss_fp4_untuned_fmoe.csv +# (GPT-OSS): D=3072, E=128, topk=4, inter_dim in {512, 1536}. +SHAPES = [ + {"name": "gptoss_t16_i512_e128_k4", "tokens": 16, "model_dim": 3072, "inter_dim": 512, "experts": 128, "topk": 4}, + {"name": "gptoss_t32_i1536_e128_k4", "tokens": 32, "model_dim": 3072, "inter_dim": 1536, "experts": 128, "topk": 4}, +] + +# Tight element-wise gate: normalized max error <= REL_TOL. +REL_TOL = 1e-2 +SEED = 20260401 +BLOCK_M, TILE_N, TILE_K, MODE = 32, 256, 256, "atomic" +# Correctness uses the deterministic "reduce" combine. The "atomic" stage-2 +# combine sums per-expert partials with order-dependent fp32 atomic-adds, so its +# result is nondeterministic run-to-run; "reduce" computes the identical math +# with a deterministic reduction (same numerics, reproducible comparison). +CORRECTNESS_MODE = "reduce" + + +def _build_model(mmod, shape, device="cuda"): + import torch + + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + model = mmod.Model( + model_dim=shape["model_dim"], inter_dim=shape["inter_dim"], + experts=shape["experts"], topk=shape["topk"], + ).to(device).eval() + hidden = torch.randn( + shape["tokens"], shape["model_dim"], dtype=torch.bfloat16, device=device + ) + return model, hidden + + +def _kernel_out(kmod, mmod, model, hidden, topk): + # Recompute the SAME routing the reference used and run the FlyDSL kernel. + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + return kmod.flydsl_moe_swiglu( + hidden, model.w1.detach(), model.w2.detach(), topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=CORRECTNESS_MODE, + ) + + +def run_correctness(verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + failures = [] + for shape in SHAPES: + model, hidden = _build_model(mmod, shape) + with torch.no_grad(): + ref = model(hidden).float() + out = _kernel_out(kmod, mmod, model, hidden, shape["topk"]).float() + torch.cuda.synchronize() + + max_abs = (ref - out).abs().max().item() + ref_scale = ref.abs().max().item() + 1e-9 + rel_err = max_abs / ref_scale + max_rel = ((ref - out).abs() / (ref.abs() + 1e-9)).max().item() + pct1e2 = torch.isclose(ref, out, atol=1e-2, rtol=1e-2).float().mean().item() * 100 + pct1e1 = torch.isclose(ref, out, atol=1e-1, rtol=1e-1).float().mean().item() * 100 + ok = rel_err <= REL_TOL + if verbose: + print( + f" {'PASS' if ok else 'FAIL'}: {shape['name']} " + f"(D{shape['model_dim']}/I{shape['inter_dim']}/E{shape['experts']}/k{shape['topk']}) " + f"norm_max_err={rel_err:.5f} (tol={REL_TOL}) " + f"max_abs={max_abs:.4f} max_rel={max_rel:.3f} " + f"close%@1e-2={pct1e2:.2f} @1e-1={pct1e1:.2f}" + ) + if not ok: + failures.append(shape["name"]) + + status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(SHAPES)})" + print(f"Status: {status}") + print(f"correctness: {'pass' if not failures else 'fail'}") + assert not failures, f"correctness FAILED for: {failures}" + return True + + +def run_benchmark(warmup=10, iters=100, verbose=True): + import torch + + kmod = _load_module(_KERNEL_DIR, KERNEL_FILE, "flydsl_kernel") + mmod = _load_module(_KERNEL_DIR, MODEL_FILE, "torch_model") + assert kmod is not None and mmod is not None, "cannot load kernel.py / model.py" + + latencies, speedups, report = [], [], [] + print(f"{'Config':<24} {'Ref':>10} {'FlyDSL':>10} {'Speedup':>10}") + print("-" * 60) + for idx, shape in enumerate(SHAPES): + model, hidden = _build_model(mmod, shape) + topk = shape["topk"] + with torch.no_grad(): + logits = model.gate(hidden) + topk_weights, topk_ids = mmod.route_topk(logits, topk) + w1, w2 = model.w1.detach(), model.w2.detach() + + def run_kernel(): + return kmod.flydsl_moe_swiglu( + hidden, w1, w2, topk_weights, topk_ids, + block_m=BLOCK_M, tile_n=TILE_N, tile_k=TILE_K, mode=MODE, + ) + + run_kernel() + torch.cuda.synchronize() + for _ in range(warmup): + run_kernel() + torch.cuda.synchronize() + ktimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); run_kernel(); e.record(); torch.cuda.synchronize() + ktimes.append(s.elapsed_time(e)) + kernel_ms = sum(ktimes) / len(ktimes) + + rtimes = [] + for _ in range(iters): + s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True) + s.record(); model(hidden); e.record(); torch.cuda.synchronize() + rtimes.append(s.elapsed_time(e)) + ref_ms = sum(rtimes) / len(rtimes) + + speedup = ref_ms / kernel_ms if kernel_ms > 0 else 1.0 + latencies.append(kernel_ms); speedups.append(speedup) + report.append({ + "test_case_id": f"test_case_{idx}", + "execution_time_ms": kernel_ms, + "shape": [shape["tokens"], shape["model_dim"], shape["inter_dim"]], + "params": {k: shape[k] for k in ("tokens", "model_dim", "inter_dim", "experts", "topk")}, + }) + if verbose: + print(f"{shape['name']:<24} {ref_ms:>8.4f}ms {kernel_ms:>8.4f}ms {speedup:>8.2f}x") + del model, hidden + torch.cuda.empty_cache() + + geomean_latency = math.exp(sum(math.log(x) for x in latencies) / len(latencies)) + geomean_speedup = math.exp(sum(math.log(x) for x in speedups) / len(speedups)) + + build_dir = Path(_KERNEL_DIR) / "build" + build_dir.mkdir(exist_ok=True) + with open(build_dir / "performance_report.json", "w") as f: + json.dump(report, f, indent=2) + + print("-" * 60) + print(f"Geometric mean latency: {geomean_latency:.4f} ms") + print(f"Geometric mean speedup: {geomean_speedup:.2f}x") + return {"geomean_latency_ms": geomean_latency, "geomean_speedup": geomean_speedup} + + +if __name__ == "__main__": + try: + import torch as _t + _arch = _t.cuda.get_device_properties(0).gcnArchName.split(":")[0] + except Exception: + _arch = "" + if _arch != "gfx950": + print(f"SKIPPED: gfx950-only task on arch={_arch or 'unknown'} (FP4/MX scaled-MFMA requires CDNA4/gfx950)") + print("correctness: skip") + sys.exit(0) + parser = argparse.ArgumentParser(description="torch2flydsl moe harness") + parser.add_argument("--correctness", action="store_true") + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--full-benchmark", action="store_true") + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + print("=" * 60) + print("torch2flydsl MoE (a4w4 swiglu, quantized reference)") + print("=" * 60) + + if args.correctness: + try: + run_correctness() + except AssertionError as exc: + print(f"ASSERTION: {exc}") + sys.exit(1) + sys.exit(0) + else: + run_benchmark(warmup=args.warmup, iters=args.iterations)