KV Cache的显存瓶颈与优化：从GQA到MQA的演进逻辑

KV Cache的显存瓶颈与优化从GQA到MQA的演进逻辑一、自回归生成中99%的显存被KV Cache占据在大模型的推理阶段自回归生成过程逐个token生成有一个显著的不对称性计算量随生成长度线性增长每生成一个新token所有之前的token都要参与注意力计算但KV Cache的显存占用也随长度线性增长——而且这个增长的系数非常惊人。对于LLaMA-70B模型使用FP16精度batch_size1生成长度2048时模型权重约140GBKV Cache约2.6GB每层约33MB × 80层看起来KV Cache远小于权重但如果batch_size32生成长度8192KV Cache约332GB——远超模型权重这就是长序列推理的显存墙显存不是被模型权重占满的是被KV Cache吃光的。flowchart TB A[标准 MHA: 每个头独立的 KV] -- B[KV Cache 2 × B × H × L × D] B -- C{KV Cache 优化方向} C -- D[GQA: 分组查询注意力] D -- D1[K/V 头数 查询头数 / G] D -- D2[KV Cache 缩小 G 倍] D -- D3[精度损失: 1%] C -- E[MQA: 多查询注意力] E -- E1[所有查询头共享 1 组 K/V] E -- E2[KV Cache 缩小 H 倍] E -- E3[精度损失: 1-3%] C -- F[MLA: 多头潜在注意力 DeepSeek] F -- F1[KV 压缩到低秩潜在空间] F -- F2[KV Cache 缩小 5-10x] F -- F3[需要从预训练阶段引入] style D fill:#e8f5e9 style E fill:#fff3e0 style F fill:#e1f5fe二、从Multi-Head到Multi-Query的权衡分析标准 Multi-Head Attention (MHA) 中每个注意力头有独立的 Q、K、V 投影。这意味着对于 H 个头KV Cache 需要存储 H 组 K 和 V。Multi-Query Attention (MQA) 的激进方案是所有头共享同一组 K 和 V只有 Q 保持多头。这样 KV Cache 缩小了 H 倍。Grouped-Query Attention (GQA) 是 MHA 和 MQA 之间的折中将 H 个查询头分为 G 组每组共享一组 K 和 V。当 G1 时退化为 MQA当 GH 时即为 MHA。关键问题是这个折中在精度和效率上究竟如何import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Optional, Tuple class GroupedQueryAttention(nn.Module): GQA: 分组查询注意力的实现。 分组策略选择 - G1 (MQA): 极致节省显存适用于推理加速 - GH (MHA): 无节省适用于训练 - GH/8 (典型GQA): LLaMA-2 70B 使用此配置 为什么不是所有层都用相同的 G 实验表明底层靠近输入对 KV 冗余更敏感 高层可以使用更激进的 G 值更小的 G。 def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: int, # K/V 头数 num_heads dropout: float 0.0 ): super().__init__() self.hidden_size hidden_size self.num_heads num_heads self.num_kv_heads num_kv_heads self.head_dim hidden_size // num_heads # GQA 的头分组比例 self.num_queries_per_kv num_heads // num_kv_heads # Q 投影每个查询头独立 self.q_proj nn.Linear(hidden_size, num_heads * self.head_dim, biasFalse) # K, V 投影仅在 num_kv_heads 个头上 self.k_proj nn.Linear(hidden_size, num_kv_heads * self.head_dim, biasFalse) self.v_proj nn.Linear(hidden_size, num_kv_heads * self.head_dim, biasFalse) self.o_proj nn.Linear(num_heads * self.head_dim, hidden_size, biasFalse) self.dropout nn.Dropout(dropout) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] None, past_key_value: Optional[Tuple[torch.Tensor]] None, ) - Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]: bsz, seq_len, _ hidden_states.shape # 投影 q self.q_proj(hidden_states) k self.k_proj(hidden_states) v self.v_proj(hidden_states) # 重塑Q 为 [bsz, num_heads, seq_len, head_dim] q q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # K, V 为 [bsz, num_kv_heads, seq_len, head_dim] k k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) # 拼接历史 KV Cache if past_key_value is not None: k torch.cat([past_key_value[0], k], dim2) v torch.cat([past_key_value[1], v], dim2) # 保存当前 KV 用于下一步的 Cache current_key_value (k, v) # GQA 关键操作将 K/V 从 num_kv_heads 扩展到 num_heads # 通过在相邻组之间重复 K/V 头实现 k_expanded k.unsqueeze(1).expand( bsz, self.num_queries_per_kv, self.num_kv_heads, k.shape[2], self.head_dim ).reshape(bsz, self.num_heads, k.shape[2], self.head_dim) v_expanded v.unsqueeze(1).expand( bsz, self.num_queries_per_kv, self.num_kv_heads, v.shape[2], self.head_dim ).reshape(bsz, self.num_heads, v.shape[2], self.head_dim) # 标准缩放点积注意力 scale 1.0 / math.sqrt(self.head_dim) attn_weights torch.matmul(q, k_expanded.transpose(-2, -1)) * scale if attention_mask is not None: attn_weights attn_weights attention_mask attn_weights F.softmax(attn_weights, dim-1) attn_weights self.dropout(attn_weights) attn_output torch.matmul(attn_weights, v_expanded) # 重塑输出 attn_output attn_output.transpose(1, 2).reshape(bsz, seq_len, -1) attn_output self.o_proj(attn_output) return attn_output, current_key_value def estimate_kv_cache_size(self, batch_size: int, seq_length: int) - int: 估算给定配置下的 KV Cache 大小字节。 bytes_per_element 2 # FP16 kv_elements ( 2 # K V * batch_size * self.num_kv_heads # ← 关键这里是 num_kv_heads 而非 num_heads * seq_length * self.head_dim ) return kv_elements * bytes_per_element三、KV Cache 的其他优化方向除了 GQA/MQA 减少 KV 头数外还有几个正交的优化方向INT8/INT4 KV Cache 量化将 KV Cache 压缩到 8-bit 或 4-bit。挑战在于 KV 值的分布是动态变化的不同层、不同token位置差异大静态量化效果有限。KIVIKV-cache quantization with 2-bit在每个token位置使用独立的量化参数。滑动窗口注意力 全局token对于长文档场景大多数token只需要关注附近的token局部依赖只有少数全局token需要关注整个序列。Mistral 的 Sliding Window Attention 和 Longformer 的混合模式都属于这个类别。Token 淘汰策略EvictionStreamingLLM 的注意力汇聚点Attention Sink发现表明前几个 token 对注意力分布有不成比例的影响。据此可以设计策略保留前K个token 最近的M个token淘汰中间的不重要token。四、不同优化路线的适用场景预训练阶段GQA 需要从预训练开始就使用分组KV结构。如果在标准MHA预训练的模型上改为GQA需要额外的微调LLaMA-2采用的就是GQA从头训练。微调阶段可以在标准MHA基础上引入GQA结构并进行少量微调~1000步以快速降低推理成本。推理阶段如果无法修改模型结构可以使用KV Cache量化INT4或token淘汰策略等后处理方法。五、总结KV Cache 是长序列推理的核心瓶颈优化方向围绕减少每个token位置存储的KV量GQA 在精度损失1%和显存节省2-8x之间取得了最佳平衡是目前工业界的主流选择。MQA 提供极致的显存节省但精度损失较大适用于对显存极度受限的边缘设备。KV Cache 量化和 Token 淘汰策略与 GQA/MQA 正交可以叠加使用。优化方案的选择取决于能否修改模型结构这一前置条件。