Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Cross-Attention › KV Caching in Cross-Attention

Michael BrenndoerferUpdated June 18, 20254 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 114Cross-Attention › KV Caching in Cross-Attention

Reading section 9 of 14 · KV Caching in Cross-Attention · 4 minSection 9 of 14 · 4 min

Section 9 of 14 · 4 min read

Cross-Attention › KV Caching in Cross-AttentionOne section · Sections in this chapter9 / 14 ↓

Sections in this chapter14

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

During autoregressive generation, cross-attention has a remarkable computational advantage over decoder self-attention: the keys and values from the encoder can be computed once and reused for every decoding step. This is the key insight behind KV caching for cross-attention.

To understand why this works, recall that cross-attention keys and values come entirely from the encoder output: $K = X_{\text{enc}} W^K$ and $V = X_{\text{enc}} W^V$ . The encoder output $X_{\text{enc}}$ is fixed throughout decoding. It does not change as the decoder generates tokens. Therefore, $K$ and $V$ are the same matrices at every decoding step. Computing them repeatedly would be wasteful. Instead, we compute them once after the encoder finishes and cache them for reuse.

Contrast this with decoder self-attention, where caching is more complex. The decoder's self-attention keys and values depend on all previously generated tokens, so as the decoder generates more tokens, the key and value matrices grow. This is the KV cache that grows during generation and is the primary memory bottleneck for long generations. Cross-attention's cache, by contrast, has a fixed size equal to the encoder length from the very first decoding step.

In[25]:

Code

class CachedCrossAttention:
    """
    Cross-attention with KV caching for efficient inference.

    The encoder's keys and values are computed once and reused
    for all decoder steps.
    """

    def __init__(self, d_model, d_k, d_v):
        self.d_k = d_k
        self.scale = 1.0 / np.sqrt(d_k)

        self.W_Q = np.random.randn(d_model, d_k) * np.sqrt(
            2.0 / (d_model + d_k)
        )
        self.W_K = np.random.randn(d_model, d_k) * np.sqrt(
            2.0 / (d_model + d_k)
        )
        self.W_V = np.random.randn(d_model, d_v) * np.sqrt(
            2.0 / (d_model + d_v)
        )

        # Cache for encoder K and V
        self.cached_K = None
        self.cached_V = None

    def cache_encoder(self, encoder_output):
        """
        Compute and cache encoder keys and values.
        Called once after encoding.
        """
        self.cached_K = encoder_output @ self.W_K
        self.cached_V = encoder_output @ self.W_V

    def forward(self, decoder_state, encoder_mask=None):
        """
        Apply cross-attention using cached encoder KV.

        Args:
            decoder_state: Current decoder position(s) (n_new, d_model)

        Returns:
            output: (n_new, d_v)
        """
        if self.cached_K is None:
            raise ValueError("Must call cache_encoder() first")

        # Only compute Q for new decoder positions
        Q = decoder_state @ self.W_Q

        # Use cached K and V
        scores = Q @ self.cached_K.T * self.scale

        if encoder_mask is not None:
            mask = encoder_mask.reshape(1, -1)
            scores = np.where(mask, scores, -1e9)

        attention_weights = softmax(scores)
        output = attention_weights @ self.cached_V

        return output, attention_weights

In[26]:

Code

# Simulate incremental decoding
d_model = 16
d_k = d_v = 16

# Encoder processes source once
source_len = 8
encoder_output = np.random.randn(source_len, d_model)

# Create cached cross-attention and cache encoder KV
cached_cross_attn = CachedCrossAttention(d_model, d_k, d_v)
cached_cross_attn.cache_encoder(encoder_output)

# Generate tokens one at a time
generated_tokens = []
for step in range(4):
    # Get representation for current position only
    current_state = np.random.randn(1, d_model)

    # Cross-attention uses cached K, V from encoder
    output, weights = cached_cross_attn.forward(current_state)
    generated_tokens.append(f"token_{step}")

Out[27]:

Console

Incremental generation with KV caching:
  Encoder length: 8
  Cached K shape: (8, 16)
  Cached V shape: (8, 16)

At each step, only Q is computed for the new token.
K and V are reused from cache, avoiding redundant computation.

This caching is particularly important for long source sequences. Without caching, generating 100 tokens from a 1000-token source would require recomputing the encoder's KV projections 100 times. With caching, we compute them once, and each decoding step only needs to compute the query projection for the new token and multiply it against the fixed, cached keys. This reduces the per-step computational cost of cross-attention from $O(n_{\text{enc}} \cdot d)$ matrix multiplications to just the query projection (one multiplication of size $d \times d_k$ ) plus the attention computation.

In memory terms, the cross-attention KV cache occupies $2 \cdot n_{\text{enc}} \cdot d_v \cdot n_{\text{layers}}$ values (two matrices of size $n_{\text{enc}} \times d_v$ per decoder layer). For a model with $d_v = 64$ , $n_{\text{enc}} = 512$ , and 12 layers, this is about 800 thousand values, which is very manageable. The decoder self-attention KV cache, by contrast, grows linearly with the number of generated tokens and can become a significant memory constraint for very long outputs.