← Back to Table of Contents

Chapter 17 — KV-Cache: Mechanics & Memory

“The KV-cache is the single most important optimization in LLM inference — without it, generating each token would require recomputing attention over the entire sequence from scratch.”

Why KV-Cache Exists

During autoregressive generation, each new token needs to attend to all previous tokens. Without caching, we’d recompute Q, K, V for every previous token at every step — O(T²) total computation to generate T tokens.

The KV-cache stores the Key and Value projections from all previous tokens, so each decode step only computes Q, K, V for the new token and reuses everything else.

Step-by-Step Visualization

At each decode step, the computation is:

• Query: only the new token → [B, H, 1, d_k]
• Key/Value: read from cache → [B, H, T_total, d_k]
• Attention: Q × K^T → [B, H, 1, T_total] → softmax → × V → [B, H, 1, d_k]

This is a matrix-vector product (not matrix-matrix), making decode steps memory-bandwidth bound rather than compute-bound.

Tensor Shapes in Detail

Memory Calculation

The KV-cache memory for the entire model:

where: L = layers, G = KV heads, d_k = head dimension, T = sequence length, B = batch size, bytes = 2 (FP16/BF16).

Worked Examples

Key insight: The KV-cache for LLaMA-3.1 405B at 128K context exceeds 63 GB per request in FP16. This is why KV-cache optimization (Chapter 18) is critical.

GQA/MQA Impact on KV-Cache

The attention variant directly determines KV-cache size (see Chapter 5):

Sliding Window KV-Cache

With Sliding Window Attention (Chapter 5), the KV-cache is a fixed-size rolling buffer:

Inspecting KV-Cache in Transformers

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B", torch_dtype=torch.bfloat16, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

input_ids = tokenizer("Hello, world!", return_tensors="pt").input_ids.to(model.device)

# First forward pass (prefill)
outputs = model(input_ids, use_cache=True)
past_kv = outputs.past_key_values

# Inspect shapes
for layer_idx, (k, v) in enumerate(past_kv):
    if layer_idx == 0:
        print(f"Layer {layer_idx}:")
        print(f"  K shape: {k.shape}")  # [1, 8, T, 128] for GQA with 8 KV heads
        print(f"  V shape: {v.shape}")  # [1, 8, T, 128]

# Total KV-cache memory
total_bytes = sum(k.nbytes + v.nbytes for k, v in past_kv)
print(f"Total KV-cache: {total_bytes / 1024**2:.1f} MB")

# Decode one more token — only feed the last token
next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1, keepdim=True)
outputs2 = model(next_token, past_key_values=past_kv, use_cache=True)
past_kv2 = outputs2.past_key_values

# Cache grew by 1 token
k0_old, k0_new = past_kv[0][0], past_kv2[0][0]
print(f"Cache grew: {k0_old.shape[2]} → {k0_new.shape[2]} tokens")

KV-Cache vs Model Weights Memory

At long context lengths, KV-cache can exceed the model weights:

At 128K context, the KV-cache equals the model size. At 512K, it’s 4× the model. This is why context extension (Chapter 12) and KV-cache optimization (Chapter 18) go hand in hand.

Batching and KV-Cache

When serving multiple requests, each request has its own KV-cache of potentially different length. Managing memory across a dynamic batch is non-trivial — this is the core problem that PagedAttention and other optimizations solve (see Chapter 18).

What’s Next

The KV-cache memory problem is the central challenge of LLM serving. The next chapter covers the full landscape of KV-cache optimization — from PagedAttention to quantized KV-cache to token eviction strategies.

← Previous: Chapter 16 — Inference & Sampling · Next: Chapter 18 — KV-Cache — Optimization →

Last updated: April 2026