TechByteByByte

Transformer Inference and KV Cache

Understand why autoregressive inference recomputes redundant work without caching, how the KV cache eliminates it, and the prefill/decode phases — with a verified proof of both correctness and a real efficiency gain.

#Transformers#KV Cache#Inference#Prefill#Decode#AI#LLM

Begin with the central question

Why would generation be painfully repetitive without a memory of earlier attention calculations?

Essential words

Inference uses fixed learned parameters. A KV cache stores past keys and values per layer. Prefill processes the prompt; decode adds new tokens one at a time.

What You Will Understand

Why naive autoregressive inference wastes enormous amounts of redundant computation, and how the KV cache eliminates it — with a direct, verified proof that caching produces mathematically identical results to recomputation, while doing meaningfully less work. This is one of the most practically important concepts in this course for real-world LLM serving.

prompt prefill -> cache K/V -> decode new query -> append new K/V -> repeat

The problem this module solves

Module 14 traced generation as: run the full forward pass, get a token, append it, repeat. Taken literally, this means recomputing attention for every previous token at every single generation step — the same Key and Value projections, computed over and over, for tokens that haven’t changed. The KV cache exists to eliminate this waste entirely.

Build the intuition

imagine re-reading an entire book from page 1 every time you want to write the next sentence of your own commentary on it — even though the book’s content never changes between commentary sentences. The KV cache is simply remembering what you already read, instead of re-reading it every single time.

4. Real-World Analogy

Think of a librarian who, every time a new patron asks a question, would have to re-catalog the entire library from scratch before answering — even though the books haven’t moved since the last question. Obviously wasteful: a sensible librarian catalogs the library once, and reuses that catalog for every subsequent question.

The KV cache is exactly this catalog — computed once per token, reused for every future generation step.

Analogy: The Repeating Library Cataloging vs. Stored Index Sheets Think of running generation steps with and without a KV cache:

  • Without a Cache (Redundant Cataloging): Every time you write a new page of notes about a book, you must read the book again from page 1 to make sure your summary flows cleanly. If the book is 100 pages, you read page 1 exactly 100 times, page 2 exactly 99 times, and so on (quadratic execution waste).
  • With a Cache (Stored Index Sheets): As you read each page the first time, you write key tags and summary points on index cards (Key/Value vectors).
  • When writing page 101’s commentary, you do not touch pages 1-100 of the book. You simply pull cards 1-100 from your folder (the KV cache), read page 101 once, write card 101, and append it. You save hours of re-reading.

📊 Visual Flowchart: KV Cache Prefill vs. Decode Lifecycles

Here is how keys and values are stored during the prompt step and appended during decode steps:

graph TD
    classDef prefill fill:#3498db,stroke:#333,stroke-width:1px,color:#fff;
    classDef decode fill:#2ecc71,stroke:#333,stroke-width:1px,color:#fff;
    classDef storage fill:#f1c40f,stroke:#333,stroke-width:1px,color:#fff;




    subgraph Prefill ["1. Prefill Phase (Prompt Step)"]
        Prompt["Input Prompt:<br>'Explain RAG'"] --> ParallelKV["Calculate K, V for all tokens in parallel"]:::prefill
        ParallelKV --> StoreCache["Save K, V to Memory"]:::storage
    end




    subgraph Decode ["2. Decode Phase (Incremental Token Steps)"]
        StoreCache --> NextQuery["Calculate Q for Newest Token ONLY"]:::decode
        GetCache["Retrieve History K, V from Cache"]:::storage --> AttentionCalc["Compute Attention: Q_new @ [K_cached + K_new]"]
        NextQuery --> AttentionCalc




        NewKV["Calculate K, V for Newest Token ONLY"]:::decode --> UpdateCache["Append New K, V to Cache Store"]:::storage




        AttentionCalc --> GenerateToken["Output Logits -> Select next token"]:::decode
    end

5. Core Concept

Prompt
  ↓
PREFILL: process the entire prompt in ONE forward pass,
         computing and CACHING every token's Key and Value
         vectors at every layer
  ↓
KV Cache (stores every prompt token's K, V -- ready to reuse)
  ↓
DECODE: generate one token at a time. At each step:
         - compute Q, K, V for ONLY the newest token
         - retrieve ALL previous tokens' K, V from the CACHE
           (no recomputation)
         - compute attention using the cached K/V plus the
           new token's K/V
  ↓
Update cache with the new token's K, V
  ↓
Generate next token
  ↓
...
TermDefinition
KV cacheStored Key and Value vectors (per layer, per token) from all previously processed tokens, reused instead of recomputed
PrefillThe initial forward pass over the full prompt, populating the cache
DecodeEach subsequent step generating one new token, using the cache

6. How It Works — Step by Step

WITHOUT KV cache (naive, wasteful):
  Step 1: compute K, V for token 1                      (1 computation)
  Step 2: compute K, V for tokens 1, 2                    (2 computations
                                                          -- token 1's K,V
                                                          recomputed!)
  Step 3: compute K, V for tokens 1, 2, 3                 (3 computations
                                                          -- tokens 1,2
                                                          recomputed again!)
  ... total work grows QUADRATICALLY with sequence length




WITH KV cache:
  Step 1: compute K, V for token 1, STORE in cache          (1 computation)
  Step 2: compute K, V for token 2 ONLY, APPEND to cache     (1 computation)
  Step 3: compute K, V for token 3 ONLY, APPEND to cache     (1 computation)
  ... total work grows LINEARLY with sequence length

Crucially: Query is never cached — only Key and Value. A new token’s Query only ever needs to be computed once, for that token’s own generation step, and is never reused later (unlike K/V, which every future token’s attention computation needs to reference).

7. Mathematical Intuition

Read the mathematics as a story

During causal generation, old token Keys and Values do not change. A KV cache stores them so the model only projects the new token and combines it with cached history.

prefill prompt -> cache K₁..Kₙ and V₁..Vₙ
new token -> new Q,K,V + cached K/V -> next token

Total K/V computations without caching, across generating n tokens: 1 + 2 + 3 + ... + n = n(n+1)/2 — quadratic growth. With caching: exactly n — linear growth. For a modest 5-token example (verified below), this is 15 vs. 5 computations — a 3x reduction; for a realistic 1000-token generation, it would be 500,500 vs. 1,000 — a roughly 500x reduction.

This quadratic-vs-linear gap is precisely why the KV cache matters so much more as sequences get longer.

8. Small Worked Example

Walk through the example

  1. Generate five positions without caching and count repeated K/V work. 2. Repeat with caching. 3. Compare 15 computations with 5 in the toy count.

Generating a 5-token sequence: without caching, by the 5th step, you’ve recomputed token 1’s Key and Value five separate times — once at every step, even though token 1 never changes after it’s first processed. With caching, token 1’s Key and Value are computed exactly once, then simply retrieved (not recomputed) at every later step.

9. Python / NumPy Example

What the code will demonstrate

This small NumPy example makes Transformer Inference and KV Cache visible with inspectable numbers and shapes. Read it in three passes: identify each input, follow the transformation line by line, and connect the printed output to the diagram above. The arrays are intentionally tiny teaching values; unless the text explicitly says otherwise, they are not weights or measurements from GPT, Gemini, or another trained model.

# The arrays are intentionally small so each transformation can be inspected.
# Printed values illustrate the mechanism; they are not trained-model measurements.
import numpy as np




np.random.seed(16)
d_model = 4




Wq = np.round(np.random.randn(d_model, d_model) * 0.4, 2)
Wk = np.round(np.random.randn(d_model, d_model) * 0.4, 2)
Wv = np.round(np.random.randn(d_model, d_model) * 0.4, 2)




# Simulate a growing sequence: tokens generated one at a time
all_tokens = np.round(np.random.randn(5, d_model) * 0.5, 3)




# --- WITHOUT KV cache: recompute K, V for the ENTIRE sequence at EVERY step ---
kv_computations_no_cache = 0
for step in range(1, len(all_tokens) + 1):
    current_seq = all_tokens[:step]
    K = current_seq @ Wk
    V = current_seq @ Wv
    kv_computations_no_cache += step




print(f"WITHOUT KV cache: total K/V projections computed = {kv_computations_no_cache}")




# --- WITH KV cache: compute K, V ONCE per token, reuse for all future steps ---
kv_cache_K, kv_cache_V = [], []
kv_computations_with_cache = 0
for step in range(1, len(all_tokens) + 1):
    new_token = all_tokens[step - 1:step]
    new_K = new_token @ Wk
    new_V = new_token @ Wv
    kv_cache_K.append(new_K)
    kv_cache_V.append(new_V)
    kv_computations_with_cache += 1




    full_K = np.vstack(kv_cache_K)
    full_V = np.vstack(kv_cache_V)




print(f"WITH KV cache: total K/V projections computed = {kv_computations_with_cache}")




# --- Verify correctness: cached K, V match fresh computation on the full sequence ---
K_fresh = all_tokens @ Wk
V_fresh = all_tokens @ Wv
print("\nDo cached K match freshly-computed K?", np.allclose(full_K, K_fresh))
print("Do cached V match freshly-computed V?", np.allclose(full_V, V_fresh))




print(f"\nRedundant computations avoided: {kv_computations_no_cache - kv_computations_with_cache}")
print(f"Reduction factor: {kv_computations_no_cache / kv_computations_with_cache:.1f}x")

Expected Output:

WITHOUT KV cache: total K/V projections computed = 15
WITH KV cache: total K/V projections computed = 5




Do cached K match freshly-computed K? True
Do cached V match freshly-computed V? True




Redundant computations avoided: 10
Reduction factor: 3.0x

10. How It Works

  • Correctness is proven, not assumed: the cached K and V — built incrementally, one token at a time, never recomputing earlier tokens — are exactly identical (True, via np.allclose) to computing K and V fresh on the entire sequence at once. Caching isn’t an approximation; it’s a genuine optimization that changes how much work is done, not what the result is.
  • The efficiency gain is concrete: 15 computations without caching vs. 5 with — a 3.0x reduction for just 5 tokens. As derived mathematically in Section 7, this gap grows quadratically worse without caching as sequences get longer — a 1000-token generation would see roughly a 500x difference.

11. How Is This Used in Modern AI?

Where this concept lives

Follow the concept at three levels: inside the model, where the computation happens; inside the AI product, where that computation supports a visible feature; and inside production, where engineers measure speed, memory, quality, and failure cases. The details below connect those levels.

🤖 How Is This Used in Modern AI?

Every production LLM serving system uses KV caching — it’s not an optional optimization but a practical necessity for serving generation requests at reasonable latency and cost. Without it, generating long responses would become prohibitively slow as context grows, exactly matching the quadratic-vs-linear gap demonstrated above.

ConceptAI application
KV cacheStandard, essential component of every real LLM inference server
PrefillThe (often relatively fast, highly parallel) initial processing of your prompt
DecodeThe (comparatively slower, one-token-at-a-time) generation phase — where KV cache savings matter most

12. How Is This Used in LLMs?

Trace one model call

User text → tokens → Transformer computation → output-token probabilities
          this topic affects one part of that computation

An LLM does not apply this idea as a separate magic step. It uses it as part of the repeated numerical pipeline that transforms token vectors and produces the next-token probabilities.

The KV cache directly explains a real, observable phenomenon: the first token of an LLM’s response (after prefill) often has different latency characteristics than each subsequent token (decode) — prefill processes the whole prompt in one parallelizable pass, while decode generates tokens one at a time, each step benefiting from (and also being partially bottlenecked by) the growing KV cache.

Real systems you can recognize

Hugging Face explains that KV caching avoids recomputing past keys and values and is intended for inference, not training; see cache explanation. Gemini separately offers context caching for repeated request prefixes, which is an API optimization and should not be confused with an internal KV cache.

13. How Is This Used in Agentic AI?

Trace one agent step

Goal + history + tool results
              ↓
     LLM processes the context
              ↓
Suggested answer or tool call
              ↓
Agent runtime validates and executes it

This distinction matters: the Transformer helps produce the proposal, while the surrounding agent software controls tools, permissions, retries, memory, and execution.

Direct relevance to Agentic AI: Very High. Long agent conversations — with extensive system prompts, conversation history, and repeated tool results — directly stress-test KV cache size and management.

This is precisely why agent applications with very long, growing contexts incur real, increasing serving costs and latency over a conversation’s lifetime — a direct, practical consequence of what this module demonstrated.

When this knowledge is useful

Use Transformer Inference and KV Cache when you need to explain, implement, debug, evaluate, or optimize the corresponding part of a Transformer pipeline. It is also useful when a model API behaves unexpectedly and you need to trace the behavior back to tokens, tensor shapes, attention visibility, training, or inference mechanics.

When it is not enough

Understanding this mechanism does not by itself prove that a complete model or application is accurate, safe, fast, or cost-effective. Production decisions still require representative evaluation data, latency and memory measurements, model-specific documentation, and tests of the surrounding retrieval or agent code.

14. Common Beginner Mistakes

⚠️ Mistake

Incorrect idea: assuming Query vectors are cached too.

Why it is incorrect: Only Key and Value are cached — a new token’s Query is computed fresh at its own generation step and never reused afterward, since future tokens need their own Query, not a past token’s.

⚠️ Mistake

Incorrect idea: thinking KV caching changes the model’s output.

Why it is incorrect: As proven directly, cached and freshly-recomputed K/V are mathematically identical — caching is purely a computational efficiency technique, with zero effect on correctness or the model’s actual predictions.

⚠️ Mistake

Incorrect idea: believing KV cache size is free or unlimited.

Why it is incorrect: The cache must store K and V for every token, at every layer, at every attention head — this consumes real, often substantial GPU memory, growing linearly with context length (Module 17 covers this cost directly).

15. Important Distinctions

Without KV CacheWith KV Cache
Recomputes ALL previous tokens’ K/V at every stepComputes each token’s K/V exactly ONCE
Total work grows QUADRATICALLY with sequence lengthTotal work grows LINEARLY with sequence length
Verified: 15 computations for 5 tokensVerified: 5 computations for 5 tokens
PrefillDecode
Processes the ENTIRE prompt in one passGenerates ONE token at a time
Highly parallelizableInherently sequential, one step per token

16. Production / Engineering Considerations

  • KV cache memory scales with context length, number of layers, and number of attention heads — this is a genuine, often dominant memory cost in LLM serving, directly motivating techniques like grouped-query and multi-query attention (Module 17), which reduce the number of independent K/V projections needing caching.
  • Longer conversations mean larger caches — a real, growing memory and cost factor across a long agent conversation, not a fixed, one-time cost.
  • Cache eviction/management becomes a genuine engineering concern at scale — serving many concurrent users each with their own growing KV cache requires careful memory management infrastructure.

17. Interview Questions

Beginner

Q: What problem does the KV cache solve?

Ans: Without it, generating each new token in autoregressive inference would require recomputing Key and Value vectors for every previous token in the sequence — even though those tokens haven’t changed. The KV cache stores each token’s Key and Value once, computed at the step it was first processed, and reuses them for every subsequent generation step, avoiding this redundant recomputation entirely.

Intermediate

Q: Why is only Key and Value cached, not Query?

Ans: Each new token’s Query is specific to that token’s own generation step — it’s used once, to compute that step’s attention, and never needed again by future steps. Key and Value, by contrast, represent “what information does this token offer,” which every future token’s attention computation needs to reference — making them worth storing and reusing, while Query has no such reuse value.

Advanced

Q: Explain, with reference to computational complexity, why KV caching matters more as sequence length grows.

Ans: Without caching, generating n tokens requires recomputing K/V for increasingly large prefixes at each step, totaling 1 + 2 + ... + n = n(n+1)/2 computations — quadratic growth in n. With caching, exactly n computations are needed — linear growth. This gap, demonstrated directly in this module (15 vs.

5 for a 5-token example, a 3x difference), widens dramatically for longer sequences — a 1000-token generation would see roughly a 500x difference between cached and uncached approaches, which is precisely why KV caching is not optional for any practically-sized LLM serving system.

Scenario

Q: A user has an extremely long conversation with an AI agent, spanning thousands of tokens of history. What practical serving consequence does this module predict, and why?

Ans: The KV cache for this conversation grows linearly with the conversation’s total token count, across every layer and attention head of the model — consuming increasing amounts of GPU memory the longer the conversation continues.

This is a genuine, predictable resource cost (not a hypothetical concern), directly motivating context-length limits and efficiency techniques (Module 17) in real production LLM serving systems.

Architecture

Q: Does using a KV cache change what tokens a model generates, compared to not using one?

Ans: No — as verified directly in this module, cached K/V values are mathematically identical to freshly recomputing them on the full sequence. The KV cache is purely a computational efficiency optimization that reduces redundant work; it has zero effect on the actual attention computation’s result or the model’s generated output.

Engineering

Q: Why might a team choose to invest in KV cache memory optimization techniques (like the ones previewed for Module 17) rather than simply adding more GPU memory?

Ans: KV cache memory usage scales directly with context length, number of layers, and number of attention heads/K-V projections — for large-context, high-throughput serving, this can become a dominant, rapidly-growing cost.

While adding more memory is one lever, techniques that reduce the amount of K/V data that needs caching in the first place (without sacrificing model quality) provide a more scalable, cost-effective solution, especially at the scale of serving many concurrent users with long contexts simultaneously — exactly the motivation behind grouped-query and multi-query attention.

18. What You Should Remember

  • The KV cache stores each token’s Key and Value (never Query), computed once, reused for every future generation step.
  • This reduces total computation from quadratic to linear in sequence length — verified directly: 15 vs. 5 computations for a 5-token example, a proven 3x reduction that grows far larger for longer sequences.
  • Caching is provably lossless — cached K/V values are mathematically identical to freshly recomputed ones, verified directly with np.allclose.
  • Prefill (processing the full prompt) and decode (generating one token at a time) are the two distinct phases this cache spans.

19. How This Helps Me Build AI Systems

You now understand, with direct proof, exactly why LLM API latency and cost behave the way they do — why longer contexts and longer generations cost more, and why this cost is a genuine architectural consequence, not an arbitrary pricing decision. This is essential context for reasoning about agent system costs, especially for long-running conversations.

Next: Module 17 — Transformer Scaling, Context Windows and Efficiency — the broader practical engineering picture: why inference is expensive, and what techniques (FlashAttention, GQA, MQA) reduce that cost.

Author
TechByteByByte Editorial Team
Reviewed by
TechByteByByte Admin
Published
Last reviewed