MLA

前言 DeepSeek 提出的 MLA（Multi-head Latent Attention）通过将 KV 压缩到低维 latent 空间，大幅降低了推理时的 KV cache 开销。但在 DCP（Decode Context Parallel，即上下文并行）分布式环境下，MLA 的 prefill 阶段设计与 decode 阶段有显著差异。本文从实现角度展开分析。 什么是 DCP DCP（Decode Context Parallel）是一种将 KV cache 按序列维度切分到多个 GPU 的分布式策略。每个 rank 只持有完整 KV cache 的 1/dcp_world_size，从而减少单卡显存占用，支持更长的上下文。 与更常见的 DP（Data Parallel，扛并发）和 EP（Expert Parallel，分摊 MoE 参数显存）不同，DCP 解决的是单请求长上下文场景下 KV cache 放不下的问题。 MLA Prefill vs Decode：两条不同的路径 MLA 在 prefill 和 decode 阶段走了截然不同的计算路径： Prefill (forward_mha) Decode (forward_mqa) KV 形态 完整 MHA（N 头） Latent（1 头） Head dim P+R（~192） Lkv+R（~576） 计算特性 Sq ≈ Skv，计算密集 Sq ≪ Skv，避免显存搬运 Prefill 走 MHA 路径：kv_c 通过 W_UK/W_UV 解压成完整多头 K/V（N 个头），然后做标准的多头注意力。因为 prefill 时新 token 数和 context 长度在同一量级，展开 KV 做计算密集的 attention 是划算的。 ...

引言 Multi-head Latent Attention（MLA）是 DeepSeek-V2/V3 系列模型中最核心的架构创新之一。它在标准 Multi-Head Attention（MHA）的基础上引入了低秩压缩，大幅降低了 KV cache 的显存占用，同时保持了与 MHA 相当的模型质量。 本文从 QKV head dimension 处理差异 这一视角切入，深入分析 MLA 的设计原理、vLLM 中的具体实现，以及这种设计带来的性能收益。 1. 标准 MHA 回顾 在标准 MHA 中，给定输入序列 X ∈ ℝ^{S×D}，Q、K、V 通过三个独立的线性变换得到： Q = X @ W_Q → [S, H, d_head] K = X @ W_K → [S, H, d_head] V = X @ W_V → [S, H, d_head] 其中 d_head = D / H，三者完全相等。注意力计算为： Attention(Q, K, V) = softmax(Q @ K^T / √d_head) @ V → [S, H, d_head] 输出再经 W_O 映射回 D 维。这种 Q/K/V 共享同一 head dim 的设计深入人心，以至于很多人默认这是注意力机制的"必须要求"。 ...