KV Cache for Efficient Generation - LLM Generation & Decoding

KV Cache for Efficient Generation

What you'll build

One decoding step of single-head attention with a KV cache: project only the new token, append its key and value to the cache, and let its query attend over every cached position.

Background

Generation is autoregressive: to produce token t+1 the model needs attention over tokens 0..t. The keys and values of tokens 0..t-1 don't change once computed (causal attention never lets old tokens see new ones), so recomputing them every step is wasted work. Caching them turns each step from "re-run the whole prefix" into "process one token". This is why LLM serving has two phases: prefill (run the prompt in parallel and fill the cache) and decode (one token per step, reading the whole cache). The cache, 2 · n_layers · n_kv_heads · d_head · seq_len numbers per sequence, is often bigger than the model's activations; GQA (Transformer Architecture) and sliding windows exist to shrink it.

Algorithm

For the new token embedding x_new (1, d_model):

q = x_new @ W_Q            # (1, d_k)   only the new token is projected
k = x_new @ W_K            # (1, d_k)
v = x_new @ W_V            # (1, d_v)
k_cache = concat(k_cache, k)  -> (t + 1, d_k)     (start from k if the cache is None)
v_cache = concat(v_cache, v)  -> (t + 1, d_v)
output  = softmax(q @ k_cache.T / sqrt(d_k)) @ v_cache    # (1, d_v)

No mask is needed: the cache only contains the past and the current token, so this is exactly causal attention for the last position.

Inputs

• x_new: (1, d_model)
• W_Q, W_K: (d_model, d_k); W_V: (d_model, d_v)
• k_cache: (t, d_k) or None; v_cache: (t, d_v) or None (None on the first step)

Expected Return Format

A tuple (output, k_cache, v_cache):

• output: (1, d_v)
• k_cache: (t + 1, d_k), v_cache: (t + 1, d_v), with the new row appended at the end

Edge cases and conventions

• First step (None caches): the new token attends only to itself, so output = v.
• Stable softmax; scale by sqrt(d_k) where d_k = W_Q.shape[1].

Constraints

NumPy only. Don't recompute keys or values for cached tokens.

Follow-up

• Per generated token, attention with a cache costs O(t · d) instead of O(t^2 · d) (or O(t · d^2) for the projections). Over n tokens: O(n^2 d) vs O(n^3 d).
• How big is the KV cache of a 32-layer, 8-KV-head, d_head = 128 model at 32k tokens in fp16? (About 4 GiB per sequence.)

Input: Step 1: kv_cache_attention(np.array([[1.0, 0.0]]), I, I, I); step 2 with x = [[0.0, 1.0]] and the returned caches (I = np.eye(2))
Output: step 1 output [[1.0, 0.0]]; step 2 output [[0.3302, 0.6698]], caches of length 2

Input: x_new (1, 8), cache of 5 previous tokens, d_k = d_v = 4
Output: output (1, 4), caches (6, 4)

import numpy as np

def kv_cache_attention(x_new, W_Q, W_K, W_V, k_cache=None, v_cache=None):
    """
    One decoding step of attention with a KV cache.

    Args:
        x_new: New token embedding (1, d_model)
        W_Q: Query weights (d_model, d_k)
        W_K: Key weights (d_model, d_k)
        W_V: Value weights (d_model, d_v)
        k_cache: Cached keys (t, d_k) or None on the first step
        v_cache: Cached values (t, d_v) or None on the first step

    Returns:
        output: Attention output for the new token (1, d_v)
        k_cache: Updated key cache (t + 1, d_k), new row last
        v_cache: Updated value cache (t + 1, d_v), new row last
    """
    # Your code here
    pass