PaperlayerView exact revision on arXiv

Source-linked AI summary

ThunderKittens: Simple, Fast, and Adorable AI Kernels

Benjamin F. Spector, Simran Arora, Aaryan Singhal, Daniel Y. Fu, Christopher Ré

arXiv:2410.20399v1cs.LGcs.AI

TL;DR

AI architectures are difficult to map efficiently onto GPU hardware, and custom kernels can remain below theoretical performance despite substantial effort. ThunderKittens provides a small set of abstractions across GPU hierarchy levels for writing AI kernels. Its kernels match or outperform prior kernels across a range of operations, including up to 14× on linear attention.

  • Problem

    Mapping AI architectures efficiently to GPU hardware remains difficult, and kernel performance can fall substantially below theoretical potential despite development effort.

    Source: §1 · p. 1

  • Method

    ThunderKittens combines managed-layout tiles and tile operations, a template for asynchronous warp coordination, and grid-level support for navigating scheduling and reuse tradeoffs.

    Source: §1 · pp. 2–3

  • Results

    TK matches CuBLAS GEMMs and FlashAttention-3 attention inference, and outperforms strongest baselines by 10–40% on attention backwards, up to 8× on state space models, and up to 14× on linear attention.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors conclude that TK’s abstractions are general and consistently meet or exceed state-of-the-art across a breadth of AI kernels.

    Source: §5 · p. 11

    The discussion focuses on NVIDIA, though the authors say the parallelism types also hold across AMD and Apple GPUs.

    Source: §1 · p. 1

Abstract

from arXiv · show

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.

1 Introduction

ThunderKittens asks whether a small set of abstractions can make high-performance AI kernels simpler to write amid the challenge of mapping diverse architectures to GPUs. It introduces abstractions across GPU hierarchy levels and reports kernels that match or outperform prior work across several AI operations.

1 Introduction

  • AI kernel performance remains below theoretical potential despite substantial effort, with even widely used softmax attention showing poor kernel support.FlashAttention-2 [12] had a 47% performance degradation when translated to H100, and FlashAttention-3 took over two years to develop after H100’s release.Source: §1 · p. 1
  • TK explores a small, opinionated abstraction set as an alternative to CUTLASS/CuTe ’s nested templates and Triton ’s simpler interface with fewer optimizations.The stated goal is a framework that combines high performance across many primitives with ease of use, learning, and maintenance.Source: §1 · p. 1
  • TK’s three-level abstractions use managed-layout tiles and tile operations, asynchronous warp coordination, and grid scheduling to overlap loads across thread-block boundaries.The warp-level tile is 16 × 16; the block template uses a producer-consumer model, while persistent grids can reduce pipeline bubbles and improve L2 cache hit rates.Source: §1 · p. 2
  • TK exposes tradeoffs in tile layouts, occupancy, and block launch order, while its simplified interface gives users control to navigate them.The authors also report TK kernels in production at ML inference providers and high-frequency trading firms.Source: §1 · p. 3
  • TK kernels match CuBLAS GEMMs and FlashAttention-3 attention inference, and outperform strongest baselines by 10–40% on attention backwards, up to 8× on state space models, and up to 14× on linear attention.The kernels were written by a small academic team, including undergraduates with no prior CUDA experience.Source: §1 · p. 3

2 GPU fundamentals

The section frames GPU execution through warp, block, and grid levels, then uses a cost model to motivate overlapping work. It contrasts programming frameworks by how they balance performance, control, and ease of use.

2 GPU fundamentals

  • GPU kernels move data from HBM, perform work, and write outputs back; the discussion focuses on H100 SXM while stating that the principles apply across GPU vendors and generations.The hierarchy follows physical hardware organization, from warps through thread blocks to grids.Source: §2 · p. 3

2.1 GPU hierarchy

  • Thread blocks group warps that execute a kernel on an SM; warps may contend for registers, shared memory, issue slots, and compute pipelines, while grids launch blocks sharing HBM and L2.Block scheduling incurs setup costs and tail effects when grid size is poor: a 133-block grid on an H100 with 132 physical SMs takes two waves, with the second running at < 1% efficiency.Source: §2.1 · p. 4

2.2 Cost model

  • The cost model groups execution time into memory, compute, and overhead costs, with perfect overlap represented by their maximum.Actual performance lies between the maximum and sum of the components, depending on workload properties and implementation efficiency.Source: §2.2 · p. 4

2.3 GPU programming frameworks

  • CUTLASS/CuTe offers nested CUDA templates and C++ control, while compiler frameworks such as Triton offer simpler interfaces but can make specialized instructions and asynchronous execution harder to manage.TK takes an opinionated abstraction approach and seeks a PyTorch-like feel alongside high performance.Source: §2.3 · pp. 4–5

3 ThunderKittens

ThunderKittens (TK) aims to simplify high-performance AI kernel development with abstractions spanning warp, thread-block, and grid parallelism. Its tile operations and asynchronous template help developers navigate hardware tradeoffs, while grid-level scheduling addresses block-launch costs and memory reuse.

3.1 Warp parallelism with familiar data structures and operations

  • TK uses 16 × 16 matrix tiles and PyTorch-like parallel operations, with layouts selected to reduce bank conflicts while supporting specialized hardware.Tile abstractions also enable compile-time checks of layout requirements for operations such as tensor-core matrix multiplication.Source: §3.1 · pp. 5–6

3.2 Block parallelism with a generalized asynchronous template

  • The LCSF template coordinates asynchronous thread-block work through load, compute, store, and finish functions, with pipelined buffers and synchronization managed by TK.Developers provide the four functions; multi-stage buffers can overlap loading the next tile with computation on the current tile.Source: §3.2 · pp. 6–7
  • LCSF expands the attention kernel’s Pareto frontier beyond a warp-level parallel kernel as occupancy varies.Both kernels improve until resource contention dominates.Source: §3.2 · p. 7, Figure 6 · p. 7

3.3 Grid parallelism with block launch scheduling

  • 2,390 GB/s HBM bandwidth and 494 TFLOPS correspond to block order {B, H, N}, versus 213 GB/s and 600 TFLOPS for {N, H, B}.The study measures HBM bandwidth and efficiency across block orders; the authors report that order substantially affects L2 reuse and can control kernel performance.Source: §3.2 · p. 8, Table 3 · p. 8, §3.3 · p. 8
  • TK explores persistent grids that reuse existing blocks for the next work chunk and overlap preparation of the next input with finishing the prior chunk.The grid-level tradeoff study targets block setup and tear-down costs as well as memory reuse between blocks.Source: §3.3 · p. 8, Table 2 · p. 8

4 Experiments

The experiments benchmark ThunderKittens kernels against optimized alternatives across workhorse and emerging AI operations, then profile selected implementations to examine their GPU utilization. TK matches or outperforms baselines across these workloads, with especially large gains on linear attention and long convolution.

4.1 TK enables simple and performant AI kernels

  • Over 40% at short sequences and 10% at longer sequences: TK outperforms FlashAttention-3 on causal and non-causal attention backward passes.TK competes with FlashAttention-3 on non-causal forward attention across sequence lengths.Source: Figure 8 · p. 9, §4.1 · p. 9
  • 14× on polynomial-based and 6.5× on learned-map linear attention: TK outperforms Flash Linear Attention’s Triton kernels.The experiments optimize both polynomial-based and learned feature-map architectures.Source: Figure 9 · p. 10, §4.1 · p. 10
  • 7.9× at sequence length 1024 and 4.7× at 4096: TK outperforms FlashFFTConv on long convolution.TK also beats PyTorch FFT operations by up to 8.7×.Source: Figure 9 · p. 10, §4.1 · p. 10
  • TK matches CuBLAS on demonstrated GEMM sizes and competes with FlashAttention-3 on non-causal forward attention, while outperforming it on causal and non-causal backward passes.The GEMM comparison uses a single kernel with 40 lines of device code.Source: §4.1 · pp. 9–10

4.2 Comparing kernel implementations

  • 4.1× higher tensor-core utilization: TK’s long-convolution profile also shows higher issue slots and fewer memory stalls than FlashFFTConv.The analysis attributes these differences to the TK template and warpgroup operations, including WGMMA operations that save registers and establish a shared-memory-to-register pipeline.Source: Table 4 · p. 11, §4.2 · p. 11
  • 10% fewer HBM-wait stalled cycles and 85% fewer shared-memory stalled cycles: TK’s attention-backward profile shows higher memory throughput and no bank conflicts.The compared FlashAttention-3 profile reports up to 9.6-way shared-memory bank conflicts; both methods have matching tensor-core utilization.Source: Table 4 · p. 11, §4.2 · p. 11

4.1 TK enables simple and performant AI kernels

  • TK extends to fused dropout-residual-layernorm and rotary memory operations, with comparisons against popular Triton kernels.The experiments also cover FFT convolutions and state space models.Source: §4.1 · pp. 10–11

5 Conclusion

ThunderKittens asks how far a few accessible GPU abstractions can go and concludes they support kernels that meet or exceed state-of-the-art performance across a breadth of AI operations. The framework and kernels are open sourced.

5 Conclusion

  • A breadth of ThunderKittens kernels consistently meets or exceeds state-of-the-art performance, supporting the conclusion that its abstractions are both general and performant.The abstractions cover tile layouts at the worker level, asynchronous execution at the thread-block level, and persistent-launch and L2-reuse options at the grid level.Source: §5 · p. 11
  • ThunderKittens provides abstractions at worker, thread-block, and grid levels, including managed-layout tiles and an asynchronous execution template.The conclusion also highlights tradeoffs for persistent block launches and L2 reuse.Source: §5 · p. 11
  • The ThunderKittens framework and kernels are open sourced for future development.The passage provides the project repository URL.Source: §5 · p. 12

A Related work

The related-work discussion positions ThunderKittens among embedded CUDA libraries, compiler-based frameworks, and hardware-aware kernel efforts. It distinguishes TK by focusing on a small set of opinionated abstractions for kernel-level optimization.

A Related work

  • CUTLASS and TK can express the same kernels as embedded C++ libraries, but TK focuses on how broadly and quickly a small set of opinionated abstractions can go.The passage contrasts this focus with CUTLASS’s extensive nested templates.Source: A Related work · p. 16
  • Compiler-based frameworks such as TVM provide an end-to-end stack for graph- and operator-level transformations, whereas TK concentrates on kernel-level optimization.The passage also names TensorFlow XLA, Glow, and DLVM as compiler-based frameworks.Source: A Related work · p. 16
  • Triton offers tiled computations and a JIT compiler with tile optimizations and auto-tuning; TK is embedded in CUDA and its abstractions fail gracefully.The passage contrasts this with Triton operations, which must be written entirely within its framework.Source: A Related work · p. 16
  • Prior hardware-aware work includes optimized attention, FlashFFTConv, and linear-attention kernels; TK asks whether a small abstraction set can support kernels across many AI operations.The discussion frames this as an alternative to independently optimizing each new architecture.Source: A Related work · p. 16

B ThunderKittens kernels

The appendix introduces its kernel template through a detailed GEMM example and provides listings for several kernel types. It also describes the benchmark setup used for Section 4 results.

B ThunderKittens kernels

  • The appendix provides kernel listings for GEMM, long convolution, attention, and rotary operations.These listings use the TK LCSF template and tile abstractions.Source: B ThunderKittens kernels · p. 17
  • The GEMM kernel is the detailed example used to introduce the template components.Its listing appears in Appendix B.1.Source: B ThunderKittens kernels · p. 17
  • Section 4 kernels are benchmarked on an NVIDIA H100 80GB SXM GPU using 10 warmup and 10 timed iterations, with C++ timings.Python-bound kernels and benchmarking infrastructure are also provided for reference.Source: B ThunderKittens kernels · p. 17

B.1 Matrix multiply

The matrix-multiply kernel uses ThunderKittens’ LCSF template to divide output tiles among compute warpgroups while load and store workers handle data movement. Its 3D grid order preserves cache locality, especially for large matrices.

Matrix multiply kernel

  • Compute warpgroups handle 64M-row by 64N-column output chunks, while load workers fetch input chunks and store workers write completed results asynchronously.Compute workers initialize accumulators, run asynchronous matrix multiplies, and store their output tile when finished.Source: §B.1 · p. 17

Tuning the number of workers and pipeline stages

  • Each pipeline stage processes 64 elements along the reduction dimensions, and ThunderKittens automatically sizes the pipeline when the user omits a value.Common 128 × 256 and 192 × 192 output tiles use a 4-stage pipeline.Source: §B.1 · p. 17

Tuning the grid order

  • A naive row-major grid order cuts 16384 × 16384 × 16384 matrix-multiply performance from 797 TFLOPs to 387 TFLOPs, a > 50% degradation.At 4096 × 4096 × 4096, the same comparison drops from 767 TFLOPs to 735 TFLOPs.Source: §B.1 · p. 17

Matrix multiply kernel

  • The templated matrix-multiply kernel is described as reasonably competitive with CuBLAS.The section presents the kernel in the LCSF template.Source: §B.1 · p. 17, Figure 10 · p. 18

B.2 Long convolution

The long-convolution kernel uses an FFT convolution algorithm computed via Monarch Matrices for sequence length 4096. Its template assigns heads and batches to blocks and uses asynchronous data movement.

Long convolution

  • The kernel assigns each block a head and calculates how many heads it handles, with the code stating that a batch is handled by just one block.It divides batch iterations across consumer warpgroups.Source: §B.2 · p. 19
  • For sequence length 4096, the templated convolution kernel outperforms FlashFFTConv [20].The kernel uses the FFT convolution algorithm computed via Monarch Matrices [11, 14, 18].Source: §B.2 · p. 19, Figure 11 · p. 20

B.3 Attention

The section presents ThunderKittens’ templated non-causal attention kernel for head dimensions 64 and 128. The kernel loads key and value tiles asynchronously and computes attention scores from query and key tiles.

Attention

  • The compute path multiplies query tiles by transposed key tiles to form attention scores.The template uses query tiles of 64 rows and key/value tiles of 192 rows for dimension 64 or 128 rows for dimension 128.Source: §B.3 · p. 21
  • A producer warp asynchronously loads key and value tiles for the selected batch, head, and sequence position.The input block contains key and value tiles.Source: §B.3 · p. 21
  • The templated non-causal attention kernel for head dimensions 64 and 128 competes with FlashAttention-3.The figure caption identifies the evaluated head dimensions.Source: §B.3 · p. 21, Figure 12 · p. 22

B.4 Rotary positional encodings

The rotary kernel applies ThunderKittens abstractions to head dimension 128 and is reported to outperform popular Triton baselines. The supplied shared-memory-layout discussion compares six layouts by where tile elements map to memory banks.

B.4 Rotary positional encodings

  • The rotary template processes sequence tiles of 16 rows and keeps sine and cosine tiles resident in registers.The producer asynchronously loads input tiles and stores output tiles.Source: §B.4 · p. 23
  • For head dimension 128, the templated rotary kernel outperforms popular Triton baselines.The section describes the kernel as written in ThunderKittens abstractions.Source: §B.4 · p. 23, Figure 13 · p. 24

C Shared memory layouts

  • Six shared-memory layouts are illustrated on a 32 × 64 16-bit tile, with colors and labels identifying which of 32 banks stores each element.The layouts include naive row-major, padded, simple swizzled, and three specialized swizzled layouts.Source: C Shared memory layouts · p. 25

C.1 Naive layout

The appendix compares shared-memory layouts by their bank conflicts, alignment, instruction support, and tile-width constraints. ThunderKittens selects the strongest supported swizzle for a tile’s compile-time width.

C.1 Naive layout

  • Row-major layout preserves HBM tensor layout and conflict-free row access, but incurs 8-way bank conflicts for tensor-core register loads and stores.Source: §C.1 · p. 25

C.2 Padded layout

  • Padding each row by one memory bank removes conflicts but misaligns addresses required by fast instructions.Source: §C.2 · p. 26

C.3 Naive Swizzled Layout

  • XORing the index with the row removes bank conflicts and preserves alignment, but lacks HGMMA and UTMA hardware support.Source: §C.3 · p. 27

C.4 32 byte swizzling

  • 32-byte swizzling directly swizzles the address with itself; it has 4-way bank conflicts, supports tiles with widths divisible by 16, and supports HGMMA and UTMA.Source: §C.4 · p. 28

C.5 64 byte swizzling

  • 64-byte swizzling uses a more aggressive pattern than 32-byte swizzling, reducing conflicts to 2-way but requiring widths divisible by 32 for half precision or 16 for full precision.Source: §C.5 · p. 29

C.6 128 byte swizzling.

  • 128-byte swizzling eliminates bank conflicts, but is valid only for half-precision tiles whose widths are divisible by 64.Source: §C.6 · p. 30

C.7 ThunderKittens

  • ThunderKittens selects the highest swizzling level valid for a tile’s compile-time width, prioritizing HGMMA and UTMA support while minimizing bank conflicts.The authors report these three hardware-supported layouts are sufficient for high performance across many kernels.Source: §C.7 · p. 30

Related papers

Loading 2410.20399v1…