从 Attention 到 RoPE 与 GQA

解析 RoPE 与 GQA 如何解决顺序建模与 KV Cache 显存压力问题。

预计阅读 7 分钟
  • AI
  • 大模型

从 Attention 到 RoPE 与 GQA：现代大模型核心结构全景解析

关键词：Self-Attention，RoPE，KV Cache，GQA，长上下文，推理优化 适合人群：已经理解 Transformer 基础公式，希望真正掌握大模型底层结构的人

摘要

Transformer 的核心是 Attention 机制，但“原始 Attention”在真实工程环境中存在两个致命问题：

  1. 不具备顺序建模能力（Permutation Invariance）
  2. 推理阶段 KV Cache 显存线性爆炸

现代大模型（如 LLaMA、Qwen）通过两项关键技术完成了结构进化：

  • RoPE（Rotary Positional Embedding） —— 用旋转将相对位置信息嵌入点积计算
  • GQA（Grouped Query Attention） —— 通过分组共享 KV 显著降低显存消耗

本文将从数学本质、几何直觉、工程实现三个层面完整推导其工作机制。

第一部分：Attention 的结构性缺陷

1.1 标准公式回顾

[ Attention(Q,K,V) = Softmax\left(\frac{QK^T}{\sqrt{d}}\right)V ]

其中：

  • Q：Query
  • K：Key
  • V：Value
  • d：Head 维度

核心计算是： [ QK^T ] 即向量内积。

1.2 缺陷一：不具备顺序感

内积的性质是什么？

  • 只与向量内容相关（方向、长度）
  • 与“这个向量来自句子的第几个位置”无关

所以如果你不给它位置，它看到的就是一堆 token embedding 的集合。 这就是常说的：Attention 对输入是置换不变的（Permutation Invariant）。

不是说 Transformer 完全不能学顺序，而是：Attention 这一层算子本身不携带顺序信息。 顺序必须由你“外加”。

若两个 token embedding 相同，则：

  • “我 爱 你”
  • “你 爱 我”

在数学上仅是向量集合不同排列。

Attention 对输入顺序是 置换不变的。

结论：

原始 Attention 不知道 token 的位置。

必须人为注入位置信息。

1.3 缺陷二：KV Cache 显存爆炸

在自回归生成中：

第 t 步只需计算： [ Q_t K_{1:t}^T ] 但必须缓存： [ K_{1:t}, V_{1:t} ] KV Cache 形状：

[Batch, SeqLen, NumHeads, HeadDim]

显存占用近似为： [ Memory \propto SeqLen \times NumHeads \times HeadDim ] 当：

  • SeqLen = 128K
  • NumHeads = 32

KV Cache 成为主要显存开销。

第二部分：RoPE —— 用旋转编码相对位置,把“相对距离”塞进点积里（而不是塞进向量里）

def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
        # 自动检测设备（是 CPU 还是某张显卡）
        if device is None:
            device = self.transformer.wte.weight.device
            
        # 1. 遍历通道 (channels) 算出不同的旋转频率
        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (base ** (channel_range / head_dim))
        
        # 2. 遍历时间步 (位置)
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        
        # 3. 计算每个位置、每个频率对应的旋转角度
        freqs = torch.outer(t, inv_freq)
        cos, sin = freqs.cos(), freqs.sin()
        
        # 4. 转换精度并调整形状，为了后续计算方便
        cos, sin = cos.bfloat16(), sin.bfloat16() 
        cos, sin = cos[None, :, None, :], sin[None, :, None, :] 
        return cos, sin

首先我们将每个head 的输入维度进行两两分组，因为只有二维的平面才能方便旋转。这样我们就得到了得到了 32 个二维平面（或 32 个频率通道）（如果head_dim是64的话）。这32位的channel其实就是32个平面，接下来我们计算每一组平面的旋转角频率，根据channel的值和head_dim去计算，这样channel值小的时候角频率大一些（分母小），channel打的时候也就是后面的分组，角频率就小一些。这个时候角频率的shape就是（32，）。有了角频率 我们需要计算每个位置要旋转的角度，这个很类似计算角速度。就是根据输入的序列的绝对位置（0，1，2，3，4…）他的shape就是（sep_len，），与我们的角频率做外积,最后得到（seq_len,32）的矩阵J，这个矩阵J其实就是每一个位置所对应的向量（我们之前提到的32位的channel）每一个chgannel应该旋转的角度。也就是第 pos 个 token 在第 i 个平面上应该旋转多少角。为了方便计算旋转后的值，我们直接计算这个矩阵J的co s和sin，并将他们和我们的输入维度对齐 方便运算。

2.1 设计目标

经典绝对位置编码是： $$ x_m’ = x_m + p_m $$ 它能注入位置信息，但有两个天然问题：

  1. 相对位置不是显式结构：模型要自己从 “$p_m$” 和 “$p_n$” 的差异里学“距离感”
  2. 外推困难：训练没见过更长位置，推理长度一变容易崩

所以 RoPE 换了一个思路： 不要把位置当作“额外信息加进去”，而是把位置当作“作用在向量上的变换”。

理想的位置编码应该：

  • 表达相对距离
  • 不增加参数规模
  • 不破坏 Attention 结构

RoPE 满足以上三点。

2.2 二维几何直觉

二维旋转矩阵： [ R(\theta)= \begin{bmatrix} \cos\theta & -\sin\theta \ \sin\theta & \cos\theta \end{bmatrix} ] RoPE 规则：

位置 m → 旋转 mθ

位置 n → 旋转 nθ

也就是： $$ q_m’ = R(m\theta)q,\quad k_n’ = R(n\theta)k $$ 注意：真实模型不是 2D，而是把维度两两分组成很多个 2D 子空间分别旋转（这个后面讲）。

2.3 关键数学推导

设原始向量 q, k

旋转后： [ q’ = R(m\theta) q ]

[ k’ = R(n\theta) k ]

计算内积： [ q’^T k’ = q^T R(m\theta)^T R(n\theta) k ] 旋转矩阵性质： [ R(\alpha)^T = R(-\alpha) ] 因此： [ R(m\theta)^T R(n\theta) = R((n-m)\theta) ] 得到： [ q’^T k’ = q^T R((n-m)\theta) k ] 结论：

旋转后的 QK 点积 只依赖相对距离 $n-m$。 这等价于把“相对位置”直接嵌入了注意力打分函数里。

这正是我们希望得到的性质。

2.4 高维实现机制

真实模型维度为 d。

实现方式：

  • 每两维一组
  • 每组构成二维平面
  • 每个平面使用不同频率

角频率公式： [ \theta_i = 10000^{-2i/d} ] 低频 → 建模长距离 高频 → 建模短距离

本质上是一种多尺度频率编码。

2.5 复数形式理解（更优雅视角）

将二维向量视作复数： [ x+iy ] 旋转等价于： [ e^{im\theta} ] 内积变为： [ e^{i(m-n)\theta} ] 自然只依赖相对距离。

这说明 RoPE 本质是：

在复平面上进行相位编码。

第三部分：GQA —— 控制 KV Cache 规模

3.1 Multi-Head Attention 的问题

假设：

  • Query 头数 H = 32
  • 每个头都有独立 K、V

KV Cache 规模：

32 组 K + 32 组 V

显存随头数线性增长。

3.2 MQA：极端压缩

  • Query：32 头
  • KV：1 头

显存减少 32 倍

但表达能力下降明显。

原因：

所有 Query 在同一语义子空间检索。

3.3 GQA：分组共享

假设：

  • 32 个 Query 头
  • 4 个 KV 头

每 8 个 Query 共享 1 个 KV。

显存节省比例： [ \frac{4}{32} = 12.5% ] 但仍保留多子空间表达能力。

3.4 张量结构变化

MHA：

Q: [B,S,H,D] K: [B,S,H,D]

GQA：

Q: [B,S,H,D] K: [B,S,H_kv,D]

其中： [ H_{kv} < H ] 计算时通过 broadcast 对齐 head 维度。

实际实现中：

不会真正复制内存，而是共享 stride。

第四部分：RoPE + GQA 的系统意义

技术解决问题影响
RoPE无顺序建模能力引入相对位置
GQAKV Cache 爆炸显著降低显存

组合效果：

  • 支持长上下文
  • 推理成本可控
  • 几乎不损失性能

这构成现代大模型的基础结构。

第五部分：工程视角总结

RoPE 的价值：

  • 数学优雅
  • 参数零增加
  • 外推能力强

GQA 的价值：

  • 显存线性压缩
  • 推理速度提升
  • 可扩展到超长上下文

一句话理解：

RoPE 让点积拥有空间感。 GQA 让多头机制更节省。

当你真正理解这两者，现代 LLM 的底层结构就不再神秘。

留言讨论

0 条留言

正在加载留言...