考察点
这道题直接指向推理成本的核心矛盾——KV cache。面试官想确认你知道自回归推理的瓶颈不是算力而是显存带宽和容量,四种结构都是围绕「压缩 KV cache」做的不同取舍。能算出各变体的 KV cache 大小、说出 MLA 的低秩压缩思路,就算答到位了。追问常往 KV cache 显存计算、为什么 GQA 成为主流折中、MLA 的实现细节上走。
参考答案
背景:KV cache 才是推理瓶颈
自回归推理时每生成一个 token,都要把历史所有 token 的 K、V 读一遍参与注意力计算。为了不重复计算,这些 K、V 缓存在显存里,这就是 KV cache。它的体积 = 2(K 和 V)× 层数 × KV head 数 × head 维度 × 序列长度 × 批大小 × 每参数字节数。长上下文 + 高并发场景下,KV cache 轻松超过模型权重本身的显存占用,而且每步都要从 HBM 全量读一遍,直接决定解码速度。四种注意力变体的差别,就在于 K、V 怎么存。
四种结构
MHA(Multi-Head Attention):标准多头注意力,Q、K、V 的头数相同且一一对应。假设 32 层、32 个头、head 维度 128,每个 token 的 KV cache 是 2 × 32 层 × 32 头 × 128 维。效果基准线,但 KV cache 最大。
MQA(Multi-Query Attention):所有 Q 头共享同一组 K、V,即 KV head 数 = 1。KV cache 直接缩小到 1/32(按上面配置),推理带宽收益巨大。代价是质量明显下降——K、V 的表达能力被压到极限,长序列上尤其吃亏。Falcon、PaLM 早期版本用过。
GQA(Grouped-Query Attention):MHA 和 MQA 的折中,把 Q 头分组,每组共享一组 K、V。比如 32 个 Q 头配 8 个 KV 头(每 4 个 Q 头共享),KV cache 缩到 1/4。实测质量几乎对齐 MHA,是收益/代价比最好的方案,LLaMA-2/3、Qwen2/3、Mistral 全线采用,已成事实标准。
MLA(Multi-head Latent Attention):DeepSeek-V2/V3 的方案,思路更进一步——不砍 head 数,而是对 KV 做低秩压缩。把 K、V 联合压缩到一个低维的 latent 向量(比如 512 维)存进 cache,用时再通过上投影矩阵还原出各头的 K、V。KV cache 从「头数 × head 维度」降到「latent 维度」,压缩率比 GQA 还高(DeepSeek-V2 官方口径 KV cache 降到原来的十几分之一),而且因为每个头仍能生成不同的 K、V,质量反而优于 GQA。工程上还配了个技巧:RoPE 直接加在低维分量上,避免还原大矩阵时再算旋转。
对比一览
| 方案 | KV head | KV cache 相对大小 | 质量 | 代表模型 |
|---|---|---|---|---|
| MHA | = Q 头数 | 1(基准) | 最好 | 早期 GPT、BERT |
| GQA | Q 头数的 1/4~1/8 | 1/4 ~ 1/8 | ≈ MHA | LLaMA、Qwen |
| MQA | 1 | 1/头数 | 略降 | Falcon |
| MLA | 压缩为低秩 latent | 可低至 ~1/16 | ≈ 或优于 GQA | DeepSeek-V2/V3 |
优化目标的本质
这一族技术都在回答同一个问题:自回归生成是 memory-bound 的,KV cache 的「容量占用」决定能开多大并发和上下文,「读取带宽」决定每 token 延迟。压缩 KV cache 换来的不只是显存,还有吞吐量——同样的卡能服务更多并发请求,这是在线推理服务成本的大头。训练时这些变体的计算量差别不大,收益全在推理侧。面试时把这个「容量 + 带宽」双目标讲出来,说明你是从系统视角理解的,不是背结构图。
可能的追问
- GQA 为什么质量几乎不掉? Q 的多样性保留了,只有 K、V 被共享;实证上 KV 头数砍到 1/8 内,下游指标差异在噪声范围内。
- MLA 的上投影会不会抵消推理收益? 还原 K、V 增加计算,但推理是 memory-bound,减少显存读取的收益远大于额外 FLOPs;且可用矩阵吸收技巧把上投影并进 Q 侧,不实际物化大矩阵。
- 除了改注意力结构,还有什么压 KV cache 的手段? KV cache 量化(INT8/INT4)、分层/逐出策略(H2O、StreamingLLM 保留 sink + 滑窗)、prefix cache 复用(多请求共享 system prompt 的 KV)。