Pith. sign in

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

Signed reviews

No signed human review yet.

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 10 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. Flash3D: Super-scaling Point Transformers through Joint Hardware-Geometry Locality

    cs.CV 2024-12 conditional novelty 6.0 of 10

    A point transformer that combines perfect spatial hashing with FlashAttention to align geometric neighborhoods with GPU memory tiles achieves 2.25x faster inference and better semantic segmentation than PTv3.

  9. 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.

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