RoPERotary Position Embedding在 Llama、Qwen 这些主流大模型里是标配。大部分人写模型的时候RoPE 是一个独立的算子先算旋转角度再对 Q、K 做复数旋转然后才进 Attention。在昇腾NPU上跑这个独立算子每次都要把 Q、K 从 Cube 单元搬出来、做完旋转再搬回去——白跑两趟显存。ops-transformer 仓库里的 RotaryEmbedding 融合算子直接把旋转操作塞进 Attention kernel省掉这两趟搬运。RoPE 的标准实现为什么慢标准 RoPE 的计算过程1. 根据 position 和 freq 计算旋转角度 cos/sin → Vector 单元 2. 把 Q 拆成相邻两两一对 → 内存搬运 3. 对每对做复数旋转[q0, q1] → [q0·cos-q1·sin, q0·sinq1·cos] → Vector 单元 4. 对 K 重复步骤 2-3 → 内存搬运 5. 把旋转后的 Q、K 送入 Attention → Cube 单元步骤 2 和 4 的拆对操作以及 Q/K 在 Vector 和 Cube 之间的来回搬运是性能黑洞。在序列长度 8K、hidden_dim 4096 的场景下Q 和 K 各有 32MB来回搬运一遍就是 128MB 的显存读写。昇腾NPU的 HBM 带宽虽然不低但跟 Cube 单元的计算速度比起来搬运永远是瓶颈。融合算子做了什么把旋转操作嵌入 Attention 的分块计算中。FlashAttention 本身就是把 Q、K、V 切块后在片上缓存计算。RotaryEmbedding 融合算子的做法是在 Q、K 的分块进入 Softmax 之前直接在片上缓存里完成旋转。原来Q_tile → 写回显存 → 读出 → 旋转 → 写回 → 再读出 → Attention 现在Q_tile → 片上缓存直接旋转 → Attention旋转计算本身很简单就是几组乘加操作Vector 单元一个 cycle 就搞定。关键是省掉了那两次显存读写。另外cos/sin 的预计算也做了优化。标准实现里每个序列位置的旋转角度是θ_i 1 / (10000^(2i/d))需要逐位置计算。融合算子把这些值预计算好存在常量缓存里不同 batch 共享同一份。这在小 batch 推理时特别有用——batch1 的场景下预计算的开销从 O(seq_len × head_dim) 降到几乎为零。收益数据在 Atlas 800I A2 上Llama2-7B 的推理性能对比配置首 token 延迟 (ms)吞吐 (tokens/s)RoPE 独立算子 FlashAttention862,840RotaryEmbedding 融合 FlashAttention723,190首 token 延迟降了 16%。这个数字看起来不算炸裂但要知道 RoPE 本身只占 Attention 计算的 5-8%融合带来的收益主要来自减少显存搬运而不是减少计算量。长序列场景收益更明显。序列长度从 2K 拉到 16K融合算子的延迟增长是线性的独立算子因为显存搬运量翻倍延迟增长接近 O(N^1.3)。怎么用PyTorch 场景下不需要改模型代码CANN 的框架适配器自动处理importtorch_npu# CANN 自动把 RoPE SDPA 路由到融合算子# 前提模型用 F.scaled_dot_product_attention 接口q_rotatedapply_rotary_emb(q,cos,sin)# 这里会被融合进 Attentionouttorch.nn.functional.scaled_dot_product_attention(q_rotated,k_rotated,v)如果你手动调用了torch_npu.npu.flash_attention需要显式传入 cos/sinouttorch_npu.npu.flash_attention(q,k,v,coscos_tensor,# 旋转角度余弦提前算好sinsin_tensor,# 旋转角度正弦rotary_mode1# 1融合旋转0不融合)踩坑提醒融合算子要求 cos/sin 的 dtype 和 Q 一致。Llama 官方权重里的 RoPE 参数是 float32但推理时 Q 通常转成了 float16。dtype 不匹配不会报错——它会默默 fallback 到独立算子。检查日志里有没有rotary_fusion_fallback关键字。还有一个融合算子目前不支持动态序列长度。如果你的推理服务支持变长输入比如 batch 里不同样本长度不同需要 padding 到统一长度再传入。CANN 8.5 加了变长支持但只限 ATB 路径PyTorch 路径还没跟上。如果你的 Llama 推理服务首 token 延迟偏高先看 RoPE 是不是独立算子在跑。融合方式很简单但 fallback 是静默的不查日志你根本不知道。仓库在这里https://atomgit.com/cann/ops-transformer