Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Decoder Architecture › Key-Value Caching for Efficient Generation

Michael BrenndoerferUpdated June 19, 20254 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 112Decoder Architecture › Key-Value Caching for Efficient Generation

Reading section 10 of 16 · Key-Value Caching for Efficient Generation · 4 minSection 10 of 16 · 4 min

Section 10 of 16 · 4 min read

Decoder Architecture › Key-Value Caching for Efficient GenerationOne section · Sections in this chapter10 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

A naive implementation recomputes attention for all positions at every generation step. This is wasteful: the representations for tokens 0 through $t-1$ don't change when we add token $t$ . KV caching stores the key and value vectors from previous positions, avoiding redundant computation.

In[24]:

Code

class CausalAttentionWithCache:
    """
    Multi-head attention with KV cache for efficient generation.
    """

    def __init__(self, d_model, n_heads):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        scale = np.sqrt(2.0 / (d_model + self.d_k))
        self.W_q = np.random.randn(d_model, d_model) * scale
        self.W_k = np.random.randn(d_model, d_model) * scale
        self.W_v = np.random.randn(d_model, d_model) * scale
        self.W_o = np.random.randn(d_model, d_model) * scale

        # Cache for past keys and values
        self.cache_k = None
        self.cache_v = None

    def clear_cache(self):
        """Reset the KV cache."""
        self.cache_k = None
        self.cache_v = None

    def __call__(self, x, use_cache=False):
        """
        Apply attention, optionally using and updating cache.

        When use_cache=True:
        - x should be just the new token(s)
        - Keys/values are appended to cache
        - Queries attend to all cached positions
        """
        seq_len = x.shape[0]

        # Compute Q, K, V for new positions
        Q = x @ self.W_q
        K = x @ self.W_k
        V = x @ self.W_v

        # Reshape for multi-head
        Q = Q.reshape(seq_len, self.n_heads, self.d_k).transpose(1, 0, 2)
        K = K.reshape(seq_len, self.n_heads, self.d_k).transpose(1, 0, 2)
        V = V.reshape(seq_len, self.n_heads, self.d_k).transpose(1, 0, 2)

        if use_cache:
            # Append new K, V to cache
            if self.cache_k is None:
                self.cache_k = K
                self.cache_v = V
            else:
                self.cache_k = np.concatenate([self.cache_k, K], axis=1)
                self.cache_v = np.concatenate([self.cache_v, V], axis=1)

            # Use full cached K, V
            K_full = self.cache_k
            V_full = self.cache_v
        else:
            K_full = K
            V_full = V

        total_len = K_full.shape[1]

        # Compute attention scores
        scores = Q @ K_full.transpose(0, 2, 1) / np.sqrt(self.d_k)

        # Apply causal mask (only for positions we're querying)
        # Shape: (seq_len, total_len)
        mask = np.zeros((seq_len, total_len))
        start_pos = total_len - seq_len
        for i in range(seq_len):
            # Position i in query corresponds to position start_pos + i in full sequence
            full_pos = start_pos + i
            # Can attend to positions 0 to full_pos
            mask[i, full_pos + 1 :] = float("-inf")

        scores = scores + mask

        # Softmax and apply to values
        attn_weights = softmax(scores, axis=-1)
        attn_output = attn_weights @ V_full

        # Reshape and project
        attn_output = attn_output.transpose(1, 0, 2).reshape(
            seq_len, self.d_model
        )
        output = attn_output @ self.W_o

        return output

Out[25]:

Console

KV Cache Demo
=============================================
Initial sequence: 5 tokens
Cache K shape after init: (4, 5, 16)
Added 1 token
Cache K shape after addition: (4, 6, 16)
Output shape for new token: (1, 64)

The cache grows from shape (4, 5, 16) to (4, 6, 16) after adding one token. Here, 4 is the number of heads, the middle dimension is sequence length (growing from 5 to 6), and 16 is the head dimension (64 / 4 heads). When processing the new token, we only compute queries for that single position, but the keys and values span all 6 positions. This asymmetry is what makes caching efficient: the expensive key-value computation for old tokens is reused.

With KV caching, generating $n$ new tokens requires $O(n \cdot L)$ compute instead of $O(n^2 \cdot L)$ , where:

$n$ : the number of new tokens to generate
$L$ : the total sequence length (prompt plus generated tokens)
$O(n \cdot L)$ : linear scaling because each new token only computes attention against the cached keys/values
$O(n^2 \cdot L)$ : quadratic scaling without caching, since we recompute all positions at each step

This optimization is critical for efficient inference in production systems. For a 1000-token prompt with 100 generated tokens, caching reduces attention computation by roughly 50x.