ThunderKittens: Simple, Fast, and Adorable AI Kernels (2410.20399v1)

Abstract: 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\times$ on state space models, and $14\times$ on linear attention.

• The paper introduces a novel framework that simplifies GPU kernel development for AI by using three levels of GPU abstraction.
• It employs warp-level, thread-block, and grid-level optimizations to achieve competitive performance with state-of-the-art libraries.
• Empirical results demonstrate over 40% improvement in attention computations and up to 14x speedup in linear attention.

ThunderKittens: Efficient AI Kernel Development

The paper "ThunderKittens: Simple, Fast, and Adorable AI Kernels," authored by Spector et al., addresses the critical challenge of efficiently mapping AI architectures to GPU hardware. This task has proven to be a bottleneck for AI advancement due to the substantial gap between theoretical and actual performance of hand-written custom kernels. ThunderKittens (TK) introduces a novel framework aimed at simplifying the development of high-performance AI kernels while maintaining ease of use and maintenance.

Key Abstractions

ThunderKittens employs a structured approach through three levels of GPU hierarchy:

1. Warp-level Abstraction: The framework provides 16x16 matrix tiles as basic data structures. These tiles are compatible with tensor cores, aiming to maximize their utilization. TK automatically optimizes memory layouts to minimize bank conflicts, drawing from familiar interfaces like PyTorch.
2. Thread-block Level Coordination: At this level, TK introduces a template for asynchronous operations using the Load-Compute-Store-Finish (LCSF) model. This template manages synchronization and overlapping of workloads across parallel warps within a thread block.
3. Grid-level Optimization: TK aids in hiding the costs of block launch and tear-down, facilitating pipeline efficiency. It explores the use of persistent grid blocks to reduce block launch overheads and improve cache reuse by controlling block execution order.

Performance and Benefits

The framework is empirically validated through extensive experimentation, demonstrating its ability to match or exceed the performance of existing optimized kernels:

• Matrix Multiplication (GEMM): TK's kernels exhibit competitive performance with NVIDIA's CuBLAS library, achieved through a straightforward 40-line implementation.
• Attention Mechanisms: For both causal and non-causal attention, TK's kernels outperform the state-of-the-art FlashAttention-3 by over 40% in the backward pass computations, highlighting the efficacy of its abstractions in different sequence lengths and head dimensions.
• Emerging AI Architectures: TK shows remarkable performance improvements on newer machine learning primitives like linear attention and state space models, outperforming the respective strongest baselines by a substantial margin. For instance, TK achieves up to 14x performance improvement on linear attention compared to Triton's Flash Linear Attention (FLA).

Insights and Implications

The findings suggest that a limited set of well-designed abstractions can significantly alleviate the complexity of writing performant kernels. By reducing the scope of choices—especially in memory layout—TK minimizes the overhead for developers while allowing them to harness the full power of modern GPU architectures. The embedded C++ design ensures that the framework doesn't impede optimization potential, providing seamless integration with existing CUDA capabilities.

Future Prospects

ThunderKittens sets a precedent for kernel development by demonstrating the feasibility of simplifying AI kernel programming without compromising performance. The framework has already seen adoption in industrial scenarios, including ML inference providers and high-frequency trading firms. Future enhancements may focus on expanding support for more complex architectures and further optimizing grid-level execution strategies.

This paper contributes a significant step forward in making high-performance GPU programming more approachable and efficient, paving the way for broader accessibility and potentially catalyzing advancements in AI hardware utilization.