HipKittens：快速而猛烈的 AMD 内核
机器学习
2025-11-12 v1
摘要
AMD GPU 提供了领先的计算和内存带宽；然而，峰值性能 AMD 内核是用原始汇编编写的。为解决将 AI 算法映射到硬件的困难，近期工作提出使用类似 PyTorch 的 C++ 嵌入式和 PyTorch 风格的领域特定语言（DSL），如 ThunderKittens（TK），以简化在 NVIDIA 硬件上开发高性能 AI 内核的工作。我们探讨了此类原语——用于显式基于块的编程，以优化内存访问并在工人之间进行细粒度异步执行——是否具有 NVIDIA 特定性或通用性。我们提供了首个详尽的 AMD AI 内核编程原语研究，并将这些见解封装在 HipKittens（HK）编程框架中。我们发现，先前 DSL 中使用的基于块的抽象可推广到 AMD GPU，但需要重新构思用于实现这些抽象的算法。我们在 CDNA3 和 CDNA4 AMD 平台上验证了 HK 原语。在评估中，HK 内核与 AMD 手优化汇编内核在 GEMM 和注意力方面持久竞争，并始终超过编译器基准。此外，汇编难以扩展到 AI 工作负载的广度；反映这一点，在某些情况下，HK 超过所有可用的内核基准表现为 （例如 注意力、GQA 反向传播、内存受限内核）。这些发现为单一的、基于块的高性能 AI 内核软件层铺平了道路，该层可跨 GPU 供应商进行翻译。HipKittens已发布于：https://github.com/HazyResearch/HipKittens。
引用
@article{arxiv.2511.08083,
title = {HipKittens: Fast and Furious AMD Kernels},
author = {William Hu and Drew Wadsworth and Sean Siddens and Stanley Winata and Daniel Y. Fu and Ryann Swann and Muhammad Osama and Christopher Ré and Simran Arora},
journal= {arXiv preprint arXiv:2511.08083},
year = {2025}
}