RoPE 旋转位置编码原理与实现详解

从二维旋转直觉到高维推广，详解 RoPE 的数学原理、NeoX/GPT-J 两种实现风格、Triton GPU Kernel 优化，以及 Linear Scaling / NTK / YaRN 三种长度外推策略。

1. 为什么需要位置编码

Transformer 的 Self-Attention 本身是置换不变的——打乱输入顺序，输出也跟着打乱但数值不变。模型无法区分 "猫吃鱼" 和 "鱼吃猫"。位置编码就是给每个 token 注入位置信息。

位置编码的演进

方案	代表模型	特点
绝对正弦编码	原始 Transformer	固定函数，不可学习
可学习绝对编码	BERT, GPT-2	受限于训练长度
相对位置编码	T5 (bias)	Attention 中加偏置
RoPE	LLaMA, Qwen, GPT-NeoX	旋转Q/K，天然相对位置
ALiBi	BLOOM	线性衰减 bias

RoPE 是当前主流 LLM 的标配，被 LLaMA、Qwen、Mistral、DeepSeek 等广泛采用。

2. 数学原理

2.1 核心思想：用旋转编码位置

对 Attention 中的 Q 和 K 向量施加与位置相关的旋转，使得 Q·K 的内积自然包含相对位置信息。

2.2 二维旋转

给定二维向量 $[x_1, x_2]$ 和旋转角度 $\theta$：

$$R(\theta) \cdot \begin{bmatrix} x_1 \ x_2 \end{bmatrix} = \begin{bmatrix} x_1 \cos\theta - x_2 \sin\theta \ x_1 \sin\theta + x_2 \cos\theta \end{bmatrix}$$

关键性质：两个旋转后向量的内积只取决于旋转角度之差：

$$\langle R(m\theta) \cdot q, ; R(n\theta) \cdot k \rangle = \langle R((m-n)\theta) \cdot q, ; k \rangle$$

这意味着 Attention score 自动编码了相对位置 $(m - n)$，而非绝对位置。

2.3 推广到高维

head_dim 为 $d$ 时，将其分成 $d/2$ 个维度对，每对使用不同频率的旋转：

$$\theta_i = \frac{1}{\text{base}^{2i/d}}, \quad i = 0, 1, \ldots, d/2 - 1$$

其中 $\text{base} = 10000$（默认值）。

对位置 $m$ 的完整旋转：

$$\text{RoPE}(x, m) = \begin{bmatrix} x_1 \cos(m\theta_0) - x_2 \sin(m\theta_0) \ x_2 \cos(m\theta_0) + x_1 \sin(m\theta_0) \ x_3 \cos(m\theta_1) - x_4 \sin(m\theta_1) \ x_4 \cos(m\theta_1) + x_3 \sin(m\theta_1) \ \vdots \end{bmatrix}$$

2.4 频率的直觉

维度对 i=0:  θ = 1.0      → 波长 2π ≈ 6     → 变化快，区分相邻位置
维度对 i=1:  θ = 0.01     → 波长 ≈ 628      → 变化中等
维度对 i=N:  θ = 0.0001   → 波长 ≈ 62832    → 变化慢，编码长距离

低维度对 → 高频 → 局部位置信息（区分相邻 token） 高维度对 → 低频 → 全局位置信息（区分远距离 token）

3. 工程实现

3.1 预计算 cos/sin 缓存

位置和频率是确定的，可以提前算好所有 $\cos(m \cdot \theta_i)$ 和 $\sin(m \cdot \theta_i)$：

inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))  # [dim/2]
t = torch.arange(max_position)                                 # [max_pos]
freqs = torch.einsum("i,j->ij", t, inv_freq)                  # [max_pos, dim/2]
cos_cache = freqs.cos()
sin_cache = freqs.sin()

对应源码：nano-vllm/nanovllm/layers/rotary_embedding.py 的 __init__

3.2 两种旋转风格

NeoX-style（LLaMA、Qwen 使用）：

将 head_dim 对半切分：

x1, x2 = x.chunk(2, dim=-1)        # 前半 & 后半
y1 = x1 * cos - x2 * sin
y2 = x2 * cos + x1 * sin
y = cat([y1, y2], dim=-1)

GPT-J-style（GPT-J、GPT-NeoX 使用）：

交错取维度对：

x1 = x[..., 0::2]                  # 偶数维度
x2 = x[..., 1::2]                  # 奇数维度
y1 = x1 * cos - x2 * sin
y2 = x1 * sin + x2 * cos
y = stack([y1, y2], dim=-1).flatten(-2)

两者数学等价，只是维度排列不同。NeoX-style 更常见，也对 GPU 更友好（连续内存访问）。

3.3 完整流程

推理时:
1. 预计算 cos/sin 缓存 (一次性)
2. 对每个 token 的位置，查表获取 cos[pos], sin[pos]
3. 对 Q 和 K 分别应用旋转（V 不旋转！）
4. 用旋转后的 Q, K 计算 Attention

4. Triton GPU Kernel 实现

4.1 为什么需要 CUDA/Triton Kernel

PyTorch 实现的问题：

chunk()  → 创建临时张量
x1*cos   → 创建临时张量
x2*sin   → 创建临时张量
cat()    → 又一次创建 + 拷贝
= 5+ 次 kernel launch, 3+ 次临时显存分配

Triton Kernel：1 次 kernel launch, 0 次临时分配。

4.2 并行策略

grid = (num_tokens * num_heads,)

每个 Triton program 处理一个 (token, head) 对:
  1. 从 positions 读取当前 token 的位置 pos
  2. 从 cos/sin 缓存加载 cos[pos], sin[pos]
  3. 从 x 加载 x1, x2 (head_dim 个元素)
  4. 计算 y1 = x1*cos - x2*sin, y2 = x2*cos + x1*sin
  5. 写回 output

head_dim 通常为 64~256，一个 program 内可以用向量化一次处理完。

4.3 FMA 优化

SGLang 的 Triton kernel 使用 FMA（Fused Multiply-Add）：

# 普通写法: 两次运算，两次舍入
y1 = x1 * cos - x2 * sin

# FMA: 一条指令，一次舍入，更精确
y1 = tl.fma(-x2, sin, x1 * cos)   # = -x2*sin + x1*cos

4.4 性能对比

场景	PyTorch	Triton	加速比
256 tokens, dim=128	基准	~1.5-2x	小 batch kernel launch 开销显著
2048 tokens, dim=128	基准	~2-3x	中 batch 收益明显
8192 tokens, dim=128	基准	~2-4x	大 batch 显存节省显著

主要收益来自消除临时张量分配和减少 kernel launch 次数。

5. 长度外推策略

RoPE 的一个挑战：模型在长度 $L$ 上训练，推理时遇到 $>L$ 的序列怎么办？

5.1 Linear Scaling

最简单的方法——等比压缩位置：

$$m' = m / s$$

训练: 位置 0, 1, 2, ..., 4095
推理: 位置 0, 0.25, 0.5, ..., 4095 (4x 外推)

问题: 所有频率都被压缩，局部分辨率下降。

对应源码：vllm/.../linear_scaling_rope.py

5.2 NTK-aware Scaling

修改 base 而非位置：

$$\text{base'} = \text{base} \times s^{d/(d-2)}$$

原始: θᵢ = 1/(10000^(2i/d))
NTK:  θᵢ = 1/(new_base^(2i/d))  ← new_base 更大

优势: 高频维度变化小（保留局部区分能力），低频维度拉伸大（扩展范围）。

对应源码：vllm/.../dynamic_ntk_scaling_rope.py

5.3 YaRN Scaling

最先进的方案——分频段处理 + 幅度修正：

高频维度 → 保持不变（外推）     不压缩，保留局部精度
中频维度 → 平滑过渡             blend(插值, 外推)
低频维度 → 等比压缩（插值）     ÷ scaling_factor
+ 幅度修正: mscale = 0.1 * log(s) + 1.0

为什么需要幅度修正？ 外推时 cos/sin 值的统计分布变了，注意力分数的量级会偏移，mscale 修正这个偏移。

对应源码：vllm/.../yarn_scaling_rope.py

5.4 三种策略对比

频率维度:  高频 ─────────────────────────── 低频
           (局部位置)                    (全局位置)

Linear:    压缩 ──────────────────────── 压缩    (全部等比压缩)
NTK:       微调 ──────────────────────── 大幅调   (只改 base)
YaRN:      不变 ── 平滑过渡 ──────────── 压缩    (分段处理)

6. 高级变体

6.1 MRoPE（多维 RoPE）

Qwen2-VL 使用，将 head_dim 分成 3 段，分别编码 时间/高度/宽度 位置：

head_dim = [T 段 | H 段 | W 段]
T 段用时间位置的 cos/sin
H 段用图像高度位置的 cos/sin
W 段用图像宽度位置的 cos/sin

对应源码：vllm/.../mrope.py

6.2 FoPE（傅里叶 RoPE）

可学习的投影矩阵替代固定频率：

sin = einsum("tD, Dd -> td", pos_sin, sin_coef)  # 可学习投影
cos = einsum("tD, Dd -> td", pos_cos, cos_coef)

对应源码：vllm/.../fope.py

7. RoPE 的重要性质

性质	说明
相对位置编码	Q·K 内积只依赖 $(m-n)$，非绝对位置
模长不变	旋转是正交变换，$|Rx| = |x|$
远程衰减	相对距离越大，期望内积越小
无额外参数	不需要学习，完全由公式确定
可外推	配合缩放策略可推广到训练长度之外

8. 与 vLLM/SGLang 源码的对应

组件	源码位置
基础 RoPE	vllm/.../rotary_embedding/base.py
Linear Scaling	vllm/.../rotary_embedding/linear_scaling_rope.py
NTK Scaling	vllm/.../rotary_embedding/dynamic_ntk_scaling_rope.py
YaRN	vllm/.../rotary_embedding/yarn_scaling_rope.py
DeepSeek 扩展	vllm/.../rotary_embedding/deepseek_scaling_rope.py
Llama3 RoPE	vllm/.../rotary_embedding/llama3_rope.py
MRoPE (多维)	vllm/.../rotary_embedding/mrope.py
FoPE (傅里叶)	vllm/.../rotary_embedding/fope.py
CUDA 融合 kernel	vllm/csrc/fused_qknorm_rope_kernel.cu
SGLang Triton kernel	sglang/.../jit_kernel/diffusion/triton/rotary.py
nano-vllm 简化实现	nano-vllm/nanovllm/layers/rotary_embedding.py

9. 教学代码

配套代码在 vllm_learn/code/course_rope/：

• rope_torch.py：PyTorch 实现，从 2D 旋转直觉到完整模块，含 3 种缩放策略
• rope_triton.py：Triton GPU Kernel 实现，NeoX/GPT-J 双风格，含性能对比