CANN ops-transformer 旋转位置编码融合算子 aclnnApplyRotaryPosEmbV2 使用指南

• 算子库
• 人工智能
• 深度学习
• Ascend

导读

aclnnApplyRotaryPosEmbV2 是 CANN ops-transformer 仓库中面向推理场景的旋转位置编码（RoPE，Rotary Position Embedding）融合算子接口，它将 Transformer 网络中的 query 与 key 两路位置编码计算融合为一路，并在原内存上执行原地更新，从而减少 kernel 启动与内存搬运开销。本文基于 aclnnApplyRotaryPosEmbV2.md 展开，结合仓库内的算子定义、tiling 实现与可编译示例源码，完整讲解该接口的产品支持情况、三种旋转模式的计算公式、两段式接口的完整参数约束、错误码语义以及可直接编译运行的调用示例。读完本文，你将掌握如何在 Ascend 推理与训练系列产品上正确选择并调用该接口完成 RoPE 计算，并理解其底层 UB 空间约束的由来。

接口定位与适用场景

功能定位

推理网络中，attention 计算前需要分别对 query 和 key 施加旋转位置编码。常规做法是调用两个独立的算子分别完成，而本接口将两路算子融合成一路：一次 kernel 启动同时处理 query 与 key，并将计算结果写回原张量（原地更新），显著减少中间张量生成与内存读写。

• 接口功能：执行旋转位置编码计算，query、key 两路融合，计算结果原地更新。
• 原始输入 query、key 同时作为输出承载计算结果（公式中的q_embed、k_embed）。

与 aclnnApplyRotaryPosEmb 的差异

本接口是 aclnnApplyRotaryPosEmb 的功能扩展版本，核心差异为新增 rotaryMode 参数，用于控制不同的旋转编码方式，使同一算子可以适配不同模型对 RoPE 切分方式的实现差异：

• 新增rotaryMode（旋转模式）参数，取值范围"half"、"interleave"、"quarter"；
• 原接口仅支持 "half" 类按半切分旋转，V2 版本在同一实现上支持三种切分/交织模式。

选择建议：如果你的模型只使用标准的 half 式旋转且不需要新模式，可继续使用原接口；需要 interleave/quarter 模式或在 Ascend 950 上使用更多 layout 时，选择本 V2 接口。

产品支持情况

补充说明（依据 op_host/apply_rotary_pos_emb_def.cpp 的 AICore 配置）：除上述产品外，仓库还针对kirinx90、kirin9030处理器系列注册了 AICore 配置（其约束与 Atlas 推理系列产品一致，不支持 BFLOAT16）；而ascend910b（Atlas A2）、ascend910_93（Atlas A3）与ascend950、mc62使用不同的 tiling/regbase 配置。可结合实际目标硬件选择对应配置。

旋转位置编码的计算公式

接口通过rotaryMode区分三种旋转编码方式。以下公式中，q_embed = (query * cos) + query_rotate * sin，k_embed = (key * cos) + key_rotate * sin，只是query_rotate/key_rotate的构造方式不同。三种模式的具体切分规则如下。

rotaryMode 为 "half"：按半切分后交叉

将最后一维从中间一分为二，前一半与后一半取负后交叉拼接：

$$query_q1 = query[..., : query.shape[-1] // 2]$$

$$query_q2 = query[..., query.shape[-1] // 2 :]$$

$$query_rotate = torch.cat((-query_q2, query_q1), dim=-1)$$

$$key_k1 = key[..., : key.shape[-1] // 2]$$

$$key_k2 = key[..., key.shape[-1] // 2 :]$$

$$key_rotate = torch.cat((-key_k2, key_k1), dim=-1)$$

$$q_embed = (query * cos) + query_rotate * sin$$

$$k_embed = (key * cos) + key_rotate * sin$$

这是大多数 RoPE 实现（如 HF 风格的rotate_half）采用的模式。

rotaryMode 为 "quarter"：四等分交叉

将最后一维均分为四段，第 2、4 段取负后按 (段2, 段1, 段4, 段3) 顺序拼接：

$$query_q1 = query[..., : query.shape[-1] // 4]$$

$$query_q2 = query[..., query.shape[-1] // 4 : query.shape[-1] // 2]$$

$$query_q3 = query[..., query.shape[-1] // 2 : query.shape[-1] // 4 * 3]$$

$$query_q4 = query[..., query.shape[-1] // 4 * 3 :]$$

$$query_rotate = torch.cat((-query_q2, query_q1, -query_q4, query_q3), dim=-1)$$

$$key_q1 = key[..., : key.shape[-1] // 4]$$

$$key_q2 = key[..., key.shape[-1] // 4 : key.shape[-1] // 2]$$

$$key_q3 = key[..., key.shape[-1] // 2 : key.shape[-1] // 4 * 3]$$

$$key_q4 = key[..., key.shape[-1] // 4 * 3 :]$$

$$key_rotate = torch.cat((-key_q2, key_q1, -key_q4, key_q3), dim=-1)$$

$$q_embed = (query * cos) + query_rotate * sin$$

$$k_embed = (key * cos) + key_rotate * sin$$

rotaryMode 为 "interleave"：奇偶交织

按奇偶位置取元素，负的偶数位元素与奇数位元素交错拼接回原 shape：

$$query_q1 = query[..., ::2].view(-1, 1)$$

$$query_q2 = query[..., 1::2].view(-1, 1)$$

$$query_rotate = torch.cat((-query_q2, query_q1), dim=-1).view(query.shape[0], query.shape[1], query.shape[2], query.shape[3])$$

$$key_q1 = key[..., ::2].view(-1, 1)$$

$$key_q2 = key[..., 1::2].view(-1, 1)$$

$$key_rotate = torch.cat((-key_q2, key_q1), dim=-1).view(key.shape[0], key.shape[1], key.shape[2], key.shape[3])$$

$$q_embed = (query * cos) + query_rotate * sin$$

$$k_embed = (key * cos) + key_rotate * sin$$

注意：以上公式以 PyTorch 语法表达切分与拼接语义，帮助理解三种模式的区别；kernel 内部通过向量指令实现等价计算，并非真的逐段 view。

三种模式支持的硬件范围

• Atlas 推理系列产品 / Atlas A2 / Atlas A3：仅支持"half"模式；
• Ascend 950PR / Ascend 950DT：支持"half"、"interleave"、"quarter"全部三种模式。

同时，"half"与"interleave"要求输入最后一维能被 2 整除，"quarter"要求最后一维能被 4 整除。

两段式接口与函数原型

与 CANN 大多数单算子接口一致，本算子采用两段式接口调用模型：先调用aclnnApplyRotaryPosEmbV2GetWorkspaceSize完成入参校验并计算所需 workspace 大小，再调用aclnnApplyRotaryPosEmbV2执行计算。

aclnnStatus aclnnApplyRotaryPosEmbV2GetWorkspaceSize( aclTensor *queryRef, aclTensor *keyRef, const aclTensor *cos, const aclTensor *sin, int64_t layout, char *rotaryMode, uint64_t *workspaceSize, aclOpExecutor **executor)

aclnnStatus aclnnApplyRotaryPosEmbV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

从源码看，V2 接口是对内部统一实现的薄封装：op_host/op_api/aclnn_apply_rotary_pos_emb_v2.cpp 中，两段接口分别透传给aclnnInnerApplyRotaryPosEmbGetWorkspaceSize与aclnnInnerApplyRotaryPosEmb，算子注册在 apply_rotary_pos_emb_def.cpp 中完成（算子名 ApplyRotaryPosEmb，输入 query/key/cos/sin，输出 query/key，属性 layout 默认 1、rotary_mode 默认 "half"）。

aclnnApplyRotaryPosEmbV2GetWorkspaceSize 参数详解

参数说明

四个输入 Tensor 均允许非连续（strides 非紧凑）传入，算子内部按 ND 格式解析。

分产品输入约束

queryRef / keyRef：

• Atlas 推理系列产品、Atlas A2、Atlas A3：不支持空 Tensor；shape 最后一维（D）必须等于 128 或 64。
• Ascend 950PR / Ascend 950DT：支持空 Tensor；shape 最后一维（D）小于等于 1024。

cos / sin：

• Atlas 推理系列产品、Atlas A2、Atlas A3：不支持空 Tensor；B 维度与 queryRef、keyRef 一致；第 3 维（N）必须等于 1；最后一维（D）必须等于 128 或 64。
• Ascend 950PR / Ascend 950DT：支持空 Tensor；B 维度与 queryRef、keyRef 一致或等于 1（支持广播）；N 维度必须等于 1；最后一维（D）小于等于 1024。

layout 参数：

• 取值范围：1-BSND、2-SBND、3-BNSD、4-TND。
• Atlas 推理系列产品、Atlas A2、Atlas A3：支持 1-BSND（4 维）与 4-TND（3 维）。
• Ascend 950PR / Ascend 950DT：支持 1-BSND、2-SBND、3-BNSD 的 4 维 Tensor 与 4-TND 的 3 维 Tensor。

rotaryMode 参数：取值范围"half"、"interleave"、"quarter"；Atlas 推理系列产品、Atlas A2、Atlas A3 仅支持"half"，Ascend 950PR / Ascend 950DT 支持全部三种。

注意：Atlas 推理系列产品不支持 BFLOAT16 输入。

返回值与错误码

第一段接口返回aclnnStatus，出现以下场景时报错：

完整返回码语义参见 aclnn 返回码。

aclnnApplyRotaryPosEmbV2 参数详解

第二段接口参数说明：

返回aclnnStatus状态码，语义见 aclnn 返回码。

约束说明与 UB 空间预算

确定性计算

aclnnApplyRotaryPosEmbV2 为默认确定性实现，即相同输入多次执行结果一致，便于调试与精度对齐（相关概念可参考 确定性计算）。

Atlas 推理系列产品 / Atlas A2 / Atlas A3 约束

• layout 为 1 时，queryRef、keyRef、cos、sin 前 2 维（B、S）必须相等；layout 为 4 时，第 1 维（T）必须相等。
• queryRef、keyRef 最后一维（D）必须相等；cos、sin 最后一维（D）必须相等。
• 四个输入 Tensor 的 dtype 必须相同。

shape 记法：layout=1 时 queryRef 为(q_b, q_s, q_n, q_d)，keyRef 为(q_b, q_s, k_n, q_d)，cos/sin 为(q_b, q_s, 1, cos_d)；layout=4 时 queryRef 为(q_t, q_n, q_d)，keyRef 为(q_t, k_n, q_d)，cos/sin 为(q_t, 1, cos_d)。其中 b 为 batch_size、s 为 seq_length、t 为 b 与 s 合轴、n 为 head_num、d 为 head_dim。

UB（Unified Buffer）空间预算：设lastDim为需要旋转位置编码的维度长度，根据 dtype 决定 cast 参数：

• 输入为 BFLOAT16 时，cast 为 1，castSize 为 4，DtypeSize 为 2；
• 输入为 FLOAT16 或 FLOAT32 时，cast 为 0，castSize = DtypeSize（FLOAT16 为 2，FLOAT32 为 4）。

所需 UB 空间估算公式：

ub_required = (q_n + k_n) * lastDim * castSize * 2 // 搬运缓冲（ping-pong） + lastDim * DtypeSize * 4 + (q_n + k_n) * lastDim * castSize // 计算缓冲 + (q_n + k_n) * lastDim * castSize * 2 + cast * (lastDim * 4 * 2) // BF16 转 FP32 的额外开销

当ub_required超过当前 AI 处理器的 UB 空间总大小时，不支持使用该融合算子。这与 tiling 实现中 apply_rotary_pos_emb_tiling.cpp 的 UB 分配逻辑一致：源码中按qPart1Ub（搬运(Q_n + K_n) * D * castDtypeSize）、cosPart1Ub、q2q1Part1Ub（计算）、sin1UbSize（BF16 cast 场景）等分量累加得到speUb，再与平台totalUbSize比较，决定是否采用小 shape 单核方案或进入多核切分路径；平台 UB 大小通过platform_ascendc接口查询（如GetCoreMemSize(UB)）。

Ascend 950PR / Ascend 950DT 约束

• 任意 layout 下，queryRef 与 keyRef 除 N 维度外其他维度必须相同。
• queryRef、keyRef 最后一维（D）必须相等；cos、sin 最后一维（D）必须相等，且小于等于 queryRef、keyRef 的最后一维（D）。
• 四个输入 Tensor 的 dtype 必须相同。
• rotaryMode 为 "half"/"interleave" 时，最后一维必须能被 2 整除；"quarter" 时最后一维必须能被 4 整除。

调用示例（完整可编译）

以下示例完整展示从 acl 初始化、构造 aclTensor、两段式调用到结果回拷与资源释放的完整流程，代码可直接参照 examples/test_aclnn_apply_rotary_pos_emb_v2.cpp（仓库中另有 test_aclnn_apply_rotary_pos_emb.cpp 展示原接口调用方式）。编译与运行步骤参考 编译与运行样例。

#include "acl/acl.h" #include "aclnnop/aclnn_apply_rotary_pos_emb_v2.h" #include <iostream> #include <vector> #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t>& shape) { int64_t shape_size = 1; for (auto i : shape) { shape_size *= i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法，资源初始化 auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); ret = aclrtSetDevice(deviceId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return 0; } template <typename T> int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size = GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); // 计算连续tensor的strides std::vector<int64_t> strides(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法，device/stream初始化，参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId = 0; aclrtStream stream; auto ret = Init(deviceId, &stream); CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); // 2. 构造输入与输出，需要根据API的接口定义构造 // 注意：Atlas 推理/A2/A3 上 D 必须为 128 或 64；Ascend 950 上 D <= 1024 std::vector<int64_t> queryShape = {1, 1, 1, 128}; // BSND std::vector<int64_t> keyShape = {1, 1, 1, 128}; std::vector<int64_t> cosShape = {1, 1, 1, 128}; // N 维必须为 1 std::vector<int64_t> sinShape = {1, 1, 1, 128}; int64_t layout = 1; // 1-BSND char rotaryMode[] = "half"; // "half" / "interleave" / "quarter" void* queryDeviceAddr = nullptr; void* keyDeviceAddr = nullptr; void* cosDeviceAddr = nullptr; void* sinDeviceAddr = nullptr; aclTensor* query = nullptr; aclTensor* key = nullptr; aclTensor* cos = nullptr; aclTensor* sin = nullptr; // host 侧测试数据（长度 128 的 FLOAT 序列，此处省略中间值，完整数据见仓库示例文件） std::vector<float> queryHostData(128, 1.0f); std::vector<float> keyHostData(128, 1.0f); std::vector<float> cosHostData(128, 1.0f); std::vector<float> sinHostData(128, 1.0f); // 创建 query / key / cos / sin 四个 aclTensor ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT, &query); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(keyHostData, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT, &key); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(cosHostData, cosShape, &cosDeviceAddr, aclDataType::ACL_FLOAT, &cos); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(sinHostData, sinShape, &sinDeviceAddr, aclDataType::ACL_FLOAT, &sin); CHECK_RET(ret == ACL_SUCCESS, return ret); // 3. 调用CANN算子库API（两段式） uint64_t workspaceSize = 0; aclOpExecutor* executor; // 第一段接口：入参校验 + 计算workspace大小 ret = aclnnApplyRotaryPosEmbV2GetWorkspaceSize(query, key, cos, sin, layout, rotaryMode, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnApplyRotaryPosEmbV2GetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret;); } // 第二段接口：执行计算 ret = aclnnApplyRotaryPosEmbV2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnApplyRotaryPosEmbV2 failed. ERROR: %d\n", ret); return ret); // 4. 固定写法，同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 5. 获取输出：原地更新，结果在 query/key 的 device 内存上，回拷到 host 侧 auto size = GetShapeSize(queryShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), queryDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("query result[%ld] is: %f\n", i, resultData[i]); } auto size1 = GetShapeSize(keyShape); std::vector<float> resultData1(size1, 0); ret = aclrtMemcpy(resultData1.data(), resultData1.size() * sizeof(resultData1[0]), keyDeviceAddr, size1 * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size1; i++) { LOG_PRINT("key result[%ld] is: %f\n", i, resultData1[i]); } // 6. 释放aclTensor aclDestroyTensor(query); aclDestroyTensor(key); aclDestroyTensor(cos); aclDestroyTensor(sin); // 7. 释放device资源 aclrtFree(queryDeviceAddr); aclrtFree(keyDeviceAddr); aclrtFree(cosDeviceAddr); aclrtFree(sinDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例要点：

• 示例采用 BSND 布局、FLOAT 类型、"half"模式，shape 为{1, 1, 1, 128}，满足 Atlas 推理/A2/A3 与 Ascend 950 两侧的约束；
• 结果是原地更新的，因此查询结果时直接回拷 queryDeviceAddr / keyDeviceAddr 即可，无需额外申请输出张量；
• 若要在 Ascend 950 上使用"quarter"/"interleave"模式，只需修改rotaryMode，并确保 D 能被 2/4 整除；
• 测试数据、随机初始化细节与完整 128 元素数据序列可在仓库示例文件中查看。

仓库中的配套实现与验证资源

• 接口实现：op_host/op_api/aclnn_apply_rotary_pos_emb_v2.cpp 与 op_host/op_api/aclnn_apply_rotary_pos_emb_v2.h；
• 算子定义：op_host/apply_rotary_pos_emb_def.cpp（属性 layout 默认 1、rotary_mode 默认 "half"，各产品 AICore 配置）；
• shape 推导：op_host/apply_rotary_pos_emb_infershape.cpp（含 cos/sin 与 query/key 之间的广播维度推导逻辑）；
• tiling 实现：op_host/apply_rotary_pos_emb_tiling.cpp 及 arch35 系列分路径实现（apply_rotary_pos_emb_tiling_ab_arch35.cpp、apply_rotary_pos_emb_tiling_aba_and_ba_arch35.cpp、apply_rotary_pos_emb_tiling_bab_arch35.cpp）；
• kernel 实现：op_kernel/apply_rotary_pos_emb.cpp、op_kernel/apply_rotary_pos_emb_base.h 及 arch35 下的分模式 kernel；
• 图融合 pass：op_graph/fusion_pass/apply_rotary_pos_emb_tensormove_pass.cpp，体现该算子在图优化层面与其他算子的融合路径；
• 测试与数据生成：tests/ut/op_kernel/apply_rotary_pos_emb_data/gen_data.py、tests/ut/op_host/test_apply_rotary_pos_emb_infershape.cpp、tests/ut/op_host/test_apply_rotary_pos_emb_tiling.cpp、tests/ut/op_host/op_api/test_aclnn_apply_rotary_pos_emb.cpp。

总结

aclnnApplyRotaryPosEmbV2 将 query/key 两路旋转位置编码融合为单次 kernel 执行并原地更新，通过新增的 rotaryMode 参数覆盖 "half"、"interleave"、"quarter" 三种主流 RoPE 实现方式。调用前务必核对目标产品的支持矩阵、shape 与 dtype 约束（尤其 D 的取值、N 维为 1、四个输入 dtype 一致），并遵循两段式接口先计算 workspace 再执行计算的流程。在 Atlas 推理/A2/A3 上仅支持 half 模式且 D 限定为 64/128；Ascend 950 上支持全部三种模式、四种 layout 与空 Tensor 场景，使用时还需关注ub_required是否超出硬件 UB 空间。结合仓库提供的示例与测试用例，可以快速完成该算子在 NPU 上的集成与验证。

• 算子库
• 人工智能
• 深度学习
• Ascend