从注意力公式出发，理解 KV Cache

示例文章：用来展示公式、代码、表格和脚注的排版，换成你自己的文章后删掉即可。

为什么会有重复计算

自回归生成时，第 $t$ 步的 query 要和前面所有位置的 key、value 做注意力：

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V

关键在于：前面位置的 $K$ 和 $V$ 只取决于那些位置自己的输入，生成新 token 时它们不会变¹。既然不变，就没必要每一步都重算一遍。

代码里长什么样

去掉多头和批次维度，核心只有几行：

def attend(q, k_new, v_new, cache):
    # 把这一步新算出的 K、V 追加进缓存
    cache.k = torch.cat([cache.k, k_new], dim=-2)
    cache.v = torch.cat([cache.v, v_new], dim=-2)
    scores = q @ cache.k.transpose(-2, -1) / math.sqrt(q.size(-1))
    return scores.softmax(dim=-1) @ cache.v

代价：显存

缓存省下了计算，但要占显存。每个 token 需要存的字节数是：

\text{bytes/token} = 2 \cdot L \cdot h \cdot d_h \cdot b

以 Llama 2 7B、FP16 为例（ $L = 32$ ， $h = 32$ ， $d_h = 128$ ， $b = 2$ ）：

上下文长度	KV Cache 大小	相对模型权重
1 token	0.5 MiB	≈ 0%
4,096 tokens	2 GiB	≈ 16%
4,096 × 批次 32	64 GiB	≈ 5 倍

旁注　这就是 MQA、GQA 要解决的问题：让多个 query 头共享同一组 $K$ 、 $V$ ，把上式里的 $h$ 变小。下一篇会讲。

在使用因果掩码（causal mask）的解码器里成立；编码器的双向注意力不适用。 ↩

Sample article: it shows how formulas, code, tables and footnotes are typeset. Replace it with your own and delete this one.

Where the repeated work comes from

In autoregressive generation, the query at step $t$ attends to the keys and values at every earlier position:

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V

The key point: the $K$ and $V$ of earlier positions depend only on those positions’ own inputs, so they don’t change when a new token is generated¹. If they don’t change, there’s no reason to recompute them at every step.

What it looks like in code

Strip away the heads and batch dimensions, and the core is only a few lines:

def attend(q, k_new, v_new, cache):
    # Append this step's new K and V to the cache
    cache.k = torch.cat([cache.k, k_new], dim=-2)
    cache.v = torch.cat([cache.v, v_new], dim=-2)
    scores = q @ cache.k.transpose(-2, -1) / math.sqrt(q.size(-1))
    return scores.softmax(dim=-1) @ cache.v

The cost: memory

The cache saves compute, but it takes memory. Each token needs this many bytes:

\text{bytes/token} = 2 \cdot L \cdot h \cdot d_h \cdot b

Take Llama 2 7B in FP16 as an example ( $L = 32$ , $h = 32$ , $d_h = 128$ , $b = 2$ ):

Context length	KV cache size	Relative to model weights
1 token	0.5 MiB	≈ 0%
4,096 tokens	2 GiB	≈ 16%
4,096 × batch of 32	64 GiB	≈ 5×

Aside　This is the problem MQA and GQA solve: several query heads share one set of $K$ and $V$ , which shrinks $h$ in the formula above. That’s the next post.

This holds for decoders with a causal mask; it doesn’t apply to an encoder’s bidirectional attention. ↩