注意力效率

多头潜在注意力(MLA)

DeepSeek 如何将 KV 缓存压缩为单一潜在向量——在不损失质量的前提下缩减 9 倍显存

KV 缓存的显存之墙

现代语言模型逐个 token 生成文本。每一步都要为之前的每个 token 重新计算 Key 和 Value——除非把它们缓存起来。KV 缓存存储这些 Key 和 Value,避免重复劳动,正是它让自回归生成如此之快。

但 KV 缓存随上下文长度线性增长,也随注意力头数增长。以 Llama-2 70B 为例,它使用分组查询注意力(GQA),仅有 8 个 KV 头(而 Query 头有 64 个)。即使有这个优化,每个 token 的 KV 缓存仍需 2 × 8 × 128 × 80 × 2 字节 ≈ 320 KB(全部 80 层合计)。在 128K 上下文窗口下,仅 KV 缓存就高达 40 GB——而假设的 MHA 等价模型(全部 64 个头)代价会高出 10 倍:每 token 约 2.5 MB,128K 上下文下达约 320 GB。这就是 KV 缓存的显存之墙,也是长上下文模型服务时最大的瓶颈。

DeepSeek-V2 提出的多头潜在注意力(MLA)从源头攻击这个瓶颈。它不再按头缓存 Key 和 Value,而是将它们压缩成一个共享的潜在向量。结果:KV 缓存缩减最多 9 倍,让 DeepSeek 得以低成本服务 128K token 的超长上下文。

回顾:标准 MHA 为何如此吃显存

在标准多头注意力(MHA)中,每一层有 h 个注意力头。对每个 token,模型把隐藏状态投影成 h 个独立的 Query 向量、h 个 Key 向量和 h 个 Value 向量。生成时,所有过去 token 的 Key 和 Value 都被缓存——每个头各存一份。

# 每 token 的 KV 缓存(FP16,每个元素 2 字节) MHA: 2 × n_heads × head_dim × n_layers × 2 字节 GQA: 2 × n_kv_heads × head_dim × n_layers × 2 字节 (n_kv_heads < n_heads) # Llama-2 70B 实际配置(GQA):8 个 KV 头,128 维,80 层 GQA: 2 × 8 × 128 × 80 × 2 = 327,680 字节 ≈ 320 KB / token → 320 KB × 131,072 (128K 上下文) ≈ 40 GB # 假设的 MHA 等价模型(64 个头): MHA: 2 × 64 × 128 × 80 × 2 = 2,621,440 字节 ≈ 2.5 MB / token → 2.5 MB × 131,072 ≈ 320 GB ← 比 GQA 差 10 倍

分组查询注意力(GQA)和多查询注意力(MQA)通过在多个 Query 头之间共享 KV 头来缓解这个问题。GQA 使用少量 KV 头组(例如用 8 组替代 64 个),按比例削减缓存。但 GQA 面临质量权衡:KV 组越少,信息损失越大,模型表现越差。MLA 的洞见在于:你不必在质量和显存之间二选一。

核心思想:压缩,而非共享

MLA 押注了一条不同于 GQA 的路线。它不共享 KV 头,而是反问:如果一个 token 的 Key 和 Value 能从一个小得多的压缩表示中重建出来呢?如果能通过上投影从一个小潜在向量恢复完整的 K 和 V 矩阵,那么你只需缓存那个潜在向量。

机制如下。MLA 为每个 token 计算一个低维潜在向量 c_KV(通常仅 512 个数值,与头数无关)。这个潜在向量是唯一被缓存的东西。当模型需要注意力所需的 Key 和 Value 时,它通过学习到的权重矩阵将 c_KV 上投影回完整的按头 K 和 V 张量。昂贵的按头存储消失了,缓存里只剩紧凑的潜在向量。

# MLA 前向传播(每 token、每层) c_KV = W_DKV · h # 下投影隐藏状态 → 潜在向量(被缓存!) K = W_UK · c_KV # 上投影潜在向量 → 全部头的 Key(实时计算) V = W_UV · c_KV # 上投影潜在向量 → 全部头的 Value(实时计算) Q = W_UQ · c_Q # Query 同样从潜在向量 c_Q 上投影得到 # KV 缓存中只存 c_KV。K 和 V 每步都重新计算。

关键洞察在于:上投影是廉价的计算——一次矩阵乘法,微秒级完成。你省下的是显存带宽和显存容量,而这些才是推理时真正稀缺的资源。你用一点点额外的算力,换来了缓存的大幅缩减。

MLA 与 MHA 的 3D 对比

下面的可视化并排对比了一个标准多头注意力层和一个多头潜在注意力层。左侧,每个头维护各自的 Key 和 Value 流,在缓存中不断累积。右侧,所有头共享一个压缩的潜在向量——只有这个潜在向量被缓存,而按头的 K 和 V 则按需重建。

左侧 MHA:4 个头各自缓存 K(金色)和 V(蓝色),共 8 个缓存张量。右侧 MLA:仅缓存一个潜在向量 c_KV(青色),K 和 V 在运行时按头上投影得到。拖动旋转,对比缓存占用的差异。

显存节省:用数字说话

数字惊人。DeepSeek-V2 使用 128 个注意力头,每个头的 Key 拆分为 128 维无 RoPE 部分和 64 维携带 RoPE 部分,另有 128 维 Value——因此等价 MHA 每 token 存储为 128 × (128 + 64 + 128) = 40,960 个元素。使用 MLA,只缓存 512 维潜在向量 c_KV 和 64 维 RoPE 向量——每层仅 576 个元素。这实现了每层 KV 缓存约 71 倍的缩减。

# DeepSeek-V2 每层、每 token 的 KV 缓存 (FP16) MHA 等价: 128 头 × (128 K_nope + 64 K_rope + 128 V) = 40,960 个元素 MLA 缓存: 512 (潜在 c_KV) + 64 (RoPE 向量) = 576 个元素 $$\text{Reduction} = \frac{40{,}960}{576} \approx 71\times \text{per layer}$$ # 实际中,MLA 通常与同等质量的 GQA 基线比较 # (GQA 已经共享 KV 头,差距缩小)。在同等质量下与 GQA 对比, # 净 KV 缓存缩减约 9 倍到 14 倍——仍然非常可观。

与同等质量的 GQA 基线对比后(GQA 本身已共享 KV 头,差距缩小),净有效 KV 缓存缩减约为 9 倍到 14 倍。上投影矩阵 W_UK 和 W_UV 是模型权重,不是逐 token 缓存——它们驻留在显存中一次,在所有 token 间复用。MLA 削减的是随上下文长度增长的逐 token、逐序列缓存,这才是真正撑爆服务预算的成本。

RoPE 的麻烦(以及巧妙的解法)

还有一个微妙之处,也是 MLA 并不简单的原因。Transformer 中的位置信息通过旋转位置编码(RoPE)注入,它会根据位置旋转 Query 和 Key 向量。RoPE 必须在 Key 缓存之前应用——否则注意力计算时模型无法比较位置。

这给 MLA 带来了难题。如果 Key 在运行时从潜在向量重建,RoPE 该在哪里应用?你不能对潜在向量应用 RoPE(维度不同),也不能缓存已应用 RoPE 的 Key(那样就失去了压缩的意义)。DeepSeek 的解法很优雅:把每个 Key 拆成对 RoPE 敏感的部分和与 RoPE 无关的部分。与 RoPE 无关的部分被压缩进潜在向量并缓存;对 RoPE 敏感的部分是一个单独携带的小额外向量。两者在注意力计算时合并。

# MLA 中的解耦 RoPE K = [ K_no_rope ; K_rope ] K_no_rope = W_UK · c_KV # 压缩在潜在向量中,上投影,不带 RoPE K_rope = RoPE(W_KR · h) # 独立的小分支,携带位置信息 # 注意力使用拼接后的完整带位置 Key, # 但主体(K_no_rope)在缓存中保持压缩状态。

这种解耦 RoPE 设计是 MLA 能在工程上落地的关键细节。它保留了注意力的位置敏感性,同时保住了压缩收益。这也是 MLA 不能直接替换 MHA 的原因——注意力内核必须重写,以处理这种拆分的 Key 结构。

RoPE 解耦可视化

这个场景展示了 MLA 最重要的架构洞见:每个 token 的 Key 如何分成两条路径。主体(K_no_rope)被压缩进潜在向量并缓存;一个小的携带 RoPE 的向量(K_rope)走另一条独立路径,在缓存前应用位置编码。两者都被缓存——只是路径不同。

观察 Key 分成两条路径:主体(K_no_rope)被压缩进潜在向量;一个小的携带 RoPE 的向量独立行走。两者都被缓存,只是路径不同。

压缩会损害质量吗?

自然的担忧:如果 Key 和 Value 被压缩进潜在向量再重建,信息会丢失吗?实测下来,几乎没有。DeepSeek-V2 在基准测试上持平甚至超越同等规模的 GQA 模型,同时只用了其零头的 KV 缓存。原因是潜在向量并非任意压缩——它是一个学习到的低秩投影。模型联合训练下投影和上投影矩阵,使重建出的 K 和 V 恰好保留注意力所需的信息。

事实上,在同等缓存预算下,MLA 往往略胜 GQA。GQA 笨拙地共享 KV 头——同一组内的所有 Query 头看到完全相同的 Key。MLA 则给每个 Query 头各自的重建 Key,保留了头的多样性,同时通过低秩瓶颈而非头共享来实现压缩。你以共享头的代价,获得了按头的丰富性。

MLA 在生产中:DeepSeek-V2 与 V3

MLA 是 DeepSeek 模型家族标志性的架构创新。DeepSeek-V2(2024)首次引入它,取代了第一代使用的 GQA。DeepSeek-V3(2024 年底)保留 MLA,并与 DeepSeekMoE 的细粒度专家结合,产出一个 6710 亿参数的模型(每 token 激活 370 亿),在远低于前沿闭源模型的服务成本下与之比肩。

实际影响:DeepSeek-V3 能在 8 张 H800 上服务 128K token 上下文,KV 缓存小到足以装下。同等规模的 MHA 模型仅缓存就需要在更多 GPU 间分片,成倍推高服务成本。MLA 很大程度上解释了为何 DeepSeek 能在 2024-2025 年提供比竞争对手低一个数量级的 API 定价。这是架构创新直接转化为经济优势的最清晰案例之一。

核心要点

1

KV 缓存随上下文长度和头数增长,是长上下文模型的主要显存开销——往往比权重本身还大。这就是为什么扩展上下文本质上是内存问题,而非计算问题。

2

GQA 和 MQA 通过共享 KV 头来削减缓存,但面临质量权衡。MLA 走了不同的路:压缩,而非共享。这个区别很重要——压缩保留了每个头的独特性,而共享会模糊它。

3

MLA 每 token 只缓存一个低维潜在向量,通过学习到的上投影实时重建按头的 Key 和 Value——缓存缩减 9-14 倍。赌注在于计算比显存带宽便宜,而这在整个 GPU 时代一直是成立的。

4

解耦 RoPE 把每个 Key 拆成压缩的无关 RoPE 部分和小的带位置部分,在压缩方案内保留位置敏感性。这是一个优雅的绕行方案:RoPE 太重要不能压缩,所以 MLA 选择绕过它。

5

MLA 驱动 DeepSeek-V2 和 V3,以同等 MHA 模型零头的成本服务 128K 上下文——架构创新直接转化为经济优势。DeepSeek 在 2024-2025 年比竞争对手低一个数量级的 API 定价就是证据。

探索相关主题:

深入了解周边技术: