Flash Attention: Fast, Memory-Efficient Attention for LLMs

Self-attention is the core operation in every transformer. It is also the operation that breaks at long context lengths.

For a sequence of N tokens, standard attention materializes an N × N matrix in GPU HBM. At N=128K (a typical long-context window), that's 128,000 × 128,000 × 2 bytes = 33 GB — just for the attention scores of a single head. Multiply by 32 heads and you've exceeded an H100's 80 GB HBM before the FFN layers even run.

Flash Attention (2022) solved this by observing that the N × N matrix never needs to exist all at once. The final attention output can be computed in tiles that fit in SRAM, reading from and writing to HBM once — not repeatedly.

---

Standard Attention: The Memory Problem

The self-attention formula:

Attention(Q, K, V) = softmax(QK^T / √d_k) V

Broken into steps:

flowchart LR

Q["Q ∈ R^(N×d)"] --> S["S = QK^T <br/> N × N matrix <br/> ← BOTTLENECK"]
K["K ∈ R^(N×d)"] --> S
S --> P["P = softmax(S) <br/> N × N matrix"]
P --> O["O = PV <br/> N × d output"]
V["V ∈ R^(N×d)"] --> O

Both S and P are N × N. For N=32,768 in BF16:

S = 32768 × 32768 × 2 bytes = 2.1 GB  (per head)
P = 32768 × 32768 × 2 bytes = 2.1 GB  (per head)

Writing 2.1 GB to HBM and reading it back for softmax, then reading again for the PV product — this is three full HBM round trips per head, per layer. On H100, HBM bandwidth is 3.35 TB/s. Every unnecessary HBM access is wasted time.

---

GPU Memory Hierarchy

Understanding Flash Attention requires understanding where computation happens:

Memory	Size (H100)	Bandwidth	Latency
SRAM (per SM, registers + shared mem)	~192 KB	~19 TB/s	~10 cycles
HBM (High Bandwidth Memory)	80 GB	3.35 TB/s	~500 cycles
System RAM	Host	PCIe ~64 GB/s	~μs

The ratio: SRAM is 6× faster than HBM but 400,000× smaller. The goal is to keep computation in SRAM and minimize HBM reads/writes.

---

Flash Attention: IO-Aware Tiling

Flash Attention computes the exact same result as standard attention but never materializes the full N × N matrix. Instead, it processes Q, K, V in tiles that fit in SRAM.

The Algorithm

flowchart TB

HBM["HBM <br/> Q, K, V stored in tiles"]

-->|"load Q_block, K_block, V_block <br/> (one tile at a time)"| SRAM["SRAM (per SM) <br/> Q_i, K_j, V_j <br/> + running softmax statistics <br/> (m_i, ℓ_i)"]

SRAM

-->|"compute attention for this tile <br/> update running output O_i"| Compute["Tile computation <br/> (QK^T, softmax, PV) <br/> in SRAM — never spills to HBM"]

Compute

-->|"write O_i back to HBM <br/> (once, at the end)"| HBM

Key insight: softmax requires knowing all scores s_ij = Q_i · K_j / √d before computing exp(s) / Σexp(s). Flash Attention uses an online softmax trick: maintain a running maximum m and running sum ℓ across tiles, then rescale the accumulated output as each new tile arrives.

For each tile K_j, V_j:
  s_ij = Q_i @ K_j^T / √d           # scores for this tile
  m_ij = max(m_{i,j-1}, rowmax(s_ij))  # update running max
  P̃_ij = exp(s_ij - m_ij)          # numerically stable softmax scores
  ℓ_ij = exp(m_{i,j-1} - m_ij) × ℓ_{i,j-1} + rowsum(P̃_ij)  # update sum
  O_i  = (ℓ_{i,j-1}/ℓ_ij) × exp(m_{i,j-1}-m_ij) × O_{i,j-1} + P̃_ij @ V_j

The final O_i = O_i / ℓ_ij gives the exact same result as computing softmax on the full matrix.

Complexity Comparison

	Standard Attention	Flash Attention
Memory	O(N²)	O(N)
HBM reads/writes	O(N² / M) where M = SRAM size	O(N² × d / M) — same FLOPs, fewer passes
N × N matrix materialized	Yes — stored in HBM	No — computed in SRAM tiles
Exact or approximate	Exact	Exact

Flash Attention is not an approximation. It computes mathematically identical output.

---

Flash Attention 2 (2023)

Flash Attention 2 improved on FA1 with:

1. Better work partitioning across warps — FA1 had warp-level synchronization overhead; FA2 assigns Q to outer loop (parallelized across warps) and K,V to inner loop
2. Reduced non-GEMM FLOPs — minimizes rescaling operations in the online softmax
3. Causal mask optimization — skips computation for masked (future) tokens in the lower triangle

Result: ~2× faster than FA1 on A100/H100.

---

Flash Attention 3 (2024 — H100-specific)

Flash Attention 3 exploits H100 hardware features unavailable on previous GPUs:

Feature	What it enables
WGMMA (Warpgroup Matrix Multiply-Accumulate)	Larger matrix operations in a single instruction
TMA (Tensor Memory Accelerator)	Asynchronous HBM→SRAM data movement — overlaps data loading with computation
FP8 Tensor Cores	2× throughput vs BF16 for the QK^T GEMM

Result: FA3 achieves 75% of H100 theoretical FlOPs — near peak utilization. FA2 achieves ~35%.

---

Attention Variants: MHA, GQA, MQA

Standard multi-head attention replicates K and V for every head. Newer architectures use shared K,V to reduce KV cache memory — the dominant inference memory bottleneck.

Multi-Head Attention (MHA)

n_q_heads = n_k_heads = n_v_heads = H
KV cache per token = 2 × H × d_head × bytes

For Llama 2 70B: H=64 heads, d_head=128, BF16 → 64 × 128 × 2 × 2 = 32 KB per token per layer. At 32 layers, 4096 context: 32 KB × 32 × 4096 = 4 GB per sequence.

Grouped Query Attention (GQA)

GQA (used in Llama 3, Mistral, Gemma) groups multiple Q heads to share one K,V pair:

flowchart LR

Q1["Q head 1"] --> KV1["K,V group 1"]
Q2["Q head 2"] --> KV1
Q3["Q head 3"] --> KV1
Q4["Q head 4"] --> KV1

Q5["Q head 5"] --> KV2["K,V group 2"]
Q6["Q head 6"] --> KV2
Q7["Q head 7"] --> KV2
Q8["Q head 8"] --> KV2

With 8 Q heads grouped into 2 K,V groups (n_kv_heads = 2):

KV cache = 2 × 2 × d_head × bytes  (4× smaller than MHA with 8 heads)

Llama 3.1 70B uses GQA with n_heads=64, n_kv_heads=8 — an 8× KV cache reduction vs MHA.

Multi-Query Attention (MQA)

All Q heads share a single K,V pair (n_kv_heads=1). Maximum KV cache reduction, some quality loss. Used in older models (PaLM, Falcon).

Variant	n_kv_heads	KV cache size	Quality
MHA	= n_heads	1× (baseline)	Highest
GQA	1 < n < n_heads	1/g × (baseline)	Near-MHA
MQA	1	1/n_heads	Lowest

GQA is the current default for large models — the sweet spot between KV cache reduction and output quality.

---

KV Cache

During inference, attention recomputes the same key and value vectors for every token it generates. The KV cache stores these vectors so they're computed once per token, not N times.

flowchart LR

Prefill["Prefill phase <br/> (process N input tokens) <br/> Compute K,V for all tokens <br/> Store in KV cache"]

-->Decode["Decode phase <br/> (generate token t) <br/> Reuse K,V for tokens 0..t-1 <br/> Append new K_t, V_t to cache"]

Decode

-->|"next token"| Decode

KV cache memory for one H100 serving Llama 3.1 70B (GQA: n_kv_heads=8):

per token per layer: 2 × 8 × 128 × 2 bytes = 4 KB
80 layers:           4 KB × 80 = 320 KB per token
4096 context:        320 KB × 4096 ≈ 1.3 GB
32 concurrent seqs:  1.3 GB × 32 ≈ 40 GB (of 80 GB HBM)

GQA's 8× reduction is the difference between supporting 32 concurrent sequences and 4.

---

PagedAttention (vLLM)

Flash Attention eliminates the N×N matrix from GPU memory. PagedAttention (vLLM, 2023) solves a different problem: KV cache memory fragmentation.

Without PagedAttention:

• Allocate a contiguous KV cache block for the maximum sequence length at request start
• A 128-token sequence and a 4096-token sequence both reserve 4096 tokens of KV cache
• Fragmentation wastes 50–60% of KV cache memory

PagedAttention divides KV cache into fixed-size pages (e.g., 16 tokens), allocated on demand:

flowchart LR

Req1["Request 1 <br/> (200 tokens)"]

-->|"allocates"| Page1["Page 1 <br/> (16 tokens)"]

-->Page2["Page 2 <br/> ..."]

-->PageN["Page 13 <br/> (8 tokens used)"]

Req2["Request 2 <br/> (32 tokens)"]

-->|"allocates"| Page14["Page 14"]

-->Page15["Page 15"]

Pages from completed sequences are freed and reused. This reduces KV cache waste from 60% to <5%.

Flash Attention and PagedAttention are complementary:

• Flash Attention: eliminates the O(N²) attention matrix from HBM during computation
• PagedAttention: manages KV cache memory allocation across concurrent sequences

Both are active simultaneously in production inference stacks (TRT-LLM, vLLM, NIM).

---

Integration in Production Stacks

Stack	Flash Attention	GQA support	PagedAttention	Notes
TensorRT-LLM	FA2 + FA3 (H100)	✅	✅ (Paged KV cache)	Used inside NIM profiles
vLLM	FA2 + Triton FA kernels	✅	✅ (native)	Open-source alternative
PyTorch (torch.compile)	scaled_dot_product_attention	✅	❌	Auto-selects FA kernel
NIM (H100 profile)	FA3	✅	✅	All optimizations active

PyTorch Usage

import torch
import torch.nn.functional as F

# PyTorch automatically dispatches to Flash Attention when available
# No code changes needed from standard attention
output = F.scaled_dot_product_attention(
    query,    # (batch, n_heads, seq, d_head)
    key,
    value,
    attn_mask=None,
    dropout_p=0.0,
    is_causal=True     # enables causal mask optimization
)

TRT-LLM Configuration

NIM selects the correct FA variant based on GPU type automatically. For manual TRT-LLM builds:

from tensorrt_llm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=8,
    kv_cache_config=KvCacheConfig(
        free_gpu_memory_fraction=0.9,   # 90% of HBM for KV cache
        enable_chunked_context=True     # chunked prefill for long contexts
    ),
    # FA3 selected automatically on H100
)

---

Performance Numbers on H100

Sequence length	Standard attention (ms)	FA2 (ms)	FA3 (ms)	Speedup (FA3 vs standard)
512	0.8	0.5	0.4	2×
2048	8.2	2.1	1.4	5.9×
8192	130	14.3	8.7	15×
32768	OOM	56	33	—

At long context lengths, Flash Attention is not a marginal optimization — standard attention simply cannot run.

---

The Inference Optimization Chain

flowchart LR

Model["Trained LLM <br/> (BF16 weights)"]

-->FA["Flash Attention <br/> (O(N) memory <br/> SRAM tiling)"]

-->GQA["GQA <br/> (8× smaller KV cache)"]

-->PA["PagedAttention <br/> (<5% KV fragmentation)"]

-->TRT["TensorRT-LLM <br/> (FP8, kernel fusion <br/> FA3 kernel)"]

-->NIM["NIM Profile <br/> (h100x8-fp8-tp8) <br/> auto-configured"]

Flash Attention is the foundation. Every other inference optimization (GQA, PagedAttention, TRT-LLM kernel fusion) is layered on top of it.

---

Key Takeaways

Concept	Purpose
Standard attention	O(N²) HBM memory — hits 33 GB per head at N=128K
Flash Attention	Tiles Q,K,V to fit in SRAM; never materializes N×N; exact result
Flash Attention 2	Better warp partitioning; ~2× faster than FA1
Flash Attention 3	WGMMA + TMA + FP8; 75% of H100 peak FLOPs
GQA	Multiple Q heads share K,V pairs; 8× KV cache reduction in Llama 3
KV cache	Stores K,V per token so decode reuses prefill computation
PagedAttention	Eliminates KV cache fragmentation; complementary to Flash Attention
NIM / TRT-LLM	FA3 + GQA + paged KV cache selected automatically for H100 profiles

The practical implication: at context lengths above 4K, Flash Attention is not optional. Standard attention fails or degrades severely. Every production LLM inference system — NIM, vLLM, TGI — ships Flash Attention by default. The question for a DGX Cloud deployment is only whether you're running FA2 (A100) or FA3 (H100).

---

Related Posts

• Optimizing AI Inference at Scale: The Full Stack — KV cache, continuous batching, tensor parallelism, and speculative decoding; Flash Attention is the foundational layer the rest builds on
• NVIDIA NIM: Optimized Inference Microservices — NIM H100 profiles include FA3, GQA, and paged KV cache out of the box; understanding Flash Attention explains why H100 profiles outperform A100 profiles
• TensorRT and High-Performance AI Inference — TRT-LLM compiles the Flash Attention kernel alongside graph fusion and FP8 quantization; the compiled engine is what NIM serves
• What are Transformer Models? — self-attention is the operation Flash Attention optimizes; understanding the QK^T·V computation is the prerequisite
• AI Programming Model (CUDA) — SRAM (shared memory), HBM bandwidth, and SM architecture that make the FA tiling strategy possible