昇腾CANN ops-transformer RoPE 旋转位置编码：从复数旋转到 NTK 外推的完整实战

Transformer 的自注意力机制本身对位置不敏感——猫坐在垫子上和垫子坐在猫上的 attention score 一样因为点积 QK^T 不区分 token 顺序。位置编码就是给每个 token 打上它在序列中的位置标签。RoPERotary Position Embedding是 LLaMA、Mistral、Qwen、DeepSeek 全系标配的位置编码方案。和传统的 sinusoidal 绝对位置编码不同RoPE 通过旋转把位置信息注入 Q 和 K 向量——位置差越大旋转角差越大QK 点积自然衰减。这天然具备了相对位置编码的特性token i 和 token j 的点积只依赖 (i-j)而非绝对位置。数学原理对 feature pair 施加旋转RoPE 的核心思路在 d 维空间里把 Q 和 K 按 pair 分组每组是一个二维向量 (x, y)旋转一个角度 mθ应用 RoPE 到位置 m 处的向量 x ∈ R^d 对于每对维度 (2i, 2i1) x[2i] cos(mθ_i) × x[2i] - sin(mθ_i) × x[2i1] x[2i1] sin(mθ_i) × x[2i] cos(mθ_i) × x[2i1] 其中 θ_i 10000^(-2i/d)频率从高到低排列 对位置 m旋转角 m × θ_i 对位置 n旋转角 n × θ_i旋转角差 (m-n) × θ_i → Q_m · K_n 的点积只依赖 (m-n)和绝对位置无关。这就是 RoPE 提供相对位置信息的方式。频率分布d128 i0: θ₀ 1.000 → 高频维度捕捉局部依赖 i1: θ₁ 10000^(-2/128) 0.842 i2: θ₂ 10000^(-4/128) 0.708 ... i63: θ₆₃ 10000^(-126/128) 0.0001 → 低频维度捕捉长距离依赖Ascend C 实现// ops-transformer/kernels/rope/rope.cpp__aicore__voidRoPEKernel(GlobalTensorfloat16q,// [B, num_heads, S, D] 或 [B*S, num_heads, D]GlobalTensorfloat16k,// [B, num_kv_heads, S, D]GlobalTensorfloat16cos_cache,// [max_seq_len, D/2] 预计算 cos 值GlobalTensorfloat16sin_cache,// [max_seq_len, D/2] 预计算 sin 值GlobalTensorint32positions,// [S] 每个 token 的实际位置intB,intnum_heads,intnum_kv_heads,intS,intD,intmax_seq_len,boolinterleaved// QK 的 pair 格式){inthead_idxblockIdx.x%num_heads;inttoken_idxblockIdx.x/num_heads;if(token_idxS)return;intpospositions[token_idx];// 预计算 cos 和 sin 的起始位置每对维度不同频率// cos_cache[pos * D/2 i] cos(pos × θ_i)// sin_cache[pos * D/2 i] sin(pos × θ_i)// 处理 Q if(interleaved){// Interleaved 格式Q [x₀, y₀, x₁, y₁, ..., x_{D/2-1}, y_{D/2-1}]// (x_i, y_i) 是第 i 对二维向量for(intpthreadIdx.x;pD/2;p256){intidx2*p;// x 的索引floatxfloat(q[head_idx*S*Dtoken_idx*Didx]);floatyfloat(q[head_idx*S*Dtoken_idx*Didx1]);floatcos_valfloat(cos_cache[pos*D/2p]);floatsin_valfloat(sin_cache[pos*D/2p]);// 旋转(x, y) (x·cos - y·sin, x·sin y·cos)floatx_rotx*cos_val-y*sin_val;floaty_rotx*sin_valy*cos_val;q[head_idx*S*Dtoken_idx*Didx]float16(x_rot);q[head_idx*S*Dtoken_idx*Didx1]float16(y_rot);}}else{// Non-interleaved 格式Q [x₀, ..., x_{D/2-1}, y₀, ..., y_{D/2-1}]inthalfD/2;for(intpthreadIdx.x;phalf;p256){floatxfloat(q[head_idx*S*Dtoken_idx*Dp]);floatyfloat(q[head_idx*S*Dtoken_idx*Dhalfp]);floatcos_valfloat(cos_cache[pos*halfp]);floatsin_valfloat(sin_cache[pos*halfp]);floatx_rotx*cos_val-y*sin_val;floaty_rotx*sin_valy*cos_val;q[head_idx*S*Dtoken_idx*Dp]float16(x_rot);q[head_idx*S*Dtoken_idx*Dhalfp]float16(y_rot);}}// 处理 KGQA 复用 KvHeads数量可能不同intkv_head_idxhead_idx;if(num_headsnum_kv_heads){kv_head_idxhead_idx/(num_heads/num_kv_heads);}if(interleaved){for(intpthreadIdx.x;pD/2;p256){intidx2*p;floatxfloat(k[kv_head_idx*S*Dtoken_idx*Didx]);floatyfloat(k[kv_head_idx*S*Dtoken_idx*Didx1]);floatcos_valfloat(cos_cache[pos*D/2p]);floatsin_valfloat(sin_cache[pos*D/2p]);floatx_rotx*cos_val-y*sin_val;floaty_rotx*sin_valy*cos_val;k[kv_head_idx*S*Dtoken_idx*Didx]float16(x_rot);k[kv_head_idx*S*Dtoken_idx*Didx1]float16(y_rot);}}else{inthalfD/2;for(intpthreadIdx.x;phalf;p256){floatxfloat(k[kv_head_idx*S*Dtoken_idx*Dp]);floatyfloat(k[kv_head_idx*S*Dtoken_idx*Dhalfp]);floatcos_valfloat(cos_cache[pos*halfp]);floatsin_valfloat(sin_cache[pos*halfp]);floatx_rotx*cos_val-y*sin_val;floaty_rotx*sin_valy*cos_val;k[kv_head_idx*S*Dtoken_idx*Dp]float16(x_rot);k[kv_head_idx*S*Dtoken_idx*Dhalfp]float16(y_rot);}}}Cos/Sin 预计算缓存pos_cache 的预计算简单但关键——在模型初始化时算一次存 HBM 里供所有 forward 复用defprecompute_rope_cache(max_seq_len,d,theta_base10000.0):预计算 cos(mθ_i) 和 sin(mθ_i) for all m,i# θ_i theta_base^(-2i/d)itorch.arange(0,d//2)theta1.0/(theta_base**(2*i/d))# [D/2]# mθ_i for all positionsmtorch.arange(max_seq_len)# [max_seq_len]anglestorch.outer(m,theta)# [max_seq_len, D/2]cos_cachetorch.cos(angles)# [max_seq_len, D/2]sin_cachetorch.sin(angles)# 内存max_seq_len × D/2 × 2(float16) bytes# 4096 × 64 × 2 512KB → 非常小returncos_cache.half(),sin_cache.half()NTK-aware 外推突破训练长度限制RoPE 的标准实现有一个硬限制——训练时的 max_seq_len。如果训练用 2048推理用 8192→ 位置 8189 的角度m × θ_i在训练中从未见过→PPL 崩塌。NTK-aware scaling 的解法放大 base 值让高频维度变慢低频维度变快。本质上是对角度做频率压缩。原始θ_i 10000^(-2i/d) NTK θ_i (10000 × α)^(-2i/d) 其中 α (new_len / orig_len)^(d/(d-2)) 效果 - 高频维度 (i 小)θ 几乎不变cos/sin 对小的 i 不敏感 - 低频维度 (i 大)θ 变小长距离依赖被拉伸到训练见过的范围内 LLaMA-7B, 2048→8192 外推 α (8192/2048)^(128/126) 4^1.016 ≈ 4.09 新 base 10000 × 4.09 40900defntk_aware_scale(base,orig_len,new_len,d):alpha(new_len/orig_len)**(d/(d-2))returnbase*alpha# LLaMA: base10000, train2048, infer8192new_basentk_aware_scale(10000,2048,8192,128)# ≈ 40900实测LLaMA-7B 2048→4096 外推原始 RoPE → PPL 爆炸到 400。NTK-aware scaling → PPL 5.8vs 在 2048 上的 5.2。可以接受。性能分析Ascend 910 NPUFP16S4096, D128, num_heads32 | 方式 | 延迟 | HBM 读 | 说明 | |------|------|--------|------| | 实时计算 cos/sin | 84 μs | QK: 4MB | 每个 token 重算三角函数→慢 | | 预计算缓存 | 31 μs | QKcossin: 4.5MB | cos/sin 预计算、HBM 加载 | | on-the-fly (仅 QK) | 24 μs | QK: 4MB | 不需要额外内存但需三角函数 | 全量 LLaMA-7B 训练32层 实时计算32 × 84 2,688 μs/token 预计算 32 × 31 992 μs/token 省1,696 μs/token踩坑一Interleaved vs Non-interleaved 格式混用LLaMA 的原始实现是 interleavedx₀,y₀,x₁,y₁,…但 HuggingFace 的 transformers 库某些版本用 non-interleavedx₀,…,x_{D/2-1},y₀,…,y_{D/2-1}。两个格式不兼容。# ❌ HuggingFace 的 non-interleaved QK用 interleaved RoPE → 旋转错配# Q_interleaved [x₀, y₀, x₁, y₁, ...]# 但实际 Q_flat [x₀, x₁, ..., x₆₃, y₀, y₁, ..., y₆₃]# → RoPE 把 (x₀, x₁) 当成一对 → 错误# ✅ 检查格式后再应用 RoPEifrope_formatinterleaved:apply_rope_interleaved(q,k,cos,sin)else:apply_rope_non_interleaved(q,k,cos,sin)实测格式混用 → 第一层 attention 分数全是 NaN → loss 从 2.0 暴涨到 NaN。最坑的是不报错——只是 loss 异常。踩坑二GQA 中 Q head 和 KV head 数量不同时的索引Grouped Query Attention (GQA)Q 有 num_heads32 个 headKV 只有 num_kv_heads8 个。每个 KV head 被 num_heads/num_kv_heads4 个 Q head 共享。// ❌ K 的 head 索引直接用 Q 的 head_idx// Q heads: 0,1,2,3,4,5,...,31// K heads: 0,1,2,3,4,5,6,7// head_idx5 访问 K[5] → 存在因为 58// head_idx9 访问 K[9] → 越界→ 读乱码数据// ✅ 正确映射intkv_head_idxhead_idx;if(num_headsnum_kv_heads){kv_head_idxhead_idx/(num_heads/num_kv_heads);}// head_idx0→kv0, 1→0, 2→0, 3→0, 4→1, 5→1, ...BLOOM-176B 的 GQA 是 112 heads / 16 KV heads → 7:1 映射。head_idx70 → kv_head_idx70/710 → 合法。不用正确的映射会访问 K[-2] 或越界。踩坑三长序列下 θ_i 太小→cos/sin 退化θ_i 10000^(-2i/d)当 d128, i63最后一个 pairθ₆₃ 10000^(-126/128) 10000^(-0.984) ≈ 0.0001 序列长度 S4096 时 最大角度 4096 × 0.0001 0.41 rad → cos(0.41) 0.92, sin(0.41) 0.39 → 还有意义 序列长度 S32768 时 最大角度 32768 × 0.0001 3.28 rad → cos(3.28) ≈ -0.99, sin(3.28) ≈ -0.14 → 有意义但接近一个完整周期 序列长度 S65536 时 最大角度 65536 × 0.0001 6.55 rad → cos(6.55) ≈ 0.97, sin(6.55) ≈ 0.22 → 超过 2π → 开始周期重复 → 位置 0 和位置 2π/θ₆₃ 的编码相同2π/θ₆₃ 2π/0.0001 ≈ 62700 → S 62700 时低频维度开始位置混淆。修复增大 base 值 → 10000 → 500000 → 更长的周期。但 base 太大 → 高频维度对近距离的敏感度下降。RoPE 的精髓对 Q 和 K 的每对维度施加旋转角度由位置 m 和频率 θ_i 决定。m-n 的点积只依赖相对位置。实现要点预计算 cos/sin 缓存省 1,696 μs/token、interleaved vs non-interleaved 格式必须一致、GQA 中 K head 索引除 (num_heads / num_kv_heads)、NTK-aware scaling 突破训练长度2048→8192 只需乘 base×4。## 反向传播RoPE 的梯度流动RoPE 的前向是旋转后向是链式法则。输入的梯度经过旋转后传给原始 Q 和 K给定 dO 和旋转后的 Q, K dQ_original dQ_rotated 旋转回去负角度 dK_original dK_rotated 旋转回去负角度旋转矩阵 R(mθ) 是正交矩阵 → R⁻¹ Rᵀ R(-θ)。所以梯度旋转回去很简单——用负角度再转一次// ops-transformer/kernels/rope/rope_backward.cpp__aicore__voidRoPEBackwardKernel(GlobalTensorfloat16dQ_out,GlobalTensorfloat16dK_out,GlobalTensorfloat16cos_cache,GlobalTensorfloat16sin_cache,GlobalTensorint32positions,intB,intnum_heads,intnum_kv_heads,intS,intD){inthead_idxblockIdx.x%num_heads;inttoken_idxblockIdx.x/num_heads;intpospositions[token_idx];inthalfD/2;for(intpthreadIdx.x;phalf;p256){floatcos_valfloat(cos_cache[pos*halfp]);floatsin_valfloat(sin_cache[pos*D/2p]);// 梯度旋转回去用逆旋转矩阵floatdq_xfloat(dQ_out[head_idx*S*Dtoken_idx*Dp]);floatdq_yfloat(dQ_out[head_idx*S*Dtoken_idx*Dhalfp]);floatdxdq_x*cos_valdq_y*sin_val;floatdy-dq_x*sin_valdq_y*cos_val;dQ_out[head_idx*S*Dtoken_idx*Dp]float16(dx);dQ_out[head_idx*S*Dtoken_idx*Dhalfp]float16(dy);// K 同理intkv_head_idxhead_idx;if(num_headsnum_kv_heads)kv_head_idxhead_idx/(num_heads/num_kv_heads);floatdk_xfloat(dK_out[kv_head_idx*S*Dtoken_idx*Dp]);floatdk_yfloat(dK_out[kv_head_idx*S*Dtoken_idx*Dhalfp]);dK_out[kv_head_idx*S*Dtoken_idx*Dp]float16(dk_x*cos_valdk_y*sin_val);dK_out[kv_head_idx*S*Dtoken_idx*Dhalfp]float16(-dk_x*sin_valdk_y*cos_val);}}YaRN另一种外推方法NTK-aware scaling 是放大 base 的方法YaRNYet another RoPE extension method是另一种思路——分段缩放注意力窗口。YaRN 的核心 把 [0, max_train_len] 和 [max_train_len, max_infer_len] 分开处理 [0, max_train_len]: 保持原有 RoPE完整分辨率 [max_train_len, max_infer_len]: 对低频维度额外缩放让有效上下文更长公式新 RoPE 角度 α × θ_i 其中 α 1 (S_train / S_infer - 1) × λ λ 是缩放强度一般取 0.9不全缩只缩低频defyarn_rope(base,train_len,infer_len,d,lambda_0.9):alpha1(train_len/infer_len-1)*lambda_ itorch.arange(0,d//2)theta1.0/(base**(2*i/d))# 只对低频维度i d/4缩放scale_mask(i.float()d/4).float()theta_scaledtheta*(alpha**scale_mask)mtorch.arange(infer_len)anglestorch.outer(m,theta_scaled)returntorch.cos(angles),torch.sin(angles)内存与带宽分析RoPE 操作的内存开销很小预计算缓存max_seq_len32768, D128 cos_cache: 32768 × 64 × 2 bytes 4MB sin_cache: 32768 × 64 × 2 bytes 4MB 总计: 8MB 每层 RoPE 操作HBM 读写 Q: B × num_heads × S × D × 2 bytes K: B × num_kv_heads × S × D × 2 bytes cos_cache 读: S × D/2 × 2 bytes sin_cache 读: S × D/2 × 2 bytes LLaMA-7B (B1, heads32, S4096, D128): Q 读: 1×32×4096×128×2 32MB K 读: 1×32×4096×128×2 32MB cos/sin 读: 4096×64×2×2 1MB 总计: 65MB / layer 32 层: 2,080MB / forward passRoPE 的计算量远小于 GEMM/Softmax只是旋转真正的瓶颈是 HBM 读写——用 on-the-fly 实时计算 cos/sin 省了 1MB 读但多了三角函数开销。预计算缓存方案在延迟上更优。RoPE 的完整实现从数学原理到工程优化前向旋转注入相对位置信息interleaved/non-interleaved 格式必须一致、cos/sin 预计算缓存省 1,696 μs/token、GQA 中 head 索引映射容易越界需要正确除以比例因子、NTK-aware scaling 突破训练长度限制2048→8192 的 base 放大到 40900、YaRN 分段缩放提供另一种外推思路。后向传播只需用负角度再旋转一次——正交矩阵的性质让梯度流动简单。