前言

学习 DeepSeek 中的 MLA 模块，究极缝合怪，东抄抄西抄抄，主要 copy 自苏神的文章，仅供自己参考😄

refer1：缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA

refer2: 博客分享：从MHA、MQA、GQA到MLA

1. 从MHA、MQA、GQA到MLA

以下内容均来自于苏神的文章：缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA

1.1 MHA

MHA（Multi-Head Attention），也就是多头注意力，是开山之作 《Attention is all you need》 所提出的一种 Attention 的形式

在数学上，多头注意力 MHA 等价于多个独立的单头注意力的拼接，假设输入的（行）向量序列为 ${\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_l$，其中 ${\mathbf{x}}_i\in {\mathbb{R}}^d$，那么 MHA 可以形式地记为：

$$\begin{gathered} {\mathbf{o}}_{\mathbf{t}}=\left[{\mathbf{o}}_{\mathbf{t}}^{(1)},{\mathbf{o}}_{\mathbf{t}}^{(2)},\cdots ,{\mathbf{o}}_{\mathbf{t}}^{(\mathbf{h})}\right] \\ {\mathbf{o}}_t^{(s)}=\mathbf{Attention}\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)} \\ \begin{array}{l}{\mathbf{q}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{k}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{v}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{v}}},{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{v}}}\end{array} \end{gathered}$$

简单起见，这里省略了 Attention 矩阵的缩放因子 $\frac{1}{\sqrt{d_k}}$

实践上，常见的设置是 $d_k=d_v=d/h$，例如对于 LLaMA2-7B 有 $d=4096,h=32,d_k=d_v=128$，LLaMa2-70B 则是 $d=8192,h=64,d_k=d_v=128$

这里只考虑主流自回归 LLM 所用的 Causal Attention，在 token by token 递归生成时，新预测出来的第 $t+1$ 个 token，并不会影响到已经算好的 $k_{\le t}^{(s)},v_{\le t}^{(s)}$，因此这部分结果我们可以缓存下来，供后续生成调用，避免不必要的重复计算，这就是所谓的 KV Cache

关于 KV Cache 大家感兴趣的可以看看：KV Cache的原理与实现

后面的 MQA、GQA、MLA 都是围绕“如何减少 KV Cache 同时尽可能地保证效果”这个主题发展而来的产物

上图展示了标准 MHA 下的 KV Cache 是多大，它和注意力头数、序列长度等相关，此时 KV Cache 的大小是 $2\ast seq\_len\ast num\_head\ast head\_dim$

Note：该图片来自于 https://github.com/preacher-1/MLA_tutorial

代码实现如下：

import math
import torch
import torch.nn as nn# Multi-Head Attention
class MHA(nn.Module):def __init__(self, d_model, num_heads):super().__init__()self.d_model = d_modelself.num_heads = num_headsassert d_model % num_heads == 0self.head_dim = d_model // num_headsself.q_linear = nn.ModuleList([nn.Linear(d_model, self.head_dim, bias=False) for _ in range(num_heads)])self.k_linear = nn.ModuleList([nn.Linear(d_model, self.head_dim, bias=False) for _ in range(num_heads)])self.v_linear = nn.ModuleList([nn.Linear(d_model, self.head_dim, bias=False) for _ in range(num_heads)])self.out_linear = nn.Linear(d_model, d_model, bias=False) def forward(self, x):bsz, seq_len, _ = x.shapeoutputs = []# Parallelfor i in range(self.num_heads):q = self.q_linear[i](x) # (bsz, seq_len, head_dim)k = self.k_linear[i](x) # (bsz, seq_len, head_dim)v = self.v_linear[i](x) # (bsz, seq_len, head_dim)# RoPE# TODO: Implement RoPE# Attentionattention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # (bsz, seq_len, seq_len)# Casual maskmask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()mask = mask.unsqueeze(0).to(x.device)   # (1, seq_len, seq_len)attention = attention.masked_fill(mask, float('-inf'))attention = torch.softmax(attention, dim=-1)# Outputoutput = torch.matmul(attention, v) # (bsz, seq_len, seq_len)outputs.append(output)# Linear projectionoutput = torch.cat(outputs, dim=-1) # (bsz, seq_len, d_model)output = self.out_linear(output)return output# Another implement for Multi-Head Attention
class MHA2(nn.Module):def __init__(self, d_model, num_heads):super().__init__()self.d_model = d_modelself.num_heads = num_headsassert d_model % num_heads == 0self.head_dim = d_model // num_headsself.q_linear = nn.Linear(d_model, num_heads * self.head_dim, bias=False)self.k_linear = nn.Linear(d_model, num_heads * self.head_dim, bias=False)self.v_linear = nn.Linear(d_model, num_heads * self.head_dim, bias=False)self.out_linear = nn.Linear(d_model, d_model, bias=False)def forward(self, x):bsz, seq_len, _ = x.shapeq = self.q_linear(x) # (bsz, seq_len, num_heads * head_dim)k = self.k_linear(x) # (bsz, seq_len, num_heads * head_dim)v = self.v_linear(x) # (bsz, seq_len, num_heads * head_dim)# matmul 只能在最后两个维度相乘, 需要对 NxD 的矩阵相乘, 做 1,2 维度的交换# (bsz, num_heads, seq_len, head_dim)q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)k = k.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)v = v.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)# RoPE# TODO: Implement RoPE# Attentionattention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # (bsz, num_heads, seq_len, seq_len)# Casual maskmask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()mask = mask.unsqueeze(0).to(x.device)   # (1, seq_len, seq_len)attention = attention.masked_fill(mask, float('-inf'))attention = torch.softmax(attention, dim=-1)# Outputoutput = torch.matmul(attention, v) # (bsz, num_heads, seq_len, head_dim)output = output.transpose(1, 2) # (bsz, seq_len, num_heads, head_dim)  output = output.contiguous().view(bsz, seq_len, -1) # (bsz, seq_len, d_model)output = self.out_linear(output)return output# Example usage
torch.manual_seed(10)
d_model = 512
num_heads = 8
mha = MHA2(d_model, num_heads)
x = torch.randn(10, 20, d_model)  # (bsz, seq_len, d_mdeol)
output = mha(x)
print(output.shape) # (10, 20, 512)

Note：代码参考自：https://github.com/preacher-1/MLA_tutorial

此外 MHA 还有另外一种实现，一次性将 $Q,K,V$ 投影，具体代码可以参考 https://github.com/karpathy/minGPT/tree/master/mingpt

1.2 瓶颈

为什么降低 KV Cache 的大小如此重要呢？🤔

众所周知，一般情况下 LLM 的推理都是在 GPU 上进行的，而单张 GPU 的显存是有限的，一部分我们要用来存放模型的参数和前向计算的激活值，这部分依赖于模型的体量，选定模型后它就是个常数；另外一部分我们要用来存放模型的 KV Cache，这部分不仅依赖于模型的体量，还依赖于模型的输入长度，也就是在推理过程中是动态增长的，当 Context 长度足够长时，它的大小就会占主导地位，可能超过一张卡甚至一台机（8张卡）的总显存量

在 GPU 上部署模型的原则是：能一张卡部署的，就不要跨多张卡；能一台机部署的，就不要跨多台机。这是因为“卡内通信带宽 > 卡间通信带宽 > 机间通信带宽”，由于“木桶效应”，模型部署时跨的设备越多，受设备间通信带宽的的“拖累”就越大，事实上即便是单卡 H100 内 SRAM 与 HBM 的带宽已经达到了 3TB/s，但对于 Short Context 来说这个速度依然还是推理的瓶颈，更不用说更慢的卡间、机间通信了

所以，减少 KV Cache 的目的就是要实现在更少的设备上推理更长的 Context，或者在相同的 Context 长度下让推理的 batch size 更大，从而实现更快的推理速度或者更大的吞吐量。当然，最终的目的都是为了实现更低的推理成本

1.3 MQA

MQA（Multi-Query Attention）是减少 KV Cache 的一次非常朴素的尝试，首次提出自 《Fast Transformer Decoding: One Write-Head is All You Need》

MQA 的思路很简单，直接让所有 Attention Head 共享同一个 K、V，用公式来说就是取消 MHA 所有的 $\mathbf{k},\mathbf{v}$ 的上标 ${}^{(s)}$

$$\begin{gathered} {\mathbf{o}}_{\mathbf{t}}=\left[{\mathbf{o}}_{\mathbf{t}}^{(1)},{\mathbf{o}}_{\mathbf{t}}^{(2)},\cdots ,{\mathbf{o}}_{\mathbf{t}}^{(\mathbf{h})}\right] \\ {\mathbf{o}}_t^{(s)}=\mathbf{Attention}\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{\bcancel{(s)}},{\mathbf{v}}_{\le t}^{\bcancel{(s)}}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{\bcancel{(s)}}{}^{\top }\right){\mathbf{v}}_i^{\bcancel{(s)}}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{\bcancel{(s)}}{}^{\top }\right)} \\ \begin{array}{l}{\mathbf{q}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{k}}_{\mathbf{i}}^{\bcancel{(\mathbf{s})}}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{\bcancel{(\mathbf{s})}}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{k}}^{\bcancel{(\mathbf{s})}}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{v}}_{\mathbf{i}}^{\bcancel{(\mathbf{s})}}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{\bcancel{(\mathbf{s})}}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{v}}},{\mathbf{W}}_{\mathbf{v}}^{\bcancel{(\mathbf{s})}}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{v}}}\end{array} \end{gathered}$$

使用 MQA 的模型包括 PaLM、StarCoder、Gemini 等。很明显，MQA 直接将 KV Cache 减少到了原来的 $1/h$，这是非常可观的，单从节省显存角度看已经是天花板了

效果方面，目前看来大部分任务的损失都比较有限，且 MQA 的支持者相信这部分损失可以通过进一步训练来弥补回。此外，注意到 MQA 由于共享了 K、V，将会导致 Attention 的参数量减少了将近一半，而为了模型总参数量的不变，通常会相应地增大 FFN/GLU 的规模，这也能弥补一部分效果损失

Note：该图片来自于 https://github.com/preacher-1/MLA_tutorial

上图展示了标准 MQA 下的 KV Cache 是多大，和标准的 MHA 相比 $Q$ 保持不变，但所有头的 $K,V$ 共享，此时 KV Cache 的大小是 $2\ast seq\_len\ast 1\ast head\_dim$

代码实现如下：

import math
import torch
import torch.nn as nn# Multi-Query Attention
class MQA(nn.Module):def __init__(self, d_model, num_heads):super().__init__()self.d_model = d_modelself.num_heads = num_headsassert d_model % num_heads == 0self.head_dim = d_model // num_headsself.q_linear = nn.Linear(d_model, d_model)self.k_linear = nn.Linear(d_model, self.head_dim)self.v_linear = nn.Linear(d_model, self.head_dim)self.out_linear = nn.Linear(d_model, d_model)def forward(self, x):bsz, _, _ = x.shape# Linear projections, all heads share the same K, Vq = self.q_linear(x) # (bsz, seq_len, d_model)k = self.k_linear(x) # (bsz, seq_len, head_dim)v = self.v_linear(x) # (bsz, seq_len, head_dim)# Reshape for multi-head attentionq = q.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2)# (bsz, num_heads, seq_len, head_dim)k = torch.unsqueeze(k, 1) # (bsz, 1, seq_len, head_dim)v = torch.unsqueeze(v, 1) # (bsz, 1, seq_len, head_dim)# Attentionattention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # (bsz, num_heads, seq_len, seq_len)attention = torch.softmax(attention, dim=-1)# Outputoutput = torch.matmul(attention, v) # (bsz, num_heads, seq_len, head_dim)output = output.transpose(1, 2) # (bsz, seq_len, num_heads, head_dim)  output = output.contiguous().view(bsz, -1, d_model)# Linear projectionoutput = self.out_linear(output)return output# Example usage
torch.manual_seed(10)
d_model = 512
num_heads = 8
mqa = MQA(d_model, num_heads)
x = torch.randn(10, 20, d_model)  # (bsz, seq_len, d_mdeol)
output = mqa(x)
print(output.shape) # (10, 20, 512)

Note：代码参考自：https://github.com/preacher-1/MLA_tutorial

这个代码和 MHA 实现类似，不同的是由于 MQA 所有头共享同一个 $K,V$，因此这里的 $W_k,W_v$ 投影矩阵的维度是 head_dim 而不再是 num_heads * head_dim，在 forward 时通过广播机制将 $W_k,W_v$ 共享到其他头即可

1.4 GQA

然而，也有人担心 MQA 对 KV Cache 的压缩太严重，以至于会影响模型的学习效率以及最终结果。为此，一个 MHA 和 MQA 之间的过渡版本 GQA（Grouped-Query Attention）应运而生，出自论文 《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》，是 23 年的工作

GQA 的思想也很朴素，它就是将所有 Head 分为 $g$ 个组（$g$ 可以整除 $h$），每组共享同一对 K、V，用数学公式表示为：

$$\begin{gathered} {\mathbf{o}}_{\mathbf{t}}=\left[{\mathbf{o}}_{\mathbf{t}}^{(1)},{\mathbf{o}}_{\mathbf{t}}^{(2)},\cdots ,{\mathbf{o}}_{\mathbf{t}}^{(\mathbf{h})}\right] \\ {\mathbf{o}}_t^{(s)}=\mathbf{Attention}\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(\lceil sg/h\rceil )},{\mathbf{v}}_{\le t}^{(\lceil sg/h\rceil )}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(\lceil sg/h\rceil )}{}^{\top }\right){\mathbf{v}}_i^{(\lceil sg/h\rceil )}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(\lceil sg/h\rceil )}{}^{\top }\right)} \\ {\mathbf{q}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}} \\ {\mathbf{k}}_{\mathbf{i}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{k}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}} \\ {\mathbf{v}}_{\mathbf{i}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{v}}},{\mathbf{W}}_{\mathbf{v}}^{(\lceil \mathbf{sg}/\mathbf{h}\rceil )}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{v}}} \end{gathered}$$

其中 $\lceil \cdot \rceil$ 是上取整符号

GQA 提供了从 MHA 到 MQA 的自然过渡，当 $g=h$ 时就是 MHA；当 $g=1$ 时就是 MQA；当 $1<g<h$ 时，它只将 KV Cache 压缩到 $g/h$，压缩率不如 MQA，但同时也提供了更大的自由度，效果上更有保证。GQA 最知名的使用者，大概是 Meta 开源的 LLAMA2-70B，以及 LLAMA3 全系列，此外使用 GQA 的模型还有 TigerBot、DeepSeek-V1、StarCoder2、Yi、ChatGLM2、ChatGLM3 等，相比使用 MQA 的模型更多

Note：该图片来自于 https://github.com/preacher-1/MLA_tutorial

上图展示了标准 GQA 下的 KV Cache 是多大，GQA 是 MHA 和 MQA 的一种折中，它将 $K,V$ 分成 group 组，每组共享同一个 $K,V$，此时 KV Cache 的大小是 $2\ast seq\_len\ast n\_groups\ast head\_dim$

代码实现如下：

import math
import torch
import torch.nn as nn# Grouped-Query Attention
class GQA(torch.nn.Module):def __init__(self, d_model, num_heads, num_groups):super().__init__()self.d_model = d_modelself.num_heads = num_headsself.num_groups = num_groupsself.group_heads = num_heads // num_groupsself.head_dim = d_model // num_headsself.W_q = nn.Linear(d_model, d_model)self.W_k = nn.Linear(d_model, self.head_dim * num_groups)self.W_v = nn.Linear(d_model, self.head_dim * num_groups)self.out_linear = nn.Linear(d_model, d_model)def forward(self, x):bsz, seq_len, _ = x.shape# Linear projections, each group share the same K, Vq = self.W_q(x) # (bsz, seq_len, d_model)k = self.W_k(x) # (bsz, seq_len, head_dim * num_groups)v = self.W_v(x) # (bsz, seq_len, head_dim * num_groups)# Reshape for multi-head attention# (bsz, num_groups, gropus_head, seq_len, head_dim)q = q.view(bsz, seq_len, self.num_groups, self.group_heads, self.head_dim).permute(0, 2, 3, 1, 4)k = k.view(bsz, seq_len, self.num_groups, self.head_dim).transpose(1, 2) # (bsz, num_groups, seq_len, head_dim)v = v.view(bsz, seq_len, self.num_groups, self.head_dim).transpose(1, 2) # (bsz, num_groups, seq_len, head_dim)k = torch.unsqueeze(k, 2) # (bsz, num_groups, 1, seq_len, head_dim)v = torch.unsqueeze(v, 2) # (bsz, num_groups, 1, seq_len, head_dim)# Attentionattention = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # (bsz, num_gropus, gropus_head, seq_len, seq_len)attention = torch.softmax(attention, dim=-1)# Outputoutput = torch.matmul(attention, v) # (bsz, num_groups, gropus_head, seq_len, head_dim)output = output.permute(0, 3, 1, 2, 4).contiguous().view(bsz, -1, self.d_model)# Linear projectionoutput = self.out_linear(output)return output# Example usage
torch.manual_seed(10)
d_model = 512
num_heads = 8
num_groups = 4
gqa = GQA(d_model, num_heads, num_groups)
x = torch.randn(32, 10, d_model)  # (bsz, seq_len, d_mdeol)
output = gqa(x)
print(output.shape) # (32, 10, 512)

Note：代码参考自：https://github.com/preacher-1/MLA_tutorial

1.5 MLA

有了 MHA、MQA、GQA 的铺垫，我们理解起 MLA（Multi-head Latent Attention）就相对容易一些了。DeepSeek-V2 的技术报告里是从低秩投影（类似于 LoRA）的角度引入 MLA 的，但苏神认为低秩投影这个角度并不贴近本质，MLA 的本质是低秩投影之后的工作

1.5.1 Part 1

GQA 在投影之后做了什么呢？首先它将向量对半分两份分别作为 K、V，然后每一份又均分为 $g$ 份，每一份复制 $h/g$ 次，以此来“凑”够 $h$ 个 Attention Head 所需要的 K、V

我们知道分割、复制都是简单的线性变换，所以 MLA 的第一个想法是将这些简单的线性变换换成一般的线性变换，以增强模型的能力：

$$\begin{gathered} {\mathbf{o}}_{\mathbf{t}}=\left[{\mathbf{o}}_{\mathbf{t}}^{(1)},{\mathbf{o}}_{\mathbf{t}}^{(2)},\cdots ,{\mathbf{o}}_{\mathbf{t}}^{(\mathbf{h})}\right] \\ {\mathbf{o}}_t^{(s)}=\mathbf{Attention}\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)} \\ \begin{array}{l}{\mathbf{q}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{k}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{k}}},{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{c}}\times {\mathbf{d}}_{\mathbf{k}}}\\ {\mathbf{v}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{v}}},{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{c}}\times {\mathbf{d}}_{\mathbf{v}}}\end{array} \\ {\mathbf{c}}_{\mathbf{i}}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{c}}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{c}}},{\mathbf{W}}_{\mathbf{c}}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{c}}} \end{gathered}$$

这里博主有些困惑，那看了 博客分享：从MHA、MQA、GQA到MLA 视频之后大概理解了 MLA 想要做的事情，前面我们提到 GQA 的实现中有一些分割、复制的变换，MLA 出于增强 GQA 性能的目的，想要将这些简单的线性变换加上一些可学习的参数让其变成一般的线性变换

以下分析内容来自于：博客分享：从MHA、MQA、GQA到MLA

MQA 和 GQA 的“升维”投影矩阵

原始的 GQA 先将输入 $x_i$ 分别压缩到 $g{d}_k$ 和 $g{d}_v$ 维，再复制 $g$ 份得到可以直接和 $q_i$ 相乘的 $k_i$ 和 $v_i$。将 GQA 的投影矩阵记为 ${\mathbf{W}}_c$，那么有：

$$\begin{gathered} {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{g(d_k+d_v)},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times g(d_k+d_v)} \\ {\mathbf{c}}_i=[{\mathbf{k}}_i^{(1)},\cdots ,{\mathbf{k}}_i^{(g)},{\mathbf{v}}_i^{(1)},\cdots ,{\mathbf{v}}_i^{(g)}]=[{\mathbf{ck}}_i,{\mathbf{cv}}_i] \\ {\mathbf{W}}_{split}^k=\left[\begin{array}{c}{\mathbf{I}}_{g{d}_k}\\ 0_{g{d}_k}\end{array}\right]{\mathbf{W}}_{split}^v=\left[\begin{array}{c}{0}_{g{d}_v}\\ {\mathbf{I}}_{g{d}_v}\end{array}\right] \\ {\mathbf{ck}}_i={\mathbf{c}}_i{\mathbf{W}}_{split}^k\in {\mathbb{R}}^{g{d}_k},{\mathbf{cv}}_i={\mathbf{c}}_i{\mathbf{W}}_{split}^v\in {\mathbb{R}}^{g{d}_v} \end{gathered}$$
这里 ${\mathbf{W}}_{split}^k,{\mathbf{W}}_{split}^v$ 实现了形式上的分割操作，得到 ${\mathbf{ck}}_i,{\mathbf{cv}}_i$。下面我们将构造“复制”操作的投影矩阵：
$${\mathbf{W}}_k\in {\mathbb{R}}^{g{d}_k\times h{d}_k}=\left[\begin{array}{ccccccccccccc}{\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}& 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}\\ 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& {\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}\\ ⋮& ⋮& \ddots & ⋮& ⋮& ⋮& \ddots & ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& \cdots & {\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}\end{array}\right]$$
其中每行 ${\mathbf{I}}_{d_k}$ 重复 $h/g$ 遍，代表从 groups 到 heads 的“放缩”倍数或“复制”次数；一共有 $g$ 行，对应原来的 $d_i^{(s)}$，共 $g$个。${\mathbf{W}}_v\in {\mathbb{R}}^{g{d}_v\times h{d}_v}$ 的形式与 ${\mathbf{W}}_k$ 相同，故不赘述。将前者左乘 ${\mathbf{ck}}_i$，则有：

$$\begin{array}{rl}{\mathbf{ck}}_i{\mathbf{W}}_k& =[{\mathbf{k}}_i^{(1)},\cdots ,{\mathbf{k}}_i^{(g)}]\cdot \left[\begin{array}{ccccccccccccc}{\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}& 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}\\ 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& {\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}\\ ⋮& ⋮& \ddots & ⋮& ⋮& ⋮& \ddots & ⋮& \ddots & ⋮& ⋮& \ddots & ⋮\\ 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& 0_{d_k}& 0_{d_k}& \cdots & 0_{d_k}& \cdots & {\mathbf{I}}_{d_k}& {\mathbf{I}}_{d_k}& \cdots & {\mathbf{I}}_{d_k}\end{array}\right]\\ & =[{\mathbf{k}}_i^{(1)}{\mathbf{I}}_{d_k},{\mathbf{k}}_i^{(1)}{\mathbf{I}}_{d_k},\cdots ,{\mathbf{k}}_i^{(1)}{\mathbf{I}}_{d_k},{\mathbf{k}}_i^{(2)}{\mathbf{I}}_{d_k},{\mathbf{k}}_i^{(2)}{\mathbf{I}}_{d_k},\cdots ,{\mathbf{k}}_i^{(2)}{\mathbf{I}}_{d_k},\cdots ,{\mathbf{k}}_i^{(g)}{\mathbf{I}}_{d_k},{\mathbf{k}}_i^{(g)}{\mathbf{I}}_{d_k},\cdots ,{\mathbf{k}}_i^{(g)}{\mathbf{I}}_{d_k}]\\ & =[{\mathbf{k}}_i^{(1)},\cdots ,{\mathbf{k}}_i^{(1)},{\mathbf{k}}_i^{(2)},\cdots ,{\mathbf{k}}_i^{(2)},\cdots ,{\mathbf{k}}_i^{(g)},\cdots ,{\mathbf{k}}_i^{(g)}]\in {\mathbb{R}}^{h{d}_k}\end{array}$$
于是我们就得到了维度为 $h{d}_k$ 的 ${\mathbf{k}}_i$，其中每个 ${\mathbf{k}}_i^{(s)}$ 都被复制了 $h/g$ 次，实现了“复制”操作。同理，${\mathbf{cv}}_i{\mathbf{W}}_v$ 得到${\mathbf{v}}_i$，维度为 $h{d}_v$。
这里我们讨论的都是单个 token 的行向量，而对于实际输入序列，其（最后两个）维度为 (seq_len, d)，同样可以直接替换上面的单一向量。

在上面所构造的所有矩阵中，最重要的是 ${\mathbf{W}}_k$ 和 ${\mathbf{W}}_v$，可以看出两者都是由若干单位矩阵组成的稀疏矩阵，是 GQA 的分割、复制操作的矩阵形式描述，那么正如苏神文章中所述，我们可以将让这两个矩阵变成可学习的参数，比如在“复制”过程中给每个头一个不同的权重，这样理论上可以增强 GQA 的能力。

这就是 MLA 的思想，它可以看作是 GQA 的一种改进，在压缩到 $\mathbf{c}$ 维之后又用一个上投影矩阵来恢复到更高的维度，在 DeepSeek-V2 的技术报告中是先利用下投影矩阵 $W^{DKV}$ 将隐藏层输入 ${\mathbf{h}}_t$ 投影得到 ${\mathbf{c}}_t^{KV}$，然后再用两个上投影矩阵 $W^{UK},W^{UV}$ 将 ${\mathbf{c}}_t^{KV}$ 还原得到 ${\mathbf{k}}_t^C,{\mathbf{v}}_t^C$

然而，理论上这样是能增加模型能力，但别忘了 GQA 的主要目的是减少 KV Cache，出于节省计算和通信成本的考虑，我们一般缓存的是投影后的 ${\mathbf{k}}_{\mathbf{i}},{\mathbf{v}}_{\mathbf{i}}$ 而不是投影前的 ${\mathbf{c}}_{\mathbf{i}}$ 或 ${\mathbf{x}}_{\mathbf{i}}$，而 MLA 的这个做法，通过不同的投影矩阵再次让所有的 K、V Head 都变得各不相同，那么 KV Cache 的大小就恢复成跟 MHA 一样大了，违背了 GQA 的初衷。

对此，MLA 发现，我们可以结合 Dot-Attention 的具体形式，通过一个简单但不失巧妙的恒等变换来规避这个问题。首先，在训练阶段还是照常进行，此时优化空间不大；然后，在推理阶段，我们利用：

$${\mathbf{q}}_{\mathbf{t}}^{(\mathbf{s})}{\mathbf{k}}_{\mathbf{i}}^{(\mathbf{s})\top }=\left({\mathbf{x}}_{\mathbf{t}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}\right)\left({\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}\right){}^{\top }={\mathbf{x}}_{\mathbf{t}}\left({\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})\top }\right){\mathbf{c}}_{\mathbf{i}}^{\top }$$

这意味着推理阶段，我们可以将 ${\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})\top }$ 合并起来作为 Q 的投影矩阵，那么 ${\mathbf{c}}_{\mathbf{i}}$ 则取代了原本的 ${\mathbf{k}}_{\mathbf{i}}$

同理，在 ${\mathbf{o}}_{\mathbf{t}}$ 后面我们还有一个投影矩阵，于是 ${\mathbf{v}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}$ 的 ${\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}$ 也可以吸收到后面的投影矩阵中去，于是等效地 ${\mathbf{v}}_{\mathbf{i}}$ 也可以用 ${\mathbf{c}}_{\mathbf{i}}$ 代替，也就是说此时 KV Cache 只需要存下所有的 ${\mathbf{c}}_{\mathbf{i}}$ 就行，而不至于存下所有的 ${\mathbf{k}}_{\mathbf{i}}^{(\mathbf{s})}$、${\mathbf{v}}_{\mathbf{i}}^{(\mathbf{s})}$。注意到 ${\mathbf{c}}_{\mathbf{i}}$ 跟 ${}^{(s)}$ 无关，也就是说所有的头共享的，即 MLA 在推理阶段它可以恒等变换为一个 MQA

再次强调，我们的主题一直都是减少 KV Cache，那到目前为止，MLA 做到了什么呢？答案是通过不同的投影矩阵来增强 GQA 的能力，并且推理时可以保持同样大小的 KV Cache。那么反过来，如果我们只需要跟 GQA 相近的能力，那么是不是就可以再次减少 KV Cache 了？换言之，$d_c$ 没必要取 $g(d_k+d_v)$，而是取更小的值（DeepSeek-V2 取了 512），从而进一步压缩 KV Cache，这就是 MLA 的核心思想

1.5.2 Part 2

一切似乎都很完美，但到目前为止的 MLA 有一个难以绕开的缺陷—不兼容 RoPE（旋转位置编码）

关于 RoPE 大家感兴趣的可以看看：RoPE旋转位置编码原理浅析

前面我们说了，MLA 之所以能保持跟 GQA 一样大小的 KV Cache，其关键一步是将 ${\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})\top }$ 合并成一个（跟位置无关的）矩阵作为 Q 的投影矩阵，但如果加了 RoPE 的话，这一步就无法实现了。这是因为 RoPE 是一个跟位置相关的、$d_k\times d_k$ 的分块对角矩阵 ${\mathcal{R}}_{\mathbf{m}}$，满足 ${\mathcal{R}}_{\mathbf{m}}{{\mathcal{R}}_{\mathbf{n}}}^{\top }={\mathcal{R}}_{\mathbf{m}-\mathbf{n}}$，MLA 加入 RoPE 之后会让 ${\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})\top }$ 之间多插入了一项 ${\mathcal{R}}_{\mathbf{t}-\mathbf{i}}$：

$$\begin{gathered} {\mathbf{q}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{q}}^{(\mathbf{s})}{\mathcal{R}}_{\mathbf{i}},{\mathbf{k}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{k}}^{(\mathbf{s})}{\mathcal{R}}_{\mathbf{i}} \\ {\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }=\left({\mathbf{x}}_t{\mathbf{W}}_q^{(s)}{\mathcal{R}}_{\mathbf{i}}\right)\left({\mathbf{c}}_i{\mathbf{W}}_k^{(s)}{\mathcal{R}}_{\mathbf{i}}\right){}^{\top }={\mathbf{x}}_t\left({\mathbf{W}}_q^{(s)}{\mathcal{R}}_{\mathbf{t}-\mathbf{i}}{\mathbf{W}}_k^{(s)}{}^{\top }\right){\mathbf{c}}_i^{\top } \end{gathered}$$

这里的 ${\mathbf{W}}_q^{(s)}{\mathcal{R}}_{\mathbf{t}-\mathbf{i}}{\mathbf{W}}_k^{(s)}{}^{\top }$ 就无法合并为一个固定的投影矩阵了（跟位置差 $t-i$ 相关），从而 MLA 的想法无法结合 RoPE 实现

最后发布的 MLA 通过将每个 Attention Head 的 Q、K 新增 $d_r$ 个维度用来添加 RoPE，其中 K 新增的维度每个 Head 共享：

$$\begin{gathered} {\mathbf{o}}_{\mathbf{t}}=\left[{\mathbf{o}}_{\mathbf{t}}^{(1)},{\mathbf{o}}_{\mathbf{t}}^{(2)},\cdots ,{\mathbf{o}}_{\mathbf{t}}^{(\mathbf{h})}\right] \\ {\mathbf{o}}_t^{(s)}=\mathbf{Attention}\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)} \\ {\mathbf{q}}_i^{(s)}=\left[{\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{qc}^{(s)},{\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{qr}^{(s)}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d\times d_k}{,\mathbf{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d\times d_r} \\ {\mathbf{k}}_i^{(s)}=\left[{\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{kc}^{(s)},{\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k}{,\mathbf{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r} \\ {\mathbf{v}}_{\mathbf{i}}^{(\mathbf{s})}={\mathbf{c}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{v}}},{\mathbf{W}}_{\mathbf{v}}^{(\mathbf{s})}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{c}}\times {\mathbf{d}}_{\mathbf{v}}} \\ {\mathbf{c}}_{\mathbf{i}}={\mathbf{x}}_{\mathbf{i}}{\mathbf{W}}_{\mathbf{c}}\in {\mathbb{R}}^{{\mathbf{d}}_{\mathbf{c}}},{\mathbf{W}}_{\mathbf{c}}\in {\mathbb{R}}^{\mathbf{d}\times {\mathbf{d}}_{\mathbf{c}}} \end{gathered}$$

这样一来，没有 RoPE 的维度就可以重复 “Part 1” 的操作，在推理时 KV Cache 只需要存 ${\mathbf{c}}_{\mathbf{i}}$，新增的带 RoPE 的维度就可以用来补充位置信息，并且由于所有 Head 共享，所以也就只有在 K Cache 这里增加了 $d_r$ 个维度，原论文取了 $d_r=d_k/2=64$，相比原本的 $d_c=512$，增加的幅度不大

1.5.3 Part 3

最后有一个细节，就是 MLA 的最终版本，还将 Q 的输入也改为了低秩投影形式，这与减少 KV Cache 无关，主要是为了减少训练期间参数量和相应的梯度所占的显存：

$$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)},{\mathbf{o}}_t^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}=\left[{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qc}^{(s)},{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qr}^{(s)}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }},{\mathbf{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_r}\\ {\mathbf{k}}_i^{(s)}=\left[{\mathbf{c}}_i{\mathbf{W}}_{kc}^{(s)},{\mathbf{x}}_i{\mathbf{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c},{\mathbf{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathbf{v}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathbf{c}}_i^{\prime }={\mathbf{x}}_i{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d_c^{\prime }},{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d\times d_c^{\prime }}\\ {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{d_c},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$$

注意 ${\mathbf{k}}_i^{(s)}$ 中的第二项，带 RoPE 的部分，其输入还是 ${\mathbf{x}}_i$ 而不是 ${\mathbf{c}}_i$，这里保持了原论文的设置，不是笔误，$d_c^{\prime }$ 原论文的取值是 1536，跟 $d_c=512$ 不同。

同时，我们把带 RoPE 的 MHA 放在下面，方便大家对比：

$$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)},{\mathbf{o}}_t^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}={\mathbf{x}}_i{\mathbf{W}}_q^{(s)}{\mathcal{R}}_i\in {\mathbb{R}}^{d_k},{\mathbf{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathbf{k}}_i^{(s)}={\mathbf{x}}_i{\mathbf{W}}_k^{(s)}{\mathcal{R}}_i\in {\mathbb{R}}^{d_k},{\mathbf{W}}_k^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathbf{v}}_i^{(s)}={\mathbf{x}}_i{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d\times d_v}\end{array}$$

可以发现，其实在训练阶段，除了多了一步低秩投影以及只在部分维度加 RoPE 外，MLA 与 Q、K 的 Head Size 由 $d_k$ 换成 $d_k+d_r$ 的 MHA 基本无异

推理阶段的 MLA 则改为：

$$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)}{\mathbf{W}}_v^{(1)},{\mathbf{o}}_t^{(2)}{\mathbf{W}}_v^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}{\mathbf{W}}_v^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{\bcancel{(s)}},{\mathbf{c}}_{\le t}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{\bcancel{(s)}}{}^{\top }\right){\mathbf{c}}_i}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{\bcancel{(s)}}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}=\left[{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qc}^{(s)}{\mathbf{W}}_{kc}^{(s)}{}^{\top },{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qr}^{(s)}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_c+d_r}\\ {\mathbf{k}}_i^{\bcancel{(s)}}=\left[{\mathbf{c}}_i,{\mathbf{x}}_i{\mathbf{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_c+d_r}\\ {\mathbf{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_k},{\mathbf{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathbf{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_r},{\mathbf{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathbf{c}}_i^{\prime }={\mathbf{x}}_i{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d_c^{\prime }},{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d\times d_c^{\prime }}\\ {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{d_c},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$$

此时 Q、K 的 Head Size 变成了 $d_c+d_r$，V 的 Head Size 则变成了 $d_c$，按照原论文的设置，这是 $d_k$、$d_v$ 的 4 倍。所以实际上 MLA 在推理阶段做的这个转换，虽然能有效减少KV Cache，但其推理的计算量是增加的。

那为什么还能提高推理效率呢？这又回到“瓶颈”一节所讨论的问题了，我们可以将LLM的推理分两部分：第一个 Token 的生成（Prefill）和后续每个 Token 的生成（Generation），Prefill 阶段涉及到对输入所有 Token 的并行计算，然后把对应的 KV Cache 存下来，这部分对于计算、带宽和显存都是瓶颈，MLA 虽然增大了计算量，但 KV Cache 的减少也降低了显存和带宽的压力，大家半斤八两；但是 Generation 阶段由于每步只计算一个 Token，实际上它更多的是带宽瓶颈和显存瓶颈，因此 MLA 的引入理论上能明显提高 Generation 的速度。

还有一个细节充分体现了这个特性。一般的 LLM 架构参数满足 $h=d$，即 num_heads * head_size = hidden_size，但 DeepSeek-V2 不一样，它 $d_k=128,d=5120$，但 $h=128$，是一般设置的 3 倍！这是因为 MLA 的 KV Cache 大小跟 $h$ 无关，增大 $h$ 只会增加计算量和提升模型能力，但不会增加 KV Cache，所以不会带来速度瓶颈。

由于篇幅原因（CSDN 对正文字数有限制，真服了😒），MLA 的代码实现我们放在下篇文章

结语

MLA 可以看作是 GQA 的优化，通过投影矩阵的方式替换 GQA 中的分割、复制等线性变换操作，并引入了一个恒等变换在推理阶段通过矩阵吸收来进一步压缩 KV Cache，同时采用了一种混合方法通过新增维度来兼容 RoPE 旋转位置编码，总的来说，MLA 算得上一种非常实用的注意力变体

大家可以多看看苏神的文章来加深理解🤗

参考

• 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA
• 博客分享：从MHA、MQA、GQA到MLA
• https://github.com/preacher-1/MLA_tutorial
• https://github.com/deepseek-ai/DeepSeek-V3