MoonMath AI has released an open-source bf16 forward attention kernel for AMD’s MI300X GPU that outperforms AMD’s own AITER v3 library on every tested shape and rounding mode. The kernel, written in HIP and released under the MIT licence, achieves geometric mean speedups of 1.18 times with RTNE rounding, 1.15 times with RTNA, and 1.08 times with RTZ, with individual shapes showing up to 1.26 times improvement.
The kernel is already in production use: a pull request integrating it into SGLang sped up Wan2.1 video diffusion by 1.23 times with no quality regression.
The kernel computes fused softmax(QK^T / sqrt(d)) dot V for bf16 data with a head dimension of 128. It supports both BSHD and BHSD memory layouts and any sequence length, including cross-attention.
The key innovation is a technical one that avoids a common trade-off in GPU kernel development. When writing high-performance GPU kernels, developers typically face a choice between compiler intrinsics, which let the compiler manage register allocation but can reorder or rename instructions in ways that subvert performance tuning, and raw inline assembly, which gives full opcode control but forces manual register management that is error-prone and hard to maintain.
MoonMath’s solution uses one-instruction assembly wrappers with the `__device__ __forceinline__` directive. The critical trick is the `”+v”(c)` constraint, which ties the accumulator input and output to the same vector general-purpose register, ensuring the compiler does not emit a separate copy instruction.
Architecture: eight waves, two groups
The kernel uses eight waves per block in two groups of four, rather than the textbook four-wave layout. The two groups run the same Q dot K, softmax, output-equals-P-dot-V sequence but offset by a phase: one group saturates the matrix core while the other runs softmax and issues memory loads. Two barriers per iteration synchronise the handoff.
The approach echoes FlashAttention-3’s alternation between matrix multiply and softmax, but does not copy FlashAttention-3’s producer-consumer warp split because CDNA3 architecture already handles memory moves asynchronously.
Data placement decisions
The team made deliberate choices about where to place each data structure:
- K goes in LDS (shared memory), double-buffered at 32 KiB, shared by all eight waves
- V stays hot in L1 cache, resident and prefetched, since it is reread across the PV step
- Q and accumulators live in vector general-purpose registers, read every iteration and never reloaded
Two later refinements improved performance further. A third Q tile increased data reuse per loaded K-V tile. A Flash-Decoding-style tail KV split rescued the stranded fractional round across the MI300X’s 304 compute units, which fixed the RTZ-mode performance that had initially lagged behind AITER.
Benchmark results
The kernel was tested on MI300X across three rounding modes. Representative results at the (batch, heads, sequence, head-dim) shapes most relevant to inference workloads:
| Shape | Round | MoonMath (ms) | AITER v3 (ms) | Speedup |
|—|—|—|—|—|
| (2, 24, 8192, 128) | RTNE | 3.083 | 3.792 | 1.23x |
| (2, 24, 16384, 128) | RTNE | 11.670 | 14.691 | 1.26x |
| (2, 24, 32768, 128) | RTNA | 44.440 | 52.363 | 1.18x |
| (1, 16, 131072, 128) | RTNE | 232.517 | 269.278 | 1.16x |
Every finite output is within one bf16 unit in the last place of AITER’s output. NaN and Infinity results are bit-identical. The kernel is deterministic.
Limits
The kernel has constraints. It does not support causal masking, grouped-query attention, or variable-length batching. It outputs bf16 only and runs exclusively on the gfx942 architecture (MI300X). These limitations are consistent with its stated purpose: a focused, high-performance kernel for the most common diffusion and inference attention patterns on AMD hardware.
The kernel is installable via pip and launches on the caller’s CUDA stream for overlap with other operations. The code and full benchmark suite are available on GitHub under the MIT licence.
Sources: MarkTechPost (June 22)