Attention & Transformers

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.

status: publishedimportance: importantdifficulty 4/5math: graduateread: 20mlive demo

Back to Attention & Transformers Next: LLM Serving at Scale: Prefill, Decode & Continuous Batching

Editorial transformer-systems illustration of KV-cache sharing, grouped query lanes, and reduced memory bandwidth.

Concept Structure

Efficient Attention at Scale: KV Cache, GQA & FlashAttention

01Intuition

Start with the picture, metaphor, or geometric mechanism.

02Math

Make the objects explicit and connect them with notation.

03Code

Mirror the equations with runnable implementation details.

04Interactive Demo

Manipulate the mechanism and watch the idea respond.

3prerequisites

4next concepts

4related links

Learning map

Efficient Attention at Scale: KV Cache, GQA & FlashAttention

BeforeScaled Dot-Product Attention & Transformer LayersNow4/4 sections readyTryManipulate one control and predict the visible change.NextLLM Serving at Scale: Prefill, Decode & Continuous Batching

Object flow

4/4 sections readyAsk about this Research room

ConceptEfficient Attention at Scale: KV Cache, GQA & FlashAttentionAttention & Transformers EquationEfficient Attention at Scale: KV Cache, GQA & FlashAttention equation 1Exact equation object CodeEfficient Attention at Scale: KV Cache, GQA & FlashAttention code wit...Exact code witness DemoEfficient Attention at Scale: KV Cache, GQA & FlashAttention interact...Visualization object ClaimAt incremental decoding time, KV cache reads can become a memory-band...Exact claim check SourceFast Transformer Decoding: One Write-Head is All You NeedExact source object

ConceptEfficient Attention at Scale: KV Cache, GQA & FlashAttentionAttention & Transformers

3 sources attachedLocal snapshot ready

concept:attention-transformers/efficient-attention

Codewitness nearby Predictbefore reveal Roomobject handoff

Conceptual Bridge

What should feel connected as you move through this page.

Carry inScaled Dot-Product Attention & Transformer Layers

Bring the mental model from Scaled Dot-Product Attention & Transformer Layers; this page will reuse it instead of restarting from zero.

Work hereEfficient 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.

Carry outLLM Serving at Scale: Prefill, Decode & Continuous Batching

The next edge should feel earned: use the demo prediction here before following LLM Serving at Scale: Prefill, Decode & Continuous Batching.

Test the linkManipulate one control and predict the visible change.Then continue to LLM Serving at Scale: Prefill, Decode & Continuous Batching

01IntuitionStart with the picture, metaphor, or geometric mechanism.02MathMake the objects explicit and connect them with notation.03CodeMirror the equations with runnable implementation details.04Interactive DemoManipulate the mechanism and watch the idea respond.

Intuition

Build the mental picture first so the rest of the page has something to attach to.

Section prompt

"Attention is $O(T^2)$ " is true on paper, but in production the pain often feels different:

During decoding, you only compute attention for the new token, but you must attend over all past tokens.
You therefore store past keys/values in a KV cache. At long context, the KV cache can dominate memory.

Modern serving is mostly a game of managing memory bandwidth and cache size:

KV cache avoids recomputing keys/values for all past tokens.
Grouped-query attention (GQA) reduces KV cache size by sharing keys/values across multiple query heads.
FlashAttention keeps the exact attention math but changes how it is computed so the attention matrix is never materialized in slow memory.

Math

Translate the story into symbols, assumptions, and a derivation you can inspect.

Section prompt

Equation 1o_t = \mathrm{softmax}\!\left(\frac{q_t K_{1:t}^{\top}}{\sqrt{d_k}}\right) V_{1:t},Equation 2\mathrm{Mem}_{KV} \approx B\cdot L\cdot T\cdot H_{kv}\cdot d_{\mathrm{head}}\cdot 2 \cdot \ma...

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.

KV cache memory scaling

Across a batch of $B$ sequences and $L$ layers, the KV cache stores keys and values for all $T$ positions:

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

The factor of 2 is for storing both $K$ and $V$ .

Grouped-query attention (GQA)

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)}.

This reduces the KV cache to roughly $H_{kv}/H_q$ of full multi-head attention. Equivalently: the KV cache becomes about $(H_{kv}/H_q)$ of full multi-head attention (so it’s reduced by roughly $H_q/H_{kv}\times$ when $H_{kv} < H_q$ ).

For the dedicated head-sharing invariant, selected Q-to-KV mapping, and cache-ratio prediction exercise, see Grouped-Query Attention.

Code

Keep the implementation aligned with the notation so the algorithm is legible.

Section prompt

Code witness 1def kv_cache_gb(T, layers, h_kv, d_head, batch=1, bytes_per_elem=2): elems = batch * layers *...python

def kv_cache_gb(T, layers, h_kv, d_head, batch=1, bytes_per_elem=2):
    elems = batch * layers * T * h_kv * d_head * 2  # keys + values
    return elems * bytes_per_elem / 1e9

T, layers, d_head, batch = 128_000, 80, 128, 1
full = kv_cache_gb(T, layers, h_kv=64, d_head=d_head, batch=batch)

for h_q, h_kv in [(64, 64), (64, 8), (64, 1)]:
    gb = kv_cache_gb(T, layers, h_kv=h_kv, d_head=d_head, batch=batch)  # fp16
    print(f"Hq={h_q:>2} Hkv={h_kv:>2}  KV ~ {gb:6.1f} GB  ({gb/full:.1%} of full)")

Interactive Demo

Use direct manipulation to connect the explanation to a moving system.

Section prompt

Use the demo to explore how KV cache size grows with context length, and how GQA changes the memory/latency tradeoff without changing the basic attention mechanism.

Live Concept Demo

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

The stage is code-native and interactive. Use it to test the explanation against the mechanism.

difficulty 4/5graduatecode-aligned

Demo Prediction Checkpoint

Manipulate one control and predict the visible change.

Commit to what Efficient Attention at Scale: KV Cache, GQA & FlashAttention should make visible before reading the result.

After The First Pass

Turn the concept into an inspected object.

Once the invariant is visible in the intuition, math, code, and demo, use these panels to inspect the mechanism visually, check source support, practice the idea, and attach a grounded research question.

Mechanism Storyboard

See the idea move before the page explains it

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

Prediction open01 / Intuition

Prediction lens

Start with the picture, metaphor, or geometric mechanism.

Commit first

Before reading further, choose the kind of change Efficient Attention at Scale: KV Cache, GQA & FlashAttention should make visible.

Visual Inquiry

Make the image answer a mathematical question

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

4/4 stages readyLive demo connected

Visual cueWhich visible object should carry the first intuition?

Inspection depth2/4

Prediction

Which visible object should carry the first intuition?

Commit first

Pick the cue that should make Efficient Attention at Scale: KV Cache, GQA & FlashAttention easier to reason about before the page gives the answer.

Source Grounding

Canonical references for the mechanism on this page.

paper · 2019Fast Transformer Decoding: One Write-Head is All You NeedShazeer

Grounds the KV-cache bandwidth motivation for sharing keys and values across query heads during decoding.

Open source

paper · 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie et al.

Introduces grouped-query attention as the practical middle point between multi-head quality and MQA memory savings.

Open source

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

Grounds the IO-aware attention kernel: tiling plus online softmax avoids materializing the full attention matrix.

Open source

Claim Review

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

Status1 substantive review recorded

Claims without a substantive review badge still need exact source-support review.

Sources3 references

shazeer-2019-mqa, ainslie-2023-gqa, dao-2022-flashattention

Witnesses4 local objects

Use equation, code, and demo objects to check whether the source support is operational.

Source support candidates

paper 2019Fast Transformer Decoding: One Write-Head is All You Need

Grounds the KV-cache bandwidth motivation for sharing keys and values across query heads during decoding.

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

Introduces grouped-query attention as the practical middle point between multi-head quality and MQA memory savings.

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

Grounds the IO-aware attention kernel: tiling plus online softmax avoids materializing the full attention matrix.

Mechanism witnesses

Equation 1

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

Equation 2

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

Code witness 1def kv_cache_gb(T, layers, h_kv, d_head, batch=1, bytes_per_elem=2): elems = batch * layers *...Demo stateLive mechanism probe

Practice Loop

Try the idea before it explains itself

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

Readiness0/3 checks ready

Predict

Before touching the demo, predict one visible change that should happen in Efficient Attention at Scale: KV Cache, GQA & FlashAttention.

Hint 1

Reveal when your model needs a nudge.

Hint 2

Reveal when your model needs a nudge.

Hint 3

Reveal when your model needs a nudge.

Your answer canvas

Local checks

Claim

A concrete answer is on the canvas.

Mechanism

The answer names why the claim should hold.

Bridge

It touches the page context or a neighboring idea.

Misconception check

Object research drawerClose

ConceptEfficient Attention at Scale: KV Cache, GQA & FlashAttentionAttention & Transformers

Code witness comparisonEfficient Attention at Scale: KV Cache, GQA & FlashAttention code witness 1def kv_cache_gb(T, layers, h_kv, d_head, batch=1, bytes_per_elem=2): elems = batch * layers *...Prediction before revealEfficient Attention at Scale: KV Cache, GQA & FlashAttention interactive demoManipulate one control and predict the visible change.

Grounded room questionWhat is the smallest example that makes Efficient Attention at Scale: KV Cache, GQA & FlashAttention click without losing the math?Local snapshot ready

Research Room

Attach the question to an exact object

Pick the concept, equation, source, code witness, claim, misconception, or demo state before asking for help. The handoff stays grounded to that object.

Next local actionNo local draft saved yet

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

conceptAttention & Transformers

Efficient Attention at Scale: KV Cache, GQA & FlashAttention

Anchored question

What is the smallest example that makes Efficient Attention at Scale: KV Cache, GQA & FlashAttention click without losing the math?

Local action draftNo local draft saved yetExpand only when ready to capture one local next action

Local action draft

This draft stays locally in this browser for concept:attention-transformers/efficient-attention.

Draft noteNext action

No local draft saved.

Evidence to inspect

Source ids to inspect: shazeer-2019-mqa, ainslie-2023-gqa, dao-2022-flashattention
Definition, prerequisite, and contrast concept links
The equation or code witness that makes the concept operational
One demo state that shows the invariant instead of a slogan

What would resolve this

The learner can state the mechanism in their own words
The learner can name the prerequisite that would repair confusion
The learner can predict how the mechanism changes under one perturbation

Grounded AI handoff

I am working in Continuous Function's research reading room. Object: concept - Efficient Attention at Scale: KV Cache, GQA & FlashAttention Object key: concept:attention-transformers/efficient-attention Context: Attention & Transformers Anchor id: concept/concept-notebook/attention-transformers/efficient-attention Open question: What is the smallest example that makes Efficient Attention at Scale: KV Cache, GQA & FlashAttention click without losing the math? Evidence to inspect: - Source ids to inspect: shazeer-2019-mqa, ainslie-2023-gqa, dao-2022-flashattention - Definition, prerequisite, and contrast concept links - The equation or code witness that makes the concept operational - One demo state that shows the invariant instead of a slogan What would resolve this: - The learner can state the mechanism in their own words - The learner can name the prerequisite that would repair confusion - The learner can predict how the mechanism changes under one perturbation Answer as a careful research tutor: stay source-grounded, separate verified evidence from assumptions, name the relevant math objects, and end with one next action.

Open source object

concept/concept-notebook/attention-transformers/efficient-attention
concept:attention-transformers/efficient-attention

Learning Map

Before / Now / Try / Next

BeforeScaled Dot-Product Attention & Transformer Layers

NowIntuition → Math → Code → Demo

TryManipulate one control and predict the visible change.

NextLLM Serving at Scale: Prefill, Decode & Continuous Batching

Intuitionready
Mathready
Codeready
Interactive Demoready

Object Companion

Ask beside the selected object

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

Your question

GoalComfortStyleStuck on

Context prompt

You are my AI learning companion for Continuous Function. Current context: Attention & Transformers concept. Learning surface: Efficient Attention at Scale: KV Cache, GQA & FlashAttention. What this page says: How attention becomes practical at long context: KV caching for decoding, grouped-query attention, and IO-aware kernels like FlashAttention. Current section: Intuition, math, code, and interactive demo. Suggested next step: Manipulate one control and predict the visible change.. Learner goal: Understand the idea. Learner comfort level: New to this. Preferred explanation style: Visual first. Task: Explain the central idea in plain language, then restate it with the exact math objects from the page. Answer in a way that helps me learn: ask one clarifying question only if needed, use intuition before notation, and end with one thing I should try on the page.

Domain

Attention & Transformers

attention-transformersefficient-attentionkv-cacheflashattentiongqa

Prerequisites

Scaled Dot-Product Attention & Transformer Layers Rotary Position Embeddings (RoPE)Efficiency: Quantization, Distillation, LoRA & Sparse MoE

Leads To

LLM Serving at Scale: Prefill, Decode & Continuous Batching Speculative Decoding: Lossless Multi-Token Generation Long Context Engineering: RoPE Scaling, KV Compression & Memory Optimization Grouped-Query Attention: Sharing KV Heads

LLM Serving at Scale: Prefill, Decode & Continuous Batching Long Context Engineering: RoPE Scaling, KV Compression & Memory Optimization Speculative Decoding: Lossless Multi-Token Generation Grouped-Query Attention: Sharing KV Heads

Within this domain