KV Cache › Implementation
Language AI HandbookAll chapters KV Cache › ImplementationOne section · Sections in this chapter8 / 14 ↓
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
Let's build a simple KV cache implementation to see how these concepts work in practice. The implementation below prioritizes clarity over production-level efficiency, but it captures all the needed logic.
import torch
class KVCache:
"""
Simple KV cache for a single attention layer.
"""
def __init__(
self,
batch_size: int,
num_heads: int,
head_dim: int,
max_seq_len: int,
dtype=torch.float32,
device="cpu",
):
self.batch_size = batch_size
self.num_heads = num_heads
self.head_dim = head_dim
self.max_seq_len = max_seq_len
# Pre-allocate cache tensors
cache_shape = (batch_size, num_heads, max_seq_len, head_dim)
self.k_cache = torch.zeros(cache_shape, dtype=dtype, device=device)
self.v_cache = torch.zeros(cache_shape, dtype=dtype, device=device)
# Track current sequence length
self.seq_len = 0
def update(self, k_new: torch.Tensor, v_new: torch.Tensor) -> tuple:
"""
Append new key-value pairs to cache and return full cache.
Args:
k_new: New keys of shape (batch, num_heads, new_len, head_dim)
v_new: New values of shape (batch, num_heads, new_len, head_dim)
Returns:
Tuple of (all_keys, all_values) including new additions
"""
new_len = k_new.shape[2]
# Store new keys and values
self.k_cache[:, :, self.seq_len : self.seq_len + new_len, :] = k_new
self.v_cache[:, :, self.seq_len : self.seq_len + new_len, :] = v_new
self.seq_len += new_len
# Return only the valid portion of the cache
return (
self.k_cache[:, :, : self.seq_len, :],
self.v_cache[:, :, : self.seq_len, :],
)
def get_seq_len(self) -> int:
return self.seq_lenThe KVCache class pre-allocates tensors at initialization and uses an integer counter to track how much of the pre-allocated space has been used. The update method writes new entries into the pre-allocated buffer using slice assignment, which is an in-place operation that avoids memory allocation during the hot loop of generation. Returning only the valid slice [:, :, :self.seq_len, :] ensures the attention computation only sees actual cached entries, not the zero-padding at the end.
from __future__ import annotations
import math
from typing import TYPE_CHECKING
import torch
import torch.nn as nn
import torch.nn.functional as F
if TYPE_CHECKING:
pass
class CausalSelfAttentionWithCache(nn.Module):
"""
Causal self-attention module that supports KV caching.
"""
def __init__(self, d_model: int, num_heads: int):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
# Combined QKV projection for efficiency
self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
self.scale = 1.0 / math.sqrt(self.head_dim)
def forward(
self,
x: torch.Tensor,
kv_cache: "KVCache | None" = None,
use_cache: bool = False,
) -> "tuple[torch.Tensor, KVCache | None]":
"""
Forward pass with optional KV caching.
Args:
x: Input tensor of shape (batch, seq_len, d_model)
kv_cache: Optional KVCache object for incremental decoding
use_cache: Whether to use/update the cache
Returns:
Tuple of (output, updated_cache)
"""
batch_size, seq_len, _ = x.shape
# Compute Q, K, V projections
qkv = self.qkv_proj(x)
qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4) # (3, batch, heads, seq, head_dim)
q, k, v = qkv[0], qkv[1], qkv[2]
# Handle KV cache
if use_cache and kv_cache is not None:
k, v = kv_cache.update(k, v)
cache_len = kv_cache.get_seq_len()
else:
cache_len = seq_len
# Compute attention scores
# q: (batch, heads, seq_len, head_dim)
# k: (batch, heads, cache_len, head_dim)
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
# Apply causal mask
# For cached attention: new queries can attend to all cached keys
if seq_len == 1 and cache_len > 1:
# Single token attending to full cache - no masking needed
pass
else:
# Create causal mask for the needed portion
mask = torch.triu(
torch.ones(
seq_len, cache_len, device=x.device, dtype=torch.bool
),
diagonal=cache_len - seq_len + 1,
)
scores = scores.masked_fill(
mask.unsqueeze(0).unsqueeze(0), float("-inf")
)
attn_weights = F.softmax(scores, dim=-1)
# Apply attention to values
out = torch.matmul(attn_weights, v)
# Reshape and project output
out = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
out = self.out_proj(out)
return out, kv_cacheThe combined QKV projection (qkv_proj) computes all three projections in a single matrix multiply, which is faster than three separate operations on modern hardware due to better memory bandwidth utilization. The reshaped qkv tensor is then split into separate query, key, and value tensors before the cache update.
Now let's show how the cache accelerates generation:
# Set up a small model setup
d_model = 64
num_heads = 4
batch_size = 1
max_seq_len = 32
# Create attention module
attention = CausalSelfAttentionWithCache(d_model, num_heads)
# Simulate a prompt of 8 tokens
prompt_len = 8
prompt = torch.randn(batch_size, prompt_len, d_model)
# PREFILL PHASE: Process entire prompt at once
cache = KVCache(batch_size, num_heads, d_model // num_heads, max_seq_len)
output, cache = attention(prompt, cache, use_cache=True)After prefill phase: Cache sequence length: 8 K cache shape (used portion): torch.Size([1, 4, 8, 16])
The output confirms that the cache has been initialized with the 8 prompt tokens. The key cache shape shows we have stored the projection results for these tokens, ready to be attended to by future generation steps.
# DECODE PHASE: Generate tokens one at a time
num_new_tokens = 5
generation_outputs = []
for step in range(num_new_tokens):
# Simulate embedding of newly generated token
new_token = torch.randn(batch_size, 1, d_model)
# Process only the new token, reusing cached K,V
output, cache = attention(new_token, cache, use_cache=True)
generation_outputs.append(output)After generating 5 tokens: Cache sequence length: 13 Total tokens processed: 13 K cache used: torch.Size([1, 4, 13, 16])
The cache has grown by 5 tokens, now containing the history for 13 tokens total.
Let's compare the computational savings. The important detail is that we only processed the 5 new tokens through the model's projection layers, yet the attention mechanism had access to the full 13-token history via the cache.
def count_qkv_projections(
prompt_len: int, gen_len: int, use_cache: bool
) -> dict:
"""
Count Q, K, V projection operations with and without caching.
"""
if use_cache:
# Prefill: compute Q, K, V for all prompt tokens
prefill_projections = prompt_len * 3 # Q, K, V for each token
# Decode: compute Q, K, V only for new tokens
decode_projections = gen_len * 3
total = prefill_projections + decode_projections
else:
# Without cache: recompute all projections at each step
total = 0
for t in range(1, gen_len + 1):
seq_len = prompt_len + t
total += seq_len * 3 # Q, K, V for entire sequence
return {"total_projections": total, "with_cache": use_cache}
# Compare for realistic generation scenario
prompt_len = 100
gen_len = 200
with_cache = count_qkv_projections(prompt_len, gen_len, use_cache=True)
without_cache = count_qkv_projections(prompt_len, gen_len, use_cache=False)
speedup = without_cache["total_projections"] / with_cache["total_projections"]Generating 200 tokens from 100-token prompt: With KV cache: 900 projection operations Without cache: 120,300 projection operations Reduction factor: 133.7x fewer operations
The savings are substantial. For generation, the projection computation savings grow quadratically with sequence length.
