Efficient Attention at Scale: KV Cache, GQA & FlashAttention

How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention.

Advanced · Graduate mathematics · about 26 minutes

Reading map and next steps

Intuition

The idea in plain words, before the symbols.

KV caching, grouped-query attention, and FlashAttention reduce different costs of attention. The usual summary, that attention costs $O(T^2)$ , does not say which of these costs a method changes. Using the generation of one new token as the main case, we compare what each method stores, reads, or recomputes. The cache calculation checks how changing the number of key/value heads changes stored bytes. You will need scaled dot-product attention and tensor shapes.

During decoding, the model generates one new token at a time. For that new query, it must still compare against all earlier keys and read all earlier values. Recomputing keys and values for the whole prefix would be wasteful, so serving systems keep a KV cache: the keys and values already produced by every layer for every cached token.

So a more useful question than "is attention quadratic?" is this one:

If the model keeps the same context length and same head dimension, which symbol changes when full multi-head attention becomes grouped-query attention?

The answer is not the number of query heads. It is the number of stored key/value heads, $H_{kv}$ .

Three mechanisms often get blurred together:

KV caching avoids recomputing past keys and values during autoregressive decoding.
Multi-query/grouped-query attention (MQA/GQA) reduces how many key/value heads are stored and read.
FlashAttention computes exact attention but changes how data moves during the computation; in particular it never writes the full attention matrix to the GPU's main (high-bandwidth) memory.

The first two change what is stored and read during decoding. FlashAttention changes how the attention computation moves data between memories, and does not shrink the cache. Real systems combine them, but they are separate changes.

Study prompts for this section: Intuition

Study prompts

Each button copies a prompt about this section to your clipboard, for an AI assistant you already use. Nothing is sent by this site.

Mathematics

Definitions, assumptions and the derivation.

Equation 1 Equation 2 Equation 3

With a KV cache, the attention output for the new token at time $t$ is:

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 cached.

This equation says what the new token reads. The next equation says what the server must keep resident so that read is possible.

KV Cache Memory Scaling

Across a batch of $B$ sequences and $N_{\mathrm{layers}}$ layers, the KV cache stores keys and values for all $T$ positions:

\mathrm{Mem}_{KV} \approx B\cdot N_{\mathrm{layers}}\cdot T\cdot H_{kv}\cdot d_{\mathrm{head}}\cdot 2 \cdot \mathrm{bytes}.

What each symbol counts:

$B$ : active sequences in the batch.
$N_{\mathrm{layers}}$ : transformer layers that each keep their own cache.
$T$ : cached tokens per sequence.
$H_{kv}$ : stored key/value heads per layer.
$d_{\mathrm{head}}$ : width of each head.
$2$ : one tensor for $K$ , one tensor for $V$ .
$\mathrm{bytes}$ : bytes per stored scalar, such as 2 for fp16/bf16.

This is an estimate from the formula alone. It leaves out memory lost to fragmentation, paging metadata, scheduling, shared prefixes and padding, and it says nothing about how changing the attention architecture affects model quality. The KV-cache reference works one case through every factor: one sequence of 4,096 tokens in a 32-layer model with 32 heads of 128 dimensions in fp16.

Grouped-Query Attention

Let there be $H_q$ query heads but only $H_{kv}$ key/value heads, with a mapping $g(h)$ from query head $h$ to a KV head:

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

If full multi-head attention has $H_q = H_{kv}=32$ and GQA keeps $H_q=32$ but uses $H_{kv}=8$ , then the cache-width term is divided by 4:

\frac{\mathrm{Mem}_{KV}^{\mathrm{GQA}}}{\mathrm{Mem}_{KV}^{\mathrm{MHA}}} \approx \frac{H_{kv}^{\mathrm{GQA}}}{H_{kv}^{\mathrm{MHA}}} = \frac{8}{32} = \frac{1}{4}.

The rule to remember:

Holding $B$ , $N_{\mathrm{layers}}$ , $T$ , $d_{\mathrm{head}}$ , and bytes fixed, KV-cache memory scales linearly with $H_{kv}$ .

And the catch for long context:

After GQA narrows $H_{kv}$ , increasing $T$ still grows the cache linearly.

For the mapping from query heads to shared key/value heads, and an exercise on the cache ratio, see Grouped-Query Attention.

What the sources support

Shazeer (2019) shows that during decoding, reading the large key and value tensors again for every new token can dominate the time per token, and proposes multi-query attention (MQA): one key/value head shared by all query heads. Ainslie et al. (2023) introduce grouped-query attention, which sits between full multi-head attention and MQA, so that $H_{kv}$ becomes a design choice between 1 and $H_q$ . Dao et al. (2022) support a different claim: exact attention runs faster when the implementation moves less data.

None of this shows that GQA always preserves quality, that FlashAttention shrinks the KV cache, or that the formula predicts serving capacity exactly; each of those would need its own evidence.

Study prompts for this section: Mathematics

Study prompts

Each button copies a prompt about this section to your clipboard, for an AI assistant you already use. Nothing is sent by this site.

Code

The same calculation as runnable code, in the notation of the derivation.

Code example 1def kv_cache_gb( batch=1, layers=32, tokens=4096, kv_heads=32, head_dim=128, bytes_per_value=...python

def kv_cache_gb(
    batch=1,
    layers=32,
    tokens=4096,
    kv_heads=32,
    head_dim=128,
    bytes_per_value=2,
):
    scalars = batch * layers * tokens * kv_heads * head_dim * 2
    return scalars * bytes_per_value / 1e9

baseline = dict(batch=1, layers=32, tokens=4096,
                head_dim=128, bytes_per_value=2)

full_mha = kv_cache_gb(**baseline, kv_heads=32)
for label, kv_heads, tokens in [
    ("MHA same context", 32, 4096),
    ("GQA same context", 8, 4096),
    ("GQA long context", 8, 32768),
]:
    gb = kv_cache_gb(**{**baseline, "tokens": tokens}, kv_heads=kv_heads)
    print(f"{label:17s} H_kv={kv_heads:2d} T={tokens:6d} "
          f"KV={gb:6.2f} GB  baseline_mha/this={full_mha/gb:4.1f}x")

The code divides by $10^9$ , so it reports decimal gigabytes: 2.15 GB for full multi-head attention at 4,096 tokens, 0.54 GB with 8 KV heads, and 4.29 GB with 8 KV heads at 32,768 tokens. The figure below reports binary gibibytes ( $1\ \mathrm{GiB} = 2^{30}$ bytes $\approx 1.074$ GB): set it to 4,096 tokens and the same 2,147,483,648 bytes appear as 2.00 GiB.

Study prompts for this section: Code

Study prompts

Each button copies a prompt about this section to your clipboard, for an AI assistant you already use. Nothing is sent by this site.

Demo

Change an input, predict the result, and see which quantities respond.

ExploreChange one control at a time and name what responds.An optional observation guide follows the demo.

Explore Efficient Attention at Scale: KV Cache, GQA & FlashAttention

Loading the interactive demo…

Observation guide

Optional and ungraded. Opening a guide saves nothing.

Choose what to inspect in Efficient Attention at Scale: KV Cache, GQA & FlashAttention, then open its guide.

Use the figure to make a prediction before you calculate:

Start from the memory formula above.
Predict what changes when $H_q=32$ query heads move from full multi-head attention to GQA with $H_{kv}=8$ .
Reveal the result at the same context length.
Increase $T$ and watch the memory grow again.
State the rule: for a fixed model shape and precision, $\mathrm{Mem}_{KV}$ scales with $T\cdot H_{kv}$ .
Say which term each change touches: GQA changes $H_{kv}$ , a longer context changes $T$ , an 8-bit cache changes the bytes, and FlashAttention changes none of them. At $g=1$ nothing is shared, so the cache is the same size as full multi-head attention.

Where to go next:

Long Context for what happens as $T$ grows: position handling and retrieval quality.
LLM Serving for memory allocation, batching, latency and throughput.
FlashAttention for exact attention that moves less data, rather than for cache storage.

Study prompts for this section: Demo

Study prompts

Each button copies a prompt about this section to your clipboard, for an AI assistant you already use. Nothing is sent by this site.

Focus on

In focus

Equation: Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2

For Equation 2 in Efficient Attention at Scale: KV Cache, GQA & FlashAttention, what does each symbol mean, what assumption makes it valid, and which source or code example supp...

BeforeScaled Dot-Product Attention & Transformer LayersNextLLM Serving at Scale: Prefill, Decode & Continuous Batching

DetailsEquation in the derivation

About this itemEquation in the derivation

2 sources citedChecked against the cited passagesChecking this browser

Reference id

equation:attention-transformers/efficient-attention#math-object-2

Math3 equationso_t = softmax (\frac{q_t K_{1:t}^{^T}}{sqrt(d_k)}) V_{1:t},CodeEfficient Attention at Scale: KV Cache, GQA & FlashAttention code exa...def kv_cache_gb( batch=1, layers=32, tokens=4096, kv_heads=32, head_dim=128, bytes_per_value=...DemoEfficient Attention at Scale: KV Cache, GQA & FlashAttention interact...Change one control at a time and name what responds.

Worked calculation

Predict how sharing key–value heads changes the cache.

Keep the equation fixed, choose a point of view, predict what grouping changes, then reveal the calculated memory.

Selected equationMem_KV = B * N_layers * T * H_kv * d_head * 2 * bytesLearner view

1Question

What symbol changed, and what stayed fixed?

Learner

2Equation

KV-cache memory equation

Mem_KV = B * N_layers * T * H_kv * d_head * 2 * bytes

3Prediction

Choose a prediction below.

No prediction yet.

4Change

T = 32k, g = 4, bytes = 2

H_kv = 8

5Evidence

Shown after you predict.

8 stored K/V heads

6What stays the same

KV-cache memory scales linearly with stored K/V heads.

The number of query heads does not change.

7Next

Move g from 1 to 4, then say why the ratio changes before reading the result.

Long context and FlashAttention build on this.

Predict first

If 32 query heads share K/V heads in groups of 4, what happens to KV-cache memory compared with ordinary multi-head attention?

What stays the sameFor fixed B, N_layers, T, d_head and bytes, KV-cache memory scales linearly with stored K/V heads, not query heads.

Make a prediction to see the calculated cache size.

Open the full KV cache demo Continue in the KV-memory Lab Ask in the discussion

Learner view

What symbol changed, and what stayed fixed?Move g from 1 to 4, then say why the ratio changes before reading the result.

Mem_KV = B * N_layers * T * H_kv * d_head * 2 * bytesH_kv = H_q / g

B = 1N_layers = 32H_q = 32d_head = 128bytes = 2

Sequence length T

KV sharing group size g

Bytes per value

32 query heads (unchanged)

8 stored K/V heads after grouping

8KV headshiddencache at 32khiddenreduction vs MHAhiddenKV-head share vs MHA

Calculated from the formula above; nothing here was measured on a model or GPU.

What the calculation showsGQA changes H_kv, not H_q.

The cache equation and the head count above separate memory pressure from the number of query heads.

Your observation

Shown after you make a prediction.

Next stepMove g from 1 to 4, then say why the ratio changes before reading the result.

Discussion notes: the KV cache

Watch which term changes before comparing efficiency methods.

Prepared by the site, not reader comments

EquationEfficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2Question

For Equation 2 in Efficient Attention at Scale: KV Cache, GQA & FlashAttention, what does each symbol mean, what assumption makes it valid, and which source or code example supports it?

Equation in the derivation

Reference id

equation:attention-transformers/efficient-attention#math-object-2

Sources

Checked against the cited passages · Sources: Fast Transformer Decoding: One Write-Head is All You Need; GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Check

Change one variable at a time in the KV cache demo.

Evidence

Separate cache writes, cache reads, and IO-aware attention.

Next step

First double T and predict the memory change. Then reduce Hkv and predict the cache ratio. Leave FlashAttention out of these predictions: it changes how attention is computed, not how much the cache stores.

Best explanationSeparate cache writes, cache reads, and IO-aware attention.

KV caching saves recomputing old keys and values, but every decode step still reads the stored K and V tensors. GQA changes the cache-size term Hkv; FlashAttention changes how exact attention is tiled so the full attention matrix is not materialized.

About Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2

Common questionIf keys and values are already cached, why does long context still hurt?

The cache removes repeated projection work, not the need to read prior K/V vectors. The first thing to settle is whether compute was saved, memory traffic was saved, or both.

Common confusion

Teaching noteState the shapes before making a claim about speed.

Name B, L, T, Hkv, d_head, and bytes before saying “efficient.” Then ask which symbol moves when context grows, when KV heads are shared, and when the kernel avoids materializing attention.

For teachers

In practiceA serving engineer cares about the term that must be read on every new token.

Hold layers, batch, and head size fixed. Doubling T doubles the KV cache. Moving from full MHA to GQA reduces the KV-head factor, but quality, batching, paging, and hardware still decide the final latency.

Production caveat

Source correctionDo not credit FlashAttention for shrinking the decode KV cache.

Use Shazeer and Ainslie for MQA/GQA decode bandwidth and KV-head sharing. Use Dao et al. for IO-aware exact attention and avoiding the full attention matrix. These are related efficiency stories, not the same lever.

Checked against the cited passages · Sources: Fast Transformer Decoding: One Write-Head is All You Need; GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Using an AI assistantAsk about this equation, not about attention in general.

Ask for one formula, one small code example, the source that supports it, one common misconception and one next experiment. An answer that treats GQA and FlashAttention as the same saving has gone wrong.

Prepared advice

Open questionWhere does the quality tradeoff appear when fewer KV heads serve more query heads?

The cache equation says nothing about model quality, so this question stays open here. Answering it needs a named model, workload and quality metric.

Not settled here

Next experimentChange one variable at a time in the KV cache demo.

Predict, then reveal

Ask an AI assistantThe copied prompt includes this item, the notes above and the cited sources.

Ask about this in the discussion Go to the demo See the code

After the first pass

Go deeper

Optional material for a second reading. Sources are always shown; the other panels open from their headings or from these links, and closing one keeps what you wrote or revealed.

What to look for in each sectionChoose a section and what you expect to see in it, then read short advice.

Section by section

Choose what to look for in each section

How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention.

Choose what to look for (no demo yet)01 / Intuition

Illustration of KV-cache sharing, grouped query lanes, and reduced memory bandwidth.

What to look for

Start with the picture, metaphor, or geometric mechanism.

Choose first

Choose what you expect to see in this section of Efficient Attention at Scale: KV Cache, GQA & FlashAttention.

Questions to ask of the pictureChoose a part of the picture and what you expect, then read short advice.

Looking at the picture

Questions to ask of the picture

How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention.

4 of 4 sections writtenNo live demo yet

Visual cueWhich part of the picture should you look at first?

Highlight size2/4

Question

Which part of the picture should you look at first?

Choose first

Pick the part of the picture you expect to explain Efficient Attention at Scale: KV Cache, GQA & FlashAttention.

Sources

References for this page

A reference shows where an idea comes from; it does not vouch for every step on this page.

paper · 2019Fast Transformer Decoding: One Write-Head is All You NeedShazeerSee arXiv:1911.02150v1, §2.4.1 (pp. 3–4) and §3–3.1 (pp. 4–6)

Why it is cited

arXiv preprint. Shows that reloading the key and value tensors at every step can dominate incremental decoding, and proposes one key/value head shared by all query heads (multi-query attention).

Open source

paper · 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie et al.See EMNLP 2023, §2.2, pp. 4895–4896 (PDF pp. 1–2)

Why it is cited

EMNLP 2023. Introduces grouped-query attention, between multi-head attention and MQA: query heads share key/value heads in groups.

Open source

paper · 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao et al.

Why it is cited

NeurIPS 2022. Tiling plus a running (online) softmax computes exact attention without storing the full attention matrix in GPU main memory.

Open source

What each claim rests on1 claim, each with its sources, where it appears on this page, and its caveats.

Claims and sources

How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention.

Checked against the cited passages

The checks were made by this site, not by an independent reviewer.

Explain it without the pageWrite your own explanation before asking for feedback. Your draft and any hints stay while this page is open.

Practice · Efficient Attention at Scale: KV Cache, GQA & FlashAttention

Try the idea in your own words

How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention.

Equation · Selected for practice

Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2

What it rests on: Checked against the cited passages · Sources: Fast Transformer Decoding: One Write-Head is All You Need; GQA: Training Generalized Multi-Query Tran...

Context and links

Equation in the derivation

Choose a task

Explain the mechanism

For Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2: For Equation 2 in Efficient Attention at Scale: KV Cache, GQA & FlashAttention, what does each symbol mean, what assumption makes it valid, and which source or code example supp... Explain your answer, including what changes, why, and which assumption matters.

Your answerNo answer yet

A rough first thought is enough. Your draft stays when you change tasks.

Your draft stays on this page and is cleared when you leave or reload.

A little help · Explain

Open one hint at a time. These are suggestions, not your answer or a grade.

0 of 3 hints shown for this question.

Clearing your answer keeps this help record. Outside help is your own report; this page cannot check it.

Where am I stuck? (optional)

Your own description, not an automatic diagnosis

Choose one, or leave this unspecified. Select it again to clear it.

Take your draft to a feedback conversation

No AI feedback runs here. You can copy a prompt to use elsewhere; nothing is sent automatically. Review the text before sharing, and leave out private information.

Write an attempt before copying a feedback prompt.

A draft, or an AI reply to it, is not a test of what you have learned. To check that, try a different case later without help.

Ask about this pagePick one item on this page and copy a prompt about it into an AI assistant you already use. Notes you write here stay in this browser.

Ask about this pageClose

EquationEfficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2Checked against the cited passages · Sources: Fast Transformer Decoding: One Write-Head is All You Need; GQA: Training Generalized Multi-Query Tran...def kv_cache_gb( batch=1, layers=32, tokens=4096, kv_heads=32, head_dim=128, bytes_per_value=...

Your question

Choose what your question is about

Pick the idea, equation, source, code, claim, misconception or demo state you want to ask about. The prompt you copy includes it, so the answer can stay on that item.

Next local actionNo local draft saved yet

Open the draft below to save one note and next action in this browser.

equationEquation snippet: Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2. Source notation: Mem_{KV} \approx B\cdot N_{layers}\cdot T\cdot H_{kv}\cdot d_{head}\cdot 2 \cdot bytes.

Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2

Starting question

For Equation 2 in Efficient Attention at Scale: KV Cache, GQA & FlashAttention, what does each symbol mean, what assumption makes it valid, and which source or code example supports it?

SourcesCheck the 3 cited sources listed under Sources.Notes can be saved for this item

Points of view on this item

Fixed prompts generated from this item. They are not comments from people or independent reviews.

Learner’s first stepAsk what would make "Efficient Attention at Scale: KV Cache, GQA & FlashAttention equation 2" feel predictable rather than familiar.

Assumption

The 3 cited sources must support this exact item, not just the surrounding topic.

Source-checking summary

The question is settled only when each term has a stated role, shape and assumption, and a check in the code or the demo.

Proposed experiment

Change the term you think matters most and compare the predicted effect with the code or the demo.

Next action

Every symbol has a clear role and shape

Evidence4 checks

ObservationChecking for a saved observation

ActionReady for one action

PromptLearner prompt ready to copy

Go to this item

01ObservationChecking what this browser has saved

02EvidenceChecking for a saved observation

03SourcesCheck the 3 cited sources listed under Sources.

04Next stepSave one next action

Local action draftNo local draft saved yetOpen when you are ready to write one next action

Local action draft

This draft stays in this browser and is attached to this item only.

Draft noteNext action

No local draft saved.

Evidence to inspect

What the 3 cited sources say about this exact item
Symbol meanings, shapes, units, and hidden assumptions
The source line, page, or concept section where the equation is justified
Runnable code that computes the same terms

What would resolve this

Every symbol has a clear role and shape
You know which term changes in the paper or the demo
The equation predicts the observed behavior without extra hand-waving

Reference id

equation/concept-notebook/attention-transformers/efficient-attention/math/equation-2
equation:attention-transformers/efficient-attention#math-object-2