diff --git a/README.md b/README.md index 4ef11906b..3009fa863 100644 --- a/README.md +++ b/README.md @@ -368,7 +368,6 @@ See `examples/` for more examples including tiled copy (`02-tiledCopy.py`), tile | **Blockscale GEMM** | `test_blockscale_preshuffle_gemm.py` | Blockscale preshuffle GEMM | | **HGEMM Split-K** | `test_hgemm_splitk.py` | FP16 GEMM split-K | | **MoE GEMM** | `test_moe_gemm.py` | MoE 2-stage (gate/up + reduce) | -| **MoE Blockscale** | `test_moe_blockscale.py` | MoE blockscale 2-stage | | **MoE Reduce** | `test_moe_reduce.py` | MoE reduce kernel | | **PagedAttention** | `test_pa.py` | Paged attention decode (FP8) — *WIP perf tuning* | | **FlashAttention** | `test_flash_attn_fwd.py` | Flash attention — *WIP perf tuning* | diff --git a/docs/api/kernels.rst b/docs/api/kernels.rst index 8a3fd2627..83246ca4e 100644 --- a/docs/api/kernels.rst +++ b/docs/api/kernels.rst @@ -21,8 +21,7 @@ MoE (Mixture-of-Experts) Kernels - ``kernels.moe.moe_gemm_2stage`` -- MoE GEMM with 2-stage pipeline (stage1 + stage2). Also provides the MoE reduction (sum over the topk dimension, ``Y[t, d] = sum(X[t, :, d])``), compiled via ``compile_moe_reduction()``. -- ``kernels.moe.mixed_moe_gemm_2stage`` -- Mixed-precision MoE GEMM -- ``kernels.moe.moe_blockscale_2stage`` -- MoE with block-scale quantization (MXFP4) +- ``kernels.moe.mxfp_moe`` -- Fused a4w4 / a8w4 MoE 2-stage GEMM (device-side fp4 re-quant) Paged Attention ---------------- diff --git a/docs/architecture_guide.md b/docs/architecture_guide.md index affeb645b..e52155249 100644 --- a/docs/architecture_guide.md +++ b/docs/architecture_guide.md @@ -103,8 +103,7 @@ FlyDSL/ │ │ └── fused_rope_cache_kernel.py # Fused RoPE + KV cache │ ├── moe/ # Mixture-of-experts kernels │ │ ├── moe_gemm_2stage.py # MoE GEMM (2-stage gate/up + reduce) -│ │ ├── moe_blockscale_2stage.py # MoE Blockscale GEMM -│ │ ├── mixed_moe_gemm_2stage.py # Mixed-precision MoE GEMM +│ │ ├── mxfp_moe/ # Fused a4w4/a8w4 MoE 2-stage (device fp4 re-quant) │ │ ├── moe_sorting_kernel.py # MoE token sorting │ │ └── topk_gating_softmax_kernel.py # Top-k gating softmax │ ├── mma/ # Shared MMA pipeline helpers diff --git a/docs/prebuilt_kernels_guide.md b/docs/prebuilt_kernels_guide.md index d248a3924..d92890126 100644 --- a/docs/prebuilt_kernels_guide.md +++ b/docs/prebuilt_kernels_guide.md @@ -348,7 +348,6 @@ What operation do you need? │ ├── MoE (Mixture of Experts) │ ├── Blockscale MoE (gate+up+reduce) -│ │ └── → kernels/moe/moe_blockscale_2stage.py │ └── Standard MoE (fp8/f16/bf16/int8/int4) │ └── → kernels/moe/moe_gemm_2stage.py │ @@ -368,8 +367,7 @@ What operation do you need? | `kernels/gemm/blockscale_preshuffle_gemm.py` | Blockscale GEMM | | `kernels/gemm/hgemm_splitk.py` | FP16 GEMM split-K | | `kernels/moe/moe_gemm_2stage.py` | MoE GEMM 2-stage (gate/up + reduce) | -| `kernels/moe/moe_blockscale_2stage.py` | MoE Blockscale 2-stage | -| `kernels/moe/mixed_moe_gemm_2stage.py` | Mixed-precision MoE GEMM | +| `kernels/moe/mxfp_moe/` | Fused a4w4/a8w4 MoE 2-stage GEMM (device fp4 re-quant) | | `kernels/attention/pa_decode_fp8.py` | Paged attention decode (FP8) | | `kernels/attention/flash_attn_generic.py` | FlashAttention generic fallback | | `kernels/attention/flash_attn_gfx950.py` | FlashAttention gfx950 bf16/f16 fast path | @@ -398,7 +396,6 @@ What operation do you need? | `tests/kernels/test_blockscale_preshuffle_gemm.py` | Blockscale GEMM | | `tests/kernels/test_hgemm_splitk.py` | FP16 GEMM split-K | | `tests/kernels/test_moe_gemm.py` | MoE GEMM | -| `tests/kernels/test_moe_blockscale.py` | MoE Blockscale GEMM | | `tests/kernels/test_moe_reduce.py` | MoE reduce kernel | | `tests/kernels/test_pa.py` | Paged attention decode | | `tests/kernels/test_flash_attn_fwd.py` | FlashAttention | diff --git a/docs/testing_benchmarking_guide.md b/docs/testing_benchmarking_guide.md index 3ceb8d7eb..4a2b282e0 100644 --- a/docs/testing_benchmarking_guide.md +++ b/docs/testing_benchmarking_guide.md @@ -69,7 +69,6 @@ Full end-to-end tests: compile FlyDSL kernels, execute on GPU, validate against | `test_preshuffle_gemm.py` | GEMM | Preshuffle MFMA GEMM (fp8/int8/fp16/bf16) | | `test_blockscale_preshuffle_gemm.py` | GEMM | Block-scale (MXFP4) preshuffle GEMM | | `test_moe_gemm.py` | MoE GEMM | Mixture-of-Experts GEMM | -| `test_moe_blockscale.py` | MoE | MoE with block-scale quantization | | `test_moe_reduce.py` | MoE Reduce | MoE reduction kernel | | `test_pa.py` | Paged Attn | Paged attention decode | | `test_quant.py` | Quantization | Quantization ops | diff --git a/kernels/common/tensor_shim.py b/kernels/common/tensor_shim.py index a47cd2a2c..7509b77e4 100644 --- a/kernels/common/tensor_shim.py +++ b/kernels/common/tensor_shim.py @@ -19,13 +19,23 @@ def _run_compiled(exe, *args): """First call: ``flyc.compile(exe, *args)`` compiles **and** executes the kernel. Subsequent calls: fast dispatch via the cached ``CompiledFunction``. + + A failed cold compile can leave an MLIR ``Context`` open on the stack; unwind + it before re-raising so the next attempt starts clean. """ cf = getattr(exe, "_cf", None) - if cf is None: - cf = flyc.compile(exe, *args) - exe._cf = cf - else: + if cf is not None: cf(*args) + return + try: + exe._cf = flyc.compile(exe, *args) + except Exception: + try: + while ir.Context.current is not None: + ir.Context.current.__exit__(None, None, None) + except Exception: + pass + raise def _to_raw(v): diff --git a/kernels/moe/mixed_moe_gemm_2stage/__init__.py b/kernels/moe/mixed_moe_gemm_2stage/__init__.py deleted file mode 100644 index 1d59653a8..000000000 --- a/kernels/moe/mixed_moe_gemm_2stage/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Mixed-dtype MoE 2-stage MFMA kernels (stage1 / stage2). - -Split from the former monolith; public API unchanged. -""" - -from kernels.moe.mixed_moe_gemm_2stage.gemm1 import compile_mixed_moe_gemm1 -from kernels.moe.mixed_moe_gemm_2stage.gemm2 import compile_mixed_moe_gemm2 - -__all__ = [ - "compile_mixed_moe_gemm1", - "compile_mixed_moe_gemm2", -] diff --git a/kernels/moe/mixed_moe_gemm_2stage/common.py b/kernels/moe/mixed_moe_gemm_2stage/common.py deleted file mode 100644 index b8c71850e..000000000 --- a/kernels/moe/mixed_moe_gemm_2stage/common.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Mixed-dtype MoE 2-stage MFMA kernels (stage1 / stage2). :: _common""" - -import os - -from flydsl._mlir.dialects import llvm - - -def _get_cu_num() -> int: - env = os.environ.get("CU_NUM") - if env: - return int(env) - try: - import torch - - return int(torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count) - except Exception: - return 304 - - -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, - ) diff --git a/kernels/moe/mixed_moe_gemm_2stage/gemm1.py b/kernels/moe/mixed_moe_gemm_2stage/gemm1.py deleted file mode 100644 index 02b46dd2a..000000000 --- a/kernels/moe/mixed_moe_gemm_2stage/gemm1.py +++ /dev/null @@ -1,2366 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Mixed-dtype MoE 2-stage MFMA kernels (stage1 / stage2). :: _gemm1""" - -import functools -import os - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm, memref, scf, vector -from flydsl._mlir.dialects.arith import CmpIPredicate -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr.typing import T -from flydsl.runtime.device import get_rocm_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr -from kernels.common import buffer_ops -from kernels.common.layout_utils import crd2idx, idx2crd -from kernels.common.layout_utils import get as layout_get -from kernels.common.mma.mfma_epilogues import c_shuffle_epilog -from kernels.common.mma.mfma_preshuffle_pipeline import ( - _buffer_load_vec, - buffer_copy_gmem16_dwordx4, - lds_store_16b_xor16, - make_preshuffle_b_layout, - make_preshuffle_scale_layout, - swizzle_xor16, - tile_chunk_coord_i32, -) -from kernels.moe.mixed_moe_gemm_2stage.common import _barrier -from kernels.moe.moe_common import ( - GateMode, -) # noqa: F401 re-exported for back-compat -from kernels.moe.moe_common import ( - i64x4_to_i32x8 as _pack, -) -from kernels.moe.moe_common import ( - i64x4_to_i32x8 as pack_i64x4_to_i32x8, -) - - -@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_rocm_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 - - # 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 - - _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.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_num_valid_ids: fx.Tensor, - arg_bias: fx.Tensor, - arg_out_scale_sorted: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_n_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - ): - - tokens_in = arith.index_cast(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) - - 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 = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_i32) - - w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False) - - # Out: [tokens*topk, inter_dim] - numids_rsrc = buffer_ops.create_buffer_resource( - arg_num_valid_ids, - max_size=False, - num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_scale_x, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_scale_w, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_sorted_token_ids, - max_size=False, - num_records_bytes=sorted_nbytes_i32, - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_weights, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_expert_ids, max_size=False, num_records_bytes=eid_nbytes_i32 - ) - bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=False) 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): - sorted_scale_rsrc = buffer_ops.create_buffer_resource(arg_out_scale_sorted, max_size=False) - - # ---- 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), as_ir_value(x_vec))) - return parts - - # Wave/lane decomposition (identical to stage2) - coord_wl = idx2crd(fx.Int32(tx), layout_tx_wave_lane) - wave_id = layout_get(coord_wl, 0) - lane_id = layout_get(coord_wl, 1) - coord_l16 = idx2crd(fx.Int32(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(fx.Int32(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(fx.Int32(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(tuple(fx.Int32(c) for c in 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, as_ir_value(b16)) - b0 = vector.extract(as_ir_value(b_i64x2), static_position=[0], dynamic_position=[]) - b1 = vector.extract(as_ir_value(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), [as_ir_value(_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), [as_ir_value(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), [as_ir_value(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), [as_ir_value(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([fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)], layout_lds) - loaded_a16 = vector.load(vec16_x, as_ir_value(lds_buffer), [as_ir_value(idx_a16)]) - a_i64x2 = vector.bitcast(vec2_i64, as_ir_value(loaded_a16)) - a0 = vector.extract(as_ir_value(a_i64x2), static_position=[0], dynamic_position=[]) - a1 = vector.extract(as_ir_value(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) - - _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( - as_ir_value(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( - as_ir_value(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( - as_ir_value(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) - - 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( - as_ir_value(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( - as_ir_value(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( - as_ir_value(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), - [as_ir_value(_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), [as_ir_value(_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), [as_ir_value(_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), [as_ir_value(_tid_val)]) - vector.store(as_ir_value(_tid_vec1), as_ir_value(lds_tid), [as_ir_value(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(as_ir_value(gate_v4), static_position=[ei], dynamic_position=[]) - u = vector.extract(as_ir_value(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, [as_ir_value(e) for e in 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(as_ir_value(gate_v4), static_position=[ei], dynamic_position=[]) - u = vector.extract(as_ir_value(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, [as_ir_value(e) for e in 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, [as_ir_value(e) for e in [_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, [as_ir_value(e) for e in [_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 - - from flydsl._mlir.dialects import fly as _fly - - _llvm_ptr_ty = ir.Type.parse("!llvm.ptr") - out_base_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty, arg_out) - out_base_i64 = llvm.ptrtoint(T.i64, out_base_ptr) - out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) - - 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(as_ir_value(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, [as_ir_value(v)]) - vector.store( - as_ir_value(v1), - as_ir_value(lds_out), - [as_ir_value(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, [as_ir_value(v_out)]) - vector.store( - as_ir_value(v1), - as_ir_value(lds_out), - [as_ir_value(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 gemm_common_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(as_ir_value(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 gemm_common_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.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_max_token_ids: fx.Tensor, - arg_bias: fx.Tensor, - arg_out_scale_sorted: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_inter_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - stream: fx.Stream, - ): - _ = _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_in = arith.index_cast(ir.IndexType.get(), i32_inter_in.ir_value()) - tile_n_index = arith.constant(tile_n, index=True) - inter_dim_pad_total = arith.constant(2 * inter_dim_pad, index=True) - if const_expr(mock_gate_only or gate_up_interleave): - gx = (inter_in - inter_dim_pad_total + tile_n_index - 1) / tile_n_index - else: - gx = (inter_in - inter_dim_pad_total + 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 diff --git a/kernels/moe/mixed_moe_gemm_2stage/gemm2.py b/kernels/moe/mixed_moe_gemm_2stage/gemm2.py deleted file mode 100644 index a1bd89f82..000000000 --- a/kernels/moe/mixed_moe_gemm_2stage/gemm2.py +++ /dev/null @@ -1,1623 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""Mixed-dtype MoE 2-stage MFMA kernels (stage1 / stage2). :: _gemm2""" - -import functools -import os - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm, memref, scf, vector -from flydsl._mlir.dialects.arith import CmpIPredicate -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr.typing import T -from flydsl.runtime.device import get_rocm_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr -from kernels.common import buffer_ops -from kernels.common.layout_utils import crd2idx, idx2crd -from kernels.common.layout_utils import get as layout_get -from kernels.common.mem_ops import buffer_atomic_add -from kernels.common.mma.mfma_epilogues import c_shuffle_epilog -from kernels.common.mma.mfma_preshuffle_pipeline import ( - _buffer_load_vec, - buffer_copy_gmem16_dwordx4, - lds_store_4b_xor16, - lds_store_8b_xor16, - lds_store_16b_xor16, - make_preshuffle_scale_layout, - swizzle_xor16, - tile_chunk_coord_i32, -) -from kernels.moe.mixed_moe_gemm_2stage.common import _get_cu_num -from kernels.moe.moe_common import ( - i64x4_to_i32x8 as pack_i64x4_to_i32x8, -) - - -@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_rocm_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" - # 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: - _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.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_num_valid_ids: fx.Tensor, - arg_bias: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_n_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - ): - - tokens_in = arith.index_cast(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) - - 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 - from kernels.common.layout_utils import _div_pow2, _mod_pow2 - - 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 = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_i32) - - w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False) - - # OUT: [tokens, model_dim] -> clamp to descriptor max (i32 bytes) to avoid overflow on huge tokens. - out_elem_bytes = 4 if out_is_f32 else 2 - out_nbytes_idx = tokens_in * n_in * 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 = buffer_ops.create_buffer_resource(arg_out, max_size=False, num_records_bytes=out_nbytes_i32) - - # num_valid_ids (sorted padded MN) for scale sizing / guards. - numids_rsrc = buffer_ops.create_buffer_resource( - arg_num_valid_ids, - max_size=False, - num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_scale_x, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_scale_x, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_scale_w, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_sorted_token_ids, - max_size=False, - num_records_bytes=sorted_nbytes_i32, - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_weights, max_size=False, num_records_bytes=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 = buffer_ops.create_buffer_resource( - arg_expert_ids, max_size=False, num_records_bytes=eid_nbytes_i32 - ) - bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=False) 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, as_ir_value(x_vec))) - elif const_expr(x_load_bytes == 8): - parts.append(vector.bitcast(vec2_i32, as_ir_value(x_vec))) - else: - parts.append(vector.bitcast(vec1_i32, as_ir_value(x_vec))) - return parts - - # tx -> wave/lane (GEMM-style decomposition). - coord_wl = idx2crd(fx.Int32(tx), layout_tx_wave_lane) - wave_id = layout_get(coord_wl, 0) - lane_id = layout_get(coord_wl, 1) - coord_l16 = idx2crd(fx.Int32(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, as_ir_value(b16)) - b0 = vector.extract(as_ir_value(b_i64x2), static_position=[0], dynamic_position=[]) - b1 = vector.extract(as_ir_value(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), [as_ir_value(s)]) - - def _apply_k_shift(scale_vec, k_shift_bits): - if const_expr(k_shift_bits > 0): - val = vector.extract(as_ir_value(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), [as_ir_value(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([fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)], layout_lds) - loaded_a16 = vector.load(vec16_x, as_ir_value(lds_buffer), [as_ir_value(idx_a16)]) - a_i64x2 = vector.bitcast(vec2_i64, as_ir_value(loaded_a16)) - a0 = vector.extract(as_ir_value(a_i64x2), static_position=[0], dynamic_position=[]) - a1 = vector.extract(as_ir_value(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) - - # 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( - as_ir_value(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( - as_ir_value(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), [as_ir_value(_tid_val)]) - vector.store(as_ir_value(_tid_vec1), as_ir_value(lds_tid), [as_ir_value(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): - buffer_atomic_add(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). - from flydsl._mlir.dialects import fly as _fly - - _llvm_ptr_ty = ir.Type.parse("!llvm.ptr") - out_base_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty, arg_out) - out_base_i64 = llvm.ptrtoint(T.i64, out_base_ptr) - out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) - - 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(as_ir_value(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, [as_ir_value(v_out)]) - - vector.store(as_ir_value(v1), as_ir_value(lds_out), [as_ir_value(lds_idx)], alignment=2) - - def precompute_row(*, row_local, row): - # 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.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_num_valid_ids: fx.Tensor, - arg_bias: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_n_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - stream: fx.Stream, - ): - _ = _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 diff --git a/kernels/moe/moe_blockscale_2stage/__init__.py b/kernels/moe/moe_blockscale_2stage/__init__.py deleted file mode 100644 index 111e73053..000000000 --- a/kernels/moe/moe_blockscale_2stage/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""MoE block-scale 2-stage MFMA kernels (stage1 / stage2 / reduction). - -Split from the former monolith; public API unchanged. -""" - -from kernels.moe.moe_blockscale_2stage.gemm1 import compile_moe_blockscale_gemm1 -from kernels.moe.moe_blockscale_2stage.gemm2 import ( - MoeGemm2Mode, - _MoeGemm2ReduceWrapper, - compile_moe_blockscale_gemm2, - compile_moe_blockscale_gemm2_ex, -) -from kernels.moe.moe_blockscale_2stage.reduction import compile_moe_reduction - -__all__ = [ - "compile_moe_blockscale_gemm1", - "compile_moe_blockscale_gemm2", - "compile_moe_blockscale_gemm2_ex", - "compile_moe_reduction", - "MoeGemm2Mode", - "_MoeGemm2ReduceWrapper", -] diff --git a/kernels/moe/moe_blockscale_2stage/gemm1.py b/kernels/moe/moe_blockscale_2stage/gemm1.py deleted file mode 100644 index 9373af18f..000000000 --- a/kernels/moe/moe_blockscale_2stage/gemm1.py +++ /dev/null @@ -1,1146 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""MoE block-scale 2-stage MFMA kernels (stage1 / stage2 / reduction). :: _gemm1""" - -import functools -import os - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import math as math_dialect -from flydsl._mlir.dialects import scf, vector -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr.arith import ArithValue -from flydsl.expr.typing import T -from flydsl.runtime.device import get_rocm_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr -from kernels.common import buffer_ops -from kernels.common.kernels_common import _if_then -from kernels.common.mma.mfma_epilogues import mfma_epilog -from kernels.common.mma.mfma_preshuffle_pipeline import ( - buffer_copy_gmem16_dwordx4, - lds_store_4b_xor16, - lds_store_8b_xor16, - lds_store_16b_xor16, - load_b_pack_k32, - make_preshuffle_b_layout, - preshuffle_crd2idx, - swizzle_xor16, - tile_chunk_coord_i32, -) -from kernels.moe.moe_common import ( - i64_to_v4f16 as _i64_to_v4f16, -) -from kernels.moe.moe_common import ( - i64x4_to_i32x8 as _pack128, -) - - -@functools.lru_cache(maxsize=1024) -def compile_moe_blockscale_gemm1( - *, - model_dim: int, - inter_dim: int, - experts: int, - topk: int, - tile_m: int, - tile_n: int, - tile_k: int, - doweight_stage1: bool, - scale_block_k: int = 128, - out_dtype: str = "f16", - use_cshuffle_epilog: bool | None = None, - waves_per_eu: int | None = None, -): - """Compile stage1 kernel (`moe_gemm1`) and return the compiled executable. - - in_dtype: - - "fp8": X/W are fp8 - - "fp16": X/W are fp16 - - "int8": X/W are int8 (X is [tokens, K]) - - "int8smooth": X/W are int8, but X is pre-expanded to [tokens*topk, K] with per-(token,slot) - quant scales (used to emulate MoE smoothquant behavior where each (token,slot)->expert route can - have a distinct input scaling before quantization). - - "int4": W4A8 path: X is int8, W is packed int4 (2 values per byte) unpacked to int8 in-kernel - """ - - gpu_arch = get_rocm_arch() - _is_gfx950 = str(gpu_arch).startswith("gfx95") - allocator = SmemAllocator(None, arch=gpu_arch) - _state = {} - - in_dtype = "fp8" # blockscale is FP8-only - is_f16 = in_dtype == "fp16" - elem_bytes = 2 if is_f16 else 1 - if out_dtype not in ("f16", "bf16"): - raise ValueError(f"out_dtype must be 'f16' or 'bf16', got {out_dtype!r}") - # NOTE: don't materialize MLIR types outside an active MLIR Context. - out_mlir = lambda: (lambda ty: ty() if callable(ty) else ty)(T.f16 if out_dtype == "f16" else T.bf16) - tile_k_bytes = int(tile_k) * int(elem_bytes) - # K64-byte micro-step: always 64 bytes per `ku`. For fp16 this is 32 elements. - if (tile_k_bytes % 64) != 0: - raise ValueError( - f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " - f"(tile_k={tile_k}, elem_bytes={elem_bytes})" - ) - is_int4 = in_dtype == "int4" - # INT4 here means W4A8: X is int8, W is packed int4 and unpacked to int8 in-kernel. - is_int8 = (in_dtype == "int8") or is_int4 - x_is_token_slot = in_dtype == "int8smooth" - # "int8smooth" still uses int8 MFMA, but X/scale_x are provided per (token,slot). - is_int8 = is_int8 or x_is_token_slot - - # Blockscale compile-time constants (K=model_dim for stage1) - if model_dim % scale_block_k != 0: - raise ValueError(f"model_dim ({model_dim}) must be divisible by scale_block_k ({scale_block_k})") - if (2 * inter_dim) % 128 != 0: - raise ValueError(f"2*inter_dim ({2 * inter_dim}) must be divisible by 128 (ScaleBlockN)") - sb_per_tile_s1 = tile_k // scale_block_k # scale blocks per tile (in K dim) - ku_per_sb_s1 = scale_block_k // 64 # K64-steps per scale block = 2 - nblk_k_w1 = model_dim // scale_block_k # K-blocks in W1 (=scale_k) - nblk_n_w1 = (2 * inter_dim) // 128 # N-blocks in W1 (ScaleBlockN=128) - # scale_w: [experts, nblk_n_w1, nblk_k_w1] f32 (per-block scale) - sw_nbytes = experts * nblk_n_w1 * nblk_k_w1 * 4 - - 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`)." - ) - - ir.ShapedType.get_dynamic_size() - # W is packed int4 for W4A8: 2 values per byte. - w_nbytes = ( - (experts * (2 * inter_dim) * model_dim) // 2 - if is_int4 - else (experts * (2 * inter_dim) * model_dim * elem_bytes) - ) - - total_threads = 256 - bytes_x_per_tile = int(tile_m) * int(tile_k) * int(elem_bytes) - if bytes_x_per_tile % total_threads != 0: - raise ValueError( - "tile_m*tile_k*elem_bytes must be divisible by " - f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={elem_bytes}" - ) - bytes_per_thread_x = bytes_x_per_tile // total_threads - # Keep MoE stage1 X gmem->LDS pipeline consistent with the optimized GEMM kernel: - # split into <=16B pieces and use `fly.copy(load-only)` for buffer_load_dwordx4. - # (Compute the split lens inside the kernel so the code matches GEMM structure.) - - # LDS128 mode (same idea as test_preshuffle_gemm.py): - # - LDS stride == tile_k (no extra padding) + XOR16 swizzle - # - Use ds_{read,write}_b128 (16B) and extract 8B halves for MFMA steps - _ck_lds128 = os.environ.get("FLYDSL_CK_LDS128", "1") in ("1", "true", "True", "YES", "yes") - pad_k = 0 if _ck_lds128 else 8 - lds_stride = tile_k + pad_k - if use_cshuffle_epilog is None: - use_cshuffle_epilog = os.environ.get("FLYDSL_MOE_STAGE1_CSHUFFLE", "1") in ("1", "true", "True", "YES", "yes") - use_cshuffle_epilog = bool(use_cshuffle_epilog) - if out_dtype != "f16" and use_cshuffle_epilog: - raise ValueError("stage1 cshuffle epilog currently supports only f16 output (out_dtype='f16')") - - epilog_tag = "cshuffle" if use_cshuffle_epilog else "direct" - # IMPORTANT: module name participates in FlyDSL's compile cache key. - # Keep an explicit ABI tag so signature changes can't accidentally reuse an old binary. - _wpe_tag = f"_wpe{waves_per_eu}" if waves_per_eu is not None else "" - module_name = ( - f"mfma_moe1_bs_{in_dtype}_{out_dtype}_{epilog_tag}" - f"_t{tile_m}x{tile_n}x{tile_k}{_wpe_tag}" - f"_abi8" # scf.for main loop (reduced ISA size) - ).replace("-", "_") - - # ── LDS sizing (pure Python; no MLIR Context needed) ───────────────────── - _use_cshuffle_epilog = bool(use_cshuffle_epilog) - lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(elem_bytes) - lds_out_bytes = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 - lds_total_bytes = max(lds_x_bytes, lds_out_bytes) - lds_total_elems = lds_total_bytes if elem_bytes == 1 else (lds_total_bytes // 2) - - lds_alloc_bytes = int(lds_total_elems) * int(elem_bytes) - lds_alloc_offset = allocator._align(allocator.ptr, 16) - allocator.ptr = lds_alloc_offset + lds_alloc_bytes - - @flyc.kernel(name=module_name) - def moe_blockscale_gemm1( - arg_out: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_max_token_ids: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_inter_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - ): - tokens_in = arith.index_cast(T.index, i32_tokens_in) - inter_in = arith.index_cast(T.index, i32_inter_in) - k_in = arith.index_cast(T.index, i32_k_in) - size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) - tokens_i32_v = i32_tokens_in - k_i32_v = i32_k_in - x_elem = T.f16 if is_f16 else (T.i8 if is_int8 else T.f8) - # For int4, weights are stored as packed bytes (i8) and unpacked to i8 packs. - w_elem = T.f16 if is_f16 else (T.i8 if is_int8 else T.f8) - vec16_elems = 16 if elem_bytes == 1 else 8 - vec8_elems = 8 if elem_bytes == 1 else 4 - vec8_x = T.vec(vec8_elems, x_elem) - vec16_x = T.vec(vec16_elems, x_elem) - - def silu(x): - # device fast path: - # emu = exp(-x) ~= exp2(log2e * (-x)) -> v_exp_f32 - # sig = rcp(1 + emu) -> v_rcp_f32 - # y = x * sig - # - # Using llvm.amdgcn intrinsics prevents lowering to the div_scale/div_fixup - # sequences that introduce extra compares/cndmasks. - t = x * (-1.4426950408889634) # -log2(e) - emu = rocdl.exp2(T.f32, t) - den = 1.0 + emu - sig = rocdl.rcp(T.f32, den) - return x * sig - - acc_init = arith.constant_vector(0, T.i32x4) if is_int8 else arith.constant_vector(0.0, T.f32x4) - - # Layouts - fx.make_layout((tokens_i32_v, k_i32_v), stride=(k_i32_v, 1)) - - # B preshuffle layout: match GEMM test helper exactly. - c_n_total = arith.index(experts * (2 * inter_dim)) - kpack_bytes = 8 if is_int4 else 16 - b_layout = make_preshuffle_b_layout( - arith, c_n=c_n_total, c_k=k_in, kpack_bytes=kpack_bytes, elem_bytes=elem_bytes - ) - layout_b = b_layout.layout_b - (k_in * arith.index(int(elem_bytes))) // fx.Index(64) - - shape_lds = fx.make_shape(tile_m, tile_k) - stride_lds = fx.make_stride(lds_stride, 1) - layout_lds = fx.make_layout(shape_lds, stride_lds) - - tx = gpu.thread_id("x") - # Align with Aiter launch mapping (NSwizzle==false): - # - blockIdx.x -> N dimension (tile along inter_dim) - # - blockIdx.y -> expert-block id / M dimension (tile along sorted M) - by = gpu.block_id("x") # tile along inter_dim - bx = gpu.block_id("y") # tile along sorted M - - # Block validity: compute as early as possible so invalid blocks skip all buffer-resource - # setup, LDS pointer math, and gmem prefetch work. - bx_m = bx * fx.Index(tile_m) - maxids_rsrc = buffer_ops.create_buffer_resource( - arg_max_token_ids, max_size=False, num_records_bytes=fx.Index(4) - ) - max_token_id_i32 = buffer_ops.buffer_load(maxids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32) - bx_m_i32 = arith.index_cast(T.i32, bx_m) - blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, max_token_id_i32) - # Common constants/atoms (hoisted): keep IR small like GEMM. - # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). - k_blocks16 = arith.index(tile_k_bytes // 16) - layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) - layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) - - # Everything below is gated by `blk_valid` to avoid doing buffer-resource setup and - # gmem work for padding blocks. - _if_blk = scf.IfOp(blk_valid) - with _if_then(_if_blk): - base_ptr = allocator.get_base() - lds_x_ptr = SmemPtr( - base_ptr, - lds_alloc_offset, - (T.f16 if is_f16 else (T.i8 if is_int8 else T.f8)), - shape=(lds_total_elems,), - ) - lds_x = lds_x_ptr.get() - # Alias LDS bytes as fp16 for optional CShuffle epilogue. - lds_out = ( - SmemPtr(base_ptr, lds_x_ptr.byte_offset, T.f16, shape=(tile_m * tile_n,)).get() - if _use_cshuffle_epilog - else None - ) - - # Buffer resources: for dynamic memrefs, provide `num_records_bytes` explicitly so - # hardware OOB behavior is stable (otherwise it falls back to a large max size). - c_topk = fx.Index(topk) - - # X: [tokens, k] bytes = tokens*k*elem_bytes - x_rows = tokens_in * (c_topk if x_is_token_slot else fx.Index(1)) - x_nbytes_idx = x_rows * k_in * arith.index(int(elem_bytes)) - x_rsrc = buffer_ops.create_buffer_resource( - arg_x, max_size=False, num_records_bytes=arith.index_cast(T.i64, x_nbytes_idx) - ) - - w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False, num_records_bytes=w_nbytes) - - # OUT: [tokens, topk, inter] f16/bf16 -> bytes = tokens*topk*inter*out_elem_bytes - out_elem_bytes = 2 # f16/bf16 - out_nbytes_idx = tokens_in * c_topk * inter_in * fx.Index(out_elem_bytes) - out_rsrc = buffer_ops.create_buffer_resource( - arg_out, max_size=False, num_records_bytes=arith.index_cast(T.i64, out_nbytes_idx) - ) - - # fp16 path ignores scales completely (implicit scale=1.0). - x_load_bytes = 16 - - sx_rsrc = -1 - sw_rsrc = -1 - if const_expr(not is_f16): - # scale_x: [nblk_k_w1, tokens] f32 transposed -> total = nblk_k_w1 * tokens - sx_nbytes_idx = arith.index(nblk_k_w1) * tokens_in * fx.Index(4) - sx_rsrc = buffer_ops.create_buffer_resource( - arg_scale_x, max_size=False, num_records_bytes=arith.index_cast(T.i64, sx_nbytes_idx) - ) - sw_rsrc = buffer_ops.create_buffer_resource(arg_scale_w, max_size=False, num_records_bytes=sw_nbytes) - - sorted_nbytes_idx = size_expert_ids_in * fx.Index(tile_m) * fx.Index(4) - sorted_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_token_ids, max_size=False, num_records_bytes=sorted_nbytes_idx - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_weights, max_size=False, num_records_bytes=sorted_nbytes_idx - ) - - # expert ids: [blocks] i32 -> bytes = size_expert_ids_in*4 - expert_rsrc = buffer_ops.create_buffer_resource( - arg_expert_ids, - max_size=False, - num_records_bytes=arith.index_cast(T.i64, size_expert_ids_in * fx.Index(4)), - ) - - # Expert id for this M tile (keep address math in `index`) - expert_i32 = buffer_ops.buffer_load(expert_rsrc, bx, vec_width=1, dtype=T.i32) - expert_idx = arith.index_cast(T.index, expert_i32) - inter2_idx = arith.index(2 * inter_dim) - expert_off_idx = expert_idx * inter2_idx # index - - # ---- X gmem->reg prefetch (match preshuffle GEMM mapping) ---- - # Prefer 16B buffer-load (dwordx4). If the per-thread byte count isn't divisible by - # 16, fall back to 8B (dwordx2) or 4B (dword) loads. For fp16 we require 16B. - x_load_bytes = 16 - if const_expr(is_f16): - if const_expr(bytes_per_thread_x % 16 != 0): - raise ValueError(f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16") - x_load_bytes = 16 - else: - if const_expr(bytes_per_thread_x % 16 == 0): - x_load_bytes = 16 - elif const_expr(bytes_per_thread_x % 8 == 0): - x_load_bytes = 8 - elif const_expr(bytes_per_thread_x % 4 == 0): - x_load_bytes = 4 - else: - raise ValueError( - f"bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 4 to use the dword-indexed load mapping." - ) - num_x_loads = bytes_per_thread_x // x_load_bytes - chunk_i32 = x_load_bytes // 4 # dwords per chunk (1/2/4) - - c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4) - c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4) - fx.make_layout((tokens_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1)) - tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4 - layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1)) - c_chunk_i32 = fx.Index(chunk_i32) - tx_i32_base = tx * c_chunk_i32 - mask24 = fx.Int32(0xFFFFFF) - # Keep i32 constants available for epilogue index math. - tokens_i32 = arith.index_cast(T.i32, tokens_in) - topk_i32 = fx.Int32(topk) - - def x_tile_chunk_coord_i32(i: int): - return tile_chunk_coord_i32( - arith, - tx_i32_base=tx_i32_base, - i=i, - total_threads=total_threads, - layout_tile_div4=layout_x_tile_div4, - chunk_i32=chunk_i32, - ) - - # decode token once (per thread's M-slice) and build a base row offset. - x_row_base_div4 = [] - x_col_local_i32 = [] - x_row_local = [] - for i in range_constexpr(num_x_loads): - row_local, col_local_i32 = x_tile_chunk_coord_i32(i) - x_row_local.append(row_local) - x_col_local_i32.append(col_local_i32) - - sorted_row_i = bx_m + row_local - # NOTE: rows beyond `num_valid_ids` can contain garbage (within the allocated - # buffer). That's OK as long as we never use an out-of-range token id to index X. - fused_i = buffer_ops.buffer_load(sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32) - t_raw = fused_i & mask24 - # NOTE: aiter moe_sorting uses sentinel token_id == tokens for padding. - # Do NOT rely on buffer OOB semantics for X loads; explicitly mask to a safe row. - t_valid_i32 = arith.cmpi(arith.CmpIPredicate.ult, t_raw, tokens_i32) - if const_expr(x_is_token_slot): - s_raw = fused_i >> 24 - # X is indexed by token-slot in **slot-major** order: - # row_ts = slot * tokens + token - # This matches CK's moe_smoothquant output layout. - row_ts_i32 = s_raw * tokens_i32 + t_raw - row_ts_idx = arith.index_cast(T.index, row_ts_i32) - # Apply bounds check to token-slot index - row_ts_safe = t_valid_i32.select(row_ts_idx, fx.Index(0)) - x_row_base_div4.append(row_ts_safe * c_k_div4) - else: - t_idx = arith.index_cast(T.index, t_raw) - t_safe = t_valid_i32.select(t_idx, fx.Index(0)) - x_row_base_div4.append(t_safe * c_k_div4) - - T.vec(1, T.i32) - T.vec(2, T.i32) - vec4_x = T.vec(4, x_elem) - - def load_x(idx_i32, x_load_bytes_v): - """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_v == 16): - idx_elem = idx_i32 if elem_bytes == 1 else (idx_i32 * fx.Index(2)) - return buffer_copy_gmem16_dwordx4( - buffer_ops, - vector, - elem_type=x_elem, - idx_i32=idx_elem, - rsrc=x_rsrc, - vec_elems=vec16_elems, - elem_bytes=elem_bytes, - ) - if const_expr(x_load_bytes_v == 8): - return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=2, dtype=T.i32) - return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=1, dtype=T.i32) - - def load_x_tile(base_k, x_load_bytes_v): - """Prefetch the per-thread X tile portion (gmem -> regs) for a given K base (in elements).""" - base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4) - parts = [] - for i in range_constexpr(num_x_loads): - idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i] - x_vec = load_x(idx_i32, x_load_bytes_v) - if const_expr(x_load_bytes_v == 16): - parts.append(vector.bitcast(T.i32x4, as_ir_value(x_vec))) - elif const_expr(x_load_bytes_v == 8): - parts.append(x_vec) - else: - parts.append(x_vec) - return parts - - # tx -> wave/lane (GEMM-style decomposition). - coord_wl = fx.idx2crd(fx.Int32(tx), layout_tx_wave_lane) - wave_id = fx.get(coord_wl, 0) - lane_id = fx.get(coord_wl, 1) - coord_l16 = fx.idx2crd(fx.Int32(lane_id), layout_lane16) - lane_div_16 = fx.get(coord_l16, 0) - lane_mod_16 = fx.get(coord_l16, 1) - - # Match GEMM naming/pattern: row in LDS is lane_mod_16, and col base is lane_div_16*16. - row_a_lds = lane_mod_16 - a_kpack_elems = 16 // elem_bytes - col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems)) - col_offset_base_bytes = ( - col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes))) - ) - - # Dynamic N tiling within block (same as existing kernels) - by_n = by * fx.Index(tile_n) - num_waves = 4 - n_per_wave = tile_n // num_waves - num_acc_n = n_per_wave // 16 - c_n_per_wave = fx.Index(n_per_wave) - wave_mod_4 = wave_id % fx.Index(4) - n_tile_base = wave_mod_4 * c_n_per_wave - - # Precompute n_blk/n_intra for gate and up rows (GEMM-style: idx2crd/get) - n_intra_gate = [] - n_blk_gate = [] - n_intra_up = [] - n_blk_up = [] - col_g_list = [] - inter_idx = fx.Index(inter_dim) - # layout for (row -> (blk,intra)) where intra is 0..15 - c_n0 = c_n_total // fx.Index(16) - c_n0_i32 = arith.index_cast(T.i32, c_n0) - layout_n_blk_intra = fx.make_layout((c_n0_i32, 16), stride=(16, 1)) - for ni in range_constexpr(num_acc_n): - offset = arith.index(ni * 16) - col_g = by_n + n_tile_base - col_g = col_g + offset - col_g = col_g + lane_mod_16 - col_g_list.append(col_g) - - row_gate = expert_off_idx + col_g - row_up = row_gate + inter_idx - - coord_gate = fx.idx2crd(fx.Int32(row_gate), layout_n_blk_intra) - n_blk_gate.append(fx.get(coord_gate, 0)) - n_intra_gate.append(fx.get(coord_gate, 1)) - - coord_up = fx.idx2crd(fx.Int32(row_up), layout_n_blk_intra) - n_blk_up.append(fx.get(coord_up, 0)) - n_intra_up.append(fx.get(coord_up, 1)) - - m_repeat = tile_m // 16 - k_unroll = tile_k_bytes // 64 # K64-byte micro-step (2x MFMA) - - # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- - def load_b_pack(base_k, ki_step, ni, blk_list, intra_list): - return load_b_pack_k32( - buffer_ops, - arith, - vector, - arg_b=arg_w, - b_rsrc=w_rsrc, - layout_b=layout_b, - base_k=base_k, - ki_step=ki_step, - n_blk=blk_list[ni], - n_intra=intra_list[ni], - lane_div_16=lane_div_16, # 0..3 - elem_type=w_elem, - kpack_bytes=kpack_bytes, - elem_bytes=elem_bytes, - unpack_int4=is_int4, - ) - - def load_b_tile(base_k, blk_list, intra_list): - """Prefetch the entire per-thread B tile (gmem -> regs) for a given K base. - - Returns a list of length `k_unroll`, where each entry is a tuple: - (packs_half0[ni], packs_half1[ni]) for the K64 micro-step. - """ - b_tile = [] - for ku in range_constexpr(k_unroll): - packs0 = [] - packs1 = [] - for ni in range_constexpr(num_acc_n): - ki0 = (ku * 2) + 0 - ki1 = (ku * 2) + 1 - b0 = load_b_pack(base_k, ki0, ni, blk_list, intra_list) - b1 = load_b_pack(base_k, ki1, ni, blk_list, intra_list) - packs0.append(b0) - packs1.append(b1) - b_tile.append((packs0, packs1)) - return b_tile - - acc_gate = [arith.constant_vector(0.0, T.f32x4)] * (num_acc_n * m_repeat) - acc_up = [arith.constant_vector(0.0, T.f32x4)] * (num_acc_n * m_repeat) - - # ---- Pipeline helpers: store X tile to LDS with ping-pong base ---- - def store_x_tile_to_lds(vec_x_in_parts, lds_base, x_load_bytes_v): - 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_v == 16): - lds_store_16b_xor16( - arith, - vector, - lds_memref=lds_x, - vec16_ty=vec16_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x4=vec_x_in_parts[i], - elem_bytes=elem_bytes, - ) - elif const_expr(x_load_bytes_v == 8): - lds_store_8b_xor16( - arith, - vector, - lds_memref=lds_x, - vec8_ty=vec8_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x2=vec_x_in_parts[i], - ) - else: - lds_store_4b_xor16( - arith, - vector, - lds_memref=lds_x, - vec4_ty=vec4_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x1=vec_x_in_parts[i], - ) - - # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- - def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base): - col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16) - col_base_swz = ( - col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes))) - ) - idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds) - idx_a16 = idx_a16 + lds_base - loaded_a16 = vector.load(vec16_x, as_ir_value(lds_x), [as_ir_value(idx_a16)]) - a_i64x2 = vector.bitcast(T.i64x2, as_ir_value(loaded_a16)) - a0 = vector.extract(as_ir_value(a_i64x2), static_position=[0], dynamic_position=[]) - a1 = vector.extract(as_ir_value(a_i64x2), static_position=[1], dynamic_position=[]) - return a0, a1 - - # --- Blockscale pre-decode and helpers --- - c_scale_block_k = fx.Index(scale_block_k) - c_128 = fx.Index(128) - c_nblk_k_w1 = fx.Index(nblk_k_w1) - row_off_base = lane_div_16 * fx.Index(4) - - # Pre-decode sorted token IDs as i32 (constant across all K-tiles). - # OOB buffer loads return 0, so no validity masking needed for scale values. - _pre_t_safe_i32 = [] - for _mi in range_constexpr(m_repeat): - _mi_safe = [] - for _ii in range_constexpr(4): - _row_in_tile = arith.index(_mi * 16) + row_off_base + fx.Index(_ii) - _sorted_row = bx_m + _row_in_tile - _fused_pre = buffer_ops.buffer_load(sorted_rsrc, _sorted_row, vec_width=1, dtype=T.i32) - _t_id_pre = _fused_pre & mask24 - _t_valid_pre = arith.cmpi(arith.CmpIPredicate.ult, _t_id_pre, tokens_i32) - _t_safe_pre = _t_valid_pre.select(_t_id_pre, fx.Int32(0)) - _mi_safe.append(_t_safe_pre) - _pre_t_safe_i32.append(_mi_safe) - - # Pre-compute N-block indices for scale_w (constant per CTA) - _pre_n_block_gate = [] - _pre_n_block_up = [] - for _ni in range_constexpr(num_acc_n): - _col_base_ni_pre = by_n + n_tile_base + arith.index(_ni * 16) - _pre_n_block_gate.append((expert_off_idx + _col_base_ni_pre) // c_128) - _pre_n_block_up.append((expert_off_idx + inter_idx + _col_base_ni_pre) // c_128) - - def load_scales_s1(k_base): - all_combined = [] - for sb in range_constexpr(sb_per_tile_s1): - kb = k_base // c_scale_block_k + fx.Index(sb) - sa_base_offset = kb * tokens_in - - s_a_vecs = [] - sa_base_i32 = arith.index_cast(T.i32, sa_base_offset) - for mi in range_constexpr(m_repeat): - s_a_row = [] - for ii in range_constexpr(4): - t_safe_i32 = _pre_t_safe_i32[mi][ii] - sa_idx_i32 = sa_base_i32 + t_safe_i32 - sa_idx = arith.index_cast(T.index, sa_idx_i32) - s_a_val = buffer_ops.buffer_load(sx_rsrc, sa_idx, vec_width=1, dtype=T.f32) - s_a_row.append(s_a_val) - s_a_vecs.append(s_a_row) - - _sw_shared_n = n_per_wave <= 128 - s_w_gate_vals = [] - s_w_up_vals = [] - s_w_gate = fx.Float32(1.0) - s_w_up = fx.Float32(1.0) - for ni in range_constexpr(num_acc_n): - if const_expr(ni == 0 or not _sw_shared_n): - sw_gate_idx = _pre_n_block_gate[ni] * c_nblk_k_w1 + kb - s_w_gate = buffer_ops.buffer_load(sw_rsrc, sw_gate_idx, vec_width=1, dtype=T.f32) - sw_up_idx = _pre_n_block_up[ni] * c_nblk_k_w1 + kb - s_w_up = buffer_ops.buffer_load(sw_rsrc, sw_up_idx, vec_width=1, dtype=T.f32) - s_w_gate_vals.append(s_w_gate) - s_w_up_vals.append(s_w_up) - - s_a_vec4_list = [] - for mi in range_constexpr(m_repeat): - s_a_vec4_list.append(vector.from_elements(T.f32x4, [as_ir_value(e) for e in s_a_vecs[mi]])) - all_combined.append((s_a_vec4_list, s_w_gate_vals, s_w_up_vals)) - return all_combined - - def compute_tile_bs_s1( - acc_gate_in, acc_up_in, b_gate_tile_in, b_up_tile_in, lds_base, pre_scales, *, a0_prefetch=None - ): - current_gate = list(acc_gate_in) - current_up = list(acc_up_in) - mfma_res_ty = T.f32x4 - - if const_expr(_is_gfx950): - - for sb in range_constexpr(sb_per_tile_s1): - s_a_vec4_list, s_w_gate_vals, s_w_up_vals = pre_scales[sb] - ku0 = sb * ku_per_sb_s1 - ku1 = ku0 + 1 - bg0_p0, bg0_p1 = b_gate_tile_in[ku0] - bg1_p0, bg1_p1 = b_gate_tile_in[ku1] - bu0_p0, bu0_p1 = b_up_tile_in[ku0] - bu1_p0, bu1_p1 = b_up_tile_in[ku1] - col0 = col_offset_base_bytes + arith.index(ku0 * 64) - col1 = col_offset_base_bytes + arith.index(ku1 * 64) - for mi in range_constexpr(m_repeat): - curr_row = row_a_lds + arith.index(mi * 16) - a0 = arith.constant(0, type=T.i64) - a1 = arith.constant(0, type=T.i64) - if const_expr(a0_prefetch is not None and sb == 0 and mi == 0): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64(curr_row, col0, lds_base) - a2, a3 = lds_load_packs_k64(curr_row, col1, lds_base) - a128 = _pack128(a0, a1, a2, a3) - s_a_v4 = s_a_vec4_list[mi] - pending_gate_up = None - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - bg128 = _pack128(bg0_p0[ni], bg0_p1[ni], bg1_p0[ni], bg1_p1[ni]) - bu128 = _pack128(bu0_p0[ni], bu0_p1[ni], bu1_p0[ni], bu1_p1[ni]) - blk_g = rocdl.mfma_scale_f32_16x16x128_f8f6f4( - mfma_res_ty, [a128, bg128, acc_init, 0, 0, 0, 0x7F7F7F7F, 0, 0x7F7F7F7F] - ) - blk_u = rocdl.mfma_scale_f32_16x16x128_f8f6f4( - mfma_res_ty, [a128, bu128, acc_init, 0, 0, 0, 0x7F7F7F7F, 0, 0x7F7F7F7F] - ) - rocdl.sched_barrier(0) - if const_expr(pending_gate_up is not None): - prev_acc_idx, prev_blk_g, prev_blk_u, prev_ni = pending_gate_up - s_wg_bc = vector.broadcast(T.f32x4, s_w_gate_vals[prev_ni]) - s_wu_bc = vector.broadcast(T.f32x4, s_w_up_vals[prev_ni]) - scale_g = ArithValue(s_a_v4) * ArithValue(s_wg_bc) - scale_u = ArithValue(s_a_v4) * ArithValue(s_wu_bc) - current_gate[prev_acc_idx] = math_dialect.fma( - prev_blk_g, scale_g, current_gate[prev_acc_idx] - ) - current_up[prev_acc_idx] = math_dialect.fma( - prev_blk_u, scale_u, current_up[prev_acc_idx] - ) - pending_gate_up = (acc_idx, blk_g, blk_u, ni) - if const_expr(pending_gate_up is not None): - prev_acc_idx, prev_blk_g, prev_blk_u, prev_ni = pending_gate_up - s_wg_bc = vector.broadcast(T.f32x4, s_w_gate_vals[prev_ni]) - s_wu_bc = vector.broadcast(T.f32x4, s_w_up_vals[prev_ni]) - scale_g = ArithValue(s_a_v4) * ArithValue(s_wg_bc) - scale_u = ArithValue(s_a_v4) * ArithValue(s_wu_bc) - current_gate[prev_acc_idx] = math_dialect.fma( - prev_blk_g, scale_g, current_gate[prev_acc_idx] - ) - current_up[prev_acc_idx] = math_dialect.fma( - prev_blk_u, scale_u, current_up[prev_acc_idx] - ) - else: - mfma_fn = ( - mfma_i32_k32 - if const_expr(is_int8) - else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) - ) - - def mfma_k64(acc_in, a0, a1, b0, b1): - if const_expr(is_f16): - a0v = _i64_to_v4f16(a0) - a1v = _i64_to_v4f16(a1) - b0v = _i64_to_v4f16(b0) - b1v = _i64_to_v4f16(b1) - acc_mid = mfma_fn(mfma_res_ty, [a0v, b0v, acc_in, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1v, b1v, acc_mid, 0, 0, 0]) - acc_mid = mfma_fn(mfma_res_ty, [a0, b0, acc_in, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1, b1, acc_mid, 0, 0, 0]) - - for sb in range_constexpr(sb_per_tile_s1): - s_a_vec4_list, s_w_gate_vals, s_w_up_vals = pre_scales[sb] - for mi in range_constexpr(m_repeat): - s_a_v4 = s_a_vec4_list[mi] - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - blk_g = acc_init - blk_u = acc_init - for ku_local in range_constexpr(ku_per_sb_s1): - ku = sb * ku_per_sb_s1 + ku_local - b_gate_packs0, b_gate_packs1 = b_gate_tile_in[ku] - b_up_packs0, b_up_packs1 = b_up_tile_in[ku] - ki64 = arith.index(ku * 64) - col_base = col_offset_base_bytes + ki64 - a0 = arith.constant(-1, type=T.i64) - a1 = arith.constant(-1, type=T.i64) - if const_expr( - (a0_prefetch is not None) and (sb == 0) and (ku_local == 0) and (mi == 0) - ): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64( - row_a_lds + arith.index(mi * 16), col_base, lds_base - ) - blk_g = mfma_k64(blk_g, a0, a1, b_gate_packs0[ni], b_gate_packs1[ni]) - blk_u = mfma_k64(blk_u, a0, a1, b_up_packs0[ni], b_up_packs1[ni]) - s_wg_bc = vector.broadcast(T.f32x4, s_w_gate_vals[ni]) - s_wu_bc = vector.broadcast(T.f32x4, s_w_up_vals[ni]) - scale_g = ArithValue(s_a_v4) * ArithValue(s_wg_bc) - scale_u = ArithValue(s_a_v4) * ArithValue(s_wu_bc) - current_gate[acc_idx] = math_dialect.fma(blk_g, scale_g, current_gate[acc_idx]) - current_up[acc_idx] = math_dialect.fma(blk_u, scale_u, current_up[acc_idx]) - return current_gate, current_up - - def compute_tile( - acc_gate_in, - acc_up_in, - b_gate_tile_in, - b_up_tile_in, - lds_base, - *, - prefetch_epilogue: bool = False, - a0_prefetch=None, - ): - gate_list = list(acc_gate_in) - up_list = list(acc_up_in) - mfma_res_ty = T.i32x4 if is_int8 else T.f32x4 - mfma_fn = ( - mfma_i32_k32 - if const_expr(is_int8) - else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) - ) - - # Optional: prefetch epilogue scales while we are about to run the last MFMA tile, - # matching the preshuffle GEMM pattern of overlapping scale loads with MFMA. - epilogue_pf = None - if const_expr(prefetch_epilogue): - expert_off_pf = expert_off_idx - sw_gate_pf = [] - sw_up_pf = [] - for ni in range_constexpr(num_acc_n): - col_g = col_g_list[ni] - row_gate_idx = expert_off_pf + col_g - row_up_idx = row_gate_idx + inter_idx - sw_gate_pf.append( - fx.Float32(1.0) - if const_expr(is_f16) - else buffer_ops.buffer_load(sw_rsrc, row_gate_idx, vec_width=1, dtype=T.f32) - ) - sw_up_pf.append( - fx.Float32(1.0) - if const_expr(is_f16) - else buffer_ops.buffer_load(sw_rsrc, row_up_idx, vec_width=1, dtype=T.f32) - ) - epilogue_pf = (sw_gate_pf, sw_up_pf) - - def mfma_k64(acc_in, a0, a1, b0, b1): - if const_expr(is_f16): - a0v = _i64_to_v4f16(a0) - a1v = _i64_to_v4f16(a1) - b0v = _i64_to_v4f16(b0) - b1v = _i64_to_v4f16(b1) - acc_mid = mfma_fn(mfma_res_ty, [a0v, b0v, acc_in, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1v, b1v, acc_mid, 0, 0, 0]) - acc_mid = mfma_fn(mfma_res_ty, [a0, b0, acc_in, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1, b1, acc_mid, 0, 0, 0]) - - for ku in range_constexpr(k_unroll): - b_gate_packs0, b_gate_packs1 = b_gate_tile_in[ku] - b_up_packs0, b_up_packs1 = b_up_tile_in[ku] - ki64 = arith.index(ku * 64) - col_base = col_offset_base_bytes + ki64 - - for mi in range_constexpr(m_repeat): - mi_val = arith.index(mi * 16) - curr_row_a_lds = row_a_lds + mi_val - - a0 = arith.constant(-1, type=T.i64) - a1 = arith.constant(-1, type=T.i64) - if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) - - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - gate_list[acc_idx] = mfma_k64( - gate_list[acc_idx], - a0, - a1, - b_gate_packs0[ni], - b_gate_packs1[ni], - ) - up_list[acc_idx] = mfma_k64( - up_list[acc_idx], - a0, - a1, - b_up_packs0[ni], - b_up_packs1[ni], - ) - return gate_list, up_list, epilogue_pf - - # ── scf.for loop helpers (acc-only loop state, CK-style) ────── - n_accs_half = m_repeat * num_acc_n - - # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- - lds_tile_elems = arith.index(tile_m * lds_stride) - lds_base_cur = fx.Index(0) - lds_base_nxt = lds_tile_elems - - rocdl.sched_barrier(0) - - def hot_loop_scheduler(): - mfma_per_ku = m_repeat * num_acc_n * 2 * 2 - total_mfma = k_unroll * mfma_per_ku - rocdl.sched_group_barrier(rocdl.mask_dsrd, ku_per_sb_s1 * m_repeat, 0) - rocdl.sched_group_barrier(rocdl.mask_mfma, total_mfma, 1) - rocdl.sched_group_barrier(rocdl.mask_vmem_rd, num_x_loads, 2) - rocdl.sched_group_barrier(rocdl.mask_dswr, num_x_loads, 3) - rocdl.sched_barrier(0) - - def do_one_stage(acc_gate_in, acc_up_in, k_compute, k_next, lds_compute, lds_store): - """One pipeline stage: load next tile data, compute current tile, store X to LDS.""" - scale_fn = load_scales_s1 - pre_scales = scale_fn(k_compute) - x_regs_next = load_x_tile(k_next, x_load_bytes) - b_gate_cur = load_b_tile(k_compute, n_blk_gate, n_intra_gate) - b_up_cur = load_b_tile(k_compute, n_blk_up, n_intra_up) - - ag, au = compute_tile_bs_s1(acc_gate_in, acc_up_in, b_gate_cur, b_up_cur, lds_compute, pre_scales) - store_x_tile_to_lds(x_regs_next, lds_store, x_load_bytes) - hot_loop_scheduler() - gpu.barrier() - return ag, au - - # Prologue: prefetch tile0 X into LDS, sync. - k0 = fx.Index(0) - x_regs0 = load_x_tile(k0, x_load_bytes) - store_x_tile_to_lds(x_regs0, lds_base_cur, x_load_bytes) - gpu.barrier() - - lds_base_pong = lds_base_cur - lds_base_ping = lds_base_nxt - - c2_tile_k = arith.index(tile_k * 2) - c_tile_k = arith.index(tile_k) - total_tiles = int(model_dim) // int(tile_k) - pair_iters = max((total_tiles - 2) // 2, 0) - c_k_main = pair_iters * tile_k * 2 - - init_state = list(acc_gate) + list(acc_up) - - for k_iv, inner in range(0, c_k_main, tile_k * 2, init=init_state): - n = n_accs_half - acc_gate_in = list(inner[:n]) - acc_up_in = list(inner[n : 2 * n]) - - next_k1 = k_iv + c_tile_k - - acc_gate_s0, acc_up_s0 = do_one_stage( - acc_gate_in, acc_up_in, k_iv, next_k1, lds_base_pong, lds_base_ping - ) - - next_k2 = k_iv + c2_tile_k - - acc_gate_s1, acc_up_s1 = do_one_stage( - acc_gate_s0, acc_up_s0, next_k1, next_k2, lds_base_ping, lds_base_pong - ) - - results = yield list(acc_gate_s1) + list(acc_up_s1) - - n = n_accs_half - acc_gate = list(results[:n]) - acc_up = list(results[n : 2 * n]) - - # Tail: use fresh scale decode (no dependency on prologue _pre_t_safe_idx) - k_tail0 = k_in - c2_tile_k - k_tail1 = k_in - c_tile_k - - acc_gate, acc_up = do_one_stage(acc_gate, acc_up, k_tail0, k_tail1, lds_base_pong, lds_base_ping) - - pre_scales_tail1 = load_scales_s1(k_tail1) - b_gate_last = load_b_tile(k_tail1, n_blk_gate, n_intra_gate) - b_up_last = load_b_tile(k_tail1, n_blk_up, n_intra_up) - acc_gate, acc_up = compute_tile_bs_s1( - acc_gate, acc_up, b_gate_last, b_up_last, lds_base_ping, pre_scales_tail1 - ) - - # Store epilogue to out[t, slot, inter] - tokens_i32_v = tokens_i32 - topk_i32_v = topk_i32 - inter_i32_v = fx.Int32(inter_dim) - mask24_i32 = fx.Int32(0xFFFFFF) - - # Blockscale: dequant already done in compute_tile_bs_s1, no sw/sx needed here. - - # Epilogue hoists to keep IR + Python build time small: - col_i32_list = [] - for ni in range_constexpr(num_acc_n): - col_i32_list.append(arith.index_cast(T.i32, col_g_list[ni])) - - lane_div_16 * fx.Index(4) - inter_i32_local = inter_i32_v - - if const_expr(use_cshuffle_epilog): - if const_expr(lds_out is None): - raise RuntimeError("CShuffle epilogue enabled but lds_out is not allocated/aliased.") - - def write_row_to_lds( - *, - mi: int, - ii: int, - row_in_tile, - row, - row_base_lds, - col_base_local, - num_acc_n: int, - lds_out, - ): - # Blockscale: dequant already done in compute_tile_bs_s1. - # Just apply silu + optional sorted weight. - if const_expr(doweight_stage1): - tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) - - for ni in range_constexpr(num_acc_n): - col_local = col_base_local + (ni * 16) - - acc_idx = mi * num_acc_n + ni - vg = vector.extract(as_ir_value(acc_gate[acc_idx]), static_position=[ii], dynamic_position=[]) - vu = vector.extract(as_ir_value(acc_up[acc_idx]), static_position=[ii], dynamic_position=[]) - - y = silu(vg) * vu - if const_expr(doweight_stage1): - y = y * tw - y16 = arith.trunc_f(T.f16, y) - - lds_idx = row_base_lds + col_local - v1 = vector.from_elements(T.vec(1, T.f16), [as_ir_value(y16)]) - vector.store(as_ir_value(v1), as_ir_value(lds_out), [as_ir_value(lds_idx)], alignment=2) - - def precompute_row(*, row_local, row): - fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) - t2 = fused2 & mask24_i32 - s2 = fused2 >> 24 - return (t2 * topk_i32_v + s2) * inter_i32_local - - def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): - # Guard against sentinel token ids (t == tokens) produced by aiter moe_sorting padding. - # OOB buffer stores are not guaranteed to be safe on all paths, so predicate explicitly. - fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) - t2 = fused2 & mask24_i32 - t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) - _if_valid = scf.IfOp(t_valid) - with _if_then(_if_valid): - idx0 = row_ctx - col_i32 = arith.index_cast(T.i32, col_g0) - idx_out = idx0 + col_i32 - # Vectorized fp16 store (EVec=4). - buffer_ops.buffer_store(frag, out_rsrc, idx_out) - - mfma_epilog( - use_cshuffle=True, - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, - tile_m=tile_m, - tile_n=tile_n, - e_vec=4, - m_repeat=m_repeat, - num_acc_n=num_acc_n, - tx=tx, - lane_div_16=lane_div_16, - lane_mod_16=lane_mod_16, - bx_m=bx_m, - by_n=by_n, - n_tile_base=n_tile_base, - lds_out=lds_out, - write_row_to_lds=write_row_to_lds, - precompute_row=precompute_row, - store_pair=store_pair, - ) - return - - def _stage1_store_row(*, mi: int, ii: int, row_in_tile, row): - # Blockscale: dequant already done in compute_tile_bs_s1. - fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) - t2 = fused2 & mask24_i32 - s2 = fused2 >> 24 - t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v) - - # out linear index base = ((t*topk + s)*inter_dim) (invariant across ni) - idx0 = (t2 * topk_i32_v + s2) * inter_i32_local - - # Sorted weight aligned with `row` (matches aiter moe_sorting output). - if const_expr(doweight_stage1): - tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) - - _if_valid = scf.IfOp(t_valid) - with _if_then(_if_valid): - for ni in range_constexpr(num_acc_n): - col_i32 = col_i32_list[ni] - - acc_idx = mi * num_acc_n + ni - vg = vector.extract(as_ir_value(acc_gate[acc_idx]), static_position=[ii], dynamic_position=[]) - vu = vector.extract(as_ir_value(acc_up[acc_idx]), static_position=[ii], dynamic_position=[]) - - y = silu(vg) * vu - if const_expr(doweight_stage1): - y = y * tw - y = arith.trunc_f(out_mlir(), y) - idx_out0 = idx0 + col_i32 - buffer_ops.buffer_store(y, out_rsrc, idx_out0) - - mfma_epilog( - use_cshuffle=False, - arith=arith, - range_constexpr=range_constexpr, - m_repeat=m_repeat, - lane_div_16=lane_div_16, - bx_m=bx_m, - body_row=_stage1_store_row, - ) - - # ── Host launcher (flyc.jit + .launch) ──────────────────────────────── - @flyc.jit - def launch_moe_blockscale_gemm1( - arg_out: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_max_token_ids: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_inter_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - stream: fx.Stream, - ): - allocator.finalized = False - ctx = CompilationContext.get_current() - with ir.InsertionPoint(ctx.gpu_module_body): - allocator.finalize() - - inter_in = arith.index_cast(T.index, i32_inter_in) - size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) - gx = inter_in // fx.Index(tile_n) - gy = size_expert_ids_in - - moe_blockscale_gemm1( - arg_out, - arg_x, - arg_w, - arg_scale_x, - arg_scale_w, - arg_sorted_token_ids, - arg_expert_ids, - arg_sorted_weights, - arg_max_token_ids, - i32_tokens_in, - i32_inter_in, - i32_k_in, - i32_size_expert_ids_in, - value_attrs={"rocdl.waves_per_eu": waves_per_eu}, - ).launch(grid=(gx, gy, 1), block=(256, 1, 1), stream=stream) - - return launch_moe_blockscale_gemm1 diff --git a/kernels/moe/moe_blockscale_2stage/gemm2.py b/kernels/moe/moe_blockscale_2stage/gemm2.py deleted file mode 100644 index a315d81a7..000000000 --- a/kernels/moe/moe_blockscale_2stage/gemm2.py +++ /dev/null @@ -1,1412 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""MoE block-scale 2-stage MFMA kernels (stage1 / stage2 / reduction). :: _gemm2""" - -import functools -import logging -import os - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm, scf, vector -from flydsl._mlir.dialects import math as math_dialect -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr.arith import ArithValue -from flydsl.expr.typing import T -from flydsl.runtime.device import get_rocm_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr -from kernels.common import buffer_ops -from kernels.common.kernels_common import _if_then -from kernels.common.mem_ops import buffer_atomic_add -from kernels.common.mma.mfma_epilogues import c_shuffle_epilog, default_epilog -from kernels.common.mma.mfma_preshuffle_pipeline import ( - buffer_copy_gmem16_dwordx4, - lds_store_4b_xor16, - lds_store_8b_xor16, - lds_store_16b_xor16, - load_b_pack_k32, - make_preshuffle_b_layout, - preshuffle_crd2idx, - swizzle_xor16, - tile_chunk_coord_i32, -) -from kernels.moe.moe_blockscale_2stage.reduction import compile_moe_reduction -from kernels.moe.moe_common import ( - i64_to_v4f16 as _i64_to_v4f16, -) -from kernels.moe.moe_common import ( - i64x4_to_i32x8 as _pack128, -) - - -@functools.lru_cache(maxsize=1024) -def compile_moe_blockscale_gemm2( - *, - model_dim: int, - inter_dim: int, - experts: int, - topk: int, - tile_m: int, - tile_n: int, - tile_k: int, - doweight_stage2: bool, - scale_block_k: int = 128, - 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, - waves_per_eu: int | None = None, -): - """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. - - in_dtype: - - "fp8": A2/W are fp8 - - "fp16": A2/W are fp16 - - "int8": A2/W are int8 - - "int4": W4A8 path: A2 is int8, W is packed int4 unpacked to int8 in-kernel - - 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). - """ - gpu_arch = get_rocm_arch() - _is_gfx950 = str(gpu_arch).startswith("gfx95") - allocator = SmemAllocator(None, arch=gpu_arch) - _state = {} - - in_dtype = "fp8" # blockscale is FP8-only - is_f16 = in_dtype == "fp16" - elem_bytes = 2 if is_f16 else 1 - out_s = str(out_dtype).strip().lower() - if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): - raise ValueError(f"out_dtype must be 'f16', 'bf16', or 'f32', got {out_dtype!r}") - out_is_f32 = out_s in ("f32", "fp32", "float") - out_is_bf16 = out_s in ("bf16", "bfloat16") - if (not bool(accumulate)) and out_is_f32: - raise ValueError("compile_moe_blockscale_gemm2(accumulate=False) only supports out_dtype in {'f16','bf16'}") - is_int4 = in_dtype == "int4" - # INT4 here means W4A8: A2 is int8, W is packed int4 and unpacked to int8 in-kernel. - is_int8 = (in_dtype in ("int8", "int8smooth")) or is_int4 - - # Blockscale compile-time constants (K=inter_dim for stage2) - if inter_dim % scale_block_k != 0: - raise ValueError(f"inter_dim ({inter_dim}) must be divisible by scale_block_k ({scale_block_k})") - if model_dim % 128 != 0: - raise ValueError(f"model_dim ({model_dim}) must be divisible by 128 (ScaleBlockN)") - sb_per_tile_s2 = tile_k // scale_block_k # scale blocks per tile (in K dim) - ku_per_sb_s2 = scale_block_k // 64 # K64-steps per scale block = 2 - nblk_k_w2 = inter_dim // scale_block_k # K-blocks in W2 (=scale_k) - nblk_n_w2 = model_dim // 128 # N-blocks in W2 (ScaleBlockN=128) - # scale_w: [experts, nblk_n_w2, nblk_k_w2] f32 (per-block scale) - sw_nbytes = experts * nblk_n_w2 * nblk_k_w2 * 4 - - 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`)." - ) - - ir.ShapedType.get_dynamic_size() - # W is packed int4 for W4A8: 2 values per byte. - w_nbytes = (experts * model_dim * inter_dim) // 2 if is_int4 else (experts * model_dim * inter_dim * elem_bytes) - - total_threads = 256 - tile_k_bytes = int(tile_k) * int(elem_bytes) - if (tile_k_bytes % 64) != 0: - raise ValueError( - f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " - f"(tile_k={tile_k}, elem_bytes={elem_bytes})" - ) - bytes_x_per_tile = int(tile_m) * int(tile_k) * int(elem_bytes) - if bytes_x_per_tile % total_threads != 0: - raise ValueError( - "tile_m*tile_k*elem_bytes must be divisible by " - f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={elem_bytes}" - ) - bytes_per_thread_x = bytes_x_per_tile // total_threads - - _ck_lds128 = os.environ.get("FLYDSL_CK_LDS128", "1") in ("1", "true", "True", "YES", "yes") - pad_k = 0 if _ck_lds128 else 8 - lds_stride = tile_k + pad_k - # gfx950+ has buffer_atomic_pk_add_bf16 → bf16 can use buffer atomics (same as f16). - # gfx942 only has global_atomic_pk_add_bf16 → must use global atomics with raw pointer. - _has_buffer_atomic_bf16 = str(gpu_arch).startswith(("gfx95", "gfx12")) - _needs_global_atomic_bf16 = out_is_bf16 and not _has_buffer_atomic_bf16 - if out_is_bf16: - if not (gpu_arch.startswith("gfx942") or gpu_arch.startswith("gfx950") or gpu_arch.startswith("gfx12")): - raise ValueError( - f"out_dtype='bf16' requires bf16 global atomics (gfx942/gfx950/gfx12), got arch={gpu_arch!r}" - ) - - if out_is_f32: - # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. - _use_cshuffle_epilog = False if use_cshuffle_epilog is None else bool(use_cshuffle_epilog) - if _use_cshuffle_epilog: - raise ValueError("out_dtype='f32' does not support CShuffle epilogue (set use_cshuffle_epilog=False).") - else: - if use_cshuffle_epilog is None: - _use_cshuffle_epilog = os.environ.get("FLYDSL_MOE_STAGE2_CSHUFFLE", "1") in ( - "1", - "true", - "True", - "YES", - "yes", - ) - else: - _use_cshuffle_epilog = bool(use_cshuffle_epilog) - if not _use_cshuffle_epilog: - raise ValueError("stage2 f16 output currently requires CShuffle epilogue (FLYDSL_MOE_STAGE2_CSHUFFLE=1).") - - # NOTE: Keep this as a callable so we don't require an MLIR Context at Python-time. - def out_elem(): - ty = T.f32 if out_is_f32 else (T.bf16 if out_is_bf16 else T.f16) - return ty() if callable(ty) else ty - - epilog_tag = "cshuffle" - # IMPORTANT: include tiling in the module name to avoid accidentally reusing a compiled - # binary for a different (tile_m, tile_n, tile_k) configuration. - # See stage1 note: include ABI tag to prevent binary reuse across signature changes. - # IMPORTANT: module name participates in FlyDSL's compile cache key. - # Dynamic-shape variant: safe to reuse across (tokens/sorted_size/size_expert_ids) at runtime. - # Keep a distinct ABI tag so the compile cache never mixes with historical signatures. - _wpe_tag2 = f"_wpe{waves_per_eu}" if waves_per_eu is not None else "" - module_name = ( - f"mfma_moe2_{in_dtype}_{out_s}_{epilog_tag}" - f"_t{tile_m}x{tile_n}x{tile_k}{_wpe_tag2}" - f"_abi6" # scale prefetch before VMEM tile loads - ).replace("-", "_") - - # ── LDS sizing (pure Python; no MLIR Context needed) ───────────────────── - lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(elem_bytes) - lds_out_bytes = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 - lds_total_bytes = max(lds_x_bytes, lds_out_bytes) - lds_total_elems = lds_total_bytes if elem_bytes == 1 else (lds_total_bytes // 2) - - lds_alloc_bytes = int(lds_total_elems) * int(elem_bytes) - lds_alloc_offset = allocator._align(allocator.ptr, 16) - allocator.ptr = lds_alloc_offset + lds_alloc_bytes - - _cshuffle_nlane = 32 - if bool(accumulate): - _e_vec = 2 - else: - _e_vec = 8 if int(tile_n) % (_cshuffle_nlane * 8) == 0 else 2 - _cshuffle_stride = _cshuffle_nlane * _e_vec - if int(tile_n) % _cshuffle_stride != 0: - raise ValueError(f"tile_n={tile_n} must be divisible by {_cshuffle_stride} when accumulate=False") - - if True: - - @flyc.kernel(name=module_name) - def moe_blockscale_gemm2( - arg_out: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_num_valid_ids: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_n_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - ): - tokens_in = arith.index_cast(T.index, i32_tokens_in) - n_in = arith.index_cast(T.index, i32_n_in) - k_in = arith.index_cast(T.index, i32_k_in) - size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) - k_i32_v = i32_k_in - x_elem = T.f16 if is_f16 else (T.i8 if is_int8 else T.f8) - # For int4, weights are stored as packed bytes (i8) and unpacked to i8 packs. - w_elem = T.f16 if is_f16 else (T.i8 if is_int8 else T.f8) - vec16_elems = 16 if elem_bytes == 1 else 8 - vec8_elems = 8 if elem_bytes == 1 else 4 - vec8_x = T.vec(vec8_elems, x_elem) - vec16_x = T.vec(vec16_elems, x_elem) - - acc_init = arith.constant_vector(0, T.i32x4) if is_int8 else arith.constant_vector(0.0, T.f32x4) - - # A2 layout (flatten token-slot -> M). - topk_idx = fx.Index(topk) - m_in = tokens_in * topk_idx - m_i32_v = arith.index_cast(T.i32, m_in) - fx.make_layout((m_i32_v, k_i32_v), stride=(k_i32_v, 1)) - - # B preshuffle layout: [experts*model_dim, inter_dim] - c_n_total = arith.index(experts * model_dim) - kpack_bytes = 8 if is_int4 else 16 - b_layout = make_preshuffle_b_layout( - arith, c_n=c_n_total, c_k=k_in, kpack_bytes=kpack_bytes, elem_bytes=elem_bytes - ) - layout_b = b_layout.layout_b - (k_in * arith.index(int(elem_bytes))) // fx.Index(64) - - shape_lds = fx.make_shape(tile_m, tile_k) - stride_lds = fx.make_stride(lds_stride, 1) - layout_lds = fx.make_layout(shape_lds, stride_lds) - - tx = gpu.thread_id("x") - # Align with Aiter launch mapping: - # - blockIdx.x -> N dimension (tile along model_dim) - # - blockIdx.y -> expert-block id / M dimension (tile along sorted M) - by = gpu.block_id("x") # tile along model_dim - bx = gpu.block_id("y") # tile along sorted M - - # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). - k_blocks16 = arith.index(tile_k_bytes // 16) - layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) - layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) - fx.make_layout((tile_m, tile_k), stride=(tile_k, 1)) - - base_ptr = allocator.get_base() - lds_x_ptr = SmemPtr( - base_ptr, - lds_alloc_offset, - (T.f16 if is_f16 else (T.i8 if is_int8 else T.f8)), - shape=(lds_total_elems,), - ) - lds_x = lds_x_ptr.get() - # Alias the same underlying LDS bytes as f16/bf16 for epilogue shuffle. - lds_out = ( - SmemPtr( - base_ptr, - lds_x_ptr.byte_offset, - (T.bf16 if out_is_bf16 else T.f16), - shape=(tile_m * tile_n,), - ).get() - if _use_cshuffle_epilog - else None - ) - - # Buffer resources. - # For dynamic memrefs, `max_size=False` cannot infer the logical size from the memref *type*, - # so we should pass `num_records_bytes` explicitly for stable hardware OOB behavior. - c_topk = fx.Index(topk) - - # X(A2): [tokens*topk, inter_dim] bytes = tokens*topk*k*elem_bytes - x_nbytes_idx = (tokens_in * c_topk) * k_in * arith.index(int(elem_bytes)) - x_rsrc = buffer_ops.create_buffer_resource( - arg_x, max_size=False, num_records_bytes=arith.index_cast(T.i64, x_nbytes_idx) - ) - - w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False, num_records_bytes=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 * fx.Index(out_elem_bytes) - if const_expr(not bool(accumulate)): - out_nbytes_idx = tokens_in * fx.Index(topk) * n_in * fx.Index(out_elem_bytes) - out_rsrc = buffer_ops.create_buffer_resource( - arg_out, max_size=False, num_records_bytes=arith.index_cast(T.i64, out_nbytes_idx) - ) - # fp16 path ignores scales completely (implicit scale=1.0). - sx_rsrc = -1 - sw_rsrc = -1 - if const_expr(not is_f16): - # scale_x (A2 scale): [nblk_k_w2, tokens*topk] f32 transposed -> total = nblk_k_w2 * tokens * topk - sx_nbytes_idx = arith.index(nblk_k_w2) * (tokens_in * c_topk) * fx.Index(4) - sx_rsrc = buffer_ops.create_buffer_resource( - arg_scale_x, max_size=False, num_records_bytes=arith.index_cast(T.i64, sx_nbytes_idx) - ) - # scale_w: [experts, nblk_n_w2, nblk_k_w2] f32 (per-block scale) - sw_rsrc = buffer_ops.create_buffer_resource(arg_scale_w, max_size=False, num_records_bytes=sw_nbytes) - - # sorted_token_ids / sorted_weights: [blocks*tile_m] (CK-style padded length) - sorted_nbytes_idx = size_expert_ids_in * fx.Index(tile_m) * fx.Index(4) - sorted_nbytes_i64 = arith.index_cast(T.i64, sorted_nbytes_idx) - sorted_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_token_ids, max_size=False, num_records_bytes=sorted_nbytes_i64 - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - arg_sorted_weights, max_size=False, num_records_bytes=sorted_nbytes_i64 - ) - - # expert ids: [blocks] i32 -> bytes = size_expert_ids_in*4 - eid_nbytes_idx = size_expert_ids_in * fx.Index(4) - expert_rsrc = buffer_ops.create_buffer_resource( - arg_expert_ids, max_size=False, num_records_bytes=arith.index_cast(T.i64, eid_nbytes_idx) - ) - bx_m = bx * fx.Index(tile_m) - - # Early-exit guard (as in 2ce65fb): some routing paths can produce extra/garbage - # expert blocks beyond `num_valid_ids`. Skip those blocks entirely to avoid OOB. - numids_rsrc = buffer_ops.create_buffer_resource( - arg_num_valid_ids, max_size=False, num_records_bytes=fx.Index(4) - ) - num_valid_i32 = buffer_ops.buffer_load(numids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32) - bx_m_i32 = arith.index_cast(T.i32, bx_m) - blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, num_valid_i32) - - def _moe_gemm2_then_body(): - # Expert id for this M tile. - expert_i32 = buffer_ops.buffer_load(expert_rsrc, bx, vec_width=1, dtype=T.i32) - expert_idx = arith.index_cast(T.index, expert_i32) - n_idx = fx.Index(model_dim) - expert_off_idx = expert_idx * n_idx # index - - # ---- X gmem->reg prefetch (match preshuffle GEMM mapping) ---- - # Prefer 16B buffer-load (dwordx4). If the per-thread byte count isn't divisible by - # 16, fall back to 8B (dwordx2) or 4B (dword) loads. For fp16 we require 16B. - x_load_bytes = 0 - if const_expr(is_f16): - if const_expr(bytes_per_thread_x % 16 != 0): - raise ValueError(f"[fp16] bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 16") - x_load_bytes = 16 - else: - if const_expr(bytes_per_thread_x % 16 == 0): - x_load_bytes = 16 - elif const_expr(bytes_per_thread_x % 8 == 0): - x_load_bytes = 8 - elif const_expr(bytes_per_thread_x % 4 == 0): - x_load_bytes = 4 - else: - raise ValueError( - f"bytes_per_thread_x ({bytes_per_thread_x}) must be divisible by 4 to use the dword-indexed load mapping." - ) - num_x_loads = bytes_per_thread_x // x_load_bytes - chunk_i32 = x_load_bytes // 4 # dwords per chunk (1/2/4) - - c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4) - c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4) - fx.make_layout((m_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1)) - tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4 - layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1)) - c_chunk_i32 = fx.Index(chunk_i32) - tx_i32_base = tx * c_chunk_i32 - - topk_i32 = fx.Int32(topk) - mask24 = fx.Int32(0xFFFFFF) - # Sentinel clamp uses `tokens` as the upper bound: t_valid = (t < tokens). - tokens_i32 = arith.index_cast(T.i32, tokens_in) - - def x_tile_chunk_coord_i32(i: int): - return tile_chunk_coord_i32( - arith, - tx_i32_base=tx_i32_base, - i=i, - total_threads=total_threads, - layout_tile_div4=layout_x_tile_div4, - chunk_i32=chunk_i32, - ) - - T.vec(1, T.i32) - T.vec(2, T.i32) - vec4_x = T.vec(4, x_elem) - - def load_x(idx_i32, x_load_bytes_v): - if const_expr(x_load_bytes_v == 16): - idx_elem = idx_i32 if elem_bytes == 1 else (idx_i32 * fx.Index(2)) - return buffer_copy_gmem16_dwordx4( - buffer_ops, - vector, - elem_type=x_elem, - idx_i32=idx_elem, - rsrc=x_rsrc, - vec_elems=vec16_elems, - elem_bytes=elem_bytes, - ) - if const_expr(x_load_bytes_v == 8): - return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=2, dtype=T.i32) - return buffer_ops.buffer_load(x_rsrc, idx_i32, vec_width=1, dtype=T.i32) - - # decode routed token once (per thread's M-slice) and build a base offset. - x_row_base_div4 = [] - x_col_local_i32 = [] - x_row_local = [] - for i in range_constexpr(num_x_loads): - row_local, col_local_i32 = x_tile_chunk_coord_i32(i) - x_row_local.append(row_local) - x_col_local_i32.append(col_local_i32) - - sorted_row_i = bx_m + row_local - fused_i = buffer_ops.buffer_load(sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32) - t_i32 = fused_i & mask24 - s_i32 = fused_i >> 24 - # aiter moe_sorting uses sentinel token_id == tokens for padding. - # Do NOT rely on buffer OOB semantics for A2/scale loads; explicitly mask. - t_valid = arith.cmpi(arith.CmpIPredicate.ult, t_i32, tokens_i32) - s_valid = arith.cmpi(arith.CmpIPredicate.ult, s_i32, topk_i32) - ts_valid = t_valid & s_valid - t_safe = ts_valid.select(t_i32, fx.Int32(0)) - s_safe = ts_valid.select(s_i32, fx.Int32(0)) - row_ts_i32 = t_safe * topk_i32 + s_safe - row_ts_idx = arith.index_cast(T.index, row_ts_i32) - # Base row offset in dword units: row_ts_idx * (k_in/4) - x_row_base_div4.append(row_ts_idx * c_k_div4) - - def load_x_tile(base_k, x_load_bytes_v): - base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4) - parts = [] - for i in range_constexpr(num_x_loads): - idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i] - x_vec = load_x(idx_i32, x_load_bytes_v) - if const_expr(x_load_bytes_v == 16): - parts.append(vector.bitcast(T.i32x4, as_ir_value(x_vec))) - elif const_expr(x_load_bytes_v == 8): - parts.append(x_vec) - else: - parts.append(x_vec) - return parts - - # tx -> wave/lane (GEMM-style decomposition). - coord_wl = fx.idx2crd(fx.Int32(tx), layout_tx_wave_lane) - wave_id = fx.get(coord_wl, 0) - lane_id = fx.get(coord_wl, 1) - coord_l16 = fx.idx2crd(fx.Int32(lane_id), layout_lane16) - lane_div_16 = fx.get(coord_l16, 0) - lane_mod_16 = fx.get(coord_l16, 1) - - row_a_lds = lane_mod_16 - a_kpack_elems = 16 // elem_bytes - col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems)) - col_offset_base_bytes = ( - col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes))) - ) - - # Dynamic N tiling within block. - by_n = by * fx.Index(tile_n) - num_waves = 4 - n_per_wave = tile_n // num_waves - num_acc_n = n_per_wave // 16 - c_n_per_wave = fx.Index(n_per_wave) - wave_mod_4 = wave_id % fx.Index(4) - n_tile_base = wave_mod_4 * c_n_per_wave - - # Precompute (n_blk, n_intra) for B, and col indices for output. - n_intra_list = [] - n_blk_list = [] - col_g_list = [] - c_n0 = c_n_total // fx.Index(16) - c_n0_i32 = arith.index_cast(T.i32, c_n0) - layout_n_blk_intra = fx.make_layout((c_n0_i32, 16), stride=(16, 1)) - for ni in range_constexpr(num_acc_n): - offset = arith.index(ni * 16) - col_g = by_n + n_tile_base + offset + lane_mod_16 - col_g_list.append(col_g) - - row_w = expert_off_idx + col_g - coord_w = fx.idx2crd(fx.Int32(row_w), layout_n_blk_intra) - n_blk_list.append(fx.get(coord_w, 0)) - n_intra_list.append(fx.get(coord_w, 1)) - - m_repeat = tile_m // 16 - k_unroll = tile_k_bytes // 64 # K64-byte micro-step (2x MFMA) - - # --- B Load Logic (K64) --- - def load_b_pack(base_k, ki_step, ni): - return load_b_pack_k32( - buffer_ops, - arith, - vector, - arg_b=arg_w, - b_rsrc=w_rsrc, - layout_b=layout_b, - base_k=base_k, - ki_step=ki_step, - n_blk=n_blk_list[ni], - n_intra=n_intra_list[ni], - lane_div_16=lane_div_16, # 0..3 - elem_type=w_elem, - kpack_bytes=kpack_bytes, - elem_bytes=elem_bytes, - unpack_int4=is_int4, - ) - - def load_b_tile(base_k): - """Prefetch the entire per-thread B tile (gmem -> regs) for a given K base. - - Returns a list of length `k_unroll`, where each entry is a tuple: - (packs_half0[ni], packs_half1[ni]) for the K64 micro-step. - """ - b_tile = [] - for ku in range_constexpr(k_unroll): - packs0 = [] - packs1 = [] - for ni in range_constexpr(num_acc_n): - ki0 = (ku * 2) + 0 - ki1 = (ku * 2) + 1 - b0 = load_b_pack(base_k, ki0, ni) - b1 = load_b_pack(base_k, ki1, ni) - packs0.append(b0) - packs1.append(b1) - b_tile.append((packs0, packs1)) - return b_tile - - # ---- Pipeline helpers: store X tile to LDS with ping-pong base ---- - def store_x_tile_to_lds(vec_x_in_parts, lds_base, x_load_bytes_v): - 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_v == 16): - lds_store_16b_xor16( - arith, - vector, - lds_memref=lds_x, - vec16_ty=vec16_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x4=vec_x_in_parts[i], - elem_bytes=elem_bytes, - ) - elif const_expr(x_load_bytes_v == 8): - lds_store_8b_xor16( - arith, - vector, - lds_memref=lds_x, - vec8_ty=vec8_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x2=vec_x_in_parts[i], - ) - else: - lds_store_4b_xor16( - arith, - vector, - lds_memref=lds_x, - vec4_ty=vec4_x, - layout_lds=layout_lds, - row_local=row_local, - col_local_i32=col_local_i32, - tx_c4=fx.Index(4), - k_blocks16=k_blocks16, - lds_base=lds_base, - vec_part_i32x1=vec_x_in_parts[i], - ) - - # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- - def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base): - col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16) - col_base_swz = ( - col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes))) - ) - idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds) - idx_a16 = idx_a16 + lds_base - loaded_a16 = vector.load(vec16_x, as_ir_value(lds_x), [as_ir_value(idx_a16)]) - a_i64x2 = vector.bitcast(T.i64x2, as_ir_value(loaded_a16)) - a0 = vector.extract(as_ir_value(a_i64x2), static_position=[0], dynamic_position=[]) - a1 = vector.extract(as_ir_value(a_i64x2), static_position=[1], dynamic_position=[]) - return a0, a1 - - # --- Blockscale pre-decode and helpers (stage2) --- - c_scale_block_k_s2 = fx.Index(scale_block_k) - c_128_s2 = fx.Index(128) - c_nblk_k_w2 = fx.Index(nblk_k_w2) - row_off_base_s2 = lane_div_16 * fx.Index(4) - fx.Index(model_dim) - - # Pre-decode sorted token IDs for stage2 (constant across all K-tiles). - # OOB buffer loads return 0, so no validity masking needed for scale values. - _pre_ts_safe_i32_s2 = [] - for _mi in range_constexpr(m_repeat): - _mi_safe = [] - for _ii in range_constexpr(4): - _row_in_tile = arith.index(_mi * 16) + row_off_base_s2 + fx.Index(_ii) - _sorted_row = bx_m + _row_in_tile - _fused_pre = buffer_ops.buffer_load(sorted_rsrc, _sorted_row, vec_width=1, dtype=T.i32) - _t_id_pre = _fused_pre & mask24 - _s_id_pre = _fused_pre >> 24 - _t_valid_pre = arith.cmpi(arith.CmpIPredicate.ult, _t_id_pre, tokens_i32) - _s_valid_pre = arith.cmpi(arith.CmpIPredicate.ult, _s_id_pre, topk_i32) - _ts_valid_pre = _t_valid_pre & _s_valid_pre - _t_safe_pre = _ts_valid_pre.select(_t_id_pre, fx.Int32(0)) - _s_safe_pre = _ts_valid_pre.select(_s_id_pre, fx.Int32(0)) - _ts_i32_pre = _t_safe_pre * topk_i32 + _s_safe_pre - _mi_safe.append(_ts_i32_pre) - _pre_ts_safe_i32_s2.append(_mi_safe) - - # Pre-compute N-block indices for scale_w (constant per CTA) - _pre_n_block_s2 = [] - for _ni in range_constexpr(num_acc_n): - _col_base_ni_pre = by_n + n_tile_base + arith.index(_ni * 16) - _pre_n_block_s2.append((expert_off_idx + _col_base_ni_pre) // c_128_s2) - - m_in_s2 = tokens_in * fx.Index(topk) - - def load_scales_s2(k_base): - all_combined = [] - for sb in range_constexpr(sb_per_tile_s2): - kb = k_base // c_scale_block_k_s2 + fx.Index(sb) - sa_base_offset = kb * m_in_s2 - - s_a_vecs = [] - sa_base_i32 = arith.index_cast(T.i32, sa_base_offset) - for mi in range_constexpr(m_repeat): - s_a_row = [] - for ii in range_constexpr(4): - ts_safe_i32 = _pre_ts_safe_i32_s2[mi][ii] - sa_idx_i32 = sa_base_i32 + ts_safe_i32 - sa_idx = arith.index_cast(T.index, sa_idx_i32) - s_a_val = buffer_ops.buffer_load(sx_rsrc, sa_idx, vec_width=1, dtype=T.f32) - s_a_row.append(s_a_val) - s_a_vecs.append(s_a_row) - - _sw_shared_n_s2 = n_per_wave <= 128 - s_w_vals = [] - s_w = arith.constant(1.0, type=T.f32) - for ni in range_constexpr(num_acc_n): - if const_expr(ni == 0 or not _sw_shared_n_s2): - sw_idx = _pre_n_block_s2[ni] * c_nblk_k_w2 + kb - s_w = buffer_ops.buffer_load(sw_rsrc, sw_idx, vec_width=1, dtype=T.f32) - s_w_vals.append(s_w) - - s_a_vec4_list = [] - for mi in range_constexpr(m_repeat): - s_a_vec4_list.append(vector.from_elements(T.f32x4, [as_ir_value(e) for e in s_a_vecs[mi]])) - all_combined.append((s_a_vec4_list, s_w_vals)) - return all_combined - - def compute_tile_bs_s2(acc_in, b_tile_in, lds_base, pre_scales, *, a0_prefetch=None): - current_acc = list(acc_in) - mfma_res_ty = T.f32x4 - - if const_expr(_is_gfx950): - - for sb in range_constexpr(sb_per_tile_s2): - s_a_vec4_list, s_w_vals = pre_scales[sb] - ku0 = sb * ku_per_sb_s2 - ku1 = ku0 + 1 - b0_p0, b0_p1 = b_tile_in[ku0] - b1_p0, b1_p1 = b_tile_in[ku1] - col0 = col_offset_base_bytes + arith.index(ku0 * 64) - col1 = col_offset_base_bytes + arith.index(ku1 * 64) - for mi in range_constexpr(m_repeat): - curr_row = row_a_lds + arith.index(mi * 16) - a0 = arith.constant(0, type=T.i64) - a1 = arith.constant(0, type=T.i64) - if const_expr(a0_prefetch is not None and sb == 0 and mi == 0): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64(curr_row, col0, lds_base) - a2, a3 = lds_load_packs_k64(curr_row, col1, lds_base) - a128 = _pack128(a0, a1, a2, a3) - s_a_v4 = s_a_vec4_list[mi] - pending_acc = None - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - b128 = _pack128(b0_p0[ni], b0_p1[ni], b1_p0[ni], b1_p1[ni]) - blk = rocdl.mfma_scale_f32_16x16x128_f8f6f4( - mfma_res_ty, [a128, b128, acc_init, 0, 0, 0, 0x7F7F7F7F, 0, 0x7F7F7F7F] - ) - rocdl.sched_barrier(0) - if const_expr(pending_acc is not None): - prev_acc_idx, prev_blk, prev_ni = pending_acc - s_w_bc = vector.broadcast(T.f32x4, s_w_vals[prev_ni]) - scale = ArithValue(s_a_v4) * ArithValue(s_w_bc) - current_acc[prev_acc_idx] = math_dialect.fma( - prev_blk, scale, current_acc[prev_acc_idx] - ) - pending_acc = (acc_idx, blk, ni) - if const_expr(pending_acc is not None): - prev_acc_idx, prev_blk, prev_ni = pending_acc - s_w_bc = vector.broadcast(T.f32x4, s_w_vals[prev_ni]) - scale = ArithValue(s_a_v4) * ArithValue(s_w_bc) - current_acc[prev_acc_idx] = math_dialect.fma( - prev_blk, scale, current_acc[prev_acc_idx] - ) - else: - mfma_fn = ( - mfma_i32_k32 - if const_expr(is_int8) - else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) - ) - - def mfma_k64(acc0, a0, a1, b0, b1): - if const_expr(is_f16): - a0v = _i64_to_v4f16(a0) - a1v = _i64_to_v4f16(a1) - b0v = _i64_to_v4f16(b0) - b1v = _i64_to_v4f16(b1) - acc1 = mfma_fn(mfma_res_ty, [a0v, b0v, acc0, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1v, b1v, acc1, 0, 0, 0]) - acc1 = mfma_fn(mfma_res_ty, [a0, b0, acc0, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1, b1, acc1, 0, 0, 0]) - - for sb in range_constexpr(sb_per_tile_s2): - s_a_vec4_list, s_w_vals = pre_scales[sb] - for mi in range_constexpr(m_repeat): - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - blk = acc_init - for ku_local in range_constexpr(ku_per_sb_s2): - ku = sb * ku_per_sb_s2 + ku_local - b_packs0, b_packs1 = b_tile_in[ku] - ki64 = arith.index(ku * 64) - col_base = col_offset_base_bytes + ki64 - a0 = arith.constant(-1, type=T.i64) - a1 = arith.constant(-1, type=T.i64) - if const_expr( - (a0_prefetch is not None) and (sb == 0) and (ku_local == 0) and (mi == 0) - ): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64( - row_a_lds + arith.index(mi * 16), col_base, lds_base - ) - blk = mfma_k64(blk, a0, a1, b_packs0[ni], b_packs1[ni]) - s_w_bc = vector.broadcast(T.f32x4, s_w_vals[ni]) - scale = ArithValue(s_a_vec4_list[mi]) * ArithValue(s_w_bc) - current_acc[acc_idx] = math_dialect.fma(blk, scale, current_acc[acc_idx]) - return current_acc - - def compute_tile(acc_in, b_tile_in, lds_base, *, prefetch_epilogue: bool = False, a0_prefetch=None): - acc_list = list(acc_in) - mfma_res_ty = T.i32x4 if is_int8 else T.f32x4 - mfma_fn = ( - mfma_i32_k32 - if is_int8 - else (rocdl.mfma_f32_16x16x16f16 if is_f16 else rocdl.mfma_f32_16x16x32_fp8_fp8) - ) - - epilogue_pf = None - if const_expr(prefetch_epilogue): - expert_off_pf = expert_off_idx - sw_pf = [] - for ni in range_constexpr(num_acc_n): - col_g = col_g_list[ni] - row_w_idx = expert_off_pf + col_g - sw_pf.append( - fx.Float32(1.0) - if is_f16 - else buffer_ops.buffer_load(sw_rsrc, row_w_idx, vec_width=1, dtype=T.f32) - ) - # Also prefetch per-row routed/topk weights (sorted_weights) when enabled. - tw_pf = None - if const_expr(doweight_stage2): - tw_pf = [] - lane_div_16_mul4_pf = lane_div_16 * fx.Index(4) - ii_idx_list_pf = [fx.Index(ii) for ii in range(4)] - for mi in range_constexpr(m_repeat): - mi_base_pf = arith.index(mi * 16) - for ii in range_constexpr(4): - row_off_pf = lane_div_16_mul4_pf + ii_idx_list_pf[ii] - row_in_tile_pf = mi_base_pf + row_off_pf - sorted_row_pf = bx_m + row_in_tile_pf - tw_pf.append( - buffer_ops.buffer_load(sorted_w_rsrc, sorted_row_pf, vec_width=1, dtype=T.f32) - ) - epilogue_pf = (sw_pf, tw_pf) - - def mfma_k64(acc0, a0, a1, b0, b1): - if const_expr(is_f16): - a0v = _i64_to_v4f16(a0) - a1v = _i64_to_v4f16(a1) - b0v = _i64_to_v4f16(b0) - b1v = _i64_to_v4f16(b1) - acc1 = mfma_fn(mfma_res_ty, [a0v, b0v, acc0, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1v, b1v, acc1, 0, 0, 0]) - acc1 = mfma_fn(mfma_res_ty, [a0, b0, acc0, 0, 0, 0]) - return mfma_fn(mfma_res_ty, [a1, b1, acc1, 0, 0, 0]) - - for ku in range_constexpr(k_unroll): - b_packs0, b_packs1 = b_tile_in[ku] - ki64 = arith.index(ku * 64) - col_base = col_offset_base_bytes + ki64 - - for mi in range_constexpr(m_repeat): - mi_val = arith.index(mi * 16) - curr_row_a_lds = row_a_lds + mi_val - - a0 = arith.constant(-1, type=T.i64) - a1 = arith.constant(-1, type=T.i64) - if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)): - a0, a1 = a0_prefetch - else: - a0, a1 = lds_load_packs_k64(curr_row_a_lds, col_base, lds_base) - - for ni in range_constexpr(num_acc_n): - acc_idx = mi * num_acc_n + ni - acc_list[acc_idx] = mfma_k64( - acc_list[acc_idx], - a0, - a1, - b_packs0[ni], - b_packs1[ni], - ) - return acc_list, epilogue_pf - - # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- - lds_tile_elems = arith.index(tile_m * lds_stride) - lds_base_cur = fx.Index(0) - lds_base_nxt = lds_tile_elems - - rocdl.sched_barrier(0) - - # def hot_loop_scheduler(): - # mfma_group = num_acc_n - # # K64 micro-step: 2x K32 MFMA per accumulator update. - # mfma_total = (k_unroll * 2) * m_repeat * mfma_group - # mfma_per_iter = 2 * mfma_group - # sche_iters = 0 if mfma_per_iter == 0 else (mfma_total // mfma_per_iter) - # rocdl.sched_dsrd(2) - # rocdl.sched_mfma(1) - # rocdl.sched_mfma(1) - # if num_acc_n < 4: - # rocdl.sched_dsrd(1) - # rocdl.sched_mfma(1) - # rocdl.sched_dsrd(1) - # rocdl.sched_mfma(1) - # rocdl.sched_vmem(1) - # rocdl.sched_mfma(1) - # rocdl.sched_vmem(1) - # rocdl.sched_mfma(2) - # rocdl.sched_dsrd(1) - # rocdl.sched_mfma(2) - # rocdl.sched_vmem(1) - - # dswr_tail = num_x_loads - # if dswr_tail > sche_iters: - # dswr_tail = sche_iters - # dswr_start = sche_iters - dswr_tail - # for sche_i in range_constexpr(sche_iters): - # rocdl.sched_mfma(mfma_group // 2) - # rocdl.sched_dsrd(1) - # rocdl.sched_mfma(mfma_group // 2) - # rocdl.sched_vmem(1) - # rocdl.sched_mfma(mfma_group) - # if sche_i >= dswr_start - 1: - # rocdl.sched_dswr(1) - # rocdl.sched_barrier(0) - - def hot_loop_scheduler(): - mfma_per_ku = m_repeat * num_acc_n * 2 # m * n_acc * 2(k32) - total_mfma = k_unroll * mfma_per_ku - rocdl.sched_group_barrier(rocdl.mask_dsrd, ku_per_sb_s2 * m_repeat, 0) - rocdl.sched_group_barrier(rocdl.mask_mfma, total_mfma, 1) - rocdl.sched_group_barrier(rocdl.mask_vmem_rd, num_x_loads, 2) - rocdl.sched_group_barrier(rocdl.mask_dswr, num_x_loads, 3) - rocdl.sched_barrier(0) - - # Prologue. - k0 = fx.Index(0) - x_regs0 = load_x_tile(k0, x_load_bytes) - b_cur = load_b_tile(k0) - store_x_tile_to_lds(x_regs0, lds_base_cur, x_load_bytes) - gpu.barrier() - - acc = [arith.constant_vector(0.0, T.f32x4)] * (num_acc_n * m_repeat) - lds_base_pong = lds_base_cur - lds_base_ping = lds_base_nxt - - # Cross-tile A0 LDS prefetch (default-on): prefetch the first A-pack (K64) for the - # tile we are about to compute from LDS, to overlap with upcoming VMEM. - a0_prefetch_pong = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_pong) - - # Main loop: process K tiles in 2-tile ping-pong steps. - # - # IMPORTANT: for odd number of K tiles, leave **1** tail tile; for even, leave **2**. - # Otherwise the 2-tile tail below would double-count the last tile when num_tiles is odd - # (e.g. inter_dim=192, tile_k=64 -> 3 tiles). - num_k_tiles_py = int(inter_dim) // int(tile_k) - odd_k_tiles = (num_k_tiles_py % 2) == 1 - tail_tiles = 1 if odd_k_tiles else 2 - k_main2_py = (num_k_tiles_py - tail_tiles) * int(tile_k) - if const_expr(k_main2_py < 0): - k_main2_py = 0 - - c2_tile_k = arith.index(tile_k * 2) - pair_iters = k_main2_py // (int(tile_k) * 2) - for pair_i in range_constexpr(pair_iters): - k_iv = arith.index(pair_i * (tile_k * 2)) - # Issue scale loads FIRST so their latency hides behind heavy tile VMEM. - pre_scales_pong = load_scales_s2(k_iv) - next_k1 = k_iv + tile_k - x_regs_ping = load_x_tile(next_k1, x_load_bytes) - b_ping = load_b_tile(next_k1) - - acc = compute_tile_bs_s2(acc, b_cur, lds_base_pong, pre_scales_pong, a0_prefetch=a0_prefetch_pong) - a0_prefetch_pong = None - store_x_tile_to_lds(x_regs_ping, lds_base_ping, x_load_bytes) - hot_loop_scheduler() - gpu.barrier() - - # Cross-tile prefetch for the ping tile we are about to compute. - a0_prefetch_ping = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) - - # Issue scale loads FIRST so their latency hides behind heavy tile VMEM. - pre_scales_ping = load_scales_s2(next_k1) - next_k2 = k_iv + c2_tile_k - x_regs_pong = load_x_tile(next_k2, x_load_bytes) - b_next = load_b_tile(next_k2) - - acc = compute_tile_bs_s2(acc, b_ping, lds_base_ping, pre_scales_ping, a0_prefetch=a0_prefetch_ping) - a0_prefetch_ping = None - store_x_tile_to_lds(x_regs_pong, lds_base_pong, x_load_bytes) - 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_bytes, lds_base_pong) - - b_cur = b_next - - if const_expr(odd_k_tiles): - # Tail: single remaining tile (already in `b_cur` / `lds_base_pong`). - k_last = arith.index((num_k_tiles_py - 1) * int(tile_k)) - pre_scales_last = load_scales_s2(k_last) - acc = compute_tile_bs_s2( - acc, - b_cur, - lds_base_pong, - pre_scales_last, - a0_prefetch=a0_prefetch_pong, - ) - else: - # Tail: 2 remaining tiles. - k_tail0 = k_in - c2_tile_k - k_tail1 = k_in - tile_k - # Issue scale loads FIRST so their latency hides behind heavy tile VMEM. - pre_scales_tail0 = load_scales_s2(k_tail0) - x_regs_ping = load_x_tile(k_tail1, x_load_bytes) - b_ping = load_b_tile(k_tail1) - - acc = compute_tile_bs_s2(acc, b_cur, lds_base_pong, pre_scales_tail0, a0_prefetch=a0_prefetch_pong) - a0_prefetch_pong = None - store_x_tile_to_lds(x_regs_ping, lds_base_ping, x_load_bytes) - hot_loop_scheduler() - gpu.barrier() - - # Epilogue tile (blockscale already dequantized). - a0_prefetch_ping = lds_load_packs_k64(row_a_lds, col_offset_base_bytes, lds_base_ping) - pre_scales_tail1 = load_scales_s2(k_tail1) - acc = compute_tile_bs_s2(acc, b_ping, lds_base_ping, pre_scales_tail1, a0_prefetch=a0_prefetch_ping) - - # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- - # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. - mask24_i32 = fx.Int32(0xFFFFFF) - model_i32 = fx.Int32(model_dim) - topk_i32_v = topk_i32 - - zero_i32 = fx.Int32(0) - c2_i32 = fx.Int32(2) # 2B element size for f16/bf16 - mask_even_i32 = fx.Int32(0xFFFFFFFE) # align element index to even for half2 atomics - - e_vec = _e_vec - - def atomic_add_f16x2(val_f16x2, byte_off_i32): - buffer_atomic_add(val_f16x2, out_rsrc, byte_off_i32, zero_i32, zero_i32) - - # Blockscale: dequant already done in compute_tile_bs_s2, no sw/sx needed here. - - if const_expr(out_is_f32): - # origin/dev_a16w4: f32 output uses scalar f32 atomics and skips CShuffle/LDS. - c4_i32 = fx.Int32(4) - - def atomic_add_f32(val_f32, byte_off_i32): - buffer_atomic_add(val_f32, out_rsrc, byte_off_i32, zero_i32, zero_i32) - - def _stage2_row_atomic(*, mi: int, ii: int, row_in_tile, row): - # Blockscale: dequant already done in compute_tile_bs_s2. - fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) - t2 = fused2 & mask24_i32 - fused2 >> 24 - - if const_expr(doweight_stage2): - tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) - - idx0 = t2 * model_i32 # i32 element index base - - for ni in range_constexpr(num_acc_n): - col_g = col_g_list[ni] - acc_idx = mi * num_acc_n + ni - v = vector.extract(as_ir_value(acc[acc_idx]), static_position=[ii], dynamic_position=[]) - if const_expr(doweight_stage2): - v = v * tw - col_i32 = arith.index_cast(T.i32, col_g) - idx_elem = idx0 + col_i32 - byte_off = idx_elem * c4_i32 - atomic_add_f32(v, byte_off) - - default_epilog( - arith=arith, - range_constexpr=range_constexpr, - m_repeat=m_repeat, - lane_div_16=lane_div_16, - bx_m=bx_m, - body_row=_stage2_row_atomic, - ) - else: - if const_expr(lds_out is None): - raise RuntimeError("FLYDSL_MOE_STAGE2_CSHUFFLE=1 but lds_out is not allocated/aliased.") - - # For bf16 global atomics (gfx942 only), precompute the output base address. - # gfx950+ has buffer_atomic_pk_add_bf16, so bf16 uses buffer atomics there. - out_base_idx = None - if const_expr(_needs_global_atomic_bf16): - out_base_idx = buffer_ops.extract_base_index(arg_out) - - def write_row_to_lds( - *, - mi: int, - ii: int, - row_in_tile, - row, - row_base_lds, - col_base_local, - num_acc_n: int, - lds_out, - ): - # Blockscale: dequant already done in compute_tile_bs_s2. - tw = arith.constant(1.0, type=T.f32) - if const_expr(doweight_stage2): - tw = buffer_ops.buffer_load(sorted_w_rsrc, row, vec_width=1, dtype=T.f32) - - for ni in range_constexpr(num_acc_n): - col_local = col_base_local + (ni * 16) - acc_idx = mi * num_acc_n + ni - v = vector.extract(as_ir_value(acc[acc_idx]), static_position=[ii], dynamic_position=[]) - if const_expr(doweight_stage2): - v = v * tw - v_out = arith.trunc_f(out_elem(), v) - - lds_idx = row_base_lds + col_local - vec1_out = T.vec(1, out_elem()) - v1 = vector.from_elements(vec1_out, [as_ir_value(v_out)]) - vector.store( - as_ir_value(v1), - as_ir_value(lds_out), - [as_ir_value(lds_idx)], - alignment=2, - ) - - def precompute_row(*, row_local, row): - # Precompute row context for cshuffle stores. - # Return (fused_i32, row_valid_i1) so the epilogue can skip the entire row - # for invalid tail rows (CK-style), avoiding per-store branching. - fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32) - row_i32 = arith.index_cast(T.i32, row) - row_valid0 = arith.cmpi(arith.CmpIPredicate.ult, row_i32, num_valid_i32) - t = fused2 & mask24_i32 - s = fused2 >> 24 - t_ok = arith.cmpi(arith.CmpIPredicate.ult, t, tokens_i32) - s_ok = arith.cmpi(arith.CmpIPredicate.ult, s, topk_i32_v) - row_valid = row_valid0 & t_ok & s_ok - return (fused2, row_valid) - - def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): - fused = row_ctx - t = fused & mask24_i32 - s = fused >> 24 - idx0 = t * model_i32 - if const_expr(not bool(accumulate)): - ts = t * topk_i32_v + s - idx0 = ts * model_i32 - col_i32 = arith.index_cast(T.i32, col_g0) - idx_elem = idx0 + col_i32 - idx_elem_even = idx_elem & mask_even_i32 - if const_expr(_needs_global_atomic_bf16): - # gfx942: no buffer_atomic_pk_add_bf16, use global atomicrmw fadd - if const_expr(bool(accumulate)): - byte_off = idx_elem_even * c2_i32 - byte_off_idx = arith.index_cast(T.index, byte_off) - ptr_addr_idx = out_base_idx + byte_off_idx - out_ptr = buffer_ops.create_llvm_ptr(ptr_addr_idx, address_space=1) - out_ptr_v = out_ptr._value if hasattr(out_ptr, "_value") else out_ptr - frag_v = frag._value if hasattr(frag, "_value") else frag - llvm.AtomicRMWOp( - llvm.AtomicBinOp.fadd, - out_ptr_v, - frag_v, - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ) - else: - buffer_ops.buffer_store(frag, out_rsrc, idx_elem_even) - else: - # f16, or bf16 on gfx950+ (has buffer_atomic_pk_add_bf16) - byte_off = idx_elem_even * c2_i32 - if const_expr(bool(accumulate)): - atomic_add_f16x2(frag, byte_off) - else: - buffer_ops.buffer_store(frag, out_rsrc, idx_elem_even) - - c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, - tile_m=tile_m, - tile_n=tile_n, - e_vec=e_vec, - m_repeat=m_repeat, - num_acc_n=num_acc_n, - tx=tx, - lane_div_16=lane_div_16, - lane_mod_16=lane_mod_16, - bx_m=bx_m, - by_n=by_n, - n_tile_base=n_tile_base, - lds_out=lds_out, - frag_elem_type=(T.bf16 if out_is_bf16 else T.f16), - write_row_to_lds=write_row_to_lds, - precompute_row=precompute_row, - store_pair=store_pair, - ) - - _if_blk = scf.IfOp(blk_valid) - with _if_then(_if_blk): - _moe_gemm2_then_body() - - # ── Host launcher (flyc.jit + .launch) ──────────────────────────────── - @flyc.jit - def launch_moe_blockscale_gemm2( - arg_out: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_scale_x: fx.Tensor, - arg_scale_w: fx.Tensor, - arg_sorted_token_ids: fx.Tensor, - arg_expert_ids: fx.Tensor, - arg_sorted_weights: fx.Tensor, - arg_num_valid_ids: fx.Tensor, - i32_tokens_in: fx.Int32, - i32_n_in: fx.Int32, - i32_k_in: fx.Int32, - i32_size_expert_ids_in: fx.Int32, - stream: fx.Stream, - ): - allocator.finalized = False - ctx = CompilationContext.get_current() - with ir.InsertionPoint(ctx.gpu_module_body): - allocator.finalize() - - n_in = arith.index_cast(T.index, i32_n_in) - size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) - gx = n_in // fx.Index(tile_n) - gy = size_expert_ids_in - - moe_blockscale_gemm2( - arg_out, - arg_x, - arg_w, - arg_scale_x, - arg_scale_w, - arg_sorted_token_ids, - arg_expert_ids, - arg_sorted_weights, - arg_num_valid_ids, - i32_tokens_in, - i32_n_in, - i32_k_in, - i32_size_expert_ids_in, - value_attrs={"rocdl.waves_per_eu": waves_per_eu}, - ).launch(grid=(gx, gy, 1), block=(256, 1, 1), stream=stream) - - return launch_moe_blockscale_gemm2 - - -class MoeGemm2Mode: - """Execution mode for MoE GEMM2.""" - - ATOMIC = "atomic" # Use atomic accumulation (default) - REDUCE = "reduce" # Use non-atomic write + reduce kernel - - -class _MoeGemm2ReduceWrapper: - """Wrapper combining GEMM2 (no atomics) with reduction kernel. - - This wrapper handles the intermediate buffer allocation and orchestrates - the two-phase computation: - 1. GEMM2 outputs to [tokens*topk, model_dim] without atomics - 2. Reduce sums over topk to produce [tokens, model_dim] - """ - - def __init__( - self, - gemm2_exe, - reduce_exe, - topk: int, - model_dim: int, - out_dtype_str: str = "f16", - use_mask: bool = False, - ): - self._gemm2_exe = gemm2_exe - self._reduce_exe = reduce_exe - self._topk = topk - self._model_dim = model_dim - self._out_dtype_str = out_dtype_str - self._use_mask = use_mask - - def _get_torch_dtype(self): - """Convert dtype string to torch dtype.""" - import torch - - dtype_map = { - "f16": torch.float16, - "fp16": torch.float16, - "bf16": torch.bfloat16, - "f32": torch.float32, - } - return dtype_map.get(self._out_dtype_str, torch.float16) - - def __call__( - self, - arg_out, - arg_x, - arg_w, - arg_scale_x, - arg_scale_w, - arg_sorted_token_ids, - arg_expert_ids, - arg_sorted_weights, - arg_num_valid_ids, - tokens_in, - n_in, - k_in, - size_expert_ids_in, - valid_mask=None, - stream=None, - ): - """Execute GEMM2 + reduce. - - Args match moe_gemm2 kernel signature (see compile_moe_blockscale_gemm2). - """ - import torch - - if stream is None: - stream = torch.cuda.current_stream() - intermediate = torch.empty( - tokens_in * self._topk, self._model_dim, device=arg_out.device, dtype=self._get_torch_dtype() - ) - if not self._use_mask: - intermediate.zero_() - # Phase 1: GEMM2 (no atomics) -> [tokens*topk, model_dim] - self._gemm2_exe( - intermediate.view(-1), - arg_x, - arg_w, - arg_scale_x, - arg_scale_w, - arg_sorted_token_ids, - arg_expert_ids, - arg_sorted_weights, - arg_num_valid_ids, - tokens_in, - n_in, - k_in, - size_expert_ids_in, - stream, - ) - # Phase 2: Reduce over topk -> [tokens, model_dim] - X = intermediate.view(tokens_in, self._topk, self._model_dim) - Y = arg_out.view(tokens_in, self._model_dim) - if not self._use_mask: - if valid_mask is not None: - logging.warning("valid_mask provided but use_mask=False; ignoring valid_mask") - valid_mask = torch.empty((0, self._topk), device=arg_out.device, dtype=torch.uint8) - self._reduce_exe(X, Y, valid_mask, tokens_in, stream) - - @property - def mode(self) -> str: - """Return the execution mode.""" - return MoeGemm2Mode.REDUCE - - -def compile_moe_blockscale_gemm2_ex( - *, - model_dim: int, - inter_dim: int, - experts: int, - topk: int, - tile_m: int, - tile_n: int, - tile_k: int, - doweight_stage2: bool, - in_dtype: str = "fp8", - out_dtype: str = "f16", - use_cshuffle_epilog: bool | None = None, - # Extended parameters for mode control - mode: str = MoeGemm2Mode.ATOMIC, - valid_mask=None, -): - """Compile MoE GEMM2 kernel with optional reduction. - - This is the extended interface that supports explicit mode control. - - Args: - mode: Execution mode selection: - - "atomic": Use atomic accumulation (original behavior) - - "reduce": Use non-atomic write + reduce kernel - - Returns: - Compiled executable (either wrapped or raw depending on mode). - """ - # Compile based on mode - if mode == MoeGemm2Mode.REDUCE: - # Determine if we need masked reduction - use_mask = valid_mask is not None - - # Compile GEMM2 with accumulate=False - gemm2_exe = compile_moe_blockscale_gemm2( - model_dim=model_dim, - inter_dim=inter_dim, - experts=experts, - topk=topk, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage2=doweight_stage2, - in_dtype=in_dtype, - out_dtype=out_dtype, - use_cshuffle_epilog=use_cshuffle_epilog, - accumulate=False, - ) - # Compile reduction kernel with masking support - out_s = str(out_dtype).strip().lower() - if out_s in ("f16", "fp16", "half"): - dtype_str = "f16" - elif out_s in ("bf16", "bfloat16"): - dtype_str = "bf16" - else: - dtype_str = "f32" - reduce_exe = compile_moe_reduction( - topk=topk, - model_dim=model_dim, - dtype_str=dtype_str, - use_mask=use_mask, - ) - return _MoeGemm2ReduceWrapper( - gemm2_exe=gemm2_exe, - reduce_exe=reduce_exe, - topk=topk, - model_dim=model_dim, - out_dtype_str=dtype_str, - use_mask=use_mask, - ) - else: - # Compile GEMM2 with accumulate=True (atomic mode) - return compile_moe_blockscale_gemm2( - model_dim=model_dim, - inter_dim=inter_dim, - experts=experts, - topk=topk, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage2=doweight_stage2, - in_dtype=in_dtype, - out_dtype=out_dtype, - use_cshuffle_epilog=use_cshuffle_epilog, - accumulate=True, - ) diff --git a/kernels/moe/moe_blockscale_2stage/reduction.py b/kernels/moe/moe_blockscale_2stage/reduction.py deleted file mode 100644 index 5ec04887d..000000000 --- a/kernels/moe/moe_blockscale_2stage/reduction.py +++ /dev/null @@ -1,226 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""MoE block-scale 2-stage MFMA kernels (stage1 / stage2 / reduction). :: _reduction""" - -import functools - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import scf, vector -from flydsl.expr import arith, const_expr, gpu, range_constexpr -from flydsl.expr.typing import T -from flydsl.runtime.device import get_rocm_arch -from kernels.common import buffer_ops -from kernels.common.kernels_common import _if_else, _if_then - - -@functools.lru_cache(maxsize=1024) -def compile_moe_reduction( - *, - topk: int, - model_dim: int, - dtype_str: str = "f16", - use_mask: bool = False, -): - """Compile a reduction kernel that sums over the topk dimension. - - Input: X [tokens, topk, model_dim] - valid_mask [tokens, topk] (optional, if use_mask=True) - Output: Y [tokens, model_dim] - - This kernel performs: Y[t, d] = sum(X[t, :, d]) for all t, d. - When use_mask=True, only sums slots where valid_mask[t,k]=1. - Used in conjunction with compile_moe_blockscale_gemm2(accumulate=False) to avoid atomic contention. - """ - get_rocm_arch() - ir.ShapedType.get_dynamic_size() - - # Kernel Config - BLOCK_SIZE = 256 - VEC_WIDTH = 8 - - masked = "masked" if use_mask else "" - - module_name = f"bs_moe_reduce_topk{topk}_{dtype_str}{masked}" - - if dtype_str == "f32": - elem_type_tag = "f32" - elif dtype_str == "f16": - elem_type_tag = "f16" - elif dtype_str == "bf16": - elem_type_tag = "bf16" - else: - raise ValueError(f"Unsupported dtype: {dtype_str}") - compute_type = lambda: T.f32 - i8_type = lambda: T.i8 - - def elem_type(): - ty = T.f32 if elem_type_tag == "f32" else (T.f16 if elem_type_tag == "f16" else T.bf16) - return ty() if callable(ty) else ty - - if True: - - @flyc.kernel(name=module_name) - def moe_reduction_kernel( - X: fx.Tensor, - Y: fx.Tensor, - valid_mask: fx.Tensor, - i32_m_tokens: fx.Int32, - ): - m_tokens = fx.Index(i32_m_tokens) - c_topk = fx.Index(topk) - c_model_dim = fx.Index(model_dim) - mask_nbytes_idx = m_tokens * c_topk - elem_bits = 32 if dtype_str == "f32" else 16 - copy_vec_width = 128 // elem_bits # 8 for f16/bf16, 4 for f32 - n_sub = VEC_WIDTH // copy_vec_width # 1 for f16/bf16, 2 for f32 - # Buffer-backed tensors via layout API (all dtypes) - X_buf = fx.rocdl.make_buffer_tensor(X) - Y_buf = fx.rocdl.make_buffer_tensor(Y) - # Scalar buffer resources for tail path and mask - x_rsrc = buffer_ops.create_buffer_resource(X, max_size=True) - y_rsrc = buffer_ops.create_buffer_resource(Y, max_size=True) - mask_rsrc = buffer_ops.create_buffer_resource(valid_mask, max_size=False, num_records_bytes=mask_nbytes_idx) - - token_idx = gpu.block_id("x") - tile_idx = gpu.block_id("y") - tid = gpu.thread_id("x") - - # Guard: token in range (Index is unsigned → auto ult) - tok_ok = token_idx < m_tokens - _if_tok = scf.IfOp(tok_ok) - with _if_then(_if_tok): - tile_cols = BLOCK_SIZE * VEC_WIDTH - c_tile_cols = fx.Index(tile_cols) - c_vecw = fx.Index(VEC_WIDTH) - - col_base = tile_idx * c_tile_cols + tid * c_vecw - - # Guard: any work in bounds (Index < → ult) - col_ok = col_base < c_model_dim - _if_col = scf.IfOp(col_ok) - with _if_then(_if_col): - # Fast path: full vector in-bounds (Index <= → ule) - end_ok = col_base + c_vecw <= c_model_dim - _if_full = scf.IfOp(end_ok, has_else=True) - with _if_then(_if_full): - # ── Vector path via layout API (all dtypes) ── - # fx.copy auto-iterates when atom width < VEC_WIDTH - # (e.g. f32: BufferCopy128b handles 4, fx.copy issues 2 calls for 8) - copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) - vec_type_c = T.vec(copy_vec_width, compute_type()) - vec_type_e = T.vec(copy_vec_width, elem_type()) - - acc_vecs = [vector.broadcast(vec_type_c, fx.Float32(0.0).ir_value()) for _ in range(n_sub)] - elem_dtype = fx.Numeric.from_ir_type(elem_type()) - - tok_i32 = fx.Int32(token_idx) - tile_i32 = fx.Int32(tile_idx) - tid_i32 = fx.Int32(tid) - - for k in range_constexpr(topk): - # X[token, k, :] → tile → thread's VEC_WIDTH slice - x_row = X_buf[tok_i32, fx.Int32(k), None] - x_tiled = fx.logical_divide(x_row, fx.make_layout(tile_cols, 1)) - x_div = fx.logical_divide(x_tiled[None, tile_i32], fx.make_layout(VEC_WIDTH, 1)) - x_thread = x_div[None, tid_i32] - - if const_expr(use_mask): - m_idx_i32 = fx.Int32(token_idx * c_topk + fx.Index(k)) - mv = buffer_ops.buffer_load(mask_rsrc, m_idx_i32, vec_width=1, dtype=i8_type()) - mv_ok = mv != fx.Int8(0) - - if const_expr(n_sub > 1): - x_inner = fx.logical_divide(x_thread, fx.make_layout(copy_vec_width, 1)) - for si in range_constexpr(n_sub): - src = x_inner[None, fx.Int32(si)] if n_sub > 1 else x_thread - r = fx.make_rmem_tensor(copy_vec_width, elem_dtype) - fx.copy_atom_call(copy_atom, src, r) - vec_e = fx.memref_load_vec(r) - - if const_expr(use_mask): - zero_e = vector.broadcast(vec_type_e, arith.constant(0.0, type=elem_type())) - vec_e = mv_ok.select(vec_e, zero_e) - - if const_expr(elem_bits < 32): - vec_c = vec_e.extf(vec_type_c) - else: - vec_c = vec_e - acc_vecs[si] = acc_vecs[si] + vec_c - - # ── Store results ── - if const_expr(n_sub > 1): - y_row = Y_buf[tok_i32, None] - y_tiled = fx.logical_divide(y_row, fx.make_layout(tile_cols, 1)) - y_div = fx.logical_divide(y_tiled[None, tile_i32], fx.make_layout(VEC_WIDTH, 1)) - y_inner = fx.logical_divide(y_div[None, tid_i32], fx.make_layout(copy_vec_width, 1)) - - for si in range_constexpr(n_sub): - out_vec = acc_vecs[si] - if const_expr(elem_bits < 32): - out_vec = out_vec.truncf(vec_type_e) - - if const_expr(n_sub > 1): - dst = y_inner[None, fx.Int32(si)] - else: - y_row = Y_buf[tok_i32, None] - y_tiled = fx.logical_divide(y_row, fx.make_layout(tile_cols, 1)) - y_div = fx.logical_divide(y_tiled[None, tile_i32], fx.make_layout(VEC_WIDTH, 1)) - dst = y_div[None, tid_i32] - - r_out = fx.make_rmem_tensor(copy_vec_width, elem_dtype) - fx.memref_store_vec(out_vec, r_out) - fx.copy_atom_call(copy_atom, r_out, dst) - - with _if_else(_if_full): - for lane in range_constexpr(VEC_WIDTH): - col = col_base + fx.Index(lane) - lane_ok = col < c_model_dim - _if_lane = scf.IfOp(lane_ok) - with _if_then(_if_lane): - a = arith.constant(0.0, type=compute_type()) - token_base = token_idx * c_topk - for k in range_constexpr(topk): - k_idx = fx.Index(k) - x_idx_i32 = fx.Int32((token_base + k_idx) * c_model_dim + col) - if const_expr(use_mask): - m_idx_i32 = fx.Int32(token_base + k_idx) - mv = buffer_ops.buffer_load(mask_rsrc, m_idx_i32, vec_width=1, dtype=i8_type()) - v = (mv != fx.Int8(0)).select( - buffer_ops.buffer_load(x_rsrc, x_idx_i32, vec_width=1, dtype=elem_type()), - arith.constant(0.0, type=elem_type()), - ) - else: - v = buffer_ops.buffer_load(x_rsrc, x_idx_i32, vec_width=1, dtype=elem_type()) - if const_expr(dtype_str in ("f16", "bf16")): - v = v.extf(compute_type()) - a = a + v - - out = a - if const_expr(dtype_str in ("f16", "bf16")): - out = out.truncf(elem_type()) - y_idx_i32 = fx.Int32(token_idx * c_model_dim + col) - buffer_ops.buffer_store(out, y_rsrc, y_idx_i32) - - # ── Host launcher (flyc.jit + .launch) ──────────────────────────────── - tile_size = BLOCK_SIZE * VEC_WIDTH - gy_static = (model_dim + tile_size - 1) // tile_size - - @flyc.jit - def launch_moe_reduction( - X: fx.Tensor, - Y: fx.Tensor, - valid_mask: fx.Tensor, - i32_m_tokens: fx.Int32, - stream: fx.Stream, - ): - gx = fx.Index(i32_m_tokens) - moe_reduction_kernel(X, Y, valid_mask, i32_m_tokens).launch( - grid=(gx, gy_static, 1), - block=(BLOCK_SIZE, 1, 1), - stream=stream, - ) - - return launch_moe_reduction diff --git a/kernels/moe/moe_gemm_2stage/__init__.py b/kernels/moe/moe_gemm_2stage/__init__.py index ba4f617e9..60667fe3f 100644 --- a/kernels/moe/moe_gemm_2stage/__init__.py +++ b/kernels/moe/moe_gemm_2stage/__init__.py @@ -4,6 +4,14 @@ """MoE 2-stage MFMA kernels (stage1 / stage2 / reduction). Split from the former monolithic ``moe_gemm_2stage.py``; public API unchanged. + +.. deprecated:: + These kernels use the legacy FlyDSL authoring API (``SmemAllocator`` / + ``SmemPtr`` + raw ``buffer_ops``) and will be deprecated soon. New MoE work + should use the current ``fx.*`` surface (``make_buffer_tensor`` + + ``SharedAllocator`` + ``fx.copy`` / ``fx.gemm``); see ``kernels/moe/mxfp_moe`` + for the fused a4w4/a8w4 pipeline and the ``kernel-code-cleanup`` skill for the + migration map. """ from kernels.moe.moe_gemm_2stage.gemm1 import compile_moe_gemm1 diff --git a/kernels/moe/moe_gemm_2stage/gemm1.py b/kernels/moe/moe_gemm_2stage/gemm1.py index c0107c3a3..d54cb590c 100644 --- a/kernels/moe/moe_gemm_2stage/gemm1.py +++ b/kernels/moe/moe_gemm_2stage/gemm1.py @@ -1,7 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""MoE GEMM stage1 (MFMA) kernel builder.""" +"""MoE GEMM stage1 (MFMA) kernel builder. + +Legacy authoring API (SmemAllocator/SmemPtr + raw buffer_ops); slated for +deprecation -- refactor to the current fx.* surface (make_buffer_tensor + +SharedAllocator + fx.copy/fx.gemm). See kernels/moe/mxfp_moe and the +kernel-code-cleanup skill. +""" import functools import os @@ -462,7 +468,8 @@ def silu(x): 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." + 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) diff --git a/kernels/moe/moe_gemm_2stage/gemm2.py b/kernels/moe/moe_gemm_2stage/gemm2.py index 4d46944ff..481483997 100644 --- a/kernels/moe/moe_gemm_2stage/gemm2.py +++ b/kernels/moe/moe_gemm_2stage/gemm2.py @@ -1,7 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""MoE GEMM stage2 (MFMA) kernel builder + reduce-mode dispatch.""" +"""MoE GEMM stage2 (MFMA) kernel builder + reduce-mode dispatch. + +Legacy authoring API (SmemAllocator/SmemPtr + raw buffer_ops); slated for +deprecation -- refactor to the current fx.* surface (make_buffer_tensor + +SharedAllocator + fx.copy/fx.gemm). See kernels/moe/mxfp_moe and the +kernel-code-cleanup skill. +""" import functools import logging @@ -215,7 +221,8 @@ def compile_moe_gemm2( if out_is_bf16: if not supports_bf16_global_atomics(gpu_arch): raise ValueError( - f"out_dtype='bf16' requires bf16 global atomics ({bf16_global_atomics_arch_description()}), got arch={gpu_arch!r}" + f"out_dtype='bf16' requires bf16 global atomics " + f"({bf16_global_atomics_arch_description()}), got arch={gpu_arch!r}" ) if out_is_f32: @@ -459,7 +466,8 @@ def _moe_gemm2_then_body(): 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." + 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) diff --git a/kernels/moe/moe_gemm_2stage/reduction.py b/kernels/moe/moe_gemm_2stage/reduction.py index 9031317a2..73db2cce7 100644 --- a/kernels/moe/moe_gemm_2stage/reduction.py +++ b/kernels/moe/moe_gemm_2stage/reduction.py @@ -1,7 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""MoE stage2 reduction kernel builder.""" +"""MoE stage2 reduction kernel builder. + +Legacy authoring API (SmemAllocator/SmemPtr + raw buffer_ops); slated for +deprecation -- refactor to the current fx.* surface. See kernels/moe/mxfp_moe and +the kernel-code-cleanup skill. +""" import functools diff --git a/kernels/moe/mxfp_moe/__init__.py b/kernels/moe/mxfp_moe/__init__.py new file mode 100644 index 000000000..5f82f9127 --- /dev/null +++ b/kernels/moe/mxfp_moe/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Fused a4w4 / a8w4 MoE 2-stage kernels (mxfp_moe). + +Hand-specialized CDNA4 (gfx950) MFMA pipeline ported back from aiter: + + - stage1 (:mod:`gemm1`): fused gate+up GEMM + SiLU + on-device fp4 re-quant, + emitting a sorted fp4 intermediate (``aqout``/``ascaleout``) directly consumed + by stage2 (no host re-quant between stages). ``a_dtype`` selects the activation: + "fp4" (a4w4, mxfp4 A) or "fp8" (a8w4, fp8 e4m3 A x mxfp4 W1); W1/W2 are always + mxfp4 and the intermediate is fp4, so stage2 is identical for both. + - stage2 (:mod:`gemm2`): down-projection GEMM with atomic / reduce / cshuffle / + mxfp4-out epilogues. + +Distinct from the parametric ``mixed_moe_gemm_2stage`` it replaces: this variant +does device-side re-quant and uses a ``cumsum`` + ``m_indices`` sorting contract. +""" + +from kernels.moe.mxfp_moe.gemm1 import compile_gemm1_a4w4_port, gemm1_grid +from kernels.moe.mxfp_moe.gemm2 import compile_gemm2_a4w4_port +from kernels.moe.mxfp_moe.host import flydsl_mxfp4_gemm1, flydsl_mxfp4_gemm2 + +__all__ = [ + "compile_gemm1_a4w4_port", + "gemm1_grid", + "compile_gemm2_a4w4_port", + "flydsl_mxfp4_gemm1", + "flydsl_mxfp4_gemm2", +] diff --git a/kernels/moe/mxfp_moe/dpp_utils.py b/kernels/moe/mxfp_moe/dpp_utils.py new file mode 100644 index 000000000..620d68a25 --- /dev/null +++ b/kernels/moe/mxfp_moe/dpp_utils.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + + +def _to_ir(v): + """Coerce DSL Numeric values to raw MLIR values.""" + from flydsl._mlir import ir as _ir + from flydsl.expr import arith as _arith_ext + + if isinstance(v, int): + return _arith_ext.unwrap(_arith_ext.constant(v, type=_ir.IntegerType.get_signless(32))) + if isinstance(v, float): + return _arith_ext.unwrap(_arith_ext.constant(v, type=_ir.F32Type.get())) + if not isinstance(v, _ir.Value) and hasattr(v, "ir_value"): + return v.ir_value() + return v + + +def update_dpp_i32( + old, + src, + dpp_ctrl: int, + row_mask: int = 0xF, + bank_mask: int = 0xF, + bound_ctrl: bool = False, + **kw, +): + """Wrapper for ``llvm.amdgcn.update.dpp.i32``. + + DPP controls are immediate operands. Common CDNA values: + 280/264 for row xor-8, 276/260 for row xor-4, 78 for xor-2, + and 177 for xor-1 within a 16-lane row. + """ + from flydsl._mlir import ir as _ir + from flydsl._mlir.dialects import llvm as _llvm + from flydsl.expr import arith as _arith_ext + from flydsl.expr.typing import T + + return _llvm.call_intrinsic( + T.i32, + "llvm.amdgcn.update.dpp.i32", + [ + _to_ir(old), + _to_ir(src), + _arith_ext.unwrap(_arith_ext.constant(dpp_ctrl, type=T.i32)), + _arith_ext.unwrap(_arith_ext.constant(row_mask, type=T.i32)), + _arith_ext.unwrap(_arith_ext.constant(bank_mask, type=T.i32)), + _arith_ext.unwrap(_arith_ext.constant(bound_ctrl, type=_ir.IntegerType.get_signless(1))), + ], + [], + [], + **kw, + ) diff --git a/kernels/moe/mxfp_moe/gemm1.py b/kernels/moe/mxfp_moe/gemm1.py new file mode 100644 index 000000000..23ee69047 --- /dev/null +++ b/kernels/moe/mxfp_moe/gemm1.py @@ -0,0 +1,833 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from . import dpp_utils +from .mxfp4_gemm_common import ( + _e8m0_from_amax, + _fabs_f32, + _global_i32_at, + _global_i32_buffer_tiles, + _global_i32_buffer_view, + _global_i32_load, + _global_scalar_tiles, + _inline_dpp_quad_amax, + _inline_e8m0, + _layout_idx, + _lds_swizzle_mask, + _pkmax_u16, + _raw, + _scalar_store, + _scale_mma_atoms, + _silu_mul_batch, + _udiv, + _umax_i32, + _umod, + bq_bytes_for, + bscale_bytes_for, + k_half_for, + k_tiles_total_for, + kas_per_chunk_dw_for, + kbs_per_expert_dw_for, + kBS_stride_k0_dw, + kbs_stride_n0_dw_for, + kmchunks_for, + kStages, + kunroll_for, + lds_acc_bytes_for, + num_n_blocks_for, +) + + +def n_out_for(inter): + return 2 * inter + + +def k_g2_half_for(inter): + return inter // 2 + + +def out_as_per_chunk_dw_for(inter): + return ((inter // 32) // 4 // 2) * 64 + + +def gemm1_grid(n_tokens, BM, *, NE, TOPK, INTER, BN=256): + num_n_blocks = num_n_blocks_for(n_out_for(INTER), BN) + if BM == 128: + max_m_blocks = (n_tokens * TOPK + NE * (BM - 1) + BM - 1) // BM + else: + active = min(n_tokens * TOPK, NE) + max_m_blocks = (n_tokens * TOPK + active * (BM - 1) + BM - 1) // BM + return max_m_blocks * num_n_blocks + + +@flyc.jit +def _gemm1_body( + lds_raw_ptr, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_mind, + arg_aqout, + arg_ascaleout, + arg_hidden, + bx_i32, + lane, + wave, + use_nt, + i32_ntok, + i32_total_m_blocks, + *, + BM, + BN, + BK, + inline_quant=False, + a_dtype="fp4", + K, + N_OUT, + NE, + interleave=False, +): + # A-code tile bytes/row: fp4 packs 2 codes/byte (BK/2); fp8 is 1 B/elem (BK). + KH_TILE = BK if a_dtype == "fp8" else BK // 2 + K_HALF = k_half_for(K) + # A row bytes: fp4 = K/2, fp8 = K (B is always mxfp4 -> keeps K_HALF). + A_ROW_BYTES = K if a_dtype == "fp8" else K_HALF + K_TILES_TOTAL = k_tiles_total_for(K, BK) + kUnroll = kunroll_for(K, BK) + kAS_per_chunk_dw = kas_per_chunk_dw_for(K) + kBS_stride_n0_dw = kbs_stride_n0_dw_for(K) + kBS_per_expert_dw = kbs_per_expert_dw_for(N_OUT, K) + BQ_BYTES = bq_bytes_for(NE, N_OUT, K) + BSCALE_BYTES = bscale_bytes_for(NE, N_OUT, K) + NUM_N_BLOCKS = num_n_blocks_for(N_OUT, BN) + inter = N_OUT // 2 + OUT_AS_PER_CHUNK_DW = out_as_per_chunk_dw_for(inter) + K_G2_HALF = k_g2_half_for(inter) + kAStages, kSubBlocks, kMChunks, _ = _bm_constants(BM, BN, KH_TILE, K_TILES_TOTAL) + + BN_INT = BN // 2 + b_aux = 2 if use_nt else 0 + M_REPS = BM // 16 + + n_block_idx = bx_i32 % fx.Int32(NUM_N_BLOCKS) + m_block_idx = bx_i32 // fx.Int32(NUM_N_BLOCKS) + e = rocdl.readfirstlane(T.i32, _raw(_global_i32_at(arg_eids, m_block_idx))) + m_row = m_block_idx * fx.Int32(BM) + + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + lane_div_8 = lane // fx.Int32(8) + lane_mod_8 = lane % fx.Int32(8) + + aq_num_records = fx.Int64(i32_ntok * fx.Int32(A_ROW_BYTES)) + _asc_per_mb = max(BM // 32, 1) * kAS_per_chunk_dw * 4 + ascale_num = fx.Int64(i32_total_m_blocks) * fx.Int64(_asc_per_mb) + + bq_tiles = _global_i32_buffer_tiles(arg_bq, BQ_BYTES, 4) + bq_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(b_aux), fx.Int32) + bq_reg_lay = fx.make_layout(4, 1) + + bscale_tiles = _global_i32_buffer_tiles(arg_bscale, BSCALE_BYTES, 1) + bscale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + bscale_reg_lay = fx.make_layout(1, 1) + + # aq/ascale: global->LDS async DMA (no register fragment), via BufferCopyLDS. + aq_buf = _global_i32_buffer_view(arg_aq, aq_num_records) + aq_dma_tiles4 = fx.logical_divide(aq_buf, fx.make_layout(4, 1)) + aq_dma_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), fx.Int32) + + ascale_buf = _global_i32_buffer_view(arg_ascale, ascale_num) + ascale_dma_tiles4 = fx.logical_divide(ascale_buf, fx.make_layout(4, 1)) + ascale_dma_tiles1 = fx.logical_divide(ascale_buf, fx.make_layout(1, 1)) + ascale_dma_atom16 = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), fx.Int32) + ascale_dma_atom4 = fx.make_copy_atom(fx.rocdl.BufferCopyLDS32b(), fx.Int32) + + if const_expr(inline_quant): + hidden_num = fx.Int64(i32_ntok * fx.Int32(K * 2)) + hidden_tiles = _global_i32_buffer_tiles(arg_hidden, hidden_num, 4) + hidden_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Int32) + hidden_reg_lay = fx.make_layout(4, 1) + + # Union LDS region [s_aq | s_asc], reused as lds_acc (f32 accumulator) in the + # epilogue. s_aq/lds_acc start at lds_raw_ptr; s_asc follows at + # +kAStages*BM*KH_TILE. + + cached_actual_row = [] + cached_row_inline = None + if const_expr(inline_quant): + rcls = wave * fx.Int32(4) + lane_div_16 + cached_row_inline = _global_i32_at(arg_mind, m_row + rcls) + else: + for sub in range_constexpr(kSubBlocks): + idx = m_row + wave * fx.Int32(BM // 4) + fx.Int32(sub * 8) + lane_div_8 + cached_actual_row.append(_global_i32_at(arg_mind, idx)) + + # -- b_load_s_base[j], readfirstlane'd uniform per wave -------------------- + N0_HALF = N_OUT // 32 + b_load_s_base = [] + for j in range_constexpr(4): + if const_expr(interleave): + col = n_block_idx * fx.Int32(BN) + wave * fx.Int32(BN // 4) + fx.Int32(j * 16) + else: + tile_il = n_block_idx * fx.Int32(16) + wave * fx.Int32(4) + fx.Int32(j) + g = tile_il & fx.Int32(1) + n0 = tile_il >> fx.Int32(1) + col = (g * fx.Int32(N0_HALF) + n0) * fx.Int32(16) + v = (e * fx.Int32(N_OUT) + col) * fx.Int32(K_HALF) + b_load_s_base.append(rocdl.readfirstlane(T.i32, v)) + + # -- b_scale_s_base / _hi -------------------------------------------------- + if const_expr(interleave): + mni_base = n_block_idx * fx.Int32(BN // 32) + wave * fx.Int32(BN // 128) + np_list = [mni_base, mni_base + fx.Int32(1)] + else: + np_gate = n_block_idx * fx.Int32(BN // 64) + wave + np_list = [np_gate, np_gate + fx.Int32(N_OUT // 64)] + b_scale_s_base, b_scale_s_base_hi = [], [] + for mw in range_constexpr(2): + base = (e * fx.Int32(kBS_per_expert_dw) + np_list[mw] * fx.Int32(kBS_stride_n0_dw)) * fx.Int32(4) + base = rocdl.readfirstlane(T.i32, base) + b_scale_s_base.append(base) + b_scale_s_base_hi.append(base + fx.Int32(16 * kBS_stride_k0_dw * 4)) + + # f32 accumulator fragments (one i32[4]->f32[4] per (mchunk, J)); seeded to 0 + # so the first K-iter's fx.gemm computes A*B+0 (matches the raw init path). + scale_atoms = _scale_mma_atoms(a_dtype) + accm = [[fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) for _ in range(4)] for _ in range(kMChunks)] + for i in range_constexpr(kMChunks): + for J in range_constexpr(4): + accm[i][J].store(fx.Vector.filled(4, 0.0, fx.Float32)) + b = [[[None, None] for _ in range(4)] for _ in range(kStages)] + b_scale_v = [[None, None] for _ in range(kStages)] + + # s_aq as flat i32, divided into 4-element (128-bit) and 1-element tiles. + s_aq_i32_flat = fx.make_view( + fx.recast_iter(fx.Int32, lds_raw_ptr), + fx.make_layout(kAStages * BM * KH_TILE // 4, 1), + ) + s_aq_i32x4_tiles = fx.logical_divide(s_aq_i32_flat, fx.make_layout(4, 1)) + s_aq_i32x1_tiles = fx.logical_divide(s_aq_i32_flat, fx.make_layout(1, 1)) + i32x4_copy_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + i32x4_reg_lay = fx.make_layout(4, 1) + + def issue_a_load_lds(slot, kt): + # fp4: 128 B/row/tile in one cooperative pass. fp8: A is 1 B/elem so the tile + # is 256 B/row (KH_TILE) -> a second pass fills the upper 128 B (gmem col +128, + # LDS off +128); the row-swizzle mask stays within each 128 B half. + for sub in range_constexpr(kSubBlocks): + lds_row = wave * fx.Int32(BM // 4) + fx.Int32(sub * 8) + mask = _lds_swizzle_mask(lds_row + lane_div_8) + voffset = ((lane_mod_8 * fx.Int32(16)) ^ mask) + cached_actual_row[sub] * fx.Int32(A_ROW_BYTES) + off = fx.Int32(slot * (BM * KH_TILE)) + lds_row * fx.Int32(KH_TILE) + fx.copy( + aq_dma_atom, + fx.slice(aq_dma_tiles4, (None, voffset // fx.Int32(16))), + fx.slice(s_aq_i32x4_tiles, (None, off // fx.Int32(16))), + soffset=fx.Int32(kt * KH_TILE) // fx.Int32(4), + ) + if const_expr(a_dtype == "fp8"): + fx.copy( + aq_dma_atom, + fx.slice(aq_dma_tiles4, (None, (voffset + fx.Int32(128)) // fx.Int32(16))), + fx.slice(s_aq_i32x4_tiles, (None, (off + fx.Int32(128)) // fx.Int32(16))), + soffset=fx.Int32(kt * KH_TILE) // fx.Int32(4), + ) + + def _lds_i32x4_frag(tile_idx): + # ds_read_b128 straight into an i32[4] register fragment (kept as a tensor + # so it can feed fx.gemm directly). + r = fx.make_rmem_tensor(i32x4_reg_lay, fx.Int32) + fx.copy_atom_call(i32x4_copy_atom, fx.slice(s_aq_i32x4_tiles, (None, tile_idx)), r) + return r + + def issue_a_ds_read(slot): + mask = _lds_swizzle_mask(lane_mod_16) + a = [[None, None] for _ in range(kMChunks)] + for k in range_constexpr(2): + if const_expr(a_dtype == "fp8"): + # fp8 128-K operand (v8i32) = two 16 B halves 64 B apart in the row + # (f8f6f4 ABI), packed lo++hi. k selects the 128-K half (upper = +128 B). + kbase = fx.Int32(k * 128) + lo_col = (lane_div_16 * fx.Int32(16) + kbase) ^ mask + hi_col = (lane_div_16 * fx.Int32(16) + kbase + fx.Int32(64)) ^ mask + for i in range_constexpr(kMChunks): + lds_row = lane_mod_16 + fx.Int32(i * 16) + boff = fx.Int32(slot * (BM * KH_TILE)) + lds_row * fx.Int32(KH_TILE) + lo = fx.Vector(fx.memref_load_vec(_lds_i32x4_frag((boff + lo_col) // fx.Int32(16)))) + hi = fx.Vector(fx.memref_load_vec(_lds_i32x4_frag((boff + hi_col) // fx.Int32(16)))) + t = fx.make_rmem_tensor(fx.make_layout(8, 1), fx.Int32) + t.store(lo.shuffle(hi, list(range(8)))) + a[i][k] = t + else: + lds_col = (lane_div_16 * fx.Int32(16) + fx.Int32(k * 64)) ^ mask + for i in range_constexpr(kMChunks): + lds_row = lane_mod_16 + fx.Int32(i * 16) + off = fx.Int32(slot * (BM * KH_TILE)) + lds_row * fx.Int32(KH_TILE) + lds_col + a[i][k] = _lds_i32x4_frag(off // fx.Int32(16)) + return a + + def issue_a_scale_load(): + chunk_base = m_row // fx.Int32(32) + v16 = (wave * fx.Int32(64) + lane) * fx.Int32(16) + v4 = (wave * fx.Int32(64) + lane) * fx.Int32(4) + for sub in range_constexpr(kSubBlocks): + s_chunk = rocdl.readfirstlane(T.i32, (chunk_base + fx.Int32(sub)) * fx.Int32(kAS_per_chunk_dw * 4)) + lds_sub = fx.Int32(sub * kAS_per_chunk_dw * 4) + fx.copy( + ascale_dma_atom16, + fx.slice(ascale_dma_tiles4, (None, v16 // fx.Int32(16))), + fx.slice(asc_i32x4_tiles, (None, (lds_sub + wave * fx.Int32(1024)) // fx.Int32(16))), + soffset=s_chunk // fx.Int32(4), + ) + for d in range_constexpr(3): + byte_off = 4096 + d * 1024 + s_off = rocdl.readfirstlane(T.i32, s_chunk + fx.Int32(byte_off)) + fx.copy( + ascale_dma_atom4, + fx.slice(ascale_dma_tiles1, (None, v4 // fx.Int32(4))), + fx.slice( + asc_i32_tiles, + (None, (lds_sub + fx.Int32(byte_off) + wave * fx.Int32(256)) // fx.Int32(4)), + ), + soffset=s_off // fx.Int32(4), + ) + + # s_asc as flat i32. + s_asc_i32_flat = fx.make_view( + fx.recast_iter(fx.Int32, fx.add_offset(lds_raw_ptr, kAStages * BM * KH_TILE)), + fx.make_layout(kSubBlocks * K_TILES_TOTAL * 64, 1), + ) + asc_i32_tiles = fx.logical_divide(s_asc_i32_flat, fx.make_layout(1, 1)) + asc_i32x4_tiles = fx.logical_divide(s_asc_i32_flat, fx.make_layout(4, 1)) + + def issue_a_scale_ds_read(kt): + out = [] + for sub in range_constexpr(kSubBlocks): + lds_dw = fx.Int32(sub * kAS_per_chunk_dw) + fx.Int32(kt * 64) + lane_div_16 * fx.Int32(16) + lane_mod_16 + out.append(_raw(_global_i32_load(asc_i32_tiles, lds_dw))) + return out + + lib = lane & fx.Int32(3) + lane_shr2_and3 = (lane >> fx.Int32(2)) & fx.Int32(3) + r_in_chunk = wave * fx.Int32(4) + lane_div_16 + + def inline_quant_load_kt(B128_IDX, kt, row_token): + v_voff = row_token * fx.Int32(K * 2) + lane_shr2_and3 * fx.Int32(64) + lib * fx.Int32(16) + s_soff = rocdl.readfirstlane(T.i32, fx.Int32(kt * (BK * 2) + B128_IDX * 256)) + r = fx.make_rmem_tensor(hidden_reg_lay, fx.Int32) + fx.copy( + hidden_copy_atom, + fx.slice(hidden_tiles, (None, v_voff // fx.Int32(16))), + r, + soffset=s_soff // fx.Int32(4), + ) + return r.load() + + def _bf16x2(dw): + return _raw(fx.Vector.from_elements([dw], fx.Int32).bitcast(fx.BFloat16)) + + def _iq_block_amax(h_dw_i): + # Per-32 amax over the 8 bf16 in this lane's block (4 bf16x2 dwords). + hm = [h_dw_i[j] & fx.Int32(0x7FFF7FFF) for j in range_constexpr(4)] + m01 = _pkmax_u16(hm[0], hm[1]) + m23 = _pkmax_u16(hm[2], hm[3]) + m0123 = _pkmax_u16(m01, m23) + lo = m0123 & fx.Int32(0xFFFF) + hi = m0123.shrui(fx.Int32(16)) & fx.Int32(0xFFFF) + return _umax_i32(lo, hi) + + def _iq_pack_store(h_dw_i, qs_raw, B128_IDX, SUB, slot): + # Quantize this lane's 8 elements (4 bf16x2 dwords) and store into the A-LDS + # slot at the position issue_a_ds_read expects. fp4: 8 fp4 -> one i32 in a + # 16 B swizzled block. fp8: 8 fp8 -> two i32; the 16 B block is + # B128_IDX*8 + lsa3*2 + lib//2 and the byte within it is (lib%2)*8 (matches + # the split-64 fp8 read). + r = fx.Int32(SUB * 16) + r_in_chunk + mask_r = _lds_swizzle_mask(r) + row_off = fx.Int32(slot * (BM * KH_TILE)) + r * fx.Int32(KH_TILE) + if const_expr(a_dtype == "fp8"): + blk_byte = fx.Int32(B128_IDX * 128) + lane_shr2_and3 * fx.Int32(32) + (lib >> fx.Int32(1)) * fx.Int32(16) + off = row_off + (blk_byte ^ mask_r) + (lib & fx.Int32(1)) * fx.Int32(8) + # cvt.pk.fp8 packs 2 fp8 into a 16b lane of a vector<2xi16> accumulator + # (dstLoHiSel = lo/hi); two lanes -> 4 fp8 = one i32 stored to LDS. + i16x2 = ir.Type.parse("vector<2xi16>") + zero16 = llvm.BitcastOp(i16x2, _raw(fx.Int32(0))).result + pk0 = rocdl.cvt_scalef32_pk_fp8_bf16(i16x2, zero16, _bf16x2(h_dw_i[0]), qs_raw, 0) + pk0 = rocdl.cvt_scalef32_pk_fp8_bf16(i16x2, pk0, _bf16x2(h_dw_i[1]), qs_raw, 1) + pk1 = rocdl.cvt_scalef32_pk_fp8_bf16(i16x2, zero16, _bf16x2(h_dw_i[2]), qs_raw, 0) + pk1 = rocdl.cvt_scalef32_pk_fp8_bf16(i16x2, pk1, _bf16x2(h_dw_i[3]), qs_raw, 1) + _scalar_store(s_aq_i32x1_tiles, off // fx.Int32(4), fx.Int32(llvm.BitcastOp(T.i32, pk0).result), fx.Int32) + _scalar_store( + s_aq_i32x1_tiles, + (off + fx.Int32(4)) // fx.Int32(4), + fx.Int32(llvm.BitcastOp(T.i32, pk1).result), + fx.Int32, + ) + else: + kb_in_kt = fx.Int32(B128_IDX * 4) + lane_shr2_and3 + off = row_off + ((kb_in_kt * fx.Int32(16)) ^ mask_r) + lib * fx.Int32(4) + pk = _raw(fx.Int32(0)) + for j in range_constexpr(4): + pk = rocdl.cvt_scalef32_pk_fp4_bf16(T.i32, pk, _bf16x2(h_dw_i[j]), qs_raw, j) + _scalar_store(s_aq_i32x1_tiles, off // fx.Int32(4), fx.Int32(pk), fx.Int32) + + def _inline_quant_core_batch(specs, slot, scale_accum): + n = len(specs) + h_dw = [[fx.Int32(_raw(h_v[j])) for j in range_constexpr(4)] for (_b, _s, h_v) in specs] + a = [_iq_block_amax(h_dw[i]) for i in range_constexpr(n)] + s1 = [ + fx.Int32(dpp_utils.update_dpp_i32(_raw(a[i]), _raw(a[i]), 0xB1, 0xF, 0xF, True)) for i in range_constexpr(n) + ] + a = [_umax_i32(a[i], s1[i]) for i in range_constexpr(n)] + s2 = [ + fx.Int32(dpp_utils.update_dpp_i32(_raw(a[i]), _raw(a[i]), 0x4E, 0xF, 0xF, True)) for i in range_constexpr(n) + ] + a = [_umax_i32(a[i], s2[i]) for i in range_constexpr(n)] + e8 = [_inline_e8m0(a[i]) for i in range_constexpr(n)] + for i in range_constexpr(n): + B128_IDX, SUB, _hv = specs[i] + qs_raw = _raw(fx.Float32(_raw(e8[i] << fx.Int32(23)).bitcast(T.f32))) + _iq_pack_store(h_dw[i], qs_raw, B128_IDX, SUB, slot) + pack_byte = B128_IDX * 2 + SUB + scale_accum = scale_accum | (e8[i] << fx.Int32(pack_byte * 8)) + return scale_accum + + def inline_quant_kt(B128_IDX, SUB, slot, kt, row_token, scale_accum): + h_v = inline_quant_load_kt(B128_IDX, kt, row_token) + return _inline_quant_core_batch([(B128_IDX, SUB, h_v)], slot, scale_accum) + + def inline_quant_pack_write(kt, scale_accum): + lane_tgt = lane_shr2_and3 * fx.Int32(16) + r_in_chunk + off = fx.Int32(kt * 256) + lane_tgt * fx.Int32(4) + _scalar_store(asc_i32_tiles, off // fx.Int32(4), scale_accum, fx.Int32) + + def issue_b_load_j(b_slot, K_C, j): + v = (lane_div_16 * fx.Int32(256)) + (lane_mod_16 * fx.Int32(16)) + fx.Int32(K_C * 2048) + for half in range_constexpr(2): + tile_idx = (v + fx.Int32(half * 1024)) // fx.Int32(16) + r = fx.make_rmem_tensor(bq_reg_lay, fx.Int32) + fx.copy( + bq_copy_atom, + fx.slice(bq_tiles, (None, tile_idx)), + r, + soffset=b_load_s_base[j] // fx.Int32(4), + ) + b_slot[j][half] = r + + def issue_b_scale_load(bs_slot, K_C): + v = ((lane_div_16 * fx.Int32(16)) + lane_mod_16) * fx.Int32(4) + K_C_HI = K_C // 16 + imm = (K_C - K_C_HI * 16) * (kBS_stride_k0_dw * 4) + for mw in range_constexpr(2): + s_off = b_scale_s_base[mw] if K_C_HI == 0 else b_scale_s_base_hi[mw] + idx = (v + fx.Int32(imm)) // fx.Int32(4) + r = fx.make_rmem_tensor(bscale_reg_lay, fx.Int32) + fx.copy( + bscale_copy_atom, + fx.slice(bscale_tiles, (None, idx)), + r, + soffset=s_off // fx.Int32(4), + ) + bs_slot[mw] = r.load()[0] + + def _mma(ci, opsel_a, opsel_b, a_frag, b_frag, sa, sb): + # Scaled 16x16x128 MFMA via fx.gemm; accumulate in place (d == c == ci). + # opsel_a/opsel_b select the e8m0 scale byte in the shared 256-K word and are + # baked into the atom, exactly mirroring the raw intrinsic's opsel operands. + fx.gemm(scale_atoms[(opsel_a, opsel_b)], ci, a_frag, b_frag, ci, scale_a=sa, scale_b=sb) + + def mfma_cluster(b_slot, a, a_scale, bs_slot, J): + # Accumulators are pre-seeded to 0, so every issue accumulates (A*B + C). + if const_expr(interleave): + mni = J // 2 + in_b = J % 2 + else: + mni = J % 2 + in_b = J // 2 + sb = bs_slot[mni] + bJ0, bJ1 = b_slot[J][0], b_slot[J][1] + if const_expr(kMChunks == 1): + sa = a_scale[0] + _mma(accm[0][J], 0, 0 + in_b, a[0][0], bJ0, sa, sb) + _mma(accm[0][J], 2, 2 + in_b, a[0][1], bJ1, sa, sb) + else: + for sub in range_constexpr(kSubBlocks): + i0 = sub * 2 + 0 + i1 = sub * 2 + 1 + sa = a_scale[sub] + _mma(accm[i0][J], 0, 0 + in_b, a[i0][0], bJ0, sa, sb) + _mma(accm[i1][J], 1, 0 + in_b, a[i1][0], bJ0, sa, sb) + _mma(accm[i0][J], 2, 2 + in_b, a[i0][1], bJ1, sa, sb) + _mma(accm[i1][J], 3, 2 + in_b, a[i1][1], bJ1, sa, sb) + + _relax_prologue = (BM == 128) and not inline_quant + if const_expr(not inline_quant): + issue_a_scale_load() + for K_C in range_constexpr(kStages): + if const_expr(inline_quant): + scale_accum = fx.Int32(0) + scale_accum = inline_quant_kt(0, 0, K_C, K_C, cached_row_inline, scale_accum) + issue_b_load_j(b[K_C], K_C, 0) + issue_b_load_j(b[K_C], K_C, 1) + scale_accum = inline_quant_kt(1, 0, K_C, K_C, cached_row_inline, scale_accum) + issue_b_load_j(b[K_C], K_C, 2) + issue_b_load_j(b[K_C], K_C, 3) + inline_quant_pack_write(K_C, scale_accum) + else: + issue_a_load_lds(K_C, K_C) + if const_expr(not _relax_prologue): + for j in range_constexpr(4): + issue_b_load_j(b[K_C], K_C, j) + if const_expr(not _relax_prologue): + issue_b_scale_load(b_scale_v[K_C], K_C) + if const_expr(_relax_prologue): + rocdl.sched_barrier(0) + for K_C in range_constexpr(kStages): + for j in range_constexpr(4): + issue_b_load_j(b[K_C], K_C, j) + issue_b_scale_load(b_scale_v[K_C], K_C) + + for OFFSET in range_constexpr(kUnroll): + K_C = kStages + OFFSET + read_slot = OFFSET % kAStages + write_slot = K_C % kAStages + slot_b = OFFSET % kStages + gpu.barrier() + if const_expr(BM == 128): + asc_cur = issue_a_scale_ds_read(K_C - kStages) + a_cur = issue_a_ds_read(read_slot) + else: + a_cur = issue_a_ds_read(read_slot) + asc_cur = issue_a_scale_ds_read(K_C - kStages) + if const_expr(not inline_quant): + issue_a_load_lds(write_slot, K_C) + if const_expr(inline_quant): + h_v0 = inline_quant_load_kt(0, K_C, cached_row_inline) + h_v1 = inline_quant_load_kt(1, K_C, cached_row_inline) + rocdl.sched_barrier(0) + for J in range_constexpr(4): + if const_expr(BM != 128): + rocdl.sched_barrier(0) + rocdl.s_setprio(1) + mfma_cluster(b[slot_b], a_cur, asc_cur, b_scale_v[slot_b], J) + if const_expr(BM != 128): + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + issue_b_load_j(b[slot_b], K_C, J) + rocdl.sched_barrier(0) + issue_b_scale_load(b_scale_v[slot_b], K_C) + if const_expr(inline_quant): + scale_accum = _inline_quant_core_batch([(0, 0, h_v0), (1, 0, h_v1)], write_slot, fx.Int32(0)) + inline_quant_pack_write(K_C, scale_accum) + + for S in range_constexpr(kStages): + kt = K_TILES_TOTAL - kStages + S + gpu.barrier() + if const_expr(BM == 128): + asc_cur = issue_a_scale_ds_read(kt) + a_cur = issue_a_ds_read(kt % kAStages) + else: + a_cur = issue_a_ds_read(kt % kAStages) + asc_cur = issue_a_scale_ds_read(kt) + for J in range_constexpr(4): + mfma_cluster(b[kt % kStages], a_cur, asc_cur, b_scale_v[kt % kStages], J) + + gpu.barrier() + + # lds_acc reuses the s_aq region (offset 0) as an f32 accumulator. + acc_layout = fx.make_layout((BM, BN), (BN, 1)) + + def acc_idx(row, col): + return _layout_idx(acc_layout, row, col) + + acc_copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32) + acc_reg_lay = fx.make_layout(1, 1) + acc_flat_view = fx.make_view(fx.recast_iter(fx.Float32, lds_raw_ptr), fx.make_layout(BM * BN, 1)) + acc_flat_tiles = fx.logical_divide(acc_flat_view, fx.make_layout(1, 1)) + + def acc_store(idx, value): + r = fx.make_rmem_tensor(acc_reg_lay, fx.Float32) + r.store(fx.Vector.from_elements([fx.Float32(value)], fx.Float32)) + fx.copy_atom_call(acc_copy_atom, r, fx.slice(acc_flat_tiles, (None, idx))) + + def acc_load(idx): + r = fx.make_rmem_tensor(acc_reg_lay, fx.Float32) + fx.copy_atom_call(acc_copy_atom, fx.slice(acc_flat_tiles, (None, idx)), r) + return r.load()[0] + + for i in range_constexpr(kMChunks): + row_base = fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + for J in range_constexpr(4): + is_up = (J % 2) == 1 + J_local = J // 2 + col_local = wave * fx.Int32(32) + fx.Int32(J_local * 16) + lane_mod_16 + lds_col = (fx.Int32(128) + col_local) if is_up else col_local + vec = fx.Vector(fx.memref_load_vec(accm[i][J])) + for v in range_constexpr(4): + idx = acc_idx(row_base + fx.Int32(v), lds_col) + acc_store(idx, vec[v]) + + gpu.barrier() + + tx_i32 = fx.Int32(gpu.thread_id("x")) + m_lane = tx_i32 // fx.Int32(16) + n_lane = tx_i32 % fx.Int32(16) + wave_grp = n_lane // fx.Int32(4) + kk = n_lane % fx.Int32(4) + + aqout_layout = fx.make_layout((BM, K_G2_HALF), (K_G2_HALF, 1)) + # UniversalCopy has no nontemporal/cache-hint knob; dropped (perf-neutral). + aqout_tiles = _global_scalar_tiles(arg_aqout, fx.Int32, 1 << 24) + scales_per_mr = [None] * M_REPS + + for mr in range_constexpr(M_REPS): + row_local = fx.Int32(mr * 16) + m_lane + + gate_vs = [None] * 8 + up_vs = [None] * 8 + for ee in range_constexpr(8): + col_in_grp = fx.Int32(8) * kk + fx.Int32(ee) + gate_col = wave_grp * fx.Int32(32) + col_in_grp + up_col = fx.Int32(128) + gate_col + gate_vs[ee] = acc_load(acc_idx(row_local, gate_col)) + up_vs[ee] = acc_load(acc_idx(row_local, up_col)) + result = _silu_mul_batch(gate_vs, up_vs) + + local_max = _fabs_f32(result[0]) + for ee in range_constexpr(1, 8): + local_max = local_max.maximumf(_fabs_f32(result[ee])) + lm_i = _inline_dpp_quad_amax(fx.Int32(_raw(local_max).bitcast(T.i32))) + local_max = fx.Float32(_raw(lm_i).bitcast(T.f32)) + + e8m0, qscale = _e8m0_from_amax(local_max) + scales_per_mr[mr] = e8m0 + + packed_i32 = _raw(fx.Int32(0)) + qscale_raw = _raw(qscale) + for w in range_constexpr(4): + packed_i32 = rocdl.cvt_scalef32_pk_fp4_f32( + T.i32, packed_i32, _raw(result[2 * w]), _raw(result[2 * w + 1]), qscale_raw, w + ) + packed = fx.Int32(packed_i32) + + byte_pos = n_block_idx * fx.Int32(BN_INT // 2) + wave_grp * fx.Int32(16) + kk * fx.Int32(4) + out_row = m_row + row_local + store_off = _layout_idx(aqout_layout, out_row, byte_pos) + _scalar_store(aqout_tiles, store_off // fx.Int32(4), packed, fx.Int32) + + # (chunk, ku, wave_grp, m_lane) -> dword index; shape is a placeholder. + ascaleout_layout = fx.make_layout((1 << 20, 2, 4, 16), (OUT_AS_PER_CHUNK_DW, 64, 16, 1)) + ascaleout_i8_tiles = _global_scalar_tiles(arg_ascaleout, fx.Int8, 1 << 26) + ascaleout_i16_tiles = _global_scalar_tiles(arg_ascaleout, fx.Int16, 1 << 25) + if kk == fx.Int32(0): + ku = n_block_idx >> fx.Int32(1) + ikxdl = n_block_idx & fx.Int32(1) + if const_expr(BM == 16): + chunk = m_block_idx + dword_off = _layout_idx(ascaleout_layout, chunk, ku, wave_grp, m_lane) + addr = dword_off * fx.Int32(4) + ikxdl * fx.Int32(2) + _scalar_store(ascaleout_i8_tiles, addr, scales_per_mr[0], fx.Int8) + else: + for sub in range_constexpr(kSubBlocks): + chunk = m_block_idx * fx.Int32(kSubBlocks) + fx.Int32(sub) + dword_off = _layout_idx(ascaleout_layout, chunk, ku, wave_grp, m_lane) + pair_i32 = scales_per_mr[sub * 2 + 0] | (scales_per_mr[sub * 2 + 1] << fx.Int32(8)) + addr = dword_off * fx.Int32(4) + ikxdl * fx.Int32(2) + _scalar_store(ascaleout_i16_tiles, addr // fx.Int32(2), pair_i32, fx.Int16) + + +def _bm_constants(BM, BN, KH_TILE, K_TILES_TOTAL): + kAStages = 2 if BM == 128 else 3 + kSubBlocks = 1 if BM < 32 else BM // 32 + kMChunks = kmchunks_for(BM) + s_aq_bytes = kAStages * BM * KH_TILE + s_asc_bytes = kSubBlocks * K_TILES_TOTAL * 256 + lds_acc_bytes = lds_acc_bytes_for(BM, BN) + lds_bytes = max(s_aq_bytes + s_asc_bytes, lds_acc_bytes) + return kAStages, kSubBlocks, kMChunks, lds_bytes + + +_G1_VARIANTS = { + "fp4": { + (32, True, False), + (32, False, False), + (64, False, False), + (128, False, False), + (16, True, True), + }, + # a8w4: fp8 (e4m3) A x mxfp4 W1. Same tiling as fp4; A is 1 B/elem so the + # A-LDS tile doubles. inline_quant (bf16 hidden -> fp8) supported at (16,True). + "fp8": { + (32, True, False), + (32, False, False), + (64, False, False), + (128, False, False), + (16, True, True), + }, +} + + +def compile_gemm1_a4w4_port( + BM=32, + use_nt=True, + inline_quant=False, + *, + D_HIDDEN, + D_INTER, + NE, + TOPK, + BN=256, + BK=256, + interleave=False, + xcd_swizzle=0, + a_dtype="fp4", +): + if a_dtype not in ("fp4", "fp8"): + raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + if (BM, use_nt, inline_quant) not in _G1_VARIANTS[a_dtype]: + raise AssertionError( + f"unsupported gemm1 variant (a_dtype={a_dtype}, BM={BM}, use_nt={use_nt}, inline_quant={inline_quant})" + ) + + assert BN == 256 and BK == 256, f"only BN==BK==256 supported, got BN={BN} BK={BK}" + KH_TILE = BK if a_dtype == "fp8" else BK // 2 + _K = D_HIDDEN + assert _K % BK == 0, f"D_HIDDEN (K) must be a multiple of {BK}, got {_K}" + _INTER = D_INTER + _N_OUT = n_out_for(_INTER) + assert _N_OUT % BN == 0, f"2*D_INTER (N_OUT) must be a multiple of {BN}, got {_N_OUT}" + _NE = NE + _K_TILES_TOTAL = k_tiles_total_for(_K, BK) + _NUM_N_BLOCKS = num_n_blocks_for(_N_OUT, BN) + + _, _, _, lds_bytes = _bm_constants(BM, BN, KH_TILE, _K_TILES_TOTAL) + + variant_tag = "iq" if inline_quant else ("nt" if use_nt else "cached") + # Tag with H/INTER/NE so different shape specializations get distinct + # kernel/smem symbols (so KIMI and non-KIMI instances never collide). + gu_tag = "il" if interleave else "sep" + name_suffix = f"{a_dtype}_h{_K}_i{_INTER}_ne{_NE}_bm{BM}_{variant_tag}_{gu_tag}" + if xcd_swizzle > 0: + name_suffix += f"_xcd{xcd_swizzle}" + + @fx.struct + class SharedStorage: + raw: fx.Array[fx.Uint8, lds_bytes, 16] + + @flyc.kernel(name=f"gemm1_a4w4_port_{name_suffix}", known_block_size=[256, 1, 1]) + def gemm1_kernel( + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_mind: fx.Int64, + i32_ntok: fx.Int32, + arg_aqout: fx.Int64, + arg_ascaleout: fx.Int64, + arg_hidden: fx.Int64, + ): + lds_raw_ptr = fx.SharedAllocator().allocate(SharedStorage).peek().raw.ptr + tx = gpu.thread_id("x") + bx = gpu.block_id("x") + tx_i32 = fx.Int32(tx) + bx_i32 = fx.Int32(bx) + lane = tx_i32 % fx.Int32(64) + wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) + cumsum0 = _global_i32_at(arg_cumsum, fx.Int32(0)) + total_m_blocks = cumsum0 // fx.Int32(BM) + bound = total_m_blocks * fx.Int32(_NUM_N_BLOCKS) + + _NXCD = 8 + _xq = _udiv(bound, _NXCD) + _xr = _umod(bound, _NXCD) + _SW = xcd_swizzle + + def _xcd(pid): + xc = _umod(pid, _NXCD) + wgid = xc * _xq + fx.Int32(arith.minsi(_raw(xc), _raw(_xr))) + _udiv(pid, _NXCD) + _ng = fx.Int32(_SW * _NUM_N_BLOCKS) + group_id = wgid // _ng + first_pid_m = group_id * fx.Int32(_SW) + remaining_m = total_m_blocks - first_pid_m + group_size_m = fx.Int32(arith.minsi(_raw(remaining_m), _raw(fx.Int32(_SW)))) + wig = wgid % _ng + m_block = first_pid_m + (wig % group_size_m) + n_block = wig // group_size_m + return m_block * fx.Int32(_NUM_N_BLOCKS) + n_block + + if bx_i32 < bound: + if const_expr(_SW > 0): + _tile = _xcd(bx_i32) + else: + _tile = bx_i32 + _gemm1_body( + lds_raw_ptr, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_mind, + arg_aqout, + arg_ascaleout, + arg_hidden, + _tile, + lane, + wave, + use_nt, + i32_ntok, + total_m_blocks, + BM=BM, + BN=BN, + BK=BK, + inline_quant=inline_quant, + a_dtype=a_dtype, + K=_K, + N_OUT=_N_OUT, + NE=_NE, + interleave=interleave, + ) + + @flyc.jit + def launch_gemm1( + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_mind: fx.Int64, + i32_ntok: fx.Int32, + i32_grid: fx.Int32, + arg_aqout: fx.Int64, + arg_ascaleout: fx.Int64, + arg_hidden: fx.Int64, + stream: fx.Stream, + ): + grid_x = fx.Int64(i32_grid) + gemm1_kernel( + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_cumsum, + arg_mind, + i32_ntok, + arg_aqout, + arg_ascaleout, + arg_hidden, + ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) + + return launch_gemm1 diff --git a/kernels/moe/mxfp_moe/gemm2.py b/kernels/moe/mxfp_moe/gemm2.py new file mode 100644 index 000000000..20e4f7fc2 --- /dev/null +++ b/kernels/moe/mxfp_moe/gemm2.py @@ -0,0 +1,858 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from .mxfp4_gemm_common import ( + _e8m0_from_amax, + _fabs_f32, + _gep1, + _gep3, + _global_base_ptr1, + _global_i32_at, + _global_i32_buffer_tiles, + _inline_dpp_quad_amax, + _lds_ptr3, + _lds_swizzle_mask, + _raw, + _scale_mma_atoms, + _udiv, + _umod, + bq_bytes_for, + bscale_bytes_for, + k_half_for, + k_tiles_total_for, + kas_per_chunk_dw_for, + kbs_per_expert_dw_for, + kBS_stride_k0_dw, + kbs_stride_n0_dw_for, + kmchunks_for, + kStages, + kunroll_for, + lds_acc_bytes_for, + num_n_blocks_for, +) + +NUM_CU = 256 + + +def aq_bytes_for(max_m, k): + return max_m * k_half_for(k) + + +def saq_slot_bytes(BM, KH_TILE): + return BM * KH_TILE + + +def tiling(BM): + n_load_waves = min(4, BM // 8) + rows_per_wave = BM // n_load_waves + return n_load_waves, rows_per_wave, rows_per_wave // 8 + + +def _issue_a_load_lds(aq_dma_tiles4, s_aq_i32x4_tiles, slot, kt, car, lane, slot_bytes, lds_row, KH_TILE, k_half): + # A global->LDS async DMA (no register fragment), via BufferCopyLDS128b. Mirrors + # gemm1's issue_a_load_lds: the BufferCopyLDS atom's soffset is an element count. + lane_mod_8 = lane % fx.Int32(8) + mask = _lds_swizzle_mask(lds_row + (lane // fx.Int32(8))) + voffset = ((lane_mod_8 * fx.Int32(16)) ^ mask) + car * fx.Int32(k_half) + off_i32 = fx.Int32(slot * slot_bytes) + lds_row * fx.Int32(KH_TILE) + aq_dma_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), fx.Int32) + fx.copy( + aq_dma_atom, + fx.slice(aq_dma_tiles4, (None, voffset // fx.Int32(16))), + fx.slice(s_aq_i32x4_tiles, (None, off_i32 // fx.Int32(16))), + soffset=fx.Int32(kt * KH_TILE) // fx.Int32(4), + ) + + +def compile_gemm2_a4w4_port( + BM=32, + use_nt=False, + *, + NE, + N_OUT, + epilog="atomic", + D_INTER, + D_INTER_REAL=None, + BN=256, + BK=256, + xcd_swizzle=0, +): + assert BN == 256 and BK == 256, f"only BN==BK==256 supported, got BN={BN} BK={BK}" + KH_TILE = BK // 2 + _K = D_INTER + _K_REAL = D_INTER if D_INTER_REAL is None else D_INTER_REAL + assert _K % BK == 0, ( + f"D_INTER (gemm2 contraction K = inter_dim) must be a multiple of {BK}, " + f"got {_K}; inter_dim not divisible by {BK} (e.g. 384/192) is not " + f"supported by this BK={BK} kernel" + ) + assert ( + _K_REAL % 128 == 0 and 0 < _K_REAL <= _K + ), f"D_INTER_REAL={_K_REAL} must be a multiple of 128 and in (0, {_K}]" + _K_HALF = k_half_for(_K) + _K_TILES_TOTAL = k_tiles_total_for(_K, BK) + _persistent = epilog in ("nonatomic", "nonatomic_mxfp4") + _slot_bytes = saq_slot_bytes(BM, KH_TILE) + _aStages = kStages if _K_TILES_TOTAL <= kStages else 3 + _acc_rows = min(BM, 64) if epilog == "nonatomic_cshuffle" else BM + _lds_bytes = ( + lds_acc_bytes_for(_acc_rows, BN) + _aStages * _slot_bytes if epilog != "nonatomic" else _aStages * _slot_bytes + ) + _num_n_blocks = num_n_blocks_for(N_OUT, BN) + _n_load_waves, _rows_per_wave, _kSubBlocks = tiling(BM) + _epi_tag = { + "atomic": "atomic", + "nonatomic": "nonatomic", + "nonatomic_mxfp4": "nonatomic_mxfp4", + "nonatomic_cshuffle": "nonatomic_cshuffle", + }[epilog] + _rtag = "" if _K_REAL == _K else f"r{_K_REAL}" + _tag = f"ne{NE}_h{N_OUT}_i{_K}{_rtag}_bm{BM}{'_nt' if use_nt else ''}_{_epi_tag}" + if xcd_swizzle > 0: + _tag += f"_xcd{xcd_swizzle}" + _name = f"gemm2_a4w4_port_{_tag}" + + @fx.struct + class SharedStorage: + raw: fx.Array[fx.Uint8, _lds_bytes, 16] + + @flyc.kernel(name=_name, known_block_size=[256, 1, 1]) + def gemm2_kernel( + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_stids: fx.Int64, + arg_sweights: fx.Int64, + i32_M: fx.Int32, + i32_max_m_blocks: fx.Int32, + arg_out: fx.Int64, + arg_out_scale: fx.Int64, + ): + tx = gpu.thread_id("x") + bx = gpu.block_id("x") + tx_i32 = fx.Int32(tx) + bx_i32 = fx.Int32(bx) + + lane = tx_i32 % fx.Int32(64) + wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) + + _aq_num_bytes = fx.Int64(i32_max_m_blocks) * fx.Int64(BM * _K_HALF) + aq_dma_tiles4 = _global_i32_buffer_tiles(arg_aq, _aq_num_bytes, 4) + lds_raw_ptr = fx.SharedAllocator().allocate(SharedStorage).peek().raw.ptr + # s_aq as flat i32, divided into 4-element (128-bit) tiles for the LDS DMA dst. + s_aq_i32_flat = fx.make_view( + fx.recast_iter(fx.Int32, lds_raw_ptr), + fx.make_layout(kStages * _slot_bytes // 4, 1), + ) + s_aq_i32x4_tiles = fx.logical_divide(s_aq_i32_flat, fx.make_layout(4, 1)) + + def _issue_all_a_loads(m_row0): + for slot in range_constexpr(kStages): + for sub in range_constexpr(_kSubBlocks): + lds_row = wave * fx.Int32(_rows_per_wave) + fx.Int32(sub * 8) + car = m_row0 + lds_row + (lane // fx.Int32(8)) + _issue_a_load_lds( + aq_dma_tiles4, + s_aq_i32x4_tiles, + slot, + slot, + car, + lane, + _slot_bytes, + lds_row, + KH_TILE=KH_TILE, + k_half=_K_HALF, + ) + + def _run_tile(tile_i32): + _gemm2_body( + lds_raw_ptr, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_stids, + arg_sweights, + i32_M, + i32_max_m_blocks, + arg_out, + arg_out_scale, + tile_i32, + lane, + wave, + BM, + use_nt, + NE, + N_OUT, + epilog, + D_INTER=_K, + D_INTER_REAL=_K_REAL, + aStages=_aStages, + BN=BN, + BK=BK, + KH_TILE=KH_TILE, + ) + + if const_expr(_persistent): + cumsum0 = _global_i32_at(arg_cumsum, fx.Int32(0)) + total_m_blocks = _udiv(cumsum0, BM) + bound = total_m_blocks * fx.Int32(_num_n_blocks) + grid_nb = fx.Int32(gpu.grid_dim.x) + + _NXCD = 8 + _xq = _udiv(bound, _NXCD) + _xr = _umod(bound, _NXCD) + _SW = xcd_swizzle + + def _xcd(pid): + xc = _umod(pid, _NXCD) + wgid = xc * _xq + fx.Int32(arith.minsi(_raw(xc), _raw(_xr))) + _udiv(pid, _NXCD) + if const_expr(_SW <= 0): + return wgid + _ng = fx.Int32(_SW * _num_n_blocks) + group_id = wgid // _ng + first_pid_m = group_id * fx.Int32(_SW) + remaining_m = total_m_blocks - first_pid_m + group_size_m = fx.Int32(arith.minsi(_raw(remaining_m), _raw(fx.Int32(_SW)))) + wig = wgid % _ng + m_block = first_pid_m + (wig % group_size_m) + n_block = wig // group_size_m + return m_block * fx.Int32(_num_n_blocks) + n_block + + if bx_i32 < bound: + tile = _xcd(bx_i32) + _issue_all_a_loads(_udiv(tile, _num_n_blocks) * fx.Int32(BM)) + rocdl.sched_barrier(0) + _run_tile(tile) + + for iv in range(bx_i32 + grid_nb, bound, gpu.grid_dim.x): + wu = fx.Int32(iv) + gpu.barrier() + tile = _xcd(wu) + _issue_all_a_loads(_udiv(tile, _num_n_blocks) * fx.Int32(BM)) + _run_tile(tile) + else: + cumsum0 = _global_i32_at(arg_cumsum, fx.Int32(0)) + total_m_blocks = _udiv(cumsum0, BM) + bound = total_m_blocks * fx.Int32(_num_n_blocks) + + # Non-persistent atomic path is HBM-bandwidth-bound (down-proj reads the + # full fp4 weight column-block per tile, ~4% L2 reuse). A plain m-major + # linear grid clusters consecutive tiles onto the same XCD/HBM channels; + # round-robin the launch index across the 8 XCDs (bijective over [0,bound)) + # to balance channel utilization. Optional group swizzle (xcd_swizzle>0) + # further improves per-XCD L2 locality along M. + _NXCD = 8 + _xq = _udiv(bound, _NXCD) + _xr = _umod(bound, _NXCD) + _SW = xcd_swizzle + + def _xcd_np(pid): + xc = _umod(pid, _NXCD) + wgid = xc * _xq + fx.Int32(arith.minsi(_raw(xc), _raw(_xr))) + _udiv(pid, _NXCD) + if const_expr(_SW <= 0): + return wgid + _ng = fx.Int32(_SW * _num_n_blocks) + group_id = wgid // _ng + first_pid_m = group_id * fx.Int32(_SW) + remaining_m = total_m_blocks - first_pid_m + group_size_m = fx.Int32(arith.minsi(_raw(remaining_m), _raw(fx.Int32(_SW)))) + wig = wgid % _ng + m_block = first_pid_m + (wig % group_size_m) + n_block = wig // group_size_m + return m_block * fx.Int32(_num_n_blocks) + n_block + + if bx_i32 < bound: + tile = _xcd_np(bx_i32) + m_row0 = _udiv(tile, _num_n_blocks) * fx.Int32(BM) + if const_expr(_n_load_waves < 4): + if wave < fx.Int32(_n_load_waves): + _issue_all_a_loads(m_row0) + else: + _issue_all_a_loads(m_row0) + rocdl.sched_barrier(0) + _run_tile(tile) + + @flyc.jit + def launch_gemm2( + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_stids: fx.Int64, + arg_sweights: fx.Int64, + i32_M: fx.Int32, + i32_max_m_blocks: fx.Int32, + arg_out: fx.Int64, + arg_out_scale: fx.Int64, + stream: fx.Stream, + ): + if const_expr(_persistent): + tw = i32_max_m_blocks * fx.Int32(_num_n_blocks) + persist = _raw(tw > fx.Int32(NUM_CU * 4)) + grid_i32 = arith.select(persist, _raw(fx.Int32(NUM_CU)), _raw(tw)) + grid_x = arith.index_cast(T.index, grid_i32) + else: + grid_x = arith.index_cast(T.index, i32_max_m_blocks) * fx.Index(_num_n_blocks) + gemm2_kernel( + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_cumsum, + arg_stids, + arg_sweights, + i32_M, + i32_max_m_blocks, + arg_out, + arg_out_scale, + ).launch(grid=(grid_x, 1, 1), block=(256, 1, 1), stream=stream) + + if BM == 16: + launch_gemm2.compile_hints["llvm_options"] = {"enable-post-misched": False} + + return launch_gemm2 + + +@flyc.jit +def _gemm2_body( + lds_raw_ptr, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_stids, + arg_sweights, + i32_M, + i32_max_m_blocks, + arg_out, + arg_out_scale, + bx_i32, + lane, + wave, + BM, + use_nt, + NE, + N_OUT, + epilog, + *, + D_INTER, + D_INTER_REAL=None, + aStages=kStages, + BN, + BK, + KH_TILE, +): + _aStages = aStages + _kMChunks = kmchunks_for(BM) + _slot_bytes = saq_slot_bytes(BM, KH_TILE) + _K = D_INTER + _K_HALF = k_half_for(_K) + _K_TILES_TOTAL = k_tiles_total_for(_K, BK) + _K_REAL = D_INTER if D_INTER_REAL is None else D_INTER_REAL + _n_real_half = (_K_REAL + 127) // 128 + _kUnroll = kunroll_for(_K, BK) + _kAS_per_chunk_dw = kas_per_chunk_dw_for(_K) + _kBS_stride_n0_dw = kbs_stride_n0_dw_for(_K) + _asc_chunk_div = 16 if const_expr(BM == 16) else 32 + _asc_per_mb = (BM // _asc_chunk_div) * _kAS_per_chunk_dw * 4 + _bq_bytes = bq_bytes_for(NE, N_OUT, _K) + _bscale_bytes = bscale_bytes_for(NE, N_OUT, _K) + _kbs_per_expert_dw = kbs_per_expert_dw_for(N_OUT, _K) + _num_n_blocks = num_n_blocks_for(N_OUT, BN) + _n_load_waves, _rows_per_wave, _kSubBlocks = tiling(BM) + b_aux = 2 if use_nt else 0 + + m_block_idx = _udiv(bx_i32, _num_n_blocks) + n_block_idx = bx_i32 - m_block_idx * fx.Int32(_num_n_blocks) + e = rocdl.readfirstlane(T.i32, _raw(_global_i32_at(arg_eids, m_block_idx))) + m_row = m_block_idx * fx.Int32(BM) + + _asc_num_bytes = fx.Int64(i32_max_m_blocks) * fx.Int64(_asc_per_mb) + ascale_tiles = _global_i32_buffer_tiles(arg_ascale, _asc_num_bytes, 1) + bq_tiles = _global_i32_buffer_tiles(arg_bq, _bq_bytes, 4) + bscale_tiles = _global_i32_buffer_tiles(arg_bscale, _bscale_bytes, 1) + ascale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + bscale_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Int32) + bq_copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(b_aux), fx.Int32) + bq_reg_lay = fx.make_layout(4, 1) + scalar_reg_lay = fx.make_layout(1, 1) + + # Sequential LDS layout: saq bytes at offset 0, f32 accumulator after them. + saq_base_i32 = fx.Int32(fx.ptrtoint(lds_raw_ptr)) + lds_acc_base_i32 = saq_base_i32 + fx.Int32(_aStages * _slot_bytes) + + # A global->LDS DMA source (buffer tensor) + s_aq LDS dst tiles (flat i32, 128-bit). + _aq_num_bytes = fx.Int64(i32_max_m_blocks) * fx.Int64(BM * _K_HALF) + aq_dma_tiles4 = _global_i32_buffer_tiles(arg_aq, _aq_num_bytes, 4) + s_aq_i32_flat = fx.make_view( + fx.recast_iter(fx.Int32, lds_raw_ptr), + fx.make_layout(_aStages * _slot_bytes // 4, 1), + ) + s_aq_i32x4_tiles = fx.logical_divide(s_aq_i32_flat, fx.make_layout(4, 1)) + lds_a_read_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + lds_a_read_lay = fx.make_layout(4, 1) + + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + + b_load_s_base = [] + for j in range_constexpr(4): + v = (e * fx.Int32(N_OUT) + n_block_idx * fx.Int32(BN) + wave * fx.Int32(BN // 4) + fx.Int32(j * 16)) * fx.Int32( + _K_HALF + ) + b_load_s_base.append(rocdl.readfirstlane(T.i32, v)) + + mni_base = n_block_idx * fx.Int32(BN // 16 // 2) + wave * fx.Int32(BN // 64 // 2) + b_scale_s_base = [] + for mw in range_constexpr(2): + v = (e * fx.Int32(_kbs_per_expert_dw) + (mni_base + fx.Int32(mw)) * fx.Int32(_kBS_stride_n0_dw)) * fx.Int32(4) + b_scale_s_base.append(rocdl.readfirstlane(T.i32, v)) + + chunk_base = m_row // fx.Int32(16 if const_expr(BM == 16) else 32) + a_scale_s_base = [ + rocdl.readfirstlane( + T.i32, + (chunk_base + fx.Int32(sub)) * fx.Int32(_kAS_per_chunk_dw) * fx.Int32(4), + ) + for sub in range_constexpr(_kSubBlocks) + ] + + v_voff_scale = ((lane_div_16 * fx.Int32(16)) + lane_mod_16) * fx.Int32(4) + + def load_a_scale_tile(kt): + out = [None] * _kSubBlocks + for sub in range_constexpr(_kSubBlocks): + idx = (v_voff_scale + fx.Int32(kt * 256)) // fx.Int32(4) + r = fx.make_rmem_tensor(scalar_reg_lay, fx.Int32) + fx.copy( + ascale_copy_atom, + fx.slice(ascale_tiles, (None, idx)), + r, + soffset=a_scale_s_base[sub] // fx.Int32(4), + ) + out[sub] = r.load()[0] + return out + + def load_b_scale_tile(kt): + imm = kt * (kBS_stride_k0_dw * 4) + out = [None, None] + for mw in range_constexpr(2): + idx = (v_voff_scale + fx.Int32(imm)) // fx.Int32(4) + r = fx.make_rmem_tensor(scalar_reg_lay, fx.Int32) + fx.copy( + bscale_copy_atom, + fx.slice(bscale_tiles, (None, idx)), + r, + soffset=b_scale_s_base[mw] // fx.Int32(4), + ) + out[mw] = r.load()[0] + return out + + def load_b_tile(kt): + v_voff_b = (lane_div_16 * fx.Int32(256)) + (lane_mod_16 * fx.Int32(16)) + fx.Int32(kt * 2048) + out = [[None, None] for _ in range(4)] + for j in range_constexpr(4): + for half in range_constexpr(2): + if const_expr(kt * 2 + half >= _n_real_half): + continue + idx = (v_voff_b + fx.Int32(half * 1024)) // fx.Int32(16) + r = fx.make_rmem_tensor(bq_reg_lay, fx.Int32) + fx.copy( + bq_copy_atom, + fx.slice(bq_tiles, (None, idx)), + r, + soffset=b_load_s_base[j] // fx.Int32(4), + ) + out[j][half] = r + return out + + def issue_a_load_lds(slot, kt): + for sub in range_constexpr(_kSubBlocks): + lds_row = wave * fx.Int32(_rows_per_wave) + fx.Int32(sub * 8) + car = m_row + lds_row + (lane // fx.Int32(8)) + _issue_a_load_lds( + aq_dma_tiles4, + s_aq_i32x4_tiles, + slot, + kt, + car, + lane, + _slot_bytes, + lds_row, + KH_TILE=KH_TILE, + k_half=_K_HALF, + ) + + def issue_a_ds_read(slot): + lane_row = lane_mod_16 + lane_col = lane_div_16 * fx.Int32(16) + mask = _lds_swizzle_mask(lane_row) + a = [[None, None] for _ in range(_kMChunks)] + for k in range_constexpr(2): + lds_col = (lane_col + fx.Int32(k * 64)) ^ mask + for i in range_constexpr(_kMChunks): + lds_row = lane_row + fx.Int32(i * 16) + byte_off = fx.Int32(slot * _slot_bytes) + lds_row * fx.Int32(KH_TILE) + lds_col + r = fx.make_rmem_tensor(lds_a_read_lay, fx.Int32) + fx.copy_atom_call(lds_a_read_atom, fx.slice(s_aq_i32x4_tiles, (None, byte_off // fx.Int32(16))), r) + a[i][k] = r + return a + + zero4 = Vec.filled(4, 0.0, fx.Float32) + # Scaled down-proj MMA via fx.gemm + CDNA4 MFMA_Scale atoms (fp4 x fp4, e8m0 + # scales). opsel_a/opsel_b select the active 128-K half of the shared operand; + # scale_a/scale_b carry the e8m0 words. Perf-neutral vs the raw + # mfma_scale_f32_16x16x128_f8f6f4 intrinsic on gfx950. + scale_atoms = _scale_mma_atoms("fp4") + # accm[i][J] holds the running f32[4] accumulator as an rmem tensor. + accm = [[None, None, None, None] for _ in range(_kMChunks)] + + def _mma(atom, cf, a_frag, b_frag, sa, sb): + fx.gemm(atom, cf, a_frag, b_frag, cf, scale_a=sa, scale_b=sb) + + def mfma_cluster(b_tile, a, a_scale_sub, b_scale_slot, init, kt=0): + _skip_h1 = (kt * 2 + 1) >= _n_real_half + for J in range_constexpr(4): + mni = J // 2 + in_b = J % 2 + sb = b_scale_slot[mni] + b_J0 = b_tile[J][0] + b_J1 = None if const_expr(_skip_h1) else b_tile[J][1] + for sub in range_constexpr(_kSubBlocks): + sa = a_scale_sub[sub] + i0 = sub * 2 + i1 = sub * 2 + 1 + if const_expr(init): + accm[i0][J] = fx.make_rmem_tensor(4, fx.Float32) + accm[i0][J].store(zero4) + if const_expr(_kMChunks > 1): + accm[i1][J] = fx.make_rmem_tensor(4, fx.Float32) + accm[i1][J].store(zero4) + _mma(scale_atoms[(0, 0 + in_b)], accm[i0][J], a[i0][0], b_J0, sa, sb) + if const_expr(_kMChunks > 1): + _mma(scale_atoms[(1, 0 + in_b)], accm[i1][J], a[i1][0], b_J0, sa, sb) + if const_expr(not _skip_h1): + _mma(scale_atoms[(2, 2 + in_b)], accm[i0][J], a[i0][1], b_J1, sa, sb) + if const_expr(_kMChunks > 1): + _mma(scale_atoms[(3, 2 + in_b)], accm[i1][J], a[i1][1], b_J1, sa, sb) + + def _kloop_fence(): + gpu.barrier() + + if const_expr(_K_TILES_TOTAL <= kStages): + a_scale_v = [load_a_scale_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + b_scale_v = [load_b_scale_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + b = [load_b_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + for S in range_constexpr(_K_TILES_TOTAL): + kt = S + slot = kt % kStages + _kloop_fence() + a = issue_a_ds_read(slot) + a_scale_sub = [a_scale_v[kt][sub] for sub in range_constexpr(_kSubBlocks)] + mfma_cluster(b[slot], a, a_scale_sub, b_scale_v[slot], init=(S == 0), kt=kt) + else: + a_scale_v = [load_a_scale_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + b_scale_v = [load_b_scale_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + # Software-pipeline the B tiles instead of preloading all _K_TILES_TOTAL of + # them: preloading all K tiles keeps every B fragment live across the whole + # k-loop (>=384 VGPR for K=3072/BK=256), which forces the f32 accumulators + # into AGPRs and drops occupancy to 1 wave/SIMD. Keeping only _bPF B tiles + # resident (one-ahead prefetch) lets the accumulators stay in ArchVGPRs and + # restores 2 waves/SIMD. A stays LDS-double-buffered as before. + _bPF = 2 + b_pf = [load_b_tile(kt) for kt in range_constexpr(_bPF)] + + for OFFSET in range_constexpr(_kUnroll): + kt = OFFSET + slot = kt % _aStages + next_kt = kStages + OFFSET + write_slot = next_kt % _aStages + _kloop_fence() + a = issue_a_ds_read(slot) + issue_a_load_lds(write_slot, next_kt) + # Prefetch the B tile _bPF iterations ahead so at most _bPF B tiles are + # live at once. + b_cur = b_pf[kt % _bPF] + b_next_kt = kt + _bPF + if const_expr(b_next_kt < _K_TILES_TOTAL): + b_pf[kt % _bPF] = load_b_tile(b_next_kt) + a_scale_sub = [a_scale_v[kt][sub] for sub in range_constexpr(_kSubBlocks)] + mfma_cluster(b_cur, a, a_scale_sub, b_scale_v[kt], init=(OFFSET == 0)) + + for S in range_constexpr(kStages): + kt = _K_TILES_TOTAL - kStages + S + slot = kt % _aStages + _kloop_fence() + a = issue_a_ds_read(slot) + b_cur = b_pf[kt % _bPF] + a_scale_sub = [a_scale_v[kt][sub] for sub in range_constexpr(_kSubBlocks)] + mfma_cluster(b_cur, a, a_scale_sub, b_scale_v[kt], init=False) + + # Materialize the f32[4] accumulators as raw vector values for the (raw) epilogs. + accm = [[accm[i][J].load().ir_value() for J in range(4)] for i in range(_kMChunks)] + + if epilog == "nonatomic": + out_base = _global_base_ptr1(arg_out) + _flat_bf16_epilog(accm, out_base, m_row, n_block_idx, wave, lane, N_OUT, BN, _kMChunks) + elif epilog == "nonatomic_cshuffle": + _cshuffle_flat_bf16_epilog( + lds_acc_base_i32, + accm, + arg_out, + m_row, + n_block_idx, + wave, + lane, + BM, + N_OUT, + BN, + ) + elif epilog == "nonatomic_mxfp4": + out_q_base = _global_base_ptr1(arg_out) + out_scale_base = _global_base_ptr1(arg_out_scale) + tid_i32 = fx.Int32(gpu.thread_id("x")) + _flat_mxfp4_epilog( + accm, + out_q_base, + out_scale_base, + m_row, + n_block_idx, + wave, + lane, + tid_i32, + N_OUT, + BN, + lds_acc_base_i32, + _kMChunks, + ) + else: + _atomic_bf16_epilog( + lds_acc_base_i32, + accm, + arg_out, + arg_stids, + arg_sweights, + m_row, + n_block_idx, + wave, + lane, + i32_M, + BM, + N_OUT, + BN, + ) + + +def _flat_bf16_epilog(accm, out_base, m_row, n_block_idx, wave, lane, N_OUT, BN, kMChunks): + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + row_base = m_row + lane_div_16 * fx.Int32(4) + gn_base = n_block_idx * fx.Int32(BN) + wave * fx.Int32(BN // 4) + lane_mod_16 + byte_base = (fx.Int64(row_base) * fx.Int64(N_OUT) + fx.Int64(gn_base)) * fx.Int64(2) + for i in range_constexpr(kMChunks): + for J in range_constexpr(4): + vec = Vec(accm[i][J]) + for v in range_constexpr(4): + const_off = ((i * 16 + v) * N_OUT + J * 16) * 2 + bf = Vec.from_elements([vec[v]], fx.Float32).to(fx.BFloat16) + llvm.StoreOp(_raw(bf), _gep1(out_base, byte_base + fx.Int64(const_off))) + + +def _cshuffle_flat_bf16_epilog(lds_acc_base_i32, accm, arg_out, m_row, n_block_idx, wave, lane, BM, N_OUT, BN): + _iC = BM // 16 + _REPS = BM // 8 + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + lds_base = _lds_ptr3(lds_acc_base_i32, fx.Int32(0)) + tx_i32 = fx.Int32(gpu.thread_id("x")) + m_lane = tx_i32 // fx.Int32(32) + n_lane = tx_i32 % fx.Int32(32) + col_start = n_lane * fx.Int32(2) + out_base = _global_base_ptr1(arg_out) + + for i in range_constexpr(_iC): + row_base = fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + for J in range_constexpr(4): + col = wave * fx.Int32(64) + fx.Int32(J * 16) + lane_mod_16 + bf4 = Vec(accm[i][J]).to(fx.BFloat16) + for v in range_constexpr(4): + idx = (row_base + fx.Int32(v)) * fx.Int32(BN) + col + llvm.StoreOp(_raw(bf4[v]), _gep3(lds_base, idx * fx.Int32(2))) + gpu.barrier() + for mr in range_constexpr(_REPS): + row_local = fx.Int32(mr * 8) + m_lane + sorted_row = m_row + row_local + for s in range_constexpr(4): + idx0 = row_local * fx.Int32(BN) + col_start + fx.Int32(s * 64) + pk = Vec(llvm.load(T.vec(2, T.bf16), _gep3(lds_base, idx0 * fx.Int32(2)))) + n_col = n_block_idx * fx.Int32(BN) + col_start + fx.Int32(s * 64) + elem = fx.Int64(sorted_row) * fx.Int64(N_OUT) + fx.Int64(n_col) + llvm.StoreOp(_raw(pk), _gep1(out_base, elem * fx.Int64(2))) + + +@flyc.jit +def _flat_mxfp4_epilog( + accm, + out_q_base, + out_scale_base, + m_row, + n_block_idx, + wave, + lane, + tid_i32, + N_OUT, + BN, + lds_acc_base_i32, + kMChunks, +): + lds_base = _lds_ptr3(lds_acc_base_i32, fx.Int32(0)) + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + for i in range_constexpr(kMChunks): + row_base = fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + for J in range_constexpr(4): + col = wave * fx.Int32(BN // 4) + fx.Int32(J * 16) + lane_mod_16 + vec = Vec(accm[i][J]) + for v in range_constexpr(4): + idx = (row_base + fx.Int32(v)) * fx.Int32(BN) + col + llvm.StoreOp(_raw(vec[v]), _gep3(lds_base, idx * fx.Int32(4))) + gpu.barrier() + + NBLK = BN // 32 + m_lane = tid_i32 // fx.Int32(16) + n_lane = tid_i32 % fx.Int32(16) + wave_grp = n_lane // fx.Int32(4) + kk = n_lane % fx.Int32(4) + _m_base = m_row + m_lane + _q_row0 = fx.Int64(_m_base) * fx.Int64(N_OUT // 2) + _s_row0 = fx.Int64(_m_base) * fx.Int64(N_OUT // 32) + _blocks = [(mr, half) for mr in range(kMChunks) for half in range(NBLK // 4)] + + def _issue_load(mr, half): + row_local = fx.Int32(mr * 16) + m_lane + group = wave_grp + fx.Int32(half * 4) + col0 = group * fx.Int32(32) + kk * fx.Int32(8) + base_idx = row_local * fx.Int32(BN) + col0 + v0 = Vec(llvm.load(T.vec(4, T.f32), _gep3(lds_base, base_idx * fx.Int32(4)))) + v1 = Vec( + llvm.load( + T.vec(4, T.f32), + _gep3(lds_base, (base_idx + fx.Int32(4)) * fx.Int32(4)), + ) + ) + return [v0[0], v0[1], v0[2], v0[3], v1[0], v1[1], v1[2], v1[3]], group, col0 + + _r_next, _grp_next, _col0_next = _issue_load(*_blocks[0]) + for _bi in range_constexpr(len(_blocks)): + mr, half = _blocks[_bi] + r, group, col0 = _r_next, _grp_next, _col0_next + if _bi + 1 < len(_blocks): + _r_next, _grp_next, _col0_next = _issue_load(*_blocks[_bi + 1]) + if True: + amax_f = _raw(_fabs_f32(r[0])) + for e in range_constexpr(1, 8): + abs_e = _raw(_fabs_f32(r[e])) + amax_f = arith.maxnumf(amax_f, abs_e) + amax = arith.shrui(arith.bitcast(T.i32, amax_f), _raw(fx.Int32(16))) + amax_dpp = _raw(_inline_dpp_quad_amax(amax)) + f32b = arith.shli(amax_dpp, _raw(fx.Int32(16))) + e8m0, qscale_f = _e8m0_from_amax(fx.Float32(arith.bitcast(T.f32, f32b))) + e8 = _raw(e8m0) + qscale = _raw(qscale_f) + packed = _raw(fx.Int32(0)) + packed = rocdl.cvt_scalef32_pk_fp4_f32(T.i32, packed, _raw(r[0]), _raw(r[1]), qscale, 0) + packed = rocdl.cvt_scalef32_pk_fp4_f32(T.i32, packed, _raw(r[2]), _raw(r[3]), qscale, 1) + packed = rocdl.cvt_scalef32_pk_fp4_f32(T.i32, packed, _raw(r[4]), _raw(r[5]), qscale, 2) + packed = rocdl.cvt_scalef32_pk_fp4_f32(T.i32, packed, _raw(r[6]), _raw(r[7]), qscale, 3) + global_col = n_block_idx * fx.Int32(BN) + col0 + blk = n_block_idx * fx.Int32(NBLK) + group + q_byte = _q_row0 + fx.Int64(mr * 16 * (N_OUT // 2)) + fx.Int64(global_col // fx.Int32(2)) + s_byte = _s_row0 + fx.Int64(mr * 16 * (N_OUT // 32)) + fx.Int64(blk) + llvm.StoreOp(packed, _gep1(out_q_base, q_byte), nontemporal=True) + if kk == fx.Int32(0): + llvm.StoreOp(arith.trunci(T.i8, e8), _gep1(out_scale_base, s_byte)) + + +@flyc.jit +def _atomic_bf16_epilog( + lds_acc_base_i32, + accm, + arg_out, + arg_stids, + arg_sweights, + m_row, + n_block_idx, + wave, + lane, + i32_M, + BM, + N_OUT, + BN, +): + _kMChunks = kmchunks_for(BM) + M_REPS = BM // 8 + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + lds_base = _lds_ptr3(lds_acc_base_i32, fx.Int32(0)) + + tx_i32 = fx.Int32(gpu.thread_id("x")) + m_lane = tx_i32 // fx.Int32(32) + n_lane = tx_i32 % fx.Int32(32) + col_start = n_lane * fx.Int32(2) + stids_base = _global_base_ptr1(arg_stids) + sweights_base = _global_base_ptr1(arg_sweights) + out_base = _global_base_ptr1(arg_out) + + packed = [] + weight = [] + for mr in range_constexpr(M_REPS): + sorted_pos = m_row + fx.Int32(mr * 8) + m_lane + packed.append(llvm.load(T.i32, _gep1(stids_base, sorted_pos * fx.Int32(4)), invariant=True)) + weight.append(llvm.load(T.f32, _gep1(sweights_base, sorted_pos * fx.Int32(4)), invariant=True)) + + for i in range_constexpr(_kMChunks): + row_base = fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + for J in range_constexpr(4): + col = wave * fx.Int32(64) + fx.Int32(J * 16) + lane_mod_16 + vec = Vec(accm[i][J]) + for v in range_constexpr(4): + idx = (row_base + fx.Int32(v)) * fx.Int32(BN) + col + llvm.StoreOp(_raw(vec[v]), _gep3(lds_base, idx * fx.Int32(4))) + + gpu.barrier() + + for mr in range_constexpr(M_REPS): + row_in_block = fx.Int32(mr * 8) + m_lane + token_id = packed[mr] & fx.Int32(0x00FFFFFF) + if token_id < i32_M: + row_base_addr = token_id * fx.Int32(N_OUT) + n_block_idx * fx.Int32(BN) + col_start + for s in range_constexpr(4): + idx0 = row_in_block * fx.Int32(BN) + col_start + fx.Int32(s * 64) + v2 = Vec(llvm.load(T.vec(2, T.f32), _gep3(lds_base, idx0 * fx.Int32(4)))) + pk = Vec.from_elements([v2[0] * weight[mr], v2[1] * weight[mr]], fx.Float32).to(fx.BFloat16) + off = (row_base_addr + fx.Int32(s * 64)) * fx.Int32(2) + out_ptr = _gep1(out_base, off) + llvm.AtomicRMWOp( + llvm.AtomicBinOp.fadd, + out_ptr, + _raw(pk), + llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ) diff --git a/kernels/moe/mxfp_moe/host.py b/kernels/moe/mxfp_moe/host.py new file mode 100644 index 000000000..e884cedfe --- /dev/null +++ b/kernels/moe/mxfp_moe/host.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Host-side launch glue for the fused a4w4 mxfp_moe kernels. + +Ported (self-contained) from aiter's ``mxfp4_gemm1_kernels.py`` / +``mxfp4_gemm2_kernels.py``. Kernel launch args are raw device pointers +(``fx.Int64``); tensors are passed as ``.data_ptr()``. +""" + +import functools + +import torch + +from kernels.common.tensor_shim import _run_compiled +from kernels.moe.mxfp_moe.gemm1 import compile_gemm1_a4w4_port, gemm1_grid +from kernels.moe.mxfp_moe.gemm2 import compile_gemm2_a4w4_port + +# gemm1 (BM, use_nt, inline_quant, a_dtype) variants the kernel supports. +# a_dtype="fp4" is a4w4 (mxfp4 A); "fp8" is a8w4 (fp8 e4m3 A x mxfp4 W1). +_G1_SUPPORTED = { + (32, True, False, "fp4"), + (32, False, False, "fp4"), + (64, False, False, "fp4"), + (128, False, False, "fp4"), + (16, True, True, "fp4"), + (32, True, False, "fp8"), + (32, False, False, "fp8"), + (64, False, False, "fp8"), + (128, False, False, "fp8"), + (16, True, True, "fp8"), +} + +# gemm2 (BM, use_nt, epilog) variants the kernel supports. +_G2_SUPPORTED = { + (16, False, "atomic"), + (16, True, "atomic"), + (32, False, "atomic"), + (32, True, "atomic"), + (64, False, "atomic"), + (64, True, "atomic"), + (128, False, "nonatomic"), + (128, False, "nonatomic_mxfp4"), + (32, False, "nonatomic_cshuffle"), + (64, False, "nonatomic_cshuffle"), + (128, False, "nonatomic_cshuffle"), +} + + +@functools.cache +def _get_compiled_gemm1( + BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk, BN, BK, interleave, xcd_swizzle, a_dtype +): + return compile_gemm1_a4w4_port( + BM, + use_nt, + inline_quant, + D_HIDDEN=D_HIDDEN, + D_INTER=D_INTER, + NE=NE, + TOPK=topk, + BN=BN, + BK=BK, + interleave=interleave, + xcd_swizzle=xcd_swizzle, + a_dtype=a_dtype, + ) + + +@functools.cache +def _get_compiled_gemm2(BM, use_nt, NE, N_OUT, epilog, D_INTER, D_INTER_REAL, BN, BK, xcd_swizzle): + return compile_gemm2_a4w4_port( + BM=BM, + use_nt=use_nt, + NE=NE, + N_OUT=N_OUT, + epilog=epilog, + D_INTER=D_INTER, + D_INTER_REAL=D_INTER_REAL, + BN=BN, + BK=BK, + xcd_swizzle=xcd_swizzle, + ) + + +def flydsl_mxfp4_gemm1( + *, + a_quant, + a_scale_sorted_shuffled, + w1_u8, + w1_scale_u8, + sorted_expert_ids, + cumsum_tensor, + m_indices, + inter_sorted_quant, + inter_sorted_shuffled_scale, + hidden_states, + n_tokens, + BM, + use_nt, + inline_quant, + NE, + D_HIDDEN, + D_INTER, + topk, + BN=256, + BK=256, + interleave=False, + xcd_swizzle=0, + a_dtype="fp4", + stream=None, +): + """Fused stage1: gate+up GEMM + SiLU + fp4 re-quant. + + ``a_dtype`` selects the activation format: "fp4" (a4w4, mxfp4 A) or "fp8" + (a8w4, fp8 e4m3 A x mxfp4 W1). Writes the sorted fp4 intermediate into + ``inter_sorted_quant`` / ``inter_sorted_shuffled_scale`` (both pre-allocated). + """ + if D_HIDDEN % BK != 0: + raise NotImplementedError(f"mxfp_moe gemm1 requires D_HIDDEN (K) % {BK} == 0, got H={D_HIDDEN}") + if (2 * D_INTER) % BN != 0: + raise NotImplementedError(f"mxfp_moe gemm1 requires 2*D_INTER (N_OUT) % {BN} == 0, got D_INTER={D_INTER}") + + # Non-temporal (streaming) B loads help at small M -- there is no cross-tile + # weight reuse, so streaming avoids polluting L2. But once M is large enough + # that each expert owns more than one M-block, consecutive M-tiles of the + # same expert reuse the same W1 columns: streaming then discards reusable B + # and the kernel becomes HBM-read-bound (L2 hit ~34% vs ~59%). Switch to the + # cached (non-nt) variant when total M-blocks >= experts (avg >= 1 padded + # block/expert), which measured ~22% faster at M_eff=8192 while the small-M + # streaming win is preserved below the crossover. Only auto-relax the BM==32 + # streaming default; never force streaming on when the caller disabled it. + if use_nt and not inline_quant and BM == 32: + total_m_blocks = (int(n_tokens) * int(topk) + BM - 1) // BM + if total_m_blocks >= int(NE): + use_nt = False + + if (BM, use_nt, inline_quant, a_dtype) not in _G1_SUPPORTED: + raise NotImplementedError( + f"mxfp_moe gemm1 unsupported variant " + f"(BM={BM}, use_nt={use_nt}, inline_quant={inline_quant}, a_dtype={a_dtype})" + ) + + launch = _get_compiled_gemm1( + BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk, BN, BK, interleave, xcd_swizzle, a_dtype + ) + grid = gemm1_grid(n_tokens, BM, NE=NE, TOPK=topk, INTER=D_INTER, BN=BN) + _run_compiled( + launch, + a_quant.data_ptr(), + a_scale_sorted_shuffled.data_ptr(), + w1_u8.data_ptr(), + w1_scale_u8.data_ptr(), + sorted_expert_ids.data_ptr(), + cumsum_tensor.data_ptr(), + m_indices.data_ptr(), + int(n_tokens), + int(grid), + inter_sorted_quant.data_ptr(), + inter_sorted_shuffled_scale.data_ptr(), + hidden_states.data_ptr(), + torch.cuda.current_stream() if stream is None else stream, + ) + return inter_sorted_quant, inter_sorted_shuffled_scale + + +def flydsl_mxfp4_gemm2( + *, + inter_sorted_quant, + inter_sorted_shuffled_scale, + w2_u8, + w2_scale_u8, + sorted_expert_ids, + cumsum_tensor, + sorted_token_ids, + sorted_weights, + flat_out, + M_logical, + max_sorted, + BM, + use_nt, + epilog, + NE, + D_HIDDEN, + D_INTER, + topk, + flat_out_scale=None, + D_INTER_REAL=None, + BN=256, + BK=256, + xcd_swizzle=0, + stream=None, +): + """Down-projection stage2. Consumes the stage1 sorted fp4 intermediate.""" + if D_INTER % BK != 0: + raise NotImplementedError( + f"mxfp_moe gemm2 contraction D_INTER (inter_dim) must be a multiple of " f"{BK}, got D_INTER={D_INTER}" + ) + if D_HIDDEN % BN != 0: + raise NotImplementedError(f"mxfp_moe gemm2 requires D_HIDDEN (N_OUT=model_dim) % {BN} == 0, got H={D_HIDDEN}") + if (BM, use_nt, epilog) not in _G2_SUPPORTED: + raise NotImplementedError(f"mxfp_moe gemm2 unsupported variant (BM={BM}, use_nt={use_nt}, epilog={epilog})") + + launch = _get_compiled_gemm2(BM, use_nt, NE, D_HIDDEN, epilog, D_INTER, D_INTER_REAL, BN, BK, xcd_swizzle) + max_m_blocks = (max_sorted + BM - 1) // BM + if flat_out_scale is None: + flat_out_scale = torch.empty(1, dtype=torch.uint8, device=flat_out.device) + + _run_compiled( + launch, + inter_sorted_quant.data_ptr(), + inter_sorted_shuffled_scale.data_ptr(), + w2_u8.data_ptr(), + w2_scale_u8.data_ptr(), + sorted_expert_ids.data_ptr(), + cumsum_tensor.data_ptr(), + sorted_token_ids.data_ptr(), + sorted_weights.data_ptr(), + int(M_logical), + int(max_m_blocks), + flat_out.data_ptr(), + flat_out_scale.data_ptr(), + torch.cuda.current_stream() if stream is None else stream, + ) + return flat_out diff --git a/kernels/moe/mxfp_moe/mxfp4_gemm_common.py b/kernels/moe/mxfp_moe/mxfp4_gemm_common.py new file mode 100644 index 000000000..82a2d9dac --- /dev/null +++ b/kernels/moe/mxfp_moe/mxfp4_gemm_common.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl._mlir.dialects import memref as memref_dialect +from flydsl.expr import arith, rocdl +from flydsl.expr.typing import Float4E2M1FN, Float8E4M3FN, T +from kernels.common import buffer_ops +from kernels.common.layout_utils import crd2idx + +from . import dpp_utils + +_PTR3 = "!llvm.ptr<3>" +kStages = 2 +kBS_stride_k0_dw = 64 +LOG2E = 1.4426950408889634 + + +def _raw(v): + if not isinstance(v, ir.Value) and hasattr(v, "ir_value"): + return v.ir_value() + return v + + +def _udiv(a, c): + cc = fx.Int32(c) if isinstance(c, int) else c + return fx.Int32(arith.divui(_raw(a), _raw(cc))) + + +def _umod(a, c): + cc = fx.Int32(c) if isinstance(c, int) else c + return fx.Int32(arith.remui(_raw(a), _raw(cc))) + + +# A elem selects the f8f6f4 cbsz for the scaled MFMA atom: fp4 (e2m1) -> cbsz 4, +# fp8 (e4m3) -> cbsz 0. B (blgp) is always mxfp4. +_A_ELEM = {"fp4": Float4E2M1FN, "fp8": Float8E4M3FN} + + +def _scale_mma_atoms(a_dtype): + """16 (opsel_a, opsel_b) scaled 16x16x128 MFMA atoms; A elem = fp4/fp8, B = fp4.""" + elem_a = _A_ELEM[a_dtype] + return { + (osa, osb): fx.make_mma_atom( + fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, elem_a, Float4E2M1FN, opsel_a=osa, opsel_b=osb) + ) + for osa in range(4) + for osb in range(4) + } + + +def _global_i32_buffer_view(addr_i64, num_bytes): + # fx.copy's BufferCopy/BufferCopyLDS atoms take soffset as an element count, not + # the bytes buffer_ops.buffer_load's soffset_bytes expected. + # make_layout's dynamic-shape leaf must be i32/i64, not fx.Index. + num_bytes_i64 = fx.Int64(num_bytes) + ptr_ty = fx.PointerType.get(T.i32, address_space=fx.AddressSpace.Global, alignment=4) + ptr = fx.inttoptr(ptr_ty, fx.Int64(addr_i64)) + view = fx.Tensor(fx.make_view(ptr, fx.make_layout(num_bytes_i64 // fx.Int64(4), 1))) + return fx.rocdl.make_buffer_tensor(view, max_size=False, num_records_bytes=num_bytes_i64) + + +def _global_i32_buffer_tiles(addr_i64, num_bytes, tile_elems): + return fx.logical_divide(_global_i32_buffer_view(addr_i64, num_bytes), fx.make_layout(tile_elems, 1)) + + +def _lds_ptr3(base_i32, byte_off_i32): + addr_i64 = fx.Int64(base_i32 + byte_off_i32) + return llvm.inttoptr(ir.Type.parse(_PTR3), _raw(addr_i64)) + + +def _lds_base_ptr3(lds_view): + base_i32 = fx.Int32(memref_dialect.extract_aligned_pointer_as_index(lds_view)) + return llvm.inttoptr(ir.Type.parse(_PTR3), _raw(fx.Int64(base_i32))) + + +def _gep3(base_ptr, byte_off_i32): + return buffer_ops.get_element_ptr(base_ptr, byte_offset=_raw(byte_off_i32), elem_type=T.i8) + + +def _global_base_ptr1(addr_i64): + return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(addr_i64))) + + +def _gep1(base_ptr, byte_off_i32): + return buffer_ops.get_element_ptr(base_ptr, byte_offset=_raw(byte_off_i32), elem_type=T.i8) + + +def _global_i32_ptr(addr_i64): + ptr_ty = fx.PointerType.get(T.i32, address_space=fx.AddressSpace.Global, alignment=4) + return fx.inttoptr(ptr_ty, fx.Int64(addr_i64)) + + +def _global_i32_at(addr_i64, idx): + # Plain scalar read: fx pointer index, no tiling/register-fragment machinery. + return _global_i32_ptr(addr_i64)[idx] + + +def _global_i32_load(tiles, idx): + # Atom/types must be built with an active MLIR trace context, not as globals. + atom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32) + r = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32) + fx.copy_atom_call(atom, fx.slice(tiles, (None, idx)), r) + return r.load()[0] + + +def _global_scalar_tiles(addr_i64, numeric_cls, num_elems): + ptr_ty = fx.PointerType.get( + numeric_cls.ir_type, + address_space=fx.AddressSpace.Global, + alignment=numeric_cls.width // 8, + ) + ptr = fx.inttoptr(ptr_ty, fx.Int64(addr_i64)) + flat = fx.make_view(ptr, fx.make_layout(num_elems, 1)) + return fx.logical_divide(flat, fx.make_layout(1, 1)) + + +def _scalar_store(tiles, idx, value, numeric_cls): + atom = fx.make_copy_atom(fx.UniversalCopy(numeric_cls.width), numeric_cls) + r = fx.make_rmem_tensor(fx.make_layout(1, 1), numeric_cls) + r.store(fx.Vector.from_elements([numeric_cls(value)], numeric_cls)) + fx.copy_atom_call(atom, r, fx.slice(tiles, (None, idx))) + + +def _layout_idx(layout, *coords): + idx_coords = [fx.Int64(c) for c in coords] + return fx.Int32(crd2idx(idx_coords, layout)) + + +def _buffer_rsrc(addr_i64, num_records_bytes): + return buffer_ops.create_buffer_resource_from_addr(_raw(fx.Int64(addr_i64)), num_records_bytes=num_records_bytes) + + +def _lds_swizzle_mask(row): + return (row & fx.Int32(14)) << fx.Int32(3) + + +def _fabs_f32(x): + return fx.Float32(llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(x)], [], [])) + + +def _e8m0_roundup(amax_f32): + wi = fx.Int32(_raw(amax_f32 * fx.Float32(1.0 / 6.0)).bitcast(T.i32)) + bexp = (wi + fx.Int32(0x7FFFFF)).shrui(fx.Int32(23)) & fx.Int32(0xFF) + lt = arith.cmpi(arith.CmpIPredicate.ult, _raw(bexp), _raw(fx.Int32(254))) + return fx.Int32(arith.select(lt, _raw(bexp), _raw(fx.Int32(254)))) + + +def _e8m0_from_amax(amax_f32): + e8m0 = _e8m0_roundup(amax_f32) + qscale = fx.Float32(_raw(e8m0 << fx.Int32(23)).bitcast(T.f32)) + return e8m0, qscale + + +def _inline_e8m0(amax_u16_i32): + f32 = fx.Float32(_raw((fx.Int32(_raw(amax_u16_i32)) & fx.Int32(0xFFFF)) << fx.Int32(16)).bitcast(T.f32)) + return _e8m0_roundup(f32) + + +def _pkmax_u16(a_i32, b_i32): + _v2i16 = T.vec(2, T.i16) + va = llvm.BitcastOp(_v2i16, _raw(a_i32)).result + vb = llvm.BitcastOp(_v2i16, _raw(b_i32)).result + vm = arith.MaxUIOp(va, vb).result + out = llvm.BitcastOp(T.i32, vm).result + return fx.Int32(out) + + +def _silu_mul_batch(gs, us): + e = [fx.Float32(rocdl.exp2(T.f32, _raw(g * fx.Float32(-LOG2E)))) for g in gs] + sig = [fx.Float32(rocdl.rcp(T.f32, _raw(fx.Float32(1.0) + ei))) for ei in e] + return [gs[i] * sig[i] * us[i] for i in range(len(gs))] + + +def _umax_i32(a, b): + is_gt = arith.cmpi(arith.CmpIPredicate.ugt, _raw(a), _raw(b)) + return fx.Int32(arith.select(is_gt, _raw(a), _raw(b))) + + +def _inline_dpp_quad_amax(a32): + a32 = fx.Int32(_raw(a32)) + s1 = fx.Int32(dpp_utils.update_dpp_i32(_raw(a32), _raw(a32), 0xB1, 0xF, 0xF, True)) + a32 = _umax_i32(a32, s1) + s2 = fx.Int32(dpp_utils.update_dpp_i32(_raw(a32), _raw(a32), 0x4E, 0xF, 0xF, True)) + return _umax_i32(a32, s2) + + +def k_half_for(k): + return k // 2 + + +def k_tiles_total_for(k, BK): + return k // BK + + +def kunroll_for(k, BK): + return k_tiles_total_for(k, BK) - kStages + + +def kas_c_k1_for(k): + return (k // 32) // 4 // 2 + + +def kbs_c_k1_for(k): + return (k // 32) // 4 // 2 + + +def kbs_stride_n0_dw_for(k): + return kbs_c_k1_for(k) * 64 + + +def kas_per_chunk_dw_for(k): + return kas_c_k1_for(k) * 64 + + +def num_n_blocks_for(n, BN): + return n // BN + + +def kbs_c_n1_for(n): + return n // 16 // 2 + + +def kbs_per_expert_dw_for(n, k): + return kbs_c_n1_for(n) * kbs_stride_n0_dw_for(k) + + +def bq_bytes_for(ne, n, k): + return ne * n * k_half_for(k) + + +def bscale_bytes_for(ne, n, k): + return ne * kbs_per_expert_dw_for(n, k) * 4 + + +def kmchunks_for(BM): + return BM // 16 + + +def lds_acc_bytes_for(rows, BN): + return rows * BN * 4 diff --git a/tests/arch_compat.py b/tests/arch_compat.py index 9361ca59d..46e610204 100644 --- a/tests/arch_compat.py +++ b/tests/arch_compat.py @@ -13,7 +13,6 @@ "test_preshuffle_gemm.py", "test_blockscale_preshuffle_gemm.py", "test_moe_gemm.py", - "test_moe_blockscale.py", "test_moe_reduce.py", "test_pa.py", "test_quant.py", diff --git a/tests/kernels/test_moe_blockscale.py b/tests/kernels/test_moe_blockscale.py deleted file mode 100644 index bbd9f7576..000000000 --- a/tests/kernels/test_moe_blockscale.py +++ /dev/null @@ -1,763 +0,0 @@ -#!/usr/bin/env python3 - -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2025 FlyDSL Project Contributors - -"""MoE Blockscale GEMM Test - flydsl 2-stage vs aiter fused kernel. - -Tests the MoE flow end-to-end: - Stage 1: X @ W1[expert].T -> SiLU(gate)*up -> fp16 intermediate - Stage 2: act @ W2[expert].T -> atomic accumulate to output - -Compares flydsl 2-stage (per-token-scale) vs aiter fused blockscale kernel. -Both do the same FP8 MFMA compute; difference is scale granularity. -""" - -import logging -import math -import os -import sys - -import pytest -import torch -import torch.nn.functional as F - -import flydsl.compiler as flyc - -pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] - -_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) -_PYTHON_CANDIDATES = [ - os.path.join(_REPO_ROOT, "build", "python_packages"), - _REPO_ROOT, -] -for _p in reversed(_PYTHON_CANDIDATES): - if os.path.isdir(_p) and _p not in sys.path: - sys.path.insert(0, _p) - -from flydsl.runtime.device import get_rocm_arch # noqa: E402 -from kernels.moe.moe_blockscale_2stage import compile_moe_blockscale_gemm1, compile_moe_blockscale_gemm2 # noqa: E402 -from tests.test_common import checkAllclose, run_perftest # noqa: E402 -from tests.utils import pertoken_quant, shuffle_weight # noqa: E402 - -logging.basicConfig(level=logging.INFO) -# Quiet aiter spam -logging.getLogger("aiter").setLevel(logging.ERROR) - -ARCH = get_rocm_arch() -DTYPE_FP8 = torch.float8_e4m3fn if "gfx95" in ARCH else torch.float8_e4m3fnuz - -try: - import aiter - from aiter.fused_moe import moe_sorting as aiter_moe_sorting - from aiter.ops.quant import per_group_quant_hip - - HAS_AITER = True -except Exception: - HAS_AITER = False - -# Use aiter or torch for routing -from tests.kernels.test_moe_gemm import ( # noqa: E402 - build_routing_buffers, -) - - -def _expand_blockscale(scale_flat, E, nblk_n, nblk_k, blk_n, blk_k): - """Expand [E, nblk_n*nblk_k] flat block scales to [E, N, K] element-wise.""" - return ( - scale_flat.view(-1, 1) - .repeat(1, blk_n * blk_k) - .view(E, nblk_n, nblk_k, blk_n, blk_k) - .permute(0, 1, 3, 2, 4) - .reshape(E, nblk_n * blk_n, nblk_k * blk_k) - ) - - -def torch_stage1_blockscale_ref( - hidden_states, - w1, - topk_ids, - a_scale, - w1_scale, - scale_blks, - inter_dim, -): - """Torch reference for stage 1 only: returns [B, topk, inter_dim] (SiLU(gate)*up).""" - computeType = torch.float32 - hidden_states = hidden_states.to(computeType) - w1 = w1.to(computeType) - B, D = hidden_states.shape - topk = topk_ids.shape[1] - E = w1.shape[0] - blk_n, blk_k = scale_blks - - if a_scale is not None: - hidden_states = hidden_states.view(B, -1, blk_k) * a_scale.unsqueeze(-1) - hidden_states = hidden_states.view(B, -1) - - nblk_n_w1 = (2 * inter_dim) // blk_n - nblk_k_w1 = D // blk_k - if w1_scale is not None: - w1 = w1 * _expand_blockscale(w1_scale, E, nblk_n_w1, nblk_k_w1, blk_n, blk_k) - - hidden_states = hidden_states.view(B, 1, D).expand(-1, topk, -1) - out = torch.zeros((B, topk, inter_dim), dtype=computeType, device=hidden_states.device) - for e in range(E): - mask = topk_ids == e - if mask.sum(): - sub = hidden_states[mask] - act = sub @ w1[e].T - gate, up = act.split([inter_dim, inter_dim], dim=-1) - out[mask] = F.silu(gate) * up - return out - - -def torch_stage2_blockscale_ref( - act_q, - w2, - topk_ids, - topk_weights, - a_scale, - w2_scale, - scale_blks, - tokens, - model_dim, - inter_dim, - topk, -): - """Torch reference for stage 2 only: takes re-quantized intermediate, returns [B, model_dim].""" - computeType = torch.float32 - blk_n, blk_k = scale_blks - E = w2.shape[0] - - act = act_q.to(computeType) - nblk_k_w2 = inter_dim // blk_k - if a_scale is not None: - act = act.view(-1, nblk_k_w2, blk_k) * a_scale.unsqueeze(-1) - act = act.view(-1, inter_dim) - - w2f = w2.to(computeType) - nblk_n_w2 = model_dim // blk_n - if w2_scale is not None: - w2f = w2f * _expand_blockscale(w2_scale, E, nblk_n_w2, nblk_k_w2, blk_n, blk_k) - - act_3d = act.view(tokens, topk, inter_dim) - out = torch.zeros((tokens, topk, model_dim), dtype=computeType, device=act.device) - for e in range(E): - mask = topk_ids == e - if mask.sum(): - out[mask] = act_3d[mask] @ w2f[e].T - return (out * topk_weights.view(tokens, -1, 1)).sum(dim=1) - - -def torch_moe_blockscale_ref( - hidden_states, - w1, - w2, - topk_weight, - topk_ids, - dtype, - scale_blks=(128, 128), - a_scale=None, - fc1_scale=None, - fc2_scale=None, -): - """Torch reference for blockscale MoE (from aiter test).""" - computeType = torch.float32 - hidden_states = hidden_states.to(computeType) - w1 = w1.to(computeType) - w2 = w2.to(computeType) - token_num = hidden_states.shape[0] - topk = topk_weight.shape[1] - expert, model_dim, inter_dim = w2.shape - B, D = hidden_states.shape - - blk_n, blk_k = scale_blks - if a_scale is not None: - hidden_states = hidden_states.view(token_num, -1, blk_k) * a_scale.unsqueeze(-1) - hidden_states = hidden_states.view(token_num, -1) - - hidden_states = hidden_states.view(token_num, 1, model_dim).repeat(1, topk, 1) - out = torch.zeros((B, topk, D), dtype=computeType, device=hidden_states.device) - - nblk_n_w1 = (2 * inter_dim) // blk_n - nblk_k = model_dim // blk_k - nblk_n_w2 = model_dim // blk_n - nblk_k_w2 = inter_dim // blk_k - - if fc1_scale is not None: - # fc1_scale: [E, nblk_n_w1 * nblk_k] -> expand to [E, 2*inter_dim, model_dim] - # "e n k bn bk -> e (n bn) (k bk)" = permute(0,1,3,2,4).reshape - nblk_n_w1 = (2 * inter_dim) // blk_n - fc1_s = ( - fc1_scale.view(-1, 1) - .repeat(1, blk_n * blk_k) - .view(expert, nblk_n_w1, nblk_k, blk_n, blk_k) - .permute(0, 1, 3, 2, 4) - .reshape(expert, 2 * inter_dim, model_dim) - ) - # fc2_scale: [E, nblk_n_w2 * nblk_k_w2] -> expand to [E, model_dim, inter_dim] - # "e n k bn bk -> e (n bn) (k bk)" = permute(0,1,3,2,4).reshape - fc2_s = ( - fc2_scale.view(-1, 1) - .repeat(1, blk_n * blk_k) - .view(expert, nblk_n_w2, nblk_k_w2, blk_n, blk_k) - .permute(0, 1, 3, 2, 4) - .reshape(expert, model_dim, inter_dim) - ) - w1 = w1 * fc1_s - w2 = w2 * fc2_s - - for E_id in range(w1.shape[0]): - mask = topk_ids == E_id - if mask.sum(): - sub = hidden_states[mask] - act = sub @ w1[E_id].T - gate, up = act.split([inter_dim, inter_dim], dim=-1) - out[mask] = (F.silu(gate) * up) @ w2[E_id].T - - return (out * topk_weight.view(B, -1, 1)).sum(dim=1).to(dtype) - - -@pytest.mark.parametrize( - "token, model_dim, inter_dim, E, topk", - [ - pytest.param(16, 7168, 256, 8, 2, id="small-E8"), - pytest.param(32, 7168, 256, 8, 2, id="medium-E8"), - pytest.param(32, 7168, 256, 256, 8, id="DS-V3-M32", marks=pytest.mark.large_shape), - pytest.param(128, 7168, 256, 256, 8, id="DS-V3-M128", marks=pytest.mark.large_shape), - ], -) -def test_moe_blockscale_e2e( - token, - model_dim, - inter_dim, - E, - topk, - tile_m=16, - tile_n=256, - tile_k=128, - tile_n2=None, - waves_per_eu=2, - num_iters=10, - num_warmup=3, -): - """End-to-end MoE blockscale test: flydsl 2-stage vs aiter fused.""" - dtype = torch.bfloat16 - scale_blks = (128, 128) - scale_blk_n, scale_blk_k = scale_blks - device = torch.device("cuda") - - print("=" * 80) - print( - f"MoE Blockscale E2E: tokens={token}, dim={model_dim}, " - f"inter={inter_dim}, E={E}, topk={topk}, tile={tile_m}x{tile_n}x{tile_k}" - ) - print("=" * 80) - - # ---- Generate data ---- - s = 0.2 - x_fp32 = torch.randn((token, model_dim), device=device, dtype=torch.float32) * s - x_q, x_scale = pertoken_quant(x_fp32, quant_dtype=DTYPE_FP8) - del x_fp32 - - score = torch.rand((token, E), device=device, dtype=torch.float32) - topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) - topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) - - # Per-expert block-quantize: generate fp32 one expert at a time, block-quantize, discard fp32. - def block_quant_expert(w_fp32, blk_n, blk_k): - """Block-quantize single expert weight [N, K] -> (w_q fp8 [N,K], scale [flat]).""" - N_w, K_w = w_fp32.shape - nbn, nbk = N_w // blk_n, K_w // blk_k - tmp = w_fp32.float().view(nbn, blk_n, nbk, blk_k).permute(0, 2, 1, 3).reshape(nbn * nbk, blk_n * blk_k) - q, sc = pertoken_quant(tmp, quant_dtype=DTYPE_FP8) - q = q.view(nbn, nbk, blk_n, blk_k).permute(0, 2, 1, 3).reshape(N_w, K_w) - return q, sc.view(-1) - - w1_bq_list, w1_bscale_list = [], [] - w2_bq_list, w2_bscale_list = [], [] - for e_i in range(E): - w1e = torch.randn((2 * inter_dim, model_dim), device=device, dtype=torch.float32) * s - q, sc = block_quant_expert(w1e, scale_blk_n, scale_blk_k) - w1_bq_list.append(q) - w1_bscale_list.append(sc) - del w1e - - w2e = torch.randn((model_dim, inter_dim), device=device, dtype=torch.float32) * (s / math.sqrt(inter_dim)) - q2, sc2 = block_quant_expert(w2e, scale_blk_n, scale_blk_k) - w2_bq_list.append(q2) - w2_bscale_list.append(sc2) - del w2e - - w1_bq = torch.stack(w1_bq_list) - w1_bscale = torch.stack(w1_bscale_list) - w2_bq = torch.stack(w2_bq_list) - w2_bscale = torch.stack(w2_bscale_list) - del w1_bq_list, w1_bscale_list, w2_bq_list, w2_bscale_list - torch.cuda.empty_cache() - - # Preshuffle block-quantized weights (shared by FlyDSL and aiter) - w1_bq_shuf = shuffle_weight(w1_bq, layout=(16, 16)) - w2_bq_shuf = shuffle_weight(w2_bq, layout=(16, 16)) - w1_shuf = w1_bq_shuf.view(-1) - w2_shuf = w2_bq_shuf.view(-1) - - # Input block quantize: dequantize x_q first (raw fp8 codes overflow the f16 pipeline; #642). - x_f32_for_blk = x_q.float() * x_scale - a1_bq, a1_bscale = pertoken_quant( - x_f32_for_blk.view(-1, model_dim // scale_blk_k, scale_blk_k), quant_dtype=DTYPE_FP8 - ) - a1_bq = a1_bq.view(-1, model_dim) - a1_bscale = a1_bscale.squeeze(-1) - del x_f32_for_blk - - # ---- Torch reference (2-stage with intermediate re-quant, matching FlyDSL pipeline) ---- - out1_torch_ref = torch_stage1_blockscale_ref(a1_bq, w1_bq, topk_ids, a1_bscale, w1_bscale, scale_blks, inter_dim) - out1_torch_fp16 = out1_torch_ref.to(torch.float16) - a2_ref_bq, a2_ref_bscale = pertoken_quant( - out1_torch_fp16.float().view(-1, inter_dim // scale_blk_k, scale_blk_k), quant_dtype=DTYPE_FP8 - ) - a2_ref_bq = a2_ref_bq.view(-1, inter_dim) - a2_ref_bscale = a2_ref_bscale.squeeze(-1) - out_ref = torch_stage2_blockscale_ref( - a2_ref_bq, - w2_bq, - topk_ids, - topk_weights, - a2_ref_bscale, - w2_bscale, - scale_blks, - token, - model_dim, - inter_dim, - topk, - ) - - # ---- flydsl 2-stage (true blockscale) ---- - routing = build_routing_buffers( - topk_ids=topk_ids, - topk_weights=topk_weights, - experts=E, - model_dim=model_dim, - tile_m=tile_m, - moe_sort_mode="aiter" if HAS_AITER else "torch", - ) - - sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, _sorted_size, _blocks = routing - size_expert_ids = sorted_expert_ids.numel() - - # Prepare blockscale tensors for FlyDSL kernels. - # Use per_group_quant_hip(transpose_scale=True) so the quant kernel writes scale - # directly in [nblk_k, token] memory layout — no separate .t().contiguous() needed. - # Fallback to pertoken_quant + manual transpose when aiter is unavailable. - nblk_k_w1 = model_dim // scale_blk_k - w1_scale_fly = w1_bscale.view(-1) # [E, nblk_n_w1 * nblk_k_w1] flat - - if HAS_AITER: - # quant kernel writes transposed scale layout directly - a1_bq, a1_scale_fly = per_group_quant_hip( - (x_q.float() * x_scale).to(torch.bfloat16), - quant_dtype=DTYPE_FP8, - group_size=scale_blk_k, - transpose_scale=True, - ) - a1_scale_fly = a1_scale_fly.view(-1) # [nblk_k_w1 * token] flat - else: - a1_scale_fly = a1_bscale.t().contiguous().view(-1) - - # Compile stage 1 - exe1 = compile_moe_blockscale_gemm1( - 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=False, - scale_block_k=scale_blk_k, - out_dtype="f16", - waves_per_eu=waves_per_eu, - ) - - out1 = torch.zeros((token, topk, inter_dim), dtype=torch.float16, device=device) - stream = torch.cuda.current_stream() - - compiled_exe1 = flyc.compile( - exe1, - out1.view(-1), - a1_bq.view(-1), - w1_shuf, - a1_scale_fly, - w1_scale_fly, - sorted_ids, - sorted_expert_ids, - sorted_weights, - num_valid_ids, - token, - inter_dim, - model_dim, - size_expert_ids, - stream, - ) - - def launch_stage1(): - compiled_exe1( - out1.view(-1), - a1_bq.view(-1), - w1_shuf, - a1_scale_fly, - w1_scale_fly, - sorted_ids, - sorted_expert_ids, - sorted_weights, - num_valid_ids, - token, - inter_dim, - model_dim, - size_expert_ids, - stream, - ) - - _, us1 = run_perftest(launch_stage1, num_iters=num_iters, num_warmup=num_warmup) - torch.cuda.synchronize() - - # Verify stage 1 independently (ref uses a1_bscale in original [token, nblk_k] layout) - out1_ref = torch_stage1_blockscale_ref(a1_bq, w1_bq, topk_ids, a1_bscale, w1_bscale, scale_blks, inter_dim) - err_s1 = checkAllclose( - out1_ref.to(out1.dtype), out1, rtol=0.1, atol=0.1, msg="flydsl-stage1 vs ref", printLog=False - ) - print(f" stage1 err_ratio = {err_s1:.4f}") - - # Re-quantize stage1 output for stage 2 - nblk_k_w2 = inter_dim // scale_blk_k - w2_scale_fly = w2_bscale.view(-1) # [E, nblk_n_w2 * nblk_k_w2] flat - - if HAS_AITER: - a2_bq, a2_scale_fly = per_group_quant_hip( - out1.to(torch.bfloat16).view(-1, inter_dim), - quant_dtype=DTYPE_FP8, - group_size=scale_blk_k, - transpose_scale=True, - ) - a2_scale_fly = a2_scale_fly.view(-1) # [nblk_k_w2 * token*topk] flat - a2_bscale_2d = a2_scale_fly.view(nblk_k_w2, -1).t().contiguous() # [token*topk, nblk_k_w2] for ref - else: - a2_bq, a2_bscale_2d = pertoken_quant( - out1.float().view(-1, inter_dim // scale_blk_k, scale_blk_k), quant_dtype=DTYPE_FP8 - ) - a2_bq = a2_bq.view(-1, inter_dim) - a2_bscale_2d = a2_bscale_2d.squeeze(-1) # [token*topk, nblk_k_w2] - a2_scale_fly = a2_bscale_2d.t().contiguous().view(-1) - - # Compile stage 2 (optionally with different tile_n) - tile_n_s2 = tile_n2 if tile_n2 is not None else tile_n - exe2 = compile_moe_blockscale_gemm2( - model_dim=model_dim, - inter_dim=inter_dim, - experts=E, - topk=topk, - tile_m=tile_m, - tile_n=tile_n_s2, - tile_k=tile_k, - doweight_stage2=True, - scale_block_k=scale_blk_k, - out_dtype="f16", - waves_per_eu=waves_per_eu, - ) - - out2 = torch.zeros((token, model_dim), dtype=torch.float16, device=device) - - compiled_exe2 = flyc.compile( - exe2, - out2.view(-1), - a2_bq.view(-1), - w2_shuf, - a2_scale_fly, - w2_scale_fly, - sorted_ids, - sorted_expert_ids, - sorted_weights, - num_valid_ids, - token, - model_dim, - inter_dim, - size_expert_ids, - stream, - ) - - def launch_stage2(): - compiled_exe2( - out2.view(-1), - a2_bq.view(-1), - w2_shuf, - a2_scale_fly, - w2_scale_fly, - sorted_ids, - sorted_expert_ids, - sorted_weights, - num_valid_ids, - token, - model_dim, - inter_dim, - size_expert_ids, - stream, - ) - - _, us2 = run_perftest(launch_stage2, num_iters=num_iters, num_warmup=num_warmup) - torch.cuda.synchronize() - - # Re-run once with clean output for verification (atomic accumulation requires pre-zeroed buffer) - out2.zero_() - launch_stage2() - torch.cuda.synchronize() - - # Verify stage 2 independently - out2_ref_s2 = torch_stage2_blockscale_ref( - a2_bq, w2_bq, topk_ids, topk_weights, a2_bscale_2d, w2_bscale, scale_blks, token, model_dim, inter_dim, topk - ) - err_s2 = checkAllclose( - out2_ref_s2.to(out2.dtype), out2, rtol=0.1, atol=0.1, msg="flydsl-stage2 vs ref", printLog=False - ) - print(f" stage2 err_ratio = {err_s2:.4f}") - - us_fly_total = us1 + us2 - print(f" flydsl: stage1={us1:.1f}us + stage2={us2:.1f}us = total {us_fly_total:.1f}us") - - # Verify flydsl full pipeline vs blockscale torch reference - err_fly = checkAllclose(out_ref.to(out2.dtype), out2, rtol=0.1, atol=0.1, msg="flydsl vs ref", printLog=False) - print(f" flydsl: pipeline err_ratio vs ref = {err_fly:.4f}") - # Assert instead of returning timings, so a wrong pipeline can't pass silently (#642). - assert torch.isfinite(out2).all(), "FlyDSL block-scale pipeline produced non-finite output" - assert err_fly <= 0.05, f"FlyDSL block-scale pipeline diverges from reference: err_ratio={err_fly:.4f}" - - # ---- aiter comparisons ---- - us_aiter_fused = 0 - us_ck_s1 = 0 - us_ck_s2 = 0 - if HAS_AITER: - from aiter import ActivationType, QuantType - from aiter.ops.moe_op import ck_moe_stage1_fwd, ck_moe_stage2_fwd - - nblk_n_w1 = (2 * inter_dim) // scale_blk_n - nblk_n_w2 = model_dim // scale_blk_n - - ck_block_m = 32 - sorted_ids_a, sorted_weights_a, sorted_expert_ids_a, num_valid_ids_a, out_aiter = aiter_moe_sorting( - topk_ids.to(torch.int32), topk_weights.float(), E, model_dim, dtype, ck_block_m - ) - - # --- aiter fused 1-stage --- - out_ref_aiter = torch_moe_blockscale_ref( - a1_bq, - w1_bq, - w2_bq, - topk_weights, - topk_ids, - dtype, - scale_blks=scale_blks, - a_scale=a1_bscale, - fc1_scale=w1_bscale, - fc2_scale=w2_bscale, - ) - try: - # a1_scale_fly is [nblk_k_w1, token] contiguous — exactly what fmoe_fp8_blockscale_g1u1 expects - a1_scale_aiter = a1_scale_fly.view(nblk_k_w1, token) - - def launch_aiter_fused(): - out_aiter.zero_() - aiter.fmoe_fp8_blockscale_g1u1( - out_aiter, - a1_bq, - w1_bq_shuf, - w2_bq_shuf, - sorted_ids_a, - sorted_weights_a, - sorted_expert_ids_a, - num_valid_ids_a, - topk, - a1_scale_aiter, - w1_bscale, - w2_bscale, - "", - scale_blk_n, - scale_blk_k, - None, - ) - - _, us_aiter_fused = run_perftest(launch_aiter_fused, num_iters=num_iters, num_warmup=num_warmup) - torch.cuda.synchronize() - err_fused = checkAllclose( - out_ref_aiter, out_aiter, rtol=0.1, atol=0.1, msg="aiter-fused vs ref", printLog=False - ) - print(f" aiter fused: {us_aiter_fused:.1f}us (err_ratio={err_fused:.4f})") - except Exception as e: - print(f" aiter fused: FAILED - {str(e)[:80]}") - - # --- aiter CK stage1 (blockscale) --- - try: - out_ck_s1 = torch.zeros((token, topk, inter_dim), dtype=torch.float16, device=device) - w1_scale_ck = w1_bscale.view(E, nblk_n_w1, nblk_k_w1).contiguous() - - def launch_ck_s1(): - ck_moe_stage1_fwd( - hidden_states=a1_bq, - w1=w1_bq_shuf, - w2=w2_bq_shuf, - sorted_token_ids=sorted_ids_a, - sorted_expert_ids=sorted_expert_ids_a, - num_valid_ids=num_valid_ids_a, - out=out_ck_s1, - topk=topk, - kernelName="", - w1_scale=w1_scale_ck, - a1_scale=a1_bscale, - block_m=ck_block_m, - sorted_weights=None, - quant_type=QuantType.per_1x128, - activation=ActivationType.Silu, - dst_type=torch.float16, - ) - - _, us_ck_s1 = run_perftest(launch_ck_s1, num_iters=num_iters, num_warmup=num_warmup) - torch.cuda.synchronize() - err_ck_s1 = checkAllclose( - out1_ref.to(out_ck_s1.dtype), out_ck_s1, rtol=0.1, atol=0.1, msg="ck-stage1 vs ref", printLog=False - ) - err_ck_s1_vs_fly = checkAllclose( - out1.float(), out_ck_s1.float(), rtol=0.1, atol=0.1, msg="ck-stage1 vs flydsl", printLog=False - ) - print(f" ck stage1: {us_ck_s1:.1f}us (err_vs_ref={err_ck_s1:.4f}, err_vs_fly={err_ck_s1_vs_fly:.4f})") - except Exception as e: - import traceback - - traceback.print_exc() - print(f" ck stage1: FAILED - {str(e)[:120]}") - - # --- aiter CK stage2 (blockscale) --- - try: - out_ck_s2 = torch.zeros((token, model_dim), dtype=torch.float16, device=device) - w2_scale_ck = w2_bscale.view(E, nblk_n_w2, nblk_k_w2).contiguous() - - def launch_ck_s2(): - out_ck_s2.zero_() - ck_moe_stage2_fwd( - inter_states=a2_bq.view(token, topk, inter_dim), - w1=w1_bq_shuf, - w2=w2_bq_shuf, - sorted_token_ids=sorted_ids_a, - sorted_expert_ids=sorted_expert_ids_a, - num_valid_ids=num_valid_ids_a, - out=out_ck_s2, - topk=topk, - kernelName="", - w2_scale=w2_scale_ck, - a2_scale=a2_bscale_2d, - block_m=ck_block_m, - sorted_weights=sorted_weights_a, - quant_type=QuantType.per_1x128, - activation=ActivationType.Silu, - ) - - _, us_ck_s2 = run_perftest(launch_ck_s2, num_iters=num_iters, num_warmup=num_warmup) - torch.cuda.synchronize() - # Re-run once for clean verification - out_ck_s2.zero_() - launch_ck_s2() - torch.cuda.synchronize() - err_ck_s2 = checkAllclose( - out2_ref_s2.to(out_ck_s2.dtype), out_ck_s2, rtol=0.1, atol=0.1, msg="ck-stage2 vs ref", printLog=False - ) - err_ck_s2_vs_fly = checkAllclose( - out2.float(), out_ck_s2.float(), rtol=0.1, atol=0.1, msg="ck-stage2 vs flydsl", printLog=False - ) - print(f" ck stage2: {us_ck_s2:.1f}us (err_vs_ref={err_ck_s2:.4f}, err_vs_fly={err_ck_s2_vs_fly:.4f})") - except Exception as e: - import traceback - - traceback.print_exc() - print(f" ck stage2: FAILED - {str(e)[:120]}") - - # ---- Summary ---- - flops_stage1 = 2 * token * topk * (2 * inter_dim) * model_dim - flops_stage2 = 2 * token * topk * model_dim * inter_dim - flops_total = flops_stage1 + flops_stage2 - tflops_fly = flops_total / (us_fly_total / 1e6) / 1e12 - tflops_fused = flops_total / (us_aiter_fused / 1e6) / 1e12 if us_aiter_fused > 0 else 0 - tflops_ck_s1 = flops_stage1 / (us_ck_s1 / 1e6) / 1e12 if us_ck_s1 > 0 else 0 - tflops_ck_s2 = flops_stage2 / (us_ck_s2 / 1e6) / 1e12 if us_ck_s2 > 0 else 0 - - print(f"\n {'':>14} | {'us':>8} | {'TFLOPS':>8} | vs flydsl") - print(f" {'flydsl-s1':>14} | {us1:>8.1f} | {flops_stage1/(us1/1e6)/1e12:>8.2f} |") - print(f" {'flydsl-s2':>14} | {us2:>8.1f} | {flops_stage2/(us2/1e6)/1e12:>8.2f} |") - print(f" {'flydsl-total':>14} | {us_fly_total:>8.1f} | {tflops_fly:>8.2f} |") - if us_aiter_fused > 0: - r = us_aiter_fused / us_fly_total - print(f" {'aiter-fused':>14} | {us_aiter_fused:>8.1f} | {tflops_fused:>8.2f} | {r:.2f}x") - if us_ck_s1 > 0: - r1 = us_ck_s1 / us1 - print(f" {'ck-stage1':>14} | {us_ck_s1:>8.1f} | {tflops_ck_s1:>8.2f} | {r1:.2f}x vs fly-s1") - if us_ck_s2 > 0: - r2 = us_ck_s2 / us2 - print(f" {'ck-stage2':>14} | {us_ck_s2:>8.1f} | {tflops_ck_s2:>8.2f} | {r2:.2f}x vs fly-s2") - if us_ck_s1 > 0 and us_ck_s2 > 0: - ck_total = us_ck_s1 + us_ck_s2 - tflops_ck_total = flops_total / (ck_total / 1e6) / 1e12 - r_total = ck_total / us_fly_total - print(f" {'ck-total':>14} | {ck_total:>8.1f} | {tflops_ck_total:>8.2f} | {r_total:.2f}x") - print() - - return us_fly_total, us_aiter_fused - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser( - description="MoE Blockscale Benchmark (FlyDSL vs aiter-fused vs aiter-ck2s)", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument("-m", type=int, default=0, help="Token count (0 = sweep [16,32,128,256,1024])") - parser.add_argument("-dim", type=int, default=7168) - parser.add_argument("-idim", type=int, default=2048) - parser.add_argument("-e", "--expert", type=int, default=32) - parser.add_argument("-k", "--topk", type=int, default=8) - parser.add_argument("--tile_m", type=int, default=64) - parser.add_argument("--tile_n", type=int, default=128) - parser.add_argument("--tile_n2", type=int, default=0, help="tile_n for stage2 (0 = same as tile_n)") - parser.add_argument("--tile_k", type=int, default=128) - parser.add_argument("--wpe", type=int, default=2, help="waves_per_eu (0 = disabled)") - parser.add_argument("--num_iters", type=int, default=20) - parser.add_argument("--num_warmup", type=int, default=5) - args = parser.parse_args() - - torch.set_default_device("cuda") - wpe = args.wpe if args.wpe > 0 else None - - token_list = [args.m] if args.m > 0 else [128, 16384] - - print( - f"\nMoE Blockscale: dim={args.dim}, inter={args.idim}, E={args.expert}, " - f"topk={args.topk}, tile={args.tile_m}x{args.tile_n}x{args.tile_k}, wpe={wpe}" - ) - print(f"{'tokens':>8} | {'flydsl':>10} | {'aiter-fused':>12} | {'aiter-ck2s':>12}") - print("-" * 52) - - tn2 = args.tile_n2 if args.tile_n2 > 0 else None - for m in token_list: - us_fly, us_aiter = test_moe_blockscale_e2e( - token=m, - model_dim=args.dim, - inter_dim=args.idim, - E=args.expert, - topk=args.topk, - tile_m=args.tile_m, - tile_n=args.tile_n, - tile_k=args.tile_k, - tile_n2=tn2, - waves_per_eu=wpe, - num_iters=args.num_iters, - num_warmup=args.num_warmup, - ) - print(f"{m:>8} | {us_fly:>10.1f} |") diff --git a/tests/kernels/test_moe_gemm.py b/tests/kernels/test_moe_gemm.py index 0c36340a3..e0b486c65 100644 --- a/tests/kernels/test_moe_gemm.py +++ b/tests/kernels/test_moe_gemm.py @@ -75,16 +75,19 @@ def _pack_shuffled_int8_to_packed_int4_no_perm(x_shuf_i8: torch.Tensor) -> torch HAS_AITER = False # Kernel implementations live under `kernels/`; this test file is the harness. -from kernels.moe.mixed_moe_gemm_2stage import ( # noqa: E402 - compile_mixed_moe_gemm1, - compile_mixed_moe_gemm2, -) +# The a4w4 (MX-FP4) and a8w4 (MX-FP8 activation) paths run through the fused mxfp_moe +# pipeline (device-side re-quant, sorted fp4 intermediate) via `_run_mxfp_moe_e2e`, +# which replaced the parametric mixed_moe_gemm_2stage builders. from kernels.moe.moe_gemm_2stage import ( # noqa: E402 MoeGemm2Mode, compile_moe_gemm1, compile_moe_gemm2, compile_moe_gemm2_ex, ) +from kernels.moe.mxfp_moe import ( # noqa: E402 + flydsl_mxfp4_gemm1, + flydsl_mxfp4_gemm2, +) logging.basicConfig(level=logging.INFO) @@ -418,32 +421,44 @@ def run_moe_stage1( if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4"): raise ValueError( - f"in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth','int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" + "in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth'," + f"'int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" ) is_int4 = in_dtype == "int4" is_int4_bf16 = in_dtype == "int4_bf16" # W4A16: bf16 activations, packed int4 weights is_int8smooth = in_dtype == "int8smooth" - is_fp4 = in_dtype == "fp4" is_a8w4 = in_dtype == "a8w4" # MX-FP8 activation + MX-FP4 weight - is_fp4_path = is_fp4 or is_a8w4 # shared weight/shuffle pipeline + is_fp4_path = in_dtype in ("fp4", "a8w4") # both drive the fused mxfp_moe pipeline + if is_fp4_path: + # a4w4/a8w4 run through the fused mxfp_moe kernels (flydsl_mxfp4_gemm1/2), which + # replaced the parametric mixed builders. The fused stage1 emits a sorted fp4 + # intermediate only the fused stage2 can consume, so route through the full e2e + # helper rather than a stage1-only launch. + _run_mxfp_moe_e2e( + tokens=tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + use_reduce=(tile_m == 128), + x_fp32=x_fp32, + w1_fp32=w1_fp32, + w2_fp32=w2_fp32, + topk_ids=topk_ids, + topk_weights=topk_weights, + routing=routing, + a_dtype="fp8" if is_a8w4 else "fp4", + skip_ref=bool(skip_ref), + num_iters=num_iters, + num_warmup=num_warmup, + in_dtype_label=in_dtype, + ) + return (None, None) if return_outputs else None use_packed_int4 = is_int4 or is_int4_bf16 # Quantize inputs / weights. - if is_fp4_path: - from tests.kernels.utils import gemm_common_utils - - # x: MX-FP8 (e4m3fn, 1 B/elem) for a8w4; MX-FP4 (packed, 0.5 B/elem) for fp4. - if is_a8w4: - x_q, x_scale_raw = _per_1x32_mxfp8_quant(x_fp32) - else: - x_q, x_scale_raw = _per_1x32_fp4_quant(x_fp32) - w1_flat_fp32 = w1_fp32.view(experts * (2 * inter_dim), model_dim) - w1_fp4, w1_scale_raw = _per_1x32_fp4_quant(w1_flat_fp32) - w1_q = w1_fp4 - w2_q = None # not needed for stage1 - scale_x = x_scale_raw # raw e8m0 scale (K dim divided by 32) - scale_w1 = w1_scale_raw # raw e8m0 [E*2N, K//32] - elif in_dtype == "fp8": + if in_dtype == "fp8": x_q, scale_x = pertoken_quant(x_fp32, quant_dtype=DTYPE_FP8) # [tokens,K], [tokens,1] w1_q, scale_w1 = pertoken_quant(w1_fp32, quant_dtype=DTYPE_FP8) # [E,2*inter,K], [E,2*inter,1] # w2 is not used by our kernel, but required by aiter stage1 API @@ -518,65 +533,34 @@ def run_moe_stage1( scale_w1 = None # Preshuffle weights and prepare scale tensors. - if is_fp4_path: - # FP4 path (fp4 or a8w4): W1 is always MX-FP4; x is MX-FP4 (fp4) or MX-FP8 (a8w4). - w1_shuffled = shuffle_weight(w1_q.view(torch.float4_e2m1fn_x2)) - w_kernel = w1_shuffled.view(torch.uint8).contiguous() - w1_q_flat = w1_q.view(experts * (2 * inter_dim), model_dim // 2).contiguous() - scale_w1_flat = scale_w1.view(experts * (2 * inter_dim), model_dim // 32).contiguous() - # Weight scale: e8m0_shuffle (no MoE sorting needed for weights). - scale_w1_1d = gemm_common_utils.e8m0_shuffle(scale_w1).view(torch.uint8).contiguous() - # Activation scale: must be sorted by MoE routing order. - scale_x_1d = ( - gemm_common_utils.moe_mxfp4_sort( - scale_x[:tokens, :].view(tokens, 1, -1), - sorted_ids=sorted_token_ids, - num_valid_ids=num_valid_ids, - token_num=tokens, - block_size=tile_m, - ) - .view(torch.uint8) - .contiguous() - ) - # Kernel consumes raw bytes: fp4 packs 2 elems/byte (K//2 cols); a8w4 is 1 B/elem (K cols). - # Preserve the original dtype view for the torch reference path: mxfp4 detection in - # ``_detect_scale_kind`` needs the packed fp4 layout (shape[-1]*2 == scale*32), while - # mxfp8 detection needs dtype=float8_e4m3fn. The kernel itself always reads uint8. - x_q_ref = x_q - x_q = x_q.view(torch.uint8).contiguous().view(tokens, -1) + w1_shuffled = shuffle_weight(w1_q) + w2_shuffled = shuffle_weight(w2_q) if in_dtype == "fp8" else None + + # Flatten W1 for our FlyDSL kernel (treat expert dim as part of N). + w1_shuffled_flat = w1_shuffled.view(experts * (2 * inter_dim), model_dim) + w1_q_flat = w1_q.view(experts * (2 * inter_dim), model_dim) + scale_w1_flat = None if scale_w1 is None else scale_w1.view(experts * (2 * inter_dim), 1) + + # No host-side padding: keep tensors contiguous and rely on kernel-side resource sizes / early-exit. + x_q = x_q.contiguous().view(tokens * topk, model_dim) if is_int8smooth else x_q.contiguous().view(tokens, model_dim) + # Pack weights for int4 variants (W4A8 and W4A16). + w_kernel = ( + _pack_shuffled_int8_to_packed_int4_no_perm(w1_shuffled_flat) if use_packed_int4 else w1_shuffled_flat + ).contiguous() + if not use_packed_int4: + w_kernel = w_kernel.view(experts * (2 * inter_dim), model_dim) + + # Flatten scales to 1D memrefs (fp16 path uses 0-sized scale tensors; kernel ignores them). + if scale_x is None: + scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32) else: - w1_shuffled = shuffle_weight(w1_q) - w2_shuffled = shuffle_weight(w2_q) if in_dtype == "fp8" else None - - # Flatten W1 for our FlyDSL kernel (treat expert dim as part of N). - w1_shuffled_flat = w1_shuffled.view(experts * (2 * inter_dim), model_dim) - w1_q_flat = w1_q.view(experts * (2 * inter_dim), model_dim) - scale_w1_flat = None if scale_w1 is None else scale_w1.view(experts * (2 * inter_dim), 1) - - # No host-side padding: keep tensors contiguous and rely on kernel-side resource sizes / early-exit. - x_q = ( - x_q.contiguous().view(tokens * topk, model_dim) - if is_int8smooth - else x_q.contiguous().view(tokens, model_dim) - ) - # Pack weights for int4 variants (W4A8 and W4A16). - w_kernel = ( - _pack_shuffled_int8_to_packed_int4_no_perm(w1_shuffled_flat) if use_packed_int4 else w1_shuffled_flat - ).contiguous() - if not use_packed_int4: - w_kernel = w_kernel.view(experts * (2 * inter_dim), model_dim) - - # Flatten scales to 1D memrefs (fp16 path uses 0-sized scale tensors; kernel ignores them). - if scale_x is None: - scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32) - else: - scale_x_1d = scale_x.view(-1).contiguous() - if use_groupwise_scale: - scale_w1_1d = scale_w1_prepared.view(-1).contiguous() - elif scale_w1_flat is None: - scale_w1_1d = torch.empty((0,), device=device, dtype=torch.float32) - else: - scale_w1_1d = scale_w1_flat.view(-1).contiguous() + scale_x_1d = scale_x.view(-1).contiguous() + if use_groupwise_scale: + scale_w1_1d = scale_w1_prepared.view(-1).contiguous() + elif scale_w1_flat is None: + scale_w1_1d = torch.empty((0,), device=device, dtype=torch.float32) + else: + scale_w1_1d = scale_w1_flat.view(-1).contiguous() sorted_weights_1d = sorted_weights.contiguous().view(-1) # [sorted_size] # Output: normal=[tokens, topk, inter_dim] f16/bf16, split-K=[tokens*topk, 2*inter_dim] f32 @@ -587,104 +571,50 @@ def run_moe_stage1( else: out = torch.empty((tokens, topk, inter_dim), device=device, dtype=_out_torch_dtype) - if is_fp4_path: - exe = 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=bool(doweight_stage1), - a_dtype="fp8" if is_a8w4 else "fp4", - b_dtype="fp4", - out_dtype="f16", - act="silu", - ) - bias_dummy = torch.empty((0,), device=device, dtype=torch.float32) - # Empty placeholder: stage1 writes a sorted E8M0 scale buffer only - # when gate_mode=INTERLEAVE + fp8/fp4 fused quant. For fp4-fp4 we - # still have to pass the arg slot (launcher signature is fixed). - out_scale_sorted_dummy = torch.empty((0,), device=device, dtype=torch.uint8) - - def _s1_args_fp4(o, x, w, sx, sw, st, eids, sw_sorted): - return ( - o, - x, - w, - sx, - sw, - st, - eids, - sw_sorted, - num_valid_ids, - bias_dummy, - out_scale_sorted_dummy, - tokens, - inter_dim * 2, - model_dim, - int(blocks), - torch.cuda.current_stream(), - ) - - compiled_exe = flyc.compile( - exe, - *_s1_args_fp4( - out, x_q, w_kernel, scale_x_1d, scale_w1_1d, sorted_token_ids, sorted_expert_ids, sorted_weights_1d - ), - ) - - def launch(o, x, w, sx, sw, st, eids, sw_sorted): - compiled_exe(*_s1_args_fp4(o, x, w, sx, sw, st, eids, sw_sorted)) + exe = compile_moe_gemm1( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + in_dtype=in_dtype, + group_size=group_size, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=bool(doweight_stage1), + use_cshuffle_epilog=None if _is_splitk else False, + scale_is_bf16=(scale_dtype == "bf16"), + out_dtype=out_dtype, + k_batch=k_batch, + ) - else: - exe = compile_moe_gemm1( - model_dim=model_dim, - inter_dim=inter_dim, - experts=experts, - topk=topk, - in_dtype=in_dtype, - group_size=group_size, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage1=bool(doweight_stage1), - use_cshuffle_epilog=None if _is_splitk else False, - scale_is_bf16=(scale_dtype == "bf16"), - out_dtype=out_dtype, - k_batch=k_batch, + def _s1_args(o, x, w, sx, sw, st, eids, sw_sorted): + return ( + o, + x, + w, + sx, + sw, + st, + eids, + sw_sorted, + num_valid_ids, + tokens, + inter_dim, + model_dim, + int(blocks), + torch.cuda.current_stream(), ) - def _s1_args(o, x, w, sx, sw, st, eids, sw_sorted): - return ( - o, - x, - w, - sx, - sw, - st, - eids, - sw_sorted, - num_valid_ids, - tokens, - inter_dim, - model_dim, - int(blocks), - torch.cuda.current_stream(), - ) - - compiled_exe = flyc.compile( - exe, - *_s1_args( - out, x_q, w_kernel, scale_x_1d, scale_w1_1d, sorted_token_ids, sorted_expert_ids, sorted_weights_1d - ), - ) + compiled_exe = flyc.compile( + exe, + *_s1_args(out, x_q, w_kernel, scale_x_1d, scale_w1_1d, sorted_token_ids, sorted_expert_ids, sorted_weights_1d), + ) - def launch(o, x, w, sx, sw, st, eids, sw_sorted): - if _is_splitk: - o.zero_() - compiled_exe(*_s1_args(o, x, w, sx, sw, st, eids, sw_sorted)) + def launch(o, x, w, sx, sw, st, eids, sw_sorted): + if _is_splitk: + o.zero_() + compiled_exe(*_s1_args(o, x, w, sx, sw, st, eids, sw_sorted)) _, us = run_perftest( launch, @@ -731,8 +661,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): atol = 0.5 if (is_int4 or is_int4_bf16) else 0.25 assert verify_output(out.to(torch.float32), ref, rtol=rtol, atol=atol) else: - # Use original-dtype view for fp4/a8w4 so `_detect_scale_kind` picks the right branch. - x_ref = x_q_ref if is_fp4_path else x_q + x_ref = x_q sx_ref = scale_x ref = torch_moe_gemm1( x_ref, @@ -746,14 +675,14 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): group_size=group_size, scale_w1_groups=scale_w1_groups, ) - rtol = 0.5 if (is_int4 or is_int4_bf16 or is_fp4_path) else 0.25 - atol = 0.5 if (is_int4 or is_int4_bf16 or is_fp4_path) else 0.25 + rtol = 0.5 if (is_int4 or is_int4_bf16) else 0.25 + atol = 0.5 if (is_int4 or is_int4_bf16) else 0.25 assert verify_output( out.to(torch.float32), ref, rtol=rtol, atol=atol, - logits_diff_threshold=1 if is_fp4_path else 2e-3, + logits_diff_threshold=2e-3, ) # Note: kernel launches full expert-block range; effective work is gated by num_valid_ids. @@ -765,12 +694,10 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): active_experts = min(experts, tokens * topk) is_f16_or_bf16_s1 = is_int4_bf16 or in_dtype in ("bf16", "fp16") # Per-element bits for X and W (W1 total cols = 2*inter_dim). - # fp4: x=4b, w=4b (both MX-FP4) - # a8w4: x=8b, w=4b (MX-FP8 act + MX-FP4 weight) # f16/bf16/int4_bf16: x=16b # int4/int8/fp8 etc.: x=8b, w packed to 4b when W4A*. - x_bits = 4 if is_fp4 else (16 if is_f16_or_bf16_s1 else 8) - w_bits = 4 if (is_fp4_path or use_packed_int4) else (16 if is_f16_or_bf16_s1 else 8) + x_bits = 16 if is_f16_or_bf16_s1 else 8 + w_bits = 4 if use_packed_int4 else (16 if is_f16_or_bf16_s1 else 8) x_rows = tokens * topk if is_int8smooth else tokens x_elems = x_rows * model_dim w_elems = active_experts * (2 * inter_dim) * model_dim @@ -781,11 +708,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): bytes_moved += tokens * topk * inter_dim * 2 # out fp16 (logical, post-silu) # scale bytes - if is_fp4_path: - # Per-1x32 E8M0 scale: 1 byte per 32 logical elements for both x and w1. - bytes_moved += x_elems // 32 - bytes_moved += w_elems // 32 - elif use_groupwise_scale: + if use_groupwise_scale: num_groups_s1 = model_dim // group_size _scale_bytes = 2 if scale_dtype == "bf16" else 4 bytes_moved += active_experts * num_groups_s1 * (2 * inter_dim) * _scale_bytes # groupwise scale @@ -873,7 +796,8 @@ def launch_ck(o, x, w1_, w2_, sorted_ids_, sorted_eids_, num_valid_, w1_scale_, flops = 2 * tokens * topk * (2 * inter_dim) * model_dim tflops_ck = flops / (us_ck / 1e6) / 1e12 print( - f"[aiter] stage1: {us_ck:.1f} us, {tflops_ck:.2f} TFLOPS, FlyDSL vs aiter speedups: {tflops / tflops_ck:.2f}x" + f"[aiter] stage1: {us_ck:.1f} us, {tflops_ck:.2f} TFLOPS, " + f"FlyDSL vs aiter speedups: {tflops / tflops_ck:.2f}x" ) except Exception as e: # Treat aiter compare as best-effort: many environments can import `aiter` but can't load @@ -941,9 +865,9 @@ def run_moe_stage2( raise ValueError( "Invalid stage2 tiling: inter_dim ({inter_dim}) must be divisible by tile_k2 ({tile_k}). " "Try setting `--tile_k2` to a divisor of inter_dim. " - "Tip: stage2 splits A2 loads across 256 threads; if you want smaller tile_k2, you may need a larger tile_m so (tile_m*tile_k2) stays divisible by 1024.".format( - inter_dim=inter_dim, tile_k=tile_k - ) + "Tip: stage2 splits A2 loads across 256 threads; if you want smaller " + "tile_k2, you may need a larger tile_m so (tile_m*tile_k2) stays " + "divisible by 1024.".format(inter_dim=inter_dim, tile_k=tile_k) ) # Enforce the kernel's stage2 gmem->reg load mapping constraints. # See: kernels/moe_gemm_2stage.py::compile_moe_gemm2 (x_load_bytes selection). @@ -1044,15 +968,39 @@ def run_moe_stage2( if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4"): raise ValueError( - f"in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth','int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" + "in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth'," + f"'int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" ) is_int4 = in_dtype == "int4" is_int4_bf16 = in_dtype == "int4_bf16" # W4A16: bf16 activations, packed int4 weights is_int8smooth = in_dtype == "int8smooth" - is_fp4 = in_dtype == "fp4" is_a8w4 = in_dtype == "a8w4" # MX-FP8 activation + MX-FP4 weight - # Share the FP4 stage2 path (W2 shuffle / scale sort / mixed kernel). - is_fp4_path = is_fp4 or is_a8w4 + is_fp4_path = in_dtype in ("fp4", "a8w4") # both drive the fused mxfp_moe pipeline + if is_fp4_path: + # a4w4/a8w4 run through the fused mxfp_moe kernels (flydsl_mxfp4_gemm1/2). The + # fused stage2 consumes the sorted fp4 intermediate emitted by the fused stage1, + # so route through the full e2e helper rather than a stage2-only launch. + _run_mxfp_moe_e2e( + tokens=tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + use_reduce=bool(use_reduce), + x_fp32=x_fp32, + w1_fp32=w1_fp32, + w2_fp32=w2_fp32, + topk_ids=topk_ids, + topk_weights=topk_weights, + routing=routing, + a_dtype="fp8" if is_a8w4 else "fp4", + skip_ref=bool(skip_ref), + num_iters=num_iters, + num_warmup=num_warmup, + in_dtype_label=in_dtype, + ) + return (None, None) if return_outputs else None use_packed_int4 = is_int4 or is_int4_bf16 # Quantize inputs / weights. @@ -1090,24 +1038,6 @@ def run_moe_stage2( w1_q, scale_w1 = pertoken_quant(w1_fp32, quant_dtype=torch.int8, dtypeMax=7) w2_q, scale_w2 = pertoken_quant(w2_fp32, quant_dtype=torch.int8, dtypeMax=7) scale_x = None - elif in_dtype in ("fp4", "a8w4"): - from tests.kernels.utils import gemm_common_utils - - if gemm_common_utils is None: - pytest.skip("gemm_common_utils not available (triton not installed)") - if "gfx95" not in ARCH: - pytest.skip(f"FP4 MFMA requires gfx950+, got {ARCH}") - # FP4 / A8W4 share the MXFP4 W2 path; quantize W2 only here. A2 comes - # from `a2_fp8_in` (FP4 packed bytes for 'fp4', MX-FP8 e4m3fn for 'a8w4'). - w2_flat_fp32 = w2_fp32.view(experts * model_dim, inter_dim) - w2_fp4, w2_scale_raw = _per_1x32_fp4_quant(w2_flat_fp32) - w2_q = w2_fp4 - scale_w2 = w2_scale_raw - # x_q, w1_q, scale_x, scale_w1 not used for stage2 (A2 comes from a2_fp8_in) - x_q = None - w1_q = None - scale_x = None - scale_w1 = None else: # W4A8: A2 is int8, W2 is int4 packed (host packs from int8 values in [-8,7]). x_q, scale_x = pertoken_quant(x_fp32, quant_dtype=torch.int8) @@ -1132,119 +1062,80 @@ def run_moe_stage2( # Override per-row scale (kernel uses groupwise scale instead). scale_w2 = None - if is_fp4_path: - # FP4 / A8W4: preshuffle W2 and prepare MXFP4 weight scales - w2_shuffled = shuffle_weight(w2_q.view(torch.float4_e2m1fn_x2)) - w2_kernel = w2_shuffled.view(torch.uint8).contiguous() - w2_scale_1d = gemm_common_utils.e8m0_shuffle(scale_w2).view(torch.uint8).contiguous() - - # A2 input: provided by the caller. For 'fp4' it is packed MXFP4 - # [tokens*topk, inter_dim//2]; for 'a8w4' it is MX-FP8 e4m3fn bytes - # [tokens*topk, inter_dim]. The companion scale is per-1x32 E8M0 - # [tokens*topk, inter_dim//32] for both paths. - if a2_fp8_in is not None and a2_scale_in is not None: - a2_q = a2_fp8_in - a2_scale_raw = a2_scale_in - a2_scale = a2_scale_raw - else: - raise RuntimeError( - f"run_moe_stage2(in_dtype={in_dtype!r}) requires a2_fp8_in and a2_scale_in " - "(A2 must be quantized externally for the FP4/A8W4 path)." - ) - # Sort A2 scale by MoE routing order (dtype-agnostic on the E8M0 bytes). - a2_scale_1d = ( - gemm_common_utils.moe_mxfp4_sort( - a2_scale_raw.view(tokens, topk, -1), - sorted_ids=sorted_token_ids, - num_valid_ids=num_valid_ids, - token_num=tokens, - block_size=tile_m, - ) - .view(torch.uint8) - .contiguous() - ) - # The kernel consumes A2 as a flat byte stream. For 'fp4' that's the - # packed fp4x2 bytes; for 'a8w4' the e4m3fn bytes. Keep a2_q as its - # logical dtype for the reference path below. - a2_q_kernel = a2_q.view(torch.uint8).contiguous() + # Preshuffle weights on the *unpacked* tensor. + w1_shuffled = shuffle_weight(w1_q) + w2_shuffled = shuffle_weight(w2_q) - w1_shuffled = None - sorted_weights_1d = sorted_weights.contiguous().view(-1) + # Stage2 input (A2): either provided (gemm1->quantize chaining) or built from stage1 reference. + # For int4_bf16, A2 is bf16 (same as fp16 for scale handling). + if a2_fp8_in is not None and (a2_scale_in is not None or in_dtype in ("fp16", "bf16", "int4_bf16")): + a2_q = a2_fp8_in + a2_scale = a2_scale_in else: - # Preshuffle weights on the *unpacked* tensor. - w1_shuffled = shuffle_weight(w1_q) - w2_shuffled = shuffle_weight(w2_q) - - # Stage2 input (A2): either provided (gemm1->quantize chaining) or built from stage1 reference. - # For int4_bf16, A2 is bf16 (same as fp16 for scale handling). - if a2_fp8_in is not None and (a2_scale_in is not None or in_dtype in ("fp16", "bf16", "int4_bf16")): - a2_q = a2_fp8_in - a2_scale = a2_scale_in - else: - w1_q_flat = w1_q.view(experts * (2 * inter_dim), model_dim) - scale_w1_flat = None if scale_w1 is None else scale_w1.view(experts * (2 * inter_dim), 1) - # Build stage2 input via reference stage1 only when correctness is enabled. - if bool(skip_ref): - raise RuntimeError( - "run_moe_stage2(skip_ref=True) requires providing a2_fp8_in and a2_scale_in " - "(so we don't have to run the huge torch reference stage1)." - ) - out1_ref = torch_moe_gemm1( - x_q, - w1_q_flat, - scale_x, - scale_w1_flat, - topk_ids.to(torch.int64), - topk_weights, - inter_dim=inter_dim, - doweight_stage1=bool(doweight_stage1), - ) # [tokens, topk, inter] fp32 - if in_dtype == "fp8": - a2_q, a2_scale = pertoken_quant(out1_ref, quant_dtype=DTYPE_FP8) - elif in_dtype == "fp16": - a2_q = out1_ref.to(torch.float16) - a2_scale = None - elif in_dtype == "bf16": - a2_q = out1_ref.to(torch.bfloat16) - a2_scale = None - elif in_dtype == "int4_bf16": - # W4A16: A2 is bf16 (no quant). - a2_q = out1_ref.to(torch.bfloat16) - a2_scale = None - else: - if is_int8smooth: - # Apply a per-expert smooth scale to A2 before W8A8 quantization. - smooth_scale2 = 0.75 + 0.5 * torch.rand((experts, inter_dim), device=device, dtype=torch.float32) - out1_ref = out1_ref * smooth_scale2[topk_ids.to(torch.int64)] - a2_q, a2_scale = pertoken_quant(out1_ref, quant_dtype=torch.int8) - - # Flatten weights/scales for the kernel. - w2_shuffled_flat = w2_shuffled.view(experts * model_dim, inter_dim) - scale_w2_flat = None if scale_w2 is None else scale_w2.view(experts * model_dim, 1) - - # For W4A8 and W4A16, pack preshuffled int8 weights into packed int4 bytes. - # Both use the same interleaved packing: [ (v4<<4)|v0, (v5<<4)|v1, (v6<<4)|v2, (v7<<4)|v3 ] - w2_kernel = w2_shuffled_flat - if use_packed_int4: - w2_kernel = _pack_shuffled_int8_to_packed_int4_no_perm(w2_shuffled_flat) - - w2_flat = w2_kernel.contiguous().view(-1) - w2_kernel = w2_flat - if not use_packed_int4: - w2_kernel = w2_kernel.view(experts * model_dim, inter_dim) - - # Flatten scales to 1D memrefs (fp16 path uses 0-sized scale tensors; kernel ignores them). - if a2_scale is None: - a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) - else: - a2_scale_1d = a2_scale.view(-1).contiguous() # [tokens*topk] - if use_groupwise_scale: - w2_scale_1d = scale_w2_prepared.view(-1).contiguous() - elif scale_w2_flat is None: - w2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) + w1_q_flat = w1_q.view(experts * (2 * inter_dim), model_dim) + scale_w1_flat = None if scale_w1 is None else scale_w1.view(experts * (2 * inter_dim), 1) + # Build stage2 input via reference stage1 only when correctness is enabled. + if bool(skip_ref): + raise RuntimeError( + "run_moe_stage2(skip_ref=True) requires providing a2_fp8_in and a2_scale_in " + "(so we don't have to run the huge torch reference stage1)." + ) + out1_ref = torch_moe_gemm1( + x_q, + w1_q_flat, + scale_x, + scale_w1_flat, + topk_ids.to(torch.int64), + topk_weights, + inter_dim=inter_dim, + doweight_stage1=bool(doweight_stage1), + ) # [tokens, topk, inter] fp32 + if in_dtype == "fp8": + a2_q, a2_scale = pertoken_quant(out1_ref, quant_dtype=DTYPE_FP8) + elif in_dtype == "fp16": + a2_q = out1_ref.to(torch.float16) + a2_scale = None + elif in_dtype == "bf16": + a2_q = out1_ref.to(torch.bfloat16) + a2_scale = None + elif in_dtype == "int4_bf16": + # W4A16: A2 is bf16 (no quant). + a2_q = out1_ref.to(torch.bfloat16) + a2_scale = None else: - w2_scale_1d = scale_w2_flat.view(-1).contiguous() # [experts*model_dim] - sorted_weights_1d = sorted_weights.contiguous().view(-1) # [sorted_size] + if is_int8smooth: + # Apply a per-expert smooth scale to A2 before W8A8 quantization. + smooth_scale2 = 0.75 + 0.5 * torch.rand((experts, inter_dim), device=device, dtype=torch.float32) + out1_ref = out1_ref * smooth_scale2[topk_ids.to(torch.int64)] + a2_q, a2_scale = pertoken_quant(out1_ref, quant_dtype=torch.int8) + + # Flatten weights/scales for the kernel. + w2_shuffled_flat = w2_shuffled.view(experts * model_dim, inter_dim) + scale_w2_flat = None if scale_w2 is None else scale_w2.view(experts * model_dim, 1) + + # For W4A8 and W4A16, pack preshuffled int8 weights into packed int4 bytes. + # Both use the same interleaved packing: [ (v4<<4)|v0, (v5<<4)|v1, (v6<<4)|v2, (v7<<4)|v3 ] + w2_kernel = w2_shuffled_flat + if use_packed_int4: + w2_kernel = _pack_shuffled_int8_to_packed_int4_no_perm(w2_shuffled_flat) + + w2_flat = w2_kernel.contiguous().view(-1) + w2_kernel = w2_flat + if not use_packed_int4: + w2_kernel = w2_kernel.view(experts * model_dim, inter_dim) + + # Flatten scales to 1D memrefs (fp16 path uses 0-sized scale tensors; kernel ignores them). + if a2_scale is None: + a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) + else: + a2_scale_1d = a2_scale.view(-1).contiguous() # [tokens*topk] + if use_groupwise_scale: + w2_scale_1d = scale_w2_prepared.view(-1).contiguous() + elif scale_w2_flat is None: + w2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) + else: + w2_scale_1d = scale_w2_flat.view(-1).contiguous() # [experts*model_dim] + sorted_weights_1d = sorted_weights.contiguous().view(-1) # [sorted_size] out_s = str(out_dtype).strip().lower() if out_s in ("f16", "fp16", "half"): @@ -1261,126 +1152,67 @@ def run_moe_stage2( doweight_stage2 = not bool(doweight_stage1) - if is_fp4_path: - fp4_accumulate = not bool(use_reduce) - a_dtype_kernel = "fp8" if is_a8w4 else "fp4" - exe = 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=bool(doweight_stage2), - a_dtype=a_dtype_kernel, - b_dtype="fp4", - out_dtype="f16", - accumulate=fp4_accumulate, - ) - bias_dummy = torch.empty((0,), device=device, dtype=torch.float32) - - if bool(use_reduce): - - def _s2_args_fp4_interm(interm, x, w, sx, sw, st, eids, sw_sorted): - return ( - interm.view(-1), - x, - w, - sx, - sw, - st, - eids, - sw_sorted, - num_valid_ids, - bias_dummy, - tokens, - model_dim, - inter_dim, - int(blocks), - torch.cuda.current_stream(), - ) - - _dummy_interm = torch.empty(tokens * topk, model_dim, device=device, dtype=torch.float16) - compiled_exe = flyc.compile( - exe, - *_s2_args_fp4_interm( - _dummy_interm, - a2_q_kernel.view(-1), - w2_kernel.view(-1), - a2_scale_1d, - w2_scale_1d, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - ), - ) - - def launch(o, x, w, sx, sw, st, eids, sw_sorted): - intermediate = torch.empty(tokens * topk, model_dim, device=device, dtype=torch.float16) - compiled_exe(*_s2_args_fp4_interm(intermediate, x, w, sx, sw, st, eids, sw_sorted)) - X = intermediate.view(tokens, topk, model_dim) - torch.sum(X, dim=1, out=o.view(tokens, model_dim)) - - else: - - def _s2_args_fp4(o, x, w, sx, sw, st, eids, sw_sorted): - return ( - o, - x, - w, - sx, - sw, - st, - eids, - sw_sorted, - num_valid_ids, - bias_dummy, - tokens, - model_dim, - inter_dim, - int(blocks), - torch.cuda.current_stream(), - ) - - compiled_exe = flyc.compile( - exe, - *_s2_args_fp4( - out_perf, - a2_q_kernel.view(-1), - w2_kernel.view(-1), - a2_scale_1d, - w2_scale_1d, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - ), - ) - - def launch(o, x, w, sx, sw, st, eids, sw_sorted): - compiled_exe(*_s2_args_fp4(o, x, w, sx, sw, st, eids, sw_sorted)) - - else: - exe = compile_fn( - model_dim=model_dim, - inter_dim=inter_dim, - experts=experts, - topk=topk, - in_dtype=in_dtype, - out_dtype=out_dtype, - group_size=group_size, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage2=bool(doweight_stage2), - scale_is_bf16=(scale_dtype == "bf16"), - ) + exe = compile_fn( + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + in_dtype=in_dtype, + out_dtype=out_dtype, + group_size=group_size, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage2=bool(doweight_stage2), + scale_is_bf16=(scale_dtype == "bf16"), + ) is_reduce_exe = (getattr(exe, "mode", None) == MoeGemm2Mode.REDUCE) or bool(use_reduce) - if not is_fp4_path: + def _s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted): + return ( + o, + x, + w, + sx, + sw, + st, + eids, + sw_sorted, + num_valid_ids, + tokens, + model_dim, + inter_dim, + int(blocks), + torch.cuda.current_stream(), + ) - def _s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted): - return ( + # In reduce mode, exe is a _MoeGemm2ReduceWrapper (not a JitFunction), + # so flyc.compile is not applicable. The wrapper internally dispatches + # to two separate JitFunctions (gemm2 + reduce). + if not is_reduce_exe and hasattr(flyc, "compile"): + compiled_exe = flyc.compile( + exe, + *_s2_args_atomic( + out_perf, + a2_q.view(-1), + w2_kernel.view(-1), + a2_scale_1d, + w2_scale_1d, + sorted_token_ids, + sorted_expert_ids, + sorted_weights_1d, + ), + ) + elif not is_reduce_exe: + compiled_exe = exe + + def launch(o, x, w, sx, sw, st, eids, sw_sorted): + if is_reduce_exe: + stream = torch.cuda.current_stream() + valid_mask = None + if bool(use_valid_mask): + valid_mask = get_topk_valid_mask(topk_ids, expert_mask=None).contiguous() + exe( o, x, w, @@ -1394,62 +1226,17 @@ def _s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted): model_dim, inter_dim, int(blocks), - torch.cuda.current_stream(), + valid_mask, + stream, ) - - # In reduce mode, exe is a _MoeGemm2ReduceWrapper (not a JitFunction), - # so flyc.compile is not applicable. The wrapper internally dispatches - # to two separate JitFunctions (gemm2 + reduce). - if not is_reduce_exe and hasattr(flyc, "compile"): - compiled_exe = flyc.compile( - exe, - *_s2_args_atomic( - out_perf, - a2_q.view(-1), - w2_kernel.view(-1), - a2_scale_1d, - w2_scale_1d, - sorted_token_ids, - sorted_expert_ids, - sorted_weights_1d, - ), - ) - elif not is_reduce_exe: - compiled_exe = exe - - def launch(o, x, w, sx, sw, st, eids, sw_sorted): - if is_reduce_exe: - stream = torch.cuda.current_stream() - valid_mask = None - if bool(use_valid_mask): - valid_mask = get_topk_valid_mask(topk_ids, expert_mask=None).contiguous() - exe( - o, - x, - w, - sx, - sw, - st, - eids, - sw_sorted, - num_valid_ids, - tokens, - model_dim, - inter_dim, - int(blocks), - valid_mask, - stream, - ) + else: + if hasattr(flyc, "compile"): + compiled_exe(*_s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted)) else: - if hasattr(flyc, "compile"): - compiled_exe(*_s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted)) - else: - exe(*_s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted)) + exe(*_s2_args_atomic(o, x, w, sx, sw, st, eids, sw_sorted)) - # Flat byte view of A2 consumed by the kernel. For the FP4/A8W4 path we - # already built a dedicated uint8 `a2_q_kernel`; for the other paths the - # kernel accepts the original a2_q dtype directly. - a2_launch_buf = (a2_q_kernel if is_fp4_path else a2_q).view(-1) + # Flat byte view of A2 consumed by the kernel. + a2_launch_buf = a2_q.view(-1) # NOTE: stage2 uses atomic-add into `out`, so we cannot reuse the same output buffer # across perf iterations for correctness. Time into a dedicated buffer, then run @@ -1523,8 +1310,6 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): "int8smooth": (8, 8, "per_token_f32", "per_row_f32"), "int4": (8, 4, "per_token_f32", "per_row_f32"), "int4_bf16": (16, 4, "none", "per_row_f32"), - "fp4": (4, 4, "per_block_e8m0_32", "per_block_e8m0_32"), - "a8w4": (8, 4, "per_block_e8m0_32", "per_block_e8m0_32"), } a_bits, w_bits, a_scale_mode, w_scale_mode = _BYTES_SPEC[in_dtype] if use_groupwise_scale: @@ -1658,6 +1443,335 @@ def launch_ck(o, a2_, w1_, w2_, sorted_ids_, sorted_eids_, num_valid_, w2_scale_ return None +def _print_mxfp_moe_perf( + stage, us, label, tokens, topk, model_dim, inter_dim, experts, *, a_dtype="fp4", use_reduce=False +): + """Emit the stage1/stage2 timing line the benchmark harness (run_benchmark.sh) greps for. + + Mirrors the inline ``FlyDSL MoE stage{1,2}`` prints in run_moe_stage1/run_moe_stage2 so + the fused a4w4/a8w4 path reports through the same table rows. + """ + active = min(experts, tokens * topk) # only routed experts move weights/scales + if stage == 1: + flops = 2 * tokens * topk * (2 * inter_dim) * model_dim + a_bits = 8 if a_dtype == "fp8" else 4 # MX-FP8 activation for a8w4, else MX-FP4 + x_elems = tokens * model_dim + w_elems = active * (2 * inter_dim) * model_dim + # A + W1 (MX-FP4) + per-1x32 E8M0 scales + sorted fp4 intermediate out. + nbytes = (x_elems * a_bits) // 8 + (w_elems * 4) // 8 + x_elems // 32 + w_elems // 32 + nbytes += tokens * topk * inter_dim // 2 + else: + flops = 2 * tokens * topk * model_dim * inter_dim + a2_elems = tokens * topk * inter_dim + w2_elems = active * model_dim * inter_dim + # A2 (MX-FP4) + W2 (MX-FP4) + per-1x32 E8M0 scales + bf16 out. + nbytes = a2_elems // 2 + (w2_elems * 4) // 8 + a2_elems // 32 + w2_elems // 32 + nbytes += tokens * model_dim * 2 + tflops = float("nan") if us <= 0 else flops / (us / 1e6) / 1e12 + tbps = float("nan") if us <= 0 else nbytes / 1e12 / (us / 1e6) + if stage == 1: + print( + f"FlyDSL MoE stage1[{label}]: " + f"{us:.1f} us, {tflops:.2f} TFLOPS(logical, M={tokens*topk}), {tbps:.3f} TB/s" + ) + else: + print( + f"FlyDSL MoE stage2 [mxfp_moe] {label} {'reduce' if use_reduce else 'atomic'} | " + f"{model_dim}x{inter_dim}, E={experts}, K={topk}, M_eff={tokens*topk} | " + f"{us:.1f} us, {tflops:.2f} TFLOPS, {tbps:.3f} TB/s" + ) + + +def _run_mxfp_moe_e2e( + *, + tokens: int, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + tile_m: int, + use_reduce: bool, + x_fp32: torch.Tensor, + w1_fp32: torch.Tensor, + w2_fp32: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + routing: RoutingBuffers, + a_dtype: str = "fp4", + inline_quant: bool = False, + interleave: bool = False, + skip_ref: bool = False, + num_iters: int = 0, + num_warmup: int = 0, + in_dtype_label: Optional[str] = None, +): + """End-to-end a4w4 / a8w4 correctness via the fused mxfp_moe pipeline. + + ``a_dtype`` selects the stage1 activation: "fp4" (a4w4, MX-FP4 A) or "fp8" + (a8w4, MX-FP8 e4m3 A). W1/W2 are always MX-FP4 and the stage1->stage2 + intermediate is re-quantized to fp4 on-device, so stage2 is identical for both. + Stage1 (fused gate+up GEMM + SiLU + on-device fp4 re-quant) writes a sorted + fp4 intermediate consumed directly by stage2 (down-proj). Compared against the + torch reference (torch_moe_gemm1 -> host re-quant -> torch_moe_gemm2). + """ + from tests.kernels.utils import gemm_common_utils as gcu + + dev = x_fp32.device + BM = int(tile_m) + N_OUT = 2 * inter_dim + sorted_token_ids, sorted_weights, sorted_expert_ids, num_valid_ids, sorted_size, _blocks = routing + sorted_token_ids = sorted_token_ids.to(dev) + sorted_weights = sorted_weights.to(dev) + sorted_expert_ids = sorted_expert_ids.to(dev) + num_valid_ids = num_valid_ids.to(dev) + + # --- quantize activations / weights (per-1x32 e8m0) ------------------------ + # A: MX-FP8 (e4m3, 1 B/elem) for a8w4; MX-FP4 (0.5 B/elem) for a4w4. W is MX-FP4. + if a_dtype == "fp8": + x_q, x_scale = _per_1x32_mxfp8_quant(x_fp32) # [T, K] fp8, [T, K/32] u8 + else: + x_q, x_scale = _per_1x32_fp4_quant(x_fp32) # [T, K/2] u8, [T, K/32] u8 + w1_q, w1_scale = _per_1x32_fp4_quant(w1_fp32.reshape(experts * N_OUT, model_dim)) + w2_q, w2_scale = _per_1x32_fp4_quant(w2_fp32.reshape(experts * model_dim, inter_dim)) + + # --- fused-kernel host contract ------------------------------------------- + # cumsum[0] == total padded sorted rows (kernel derives m-blocks = cumsum[0]//BM). + cumsum = num_valid_ids.to(torch.int32).contiguous() + # A rows are gathered by unpacked token id (padding entries -> OOB, clamped to 0). + m_indices = (sorted_token_ids & 0x00FFFFFF).to(torch.int32).contiguous() + + w1_shuf = shuffle_weight(w1_q.view(torch.float4_e2m1fn_x2)).view(torch.uint8).contiguous() + w1_scale_1d = gcu.e8m0_shuffle(w1_scale.view(experts * N_OUT, model_dim // 32)).view(torch.uint8).contiguous() + if inline_quant: + # inline_quant computes A + A-scale on-device from bf16 hidden; the sorted + # A-scale (and a_quant) are unused, so skip the moe sort (also undefined at BM<32). + x_scale_sort = x_scale.view(torch.uint8).contiguous() + else: + x_scale_sort = ( + gcu.moe_mxfp4_sort( + x_scale[:tokens].view(tokens, 1, -1), + sorted_ids=sorted_token_ids, + num_valid_ids=num_valid_ids, + token_num=tokens, + block_size=BM, + ) + .view(torch.uint8) + .contiguous() + ) + + # sorted fp4 intermediate (stage1 output -> stage2 input) + scale_cols = inter_dim // 32 + padded_rows = (sorted_size + 255) // 256 * 256 + padded_cols = (scale_cols + 7) // 8 * 8 + aqout = torch.zeros(sorted_size, inter_dim // 2, dtype=torch.uint8, device=dev) + ascaleout = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=dev) + # inline_quant reads bf16 hidden and quantizes A on-device (a_quant/a_scale unused); + # otherwise A is pre-quantized above and hidden is a dummy. + if inline_quant: + hidden = x_fp32.to(torch.bfloat16).contiguous() + else: + hidden = torch.zeros(tokens, model_dim, dtype=torch.bfloat16, device=dev) + + # gemm1 uses the non-temporal weight load at BM == 32 and for inline (BM == 16); + # larger cached tiles reuse weights across m-blocks. + g1_use_nt = True if inline_quant else (BM == 32) + + def _g1_launch(): + flydsl_mxfp4_gemm1( + a_quant=x_q.view(torch.uint8).contiguous(), + a_scale_sorted_shuffled=x_scale_sort, + w1_u8=w1_shuf, + w1_scale_u8=w1_scale_1d, + sorted_expert_ids=sorted_expert_ids, + cumsum_tensor=cumsum, + m_indices=m_indices, + inter_sorted_quant=aqout, + inter_sorted_shuffled_scale=ascaleout, + hidden_states=hidden, + n_tokens=tokens, + BM=BM, + use_nt=g1_use_nt, + inline_quant=inline_quant, + interleave=interleave, + NE=experts, + D_HIDDEN=model_dim, + D_INTER=inter_dim, + topk=topk, + a_dtype=a_dtype, + ) + + # gemm1 writes (not accumulates) the sorted fp4 intermediate, so repeated + # timed launches leave a valid `aqout`/`ascaleout` for stage2 below. + label = in_dtype_label or ("a8w4" if a_dtype == "fp8" else "fp4") + if num_iters > 0: + _, us1 = run_perftest(_g1_launch, num_iters=int(num_iters), num_warmup=int(num_warmup)) + _print_mxfp_moe_perf(1, us1, label, tokens, topk, model_dim, inter_dim, experts, a_dtype=a_dtype) + else: + _g1_launch() + torch.cuda.synchronize() + + # --- stage2 --------------------------------------------------------------- + w2_shuf = shuffle_weight(w2_q.view(torch.float4_e2m1fn_x2)).view(torch.uint8).contiguous() + w2_scale_1d = gcu.e8m0_shuffle(w2_scale.view(experts * model_dim, inter_dim // 32)).view(torch.uint8).contiguous() + + if use_reduce: + # Reduce (nonatomic) epilog writes flat bf16 per sorted position; reduce on + # host. The mxfp_moe stage2 mode is coupled to tile_m: reduce is a BM == 128 + # path, atomic is the BM in {16,32,64} path. (Callers pick the mode by + # tile_m; the pytest matrix skips the mismatched combo.) + if BM != 128: + pytest.skip(f"fused mxfp_moe reduce (nonatomic) epilog requires tile_m == 128, got {BM}") + flat = torch.zeros(sorted_size * model_dim, dtype=torch.bfloat16, device=dev) + + def _g2_launch(): + flydsl_mxfp4_gemm2( + inter_sorted_quant=aqout, + inter_sorted_shuffled_scale=ascaleout, + w2_u8=w2_shuf, + w2_scale_u8=w2_scale_1d, + sorted_expert_ids=sorted_expert_ids, + cumsum_tensor=cumsum, + sorted_token_ids=sorted_token_ids, + sorted_weights=sorted_weights, + flat_out=flat, + M_logical=tokens, + max_sorted=sorted_size, + BM=BM, + use_nt=False, + epilog="nonatomic", + NE=experts, + D_HIDDEN=model_dim, + D_INTER=inter_dim, + topk=topk, + ) + + # nonatomic overwrites `flat`, so a timed loop leaves a valid result. + if num_iters > 0: + _, us2 = run_perftest(_g2_launch, num_iters=int(num_iters), num_warmup=int(num_warmup)) + _print_mxfp_moe_perf(2, us2, label, tokens, topk, model_dim, inter_dim, experts, use_reduce=True) + else: + _g2_launch() + torch.cuda.synchronize() + flat = flat.view(sorted_size, model_dim).float() + tok = (sorted_token_ids & 0x00FFFFFF).long() + valid = tok < tokens + out = torch.zeros(tokens, model_dim, dtype=torch.float32, device=dev) + out.index_add_(0, tok[valid], flat[valid] * sorted_weights[valid].unsqueeze(-1)) + else: + # The atomic (scatter-to-token) epilog is supported at BM in {16, 32, 64}. + # BM == 128 down-proj is covered by the reduce (nonatomic) path instead. + if BM == 128: + pytest.skip("fused mxfp_moe atomic epilog unsupported at tile_m == 128; use reduce mode") + out_buf = torch.zeros(tokens * model_dim, dtype=torch.bfloat16, device=dev) + + def _g2_launch(): + flydsl_mxfp4_gemm2( + inter_sorted_quant=aqout, + inter_sorted_shuffled_scale=ascaleout, + w2_u8=w2_shuf, + w2_scale_u8=w2_scale_1d, + sorted_expert_ids=sorted_expert_ids, + cumsum_tensor=cumsum, + sorted_token_ids=sorted_token_ids, + sorted_weights=sorted_weights, + flat_out=out_buf, + M_logical=tokens, + max_sorted=sorted_size, + BM=BM, + use_nt=True, + epilog="atomic", + NE=experts, + D_HIDDEN=model_dim, + D_INTER=inter_dim, + topk=topk, + ) + + # atomic accumulates into out_buf; time into it, then zero + one clean + # launch for the correctness check below. + if num_iters > 0: + _, us2 = run_perftest(_g2_launch, num_iters=int(num_iters), num_warmup=int(num_warmup)) + _print_mxfp_moe_perf(2, us2, label, tokens, topk, model_dim, inter_dim, experts, use_reduce=False) + out_buf.zero_() + _g2_launch() + torch.cuda.synchronize() + out = out_buf.view(tokens, model_dim).float() + + # --- reference (stage1 -> host re-quant -> stage2) ------------------------ + if not skip_ref: + ref1 = torch_moe_gemm1( + x_q, w1_q, x_scale, w1_scale, topk_ids.long(), topk_weights, inter_dim=inter_dim, doweight_stage1=False + ) + a2_q, a2_scale = _per_1x32_fp4_quant(ref1.reshape(tokens * topk, inter_dim)) + ref2 = torch_moe_gemm2( + a2_q.view(tokens, topk, -1), + w2_q, + a2_scale.view(tokens, topk, -1), + w2_scale, + topk_ids.long(), + topk_weights, + model_dim=model_dim, + doweight_stage2=True, + ) + verify_output(out, ref2, rtol=0.5, atol=0.5, logits_diff_threshold=1) + + +@pytest.mark.skipif("gfx95" not in ARCH, reason="mxfp_moe a4w4/a8w4 requires gfx950+") +@pytest.mark.parametrize("a_dtype", ["fp4", "fp8"]) +@pytest.mark.parametrize( + "variant", + ["bm32_atomic", "inline_bm16", "interleave_bm64"], +) +def test_mxfp_moe_variants(a_dtype, variant): + """Cover the fused mxfp_moe gemm1 variants the FP4-M/L e2e shapes don't reach: + BM==32 (atomic), inline_quant (BM==16, bf16 hidden -> on-device quant), and the + interleaved gate/up layout. Both a4w4 (fp4) and a8w4 (fp8) activations.""" + device = torch.device("cuda") + tokens, model_dim, inter_dim, experts, topk = 128, 1024, 256, 8, 2 + s = 0.2 + x_fp32 = torch.randn((tokens, model_dim), device=device, dtype=torch.float32) * s + w1_fp32 = torch.randn((experts, 2 * inter_dim, model_dim), device=device, dtype=torch.float32) * s + w2_fp32 = torch.randn((experts, model_dim, inter_dim), device=device, dtype=torch.float32) * ( + s / math.sqrt(inter_dim) + ) + score = torch.rand((tokens, experts), device=device, dtype=torch.float32) + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) + + if variant == "inline_bm16": + tile_m, inline_quant, interleave = 16, True, False + elif variant == "bm32_atomic": + tile_m, inline_quant, interleave = 32, False, False + else: # interleave_bm64 + tile_m, inline_quant, interleave = 64, False, True + + routing = build_routing_buffers( + topk_ids=topk_ids, + topk_weights=topk_weights, + experts=experts, + model_dim=model_dim, + tile_m=tile_m, + moe_sort_mode="torch", + ) + _run_mxfp_moe_e2e( + tokens=tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + use_reduce=False, + x_fp32=x_fp32, + w1_fp32=w1_fp32, + w2_fp32=w2_fp32, + topk_ids=topk_ids, + topk_weights=topk_weights, + routing=routing, + a_dtype=a_dtype, + inline_quant=inline_quant, + interleave=interleave, + ) + + @pytest.mark.parametrize( "tokens, model_dim, inter_dim, experts, topk, tile_m, tile_n1, tile_k1, tile_n2, tile_k2, doweight_stage1", [ @@ -1730,6 +1844,7 @@ def launch_ck(o, a2_, w1_, w2_, sorted_ids_, sorted_eids_, num_valid_, w2_scale_ "int4", "int4_bf16", pytest.param("fp4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+")), + pytest.param("a8w4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="A8W4 requires gfx950+")), ], ) @pytest.mark.parametrize("out_dtype", ["f16", "bf16", "f32"], ids=["out_f16", "out_bf16", "out_f32"]) @@ -1796,6 +1911,11 @@ def test_moe_gemm_2stage( pytest.skip(f"{in_dtype} stage2 requires inter_dim >= 256 and tile_k2 >= 256, got {inter_dim}, {tile_k2}") if tile_m < 32 or tile_m % 32 != 0: pytest.skip(f"{in_dtype} requires tile_m % 32 == 0 and tile_m >= 32, got {tile_m}") + # The fused mxfp_moe gemm1 unrolls the K loop as prologue(kStages=2) + main + + # drain; with model_dim == kStages*tile_k (512) the main loop that inits + # the accumulator is empty, so require model_dim > 512 (FP4-S is skipped). + if model_dim <= 512: + pytest.skip(f"fused mxfp_moe gemm1 requires model_dim > 512, got {model_dim}") device = torch.device("cuda") # torch.manual_seed(int(seed)) @@ -1838,6 +1958,31 @@ def test_moe_gemm_2stage( if compare_aiter_ck is None: compare_aiter_ck = False + if in_dtype in ("fp4", "a8w4"): + # a4w4 / a8w4 drive the fused mxfp_moe pipeline end-to-end (device-side + # re-quant, sorted fp4 intermediate) rather than the retired mixed path. + _run_mxfp_moe_e2e( + tokens=tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + tile_m=tile_m, + use_reduce=bool(use_reduce), + x_fp32=x_fp32, + w1_fp32=w1_fp32, + w2_fp32=w2_fp32, + topk_ids=topk_ids, + topk_weights=topk_weights, + routing=routing, + a_dtype="fp8" if in_dtype == "a8w4" else "fp4", + skip_ref=bool(skip_ref), + num_iters=num_iters, + num_warmup=num_warmup, + in_dtype_label=in_dtype, + ) + return + out1_fp16, _us1 = run_moe_stage1( tokens=tokens, model_dim=model_dim, @@ -1867,15 +2012,7 @@ def test_moe_gemm_2stage( test_graph=test_graph, ) - if in_dtype in ("fp4", "a8w4"): - # Re-quantize stage1 output for stage2 input: - # fp4 -> MX-FP4 (0.5 B/elem, packed) - # a8w4 -> MX-FP8 e4m3fn (1 B/elem) - # run_moe_stage2 sorts the raw E8M0 scale [tokens*topk, inter_dim//32] internally. - out1_fp32 = out1_fp16.to(torch.float32).view(tokens * topk, inter_dim) - quantize_a2 = _per_1x32_mxfp8_quant if in_dtype == "a8w4" else _per_1x32_fp4_quant - a2_q, a2_scale = quantize_a2(out1_fp32) - elif w_fp4_kernel: + if w_fp4_kernel: a2_q = out1_fp16.to(torch.float32) a2_scale = None elif in_dtype == "fp8": @@ -2304,10 +2441,6 @@ def test_moe_stage2_standalone( """ is_fp4_path = in_dtype in ("fp4", "a8w4") if is_fp4_path: - from tests.kernels.utils import gemm_common_utils - - if gemm_common_utils is None: - pytest.skip("FP4 dependencies not available (triton/mixed_moe_gemm not installed)") if "gfx95" not in ARCH: pytest.skip(f"{in_dtype} requires gfx950+, got {ARCH}") if inter_dim < 256 or tile_k < 256: @@ -2336,19 +2469,14 @@ def test_moe_stage2_standalone( ) if is_fp4_path: - # FP4 / A8W4 require pre-quantized a2 (stage1 output); torch reference - # cannot synthesize FP4/FP8 packed activations, so we build them here. - _A2_QUANTIZERS = { - "fp4": _per_1x32_fp4_quant, - "a8w4": _per_1x32_mxfp8_quant, - } - device = torch.device("cuda") - torch.manual_seed(seed) - a2_fp32 = torch.randn((tokens * topk, inter_dim), device=device, dtype=torch.float32) * 0.2 - a2_q, a2_scale = _A2_QUANTIZERS[in_dtype](a2_fp32) - fp4_args = dict(common_args, a2_fp8_in=a2_q, a2_scale_in=a2_scale) - run_moe_stage2(**fp4_args, kernel_name=f"moe_gemm2_atomic_{in_dtype}") - run_moe_stage2(**fp4_args, use_reduce=True, kernel_name=f"moe_gemm2_reduce_torch_{in_dtype}") + # a4w4/a8w4 stage2 is only defined within the fused mxfp_moe pipeline; + # run_moe_stage2 routes these dtypes through the e2e helper. The fused + # stage2 mode is tile_m-coupled: atomic for tile_m < 128, reduce at 128. + run_moe_stage2( + **common_args, + use_reduce=(tile_m == 128), + kernel_name=f"moe_gemm2_{in_dtype}", + ) return # Run baseline stage2 (atomic accumulation) @@ -2403,7 +2531,9 @@ def _str2tuple_dim(v: str) -> Tuple[int, int]: parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, - description="MoE 2-stage (FlyDSL MFMA FP8) test/benchmark (argparse subset aligned with aiter test_moe_2stage.py)", + description=( + "MoE 2-stage (FlyDSL MFMA FP8) test/benchmark " "(argparse subset aligned with aiter test_moe_2stage.py)" + ), ) parser.add_argument( "--in_dtype", @@ -2587,5 +2717,12 @@ def run_one(dt: str, use_reduce: bool): if dt in ("fp4", "a8w4") and "gfx95" not in ARCH: print(f"Skipping {dt}: requires gfx950+, got {ARCH}") continue - for use_reduce in reduce_flags: + # mxfp_moe (fp4/a8w4) stage2 mode is coupled to tile_m: atomic for tile_m<128, + # reduce (nonatomic) only at tile_m==128. Run the one applicable mode rather + # than the fp8-style atomic/reduce sweep (which would pick an invalid combo). + if dt in ("fp4", "a8w4"): + dt_reduce_flags = [int(args.tile_m) == 128] + else: + dt_reduce_flags = reduce_flags + for use_reduce in dt_reduce_flags: run_one(dt, use_reduce)