Pith. sign in

REVIEW 17 cited by

ThunderKittens: Simple, Fast, and Adorable AI Kernels

Not yet reviewed by Pith; the record is open.

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2410.20399 v1 pith:TK47WXRH submitted 2024-10-27 cs.LG cs.AI

ThunderKittens: Simple, Fast, and Adorable AI Kernels

classification cs.LG cs.AI
keywords kernelsattentionoperationsperformanceabstractionshardwarelinearmatch
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved
0 comments
read the original abstract

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

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 17 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. TileSight: A First-Principles Tile-Centric Analytical GPU Performance Model from Cores to Clusters

    cs.DC 2026-07 conditional novelty 7.0

    A tile-centric analytical model predicts GPU kernel latency and cache behavior purely from microbenchmark-calibrated hardware rates, reaching about 12% GEMM error and 13% end-to-end LLM serving error across five GPU lines.

  2. TileLens: Efficiently Using Large-Granularity Memory Systems with Transparent Two-Dimensional Memory Layout

    cs.AR 2026-07 conditional novelty 7.0

    Tile-major memory layout plus TileLens SW/HW extensions and adaptive prefetching eliminate read amplification on kilobyte-granularity GPU memories, restoring near-HBM matmul speed for LLM kernels.

  3. Characterizing Real-World Bugs in Tile Programs for Automated Bug Detection

    cs.SE 2026-05 unverdicted novelty 7.0

    A systematic analysis of 301 real-world code generation bugs in tile programs, categorizing root causes, symptoms, input patterns, test oracles, and fix strategies from curated GitHub reports.

  4. Characterizing Real-World Bugs in Tile Programs for Automated Bug Detection

    cs.SE 2026-05 accept novelty 7.0

    A taxonomy of 301 real tile-program codegen bugs: 48.8% are type/operator bugs, 19.3% are memory bugs, and detection needs shape-aware fuzzing plus differential oracles.

  5. CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

    cs.LG 2026-05 unverdicted novelty 7.0

    CODA re-expresses non-attention Transformer computations as GEMM epilogues using composable primitives for scaling, reductions, and transformations.

  6. CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

    cs.LG 2026-05 unverdicted novelty 7.0

    CODA re-expresses most non-attention Transformer computations as GEMM-plus-epilogue programs using a constrained set of composable primitives to keep intermediate results on-chip and cut global memory traffic.

  7. Sparse Forcing: Native Trainable Sparse Attention for Real-time Autoregressive Diffusion Video Generation

    cs.CV 2026-04 unverdicted novelty 7.0

    Sparse Forcing adds a native trainable sparsity mechanism and PBSA kernel to autoregressive diffusion video models, yielding higher VBench scores and 1.1-1.27x speedups on 5s to 1min generations.

  8. Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

    cs.PL 2026-04 unverdicted novelty 7.0

    Nautilus auto-compiles math-like tensor descriptions into optimized GPU kernels, delivering up to 42% higher throughput than prior compilers on transformer models across NVIDIA GPUs.

  9. AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

    cs.LG 2026-06 unverdicted novelty 6.0

    A system that auto-generates statically verified megakernels for LLM forward passes, retargets across NVIDIA architectures from one source, matches reference outputs exactly, and self-improves via an agent loop.

  10. Veda: Scalable Video Diffusion via Distilled Sparse Attention

    cs.CV 2026-05 unverdicted novelty 6.0

    Veda formulates tile selection in video diffusion attention as a reconstruction problem from full attention maps, using statistics-aware and head-aware scoring to enable high sparsity with maintained quality and hardw...

  11. TLX: Hardware-Native, Evolvable MIMW GPU Compiler for Large-scale Production Environments

    cs.AR 2026-05 unverdicted novelty 6.0

    TLX is a Triton extension that exposes multi-warp, asynchronous, and cluster-level controls for modern GPUs, delivering competitive performance with low programmer effort and production deployment.

  12. TLX: Hardware-Native, Evolvable MIMW GPU Compiler for Large-scale Production Environments

    cs.AR 2026-05 unverdicted novelty 6.0

    TLX introduces MIMW-based extensions to Triton that let developers orchestrate warp-group execution and asynchronous hardware features while preserving blocked programming productivity, with kernels deployed in large-...

  13. DAK: Direct-Access-Enabled GPU Memory Offloading with Optimal Efficiency for LLM Inference

    cs.DC 2026-04 unverdicted novelty 6.0

    DAK enables direct GPU access to remote memory for LLM inference via TMA repurposing and a greedy offloading algorithm, achieving up to 3x gains over prefetching baselines on NVLink-C2C and 1.8x on PCIe.

  14. CuTe Layout Representation and Algebra

    cs.MS 2026-03 conditional novelty 6.0

    CuTe defines a hierarchical shape-stride algebra for tensor layouts that generalizes flat layouts and supports compose, inverse, complement, and tiling operations for GPU kernels.

  15. Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants

    cs.LG 2025-11 unverdicted novelty 6.0

    Flashlight is a compiler-native PyTorch framework that generates efficient fused kernels for arbitrary and data-dependent attention variants, supporting more cases than FlexAttention with competitive performance.

  16. ShinkaEvolve: Towards Open-Ended And Sample-Efficient Program Evolution

    cs.CL 2025-09 unverdicted novelty 6.0

    ShinkaEvolve improves sample efficiency in LLM-driven program evolution via parent sampling, code novelty rejection-sampling, and bandit LLM ensemble selection, achieving new SOTA circle packing with 150 samples and g...

  17. Sakana Fugu Technical Report

    cs.LG 2026-06 unverdicted novelty 5.0

    Sakana Fugu trains LLM orchestrators using fine-tuning, evolutionary algorithms, and RL to build query-adaptive multi-agent scaffolds, claiming SOTA results on benchmarks including SWE-Bench Pro and GPQA-Diamond.