ThunderKittens: Simple, Fast, and $\textit{Adorable}$ Kernels
Authors: Benjamin Frederick Spector
Simran Arora
Aaryan Singhal
Arjun Parthasarathy
Daniel Y Fu
Christopher Re
Venue: ICLR 2025 Spotlight (2025)
Keywords: Systems
Kernels
Efficiency
Efficient Models
IO Awareness
GPUs
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: ThunderKittens (TK) is a framework that simplifies the creation of high-performance AI kernels through key abstractions, enabling efficient implementation of ML architectures on GPU hardware and surpassing previous approaches in hardware utilization.

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 capabilities of GPUs suggests we might we need a wide variety of techniques to achieve high performance. However, our work explores if 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. 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 operations, (2) at the thread-block level, we provide templates for asynchronously overlapping operations, and (3) at the grid-level, TK helps hide block launch, tear-down, and memory costs. We show the value of TK by providing simple & diverse kernels that match or outperform prior art. 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.

Published as a conference paper at ICLR 2025
THUNDERKITTENT SIMPLE, FAST, AND Adorable KERNELS
Benjamin F. Spector', Simran Arora', Aaryan Singhal', Arjun Parthasarathy", Daniel Y. Fu Christopher Ré'
Stanford University, 2 Columbia University, UCSD, Together AI
{bfs, simarora, aarsan04,chrisrre) } @ stanfoord.edu, arjun. columbia.edu edu, danfu @ @ucsd.edu
ABSTRACT
The challenge of mapping Al architectures to GPU hardware creating a critical bottleneck in Al progress. Despite substantial efforts, fail to meet their theoretical performance thresholds, even on well-established op- erations like linear attention. The diverse capabilities of GPUs suggests we might we need a wide variety of techniques to achieve high performance. However, our work explores
custom kernels
small number of key abstractions can drastically simplify the process. We present THUNDERKITTES mant Al kernels while remaining easy to use. Our abstractions map to the three levels of the GPU hierarchy: (1) at the warp-level, we provide 16x matrix tiles as basic data structures and PyTorch-like operations, (2) at the thread-block level, we provide templates for asynchronously overlapping operations, and (3) at the grid-level, TK helps hide block launch, tear-down, and memory costs. We show the value of TK by providing simple & diverse kernels that match or outperform prior art. We match CuBLAS and FlashAttention-3 on GEMM and attention infer- ence performance and outperform the strongest baselines by 10 10-40 40% attention backwards, 8x on state space models, and 14x on linear attention.
(TK), a framework for writing perfor-
INTRODUCTION
AI is bottlenecked by the problem of efficiently mapping Al architectures onto accelerated GPU hardware. There has been a Cambrian explosion of MI. architectures (Ho et al., 2020) Gu & Dao 2023}; however, the performance of these architectures remains substantially below their theoretical potential, despite substantial effort to develop kernels, or GPU implementations. Notably, kernel support has been poor even for softmax attention, which is used throughout industry. FlashAttention- 2 (Dao 2024) suffered 47% performance degradation when translated to the H100 GPU, and it took over two years from the release of the H100 to develop FlashAttention-3 (Shah et al. 2024).
We are inspired by several approaches to supporting the development of Al kernels. Ideally, we would have a framework that supports high performance for breadth of primitives, while being easy to use. learn from, and maintain. High performance C++ embedded libraries like NVIDIA CUTLASS/CUTe NVIDIA 2017) contain myriad of nested templates, while compiler based ap- proaches like Triton Tiileet et al.) 2019) provide users with simpler interfaces, but fewer optimiza- tions. We ask how broad and fast we can go by choosing a small and opinionated set of abstractions.
The main vector of growth for accelerated compute is in specialized matrix multiply units. On the NVIDIA H100 and NVIDIA A100 GPUs, BF16 tensor cores represent 15 16x the FLOPS available relative to general-purposs BF16 FP32 compute. Consequently, any high performance framework must prioritize keeping tensor cores at high utilization whenever possible. However, all kernels have non-tensor operations, too (like memory loads or the softmax in attention). and it is crucial to minimize their overhead. This proposition is at the heart of our approach.
To understand the complexities and opportunities in building a simple, yet high performance frame- work, we examine simplified model of GPU parallelism, further detailed in section 2
1. Warp-level parallelism: Modern GPUs consist of tens of thousands of hardware threads which execute in parallel. Threads are organized into small groups, "warps", which execute instructions
We discuss primarily NVIDIA, but the parallelism types hold across architectures, including AMD and Apple GPUs; we provide experiments an Apple M2 Pro in Appendix

---

Published as a conference paper at ICLR 2025
Tensor care register tiles
Load and store workers
Shared Buffers
Compute workers
130
Mt X>
Inputs
Sync/Boync leads and stores from HBM to Shared memorx
PyTorch-like operations on shared and register mamory tiles
Swizzled shared tiles
Shared
33 TB/s
KB
st-M X>;
Sync buffer leadik uven(() (0), global, 1dx03
inneegister
s- Licis ky Ji.lighh> ahin
L2 Cache 50 MB, TD/s
Asyne buffer
Output
varpgrougetss
Tennso
registers
sayae P_giobal. sed, harly
c_ssin[voccidd) [D]
HBM 80 GB TB/s
Global Tensors
#l Asyne store from buffer 2
tra istore acynoco global
o_sri,, idsi):
// Use non tensor core units
exp (aton_ attn
Figure 1: THUNDERKITTENS explores whether a small set of abstractions can enable performant AI kernels. Inspired by PyTorch, we first provide tiles with managed layouts and operations over these tiles. Second, we provide program templates for coordinating asynchronous workers c.g., workers that load and store data, while other workers perform computations in fast memory.
together. Memory layouts determine how the logical data elements are mapped to physical thread ownership. If multiple threads try to access the same region ("bank") of memory, this can create expensive serializations between the threads (called "bank conflicts").
Block-level parallelism: Warps are grouped into "blocks" of threads, which can quickly share data. Warps execute their instructions on physical execution units, and having more warps in block (called occupancy) can help run more instructions at the same time, reducing runtime. For example, one warp can run tensor cores for matmul, while another uses the ALU for max.
3. Grid-level parallelism. GPUs run many blocks threads at once, which communicate through large but slow global memory (HBM). An on-chip shared L2 cache helps reduce memory laten- cies and increase bandwidth if thread blocks reuse the same data. Thread blocks also face setup and tear-down overheads, which can introduce "pipeline bubbles" that hurt performance.
Despite the apparent need for a myriad of techniques to leverage all these hardware capabilities, our central technical finding is that indeed, for many Al kernels, small number of key abstractions exist that can simplify the process of writing high-performance kernels. Our exploration led us to develop THUNDER KITTENS KITTENS (TK), an Al kernel framework built around three key principles:
1. Tile data structures with managed layouts: Our interface is inspired by familiar ML frame- works like PyTorch and NumPy (Paszke et al. 2019), highlighted in Figure 2 At the warp level, we use 16 x 16 matrix tile as our basic data structure, maximizing compatibility with and encouraging the use of tensor cores. TK automatically picks the optimal memory layouts for the tiles to minimize bank conflicts while remaining compatible with specialized hardware instruc- tions, avoiding user effort. We provide set of parallel compute primitives over tiles, based on the suite of operations in PyTorch (e.g.. pointwise multiply, mma, exp, and cumsum over tiles).
2. Program template for asynchronous work: At the block level, TK provides general kernel template for coordinating asynchronous execution across warps in a thread block, built on the producer-consumer paradigm (Dijkstral 1968). The developer's effort reduces to populating few boilerplate functions within this model, using our Py -orchh-like operands, and the template internally hides latencies through memory pipelines and synchronization primitives (Figure DD.
3. Grid scheduling for pipelining Chread-blocks, At the grid level, we show TK can help devel- opers reduce pipeline bubbles and improve I.2 cache hit rates. Our template supports persistent grid, where we overlap memory loads across thread block boundaries.
We highlight the value of these abstractions for developers in two ways:
Through our exploration, we identify a few fundamental tradeoffs between achieving different types of parallelism including in setting the tile layouts (warp-level), occupancy (block level), and block launch order (grid level). Through our ablation studies (Section 3) we show how the simplified interface in TK gives users the control navigate the tradeoffs.
We validate the TK abstractions by providing kernels that match or outperform prior kernels for a range of AI operations. We match CuBLAS GEMMs and Attentionn attention inference,
2

---

Published as a conference paper at ICLR 2025
PyTorch attention:
THUNDERKITTENS attention:
imports
inports using nettagace kittens: te 1,,,, 64> <_reg, v_reg:
lngort Lozch
imgort torch.nr.
Functional www.ttnn..untio........ de
// load k. from shared menory register
Ik [_anen[subbiiel]
corpute GBK.T ett torch.netmell 3:1
// conpute OEK.T
zerolatt:;
stt)s
congute
corpute snferme
10
sub, row att, nex_vee);
II
att F.cortmox( d:m=-1, dcype Float:2)
11
exp (att, attl:
12
12
div rew
vecly
convert copy latt latt_mn,, attlj
toz ma_48
14
15
convert bade to 6F16 stt.zo(q.d:ypo)
15
// load v. from shared menory to register
16
16
io v_snen(subtilell
$y_reg_ evep__ayour_ingg.oa 18 // mna et onto o reg
mm stt
19
outpuc
Borch.matmu (att,
19 mna_1B att_rma, Y_rog_e, eurcar
Figure 2: A snippet of attention kernel to show the PyTorch-like operations tiles.
and outperform the strongest baselines by 10 40% on attention backwards, up to 8x on state space models, and up to 14x on linear attention.
Our contributions are (I) showing small and opinionated set of abstractions in TK that goes sur- prisingly far for writing simple and performant kernels; and (2) providing a collection of performant AI kernels. We hope that TK and its insights help improve the accessibility of kernels.
2 GPU FUNDAMENTALS
GPU tasks are divided into small programs called kernels. A kernel typically loads data from high bandwidth memory (HBM), performs work on it, and writes the outputs back to HBM before concluding. Before we explain THUNDERKIT- TENS's abstractions, we provide background on GPU paral- lelism at the warp, block and grid levels. We follow NVIDIA's terminology and focus on the H100 SXM GPU, though the principles apply across GPU vendors and generations.
32
(Lane)
132
Warp (Quadrant)
Block (SM)
Grid (GPU)
Figure 3: The software (and physical) GPU hierarchy.
2.1 GPU HIERARCHY
The GPU software hierarchy closely follows its physical hardware hierarchy (Figure) Here, illustrate several of its most important components and aspects.
Warps consist of groups of 32 consccutive threads that operate on data in small but fast reg- ister memory. These instructions run on physical execution wnits, and individual threads can frequently occupy multiple specialized execution pipelines (below) at once through instruction- level parallelism, and different warps can further occupy available execution hardware:
(a) Load and store units, to bring data into and out of registers. Advanced GPUs have also introduced dedicated hardware acceleration (c.g. H100 Tensor Memory Accelerator) for asynchronous bulk data movement between HBM and shared memory.
(b) General purpose compute pipelines, such ALU for max, min, FMA for multiplies and adds, and XU for complex operations like exp. Throughput differs across the pipelines.
(c) Accelerated matrix multiply hardware (tensor cores), which have most of the GPU compute.
Threads can temporarily stall for a variety reasons, including (but not limited to) fixed instruc- tion latencies, memory latencies, barriers, pipeline throttles, or instruction cache misses.
Thread blocks are groups of warps which together execute a kernel on a physical core, called a streaming multiprocessor (SM). Although each SM has just four physical execution units, up to 64 software warps can simultaneously run on (called "occupancy"). These collocated warps of- ten contend hardware resources: registers, shared memory, issue slots, and compute pipelines, but together they can help keep many work streams running at the same time within each execu- tion unit. Warps synchronize at barriers, during which they cannot issue work.
3

---

Published as a conference paper at ICLR 2025
Importantly, warps within the same block can quickly communicate through special shared mem- ory (SMEM, 227 KB, 33 TB/s). To improve bandwidth, SMEM is grouped into 32 physical serve memory simultancously. However, if different threads try
"banks" of memory, which to access the same bank at the same time (called a bank conflict), their accesses must he serial- ized, which both increases access latencies and reduces available bandwidth. Hopper has limit of 255 registers per thread and attempts to request more, results in spills to the LI cache. SMEM can be reallocated as an LI cache for fast access to frequently used memory like spilled registers.
3. Grids of multiple thread blocks are launched to run the kernel. The H100 SXM GPU has 132 physical SM's which run thread blocks at the same time. Although SM's are capable of collocating multiple thread blocks, most AI kernels can achieve high performance by simply collocating more warps within single thread block (increasing the occupancy).
Thread blocks on the same GPU share common memory resources: large but slow high- bandwidth memory (80 GB, TB/s), which has both the greatest latency and least bandwidth of all GPU memory, and smaller but faster hardware-manage 1.2 cache (50 MB, 12 TB/s).
There are overheads to scheduling blocks. First, the block launch incurs setup costs and although this cost must be paid at least once at the initial kernel launch, kemnlss that continuously launch many large blocks can incur further costs. Second, there are tail effect costs if the grid is sized poorly. If a kernel of 133 blocks is executed an H100 with 132 physical SMs, the kernel would require two waves to execute, the first with full efficiency, and the second with < 1% efficiency. More advanced schedules using independent CUDA streams can ameliorate these tail effects, such as recent work on asynchronous tensor-parallel schedules Chang et al. (2024).
2.2 COST MODEL
Summarizing the above components, show simplified cost model for GPU parallelism below. This cost model is inspired by the roofline model Williams et al. (2008). The overall kernel execution time Coverall is the sum of the following costs where memory costs combination of the latency and bandwidth, and compute costs are a combination of latency and throughput.
Coverall = max
CHEM, CL2, CLI, Cstared
C Vensor: CALU, CAMA, Cxu
+ Setup + Csyne
Memory
Compute
Overhead
This model represents the ideal case of perfect overlapping between memory, compute, and tensor core costs. A kernel's actual cost will lie between the max and the sum of the terms, depending on the workload properties (i.e. some operations are inherently sequential), and the implementation efficiency. We aim to (1) reduce these individual costs, and (2) improve their collective overlapping. In Section 3 where we detail TK, we connect our primitives and optimizations back to these costs.
2.3 GPU PROGRAMMING FRAMEWORKS
We are inspired by a number of related efforts to simplify the development of Al kernels, such as NVIDIA CUTLASS/CUTe (NVIDIA 2017) and Triton (Tillet et al., 2019).
CUTLASS's myriad of nested CUDA templates helps power highly optimized Al kernels (Shah et al. 2024; Bikshandi & Shah 2023a b) and fundamentally, the same kemelss are expressible in TK and CUTLASS, since both are embedded libraries, giving users the full power of C++. We take a complementary approach by being rather opinionated about the abstractions. We ask: (1) How far can we get with a small set of templates? and (2) Does concision sacrifice performance? An appealing outcome is improved accessibility to AI researchers, since it can be challenging to fully leverage the capabilities of CUTLASS (Bikshandi & Shah, 2023b). We find that even industrially popular kernels written in CUTLASS, like FlashAttention--3. struggle from preventable issues like bank conflicts. We seek abstractions that manage such issues for users. Most recent AI architectures use high level compilers instead (Dao & Gu 2024 Yang & Zhang 2024
Fu et al. 2023ch.
Triton, PyTorch (Pazkke et al., 2019p. TVM (Chen et al. 2018 TensorFlow XLA Abadi et al, 2016), and others approach the problem from compiler perspective. The frameworks are not C++ embedded, so it can be challenging to use unsupported specialized hardware instructions. It can also be difficult manage asynchronous execution and register usage in high level frameworks. We explore avenues that retain the simple, PyTorch-like feel while enabling maximum performance in the next section. An extended discussion of related work is in Appendix A
4

---

Published as a conference paper at ICLR 2025
3 THUNDERKITTENS
We present THUNDERKITTENS (TK), a framework designed to simplify the development of high- performance AI kernels while leveraging the full capabilities of modern GPUs. This section (1) introduces our key programming abstractions and (2) shows how they can help developers navigate the tradeoff's between different types of parallelism. Section 3.1 focuses on warp level, Section 3.2 on thread block level, and Section 3.3 on grid level parallelism.
As running examples in this section, we show how TK helps optimize attention (Vaswani et al.) 2017) and GEMM kernels. Section I demonstrates how the principles yield performant kernels for breadth of AI operations (e.g., attention variants, convolution, SSM, rotary).
3.1 WARP PARALLELISM WITH FAMILIAR
STRUCTURES AND OPERATIONS
At its core, THUNDERKITTENS is built on two fundamental abstractions tile data structures at each level of the memory hicrarchy and bulk operands on tiles akin to the familiar suite of operations in PyTorch and NumPy. We first define the abstractions, and then show they help developers navigate tradeoffs between the tile sizes and compute efficiency.
Programming abstractions TK is heavily inspired by PyTorch and NumPy, given their familiar- ity to ML audiences (Paszke et al., 2019). We provide a concise set of parallel compute operations, based on the suite of operations "worker" abstraction, or warp or warpgroup (4 warps) of threads that collaboratively own and operate on a piece of data. TK uses 16 16 matrix tile its basic data structure, designed to maximize compatibility with tensor cores. We provide tiles for each level of the memory hierarchy:
PyTorch (e.g., in Figure 2 The operations are executed by
Register tiles and vectors, which are templated by type, shape, and layout. Figure we initial- ize a 6 type tile with column-mayj major layout, height 16, width 64. The explicit control of register memory can help users reducing CMemory in Section
Shared tiles and vectors, which are templated by type and shape.
3. Global layout descriptors: We set up HBM loads and stores as indexing into 4D tensors, where the dimensions can be known runtime compile-ting time (saving valuable registers).
An advantage of these tile-based abstractions is that they enable TK to statically check layouts and operations, which is important because GPU kernels are often difficult to debug. For example, in-register tensor core multiply mma AB requires to be a row-major layout, and B to be in a column-majo layout, and TK can raise compile-time errors these conditions are not met.
Figure 4: Shared memory bank layouts, illustrated for a 16x64 16-bit tile; each memory bank has its own color. Top left: naive, row-major layout. Although loading rows is efficient, loading into tensor core register layout suffers 8-way bank conflicts. Top right: A padded layout, which has no bank conflicts but consumes additional memory. Bottom: Two of TK's selected layouts, with compile-time selection based on width. (Bank confficts unavoidable for some tile sizes while maintaining good hardware support.) These layouts have 2-way and no bank conflicts, respectively.
Choosing a memory layout Layouts specify how logical data elements are mapped to physical thread ownership. Different tile sizes, types, and handware-acceleratee instructions benefit from different layouts, and some layouts lead to bank conflicts Our goals are:
We want our register tiles (the fastest memory) to keep memory in the layouts used by tensor cores (the fastest compute). Shown Figure [ (Left); each color represents a different thread's ownership over the data. These formats difficult use, further highlighted in Figure
5

---

Published as a conference paper at ICLR 2025
We want to support the use of tiply and bulk copy instructions), which also require specific shared memory layouts.
instructions (e.g., asynchronous matrix mul-
In TK, simplify to 3 layouts swizzled on 32, 64, and 128 byte boundaries and automati- cally assign shared tiles with layouts that minimize bank conflicts for their size and type. Seen in Section 4.2] even the FlashAttention-3 kernels written with CUTLASS templates can face bank con- flicts, hurting performance. Our approach helps minimize conflicts, reducing Chared in Section
3.2 BLOCK PARALLELISM WITH A GENERALIZED ASYNCHRONOUS TEMPLATE
THUNDER helps developers reduce overheads by coordinating how workers in thread block asynchronously overlap execution. Though the GPU hierarchy might suggest that we need wide variety of techniques, we propose a single concise template that we find enables high perfor- mance surprisingly broad range of AI workloads. We first define the template, which has four steps cead--copue-sttee seeee--trr--iiieeee finisfiihhhhhh (LCSF for short) and builds on the classical producer-consume paradigm (Dijkstral 1968 [Bauer al. (01I). We show how the LCSF template can help navigate the tradeoffs between occupancy and efficiency (reducing CHEM, CCompute in Section
Load function:
Compute function:
15 (varparoup: (warpid!)
karpgroup: ABE (att, scratch, gistato.idl, block.k))
tmarsexpest iinputs_arrived. black.k, block.v) tma:: ima::iass_aayne
Narpgroupi ham__asyye__eieee
// cottmax (simplicted)
black.k, globals.A [batch, head, inputa_acrivedl
sub_ron (att, stt, nax_vec); (act, div_ron(att, att, aosn_vec);
tea tnarsload maa::iod_aayyee
black.9, slobals. (batch, aead, itor, anputa__rrived)
(att_ard,,,,,, att]]
karparoup: tiaa (Btatt.. ati_bfle, block.9);
12 }
warpareap: mtpp:::m__sssy_sssseee (1 12
13 cise arcive sinpttesippttteerrtvveoi
azrive srrive (iaperrtts_s__tithce 13
Figure 5: A simplified depiction of attention in the LCSF template highlight the role of different specialized workers. Left is executed by workers that manage HBM to SRAM memory movement, and right by parallel compute workers, which operate in fast memory, registers and SRAM.
Programming abstractions per Section 2 AI kernel usually load tiles of large tensors from HBM to SRAM, perform computation in fast memory, store the result for the tile back to HBM, and repeat this for the next tiles. To use the LCSF template, the developer writes four functions:
1. Load function. Specifies the data that load workers should load from HBM to shared memory, and when to signal to compute workers that this memory ready for use.
Compute function. Specifies the kernel instructions that compute workers should execute, using the tile data structure and operation primitives from Section 3.1 Store function. Specifies what data workers need store to HBM. Finish function. At the end of the kernel, the workers store any final state and exit.
TK provides abstractions to help the developer manage worker overlapping and synchronization.
M N= K
Stages
TFLOPS
4096
1
260
The template maintains N-stage pipelined buffers in shared memory, which are used for loads and stores from HBM. Load/store workers add/re- move tiles of data from the buffers, based on the status of compute workers. With single stage, load workers would need to wait for all compute workers to finish ex- ecuting before replacing the input tile. A 2-stage buffer can hide the HBM load (store) latency since the next tile can asynchronously load, while the compute work- ers execute on the current tile. Deep buffers can reduce the synchronization required across compute workers, allowing them to operate multiple tiles concurrently. TK lets the user set single num- ber to specify the number of stages, and manages the setup and use of these buffers for the user. In Section 3.2 vary the number of stages N E (1,2,3,4) 4] for GEMM kemel.
1. Multi-stage buffer:
4096
2
484
4096
3
683
4096
4
760
Pipeline buffer stages We measure efficiency in TFLOPS for our GEMM kernels as we vary the number of pipeline buffer stages in the TK template.
Table 1:

---

Published as a conference paper at ICLR 2025
2. Synchronization barriers: Load/store workers need to alert compute workers when new memory written to the imput buffer. Compute workers need to alert load/store workers when tiles are written to the output buffer, or when input tiles can be evicted from the imput buffer. Within the TK template, we provide an a arrive ve function for workers to signal that they have finished their stage.
3. Asynchronous I/O: We wrap synchronous and asynchronous load and store instructions, including cp async and TMA, in the same interface. We automate tensor map descriptor creation for TMA hardware-acceleratte address generation for our global layout descriptors (g1).
Tradeofl's between occupancy and efficiency TK parametrizes the number of load/store and compute workers (or occupancy) providing simple way for developers tune their kernels. As discussed in Section 2 higher occupancy increases overlapping, but creates contention over limited hardware resources (c.g., regis- ters). With fewer registers, workers need to op- erate on smaller tiles of data, resulting in more instruction issues, SRAM to register I/O, and potentially higher synchronization costs due to the increased data partitioning across workers.
Synchronous vs LCSF Attention
341
TFLOPS
300
300
216
200
181
123
100
Worker Warpgroups
Synchrenous
Figure 6: Occupancy tradeoff: (Left) Attention TFLOPs function of occupaney, benchmarked with head dimension 64 and context length 4096. We compare synchronous and LCSF kernels.
Figure shows the occupancy tradeoffs for at- tention kernels. We consider (1) simple kernel that only uses warp level parallelism (Listing(2) and (2) kernel written in the LCSF template (Listing 5). Although with both kernels, performance increases with occupancy until resource con- tention dominates, LCSF expands the Pareto frontier beyond the naive kernel.
We find the general LCSF template to be effective across a range of Al workloads. We keep the template lightweight and simple by making opinionated design choices. However, don't want TK to get in the way of achieving peak GPU performance - TKK is embedded, meaning developers can use the full power of CUDA to extend the library as warranted.
3.3 GRID PARALLELISM WITH BLOCK LAUNCH SCHEDULING
TK makes easier for users to quickly try varied grid layouts and coordinate thread block launches. This help reduce the setup and tear-down costs for each thread block (Csetup in Section 2). and encourage memory reuse between thread blocks, to avoid slow HBM accesses (CHBM in Section 2
Block launch costs We provide optimiza-
M-N
K
TK-No
TK- Yes
CuBLAS
tions to minimize launch costs, centered around
persistent grid, where we launch thread
4096
64
93
108
69
blocks on the full set of SMs upfront, and sim-
4096
128
161
184
133
ply load the next task for the kernel within the existing block. We further eliminate pipeline bubbles by having load/store workers anticipate the next task and pre-load memory to prepare for future work, while the compute workers run the finish stage for the prior task. Table these optimizations for matrix multiplies.
4096
256
271
309
242
4096
512
414
450
407
4096
1024
565
600
633
Table 2: Persistent block launch TFLOPS for TK GEMM kernels with (yes) persistent and without (no) persistent launch as we vary matrix dimension K.
L2 reuse and block launch order Recall that thread blocks need to communicate via HBM. As introduced in Section 2 when thread blocks reuse memory, the data is often available in L2 cache, which significantly faster than HBM. However, cache eviction means that these reuse qualities depend on the order in which blocks get launched. For our attention and GEMM kernels, we measure efficiency as we vary block order, summarized in Table Block order substantially affects L2 reuse (measured through HBM bandwidth), which in turn control kernel performance.

---

Published as a conference paper at ICLR 2025
Matrix Multiply (M=N=K=16384)
Attention Forward (D=128)
Block Order
HBM GB/s
TFLOPS
Block Order
HBM GB/s
TFLOPS
{8. N. M/8} 982
805
(N, H. }} H.
213
600
(N. M}
3,070
392
2,390
494
Table 3: L2 reuse: We vary the block orders and measure both consumed bandwidth from HBM (GB/s) and efficiency (TFLOPS). For attention, we consider an optimized kernel, with an internal tiling of 8 rows of blocks, versus naive kernel that schedules blocks in row-major order. For attention, we compare block order (I) sequence length N, heads H, and outermost batch vs. (2) innermost B, H, then outermost N. Block order has significant performance implications.
EXPERIMENTS
In experiments, we validate that THUNDERKITTENS speeds up broad range of ML primitives. We compare to well-optimized kernels from prior work, written alternate frameworks such as CUTLASS, CuBLAS, general CUDA, and Triton. We compare our kernels for the "workhorse" operations in AI, GEMM and attention, as well as kernels for emerging AI architectures, such as linear attention and state space models (Section 4.1). We profile the kernels to understand TK's role in achieving high performance in Section 4.2 Kernel listings, in the TK template, are in Appendix
4.1 TK ENABLES SIMPLE AND PERFORMANT AI KERNELS
We evaluate a suite of TK kernels. We bench- mark on an NVIDIA H100 80GB SXM GPUs using CUDA 12.6 and report average TFLOPS. We provide experiments on an NVIDIA RTX 4090 and Apple M2 Pro in Appendix B
Multiplication
CUBLAS CUTLASS
792780
B0C 0.30
703
ThunderkAtens
658
690
E00
580
Workhorse kernels for AI Industry teams and researchers have made significant invest- ments into optimizing GEMMs and attention over the past several years (NVIDIA) 2023
269
Dao et al. 2022b Bikshandi & Shah 2023a
Shah et al. 2024 inter alia.), two workhorse
Figure 7: GEMM kernel from CuBLAS and TK.
operations that power the Transformer architecture (Vaswani et al. 2017). While the baselines are strong, TK kernels match or outperform:
GEMM: We compare to the strongest baselines: CuBLAS(NVIDIA 2023), OUTLASS
(017)). A single TK matrix multiply kernel, with just 40 lines of device code, is competitive.
Attention: We support multiple variants attention: causal, non-causal, and grouped query at- tention (Ainslie et al. 2023) head dimensions 64 and 128. We compare the strongest baseline,
which is concurrent to our work: FlashAttention-3 (FA3) (Shah et al. 2024). TK competes with FA3 across sequence lengths on the non-causal forwards pass, and outperforms FA3 on the causal and non-causal backwards pass by over 40% at shorter and 10% at longer sequences.
We find that TK makes it easy to use the GPU effectively by simplifying the choice of memory layouts, exploration of grid patterns for L2 reuse, and selection of occupancy and pipeline depth. The baseline kernels successfully use specialized H100 instructions and manage memory. However, existing kernels are relatively complex: FlashAttention-3 proposes a "ping-pong scheduler" for workers, and the CuBLAS library is >60MBB in CUDA 12.6 (Table (5. containing many tuned GEMM variants and logic to select the best option at runtime (Schuetze (2024). With TK, we remove the ping-pong and maintain FA3-level efficiency, and we compete with CuBLAS the demonstrated matrix sizes, using a single GEMM kernel (entirely in Appendix C.T).
Kernels for emerging AI architectures In addition to supporting peak performance on popular operations like GEMMs and attention, TK is also designed to be extensible to emerging Al work- loads. We release kernels across recent ML primitives, including linear attention (Katharopoulos al. 2020). FFT convolutions (Cooley & Tukey 1965), and state space models (Gu al. 2021).
8

---

Published as a conference paper at ICLR 2025
Min Fwd (Non Causal, 8m16, Hm16, Dm126)
Attn 8kwd (Non Causal, Bm16, Hm16, Dm128)
Teshh-ttertion-3
700
Feshhittemion- derkittess
667
668
E35
092
500
000
45.
547
.99
500
253
50B
400
446
140
19.9
10.99
R0M4
51280
Size
12288
Atti Fwd (Causal, B-16,
16, H- H6, D=126)
Attn Bkwd (Causal, 8m16, H=16, D- 128)
800
Flastiittenttoon
Fashhittertion-
Thundunnitaas
Thundar
500
600
537
190
400
400
422
468
4.49
494
437
330
2648
300
153
100
100
Size
Figure 8: Attention causal and non causal inference and backwards pass efficiency.
Linear attention We optimize two different classes of linear attention architectures, polynomial- based feature maps as in (Arora et al. 2024 Aksenov et al.) 2024 Keles et al. 2023 Kacham
ct al. 2024) and learned feature maps as (Zhang et al. 2024b a). In Figure we compare to
the strongest available baselines: the popular Flash Linear Attention (FLA) CUDA kernels Yangg
& Zhang 2024), which are written in Triton. We show TK outperforms FLAA's polynomial-baset
linear attention by 14x. TK outperforms FLA''s learned map linear attention by 6.5x.
State space models The long convolution, implemented with Fourier transforms using the convo- lution theorem, is the key primitive in popular state space modeling architectures such as S4, H3, 2021 Poli et al. 2023 Fu al.) 2023b Hasani al. 2022 Agarwal al., 2024 inter alia.). ). In Figure 9 we compare to the strongest available baseline: the FlashFFTCom
and Hyena (Gu et al.
CUDA kernels in Fu al. (2023c) and show TK outperforms the prior work by 4.7x sequence
length 4096 and 7.9x at 1024. TK outperforms PyTorch's 's FFT operations by to 8.7x.
We also optimize the recent Mamba-2 state space model (Dao & Gu (024). We provide a TK kernel that outperforms
Triton kernels prior work Dao & Gu (2024) by > 3x (Figure 9.
This gap is primarily due to the ease of fusing complex operations in TK.
We also develop kernels for common memory- intensive Al operations fused dropout-residual- layernorm (Ba et al. 2016), and rotary (Su 2023) and show TK is effective, We compare to popular Triton kernels for these operations,
TK's s programming model is extensible. In Appendix we demonstrate that TK extends across (1) hardware platforms by providing competitive NVIDIA 4090 and Apple M2 kernels, (2) preci- sions by providing a competitive FP8 GEMM kernel, and (3) tile shapes by providing an attention kernel for arbitrary shapes.
4.2 COMPARING KERNEL IMPLEMENTATION
To understand TK's improvements, we profile the kernels using NVIDIA's NSight Compute (NCU) tool. In Table 4 give NCU profiles for both the emerging long convolution primitive and the well-optimized attention backwards pass, comparing to the strongest respective baselines.
Long convolution: We profile FlashFFTConv (FC) and TK convolution kernels B, D 16, 1024, 4096 in NCU. We find TK helps both with overlapping the workers (indicated by higher issue slots and fewer memory stalls) and in tensor core utilization (4. increase). This enabled by our TK template, and use of TK warpgroup operations (which saves registers and establishes a SMEM to register memory pipeline through warpgroup matrix muliiiiluuttppyyaaaa operations).
Reference Triton kemels from: httpps://giiii wwtp:/..commaooooo l.bb/ Flash attention

---

Published as a conference paper at ICLR 2025
Mammba22 18-16, Hm32, D-64, D-State 64)
Long Convolution (B-16, 0-1024)
160
Mamba Triton
Pylrrch
144
ThundurKttarg
128
140
ThundenSttann
FlssiPFTCnn (CLDA)
Spedd (1F000g))
44
4.3
1024
2040
4096
0192
16384
Sequence Langth
Seceines Lengh
Based Linear Attention (B-16, D-64)
Hedghog Lincar Attention (B-16, H-B, D-128)
232
175
FIA Triton
PyTorch
150
Thundenötteens
133
150-
125
Pylorch
100
200
50-
30
22
0,4
0,4
X
0.9
X
1024
2048
4096
9192
16384
1024
2048
4096
$192
16384
Sequence Length
Size
Fused droppoutresiduull-rmmm
D-1024)
Rotary
1.2
1.
Hashiorrm Triton
1.0
1.0
0.7
0.8
0.8
1.5
0.4
1.0
0.9
0.4
0.4 0.4
3.5 0.5
6.5
0.2
0.2
0.1
1401
2016
0.0
1024
2040
Sequence Langtth
4396
0192
10304
Figure 9: ThunderKittens kernels performant across wide range of kernels.
Occupancy utilizations (%)
HBM
Shared
1mpl.
Tensor core
Issue slots
TPS (GB/s)
Stalls (Cycles)
Stalls (Cycles)
FA3 Bkwd
61.2
25.1
328
1.83
0.92
TK Bkwd
58.2
34.8
490
1.63
0.14
FlashFFT 54.8
13.4
25.5
14.8
2.5
1.6
TK
40.0
31.4
0.6
0.3
Table 4: NCU profiles for 1) attention backwards pass kernels from FlashAttention-3 (Shah al.
2024) vs. TK and 2) long convolution kernels from FlashFFTConv (Fu et al. 2023c) vs. TK.
Attention backwards: We consider FA3 and TK at B, H,N, = 16, 16, 3072, 128. The meth- ods match in tensor core utilization, but TK gives higher issue slot utilization, suggesting the oc- cupancy may be better-tuned. TK gives higher HBM memory throughput and incurs 10% fewer stalled cycles on HBM waits. For shared memory, TK incurs 85% fewer stalled cycles we find TK has no bank conflicts, but NVIDA's NCU profiler reports up to 9.6-way bank conflicts in FA-3.
CONCLUSION
Given the challenge of mapping Al architectures to GPU hardware, our work asks how far we can get with few, easy to use programming abstractions. In THUNDERKITTENS we give abstractions for each level of the GPU hierarchy: tiles with managed layouts at the worker level, and an asynchronous LCSF template at the thread block level. We highlight tradeoff's for persistent block launches and L2 reuse at the grid level. The natural question is whether we sacrifice anything in performance when write kernels with few abstractions. We implement breadth of AI kernels in TK and excitingly find that our abstractions are both general and consistently meet or exceed state-of-the-aa We optimistic about the potential for accessible ways for programming Al hardware.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGEMENTS
We are grateful to Together.ai for making this work possible. We than Arjun Parthasarthy for assistance in developing the complex tile support and FlashFFT kemells. We thank Mayee Chen, Tri Dao, Kawin Ethyarajh, Sabri Eyuboglu, Neel Guha, David Hou, Jordan Juravsky, Her- mann Kumbong, Jerry Liu, Avner May, Quinn McIntyre, Jon Saad-Falcon, Vijay Thakkar, Al- bert Tseng, Michael Zhang for helpful feedback and discussions during this work. We grate- fully acknowledge the support of NIH under No. US4EB020405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-Aware), CCF1763315 (Beyond Sparsity), CCFL563078 (Volume to Veloc- ity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-0188 (Long-context) and
(Interactive Human-AI Teaming); ONR under Nos. N000142312633 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), BFS supported by a Hertz Fellowship, SA supported by SGF Fellowship, and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or rec- ommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed implied, of NIH, ONR, or the U.S. Govern- ment.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, et Tensorflow: Large-scale machine learning heterogencous distributed systems. arXiv preprint arXiv:1603.00 04467, 2016.
Naman Agarwal, Daniel Suo, Xinyi Chen, and Elad Hazan. Spectral state space models, 2024. URL https://aras org/
06837
Joshua Ainslic, James Lee- Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. The 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP), 2023.
Yaroslav Aksenov, Nikita Balagansky, Sofia Maria Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gorbatovski, and Daniil Gavrilov. Linear transformers with learnable kernel functions are better in-context models. 2024.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. International Conference Machine Learning, 2024.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
1rXiv:1607.0644 064500 2016.
Michael Bauer, Henry Cook, and Brucek Khailany. Cudadma: Optimizing gpu memory bandwidth via warp specialization. SC '11: Proceedings of 201/ International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, 2011.
Ganesh Bikshandi and Jay Shah. A case study in cuda kernel fusion: Imple- menting flashattention-2 on nvidia hopper architecture using the cutlass library, 2023a.
URL htttps:/reesaarchccm colffx www.nnnontniiii...moommmm
folffxx-ilashatttetiinn..
Ganesh Bikshandi and Jay Shah. Developing cuda kernels for accelerated matrix mul-
tiplication on nvidia hopper architecture using the cutlass library, 2023b. URL https://research. cowfax-
col F htt.aa/...ao.ooooooooeaeeoooooaa

---

Published as a conference paper at ICLR 2025
Li- Li- Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong. Xuanrun Zhang, Zuquan Song, Chengji Yao, Ziheng Jiang, Haibin Lin, Xin Jin, and Xin Liu. Flux: Fast software- based communication overlap on gpus through kernel fusion, 2024.
URL https:
06858
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, al. {TVM): An automated (End-to-End) op- timizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pp. 578-594, 2018.
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297-301, 1965.
Tri Dao. 21 Faster attention with better parallelism and work partitioning. Interna- tional Conference on Learning Representations, 2024.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. International Conference Machine Learning (ICML), 2024.
Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training, 2022a. TtLhstts::/ wwxiv. hatrg//ab///a 2204. 00595
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 10-awareness. In Advances in Neural Information Process- ing Systems, 2022b.
E. W. Dijkstra. Co-operating sequential processes. In F. Genuys (ed.), Programming Languages. NATO Advanced Study Institute, 1968. Lectures given at a three weeks Summer School held in Villard-le- Leard-le-Land Lans, 1966.
Venmugil Elango, Norm Rubin, Mahesh Ravishankar, Hariharan Sandanagobalane, and Vinod Grover. Diesel: Dsl for linear algebra and neural net computations on gpus. In Proceed- ings of the 2nd ACM SIGPLAN International Workshop on Machine Learning and Program- ming Languages, MAPL 2018, pp. 42-51, New York, NY, USA, 2018. Association for Com- doi: 10.1 00.1145/2211344.211135
puting Machinery. ISBN 9781450358347.
URL https:
//doi.org/1i 10. 1145/ 32155422114555 3..1352113
Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub- quadratic gemm- based architecture. 37th Conference on Neural Information Processing Systems (NeurIPS 2023), 2023a.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. In International Confer- ence on Learning Representations, 2023b.
Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. Flashfftconv: Efficient convo-
lutions for long sequences with tensor cores. arXiv preprint wwXiv:2311005900 2023c.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured
state spaces. International Conference on Learning Representations (ICLR), 2021.
Ramin Hasani, Mathias Lechner, Tsun-Huang Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209, 12951, 2022.
Horace He, Driss Guessous, Yanbo Liang, and Joy Dong. Flexattention: The flexibility of py- torch with the performance of flashattention, 2024. URL
cut Lass
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
12

---

Published as a conference paper at ICLR 2025
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketching polynomial kernels, 2024. URLhttps: hattps://axxiv. org.oo/////oeoo abs/2310 01655
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), 2020. URL https: http///xxiviorg.//. a...
abs http://///aab//....100000
Feyza Duman Keles, Pruchuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity self-attention. In 34th International Conference on Algorithmic Learning Theory, volume 201, pp. 1-23, 2023.
NVIDIA. Cuda templates for linear algebra subroutines, 2017. URL
com/
TVIDIA/cutlase Lass
NVIDIA. cuBLAS, 2023, URL.ttps://oos.....ia..om http:ss.w...aaa..cooocccc//baaa.
http:/udaaaa/// cuda/cublas/
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang. Zach Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, nigh-performance deep learning library, 2019. URL htthhtts: /ttp:///avvvvvv. otg.oo//////oooooooooo.ooooaea/
01703
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung. Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemys- law Kazienko, Jan Kocon, and Jiaming et al. Kong. Rwkv: Reinventing mnss for the transformer era. Findings of the Association for Computational Linguistics: EMNLP 2023, 2023.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Fu, Tri Dao, Stephen Baceus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional lan- guage models. Proceedings of the 40th International Conference on Machine Learning (ICML), 2023.
Jonathan Kagan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: language and compiler for optimizing parallelism, locality, and recom- putation in image processing pipelines. Acm Sigplan Notices, 48(6))519-55 1305 2013.
Nadav Rotem, Jordan Fix, Saleem Abdulrasool, Garret Catron, Summer Deng, Roman Dzhabarov, Nick Gibson, James Hegeman, Meghan Lele, Roman Levenstein, et al. Glow: Graph lowering compiler techniques for neural networks. arXiv preprint arXiv:/805.00007, 2018.
Fabian Schuetze, Reverse-enginnerring cublas, 2024. URL https///t/tttsss////ooooooo
over
Jay Shah, Ganesh Bikshandi. Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024. URL.
2407 08608
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En-
hanced transformer with rotary position embedding, 2023.
Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, pp. 10-19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367 196. doi: 3329973
15114//315550.33229999
StRLittts:://dd httl//iiiorgg/..// www.oogg//........ http//10..11////1 1011111433000
Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. Tensor comprehen- sions: cramework-agnostic high-performance machine learning abstractions. arXiv preprint arXiv: 1.0iv:1002.0444 04730, 2018.
13

---

Published as a conference paper at ICLR 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. 31st Conference on Neural Infor- mation Processing Systems (NIPS 2017), 2017.
A. Vyas, A. Katharopoulos, and F. Fleuret. Fast transformers with clustered attention. In Pro- ceedings of the International Conference on Neural Information Processing Systems (NeurIPS), 2020.
Richard Wei, Lane Schwartz, and Vikram Adve. Dlvm: modern compiler infrastructure for deep learning systems. arXiv preprint
2017.
Samuel Williams, Andrew Waterman, and David Patterson. Roofline: An insightful visual perfor- mance model for floating-poin programs and multicore architectures. Communications of the ACM, 2008
Songlin Yang and Yu Zhang. Fla: triton-based library for hhacdhaww--cffccccemmmmmmmmmtttt of linear attention mechanism, January 2024. URL https://ggtabb.com//sstttssssssd
flasb--reearaatteett..o
Michael Zhang, Simran Arora, Rahul Chalamala, Alan Wu, Benjamin Spector, Aaryan Singhal, Krithik Ramesh, and Christopher Ré. Lolcats: On low-rank lincarizing of large language models.
http:////gaaa/////a/// 2024a. /ttpt//////aiii.o///a///////////
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcu- pine: Expressive linear attentions with softmax mimicry. International Conference on Learning Representations (ICLR), 2024b.
14

---

Published as a conference paper at ICLR 2025
Our appendix is organized follows:
Appendix
provides an extended discussion of related work.
Appendix provides a set of kernel that demonstrate the extensibility and simplicity of TK kernels.
Appendix provides set of kernel listings written in the TK abstractions.
4. Appendix D provides extended discussion of the TK library implementation details, including approaches to constructing shared memory layouts and the set of layouts used in TK.
A EXTENDED DISCUSSION OF PRIOR WORK
We provide an initial discussion related work in Section and extended discussion here. We discuss various programming frameworks that support AI kernel development and discuss prior work to develop hardware-awar AI algorithms.
CPP embedded libraries Towards raising the level of abstraction and supporting simpler pro- grams, NVIDIA maintains the CuTe and CUTLASS libraries, which are CUDA primitives library for graphics, scientific computing, and ML. As discussed in Section both CUTLASS and TK can support the same kernels, since both are C++ embedded libraries. Developers can use the power of the full C++ library, including raw CUDA, when using these frameworks. Distinct from CUT- LASS's 's objectives, specifically explore how broad and fast we can go. just using small number of opinionated abstractions.
Compiler-based libraries Many machine learning frameworks employ high-level computational graph representations for optimizations, such as TVM Chen et al. (2018). TensorFlow XLA Abadi jet al. ( Glow Rotem et al. (2018), and DLVM Wei et al. (017). TVM, for instance, incor- porates flexible tensor expression language and automated schedule optimization, Building upon Halide's Baagan-Kellpp
al.] (2013) separation of algorithm and schedule, TVM introduces new primitives such tensorization. The approaches of these frameworks differ from THUNDERKIT KIT- TENS that they provide a full end-to-end stack that optimizes both graph-level and operator-level transformations, while THUNDERKITTEN, concentrates specifically on kernel-level optimizations.
Triton Tillet et al. (2019) builds on existing approaches to deep learning compilation while intro- ducing novel techniques for efficient tiled computations GPUs. Recently, tools such as Flex Attention also provide easy to use interfaces to write kernels for attention variants, and compile down Triton He al. (2024). Unlike XLA and Glow, which use tensor-level IRs and predefined templates, or Tensor Comprehensions Vasilache et al. (2018) and Diesel Elango et al. (0018). which rely on polyhedral models, Triton introduces C-like language (Triton-C) and LLVM-based IR (Triton-IR) centered around parametric tile variables. This approach supports non-affine tensor in- dices that polyhedral models struggle with. Triton's 's JIT compiler (Triton-JIT) implements tile-level optimizations and an auto-tuner, often enabling high performance on par with hand-tuned libraries. key difference between TK and Triton that TK is embedded within CUDA, and as a result its abstractions fail gracefully. In contrast, Triton supports inline PTX only for element-wise operations on tensors.
Hardware-aware AI architectures We are inspired by the success of several prior works that introduce systems-level innovations to improve ML efficiency such as FlashAttention (Dao et al. 2022b Dao 2024 Shah et al. 2024) and other optimized attentions (Bikshandi & Shahl 2023a). FlashFFTCon (Fu et al., 2023c). linear attention kernels Vyas et al. (2020); Yang & Zhang (2024): Peng et al. (202)); Dao Gu (2024). Given the effort required independently optimize each new architecture, ask whether a small set of opinionated GPU programming abstractions can help
researchers obtain kernels for a broad range of Al operations.
B EXTENDED RESULTS
In Section 4 we show that THUNDERKITTENS results in state of the art kernels across a broad range of AT workloads NVIDIA H100 GPUs. This section provides extended details of our
15

---

Published as a conference paper at ICLR 2025
experimental protocol, additional results to analyze the extensibility and simplicity of TK kernels, and additional results to further highlight the breadth of features.
B. EXTENDED DETAILS OF EXPERIMENTAL PROTOCOL
Benchmarking kernels In order to ensure fair performance comparisons between TK kernels and others, we run 10 warm-up iterations then use cudaEvents to measure total kernel execution time over 10 benchmarking iterations. Reported performance is the average of the 10 benchmarking iterations.
with and compile is any time the input configurations (sequence length, batch size, etc.) change. We report the maximum TFLOPS with torch.compile across settings.
Benchmarking Py PyTorch We measure baseline PyTorch without torch.compile
Tuning kernels Baseline GEMM kernels are tuned via a grid-search over the default execution parameters exposed (if any) and through auto-tuning methods (exposed via CuBLASLo) for baselines, the maximum performance achieved is reported. Furthermore, for Triton kernels, we run ariton.autotune over the default parameter configurations provided in baselines we compare TK kernels to. In order to avoid impacting performance measurements, kernel tuning is done in separate iterations prior to warmup and benchmarking.
We the following software versions for benchmarking: CUDA 12.6, Triton version 3.00, and PyTorch version 2.4.
B.2 ANALYZING THE EXTENSIBILITY OF TK ACROSS WORKLOADS
We include additional kernels to demonstrate additional TK features:
FP8 precision:
We provide an FP8 GEMM kernel TK and compare this to CuBLAS in Figure [T] The inputs and outputs are both in FP8 precision, and the accumulate is in FP32.
2. Padded tiles While TK uses 16 x tiles by default, to encourage users to utilize tensor cores and coalesced loads, it is important to support un-aligned workloads. TK handles this by padding the tiles (discussed in Appendix D). We implement attention with non- aligned dimensions the NVIDIA 4090 and H100 GPUs and find that the performance characteristics remain the same our aligned kernels.
FP8 Matrix Multiplication M-N-K)
2000
CUBLAS
ThunderKittens
1534 1532
Speed (TFLOPSIS)
1500
1439
1429
1396
1250
1121
1000
750
250
0
2048
4096
6144
8192
12288
16384
Size
Figure 10: GEMM kernel using FP8 precision on the NVIDIA Hopper GPU.
B.3 ANALYZING THE EXTENSIBILITY OF TK ACROSS HARDWARE PLATFORMS
We find that TK is extensible across hardware platforms. While we focused on the top-of-line data center hardware, NVIDIA H100, in Section 4 here additionally consider the top-of-line consumer hardware, NVIDIA 4090, and personal hardware, Apple M2 chips.
Consumer hardware: NVIDIA 4090 GPU We implement causal attention at head dimen- sions 64 and 128 using TK. We compare to FlashAttention D (2024), a popular reference kernel, and find that the TK kernel is competitive across settings.
16

---

Published as a conference paper at ICLR 2025
Attn Fwd 4090 (Non Causel, B=16, I1=16, D=64;
Attr Fwdd 4090 (Non Causal, B= II=16, D=120)
200
175
150
153
134
160
162
160
179
154
145
165
100-
100-
75-
2045
2048
Size
Size
Figure 11: Attention non causal inference at head dimensions and 128 and Based kernels, on NVIDIA 4090 chips using TK and the reference baselines.
Personal hardware: Apple M2 chip We implement non causal attention head dimensions 64 and 128, and GEMMs the M2 chip. We compare to the Apple MLX framework example kernels and find that the TK kernel competitive across settings.
Attn Fwd Apple M2 (Non Causal, B-26, H-16, D-54)
Attn Fwd Apple M2 (Non Causal, B-18, Hw16, D 128)
Apple MLX ThunderKitan
Apple MLX ThundarKiteens
(TFLOPS3)
3.5
3.8
(TFLOPU0)
3.8
4.0
4.0
512
1024
1536
2048
2560
512
1024
1536
2048
2560
Size
$ze
Matrix Multiply Apple M2 (M=N=KI
Apple MLX Thuncenkiteens
5.6
So0od ITLOOPOI
5.2
5.2
5.3
3.8
3.4
1 2
512
1024
1536
2048
2560
Sum
Figure 12: Attention non causal inference at head dimensions 64 and 128 and GEMM kernels, on Apple M2 chips using TK and the Apple MLX reference baselines.
We next include code listings for the attention kernel on NVIDIA 4090 fig. 13 and Apple M2 Fig- ure 14 highlighting the close resemblance between the implementations.
17

---

Published as a conference paper at ICLR 2025
NVIDIA 4090 attention:
templ emplatecint consteepr
16x1128/71; height of wach worker tilee Lemplatacint D, Lygemane Tebff16, Lypename Lerov_1> using gkva stct, nowS<Ds, D, 10;
cemplatecint D, typenane T-tlost> using atle rt<T, ROMS-DS, ROMS<D>>)
cemplatecint ni using shared tile (F
templstesint using tiyoai_ gl<b516, -1, -1, DS 17 N, st runtime, at compiee time templatecint D> struct glosels global L 1eyout<D> og, 8g. Vg. GE7 17
cemplatecint D> launch bounde grid constant globnlscn
globsl voiE lounch_ INUM MORKERS- TEREADS,
11
pairs of workers col cooossocsy load k, y tiles int loadid load groups1groupid). workerld cittens:swarpie // which worker 1? constexpt inz LORD BLOCKS NUX MURKES 2osd MARP8;
using Load_greup kattons:
12
14
const int batch blockisk.t, head blocksox.y, s_eeq blossIdx.x RIN MORKMORSERS Norkerid;
15
16
starrd slipansst_dermy stmili aaares_allosatee al(iiat=)6 =
18
shared ti exts (8k sment ILADD BLOCKS (PTPE, STAGES all.alloooeeeataree tilesto, toan BLOCKS, PEPR. STAGES> F) marrci_tileq (sy_emen) (LOND_ BLOCKS] [PIPE_ STAGES!
19
LOAD LOND_BLOCXS, BLOCKS, P1PE_ () snates te 1égo_ INUN WORXERS] rointerpret ctterppretccsscctrreed.... wiieee- WORKERS (eceEessss___emn))) the reqister ciles. gevo t:154,, bele> (_rco, k // and & are beth row layout, as wo use ma_.at. gevo bilet, btl6, to v_regi // is column layout, we use mna_AB.
20
Initialize
23
aive_tsis4B, float> stta_tsiss,, sloat> att_block) attention tile, flost.
//
tile.
atta t:le4,, 6t att block maj may // bt16 attention tile for secord maa We cast right betore. typoname atta_tile<2, dieat>::col_vec lex_vec_last, max_vae, aern_ves] // these aze column vezters. // each
27
tile 16x6
(aLDe2=RoM--D> a.Co.depst.)
30
isadci, foo_s=oooossooeeekeeeeee
2.02. [batch, e_seg, head, C)1 // going through shared monery
31
__ynoxarp(1/
load(e_reg, so_omon/workeridd]
syncthhooad((1:
25
- 501
q_re,, _flattabicoo1i10..1s 1.26229524499)) mcgg g_ www.tt001o0t111..0ooo..
eise co
1.8126950188917
28
nes_infty (nax_vec]]
29
zero inorm_vzc)
(o_reg);
launch the load the int ky_blocks dopth 1020 BLO3KS+BONS-DD-11 (1orp BOXD__BOCCS-ROOE--DD tic 02
42
43
www.rroroooaaaaaeeeoror. load_aroup: itt.pppppiad__sseeeeee.e false> (y_smen|laasid [0], [barch, loadid, head, itcrate over k, for those g's have been loaded
false> "Lasmen]ooddid [0], 9.Kg, [barch, loadid, head, 01);
46
Zor (auzo kv_idx = 0, kv_ids ky.blocks; kv_idktt, tiee(tie11:33)
47
Inz next_load_idx [wy_id==]) 2EAA) Liresttjoee_ideeeeeee
loadidy
g.Kg.dephh)) nexi_tie (tier1)
50
toad_group: asyaccc, faissh (k_ssen [naxt_sta], 9.80, (batchy next Load_ioadddd,, head, 01) grougg iload sayack, falueo (v_amen [next_-iel, 9.99, (butch, load__oedd head, 01); soudLasyo__uaited> //
51
stay in flight.
54
wlas waynLwayyccvvii (1s
synothreada (1:
apragna
Loan LoAn procke.
for fint fint subcile Di subcile LOAD_BLOCKS 64 (kv_idx<LOAD BLOCKS subtile) RRMS-DD> dep.ly my subtilert)
58
taadik_reg. ssssee ilel [tiel); // load k fron sbared Inco registers 15x15 atlenilen ana latc_block, a_zea, K_reg, att_block)] // QRX.I = [ky, int fill deppd right
block)
int
subt ile) -OOM // one past last RV index of Lile x RORS<D> g.Rg. Sapth iidat_index RO88<b>> att black, F111, sypes antossttttcflleee tneg. 1n
62
copy veo__veo_last, max_vec);
con inax_vee, at blosk, nax_ sub row latt block, exp2 black, att_blook); sub vec___aa max_we_llat,, nax_vee)) exp2 max vec vec_last, last, max vee last mul (oorn
66
black, nax vestz
69
70
vex laxt];
(nom_ block, vest; onpy late lock ma, att block); load torreg, s_mmeffs ils [Liel);
73
74
mit revie blocklockk me, me, v_reg, o_reg); regr
max lestlz
77
78
div reg. r_reg, norn_vec); synnshhswdda())
if ig RoM--oM < degg.desst) 1// write out o. store ligo aonaammnnia morkerirr, euregl; // going through shared memory improves coal coelescing of dran writes. synewarp
82
storecl, fal (g.00, go, enen/vorkeridd :batch, q seg.
013;
85
Figure 13: Attention implemented in TK on NVIDIA 4090 chips.
18

---

Published as a conference paper at ICLR 2025
Apple M2 attention:
namerpace custom ops
struct guboxp2 I7
tomplatestsperane T> static FUTAL_FUNC
6b)
big
complssexcyyeneme RI, typenana
staric METAL,
MENA FUNC typenare htt.aa/ee_.ddettttrrerooooeeeeee
(1 66 ducks:ris retp:rriss_rrrttteeeeeeeeeeeee woid
by >ritype
subexp2 KI idst, thread
RV irex_values)
SI, RV>-ddet, sec, roN_valoes);
11
12
complaoocyypenee RV, typenam: U>
atacio METAI_FUNC acbexo2(thzead idst, thread canst 4iha, thread const. 4tha)
:seablee irsuukkss re (eoterddi>> voidr::typa
htw.:pccuutonneccccooooo.. R2> (dst, 1h5, Ets!)
16
/oonstant
reatine
19
saplaaescnnt
20
sernel vold ateead_ter IMITEND BERTEEDD_ERREEEE FARANS)
128, to be 128113
Loing qichal__ayoo kittenssions sostena:orre:oo waw.reeicffooattt.
-1,
23
glomci_layout ql_q acliptr, B, H, matiptels
24
globel_layout g1_
auliptx, H, nallptels
acliptx, acliptr, N, nuliptels SL_DES8, DP2 Laing staake
clobst_iavous
nullptzls
gionei_layout q1_.
27
28
rt_bike,
Laing aL
=L_bbR8, day
31
csing stcact st_ B>)
32
uxing
flx8 niz
caing
si<8 Briseol_vec7
consl int. block blackTex..2; const int bead blockIds.y:
36
(blockIdx. NUN_RORKER3) wazpidj
tonel kv.bkocks at st_skvssroos
39
st_AV 4_2997
40
kt.
sL_g v_regz
st_att att_blocks
43
*t.0 o_reg:
44
nex_verc_aaty
nax_vees
reaatt norr_vecz
47
loadigg siuge (block, head, g_awg, 0), lameId);
48
infty (nax vec);
(nornueest; zero togreg]
51
constexpr const bé16 q mol - 140 - 1283 7 D. 0.080388347644 0-125bE1 1.46269504C888 mul (g_r g_reg, loop uncoll (fu12)
52
Apragna
for tauta ky idx 0;
kv. blockez kv idxt4)
Load (k re_ gi_3, Sblack, head, idd, 31, lanetd);
56
block);
block, rereg ter, block);
copy vens_ven nex_vee, laneld); max vec, stt_block, ves, lsneidi; subesp2 tnax ven last,
59
60
last, vecl;
subexp2 block, ati_bleck, nax_vesl; mit (norn vec, nomn norn_vec, vec, wex_vee_laase
63
row son (nomn vec, att block, norn vec, larefd); aLawe, nex_vec_last))
64
rus
toea reg, gl.s, (block, head, kdx,dx, laneId); me reg. att bock, v_reg, ouregt:
67
68
div. roxio.
wea);
store(g). aureg, (blockk g.eeg, 01, lareTd);
Figure 14: Attention implemented in TK on Apple M2 chips.
19

---

Published as a conference paper at ICLR 2025
B.4 ANALYZING THE SIMPLICITY OF TK
As a proxies for understanding the simplicity of the TK library, we measure (1) the size of various popular frameworks in bytes and (2) the lines of code across our kernels.
The library sizes are shown in 5 For CUTLASS and TK we report the size of the "include/" directory, and for Triton we report the combined size of the "include/" directories in Triton plus the "include/" in the core MLIR compiler dependency.
Library
Size (Bytes)
Date / Version
CUTLASS
22 MB
10/22/2024
Triton
12.6 MB
10/22/2024
TK
<1.0 MB
10/22/2024
Table 5: Sizes of various CUDA libraries.
We find that the TK kernels in Table 6 average at 200 lines of code. We compare to the lines of code in the corresponding state the art baseline kernels, and the TK speed ups over these baselines. While measuring lines of code may be difficult, we provide links in the table indicate our approach. For TK, we include many comments, all the global data descriptor, and custom functions. We exclude the python bindings and other wrapper functions for all baselines. We generally observe that TK kernels use fewer lines of code and provide speed ups.
Workload
TK kernel (LoC)
Reference kernel (LoC)
Speed up (min-max)
Attention forwards
217
2325 (CUTLASS FA3)
0.87-1.14x
GEMM
84
463 (CUTLASS)
0.98-2.05x
Convolution (NV 4096)
131
624 (CUDA FlashFF ICCCV
4.7x
Based linear attention
282
89
(Tnton)
3.7-14.5x
Hedgehog linear attention
316
104
Triton
4.0-6.5x
Mamba-2
192
532 Triton
3.0-3.7x
Rotary
Triton)
1.1-2.3x
Fused layernorm
146
124 Triton)
1.0-2.2x
Table 6: Lines of code (LoC) across TK H100 kernels, state of the art non TK kernels, and the TK
speed up over the reference across the evaluated input dimensions in Section
20

---

Published as a conference paper at ICLR 2025
C THUNDERKITTENS KERNEL LISTINGS
This section first recaps our benchmarking methodology for the results and provides a set of kernels written the TK LCSF template and tile abstractions:
1. Appendix
GEMM kernel
2. Appendix
Long convolution kernel
3. Appendix
Attention kernel
Appendix
Rotary kernel
To introduce the template components, we describe the GEMM kernel in detail in Appendix C.1
Benchmarking approach Our kernels in Section 4 are benchmarked on an NVIDIA H100 80GB SXM GPU with 10 warmup and 10 timed iterations using timings measured in C++. We also provide Python-boumd kernels and benchmarking infrastructure our repository for reference.
C.1 MATRIX MULTIPLY
First we show and describe a TK GEMM kernel in the LCSF template.
Each compute warpgroup is responsible for computing 64M-row, 64NN-column chunk of the result- ing output matrix. Each compute worker identifies the coordinates for its chunk, zeros its accumula- tor registers, repeatedly runs large asynchronous matrix multiplies (compute), and finally stores out its tile in the end (finish). The load workers also compute their coordinates, and then repeatedly load chunks of the input matrices (load), Store workers perform asynchronous stores when the compute workers are finished with the chunks (stores).
Tuning the number of workers and pipeline stages The computation is divided into stages, with each stage processing 64 elements along the reduction dimensions the input matrices. The input pipeline is automatically sized by THUNDERKITTEN if the user does not specify value. For common configurations either a (2 compute warpgroup) 128 x 256 or (3 compute warpgroups) 192 192 output tile per block, it generates 4-stage pipeline.
Tuning the grid order The greatest complexity of this kernel is in setting the grid parameters. This kernel adopts 3D stride over the input matrices, which has a significant effect for large matrices which do not fit in I.2 cache. The order in which blocks execute strongly influences cache locality and thus available memory bandwidth. To illustrate the magnitude the effect, comparing the presented scheme versus naive grid (in which blocks 4096 x 4096 matrix multiply only drops from 767 TFLOPs to 735 TFLOPs, but a 16384 16384 x 16384 matrix multiply drops from 797 TFLOPs to 387 TFLOPs, 50% performance degradation.
executed in row-major order) a 4096 x
using namespace kilLens;
using namespace kittens: sprotctype
using namespace kittens: sprpecn:sspoottttt::::e: template<int M_Rt,ok, struct matmul_layout
N_RTOCK>
using base_tile
st_bfe61, 61>; gl<bf16, 1, 1, -1, -1, bese_tile>>
using globrl_layeut
struct globels
global_layose à, D, C; 1;
struct imput_block
base_cile BLOCK] BLOCK]] 1:
struct Finish_blon
base_tile c[K_ [N_SLOCK); 1:
struct common_state
int2 coord; 1
struct consumer_stale 13 11
rt_[1<16, MLBLOCC--bas_
accunt
tamplatekint ___BOKK=2,, int _L_3L0C=44, int _SPPER_M=12>
15
struct static constexpe BLOCK BLUCK, N_BLOCK _LLLOCK SUPER_M SU using layout matmul pottml_laayoutee Layootts__aBB. SMyout BLACK, M_RTOCK>;
17
using wide_tile st_br<54, BLNCBL static constexpe int WOW.LOOOOSSOAPPPEEEEEE BARES- INPUT_PIPE_STASEESS
4 PRODUCER_ WAW.REAAAAARARRRAAAAAAAA
21

---

Published as a conference paper at ICLR 2025
Hellpe: Functions
template<beol PERISASTENT t host__ stetic inline dim3 grid int int K) return dim3 (PERISISTENT_ ? 132 MAN/
wttyyyttttttttoooooooo thase_ til:: tnum_nlemente
// TrunderKittens template funct device static inline void common_setar (comsonn aetup_argeslaoou. int Bblocks arge.glabsls.
BLLCLO Cblocks ergs.gichals. C.coll / (BLBIOCK
+641;
int super_rows (Rb) Roblock/SSPFRR
(AAPFRRMMMMSSPP M)SPPRR__M
rinal_rows
super_
10
super_repeat = SUPER
int task_id arga. argg.tesk__terrggrllD..o if (task_id arga.common, ceord SUPER (task (tast__d/sspperrrppes + task_ishSU
blockIdx.X;
Cbiecks)
else if (task_id ld Rblocka+Cbloekk 15
zemeindes_ Lask_id Cuper_k
17
common. ecord super_rows + (emminder_idhhfiaaLrooe) remainder_ id/finsi_ )
else // Id is high, 20
work to do
anga.num_ites. -17 return;
22
arge.num
args args.
www..glloosss..coom
24
int id = parpsccup:tgroppi.e. == NUM WWW_CCOSURRR
warpgroup: (ro //
4 produces sets as cood id, arga. cooré. BLOCK
25
args.cormon.com segs.cormon.core.
producer
28
device static void setup (producer. setup wttp://dweeeeesuppp.aeeeeaeaeeee azgs) {
// decrease registers producers
31
device statis void
ergst {
if (xarpgroup:: (xarpiid :xarpid teapsexpact (args.1npute hnpptt_arriieeeee angs.inputl, fint
B_BLOCK;
35
tma:: ina::load_asyne (args iapgt.aapt.. args globalbb.s..,,
arga.iter), argo.inputs_errivedd
37
fortint 0; i < N BLoCK, i++)
tma::loam load_aood_aano (arga.1nput.ar ioplt arga. glga.globll.....
(args.iter, arga. coord.ytily arga.inputa_arrivedd)
40
42
3;
struct consumer
44
device__ statie wold setup (consumer, wactnnsumerr__etteee www.ssetrgsssllaaeeeeeeeeeseesse lagostt args)
45
wazpgroupt: iit.epppiiineeeeeeeeeeee hetpeaassseeesssseseeeeeeeeeeeee
// increase registers foz conaumera
(aro(arga.attte..nou
47
device static compute (conssmer__oopptteerrreeeee
args)
49
impgroupp::ma__t
arga.state.acot // dest registers
51
inpot a [Marpgrropp (tt.pprooupprrouppiliiio sgroupid()), // matrix
(args. (args.input.by // matrix
54
wazpgroup::nma saapppppsmmmaaayc_wwlllll 1f (flaniid) 0) arrive aagg..inputs_fiiiieeee..
56
device statie woid *Inish tconsumer, www.shnnnnssseeeeeae argacclyoos args) warpgroppitatoee htw.:rattteeeeeneenneenne teeintergree__oststtieee Ltwleeiwweedeeeeeeeeee_eeeeeeeee (args, finish.c (atpoe/waaaggoogggoo () 4 .state
58
sarpgroup:tayne()
60
1E iwwappgroup::warppaee -- 3) for tint 1 0; i BLBLOCK; itt) hatppooe__.ccctrggggoooooaaaee..
smat::toreeasyyccssrr
Sesiiihh.wwwargreupp wttprruup::ggoueee groupid(() [i],
62
sergs.comman wowmon.coord...... coord.x, args. common cammonn.cade coord.yti]);
tmas:store_nn aed_aaass________aee_____ memory
// wast that store 1a rinished before reusing Firish
(args, www.sstteeccc.oo accun);
66
[f(laneid =- 0) arrive (args. wwwishhnniiiiiiieeee
68
1:
Figure 15: Templated matrix multiply kernel which reasonably competitive with CuBLAS.
22

---

Published as a conference paper at ICLR 2025
C.2 LONG CONVOLUTION
This section shows the long convolution kernel for sequence length 4096, written in the TK ab- stractions. We the FFT convolution algorithm, computed via Monarch Matrices, for long convolution kernel Coooley & Tukey 1965 Fu et al. 2023a Dao et al. 2022a)
struct consumer
device statis void setup pocnsumer_saee
arga)
pmmpgrooo::coosuumm ceoonsuer__reeeeeee seersseeraeee CONSCMMRRRRRR WARPS/4> (1;
iters_pe _ters_pes_heed
CONSUMER WARPGROUTS-1]
b MIM Www._CCOSIIER_AARPRRRRPSSS
stga.state.cuureen
(0 itera_per_headl +132 blocxIdx.X; // start for iter
using
consumers: :oad (args.
args.gicbals.f,
[0, 0, 0, 0]11
args.glcooll..imn
0, G11;
conaumera ra:: : load loa.fargs.sraccch.tt
args.gicoals. to,
[0, 0, 0, 011;
consusers: :ow.:::::oodddge..eooooo load_heee
twinv args.glcoals... twinv
11
args.globals, atgs. arg..atate.curren__neee
device_ static void compute (consumeeoe
arga)
warpgroupid warpgroup: (ttkkk/kkkkktttteeeeeeeee KARPGROUP KARPS; Int befault_bacreee warparoupid + // P
16
18
ert_fik16, 64s ama_reg; // 64 registers
ert_bf<l6, accum, tmpt warpgroup:::m (mma args. segs.ssratch.f.ree args, input. (xarpgroup: (groupid warpgroup::nm :wm (mma_rag. imag, args.ccratch. angs.inpute x (woppgroup: igrouiidd ima__yoop:ma___yy_______
registers each
21
23
copy (accum, mma_reg);;
sarggroup::loa((::: azgs.
twiddle first
25
mul (accum, accum, tmply
26
grOUp<tUKKCCONSSNER (.._::
warpgroup:::m (mma_reg,
args.scratch.flp
28
"yrpgroop:::maaayyynaal copy (accum, mma_reg);;
30
warpgroup::Leed(ttrpp args.scratch.kEl, // for filter second
(accum, accum, tmp);
32
httpr/ppp:smmm__mmaa_____._ accum, args.sorstch.fsrvv)
33
())
(acoum, mna_reg);
(oad argg..oratch..wwiny
// twiddle inverse 1s pre-transposed
(accuri, accus, imp);
37
warpgroup::atcre htr.::pppc..tceee.a.ae...aa.e..a swr.sgg.scaatc..cm.. : waspggroup:oeeoupidl)) accum) // must store for A:B
2yr (defaaull batbeere__brrerddd
39
warpgroup:::me (_ma__ee args. sorsschhhnhh,ing arg.. (rro ()
40
varpgroup::storr wwtpra.....t......o..
wwa.pppttooorrggouuuup.. wwrpgro ()), mma_reg.reall) real)
42
44
(E(laneid == 0)
arsive (args. sn #inished); npors_finichedd aor.
46
47
synowarp
49
Int. pess__pe__hed (arga. sl.ga.globsiee..baad.. NON_CONOUMER_MARROOO 1)
50
int next heat_head-- = (largs.iter-))
itter_per her_p + blockidx.x;
(next_head argg.statee.uureent...
52
head_aeeddaaaa (argaataaaraar.oo scratch, arbals, next_head);
53
ctr.ettte..rrrttt.eeeeeee next_head;
55
device
void finish iconsumer frgeehhes args) (laneid() finih. sainish__ffiini finishad);
b arrive (arga. finish
Figure 16: A convolution kernel for context length 4096, written the TK LCSF template, which outperforms FlashFFTConv (Fu et al. 2023c).
23

---

Published as a conference paper at ICLR 2025
using namespace kittens;
uning namespace srctens::prottppe::
using namespace kittens :prototypp:: loaf; templatesint _Mg2
static constexpr int ffLconv_4096_ using seg_tile seq_leyout st_bt<64, 64%;
_ggi
gl<b516, 1, 64, 64, seq_tile>; using layout_
egl L. -1, 64, 64, seg_tile>>;
using fft__ atruct globals seg_layout 9, x7
eglsgl<bfie, i, 64, 64>>;
10
12
ayout filter_ ayout kf;
IIL_layout I, rinv, tw, twiny_tz
struct input_bloc 15
segitile x[xg]; 1:
output_bboo struct bcratch_blees est_b kf, finv, tx, twinv_t, tmp:2]; };
seq_tile o[wg];
17
19
struct consuner_state int cucrent_heady ); 20
21
struct fft #ft_ 0_9.40066_ee__ template
atatic constexpr NUM CONSUMER KARFS-8, CONSUMER WARPGROUPS--UNG NARPS/A,
4 NIX_BLOCKS-1, OUTPCT_-39R-:1105---2 INPUT_ SIPE_8TAGES-4,
using layout #ft:oov_4096_aayeeeee // mine
CONEUNER WARPGROUPSS
25
26
device static inline void load_head (typename layout::scra:ch block &scratch, const
4 layeut:iglobale
int head)
27
using consumars groupINNMM CONSUMER
CrCOppcNMMCCNNSIMRRK CCNNIMMR__ARFFS> WARFS>;
coneumera::syn0
29
con chcratch. g.kf, (0, head, 0, 011 // sext chunk consumers::
31
32
//
devico static void cammen_sstupp (eommon_setup aot.oaarepp.....lyoottt... argal int teada_hand = (args.g gobb a a.x.a.xx...
by xeeett++13" -blockTox.x) / 132; // an guaranteeing
batea is wea.g..glbasss..ceo
block.
35
int pesrs_per_heeee WARFGROUES;
(AANUUNNNRAPPPRRUUUSSS)1 /
36
wrga.nom_iters. arga.task_iter =s G heads_handled iters_per_head -1;
37
producer
39
device__ static void setup stgp:oupppppccuuee___eeeeeeeeeee
args] [
41
device static load (prcduser, lor_____odd____slaaooo_o argsl |
43
int itera_per_head NARPGROUPS;
Aarga.globals. sab.la....i..... batch CONSUMER_WARPGROOEE
4 WWI__CCOMSNEELMMACCOUUUSSS
int head largs.
iters_per_heet _e___ee__eeee_e_ee +132 blockTdx. NARPGROUPS;
45
int batch tacge.ite: iters_oer_headd * WUN_COONTAERRWAAPPRRUUPP)) [tweerpgrou:::wwappddd tma: mrat:expact__bae sepeeeee_ttttteeeeeeeeee (args, wtt.ss..a.tssarrtaaaoo..o
args.iter:d]
47
sizecf (args. ingpriinuuu.... xx]) ((int)
+ NUM_CONSUNER WARPGROOPS, (int) (art) (args.globss wew.gggbblll....ee...
batch)));;
48
Tor fint b batch;
Batch=SUM_CORSUMTTE WARPGROUPS *s azga. sx..e..loolls..bbble. bte)
+
tnarsload_aaes + inpats_acrivedd
bnppu.x(b-batch), args.globals.x,
head, J, arge.
51
if (laneid)) == 0) arrive iogaa.a.tttttarrrrteeeooo 3); // excra arsivals neoded synowarp!);
53
device static void (producer_ attplucce__stooeeaaeeelllaaaal aaggellyougss args) 54
56
int peers_perheedd hargs.globsie. baabsle...batchh BEW_LCNNNOEEEAARRGRROOOO
57
int head = strrs.iter" / sters_per (oar_head)+133 +132 int batch large. iters_per_heac) MOM_CONSUNER WWW.LCNNSNER_NAARROOCCC NARPGROUPS; LE [EIearpgroup:::arppill -- itgrriter
blackTdx.
fint batch;
MARPRROUUS &6 args.globals.xx batch; b++)
61
ana:: (arge..cbbll,,.... args.output. (b-batchh, { b, head, 0, ));
63
patrssstrr___aan.._m waid
([(laneid() 0) (ergee cwt.eeeeeettteeeeiiieeeeeee 4)
65
syncwarp();
68
3
24

---

Published as a conference paper at ICLR 2025
C.3 ATTENTION
This section shows non-causal attention at head dimensions 64, 128, in the TK abstractions.
exp2 (args. atete. mex_vee__ast_scaieed arga.stete max vec__ee__laae last_pcsiedi
nul (args. atate.norn__eee norni_ees, args.atate..dm nernaLve, args.state statstttt.eee.eeeeeeeeeeeee max_vve_last__ealll wow_sum/argg.ee state.
noren.eee // accumulate
onto the norm_vec
atgs.state..__ae arga.atate.mas seteeeeee_____etttsollllll //
normalice o_seg copy att.gs..satee.t...ooc...
atg.: // convert to bE16 for mma
// AB
waspgroup:::mn largs.
www.gg.attt......oo............e state.o_seg, state.ata. attookkkmaa azge.input.vl,
())
if(laneid() == 0) arrive (args. fntpsiinntaaaaissheeeee sinis)))); // done!
10
device static inline commarg.. htwp://wwwtaawwwrrrroooooosoooe
(conumeerce argscllyuuts arge) {
12
164 >= args.globals..... rorooss return;
// out bounds?
wtw.:owwargs.satee.... args.
szgs.state.... norml
auto sc_ ccsntarpree_aarrttpppmmm. layout: top_ttllleeeeooooe.ooo.e scrach. www.acct.ttaaaaeeeeeeeccoooe qlwaggggogggarrrpppp:
groupid(1
15
rarpgroupssalozed (o_ acge.alate.o_eeg))
sarpgroup::ayye (aarggooup::grouii(11)
17
if (warpgroup: (vrppid() ==
tmas:store aacrre_sssoceaaao (arg o amen, larga, common, batthh, args. rygs.common.haet arga, 4 common. wew.....oonLLuuueesaarrrrrrr grop 0));
21
11
kernel kittens: hete
Figure 17: A templated non-causal attention kernel for head dims. 64 and 128 that competes with FlashAttention-3.
25

---

Published as a conference paper at ICLR 2025
waing namespace kittens;
using namespace kittena: sprototype;
uning namespace kittens: :prototype:
templatesint D, int NUM KORKERS> struct actn__we_laaour
using ky_tile st_bf<51, D>;
ge_tile
=t #__bfo0m=6621922123, u>;
using qo_globbl kittana::gl<bffG,
-i, -1, go_tile>;
using kv_ kv_global kittena::gl<bff1, -1, -, -1, D kv_tlle>;
struct globals qo_glebal 0r kv_globa:
atruct input_ b ook kv_tile struct coratch_bloss
1
qo_tile [[NOK_NORKERS]] 1
struct
batch, head, segi
13
struct consumer_
rt_fl<16, go_tile::0ola>> o_reg:
15
col col_vectrt_fille w_tille::row>> nax_vee, vec;
coi_vecsrt_flll16, ky_tile:srowep rax__ee__est__caale mex_vec_acaledi
17
zt st_f1<16, kr_tile::reRe> at=_block;
st_bfe15, kv_tile::sews> block
20 11
templatecint
sttn__wd__tte fed_templateeeeeeeeee
22
static comalexpe int CONSOCRR HARES - 12, NOR_HORKERS CONNUUMERNARRSSSAEE INPUT_PIPE_S200
23
using layout attn__nedlaacuutca NUK_MORKER52;
24
devieeeee___ static inline void common
args)
arga. hews.coooon.atcc. blockidx.z; azga head blockidx.y) arga, blockidx.
4 x; arge.nun_iters Lag
arga.globale :oo -1;
27
struct producer
29
devica_ static inline void setup arcecioor_stttp_aggcllpaop angs)
30
werpropppprruurr_eessssstl
device statíc inline void load iproducer_ ood_arga<layooe args)
1[(warpgroup: iwappiid tma::expect iargs. args. ingutt)
34
tma::lead__eynccarrgg + ater, 0), args. inpptt_arrrttttoeoooren
args.globai.. 58,
batch, cagmocommon. head,
36
tma::oad_esynetrrggggc largss.itts.....t....... arga. globals. ates. arga.tnputs_arriven
largs.ccmmon, common, ccmmon.batchy batch, args, common, wowmon.mmeeeee head,
37
LE(laneid() arrive (azge, bazge.iae wazggeinputee.e..meoeeo [
39
41
atruct consumer
42
device_ statis inline void satup htapsumme__otupp.rrrcaaaopo args)
warpgroup: sconssuerr stera<NMM MORKERS>
44
LEl commeeg.commmn... common. commen.ee...t......... HDU WORKERS
:oaa::))):aaaat::oo__eee:::::::: azgs.
2lobals.d.roxe) rows)) out bounds?
45
warpgroup: load jargs. scratch.g...g www.ggrrroopooprroouo....
arga.globals.U,
(args.common. batch, argg..oomon,.eeee wwrgs.coommnn MORKKRR
porpproupttggouppli. Ut);
47
zero wwwgs.stete.c. stsge.
(arga. norm_vec)s
49
neg_infty (args, state..mx__eee..eeeeeeeeeesseeee
paspgroup::syn: www.pgoopp::
51
device statio iniine void compute (consumer comput argsslayouut args)
53
conalexpr float TEXFERATURS_GCALE (D -- :28) 0.009333311111211 (18.33311T.1111122500000033 44269504089 0.1251
4 $2.662595040889
i
4 http:a/cttat.nnoove____eee_ecce args stage...tte................... .aaeeemax_vveee TEMPERATURE SCALE)
56
warpgroup: ioma :mma_async_woittg // softmax
39
xow_max wargs. www.te.mx_veee..eeeeeeeee. args.stats.att
// accumulete
onto the mex_vec
60
mul (args. saags.sttate www.:ettttteeeeeeee.eoee vex___ecsscaaed stt.gp..sttttt.....o.
stgs.state.rtg saa.tte.aax.eee COMPERATURE_SCALE
args.stato.att_bbottt TEMPERATURE, SCALE)
62.
sub row targa.araatt atate.att aat__blockk trgs.state.att_b attpbtactkkkk ata.stttee.a. max_vee_sea ed); exp2 (trgs sub (args. mex_vee__ssst__celeeed. vex_vee__esled)
args.state.de btt_kkock/
state...a..vee_lsstt stas.
26

---

Published as a conference paper at ICLR 2025
C.4
ROTARY POSITIONAL ENCODINGS
This section shows the rotary kerel for head dimension 128, written in the TK abstractions.
Load (atgaarattte.coo. args aogs.q.a www.glloolloooooo cos, 1dx1;
device statis void
azga)
rt_61< beaddim>
xt_f1k16, headdim/2> xi, x2, templ, temp2;
(xad(x, segs.input.ed xxxpt..xxxxraddtlllll
(f(lnned() == 0) arrive inputasstpaattaaiinn ())
for (int i 0z i headdim/32; 1+4)
#pragma unroll fortint x1.riles[0] [i] da)ad
1 4; 1-t)
12
x.tiles [0; [4] da]a
x2.tilea(0] (3). x.tilee [1+hoaddim/32] destal (i];
14
mul (templ, xi, args.state.cos) mul (tamp2, x2, arga.atate.coa);
17
(x2, x2, mul (xc,, x1, arga.state.sin),
19
(x2, x2, arga.state.sin))
21
add (Lempl, Lengl, x21;
22
add (temp2, temp2, x11:
for (int 0; headdim/32; i++)
24
#pragma unroll fortint x.tiles(0] dd]aaa] x.tiles [1=headdin/32 -aattitmmnaaaaaaaaa_a cata. 1]
< 1-+)
36
temp1, till [0] [4] [d]a (3]
[1] data [1];
28
29
(xarppt.((aarpd())) synowarp
Ir (1lmmiid
acrive (args..ess......ussss..... outpupts
33
device static void Sinisa www.soooonnnoiiiiieee frggeiiiirrreeeayylll azgs) (laneid) -- 0) arrive (args. finish_finishedl) // nothing to do here
35
38
11
Figure 18: A templated rotary kernel for head dim. 128 that outperforms popular Triton baselines.
27

---

Published as a conference paper at ICLR 2025
using namespace kittens; using namespace kittens: :proocttpp: using namespace kittens: :ppotot templatesint _headdim, static constexpr int headdim _headdim, using seg_tile using seg_global gl<bf16, using rope_global gisbf16, 1, 1, -I, headdim/22; struct globals seg_global
_arpss scruct sotary_layout
_Napps?
st_bfsle, headdims;
-1, headd seg_tile>;
rope_globel sin, cos;
int batches; // many batches block, for sizing grid
13 };
input_ ook
seq_tile [[arpa]];
15
struct blocut_bbock
seqliile o[warpa]; 1;
struct producer_stare
act! vettvveeee ve__aps
struct conaumer_state 11
rt_[1<16, heacdim/2> sit,
1: // Long-resident tiles
templatecint _headdim> static constexpr int me__dddm_ned, NJX NOX_ OOTPUT_PIPR__TTA 3TAGE_8 "NPUT_FIPE S1PE_STAGBS-3;
comply_template
20
C NUK [XX_BLOCK=1,
21
using layout sootary_syout-teecdie devivis statie intine sold common_secup (comnon_ wttpsseeetrppplaapoosrr args)
CTML_CONSUMRR_AA KARPS>;
23
iilargs. (islergg.gicbals.bttee.... arçs.num min blocals. crga.globalaxx.doe batches+heads hardled by black
24
(intl (args. www.ggggbbss....oooooo..
25
arge.nus
27
struct producer
29
device static void setup (ttuppppodicessstttpp.cca ppreducer_regisseer.
azgs)
30
args.atste. vvrps (inn(int)))) WOX_CONSCMER MARPS,
32
(int) (args. www.l...b.a..oooooooooo rows///1x..wo11 blookTdx.X*
33
devico static porrchaarr_oodd_rrrppoeooe
argsl
35
(w (warpid args. iterid)
36
kittens: cocord idx hho.kktt....agg.......aaaaaaaat. depth,
37
arga Stga.ilerbarga.. wlw.ertaagggggbbllee..........ee wobball.x.dee... depth, blockicx.x-fUM_COMB
39
0 1:
tmastexpact_byteee byteeyyeeeee (arga, torrs. inpptsaattttattaaaae 51 zeze
4 active_warps);
41
for lint i = 0; lma::1oad_asyne
i++) [
(xp]t.x(:)), args.globala.K, (idx.b, ddx.d, idx, rti, idx.cl, + asgs.inputs_arrivved)
44
(lanaid() == 01 arriva (wwgga.innnttaarre...r.o..a inputa 3); synowarp
47
device static void (produceerr _tr___ooeeeeaeelaaouule args) 1f 1f(werpgroup: warppid() =- args.
49
args.iteråd] iterrd)
kistens::coars
b depth,
arga.itertarga global. x...lxxx.ee depth,
53
0 1:
for aeotsstore__uanccarggggll...o.
atgs warpr warps; i11)
55
args.output.off1) (idx.b, idx. d, idx. r+1, idx. c));
57
tmaristore_aant aaw_aaaaaee_d_eeaeeee read_waatdd
if(laneid() (erive (iniiuttt__iiihedd..__ finished, 4)
synswarp
60
62
};
struct consumer
64
devieee_______ static woid setup (wonaurerreeccooc cewerrauer__setuppe wwwputppargsllayurrr.... args] wew MCOBCCCRR_MARES/1> (1
kittens: idx blockidx. www.UL_.LSSUEEE KARES warpid(l, 0
67
load(args. state. sigs.globals.sin idx1; // could be better coalesced but doing just
once
28

---

Published as a conference paper at ICLR 2025
D LIBRARY IMPLEMENTATION DETAILS
This section provides additional implementation details for THUNDERKITTENS
D.1 TILE DATA STRUCTURES
The core primitive in THUNDER KITTENS is the tile data structure as introduced in section 3.1 Tiles exist at the shared memory and register memory levels of the GPU hierarchy, and are created with multiples of 16 x 16 in dimension.
Precision THUNDERKITTENS extensible across data types: FP32, FP16, BF16, and FP8. De- signing unified tile data structure that seamlessly supports different data types is challenging for two reasons:
1. Each data type requires using different memory layouts both at the shared and register memory levels, order to use specialized hardware instructions like tensor cores.
Each data type uses different amount of space, meaning that the 16 x 16 tile that the user sees could contain fixed number of bits, fixed number of elements, some other option. Ideally, can store elements in fully packed formats (e.g., o 6 2 for b 1 6, e 4m338xx for e e4m FP8).
We let the users think in terms of the number of elements per tile. When the user defines 16 x 32 tile for instance, we store this as 16 8 packed elements of e 4m3_8x4 or 16 x 16 packed elements of bf1 6.2 2 in registers. In the library, hide this complexity from the user. Taking care of differences across data types at tile data structure level,
define and operate underlying tile width for the tiles to
then use the exact same library functions (e.g., mma, exp, cumsum um) across tiles of different data types, preserving the simplicity of the library.
Padding Some AI workloads require shapes that not multiples of 16. THUNDERKITTENS provides mechanisms to support these workloads, too, without compromising performance on hardware-friendly workloads.
1. Loads & Stores. BOUNNERRKITERSE loads and stores take an optional template arguments for whether assume tensors are multiples of 16. If not, each load store preceded by check to ensure that it is in-bounds. Out-of-bound loads are filled with zeros; out-of-bound stores are not performed. For safety, these checks are enabled by default; however, they are not free, and they can also be disabled by setting the appropriate template flag. This abides by TK's '''s philosophy of "extensible, but with good defaults." " For TMA loads and stores, we use the built-in hardware padding features (that is, oft-of-bounds accesses are automatically filled in with zeros).
2. Fills & Masks. In addition preventing illegal memory accesses, one often needs to alter data within a tile to prevent accidental computations. THUNDERKITTENS provides functionality for this, too, in the form of six functions: top fill, bottom fill, left fill, right fill, triu, and tril, which respectively fill the top, bottom, left, right, upper triangle, or lower triangle of tiles. our experience, these functions have proven sufficient for all kernels considered.
D.2 SHARED MEMORY LAYOUTS
To illustrate some of the choices available in shared memory layouts, this appendix outlines six different shared memory layouts for GPU tiles: naive row-major layout, padded layout, simple swizzled layout, and three more specialized swizzled layouts. We are particularly interested in which memory banks (numbered from 00 to 31) store each element of the tile; for each layout, we color and label the element of the tile accordingly. We illustrate all layouts using 32 x 64 16-bit tile.
D.3 NAIVE LAYOUT
row-major layout, illustrated in figure 19 is among the simplest layouts. It has the benefit of accurately reflecting tensor layouts HBM. Furthermore, for access patterns that access row-wise, it has no bank conflicts. But when loading or storing tensor core register layouts, it suffers 8-way bank conflicts, and thus extremely slow.
#asynchronous wa www.:ooonuuuwwwaagrrorgoo
29

---

Published as a conference paper at ICLR 2025
04
50 20 24 04
20 00 2400
20 04
29 20000004000
20 0000400000 04.0000000000000000
29 23
20 04 04
20:44 04:00
20:04 04
20:24 04
20:24
20:24
20:04
20:34
20:34
20:2
20024
16
20:24
16
20 20:04
16
20
20:04
10 10
20004 04000
1S
16
10
20
16
i0
20 20 04 20
16
20 20
20
10 16
00 00
20
16 I 10
20 20 00
20
16 IT in
24
20
Le 16
16 IT Sa
20 04
20 24
10
24 16
20 04
20 20
10
1G 10 10
2004
00
5035
14 36 36 AT 58 18:039
Figure 19; Row-major shared memory layout.
brls+ naive_ Layout (bt16 edata, int F, int c) I seturn 4datais colunna c]i
D.3.1 PADDED LAYOUT
D6 CRICeIO CO 1031
10 10
16
10
10 30 1P
16
18
14
10 3Y 16
20
28
20 20
37 26
20
26:6
20
20
21
24
20
20:0020
10 16 Te
28 20
1000 COLOCCO
DE 6O DO
10
To 1016
DO OO DO OE D6 COCOCDR
CO 10 BP 16 TO 10 10000000
GOLGOO BY CI 11 1O 10
DICO TO
HO 16 BY
TH 16 16 BY BT 16 TO TO
2E 20
20 SO
DO DE De Do to
3T 29 D0 20 PO
202020
2O 2O 20 DE D2 DO
be 20 12 bolo
24
20 20 20
RO DO D2 DO DO DO DO
200 12 1eleea bobrrar bobo
S0 SU DO QU WELIEO IO DO 1030 1910000 1o bo belaai
DOCEIO DN WEW
TO IF 191011 TO 10 BY 10 NO So
OF ce 50 10 12 10 20
20 NO 2O 22
20 b0
30 18
20 20201000
22 20 20 ST 2010
14
29 16 20 3Y 21 26
20
20.0 SO SO
te weleerr
2000002
20 BT
28
20:0
SO
to to
belrrar
26
50 50
50 00
De
10001: 10 La 10 bolee be Seihaiin
10
21 10000
pece
100010000
22 Le Le
20 pebrrari he 20 tobooolon
10
10
14 10
11 12 1e
32 18
20 37 28:30 10
29
be
20:20/0
20 20
20 52
24
bathe an peleeaal be bnllar
28-28
20/20/2
o0 colod col tei be pel eeleei an pal talaaaa be be an belaeii 20 be
cologioe
peice or holoooii 3203203 bellei arrar ho bolbo for beaeaain bes
Figure 20: Padded shared memory layout.
A common solution to these bank conflicts is to "pad" each row by one memory bank, thereby in- troducing offset to shift consecutive elements of a column into different memory banks. This eliminates bank conflicts, but creates misaligned addresses which interferes with fast instructions that require aligned addresses. For example, it wouldn't be possible to use TMA to store the sec- ond row of this layout due to it only having a 16-byte alignment, whereas TMA requires 128-byte alignments.
30

---

Published as a conference paper at ICLR 2025
brl5x padded_ layout (b216 vdata, int int. c) #data/r (columeet))
D.3.2
NAIVE SWIZZLED LAYOUT
00.0000
20
14
10
12
29
20:00000
10
20:0
16
28
25
2T
2424
27
10
20
24
24
2
16
24
2T
21
22
34
2424425
22
19
30
21221
10
Figure 21: Naive swizzled shared memory layout.
A third option is to "swizzle" the memory, in which progressive rows reshuffled alter their banking. This layout accomplishes this by xor 'ing the index with the row, which reduces bank conflicts. However, this layout lacks hardware support for HGMMA and UTMA instructions, which particularly important on H100 GPUs for achieving high performance. Additionally, the granularity of the swizzling must be large enough to totally prevent bank conflicts when loading into registers. We illustrate simple, naive swizzling pattern here, which used to be recommended for preventing bank conflicts before the advent of tensor cores:
S4 swissled leyout
T. vintsd_t adde (aint columns elr return 1ec1 (adde (t << 21)7
31

---

Published as a conference paper at ICLR 2025
D.3.3 32 BYTE SWIZZLING
06:00
06:00
10
06:0
30
0004
060000
Co
160606
50
16 10 se to
Co
22
Figure 22: 32 byte swizzled shared memory layout.
32 byte swizzling is the first of family of layouts (of which we will examine three), where instead of swizzling the index with the row, the memory address is instead swizzled directly with itself. This layout is defined by the following C code:
br15x Satezied__aayutt__e (DC16 edata, int lnt e) 64_t. addr juint 64 elz
oeluen (bC16.) (adds (ffadde (32.81) 7) 4b12
This layout here suffers from 4- bank confficts, but is valid for all tiles whose width is multiple of 16. However, importantly, it has (as do its siblings below) hardware support from both HGMMA and UTMA instructions.
32

---

Published as a conference paper at ICLR 2025
D.3.4 64 BYTE SWIZZLING
00000000
2031 19252-34
2020202121
20202010
26
16 be
18
160000
1031 18111111111111
20
06
22
29 1900/000011
2020
20 20203030
D00000
202661222 BT 20 10 TOP SEEE
2000
26
20:0
222222
12232
25 2020 21222222 252022200
20 10 10202 2022
15
20 PE PEEEEEE
26222222
20 20202030
2520 202000 21 2020 prpiettee
1676611 AT 101000
20 PO CEEE
20000
10:0212 22 P22220
19:00
12332
Pactc 20 20 20
20
20:0 20
20
15
De SY 20 2E PEEEEEE 2
25222222
22222222 20 202020 25
24 2o
27 2000
30
10
21 20 19
00:00:0
20
20:0
12:213
2020 20 20 20:0
27
04.00
2eleera ar 10 1019
2000
00
12:11
20
20:0020
30
04.00
20
10 10
55.05
10:02
59:0
12:13
18
15:00
30 30
20
16
05.00
30:0020
20
12:22
242422
29
13
15
100
04006 96 om
20 20:2029 30000
Figure 23: 64 byte swizzled shared memory layout.
64 byte swizzling is a layout similar to 32 byte swizzling with a more aggressive pattern
bifl Sx exiienled Layout, 64 66F16 edata, int T, fin 6) 64 adde 14 recurn Ibcl (addr o (riader 166-811 55 55 4)12
cul clz
64 byte swizzling suffers from just 2-way bank conflicts, but is only valid for tiles whose width is multiple of 32 (for half half-preeision types, or 16 for full-precision).
D.3.5 128 BYTE SWIZZLING.
202626
2010000000
18
202020
20:200
1030011
20:200
13303010
16
19
20 PERRIR 1EAVIE
0O 1031 1F 1938
10
18301010 20
2020
20
20:021
20
20
202000000
20101010 20
20
2020000
20
2121 10
202020 20
20
DO OO DO
DO 10100111
10
80
20:00
DE
DO 00 RE 19 19301416
19 10
28
MU
TU 103010
NO 103011 1239
1636 MI PO De DY
1OEMEMITO TO
DN IN BOT
1010
10 EP 19
10 10101010
30
1999 2020
20
161600 00000010
20:0 PERUE 10 1010 20
18
06 D6
20
20
20
BL 2010 20100000
20
D6
20
SO 202000 2Y 20:0 2202009
20:20
202020
20
20
DO DOROONII DO OYENIN DO 2010 123220
202020
26 2000
DE
DO 0O RE 19 19300010 08 100000000
30
101010
10:011 1738
1630 ME Dal0 Debo
2022
20 20 POR
16
DE 1010011 HI IN DN DO 1010 by
celcelor DROEMI 10T1 32
20
10101010 2030 2920 25 20000 20200200
12010 16 10100100
20:031
10
10
12
DE 00
202420
20
20,030
2010000
1010011
De DO
20.00 DE ba
20
3L 20202000 29 20,0
20
1910/3020
2022 20202222 29 29 20,000300 31
02 De DA D4 De DrorDeDe
00 201011 12322291
2522
1020
2004 be
26,266222220
27222022
30
20 10100030
10:001
19:030
16 0000030
De be 3o b0 los
Figure 24: 128 byte swizzled shared memory layout.
128 byte swizzling is a further extension of its kin:
33

---

Published as a conference paper at ICLR 2025
brl5x wniszled Layoeale___ 1288 (br16 rdata, int int
dlz oecuen (bC16x1 (edde (fjadde (128x8)) 71 4137
addr juint teseeall
Finally, 128 byte swizzling has no bank conflicts, but only valid for half-precision tiles whose width is multiple of 64.
D.3.6 THUNDERKITHN
APPROACH
After substantial evaluation of these layouts, we concluded that the three final layouts were the three most important, because HGMMA and UTMA instructions are critical high performance, and further- more that they are good enough to yield high performance across many kernels. Correspondingly, depending on the width of the tile compile time we select the highest level of swizzling possible to minimize bank conflicts.
34


Reviews:
[{"review_id": "iOBcibUeJ4", "reviewer": "Reviewer_kgD5", "summary": "The paper proposes a new programming library for implementing efficient CUDA kernels. tThe paper contains the three ideas at three different levels for CUDA kernel implementation: (1) At warp-level, the author proposes to organize the tile as the multiple of 16; (2) at thread-block level, the author devises template libraries to overlap between different asynchronous warps; (3) at grid-level, the author proposes methods for managing kernel launch kernel launching overheads. As a result, the proposed library can achieve a performance on par with the existing state-of-the-art implementations.", "questions": "Thanks for submitting the excellent paper to ICLR. While in general I enjoyed reading the paper, I have a few thoughts on the extension of the paper. Specifically, this paper proposes a new CUDA abstraction that allows users to write new kernels. However, it seems that it is built on top of the fact that all the dimensions should be a multiple of 16. This could be problematic in the context of dynamic shapes where the dimension does not divide 16. Could you please elaborate on how could the proposed technique be extended to such cases?\n\nBesides, the paper uses auto-tuning for further adjust the hyperparameters for a better performance. Could you elaborate how much the tunning overhead is?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper proposes methods at different levels that simplify CUDA kernel implementations\n* The paper can achieve a similar performance compared to the state-of-the-art implementation", "weaknesses": "* The paper has not discussed the tunning overhead with the proposed techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rx8VtsHol3", "reviewer": "Reviewer_nd8y", "summary": "This paper proposes a framework to facilitate easy writing of efficient CUDA kernels. The authors leverage the asynchronous compute capabilities of the Hopper series GPUs by following a producer-consumer paradigm, to efficiently overlap different kernel operations. Additionally, the authors investigate the impact of various memory ordering strategies, demonstrating that relatively simple strided patterns offer the best tradeoffs. Lastly, the authors demonstrate performance that is comparable to or exceeds existing methods, including Triton.\n\nOverall, the work provides a significant contribution to improving computational efficiency for common operations, though the application appears limited in scope. Additionally, minor technical and structural errors impact readability. These issues could be addressed in a revision, at which point I would be inclined to raise my score.", "questions": "1)\tIs your framework limited to the Hopper series? Can it be applied to A100s, or other GPUs such as the A40/L40?\n2)\tYou focus on the 16x16 register block level, but how can your framework be extended to smaller blocks, such as with GEMV, sparse operations, and masked operations (e.g. non-power-of-two dimensions and strided masking, such as in Natten).\n3)\tThroughout the paper, you focus on BF16 precision (with the exception of softmax); have you considered other data types, such as integer types or floating-point formats like FP8?\n4)\tHow could your framework be extended to handle multi-GPU operations, such as Fully Sharded Data Parallel (FSDP) for split operations? This seems like a natural extension of the producer-consumer model.\n5)\tYou compare yourself against Triton, which also supports AMD GPUs. Can you address this as a potential tradeoff in the paper? Alternatively, if your framework can be trivially extended to ROCm, this should be included in the paper with a demonstration, otherwise it represents a tradeoff between efficiency and portability.\n6)\tYour cost model in Section 2.2 is effectively a Roofline model; could you contextualize this in the existing literature? The results in Table 3 are expected, as reordering increases the arithmetic intensity (FLOPs/Byte) of the inner loops.\n7)\tThroughout the paper, the emphasis on industry versus academic adoption (including the use by undergraduates) feels extraneous and detracts from the main narrative. The paper’s contributions should stand on their own without reliance on external endorsements or applications.\n8)\tFigures 2 and 5 present a simplified sketch for softmax, whereas the true implementation is significantly more complex, potentially leading to a misleading comparison with PyTorch.  Furthermore, Figure 2 led me to question why you are using C at all for the API, when the listing could easily have been captured by a python trace (e.g. Triton). This design choice is only clarified upon reviewing the implementation details provided in the appendix and supplementary material.\n\nTo build on these questions, the feedback below addresses specific technical details and aims to enhance overall clarity. While this paper presents a strong contribution toward improving kernel efficiency, addressing these points will better showcase the authors’ contributions.\n\nMinor Technical Errors:\n\n-\t044: The H100 datasheet shows a 7.4x ratio between TCs and ALUs, not 16x. Additionally, my understanding is that the TCs necessarily require bubbles as the Register path cannot keep up with the TC I/O for full throughput. \n-\t136: This should be \"can load\" or \"may load\" instead of \"loads.\" In general, a kernel does not necessarily need to load data from memory. Kernels can rely solely on arguments (loaded into registers at startup) to generate new data. For example, a kernel might generate a pseudo-random noise tensor without accessing memory.\n-\t148: The 32 threads must be within the same quadrant, where “consecutive” or “adjacent” would be more appropriate than “nearby”.\n-\t150: In Ampere, a warp cannot simultaneously occupy different functional units, though separate warps can. For accuracy, please verify this claim against the Hopper documentation or micro-benchmarking paper, otherwise consider omitting if verification is unavailable.\n-\t167: Excess registers spill over into Global Memory, not L1. They can appear in L1 due to the memory hierarchy, but this is at the discretion of the hardware cache manager.\n-\t171: Multiple thread blocks can only schedule on the same SM if there is sufficient space (e.g. SMem), otherwise they would clobber each other.\n-\t173: This statement should be more precise to mention “all thread blocks” and that the L2 is hardware managed, making it distinct from the software managed SMem.\n-\t179: The tail-effect cost mentioned only applies to singular kernels. Ideally the GPU should have multiple kernels in flight, which can run concurrently.\n-\tIt would also be relevant to mention that kernels which contain too many instructions can cause slowdown as they will incur ICache misses.\n\nPresentation Issues:\n\n-\tThe abstract should be revised for clarity, with suggested improvements like “creates a”, “suggest that”, and “resembling PyTorch.”\n-\tThe paper could benefit from clarity revisions in several sections, where phrasing and word choice could make technical details easier to follow. Lines: 073, 170, 178, 205, 278, 299, 301, 328, 370, 397\n-\t325: You should not use \"[1]\" and \"[2]\" to enumerate concepts as they are easily confused with reference indicators. \n-\tTable 2 and Table 3 should probably be Figures like Figure 6. It is also unclear why these stop at 4 stages, K=1024, and what K is. (MxN)x(NxK)? \n-\tFigure 7 and 8 should use subfig captions rather than plot titles. If parameters are common among subfigures, then they should be stated in the figure caption, otherwise in the subfig caption. The fontsize for the axis and labels is too small. Finally, the batch size does not match with the titles and caption.\n-\tThe table in Section 4.2 is missing a caption and column (TK is listed twice).\n-\tThe reference links are broken in Appendix B.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors demonstrate significant improvements to computational efficiency within a clearly defined framework that appears relatively straightforward to adapt. Their framework also provides functionality for more complex resource management, which is often challenging to manage directly in CUDA.  Additionally, the authors demonstrate the impact of varying hyperparameters for several key kernel operations, most of which match or exceed standard baselines. Lastly, the results show a surprising contrast with Triton implementations, positioning their approach within the CUDA domain while achieving a similar level of complexity to Triton.", "weaknesses": "-\tThe application appears limited in scope, which should be explicitly addressed. For example, is the framework limited to Hopper GPUs and above? And the focus on 16x16 register blocks may limit extensibility to other common cases such as GEMV and sparse computations.\n-\tThe paper contains many issues with presentation, including caption errors, grammatical and awkward wording, and typos, all of which impair readability. \n-\tThe paper overlooks relevant computer architecture literature regarding performance modeling, specifically in the context of balancing compute and memory (e.g. roofline analysis). Many of the findings presented in the paper are expected from the existing literature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pWJTtQTpEL", "reviewer": "Reviewer_ANEJ", "summary": "This paper presents ThunderKittens (TK), a C++ embedded library for writing high-performance CUDA kernels for NVIDIA GPUs. It introduces warp-, thread-block-, and grid-level abstractions to facilitate mapping of kernels to the GPU hierarchy. Experimental results indicate that TK can outperform strong industrial baselines, achieving superior performance for GEMM and attention kernels.", "questions": "1. What are the fundamental challenges preventing CUTLASS from avoiding bank conflicts? Could it be that the FlashAttention3 kernel simply did not select the optimal layout?\n2. CUTLASS has implemented both ping-pong and cooperative kernel variants for GEMM, with varying performance across different scenarios. How does TK support ping-pong and cooperative kernels, and could you include a comparison with CUTLASS in Figure 7’s GEMM kernel results?\n3. TK appears designed specifically for the Hopper architecture with asynchronous features. Is it also compatible with Ampere or other GPU generations? How does TK’s performance on an A100 compare to Triton?\n4. Following Q3, if Blackwell GPUs were released, would TK’s abstractions remain applicable? How do you plan to ensure extensibility across GPU generations?\n5. What's the usage of the cost model in Section 2.2? This formula is highly simplified and does not guide any optimization or automatic search later.\n6. Section 3.1 discusses various layouts — do users need to manually manage data organization and specify layouts in TK?\n7. Figure 5 is just some wrappers of mbarriers. Any insights here?\n8. Can TK effectively handle quantized kernels, where data layout is crucial for efficient transfers from TMA and WGMMA computation? How does it perform on FP8 GEMM and FlashAttention kernels?\n9. What is TK's performance on causal attention kernels?\n10. Please provide detailed experimental configurations in the Appendix. For example, which versions of PyTorch and Triton were used? Was `torch.compile` employed to optimize those network layers? For cuBLAS, was the latest [cuBLASLt](https://developer.nvidia.com/blog/introducing-grouped-gemm-apis-in-cublas-and-more-performance-updates/) autotuning enabled? Since PyTorch also uses Triton as a backend, what distinguishes the two baselines in Figure 8?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The TK library provides a useful abstraction for writing high-performance asynchronous kernels on GPU.\n2. The presentation is clear and accessible, especially the introductory sections on GPU architecture, which provide a helpful overview for ML researchers who may lack in-depth experience with GPU programming.\n3. The experimental results are compelling, showing performance on par or better than highly optimized kernels, such as FlashAttention3. The paper also demonstrates significant speedups across different kernel types compared to state-of-the-art frameworks like Triton and PyTorch.", "weaknesses": "1. The TK library is still too low-level with too many details, which requires users to manage synchronization carefully and does not simplify the programming burden.\n2. The novelty and advantages of TK over CUTLASS are unclear. Many functionalities seem achievable with CUTLASS as well. The authors mention that TK addresses bank conflicts, but the evidence presented is minimal. There appear to be no inherent limitations in CUTLASS that would prevent it from avoiding bank conflicts.\n3. Similarly, the benefits of TK over Triton are not well established. Triton, embedded in Python with a PyTorch-like API, may offer a more accessible interface. By contrast, TK, embedded in C++, still requires explicit handling of communication with mbarrier operations like expect and arrive. No user study or lines of code comparisons are provided to demonstrate that TK improves programmer productivity.\n4. Experimental results are good, but still missing comparisons in some important cases like quantized kernels and causal attention.\n5. The work reads more like a system paper, with limited ML-focused insights, raising questions about its fit for ICLR.\n\nMinor:\n- P4: \"Since the frameworks are not C++ embedded, it can be challenging to use specialized hardware instructions\" This statement is inaccurate; TVM provides mechanisms to incorporate low-level TensorCore instructions, and Triton also has [inline](https://triton-lang.org/main/python-api/triton.language.html#inline-assembly) operation to include PTX code.\n- Section 2 does not discuss the Tensor Memory Accelerator (TMA) on Hopper, which is essential for asynchronous optimizations mentioned in Contribution 2.\n- Appendix B labels appear broken (??).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R1mEaXgIHt", "reviewer": "Reviewer_5GXe", "summary": "The paper introduces THUNDERKITTENS (TK), a framework that simplifies writing AI kernels for GPUs while still allowing for high performance. Using a few key abstractions, TK provides tools for developers to create efficient kernels without deep expertise in GPU programming. Through benchmarking, the authors show that TK performs on par with or better than other leading frameworks like CuBLAS and FlashAttention-3 for various AI tasks. TK’s accessible design, inspired by PyTorch and NumPy, aims to make high-performance kernel development more straightforward and accessible to a wider audience.", "questions": "Could the authors elaborate on the potential for cross-platform compatibility? Given the focus on NVIDIA’s H100 GPUs, it would be helpful to understand whether TK’s abstractions could be adapted to other GPU architectures, like AMD or Apple, and what challenges might arise.\n\nThe paper demonstrates TK’s strong performance on medium-sized data blocks, but could the authors provide more insights into how well TK scales with very large datasets? Are there specific limitations to consider for applications requiring high parallelization or extensive data handling?\n\nCould the authors expand on their design choice to limit TK to a few key abstractions? Are there specific reasons why additional templates or adaptive features were not incorporated, and would doing so have risked undermining the framework’s simplicity?\n\nIn scenarios with high memory demands, how does TK manage the balance between memory overhead and computational efficiency? Further detail on this balance could clarify TK’s suitability for applications with varied memory and compute requirements.\n\nLastly, could the authors clarify TK’s debugging process, especially for users who may not be familiar with GPU optimization? Since GPU kernel errors can be challenging to diagnose, any insights into how TK might support error handling and debugging would be valuable for potential adopters.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper offers a fresh and practical approach to GPU kernel programming, using only a handful of essential abstractions to make high-performance kernel writing accessible to a wider range of developers. This simplicity-oriented approach can reduce the complexity typically associated with GPU development, which could be particularly valuable for those without extensive CUDA experience. In terms of performance, THUNDERKITTENS shows impressive results, even surpassing established libraries like CuBLAS and FlashAttention-3 in several tasks, especially in backward pass operations for attention mechanisms and linear attention. The results strongly suggest that TK’s design strikes a good balance between simplicity and performance optimization. Furthermore, by aligning its design with PyTorch and NumPy, TK makes it easier for non-specialists to adopt, potentially expanding the accessibility of efficient GPU programming.", "weaknesses": "1- While the minimalistic design is a key strength, it may also limit TK’s flexibility for more specialized AI tasks that require tailored optimization strategies. As demands grow for handling complex and emerging AI workloads, the current set of abstractions could potentially fall short.\n\n2- The focus on NVIDIA’s H100 GPUs raises questions about how well TK can transfer to other platforms, such as AMD or Apple GPUs. Expanding on cross-platform compatibility would provide more clarity about TK’s broader usability.\n\n3- Though the paper demonstrates strong performance on medium-sized data, it is less clear how TK handles scalability with very large datasets or highly parallelized scenarios. Addressing its limitations in these settings could further support TK’s value in real-world applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)