4. RoPE的代码实现：PyTorch实现RoPE的前向传播
好，咱们直接进入正题。上一节我们把RoPE的数学原理讲清楚了，这一节就动手写代码。我个人习惯是，先理解公式，再动手实现，最后再集成到Attention里。这样每一步都踩实了，后面调bug的时候心里才有底。
4.1 从公式到代码：RoPE的核心逻辑
先回顾一下RoPE的核心操作。说白了，就是对Query和Key做旋转，让它们带上位置信息。旋转矩阵是分块的，每个2x2的块对应一个频率。
代码实现的关键步骤就三步：
- 生成频率向量（也就是公式里的θ）
- 计算旋转角度（位置m乘以频率θ）
- 应用旋转矩阵到Query和Key上
嗯，这里要注意，实际实现时我们不会真的构建一个巨大的旋转矩阵，那样太浪费显存了。我们用的是“按位旋转”的技巧，把矩阵乘法变成了向量运算。
4.2 逐行解析：PyTorch实现
下面是我写的一个RoPE实现，代码不长，但每一行都有讲究。我在项目中遇到过几次因为维度处理不对导致训练loss不收敛的情况，所以这次我会把维度变化也讲清楚。
import torch
import torch.nn as nn
import math
class RotaryPositionalEmbedding(nn.Module):
def __init__(self, dim, max_seq_len=2048, base=10000.0):
super().__init__()
self.dim = dim
self.max_seq_len = max_seq_len
self.base = base
# 1. 生成频率向量 theta_i = base^(-2i/dim)
# 注意：dim必须是偶数，因为每2个维度共享一个频率
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
# inv_freq shape: [dim/2]
# 2. 预计算所有位置的旋转角度
# 位置索引: [0, 1, 2, ..., max_seq_len-1]
position = torch.arange(max_seq_len).float()
# position shape: [max_seq_len]
# 外积: 每个位置乘以每个频率
freqs = torch.einsum("i,j->ij", position, inv_freq)
# freqs shape: [max_seq_len, dim/2]
# 3. 存储cos和sin值
# 注意：这里要重复一次，因为每个频率对应2个维度
self.register_buffer("cos_cached", freqs.cos().repeat_interleave(2, dim=-1))
self.register_buffer("sin_cached", freqs.sin().repeat_interleave(2, dim=-1))
# cos_cached shape: [max_seq_len, dim]
def forward(self, x, offset=0):
"""
x: [batch_size, seq_len, num_heads, dim]
offset: 用于推理时的KV Cache偏移
"""
seq_len = x.shape[1] + offset
# 从缓存中取出对应长度的cos和sin
cos = self.cos_cached[offset:seq_len].unsqueeze(0).unsqueeze(0)
sin = self.sin_cached[offset:seq_len].unsqueeze(0).unsqueeze(0)
# cos shape: [1, seq_len, 1, dim]
# 应用旋转
# 关键公式: RoPE(x) = x * cos + rotate_half(x) * sin
x_rotated = x * cos + self._rotate_half(x) * sin
return x_rotated
def _rotate_half(self, x):
"""
将x的后半部分取负，然后交换前后半部分
对应旋转矩阵中的 [-x1, x0] 操作
"""
x1, x2 = x.chunk(2, dim=-1)
# x1: 前半部分, x2: 后半部分
return torch.cat([-x2, x1], dim=-1)
代码要点解析：
- inv_freq的计算：这里用了
torch.arange(0, dim, 2)，只取偶数索引。因为每两个维度共享一个频率，所以频率向量的长度是dim/2。 - repeat_interleave的作用：把频率对应的cos/sin值复制一份，这样就能和完整的dim维度对齐。比如dim=4时，频率向量长度是2，复制后变成4。
- _rotate_half的实现：这个函数实现了旋转矩阵的核心操作——把向量分成两半，交换并取负。你想想看，这不就是二维旋转的推广吗？
4.3 与标准Attention的集成
有了RoPE模块，集成到Attention里就很简单了。我见过一些实现把RoPE放在Attention外面，但我觉得放在里面更干净，因为RoPE只影响Q和K，不影响V。
class RotaryAttention(nn.Module):
def __init__(self, dim, num_heads, max_seq_len=2048):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
# QKV投影
self.wq = nn.Linear(dim, dim, bias=False)
self.wk = nn.Linear(dim, dim, bias=False)
self.wv = nn.Linear(dim, dim, bias=False)
self.wo = nn.Linear(dim, dim, bias=False)
# RoPE模块
self.rotary = RotaryPositionalEmbedding(
dim=self.head_dim,
max_seq_len=max_seq_len
)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
# 1. 线性投影并分头
q = self.wq(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
k = self.wk(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
v = self.wv(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
# q/k/v shape: [batch_size, seq_len, num_heads, head_dim]
# 2. 应用RoPE（只对Q和K）
q = self.rotary(q)
k = self.rotary(k)
# 3. 标准Attention计算
# 转置为 [batch_size, num_heads, seq_len, head_dim]
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# 计算注意力分数
scale = math.sqrt(self.head_dim)
attn_weights = torch.matmul(q, k.transpose(-2, -1)) / scale
if mask is not None:
attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))
attn_weights = torch.softmax(attn_weights, dim=-1)
# 加权求和
out = torch.matmul(attn_weights, v)
# out shape: [batch_size, num_heads, seq_len, head_dim]
# 合并多头
out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
out = self.wo(out)
return out
避坑指南：
我曾经在集成RoPE时犯过一个低级错误——把RoPE应用到了V上。结果模型训练了三天，loss死活降不下去。后来排查才发现，V不应该带位置信息，它只是内容表示。所以记住：RoPE只作用于Q和K。
4.4 验证实现正确性
写完了代码，怎么验证对不对？我一般会做两个测试：
- 对称性测试：对于相同位置的Q和K，注意力分数应该最大。
- 距离衰减测试：位置差越大，注意力分数应该越小（在合理范围内）。
下面是一个简单的测试脚本：
def test_rope():
dim = 64
num_heads = 8
seq_len = 10
attn = RotaryAttention(dim, num_heads)
x = torch.randn(1, seq_len, dim)
# 测试1: 相同位置的注意力
out = attn(x)
print(f"Output shape: {out.shape}") # 应该是 [1, 10, 64]
# 测试2: 检查位置编码效果
# 构造两个相同的token在不同位置
token = torch.randn(1, 1, dim)
x1 = torch.cat([token, torch.zeros(1, 9, dim)], dim=1)
x2 = torch.cat([torch.zeros(1, 5, dim), token, torch.zeros(1, 4, dim)], dim=1)
out1 = attn(x1)
out2 = attn(x2)
# 相同token在不同位置的输出应该不同
assert not torch.allclose(out1[:, 0, :], out2[:, 5, :], atol=1e-6)
print("位置编码生效：相同token在不同位置输出不同 ✓")
test_rope()
性能注意事项：
在实际训练大模型时，RoPE的cos/sin缓存可能会占用不少显存。如果max_seq_len设得太大（比如32K），建议使用torch.jit.script或者torch.compile来优化。我在训练7B模型时，就因为这个缓存导致OOM，后来改成按需计算才解决。
4.5 小结
这一节我们完成了RoPE的完整实现，从频率生成到旋转应用，再到与Attention的集成。代码量不大，但每个细节都值得推敲。特别是那个_rotate_half函数，它把矩阵乘法变成了向量操作，是RoPE能高效运行的关键。
下一节我们会深入RoPE的数学性质，看看它为什么能处理长序列，以及如何扩展到更高维度。但在此之前，我建议你把这段代码跑一遍，亲手感受一下位置编码的效果。