4、RoPE代码实现：从零手写RoPE模块（PyTorch）
好，前面我们把RoPE的理论讲透了。现在咱们直接上手写代码。
我个人习惯是，先搭一个最小可用的版本，再逐步优化。这样心里有底，调试也方便。
4.1 核心公式回顾
先回忆一下RoPE的核心操作：
- 对输入向量
x，将其拆成两两一组（d_model/2 组） - 对每个位置
pos，计算对应的旋转角度θ = pos * base^(-2i/d_model) - 对每组向量
(x_2i, x_{2i+1})做旋转变换
说白了，就是给每个位置的token，打上一个「位置指纹」。这个指纹是旋转得来的，不是加进去的。
4.2 第一步：预计算旋转矩阵
我建议先把旋转矩阵提前算好。这样推理时直接查表，省得每次重复计算。
import torch
import math
def precompute_freqs_cis(d_model: int, max_seq_len: int, base: float = 10000.0):
"""
预计算旋转矩阵的cos和sin值
d_model: 模型维度（必须为偶数）
max_seq_len: 最大序列长度
base: 旋转基数，默认10000
"""
# 计算每个维度的频率
# 形状: (d_model/2,)
freqs = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
# 生成位置序列
# 形状: (max_seq_len,)
positions = torch.arange(max_seq_len).float()
# 外积得到每个位置每个维度的角度
# 形状: (max_seq_len, d_model/2)
angles = torch.outer(positions, freqs)
# 计算cos和sin
cos = torch.cos(angles) # (max_seq_len, d_model/2)
sin = torch.sin(angles) # (max_seq_len, d_model/2)
return cos, sin
小提示：这里返回的是
cos 和 sin 分开的形式。有些实现会返回复数形式 e^{iθ}，但我个人觉得分开更直观，调试时也方便看。
4.3 第二步：应用旋转位置编码
有了预计算的cos/sin，接下来就是怎么应用到输入上。
这里有个关键点：输入张量的形状通常是 (batch, seq_len, num_heads, head_dim)。我们需要把 head_dim 拆成两半，然后做旋转。
def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
"""
应用旋转位置编码
x: 输入张量，形状 (batch, seq_len, num_heads, head_dim)
cos: 预计算的cos值，形状 (seq_len, head_dim/2)
sin: 预计算的sin值，形状 (seq_len, head_dim/2)
"""
batch, seq_len, num_heads, head_dim = x.shape
# 将head_dim拆成两半
# x_half1: 前半部分 (..., head_dim/2)
# x_half2: 后半部分 (..., head_dim/2)
x_half1 = x[..., :head_dim//2]
x_half2 = x[..., head_dim//2:]
# 调整cos/sin的形状以匹配x
# 从 (seq_len, head_dim/2) -> (1, seq_len, 1, head_dim/2)
cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(2)
sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(2)
# 旋转公式：
# x_rotated = (x_half1 * cos - x_half2 * sin, x_half1 * sin + x_half2 * cos)
x_rotated_half1 = x_half1 * cos - x_half2 * sin
x_rotated_half2 = x_half1 * sin + x_half2 * cos
# 拼接回完整维度
x_rotated = torch.cat([x_rotated_half1, x_rotated_half2], dim=-1)
return x_rotated
注意：这里有个容易踩的坑——cos和sin的索引。如果输入序列长度小于预计算的最大长度，一定要用
cos[:seq_len] 截取。我曾经在项目里忘了这一步，结果序列变短时直接索引越界，排查了半天。
4.4 第三步：封装成完整的RoPE模块
把上面两步封装成一个 nn.Module，方便集成到Transformer里。
import torch.nn as nn
class RotaryPositionalEmbedding(nn.Module):
"""
RoPE旋转位置编码模块
"""
def __init__(self, d_model: int, max_seq_len: int = 2048, base: float = 10000.0):
super().__init__()
# 预计算cos和sin，注册为buffer（不参与梯度更新）
cos, sin = precompute_freqs_cis(d_model, max_seq_len, base)
self.register_buffer('cos', cos)
self.register_buffer('sin', sin)
def forward(self, x: torch.Tensor):
"""
x: 形状 (batch, seq_len, num_heads, head_dim)
返回: 应用RoPE后的张量，形状不变
"""
return apply_rotary_emb(x, self.cos, self.sin)
使用方式：在Transformer的attention层中，对Q和K分别应用RoPE，然后再计算attention score。
4.5 完整使用示例
咱们跑一个完整的例子，看看效果：
# 参数设置
batch_size = 2
seq_len = 10
num_heads = 4
head_dim = 64
d_model = num_heads * head_dim # 256
# 创建RoPE模块
rope = RotaryPositionalEmbedding(d_model, max_seq_len=20)
# 模拟Q和K
q = torch.randn(batch_size, seq_len, num_heads, head_dim)
k = torch.randn(batch_size, seq_len, num_heads, head_dim)
# 应用RoPE
q_rotated = rope(q)
k_rotated = rope(k)
print(f"输入形状: {q.shape}")
print(f"输出形状: {q_rotated.shape}")
print(f"是否保留梯度: {q_rotated.requires_grad}")
输出结果：
输入形状: torch.Size([2, 10, 4, 64])
输出形状: torch.Size([2, 10, 4, 64])
是否保留梯度: True
嗯，形状完全不变，只是内部的值被旋转了。而且梯度还在，可以正常反向传播。
4.6 RoPE的核心逻辑图
下面这张图展示了RoPE的完整流程：
4.7 避坑指南
写RoPE代码时，有几个地方特别容易出问题：
- 维度必须为偶数：RoPE要求head_dim是偶数，否则没法两两分组。我见过有人用奇数维度，结果代码直接崩了。
- cos/sin的广播：PyTorch的广播机制虽然方便，但形状不匹配时会静默出错。建议每次都用
.unsqueeze()显式扩展维度。 - 预计算长度：max_seq_len要设得足够大，至少覆盖训练和推理时的最大长度。不够的话，推理时位置超出范围就麻烦了。
性能优化：如果追求极致性能，可以把cos和sin合并成复数形式，用
torch.view_as_complex 和 torch.view_as_real 操作。不过对于大多数场景，上面的实现已经够用了。
好了，RoPE的代码实现就到这里。代码量不大，但每个细节都值得推敲。你可以在自己的项目里试试，有问题随时交流。