Fused Rotary Position Encoding (RoPE) 全方位解析

背景 / Background

Rotary Position Embedding (RoPE) 由 Su et al. 在 2021 年的 RoFormer 论文中提出，已成为现代 LLM 的事实标准位置编码方案。LLaMA、Qwen、Baichuan、Mistral、Gemma、Phi、GLM 等主流模型均采用 RoPE 的某种变体。

与绝对位置编码 (APE) 直接加到 token embedding 不同，RoPE 在 attention 计算前对 Query 和 Key 向量施加位置相关的旋转变换。这种设计使得两个 token 的 attention score 只依赖它们的相对位置，天然支持位置外推。

RoPE 数学原理

对于 head 维度 d 中的第 i 对元素，RoPE 定义旋转角度为：

为什么有多种变体

RoPE 的数学原理只定义了"对 d/2 对元素各施加一个旋转"，但没有规定这些"对"如何从 d 维向量中选取。这就是不同 RoPE 变体的根本分歧点：

四大 RoPE 配对方式详解

1. GPT-J 模式 (Interleaved / 交错配对)

GPT-J 风格将相邻的两个元素视为一对：(arr[2i], arr[2i+1]) 构成第 i 个旋转对。这是 RoPE 原始论文中最自然的实现方式。

内存布局 (head_dim=8, rot_dim=8)

配对关系：(x₀,y₀), (x₁,y₁), (x₂,y₂), (x₃,y₃) — 每对共享同一个 θ

索引公式

对应 Kernel 代码 (apply_token_rotary_embedding_kernel, IS_NEOX=false)

// GPT-J: 相邻元素配对
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos = cos_ptr[x_index / 2];
sin = sin_ptr[x_index / 2];

const T x = arr[x_index];
const T y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;CUDA

代表模型

GPT-J, CodeGen, GPT-NeoX (early), 部分 Falcon 变体

2. GPT-NeoX 模式 (Half-Split / 前后半配对)

GPT-NeoX 风格将 head 维度切成前半和后半，前半的第 i 个元素与后半的第 i 个元素配对：(arr[i], arr[i + d/2])。

内存布局 (head_dim=8, rot_dim=8)

前半 [0, d/2) 是所有 x，后半 [d/2, d) 是所有 y — 配对关系跨越半个 head

索引公式

对应 Kernel 代码 (apply_token_rotary_embedding_kernel, IS_NEOX=true)

// NeoX: 前半后半配对
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos = cos_ptr[x_index];
sin = sin_ptr[x_index];

const T x = arr[x_index];
const T y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;CUDA

对应 Kernel 代码 (FusedNeoxRopeEmbeddingKernel)

// Fused 版本：从 packed QKV 中拆分并旋转
const int base_idx_left  = ... + h_bias;            // 前半
const int base_idx_right = base_idx_left + half_lastdim; // 后半

Load<T, VecSize>(&qkv[base_idx_left],  &left_vec);
Load<T, VecSize>(&qkv[base_idx_right], &right_vec);

// 旋转：left=x, right=y
left_vec[i]  = input_left * cos_tmp - input_right * sin_tmp;
right_vec[i] = input_right * cos_tmp + input_left * sin_tmp;CUDA

代表模型

LLaMA 1/2/3, Qwen, Baichuan, Mistral, Gemma, Yi, DeepSeek

3. Partial Rotary (部分旋转)

某些模型只旋转 head 维度的前一部分（rotary_dim < head_dim），后面的维度保持不变。这给模型保留了一部分"不受位置影响"的表示空间。

内存布局 (head_dim=8, rotary_dim=4, NeoX 配对)

前 rotary_dim 个元素参与旋转（NeoX 前后半配对），后 head_dim-rotary_dim 个元素 pass-through

对应 Kernel 代码 (GQAVariableLengthNeoxPartialRotarySplitKernel)

if (hi < q_num_head + kv_num_head) {       // 只有 Q 和 K 需要旋转
  if (h_bias < rotary_dim) {                // 在旋转区域内
    if (h_bias < half_rotary_dim) {
      // 前半：加载后半作为配对
      Load(&qkv[base_idx + half_rotary_dim], &src_vec_right);
      // x' = x*cos - y*sin
      src_vec[i] = input_left * cos - input_right * sin;
    } else {
      // 后半：加载前半作为配对
      Load(&qkv[base_idx - half_rotary_dim], &src_vec_right);
      // y' = y*cos + x*sin
      src_vec[i] = input_left * cos + input_right * sin;
    }
  }
  // else: h_bias >= rotary_dim, 不旋转，保持原值
}CUDA

代表模型

Phi-2/3 (partial_rotary_factor=0.5), GLM-4/4.1V, ChatGLM 系列, Gemma (部分层)

4. 3D RoPE (多维旋转位置编码)

3D RoPE 将位置从标量扩展到多维坐标 (t, h, w)，用于多模态和视频场景。在 GQAVariableLengthRotarySplitKernel 中通过 rope_3d 参数支持：

// 3D RoPE: 每个 batch 有独立的 emb 偏移
int64_t new_emb_idx = rope_3d
    ? emb_idx + ori_bi * head_dim * max_model_len  // 按 batch 偏移
    : emb_idx;                                       // 标准 1DCUDA

3D RoPE 的核心思想是将 head 维度再细分为若干组，分别编码时间、高度、宽度三个维度的位置。旋转公式不变，只是 cos/sin cache 的组织方式变为 [batch, temporal_pos * spatial_pos, dim_group]。

代表模型

Qwen-VL, LLaVA-Video, VideoRoPE (ICML 2025), VRoPE

Kernel 实现深度分析

下面对用户提供的 4 个 CUDA Kernel 进行逐一拆解。

Kernel 1: apply_rotary_embedding_kernel (基础版)

Kernel 2: FusedNeoxRopeEmbeddingKernel (Fused QKV Split + NeoX)

Kernel 3: GQAVariableLengthRotarySplitKernel (Interleaved + RMSNorm)

Kernel 4: GQAVariableLengthNeoxPartialRotarySplitKernel (Partial NeoX)

对比分析 / Comparison

配对方式对比

Kernel 实现对比

性能特征对比

Fused 程度 (越高越好)

内存效率 (减少 Global Memory 读写)

频率缩放方法 (影响 cos/sin cache，不影响 kernel)

以下方法只改变 cos_sin_cache 的预计算方式，不需要修改旋转 kernel 本身：

结论与建议 / Conclusions

从 4 个真实 kernel 的分析中，可以总结出以下工程建议：

1. 配对方式选择跟随模型架构。 NeoX (Half-Split) 是当前主流（LLaMA 系列），Interleaved 在部分模型中使用。两者数学等价，但内存访问模式不同。不要混用。

2. Fused Kernel 是 LLM 推理优化的标准做法。 将 QKV Split + RoPE + RMSNorm 等操作融合为一个 kernel，可以将 global memory 读写次数从 N 次降为 1 次读 + 1 次写，这对 memory-bound 的 RoPE 操作收益显著。

3. 向量化 Load/Store 必不可少。 基础版 kernel（逐元素访问）与 VecSize=4 版本相比，带宽利用率差距巨大。LDG.128（float4 / 4×bf16）是标准配置。

4. Warp-level 编程是 Fused Kernel 的关键技巧。 如 Kernel 3 中用 WelfordWarpAllReduce 做 warp 内归约计算 RMSNorm，不需要 shared memory 同步，性能和代码复杂度都优于 block-level 方案。

5. 关注 SM90+ 新特性。 Kernel 4 使用的 PDL (Programmatic Dependent Launch) 可以减少 kernel chain 的 pipeline bubble，对 decode latency 敏感的场景值得采用。

6. 频率缩放方法不影响 kernel。 所有 PI / NTK / YaRN / LongRoPE 等长度外推方法只改变 cos/sin cache 的预计算，kernel 本身无需任何修改即可支持。

参考资料 / References

1. RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021) — RoPE 原始论文
2. Rotary Embeddings: A Relative Revolution (EleutherAI Blog) — GPT-NeoX/GPT-J RoPE 实现差异
3. Extending the RoPE (EleutherAI Blog, YaRN) — YaRN 与频率缩放方法综述
4. YaRN: Efficient Context Window Extension (Peng et al., 2023) — YaRN 论文
5. LongRoPE: Extending LLM Context to 2048K (Microsoft, ICML 2024) — 非均匀维度缩放
6. LongRoPE2: Near-Lossless LLM Context Window Scaling (2025) — LongRoPE 改进版
7. HuggingFace Transformers Issue #25199 — LLaMA RoPE 实现差异讨论
8. vLLM Issue #747 — GPT-J vs GPT-NeoX RoPE 实现对比
9. vLLM PR #22593 — Partial Rotary Embedding MRoPE Triton 支持
10. VideoRoPE: What Makes for Good Video Rotary Position Embedding? (2025) — 3D RoPE 在视频 LLM 中的应用
11. Inside RoPE: Rotary Magic into Position Embeddings — 图解 RoPE 原理
12. CUDA C++ Programming Guide (NVIDIA) — grid 维度上限与 SM90 PDL 参考