Transformer 架构优化学习笔记：GQA 与 MoE 如何把推理成本打下来

这是一篇学习笔记，记录我在研究大模型推理优化时的思考路径。全文围绕两个核心问题展开：GQA（分组查询注意力）怎么减少了推理时的显存压力？MoE（混合专家）怎么在参数量暴涨的同时控制住了计算量？内容从最基础的概念讲起，逐步深入到代码实现和工程权衡，适合对 Transformer 有基本了解、想进一步理解现代大模型架构优化的读者。

推理成本为什么是大模型的命门

训练一个大模型是一次性投入，推理则是持续支出。当一个 70B 参数的模型部署上线，每处理一个请求都要把模型权重从显存搬到计算单元，每生成一个 token 都要维护不断增长的 KV Cache。用户多了，GPU 账单按小时跳动——推理成本直接决定了一个大模型产品能不能活下去。

推理成本的两个核心来源：一是显存占用，决定了你需要多少张卡；二是计算量，决定了每张卡能服务多少用户。GQA 主要解决前者，MoE 主要解决后者。两者的思路截然不同，但目标一致：在不牺牲模型效果的前提下，让推理更便宜。

第一部分：GQA —— 给 KV Cache 瘦身

从 KV Cache 说起

要理解 GQA，先得搞清楚 KV Cache 是什么，以及它为什么成了瓶颈。

Transformer 的自注意力计算中，每个 token 会生成三组向量：Q（查询）、K（键）、V（值）。在自回归生成时，每生成一个新 token，都要拿它的 Q 去和之前所有 token 的 K 做点积，再用得到的注意力权重去加权之前的 V。如果不做任何缓存，每生成一个 token 都要重新计算整段序列的 K 和 V，复杂度是 O(L²)。

KV Cache 的做法很直接：把每一层、每一步算出的 K 和 V 张量存起来，下一步直接复用。这样一来，新 token 只需算自己的 Q，和缓存里的 K/V 做一次点积，单步复杂度降到 O(L)。代价是显存占用随序列长度线性增长。

来算一笔账。以 LLaMA-2 7B 为例：

模型参数：32 层 × 32 头 × 128 维 = 每层 32 组 K/V
KV Cache 大小 = 2(K和V) × 32层 × 32头 × 128维 × seq_len × 2字节(fp16)

当 seq_len=4096 时，单个请求的 KV Cache 约 2.1 GB。batch size 开到 32，光 KV Cache 就要 67 GB——比模型权重本身还大。当上下文拉到 128K（LLaMA 3 的水平），这个数字会更加夸张。

这就是 GQA 要解决的问题：KV Cache 太大了，而它的增长速度和注意力头数成正比。

MHA、MQA、GQA：三条路线

理解 GQA 之前，先看它的两个"邻居"——MHA 和 MQA。

MHA（Multi-Head Attention） 是原始 Transformer 的标准方案。假设有 32 个注意力头，每个头都有自己独立的 Q、K、V 投影。KV Cache 的大小和头数（32）成正比。效果最好，但显存开销最大。

MQA（Multi-Query Attention） 是 2019 年 Shazeer 提出的极端方案：所有查询头共享同一组 K 和 V。KV 头数从 32 降到 1，KV Cache 直接缩减到 1/32。推理速度飞快，但效果损失明显——多个查询头被迫看同一份 K/V，表达能力被压缩。

GQA（Grouped Query Attention） 是两者的折中。把 32 个查询头分成若干组，每组共享一组 K/V。比如分成 8 组，就是 8 组 K/V 头服务 32 个查询头，每组 4 个查询头共享同一组 K/V。KV Cache 缩减到 MHA 的 1/4，效果接近 MHA，远好于 MQA。

直觉上可以这样理解：MQA 像是让 32 个人共用一本书，信息太少；MHA 是每人一本，太浪费；GQA 是 4 个人共用一本，够用且省纸。

GQA 的论文发现

GQA 来自 Ainslie 等人 2023 年的论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》。论文有两个关键发现：

第一，把已有的 MHA 模型转换为 GQA（通过合并相邻的 KV 头），效果几乎无损。这意味着不需要从头训练，可以直接拿现有的 MHA 模型改造。

第二，从头训练的 GQA 模型在 KV 头数达到一定阈值后（实验中是 8 组左右），效果与 MHA 几乎一致。这说明 MHA 中多个 KV 头之间存在大量冗余——它们学到的 K/V 投影高度相似，完全可以共享。

这个发现很有启发性：很多时候，模型架构的"冗余"就是优化的空间。

代码实现

GQA 的实现分两步：投影时只生成 n_kv_heads 组 K/V，计算注意力时把 K/V 重复扩展到和 Q 头数一致。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GroupedQueryAttention(nn.Module):
    def __init__(self, dim, n_heads, n_kv_heads=None):
        super().__init__()
        self.n_heads = n_heads                  # 查询头数，如 32
        self.n_kv_heads = n_kv_heads or n_heads  # KV 头数，如 8
        self.head_dim = dim // n_heads
        self.n_rep = n_heads // self.n_kv_heads  # 每组 KV 共享几个 Q 头，如 4

        # Q 投影到全部头，K/V 只投影到 KV 头
        self.wq = nn.Linear(dim, n_heads * self.head_dim)
        self.wk = nn.Linear(dim, self.n_kv_heads * self.head_dim)
        self.wv = nn.Linear(dim, self.n_kv_heads * self.head_dim)
        self.wo = nn.Linear(n_heads * self.head_dim, dim)

    def forward(self, x, mask=None, freqs_cis=None,
                cache_k=None, cache_v=None):
        bsz, seqlen, _ = x.shape

        xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim)
        xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim)
        xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim)

        # 此处省略 RoPE 位置编码的施加

        # 更新 KV Cache
        if cache_k is not None:
            xk = torch.cat([cache_k, xk], dim=1)
            xv = torch.cat([cache_v, xv], dim=1)

        # 关键步骤：把 KV 头重复扩展到和 Q 头数一致
        xk = self._repeat_kv(xk)   # [bsz, seqlen, n_heads, head_dim]
        xv = self._repeat_kv(xv)

        # 标准注意力计算
        xq = xq.transpose(1, 2)    # [bsz, n_heads, seqlen, head_dim]
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)

        scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask
        attn = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(attn, xv)  # [bsz, n_heads, seqlen, head_dim]

        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output), xk, xv

    def _repeat_kv(self, x):
        """将 KV 头从 n_kv_heads 扩展到 n_heads"""
        bsz, seqlen, n_kv_heads, head_dim = x.shape
        if self.n_rep == 1:
            return x
        return (
            x[:, :, :, None, :]
            .expand(bsz, seqlen, n_kv_heads, self.n_rep, head_dim)
            .reshape(bsz, seqlen, n_kv_heads * self.n_rep, head_dim)
        )

_repeat_kv 是 GQA 的核心操作。它通过 expand（视图扩展，不拷贝数据）把每组 KV 头"复制"到对应的 Q 头上。在实际推理框架中，这一步往往被优化为直接在 attention kernel 中用索引读取，避免显式的数据复制。

主流模型中的应用

GQA 已经成为现代大模型的标配：

• LLaMA 3：全系列（8B/70B/405B）均采用 GQA，8 组 KV 头
• Qwen3：28 个 Q 头共享 4 组 KV 头，n_rep=7
• Mistral 7B：8 个 KV 头服务 32 个 Q 头
• DeepSeek-V3：没有用 GQA，而是用了 MLA（Multi-head Latent Attention），把 KV Cache 压缩到低秩隐空间，思路不同但目标一致

一个值得注意的趋势：KV Cache 的压缩方式正在从"减少头数"（GQA）向"低秩压缩"（MLA）演进。DeepSeek 的 MLA 把 KV Cache 压缩到了更极致的水平，但实现复杂度也更高。GQA 的优势在于简单直接，几乎不增加任何额外计算，工程落地成本低。

第二部分：MoE —— 用稀疏激活打破参数与计算的绑定

密集模型的困境

标准的 Transformer 是一个"密集模型"（Dense Model）：每个 token 都要经过所有层的所有参数。模型有 70B 参数，处理每个 token 就要做 70B 参数量的计算。参数量和计算量是绑定的——想让模型更聪明就加参数，加了参数就加了计算量。

这个绑定关系在追求更大模型时变成了死结。GPT-4 传闻有 1.8T 参数，如果用密集架构，推理时每个 token 都要跑完整 1.8T 的计算，成本高到难以商业化。

MoE 的核心洞察是：模型需要大参数量来存储知识，但处理每个 token 时并不需要用到所有知识。 就像人脑处理不同问题时激活不同区域，MoE 让每个 token 只激活模型中的一小部分"专家"。

MoE 的基本结构

在标准 Transformer 中，每个层有两个子模块：注意力层和前馈网络（FFN）。MoE 的做法是把 FFN 替换为多个并行的"专家"FFN，再加一个路由器（Router/Gate）来决定每个 token 交给哪些专家处理。

标准 FFN:
  x → Linear → GeLU → Linear → output

MoE 层:
  x → Router(x) → 选出 top-k 个专家
    → 分别过选中的 k 个专家 FFN
    → 按路由权重加权求和 → output

假设一个 MoE 层有 8 个专家，每次激活 top-2，那么每个 token 只需经过 2 个专家的计算。模型总参数量是 8 份 FFN 的参数，但计算量只有 2 份——参数量和计算量被解耦了。

一个具体的数字感受：Qwen3-235B-A22B 总参数 2350 亿，每次推理只激活 220 亿参数。相当于一个 22B 密集模型的计算成本，却拥有 235B 模型的知识容量。

路由器：MoE 的核心组件

路由器是一个小型线性层，输入是 token 的隐藏状态，输出是每个专家的得分：

class MoERouter(nn.Module):
    def __init__(self, dim, n_experts):
        super().__init__()
        self.gate = nn.Linear(dim, n_experts)

    def forward(self, x):
        # x: [batch * seq_len, dim]
        logits = self.gate(x)              # [tokens, n_experts]
        scores = F.softmax(logits, dim=-1)  # 每个专家的权重
        # 选出得分最高的 top-k 个专家
        topk_scores, topk_indices = torch.topk(scores, k=2, dim=-1)
        # 重新归一化，使选中的 k 个权重之和为 1
        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
        return topk_scores, topk_indices

路由器的输出是两个张量：topk_indices 告诉每个 token 该去哪些专家，topk_scores 告诉每个专家的输出该占多大比重。

路由策略是 MoE 设计中最关键的部分，不同模型在这里走了不同的路线。

Switch Transformer：极简主义

Google 2021 年的 Switch Transformer 把路由简化到了极致：每个 token 只路由到 1 个专家（top-1 路由）。这种设计最大化了稀疏性——8 个专家只用 1 个，计算量降到密集模型的 1/8。

Switch Transformer 还引入了辅助负载均衡损失。直觉是：如果不加约束，路由器可能偷懒把所有 token 都送到同一个"最擅长"的专家，其他专家闲置。辅助损失通过惩罚专家负载的变异系数（CV），迫使路由器把 token 尽量均匀地分配到各专家。

def load_balancing_loss(gate_scores, topk_indices, n_experts):
    """
    gate_scores: [tokens, n_experts] 路由器输出的概率
    topk_indices: [tokens, k] 选中的专家索引
    """
    # 每个专家被选中的 token 比例
    expert_mask = F.one_hot(topk_indices, n_experts).float()  # [tokens, k, n_experts]
    expert_load = expert_mask.sum(dim=1).mean(dim=0)           # [n_experts]

    # 每个专家的平均路由概率
    expert_prob = gate_scores.mean(dim=0)  # [n_experts]

    # 负载均衡损失 = n_experts * Σ(expert_load_i × expert_prob_i)
    loss = n_experts * (expert_load * expert_prob).sum()
    return loss

这个损失函数的设计很巧妙：它同时考虑了"每个专家分到了多少 token"（负载）和"路由器对每个专家的平均偏好"（概率），两者的乘积在负载完全均匀时达到最小值。

DeepSeekMoE：细粒度专家 + 共享专家

DeepSeek 在 MoE 架构上做了两个重要创新，让它成为目前开源 MoE 模型中效果最好的方案之一。

细粒度专家分割：传统 MoE 的每个专家是一个完整的 FFN，参数量较大（比如每个专家的隐藏维度是模型维度的 4 倍）。DeepSeek 把每个大专家拆成多个小专家——比如把 1 个大 FFN 拆成 4 个小 FFN，专家总数从 8 变成 32，但总参数量不变。每个 token 仍然激活 2 个小专家，但这两个小专家的参数量只有原来的 1/4。

这样做的动机：大专家容易"知识冲突"——一个专家被迫同时处理数学、代码、语言等不同领域的知识，内部参数互相干扰。小专家更容易"专业化"，每个小专家专注于更细粒度的知识领域。实验表明，细粒度分割在相同参数量和计算量下效果更好。

共享专家隔离：DeepSeek 从所有专家中划出一部分作为"共享专家"（Shared Expert），每个 token 无论路由结果如何，都会经过这些共享专家。共享专家负责处理通用知识（如语法结构、常见搭配），路由专家则负责领域特化知识。

DeepSeekMoE 层:
  x → 共享专家 (固定激活)
    → Router → 选出 top-k 个路由专家 (条件激活)
    → 共享专家输出 + 路由专家加权输出 → output

DeepSeek-V3 的配置是 1 个共享专家 + 256 个路由专家，每个 token 激活 8 个路由专家。共享专家保证通用能力不丢失，256 个细粒度路由专家提供高度专业化的知识覆盖。

无辅助损失的负载均衡

DeepSeek-V3 还引入了一个工程上很重要的创新：无辅助损失负载均衡。

传统的辅助损失方法有一个副作用：它通过梯度直接干预路由器的学习方向，导致路由器为了"均匀分配"而做出次优的路由决策。模型效果会受影响——明明某个 token 送给专家 A 最好，但为了均衡负载被强行送给了专家 B。

DeepSeek-V3 的做法是给每个专家加一个可动态调整的偏置项（bias），偏置项不参与梯度计算，而是根据每个专家的实际负载动态调整：负载高的专家偏置降低，负载低的专家偏置升高。这样路由器本身只管"把 token 送给最合适的专家"，负载均衡交给偏置项的动态调节来完成。

# 伪代码：无辅助损失的负载均衡
# 每个 expert 有一个 bias，不参与梯度更新
expert_bias = torch.zeros(n_experts)  # 动态调整

def route_with_bias(gate_logits, expert_bias):
    # 路由得分 = 路由器输出 + 动态偏置
    adjusted_logits = gate_logits + expert_bias
    scores = F.softmax(adjusted_logits, dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=2, dim=-1)
    return topk_scores, topk_indices

def update_bias(expert_bias, expert_load, target_load):
    # 简单的比例控制：负载高的降偏置，负载低的升偏置
    # 实际实现中用更平滑的更新策略
    expert_bias += 0.01 * (target_load - expert_load)

这个思路的好处是：模型训练的优化目标纯粹是"做好任务"，负载均衡完全是一个外部的、不干扰学习的控制回路。

MoE 的工程挑战

MoE 在理论上很优雅，但在工程实现上有两个硬骨头。

负载均衡：前面已经讨论过了。如果路由不均衡，有些专家被塞满、有些专家闲置，计算资源就浪费了。尤其在多 GPU 并行训练时，每个 GPU 负责一部分专家（Expert Parallelism），负载不均衡意味着有些 GPU 在等其他 GPU 算完。

通信开销：在 Expert Parallelism 下，token 可能需要被发送到其他 GPU 上的专家处理，处理完再发回来。这种 all-to-all 通信在千卡规模下会消耗大量时间。DeepSeek-V3 的解法是"节点限制路由"——限制每个 token 最多被路由到 4 个节点上的专家，控制跨节点通信量。

MoE 的推理成本优势

回到推理成本这个核心话题。MoE 模型在推理时的优势可以用一个公式概括：

推理计算量 ≈ 激活参数量 × seq_len
         << 总参数量 × seq_len

以 Qwen3-235B-A22B 为例：

注意一个细节：MoE 减少的是计算量，不是显存占用。所有专家的权重都要驻留在显存里，只是每次推理只用其中一部分。这意味着 MoE 模型仍然需要和密集模型一样多的 GPU 来放权重，但每张 GPU 的计算负载大幅降低，可以服务更多并发请求。

换句话说：MoE 省的是每 token 的计算开销（FLOPs），GQA 省的是每 token 的显存开销（KV Cache）。两者的优化维度不同，可以叠加使用。

GQA 与 MoE 的协同效应

现代大模型通常同时使用 GQA 和 MoE。以 Qwen3-235B-A22B 为例，它采用了 128 专家 / 8 激活的 MoE 架构，同时注意力层使用 GQA（28 个 Q 头共享 4 组 KV 头）。两个优化各自独立工作，分别解决不同维度的瓶颈：

• GQA 把 KV Cache 从 MHA 的 100% 压到约 14%（4/28），降低了长序列推理的显存压力
• MoE 把每 token 的计算量从 235B 压到 22B，降低了每 token 的计算开销

两者叠加的效果是：模型可以在有限的 GPU 资源下，同时处理更长的上下文（GQA 的功劳）和更多的并发请求（MoE 的功劳）。

这种"分维度优化"的思路在系统设计中很常见：当你面对一个多维度的瓶颈时，不要试图用一个方案解决所有问题，而是针对每个维度找到最直接有效的优化手段。

总结：两种思路，一个目标

GQA 和 MoE 代表了大模型推理优化的两种典型思路：

GQA 做减法。它发现 MHA 中多个 KV 头存在冗余，通过共享 KV 头来减少冗余。优化的对象是推理时的显存占用（KV Cache），不影响每 token 的计算量。实现简单，几乎没有副作用，已经成为行业标配。

MoE 做拆分。它把一个大的 FFN 拆成多个小的专家 FFN，每次只激活一部分。优化的对象是每 token 的计算量，不影响（甚至增加了）显存占用。实现复杂，涉及路由、负载均衡、并行通信等工程难题，但带来的计算节省是数量级的。

一个更宏观的观察：这两个优化都不是从理论推导出来的，而是从工程实践中"长"出来的。GQA 来源于对 MHA 冗余的实验观察，MoE 来源于"参数存知识、计算用知识"的直觉。大模型架构的演进，很多时候不是数学上的突破，而是对已有结构的精细拆解和重组。

如果你正在学习大模型架构，建议把 LLaMA 3 的 model.py（GQA 实现）和 DeepSeek 的开源代码（MoE 实现）对照着读一遍。前者不到 300 行，后者稍长但注释清晰。读懂这两段代码，比看十篇综述论文更能建立对架构优化的直觉。