LM10-L3 — Implement a KV Cache and Rotary Embeddings in PyTorch

Objective

Implement RoPE tables and rotation, a cache update and an offset causal mask so that a grouped-query attention layer decodes token by token with the same logits as a full forward pass, and so that the cache shrinks in proportion to the number of key-value heads.

Why a notebook? This exercise trains or runs a real PyTorch model. PyTorch does not run in the browser Python used by the site’s other labs, so you run this one on your own computer or in a hosted notebook service. The interactive browser labs of this module cover the surrounding ideas without PyTorch.

Setup and requirements

Hardware: CPU only, about 2 GB of RAM, a few minutes. CPU only; about 2 GB of RAM; tensors are tiny, so the self-checks take seconds. The timing experiment at the end generates 96 tokens with a model of width 128 and takes a few seconds on a laptop.

Download the notebook (.ipynb) Reference solution (try the lab first)

On your own machine (Python 3.10 or newer):

python3 -m venv llm-env
source llm-env/bin/activate        # on Windows: llm-env\Scripts\activate
pip install torch numpy jupyterlab
jupyter lab kv-cache-and-rope-in-pytorch.ipynb

In a hosted notebook service such as Google Colab or Kaggle, upload the file and run it there; PyTorch is usually preinstalled. Matplotlib is used for optional plots when it is available; the notebook runs without it.

What you will do

1. Rotary embeddings. Build the cos and sin tables and rotate the two halves of each head vector; check the norm, position 0 and the dependence on the offset only.
2. Cache and mask. Append new keys and values along the time dimension and build the causal mask for queries that are the last positions of a longer sequence.
3. Grouped-query layer. Read the given layer that uses your functions and shares each key-value head between several query heads.
4. Cached equals full. Compare an eight-token prefill plus four decode steps with one full pass, and compare cache sizes for different numbers of key-value heads.
5. Generation. Write greedy generation with and without a cache and count the tokens processed.
6. Timing. Time both ways on a longer prompt and read how the work grows.

Cells that you must complete are marked TODO in the notebook and have a self-check right after them (an assert with a message). Run the cells from the top; a failing check tells you what is wrong.

Expected behaviour

• Self-check 1: the tables have shape (n_pos, head_dim/2); position 0 is the identity; the rotation preserves the norm; the score between positions m and n depends only on m − n.
• Self-check 2: the cache grows along dimension 2 with old keys first; the offset mask is the usual causal mask without a past, empty for a single new token, and correct for a two-token chunk.
• Self-check 3: the largest difference between full and cached logits is below 1e-5, and the cache holds 12 positions of 2 key-value heads; the cache bytes are in the ratio 4 : 2 : 1 for 4, 2 and 1 key-value heads.
• Self-check 4: generation with and without a cache yields identical tokens; the model processes 105 tokens without a cache and 15 with one.

How to know you are done

• All four self-checks print “passed” when the notebook is run from the top.
• You can explain why the rotated keys are cached and why cached decoding is exact rather than approximate.
• You can compute the cache bytes per token for a configuration and relate it to Lab LM10-L1.

The site cannot run or grade the notebook for you. The self-checks inside the notebook are your verification. The reference solution was executed with PyTorch on a CPU when this module was prepared; timings depend on your machine, and the model is random, so nothing is learned.

If something goes wrong

• The cached logits differ from the full ones: check start_pos: each new token must be rotated by its absolute position, and the mask offset must be n_total − n_query.
• A shape error in the cache: the time dimension is dimension 2 of a (batch, heads, time, width) tensor, not dimension 1.
• The cache has the wrong number of heads: cache the keys and values before repeating them for the query groups.
• No PyTorch installed: run pip install torch in your virtual environment; the CPU build is enough.

Theory for this lab

• LM10-T2: Rotary Position Embeddings
• LM10-T5: KV Caching, Prefill and Decode

Reference material

Modern LLM Architecture and Inference labs

1. LM10-L1: KV-Cache Memory and Behaviour
2. LM10-L2: Compare Attention Variants
3. LM10-L3: Implement a KV Cache and Rotary Embeddings in PyTorch: this lab.

Next lab: LM11-L1: Chunking and Retrieval.