# ThunderKittens: Simple, Fast, and Adorable AI Kernels

arXiv: https://arxiv.org/abs/2410.20399 · HTML: https://mlsanity.com/paper/2410.20399/

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 50 references that are in the map.
- First posted on arXiv 27 Oct 2024.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Hyena Hierarchy: Towards Larger Convolutional Language Models](https://mlsanity.com/paper/2302.10866/) (2023)
- [Simple linear attention language models balance the recall-throughput tradeoff](https://mlsanity.com/paper/2402.18668/) (2024)
- [On The Computational Complexity of Self-Attention](https://mlsanity.com/paper/2209.04881/) (2022)
- [Liquid Structural State-Space Models](https://mlsanity.com/paper/2209.12951/) (2022)
- [Fast Transformers with Clustered Attention](https://mlsanity.com/paper/2007.04825/) (2020)
- [The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry](https://mlsanity.com/paper/2402.04347/) (2024)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture](https://mlsanity.com/paper/2310.12109/) (2023)
- [PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels](https://mlsanity.com/paper/2310.01655/) (2023)
- [FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores](https://mlsanity.com/paper/2311.05908/) (2023)
- [LoLCATs: On Low-Rank Linearizing of Large Language Models](https://mlsanity.com/paper/2410.10254/) (2024)
- [Spectral State Space Models](https://mlsanity.com/paper/2312.06837/) (2023)
- [Linear Transformers with Learnable Kernel Functions are Better In-Context Models](https://mlsanity.com/paper/2402.10644/) (2024)
- [A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library](https://mlsanity.com/paper/2312.11918/) (2023)
- [DLVM: A modern compiler infrastructure for deep learning systems](https://mlsanity.com/paper/1711.03016/) (2017)

## Built on by (most cited first, 30 of 36)

- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [VSA: Faster Video Diffusion with Trainable Sparse Attention](https://mlsanity.com/paper/2505.13389/) (2025)
- [One-Minute Video Generation with Test-Time Training](https://mlsanity.com/paper/2504.05298/) (2025)
- [Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization](https://mlsanity.com/paper/2509.14279/) (2025)
- [Astra: A Multi-Agent System for GPU Kernel Performance Optimization](https://mlsanity.com/paper/2509.07506/) (2025)
- [AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization](https://mlsanity.com/paper/2511.15915/) (2025)
- [Adaptive Self-improvement LLM Agentic System for ML Library Development](https://mlsanity.com/paper/2502.02534/) (2025)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels](https://mlsanity.com/paper/2503.14376/) (2025)
- [Sparse Forcing: Native Trainable Sparse Attention for Real-time Autoregressive Diffusion Video Generation](https://mlsanity.com/paper/2604.21221/) (2026)
- [CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe](https://mlsanity.com/paper/2604.01489/) (2026)
- [Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs](https://mlsanity.com/paper/2512.18134/) (2025)
- [CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs](https://mlsanity.com/paper/2605.19269/) (2026)
- [Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs](https://mlsanity.com/paper/2510.08726/) (2025)
- [Optimizing PyTorch Inference with LLM-Based Multi-Agent Systems](https://mlsanity.com/paper/2511.16964/) (2025)
- [GPU Kernel Optimization Beyond Full Builds: An LLM Framework with Minimal Executable Programs](https://mlsanity.com/paper/2512.22147/) (2025)
- [A Two-Stage GPU Kernel Tuner Combining Semantic Refactoring and Search-Based Optimization](https://mlsanity.com/paper/2601.12698/) (2026)
- [Sliding-window beats linear attention](https://mlsanity.com/paper/2608.28444/) (2026)
- [Flash3D: Super-scaling Point Transformers through Joint Hardware-Geometry Locality](https://mlsanity.com/paper/2412.16481/) (2024)
- [HipKittens: Fast and Furious AMD Kernels](https://mlsanity.com/paper/2511.08083/) (2025)
- [Early Discoveries of Algorithmist I: Promise of Provable Algorithm Synthesis at Scale](https://mlsanity.com/paper/2603.22363/) (2026)
- [Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels](https://mlsanity.com/paper/2604.14825/) (2026)
- [Veda: Scalable Video Diffusion via Distilled Sparse Attention](https://mlsanity.com/paper/2605.30325/) (2026)
- [DeltaLog: Deferred Materialization of Recurrent States for Linear Attention Decoding](https://mlsanity.com/paper/2608.15533/) (2026)
- [Iris: First-Class Multi-GPU Programming Experience in Triton](https://mlsanity.com/paper/2511.12500/) (2025)
- [PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks](https://mlsanity.com/paper/2603.06731/) (2026)
- [Learning When to Optimize: Verified Optimization Skills from Expert GPU-Kernel Lineages](https://mlsanity.com/paper/2605.28213/) (2026)
- [AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis](https://mlsanity.com/paper/2606.09682/) (2026)
- [SPADE: An Input-Adaptive Sparse Attention Engine for Fast Video Diffusion Models Inference](https://mlsanity.com/paper/2608.03335/) (2026)
- [HIERA: Workload-Aware Planning Across Implementation Spaces for GPU Kernel Optimization](https://mlsanity.com/paper/2608.21157/) (2026)

## Abstract (arXiv)

The challenge of mapping AI architectures to GPU hardware is creating a critical bottleneck in AI progress. Despite substantial efforts, hand-written custom kernels fail to meet their theoretical performance thresholds, even on well-established operations like linear attention. The diverse hardware capabilities of GPUs might suggest that we need a wide variety of techniques to achieve high performance. However, our work explores whether a small number of key abstractions can drastically simplify the process. We present ThunderKittens (TK), a framework for writing performant AI kernels while remaining easy to use and maintain. Our abstractions map to the three levels of the GPU hierarchy: (1) at the warp-level, we provide 16x16 matrix tiles as basic data structures and PyTorch-like parallel compute operations over tiles, (2) at the thread-block level, we provide a template for overlapping asynchronous operations across parallel warps, and (3) at the grid-level, we provide support to help hide the block launch and tear-down, and memory costs. We show the value of TK by providing kernels that match or outperform prior kernels for a range of AI operations. We match CuBLAS and FlashAttention-3 on GEMM and attention inference performance and outperform the strongest baselines by 10-40% on attention backwards, 8× on state space models, and 14× on linear attention.