ROCm: port aiter FlashAttention and fused kernels to gfx1201 (RDNA4) via rocWMMA

[kmbandy]

Overview

AMD's aiter library contains highly optimized FA3 FlashAttention, fused softmax+attention, and KV cache quantization kernels targeting MI300X (CDNA3). While MFMA-based kernels are not portable to RDNA4, the algorithmic approaches are worth adapting for gfx1201 (Radeon AI Pro R9700) using rocWMMA.

What's worth porting

High value

• FlashAttention-3 tiling logic — aiter's FA3 implementation is a clean reference. The tiling strategy, softmax accumulation, and causal masking are architecture-agnostic. Needs retuning for:
  • RDNA4 32-wide wavefronts (vs CDNA 64-wide)
  • WMMA 16x16x16 tile dimensions (vs MFMA shapes)
  • GDDR6 bandwidth characteristics (~960 GB/s vs HBM 5+ TB/s) — different tiling assumptions
• Fused softmax + attention kernels — fusion strategy is valid on WMMA, just swap MFMA calls for WMMA equivalents via rocWMMA

Usable as-is or near as-is

• KV cache quantization kernels (int8/fp8) — some don't use matrix cores at all, pure vector ops that run on RDNA4 without modification

Not worth porting

• GEMM kernels tuned for HBM bandwidth — access patterns assume HBM throughput, GDDR6 tiling assumptions are fundamentally different

Relationship to existing work

This overlaps with the planned fattn-wmma-f16.cuh rocWMMA Flash Attention improvement for gfx12 targets. aiter should be used as the reference implementation rather than building from scratch — the FA3 block structure is well-validated on hardware.

Hardware context

• gfx1201 (RDNA4), 32GB VRAM, GDDR6 ~960 GB/s
• gfx1030 (RDNA2) secondary GPU — any WMMA work is gfx1201-only, gfx1030 does not have WMMA
• ROCm 7.2

References

• https://github.com/ROCm/aiter
• rocWMMA: https://github.com/ROCm/rocWMMA

[github-actions]

This issue was closed because it has been inactive for 14 days since being marked as stale.