AIAny
AI Infra2022
Icon for item

FlashAttention

Fused CUDA kernels that compute exact attention without ever writing the full N×N score matrix to GPU memory, cutting memory from quadratic to linear and speeding up training and inference on A100/H100. Ships FlashAttention-2/3 plus KV-cache decode paths.

Introduction

The bottleneck in attention was never the math — it's the memory traffic. A standard implementation writes the full N×N score matrix out to slow GPU HBM and reads it back, so attention is IO-bound long before it is compute-bound. The trick here is to tile the computation and keep intermediate scores in fast on-chip SRAM, streaming through the sequence and never materializing that matrix in HBM — producing the numerically exact same output.

What Sets It Apart
  • Exact, not approximate. Unlike sparse or low-rank attention, outputs are bit-for-bit equivalent to naive attention, so it drops into an existing model with zero accuracy trade-off — no retraining, no quality regression.
  • Linear memory in sequence length. Removing the quadratic activation term is what made training at 32k+ context practical on ordinary hardware, rather than a research stunt.
  • Each version chases the silicon. v2 rebalanced work across GPU warps for higher occupancy; v3 exploits Hopper (H100) async copies and FP8; Triton and ROCm backends extend reach beyond hand-written CUDA and NVIDIA.
  • A dedicated decode path. KV-cache and paged variants target the memory-bound single-token generation of inference, not just the throughput-bound training case.
Great Fit If / Look Elsewhere

Great fit if you train or serve transformers on modern NVIDIA GPUs and want speed and memory wins with no model changes — it is already the default kernel inside PyTorch, vLLM, and most training stacks, so you may be using it without knowing. Look elsewhere if you are on pre-Ampere cards (support is partial or absent), need CPU inference, or run short sequences where the score matrix already fits comfortably in cache — there the payoff shrinks toward noise.

Information

  • Websitegithub.com
  • OrganizationsDao AI Lab
  • AuthorsDao-AILab, Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
  • Published date2022/05/19

Categories

More Items

GitHub
AI Infra2025

Defines a vendor-neutral JSON/YAML semantic model specification and tooling to exchange metrics, dimensions, lineage and other business semantics across analytics, AI and BI platforms; includes a core spec, validators, converters (dbt, GoodData, Salesforce) and example models.

GitHub
AI Train2025

An asynchronous, high-throughput framework for large-scale reinforcement learning and agentic training that scales to 1T+ MoE models and 1000+ GPUs, with native verifiers integration, end-to-end SFT/RL/evals, and Slurm/Kubernetes deployment; requires NVIDIA GPUs.

GitHub

Runs a self-hosted meeting bot and transcription API that joins Google Meet, Teams and Zoom and streams speaker-attributed transcripts in real time. Compiles meetings into a git-backed Markdown workspace and runs sandboxed agents on your infrastructure; Apache-2.0 and air-gap capable.