ThunderKittens: Simple, Fast, and Adorable AI Kernels Article Swipe

PDF

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

YOU? · · 2024 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2410.20399

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

Related Topics

Computer Science

Neuroscience

Artificial Intelligence

Philosophy

Epistemology

Concepts

Simple (philosophy) Computer science Neuroscience Artificial intelligence Psychology Philosophy Epistemology

Metadata

Type: preprint
Language: en
Landing Page: http://arxiv.org/abs/2410.20399
PDF: https://arxiv.org/pdf/2410.20399
OA Status: green
Related Works: 10
OpenAlex ID: https://openalex.org/W4404313951

All OpenAlex metadata

Raw OpenAlex JSON

OpenAlex ID: https://openalex.org/W4404313951

Canonical identifier for this work in OpenAlex
DOI: https://doi.org/10.48550/arxiv.2410.20399

Digital Object Identifier
Title: ThunderKittens: Simple, Fast, and Adorable AI Kernels

Work title
Type: preprint

OpenAlex work type
Language: en

Primary language
Publication year: 2024

Year of publication
Publication date: 2024-10-27

Full publication date if available
Authors: Benjamin Spector, Simran Arora, Aaryan Singhal, Daniel Y. Fu, Christopher Ré

List of authors in order
Landing page: https://arxiv.org/abs/2410.20399

Publisher landing page
PDF URL: https://arxiv.org/pdf/2410.20399

Direct link to full text PDF
Open access: Yes

Whether a free full text is available
OA status: green

Open access status per OpenAlex
OA URL: https://arxiv.org/pdf/2410.20399

Direct OA link when available
Concepts: Simple (philosophy), Computer science, Neuroscience, Artificial intelligence, Psychology, Philosophy, Epistemology

Top concepts (fields/topics) attached by OpenAlex
Cited by: 0

Total citation count in OpenAlex
Related works (count): 10

Other works algorithmically related by OpenAlex

Full payload

id	https://openalex.org/W4404313951
doi	https://doi.org/10.48550/arxiv.2410.20399
ids.doi	https://doi.org/10.48550/arxiv.2410.20399
ids.openalex	https://openalex.org/W4404313951
fwci
type	preprint
title	ThunderKittens: Simple, Fast, and Adorable AI Kernels
biblio.issue
biblio.volume
biblio.last_page
biblio.first_page
topics[0].id	https://openalex.org/T13650
topics[0].field.id	https://openalex.org/fields/17
topics[0].field.display_name	Computer Science
topics[0].score	0.5666999816894531
topics[0].domain.id	https://openalex.org/domains/3
topics[0].domain.display_name	Physical Sciences
topics[0].subfield.id	https://openalex.org/subfields/1702
topics[0].subfield.display_name	Artificial Intelligence
topics[0].display_name	Computational Physics and Python Applications
is_xpac	False
apc_list
apc_paid
concepts[0].id	https://openalex.org/C2780586882
concepts[0].level	2
concepts[0].score	0.74149489402771
concepts[0].wikidata	https://www.wikidata.org/wiki/Q7520643
concepts[0].display_name	Simple (philosophy)
concepts[1].id	https://openalex.org/C41008148
concepts[1].level	0
concepts[1].score	0.4958103597164154
concepts[1].wikidata	https://www.wikidata.org/wiki/Q21198
concepts[1].display_name	Computer science
concepts[2].id	https://openalex.org/C169760540
concepts[2].level	1
concepts[2].score	0.42376890778541565
concepts[2].wikidata	https://www.wikidata.org/wiki/Q207011
concepts[2].display_name	Neuroscience
concepts[3].id	https://openalex.org/C154945302
concepts[3].level	1
concepts[3].score	0.36253494024276733
concepts[3].wikidata	https://www.wikidata.org/wiki/Q11660
concepts[3].display_name	Artificial intelligence
concepts[4].id	https://openalex.org/C15744967
concepts[4].level	0
concepts[4].score	0.16345903277397156
concepts[4].wikidata	https://www.wikidata.org/wiki/Q9418
concepts[4].display_name	Psychology
concepts[5].id	https://openalex.org/C138885662
concepts[5].level	0
concepts[5].score	0.0709221363067627
concepts[5].wikidata	https://www.wikidata.org/wiki/Q5891
concepts[5].display_name	Philosophy
concepts[6].id	https://openalex.org/C111472728
concepts[6].level	1
concepts[6].score	0.05614888668060303
concepts[6].wikidata	https://www.wikidata.org/wiki/Q9471
concepts[6].display_name	Epistemology
keywords[0].id	https://openalex.org/keywords/simple
keywords[0].score	0.74149489402771
keywords[0].display_name	Simple (philosophy)
keywords[1].id	https://openalex.org/keywords/computer-science
keywords[1].score	0.4958103597164154
keywords[1].display_name	Computer science
keywords[2].id	https://openalex.org/keywords/neuroscience
keywords[2].score	0.42376890778541565
keywords[2].display_name	Neuroscience
keywords[3].id	https://openalex.org/keywords/artificial-intelligence
keywords[3].score	0.36253494024276733
keywords[3].display_name	Artificial intelligence
keywords[4].id	https://openalex.org/keywords/psychology
keywords[4].score	0.16345903277397156
keywords[4].display_name	Psychology
keywords[5].id	https://openalex.org/keywords/philosophy
keywords[5].score	0.0709221363067627
keywords[5].display_name	Philosophy
keywords[6].id	https://openalex.org/keywords/epistemology
keywords[6].score	0.05614888668060303
keywords[6].display_name	Epistemology
language	en
locations[0].id	pmh:oai:arXiv.org:2410.20399
locations[0].is_oa	True
locations[0].source.id	https://openalex.org/S4306400194
locations[0].source.issn
locations[0].source.type	repository
locations[0].source.is_oa	True
locations[0].source.issn_l
locations[0].source.is_core	False
locations[0].source.is_in_doaj	False
locations[0].source.display_name	arXiv (Cornell University)
locations[0].source.host_organization	https://openalex.org/I205783295
locations[0].source.host_organization_name	Cornell University
locations[0].source.host_organization_lineage	https://openalex.org/I205783295
locations[0].license
locations[0].pdf_url	https://arxiv.org/pdf/2410.20399
locations[0].version	submittedVersion
locations[0].raw_type	text
locations[0].license_id
locations[0].is_accepted	False
locations[0].is_published	False
locations[0].raw_source_name
locations[0].landing_page_url	http://arxiv.org/abs/2410.20399
locations[1].id	doi:10.48550/arxiv.2410.20399
locations[1].is_oa	True
locations[1].source.id	https://openalex.org/S4306400194
locations[1].source.issn
locations[1].source.type	repository
locations[1].source.is_oa	True
locations[1].source.issn_l
locations[1].source.is_core	False
locations[1].source.is_in_doaj	False
locations[1].source.display_name	arXiv (Cornell University)
locations[1].source.host_organization	https://openalex.org/I205783295
locations[1].source.host_organization_name	Cornell University
locations[1].source.host_organization_lineage	https://openalex.org/I205783295
locations[1].license	public-domain
locations[1].pdf_url
locations[1].version
locations[1].raw_type	article
locations[1].license_id	https://openalex.org/licenses/public-domain
locations[1].is_accepted	False
locations[1].is_published
locations[1].raw_source_name
locations[1].landing_page_url	https://doi.org/10.48550/arxiv.2410.20399
indexed_in	arxiv, datacite
authorships[0].author.id	https://openalex.org/A5004675499
authorships[0].author.orcid	https://orcid.org/0000-0003-0468-5986
authorships[0].author.display_name	Benjamin Spector
authorships[0].author_position	first
authorships[0].raw_author_name	Spector, Benjamin F.
authorships[0].is_corresponding	False
authorships[1].author.id	https://openalex.org/A5075032133
authorships[1].author.orcid	https://orcid.org/0000-0002-2087-8043
authorships[1].author.display_name	Simran Arora
authorships[1].author_position	middle
authorships[1].raw_author_name	Arora, Simran
authorships[1].is_corresponding	False
authorships[2].author.id	https://openalex.org/A5104346387
authorships[2].author.orcid
authorships[2].author.display_name	Aaryan Singhal
authorships[2].author_position	middle
authorships[2].raw_author_name	Singhal, Aaryan
authorships[2].is_corresponding	False
authorships[3].author.id	https://openalex.org/A5032865467
authorships[3].author.orcid	https://orcid.org/0000-0002-2500-2577
authorships[3].author.display_name	Daniel Y. Fu
authorships[3].author_position	middle
authorships[3].raw_author_name	Fu, Daniel Y.
authorships[3].is_corresponding	False
authorships[4].author.id	https://openalex.org/A5103852640
authorships[4].author.orcid
authorships[4].author.display_name	Christopher Ré
authorships[4].author_position	last
authorships[4].raw_author_name	Ré, Christopher
authorships[4].is_corresponding	False
has_content.pdf	False
has_content.grobid_xml	False
is_paratext	False
open_access.is_oa	True
open_access.oa_url	https://arxiv.org/pdf/2410.20399
open_access.oa_status	green
open_access.any_repository_has_fulltext	False
created_date	2025-10-10T00:00:00
display_name	ThunderKittens: Simple, Fast, and Adorable AI Kernels
has_fulltext	False
is_retracted	False
updated_date	2025-11-06T06:51:31.235846
primary_topic.id	https://openalex.org/T13650
primary_topic.field.id	https://openalex.org/fields/17
primary_topic.field.display_name	Computer Science
primary_topic.score	0.5666999816894531
primary_topic.domain.id	https://openalex.org/domains/3
primary_topic.domain.display_name	Physical Sciences
primary_topic.subfield.id	https://openalex.org/subfields/1702
primary_topic.subfield.display_name	Artificial Intelligence
primary_topic.display_name	Computational Physics and Python Applications
related_works	https://openalex.org/W4391375266, https://openalex.org/W2899084033, https://openalex.org/W2748952813, https://openalex.org/W1585007175, https://openalex.org/W2382521049, https://openalex.org/W2390279801, https://openalex.org/W4391913857, https://openalex.org/W2144385241, https://openalex.org/W2358668433, https://openalex.org/W4396701345
cited_by_count	0
locations_count	2
best_oa_location.id	pmh:oai:arXiv.org:2410.20399
best_oa_location.is_oa	True
best_oa_location.source.id	https://openalex.org/S4306400194
best_oa_location.source.issn
best_oa_location.source.type	repository
best_oa_location.source.is_oa	True
best_oa_location.source.issn_l
best_oa_location.source.is_core	False
best_oa_location.source.is_in_doaj	False
best_oa_location.source.display_name	arXiv (Cornell University)
best_oa_location.source.host_organization	https://openalex.org/I205783295
best_oa_location.source.host_organization_name	Cornell University
best_oa_location.source.host_organization_lineage	https://openalex.org/I205783295
best_oa_location.license
best_oa_location.pdf_url	https://arxiv.org/pdf/2410.20399
best_oa_location.version	submittedVersion
best_oa_location.raw_type	text
best_oa_location.license_id
best_oa_location.is_accepted	False
best_oa_location.is_published	False
best_oa_location.raw_source_name
best_oa_location.landing_page_url	http://arxiv.org/abs/2410.20399
primary_location.id	pmh:oai:arXiv.org:2410.20399
primary_location.is_oa	True
primary_location.source.id	https://openalex.org/S4306400194
primary_location.source.issn
primary_location.source.type	repository
primary_location.source.is_oa	True
primary_location.source.issn_l
primary_location.source.is_core	False
primary_location.source.is_in_doaj	False
primary_location.source.display_name	arXiv (Cornell University)
primary_location.source.host_organization	https://openalex.org/I205783295
primary_location.source.host_organization_name	Cornell University
primary_location.source.host_organization_lineage	https://openalex.org/I205783295
primary_location.license
primary_location.pdf_url	https://arxiv.org/pdf/2410.20399
primary_location.version	submittedVersion
primary_location.raw_type	text
primary_location.license_id
primary_location.is_accepted	False
primary_location.is_published	False
primary_location.raw_source_name
primary_location.landing_page_url	http://arxiv.org/abs/2410.20399
publication_date	2024-10-27
publication_year	2024
referenced_works_count	0
abstract_inverted_index.a	11, 48, 62, 77, 129, 173
abstract_inverted_index.AI	4, 15, 82, 176
abstract_inverted_index.TK	162
abstract_inverted_index.We	73, 157, 178
abstract_inverted_index.as	111
abstract_inverted_index.at	103, 123, 140
abstract_inverted_index.by	163, 194
abstract_inverted_index.in	14
abstract_inverted_index.is	9
abstract_inverted_index.of	2, 41, 51, 65, 98, 161, 175
abstract_inverted_index.on	31, 183, 196, 200, 206
abstract_inverted_index.or	168
abstract_inverted_index.to	6, 24, 53, 87, 94, 146
abstract_inverted_index.we	46, 106, 127, 143
abstract_inverted_index.(1)	102
abstract_inverted_index.(2)	122
abstract_inverted_index.(3)	139
abstract_inverted_index.GPU	7, 100
abstract_inverted_index.Our	91
abstract_inverted_index.The	0, 37
abstract_inverted_index.and	89, 115, 138, 152, 154, 181, 185, 189, 204
abstract_inverted_index.can	68
abstract_inverted_index.for	79, 131, 172
abstract_inverted_index.key	66
abstract_inverted_index.map	93
abstract_inverted_index.our	58
abstract_inverted_index.the	71, 95, 99, 104, 124, 141, 149, 159, 191
abstract_inverted_index.use	88
abstract_inverted_index.GEMM	184
abstract_inverted_index.GPUs	42
abstract_inverted_index.data	113
abstract_inverted_index.easy	86
abstract_inverted_index.even	30
abstract_inverted_index.fail	23
abstract_inverted_index.help	147
abstract_inverted_index.hide	148
abstract_inverted_index.high	55
abstract_inverted_index.like	34
abstract_inverted_index.meet	25
abstract_inverted_index.need	47
abstract_inverted_index.over	120
abstract_inverted_index.show	158
abstract_inverted_index.that	45, 166
abstract_inverted_index.wide	49
abstract_inverted_index.work	59
abstract_inverted_index.(TK),	76
abstract_inverted_index.16x16	108
abstract_inverted_index.basic	112
abstract_inverted_index.block	150
abstract_inverted_index.match	167, 179
abstract_inverted_index.might	43
abstract_inverted_index.prior	170
abstract_inverted_index.range	174
abstract_inverted_index.small	63
abstract_inverted_index.space	202
abstract_inverted_index.state	201
abstract_inverted_index.their	26
abstract_inverted_index.three	96
abstract_inverted_index.tiles	110
abstract_inverted_index.value	160
abstract_inverted_index.while	84
abstract_inverted_index.CuBLAS	180
abstract_inverted_index.across	135
abstract_inverted_index.costs.	156
abstract_inverted_index.custom	21
abstract_inverted_index.launch	151
abstract_inverted_index.level,	126
abstract_inverted_index.levels	97
abstract_inverted_index.linear	35, 207
abstract_inverted_index.matrix	109
abstract_inverted_index.memory	155
abstract_inverted_index.number	64
abstract_inverted_index.tiles,	121
abstract_inverted_index.warps,	137
abstract_inverted_index.Despite	17
abstract_inverted_index.achieve	54
abstract_inverted_index.compute	118
abstract_inverted_index.diverse	38
abstract_inverted_index.kernels	22, 83, 165, 171
abstract_inverted_index.mapping	3
abstract_inverted_index.models,	203
abstract_inverted_index.present	74
abstract_inverted_index.provide	107, 128, 144
abstract_inverted_index.suggest	44
abstract_inverted_index.support	145
abstract_inverted_index.variety	50
abstract_inverted_index.whether	61
abstract_inverted_index.writing	80
abstract_inverted_index.However,	57
abstract_inverted_index.creating	10
abstract_inverted_index.critical	12
abstract_inverted_index.efforts,	19
abstract_inverted_index.explores	60
abstract_inverted_index.hardware	8, 39
abstract_inverted_index.parallel	117, 136
abstract_inverted_index.process.	72
abstract_inverted_index.simplify	70
abstract_inverted_index.template	130
abstract_inverted_index.$10-40\%$	195
abstract_inverted_index.$8\times$	199
abstract_inverted_index.attention	186, 197
abstract_inverted_index.baselines	193
abstract_inverted_index.challenge	1
abstract_inverted_index.framework	78
abstract_inverted_index.inference	187
abstract_inverted_index.maintain.	90
abstract_inverted_index.progress.	16
abstract_inverted_index.providing	164
abstract_inverted_index.remaining	85
abstract_inverted_index.strongest	192
abstract_inverted_index.$14\times$	205
abstract_inverted_index.attention.	36, 208
abstract_inverted_index.backwards,	198
abstract_inverted_index.bottleneck	13
abstract_inverted_index.hierarchy:	101
abstract_inverted_index.operations	33, 119, 134
abstract_inverted_index.outperform	169, 190
abstract_inverted_index.performant	81
abstract_inverted_index.structures	114
abstract_inverted_index.tear-down,	153
abstract_inverted_index.techniques	52
abstract_inverted_index.drastically	69
abstract_inverted_index.grid-level,	142
abstract_inverted_index.operations.	177
abstract_inverted_index.overlapping	132
abstract_inverted_index.performance	28, 188
abstract_inverted_index.substantial	18
abstract_inverted_index.theoretical	27
abstract_inverted_index.thresholds,	29
abstract_inverted_index.warp-level,	105
abstract_inverted_index.PyTorch-like	116
abstract_inverted_index.abstractions	67, 92
abstract_inverted_index.asynchronous	133
abstract_inverted_index.capabilities	40
abstract_inverted_index.hand-written	20
abstract_inverted_index.performance.	56
abstract_inverted_index.thread-block	125
abstract_inverted_index.architectures	5
abstract_inverted_index.ThunderKittens	75
abstract_inverted_index.FlashAttention-3	182
abstract_inverted_index.well-established	32
cited_by_percentile_year
countries_distinct_count	0
institutions_distinct_count	5
citation_normalized_percentile