从 LLaMA 到 DeepSeek · EP01

多头潜在注意力（上）

MLA · Part 1

新系列开篇。上一个系列「从GPT-2到LLaMA」里，我们把 GPT-2 的零件一个个换成了 LLaMA 的。这个系列继续往前走：在 LLaMA 的骨架上，DeepSeek 又换掉了哪些零件，为什么要换。

第一个零件，是 DeepSeek 最有代表性的设计：MLA（Multi-head Latent Attention，多头潜在注意力）。

一、先回到老问题：KV Cache 太大了

推理优化篇里我们讲过 KV Cache：生成第 t 个 token 时，前面所有 token 的 K 和 V 不用重算，直接从缓存里取。

代价是显存。每个 token、每一层，要存的数字个数是：

KV Cache（每 token 每层）= 2 × n_heads × head_dim
                           ↑ K 和 V 各一份

拿一个 DeepSeek 规模的配置算一下：128 个头，head_dim = 128。

2 × 128 × 128 = 32768 个数 / token / 层

更关键的是，decode 阶段每生成一个 token，都要把整个 KV Cache 从显存里读一遍。算力往往不是瓶颈，搬数据才是。所以 KV Cache 越小，能塞的上下文越长、能同时服务的请求越多、生成越快。

二、上一季的解法：GQA，让头共享 K/V

「从GPT-2到LLaMA」EP4 讲过 GQA 的思路：既然每个头都存一份 K/V 太贵，那就让几个 Query 头共用一组 K/V。

MHA：8 个 Q 头，8 组 K/V
GQA：8 个 Q 头，2 组 K/V（每 4 个头共用一组）
MQA：8 个 Q 头，1 组 K/V（所有头共用）

缓存确实小了，但有一个代价：共用同一组的头，看到的 K 一模一样。头与头之间的差异，只能靠各自的 Q 来体现。共享得越狠，表达能力损失越多。

那有没有办法，缓存像 MQA 一样小，但每个头又能拿到属于自己的 K？

三、MLA 的思路：不共享成品，共享原料

先用一个工厂的比喻（接着原系列「装配车间」「流水线」的说法）。

GQA 的做法是：几个工位共用同一批成品零件。省了仓库，但工位之间拿到的零件完全一样。

MLA 的做法是：仓库里不存成品，只存一份很小的毛坯件。每个工位按自己的图纸，把毛坯加工成自己需要的零件。

  • 仓库只存毛坯 → 缓存很小
  • 每个工位图纸不同 → 每个头的 K 依然不同

翻译成矩阵语言：

  1. 先把每个 token 的隐藏向量 h，压缩成一个很短的潜在向量 c（latent）。这一步叫下投影（down-projection）。
  2. KV Cache 只存 c。
  3. 用的时候，每个头用自己的矩阵把 c 还原成自己的 K 和 V。这一步叫上投影（up-projection）。

这就是名字里 Latent 的含义：缓存里存的不是 K/V 本身，而是生成 K/V 的潜在表示。

四、公式 + 猫吃鱼手算

沿用老例子：「猫吃鱼」3 个 token，d_model = 4。设定：

n_heads = 2，head_dim = 2（2 × 2 = 4 = d_model）
d_c = 2（潜在向量 c 的维度）

为了专注 MLA 本身，这一期先假装没有 RoPE。RoPE 和 MLA 之间有一个很微妙的冲突，留给下一期专门讲。

4.1 三个公式

c_s     = W_DKV · h_s        下投影：(d_c × d_model)，4 维 → 2 维
k_{s,i} = W_UK_i · c_s       头 i 的 K 上投影：(head_dim × d_c)
v_{s,i} = W_UV_i · c_s       头 i 的 V 上投影：(head_dim × d_c)

Q 照常计算：q_{t,i} = W_Q_i · h_t。

旁注：DeepSeek 实际上对 Q 也做了类似的低秩压缩。但 Q 不进缓存，压缩 Q 主要是为了省训练时的激活显存，和 KV Cache 无关，这里先略过。

4.2 第一步：下投影，得到要缓存的 c

假设经过 RMSNorm 后，三个 token 的隐藏向量是：

h_猫 = [1, 0, 1, 0]
h_吃 = [0, 1, 0, 1]
h_鱼 = [1, 1, 0, 0]

下投影矩阵（2 × 4）：

W_DKV = [[0.5, 0,   0.5, 0  ],
         [0,   0.5, 0,   0.5]]

逐个算：

c_猫 = W_DKV · h_猫 = [0.5×1 + 0.5×1,  0]        = [1,   0  ]
c_吃 = W_DKV · h_吃 = [0,  0.5×1 + 0.5×1]        = [0,   1  ]
c_鱼 = W_DKV · h_鱼 = [0.5×1,  0.5×1]            = [0.5, 0.5]

shape：(4,) → (2,)。缓存里存的就是这三个 2 维向量，总共 6 个数。

4.3 第二步：上投影，每个头还原出自己的 K

两个头用不同的上投影矩阵：

W_UK_1 = [[1, 0],      头 1：原样输出
          [0, 1]]

W_UK_2 = [[0, 1],      头 2：交换两个分量
          [1, 0]]

还原出来的 K：

token c 头 1 的 K 头 2 的 K
猫 [1, 0] [1, 0] [0, 1]
吃 [0, 1] [0, 1] [1, 0]
鱼 [0.5, 0.5] [0.5, 0.5] [0.5, 0.5]

关键就在这张表：缓存里每个 token 只有一份 c，但两个头还原出来的 K 是不同的。这是 MLA 和 GQA 的本质区别：GQA 共享的是 K 本身，MLA 共享的是生成 K 的原料。

4.4 第三步：照常算注意力

看最后一个 token「鱼」（位置 2，因果掩码下能看到猫、吃、鱼）。假设它在两个头上的 Query 都是：

q_鱼,1 = [1, 0]
q_鱼,2 = [1, 0]

注意力分数 = q · k（为了手算方便，先不除 √head_dim，不影响后面的对比）：

头 1：q · [k_猫, k_吃, k_鱼] = [1, 0, 0.5]   → 最关注「猫」
头 2：q · [k_猫, k_吃, k_鱼] = [0, 1, 0.5]   → 最关注「吃」

两个头的 Query 完全一样，却因为各自的 W_UK 不同，关注了不同的 token。如果换成 MQA，两个头看到的是同一组 K，分数也会完全一样。

五、到底省了多少

先看玩具例子，每个 token 每层要缓存的数字个数：

MHA：K 4 个 + V 4 个 = 8
MQA：K 2 个 + V 2 个 = 4
MLA：c 2 个           = 2   ← K 和 V 共用同一个 c

注意最后一行：K 和 V 是从同一个 c 还原出来的，所以只需要缓存一份。

再看真实配置。DeepSeek-V2/V3 的 MLA 参数是：128 个头，head_dim = 128，d_c = 512，另外还有一个 64 维的 RoPE 分量（下期讲）。

MHA：2 × 128 × 128 = 32768 个数 / token / 层
MLA：512 + 64      =   576 个数 / token / 层

压缩了约 57 倍。DeepSeek-V2 论文给了一个直观的换算：MLA 的缓存量，大约相当于只有 2.25 组的 GQA。

放到 DeepSeek-V3 的 61 层、BF16、128K 上下文下粗算：

MHA：约 500 GB
MLA：约 9 GB

前者一张卡都放不下，后者轻轻松松。

六、问题来了：每次都要还原 K，不就白省了？

看到这里，很自然会有一个疑问：

缓存里只存 c，可每次算注意力都得用 W_UK 把所有历史 token 的 c 还原成 K。上下文 10 万个 token，每一步都还原 10 万次，这不是把省下的显存换成了更多计算吗？

这就是 MLA 最巧妙的地方：矩阵吸收（absorb）。

6.1 把 W_UK 吸收进 Query

展开一个头的注意力分数：

score = q_t · k_s
      = q_t^T · (W_UK · c_s)
      = (W_UK^T · q_t)^T · c_s

只是换了一下括号的位置。但意义完全不同：

  • 原来：先把每个历史 token 的 c 还原成 k，再和 q 点积 → 还原 T 次
  • 现在：先把当前的 q 变换成 q̃ = W_UK^T · q，直接和 c 点积 → 变换 1 次

更进一步，q̃ = W_UK^T · W_Q · h_t，而 W_UK^T · W_Q 是两个固定的权重矩阵，推理前就可以离线乘好合并成一个。这样推理时，K 从头到尾都不会被真正算出来。

6.2 手算验证

拿头 2 验证，q_鱼,2 = [1, 0]：

q̃ = W_UK_2^T · q = [[0,1],[1,0]] · [1, 0] = [0, 1]

直接和缓存里的 c 点积：

q̃ · c_猫 = [0,1] · [1,   0  ] = 0
q̃ · c_吃 = [0,1] · [0,   1  ] = 1
q̃ · c_鱼 = [0,1] · [0.5, 0.5] = 0.5

结果 [0, 1, 0.5]，和 4.4 节先还原 K 再点积的结果完全一致。

也可以验证离线合并这条路。假设头 2 的 Query 矩阵是：

W_Q_2 = [[1, 0, 0, 0],
         [0, 0, 1, 0]]        （h_鱼 = [1,1,0,0] 算出来正是 [1, 0]）

离线合并：

W_UK_2^T · W_Q_2 = [[0, 0, 1, 0],
                    [1, 0, 0, 0]]

直接作用在 h_鱼 上：[0, 1]，就是上面的 q̃。一个矩阵，一步到位。

6.3 V 这边同理：把 W_UV 吸收进输出矩阵

注意力的输出是对 V 的加权求和：

o = Σ_s a_s · v_s = Σ_s a_s · (W_UV · c_s) = W_UV · (Σ_s a_s · c_s)

先在潜在空间里对 c 加权求和，最后只乘一次 W_UV。而 W_UV 后面紧跟着输出投影 W_O，两者同样可以离线合并成一个矩阵。

手算头 2。分数 [0, 1, 0.5] 过 softmax：

e^0 = 1，e^1 ≈ 2.718，e^0.5 ≈ 1.649，总和 ≈ 5.367
a ≈ [0.186, 0.507, 0.307]

设 W_UV_2 = [[1, 1], [0, 1]]。

原路线（先还原 V 再加权）：

v_猫 = [1, 0]，v_吃 = [1, 1]，v_鱼 = [1, 0.5]
o = 0.186×[1,0] + 0.507×[1,1] + 0.307×[1,0.5] ≈ [1.000, 0.661]

吸收路线（先对 c 加权，最后乘一次 W_UV）：

Σ a·c = 0.186×[1,0] + 0.507×[0,1] + 0.307×[0.5,0.5] ≈ [0.340, 0.661]
o = W_UV_2 · [0.340, 0.661] = [0.340 + 0.661, 0.661] ≈ [1.000, 0.661]

结果一致。

七、吸收不是免费的：用计算换显存

矩阵吸收也有代价，需要说清楚。

吸收之后，q̃ 的维度从 head_dim 变成了 d_c。在 DeepSeek 的配置里，是从 128 维变成 512 维，每个注意力分数的点积计算量变成了 4 倍。

之所以划算，是因为 decode 阶段的瓶颈在显存带宽而不在算力：多算一点无所谓，少读几十倍的数据才是关键。

反过来，在训练和 prefill 阶段，一次要处理大量 token，瓶颈变成了算力。这时实际实现通常不做吸收，而是直接把 K、V 还原出来，按普通的多头注意力算。同一套权重，两种算法，看阶段选择。

八、压缩了，模型能力会不会变差？

这是最常见的担心。有两点值得澄清。

第一，MLA 不是把训练好的模型拿来事后压缩。它是从训练一开始就带着这个结构。模型从一开始就只能通过 2 维（真实配置是 512 维）的 c 来传递 K/V 的信息，它会在这个约束下学会如何利用这些维度。

第二，和 GQA 比，MLA 保留了每个头 K/V 的差异性。DeepSeek-V2 论文的消融实验报告，MLA 的效果不输甚至略好于标准 MHA，而缓存只有它的零头。当然，这是论文在其实验设置下的结论，不同规模、不同数据下是否都成立，还有待更多验证。

九、PyTorch 实现

下面写一个不含 RoPE 的 MLA，同时实现两条路线：forward_naive 先还原 K/V，forward_absorbed 用矩阵吸收、全程不生成 K/V，最后验证两者结果一致。

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


class MLA(nn.Module):
    """Multi-head Latent Attention（暂不含 RoPE，下期补上）"""

    def __init__(self, d_model, n_heads, d_c):
        super().__init__()
        self.n_heads = n_heads
        self.d_h = d_model // n_heads
        self.d_c = d_c

        self.W_Q   = nn.Linear(d_model, n_heads * self.d_h, bias=False)
        self.W_DKV = nn.Linear(d_model, d_c, bias=False)              # 下投影
        self.W_UK  = nn.Linear(d_c, n_heads * self.d_h, bias=False)   # K 上投影
        self.W_UV  = nn.Linear(d_c, n_heads * self.d_h, bias=False)   # V 上投影
        self.W_O   = nn.Linear(n_heads * self.d_h, d_model, bias=False)

    def forward_naive(self, h):
        """路线一：先把 c 还原成 K、V，再走标准多头注意力"""
        B, T, _ = h.shape
        c = self.W_DKV(h)                                    # (B, T, d_c)  ← 只缓存它

        q = self.W_Q(h).view(B, T, self.n_heads, self.d_h).transpose(1, 2)
        k = self.W_UK(c).view(B, T, self.n_heads, self.d_h).transpose(1, 2)
        v = self.W_UV(c).view(B, T, self.n_heads, self.d_h).transpose(1, 2)
        # q, k, v: (B, H, T, d_h)

        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).reshape(B, T, -1)          # (B, T, H*d_h)
        return self.W_O(out)

    def forward_absorbed(self, h):
        """路线二：矩阵吸收，全程不生成 K、V"""
        B, T, _ = h.shape
        c = self.W_DKV(h)                                    # (B, T, d_c)
        q = self.W_Q(h).view(B, T, self.n_heads, self.d_h)   # (B, T, H, d_h)

        # nn.Linear 的 weight 形状是 (out, in)，拆成每个头一块
        W_UK = self.W_UK.weight.view(self.n_heads, self.d_h, self.d_c)  # (H, d_h, d_c)
        W_UV = self.W_UV.weight.view(self.n_heads, self.d_h, self.d_c)  # (H, d_h, d_c)

        # 吸收 W_UK：q̃ = W_UK^T · q
        q_tilde = torch.einsum('bthd,hdc->bhtc', q, W_UK)    # (B, H, T, d_c)

        # q̃ 直接和缓存里的 c 点积
        scores = torch.einsum('bhtc,bsc->bhts', q_tilde, c) / self.d_h ** 0.5
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=h.device), 1)
        scores = scores.masked_fill(mask, float('-inf'))
        attn = scores.softmax(dim=-1)                        # (B, H, T, T)

        # 先在潜在空间加权求和，再乘一次 W_UV
        ctx = torch.einsum('bhts,bsc->bhtc', attn, c)        # (B, H, T, d_c)
        out = torch.einsum('bhtc,hdc->bthd', ctx, W_UV)      # (B, T, H, d_h)
        return self.W_O(out.reshape(B, T, -1))


torch.manual_seed(0)
mla = MLA(d_model=4, n_heads=2, d_c=2)
h = torch.randn(1, 3, 4)                                     # 猫 吃 鱼

out1 = mla.forward_naive(h)
out2 = mla.forward_absorbed(h)
print("两条路线结果一致：", torch.allclose(out1, out2, atol=1e-6))

# 每 token 每层的缓存量对比
n_heads, d_h, d_c = 2, 2, 2
print("MHA 缓存:", 2 * n_heads * d_h)   # 8
print("MQA 缓存:", 2 * d_h)             # 4
print("MLA 缓存:", d_c)                 # 2

输出：

两条路线结果一致： True
MHA 缓存: 8
MQA 缓存: 4
MLA 缓存: 2

几点说明：

  • 为了对照，这里的吸收版本是在运行时做的 W_UK^T · q。真实部署中，会把 W_UK^T · W_Q 和 W_O · W_UV 离线合并成两个新矩阵，推理时一步到位。
  • 真正的 decode 循环里，每生成一个新 token，只需要把它的 c 追加进缓存，和推理优化篇的 KV Cache 写法一样，只是存的东西从 (K, V) 换成了 c。
  • 权重是随机初始化的，输出数值本身没有任何语义，这里只验证两条路线在数学上等价。

十、小结

MHA GQA / MQA MLA
缓存什么 每个头的 K、V 每组共享的 K、V 一个潜在向量 c
每个头的 K 是否不同 是 组内相同 是
缓存大小 大 中 / 小 很小
推理时额外开销 无 无 点积维度变大（用计算换显存）

一句话总结：MLA 共享的不是成品，而是原料；再通过矩阵吸收，让推理时根本不需要把成品造出来。

前置阅读：

下期预告：RoPE 卡住了矩阵吸收

这一期我们刻意跳过了 RoPE。把它加回来，会发生什么？

RoPE 要按位置旋转 q 和 k。加上 RoPE 后，注意力分数变成：

score = (R_t · W_Q · h_t)^T · (R_s · W_UK · c_s)
      = h_t^T · W_Q^T · R_t^T · R_s · W_UK · c_s

注意中间那个 R_t^T · R_s，它取决于两个 token 的相对位置，每一对 token 都不一样。它卡在 W_Q 和 W_UK 之间，导致这两个矩阵再也无法离线合并，这一期最巧妙的矩阵吸收直接失效。

DeepSeek 的解法叫 decoupled RoPE：把位置信息单独拆成一路很小的向量，不参与压缩，也不参与吸收。上面缓存计算里那个多出来的 64 维，就是它。

下期见。

配套互动：在矩阵里一格一格看 MLA 和解耦 RoPE →