From ddfa0e279aeaac13092806c2d6cf632db295ddc5 Mon Sep 17 00:00:00 2001 From: Felix Li Date: Thu, 25 Jun 2026 09:52:40 +0000 Subject: [PATCH 1/2] feat(FlyROCDL): cache_modifier (nt) on CopyOpCDNA3BufferCopy atom Add an optional cacheModifier param (DefaultValuedParameter, default 0) to the cdna3.buffer_copy atom type, forwarded as the rocdl raw buffer load/store `cachepolicy`/aux operand (0=cached, 2=non-temporal). Lets the layout-API fx.copy express non-temporal B-weight loads, which previously required a raw buffer_load(cache_modifier=2) fallback. - CopyAtom.td: new param + optional assembly group (`<128>` still parses; `<128, cache = 2>` when set). - CDNA3/CopyAtom.cpp: thread getCacheModifier() into the load+store aux. - FlyROCDLExtension.cpp: BufferCopy get() gains cache_modifier=0 kwarg. - universal.py: BufferCopy/BufferCopy{8,16,32,64,128}b accept cache_modifier. Backward compatible (default 0 == prior behavior). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../flydsl/Dialect/FlyROCDL/IR/CopyAtom.td | 8 +++++-- lib/Bindings/Python/FlyROCDLExtension.cpp | 12 ++++++----- lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp | 8 ++++--- python/flydsl/expr/rocdl/universal.py | 21 ++++++++++++------- 4 files changed, 31 insertions(+), 18 deletions(-) diff --git a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td index 5f0d924b5..d774b7839 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td +++ b/include/flydsl/Dialect/FlyROCDL/IR/CopyAtom.td @@ -8,9 +8,13 @@ include "flydsl/Dialect/FlyROCDL/IR/Dialect.td" def FlyROCDL_CopyOpBufferCopy : FlyROCDL_StatefulCopyOp<"CopyOpCDNA3BufferCopy", "cdna3.buffer_copy",[]> { let parameters = (ins - "int32_t":$bitSize + "int32_t":$bitSize, + // aux/cache-policy bits forwarded to the rocdl raw buffer load/store + // `cachepolicy` operand. 0 = default (cached); 2 = non-temporal (nt). Mirrors + // the raw buffer_load `cache_modifier` so the layout-API copy can express nt. + DefaultValuedParameter<"int32_t", "0">:$cacheModifier ); - let assemblyFormat = "`<` $bitSize `>`"; + let assemblyFormat = "`<` $bitSize (`,` `cache` `=` $cacheModifier^)? `>`"; } def FlyROCDL_CopyOpBufferCopyLDS : FlyROCDL_StatefulCopyOp<"CopyOpCDNA3BufferCopyLDS", "cdna3.buffer_copy_lds",[]> { diff --git a/lib/Bindings/Python/FlyROCDLExtension.cpp b/lib/Bindings/Python/FlyROCDLExtension.cpp index 94f309f02..7a439f063 100644 --- a/lib/Bindings/Python/FlyROCDLExtension.cpp +++ b/lib/Bindings/Python/FlyROCDLExtension.cpp @@ -106,13 +106,15 @@ struct PyCopyOpCDNA3BufferCopyType : PyConcreteType static void bindDerived(ClassTy &c) { c.def_static( "get", - [](int32_t bitSize, DefaultingPyMlirContext context) { + [](int32_t bitSize, int32_t cacheModifier, DefaultingPyMlirContext context) { MLIRContext *ctx = unwrap(context.get()->get()); - return PyCopyOpCDNA3BufferCopyType(context->getRef(), - wrap(CopyOpCDNA3BufferCopyType::get(ctx, bitSize))); + return PyCopyOpCDNA3BufferCopyType( + context->getRef(), + wrap(CopyOpCDNA3BufferCopyType::get(ctx, bitSize, cacheModifier))); }, - "bit_size"_a, nb::kw_only(), "context"_a = nb::none(), - "Create a CopyOpCDNA3BufferCopyType with the given bit size"); + "bit_size"_a, "cache_modifier"_a = 0, nb::kw_only(), "context"_a = nb::none(), + "Create a CopyOpCDNA3BufferCopyType with the given bit size and " + "cache_modifier (0=cached, 2=non-temporal)"); } }; diff --git a/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp b/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp index 7f381e764..e8138d67e 100644 --- a/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp +++ b/lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp @@ -88,7 +88,9 @@ FailureOr CopyOpCDNA3BufferCopyType::emitAtomCallSSA(OpBuilder &builder, return arith::DivUIOp::create(builder, loc, bits, eight); }; - Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32); + // aux/cachepolicy bits for the raw buffer load/store (0=cached, 2=nt). Carried + // on the atom type so the layout-API copy can request non-temporal loads. + Value aux = arith::ConstantIntOp::create(builder, loc, getCacheModifier(), 32); ArrayAttr noAttrs; auto srcMemTy = srcTyArg ? dyn_cast(srcTyArg) : fly::MemRefType(); @@ -102,7 +104,7 @@ FailureOr CopyOpCDNA3BufferCopyType::emitAtomCallSSA(OpBuilder &builder, Value srcOff = bp.swizzleByteOffset(builder, loc); Value loaded = ROCDL::RawPtrBufferLoadOp::create(builder, loc, copyTy, srcRsrc, srcOff, soffset, - zero, noAttrs, noAttrs, noAttrs); + aux, noAttrs, noAttrs, noAttrs); if (resultTy && loaded.getType() != resultTy) loaded = LLVM::BitcastOp::create(builder, loc, resultTy, loaded); return loaded; @@ -118,7 +120,7 @@ FailureOr CopyOpCDNA3BufferCopyType::emitAtomCallSSA(OpBuilder &builder, Value stored = src; if (stored.getType() != copyTy) stored = LLVM::BitcastOp::create(builder, loc, copyTy, stored); - ROCDL::RawPtrBufferStoreOp::create(builder, loc, stored, dstRsrc, dstOff, soffset, zero, + ROCDL::RawPtrBufferStoreOp::create(builder, loc, stored, dstRsrc, dstOff, soffset, aux, noAttrs, noAttrs, noAttrs); return stored; } diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index 7e9e3b890..232e0373e 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -16,21 +16,26 @@ from ..typing import Int16, Int32, Int64, Tensor -def BufferCopy(bit_size): +def BufferCopy(bit_size, cache_modifier=0): """Create a CDNA3 buffer copy atom. + Args: + bit_size: copy width in bits. + cache_modifier: aux/cachepolicy bits forwarded to the raw buffer + load/store (0 = default cached, 2 = non-temporal / nt). + Current atom state: - `soffset` (`i32`), default zero """ - return CopyOpCDNA3BufferCopyType.get(bit_size) + return CopyOpCDNA3BufferCopyType.get(bit_size, cache_modifier) -# BufferCopy aliases for convenience -BufferCopy8b = lambda: CopyOpCDNA3BufferCopyType.get(8) -BufferCopy16b = lambda: CopyOpCDNA3BufferCopyType.get(16) -BufferCopy32b = lambda: CopyOpCDNA3BufferCopyType.get(32) -BufferCopy64b = lambda: CopyOpCDNA3BufferCopyType.get(64) -BufferCopy128b = lambda: CopyOpCDNA3BufferCopyType.get(128) +# BufferCopy aliases for convenience (optional cache_modifier: 0=cached, 2=nt) +BufferCopy8b = lambda cache_modifier=0: CopyOpCDNA3BufferCopyType.get(8, cache_modifier) +BufferCopy16b = lambda cache_modifier=0: CopyOpCDNA3BufferCopyType.get(16, cache_modifier) +BufferCopy32b = lambda cache_modifier=0: CopyOpCDNA3BufferCopyType.get(32, cache_modifier) +BufferCopy64b = lambda cache_modifier=0: CopyOpCDNA3BufferCopyType.get(64, cache_modifier) +BufferCopy128b = lambda cache_modifier=0: CopyOpCDNA3BufferCopyType.get(128, cache_modifier) def BufferCopyLDS(bit_size): From 9d18d299f6c7c304b686d8fcdc6c2760f3a219cd Mon Sep 17 00:00:00 2001 From: Felix Li Date: Fri, 26 Jun 2026 04:39:55 +0000 Subject: [PATCH 2/2] feat(moe): layout-API MXFP4 (a4w4/a8w4) MoE gemm, opus-sort only Add a layout-API MXFP4 MoE up/gate + down-proj gemm that consumes the standard (opus) sort contract from moe_sorting_kernel directly -- gemm1 gathers A rows via sorted_token_ids & 0xFFFFFF; gemm2 scatters per sorted row via global atomic add weighted by sorted_weights. No fused-sort extras (m_indices / reverse_sorted) are needed. kernels/mxfp4_moe_layout.py - layout-API building blocks (fx.copy B/B-scale, fx.gemm scaled-MFMA atoms) kernels/mxfp4_moe_common.py - shared raw helpers / constants / K-derived size formulas / atomic bf16 epilogue kernels/mxfp4_moe_gemm1.py - BM32 up/gate gemm (a4w4 + a8w4, interleave + separated, nt/cached, out fp4/fp8) kernels/mxfp4_moe_gemm2.py - BM32 atomic down-proj (a4w4 + a8w4 fp8 input) kernels/mxfp4_moe_gemm_2stage.py - public API + host launchers Wire a4w4/a8w4 of tests/kernels/test_moe_gemm.py::test_moe_gemm_2stage to the new pipeline (opus sort -> gemm1 -> gemm2 atomic) vs an independent dequant-MoE reference; a8w4 added to the in_dtype matrix. Validated on gfx950: chain cosine a4w4=0.988, a8w4=1.000 (interleave + separated); test_moe_gemm_2stage fp4/a8w4 over FP4-S/M/L pass. Co-Authored-By: Claude Opus 4.8 --- kernels/mxfp4_moe_common.py | 217 +++++++++ kernels/mxfp4_moe_gemm1.py | 784 +++++++++++++++++++++++++++++++ kernels/mxfp4_moe_gemm2.py | 522 ++++++++++++++++++++ kernels/mxfp4_moe_gemm_2stage.py | 205 ++++++++ kernels/mxfp4_moe_layout.py | 128 +++++ tests/kernels/test_moe_gemm.py | 281 +++++++++++ 6 files changed, 2137 insertions(+) create mode 100644 kernels/mxfp4_moe_common.py create mode 100644 kernels/mxfp4_moe_gemm1.py create mode 100644 kernels/mxfp4_moe_gemm2.py create mode 100644 kernels/mxfp4_moe_gemm_2stage.py create mode 100644 kernels/mxfp4_moe_layout.py diff --git a/kernels/mxfp4_moe_common.py b/kernels/mxfp4_moe_common.py new file mode 100644 index 000000000..f93825966 --- /dev/null +++ b/kernels/mxfp4_moe_common.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors +"""Shared raw helpers / constants / K-derived size formulas for the layout-API +MXFP4 MoE gemm (down-proj gemm2 + up/gate gemm1). + +Extracted verbatim from the original aiter port so both ``mxfp4_moe_gemm1`` and +``mxfp4_moe_gemm2`` share one definition of the pointer/LDS helpers, the e8m0 +scale-layout size formulas, and the atomic bf16 epilogue. The gemm bodies +themselves live in the sibling kernel modules. +""" + +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, buffer_ops, const_expr, gpu, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# -- shape constants (BM-independent; KIMI defaults, per-shape via compile args) -- +MAX_M = 655360 +NE = 385 +K = 512 # gemm2 contraction = inter_dim (DEFAULT / KIMI) +N_OUT = 7168 # default gemm2 output dim = model_dim +BN = 256 +BK = 256 +kStages = 2 +# e8m0 scale-layout K-independent stride. +kBS_stride_k0_dw = 64 + + +# -- K-derived sizes (parametrized over the contraction dim K = inter_dim) ----- +def k_half_for(k): + return k // 2 # packed-fp4 bytes along K (KIMI: 256) + + +def k_tiles_total_for(k): + return k // BK # KIMI: 2 + + +def kunroll_for(k): + # streaming main-loop trip count: kUnroll = K_TILES_TOTAL - kStages. + return k_tiles_total_for(k) - kStages + + +def kbs_c_k1_for(k): + return (k // 32) // 4 // 2 # KIMI: 2 + + +def kbs_stride_n0_dw_for(k): + return kbs_c_k1_for(k) * 64 # KIMI: 128 + + +def kas_c_k1_for(k): + return (k // 32) // 4 // 2 # KIMI: 2 + + +def kas_per_chunk_dw_for(k): + return kas_c_k1_for(k) * 64 # KIMI: 128 + + +# -- shape-parametrized sizes (NE/N_OUT vary per instance; N_OUT % 256 == 0) ---- +def num_n_blocks_for(n_out): + return n_out // 256 + + +def kbs_per_expert_dw_for(n_out, k=K): + return (n_out // 16 // 2) * kbs_stride_n0_dw_for(k) + + +def kmchunks(BM): + return 1 if const_expr(BM == 16) else BM // 16 + + +_PTR3 = "!llvm.ptr<3>" + + +def _raw(v): + """Unwrap an fx value to a raw ir.Value for raw llvm/arith ops.""" + 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 _lds_ptr3(base_i32, byte_off_i32): + """ptr<3> = inttoptr(i64(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): + """One ptr<3> for the LDS base; offsets via GEP.""" + 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): + """getelementptr i8, base_ptr, byte_off_i32 (ptr<3>).""" + return buffer_ops.get_element_ptr( + base_ptr, byte_offset=_raw(byte_off_i32), elem_type=T.i8 + ) + + +def _global_base_ptr1(addr_i64): + """One ptr<1> base from a raw i64 device address (bare data_ptr() kernarg).""" + return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(addr_i64))) + + +def _gep1(base_ptr, byte_off_i32): + """getelementptr i8, base_ptr, byte_off_i32 (ptr<1>).""" + return buffer_ops.get_element_ptr( + base_ptr, byte_offset=_raw(byte_off_i32), elem_type=T.i8 + ) + + +def _global_ptr1(arg, byte_off_i32): + return _gep1(_global_base_ptr1(arg), byte_off_i32) + + +def _lds_swizzle_mask(row): + """lds_swizzle_mask(row): mask = (row & 14) << 3.""" + return (row & fx.Int32(14)) << fx.Int32(3) + + +def _atomic_bf16_epilog( + lds_acc, + accm, + arg_out, + arg_stids, + arg_sweights, + m_row, + n_block_idx, + wave, + lane, + i32_M, + BM, + N_OUT, +): + _kMChunks = kmchunks(BM) + M_REPS = BM // 8 # BM32: 4, BM16: 2 + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + lds_base = _lds_base_ptr3(lds_acc.get()) + + 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) + + # Prefetch sorted_token_ids / sorted_weights BEFORE the cshuffle stores and + # both LDS barriers (invariant => freely hoistable), overlapping their global + # latency with the store + barriers instead of exposing it in the atomic loop. + 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 + ) + ) + + # pre-store fence+barrier (HIP run_one __syncthreads() before the epilog). + gpu.barrier() + + # write accm -> lds_acc cshuffle (scalar f32 stores, as HIP does) + 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() + + # read back + weighted atomic add (token_id / weight prefetched above) + 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): + # adjacent ee=0,1 are contiguous -> one <2xf32> load (as HIP vectorizes) + 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) # bf16 byte off + 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/mxfp4_moe_gemm1.py b/kernels/mxfp4_moe_gemm1.py new file mode 100644 index 000000000..4aba1be00 --- /dev/null +++ b/kernels/mxfp4_moe_gemm1.py @@ -0,0 +1,784 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors +"""FlyDSL **layout-API** port of aiter ``gemm1_a4w4`` (MXFP4 MoE up/gate-proj). + +Variant: ``(BM=32, inline_quant=False)`` for both ``use_nt`` and both gate modes +(kSubBlocks=1, kMChunks=2, kAStages=3, kStages=2). ``use_nt`` is the B-load cache +policy (tuned config BM32_NT vs BM32_CACHED): True -> non-temporal, False -> cached. +``interleave`` picks the gate/up layout (True=interleaved, False=separated); it only +changes the B-load column, the B-scale n-pack base, and the mfma opsel (the epilogue +gate/up split is the same for both, since accm[J] holds the same logical (g/u, n0)). + +Layout-API pieces (the B/B-scale views, copy atoms, and the scaled-MFMA atom set + +``gemm_mma`` are shared with gemm2 via ``mxfp4_moe_layout``): + * B load -> ``ml.bq_view`` + ``fx.copy`` into register fragments; the nt/cached + policy rides on the copy atom's ``cache_modifier``. + * B-scale load -> ``ml.bscale_view`` + ``fx.copy`` (32b, cached) into per-stage i32 + fragments; the per-K-tile offset rides the voffset (no hi/lo split). + * MMA -> one ``ml.gemm_mma`` (fx.gemm) per mfma over rank-1 register fragments + (A/B/C), with per-K-block e8m0 scales via ``scale_a=/scale_b=`` and a pre-built + opsel-specialized ``MFMA_Scale`` atom per (opselA,opselB). C accumulates in + place (d == c). mem2reg folds the fragment plumbing -> ISA == the raw intrinsic. + +Kept raw (self-contained helpers below): math/quant/pointer helpers, the A-side +LDS stage + ds-read addressing + A-scale machinery, and the epilogue math. (B and +B-scale now both ride the layout API; only the A path and epilogue stay raw.) +Acceptance: KIMI BM32 interleave numeric gate (mean_row_cos > 0.85). +""" + +import flydsl.compiler as flyc +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, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +from . import mxfp4_moe_layout as ml + +# ---- constants (KIMI defaults; per-shape values come from compile args) ------ +NE_DEFAULT, K_DEFAULT, INTER_DEFAULT, TOPK_DEFAULT = 385, 7168, 512, 9 +BN = BK = 256 +KH_TILE = BK // 2 # 128 packed bytes per K-tile +kStages = 2 +kBS_stride_k0_dw = 64 +LOG2E = 1.4426950408889634 +_PTR3 = "!llvm.ptr<3>" + +# BM32 path: fixed for the single supported variant. +BM = 32 +kAStages = 3 +kSubBlocks = 1 +kMChunks = 2 # BM // 16 +M_REPS = BM // 16 +BN_INT = BN // 2 # 128 + + +# ---- self-contained math / pointer / size helpers --------------------------- +def _raw(v): + """Unwrap an fx value to a raw ir.Value.""" + if not isinstance(v, ir.Value) and hasattr(v, "ir_value"): + return v.ir_value() + return v + + +def _lds_ptr3(base_i32, byte_off_i32): + return llvm.inttoptr(ir.Type.parse(_PTR3), _raw(fx.Int64(base_i32 + byte_off_i32))) + + +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_ptr1(arg, byte_off_i32): + return _gep1(_global_base_ptr1(arg), byte_off_i32) + + +def _lds_swizzle_mask(row): + """lds_swizzle_mask(row) = (row & 14) << 3 (fp4 A tile).""" + return (row & fx.Int32(14)) << fx.Int32(3) + + +def _lds_swizzle_mask_f8(row): + """lds_swizzle_mask(row) = (row & 15) << 4 (fp8 A tile).""" + return (row & fx.Int32(15)) << fx.Int32(4) + + +def _silu_mul_batch(gs, us): + """silu(g)*u via exp2/rcp (matches HIP silu_mul_fast).""" + 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 _fabs_f32(x): + """fabsf via sign-bit clear (FlyDSL has no arith.absf).""" + abs_bits = _raw(x).bitcast(T.i32) & _raw(fx.Int32(0x7FFFFFFF)) + return fx.Float32(abs_bits.bitcast(T.f32)) + + +def _e8m0_from_amax(amax_f32, dtype_max=6.0): + """(e8m0_i32, quant_scale_f32) = ceil_pow2(amax/dtype_max) clamped to 254. + dtype_max is the output format's max magnitude (fp4 e2m1 = 6, fp8 e4m3 = 448).""" + wi = fx.Int32(_raw(amax_f32 * fx.Float32(1.0 / dtype_max)).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))) + e8m0 = fx.Int32(arith.select(lt, _raw(bexp), _raw(fx.Int32(254)))) + qscale = fx.Float32(_raw(e8m0 << fx.Int32(23)).bitcast(T.f32)) + return e8m0, qscale + + +def gemm1_grid(n_tokens, BM=32, NE=NE_DEFAULT, TOPK=TOPK_DEFAULT, INTER=INTER_DEFAULT): + """Host-side grid size (BM=32 active-experts bound).""" + active = min(n_tokens * TOPK, NE) + max_m_blocks = (n_tokens * TOPK + active * (BM - 1) + BM - 1) // BM + return max_m_blocks * ((2 * INTER) // 256) + + +@flyc.jit +def _gemm1_body_v2( + allocator, + lds_off, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_sti, + arg_aqout, + arg_ascaleout, + bx_i32, + lane, + wave, + i32_ntok, + i32_total_m_blocks, + *, + K, + INTER, + NE, + interleave, + b_nontemporal, + a_dtype, + out_dtype, +): + # A activation dtype: fp4 (packed 2/byte) or fp8 (1 byte/elem). Only the A + # payload path differs (LDS tile size, ds-read gather, mfma A-format); the weight + # + all e8m0 scale paths are identical. fp8 A uses the raw mfma_scale intrinsic + # (cbsz=0); fp4 A keeps the fx.gemm fragment path. + is_f8_a = a_dtype == "fp8" + # Intermediate OUTPUT dtype: fp4 (8/i32, INTER//2 bytes/row) or fp8 (mxfp8: 4/i32, + # INTER bytes/row). Only the epilogue requant/pack/store differs. + is_f8_out = out_dtype == "fp8" + out_max = 448.0 if is_f8_out else 6.0 # e4m3 / e2m1 max magnitude + out_pack = 1 if is_f8_out else 2 # logical out elems per stored byte + a_pack = 1 if is_f8_a else 2 # logical A elems per stored byte + am = 2 // a_pack # A row-group calls per 8-row sub (fp8=2, fp4=1) + KH_TILE_A = BK // a_pack # A bytes per K-tile row in LDS (fp8=256, fp4=128) + cbsz_a = 0 if is_f8_a else 4 # mfma A-format selector (fp8=0, fp4=4) + # K-/INTER-derived sizes (compile-time Python ints; parametrized over the contraction dim). + _kc = (K // 32) // 4 // 2 + K_HALF = K // 2 + K_BYTES = K // a_pack # a_quant row stride in bytes (= K_HALF for fp4) + K_TILES_TOTAL = K // BK + kUnroll = K_TILES_TOTAL - kStages + kAS_per_chunk_dw = _kc * 64 + kBS_stride_n0_dw = _kc * 64 + N_OUT = 2 * INTER + kBS_per_expert_dw = (N_OUT // 16 // 2) * kBS_stride_n0_dw + NUM_N_BLOCKS = N_OUT // 256 + OUT_AS_PER_CHUNK_DW = ((INTER // 32) // 4 // 2) * 64 + K_G2_BYTES = INTER // out_pack # output intermediate row stride (fp4 INTER/2, fp8 INTER) + + # block -> (m_block_idx, n_block_idx) ; e = sorted_expert_ids[m_block_idx] + 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, llvm.load(T.i32, _global_ptr1(arg_eids, m_block_idx * fx.Int32(4))) + ) + m_row = m_block_idx * fx.Int32(BM) + + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + + # buffer resources (A-gather + scales) + aq_num_records = arith.index_cast(T.index, _raw(i32_ntok * fx.Int32(K_BYTES))) + aq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_aq)), num_records_bytes=aq_num_records + ) + _asc_per_mb = max(BM // 32, 1) * kAS_per_chunk_dw * 4 + ascale_num = arith.index_cast(T.index, _raw(i32_total_m_blocks)) * fx.Index( + _asc_per_mb + ) + ascale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_ascale)), num_records_bytes=ascale_num + ) + + # LDS views (s_aq / s_asc, union-overlapping lds_acc) + lds_base = allocator.get_base() + s_aq = SmemPtr(lds_base, lds_off, T.i8, shape=(kAStages * BM * KH_TILE_A,)) + s_asc = SmemPtr( + lds_base, + lds_off + kAStages * BM * KH_TILE_A, + T.i8, + shape=(kSubBlocks * K_TILES_TOTAL * 256,), + ) + lds_acc = SmemPtr(lds_base, lds_off, T.f32, shape=(BM * BN,)) + + # cached A rows (A-gather base offset). buffer_load_lds fills 64*16B/wave: fp4 -> + # 8 rows x 128B (lane//8), fp8 -> 4 rows x 256B (lane//16); fp8 needs `am` calls. + lanes_per_row = KH_TILE_A // 16 # 8 (fp4) / 16 (fp8) + rows_per_call = 64 // lanes_per_row # 8 (fp4) / 4 (fp8) + a_lane_row = lane // fx.Int32(lanes_per_row) + # Gather row is read from sorted_token_ids and masked to the low 24 bits + # (token_id; high byte = topk_id) -- the reference mixed_moe gather. Pad rows + # carry token_id==M (OOB) so the A_q buffer-bounds load returns 0. + mask24_i32 = arith.constant(0xFFFFFF, type=T.i32) + cached_actual_row = [] + for sub in range_constexpr(kSubBlocks): + for h in range_constexpr(am): + idx = ( + m_row + + wave * fx.Int32(BM // 4) + + fx.Int32(sub * 8 + h * rows_per_call) + + a_lane_row + ) + cached_actual_row.append( + arith.andi( + llvm.load(T.i32, _global_ptr1(arg_sti, idx * fx.Int32(4))), + mask24_i32, + ) + ) + + # B-scale n-pack words (gate/up split differs by mode); the per-(wave,mw) base + # is uniform, the per-lane + per-K-tile parts become layout axes (see views below). + 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)] + + def issue_a_load_lds(slot, kt): + # lane L -> LDS[base+L*16]: fp4 8 rows x 128B (lane//8,lane%8), fp8 4 rows x + # 256B (lane//16,lane%16); fp8 splits each 8-row sub into `am` row-groups. + lane_col = (lane % fx.Int32(lanes_per_row)) * fx.Int32(16) + base_i32 = fx.Int32(memref_dialect.extract_aligned_pointer_as_index(s_aq.get())) + for sub in range_constexpr(kSubBlocks): + for h in range_constexpr(am): + lds_row = wave * fx.Int32(BM // 4) + fx.Int32(sub * 8 + h * rows_per_call) + mask = ( + _lds_swizzle_mask_f8(lds_row + a_lane_row) + if const_expr(is_f8_a) + else _lds_swizzle_mask(lds_row + a_lane_row) + ) + voffset = (lane_col ^ mask) + cached_actual_row[sub * am + h] * fx.Int32( + K_BYTES + ) + off = fx.Int32(slot * (BM * KH_TILE_A)) + lds_row * fx.Int32(KH_TILE_A) + rocdl.raw_ptr_buffer_load_lds( + aq_rsrc, + _lds_ptr3(base_i32, off), + fx.Int32(16), + voffset, + fx.Int32(kt * KH_TILE_A), + fx.Int32(0), + fx.Int32(0), + ) + + def issue_a_ds_read(slot): + # fp4: 32 contiguous K (Vec4 i32) at col g*16+k*64 -> A fragment for fx.gemm. + # fp8: a lane's 32 K split into two 16-K halves 64B apart -> Vec8 i32 (raw, + # for the mfma_scale intrinsic; cbsz=0). + base_ptr = _lds_base_ptr3(s_aq.get()) + for k in range_constexpr(2): + for i in range_constexpr(kMChunks): + lds_row = lane_mod_16 + fx.Int32(i * 16) + row_off = fx.Int32(slot * (BM * KH_TILE_A)) + lds_row * fx.Int32(KH_TILE_A) + if const_expr(is_f8_a): + mask = _lds_swizzle_mask_f8(lane_mod_16) + col0 = lane_div_16 * fx.Int32(16) + fx.Int32(k * 128) + col_lo = col0 ^ mask + col_hi = (col0 + fx.Int32(64)) ^ mask + lo = Vec(llvm.load(T.vec(2, T.i64), _gep3(base_ptr, row_off + col_lo))) + hi = Vec(llvm.load(T.vec(2, T.i64), _gep3(base_ptr, row_off + col_hi))) + a64 = Vec.from_elements([lo[0], lo[1], hi[0], hi[1]], fx.Int64) + _a_vals[i][k] = _raw(a64.bitcast(fx.Int32)) + else: + mask = _lds_swizzle_mask(lane_mod_16) + lds_col = (lane_div_16 * fx.Int32(16) + fx.Int32(k * 64)) ^ mask + vec = llvm.load(T.vec(4, T.i32), _gep3(base_ptr, row_off + lds_col)) + _a_frags[i][k].store(Vec(vec)) + + 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) + asc_base = fx.Int32( + memref_dialect.extract_aligned_pointer_as_index(s_asc.get()) + ) + 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) + rocdl.raw_ptr_buffer_load_lds( + ascale_rsrc, + _lds_ptr3(asc_base, lds_sub + wave * fx.Int32(1024)), + fx.Int32(16), + v16, + s_chunk, + fx.Int32(0), + fx.Int32(0), + ) + for d in range_constexpr(3): + byte_off = 4096 + d * 1024 + s_off = rocdl.readfirstlane(T.i32, s_chunk + fx.Int32(byte_off)) + rocdl.raw_ptr_buffer_load_lds( + ascale_rsrc, + _lds_ptr3( + asc_base, lds_sub + fx.Int32(byte_off) + wave * fx.Int32(256) + ), + fx.Int32(4), + v4, + s_off, + fx.Int32(0), + fx.Int32(0), + ) + + def issue_a_scale_ds_read(kt): + base_ptr = _lds_base_ptr3(s_asc.get()) + 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(llvm.load(T.i32, _gep3(base_ptr, lds_dw * fx.Int32(4)))) + return out + + # B load: CK preshuffle as an fx.make_layout view over bq. The descriptor base + # MUST stay uniform per wave (folding the per-lane part in makes make_buffer_tensor + # emit a per-lane WATERFALL, ~14x slower), so the base is the uniform col offset + # and the per-lane (klane,nlane) are layout axes -> a VGPR voffset at copy time. + # nt/cached rides on the copy atom's cache_modifier (2=nt/0=cached). + KH4 = K_HALF // 4 # i32 stride for the col axis + _b_copy_atom = ml.b_copy_atom(b_nontemporal) + _bs_copy_atom = ml.bscale_copy_atom() + + N0_HALF = N_OUT // 32 # separate-mode gate/up column split + + # B-load view per j-tile (shared layout primitive). interleave / separated only + # change which logical N-row `col` maps to; the view layout is identical. + def _make_bq_view_for_jtile(j): + 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) + col = ((tile_il & fx.Int32(1)) * fx.Int32(N0_HALF) + (tile_il >> fx.Int32(1))) * fx.Int32(16) + return ml.bq_view(arg_bq, e * fx.Int32(N_OUT) + col, KH4, K_TILES_TOTAL) + + _bq_views = [_make_bq_view_for_jtile(j) for j in range_constexpr(4)] + + # B-scale view per n-pack word (shared layout primitive). + _bscale_views = [ + ml.bscale_view( + arg_bscale, + e * fx.Int32(kBS_per_expert_dw) + np_list[mw] * fx.Int32(kBS_stride_n0_dw), + K_TILES_TOTAL, + k0_stride_dw=kBS_stride_k0_dw, + ) + for mw in range_constexpr(2) + ] + + # B is loaded via fx.copy into i32<4:1> fragments (16B = 32 fp4) regardless of A + # dtype. PER-STAGE (kStages) prefetch double-buffer. + _frag_tmpl = ml.bq_frag_tmpl(_bq_views[0]) # i32<4:1> + _bq_frags = [ + [ + [fx.make_fragment_like(_frag_tmpl) for _ in range_constexpr(2)] + for _ in range_constexpr(4) + ] + for _ in range_constexpr(kStages) + ] + # A / C: fp4 uses fx.gemm register fragments (A refilled per K iter, C accumulates + # in place). fp8 uses the raw mfma_scale intrinsic, so A is a per-iter Vec8 i32 + # value (_a_vals) and C is a raw f32x4 accumulator (accm, init to zero). + zero4 = Vec.filled(4, 0.0, fx.Float32) + if const_expr(is_f8_a): + _a_vals = [[None, None] for _ in range(kMChunks)] + accm = [[zero4 for _ in range(4)] for _ in range(kMChunks)] + else: + _a_frags = [ + [fx.make_fragment_like(_frag_tmpl) for _ in range_constexpr(2)] + for _ in range_constexpr(kMChunks) + ] + _c_frags = [ + [ + fx.make_fragment_like(_frag_tmpl, dtype=fx.Float32.ir_type) + for _ in range_constexpr(4) + ] + for _ in range_constexpr(kMChunks) + ] + # B-scale fragments: i32<1:1>, PER-STAGE (kStages) double-buffer like _bq_frags. + _bs_frag_tmpl = ml.bscale_frag_tmpl(_bscale_views[0]) # i32<1:1> + _bs_frags = [ + [fx.make_fragment_like(_bs_frag_tmpl) for _ in range_constexpr(2)] + for _ in range_constexpr(kStages) + ] + + def issue_b_load_j(stage, K_C, j): + view = _bq_views[j] + for half in range_constexpr(2): + fx.copy( + _b_copy_atom, + view[lane_div_16, lane_mod_16, K_C, half, None], + _bq_frags[stage][j][half], + ) + + def issue_b_scale_load(stage, K_C): + for mw in range_constexpr(2): + fx.copy( + _bs_copy_atom, + _bscale_views[mw][lane_div_16, lane_mod_16, K_C, None], + _bs_frags[stage][mw], + ) + + # MMA. fp4: one fx.gemm per mfma over rank-1 fragments (shared layout primitive), + # scales ride scale_a=/scale_b=, C accumulates in place. fp8: the raw scaled-MFMA + # intrinsic (cbsz=0, A is the Vec8 i32 from ds-read), C accumulates via accm. + _scale_mma_atoms = ml.scale_mma_atoms() if const_expr(not is_f8_a) else None + + def _gemm_mma(a_frag, b_frag, c_frag, opsel_a, opsel_b, sa, sb): + ml.gemm_mma(_scale_mma_atoms, a_frag, b_frag, c_frag, opsel_a, opsel_b, sa, sb) + + def mfma_cluster(stage, a_scale, J): + # interleave: mni=J//2 (n0), in_b=J%2 (gate/up); separate: swapped. + if const_expr(interleave): + mni, in_b = J // 2, J % 2 + else: + mni, in_b = J % 2, J // 2 + sb = _raw(Vec(_bs_frags[stage][mni].load())[0]) + sa = a_scale[0] # kSubBlocks == 1 + if const_expr(is_f8_a): + bJ0 = Vec(_bq_frags[stage][J][0].load()) + bJ1 = Vec(_bq_frags[stage][J][1].load()) + for osa, k, i in ((0, 0, 0), (1, 0, 1), (2, 1, 0), (3, 1, 1)): + bJ = bJ0 if k == 0 else bJ1 + osb = (0 if k == 0 else 2) + in_b + accm[i][J] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.f32x4, [_a_vals[i][k], bJ, accm[i][J], cbsz_a, 4, osa, sa, osb, sb] + ) + else: + bJ0, bJ1 = _bq_frags[stage][J][0], _bq_frags[stage][J][1] + _gemm_mma(_a_frags[0][0], bJ0, _c_frags[0][J], 0, 0 + in_b, sa, sb) + _gemm_mma(_a_frags[1][0], bJ0, _c_frags[1][J], 1, 0 + in_b, sa, sb) + _gemm_mma(_a_frags[0][1], bJ1, _c_frags[0][J], 2, 2 + in_b, sa, sb) + _gemm_mma(_a_frags[1][1], bJ1, _c_frags[1][J], 3, 2 + in_b, sa, sb) + + # zero C (fp4 fragments accumulate in place thereafter; fp8 accm pre-init above). + if const_expr(not is_f8_a): + for i in range_constexpr(kMChunks): + for J in range_constexpr(4): + _c_frags[i][J].store(zero4) + + # prologue: stages 0,1 + issue_a_scale_load() + for K_C in range_constexpr(kStages): + issue_a_load_lds(K_C, K_C) + for j in range_constexpr(4): + issue_b_load_j(K_C, K_C, j) + issue_b_scale_load(K_C, K_C) + + # main loop. sched_barrier/s_setprio fence the mfma chain from the B loads (mirror + # v1's BM!=128 hints) so it stays dense -- closes the small-M gap. + 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() + issue_a_ds_read(read_slot) + asc_cur = issue_a_scale_ds_read(K_C - kStages) + issue_a_load_lds(write_slot, K_C) + for J in range_constexpr(4): + rocdl.sched_barrier(0) + rocdl.s_setprio(1) + mfma_cluster(slot_b, asc_cur, J) + rocdl.s_setprio(0) + rocdl.sched_barrier(0) + issue_b_load_j(slot_b, K_C, J) + rocdl.sched_barrier(0) + issue_b_scale_load(slot_b, K_C) + + # drain: last kStages + for S in range_constexpr(kStages): + kt = K_TILES_TOTAL - kStages + S + gpu.barrier() + issue_a_ds_read(kt % kAStages) + asc_cur = issue_a_scale_ds_read(kt) + for J in range_constexpr(4): + mfma_cluster(kt % kStages, asc_cur, J) + + gpu.barrier() + s_aq._view_cache = None + s_asc._view_cache = None + lds_acc._view_cache = None + + # epilog: cshuffle -> SwiGLU -> fp4 + e8m0 requant (raw math) + wave_n = wave + lds_acc_base = _lds_base_ptr3(lds_acc.get()) + + # Read accumulators (flat slot [i,J,v]): fp4 from the C fragments, fp8 from accm. + if const_expr(is_f8_a): + _acc_vecs = [[Vec(accm[i][J]) for J in range(4)] for i in range(kMChunks)] + else: + _acc_vecs = [[Vec(_c_frags[i][J].load()) for J in range(4)] for i in range(kMChunks)] + + def _acc(i, J, v): + return _acc_vecs[i][J][v] + + 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_n * fx.Int32(32) + fx.Int32(J_local * 16) + lane_mod_16 + lds_col = (fx.Int32(128) + col_local) if is_up else col_local + for v in range_constexpr(4): + idx = (row_base + fx.Int32(v)) * fx.Int32(BN) + lds_col + llvm.StoreOp( + _raw(fx.Float32(_acc(i, J, v))), + _gep3(lds_acc_base, idx * fx.Int32(4)), + ) + + gpu.barrier() + + tx_i32 = arith.index_cast(T.i32, 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_base = _global_base_ptr1(arg_aqout) + 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_off = (row_local * fx.Int32(BN) + gate_col) * fx.Int32(4) + up_off = (row_local * fx.Int32(BN) + up_col) * fx.Int32(4) + gate_vs[ee] = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, gate_off))) + up_vs[ee] = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, up_off))) + 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])) + local_max = local_max.maximumf(local_max.shuffle_xor(fx.Int32(1), fx.Int32(64))) + local_max = local_max.maximumf(local_max.shuffle_xor(fx.Int32(2), fx.Int32(64))) + + e8m0, qscale = _e8m0_from_amax(local_max, out_max) + scales_per_mr[mr] = e8m0 + + qscale_raw = _raw(qscale) + # fp4 byte position of this lane's 8 elems (8 fp4 = 4 bytes); fp8 doubles it + # (8 fp8 = 8 bytes), and the row stride is INTER (vs INTER//2). + byte_pos_fp4 = ( + n_block_idx * fx.Int32(BN_INT // 2) + + wave_grp * fx.Int32(16) + + kk * fx.Int32(4) + ) + out_row = m_row + row_local + if const_expr(is_f8_out): + # 8 f32 -> 8 fp8 = 2x vector<2xi16> (4 fp8 each): cvt packs 2 fp8 into the + # lo/hi 16-bit half of the running vector. lo holds elems 0..3, hi 4..7. + v2i16 = T.vec(2, T.i16) + lo = _raw(Vec.filled(2, 0, fx.Int16)) + lo = rocdl.cvt_scalef32_pk_fp8_f32(v2i16, lo, _raw(result[0]), _raw(result[1]), qscale_raw, 0) + lo = rocdl.cvt_scalef32_pk_fp8_f32(v2i16, lo, _raw(result[2]), _raw(result[3]), qscale_raw, 1) + hi = _raw(Vec.filled(2, 0, fx.Int16)) + hi = rocdl.cvt_scalef32_pk_fp8_f32(v2i16, hi, _raw(result[4]), _raw(result[5]), qscale_raw, 0) + hi = rocdl.cvt_scalef32_pk_fp8_f32(v2i16, hi, _raw(result[6]), _raw(result[7]), qscale_raw, 1) + store_off = out_row * fx.Int32(K_G2_BYTES) + byte_pos_fp4 * fx.Int32(2) + llvm.StoreOp(lo, _gep1(aqout_base, store_off), alignment=4, nontemporal=True) + llvm.StoreOp(hi, _gep1(aqout_base, store_off + fx.Int32(4)), alignment=4, nontemporal=True) + else: + packed_i32 = _raw(fx.Int32(0)) + 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, + ) + store_off = out_row * fx.Int32(K_G2_BYTES) + byte_pos_fp4 + llvm.StoreOp( + _raw(fx.Int32(packed_i32)), + _gep1(aqout_base, store_off), + alignment=4, + nontemporal=True, + ) + + ascaleout_base = _global_base_ptr1(arg_ascaleout) + if kk == fx.Int32(0): + ku = n_block_idx >> fx.Int32(1) + ikxdl = n_block_idx & fx.Int32(1) + for sub in range_constexpr(kSubBlocks): + chunk = m_block_idx * fx.Int32(kSubBlocks) + fx.Int32(sub) + dword_off = ( + chunk * fx.Int32(OUT_AS_PER_CHUNK_DW) + + ku * fx.Int32(64) + + wave_grp * fx.Int32(16) + + m_lane + ) + pair_i32 = scales_per_mr[sub * 2 + 0] | ( + scales_per_mr[sub * 2 + 1] << fx.Int32(8) + ) + pair_i16 = arith.TruncIOp(T.i16, _raw(pair_i32)).result + addr = dword_off * fx.Int32(4) + ikxdl * fx.Int32(2) + llvm.StoreOp( + pair_i16, + _gep1(ascaleout_base, addr), + alignment=2, + ) + + +def _lds_bytes_for(K_TILES_TOTAL, KH_TILE_A=KH_TILE): + s_aq_bytes = kAStages * BM * KH_TILE_A # fp8 A tile is 2x (256B vs 128B) + s_asc_bytes = kSubBlocks * K_TILES_TOTAL * 256 + lds_acc_bytes = BM * BN * 4 + return max(s_aq_bytes + s_asc_bytes, lds_acc_bytes) + + +def compile_gemm1_a4w4_port( + BM=32, + use_nt=True, + inline_quant=False, + D_HIDDEN=K_DEFAULT, + D_INTER=INTER_DEFAULT, + NE=NE_DEFAULT, + TOPK=TOPK_DEFAULT, + interleave=True, + a_dtype="fp4", + out_dtype="fp4", +): + # use_nt IS the B-load cache policy (v1's `b_aux = 2 if use_nt else 0`; + # tuned config BM32_NT vs BM32_CACHED): True -> nt (decode), False -> cached. + b_nontemporal = use_nt + if (BM, inline_quant) != (32, False): + raise AssertionError( + f"mxfp4_moe_gemm1 supports only (BM=32, inline_quant=False); " + f"got (BM={BM}, inline_quant={inline_quant})" + ) + if a_dtype not in ("fp4", "fp8"): + raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + if out_dtype not in ("fp4", "fp8"): + raise AssertionError(f"out_dtype must be 'fp4' or 'fp8', got {out_dtype!r}") + + _K, _INTER, _NE = D_HIDDEN, D_INTER, NE + assert _K % BK == 0, f"D_HIDDEN (K) must be a multiple of {BK}, got {_K}" + _N_OUT = 2 * _INTER + assert _N_OUT % BN == 0, f"2*D_INTER (N_OUT) must be a multiple of {BN}, got {_N_OUT}" + + _KH_TILE_A = BK // (1 if a_dtype == "fp8" else 2) + lds_bytes = _lds_bytes_for(_K // BK, _KH_TILE_A) # K_TILES_TOTAL + + gu_tag = "il" if interleave else "sep" + bnt_tag = "nt" if b_nontemporal else "cached" + a_tag = "a8" if a_dtype == "fp8" else "a4" + o_tag = "o8" if out_dtype == "fp8" else "o4" + name_suffix = f"h{_K}_i{_INTER}_ne{_NE}_bm{BM}_{bnt_tag}_{gu_tag}_{a_tag}{o_tag}_v2" + + allocator = SmemAllocator( + None, arch="gfx950", global_sym_name=f"gemm1port_v2_smem_{name_suffix}" + ) + lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_off + lds_bytes + + @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_sti: fx.Int64, + i32_ntok: fx.Int32, + arg_aqout: fx.Int64, + arg_ascaleout: fx.Int64, + arg_hidden: fx.Int64, + ): + tx = gpu.thread_id("x") + bx = gpu.block_id("x") + tx_i32 = arith.index_cast(T.i32, tx) + bx_i32 = arith.index_cast(T.i32, bx) + lane = tx_i32 % fx.Int32(64) + wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) + cumsum0 = llvm.load(T.i32, _global_ptr1(arg_cumsum, fx.Int32(0))) + total_m_blocks = cumsum0 // fx.Int32(BM) + bound = total_m_blocks * fx.Int32(_N_OUT // 256) # * NUM_N_BLOCKS + if fx.Int32(bx_i32) < bound: + _gemm1_body_v2( + allocator, + lds_off, + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_sti, + arg_aqout, + arg_ascaleout, + bx_i32, + lane, + wave, + i32_ntok, + total_m_blocks, + K=_K, + INTER=_INTER, + NE=_NE, + interleave=interleave, + b_nontemporal=b_nontemporal, + a_dtype=a_dtype, + out_dtype=out_dtype, + ) + + @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_sti: 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, + ): + from flydsl.compiler.kernel_function import CompilationContext + + ctx = CompilationContext.get_current() + allocator.finalized = False + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + grid_x = arith.index_cast(T.index, i32_grid) + gemm1_kernel( + arg_aq, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_cumsum, + arg_sti, + 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/mxfp4_moe_gemm2.py b/kernels/mxfp4_moe_gemm2.py new file mode 100644 index 000000000..74cad9592 --- /dev/null +++ b/kernels/mxfp4_moe_gemm2.py @@ -0,0 +1,522 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors +"""FlyDSL **layout-API** port of aiter ``gemm2_a4w4`` (MXFP4 MoE down-proj). + +Variant: ``(BM=32, epilog="atomic")`` for both ``use_nt`` (B-load nt/cached policy) +-- exactly what the tuned KIMI config's gemm2 selects (kernelName2 BM32_ATOMIC_NT for +tok 256/512, BM32_ATOMIC for tok 1024/2048 at NE385/H7168/inter512/TOPK9). Covers the +KIMI fast path (D_INTER<=512, K_TILES<=2, fully unrolled) and the streaming K-loop +(D_INTER>512). The BM32 GEMM core is identical to ``mxfp4_moe_gemm1`` (same +preshuffled-B layout + scaled-MFMA opsel), so the B-load / B-scale / MMA pieces come +straight from ``mxfp4_moe_layout``: + + * B load -> ``ml.bq_view`` + ``fx.copy`` into per-tile register fragments; + nt/cached rides the copy atom's cache_modifier. + * B-scale load -> ``ml.bscale_view`` + ``fx.copy`` into per-tile i32 fragments. + * MMA -> one ``ml.gemm_mma`` (fx.gemm) per mfma over rank-1 fragments (A/B/C), + per-K-block e8m0 scales via scale_a=/scale_b=, C accumulating in place. + +Kept raw (shared via ``mxfp4_moe_common``): pointer helpers, the A-side +LDS stage + ds-read + A-scale machinery, and the atomic-bf16 epilogue. Unlike gemm1 +(which streams B per-K), gemm2 loads ALL B tiles up front into registers (B is not +LDS-bound), so the B/B-scale fragments are PER-TILE (K_TILES_TOTAL); only the A->LDS +stage streams (triple-buffered for K_TILES>2). BM is fixed at 32, so the BM-dependent +tile sizes are module constants (mirrors mxfp4_moe_gemm1). The D_INTER_REAL pad-tail +path is not supported (KIMI never pads). +""" + +import flydsl.compiler as flyc +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, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +from . import mxfp4_moe_layout as ml + +# Raw helpers / K-derived size formulas / constants reused verbatim from v1; the +# A-side LDS stage + ds-read + A-scale and the atomic epilogue are unchanged. +from .mxfp4_moe_common import ( + BK, + BN, + K, + MAX_M, + NE, + N_OUT, + kBS_stride_k0_dw, + kStages, + _atomic_bf16_epilog, + _gep3, + _global_ptr1, + _lds_base_ptr3, + _lds_ptr3, + _lds_swizzle_mask, + _raw, + _udiv, + k_half_for, + k_tiles_total_for, + kas_per_chunk_dw_for, + kbs_per_expert_dw_for, + kbs_stride_n0_dw_for, + kunroll_for, + num_n_blocks_for, +) + +# BM32 is the only variant -> the BM-dependent tile sizes are constants. +BM = 32 +kMChunks = BM // 16 # 2 (M-subblocks of 16 rows) +N_LOAD_WAVES = 4 # all 4 waves load A rows +ROWS_PER_WAVE = BM // N_LOAD_WAVES # 8 + + +def _lds_swizzle_mask_f8(row): + """lds_swizzle_mask(row) = (row & 15) << 4 (fp8 A tile).""" + return (row & fx.Int32(15)) << fx.Int32(4) + + +def _issue_a_load_lds_dt(aq_rsrc, saq, slot, kt, m_row, wave, lane, is_f8, KH_TILE_A, K_BYTES): + """A->LDS for one K-tile (gemm2: A is the already-sorted intermediate, so the + source row is the sorted row directly -- no m_indices gather). fp4: 8 lanes/row x + 128B; fp8: 16 lanes/row x 64B x am=2 row-groups, with the 256B-row swizzle. + Mirrors mxfp4_moe_gemm1.issue_a_load_lds; BM32 -> kSubBlocks=1.""" + am = 2 if is_f8 else 1 # row-group calls per 8-row wave (fp8 4 rows/call) + lanes_per_row = KH_TILE_A // 16 # 8 (fp4) / 16 (fp8) + rows_per_call = 64 // lanes_per_row # 8 (fp4) / 4 (fp8) + a_lane_row = lane // fx.Int32(lanes_per_row) + lane_col = (lane % fx.Int32(lanes_per_row)) * fx.Int32(16) + base_i32 = fx.Int32(memref_dialect.extract_aligned_pointer_as_index(saq.get())) + for h in range_constexpr(am): + lds_row = wave * fx.Int32(ROWS_PER_WAVE) + fx.Int32(h * rows_per_call) + mask = ( + _lds_swizzle_mask_f8(lds_row + a_lane_row) + if const_expr(is_f8) + else _lds_swizzle_mask(lds_row + a_lane_row) + ) + car = m_row + lds_row + a_lane_row # direct sorted row + voffset = (lane_col ^ mask) + car * fx.Int32(K_BYTES) + off = fx.Int32(slot * (BM * KH_TILE_A)) + lds_row * fx.Int32(KH_TILE_A) + rocdl.raw_ptr_buffer_load_lds( + aq_rsrc, + _lds_ptr3(base_i32, off), + fx.Int32(16), + voffset, + fx.Int32(kt * KH_TILE_A), + fx.Int32(0), + fx.Int32(0), + ) + + +def compile_gemm2_a4w4_port( + BM=32, + use_nt=False, + NE=NE, + N_OUT=N_OUT, + MAX_M=MAX_M, + epilog="atomic", + D_INTER=K, + D_INTER_REAL=None, + a_dtype="fp4", +): + """Compile the gemm2 a4w4 layout-API down-proj. Only (BM=32, epilog="atomic") is + supported; D_INTER (= contraction K = inter_dim) must be a multiple of BK(256) + (512 keeps the fully-unrolled fast path; >512 uses the streaming K-loop). + a_dtype="fp8" reads an mxfp8 intermediate (gemm1 out_dtype="fp8"). The + D_INTER_REAL pad-tail (unpadded non-256-aligned shards) is not supported.""" + if (BM, epilog) != (32, "atomic"): + raise AssertionError( + f"mxfp4_moe_gemm2 supports only (BM=32, epilog='atomic'); " + f"got (BM={BM}, epilog={epilog})" + ) + if D_INTER_REAL is not None and D_INTER_REAL != D_INTER: + raise AssertionError( + f"mxfp4_moe_gemm2 does not support D_INTER_REAL padding " + f"(D_INTER_REAL={D_INTER_REAL}, D_INTER={D_INTER})" + ) + if a_dtype not in ("fp4", "fp8"): + raise AssertionError(f"a_dtype must be 'fp4' or 'fp8', got {a_dtype!r}") + _K = D_INTER + assert _K % BK == 0, ( + f"D_INTER (gemm2 contraction K = inter_dim) must be a multiple of {BK}, got {_K}" + ) + _is_f8 = a_dtype == "fp8" + _KH_TILE_A = BK // (1 if _is_f8 else 2) # A LDS K-tile bytes (fp8 256, fp4 128) + _K_BYTES = _K // (1 if _is_f8 else 2) # A row stride bytes (fp8 K, fp4 K//2) + _slot_bytes = BM * _KH_TILE_A + _K_TILES_TOTAL = k_tiles_total_for(_K) + _aStages = kStages if _K_TILES_TOTAL <= kStages else 3 + _lds_bytes = max(BM * BN * 4, _aStages * _slot_bytes) + _num_n_blocks = num_n_blocks_for(N_OUT) + + _atag = "_a8" if _is_f8 else "" + _tag = f"ne{NE}_h{N_OUT}_i{_K}_bm{BM}{'_nt' if use_nt else ''}_atomic{_atag}_v2" + _name = f"gemm2_a4w4_port_{_tag}" + + allocator = SmemAllocator( + None, arch="gfx950", global_sym_name=f"gemm2port_v2_smem_{_tag}" + ) + lds_off = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_off + _lds_bytes + + @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, # unused (atomic epilog); kept for signature parity + ): + 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 = arith.index_cast(T.index, _raw(i32_max_m_blocks)) * fx.Index( + BM * _K_BYTES + ) + aq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_aq)), num_records_bytes=_aq_num + ) + saq = SmemPtr( + allocator.get_base(), lds_off, T.i8, shape=(_aStages * _slot_bytes,) + ) + + # Preload the first kStages K-tiles (== ALL tiles for the K_TILES<=2 fast + # path; == prologue for the streaming path). slot == kt for the preload. + def _issue_all_a_loads(m_row0): + for slot in range_constexpr(kStages): + _issue_a_load_lds_dt( + aq_rsrc, saq, slot, slot, m_row0, wave, lane, + _is_f8, _KH_TILE_A, _K_BYTES, + ) + + # One-shot grid (atomic). Issue A->LDS BEFORE the cumsum load so the HBM + # latency overlaps the cumsum + bound check (A->LDS depends only on bx/lane). + _issue_all_a_loads(_udiv(bx_i32, _num_n_blocks) * fx.Int32(BM)) + rocdl.sched_barrier(0) + + cumsum0 = llvm.load(T.i32, _global_ptr1(arg_cumsum, fx.Int32(0))) + total_m_blocks = _udiv(cumsum0, BM) + bound = total_m_blocks * fx.Int32(_num_n_blocks) + + if fx.Int32(bx_i32) < bound: + _gemm2_body_v2( + allocator, + lds_off, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_stids, + arg_sweights, + i32_M, + i32_max_m_blocks, + arg_out, + bx_i32, + lane, + wave, + aq_rsrc, + use_nt=use_nt, + NE=NE, + N_OUT=N_OUT, + D_INTER=_K, + aStages=_aStages, + a_dtype=a_dtype, + ) + + @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, + ): + from flydsl.compiler.kernel_function import CompilationContext + + ctx = CompilationContext.get_current() + allocator.finalized = False + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + 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) + + return launch_gemm2 + + +@flyc.jit +def _gemm2_body_v2( + allocator, + lds_off, + arg_ascale, + arg_bq, + arg_bscale, + arg_eids, + arg_stids, + arg_sweights, + i32_M, + i32_max_m_blocks, + arg_out, + bx_i32, + lane, + wave, + aq_rsrc, + *, + use_nt, + NE, + N_OUT, + D_INTER, + aStages, + a_dtype, +): + _aStages = aStages + # A activation dtype: fp4 (intermediate from gemm1 fp4-out) or fp8 (mxfp8). Only + # the A LDS tile size, ds-read gather, and mfma A-format differ; B/scale identical. + is_f8_a = a_dtype == "fp8" + KH_TILE_A = BK // (1 if is_f8_a else 2) + K_BYTES = D_INTER // (1 if is_f8_a else 2) + slot_bytes = BM * KH_TILE_A + cbsz_a = 0 if is_f8_a else 4 + # K-derived sizes (parametrized over contraction K = inter_dim = D_INTER). + _K = D_INTER + _K_HALF = k_half_for(_K) + _K_TILES_TOTAL = k_tiles_total_for(_K) + _kUnroll = kunroll_for(_K) + _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) + _num_n_blocks = num_n_blocks_for(N_OUT) + KH4 = _K_HALF // 4 + + # block -> (m_block_idx, n_block_idx) ; e = sorted_expert_ids[m_block_idx] + 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, llvm.load(T.i32, _global_ptr1(arg_eids, m_block_idx * fx.Int32(4))) + ) + m_row = m_block_idx * fx.Int32(BM) + + lane_div_16 = lane // fx.Int32(16) + lane_mod_16 = lane % fx.Int32(16) + + # A-scale buffer resource + uniform base (A-scale load stays raw). BM32 -> one + # 32-row chunk, one subblock. + _asc_per_mb = (BM // 32) * _kAS_per_chunk_dw * 4 + _asc_num = arith.index_cast(T.index, _raw(i32_max_m_blocks)) * fx.Index(_asc_per_mb) + ascale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_ascale)), num_records_bytes=_asc_num + ) + a_scale_s_base = rocdl.readfirstlane( + T.i32, (m_row // fx.Int32(32)) * fx.Int32(_kAS_per_chunk_dw) * fx.Int32(4) + ) + v_voff_scale = ((lane_div_16 * fx.Int32(16)) + lane_mod_16) * fx.Int32(4) + + def load_a_scale_tile(kt): + return buffer_ops.buffer_load( + ascale_rsrc, + (v_voff_scale + fx.Int32(kt * 256)) // fx.Int32(4), + vec_width=1, + dtype=T.i32, + soffset_bytes=a_scale_s_base, + ) + + lds_base = allocator.get_base() + saq = SmemPtr(lds_base, lds_off, T.i8, shape=(_aStages * slot_bytes,)) + lds_acc = SmemPtr(lds_base, lds_off, T.f32, shape=(BM * BN,)) + + # -- B / B-scale layout-API views (shared primitives) --------------------- + _b_copy_atom = ml.b_copy_atom(use_nt) + _bs_copy_atom = ml.bscale_copy_atom() + + def _make_bq_view(j): + col = n_block_idx * fx.Int32(BN) + wave * fx.Int32(BN // 4) + fx.Int32(j * 16) + return ml.bq_view(arg_bq, e * fx.Int32(N_OUT) + col, KH4, _K_TILES_TOTAL) + + _bq_views = [_make_bq_view(j) for j in range_constexpr(4)] + + mni_base = n_block_idx * fx.Int32(BN // 16 // 2) + wave * fx.Int32(BN // 64 // 2) + _bscale_views = [ + ml.bscale_view( + arg_bscale, + e * fx.Int32(_kbs_per_expert_dw) + + (mni_base + fx.Int32(mw)) * fx.Int32(_kBS_stride_n0_dw), + _K_TILES_TOTAL, + k0_stride_dw=kBS_stride_k0_dw, + ) + for mw in range_constexpr(2) + ] + + # Fragments. B / B-scale are PER-TILE (all tiles loaded up front, as v1); A is + # refilled per K iter; C (f32) accumulates in place (fx.gemm d == c). + _frag_tmpl = ml.bq_frag_tmpl(_bq_views[0]) # i32<4:1> + _bs_frag_tmpl = ml.bscale_frag_tmpl(_bscale_views[0]) # i32<1:1> + _bq_frags = [ + [[fx.make_fragment_like(_frag_tmpl) for _ in range_constexpr(2)] for _ in range_constexpr(4)] + for _ in range_constexpr(_K_TILES_TOTAL) + ] + _bs_frags = [ + [fx.make_fragment_like(_bs_frag_tmpl) for _ in range_constexpr(2)] + for _ in range_constexpr(_K_TILES_TOTAL) + ] + # A / C: fp4 uses fx.gemm register fragments; fp8 uses the raw mfma_scale intrinsic + # (A is a per-iter Vec8 i32 value, C a raw f32x4 accumulator init to zero). + zero4 = Vec.filled(4, 0.0, fx.Float32) + if const_expr(is_f8_a): + _a_vals = [[None, None] for _ in range(kMChunks)] + accm = [[zero4 for _ in range(4)] for _ in range(kMChunks)] + else: + _a_frags = [ + [fx.make_fragment_like(_frag_tmpl) for _ in range_constexpr(2)] + for _ in range_constexpr(kMChunks) + ] + _c_frags = [ + [fx.make_fragment_like(_frag_tmpl, dtype=fx.Float32.ir_type) for _ in range_constexpr(4)] + for _ in range_constexpr(kMChunks) + ] + + def issue_b_load_tile(kt): + for j in range_constexpr(4): + for half in range_constexpr(2): + fx.copy( + _b_copy_atom, + _bq_views[j][lane_div_16, lane_mod_16, kt, half, None], + _bq_frags[kt][j][half], + ) + + def issue_b_scale_tile(kt): + for mw in range_constexpr(2): + fx.copy( + _bs_copy_atom, + _bscale_views[mw][lane_div_16, lane_mod_16, kt, None], + _bs_frags[kt][mw], + ) + + # A ds-read (raw). fp4 -> i32x4 into fragments (fx.gemm); fp8 -> Vec8 i32 (two + # i64x2 halves 64B apart) as a raw value for the mfma intrinsic. + def issue_a_ds_read(slot): + base_ptr = _lds_base_ptr3(saq.get()) + for k in range_constexpr(2): + for i in range_constexpr(kMChunks): + lds_row = lane_mod_16 + fx.Int32(i * 16) + row_off = fx.Int32(slot * slot_bytes) + lds_row * fx.Int32(KH_TILE_A) + if const_expr(is_f8_a): + mask = _lds_swizzle_mask_f8(lane_mod_16) + col0 = lane_div_16 * fx.Int32(16) + fx.Int32(k * 128) + col_lo = col0 ^ mask + col_hi = (col0 + fx.Int32(64)) ^ mask + lo = Vec(llvm.load(T.vec(2, T.i64), _gep3(base_ptr, row_off + col_lo))) + hi = Vec(llvm.load(T.vec(2, T.i64), _gep3(base_ptr, row_off + col_hi))) + a64 = Vec.from_elements([lo[0], lo[1], hi[0], hi[1]], fx.Int64) + _a_vals[i][k] = _raw(a64.bitcast(fx.Int32)) + else: + mask = _lds_swizzle_mask(lane_mod_16) + lds_col = (lane_div_16 * fx.Int32(16) + fx.Int32(k * 64)) ^ mask + vec = llvm.load(T.vec(4, T.i32), _gep3(base_ptr, row_off + lds_col)) + _a_frags[i][k].store(Vec(vec)) + + def issue_a_load_lds(slot, kt): + _issue_a_load_lds_dt( + aq_rsrc, saq, slot, kt, m_row, wave, lane, is_f8_a, KH_TILE_A, K_BYTES + ) + + _scale_mma_atoms = ml.scale_mma_atoms() if const_expr(not is_f8_a) else None + + def mfma_cluster(kt, sa): + # interleave-equivalent opsel (gemm2 has no gate/up split): mni=J//2, in_b=J%2. + # BM32: kSubBlocks=1 (sub=0), kMChunks=2 (i0=0, i1=1). + for J in range_constexpr(4): + mni, in_b = J // 2, J % 2 + sb = _raw(Vec(_bs_frags[kt][mni].load())[0]) + if const_expr(is_f8_a): + bJ0 = Vec(_bq_frags[kt][J][0].load()) + bJ1 = Vec(_bq_frags[kt][J][1].load()) + for osa, k, i in ((0, 0, 0), (1, 0, 1), (2, 1, 0), (3, 1, 1)): + bJ = bJ0 if k == 0 else bJ1 + osb = (0 if k == 0 else 2) + in_b + accm[i][J] = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.f32x4, [_a_vals[i][k], bJ, accm[i][J], cbsz_a, 4, osa, sa, osb, sb] + ) + else: + bJ0, bJ1 = _bq_frags[kt][J][0], _bq_frags[kt][J][1] + ml.gemm_mma(_scale_mma_atoms, _a_frags[0][0], bJ0, _c_frags[0][J], 0, 0 + in_b, sa, sb) + ml.gemm_mma(_scale_mma_atoms, _a_frags[1][0], bJ0, _c_frags[1][J], 1, 0 + in_b, sa, sb) + ml.gemm_mma(_scale_mma_atoms, _a_frags[0][1], bJ1, _c_frags[0][J], 2, 2 + in_b, sa, sb) + ml.gemm_mma(_scale_mma_atoms, _a_frags[1][1], bJ1, _c_frags[1][J], 3, 2 + in_b, sa, sb) + + # zero C (fp4 fragments accumulate in place; fp8 accm pre-init above). + if const_expr(not is_f8_a): + for i in range_constexpr(kMChunks): + for J in range_constexpr(4): + _c_frags[i][J].store(zero4) + + # Load ALL B-q + B-scale + A-scale tiles up front (B is not LDS-bound), as v1. + a_scale_v = [load_a_scale_tile(kt) for kt in range_constexpr(_K_TILES_TOTAL)] + for kt in range_constexpr(_K_TILES_TOTAL): + issue_b_load_tile(kt) + issue_b_scale_tile(kt) + + if const_expr(_K_TILES_TOTAL <= kStages): + # Fast path: all tiles preloaded in LDS by the kernel. + for kt in range_constexpr(_K_TILES_TOTAL): + gpu.barrier() + issue_a_ds_read(kt % kStages) + mfma_cluster(kt, a_scale_v[kt]) + else: + # Streaming double-buffered K-loop (triple-buffered LDS): process tile + # kt=OFFSET (read slot kt%aStages) and stream the next tile into its slot. + for OFFSET in range_constexpr(_kUnroll): + kt = OFFSET + gpu.barrier() + issue_a_ds_read(kt % _aStages) + issue_a_load_lds((kStages + OFFSET) % _aStages, kStages + OFFSET) + mfma_cluster(kt, a_scale_v[kt]) + for S in range_constexpr(kStages): + kt = _K_TILES_TOTAL - kStages + S + gpu.barrier() + issue_a_ds_read(kt % _aStages) + mfma_cluster(kt, a_scale_v[kt]) + + # -- epilog: atomic bf16 (raw, reused from v1). fp8 accm holds raw f32x4 results; + # fp4 loads the C fragments. --- + saq._view_cache = None + lds_acc._view_cache = None + if const_expr(is_f8_a): + accm_vecs = accm + else: + accm_vecs = [[_c_frags[i][J].load() for J in range(4)] for i in range(kMChunks)] + _atomic_bf16_epilog( + lds_acc, accm_vecs, arg_out, arg_stids, arg_sweights, + m_row, n_block_idx, wave, lane, i32_M, BM, N_OUT, + ) diff --git a/kernels/mxfp4_moe_gemm_2stage.py b/kernels/mxfp4_moe_gemm_2stage.py new file mode 100644 index 000000000..82552c9e0 --- /dev/null +++ b/kernels/mxfp4_moe_gemm_2stage.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors +"""Layout-API MXFP4 MoE gemm (2-stage), opus-sort only. + +This is the layout-API replacement for ``mixed_moe_gemm_2stage`` for the MXFP4 +a4w4 / a8w4 surface. It consumes the standard (opus-style) sort contract emitted +by ``moe_sorting_kernel`` -- ``sorted_token_ids`` packed ``(topk<<24)|token_id`` +with sentinel ``(topk<<24)|M`` -- and needs NO fused-sort extras: + + * gemm1 gathers its A rows straight from ``sorted_token_ids & 0xFFFFFF`` + (the reference ``mixed_moe_gemm_2stage`` gather; padding rows carry M -> the + buffer-bounds load returns 0). + * gemm2 uses the atomic bf16 epilogue, scattering per sorted row into the + output via ``global.atomic.fadd`` weighted by ``sorted_weights`` -- so there + is no inverse-permutation (``reverse_sorted``) dependency. + +Covered surface (BM=32): + * gemm1: a4w4 + a8w4 (fp8 act), interleave + separated gate, nt/cached B-load, + out fp4 / fp8. + * gemm2: atomic epilog, a4w4 + a8w4 (fp8 intermediate). + +The MMA + B / B-scale data movement run through the FlyDSL layout API +(``mxfp4_moe_layout``: ``fx.copy`` + ``fx.gemm``); the A-side LDS staging, the +e8m0 scale math, and the atomic epilogue are raw (shared via +``mxfp4_moe_common``). +""" + +import flydsl.compiler as flyc +from flydsl._mlir import ir + +from .mxfp4_moe_gemm1 import compile_gemm1_a4w4_port, gemm1_grid +from .mxfp4_moe_gemm2 import compile_gemm2_a4w4_port + +__all__ = [ + "compile_gemm1_a4w4_port", + "compile_gemm2_a4w4_port", + "gemm1_grid", + "mxfp4_moe_gemm1", + "mxfp4_moe_gemm2", +] + +# -- launcher cache + dispatch (compile once per config, fast-dispatch after) --- +_G1_CACHE = {} +_G2_CACHE = {} + + +def _run_compiled(exe, args): + """First call: flyc.compile (compiles + executes + caches the CompiledFunction) + on ``exe._cf``. Subsequent calls: fast dispatch via the cached function.""" + cf = getattr(exe, "_cf", None) + if cf is not None: + cf(*args) + return + try: + cf = flyc.compile(exe, *args) + exe._cf = cf + except Exception: + # JitFunction.__call__ leaks ir.Context on compile failure; clean up so a + # later call doesn't take the wrong (no-CompilationContext) code path. + try: + while ir.Context.current is not None: + ir.Context.current.__exit__(None, None, None) + except Exception: + pass + raise + + +def _get_g1(BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk, interleave, + a_dtype, out_dtype): + key = (BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk, interleave, + a_dtype, out_dtype) + launch = _G1_CACHE.get(key) + if launch is None: + launch = compile_gemm1_a4w4_port( + BM=BM, use_nt=use_nt, inline_quant=inline_quant, D_HIDDEN=D_HIDDEN, + D_INTER=D_INTER, NE=NE, TOPK=topk, interleave=interleave, + a_dtype=a_dtype, out_dtype=out_dtype, + ) + _G1_CACHE[key] = launch + return launch + + +def _get_g2(BM, use_nt, NE, D_HIDDEN, epilog, D_INTER, D_INTER_REAL, a_dtype): + key = (BM, use_nt, NE, D_HIDDEN, epilog, D_INTER, D_INTER_REAL, a_dtype) + launch = _G2_CACHE.get(key) + if launch is None: + launch = compile_gemm2_a4w4_port( + BM=BM, use_nt=use_nt, NE=NE, N_OUT=D_HIDDEN, epilog=epilog, + D_INTER=D_INTER, D_INTER_REAL=D_INTER_REAL, a_dtype=a_dtype, + ) + _G2_CACHE[key] = launch + return launch + + +def mxfp4_moe_gemm1( + *, + a_quant, + a_scale_sorted_shuffled, + w1_u8, + w1_scale_u8, + sorted_expert_ids, + cumsum_tensor, + sorted_token_ids, + inter_sorted_quant, + inter_sorted_shuffled_scale, + hidden_states, + n_tokens, + NE, + D_HIDDEN, + D_INTER, + topk, + BM=32, + use_nt=True, + inline_quant=False, + interleave=True, + a_dtype="fp4", + out_dtype="fp4", + stream=None, +): + """Stage-1 up/gate gemm: A_q x w1 -> inter (packed MXFP4 / MXFP8, sorted layout). + + Buffers are pre-allocated by the caller. w1_u8 / w1_scale_u8 must be uint8 + views. ``sorted_token_ids`` is the opus-sort output (gemm1 masks it to the + token id internally). + """ + import torch + + launch = _get_g1(BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk, + interleave, a_dtype, out_dtype) + grid = gemm1_grid(n_tokens, BM, NE=NE, TOPK=topk, INTER=D_INTER) + _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(), + sorted_token_ids.data_ptr(), + n_tokens, + 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 mxfp4_moe_gemm2( + *, + inter_sorted_quant, + inter_sorted_shuffled_scale, + w2_u8, + w2_scale_u8, + sorted_expert_ids, + cumsum_tensor, + sorted_token_ids, + sorted_weights, + out, + M_logical, + max_sorted, + NE, + D_HIDDEN, + D_INTER, + topk, + BM=32, + use_nt=False, + a_dtype="fp4", + D_INTER_REAL=None, + stream=None, +): + """Stage-2 down-proj gemm (atomic bf16 epilog): scatters per sorted row into + ``out`` via weighted ``global.atomic.fadd`` (opus-sort only, no reverse_sorted). + + ``out`` MUST be pre-zeroed ([M, D_HIDDEN] bf16) -- the opus sort zeroes its + ``moe_buf`` for exactly this accumulation. + """ + import torch + + launch = _get_g2(BM, use_nt, NE, D_HIDDEN, "atomic", D_INTER, D_INTER_REAL, + a_dtype) + max_m_blocks = (max_sorted + BM - 1) // BM + out_scale = out # unused by the atomic epilog; any valid device ptr is fine + _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(), + M_logical, + max_m_blocks, + out.data_ptr(), + out_scale.data_ptr(), + torch.cuda.current_stream() if stream is None else stream, + ), + ) + return out diff --git a/kernels/mxfp4_moe_layout.py b/kernels/mxfp4_moe_layout.py new file mode 100644 index 000000000..c3daaf6ea --- /dev/null +++ b/kernels/mxfp4_moe_layout.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2025-2026 FlyDSL Project Contributors +"""Shared FlyDSL **layout-API** primitives for the MXFP4 MoE GEMM ports. + +The BM32 GEMM core is identical between gemm1 (up/gate-proj) and gemm2 (down-proj): +the CK-preshuffled B weight and its e8m0 B-scale share the same on-disk layout, and +the scaled 16x16x128 fp4 MFMA uses the same opsel pattern. This module factors out +those layout-API building blocks so both ``mxfp4_moe_gemm1`` and ``mxfp4_moe_gemm2`` +reuse one definition: + + * ``b_copy_atom`` / ``bscale_copy_atom`` -- the BufferCopy atoms (the nt/cached + policy rides on the B atom's ``cache_modifier``). + * ``bq_view`` / ``bscale_view`` -- ``fx.make_layout`` views over the preshuffled + weight / scale, anchored at a UNIFORM per-(wave) base so the per-lane + (klane,nlane) + K-tile become layout axes (a VGPR voffset at copy time, not a + divergent-pointer waterfall). + * ``bq_frag_tmpl`` / ``bscale_frag_tmpl`` -- register-fragment templates sliced + from those views (i32<4:1> for B, i32<1:1> for B-scale). + * ``scale_mma_atoms`` / ``gemm_mma`` -- the pre-built (opselA,opselB) MFMA_Scale + atom set and the one-mfma ``fx.gemm`` wrapper (scales ride scale_a=/scale_b=). + +Callers keep their own fragment bookkeeping (per-stage vs per-tile), A-side LDS / +ds-read / A-scale, and epilogue -- only these B/B-scale/MMA pieces are shared. +""" + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl.expr import arith, rocdl +from flydsl.expr.typing import Float4E2M1FN +from flydsl.expr.typing import T + + +def _raw(v): + """Unwrap an fx value to a raw ir.Value.""" + if not isinstance(v, ir.Value) and hasattr(v, "ir_value"): + return v.ir_value() + return v + + +def b_copy_atom(nontemporal): + """BufferCopy128b (4x i32 = one 128b weight chunk). nt rides cache_modifier.""" + return fx.make_copy_atom(fx.rocdl.BufferCopy128b(2 if nontemporal else 0), 32) + + +def bscale_copy_atom(): + """BufferCopy32b (1x i32 e8m0 scale word); always cached (scales reuse heavily).""" + return fx.make_copy_atom(fx.rocdl.BufferCopy32b(0), 32) + + +def bq_view(arg_bq, row_elems, KH4, K_TILES_TOTAL): + """Layout view over the preshuffled B weight for one N-row tile. + + ``row_elems`` = (e*N_OUT + col): the logical N-row index into the weight. The + uniform per-(wave) base is ``readfirstlane(row_elems * KH4)`` (KH4 = K_HALF//4, + the i32 col stride); the per-lane (klane,nlane), K-tile, K-half, and kpack4 are + layout axes -> a VGPR voffset at copy time. The byte base is zext'd before *4 + (it can exceed a signed i32). Index ``view[lane//16, lane%16, kt, half, None]`` + -> an i32<4:1> (16B = 32 fp4) slice for fx.copy / fx.gemm. + """ + col_base = rocdl.readfirstlane(T.i32, _raw(row_elems) * fx.Int32(KH4)) + i32_ptr_ty = fx.PointerType.get( + T.i32, address_space=fx.AddressSpace.Global, alignment=16 + ) + off_i64 = fx.Int64(arith.ExtUIOp(T.i64, _raw(col_base)).result) + base_iter = fx.inttoptr(i32_ptr_ty, fx.Int64(arg_bq) + off_i64 * fx.Int64(4)) + # i32 strides: klane[0,4)->64, nlane[0,16)->4, K_tile->512, half[0,2)->256, kpack4->1 + shape = (4, 16, K_TILES_TOTAL, 2, 4) + view = fx.Tensor(fx.make_view(base_iter, fx.make_layout(shape, (64, 4, 512, 256, 1)))) + return fx.rocdl.make_buffer_tensor(view, max_size=False) + + +def bscale_view(arg_bscale, base_dw, K_TILES_TOTAL, k0_stride_dw=64): + """Layout view over the e8m0 B-scale for one n-pack word. + + ``base_dw`` = (e*kBS_per_expert_dw + np*kBS_stride_n0_dw): the uniform per-(wave) + dword base (readfirstlane'd here). The per-lane (klane,nlane) and the K-tile are + layout axes; the full K-tile rides the voffset (no hi/lo soffset split). Index + ``view[lane//16, lane%16, kt, None]`` -> an i32<1:1> scale word. + """ + base_dw = rocdl.readfirstlane(T.i32, _raw(base_dw)) + i32_ptr_ty = fx.PointerType.get( + T.i32, address_space=fx.AddressSpace.Global, alignment=4 + ) + off_i64 = fx.Int64(arith.ExtUIOp(T.i64, _raw(base_dw)).result) + base_iter = fx.inttoptr(i32_ptr_ty, fx.Int64(arg_bscale) + off_i64 * fx.Int64(4)) + shape = (4, 16, K_TILES_TOTAL, 1) + stride = (16, 1, k0_stride_dw, 1) + view = fx.Tensor(fx.make_view(base_iter, fx.make_layout(shape, stride))) + return fx.rocdl.make_buffer_tensor(view, max_size=False) + + +def bq_frag_tmpl(view): + """i32<4:1> fragment template sliced from a bq_view (16B = 32 fp4).""" + return view[0, 0, 0, 0, None] + + +def bscale_frag_tmpl(view): + """i32<1:1> fragment template sliced from a bscale_view (one e8m0 word).""" + return view[0, 0, 0, None] + + +def scale_mma_atoms(): + """Pre-build all 16 (opselA,opselB) scaled-MFMA atoms (opsel is a TYPE param, so + one atom per pair; built once at trace time). cbsz/blgp(=4 for fp4) are inferred + from Float4E2M1FN.""" + return { + (osa, osb): fx.make_mma_atom( + fx.rocdl.cdna4.MFMA_Scale( + 16, 16, 128, Float4E2M1FN, opsel_a=osa, opsel_b=osb + ) + ) + for osa in range(4) + for osb in range(4) + } + + +def gemm_mma(atoms, a_frag, b_frag, c_frag, opsel_a, opsel_b, sa, sb): + """One scaled MFMA via fx.gemm over rank-1 register fragments (-> one + MmaAtomCall). C accumulates in place (d == c); scales ride scale_a=/scale_b=.""" + fx.gemm( + atoms[(opsel_a, opsel_b)], + c_frag, + a_frag, + b_frag, + c_frag, + scale_a=sa, + scale_b=sb, + ) diff --git a/tests/kernels/test_moe_gemm.py b/tests/kernels/test_moe_gemm.py index e956f83b5..8e838768b 100644 --- a/tests/kernels/test_moe_gemm.py +++ b/tests/kernels/test_moe_gemm.py @@ -85,6 +85,11 @@ def _pack_shuffled_int8_to_packed_int4_no_perm(x_shuf_i8: torch.Tensor) -> torch compile_moe_gemm2, compile_moe_gemm2_ex, ) +from kernels.moe_sorting_kernel import moe_sorting_flydsl # noqa: E402 +from kernels.mxfp4_moe_gemm_2stage import ( # noqa: E402 + mxfp4_moe_gemm1, + mxfp4_moe_gemm2, +) logging.basicConfig(level=logging.INFO) @@ -1730,6 +1735,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 +1802,12 @@ 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 layout-API MXFP4 pipe (mxfp4_moe_gemm_2stage) is the BM32 atomic, opus-sort + # path: no reduce mode, and graph capture is out of scope here. + if bool(use_reduce): + pytest.skip(f"{in_dtype} layout-API pipe is atomic-only (no reduce mode)") + if bool(test_graph): + pytest.skip(f"{in_dtype} layout-API pipe: graph capture not covered") device = torch.device("cuda") # torch.manual_seed(int(seed)) @@ -1821,6 +1833,26 @@ def test_moe_gemm_2stage( topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) + # a4w4 / a8w4 run the layout-API MXFP4 pipeline (opus sort -> gemm1 -> gemm2 + # atomic), replacing the mixed_moe path; it fuses both stages + the topk + # reduction, so it bypasses run_moe_stage1 / run_moe_stage2. + if in_dtype in ("fp4", "a8w4"): + run_mxfp4_moe_2stage( + tokens=tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + in_dtype=in_dtype, + x_fp32=x_fp32, + w1_fp32=w1_fp32, + w2_fp32=w2_fp32, + topk_ids=topk_ids, + topk_weights=topk_weights, + seed=seed, + ) + return + routing = build_routing_buffers( topk_ids=topk_ids, topk_weights=topk_weights, @@ -1985,6 +2017,255 @@ def _per_1x32_mxfp8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: return x_q, scale_bytes +# --------------------------------------------------------------------------- +# Layout-API MXFP4 MoE pipeline (mxfp4_moe_gemm_2stage), opus-sort only. +# Drives the a4w4 / a8w4 path of test_moe_gemm_2stage instead of mixed_moe: +# moe_sorting_flydsl (opus sort) -> gemm1 -> gemm2 (atomic scatter) -> bf16 out +# gemm1 gathers from sorted_token_ids (& 0xFFFFFF); gemm2 scatters via atomic add +# weighted by sorted_weights -- no fused-sort extras (m_indices / reverse_sorted). +# --------------------------------------------------------------------------- +def _mxfp4_shuffle_weight_a16w4(x, gate_up, NLane=16, KPack=16): + """CK a16w4 weight preshuffle (is_guinterleave path).""" + x_type = x.dtype + if hasattr(torch, "float4_e2m1fn_x2") and x_type == torch.float4_e2m1fn_x2: + x = x.view(torch.uint8) + E, N, K_pk = x.shape + if gate_up: + N = N // 2 + KLane = 64 // NLane + N0 = N // NLane + K0 = K_pk // (KLane * KPack) + if gate_up: + x_ = x.view(E, 2, N0, NLane, K0, KLane, KPack).permute(0, 2, 1, 4, 5, 3, 6) + else: + x_ = x.view(E, N0, NLane, K0, KLane, KPack).permute(0, 1, 3, 4, 2, 5) + return x_.contiguous().view(*x.shape).contiguous().view(x_type) + + +def _mxfp4_shuffle_scale_a16w4(src, E, gate_up): + """CK a16w4 e8m0 scale preshuffle (is_guinterleave path).""" + n_experts, k_ = src.shape + n_ = n_experts // E + K_Pack, N_Pack, N_Lane = 2, 2, 16 + K_Lane = 64 // N_Lane + K1 = k_ // K_Pack // K_Lane + N1 = n_ // N_Lane // N_Pack + if gate_up: + s = src.view(E, N_Pack, N1, N_Lane, K1, K_Pack, K_Lane).permute(0, 2, 4, 6, 3, 5, 1) + else: + s = src.view(E, N1, N_Pack, N_Lane, K1, K_Pack, K_Lane).permute(0, 1, 4, 6, 3, 5, 2) + return s.contiguous().view(*src.shape).contiguous() + + +def _mxfp4_a_scale_sorted_shuffled(asc, sti, cumsum, max_sorted, H, BM=32, BK=256): + """Torch reconstruction of moe_sort_scales: sort + CK-shuffle the e8m0 A-scale + by sorted row, exactly as gemm1 consumes it (opus gather: sti & 0xFFFFFF).""" + device = asc.device + MN_PACK = 2 + K_PACK = BK // 128 + C_M1 = BM // (16 * MN_PACK) + C_K1 = (H // 32) // (4 * K_PACK) + K_LANE, N_LANE = 4, 16 + DWORDS_PER_CHUNK = C_M1 * C_K1 * K_LANE * N_LANE + n_chunks = max_sorted // BM + actual_sorted = int(cumsum[0].item()) + actual_n_chunks = (actual_sorted + BM - 1) // BM + total_work = n_chunks * DWORDS_PER_CHUNK + sti_c = sti & 0x00FFFFFF + out = torch.zeros((total_work, 4), dtype=torch.uint8, device=device) + wid = torch.arange(total_work, device=device) + r = wid.clone() + n_lane = r % N_LANE + r //= N_LANE + k_lane = r % K_LANE + r //= K_LANE + ku = r % C_K1 + r //= C_K1 + mi = r % C_M1 + r //= C_M1 + chunk = r + valid_chunk = chunk < actual_n_chunks + M = asc.shape[0] + for ikxdl in range(K_PACK): + for im_a in range(MN_PACK): + sorted_row = chunk * BM + (mi * MN_PACK + im_a) * 16 + n_lane + rowok = (sorted_row < actual_sorted) & valid_chunk + srow = torch.clamp(sorted_row, max=max_sorted - 1) + stiv = sti_c[srow] + tid = torch.where((stiv < M) & rowok, stiv, torch.zeros_like(stiv)) + k_idx = ku * K_PACK * 4 + ikxdl * 4 + k_lane + byte = asc[tid.long(), k_idx.long()] + out[:, ikxdl * MN_PACK + im_a] = torch.where( + rowok, byte, torch.zeros_like(byte) + ) + return out.reshape(-1).contiguous() + + +def _u8v(t): + return ( + t.view(torch.uint8) + if (t is not None and t.element_size() == 1 and t.dtype != torch.uint8) + else t + ) + + +def run_mxfp4_moe_2stage( + *, + tokens, + model_dim, + inter_dim, + experts, + topk, + in_dtype, + x_fp32, + w1_fp32, + w2_fp32, + topk_ids, + topk_weights, + interleave=True, + seed=0, +): + """Run the layout-API MXFP4 MoE (opus sort -> gemm1 -> gemm2 atomic) and verify + against an independent dequant-MoE reference. Returns the bf16 output.""" + from tests.kernels.utils import fp4_utils + + device = x_fp32.device + NE, H, INTER, TOPK = experts, model_dim, inter_dim, topk + BM = 32 + is_f8 = in_dtype == "a8w4" + + # weights (fp4) + CK a16w4 preshuffle + w1q, w1s = _per_1x32_fp4_quant(w1_fp32) + w2q, w2s = _per_1x32_fp4_quant(w2_fp32) + w1u8 = _u8v(_mxfp4_shuffle_weight_a16w4(w1q, gate_up=interleave)) + w1sc = _u8v(_mxfp4_shuffle_scale_a16w4(w1s, NE, gate_up=interleave)) + w2u8 = _u8v(_mxfp4_shuffle_weight_a16w4(w2q, gate_up=False)) + w2sc = _u8v(_mxfp4_shuffle_scale_a16w4(w2s, NE, gate_up=False)) + + # opus sort (FlyDSL moe_sorting_kernel) + topk_ids_i32 = topk_ids.to(torch.int32) + topk_w_f32 = topk_weights.to(torch.float32) + max_padded = tokens * TOPK + NE * BM - TOPK + max_sorted = ((max_padded + BM - 1) // BM) * BM + sti = torch.empty(max_sorted, dtype=torch.int32, device=device) + swt = torch.empty(max_sorted, dtype=torch.float32, device=device) + sei = torch.empty(max_sorted // BM, dtype=torch.int32, device=device) + nv = torch.empty(2, dtype=torch.int32, device=device) + moe_buf = torch.empty((tokens, H), dtype=torch.bfloat16, device=device) + moe_sorting_flydsl(topk_ids_i32, topk_w_f32, sti, swt, sei, nv, moe_buf, NE, BM) + torch.cuda.synchronize() + cumsum = nv + n = int(cumsum[0].item()) + + # A quant (+ sorted/shuffled e8m0 A-scale) + hidden = x_fp32.to(torch.bfloat16) + if is_f8: + aq, asc = _per_1x32_mxfp8_quant(hidden) + aq = aq.view(torch.uint8).view(tokens, H).contiguous() + else: + aq, asc = _per_1x32_fp4_quant(hidden) + aq = aq.view(torch.uint8).view(tokens, H // 2).contiguous() + asc = asc.view(torch.uint8).view(tokens, H // 32).contiguous() + assh = _mxfp4_a_scale_sorted_shuffled(asc, sti, cumsum, max_sorted, H, BM=BM) + + # gemm1 -> intermediate (fp4 / fp8) + shuffled scale + out_dtype = "fp8" if is_f8 else "fp4" + inter_cols = INTER if is_f8 else INTER // 2 + isq = torch.zeros((max_sorted, inter_cols), device=device, dtype=torch.uint8) + isc_cols = INTER // 32 + isr = ( + (((max_sorted * ((2 * INTER) // 64) * 4) + isc_cols - 1) // isc_cols + 31) + // 32 * 32 + ) + iss = torch.zeros((isr, isc_cols), device=device, dtype=torch.uint8) + mxfp4_moe_gemm1( + a_quant=aq, + a_scale_sorted_shuffled=assh, + w1_u8=w1u8, + w1_scale_u8=w1sc, + sorted_expert_ids=sei, + cumsum_tensor=cumsum, + sorted_token_ids=sti, + inter_sorted_quant=isq, + inter_sorted_shuffled_scale=iss, + hidden_states=hidden, + n_tokens=tokens, + NE=NE, + D_HIDDEN=H, + D_INTER=INTER, + topk=TOPK, + BM=BM, + use_nt=True, + inline_quant=False, + interleave=interleave, + a_dtype=("fp8" if is_f8 else "fp4"), + out_dtype=out_dtype, + ) + torch.cuda.synchronize() + + # gemm2 (atomic): inter x w2 -> per-token weighted topk sum + out = torch.zeros((tokens, H), dtype=torch.bfloat16, device=device) + mxfp4_moe_gemm2( + inter_sorted_quant=isq, + inter_sorted_shuffled_scale=iss, + w2_u8=w2u8, + w2_scale_u8=w2sc, + sorted_expert_ids=sei, + cumsum_tensor=cumsum, + sorted_token_ids=sti, + sorted_weights=swt, + out=out, + M_logical=tokens, + max_sorted=max_sorted, + NE=NE, + D_HIDDEN=H, + D_INTER=INTER, + topk=TOPK, + BM=BM, + use_nt=False, + a_dtype=("fp8" if is_f8 else "fp4"), + ) + torch.cuda.synchronize() + + # reference: independent dequant MoE (opus gather: tok = sti & 0xFFFFFF) + if is_f8: + A = fp4_utils.fp8_e4m3_to_f32(aq.view(torch.float8_e4m3fn)).view(tokens, H) + else: + A = fp4_utils.mxfp4_to_f32(aq.view(torch.uint8)).view(tokens, H) + Asc = fp4_utils.e8m0_to_f32(asc.view(torch.uint8)) + A = (A.view(tokens, H // 32, 32) * Asc.unsqueeze(-1)).view(tokens, H) + W1 = fp4_utils.mxfp4_to_f32(w1q.view(torch.uint8)) + W1s = fp4_utils.e8m0_to_f32(w1s.view(torch.uint8)).view(NE, 2 * INTER, H // 32) + W1 = (W1.view(NE, 2 * INTER, H // 32, 32) * W1s.unsqueeze(-1)).view(NE, 2 * INTER, H) + W2 = fp4_utils.mxfp4_to_f32(w2q.view(torch.uint8)) + W2s = fp4_utils.e8m0_to_f32(w2s.view(torch.uint8)).view(NE, H, INTER // 32) + W2 = (W2.view(NE, H, INTER // 32, 32) * W2s.unsqueeze(-1)).view(NE, H, INTER) + + sti_c, sei_c, swt_c = sti[:n].cpu(), sei.cpu(), swt[:n].cpu() + ref = torch.zeros((tokens, H), dtype=torch.float32, device=device) + for r in range(n): + tok = int(sti_c[r].item()) & 0x00FFFFFF + if tok >= tokens: + continue + e = int(sei_c[r // BM].item()) + gate = A[tok] @ W1[e, :INTER].T + up = A[tok] @ W1[e, INTER : 2 * INTER].T + inter_r = torch.nn.functional.silu(gate) * up + ref[tok] += (inter_r @ W2[e].T) * float(swt_c[r].item()) + + cos = torch.nn.functional.cosine_similarity( + ref.reshape(-1), out.float().reshape(-1), dim=0 + ).item() + thr = 0.95 if is_f8 else 0.85 + logging.info( + "[mxfp4 moe %s %s] cos=%.4f n=%d (model_dim=%d inter=%d E=%d topk=%d)", + in_dtype, "il" if interleave else "sep", cos, n, model_dim, inter_dim, experts, topk, + ) + assert verify_output(out.to(torch.float32), ref, rtol=0.5, atol=0.5, logits_diff_threshold=1) + assert cos > thr, f"{in_dtype} cos={cos:.4f} <= {thr}" + return out + + # Test Helpers for MoE GEMM2 Mode Comparison def _make_reduce_mode_compile_fn( use_flydsl_reduce: bool = True, use_valid_mask: bool = False, scale_dtype: str = "f32"