KV Cache — Module 12.3

How This Lesson Fits the Module & Volume

Volume 11 introduced the KV cache as the core trick that makes autoregressive decoding incremental: store past keys and values so each new token only computes its own K/V. Module 12.3 elevates that idea into a serving systems concern—multi-request GPUs, long contexts, and memory managers like paged attention in vLLM.

This lecture opens Inference Optimization. Everything that follows—Flash Attention, speculative decoding, prefix cache, and batching strategies—assumes you understand how KV memory dominates decode throughput.

Learning Objectives

By the end of this lesson, students should be able to:

Recall prefill vs decode and why the KV cache is essential for serving.
Estimate KV memory from layers, heads, head dim, sequence length, dtype, and batch.
Explain paged / block KV layouts used by modern engines (vLLM-style).
Distinguish contiguous per-request caches from shared block tables.
Relate cache pressure to batch size, context length, and concurrent users.
Sketch how a serving loop appends, evicts, and reuses KV blocks.

Definition

A KV cache is the per-layer store of key and value tensors for tokens already processed. In a serving system it is not just a tensor concatenate—it is a managed memory pool (often paged into fixed-size blocks) that must be allocated, referenced by block tables, shared across prefix hits, and freed when requests finish.

From Single Session to Multi-Tenant Serving

In Vol. 11 you cached one sequence. In production, dozens of conversations share one GPU. Each request needs its own growing KV footprint; naive dense allocation leaves huge holes when sequences finish at different times. That fragmentation is why engines moved to paged attention: split KV into fixed blocks and map logical positions to physical blocks via a block table.

1. Prefill

Write prompt KV into blocks

2. Decode

Append one token’s K/V

3. Schedule

Batch ready requests

4. Free

Return blocks on finish

Memory Model

Rough bytes for one request (ignoring overhead):

2 × L × H × D × T × sizeof(dtype) — keys and values, L layers, H heads (or KV heads for GQA), D head dim, T tokens.

Factor	Effect on KV RAM	Serving lever
Sequence length T	Linear growth	Context limits, sliding window
Concurrent requests	Multiplies footprint	Max batch / admission control
Dtype (FP16 vs FP8/INT8)	2–4× savings if quantized	KV cache quantization
GQA / MQA	Fewer KV heads	Architecture choice
Fragmentation	Wasted reserved slots	Paged / block allocator

Paged KV (vLLM-Style Sketch)

Logical tokens map to physical blocks. Attention gathers K/V via the block table instead of assuming one contiguous tensor per request.

# Conceptual sketch — not a full engine
from dataclasses import dataclass, field

BLOCK_SIZE = 16  # tokens per KV block

@dataclass
class BlockTable:
    blocks: list[int] = field(default_factory=list)  # physical block ids

class PagedKVPool:
    def __init__(self, num_blocks: int):
        self.free = list(range(num_blocks))
        # pool[block_id] holds K/V for BLOCK_SIZE positions

    def alloc(self, n_tokens: int) -> list[int]:
        n_blocks = (n_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE
        if len(self.free) < n_blocks:
            raise MemoryError("KV pool exhausted — reject or swap request")
        out = [self.free.pop() for _ in range(n_blocks)]
        return out

    def free_blocks(self, ids: list[int]) -> None:
        self.free.extend(ids)

# Prefill 48 tokens → 3 blocks; decode grows into a new block every 16 tokens
pool = PagedKVPool(1024)
bt = BlockTable(blocks=pool.alloc(48))
print("prefill blocks:", bt.blocks)

Serving Implications

Contiguous cache

Simple indexing.
Pre-reserve max length.
Fragmentation under load.

Paged cache

High utilization.
Prefix sharing via shared blocks.
Needs block-aware kernels.

Decode bottleneck

Often memory-bandwidth bound.
Reading huge KV each step.
Batching amortizes weight reads.

Strengths and Tradeoffs

Strengths

Makes long interactive decode feasible.
Paged layouts unlock high concurrency.
Enables prefix reuse across users.

Tradeoffs

KV RAM often exceeds weight RAM at long context.
Complex schedulers and eviction policies.
Bugs in block maps cause silent wrong outputs.

Common Misconception

“Bigger GPU batch always means higher tokens/sec.” Once KV memory fills, the scheduler must reject or pause requests. Optimal throughput balances weight reuse (larger batch) against KV footprint (smaller concurrent contexts).

Knowledge Check

Short Answer: What does the KV cache store? Answer: Per-layer key and value tensors for past tokens.
True/False: Prefill writes prompt KV; decode appends one position at a time. Answer: True.
Multiple Choice: Paged attention mainly helps: (a) tokenizer speed, (b) KV memory utilization under many requests, (c) dataset shuffling. Answer: (b).
Short Answer: Name two factors that grow KV RAM linearly. Answer: Sequence length and number of layers (also batch/concurrency, heads, dtype).
True/False: The KV cache stores detokenized chat text. Answer: False—it stores activations.
Multiple Choice: Decode is often limited by: (a) memory bandwidth reading KV, (b) only CSS, (c) BLEU. Answer: (a).
Short Answer: Why did Vol. 11’s simple concatenate fail at scale? Answer: Contiguous per-request reservation fragments under multi-tenant load.
True/False: GQA/MQA reduces KV heads and thus cache size. Answer: True.
Multiple Choice: A block table maps: (a) URLs to DNS, (b) logical token ranges to physical KV blocks, (c) ranks to GPUs only. Answer: (b).
Short Answer: Which next lecture attacks attention IO specifically? Answer: Flash Attention.

Key Takeaways

Serving KV caches are managed pools, not just torch.cat.
Memory scales with length × layers × concurrency; it often caps batch size.
Paged layouts raise utilization and enable prefix sharing.
Decode throughput is tightly coupled to how fast you can read KV.
Next: Flash Attention.

Trainer’s Guide

Hands-on idea: Estimate KV GB for a 32-layer GQA model at 8k context, FP16, for 1 vs 32 concurrent users; discuss when the GPU is “full.”

Discussion prompt: Would you rather quantize the KV cache or shorten max context to admit more users?

Recap: The KV cache is the central memory structure of LLM serving; paged management turns a single-session trick into a multi-tenant system. Continue with Flash Attention.