Context windows and the KV cache interview questions

It's the most tokens a model can read at once, shared by prompt, history and output: 1,024 for GPT-2 small. Attention alone ignores order, so GPT-2 adds a learned position vector from a table with exactly 1,024 rows. Position 1,025 has no row; on our CPU run it raised an IndexError.

The wait is prefill. Prefill runs the whole prompt in one parallel pass and sets the time to first token; decode then writes one token per step, each costing about the same.

On our test machine, GPT-2 small prefilled 16 tokens in 15.3 ms and 960 in 219.6 ms, while a decode step stayed between 8.0 and 9.6 ms. The lever is how much text goes into the prompt.

It stores every processed token's keys and values, per layer and head, so each decode step computes only the new token's.

It's exact because of the causal mask: earlier tokens never depend on later ones, and greedy output was identical with and without the cache. After a 960-token prompt on our test machine, a step took 9.6 ms with the cache against 228.8 ms without, about 24 times faster.

Multiply it out: 2 (key and value) × layers × heads × head dimension × tokens × bytes per number × batch. GPT-2 small in 32-bit floats needs 72 KiB per token; a full 1,024-token window at batch 64 is 4.50 GiB, nearly ten times the weights. Cache, not weights, can cap context length and concurrent users.

Each factor is a lever: fewer bytes (a lower-precision cache), fewer tokens (a local window, trimmed history), keys and values shared across heads, a smaller batch. I'd model the real mix of context lengths and batch sizes before buying hardware.

A. On our GPT-2 small measurements, A is about 220 ms of prefill plus 10 steps of about 10 ms, roughly 0.3 s; B is about 15 ms plus 200 sequential steps of about 8 ms, roughly 1.6 to 1.7 s. Output tokens run one after another, while prompt tokens are processed together.

Resending the history re-prefills more every turn, so time to first token and cache memory grow turn by turn; once history won't fit the window, apps drop or summarise old turns. APIs typically bill input and output tokens, so trim both.

In multi-head attention every head has its own queries, keys and values. Grouped-query attention lets a group of query heads share one key-value head.

The cache holds keys and values per key-value head, so fewer of them means a proportionally smaller KV cache: Llama 2 70B's 64 query heads share 8, an 8× cut, with little quality loss.

Usual planted bugs: the mask applied after softmax, scaling by d instead of head size, a missing transpose before merging heads. For the cache, each decode step computes q, k and v for the new token only, appends k and v, and attends over everything cached.

def step(x, w_qkv, w_o, h, cache):  # x: (B, 1, d)
    B, _, d = x.shape
    q, k, v = (x @ w_qkv).split(d, -1)
    sh = lambda t: t.view(B, 1, h, -1).transpose(1, 2)
    q, k, v = sh(q), sh(k), sh(v)   # (B, h, 1, d/h)
    if "k" in cache:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v
    s = q @ k.transpose(-2, -1) / (d // h) ** 0.5
    y = s.softmax(-1) @ v           # no mask: all past
    return y.transpose(1, 2).reshape(B, 1, d) @ w_o

Interviewers check that cached decoding matches the full forward pass and that positions continue from the cache length.

Write it once as GQA: plain MHA is the case where key-value heads equal query heads, and multi-query is one key-value head.

import torch, torch.nn as nn, torch.nn.functional as F

class GQA(nn.Module):
    def __init__(self, d, n_q, n_kv):
        super().__init__()
        self.hd, self.nq, self.nkv = d // n_q, n_q, n_kv
        self.q = nn.Linear(d, d, bias=False)
        self.kv = nn.Linear(d, 2 * n_kv * self.hd,
                            bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, cache=None):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.nq, self.hd)
        k, v = self.kv(x).view(
            B, T, 2, self.nkv, self.hd).unbind(2)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))
        if cache is not None:
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        cache = (k, v)  # n_kv heads: the memory saving
        r = self.nq // self.nkv
        k = k.repeat_interleave(r, dim=1)
        v = v.repeat_interleave(r, dim=1)
        S = k.shape[2]  # cached + new positions
        # new query i sits at position S - T + i
        m = torch.ones(T, S, dtype=torch.bool,
                       device=x.device).tril(S - T)
        a = q @ k.transpose(-2, -1) / self.hd ** 0.5
        a = a.masked_fill(~m, float("-inf"))
        y = F.softmax(a, dim=-1) @ v
        y = y.transpose(1, 2).reshape(B, T, -1)
        return self.o(y), cache

Interviewers check three things: the mask is offset by the cached length, the cache stores keys and values before repeating them, so it shrinks by the group factor, and cached decoding reproduces a full forward pass token for token.

FlashAttention computes exact attention in tiles kept in on-chip SRAM, with an online softmax, so the full score matrix never goes to HBM. Memory becomes linear in length, and it's faster because attention is memory-bound.

The KV cache keeps past keys and values so each decode step computes only the new token, trading memory for compute. RoPE rotates queries and keys by position-dependent angles, so dot products depend on relative position, parameter-free.

Attention builds each token's output as a weighted mix of other tokens. Each hidden vector is multiplied by three learned matrices to give a query, key and value. Scores are query-key dot products divided by the root of the head size, masked causally in a decoder, softmaxed, then used to average the values.

In decoding, past tokens' keys and values never change, so the KV cache stores them. Each step projects only the new token and attends over the cache, instead of reprocessing the prefix, paying in memory that grows with batch times context.

class KVCache:
    def __init__(self, L, B, H, T_max, dh, dtype, dev):
        shape = (L, B, H, T_max, dh)
        self.k = torch.zeros(shape, dtype=dtype, device=dev)
        self.v = torch.zeros_like(self.k)

    def update(self, layer, k, v, start):
        # k, v: (B, H, n, dh); n = prompt len, then 1
        end = start + k.size(2)
        self.k[layer, :, :, start:end] = k
        self.v[layer, :, :, start:end] = v
        return (self.k[layer, :, :, :end],
                self.v[layer, :, :, :end])

    def reorder(self, beam_idx):
        # beam search: rows follow surviving beams
        self.k = self.k.index_select(1, beam_idx)
        self.v = self.v.index_select(1, beam_idx)

Prefill writes the whole prompt in one pass; each decode step writes one position and attends over everything so far. Batch with left padding plus an attention mask so positions line up.

Memory per token is 2 × layers × KV heads × head dim × bytes, so cap batch times length. Beyond that: GQA, quantized caches, or paged blocks that avoid reserving T_max per sequence.