Attention Decode on AMD MI450 GPUs: A Gluon Kernel Optimization Guide
2 days ago
- MI450 GPU引入了全新硬件特性，如TDM（三角数据传输模式）、更大的LDS（本地数据共享）和工作组集群，将注意力解码工作负载的性能瓶颈从计算转移到内存。
- 注意力解码内核通过使用基于Triton的领域特定语言Gluon进行优化，该语言具有显式张量布局，允许对生成的指令进行精细控制。
- 关键优化包括：优化WMMA张量布局；利用TDM通过直接路径实现高效数据加载；采用三重缓冲流水线以重叠内存与计算；以及使用带有工作组集群的Split-k技术提升并行度。
- 优化后的内核在MI450上针对长上下文解码实现了高达峰值HBM带宽的85%（GQA为17.10 TB/s，MQA为16.65 TB/s）。
- 内核源代码已开源，可使用提供的命令配合ROCm 7.14.0和PyTorch 2.11.0进行复现。
- 未来工作包括将这些优化扩展到其他注意力变体（如MLA/DSA），并集成到生产级LLM推理框架中。