KV cache
Stored keys and values of earlier tokens, so each decode step only computes them for the new token. Exact, but costs memory.
Interview questions about KV cache
- Walk me through Transformer attention: causal and padding masks, cross-attention, why a KV cache helps at inference, and how mixed-precision training works.
Masks set scores to minus infinity before the softmax. The causal mask hides future positions; the padding mask hides pad keys, and pads are excluded from the…
- What does the KV cache store, why is using it exact rather than an approximation, and how much does it save?
It stores every processed token's keys and values, per layer and head, so each decode step computes only the new token's.
- What are multi-head and grouped-query attention, and what does GQA do to the KV cache?
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.
- Here's a small transformer with a few bugs. Fix them, then add a KV cache for decoding.
Usual planted bugs: the mask applied after softmax, scaling by d instead of head size, a missing transpose before merging heads.
- Implement causal multi-head attention, add a KV cache for token-by-token decoding, then extend it to grouped-query attention.
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.
- Explain FlashAttention, the KV cache and RoPE, and the memory and compute trade-offs each addresses.
FlashAttention computes exact attention in tiles kept in on-chip SRAM, with an online softmax, so the full score matrix never goes to HBM.
- What is attention, exactly how are the queries, keys and values computed, and how does a KV cache speed up decoding?
Attention builds each token's output as a weighted mix of other tokens.
- 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.