Published2025-10-13
Topic推理系统 · 笔记
Size1265 chars · 3 images

LLM —— attention、MHA、MQA、GQA、MLA

LLM —— attention、MHA、MQA、GQA、MLA

文章启蒙来自苏神的缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA，拜谢Orz。

attention

首先attention的公式我们都知道如下：

$$ \mathrm{Attention}(K,Q,V)=\mathrm{softmax}(\frac{QK^{\top}}{\sqrt{d_{k}}})V $$

我们尝试用数学来表述清楚这个问题，首先假设输入为一条token的$x$，输出为一条token的$o$，结合上面的公式得到QKV矩阵：

$$ q=x W_q\in\mathbb{R}^{d_k},W_q\in\mathbb{R}^{d\times d_k} $$
$$ k=x W_{k}\in\mathbb R^{d_k},W_{k}\in\mathbb R^{d\times d_k} $$
$$ v=xW_{v}\in\mathbb{R}^{d_k},W_{v}\in\mathbb{R}^{d\times d_v} $$

其中分母是缩放因子，通常可以人为设定，这里我们后面再提

MHA

那么对于多头注意力机制，也就是attention all you need的原文中提到的，是attention的一种并发形式，我们假设输入为：

$$ x_{1},x_{2},x_{3},\ldots,x_{l}.~~~~~~x_{i}\in\mathbb{R}^{d} $$

假设有h个注意力头组成一个多头注意力，那么在t时刻的输出为：

$$ o_{t}=[o_{t}^{(1)},o_{t}^{(2)},o_{t}^{(3)},\ldots,o_{t}^{(h)}] $$

那么在第s个头内，进行的计算有：

$$ q_{i}^{(s)}=x_{i}W_{q}^{(s)}\in R^{d_k},~~~~~W_{q}^{(s)}\in R^{d\times d_k} $$
$$ k_{i}^{(s)}=x_{i}W_{k}^{(s)}\in R^{d_k},~~~~~W_{k}^{(s)}\in R^{d\times d_k} $$
$$ v_{i}^{(s)}=x_{i}W_{v}^{(s)}\in R^{d_k},~~~~~W_{v}^{(s)}\in R^{d\times d_v} $$

其架构如下：

上面QKV切分之后的size都是$[seq, d_i]$，那么$QK^T$就是[seq, seq]，完整QKV经过softmax之后的size为$[seq, d_i]$，也就是说每个$head_i$产出的size都是$[seq, d_i]$，拼起来就是[seq, seq] ，最后通过一个合并头将所有头的输出合并回一个大的张量。

这里需要注意的是，在MoE架构中，多头模型的权重太大了，进行TP的时候会把权重矩阵进行切分，分配到不同的GPU上，具体可以看vllm v1 源码解析 —— 单机八卡推理

还需要注意的是，这里的qkv在存储上并非对称，详情请看为什么加速LLM推断有KV Cache而没有Q Cache？

综上所述，得出的attention如下式：

$$ \boldsymbol{o}_t^{(s)}=Attention\left(\boldsymbol{q}_t^{(s)},\boldsymbol{k}_{\leq t}^{(s)},\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)}\boldsymbol{k}_i^{(s)\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)}\boldsymbol{k}_i^{(s)\top}\right)} $$

上面公式省略了原公式分母的缩放因子，实践中常见的缩放因子设计为$d_{k}=d_{v}=d/h$。LLAMA2-7b有$d=4096,h=32,d_{k}=d_{v}=128 $，而deepseek v3.2的$h = 128，d=7168，d_{k}=56$。

MQA和GQA

MQA

MQA就是舍弃了多个KV，只用一份KV。

GQA

GQA没有那么极端，只是分成了G组，组内共享同一个KV。

GQA提供了MHA到MQA的自然过渡，当$g=h$时就是MHA，$g=1$时就是MQA，当$1<g<h$时，它只将KV Cache压缩到$g/h$，压缩率不如MQA，但同时也提供了更大的自由度，效果上更有保证。GQA最知名的使用者，大概是Meta开源的LLAMA2-70B，以及LLAMA3全系列。（苏神文章原话）

在llama2/3-70B中，GQA的$g=8$，其他用了GQA的同体量模型基本上也保持了这个设置，这并非偶然，而是同样出于推理效率的考虑。我们知道，70B这个体量的模型，如果不进行极端的量化，那么不可能部署到单卡（A100/H100 80G）上。单卡不行，那么就能单机了，一般情况下一台机可以装8张卡，刚才我们说了，Attention的每个Head实际上是独立运算然后拼接起来的，当$g=8$时，正好可以每张卡负责计算一组K、V对应的Attention Head，这样可以在尽可能保证K、V多样性的同时最大程度上减少卡间通信。（苏神文章原话）

MLA Multi-head Latent Attention

参考文献：

DeepSeek-V3.2-Exp.json