4、RoPE代码实现：Python实现RoPE编码函数，批量处理与维度适配，与PyTorch的集成

好，咱们直接进入正题。前面讲了那么多理论，现在该动手了。我个人习惯是，先写一个最朴素的版本，跑通了再考虑优化。你想想看，如果连单样本都搞不定，谈什么批量处理？

4.1 从零开始：一个朴素的RoPE实现

先别急着上PyTorch，咱们用纯NumPy把核心逻辑理清楚。RoPE说白了就是给每个位置分配一个旋转角度，然后对query和key做旋转。

import numpy as np

def rope_single_pos(pos, d_model, base=10000.0):
    """
    计算单个位置的RoPE旋转矩阵
    pos: 位置索引 (int)
    d_model: 特征维度 (必须为偶数)
    base: 频率基数，默认10000
    """
    # 计算每个维度的频率
    i = np.arange(0, d_model, 2)  # 取偶数索引
    theta = pos / (base ** (i / d_model))
    
    # 构建旋转矩阵
    cos_vals = np.cos(theta)
    sin_vals = np.sin(theta)
    
    return cos_vals, sin_vals

# 测试一下
d_model = 4
cos0, sin0 = rope_single_pos(0, d_model)
cos1, sin1 = rope_single_pos(1, d_model)
print(f"位置0: cos={cos0}, sin={sin0}")
print(f"位置1: cos={cos1}, sin={sin1}")

嗯，这里要注意。我刚开始写的时候犯过一个低级错误——忘了d_model必须是偶数。因为RoPE是把特征维度两两分组旋转的，奇数维度会出问题。后来我加了个断言，养成习惯了。

4.2 批量处理：向量化才是王道

单样本跑通了，但实际训练时我们是一次处理一个batch的。用for循环逐位置计算？太慢了。咱们得用矩阵运算一次性搞定。

def rope_batch(seq_len, d_model, base=10000.0):
    """
    批量生成所有位置的旋转矩阵
    seq_len: 序列长度
    d_model: 特征维度
    返回: cos_emb, sin_emb, 形状均为 [seq_len, d_model//2]
    """
    # 生成位置序列 [0, 1, 2, ..., seq_len-1]
    pos = np.arange(seq_len)[:, np.newaxis]  # [seq_len, 1]
    
    # 生成维度索引 [0, 2, 4, ..., d_model-2]
    i = np.arange(0, d_model, 2)[np.newaxis, :]  # [1, d_model//2]
    
    # 计算频率
    theta = pos / (base ** (i / d_model))  # [seq_len, d_model//2]
    
    cos_emb = np.cos(theta)
    sin_emb = np.sin(theta)
    
    return cos_emb, sin_emb

# 测试
cos_emb, sin_emb = rope_batch(4, 8)
print(f"cos形状: {cos_emb.shape}")  # [4, 4]
print(f"sin形状: {sin_emb.shape}")  # [4, 4]

为什么会这样设计？你看，我们把位置和维度索引都广播成矩阵，一次计算就得到所有位置的cos和sin值。这比逐位置计算快了不止一个数量级。

4.3 维度适配：把旋转应用到特征上

有了旋转矩阵，怎么把它应用到query和key上？这里有个小技巧——需要把特征向量两两分组，然后做旋转。

def apply_rope(x, cos_emb, sin_emb):
    """
    将RoPE旋转应用到输入张量x上
    x: [batch_size, seq_len, d_model]
    cos_emb: [seq_len, d_model//2]
    sin_emb: [seq_len, d_model//2]
    """
    batch_size, seq_len, d_model = x.shape
    
    # 将x分成两半：x1是偶数索引，x2是奇数索引
    x1 = x[:, :, 0::2]  # [batch, seq_len, d_model//2]
    x2 = x[:, :, 1::2]  # [batch, seq_len, d_model//2]
    
    # 应用旋转
    # [x1'; x2'] = [cos, -sin; sin, cos] @ [x1; x2]
    x1_new = x1 * cos_emb - x2 * sin_emb
    x2_new = x1 * sin_emb + x2 * cos_emb
    
    # 合并回去
    x_rotated = np.stack([x1_new, x2_new], axis=-1)  # [batch, seq_len, d_model//2, 2]
    x_rotated = x_rotated.reshape(batch_size, seq_len, d_model)
    
    return x_rotated

# 完整测试
batch_size, seq_len, d_model = 2, 4, 8
x = np.random.randn(batch_size, seq_len, d_model)
cos_emb, sin_emb = rope_batch(seq_len, d_model)
x_rotated = apply_rope(x, cos_emb, sin_emb)
print(f"输入形状: {x.shape}, 输出形状: {x_rotated.shape}")

4.4 与PyTorch集成：从NumPy到GPU

实际项目中我们肯定用PyTorch。把上面的代码迁移到PyTorch很简单，主要改动两点：用torch代替numpy，注意设备（device）管理。

import torch
import torch.nn as nn

class RotaryPositionEmbedding(nn.Module):
    def __init__(self, d_model, base=10000.0):
        super().__init__()
        self.d_model = d_model
        self.base = base
        
    def forward(self, x):
        """
        x: [batch_size, seq_len, d_model]
        """
        batch_size, seq_len, d_model = x.shape
        device = x.device
        
        # 生成位置和维度索引
        pos = torch.arange(seq_len, device=device).float()[:, None]  # [seq_len, 1]
        i = torch.arange(0, d_model, 2, device=device).float()[None, :]  # [1, d_model//2]
        
        # 计算频率
        theta = pos / (self.base ** (i / d_model))  # [seq_len, d_model//2]
        
        cos_emb = theta.cos()  # [seq_len, d_model//2]
        sin_emb = theta.sin()  # [seq_len, d_model//2]
        
        # 应用旋转
        x1 = x[:, :, 0::2]  # [batch, seq_len, d_model//2]
        x2 = x[:, :, 1::2]  # [batch, seq_len, d_model//2]
        
        x1_new = x1 * cos_emb - x2 * sin_emb
        x2_new = x1 * sin_emb + x2 * cos_emb
        
        # 合并
        x_rotated = torch.stack([x1_new, x2_new], dim=-1)  # [batch, seq_len, d_model//2, 2]
        x_rotated = x_rotated.reshape(batch_size, seq_len, d_model)
        
        return x_rotated

# 测试
rope = RotaryPositionEmbedding(d_model=8)
x = torch.randn(2, 4, 8)
output = rope(x)
print(f"PyTorch输出形状: {output.shape}")

4.5 核心逻辑流程图

下面我用一张图把RoPE的整个计算流程串起来，方便你理解数据是怎么流动的。

这张图把RoPE的流程分成了三步：先生成旋转矩阵，同时把输入特征拆成奇偶两组，最后做旋转合并。我在项目中调试时，经常对着这张图检查维度有没有对齐。

好了，代码实现部分就到这里。你可以在自己的项目里直接复用这些代码，记得根据实际情况调整base值和d_model大小。我个人习惯把base设成10000，但如果你处理的是超长序列（比如8192以上），可以试试调大到50000或100000。

---

公众号:蓝海资料掘金营，微信deep3321