逆频率计算、cos/sin 缓存与 apply_rotary_pos_emb 的向量化实现

配套代码

逆频率预计算

RoPE 实现的第一步是预计算每个维度对的频率。回顾上一章的公式：

在代码中，我们通常计算逆频率（inverse frequency），也就是 𝜃 𝑖 θ i 本身：

# Step 1: 计算逆频率
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
# inv_freq shape: (dim // 2,)
# inv_freq[0] = 1.0, inv_freq[1] ≈ 0.85, ..., inv_freq[-1] ≈ 0.00011
# Step 2: 构建位置-频率矩阵
t = torch.arange(max_seq_len).float()
freqs = torch.outer(t, inv_freq)  # (max_seq_len, dim // 2)
# freqs[m, i] = m * theta_i

我们来看看 torch.outer 做了什么：

# 假设 max_seq_len = 4, dim = 6 (3 个维度对)
t = [0, 1, 2, 3]
inv_freq = [θ₀, θ₁, θ₂]
# outer product:
# freqs[m, i] = m * θᵢ
freqs = [[0·θ₀, 0·θ₁, 0·θ₂],    # 位置 0：不旋转
         [1·θ₀, 1·θ₁, 1·θ₂],    # 位置 1
         [2·θ₀, 2·θ₁, 2·θ₂],    # 位置 2
         [3·θ₀, 3·θ₁, 3·θ₂]]    # 位置 3

有了 freqs，对每个位置 𝑚 m 取 cos 和 sin 就得到旋转所需的系数：

cos = freqs.cos()  # (max_seq_len, dim // 2)
sin = freqs.sin()  # (max_seq_len, dim // 2)
# cos[m, i] = cos(m * θ_i)，sin[m, i] = sin(m * θ_i)

到这一步，两种实现风格的计算完全相同。它们的区别在于下一步：如何把这些 cos/sin 系数应用到向量上，也就是维度怎么配对。

两种实现风格

RoPE 的数学公式只规定了”每对维度独立旋转”，但没有规定维度怎么配对。实践中有两种常见的配对方式，数学上完全等价，只是内存布局不同：

𝑥 0 , 𝑥 1 ) , ( 𝑥 2 , 𝑥 3 ) , … (x 0

,x 1

),(x 2

,x 3

),… | ( 𝑥 0 , 𝑥 𝑑 / 2 ) , ( 𝑥 1 , 𝑥 𝑑 / 2 + 1 ) , … (x 0