Pith. sign in

REVIEW 9 cited by

ThunderKittens: Simple, Fast, and Adorable AI Kernels

Not yet reviewed by Pith; the record is open.

Re-run · record.json Download PDF Read on arXiv ↗

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

Benjamin F. Spector , Simran Arora , Aaryan Singhal , Daniel Y. Fu , Christopher Ré This is my paper · ORCID
classification cs.LGcs.AI
keywords kernelsattentionoperationsperformanceabstractionshardwarelinearmatch
verification ladder T0 review T1 audit T2 compute T3 formal
0 comments
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). Continue with ORCID to comment.

Forward citations

Cited by 9 Pith papers

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

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

    cs.DC 2026-07 conditional novelty 7.0 of 10

    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 of 10

    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 of 10

    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.

  4. FG-Attn: Leveraging Fine-Grained Sparse Attention in Video Diffusion Models

    cs.CV 2025-09 conditional novelty 7.0 of 10

    FG-Attn accelerates video DiT inference by computing attention over only the relevant Mx1 key slices per query group, using an asynchronous gather-load kernel, and reports up to 1.65x end-to-end speedup with negligibl...

  5. Fast Video Generation with Sliding Tile Attention

    cs.CV 2025-02 conditional novelty 7.0 of 10

    Sliding tile attention (STA) replaces full 3D attention in video diffusion transformers with dense tile-local windows, achieving 1.89x training-free and up to 3.53x fine-tuned end-to-end speedups on HunyuanVideo with ...

  6. CuTe Layout Representation and Algebra

    cs.MS 2026-03 conditional novelty 6.0 of 10

    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.

  7. KPerfIR: Towards an Open and Compiler-centric Ecosystem for GPU Kernel Performance Tooling on Modern AI Workloads

    cs.DC 2025-05 conditional novelty 6.0 of 10

    KPerfIR is a compiler-centric profiling infrastructure for Triton, demonstrated by a region-based timing tool that improved Triton's FlashAttention-3 by 24.1%.

  8. SPADE: An Input-Adaptive Sparse Attention Engine for Fast Video Diffusion Models Inference

    cs.CV 2026-08 conditional novelty 5.0 of 10

    SPADE combines static, semi-static, and dynamic block-sparse attention with a cheap SICS-based blocking heuristic to speed up video diffusion inference by up to 1.80x end-to-end.

  9. GPU Kernel Scientist: An LLM-Driven Framework for Iterative Kernel Optimization

    cs.LG 2025-06 conditional novelty 5.0 of 10

    An LLM agent loop iteratively optimized a HIP GEMM kernel for AMD MI300 using only end-to-end timings, reaching about 450µs versus 850µs for the PyTorch reference.

Pith tools