摘要：按需备货：查询头保留独立配置，键值头按组共享，减少重复计算灵活分餐：用repeat_interleave魔法，让每个厨师都能拿到对应组的食材高效出餐：相比MHA，直接砍掉2/3的KV存储量，长文本场景显存暴降75%！
继上次聊透 MHA 后，今天和大家聊聊 ——Group Query Attention（GQA）！
曾几何时，这玩意儿简直是长文本处理的救星，用对了直接让模型跑得又快又省！
传统MHA就像12个厨师各开小灶，每个都要备全套食材（KV矩阵），但GQA直接玩起「分餐制」：
12个厨师（查询头）照常工作，但后厨只设3个食材组（键值头组），每组给4个厨师供菜！
核心原理就三步：
按需备货：查询头保留独立配置，键值头按组共享，减少重复计算灵活分餐：用repeat_interleave魔法，让每个厨师都能拿到对应组的食材高效出餐：相比MHA，直接砍掉2/3的KV存储量，长文本场景显存暴降75%！Llama3、QWen2.5等用的就是GQA，上图做了MHA、MQA、GQA的效果对比，可以看到效果还不错。
import torchimport torch.nn as nnclass GroupQueryAttention(nn.Module):def __init__(self, hidden_size, num_heads, group_num):super.__init__self.num_heads = num_heads # 12个厨师self.head_dim = hidden_size // num_heads # 每个厨师管64道菜self.group_num = group_num # 3个食材组# 四个关键操作台self.q_linear = nn.Linear(hidden_size, hidden_size) # 处理顾客订单self.k_linear = nn.Linear(hidden_size, group_num * self.head_dim) # 食材缩减备货self.v_linear = nn.Linear(hidden_size, group_num * self.head_dim) # 同上self.o_linear = nn.Linear(hidden_size, hidden_size) # 装盘上菜# 前向传播def forward(self, hidden_state, attention_mask=None):batch_size, seq_len, _ = hidden_state.size # 4桌客人，每桌256道菜# 步骤1：订单预处理query = self.q_linear(hidden_state) # 翻译订单需求key = self.k_linear(hidden_state) # 按组备货（键）value = self.v_linear(hidden_state) # 按组备货（值）# 步骤2：神奇分组术query = query.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 12个厨师各管64维key = key.view(batch_size, seq_len, self.group_num, self.head_dim).transpose(1, 2) # 3组食材key = key.repeat_interleave(self.num_heads//self.group_num, dim=1) # 每组复制4次供12人# 步骤3：计算关联性attn_scores = torch.matmul(query, key.transpose(-1, -2)) / torch.sqrt(torch.tensor(self.head_dim))# 厨师们对比256道菜的关联度，缩放评分防止数值爆炸# 步骤4：特殊要求处理（掩码）if attention_mask is not None: attn_scores += attention_mask * -1e9# 步骤5：确定优先级attn_probs = torch.softmax(attn_scores, dim=-1) # 算出每道菜参考其他菜的权重# 步骤6：组合菜品context = torch.matmul(attn_probs, value) # 按权重混合食材# 步骤7：装盘上菜context = context.transpose(1, 2).contiguous.view(batch_size, seq_len, -1)return self.o_linear(context) # 微调后上菜分组策略：合理选择键值头分组数 G 至关重要，它直接影响计算量和显存占用。分组数越多，KV 存储量越少，但可能会牺牲一定的模型表达能力。维度匹配：在进行张量重塑和维度变换时，要确保各维度之间的匹配关系正确。特别是键值向量的扩展操作 repeat_interleave，需要根据分组数和查询头数量准确设置参数，以保证每个查询头都能获取到正确的键值信息。与 MHA 对比：理解 GQA 与传统 MHA 的差异，有助于更好地把握 GQA 的优势和适用场景。GQA 在保持多查询头灵活性的同时，通过共享键值头降低了计算复杂度和显存占用，尤其适合处理长序列数据。实际应用调整：在实际应用中，需要根据具体任务和硬件资源来调整 GQA 的参数。例如，在移动端或内存受限的设备上，可以适当增加分组数以减少显存占用。✅ 优点：
显存暴减：相比传统 MHA，GQA 通过共享键值头组，直接砍掉 2/3 的 KV 存储量。在长文本场景下，比如处理万字论文时，KV 缓存能从 16MB 直降到 4MB，简直是显存不足党的福音！推理加速：减少了键值向量的重复计算，模型推理速度大幅提升，尤其在处理超长上下文时优势明显，再也不用苦等模型「龟速」输出。设备友好：对硬件资源要求更低，在手机端、边缘设备等内存受限的场景下，也能流畅运行 AI 模型，让小设备也能玩转大模型！性能平衡：保留了多头注意力的灵活性，在减少计算量的同时，仍能保持不错的模型性能，兼顾效率与效果。❌ 缺点：
表达能力受限：由于键值头分组共享，相比 MHA 可能会丢失一些细节信息，在某些对特征捕捉精度要求极高的任务（如复杂语义理解、专业领域知识建模）中，模型表达能力稍显不足。参数敏感：分组数等关键参数的设置对性能影响较大，如果选择不当，可能无法充分发挥 GQA 的优势，甚至导致性能下降，需要反复调参优化。适用场景局限：在短序列任务或对计算量要求不高的场景中，GQA 相比 MHA 的优势不明显，反而可能因为额外的分组逻辑增加不必要的计算开销 。好了，学会GQA等于掌握大模型的「Attention变体黑魔法」！点赞收藏这篇硬核干货，想看更多大模型知识？评论区告诉我，下期接着盘！
来源：py大模型飞飞