多种Attention对比

假设hidden_state为4096维

h(1*4096)变成Q(1*4096)之后会被切成Q(32*128)，然后每个Q后续只需要和自己匹配的KV进行运算，优点是我们能够对Q的每个地方特化expression。这里自己匹配的KV也可以分成很多份。共享KV会有损失。

结构	Attention 表达能力	单个token KV Cache
MHA	32 套权重，KV也独立	321282
GQA	32 套权重，部分共享KV （例如共享四套KV Cache）	41282
MQA	32 套权重，完全共享KV	11282
MLA	32 套权重，完全共享KV，且KV可以由同一个L变出来	1*128
GQLA	32 套权重，部分共享KV（例如共享四套KV Cache），且KV可以由同一个L变出来	4*128
不切Q	只有1套权重	4096*2

MLA

MLA是把K和V都变成latent,然后用不同的矩阵还原出K和V。

C: [L, 512]
W_UK: [512, 32×128]
K: [L, 32, 128]

$K_i=c^{KV}W_{UK}^{(i)}$

$V_i=c^{KV}W_{UV}^{(i)}$

原本 QK 是：

$Q_iK_i^\top = Q_i(c^{KV}W_{UK}^{(i)})^\top$

利用结合律改写为：

$(Q_iW_{UK}^{(i)\top})(c^{KV})^\top$

Decode 时会重复展开历史 KV

假设目前已有4096个 token。

第一个 decode token：

C[0:4096] @ WUK → 完整 K
C[0:4096] @ WUV → 完整 V

下一个 decode token：

C[0:4097] @ WUK → 又展开一次
C[0:4097] @ WUV → 又展开一次

之前4096个 token 的 K/V 没有发生任何变化，却被重复计算。利用K projection可以让KV一直以c的方式保存着。