从Attention到GQA：大模型Transformer架构的核心机制与显存优化演变

一、引言：为什么大模型推理这么吃显存？

如果你部署过大模型（比如LLaMA 2、Qwen），一定遇到过这样的问题：明明模型权重已经加载完毕，一跑长文本推理，显存直接OOM。很多开发者的第一反应是“加卡”，但多加几块A100的成本实在不菲。

问题的根源，其实藏在Transformer的自注意力机制里——更准确地说，藏在KV Cache里。

本文将从多头注意力（MHA）的原理出发，逐步剖析为什么大模型推理会遭遇“显存墙”，以及GQA（分组查询注意力）如何成为当前主流大模型（LLaMA 2/3、Qwen-2、Mistral等）的共同选择。全程配合PyTorch代码实现，力求讲透、讲深。

---

二、从RNN到Transformer：为什么需要注意力机制

在Transformer出现之前，NLP领域的主流是RNN（循环神经网络）及其变体LSTM、GRU。RNN的本质是串行处理：每个时间步的输出依赖前一个时间步的隐藏状态。这种设计带来了两个致命问题：

第一，长距离依赖难以捕捉。 当序列长度超过100个token时，最早的输入信息经过多层循环传递后已经衰减得所剩无几，模型很难建立起长距离的语义关联。

第二，训练无法并行。 因为每个时间步必须等前一步计算完成才能开始，GPU的并行计算能力被严重浪费，训练大规模模型的时间成本高得离谱。

Transformer提出的自注意力机制彻底改变了这一切。它让序列中的每个token都可以直接关注到所有其他token，无论距离多远，语义关联都能一步到位。更重要的是，所有token的注意力计算可以完全并行，这为后续大模型的爆发铺平了道路。

自注意力的核心公式是：

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$

其中Q（Query）代表当前token要“问什么”，K（Key）代表其他token能“提供什么”，V（Value）代表其他token携带的“实际信息”。通过Q和K的点积算出注意力权重，再用这个权重去加权V，就完成了信息聚合。

---

三、Multi-Head Attention（MHA）：多头带来的能力跃升

3.1 为什么需要“多头”

单头注意力虽然能让token关注到其他位置，但它只能从一个角度去理解语义关系。比如一个句子中，“苹果”这个词可能同时和“吃”有动作关系、和“红色”有属性关系、和“水果”有类别关系。单头注意力只能捕捉其中一种。

多头注意力的做法是：并行运行多组独立的注意力，每组（称为一个“头”）都从不同的子空间去建模语义关系，最后把所有头的结果拼接起来。数学表达如下：

$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,...,{\text{head}}_h)W^O$$

其中每个头的计算方式是：${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$。每个头有自己的Q、K、V投影矩阵，互不干扰。

假设有8个头，那么有的头可能专门关注“主谓宾”的语法结构，有的头关注指代关系（代词指向前文的名词），有的头关注相邻词汇的修饰关系。这种多角度的信息抽取能力，让大模型真正具备了理解复杂语义的基础。

3.2 MHA的“显存陷阱”：KV Cache

理解了MHA的原理，我们再来看工程落地的核心问题。

在大模型的自回归生成场景中（也就是我们平时和ChatGPT对话时的逐字输出），模型是逐个token地生成。每生成一个新token，都需要计算它和前面所有token的注意力关系。

如果每次生成新token时都重新计算前面所有token的K和V，计算量会呈平方级增长——生成第N个token时要算N次K和N次V，总复杂度是$O(N^2)$。对于动辄生成几千个token的场景，这个计算量在工程上完全不可接受。

所以业界引入了KV Cache（键值缓存）机制：把已经生成过的token对应的Key和Value缓存下来。生成新token时，只计算新token的Q、K、V，然后把新token的K和V追加到缓存中，复用缓存的K和V去计算注意力。

这个机制在计算上是高效的，但在显存上却代价高昂。

KV Cache的大小计算公式为：

$$\text{KV Cache大小}=2\times \text{batch\_size}\times \text{num\_layers}\times \text{num\_heads}\times \text{head\_dim}\times \text{seq\_len}\times \text{字节数}$$

最核心的变量是 num_heads——也就是注意力头的数量。每个头都要缓存自己的K和V，头越多，缓存越大。

以LLaMA-2 70B为例：64个Q头（在MHA中也就意味着64个KV头）、80层、head_dim=128，每个参数用FP16（2字节）存储。

• 单token单层的KV Cache = 2（K和V）× 64（头）× 128（维度）× 2（字节）= 32,768字节 ≈ 32KB
• 80层总计 = 32KB × 80 = 2.56MB
• 对于一个4096长度的输入序列：2.56MB × 4096 ≈ 10.5GB

这只是KV Cache的占用，还没算模型权重（70B参数量约140GB，用FP16存储）和中间激活值。一个推理请求就要吃掉10GB+的KV Cache，如果batch_size=8，单是KV Cache就逼近100GB。

这就是大模型长文本推理的“显存第一堵墙”。 也正是这堵墙，推动着学术界和工业界去探索更高效的注意力变体。

---

四、MQA（Multi-Query Attention）：极致压缩的先行者

4.1 MQA的暴力压缩思路

面对MHA的显存困境，Google在2019年提出了一个极其简洁的解决方案——MQA（Multi-Query Attention）。思路可以一句话概括：所有Q头共享同一组K和V，即强制 num_kv_heads = 1。

也就是说，不管你有64个Q头还是128个Q头，K和V只各有一份。这意味着模型不再允许不同头从不同子空间去捕捉K和V的语义，所有头看到的是完全相同的K和V，只是在Q的投影上保留多头差异。

4.2 收益和代价同样显著

显存收益是爆炸性的。 沿用上面的LLaMA-2 70B数据，MQA将KV头数从64压到1，KV Cache从10.5GB直接降到约164MB，压缩到原来的1/64。推理速度受益于此，可以提升8到12倍。

但代价同样沉重。 所有Q头共享K和V，相当于“多个人共用同一双眼睛来看世界”。虽然每个头还保留着自己独特的Q投影，但信息来源被严重限制了。模型在捕捉多样化语义关系时能力明显下降，具体表现为：在多任务基准测试中，MQA相比MHA平均有2%到5%的精度损失，尤其是在需要细粒度语义理解的复杂推理任务上。

更麻烦的是，MQA模型必须从头训练，不能从已训练好的MHA模型做简单转换就得到。这意味着如果一家公司已经花了几千万美金训练了一个MHA大模型，想换成MQA来省推理成本，只能重新训练——这在时间和资金上都是巨大的门槛。

MQA的困境在于：压缩得够狠，但精度掉得太多，且无法“中途上车”。

---

五、GQA（Grouped Query Attention）：折中之道，兼顾两边

5.1 核心思想：在MHA和MQA之间画一条连续谱

2023年，Google DeepMind的Ainslie等人在论文中正式提出了GQA（Grouped Query Attention）。核心思想非常优雅：不非此即彼，而是在MHA和MQA之间连续调节。

具体做法是：将Q头分成若干组，每组内的所有Q头共享一组K和V。

配置参数用两个数字描述：

• num_q_heads = h：Q头的总数
• num_kv_heads = g：KV头（即组）的总数，要求 g 能整除 h

每组包含 $h/g$ 个Q头，这 $h/g$ 个Q头共享1组KV。

在这个框架下：

• 当 $g=h$ 时，每个Q头单独配一组KV，退化为MHA
• 当 $g=1$ 时，所有Q头共享一组KV，退化为MQA
• 当 $1<g<h$ 时，就是真正的GQA，在两者之间取平衡

这意味着GQA不是一个固定的结构，而是一整个设计空间。工程师可以根据自己的显存预算和精度要求，自由选择合适的 $g$ 值。这种灵活性在工程落地中极其宝贵。

5.2 显存收益有多实在？

还是用LLaMA-2 70B做参照，该模型的实际配置就是GQA，且 num_q_heads=64，num_kv_heads=8，head_dim=128，80层。

• MHA（64个KV头）：KV Cache = 32KB/token/层，4096长度约10.5GB
• GQA（8个KV头）：KV Cache = 2 × 8 × 128 × 2字节 = 4KB/token/层，4096长度约 1.3GB
• 压缩到原来的 1/8

再看几个主流模型的实际配置：

• LLaMA-2 70B：num_q_heads=64，num_kv_heads=8
• LLaMA-3 70B：num_q_heads=64，num_kv_heads=8
• Qwen-2 72B：num_q_heads=64，num_kv_heads=8
• Mistral 7B：num_q_heads=32，num_kv_heads=8
• DeepSeek-V1：num_q_heads=64，num_kv_heads=4

你会发现，主流7B到72B级别的大模型，KV头数量几乎都落在4到8之间。这说明业界经过大量实验验证，GQA-8在精度损失（通常控制在1%以内）和显存节省（8倍压缩）之间达成了最优平衡。

更关键的是，GQA支持从MHA模型做转换：将MHA中多个头的K和V做分组平均，即可得到GQA的初始权重。这使得已训练好的MHA模型无需从头再训，通过少量继续预训练（约训练总量的5%到10%）就能转为GQA结构。对于大模型厂商来说，这是决定性的优势——省下了数千万美金的重新训练成本。

---

六、手写GQA：PyTorch完整实现（含KV Cache模拟）

原理讲透之后，代码必须跟上。下面是一个生产级可用的GQA实现，与LLaMA系列的实际代码对齐，并且增加了KV Cache的模拟逻辑，让你直观看到显存节省发生在哪里。

import torch
import torch.nn as nn
import math
from typing import Optional, Tuple

class GroupedQueryAttention(nn.Module):
    def __init__(
        self,
        d_model: int,              # 模型维度
        num_heads: int,            # Q头总数
        num_kv_heads: int,         # KV头组数（GQA核心参数）
        dropout: float = 0.0,
        max_seq_len: int = 4096    # 最大序列长度，用于预分配缓存
    ):
        super().__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        assert num_heads % num_kv_heads == 0, "num_heads必须能被num_kv_heads整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = d_model // num_heads
        
        # Q投影：完整映射到 d_model（所有Q头）
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        # K投影：只映射到 num_kv_heads * head_dim（参数量大幅缩减）
        self.W_k = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        # V投影：同K
        self.W_v = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        # 输出投影
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
        
        # 用于缓存KV（推理时使用）
        self.max_seq_len = max_seq_len
        self.cache_k = None  # 将在第一次forward时初始化
        self.cache_v = None
        self.step = 0  # 当前生成步数
        
    def _init_cache(self, batch_size: int, device: torch.device):
        """初始化KV Cache，形状为 (batch, num_kv_heads, max_seq_len, head_dim)"""
        # 注意：缓存的是GQA压缩后的KV（num_kv_heads组），而不是扩展后的
        shape = (batch_size, self.num_kv_heads, self.max_seq_len, self.head_dim)
        self.cache_k = torch.zeros(shape, dtype=torch.float16, device=device)
        self.cache_v = torch.zeros(shape, dtype=torch.float16, device=device)
        self.step = 0
        
    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        use_cache: bool = False  # 是否使用KV Cache
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        batch, seq_len, _ = x.shape
        
        # 1. 线性投影
        Q = self.W_q(x)  # (batch, seq_len, d_model)
        K = self.W_k(x)  # (batch, seq_len, num_kv_heads * head_dim)
        V = self.W_v(x)  # (batch, seq_len, num_kv_heads * head_dim)
        
        # 2. 拆分为多头
        # Q: (batch, seq_len, num_heads, head_dim) -> (batch, num_heads, seq_len, head_dim)
        Q = Q.view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # K: (batch, seq_len, num_kv_heads, head_dim) -> (batch, num_kv_heads, seq_len, head_dim)
        K = K.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        
        # 3. 【核心】KV Cache逻辑
        if use_cache:
            if self.cache_k is None:
                self._init_cache(batch, x.device)
            
            # 将当前步的K和V写入缓存
            self.cache_k[:, :, self.step:self.step + seq_len, :] = K
            self.cache_v[:, :, self.step:self.step + seq_len, :] = V
            
            # 从缓存中读取完整的K和V（包含历史 + 当前）
            K = self.cache_k[:, :, :self.step + seq_len, :]
            V = self.cache_v[:, :, :self.step + seq_len, :]
            
            self.step += seq_len
        
        # 4. 【GQA核心】将KV头扩展到Q头的数量
        # 每组内的Q头共享同一组KV，所以需要repeat KV头
        # 例如：num_heads=8, num_kv_heads=2，则每个KV头对应4个Q头
        K = K.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        V = V.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        # 此时K和V的形状都是 (batch, num_heads, total_seq_len, head_dim)
        
        # 5. 缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        # scores: (batch, num_heads, seq_len, total_seq_len)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 6. 加权求和
        attn_output = torch.matmul(attn_weights, V)
        # attn_output: (batch, num_heads, seq_len, head_dim)
        
        # 7. 合并多头 -> 输出
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch, seq_len, self.d_model)
        output = self.W_o(attn_output)
        
        # 8. 返回KV Cache供下一轮使用
        if use_cache:
            return output, (self.cache_k[:, :, :self.step, :], self.cache_v[:, :, :self.step, :])
        return output, None
    
    def reset_cache(self):
        """重置缓存（开始新序列生成时调用）"""
        self.cache_k = None
        self.cache_v = None
        self.step = 0


# 测试与显存对比
if __name__ == "__main__":
    d_model = 8192      # LLaMA-2 70B的维度
    num_heads = 64      # Q头数
    
    print("=" * 60)
    print("GQA配置与参数量对比")
    print("=" * 60)
    
    # 测试不同num_kv_heads配置
    for kv_heads in [64, 8, 1]:
        gqa = GroupedQueryAttention(d_model, num_heads, kv_heads)
        total_params = sum(p.numel() for p in gqa.parameters())
        
        # 计算KV Cache每token每层的开销
        kv_cache_per_token_per_layer = 2 * kv_heads * (d_model // num_heads) * 2  # FP16，2字节
        total_layers = 80  # LLaMA-2 70B层数
        
        print(f"\nnum_kv_heads={kv_heads} (MHA if 64, GQA if 8, MQA if 1)")
        print(f"  总参数量: {total_params / 1e6:.2f}M")
        print(f"  单token单层KV Cache: {kv_cache_per_token_per_layer / 1024:.2f}KB")
        print(f"  4096 tokens KV Cache: {kv_cache_per_token_per_layer * 4096 * total_layers / 1024 / 1024:.2f}GB")
    
    # 前向传播测试（带Cache）
    print("\n" + "=" * 60)
    print("前向传播与KV Cache测试")
    print("=" * 60)
    
    model = GroupedQueryAttention(d_model, num_heads, 8)  # GQA-8
    x = torch.randn(1, 1, d_model)  # 单步生成场景
    
    # 模拟生成5个token
    print("模拟自回归生成（逐token）：")
    for step in range(5):
        out, cache = model(x, use_cache=True)
        print(f"  第{step+1}步: 输入shape={x.shape}, 输出shape={out.shape}, "
              f"缓存K shape={cache[0].shape}, 缓存V shape={cache[1].shape}")
        # 下一步的x是上一步的输出（这里用随机数模拟）
        x = torch.randn(1, 1, d_model)
    
    model.reset_cache()
    print("\n缓存已重置，可用于下一个新序列的生成。")

代码的核心要点：

1. K/V投影的维度缩减：W_k和W_v的输出维度是 num_kv_heads * head_dim，而不是MHA中的 num_heads * head_dim。这直接减少了线性层的参数量，从模型权重层面就开始节省显存。

2. 缓存的是压缩后的KV：_init_cache中创建的缓存形状是 (batch, num_kv_heads, max_seq_len, head_dim)，而不是扩展到 num_heads 后的版本。这是GQA节省显存的根本原因——存得少。

3. 推理时才做扩展：repeat_interleave操作在每步前向中完成，将 num_kv_heads 组KV复制扩展到 num_heads 组。这一步不产生额外的持久化显存占用，只是临时计算。

4. reset_cache机制：生成完一个完整序列后，必须重置缓存，否则下一个序列会错误地复用前一个序列的KV。

---

七、GQA在实际部署中的工程考量

7.1 与Flash Attention的兼容性

Flash Attention是NVIDIA在2022年提出的高效注意力实现，通过分块计算大幅减少HBM（高带宽显存）访问，在长序列场景下可将注意力计算速度提升2到4倍。

GQA与Flash Attention完全兼容。Flash Attention在计算时会将Q、K、V的head维度统一处理，GQA先通过repeat_interleave将K、V扩展到与Q相同的头数，再送入Flash Attention。扩展操作非常轻量，几乎不增加计算开销。

7.2 服务框架中的配置示例

目前在主流推理框架vLLM中，启用GQA只需在模型配置文件中指定：

{
    "num_attention_heads": 64,
    "num_key_value_heads": 8
}

框架会自动识别并采用GQA的KV Cache管理策略。Hugging Face Transformers库也提供了类似的num_key_value_heads参数，开箱即用。

7.3 GQA适用场景的建议

从工程实践来看，不同场景对GQA的配置选择有不同倾向：

场景	推荐 num_kv_heads	理由
短文本对话（< 2K tokens）	16 或 32	显存压力不大，保留更多头保证精度
长文档分析（8K+ tokens）	4 或 8	显存是关键瓶颈，压缩收益巨大
批处理推理（batch_size > 8）	4	多批次叠加后显存压力成倍放大，必须极致压缩
代码生成/推理任务	8 或 16	这类任务对精度敏感，不宜压缩过多

这些数值来自主流模型的配置经验，实际使用时建议在自己的业务数据上做A/B测试，找到精度和成本的平衡点。

---

八、展望：GQA之后，下一代注意力是什么？

GQA目前已经成为大模型的事实标准，但技术演进从未停止。值得关注的下一代方向包括：

8.1 MLA（Multi-head Latent Attention）

DeepSeek-V2/V3引入的MLA，将KV进行低秩压缩——用低维的latent向量来表示高维的K和V。在推理时，只需要缓存这个低维latent向量，再通过一个上投影矩阵恢复出K和V。效果相当惊人：在DeepSeek-V2中，KV Cache被压缩到GQA的约1/29，同时模型精度几乎无损。

8.2 稀疏注意力（Sparse Attention）

核心思路是：不是所有token都值得被关注。比如在长文档中，相隔很远的token之间往往没有强语义关联。通过稀疏化（只让每个token关注邻近的或少数重要的token），可以将注意力计算复杂度从$O(N^2)$降到$O(N\log N)$甚至$O(N)$。Longformer、BigBird等模型已经验证了这条路的可行性。

8.3 线性注意力（Linear Attention）

将Softmax替换为核函数映射，让注意力计算从矩阵乘法变为向量点积，复杂度从$O(N^2)$直接降到$O(N)$。但目前的线性注意力方案在精度上还无法超越标准Softmax注意力，仍需进一步的算法突破。

---

九、结语

从MHA到MQA，再到GQA，我们看到的是一个精度与效率反复博弈、最终走向工程最优解的清晰技术演进路径。

对于AI应用工程师而言，理解GQA的意义远不止于面试时多一个亮点——当你在生产环境中部署大模型服务，面对单卡A100 80GB显存、用户要求支持32K上下文、同时batch_size还得大于4时，你会真切地感受到KV Cache是真正的瓶颈，而GQA就是那个让你不用再着急申请预算加卡的救命设计。

希望这篇文章能帮你真正理解：一个看似简单的结构变化，背后是如何牵引整个大模型推理基础设施演进的。 如果你在实际部署中遇到了GQA相关的配置问题或显存优化难题，欢迎在评论区交流讨论。