algorithms.live
    AI-ENGINEERING
    advanced

    KV Cache & Attention Optimization

    AI Engineer path · lesson 104 of 124
    The idea
    Reading a book aloud, where before every word you re-read every page so far. The KV cache is the margin notes you make as you go: one note per page, written once, glanced at thereafter. The notes make the reading fast — and eventually the notes outweigh the book.
    How to use
    Use the calculator to reproduce the book's Llama 2 13B example (54 GB), then drop bytes per value from 2 to 1 and watch it halve.
    At each generation step the model attends over every previous token. Those tokens' key and value vectors were already computed and will never change — so recomputing them is pure waste.
    click to toggle
    step 11 computedstep 21 computedstep 31 computedcomputed nowread from cache
    3
    1
    projections at the last step
    3
    projections in total
    With the cache, each step projects exactly one new token and reads the rest back. Total work grows linearly with the output length.
    Why training does not need one
    During training every token in the sequence is already known, so the whole thing is computed in one pass — there is no step-by-step dependency to cache across. A KV cache is used only during inference. This is also why a model that trains comfortably can be surprisingly hard to serve: the memory profile of generation is not the memory profile it was tuned on.
    The trade you just made
    The cache converts a compute problem into a memory problem, and memory is the scarcer resource. Attention computation grows exponentially with sequence length while the cache grows linearly — so the cache is unambiguously the right trade, right up until it does not fit, at which point your maximum batch size collapses and throughput goes with it.

    Key insight

    The cache scales with batch size as well as sequence length, which is what makes it a serving problem rather than a per-request one. Every extra concurrent user costs you another full sequence of keys and values, so the cache — not the weights — is usually what caps how many people you can serve at once.

    Why the cache exists

    Sequence

    Order what happens on each decoding step with a KV cache.

    Tap a piece, then tap where it goes — or drag it. Tap the tray to take one back.

    1. 1Drop a step here
    2. 2Drop a step here
    3. 3Drop a step here
    4. 4Drop a step here
    5. 5Drop a step here
    5 pieces left to place

    Test your knowledge

    20 questions
    Question 1 of 20
    Score: 0

    The KV cache exists because, across generation steps, previous tokens' key and value vectors:

    The short version

    Reading a book aloud, where before every word you re-read every page so far. The KV cache is the margin notes you make as you go: one note per page, written once, glanced at thereafter. The notes make the reading fast — and eventually the notes outweigh the book.

    In more depth

    KV Cache & Attention Optimization

    Generating a token means attending to every token before it — and those tokens' key and value vectors never change. The KV cache stores them, turning quadratic recomputation into linear.

    The trade

    Attention computation grows exponentially with sequence length; the KV cache grows linearly. So the cache is the right trade — until it does not fit.
    A KV cache is used only during inference, not training. In training all tokens are known in advance, so generation is computed in one pass.
    It scales with batch size too, which makes it a serving problem: the cache, not the weights, usually caps concurrency.

    The formula

    KV cache bytes = 2 × B × S × L × H × M

    B batch · S sequence length · L layers · H model dimension · M bytes per value. The leading 2 is because you store both a key and a value.

    Worked example — Llama 2 13B (40 layers, dimension 5,120), batch 32, sequence 2,048, 2 bytes:

    2 × 32 × 2,048 × 40 × 5,120 × 2 = **54 GB**

    The number that makes the point: Pope et al. (2022) calculated 3 TB for a 500B+ model with multi-head attention at batch 512, context 2,048 — three times the size of that model's weights.

    Three buckets of optimization

    1. Redesign the mechanism (training/finetuning only — it changes the architecture):

    TechniqueWhoCuts
    Multi-query attentionShazeer, 2019Shares KV across all query heads
    Grouped-query attentionAinslie et al., 2023Groups of heads share KV — the dial between MQA and MHA
    Cross-layer attentionBrandon et al., 2024Shares KV across adjacent layers (3 layers ⇒ 3× smaller)
    Local windowed attentionBeltagy et al., 20201,000-token window on 10,000-token sequences ⇒ 10×

    Character.AI (2024): average conversation carries 180 messages; combining multi-query, interleaved local/global and cross-layer attention cut their KV cache over 20×. They compose because each cuts a different dimension.

    2. Optimise the cache (inference time): PagedAttention (Kwon et al., 2023, vLLM) — non-contiguous blocks, removing fragmentation and enabling prefix sharing; KV cache quantization; adaptive compression; selective KV cache.

    3. Write kernels: FlashAttention (Dao et al., 2022) fuses operations so intermediates stay in SRAM. Exact, not approximate. FlashAttention-3 followed for H100s.

    Interview Questions

    NVIDIA
    Anthropic
    Meta
    Google
    Anthropic
    OpenAI
    Google
    Character.AI
    Meta
    Google
    Mistral
    Anthropic
    Databricks
    Amazon

    References & Further Reading

    LLM Inference Engines: vLLM, KV Cache, PagedAttention, BatchingYouTube
    video
    Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)arXiv
    paper
    GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)arXiv
    paper
    phoenixtype/ai-engineering-interview-prepGitHub
    docs
    All AI Engineering topics

    Love this visualization?

    Sign up to track your progress, earn achievements, use the AI assistant, and access your personal learning dashboard.

    100% free • No credit card required