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之上，讓你無需經歷通常的組合語言級繁瑣，即可撰寫生產級內核。