HazyResearch/HipKittens

What is HipKittens?

HipKittens is an open‑source C++ library that provides low‑level building blocks for writing high‑performance AI kernels on AMD GPUs (CDNA3, CDNA4, CDNA5). It sits in the ThunderKittens ecosystem and targets the same kind of tensor‑core‑style programming that projects like Triton or Torch‑Inductor use, but it is written for AMD’s ROCm stack (HIP, hipcc).

Core ideas

Idea	What it means
Tile primitives	Small, fixed‑size blocks that map directly to AMD tensor‑core units. The library supplies memory‑access helpers that are coalesced, avoid bank conflicts, and use the native tensor‑core layout.
Python‑inspired bulk compute	Thin C++ wrappers that look like Python functions (e.g., gemm, attention) but compile down to inline assembly and HIP calls.
Async loads/stores	Helpers that issue non‑blocking loads into shared memory, letting the hardware hide latency.
Scheduling patterns	Two reusable patterns – 8‑wave ping‑pong and 4‑wave interleave – that overlap compute and memory to keep the GPU pipelines busy.

What can you do with it?

• BF16 GEMM kernels (e.g., 8192×8192×8192 matrix multiply) that can be benchmarked against AITER and PyTorch.
• Attention kernels (MHA, GQA, causal/non‑causal) for both forward and backward passes, with configurable batch size, sequence length, head count, and head dimension.
• Memory‑bound kernels such as rotary‑position embeddings and fused layer‑norm.
• End‑to‑end training scripts for BERT or LLaMA that plug HipKittens attention kernels into a PyTorch training loop.

All of these are compiled with ROCm’s hipcc and can be run inside the provided AMD Docker images.

Who might use it?

• Researchers exploring AMD‑only or multi‑silicon AI performance and who need a lightweight, “close‑to‑metal” kernel library.
• Engineers building custom transformer back‑ends for AMD GPUs, especially when they want to compare against existing back‑ends like AITER, Triton, or HipBLASLT.
• Students learning how tensor‑core programming works on AMD hardware; the library’s unit tests and example make‑files serve as concrete learning material.

How to get started?

1. Clone the repo and pull the matching ROCm Docker (the README lists two images for MI35x and MI30x GPUs).
2. Enter the container, source env.src to set up environment variables, and install the AITER baseline if you want to run comparative benchmarks.
3. Build a kernel (e.g., cd kernels/cdna4/gemm/bf16fp32 && make && python bench.py). The bench.py script will run the HipKittens kernel, the AITER reference, and PyTorch’s implementation and print timing numbers.
4. Run unit tests for your GPU generation (cd HipKittens/tests/unit/cdna4 && make && ./unit_tests).
5. Explore the analysis scripts under HipKittens/analysis to reproduce the paper’s performance plots, or look at training/ for a quick BERT/LLaMA training run.

When is it useful?

• When you are training or inference on AMD GPUs (MI300X, MI350X, MI455, etc.) and need kernels that are faster than the generic ROCm libraries.
• When you want to experiment with new scheduling or memory‑access patterns without writing raw HIP assembly yourself.
• When you are building a multi‑silicon AI stack and need a reference implementation that follows the same tile‑based abstraction used by other DSLs (e.g., Triton) but works on AMD hardware.

---

Contact & community: William Hu (willhu@stanford.edu) and Simran Arora (simran@cs.stanford.edu) run a Discord channel for HipKittens users. The repo also links to the arXiv paper, blog posts, and a YouTube lecture for deeper background.