HazyResearch/ThunderKittens

ThunderKittens – A CUDA‑DSL for high‑performance AI kernels

What it is – ThunderKittens is a header‑only C++/CUDA library that lets you write deep‑learning kernels (e.g., GEMM, FlashAttention) as compact, tile‑oriented code. It abstracts low‑level GPU details (tensor‑core calls, async copies, shared‑memory banking) while staying close to the hardware, so the generated kernels run at near‑theoretical speed on modern NVIDIA GPUs (H100, Blackwell, Vera Rubin).

Why it matters – Training and inference of large language models are limited by how fast matrix‑multiply and attention kernels can be executed. Hand‑crafting such kernels is error‑prone and requires deep knowledge of NVIDIA’s programming model. ThunderKittens provides a small, extensible DSL that:

• Keeps the source under 100 lines for many kernels.
• Guarantees correct layout handling at compile time.
• Exposes the newest tensor‑core instructions (WGMMA, TCGEN05, MXFP8, NVFP4) without forcing you to write assembly.
• Works directly from C++ or can be wrapped for PyTorch via PyBind11.

Key concepts

Concept	What it does
Tile primitives	Operate on 16×16‑or‑larger blocks (register tiles, shared tiles, vectors) that map naturally to tensor‑core lanes.
Warp / Warp‑group	Functions are written for a single warp (32 threads) by default; collaborative groups of 4 warps (warp‑group) expose async matrix‑multiply‑accumulate instructions.
TMA / async copies	Built‑in helpers hide latency by loading/storing through NVIDIA’s TMA (Tensor Memory Access) mechanism.
Static layout checking	Templates encode data type, shape, and memory layout, so mismatched operations are caught at compile time.
Load‑Store‑Compute‑Finish template	A recommended pattern that overlaps memory movement with compute, maximizing occupancy.

Typical workflow

1. Clone the repo and include kittens.cuh in your CUDA source.
2. Define a kernel using the provided matmul_layout / matmul_template (or other primitives) – you only need to fill in the producer, consumer, and optional common_setup callbacks.
3. Compile the kernel with the supplied Makefile (CUDA 12.8+, C++20). Each kernel lives in its own folder under kernels/ and can be built independently.
4. Call from Python (optional) – after make, a small PyBind11 wrapper lets you invoke the kernel from PyTorch 2.8+.
5. Benchmark / test – correctness tests and performance scripts are co‑located with each kernel.

Example: 855 TFLOPs matrix multiply on an H100

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

// layout definition (tiles, global pointers, etc.)
template<int M_BLOCK, int N_BLOCK>
struct matmul_layout { … };

// kernel template that wires together producer, consumer, and common setup
template<int _M_BLOCK=2, int _N_BLOCK=4, int _SUPER_M=12>
struct matmul_template { … };

The full source (≈100 lines) is shown in the README and compiles to a kernel that reaches ~86 % of the H100’s theoretical peak.

Supported hardware – Primarily NVIDIA Hopper (H100) and Blackwell (B200) GPUs. Newer Vera Rubin GPUs are supported as of Sep 2026. Ampere works but receives no further updates. AMD users can look at the sister project HipKittens.

Installation checklist

• CUDA 12.8+ (set CUDA_HOME, update PATH/LD_LIBRARY_PATH).
• C++20 compiler (gcc‑11 or clang‑11).
• (Optional) PyTorch 2.8+ and PyBind11 if you want Python bindings.
• Clone, include the header, and compile kernels with the provided Makefiles.

Where it’s used – Production‑scale training and inference at companies such as Together AI, Jump Trading, and Cursor. The library is also used internally at Stanford’s Hazy Research Lab.

Learning resources

• ThunderKittens Manual – short guide covering tiles, scopes, and API conventions.
• Educational kernel series – kernels/gemm/educational_h100 walks through a step‑by‑step GEMM implementation.
• Deep‑dive blog – Hamza Elshafie’s dissection (May 2026) explains the DSL internals.
• Onboarding doc – a Google‑Doc (link in README) for newcomers.

Demos – The demos/ folder contains ready‑to‑run examples for LLM training and inference (e.g., Qwen, Llama, LoLCATS). They demonstrate how to plug a ThunderKittens kernel into a PyTorch workflow and run on an 8‑B model from Hugging Face.

---

Bottom line – If you need to squeeze every ounce of performance out of NVIDIA tensor cores for custom attention, GEMM, or other matrix‑heavy operations, ThunderKittens gives you a tiny, type‑safe DSL that sits directly on top of CUDA, letting you write production‑grade kernels without the usual assembly‑level hassle.