HazyResearch/HipKittens

HipKittens 是什么？

HipKittens 是一个开源的 C++ 库，提供低级构建块，用于在 AMD GPU 上编写高性能 AI 内核（CDNA3、CDNA4、CDNA5）。它属于 ThunderKittens 生态系统，并针对与 Triton 或 Torch‑Inductor 等项目所使用的同类 tensor‑core 风格编程，但它是为 AMD 的 ROCm 栈（HIP、hipcc）而编写的。

核心概念

概念	含义
Tile 原语	直接映射到 AMD tensor‑core 单元的小型固定大小块。该库提供合并、避免 bank conflict，并使用原生 tensor‑core 布局的内存访问辅助工具。
Python 风格的大规模计算	看起来像 Python 函数（例如 `gemm`、`attention`）的轻量 C++ 包装，但会编译成内联汇编和 HIP 调用。
异步加载/存储	发出非阻塞加载到共享内存的辅助工具，让硬件隐藏延迟。
调度模式	两种可重复使用的模式 – 8‑wave ping‑pong 和 4‑wave interleave – 重叠计算和内存，以保持 GPU 流水线忙碌。

你可以用它做什么？

BF16 GEMM 内核（例如 8192×8192×8192 矩阵乘法），可与 AITER 和 PyTorch 进行基准测试。
Attention 内核（MHA、GQA、因果/非因果），支持前向和反向传播，并可配置批次大小、序列长度、头数和头维度。
内存密集型内核，例如旋转位置嵌入和融合层归一化。
用于 BERT 或 LLaMA 的端到端训练脚本，可将 HipKittens attention 内核插入 PyTorch 训练循环。

所有这些都使用 ROCm 的 hipcc 编译，并可在提供的 AMD Docker 镜像中运行。

谁可能使用它？

研究人员探索仅限 AMD 或多芯片 AI 性能，并需要轻量、贴近硬件的内核库。
工程师为 AMD GPU 构建自定义 transformer 后端，尤其是当他们想要与现有后端（如 AITER、Triton 或 HipBLASLT）进行比较时。
学生学习 AMD 硬件上的 tensor‑core 编程；该库的单元测试和示例 make‑file 可作为具体的学习材料。

如何开始？

克隆仓库并拉取匹配的 ROCm Docker（README 列出了 MI35x 和 MI30x GPU 的两个镜像）。
进入容器，source env.src 以设置环境变量，并安装 AITER 基线（如果您想运行比较基准测试）。
构建内核（例如 cd kernels/cdna4/gemm/bf16fp32 && make && python bench.py）。bench.py 脚本将运行 HipKittens 内核、AITER 参考和 PyTorch 的实现，并打印计时数字。
运行单元测试以测试您的 GPU 世代（cd HipKittens/tests/unit/cdna4 && make && ./unit_tests）。
探索分析脚本在 HipKittens/analysis 下，以重现论文的性能图表，或查看 training/ 以进行快速的 BERT/LLaMA 训练运行。

何时有用？

当您在 AMD GPU 上进行训练或推理（MI300X、MI350X、MI455 等）且需要比通用 ROCm 库更快的内核时。
当您想要 实验新的调度或内存访问模式，而不想自己编写原始 HIP 汇编时。
当您正在构建 多芯片 AI 栈，且需要一个遵循其他 DSL（例如 Triton）所用相同 tile 抽象，但能在 AMD 硬件上运行的参考实现时。

联系与社区：William Hu (willhu@stanford.edu) 和 Simran Arora (simran@cs.stanford.edu) 为 HipKittens 用户运营一个 Discord 频道。该仓库还链接到 arXiv 论文、博客文章和 YouTube 讲座，以提供更深入的背景。