CookLLM文档RoPE 代码实现
会员专享逆频率计算、cos/sin 缓存与 apply_rotary_pos_emb 的向量化实现
在权益中心获取代码逆频率预计算
上一章把 RoPE 推导成了"对每个二维子空间做旋转"。落到代码上，这句话要拆成三件具体的事：算出每个子空间的旋转速度、把速度乘上位置得到旋转角、再把角度变成可直接相乘的 / 张量。
第一件事最容易被当成一行公式带过，但它决定了模型的外推行为。RoPE 实现的第一步是预计算每个维度对的频率。回顾上一章的公式：
注意这是一个指数衰减的序列： 时 ，转得最快，一个位置就走一整弧度； 取到 时 已经小到 量级，要走上万个位置才转过一圈。低维负责分辨相邻 token，高维负责编码长程距离，这个分工是后面所有长度外推方法（NTK-aware、YaRN）的作用对象。
在代码中，我们通常计算逆频率（inverse frequency），也就是 本身：
# 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，对每个位置 取 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 系数应用到向量上，也就是维度怎么配对。
登录以继续阅读
这是一篇付费内容，请登录您的账户以访问完整内容。