Lab 08 - KV Cache Step-by-Step

Overview

Lab 08: KV Cache Step-by-Step

Annotated code reading lab. Running code is optional.

Related handbook section

KV Cache Step-by-Step

This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Open handbook section Open starter file

Concept Goal

KV Cache Step-by-Step

Read prefill and decode as two phases of the same cache contract.

Mental Model

Mechanism to keep in mind

`prefill` writes K/V for the prompt.
`decode_step` appends one new K/V per layer.
The new query reads all cached K/V instead of recomputing history.

Memory estimate

Cache size variables

KV Cache memory

2 * layers * batch * seq_len * kv_heads * head_dim * bytes

2 is K plus V.
seq_len is cached context length.
bytes depends on cache precision and engine support.

Annotated Code Preview

Starter preview

Excerpt from code/lab-08-kv-cache-step-by-step/kv_cache.py. The linked starter file is the source of truth.

Open starter file

# KV Cache Step-by-Step
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Read prefill and decode as two phases of the same cache contract.

kv_cache = []
for prompt_token in ["A", "B", "C"]:
    kv_cache.append(f"K,V({prompt_token})")  # prefill

new_token = "D"
query = f"Q({new_token})"
context = f"attention({query}, cached={len(kv_cache)} tokens)"
kv_cache.append(f"K,V({new_token})")

# What to explain while reading:
# - prefill writes K/V for the prompt.
# - decode_step appends one new K/V per layer.
# - The new query reads all cached K/V instead of recomputing history.
#
# Common traps:
# - KV cache is mainly an inference serving concern.
# - Weight quantization does not automatically shrink KV cache.

Line-by-line Explanation

What each block is doing

Setup / contract: `prefill` writes K/V for the prompt.
Main transition: `decode_step` appends one new K/V per layer.
Interview hook: The new query reads all cached K/V instead of recomputing history.

What to Notice

Reading checkpoints

KV cache grows with context and concurrency.
The cache saves compute but costs memory.
GQA/MQA reduce KV heads and cache size.

Common Misunderstandings

What this lab prevents

KV cache is mainly an inference serving concern.
Weight quantization does not automatically shrink KV cache.

Interview Explanation

How to say it out loud

Read prefill and decode as two phases of the same cache contract. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.

External intuition notes

Additional intuition

Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.

InfraLens