Efficient Attention at Scale: KV Cache, GQA & FlashAttention

Why It Matters for Modern Models

A key–value (KV) cache stores the keys and values of earlier tokens so attention can reuse them during generation. Llama 2 70B and the Llama 3 models use grouped-query attention with 8 key-value heads to shrink the KV cache (Touvron et al., 2023; Grattafiori et al., 2024).
KV-cache memory grows linearly with context length and batch size. For a model with about 7 billion parameters, 32 layers, 32 heads of width 128 and no head sharing, 16-bit keys and values take 524,288 bytes (0.5 MiB) per token, so about 27,000 cached tokens (summed over the batch) take as much memory as the 14 GB of 16-bit weights; beyond that, reading the cache is most of the memory traffic in each decode step.
FlashAttention made training faster and longer contexts practical by reducing memory traffic without approximating attention (Dao et al., 2022).
KV-cache compression is an active research area: SnapKV keeps, for each head, only the prompt positions that receive the most attention from the end of the prompt (Li et al., 2024), and RazorAttention keeps the full cache only for a few "retrieval" heads (Tang et al., 2025).

Questions to Think About

When does the KV cache, rather than the weights, become the largest use of inference memory?
Why does GQA save memory and bandwidth without reducing the attention arithmetic each query head performs?
Why can some KV-compression methods drop most cached positions for many heads, and which heads need the full cache?

Interactive Demos

Core Math

If you prefer intuition first, try the demos above, then come back for the math.

Key Equation

o_t = \mathrm{softmax}\!\left(\frac{q_t K_{1:t}^{\top}}{\sqrt{d_k}}\right) V_{1:t}

At long context, the cost of attention is set less by its formula than by how keys and values are stored and moved.

Autoregressive attention with a KV cache:

During decoding, the model forms the new query $q_t$ and attends to the cached keys and values of all earlier tokens:

o_t = \mathrm{softmax}\!\left(\frac{q_t K_{1:t}^{\top}}{\sqrt{d_k}}\right) V_{1:t}

where $(K_{1:t}, V_{1:t})$ are stored, not recomputed. Without the cache, every new token would recompute keys and values for the whole prefix, so generating $T$ tokens would repeat $O(T^2)$ projection work.

Grouped-Query Attention (GQA):

Multiple query heads share fewer KV heads. Let query head $h \in \{1,\dots,H_q\}$ map to KV-head group $g(h) \in \{1,\dots,H_{kv}\}$ :

o_h = \mathrm{softmax}\!\left(\frac{Q_h K_{g(h)}^{\top}}{\sqrt{d_k}}\right) V_{g(h)}

This interpolates between multi-head attention ( $H_{kv} = H_q$ , no sharing) and multi-query attention ( $H_{kv} = 1$ , all query heads share one key-value head).

KV cache memory:

For one sequence, each layer stores keys and values for every past token:

\text{Mem}_{KV} = 2 \cdot T \cdot H_{kv} \cdot d_{\text{head}} \cdot \text{bytes per element}

where $T$ is the context length and the factor 2 counts keys and values. A GQA cache is the fraction $H_{kv}/H_q$ of the full multi-head cache, a reduction by the factor $H_q/H_{kv}$ .

FlashAttention: standard attention writes the full $T \times T$ score matrix to slow GPU memory (HBM) and reads it back. FlashAttention splits the computation into tiles that fit in fast on-chip memory (SRAM) and combines the partial softmax results exactly, so it returns the same output with memory linear in $T$ and far fewer HBM reads and writes (Dao et al., 2022).

Key Papers

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Dao et al.2022NeurIPS

Read paper →

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Ainslie et al.2023EMNLP

Read paper →

SnapKV: LLM Knows What You Are Looking for Before Generation

Li et al.2024NeurIPS

Read paper →

RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

Tang et al.2025ICLR

Read paper →

Connections

Prerequisites

Attention RoPE Efficiency

Leads to

Speculative Decoding LLM Serving MoE MoE Serving Long Context State Space Models FlashAttention GQA

Where to go next

A newer notebook covers Efficient Attention, with an interactive demo.

Continue in the newer notebook

Other connections (8)

Which limitations does it address, or leave to a later idea?

KV compression: a limitation of Efficient Attention that Long Context addresses. Open Long Context
FlashAttention removes the quadratic memory cost of attention (the computation is still quadratic), but during generation the key-value cache still grows linearly with context length T. Long-context systems add cache quantization and compression to reduce that remaining cost.
Attention costs: a limitation of Efficient Attention that Attention addresses. Open Attention
Efficient-attention methods target different costs of standard attention. FlashAttention computes exact attention without storing the n × n matrix; grouped-query attention shrinks the key-value cache; sliding-window attention limits each token to recent positions, which changes what is computed.
KV cache memory: a limitation of Attention that Efficient Attention addresses. Open Attention
During generation every attention head caches its keys and values, so the cache grows with the number of tokens, heads and layers. Multi-query and grouped-query attention shrink it by letting several query heads share one key and value head (Shazeer, 2019; Ainslie et al., 2023). FlashAttention addresses a different cost, the n × n attention matrix.

Where does it break down?

O(n²) memory: where Attention breaks down, taken up on this page. Open Attention
A standard implementation stores the full n × n matrix of attention weights for the backward pass; at n = 100,000 that is 10¹⁰ numbers per head, 20 GB in 16-bit precision. FlashAttention avoids storing it by computing attention in tiles and recomputing them during the backward pass (Dao et al., 2022).

Where else does the same idea appear?

Blocked KV memory layout: the same technique appears in LLM Serving. Open LLM Serving
Both split attention data into fixed-size blocks, for different reasons. FlashAttention tiles the computation so that each block fits in fast on-chip memory (Dao et al., 2022). PagedAttention stores the key-value cache in fixed-size blocks that need not be contiguous, so memory can be allocated as sequences grow without fragmentation (Kwon et al., 2023).
Recurrent view: the same technique appears in State Space Models. Open State Space Models
Both SSMs and linear attention can be written as recurrences with particular transition matrices. The "selective" in Mamba means that the recurrence depends on the input (Gu & Dao, 2024), the same idea as the data-dependent gates of gated linear attention.
Memory efficiency: the same technique appears in Long Context. Open Long Context
Long context requires efficient attention: GQA reduces KV cache, sliding window limits active tokens, both enable longer sequences within memory constraints.
Positions at attention time: the same technique appears in RoPE. Open RoPE
RoPE rotates each query and key by an angle set by its position just before the attention scores are computed, so no position embeddings are added to the inputs. Because only queries and keys change, it combines directly with FlashAttention and sliding-window attention.

Efficient Attention at Scale: KV Cache, GQA & FlashAttention

Where to start

Why It Matters for Modern Models

Questions to Think About

Interactive Demos

Core Math

Key Papers

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

SnapKV: LLM Knows What You Are Looking for Before Generation

RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

Connections

Prerequisites

Leads to

Where to go next

Which limitations does it address, or leave to a later idea?

Where does it break down?

Where else does the same idea appear?