KV Cache Implementation and Memory Accounting

Lab: Run the companion CPU exercise at /labs/track1/03-cache-and-gqa/ (per-layer cache update, logit parity, memory accounting).

KV Cache Implementation and Memory Accounting

Why per-layer cache updates matter

In a decoder-only transformer, every attention layer maintains its own key and value cache. During prefill, the entire prompt is processed in one parallel forward pass, and the K/V tensors for every prompt token are stored. During decode, one new token is consumed per step: its key and value are computed, appended to the cache, and the next-token logits are produced.

The critical detail is that each layer writes its own K/V column independently. If a layer’s cache write is off by one position, or if the position index is misaligned with the attention mask, the next token can attend to the wrong key/value pair, producing incorrect logits. This lesson walks through the per-layer update, verifies logit parity with an uncached reference, and quantifies the memory cost.

Token trace: prefill through two decode inputs

Consider a causal decoder with a 3-token prompt t1, t2, t3 and greedy generation of two output tokens. The trace below tracks cache length, K/V shape, and the token whose logits are produced at each step.

KV cache trace: prefill → g1 → g2

Step	Prompt / Input	Cache length	K/V shape (per layer)	What happens	Logits predict
0	Prompt tokens t1, t2, t3 available	—	—	Three tokens ready for processing	—
1	Prefill: t1, t2, t3 processed together in one forward pass. Final prompt-position logits predict the first generated token. K/V states for all 3 tokens are stored in the cache.	P = 3	[B, H, 3, D]	All three K/V pairs written to their respective cache slots.	g1 (first generated token)
2	Decode step 1: consume token g1. With the KV cache, the new token attends to cached K/V for t1, t2, t3 plus its own K/V. Cache grows from 3 to 4. Logits predict the next token.	P+1 = 4	[B, H, 4, D]	g1’s key and value are appended as column index 3. The attention mask now includes the new token position.	g2
3	Decode step 2: consume token g2. Attend to cached K/V for t1, t2, t3, g1, plus g2’s own K/V. Cache grows from 4 to 5.	P+2 = 5	[B, H, 5, D]	g2’s key and value are appended as column index 4.	— (trace end)

Three-prompt token trace with two decode steps. Cache: 3 → 4 → 5 columns.

Read this diagram as text

Prefill fills the cache with 3 columns (one per prompt token). Decode step 1 appends g1's K/V, extending the cache to 4 columns; its logits predict g2. Decode step 2 appends g2's K/V, extending the cache to 5 columns; its logits would predict g3.

Key point: Each decode step processes exactly one new token. The cache extends by one column per step. The prefill pass already produces the logits for g1; decode step 1 consumes g1 and produces logits for g2. No extra model pass is needed to produce the first generated token.

Shape annotations for the trace

• Prefill K/V: [B, H, P, D] — one key/value pair per token, per KV head, per head dimension, for all prompt tokens together. With B=1, H=8, P=3, D=64 this is [1, 8, 3, 64].
• New K/V per decode step: [B, H, 1, D] — the key and value for the single new token, computed anew and written into the cache at the next free slot.
• Usable K/V after g1: [B, H, P+1, D] — the cache now contains the original P prompt columns plus 1 generated column. Any future attention pass can read all P+1 columns.
• Usable K/V after g2: [B, H, P+2, D] — the cache contains P prompt columns plus 2 generated columns.

These shapes assume a conventional decoder-only architecture where every attention layer stores K and V for every token in the context, across all KV heads. Grouped-query attention (GQA) would change the H dimension (to the number of KV heads) but the overall structure is the same.

Per-layer cache write: mechanics

Each transformer layer l maintains a key cache and a value cache. The write operation for a single decode step can be expressed as:

# Pseudocode: write new K/V into layer l's cache
cache_key[l][:, :, write_position, :] = new_key   # shape: [B, H, 1, D]
cache_value[l][:, :, write_position, :] = new_value

where write_position is the current cache length (number of tokens already stored). After prefill with P prompt tokens, write_position = P. After each decode step, write_position increments by 1.

The read operation for attention in the same layer uses a causal mask that allows the query at position s to attend to all key/value positions 0 … s (including the newly written column). In practice, the framework’s past_key_values argument handles this: when past_key_values is passed, the attention mechanism reads from the cached K/V and combines it with the new token’s Q, K, V so that the causal mask and cached storage are consistent.

Common mistake: Writing the new K/V at slot write_position + 1 instead of write_position. This shifts every subsequent read by one, causing the attending token to see a zero column (or wrap-around) and producing silently wrong logits. The trace above guards against this by always incrementing the cache length after the write is confirmed.

Tensor axes: query heads versus KV heads

In a model with n_query_heads query heads and n_kv_heads key/value heads:

Dimension	Meaning	Typical model example
n_query_heads	Number of Q projection heads, one per attention head	Llama-2-7B: 4
n_kv_heads	Number of K/V projection heads, shared across multiple Q heads (GQA/MQA)	Llama-2-7B: 1 (MQA) / Llama-2-70B: 8
Head dimension d	Size of each head’s key/value/query vector	Llama-2-7B: 64

The cache stores only the K/V dimensions, i.e. n_kv_heads × d per layer per token. The Q projection still operates on all n_query_heads, but the K/V it attends to come from the smaller KV cache.

GQA effect on cache size: If n_query_heads = 32 and n_kv_heads = 8, the cache stores 8 KV heads instead of 32, reducing KV bytes by 4×. Query computation is unchanged; only the K/V read path is smaller.

For the numerical example in this lesson, I’ll assume a plain-full-attention model where n_query_heads = n_kv_heads = H, so the cache dimension H equals the number of attention heads. This keeps the byte formula straightforward, and the GQA reduction can be discussed as a lever in the “What you can trade” section.

Cached versus uncached logit parity

Claim: If the same causal attention computation is performed (1) with the KV cache enabled and (2) without the cache (full sequence K/V passed explicitly), the next-token logits should agree to within floating-point tolerance.

Verification setup: Run a tiny decoder-only model through two forward passes:

1. Cached pass: model(input_ids, past_key_values=initial_past, use_cache=True) — the model reads K/V from past_key_values and appends the new token’s K/V.
2. Uncached pass: model(input_ids_full, use_cache=False) — the model receives the entire sequence including the new token, so attention computes Q/K attention over all tokens without a cache.

The logits from the last position should satisfy:

‖logits_cached − logits_uncached‖_∞ < ε

where ε depends on precision:

• FP32: ε ≈ 1e-5 (machine epsilon for float32)
• FP16: ε ≈ 1e-3 (reduced precision introduces rounding in softmax and matmul)

If the difference exceeds ε, check: (a) the cache write position is correct, (b) the causal mask aligns with the extended cache length, (c) the query/key/value projections have the same shape in both paths.

Source: This parity check is a standard regression test in Transformer implementations (e.g., Hugging Face transformers cache test suite) and is mentioned in the vLLM prefix caching design.

Raw cache bytes: assumptions and a checked numerical example

The raw KV cache size depends on several stated assumptions. Below is the formula and a worked example that you can reproduce with a calculator or a small Python script.

Formula

Let:

• L = number of transformer layers
• H = number of key/value heads per layer
• D = head dimension (the size of each key/value vector)
• S = total sequence length, including prompt and every generated token
• B = batch size (concurrent sequences)
• b = bytes per stored value (2 for FP16, 4 for FP32)

In a decoder where every layer uses full attention, each token needs a key and a value for every layer and every key/value head. The raw cache size is:

KV bytes = 2 × L × H × D × S × B × b

The leading 2 accounts for the key and value pair. This formula assumes:

1. Full attention in every layer — each token’s key and value are stored for all heads.
2. No GQA/MQA compression — H equals the number of query heads; if GQA is used, replace H with the number of KV heads.
3. Dense, unpadded storage — every position from 0 to S−1 occupies exactly one slot; no block/paged allocation savings are counted.
4. One sequence per batch item — batch dimension B multiplies the per-sequence cache linearly.
5. No allocator metadata, fragmentation, or workspace overhead — the formula gives the tensor bytes, not the GPU memory bill inclusive of the allocator.

Worked example: 32-layer model, 8 heads, D=128, 1K context, 1 sequence, FP16

Take a model with L = 32 layers, H = 8 KV heads, head dimension D = 128, a context of S = 1,024 tokens (prompt + generated), B = 1 sequence, and FP16 storage (b = 2 bytes).

Step	Calculation	Result
Per-token, per-layer cache size	2 × H × D × b = 2 × 8 × 128 × 2	4,096 bytes ≈ 4 KiB
Per-token, all layers	4 KiB × L = 4 KiB × 32	128 KiB
One full 1K sequence	128 KiB × S = 128 KiB × 1,024	128 MiB
Eight concurrent sequences	128 MiB × B = 128 MiB × 8	1 GiB

Check: 2 × 32 × 8 × 128 × 1,024 × 1 × 2 = 2 × 32 × 8 × 128 × 2,048 = 2 × 32 × 8 × 262,144 = 134,217,728 bytes = 128 MiB. ✓

The formula scales linearly: doubling the context length S doubles the cache; doubling the batch B doubles it; moving from FP16 to FP32 (b = 4) doubles it.

What changes with GQA?

If the same 32-layer model uses GQA with 32 query heads but only 4 KV heads (H = 4 in the formula), the cache halves relative to the H = 8 example above, because the formula scales linearly in H:

• FP16, S = 1K, B = 1: 128 MiB to 64 MiB
• FP16, S = 8K, B = 8: 8 GiB to 4 GiB

For reference, a model with no GQA at all (H = 32, one KV head per query head) would need 512 MiB in the 1K/FP16/B=1 case: 2 x 32 x 32 x 128 x 1,024 x 1 x 2 = 536,870,912 bytes.

The quality impact of reducing KV heads should be evaluated on the target task, as fewer heads may reduce the model’s capacity to store diverse key patterns.

Reproducible: per-layer KV cache update and logit parity check

The following self-contained PyTorch program builds a tiny two-layer decoder with an explicit per-layer KV cache, runs a cached and an uncached forward pass over the same tokens, and asserts that the next-token logits agree. It was executed for this lesson with torch 2.15.0, Python 3.12, CPU; the output below is the real run, not an illustration.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
torch.set_grad_enabled(False)

D, H, D_HEAD, L, V = 32, 2, 16, 2, 50


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.q = nn.Linear(D, H * D_HEAD)
        self.k = nn.Linear(D, H * D_HEAD)
        self.v = nn.Linear(D, H * D_HEAD)
        self.o = nn.Linear(H * D_HEAD, D)
        self.ln1 = nn.LayerNorm(D)
        self.ln2 = nn.LayerNorm(D)
        self.fc1 = nn.Linear(D, 4 * D)
        self.fc2 = nn.Linear(4 * D, D)

    def forward(self, x, past=None):
        B, S, _ = x.shape
        q = self.q(x).view(B, S, H, D_HEAD).transpose(1, 2)  # [B,H,S,Dh]
        k = self.k(x).view(B, S, H, D_HEAD).transpose(1, 2)
        v = self.v(x).view(B, S, H, D_HEAD).transpose(1, 2)
        if past is not None:
            k = torch.cat([past[0], k], dim=2)
            v = torch.cat([past[1], v], dim=2)
        S_all = k.shape[2]
        scores = (q @ k.transpose(-2, -1)) / (D_HEAD ** 0.5)
        q_pos = torch.arange(S_all - S, S_all)
        k_pos = torch.arange(S_all)
        blocked = k_pos[None, :] > q_pos[:, None]          # [S, S_all]
        scores = scores.masked_fill(blocked[None, None], float("-inf"))
        attn = scores.softmax(-1)
        out = (attn @ v).transpose(1, 2).reshape(B, S, H * D_HEAD)
        x = x + self.o(out)
        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))
        return x, (k, v)


class Tiny(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D)
        self.blocks = nn.ModuleList([Block() for _ in range(L)])
        self.ln_f = nn.LayerNorm(D)
        self.head = nn.Linear(D, V, bias=False)

    def forward(self, idx, past=None):
        x = self.emb(idx)
        new_past = []
        for i, blk in enumerate(self.blocks):
            x, layer_past = blk(x, None if past is None else past[i])
            new_past.append(layer_past)
        return self.head(self.ln_f(x)), new_past


model = Tiny()
prompt = torch.tensor([[1, 2, 3, 4]])                     # [B=1, P=4]

# Cached path: prefill the 4 prompt tokens, then feed one token against the cache
logits_prefill, past = model(prompt)
prefill_cache_len = past[0][0].shape[2]
next_id = logits_prefill[:, -1, :].argmax(-1, keepdim=True)   # first generated token
logits_decode, past = model(next_id, past=past)
decode_cache_len = past[0][0].shape[2]

# Uncached reference: run the S=5 sequence (prompt + the same generated token) in one pass
full_ids = torch.cat([prompt, next_id], dim=1)
logits_full, _ = model(full_ids)

# The parity claim: last cached-decode logits == last uncached logits
diff = (logits_decode[:, -1, :] - logits_full[:, -1, :]).abs().max().item()
print(f"cache length after prefill: {prefill_cache_len}, after one decode step: {decode_cache_len}")
print(f"cached next-token logits: {logits_decode[:, -1, :].shape}")
print(f"uncached next-token logits: {logits_full[:, -1, :].shape}")
print(f"max |cached - uncached| logit diff: {diff:.3e}")
print(f"parity within 1e-5: {diff < 1e-5}")

Actual output:

cache length after prefill: 4, after one decode step: 5
cached next-token logits: torch.Size([1, 50])
uncached next-token logits: torch.Size([1, 50])
max |cached - uncached| logit diff: 1.788e-07
parity within 1e-5: True

The cache grows by exactly one position per decode step (4 to 5), and the cached path matches the uncached reference to within 1.8e-07, far inside the FP32 tolerance discussed below.

Numerical tolerance note: The cached and uncached logit difference should be < 1e-3 in FP16 and < 1e-5 in FP32. If you see larger differences, verify that (a) the cache write position increments by exactly 1 per decode step, (b) the causal mask covers the full extended cache length, and (c) no head-dimension transpose error has occurred. The above script is a minimal reproduction; a full check would use a pretrained model (e.g., GPT-2) and compare model(..., use_cache=True).logits against model(..., use_cache=False).logits on the same token sequence.

Related lessons

• LLM Internals & Architecture · Inference foundationsUnderstanding the KV Cache12 min
• LLM Internals & Architecture · Generation, KV cache, and contextPrefill and Decode: Measuring the Two Phases11 min
• LLM Internals & Architecture · Generation, KV cache, and contextPretrained Weights and Text Generation12 min