The MoonMath AI 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 cloud provider.
What the MoonMath AI Kernel Does and Why It Matters
Attention is the fused softmax(QKT/√d)·V 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.
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 — 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.
The Core Technique: One-Instruction Assembly Wrappers
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.
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.
Architecture: Eight Waves, Two Groups, Two Barriers
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·K, softmax, O += P·V 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 — one at the phase handoff and one at the iteration boundary — with per-counter waits handling the rest of the synchronization.
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.
Memory Placement Strategy: Where Data Lives and Why
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.
The team chose the 16×16×16 MFMA over the 32×32×8 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 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 — each decision cascades into the next.
Benchmark Results: Beating AITER v3 Across the Board
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).
Selected results:
- (2, 24, 8192, 128), RTNE: MoonMath 3.083 ms vs AITER 3.792 ms — 1.23× faster
- (2, 24, 16384, 128), RTNE: 11.670 ms vs 14.691 ms — 1.26× faster
- (2, 24, 32768, 128), RTNA: 44.440 ms vs 52.363 ms — 1.18× faster
- (1, 16, 131072, 128), RTNE: 232.517 ms vs 269.278 ms — 1.16× faster
Geometric means across the full sweep favor MoonMath by 1.18× (RTNE), 1.15× (RTNA), and 1.08× (RTZ) versus AITER v3. Against Modular MAX, geomeans range from 1.44× to 1.49×, with per-shape speedups reaching 1.59×. The RTZ mode is the tightest race; the (4, 16, 16384) RTZ shape improved from 0.95× to 1.07× after the tail KV split was implemented.
Real-World Validation: Wan2.1 Video Diffusion
The kernel has already been tested in production. A real SGLang pull request used it to speed up Wan2.1 video diffusion by 1.23× with no quality regression. This demonstrates that the performance gains translate beyond synthetic benchmarks into actual model workloads.
What This Means for Developers
For developers working with AMD MI300X hardware, this kernel offers an immediate performance uplift for bf16 attention with head dimension 128, available now under a permissive MIT license. The code demonstrates that HIP-based kernels can outperform hand-tuned assembly when combined with careful memory placement and pipeline design. The one-instruction asm wrapper technique is a practical pattern that other kernel developers can adopt to gain instruction-level control without sacrificing compiler register management. Anyone running transformer inference or training on MI300X hardware should evaluate this kernel against their existing attention implementation, particularly for long-sequence workloads where the memory placement advantages compound most significantly.
