第一次在昇腾NPU上跑 LLaMA 推理，发现 context 超过 2K 之后，模型开始"胡言乱语"——不是模型不行，是位置编码没处理好。

---

位置编码在干什么？

Transformer 的 Attention 是排列不变的——把输入句子的词序打乱，Attention 的输出不变。

这显然有问题：[“猫”, “吃”, “鱼”] 和 [“鱼”, “吃”, “猫”] 意思完全不同，但 Attention 看不出区别。

位置编码就是给每个 token 打上一个"我是第几个"的标签，让模型能区分词序。

---

RoPE 是怎么打标签的？

RoPE（Rotary Position Embedding，旋转位置编码）的思路很巧妙：

不用加一个位置向量，而是对 Q 和 K 做"旋转"——位置不同，旋转角度不同。

数学上长这样：

Q_m = RoPE(Q, m) = Q × rotation_matrix(m)
K_n = RoPE(K, n) = K × rotation_matrix(n)

然后算 Attention：Q_m · K_n = 内积，这个结果只和相对位置 (m-n) 有关。

也就是：RoPE 天然支持相对位置编码，这是它最大的优势。

---

标准实现的问题

RoPE 的计算看起来简单（就是个旋转矩阵乘法），但在长 context 下是瓶颈：

1. 每个头、每个位置都要算一次旋转矩阵
   LLaMA-2 7B：32 个头 × 4096 个位置 = 13 万次旋转——每次推理都要重算

2. 旋转矩阵是复数运算，标准实现用实数矩阵乘法模拟，慢

3. 不能很好地利用昇腾NPU 的 Cube 核——旋转运算是逐元素的，Cube 核（矩阵乘法）吃不下这么细的活

---

ops-transformer 里的 RoPE 算子优化

1. 预计算 + 查表（高频 cos/sin 不重算）

RoPE 的旋转角度公式是：

θ_i = 1 / (base^(2i/d))

其中 base = 10000（LLaMA 的默认配置）。

关键观察：θ_i 只和 i（维度索引）和 base 有关，和具体位置 m 无关。

ops-transformer 的实现里：

• 把 cos(mθ_i) 和 sin(mθ_i) 做成一个查找表（Look-up Table）
• 每个位置 m 来，直接从表里取，不重算
• 这个表存在 UB（片上内存） 里，不读 HBM

// RoPE 查找表（存在 UB 里，不读 HBM）
__aicore__ void RoPELookup(AscendC::LocalTensor<float> &ubCos,
                           AscendC::LocalTensor<float> &ubSin,
                           int pos, int headDim) {
    // 直接从 UB 表里取 cos(mθ_i) 和 sin(mθ_i)
    // 不重算，不读 HBM
    for (int i = 0; i < headDim / 2; ++i) {
        ubCos(i) = cosTable[pos][i];  // UB 内直接取
        ubSin(i) = sinTable[pos][i];
    }
}

2. 用 Vector 核做逐元素旋转（不占 Cube 核）

RoPE 的本质是：

Q_out[i] = Q_in[i] * cos(mθ_i) - Q_in[i + d/2] * sin(mθ_i)  // 实部
Q_out[i + d/2] = Q_in[i] * sin(mθ_i) + Q_in[i + d/2] * cos(mθ_i)  // 虚部

这是逐元素运算（element-wise），适合 Vector 核，不适合 Cube 核。

ops-transformer 里把 RoPE 的计算拆到 Vector 核上跑，和 Attention 的 Matrix Multiply（Cube 核）并行：

Cube 核：算 Q × K^T（Matrix Multiply）
Vector 核：同时算 RoPE 旋转（Element-wise）

两个核不抢资源，整体吞吐上去。

3. 长 context 的"分块 RoPE"（避免 HBM 读写）

context 很长的时候（8K、16K），RoPE 的查找表也很大（8K × 32 头 × 128 维 = 不小的显存）。

ops-transformer 的实现里：

• 把 context 按 Tile（块） 切分
• 每个 Tile 只加载自己需要的那部分 RoPE 查找表
• 查找表存在 L2 Buffer（片上缓存） 里，不写 HBM

这个优化在 8K context 以上的收益特别明显——标准实现里 RoPE 的 HBM 访问能占到 Prefill 阶段的 15-20%，ops-transformer 的版本压到 5% 以内。

---

实际收益（LLaMA-2 7B，Atlas 300I Duo，Batch=1）

配置	Prefill 吞吐 (tokens/s)	Decode 延迟 (ms/token)	RoPE 占比
标准 RoPE	~1,800	~20	18%
+ ops-transformer RoPE 优化	~2,600	~14	6%
提升幅度	+44%	-30%	-12pp

长 context（8K）下，RoPE 的优化收益更大——因为查找表更大，HBM 访问更频繁，优化效果更明显。

---

代码示例（PyTorch，调用 RoPE）

import torch
import torch_npu

# RoPE 的配置（LLaMA-2 7B）
rope_config = {
    "rope_theta": 10000.0,   # base 值
    "max_seq_len": 4096,
    "head_dim": 128,
}

# 在昇腾NPU上，RoPE 底层走的是 ops-transformer 的优化算子
# 不需要额外配置，CANN 8.0+ 自动识别
q = torch.randn(1, 32, 128, 128).npu()
k = torch.randn(1, 32, 128, 128).npu()
positions = torch.arange(128).npu()

# 调用 RoPE（底层是 ops-transformer 的优化版本）
q_rope = apply_rotary_emb(q, positions, rope_config)
k_rope = apply_rotary_emb(k, positions, rope_config)
# 上面的 apply_rotary_emb 在昇腾NPU上走的是：
#   ops-transformer 的 RoPE Kernel（查找表 + Vector 核 + Tile 切分）

---

一个容易踩的坑

RoPE 的 rope_theta（base 值）要和训练时一致。

训练时如果用了 rope_theta = 10000，推理时改成 500000（为了支持更长的 context），位置编码的尺度就变了，模型会"不认识"这些位置。

正确做法是：

• 推理时要改 context 长度，用线性插值改 RoPE 查找表，不改 rope_theta 的值
• 或者直接用 动态 NTK（Neural Tangent Kernel） 扩展——这个在 ops-transformer 的 rope_ntk_scaling.py 里有参考实现

---

如果你想在自己的模型里用 RoPE，或者想改 RoPE 支持更长的 context，去 ops-transformer 的 ops/rope/ 目录：

https://atomgit.com/cann/ops-transformer

里面有：

• rope_kernel.cpp — RoPE 的 Ascend C 实现（查找表 + Vector 核优化）
• rope_ntk_scaling.py — 动态 NTK 扩展的参考实现
• examples/rope_long_context.py — 跑这个脚本测试 8K/16K context 的 PPL

一句话总结：RoPE 不是"算得快"，而是"不重算、不读 HBM、不抢 Cube 核"——预计算、查表、Vector 核并行，三件事同时做，RoPE 就不会成为瓶颈。

昇腾NPU 上跑 LLaMA，RoPE 的优化在长 context（≥4K）下收益才明显。如果你的 context 很短（<1K），这个优化可能感觉不到——但这不是 RoPE 的问题，是你的 context 还不够长。