10 · Multi-head Latent Attention (MLA)

注意力 DeepSeek 2024
DeepSeek-AI · arXiv:2405.04434

迭代了什么

上一代的问题:GQA/MQA 通过减少 KV 头数来压缩缓存,但牺牲模型容量。

这代改了什么:低秩 KV 联合压缩:投影到低维潜在空间再恢复,推理仅缓存压缩后的向量。算子吸收消除额外计算。

效果:KV 缓存降至 MHA 的 14%,训练成本节省 42.5%,8×H800 超 50K tokens/s。质量反而超过 MHA(BBH 50.7 vs 46.6)。

一、解决的问题

MHA 的 KV 缓存是推理瓶颈。GQA/MQA 通过减少 KV 头数来压缩,但会影响模型容量。需要既能大幅压缩 KV 缓存又不损失质量的方法。

二、核心创新:低秩 KV 联合压缩

将 Q/K/V 投影到低维潜在空间再恢复:

ctKV = WDKV · ht (压缩到 dc 维)
ktC = WUK · ctKV
vtC = WUV · ctKV

推理时仅缓存 ctKV(dc·l),而非 2·nh·dh·l。

三、关键工程技巧:算子吸收

利用矩阵乘法结合律,推理时将 WUK 吸收到 WQ,WUV 吸收到 WO,消除压缩的额外开销。

四、解耦 RoPE

低秩 KV 压缩与 RoPE 不兼容,引入额外的解耦 query 和共享 key 承载 RoPE:

qtR = RoPE(WQR · ctQ)
ktR = RoPE(WKR · ht) (所有头共享)

五、实际效果

DeepSeek V2:KV 缓存降至 MHA 的 14%,训练成本节省 42.5%,吞吐量超 50K tokens/s on 8×H800。