CANN ops-transformer 旋转位置编码算子 RopeWithSinCosCache 深度解析：cache 化 RoPE 的计算原理与 aclnn 两段式调用实战

CANN ops-transformer 旋转位置编码算子 RopeWithSinCosCache 深度解析：cache 化 RoPE 的计算原理与 aclnn 两段式调用实战

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

导读：本文聚焦 CANN ops-transformer 仓库中的旋转位置编码算子RopeWithSinCosCache（位于 posembedding/rope_with_sin_cos_cache），该算子是推理网络中用于将预计算的 sin/cos 位置编码表（cache）直接应用到 query/key 上的高性能实现。读完本文，你将完整掌握该算子的 rope/mrope 两种计算模式、GPT-NeoX 与 GPT-J 两种旋转风格的计算公式、全部入参约束，以及基于aclnnRopeWithSinCosCache两段式接口的落地调用方法，并能结合仓库源码理解其参数校验、多核 tiling 与 kernel 分派原理。

产品支持情况

依据 posembedding/rope_with_sin_cos_cache/README.md，该算子对产品的支持情况如下：

需要补充说明的是：

• 在 op_host/rope_with_sin_cos_cache_def.cpp 的算子定义中，RopeWithSinCosCache实际为ascend910b、ascend910_93、ascend950、kirinx90、kirin9030五个平台配置了 AICore 算子实现（其中ascend950使用独立的rope_with_sin_cos_cache_apt实现，Kirin 平台使用独立配置），与产品支持矩阵一致。
• Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16，其算子配置中 queryIn/keyIn/cosSinCache 仅支持DT_FLOAT与DT_FLOAT16（见 rope_with_sin_cos_cache_def.cpp）。
• 带cacheMode扩展参数的 V2 接口aclnnRopeWithSinCosCacheV2目前不支持 Ascend 950PR/Ascend 950DT（见 docs/aclnnRopeWithSinCosCacheV2.md）。

功能说明：为推理性能设计的 cache 化旋转位置编码

设计动机

旋转位置编码（RoPE）在 Transformer 推理网络中广泛应用。常规实现中，每个 token 的 sin/cos 值需要在计算时临时生成；而在推理场景下，位置编码值往往是可预先计算并缓存的。RopeWithSinCosCache 算子正是将 sin 和 cos 输入通过 cache（cosSinCache）传入，直接按位置索引查表并执行旋转位置编码计算，从而省去重复计算，提升推理性能。

整体计算分三步：

1. 按位置索引positions从cosSinCache中取出对应的cosSin张量；
2. 将cosSin沿最后一维切分为cos与sin两半；
3. 将query/key拆分为需要旋转的部分queryRot（前rotaryDim维）和直接透传的部分queryPass（其余维度），对queryRot应用旋转，最后拼回。

mrope 模式：多模态位置编码

mrope 模式下，positions的 shape 为[m, numTokens]，其中 m 为mropeSection的元素个数（支持 3 或 4），用于多模态（如图像、视频、文本混合）场景下对不同位置使用不同的位置编码频率：

cosSin[i] = cosSinCache[positions[i]] cos, sin = cosSin.chunk(2, dim=-1)

当mropeSection元素个数为3时，按三段拼接cos与sin：

cos0 = cos[0, :, :mropeSection[0]] cos1 = cos[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])] cos2 = cos[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])] cos = torch.cat((cos0, cos1, cos2), dim=-1) sin0 = sin[0, :, :mropeSection[0]] sin1 = sin[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])] sin2 = sin[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])] sin = torch.cat((sin0, sin1, sin2), dim=-1)

当mropeSection元素个数为4时，使用更简洁的 split 形式表达同一语义：

cos = torch.cat([m[i] for i, m in enumerate(cos.split(mropeSection, dim=-1))], dim=-1) sin = torch.cat([m[i] for i, m in enumerate(sin.split(mropeSection, dim=-1))], dim=-1)

之后统一进入旋转计算（与 rope 模式共用）：

queryRot = query[..., :rotaryDim] queryPass = query[..., rotaryDim:]

rope 模式：标准位置编码

rope 模式下，positions的 shape 为[numTokens]（一维索引），直接查表并切分：

cosSin[i] = cosSinCache[positions[i]] cos, sin = cosSin.chunk(2, dim=-1) queryRot = query[..., :rotaryDim] queryPass = query[..., rotaryDim:]

说明：在 aclnn 接口中，mropeSection入参传nullptr即表示 rope 模式；传入非空mropeSection即开启 mrope 模式（详见 docs/aclnnRopeWithSinCosCache.md 参数说明）。从源码看，kernel 中以mropeSection[0] > 0作为进入 mrope 分支的判断条件（见 op_host/op_api/aclnn_rope_with_sin_cos_cache.cpp）。

rotate_half（GPT-NeoX style）旋转模式

当isNeoxStyle=true时采用该模式，先将queryRot沿最后一维对半切分，再交叉相乘后拼接：

x1, x2 = torch.chunk(queryRot, 2, dim=-1) o1[i] = x1[i] * cos[i] - x2[i] * sin[i] o2[i] = x2[i] * cos[i] + x1[i] * sin[i] queryRot = torch.cat((o1, o2), dim=-1) query = torch.cat((queryRot, queryPass), dim=-1)

rotate_interleaved（GPT-J style）旋转模式

当isNeoxStyle=false时采用该模式，按奇偶位交错拆分并旋转：

x1 = queryRot[..., ::2] x2 = queryRot[..., 1::2] o1[i] = x1[i] * cos[i] - x2[i] * sin[i] o2[i] = x2[i] * cos[i] + x1[i] * sin[i] queryRot = torch.stack((o1, o2), dim=-1) query = torch.cat((queryRot, queryPass), dim=-1)

对keyIn执行完全相同的旋转流程，得到keyOut。

参数说明

下表完整摘录自 posembedding/rope_with_sin_cos_cache/README.md：

补充要点：

• Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16（该说明位于 README 参数表下方）。
• 在算子定义 rope_with_sin_cos_cache_def.cpp 中，除上述接口入参外，底层算子还包含numQHeads、numKHeads、headSize、mropeSection、qStride、kStride、isNeoxStyle、cacheMode等属性，其中qStride/kStride表示 query/key 行间步长（用于支持非连续输入），cacheMode默认为 0。

约束说明

调用该算子前需满足以下约束（源自 README 与 docs/aclnnRopeWithSinCosCache.md 的约束章节）：

1. queryIn、keyIn、cosSinCache只支持2 维shape 输入；其中queryIn/keyInshape 为(numTokens, numQHeads*headSize)/(numTokens, numKHeads*headSize)，cosSinCacheshape 为(maxSeqLen, rotaryDim)，maxSeqLen表示模型处理序列的最大长度，rotaryDim表示旋转位置嵌入的维度大小。
2. queryIn、keyIn、cosSinCache的数据类型需要保持一致，且queryOut/keyOut与对应输入数据类型一致。
3. headSize：数据类型为 BFLOAT16 或 FLOAT16 时为32 的倍数，数据类型为 FLOAT32 时为16 的倍数（等价于 64 Byte 对齐，tiling 侧按headSize % (64 / 元素字节数) == 0校验）。
4. rotaryDim：始终小于等于headSize；BFLOAT16/FLOAT16 时为 32 的倍数，FLOAT32 时为 16 的倍数；mrope 模式下应满足rotaryDim = mropeSection所有元素之和的 2 倍，即mropeSection[0] + mropeSection[1] + mropeSection2 = rotaryDim / 2。
5. 输入 tensorpositions的取值应小于cosSinCache的 0 维maxSeqLen。
6. aclnnRopeWithSinCosCache为默认确定性实现。
7. mropeSection 取值限制：README 标注为[16, 24, 24]；而 docs/aclnnRopeWithSinCosCache.md 与 V2 接口文档给出更完整的白名单：[16, 24, 24]、[24, 20, 20]、[8, 12, 12]、[16, 16, 16, 16]，V2 接口另支持[11, 11, 10]。源码中的mropeSupportList校验了全部 5 种取值（见 aclnn_rope_with_sin_cos_cache.cpp）。

调用方式：aclnn 两段式接口

该算子通过 CANN 的 aclnn（AscendCL NN）两段式接口调用：必须先调用aclnnRopeWithSinCosCacheGetWorkspaceSize获取 workspace 大小与执行器，再调用aclnnRopeWithSinCosCache执行计算。

接口原型

aclnnStatus aclnnRopeWithSinCosCacheGetWorkspaceSize( const aclTensor *positions, const aclTensor *queryIn, const aclTensor *keyIn, const aclTensor *cosSinCache, const aclIntArray *mropeSection, int64_t headSize, bool isNeoxStyle, aclTensor *queryOut, aclTensor *keyOut, uint64_t *workspaceSize, aclOpExecutor **executor);

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

第一段接口参数详解

返回值与错误码

第一段接口完成入参校验，出现以下场景时报错：

完整调用示例

以下为完整可运行流程（关键步骤摘录，完整样例见 examples/test_aclnn_rope_with_sin_cos_cache.cpp）。示例为 rope 模式（mropeSection传nullptr），采用 BFLOAT16 类型：

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_rope_with_sin_cos_cache.h" #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) // 计算连续 tensor 的 strides 并用 aclCreateTensor 构造 aclTensor template <typename T> int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr, aclDataType dataType, aclTensor **tensor) { auto size = 1; for (auto i : shape) size *= i; size *= sizeof(T); auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, return ret); 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]; } *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 初始化（固定写法） int32_t deviceId = 0; aclrtStream stream; aclInit(nullptr); aclrtSetDevice(deviceId); aclrtCreateStream(&stream); // 2. 构造输入与输出（rope 模式：positions 为 1 维索引） std::vector<int64_t> positionsShape = {2}; // numTokens = 2 std::vector<int64_t> queryInShape = {2, 64}; // numQHeads*headSize = 2*32 std::vector<int64_t> keyInShape = {2, 64}; std::vector<int64_t> cosSinCacheShape = {2, 32}; // (maxSeqLen, rotaryDim) std::vector<int64_t> queryOutShape = {2, 64}; std::vector<int64_t> keyOutShape = {2, 64}; void* positionsDeviceAddr = nullptr; void* queryInDeviceAddr = nullptr; void* keyInDeviceAddr = nullptr; void* cosSinCacheDeviceAddr = nullptr; void* queryOutDeviceAddr = nullptr; void* keyOutDeviceAddr = nullptr; aclTensor* positions = nullptr; aclTensor* queryIn = nullptr; aclTensor* keyIn = nullptr; aclTensor* cosSinCache = nullptr; aclTensor* queryOut = nullptr; aclTensor* keyOut = nullptr; int64_t headSize = 32; // BF16 下需为 32 的倍数 bool isNeoxStyle = true; // true: GPT-NeoX style std::vector<int64_t> positionsHostData = {0, 1}; std::vector<float> queryInHostData(128, 1.0f); // 实际请填充真实数据 std::vector<float> keyInHostData(128, 1.0f); std::vector<float> cosSinCacheHostData(64, 1.0f); // 长度 = 2*32 std::vector<float> queryOutHostData(128, 0.0f); std::vector<float> keyOutHostData(128, 0.0f); CreateAclTensor(positionsHostData, positionsShape, &positionsDeviceAddr, aclDataType::ACL_INT64, &positions); CreateAclTensor(queryInHostData, queryInShape, &queryInDeviceAddr, aclDataType::ACL_BF16, &queryIn); CreateAclTensor(keyInHostData, keyInShape, &keyInDeviceAddr, aclDataType::ACL_BF16, &keyIn); CreateAclTensor(cosSinCacheHostData, cosSinCacheShape, &cosSinCacheDeviceAddr, aclDataType::ACL_BF16, &cosSinCache); CreateAclTensor(queryOutHostData, queryOutShape, &queryOutDeviceAddr, aclDataType::ACL_BF16, &queryOut); CreateAclTensor(keyOutHostData, keyOutShape, &keyOutDeviceAddr, aclDataType::ACL_BF16, &keyOut); // 3. 两段式调用：第一段获取 workspace 与执行器 uint64_t workspaceSize = 0; aclOpExecutor* executor = nullptr; auto ret = aclnnRopeWithSinCosCacheGetWorkspaceSize( positions, queryIn, keyIn, cosSinCache, nullptr /* mropeSection, rope模式 */, headSize, isNeoxStyle, queryOut, keyOut, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, return ret); // 4. 按需申请 workspace void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, return ret); } // 5. 第二段接口执行计算 ret = aclnnRopeWithSinCosCache(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, return ret); // 6. 同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, return ret); // 7. 释放资源 aclDestroyTensor(positions); aclDestroyTensor(queryIn); aclDestroyTensor(keyIn); aclDestroyTensor(cosSinCache); aclDestroyTensor(queryOut); aclDestroyTensor(keyOut); aclrtFree(positionsDeviceAddr); aclrtFree(queryInDeviceAddr); aclrtFree(keyInDeviceAddr); aclrtFree(cosSinCacheDeviceAddr); aclrtFree(queryOutDeviceAddr); aclrtFree(keyOutDeviceAddr); if (workspaceSize > 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

要点提醒：

• 若开启 mrope 模式，只需将示例中的nullptr替换为按aclIntArray构造的mropeSection（如{16, 24, 24}或{16, 16, 16, 16}），并将positions构造为(3, numTokens)或(4, numTokens)的二维张量。
• 示例中queryIn等输入张量支持非连续（stride 不连续）Tensor，接口内部会先做Contiguous处理；queryOut/keyOut不支持非连续（见 aclnn_rope_with_sin_cos_cache.cpp 中l0op::Contiguous与l0op::ViewCopy的组合使用）。

V2 接口：cacheMode 扩展

仓库同时提供了扩展接口aclnnRopeWithSinCosCacheV2（文档见 docs/aclnnRopeWithSinCosCacheV2.md，样例见 examples/test_aclnn_rope_with_sin_cos_cache_v2.cpp），与 V1 相比新增cacheMode参数，用于指示拼接 cos/sin 的方式：

• cacheMode = 0：与aclnnRopeWithSinCosCache实现相同，为分段式拼接 cos 和 sin（即上文 mrope 三段/四段 cat 的逻辑）。
• cacheMode = 1：为交错式拼接 cos 和 sin。以三段为例：

cosTmp = cos cos[..., 1:mropeSection[1] * 3:3] = cosTmp[1, ..., 1:mropeSection[1] * 3:3] cos[..., 2:mropeSection[1] * 3:3] = cosTmp[2, ..., 2:mropeSection[1] * 3:3] sinTmp = sin sin[..., 1:mropeSection[1] * 3:3] = sinTmp[1, ..., 1:mropeSection[1] * 3:3] sin[..., 2:mropeSection[1] * 3:3] = sinTmp[2, ..., 2:mropeSection[1] * 3:3]

V2 的约束差异：mrope 模式下cacheMode仅支持 0 和 1，且当mropeSection为[16, 16, 16, 16]时仅支持 0（tiling 侧同样做了该校验，见 rope_with_sin_cos_cache_tiling.cpp）。源码中 V1 接口内部固定传入cacheMode = 0，V2 才真正暴露该参数（见 aclnn_rope_with_sin_cos_cache.cpp）。

源码实现原理

算子定义与属性

在 op_host/rope_with_sin_cos_cache_def.cpp 中，RopeWithSinCosCache通过OpDef注册：

• 4 个必选输入（positions、queryIn、keyIn、cosSinCache）与 2 个必选输出（queryOut、keyOut），数据类型均为DT_FLOAT / DT_FLOAT16 / DT_BF16（positions 为 INT64），格式均为 ND；
• 属性包括必选的numQHeads、numKHeads、headSize，可选的mropeSection（默认{0, 0, 0}）、qStride、kStride、isNeoxStyle（默认 true）、cacheMode（默认 0）；
• 各平台配置均开启动态编译（DynamicCompileStaticFlag）、动态 rank 与动态 shape 支持。

入参校验逻辑

op_host/op_api/aclnn_rope_with_sin_cos_cache.cpp 中CheckParams依次完成四类校验：

1. 空指针检查（CheckNotNull）：6 个 tensor 均非空；
2. 数据类型检查（CheckDtypeValid）：positions 必须为 INT64，keyIn/cosSinCache/queryOut/keyOut的数据类型须与queryIn一致，且整体在支持列表内；
3. shape 检查（CheckShape）：三个输入 tensor 均为 2 维、均非空、keyIn.shape[0] == queryIn.shape[0]、输出与输入 shape 相等；
4. mropeSection 校验：元素个数必须为 3 或 4、元素非负、当mropeSection[0] > 0（mrope 模式）时要求所有元素之和的 2 倍等于rotaryDim（即cosSinCache第 1 维），且取值必须命中mropeSupportList白名单{16,24,24}、{8,12,12}、{24,20,20}、{11,11,10}、{16,16,16,16}（见 aclnn_rope_with_sin_cos_cache.cpp）。

Tiling 与多核切分

op_host/rope_with_sin_cos_cache_tiling.cpp 展示了 host 侧 tiling 策略，从中可以看到该算子典型的性能设计：

• 按 token 数切分 core：front_core与tail_core两级分配，将numTokens行数据近似均匀地分配到所有 AI Core，每个 core 内部再按 UB 容量分多轮循环搬运（loop_time_each_front_core/loop_time_each_tail_core）；
• 按 UB 空间约束分片：以maxUbSize（UB 内存）为上限，根据rotaryDim、headSize、数据类型字节数估算单轮可加载的最大 token 行数与 head 数（maxNPerLoopForUb、numHeadsForUb），超过时再沿 head 维度循环（loop_along_qheads/loop_along_kheads）；
• tilingKey 按数据类型分派：BF16 对应TILING_BF16 = 20、FP16 对应 21、FP32 对应 22（见 rope_with_sin_cos_cache_tiling.cpp），同时校验headSize/rotaryDim的 64 Byte 对齐与headSize >= rotaryDim；
• 全部 tiling 参数（约 30 个字段）定义在 op_host/rope_with_sin_cos_cache_tiling.h 中，kernel 侧通过GET_TILING_DATA读取。

Kernel 入口与数据类型分派

op_kernel/rope_with_sin_cos_cache.cpp 是 AI Core 侧入口函数rope_with_sin_cos_cache，接收position_id、query_in、key_in、cos_sin_cache、query_out、key_out、workspace、tiling八个 GM 地址参数，并按 tilingKey 实例化不同的模板实现：

• TILING_KEY_IS(20)→RopeWithSinCosCacheFP16<bfloat16_t>（BF16）；
• TILING_KEY_IS(21)→RopeWithSinCosCacheFP16<half>（FP16）；
• TILING_KEY_IS(22)→RopeWithSinCosCacheF32<float>（FP32）。

针对不同架构还提供了独立实现：op_kernel/arch35/目录下的rope_with_sin_cos_cache_f_bf16_arch35.h、rope_with_sin_cos_cache_fp32_arch35.h等文件面向 arch35 平台，另有 op_kernel/rope_with_sin_cos_cache_apt.cpp 对应 Ascend 950 的 regbase 实现。

测试与验证

仓库为算子提供了完善的测试用例：

• ST 用例：tests/st/aclnnRopeWithSinCosCache/atk_aclnnRopeWithSinCosCache.json 定义了 210 组边界与随机用例，覆盖 fp16/bf16/fp32 三种类型、headSize 32/64/96/128、rotaryDim 32/64/96/128 以及 numTokens 从 1 到 131073（2^17+1）的极大规模组合，测试数据还包含nan、inf、-inf等异常值输入；
• UT 用例：见 tests/ut/op_host 下的 infershape 与 tiling 单测、tests/ut/op_kernel 下的 kernel 单测及其 gen_data.py 数据生成脚本，可在无 NPU 环境下先行验证 tiling 与 kernel 逻辑。

小结

RopeWithSinCosCache 是 CANN ops-transformer 中面向推理场景的位置编码算子，核心价值在于：

• cache 化设计：将 sin/cos 位置编码预计算并缓存，按positions索引查表，避免推理路径上的重复计算；
• 双模式覆盖：支持标准 rope 与多模态 mrope（3/4 段mropeSection），并兼容 GPT-NeoX（rotate_half）与 GPT-J（rotate_interleaved）两种主流旋转风格；
• 产品矩阵完善：覆盖 Ascend 950、Atlas A2/A3 系列与 Kirin X90/Kirin 9030，并提供 V1/V2（cacheMode）两套 aclnn 接口；
• 实现细节完整：从算子定义、参数校验白名单、UB 感知的多核 tiling 到按数据类型分派的 kernel 模板，仓库源码均可逐一对照验证。

如需进一步深入，可继续阅读 docs/aclnnRopeWithSinCosCache.md（V1 接口完整说明）、docs/aclnnRopeWithSinCosCacheV2.md（V2 接口完整说明）以及上文列出的源码与测试文件。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer