MLA 结构

• 需要缓存
  • KV 向量共用的压缩隐特征
  • K 向量多头共享的带位置编码的向量
• 为什么带有位置信息的 Q 向量来自于隐特征向量，而带有位置的 K 向量来自于 H 向量且共享呢？

最好的方法肯定是从H向量直接计算并且不共享，但是会大大增加显存使用和降低计算效率。原理上，基于隐向量计算ROPE肯定是有损的，共享也肯定牺牲了表达能力，所以做了一些权衡：

1、Q向量都基于潜向量生成RoPE向量而不共享，主要是为了增加计算效率。因为隐向量小所以计算更快，而且每次都要计算。不共享是为了保证表达能力。

2、K向量是从缓存中取的，不用每次计算，所以直接在H中计算就好。但是如果不共享将会让每个头都有一个RoPE向量，大大增加显存占用，所以共享。

---

多头注意力机制 MHA + KV cache

在生成第三个 token 的时候，第一个 token 进行的计算已经在生成第二个 token 的时候计算过了，重复计算。–》缓存第一个 token 计算的中间变量，并且只保留生成新 token 所需要的中间变量（KV cache）

有了 KV cache 后生成第三个 token 的过程

生成第四个 token

---

GQA/MQA

这里展示的是 MQA，生成 3 个 head 的 Q 向量，只生成 1 个 head 维度的 K 和 V 向量

多头间通过复制共享 query 向量一起来计算注意力，从而减少 kv cache，但会大大影响性能

---

为了折中，提出了 GQA，每组 query 共享一个 k 和 v 向量

---

MLA-Multi-Head Latent Attention

多头潜在注意力机制

• 目的：减少 kv cache + 尽量不影响性能或者提高性能
• 原理：对 token 的特征向量进行压缩转换，缓存压缩后的向量，在计算 attention 之后再解压回原来的尺寸

• 可以提效果，很不错

压缩 KV 向量

kv cache 本意是为了减少推理时对之前 token 的 k 和 v 向量的计算

MLA 因为缓存了压缩的 kv cache，而减小了 kv cache 的显存占用，但是在取出缓存后，k 和 v 不能直接使用，需要经过解压计算才可以，引入了额外的计算，与 kv cache 初衷相悖

• 对 k 进行解压操作的矩阵可以和 Wuq 矩阵进行融合，这个融合可以在推理之前算好，这样在推理时就不用进行对 k 的额外解压计算了【利用矩阵相乘的结合律，对矩阵提前进行融合，从而规避 MLA 引入的因解压隐特征带来的额外计算】

Wuv 同理，可以和 Wo 融合

---

压缩 Q 向量

除了对 KV 向量进行压缩外，对 Q 向量也进行了压缩，好处是降低了参数量，而且可以提高模型性能

---

考虑 RoPE

RoPE 需要对每一层的 Q 向量和 K 向量进行旋转，而且根据 token 位置的不同，旋转矩阵的参数也是不同的。加入了 RoPE 的矩阵无法融合，因为中间两个矩阵与 token 位置相关。

• 解决方案：为 Q 和 K 向量额外增加一些维度来表示位置信息

对于 Q 向量，通过 WQR 为每一个头生成一些原始特征，然后通过 RoPE 增加位置信息，再把生成带有位置信息的特征拼接到每个注意力头的 Q 向量

↓拼接

对于 K 向量，通过 WKR 矩阵生成一个头共享的特征，然后通过 RoPE 增加位置信息，然后复制到多个头共享位置信息。**这里多头共享带位置编码的 K 向量，也需要被缓存，**以便在生成带有位置信息的 K 向量时用到

---

在推理时

• 不带 RoPE 的 Q 向量和 K 向量进行点积运算（结果为数值），可以用融合的矩阵来消除解压操作
• 带 RoPE 的部分进行点积运算

将两部分得到的两个值进行逐元素相加：⊕ ，就相当于对拼接了位置信息的完整的 Q 和 K 向量进行点积操作的值。

---

参考

1. https://www.bilibili.com/video/BV1BYXRYWEMj
2. https://arxiv.org/pdf/2412.19437