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

Measures generative AI inference performance with token-level metrics (TTFT, inter-token latency), latency, and throughput under realistic traffic patterns. Provides a multiprocess engine, real-time TUI dashboard, extensible plugins, and integrations for telemetry and result uploads, aimed at inference benchmarking and capacity planning.

GitHub
AI Train2019

Train and experiment with multi-billion to trillion-parameter transformer models on large GPU clusters using GPU-optimized building blocks and reference training scripts; offers advanced parallelism and mixed-precision support for research teams and ML engineers.

GitHub

Indexes full text of visited web pages and local files on a self‑hosted server so you can search your personal knowledge from a web UI, terminal, CLI, or an AI assistant. Runs without mandatory telemetry, offers a browser extension for automatic capture, and supports optional semantic search via a configurable embeddings endpoint.