Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

KV Cache › Verifying Cache Correctness

Michael BrenndoerferJanuary 6, 20262 min read

Language AI HandbookAll chapters

Part 42 of 67Inference Optimization

Chapter 313KV Cache › Verifying Cache Correctness

Reading section 9 of 14 · Verifying Cache Correctness · 2 minSection 9 of 14 · 2 min

Section 9 of 14 · 2 min read

KV Cache › Verifying Cache CorrectnessOne 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

A necessary property of KV caching is that it must produce identical outputs to non-cached attention. This guarantee makes caching a valid optimization rather than an approximation. If caching changed the output even slightly, every generated token would be slightly different, and those small errors would compound over the course of a generation, creating wildly different sequences. Let's verify this mathematically and empirically.

The algebraic argument is straightforward. In standard full-sequence attention, the output for position $t$ in a sequence of length $T$ is:

o_t = \text{softmax}\left(\frac{q_t K[1:t]^T}{\sqrt{d_k}}\right) V[1:t]

where $K[1:t]$ denotes the submatrix of $K$ containing only the first $t$ rows. In cached attention, we have:

o_t = \text{softmax}\left(\frac{q_t K_{\text{cache}}^T}{\sqrt{d_k}}\right) V_{\text{cache}}

where $K_{\text{cache}}$ at step $t$ contains exactly the first $t$ rows of $K$ . Since $k_i = x_i W_K$ is computed identically whether we process the full sequence at once or token-by-token, the two expressions are algebraically identical. The numerical values in $K_{\text{cache}}$ are bitwise identical to the corresponding rows of $K$ , so the outputs are numerically equivalent up to floating-point precision.

In[15]:

Code

def verify_cache_equivalence():
    """
    Verify that cached and non-cached attention produce identical outputs.
    """
    torch.manual_seed(42)

    d_model = 64
    num_heads = 4
    batch_size = 2

    attention = CausalSelfAttentionWithCache(d_model, num_heads)

    # Create a sequence
    prompt = torch.randn(batch_size, 5, d_model)
    next_tokens = torch.randn(batch_size, 3, d_model)
    full_sequence = torch.cat([prompt, next_tokens], dim=1)

    # Method 1: Process entire sequence without cache
    output_no_cache, _ = attention(full_sequence, use_cache=False)

    # Method 2: Process with cache (prefill + decode)
    cache = KVCache(batch_size, num_heads, d_model // num_heads, max_seq_len=32)

    # Prefill with prompt
    output_prefill, cache = attention(prompt, cache, use_cache=True)

    # Decode remaining tokens one at a time
    cached_outputs = [output_prefill]
    for i in range(next_tokens.shape[1]):
        token = next_tokens[:, i : i + 1, :]
        output_step, cache = attention(token, cache, use_cache=True)
        cached_outputs.append(output_step)

    output_with_cache = torch.cat(cached_outputs, dim=1)

    return output_no_cache, output_with_cache


output_no_cache, output_with_cache = verify_cache_equivalence()

# Compare outputs
max_diff = (output_no_cache - output_with_cache).abs().max().item()
are_equal = torch.allclose(output_no_cache, output_with_cache, atol=1e-5)

Out[16]:

Console

Maximum difference between cached and non-cached: 1.19e-07
Outputs are equivalent: True

The numerical equivalence confirms that caching is purely an optimization: it changes how we compute the result, not what we compute. The tiny residual difference (typically around $10^{-7}$ ) shows floating-point rounding from the different order of operations, not a systematic error.