# Context windows and the KV cache interview questions

> 11 context windows and the kv cache interview questions for AI and ML engineer roles, each with a model answer.

Canonical: https://llmdojo.dev/interview/topic/context-and-kv-cache
Last updated: 2026-10-23

11 questions: 1 junior, 6 mid-level and 4 senior. Taught in [Context Windows, Prefill and the KV Cache](https://llmdojo.dev/learn/model-foundations/context).

## [What is a context window, and why can't GPT-2 small take a 1,025th token even with plenty of memory?](https://llmdojo.dev/interview/q/what-is-a-context-window-and-why-can-t-gpt-2-small)

Level: Junior

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.

## [Users say your summariser 'takes ages to start' on long documents but then writes quickly. Explain what's happening in terms of prefill and decode.](https://llmdojo.dev/interview/q/users-say-your-summariser-takes-ages-to-start-on-long-documents-but)

Level: Mid-level

**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.

## [What does the KV cache store, why is using it exact rather than an approximation, and how much does it save?](https://llmdojo.dev/interview/q/what-does-the-kv-cache-store-why-is-using-it-exact-rather)

Level: Mid-level

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

**Diagram: What the KV cache saves**

1. Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
2. Each decode step computes the key and value of the new token only, and appends them to the cache.
3. The new token's query still attends to every earlier key. Those keys and values are read back from the cache, not recomputed.
4. Earlier tokens' keys and values never change, because of the causal mask, so the cache is exact. With a 4-token prompt and 3 new tokens, that is 7 key-value computations with the cache and 22 without.

Counts from the KV-cache lab in lesson 5: a 4-token prompt and 3 new tokens.

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**.

## [You're sizing GPUs for a chat service. How do you estimate KV-cache memory, and which levers would you pull if it doesn't fit?](https://llmdojo.dev/interview/q/you-re-sizing-gpus-for-a-chat-service-how-do-you-estimate)

Level: Senior

**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.

## [Request A has a 960-token prompt and 10 output tokens. Request B has a 16-token prompt and 200 output tokens. Same model and machine: which finishes first, and what does that tell you about a chat app that resends the full history every turn?](https://llmdojo.dev/interview/q/request-a-has-a-960-token-prompt-and-10-output-tokens-request)

Level: Senior

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.

## [What are multi-head and grouped-query attention, and what does GQA do to the KV cache?](https://llmdojo.dev/interview/q/what-are-multi-head-and-grouped-query-attention-and-what-does-gqa)

Level: Mid-level

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**.

**Diagram: Multi-head vs grouped-query attention**

1. In multi-head attention, each head projects its own queries, keys and values. Every key and value is cached for every token.
2. Grouped-query attention keeps all the query heads but lets each group read one shared key-value head. Multi-query attention is the extreme: one K, V for all.
3. The cache shrinks in proportion to the key-value heads, which is why most large open models use GQA.

Ainslie et al. (2023), GQA; Llama 2 70B's configuration: 64 query heads, 8 key-value heads (Touvron et al. 2023, Section 2.2).

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.

## [Here's a small transformer with a few bugs. Fix them, then add a KV cache for decoding.](https://llmdojo.dev/interview/q/here-s-a-small-transformer-with-a-few-bugs-fix-them-then)

Level: Mid-level

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.

```python
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
```

**Diagram: What the KV cache saves**

1. Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
2. Each decode step computes the key and value of the new token only, and appends them to the cache.
3. The new token's query still attends to every earlier key. Those keys and values are read back from the cache, not recomputed.
4. Earlier tokens' keys and values never change, because of the causal mask, so the cache is exact. With a 4-token prompt and 3 new tokens, that is 7 key-value computations with the cache and 22 without.

Counts from the KV-cache lab in lesson 5: a 4-token prompt and 3 new tokens.

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

## [Implement causal multi-head attention, add a KV cache for token-by-token decoding, then extend it to grouped-query attention.](https://llmdojo.dev/interview/q/implement-causal-multi-head-attention-add-a-kv-cache-for-token-by)

Level: Mid-level

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.

```python

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.

## [Explain FlashAttention, the KV cache and RoPE, and the memory and compute trade-offs each addresses.](https://llmdojo.dev/interview/q/explain-flashattention-the-kv-cache-and-rope-and-the-memory-and-compute)

Level: Mid-level

**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.

**Diagram: What the KV cache saves**

1. Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
2. Each decode step computes the key and value of the new token only, and appends them to the cache.
3. The new token's query still attends to every earlier key. Those keys and values are read back from the cache, not recomputed.
4. Earlier tokens' keys and values never change, because of the causal mask, so the cache is exact. With a 4-token prompt and 3 new tokens, that is 7 key-value computations with the cache and 22 without.

Counts from the KV-cache lab in lesson 5: a 4-token prompt and 3 new tokens.

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.

## [What is attention, exactly how are the queries, keys and values computed, and how does a KV cache speed up decoding?](https://llmdojo.dev/interview/q/what-is-attention-exactly-how-are-the-queries-keys-and-values-computed)

Level: Senior

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**.

## [Design and implement a KV cache for autoregressive decoding in a Transformer: tensor shapes for prefill and per-step decode, per-layer caches, batching, memory limits and beam search.](https://llmdojo.dev/interview/q/design-and-implement-a-kv-cache-for-autoregressive-decoding-in-a-transformer)

Level: Senior

```python
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.

**Diagram: What the KV cache saves**

1. Prefill runs the prompt once and computes the keys and values of all 4 prompt tokens, in every layer and every head. They go into the cache.
2. Each decode step computes the key and value of the new token only, and appends them to the cache.
3. The new token's query still attends to every earlier key. Those keys and values are read back from the cache, not recomputed.
4. Earlier tokens' keys and values never change, because of the causal mask, so the cache is exact. With a 4-token prompt and 3 new tokens, that is 7 key-value computations with the cache and 22 without.

Counts from the KV-cache lab in lesson 5: a 4-token prompt and 3 new tokens.

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.

## Related
- [All interview questions](https://llmdojo.dev/interview)