RoPE Kernel 可视化解析

1 整体数据流

2 RoPE 核心运算：2D 旋转

对 head 内的每一对元素 (x, y)，施加一个角度为 θ 的 2D 旋转。θ 由 position 和维度下标共同决定（通过预计算的 cos/sin cache 查表）。

x' = x × cos(θ) - y × sin(θ)
y' = y × cos(θ) + x × sin(θ)

3 四种 RoPE 配对模式

配对方式不同，但旋转公式完全相同。head_size 内只有前 rot_dim 个元素参与旋转，后面的元素不动。

NeoX：前半和后半配对。 arr[i] 和 arr[i + embed_dim] 组成旋转对。LLaMA、Qwen、Baichuan 等主流模型使用此方式。

// rot_offset = 0,1,2,...,embed_dim-1
x_index = rot_offset // 前半
y_index = embed_dim + rot_offset // 后半

GPT-J：相邻元素配对。 arr[2i] 和 arr[2i+1] 组成旋转对。GPT-J、CodeGen 使用此方式。

// rot_offset = 0,1,2,...,embed_dim-1
x_index = 2 * rot_offset // 偶数位
y_index = 2 * rot_offset + 1 // 奇数位

CUDA 优势：一次 float4 向量化 Load（4 个 fp16 元素）恰好覆盖 2 个旋转对，cos/sin 只需加载 HalfVecSize。相比 NeoX 的两次 Load，Interleaved 的内存访问效率更高。

Partial Rotary：只旋转前 rotary_dim 维度。 head_dim 中只有前 rotary_dim 个元素参与 NeoX 配对旋转，后面的维度保持原值。Phi-3、GLM-4 使用此策略。

// GQAVariableLengthNeoxPartialRotarySplitKernel
if (h_bias < half_rotary_dim) {
  src_vec[i] = left * cos - right * sin   // x' = x*cos - y*sin
} else if (h_bias < rotary_dim) {
  src_vec[i] = left * cos + right * sin   // y' = y*cos + x*sin
} else {
  // 不旋转, 保持原值
}

代表模型：Phi-2/3（partial_rotary_factor=0.5）、GLM-4/4.1V、ChatGLM 系列。保留部分维度不受位置编码干扰，让模型学习纯内容相关的 attention pattern。

Fused NeoX：从 packed QKV 中拆分并同时旋转。 输入是 Q/K/V 拼在一起的 packed tensor，kernel 一次完成 QKV Split + NeoX RoPE，减少 global memory 读写。

// FusedNeoxRopeEmbeddingKernel (VecSize=4)
Load<T, 4>(&qkv[base_idx_left], &left_vec);   // 前半
Load<T, 4>(&qkv[base_idx_right], &right_vec); // 后半 = left + half_lastdim

if (qkv_id < 2) { // Q 和 K 旋转
  left_vec[i] = left * cos - right * sin
  right_vec[i] = right * cos + left * sin
}
Store(left_vec, &out[split_left]); Store(right_vec, &out[split_right]);

精度细节：cos/sin 使用 float 精度，bf16/fp16 数据在运算前 cast 到 float，运算后 cast 回来。避免低精度累积误差。

4 GPU 线程映射

Grid 维度

grid.x = num_tokens

1 block = 1 token，所有 head 的旋转

Block 维度

blockDim = min(num_heads * embed_dim, 512)

线程在 block 内 stride-loop 处理所有 head

stride-loop 模式：for (int i = threadIdx.x; i < nq; i += blockDim.x) — 每个线程从自己的 threadIdx 开始，每次跨 blockDim 步长，直到覆盖全部 nq 个旋转对。这意味着一个线程可能跨越多个 head。同一组 cos/sin（同一 position）被所有 head 共享。

5 cos_sin_cache 查表过程

cache_ptr = cos_sin_cache + pos * rot_dim
cos_ptr = cache_ptr // 前 embed_dim 个
sin_ptr = cache_ptr + embed_dim // 后 embed_dim 个

6 head_size 内部：旋转区 vs 非旋转区

当 rot_dim < head_size 时，head 的前 rot_dim 个元素参与旋转，后面的元素保持不动。常见于部分旋转模型。

7 GQA + Variable-Length RoPE

GQA (Grouped-Query Attention) 下 Q heads 和 KV heads 数量不同。Variable-Length 模式支持 batch 内每条序列长度不等（padding_offset 压缩存储）。kernel 需要同时处理 Q 和 K 的旋转，并正确映射 GQA 的 head 分组。

// GQAVariableLengthRotarySplitKernel — 位置恢复逻辑
token_idx = blockIdx.x; // compressed 序列中的 token 位置
pos = token_idx + padding_offset[token_idx]; // 恢复原始位置
// 或者使用 position_ids[] 直接传入

// GQA head 映射:
q_head_idx = i / embed_dim; // Q 的 head 编号
kv_head_idx = q_head_idx / num_heads_per_group; // 对应的 KV head

代表实现：GQAVariableLengthRotarySplitKernel（GPT-J 配对）和 GQAVariableLengthNeoxPartialRotarySplitKernel（NeoX 配对 + Partial Rotary）。两者共享 GQA + Variable-Length 的线程映射逻辑，区别仅在配对方式和是否有 pass-through 区域。

8 Fused RMSNorm + QKV Split + RoPE Kernel

最高级的融合方案：在一个 kernel 内完成 RMSNorm 归一化 + QKV 拆分 + NeoX RoPE 旋转。消除所有中间 tensor 的 HBM 读写，将 3 次 kernel launch 合并为 1 次。

// FusedNeoxRopeEmbeddingKernel + RMSNorm 的伪代码
// Step 1: Load + Norm
vec = Load<T, VecSize>(&qkv[idx]);
float local_sum = 0;
for (int i = 0; i < VecSize; i++) local_sum += vec[i] * vec[i];
mean_sq = WelfordWarpAllReduce(local_sum) / head_dim;
rms = rsqrt(mean_sq + 1e-6f) * weight[h_offset];

// Step 2: Normalize → Split → Rotate → Store
left = (float)vec_left[i] * rms;
right = (float)vec_right[i] * rms;
out_left[i] = left * cos - right * sin;
out_right[i] = right * cos + left * sin;

WelfordWarpAllReduce 关键特性：使用 __shfl_xor_sync 在 warp 的 32 个 lane 内完成规约，无需 __syncthreads() 和 shared memory。5 轮 shuffle（log2(32)=5）即可完成全 warp 归约。这是 fused kernel 的效率核心——把本需 shared memory 的 norm 压缩到纯寄存器操作。

9 频率缩放方法：只改 cos/sin cache，不改 kernel

所有频率缩放方法都是正交于配对方式的。它们只修改预计算的 cos_sin_cache（改变 θ 的计算方式），kernel 本身的旋转逻辑完全不变。这是 RoPE 系统设计中非常重要的解耦。

工程意义：添加新的频率缩放方法（如从 NTK 切换到 YaRN）不需要修改任何 CUDA kernel 代码。只需在 Python 侧修改 cos_sin_cache 的预计算逻辑，kernel 的四种配对模式可以复用。这也是为什么 PaddlePaddle / vLLM 等框架把 RoPE kernel 和 frequency scaling 完全解耦的原因。

10 性能优化要点

11 全局对比汇总

特性	NeoX (Half-Split)	GPT-J (Interleaved)	Partial Rotary	Fused NeoX
配对规则	arr[i] ↔ arr[i+d/2]	arr[2i] ↔ arr[2i+1]	NeoX 方式 (仅前 rotary_dim)	NeoX 方式 (从 packed QKV)
代表模型	LLaMA, Qwen, Mistral, Baichuan	GPT-J, CodeGen	Phi-2/3, GLM-4, ChatGLM	PaddlePaddle 推理优化
旋转维度	全部 rot_dim	全部 rot_dim	前 rotary_dim, 后面 pass-through	全部 head_dim (Q,K); V 不旋转
cos/sin cache 读取	2 次 Load (前半+后半分开)	1 次 Load (交错对内连续)	2 次 Load (同 NeoX)	2 次 Load (同 NeoX)
HBM 读写次数	1R + 1W (in-place)	1R + 1W (in-place)	1R + 1W (in-place/split)	1R + 1W (含 QKV split)
融合程度	仅 RoPE	仅 RoPE	仅 RoPE	QKV Split + RoPE (可+RMSNorm)
GQA 支持	需额外 kernel 或参数	GQAVariableLengthRotarySplitKernel	GQAVariableLength...PartialRotary...	通过 num_heads 参数
Variable-Length	padding_offset	padding_offset + seq_lens	padding_offset + seq_lens	通常固定长度
向量化策略	VecSize 对齐 head_dim	float4 覆盖 2 个旋转对	VecSize 对齐, 分段判断	VecSize=4, 前后半分别 Load

选择建议：优先使用模型默认的配对方式（模型权重训练时确定，推理时不能更改）。在此基础上，如果需要长上下文，选择合适的频率缩放（如 YaRN、Dynamic NTK）。如果目标是推理性能，优先 Fused kernel，把 QKV Split / RMSNorm / RoPE 合并。

Fused Rotary Position Encoding (RoPE)