12.1 KV Cache Management

Autoregressive serving has two different hardware regimes. Prefill processes the prompt in parallel and is commonly limited by matrix-multiplication throughput. Decode produces one token per sequence at a time; it repeatedly reads weights and attention state, so memory bandwidth, cache capacity, and scheduling often dominate. The KV cache is the state that connects those phases.

The engineering question is not merely “does the cache fit?” A useful design must preserve model quality, keep inter-token latency predictable under mixed sequence lengths, and recover capacity when requests finish.

1. What the Cache Actually Saves

At layer $l$, attention projects token state $x_t$ into

During causal decoding, keys and values of earlier tokens do not change. Retaining $\{k_i,v_i\}_{i\le t}$ means step $t+1$ computes only the new projections and attends to cached history:

where $d_h$ is head dimension and $M$ is a causal or padding mask.

Complexity needs careful wording. Without a cache, decode step $t$ reruns a length-$t$ forward pass: self-attention score work is $O(t^2)$. Summed across $N$ newly generated tokens, that attention work is $O(N^3)$, while projections and MLP work sum to $O(N^2)$. With a cache, attention at step $t$ is $O(t)$, so its total is $O(N^2)$. These expressions suppress model width and batch dimensions; wall-clock speed is usually governed by memory traffic and kernel utilization, not asymptotic notation alone.

2. Capacity Model

For a dense cache with equal-length sequences,

where $B$ is concurrent sequences, $S$ cached tokens per sequence, $L$ layers, $H_{KV}$ KV heads, $d_h$ values per head, and $p$ bytes per value. The factor two accounts for keys and values.

For heterogeneous traffic, replace $BS$ with actual allocated token count $\sum_i S_i^{\text{allocated}}$. Padding, block metadata, allocator slack, temporary workspaces, speculative branches, beam hypotheses, and tensor-parallel replication or sharding are separate costs.

For 32 sequences, 8,192 tokens each, 80 layers, 8 KV heads, $d_h=128$, and BF16:

This is the aggregate for 32 equal-length sequences, not one request. Nor does the cache always exceed the weights. It can do so at large concurrency or context length, so both weights and peak live-cache occupancy belong in the capacity plan.

3. Reducing KV Width

Let $H_Q$ be query heads.

• Multi-Head Attention (MHA): $H_{KV}=H_Q$.
• Multi-Query Attention (MQA): one shared KV head.
• Grouped-Query Attention (GQA): $1<H_{KV}<H_Q$, with multiple query heads sharing each KV head.

Payload scales linearly with $H_{KV}$. GQA therefore exposes a quality–capacity–bandwidth trade-off. The original study reported quality close to MHA for its evaluated converted and trained models while retaining speed near MQA [1]. That evidence does not guarantee that an arbitrary checkpoint can change head structure without retraining.

Multi-Head Latent Attention

DeepSeek-V2 introduced Multi-Head Latent Attention (MLA) [2]. MLA caches a low-rank latent $c_t^{KV}$, plus the position-dependent component required by decoupled RoPE, instead of a conventional per-head KV payload. Linear key/value up-projections can be absorbed algebraically into adjacent projections during decode. RoPE is separate because a position-dependent rotation breaks simple static absorption.

The DeepSeek-V2 report states a 93.3% KV-cache reduction and 5.76-times maximum generation throughput relative to its DeepSeek 67B comparison system. Those are reported system results for that architecture and implementation—not universal MLA multipliers. Kernel support, latent width, positional encoding, parallel layout, and traffic shape determine realized gain.

4. Cache Precision Is a Systems Trade-off

FP8 storage halves dense payload relative to FP16/BF16; four-bit storage quarters it. Capacity gains are deterministic, but latency and quality gains are not. A serving engine may spend instructions converting values, calculating scales, or moving metadata. Workloads that are not KV-bandwidth-bound may see little speedup.

Quantization also changes the numerical contract. Scale granularity, calibration data, format such as E4M3 versus E5M2, outliers, and accumulation precision matter. vLLM’s documentation recommends calibrated scales when possible and asks users to evaluate workload accuracy [3].

A sound rollout compares:

• quality by context-length slice;
• prefill latency, p50/p95 time to first token, and p50/p95 inter-token latency;
• throughput at a fixed latency service-level objective;
• cache occupancy, allocation failures, and eviction or recomputation rate.

“Twice as many cache tokens fit” is not the same as “tokens are generated twice as fast.”

5. An Educational GQA Cache

The following kernel demonstrates cache shape and causal masking, including chunked decode where cached tokens precede several new tokens.

Deliberately simplified implementation. It omits RoPE, padding masks, sliding windows, quantized storage, tensor parallelism, paged allocation, prefix sharing, and an in-place writer. It uses torch.cat, so it is unsuitable for production. The enable_gqa flag also depends on the installed PyTorch backend.

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

KVCache = tuple[torch.Tensor, torch.Tensor]  # [B, H_kv, S, d_h] each


class CausalGQAAttention(nn.Module):
    def __init__(self, d_model: int = 4096, n_heads: int = 32, n_kv_heads: int = 8):
        super().__init__()
        if d_model % n_heads or n_heads % n_kv_heads:
            raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero")
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.d_head = d_model // n_heads
        self.q_proj = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    @torch.no_grad()
    def forward(self, x: torch.Tensor, cache: KVCache | None = None):
        batch, new_tokens, width = x.shape
        q = self.q_proj(x).view(batch, new_tokens, self.n_heads, self.d_head)
        k = self.k_proj(x).view(batch, new_tokens, self.n_kv_heads, self.d_head)
        v = self.v_proj(x).view(batch, new_tokens, self.n_kv_heads, self.d_head)
        q, k, v = (z.transpose(1, 2) for z in (q, k, v))

        past_tokens = 0
        if cache is not None:
            past_k, past_v = cache
            past_tokens = past_k.size(2)
            k = torch.cat((past_k, k), dim=2)
            v = torch.cat((past_v, v), dim=2)

        # True means “allowed” for PyTorch SDPA boolean masks.
        q_pos = torch.arange(past_tokens, past_tokens + new_tokens, device=x.device)
        k_pos = torch.arange(k.size(2), device=x.device)
        allowed = (k_pos[None, :] <= q_pos[:, None])[None, None, :, :]

        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=allowed, is_causal=False, enable_gqa=True
        )
        y = y.transpose(1, 2).contiguous().view(batch, new_tokens, width)
        return self.o_proj(y), (k, v)

The explicit position mask matters. Naively setting is_causal=True for a multi-token chunk with a non-empty cache can align the triangular mask to the wrong origin, depending on the kernel contract.

6. Production Cache Contract

A production engine preallocates blocks and writes new KV rows in place. PagedAttention applies virtual-memory-style block mapping so sequences do not require one contiguous allocation and blocks can be shared for common prefixes [4]. The next section develops that allocator.

Before launch, record:

• Artifact: model, tokenizer, positional encoding, cache dtype/scales, kernel/engine versions, block size, parallel layout.
• Capacity: weights, cache, activations, workspace, CUDA graphs, fragmentation, and failure reserve separately.
• Telemetry: logical versus allocated tokens, block utilization, prefix hit rate, eviction/preemption, recomputed tokens, HBM use, batch tokens, TTFT, inter-token latency, output tokens/s.
• Failure policy: stop admission or shed low-priority work before allocation failure; roll back a format when quality slices regress or latency tails exceed the declared window.
• Recovery: verify token continuation after preemption or migration; do not assume a serialized KV tensor is portable across model or kernel revisions.

7. Takeaway

KV caching exchanges redundant computation for persistent memory traffic. MQA/GQA/MLA reduce width, lower precision reduces bytes per value, and paged allocation reduces waste. None is an isolated switch: quality, kernels, scheduling, and request-length distribution decide whether a smaller cache becomes a faster, reliable service.

Quizzes

References

1. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.
2. DeepSeek-AI. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.
3. vLLM Project. Quantized KV Cache. Official documentation.
4. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.