HazyResearch/ThunderKittens

ThunderKittens – 高性能AI内核的CUDA-DSL

是什么 – ThunderKittens 是一个 仅头文件 的 C++/CUDA 库，可让你以紧凑、基于分块（tile）的方式编写深度学习内核（例如 GEMM、FlashAttention）。它抽象了低级GPU细节（张量核调用、异步复制、共享内存分组），同时保持与硬件的紧密联系，使生成的内核在现代NVIDIA GPU（H100、Blackwell、Vera Rubin）上接近理论峰值速度运行。

为何重要 – 大型语言模型的训练和推理受限于矩阵乘法和注意力内核的执行速度。手工编写此类内核容易出错，且需要对NVIDIA编程模型有深入理解。ThunderKittens 提供了一个小巧、可扩展的DSL，具备以下优势：

• 多数内核的源码可控制在100行以内。
• 编译时保证正确的布局处理。
• 无需编写汇编即可使用最新的张量核指令（WGMMA、TCGEN05、MXFP8、NVFP4）。
• 可直接从C++调用，或通过PyBind11封装为PyTorch使用。

核心概念

概念	作用
分块原语	操作16×16及以上块（寄存器分块、共享分块、向量），自然映射到张量核通道。
线程束 / 线程束组	函数默认针对单个线程束（32线程）编写；4个线程束的协作组（线程束组）可暴露异步矩阵乘加指令。
TMA / 异步复制	内置辅助函数通过NVIDIA的TMA（张量内存访问）机制隐藏延迟，实现加载/存储。
静态布局检查	模板编码数据类型、形状和内存布局，不匹配的操作在编译时即可捕获。
加载-存储-计算-完成模板	推荐模式，将内存移动与计算重叠，最大化占用率。

典型工作流程

1. 克隆 仓库并将在CUDA源码中包含 kittens.cuh。
2. 定义内核 使用提供的 matmul_layout / matmul_template（或其他原语）——你只需填写 producer、consumer 和可选的 common_setup 回调函数。
3. 编译 使用提供的Makefile（CUDA 12.8+，C++20）。每个内核位于 kernels/ 下的独立文件夹中，可独立构建。
4. 从Python调用（可选）——make 后，一个小型PyBind11包装器可让你从PyTorch 2.8+调用该内核。
5. 基准测试 / 测试 —— 正确性测试和性能脚本与每个内核同目录存放。

示例：H100上实现855 TFLOPs的矩阵乘法

#include "kittens.cuh"
#include "prototype.cuh"
using namespace kittens;
using namespace kittens::prototype;

// 布局定义（分块、全局指针等）
template<int M_BLOCK, int N_BLOCK>
struct matmul_layout { … };

// 连接producer、consumer和common_setup的内核模板
template<int _M_BLOCK=2, int _N_BLOCK=4, int _SUPER_M=12>
struct matmul_template { … };

完整源码（约100行）在README中展示，编译后生成的内核可达到H100理论峰值的约86%。

支持的硬件 – 主要支持NVIDIA Hopper（H100）和Blackwell（B200）GPU。自2026年9月起支持新型Vera Rubin GPU。Ampere可运行但不再更新。AMD用户可参考姊妹项目 HipKittens。

安装检查清单

• CUDA 12.8+（设置 CUDA_HOME，更新 PATH/LD_LIBRARY_PATH）。
• C++20编译器（gcc-11 或 clang-11）。
• （可选）若需Python绑定，安装PyTorch 2.8+ 和 PyBind11。
• 克隆仓库，包含头文件，并使用提供的Makefile编译内核。

使用场景 – 在Together AI、Jump Trading、Cursor等公司用于生产级训练和推理。斯坦福大学Hazy研究实验室也内部使用。

学习资源

• ThunderKittens手册 – 简明指南，涵盖分块、作用域和API约定。
• 教育性内核系列 – kernels/gemm/educational_h100 逐步讲解GEMM实现。
• 深度剖析博客 – Hamza Elshafie（2026年5月）对DSL内部结构的解析。
• 入门文档 – 新手用Google文档（README中链接）。

演示 – demos/ 文件夹包含可直接运行的LLM训练和推理示例（如Qwen、Llama、LoLCATS）。演示如何将ThunderKittens内核集成到PyTorch工作流中，并在Hugging Face的8B模型上运行。

---

总结 – 若你需要为自定义注意力、GEMM或其他矩阵密集型操作榨取NVIDIA张量核的每一丝性能，ThunderKittens提供了一个小巧、类型安全的DSL，直接建立在CUDA之上，让你无需经历通常的汇编级繁琐，即可编写生产级内核。