注意力变体与 KV 缓存压缩

采用全因果注意力进行自回归解码时，每个新 token 都会使用此前各 token 的 Key 和 Value；滑动窗口等注意力变体则会限制访问范围。缓存这些张量避免了重复计算，但代价是内存占用随序列长度和 Batch Size 线性增长。对于长上下文或小 Batch 解码场景，读取缓存往往成为主要的带宽瓶颈。

目前的优化手段主要围绕三个维度展开：减少 KV 头数量、压缩缓存表示的宽度，或者降低缓存 Token 的数量/精度。

查看清晰大图

从底部蓝色查询头往上看：中间每两个查询头共享一组键值头，右侧所有查询头共享一组。查询头数量没有减少；其他维度不变时，缓存节省来自需要保存的键和值变少。

Ainslie et al., GQA (2023), Figure 2. · CC BY 4.0

MHA 与 GQA

假设模型有 $L$ 层，序列长度 $S$，Query 头数 $n_h$，KV 头数 $n_{kv}$，头维度 $d_h$，每个缓存元素占用 $b$ 字节。忽略分配器和元数据开销，单条序列的 KV 缓存大小约为：

$$2L S n_{kv} d_h b.$$

系数 2 对应 Key 和 Value 两部分。

• 多头注意力 (MHA)：通常 $n_{kv}=n_h$，即每个 Query 头对应独立的 KV 头。
• 多查询注意力 (MQA)：所有 Query 头共享同一个 KV 头。
• 分组查询注意力 (GQA)：介于两者之间，多个 Query 头共享一个 KV 头。

在其他参数固定的前提下，从 MHA 切换到 GQA，缓存内存占用会按 $n_{kv}/n_h$ 的比例下降。但质量、吞吐量及内核支持是否达标，取决于具体的训练模型和服务栈。GQA 是训练阶段的架构选择，并非可以无损套用到任意权重的开关。

举一个假设配置：$L=32$、$S=4096$、$n_h=32$、$n_{kv}=8$、$d_h=128$，每个元素两字节，单序列缓存为 $536{,}870{,}912$ 字节，即 512 MiB。若 MHA 有 32 个 KV 头，则需 2 GiB；四条等长 GQA 序列也需 2 GiB，尚未计额外开销。这只是缓存，不是整个模型的内存需求。

什么条件下可以精确复用？

评估模式下的因果模型追加 token 时，之前的 token 表示不会改变。每层保存自己的 K/V，只需为新位置计算 query。精确复用要求分词后的前缀、权重（含适配器）、位置索引、注意力规则和数值表示都相同。修改早期 token 会使下游缓存状态失效。双向注意力通常无法在追加输入后原样复用旧状态，因为旧位置可能需要关注新 token。

若缓存已有 $S$ 个 token，新片段有 $q$ 个 token，片段内从零计数的 query 行 $i$ 可以访问到绝对 key 位置 $S+i$。直接在 $q\times(S+q)$ 矩阵左上角套三角掩码，会错误屏蔽有效缓存。剔除 token 后要保留原始位置索引，重新编号会改变位置注意力。先比较相同前缀下有无缓存的 logits，再测试近似压缩；比较时应容许内核形状或归约顺序变化造成的浮点误差。

树形推测解码把状态管理扩展到候选分支：需要隔离各分支，并沿接受路径保留一致的 KV 与压缩状态。

多头潜在注意力 (MLA)

DeepSeek-V2 引入了多头潜在注意力 (MLA)。它不再缓存完整的每头 Key/Value 张量，而是学习一个低维的潜在表示（Latent Representation）。由于旋转位置编码（RoPE）依赖位置信息，无法简单折叠进固定投影，因此用于 RoPE 的 Key 部分保持独立。

简化的缓存宽度对比如下：

$$\begin{aligned} \text{MHA/GQA 每层每 token 宽度} &\approx 2n_{kv}d_h,\\ \text{MLA 每层每 token 宽度} &\approx d_c+d_R, \end{aligned}$$

其中 $d_c$ 为压缩后的潜在宽度，$d_R$ 为解耦的旋转 Key 宽度。实际的内存和计算路径取决于模型具体的投影矩阵及推理实现。MLA 能显著降低缓存流量，但其公开效果仅适用于采用该架构训练的模型。

现有缓存的压缩策略

SnapKV 和 PyramidKV 等方法不改变训练好的注意力架构，而是直接操作模型缓存中保留的 Token。

• SnapKV：利用观察窗口内的注意力模式对 Prompt 位置评分，仅保留选定的历史 Key/Value 子集。
• PyramidKV：基于不同层在 Prompt 上注意力分布的差异，为不同层分配不同的缓存预算。
• KV 缓存量化：以更低精度存储缓存元素。
• 窗口/滑动注意力：仅在支持该设计的层中保留有限的近期区域。

这些技术不可互换。Token 剔除可能丢失虽小但关键的片段；量化可能改变注意力分数；自定义布局可能需要专用内核。论文中的准确率或困惑度结果仅适用于评估时的特定模型、任务、预算和实现，不能直接推广到所有长上下文工作负载。

关键测量指标

针对具体部署，需记录以下数据：

1. 确切模型、注意力架构及缓存数据类型；
2. Prompt 长度、生成长度及并发序列数；
3. GPU 峰值内存及主机内存占用；
4. Prefill 时间、首 Token 延迟 (TTFT) 及解码吞吐量；
5. 需要跨位置获取信息的长上下文任务成功率；
6. 缓存剔除、量化或卸载至 CPU 后的行为表现。

宣称支持长上下文的模型在现有硬件上可能无法使用，而节省缓存的方法虽然能装入内存，却可能导致任务失败。内存适配性和答案质量保留情况需要分别测试。

我的阅读 · 打开过的笔记会记在当前浏览器里。