arXiv:2511.08083cs.LG2025-11被引 2

HipKittens让AMD GPU也能高效运行AI算子,性能媲美手写汇编。

HipKittens: Fast and Furious AMD Kernels

  • 基于分块内存访问与细粒度异步执行,适配AMD GPU的编程框架
  • 在CDNA3/CDNA4上实现GEMM和注意力层性能接近手写汇编
  • 支持多场景通用性,比编译器基线快1.2到2.4倍

AMD GPU具备顶尖的计算能力和内存带宽;然而,其峰值性能的算子仍需用原始汇编编写。为降低将AI算法映射到硬件的难度,近期工作提出类PyTorch的领域特定语言(如ThunderKittens),以简化NVIDIA硬件上的高性能算子开发。本文首次系统研究了这些编程原语在AMD GPU上的适用性——分块编程、优化内存访问与细粒度异步执行是否仅限于NVIDIA。我们封装经验,提出HipKittens(HK)编程框架。发现分块抽象可泛化至AMD,但需重构底层算法。在CDNA3和CDNA4平台上验证,HK算子在GEMM和注意力计算中性能媲美手写汇编,并持续优于编译器基线。尤其在部分内存密集型场景下(如d=64注意力、GQA反向传播),性能提升达1.2–2.4倍。这为跨厂商统一高性能算子软件层提供可能。HipKittens已开源:https://github.com/HazyResearch/HipKittens。

原文摘要 · Abstract (English)

AMD GPUs offer state-of-the-art compute and memory bandwidth; however, peak performance AMD kernels are written in raw assembly. To address the difficulty of mapping AI algorithms to hardware, recent work proposes C++ embedded and PyTorch-inspired domain-specific languages like ThunderKittens (TK) to simplify high performance AI kernel development on NVIDIA hardware. We explore the extent to which such primitives -- for explicit tile-based programming with optimized memory accesses and fine-grained asynchronous execution across workers -- are NVIDIA-specific or general. We provide the first detailed study of the programming primitives that lead to performant AMD AI kernels, and we encapsulate these insights in the HipKittens (HK) programming framework. We find that tile-based abstractions used in prior DSLs generalize to AMD GPUs, however we need to rethink the algorithms that instantiate these abstractions for AMD. We validate the HK primitives across CDNA3 and CDNA4 AMD platforms. In evaluations, HK kernels compete with AMD's hand-optimized assembly kernels for GEMMs and attention, and consistently outperform compiler baselines. Moreover, assembly is difficult to scale to the breadth of AI workloads; reflecting this, in some settings HK outperforms all available kernel baselines by $1.2-2.4\times$ (e.g., $d=64$ attention, GQA backwards, memory-bound kernels). These findings help pave the way for a single, tile-based software layer for high-performance AI kernels that translates across GPU vendors. HipKittens is released at: https://github.com/HazyResearch/HipKittens.

AMD GPUAI算子分块编程高性能计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。