- MI450 GPU introduces new hardware features like TDM, larger LDS, and workgroup clusters that shift the performance bottleneck from compute to memory for attention decode workloads.
- Attention decode kernels are optimized using Gluon, a Triton-based DSL with explicit tensor layouts, allowing fine-grained control over generated instructions.
- Key optimizations include: optimizing WMMA tensor layouts, using TDM for efficient data loading with direct path, pipelining with triple buffering to overlap memory and computation, and using Split-k with workgroup clusters to boost parallelism.
- The optimized kernel achieves up to 85% of peak HBM bandwidth on MI450 (17.10 TB/s for GQA, 16.65 TB/s for MQA) for long-context decode.
- The kernel source code is open-sourced and can be reproduced using provided commands with ROCm 7.14.0 and PyTorch 2.11.0.
- Future work includes extending these optimizations to other attention variants (e.g., MLA/DSA) and integrating into production LLM inference frameworks.