glossary/compute/flash-attention
FlashAttention
An attention kernel that never writes the full attention matrix to HBM. It processes keys and values in tiles that fit the GPU's on-chip SRAM, keeping a running softmax so the result is exact, and reads each tile once. Standard attention materialises a matrix that grows with the square of sequence length; this reads memory linear in it. It made long contexts affordable and is the default kernel in every serving engine.
19 TB/s versus 1.5 TB/s
On-chip SRAM against HBM bandwidth on the A100, the gap the tiling exploits, per the paper.
Related
- glossary/
- Kernel fusion
Combining several operations into one kernel so intermediate results stay in registers or on-chip memory instead of making a round trip to HBM.
- High bandwidth memory
The stacked DRAM on the GPU package that holds the weights and the KV cache.
- Memory-bound
Limited by how fast bytes can be read, not by how fast they can be operated on.
- FlashInfer
A kernel library for the operations serving engines share: attention over paged KV cache, in prefill and decode shapes, plus sampling, GEMM and MoE routines.
- sources/
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
- FlashAttention-2 (Dao, 2023)
- topics/