注意力变体与 MoE

Transformer 架构 里的标准模型有两个越来越贵的地方。第一，生成文本时每个新 token 都要读一遍之前所有 token 的 K 和 V，一个 70B 模型如果用标准多头注意力，16 条 32K 长的序列光是这些缓存就要 1.28 TB，相当于 16 张 80GB 的卡。第二，想要更聪明的模型就要更多参数，而每个 token 的计算量和参数量成正比。这一篇讲两类改动：改注意力，让缓存更小、长序列更便宜；改 FFN，用 MoE 让参数变多而每个 token 的计算不变。

前置：Transformer 架构、资源核算与 Scaling Laws

生成时的瓶颈：每个 token 都要读一遍 K/V

自回归生成时，已经算过的 token 的 K、V 不会变，推理引擎把它们缓存下来，新 token 只需算自己的 q、k、v，再和缓存里的所有 K/V 做注意力。这个缓存叫 KV cache，它的管理（分页、调度、带宽分析）见 推理系统。这里只关心它有多大：

$\text{每 token 字节数}=2(\text{K 和 V})\times n_{\text{layer}}\times n_{\text{kv\_head}}\times d_{\text{head}}\times \text{每元素字节数}$

拿 LLaMA 2 70B 的形状（80 层、64 个 query 头、d_head = 128、bf16）算三种情况：

def kv_bytes(layers, kv_heads, hd, bytes_per=2):
    return 2 * layers * kv_heads * hd * bytes_per          # K 和 V 各一份
cfgs = {
    "MHA, 64 kv heads": (80, 64, 128),
    "GQA, 8 kv heads":  (80, 8, 128),
    "MQA, 1 kv head":   (80, 1, 128),
}
for name, (L, kv, hd) in cfgs.items():
    b = kv_bytes(L, kv, hd)
    print(f"{name:17s} {b/1024:8.0f} KiB/token   32K ctx x 16 seqs = {b*32768*16/2**30:6.1f} GiB")

MHA, 64 kv heads      2560 KiB/token   32K ctx x 16 seqs = 1280.0 GiB
GQA, 8 kv heads        320 KiB/token   32K ctx x 16 seqs =  160.0 GiB
MQA, 1 kv head          40 KiB/token   32K ctx x 16 seqs =   20.0 GiB

标准多头注意力（MHA）每个 token 要 2.5 MiB，16 条 32K 的序列要 1.28 TB，模型权重本身才 140 GB。

缓存大不只是放不下的问题，还是速度问题。decode 阶段每一步只有一个新 token，对每个头而言，读入长度为 T 的 K 和 V（2·T·d_head 个元素），只做 2·2·T·d_head 次浮点运算（q·K 和 权重·V）。在 MHA 下，每个 query 头独享一份 K/V，算术强度（每读一个字节做多少 FLOP）大约是 bf16 下 1 FLOP/byte。H100 的 bf16 算力约 989 TFLOP/s、显存带宽约 3.35 TB/s，要用满算力需要每字节约 300 FLOP。decode 时的注意力离这个数差了两个数量级，GPU 绝大部分时间在等显存（roofline 分析见 GPU 与 Kernel）。训练和 prefill 没有这个问题，因为它们一次处理整个序列，K/V 读进来后被 T 个 query 复用。

MQA 与 GQA：让多个 query 头共用一份 K/V

既然问题在于 K/V 被读得多、用得少，最直接的修法是让多个 query 头共享同一份 K/V。所有头共用一份叫 MQA（Multi-Query Attention，Shazeer 2019）；折中一下，把 query 头分成 G 组，每组共享一份，叫 GQA（Grouped-Query Attention，Ainslie et al. 2023）。实现时只需在头维度上把 K/V 复制 n_head / n_kv_head 份：

import torch
def attention(x, Wq, Wk, Wv, n_head, n_kv_head):
    B, T, C = x.shape; hd = C // n_head
    q = (x @ Wq).view(B, T, n_head, hd).transpose(1, 2)
    k = (x @ Wk).view(B, T, n_kv_head, hd).transpose(1, 2)   # K/V 只有 n_kv_head 个头
    v = (x @ Wv).view(B, T, n_kv_head, hd).transpose(1, 2)
    g = n_head // n_kv_head
    # 每组 g 个 query 头共用同一份 K/V：沿头维复制，不需要新参数
    k, v = k.repeat_interleave(g, dim=1), v.repeat_interleave(g, dim=1)
    s = q @ k.transpose(-2, -1) / hd**0.5
    s = s.masked_fill(torch.ones(T, T, dtype=torch.bool).triu(1), float("-inf"))
    return (s.softmax(-1) @ v).transpose(1, 2).reshape(B, T, C)
torch.manual_seed(0)
C, H = 64, 8
x = torch.randn(2, 10, C)
Wq = torch.randn(C, C) / C**0.5
for kv in (8, 2, 1):
    Wk = torch.randn(C, C // H * kv) / C**0.5; Wv = torch.randn(C, C // H * kv) / C**0.5
    y = attention(x, Wq, Wk, Wv, H, kv)
    print(f"n_kv_head={kv}: out {tuple(y.shape)}, K/V params {Wk.numel() + Wv.numel()}")

n_kv_head=8: out (2, 10, 64), K/V params 8192
n_kv_head=2: out (2, 10, 64), K/V params 2048
n_kv_head=1: out (2, 10, 64), K/V params 1024

n_kv_head = n_head 就是 MHA，n_kv_head = 1 就是 MQA，GQA 是两者之间的连续谱。输出形状不变，K/V 投影的参数和缓存都按 n_kv_head / n_head 缩小。这里的 repeat_interleave 是为了讲清楚而写的；真正的推理 kernel 不会复制数据，而是让同组的几个 query 头去读同一块 K/V，这样每读一个字节就被 g 个头复用，算术强度也提高到原来的 g 倍。

代价是表达力：MQA 在一些任务上质量有可见的下降，训练也更不稳定；GQA 用 8 个 KV 头时质量几乎和 MHA 一样，缓存却只有 1/8。所以 LLaMA 2 70B、LLaMA 3 全系、Mistral、Qwen 等主流开源模型都用 GQA。GQA 论文还给了一个省钱的办法：已经训练好的 MHA 模型，可以把每组的 K/V 头取平均合成一个，再用原预训练算力的约 5% 继续训练，就能转成 GQA。

停一下，检查理解单选

一个模型从 MHA（32 个 KV 头）改成 GQA（8 个 KV 头），其他不变。下列哪个说法正确？

训练时每个 token 的 FLOPs 减少为原来的 1/4KV cache 变为原来的 1/4，decode 时注意力的算术强度变为原来的 4 倍query 头的数量也必须降到 8模型的上下文长度上限扩大为原来的 4 倍

检查答案

MLA：缓存一个低秩的 latent

GQA 是在”头”这个维度上压缩。DeepSeek-V2 提出的 MLA（Multi-head Latent Attention）换了一个维度：把每个 token 的 K、V 联合压成一个低维向量 $c_t$，只缓存它。

$c_t=W_{DKV}{x}_t\in {\mathbb{R}}^{d_c},k_t^{(h)}=W_{UK}^{(h)}{c}_t,v_t^{(h)}=W_{UV}^{(h)}{c}_t$

在 DeepSeek-V2/V3 里 $d_c=512$，而 128 个头的 K 和 V 合起来有 2 × 128 × 128 = 32768 维。看起来每次用 K/V 都要先从 $c_t$ 升维，计算反而多了。妙处在于矩阵可以吸收：打分是 $q^{(h)\top }{k}^{(h)}=(W_{UK}^{(h)\top }{q}^{(h)}{)}^{\top }{c}_t$，把 $W_{UK}$ 并进 query 一侧，直接拿 query 和 latent 做点积；同理 $W_{UV}$ 可以并进输出投影。推理时根本不需要把每个头的 K、V 显式算出来。

RoPE 是这个吸收技巧的障碍：旋转矩阵依赖位置，夹在 $q$ 和 $W_{UK}{c}_t$ 之间就没法预先合并。DeepSeek 的解法是”解耦”：额外给每个 token 算一个 64 维、所有头共享的 key 专门承载 RoPE，和不带位置的那部分拼起来打分。于是每层每个 token 只缓存 512 + 64 = 576 个数。按 V3 的形状（61 层、128 头、d_head = 128）算：

L = 61
mha = 2 * L * 128 * 128 * 2          # 如果用标准 MHA，bf16
mla = L * (512 + 64) * 2             # MLA：latent + 共享的 RoPE key
print(f"V3 if MHA: {mha/1024:.0f} KiB/token,  MLA: {mla/1024:.1f} KiB/token,  ratio {mha/mla:.1f}x")

V3 if MHA: 3904 KiB/token,  MLA: 68.6 KiB/token,  ratio 56.9x

缓存小了 57 倍，比 LLaMA 2 70B 的 GQA-8（320 KiB）还小得多，而 DeepSeek-V2 报告里 MLA 的质量不低于 MHA。代价是实现复杂：训练时要显式升维，推理时要走吸收后的路径，RoPE 要单独处理，推理引擎需要专门的 kernel。

长上下文：滑动窗口、线性注意力与稀疏注意力

GQA 和 MLA 压的是每个 token 的缓存大小，但注意力对序列长度仍然是 $O(T^2)$ 的计算、$O(T)$ 的缓存。上下文到几十万 token 时，要换一种思路：不让每个 token 看所有前文。

滑动窗口最简单：每个 token 只看最近 W 个 token。缓存上限固定为 W，计算是 $O(TW)$。信息仍然能传得更远，因为第 $\ell$ 层的位置能看到第 $\ell -1$ 层的窗口，而那一层又看到更早的窗口，L 层叠起来感受野是 $L\times W$；但这种间接传递是有损的。Mistral 7B 用的是 4096 的窗口。更常见的做法是交错：大部分层用滑动窗口，每隔几层放一层全局注意力，例如 Gemma 3 是 5 层局部（窗口 1024）配 1 层全局。全局层负责精确的长程检索，局部层省掉大部分缓存。

线性注意力动的是 softmax。如果把 $\exp (q\cdot k)$ 换成某个特征映射的内积 $\varphi (q)\cdot \varphi (k)$，注意力就变成纯矩阵乘，可以换括号：

import torch, time
torch.manual_seed(0)
T, d = 4096, 64
phi = lambda x: torch.nn.functional.elu(x) + 1      # 保证"相似度"非负，代替 exp
Q, K, V = phi(torch.randn(T, d)), phi(torch.randn(T, d)), torch.randn(T, d)
t = time.time(); a = (Q @ K.T) @ V; t1 = time.time() - t   # 先算 T×T
t = time.time(); b = Q @ (K.T @ V); t2 = time.time() - t   # 先算 d×d
print(f"max rel diff {((a-b).abs().max()/a.abs().max()):.1e}", f"(QK^T)V {t1*1e3:.1f} ms   Q(K^T V) {t2*1e3:.1f} ms")
print(f"FLOPs: {2*2*T*T*d/1e9:.2f} G vs {2*2*T*d*d/1e9:.3f} G, ratio T/d = {T//d}")
# 因果版本 = 一个 RNN：状态 S 是 d×d，每来一个 token 加一个外积
S = torch.zeros(d, d); out = []
for t in range(T):
    S = S + K[t:t+1].T @ V[t:t+1]
    out.append(Q[t] @ S)
causal = (Q @ K.T).tril() @ V
print(torch.allclose(torch.stack(out), causal, rtol=1e-3, atol=1e-2))

max rel diff 3.2e-06 (QK^T)V 25.1 ms   Q(K^T V) 0.5 ms
FLOPs: 4.29 G vs 0.067 G, ratio T/d = 64
True

（这里省略了归一化分母，它可以用同样的方式递推。）结合律让计算从 $O(T^{2}d)$ 降到 $O(T{d}^2)$，序列越长省得越多。更重要的是因果版本：它等价于一个状态为 $d\times d$ 矩阵的 RNN，decode 时每步只更新状态，缓存大小和上下文长度无关。

代价是记忆容量。softmax 注意力可以精确地取回任意一个历史 token；线性注意力把所有历史压进一个固定大小的矩阵，只会不断累加，旧信息无法擦掉。后续工作都在给这个更新加门控：Mamba-2 让状态每步乘一个与数据相关的衰减系数，DeltaNet 用”先减去对 k 的旧预测，再写入新值”的 delta 规则做有选择的覆盖，Gated DeltaNet 把两者合起来。纯线性模型在需要精确检索的任务上仍然吃亏，所以落地的做法多是混合：大部分层用线性注意力，少数层保留完整注意力，例如 Qwen3-Next 是 3 层 Gated DeltaNet 配 1 层普通注意力。

稀疏注意力保留 softmax，但每个 query 只挑一部分 key 去算。DeepSeek-V3.2 的 DSA（DeepSeek Sparse Attention）用一个很轻量的”索引器”给所有历史 token 快速打分，每个 query 只取得分最高的 2048 个 token 做真正的 MLA 注意力。这样主注意力的计算从 $O(T^2)$ 降到 $O(Tk)$；索引器本身仍然是 $O(T^2)$，但它头少、维度低、可以用 FP8，远比完整注意力便宜。

停一下，检查理解多选

关于线性注意力，哪些说法正确？

去掉 softmax 后可以先算 K^T V，计算量从与 T^2 成正比变为与 T 成正比因果线性注意力可以写成状态大小固定的 RNN，decode 时缓存不随上下文增长它与 softmax 注意力在数学上等价，只是计算顺序不同它的状态能精确存储并取回任意一个历史 token

检查答案

MoE：参数变多，每个 token 的计算不变

换到另一个问题。资源核算与 Scaling Laws 会讲到：模型越大 loss 越低，但稠密模型每个 token 的训练计算是 $6N$，参数翻倍计算就翻倍。有没有办法让参数翻倍而计算不翻倍？

Mixture of Experts 的想法是：把每层的一个 FFN 换成 E 个并列的 FFN（专家），再加一个小的路由器（router），给每个 token 只挑其中 k 个去算。参数是 E 份，每个 token 的计算只有 k 份。

import torch, torch.nn as nn, torch.nn.functional as F
class MoE(nn.Module):
    def __init__(self, d, d_ff, n_exp, k):
        super().__init__()
        self.k = k
        self.router = nn.Linear(d, n_exp, bias=False)
        self.w1 = nn.Parameter(torch.randn(n_exp, d, d_ff) / d**0.5)
        self.w2 = nn.Parameter(torch.randn(n_exp, d_ff, d) / d_ff**0.5)
    def forward(self, x):                       # x: (N, d)，batch 和序列已展平
        probs = self.router(x).softmax(-1)      # (N, E)
        w, idx = probs.topk(self.k, dim=-1)     # 每个 token 选 k 个专家
        w = w / w.sum(-1, keepdim=True)         # 只在选中的专家之间重新归一化
        out = torch.zeros_like(x)
        for e in range(self.w1.shape[0]):
            tok, slot = (idx == e).nonzero(as_tuple=True)   # 哪些 token 选了专家 e
            if len(tok) == 0: continue
            h = F.silu(x[tok] @ self.w1[e]) @ self.w2[e]
            out.index_add_(0, tok, h * w[tok, slot, None])
        return out, probs, idx
torch.manual_seed(0)
d, d_ff, E, k = 64, 128, 8, 2
moe = MoE(d, d_ff, E, k)
x = torch.randn(1024, d)
y, probs, idx = moe(x)
total = sum(p.numel() for p in moe.parameters())
active = d * E + k * 2 * d * d_ff
print(f"total params {total:,}, active per token {active:,}")
print("tokens per expert:", torch.bincount(idx.flatten(), minlength=E).tolist())

total params 131,584, active per token 33,280
tokens per expert: [222, 264, 277, 232, 228, 285, 264, 276]

8 个专家、top-2，参数是一个 FFN 的 8 倍，每个 token 只用到约 1/4。实现的关键是”按专家分组”：先找出选了专家 e 的所有 token，把它们打包成一个矩阵乘，结果按路由权重加回原位。

top-k 的选择本身不可导，梯度怎么流到路由器？靠的是那个权重 w：被选中专家的输出乘以它的路由概率，loss 对输出的梯度会经过 w 传回 router。模型因此学会”给效果好的专家更高分”，没被选中的专家这一步收不到梯度。

MoE 的设计空间主要有三个维度。k 的大小：Switch Transformer 用 top-1，Mixtral 8x7B 用 8 选 2（总参数 46.7B，每 token 激活约 12.9B）。专家粒度：DeepSeekMoE 把每个专家切小，用更多、更细的专家，并提高 k，组合数多得多，专家可以更专精。共享专家：留一两个所有 token 都必经的专家，承载通用知识，让路由专家不必各自重复学一遍。

停一下，检查理解单选

上面的 MoE 实现中，router 是如何获得梯度的？

通过 topk 的索引，PyTorch 对离散选择自动做了可导近似通过把专家输出乘以路由权重 w，loss 对 w 的梯度传回 routerrouter 不需要梯度，它是随机初始化后固定不变的通过 load balancing loss，router 只从这个辅助损失获得梯度

检查答案

训练 MoE：路由会坍缩，所以要平衡负载

路由器有一个正反馈：某个专家早期碰巧被选得多，就得到更多训练、变得更好，于是被选得更多。最后少数专家包揽一切，其余专家闲置，参数白白浪费。系统上也有问题：专家通常分布在不同 GPU 上（专家并行，见 并行训练），负载不均意味着有的卡忙死、有的卡空转，而且每个专家的容量有上限，超出的 token 会被丢弃。

辅助平衡损失是 Switch Transformer 的做法。设 $f_i$ 是这个 batch 里实际分给专家 $i$ 的 token 比例，$P_i$ 是所有 token 给专家 $i$ 的平均路由概率：

${\mathcal{L}}_{\text{bal}}=\alpha \cdot E{\sum }_{i=1}^E{f}_i{P}_i$

$f_i$ 来自 argmax，不可导，梯度只走 $P_i$，且对 $P_i$ 的梯度正比于 $f_i$：分得越多的专家，它的概率被压得越狠。均匀时这一项等于 $\alpha$，越不均匀越大：

import torch
def switch_loss(f, P):            # f：实际分到的 token 比例；P：路由概率的平均
    return len(f) * (f * P).sum()
print(f"{switch_loss(torch.tensor([.7,.1,.1,.1]), torch.tensor([.6,.2,.1,.1])):.2f}",
      f"{switch_loss(torch.full((4,), .25), torch.full((4,), .25)):.2f}")

1.84 1.00

问题是这个损失和主任务的 loss 在打架，α 调大了会损害模型质量。DeepSeek-V3 提出了一种无辅助损失的平衡方法：给每个专家一个偏置 $b_i$，只在选 top-k 时加到分数上，不参与输出的加权；每一步之后，过载的专家把 $b_i$ 调低一点，闲置的调高一点。用一个天生偏心的路由器模拟一下：

import torch
torch.manual_seed(0)
N, E, k, gamma = 4096, 16, 2, 0.01
skew = torch.linspace(1.5, 0, E)                 # 模拟一个偏心的 router：前几个专家天生得分高
bias = torch.zeros(E)
for step in range(301):
    scores = torch.sigmoid(torch.randn(N, E) + skew)
    idx = (scores + bias).topk(k, -1).indices    # bias 只影响"选谁"，不进入输出的加权
    load = torch.bincount(idx.flatten(), minlength=E).float()
    if step in (0, 100, 300):
        print(f"step {step:3d}  max/mean load {load.max()/load.mean():.2f}  min/mean {load.min()/load.mean():.2f}")
    bias += gamma * torch.sign(load.mean() - load)   # 过载的专家降 bias，闲置的升 bias

step   0  max/mean load 2.40  min/mean 0.16
step 100  max/mean load 1.21  min/mean 0.82
step 300  max/mean load 1.16  min/mean 0.89

最忙的专家从平均负载的 2.4 倍降到 1.16 倍，最闲的从 0.16 倍升到 0.89 倍。偏置不进入梯度，不和主 loss 争；DeepSeek-V3 仍保留了一个权重很小的序列级平衡损失，防止单条序列内部极端不均。

MoE 训练还有一些其他的坑：router 的 softmax 在低精度下容易不稳定，通常用 fp32 计算并加 z-loss；微调时 MoE 比稠密模型更容易过拟合；也可以把训练好的稠密模型复制成多个专家再继续训练（upcycling），省下从头训练的成本。

停一下，检查理解单选

DeepSeek-V3 的无辅助损失负载均衡中，偏置 b_i 为什么只用于选择专家、不参与输出加权？

因为偏置太小，加进去对输出没有影响为了只改变路由的分配，而不改变被选中专家的输出贡献，从而不干扰主任务的梯度因为偏置是负数，加进去会让输出变号因为加进输出会让每个 token 激活的专家数变多

检查答案

拼起来核对一遍：DeepSeek-V3

DeepSeek-V3 把这篇的东西几乎都用上了：MLA 注意力、细粒度 MoE（每层 256 个路由专家 + 1 个共享专家，每个 token 选 8 个）、无辅助损失的负载均衡。官方给出的规模是总参数 671B、每 token 激活 37B。用公开的配置自己核对一下：

# DeepSeek-V3 配置（来自官方 config.json）
d, L, dense_L, V = 7168, 61, 3, 129280
H, hd_nope, hd_rope, hd_v, q_rank, kv_rank = 128, 128, 64, 128, 1536, 512
E, k, shared, d_exp, d_dense = 256, 8, 1, 2048, 18432
mla = (d * q_rank + q_rank * H * (hd_nope + hd_rope)      # q 先降到 1536 维再升回来
       + d * (kv_rank + hd_rope)                            # 压缩成 512 维 latent + 共享的 RoPE key
       + kv_rank * H * (hd_nope + hd_v)                     # latent 升维出每个头的 K、V
       + H * hd_v * d)                                      # 输出投影
expert = 3 * d * d_exp
moe_L = L - dense_L
total = L * mla + dense_L * 3 * d * d_dense + moe_L * (E + shared) * expert + 2 * V * d
active = L * mla + dense_L * 3 * d * d_dense + moe_L * (k + shared) * expert + 2 * V * d
print(f"attention/layer {mla/1e6:.0f}M, expert {expert/1e6:.1f}M")
print(f"total {total/1e9:.0f}B, active {active/1e9:.1f}B ({active/total:.1%})")
print(f"forward FLOPs/token ~ 2*active = {2*active/1e9:.0f} GFLOPs; FP8 weights {total/1e9:.0f} GB")

attention/layer 187M, expert 44.0M
total 671B, active 37.4B (5.6%)
forward FLOPs/token ~ 2*active = 75 GFLOPs; FP8 weights 671 GB

（这里忽略了 norm 和 router 的少量参数，也不含额外的 MTP 模块；embedding 查表其实不算 FLOPs，上面的估计略偏大。）两个数字都和官方吻合。这个结果说明了 MoE 的本质交换：每个 token 的计算相当于一个 37B 的稠密模型，但所有 671B 参数都要放在显存里，用 FP8 存也要 671 GB，至少要一整台 8 卡机器。MoE 省的是算力，不省显存，而且专家分布在多卡上带来了大量的 all-to-all 通信。

V3 还用了 MTP（Multi-Token Prediction）：主模型之外加一个小模块，在每个位置额外预测再往后一个 token，作为训练时的额外监督信号；推理时这个模块可以丢掉，也可以拿来做投机采样（见 推理系统）。

记住这三件事

1. decode 时注意力受显存带宽限制，KV cache 的大小决定了速度和并发：GQA 让多个 query 头共享 K/V，MLA 只缓存一个低秩 latent，DeepSeek-V3 每 token 每层只存 576 个数。
2. 长上下文的思路是不让每个 token 看全部前文：滑动窗口与全局层交错，线性注意力把历史压成固定状态（需要门控，常与普通注意力混合），稀疏注意力用便宜的索引器挑 top-k。
3. MoE 用路由器给每个 token 选 k 个专家，参数变多而每 token 计算不变；路由会坍缩，所以需要平衡损失或 DeepSeek-V3 式的偏置调节；它省算力不省显存。

延伸阅读

• Stanford CS336 Lecture 3（GQA/滑动窗口）与 Lecture 4（注意力替代方案与 MoE）
• Ainslie et al. 2023: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• DeepSeek-V2 技术报告（MLA、DeepSeekMoE） 与 DeepSeek-V3 技术报告
• Fedus et al. 2022: Switch Transformers
• Katharopoulos et al. 2020: Transformers are RNNs (线性注意力)
• Yang et al. 2024: Gated Delta Networks