MoE: add a16w mix fused 2-stage kernels (bf16 A × mxfp4/int4/bf16 W) (#948)

Author: coderfeli

kernels/moe/moe_2stage_a16wmix/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+"""Fused a16w-mix (bf16 A x mxfp4/int4 W) 2-stage MoE kernels.
+
+Kernel builders live in :mod:`gemm1` / :mod:`gemm2` (shared helpers in
+:mod:`utils`). Host launch/tile-config/CSV glue is a test-side concern and
+lives in ``tests/kernels/moe_a16wmix_host.py``.
+"""

kernels/moe/moe_2stage_a16wmix/gemm1.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (C) 2025-2026 FlyDSL Project Contributors
+
+"""Shared low-level helpers for the a16w4/a16wi4/a16w16 fused MoE kernels
+(:mod:`gemm1` stage1 and :mod:`gemm2` stage2). Pointer/GEP builders, buffer-tensor
+views, e8m0/int4 dequant, the A-LDS XOR swizzle, and the arch gate."""
+
+import os
+
+import flydsl.expr as fx
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import llvm
+from flydsl.expr import arith, range_constexpr, rocdl
+from flydsl.expr.typing import T
+from flydsl.runtime.device import get_rocm_arch
+from kernels.common import buffer_ops
+
+_PTR3 = "!llvm.ptr<3>"
+LOG2E = 1.4426950408889634
+
+# a16wi4 (int4 W) groupwise scale: group_size = 32 == one MFMA K32 step (one ku per
+# K-group). Scale packed bf16 pairs (E, N, G//2, 2); even/odd ku selects lo/hi half.
+A16WI4_GROUP_SIZE = 32
+
+
+def a16wmix_use_k16(arch=None):
+    """True for the gfx942 (CDNA3) codepath: K=16 MFMA + scalar int4 dequant.
+
+    Arch-gate: gfx950 (CDNA4) has K=32 mfma_f32_16x16x32_bf16 + v_cvt_pk_bf16_f32;
+    gfx942 has neither and falls back to K=16 MFMA + scalar-trunc dequant.
+    ``FLYDSL_A16WMIX_FORCE_K16=1`` forces the gfx942 path (a strict ISA subset) for
+    validation on a gfx950 box.
+    """
+    if os.environ.get("FLYDSL_A16WMIX_FORCE_K16", "0") not in ("0", "", "false", "False"):
+        return True
+    if arch is None:
+        arch = get_rocm_arch() or ""
+    return "gfx95" not in str(arch)
+
+
+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)))
+
+
+def _global_i32_buffer_view(addr_i64, num_bytes):
+    # fx.copy BufferCopy atoms take soffset as an element count (not bytes); the
+    # make_layout 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 _buffer_i32_scalar_read(tiles1, idx, atom):
+    """Read one i32 dword at element ``idx`` from a ``_global_i32_buffer_tiles(..., 1)``
+    view via the layout-API BufferCopy atom (buffer_load_dword; OOB-clamped by the
+    buffer resource). ``tiles1`` is 1-dword tiles so the tile index == ``idx``.
+    """
+    r = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Int32)
+    fx.copy(atom, fx.slice(tiles1, (None, idx)), r)
+    return fx.Int32(fx.Vector(fx.memref_load_vec(r))[0])
+
+
+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 _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):
+    return _global_i32_ptr(addr_i64)[idx]
+
+
+def _e8m0_byte_to_f32(packed_i32, byte_pos):
+    shift = byte_pos * fx.Int32(8)
+    b = packed_i32.shrui(shift) & fx.Int32(0xFF)
+    return fx.Float32(_raw(b << fx.Int32(23)).bitcast(T.f32))
+
+
+def _cvt_pk_bf16_f32_se(src_a_f32, src_b_f32):
+    # Side-effecting v_cvt_pk_bf16_f32 (pack 2 f32 -> 2xbf16 in i32). LOAD-BEARING:
+    # the stateless rocdl.cvt_pk_bf16_f32 gets CSE-merged/reordered across K steps in
+    # the a16wi4 gemm1 hot loop (garbage output); side_effects pins each call.
+    return llvm.inline_asm(
+        ir.IntegerType.get_signless(32),
+        [_raw(src_a_f32), _raw(src_b_f32)],
+        "v_cvt_pk_bf16_f32 $0, $1, $2",
+        "=v,v,v",
+        has_side_effects=True,
+    )
+
+
+def _int4_nibble_to_bf16x8(raw_i32, scale_f32, *, use_k16=False):
+    """int4 (signed) -> bf16 upconvert for one MFMA K32 step (8 nibbles -> v8bf16).
+
+    ``raw_i32`` holds 8 signed-int4 nibbles in bits[4n+3:4n] (same K order as the
+    mxfp4 sel 0..3 path). ``v_cvt_off_f32_i4`` reads the nibble unsigned, subtracts 8,
+    and scales the mantissa by 16, so the x16 is folded into eff = scale*16.
+    ``use_k16`` (gfx942): v_cvt_pk_bf16_f32 is gfx950-only -> scalar .to(BFloat16).
+    """
+    eff = fx.Float32(scale_f32 * fx.Float32(16.0))
+    raw_even = fx.Int32(raw_i32)
+    raw_odd = raw_even.shrui(fx.Int32(4))
+    if use_k16:
+        # gfx942 fallback: scalar f32 -> bf16 truncation (no v_cvt_pk_bf16_f32).
+        bf16s = []
+        for j in range_constexpr(4):
+            f_lo = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_even), byte_sel=j)) * eff
+            f_hi = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_odd), byte_sel=j)) * eff
+            bf16s.append(f_lo.to(fx.BFloat16))
+            bf16s.append(f_hi.to(fx.BFloat16))
+        return fx.Vector.from_elements([_raw(x) for x in bf16s], fx.BFloat16)  # v8bf16
+    # byte_sel loads (1 shift total); side-effecting pk-convert.
+    i32s = []
+    for j in range_constexpr(4):
+        f_lo = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_even), byte_sel=j)) * eff
+        f_hi = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_odd), byte_sel=j)) * eff
+        i32s.append(fx.Int32(_cvt_pk_bf16_f32_se(_raw(f_lo), _raw(f_hi))))
+    v4i32 = fx.Vector.from_elements([_raw(x) for x in i32s], fx.Int32)
+    return v4i32.bitcast(fx.BFloat16)  # v8bf16
+
+
+def _int4_nibble_to_bf16x8_raw(raw_i32, *, use_k16=False):
+    """int4 (signed) -> bf16 for one MFMA K32 step WITHOUT the groupwise scale.
+
+    Same as :func:`_int4_nibble_to_bf16x8` but emits the raw dequant weights
+    ``(nibble-8)/16`` (``v_cvt_off_f32_i4``'s native output -- no per-element
+    ``v_mul_f32``). The groupwise scale (and the folded x16) is applied ONCE per
+    K-group on the small MFMA accumulator instead (see the ``_acc_scale_int4`` path in
+    the stage1 body): for BM16 (m_repeat=1) that trades 8 per-nibble muls for 4
+    per-accumulator fmas and drops the long-lived scaled-f32 operand VGPRs.
+    ``(nibble-8)/16`` is bf16-exact (values in ``{-7/16..7/16}``).
+    ``use_k16`` (gfx942): v_cvt_pk_bf16_f32 is gfx950-only -> scalar .to(BFloat16).
+    """
+    raw_even = fx.Int32(raw_i32)
+    raw_odd = raw_even.shrui(fx.Int32(4))
+    if use_k16:
+        bf16s = []
+        for j in range_constexpr(4):
+            f_lo = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_even), byte_sel=j))
+            f_hi = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_odd), byte_sel=j))
+            bf16s.append(f_lo.to(fx.BFloat16))
+            bf16s.append(f_hi.to(fx.BFloat16))
+        return fx.Vector.from_elements([_raw(x) for x in bf16s], fx.BFloat16)  # v8bf16
+    i32s = []
+    for j in range_constexpr(4):
+        f_lo = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_even), byte_sel=j))
+        f_hi = fx.Float32(rocdl.cvt_off_f32_i4(_raw(raw_odd), byte_sel=j))
+        i32s.append(fx.Int32(_cvt_pk_bf16_f32_se(_raw(f_lo), _raw(f_hi))))
+    v4i32 = fx.Vector.from_elements([_raw(x) for x in i32s], fx.Int32)
+    return v4i32.bitcast(fx.BFloat16)  # v8bf16
+
+
+def kmchunks_for(BM):
+    return BM // 16
+
+
+def lds_acc_bytes_for(rows, BN):
+    return rows * BN * 4
+
+
+def _a16w4_swizzle_xor16(row, col_bytes, k_blocks16, *, enable=False):
+    """A-LDS bank-conflict XOR swizzle (aiter swizzle_xor16: col ^ ((row&(kb16-1))*16)).
+
+    Both the DMA write and the LDS read go through this helper so the physical layout
+    stays consistent. gemm1 keeps linear (enable=False); gemm2 enables it.
+    """
+    if not enable:
+        return col_bytes
+    rem = row & fx.Int32(k_blocks16 - 1)
+    return col_bytes ^ (rem * fx.Int32(16))