Together AI's ThunderKittens: A New Language for Efficient GPU Kernels

Together AI has released ThunderKittens, a domain-specific language (DSL) for writing optimized GPU code. The project is positioned as an attempt to simplify neural network development, which today demands deep knowledge of hardware architecture.

Why it was needed

Writing efficient GPU code is black magic for most ML engineers. Chips like the NVIDIA H100 have dedicated tensor cores that deliver 94% of the total compute power. But using them means writing CUDA, a complex low-level language that only a handful of people in the industry understand.

In 2023, the popular FlashAttention2 showed that working with attention in transformers could be simplified significantly. But that is only a special case: an algorithm for one specific operator. Developers need a universal way to write fast code without the hellish CUDA, one that works across different kernels.

How ThunderKittens works

ThunderKittens sits between two extremes. On one side is pure CUDA: very fast, but very complex, with a high barrier to entry. On the other is Triton, which hides GPU details and makes life easier, but sometimes cannot squeeze the maximum performance out of the hardware.

ThunderKittens offers a third path. Its API resembles PyTorch, which is familiar to every ML developer used to tensor operations. At the same time, it is transparent enough for the developer to understand what happens at the hardware level. The authors say that if you know CUDA, you can "compile" ThunderKittens in your head.

The fundamental object in ThunderKittens is the tile: a matrix sized to fit into a tensor core. It is the unit of computation that lets you fully load the specialized core and extract the full potential of modern hardware.

Performance and results

On A100 and RTX 4090 chips, ThunderKittens matches FlashAttention2 in speed, while the code is somewhat shorter and clearer. On the H100, it wins: it is faster than FlashAttention2 in both the forward and backward passes. In other words, there is no trade-off between clean code and speed.

The authors have already written several kernels for other algorithms in ThunderKittens:

• Based, an optimized version of linear attention
• Hedgehog and other specialized kernels for transformers
• Several solutions that compare favorably with Triton versions on speed

One interesting detail: engineers who had taken only a two-hour CUDA course were able to write their own code in ThunderKittens. This suggests the language really does simplify development, even though it does not hide the hardware details.

Open project and teaching

The Together AI authors say frankly that this is an art project. Do not expect regular updates or support for every complaint in the bug tracker. The project is released openly because the developers find it interesting to share ideas and tools with the community.

Alongside ThunderKittens, they released NanoGPT-TK, a version of Andrej Karpathy's iconic NanoGPT project in which the main compute kernels are rewritten in ThunderKittens. It was built specifically for teaching and to demonstrate the language's capabilities. NanoGPT has long been regarded as one of the best AI projects for understanding how transformer training works from scratch.

What this means

ThunderKittens shows that AI has a real gap between the convenience of abstractions (PyTorch, Triton) and control over the actual hardware (CUDA). It turns out developers are willing to write slightly more complex code if it gives them control and real speed in practice.

For ML engineers, this may mean that moving a trained model from one chip to another gets easier in the future: rewriting a few GPU kernels would be enough, rather than reworking half the infrastructure. For researchers, it is a tool for quickly experimenting with specialized algorithms that does not require a months-long CUDA course.