{"id":57623,"date":"2026-06-22T10:33:00","date_gmt":"2026-06-22T14:33:00","guid":{"rendered":"https:\/\/overcentral.com\/en\/?p=57623"},"modified":"2026-06-22T10:33:00","modified_gmt":"2026-06-22T14:33:00","slug":"moonmath-hip-kernel-mi300x-aiter-benchmark","status":"publish","type":"post","link":"https:\/\/overcentral.com\/en\/moonmath-hip-kernel-mi300x-aiter-benchmark\/","title":{"rendered":"MoonMath AI Open-Sources HIP Kernel That Beats AITER v3 on MI300X"},"content":{"rendered":"

The MoonMath AI<\/a> team has open-sourced a bf16 forward attention kernel for AMD’s MI300X GPU, achieving significant performance gains over AMD’s own optimized AITER v3 kernel across all tested configurations. Written entirely in HIP rather than hand-tuned assembly, the kernel is available under the MIT license and represents a notable advance in open-source GPU kernel development for AMD’s CDNA3 architecture. Bare-metal access for development and testing was provided by HotAisle, an AMD<\/a> cloud provider.<\/p>\n

What the MoonMath AI Kernel Does and Why It Matters<\/h2>\n

Attention is the fused softmax(QKT<\/sup>\/\u221ad)\u00b7V operation at the heart of every transformer model. The MI300X is AMD’s CDNA3 data-center GPU targeting the gfx942 instruction set, and this kernel runs exclusively on that hardware. It computes forward attention in bf16 with a fixed head dimension of 128, accepting inputs in either BSHD or BHSD layout without requiring a transpose. Any sequence length is supported, including cross-attention.<\/p>\n

The kernel has real limitations: no causal mask, no grouped-query attention, and no varlen batching. Outputs are in bf16, and it runs only on gfx942 hardware. Numerics are tightly controlled \u2014 all three rounding modes (RTNE, RTNA, RTZ) match AITER’s per-mode rules, every finite output sits within 1 bf16 ULP of AITER, and NaN and Inf handling is bit-identical and deterministic.<\/p>\n

The Core Technique: One-Instruction Assembly Wrappers<\/h2>\n

The central innovation avoids a familiar dilemma in GPU kernel development. Compiler intrinsics keep code tidy but let the compiler reorder or rename operands. Raw inline assembly gives full control but forces manual register and address management. MoonMath wraps exactly one instruction in a __device__ __forceinline__ function, using extended asm constraints to describe operands. The team picks the opcode; the compiler still allocates registers and tracks data flow.<\/p>\n

The key pattern uses a “+v”(c) constraint that ties the accumulator input and output to the same VGPR, eliminating copy instructions. This keeps the kernel close to ordinary HIP while still steering the machine one instruction at a time. The approach demonstrates that precise control over instruction selection does not require abandoning compiler register allocation.<\/p>\n

Architecture: Eight Waves, Two Groups, Two Barriers<\/h2>\n

A CDNA3 compute unit has four SIMD units, but MoonMath runs eight waves per block in two groups of four. Both groups execute the same Q\u00b7K, softmax, O += P\u00b7V sequence but are offset by a phase. While one group saturates the matrix core, the other runs softmax and issues loads. They then swap, keeping the matrix core continuously busy. Two s_barrier instructions bound each iteration \u2014 one at the phase handoff and one at the iteration boundary \u2014 with per-counter waits handling the rest of the synchronization.<\/p>\n

This approach echoes FlashAttention-3’s matmul and softmax alternation but without a dedicated producer warp. On CDNA3 every memory move is already asynchronous, so a separate producer wave is unnecessary.<\/p>\n

Memory Placement Strategy: Where Data Lives and Why<\/h2>\n

Most of the speedup comes from deliberate memory placement. K streams from HBM into LDS in a double-buffered 32 KiB buffer shared by all eight waves. V is kept hot in L1, reread on every PV matmul. Q and accumulators reside in registers, read every iteration and never reloaded.<\/p>\n

The team chose the 16\u00d716\u00d716 MFMA over the 32\u00d732\u00d78 variant. Both offer identical throughput, but the smaller tile accumulates into only 4 fp32 elements per lane instead of 16. Lower accumulator pressure leaves room for deeper prefetch and a third Q tile (3Q), which raises data reuse per loaded K and V tile. A Flash-Decoding-style tail KV<\/a> split rescues the stranded fractional round across MI300X’s 304 compute units. Moving V to L1 freed the LDS that the third Q tile then fills \u2014 each decision cascades into the next.<\/p>\n

Benchmark Results: Beating AITER v3 Across the Board<\/h2>\n

Tests ran on MI300X in bf16 with head dimension 128. Each shape was measured at three rounding modes: RTNE (round to nearest even), RTNA (round to nearest, ties away from zero, AITER’s default), and RTZ (truncate toward zero, AITER’s fastest mode).<\/p>\n

Selected results:<\/p>\n