算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载本篇文章围绕 CANN ops-transformer 开源仓库中的inplace_partial_rotary_mul算子展开系统讲解其在 NPU 上以 Inplace 方式执行单路旋转位置编码RoPE的接口设计、interleave 旋转计算公式、partial_slice局部切片语义、自动微分链路以及单算子模式、训练模式与图模式三种调用方式。读完本文你将能够正确配置该算子的输入与约束理解其与反向算子inplace_partial_rotary_mul_backward的联动机制并直接复用文中完整示例完成 NPU 上的 RoPE 编码。算子概览与产品支持情况inplace_partial_rotary_mul是 CANN transformer 类大模型算子库项目主页中 posembedding 目录下的核心位置编码算子源码位于 posembedding/inplace_partial_rotary_mul。其最显著的特点是原地Inplace计算执行单路旋转位置编码时直接修改输入张量x不产生新的输出张量同时支持通过partial_slice参数指定输入张量最后一维上的局部范围仅对该范围内的数据执行旋转位置编码其余位置保持原值。这一设计特别适合在多头注意力模型中对隐藏维度做部分旋转如仅旋转前 D/2 维的常见 RoPE 用法可显著节省显存与访存开销。该算子在不同 NPU 产品上的支持情况如下依据 torchapi_inplace_partial_rotary_mul.md产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册代码 inplace_partial_rotary_mul_def.cpp 可以看到底层 AICore 侧分别针对ascend910b、ascend910_93即 Atlas A2/A3 系列与ascend950配置了不同的算子实现文件其中 ascend950 走inplace_partial_rotary_mul_apt实现路径与上文产品支持矩阵一一对应。功能与计算公式接口功能执行单路旋转位置编码的 Inplace 计算直接修改输入张量x不产生新的输出张量。该接口支持通过partial_slice参数指定输入张量最后一维上的局部范围仅对该范围内的数据执行旋转位置编码其余位置保持原值。输入x采用BSND维度格式BBatch批量大小SSeq-Length序列长度NHead-Num多头数DHead-Dim每个头的隐藏维度大小。partial_slice作用于输入x的最后一维 D 维取值范围为左闭右开区间[start, end)。Python 接口中partial_slice默认值为None内部按[0, 0]处理对应空切片、不执行旋转的语义见 inplace_partial_rotary_mul.py。计算公式interleave 模式rotary_mode为interleave下设partial_slice[start, end]被旋转的局部张量为$$x_{slice} x[..., start:end]$$计算过程如下$$x_1 x_{slice}[..., ::2]$$$$x_2 x_{slice}[..., 1::2]$$$$x_{rotate} \text{cat}(-x_2, x_1)$$$$x_{out} x_{slice} \cdot \cos x_{rotate} \cdot \sin$$最终将 $x_{out}$ 原地写回x[..., start:end]。其中$x$ 表示参数x$\cos$ 表示参数r1$\sin$ 表示参数r2。当start与end相等时不执行旋转位置编码x保持不变。反向算子inplace_partial_rotary_mul_backward同样支持空 Tensor 和切片长度为零的场景执行 no-op。从底层 Kernel 实现可以印证上述公式在 inplace_partial_rotary_mul.h 的Process中先通过SetGatherSrcOffset构造奇偶索引交换的 Gather 表idsUb.SetValue(i, i ^ 1)即以 XOR 1 交换偶数位与奇数位再依次完成x * cosComputeMul、奇偶交换取数Gather、x_rotate * sinComputeMul以及InterleavedInversion对奇数位乘以 -1mask 为0x5555555555555555和最终的Add求和恰好对应公式中拆奇偶、取负拼接、乘 cos/sin 相加的完整计算链。函数原型与参数说明函数原型cann_ops_transformer.inplace_partial_rotary_mul(x, r1, r2, *, rotary_modeinterleave, partial_sliceNone) - None该接口位于cann_ops_transformer.ops包中实际源码在 torch_extension/inplace_partial_rotary_mul.py并由 torch_extension/__init__.py 导出。接口的算子 schema 为inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, *, str rotary_modeinterleave, int[2] partial_slice[0, 0]) - ()其中Tensor(a!)表示x为可写别名inplace 修改语义。参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选待执行旋转位置编码的张量对应公式中的 x。Inplace 模式下计算结果直接写回该 Tensor。bfloat16、float16、float32(B, S, N, D)r1Tensor必选位置编码张量对应公式中的 cos 分量。bfloat16、float16、float324 维需与 x 满足广播关系r2Tensor必选位置编码张量对应公式中的 sin 分量r2 和 r1 的数据类型必须一致。bfloat16、float16、float324 维需与 x 满足广播关系rotary_modestr可选旋转模式。当前仅支持 interleave默认值为 interleave。--partial_sliceList[int]可选部分旋转的切片范围 [start, end)作用于 x 的最后一维 D 维。默认值为 None接口内部按 [0, 0] 处理。--关于数据类型与格式算子注册代码 inplace_partial_rotary_mul_def.cpp 明确了x支持DT_FLOAT16 / DT_FLOAT / DT_BF16cosr1与sinr2支持DT_FLOAT16 / DT_FLOAT / DT_BF16格式统一为 ND即内存连续布局并带有AutoContiguous声明——这解释了文档中不支持非连续 Tensor的约束来源。返回值说明该接口无返回值None。计算结果直接 inplace 写回输入张量xx在计算后 shape 和数据类型保持不变partial_slice指定范围以外的数据保持原值。这一行为在 inplace_partial_rotary_mul_infershape.cpp 中有直接体现InferShape 将输出 shape 直接赋值为输入x的 shape*yShape *xShapeInferDataType 同样透传输入数据类型确认输出即输入、不产生新张量。自动微分说明当x.requires_grad为 True 时对 loss 执行.backward()即自动触发inplace_partial_rotary_mul_backward无需手动调用反向算子。r1cos、r2sin的梯度不计算始终为 None。反向算子inplace_partial_rotary_mul_backward支持空 Tensor 和切片长度为零的场景执行 no-op自动微分可正常使用。自动微分链路的实现位于 inplace_partial_rotary_mul.pyInplacePartialRotaryMulFn继承torch.autograd.Functionforward 中通过ctx.mark_dirty(x)标记 inplace 修改并保存r1、r2用于反向backward 中调用torch.ops.cann_ops_transformer.inplace_partial_rotary_mul_backward就地计算grad_input且返回值中除grad_input外其余均为None——这正是r1、r2 不计算梯度的代码级证据。此外Python 入口 inplace_partial_rotary_mul.py 会按x.requires_grad自动分派需要梯度时走 autograd Function 包装否则直接调用底层算子避免不必要的计算图开销。约束说明该接口支持训练、推理场景下使用。该接口支持单算子模式和图模式调用。不支持非连续 Tensor。x最后一维 D 大小不超过 1024且 D 必须为 2 的倍数。该上限在 tiling 代码 inplace_partial_rotary_mul_tiling.cpp 中以D_LIMIT 1024常量直接体现。partial_slice必须包含两个整数满足start 0、end 0、end D、end - start 0。partial_slice切片长度即end - start必须为 2 的倍数。当end和start相等时切片长度为零正向和反向计算均执行 no-op直接返回。tiling 代码中以TILING_KEY_NOOP 1sliceLength 0, no computation标记该分支见 inplace_partial_rotary_mul_tiling.cpp。r1、r2最后一维大小必须相同且必须等于partial_slice的切片长度即end - start。r1、r2的 shape 必须与x[..., start:end]满足广播关系且存在如下约束Ascend 950PR / Ascend 950DTr1、r2的 shape 当前只支持 BSND、B1ND、B11D、111D 排布。Atlas A3 训练系列产品 / Atlas A3 推理系列产品、Atlas A2 训练系列产品 / Atlas A2 推理系列产品r1、r2的 shape 当前只支持 BS1D、B11D 排布。x的各维度值必须大于 0当partial_slice不是空切片时r1、r2参与计算的维度值必须大于 0。自动微分约束仅计算x的梯度r1、r2的梯度不计算始终为 None。因算子为输入输出同地址操作x不能是requires_gradTrue的叶子张量。关于最后一条约束从代码实现看autograd 包装在 backward 阶段直接对grad_output执行就地反向inplace_partial_rotary_mul.py因此若x本身是requires_gradTrue的叶子张量inplace 修改会被 PyTorch 的版本计数器version counter机制判定为非法改写这也是文档要求训练时对x先做一次拷贝如y x * 1.0再传入的原因。确定性计算默认支持确定性计算。调用说明单算子模式调用单算子模式下接口直接完成 RoPE 编码无需手动管理计算图import torch import torch_npu from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B 2 S 32 N 8 D 128 slice_start 0 slice_end 64 x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) inplace_partial_rotary_mul( x, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], )示例要点说明x为 BSND 布局的(2, 32, 8, 128)D128partial_slice[0, 64]表示仅旋转最后一维的前 64 个元素后 64 个元素保持原值。r1、r2取 BS1D 排布(2, 32, 1, 64)其最后一维 64 恰等于切片长度end - start满足上文约束同时1维与x的 N 维满足广播关系Atlas A2/A3 系列支持的 BS1D、B11D 排布之一。调用后无需接收返回值x本身即被修改可在调用后直接打印x[0, 0, 0, :8]对比partial_slice内外的值验证效果。训练模式调用自动微分训练场景下接口通过 autograd 包装自动衔接反向计算。注意由于 inplace 特性需先对叶子张量x做一次拷贝import torch import torch_npu from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B, S, N, D 2, 32, 8, 128 slice_start, slice_end 0, 64 x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16, requires_gradTrue) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) y x * 1.0 y.retain_grad() # 正向自动追踪计算图y被inplace修改无需接收返回值 inplace_partial_rotary_mul( y, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], ) # 继续前向计算 loss y.sum() loss.backward() # 自动调用inplace_partial_rotary_mul_backward print(y.grad.shape) print(x.grad.shape) # r1.grad, r2.grad始终为Nonecos/sin不计算梯度此处的关键点在于y x * 1.0构造出非叶子的中间张量作为 inplace 修改对象规避x不能是requires_gradTrue的叶子张量的约束loss.backward()会经由InplacePartialRotaryMulFn.backward自动调用反向算子inplace_partial_rotary_mul_backward源码见 inplace_partial_rotary_mul.py并在反向中对r1、r2返回None梯度。图模式调用图模式torch.compiletorchairNPU 后端下将算子封装进nn.Module后编译执行import torch import torch_npu import torchair from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B 2 S 32 N 8 D 128 slice_start 0 slice_end 64 class InplacePartialRotaryMulModel(torch.nn.Module): def forward(self, x, r1, r2): inplace_partial_rotary_mul( x, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], ) return x model InplacePartialRotaryMulModel().npu() npu_backend torchair.get_npu_backend() model torch.compile(model, backendnpu_backend, dynamicFalse) x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) output model(x, r1, r2)从实现角度看图模式之所以可行得益于 torch 扩展侧对算子做了两层准备一是注册了Meta实现inplace_partial_rotary_mul.py对x.dim() ! 4、rotary_mode ! interleave、partial_slice长度不等于 2 等非法输入在编译期即抛出明确错误二是通过_ensure_initialized()与load()预加载算子模块避免 dynamo 跟踪期才触发torch.utils.cpp_extension.load()见 inplace_partial_rotary_mul.py。同时 graph_convert_inplace_partial_rotary_mul.py 负责在构图阶段完成算子到图 IR 的转换。底层实现与调用链Python 到 AICore 的完整调用链该算子的完整调用链可以概括为四层Python 接口层inplace_partial_rotary_mul.py 负责默认参数归一化partial_sliceNone - [0, 0]、autograd 分派最终下发到torch.ops.cann_ops_transformer.inplace_partial_rotary_mulC torch 扩展层csrc/inplace_partial_rotary_mul.cpp 校验x.dim() 4与rotary_mode将字符串模式映射为整数mode_map中interleave - 1再调用aclnnInplacePartialRotaryMulACLNN 计算接口层op_api/aclnn_inplace_partial_rotary_mul.h 提供标准的GetWorkspaceSizeExecute两段式接口rotary_mode以int64_t传入当前仅支持 interleave即 rotary_mode1Host/Kernel 层inplace_partial_rotary_mul_def.cpp 完成算子注册与 AICore 配置inplace_partial_rotary_mul_tiling.cpp 根据输入 shape、partial_slice、dtype 计算分核与分块策略最终由 op_kernel 下的 Kernel 代码在 NPU AICore 上执行旋转计算。Tiling 与 Kernel 的工程细节从 tiling 源码可以看出该算子在工程实现上的几个关键决策多种旋转布局支持Kernel 头文件 inplace_rotate_half.h 中定义了LAYOUT_BNSD / LAYOUT_BSND / LAYOUT_SBND / LAYOUT_R_B1SD / LAYOUT_BND / LAYOUT_NO_BROADCAST等多种布局路径分别对应r1/r2广播形态不同时的访存策略如 B1ND 广播需要按 batch 复用旋转系数对应RB1sdProcess混合精度计算非 FP32 输入会在 Kernel 内先Cast到 FP32 参与乘加最终以CAST_RINT模式写回原精度见 inplace_partial_rotary_mul.h对应 tiling 中的TILING_KEY_BFLOAT16_FLOAT32_MIXED等混合精度 tiling 分支no-op 短路切片长度为零时 tiling 直接标记TILING_KEY_NOOP正向与反向均跳过实际计算多核并行分片tiling 依据 batch、seq、head 维度做 Split 分核TilingSplitN / TilingSplitB / TilingSplitS每个核处理独立的 batch/head 分片平衡负载。示例与测试验证仓库为该算子提供了完整的示例与单测C 单算子示例examples/test_aclnn_inplace_partial_rotary_mul.cpp 与 examples/test_geir_inplace_partial_rotary_mul.cpp分别演示 ACLNN 接口与 GEIR 图接口的直接调用Host 侧单测tests/ut/op_host/test_inplace_partial_rotary_mul_a3_tiling.cpp 与 test_inplace_partial_rotary_mul_a5_tiling.cpp覆盖 Atlas A3 与 A5 平台的 tiling 策略Kernel 侧单测tests/ut/op_kernel/test_inplace_partial_rotary_mul_kernel.cpp配套 tests/ut/op_kernel/inplace_partial_rotary_mul_data 下的gen_data.py构造输入、gen_golden.py生成 golden 结果、gen_tiling.py生成 tiling 参数三个脚本构成完整的数据生成—计算—校验闭环。总结inplace_partial_rotary_mul通过输入输出同地址的 inplace 设计配合partial_slice局部切片能力为 transformer 类大模型在 NPU 上的旋转位置编码提供了一种低显存开销、支持自动微分的实现方案。其核心价值可以归纳为三点省显存、免拷贝直接改写输入张量x不产生额外输出适合长序列、大 batch 场景局部旋转partial_slice[start, end)精确控制旋转范围支持仅旋转头维度前 D/2等常见 RoPE 变体start end时自动 no-op全链路支持单算子、训练autograd 自动衔接反向、图模式torch.compile torchair三种调用方式覆盖推理与训练全场景且r1/r2作为常量参与计算、不产生梯度符合 RoPE 的标准用法。在接入该算子时务必重点核对x的 BSND 布局与 D 维约束、r1/r2最后一维等于切片长度且满足目标产品的广播排布A2/A3 系列为 BS1D、B11D950 系列为 BSND、B1ND、B11D、111D、以及训练场景下对叶子张量的拷贝处理即可稳定获得正确的部分旋转位置编码结果。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 局部旋转位置编码反向算子 aclnnInplacePartialRotaryMulGrad 原理与调用指南CANN ops transformer 局部旋转位置编码反向算子 aclnnInplacePartialRotaryMulGrad 原理与调用指南 本文以 C算子库人工智能深度学习Ascend使用Next-on-Netlify实现服务器端渲染的10个实用技巧使用Next on Netlify实现服务器端渲染的10个实用技巧 Next on Netlify是一个强大的工具它允许开发者在Netlify平台上构建和部署算子库人工智能深度学习AscendCANN ops-transformer 反向旋转位置编码算子 aclnnRotaryPositionEmbeddingGrad 使用与实现解析CANN ops transformer 反向旋转位置编码算子 aclnnRotaryPositionEmbeddingGrad 使用与实现解析 本文围绕 CA算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer 算子详解：inplace_partial_rotary_mul 原地部分旋转位置编码
算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载本篇文章围绕 CANN ops-transformer 开源仓库中的inplace_partial_rotary_mul算子展开系统讲解其在 NPU 上以 Inplace 方式执行单路旋转位置编码RoPE的接口设计、interleave 旋转计算公式、partial_slice局部切片语义、自动微分链路以及单算子模式、训练模式与图模式三种调用方式。读完本文你将能够正确配置该算子的输入与约束理解其与反向算子inplace_partial_rotary_mul_backward的联动机制并直接复用文中完整示例完成 NPU 上的 RoPE 编码。算子概览与产品支持情况inplace_partial_rotary_mul是 CANN transformer 类大模型算子库项目主页中 posembedding 目录下的核心位置编码算子源码位于 posembedding/inplace_partial_rotary_mul。其最显著的特点是原地Inplace计算执行单路旋转位置编码时直接修改输入张量x不产生新的输出张量同时支持通过partial_slice参数指定输入张量最后一维上的局部范围仅对该范围内的数据执行旋转位置编码其余位置保持原值。这一设计特别适合在多头注意力模型中对隐藏维度做部分旋转如仅旋转前 D/2 维的常见 RoPE 用法可显著节省显存与访存开销。该算子在不同 NPU 产品上的支持情况如下依据 torchapi_inplace_partial_rotary_mul.md产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册代码 inplace_partial_rotary_mul_def.cpp 可以看到底层 AICore 侧分别针对ascend910b、ascend910_93即 Atlas A2/A3 系列与ascend950配置了不同的算子实现文件其中 ascend950 走inplace_partial_rotary_mul_apt实现路径与上文产品支持矩阵一一对应。功能与计算公式接口功能执行单路旋转位置编码的 Inplace 计算直接修改输入张量x不产生新的输出张量。该接口支持通过partial_slice参数指定输入张量最后一维上的局部范围仅对该范围内的数据执行旋转位置编码其余位置保持原值。输入x采用BSND维度格式BBatch批量大小SSeq-Length序列长度NHead-Num多头数DHead-Dim每个头的隐藏维度大小。partial_slice作用于输入x的最后一维 D 维取值范围为左闭右开区间[start, end)。Python 接口中partial_slice默认值为None内部按[0, 0]处理对应空切片、不执行旋转的语义见 inplace_partial_rotary_mul.py。计算公式interleave 模式rotary_mode为interleave下设partial_slice[start, end]被旋转的局部张量为$$x_{slice} x[..., start:end]$$计算过程如下$$x_1 x_{slice}[..., ::2]$$$$x_2 x_{slice}[..., 1::2]$$$$x_{rotate} \text{cat}(-x_2, x_1)$$$$x_{out} x_{slice} \cdot \cos x_{rotate} \cdot \sin$$最终将 $x_{out}$ 原地写回x[..., start:end]。其中$x$ 表示参数x$\cos$ 表示参数r1$\sin$ 表示参数r2。当start与end相等时不执行旋转位置编码x保持不变。反向算子inplace_partial_rotary_mul_backward同样支持空 Tensor 和切片长度为零的场景执行 no-op。从底层 Kernel 实现可以印证上述公式在 inplace_partial_rotary_mul.h 的Process中先通过SetGatherSrcOffset构造奇偶索引交换的 Gather 表idsUb.SetValue(i, i ^ 1)即以 XOR 1 交换偶数位与奇数位再依次完成x * cosComputeMul、奇偶交换取数Gather、x_rotate * sinComputeMul以及InterleavedInversion对奇数位乘以 -1mask 为0x5555555555555555和最终的Add求和恰好对应公式中拆奇偶、取负拼接、乘 cos/sin 相加的完整计算链。函数原型与参数说明函数原型cann_ops_transformer.inplace_partial_rotary_mul(x, r1, r2, *, rotary_modeinterleave, partial_sliceNone) - None该接口位于cann_ops_transformer.ops包中实际源码在 torch_extension/inplace_partial_rotary_mul.py并由 torch_extension/__init__.py 导出。接口的算子 schema 为inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, *, str rotary_modeinterleave, int[2] partial_slice[0, 0]) - ()其中Tensor(a!)表示x为可写别名inplace 修改语义。参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选待执行旋转位置编码的张量对应公式中的 x。Inplace 模式下计算结果直接写回该 Tensor。bfloat16、float16、float32(B, S, N, D)r1Tensor必选位置编码张量对应公式中的 cos 分量。bfloat16、float16、float324 维需与 x 满足广播关系r2Tensor必选位置编码张量对应公式中的 sin 分量r2 和 r1 的数据类型必须一致。bfloat16、float16、float324 维需与 x 满足广播关系rotary_modestr可选旋转模式。当前仅支持 interleave默认值为 interleave。--partial_sliceList[int]可选部分旋转的切片范围 [start, end)作用于 x 的最后一维 D 维。默认值为 None接口内部按 [0, 0] 处理。--关于数据类型与格式算子注册代码 inplace_partial_rotary_mul_def.cpp 明确了x支持DT_FLOAT16 / DT_FLOAT / DT_BF16cosr1与sinr2支持DT_FLOAT16 / DT_FLOAT / DT_BF16格式统一为 ND即内存连续布局并带有AutoContiguous声明——这解释了文档中不支持非连续 Tensor的约束来源。返回值说明该接口无返回值None。计算结果直接 inplace 写回输入张量xx在计算后 shape 和数据类型保持不变partial_slice指定范围以外的数据保持原值。这一行为在 inplace_partial_rotary_mul_infershape.cpp 中有直接体现InferShape 将输出 shape 直接赋值为输入x的 shape*yShape *xShapeInferDataType 同样透传输入数据类型确认输出即输入、不产生新张量。自动微分说明当x.requires_grad为 True 时对 loss 执行.backward()即自动触发inplace_partial_rotary_mul_backward无需手动调用反向算子。r1cos、r2sin的梯度不计算始终为 None。反向算子inplace_partial_rotary_mul_backward支持空 Tensor 和切片长度为零的场景执行 no-op自动微分可正常使用。自动微分链路的实现位于 inplace_partial_rotary_mul.pyInplacePartialRotaryMulFn继承torch.autograd.Functionforward 中通过ctx.mark_dirty(x)标记 inplace 修改并保存r1、r2用于反向backward 中调用torch.ops.cann_ops_transformer.inplace_partial_rotary_mul_backward就地计算grad_input且返回值中除grad_input外其余均为None——这正是r1、r2 不计算梯度的代码级证据。此外Python 入口 inplace_partial_rotary_mul.py 会按x.requires_grad自动分派需要梯度时走 autograd Function 包装否则直接调用底层算子避免不必要的计算图开销。约束说明该接口支持训练、推理场景下使用。该接口支持单算子模式和图模式调用。不支持非连续 Tensor。x最后一维 D 大小不超过 1024且 D 必须为 2 的倍数。该上限在 tiling 代码 inplace_partial_rotary_mul_tiling.cpp 中以D_LIMIT 1024常量直接体现。partial_slice必须包含两个整数满足start 0、end 0、end D、end - start 0。partial_slice切片长度即end - start必须为 2 的倍数。当end和start相等时切片长度为零正向和反向计算均执行 no-op直接返回。tiling 代码中以TILING_KEY_NOOP 1sliceLength 0, no computation标记该分支见 inplace_partial_rotary_mul_tiling.cpp。r1、r2最后一维大小必须相同且必须等于partial_slice的切片长度即end - start。r1、r2的 shape 必须与x[..., start:end]满足广播关系且存在如下约束Ascend 950PR / Ascend 950DTr1、r2的 shape 当前只支持 BSND、B1ND、B11D、111D 排布。Atlas A3 训练系列产品 / Atlas A3 推理系列产品、Atlas A2 训练系列产品 / Atlas A2 推理系列产品r1、r2的 shape 当前只支持 BS1D、B11D 排布。x的各维度值必须大于 0当partial_slice不是空切片时r1、r2参与计算的维度值必须大于 0。自动微分约束仅计算x的梯度r1、r2的梯度不计算始终为 None。因算子为输入输出同地址操作x不能是requires_gradTrue的叶子张量。关于最后一条约束从代码实现看autograd 包装在 backward 阶段直接对grad_output执行就地反向inplace_partial_rotary_mul.py因此若x本身是requires_gradTrue的叶子张量inplace 修改会被 PyTorch 的版本计数器version counter机制判定为非法改写这也是文档要求训练时对x先做一次拷贝如y x * 1.0再传入的原因。确定性计算默认支持确定性计算。调用说明单算子模式调用单算子模式下接口直接完成 RoPE 编码无需手动管理计算图import torch import torch_npu from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B 2 S 32 N 8 D 128 slice_start 0 slice_end 64 x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) inplace_partial_rotary_mul( x, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], )示例要点说明x为 BSND 布局的(2, 32, 8, 128)D128partial_slice[0, 64]表示仅旋转最后一维的前 64 个元素后 64 个元素保持原值。r1、r2取 BS1D 排布(2, 32, 1, 64)其最后一维 64 恰等于切片长度end - start满足上文约束同时1维与x的 N 维满足广播关系Atlas A2/A3 系列支持的 BS1D、B11D 排布之一。调用后无需接收返回值x本身即被修改可在调用后直接打印x[0, 0, 0, :8]对比partial_slice内外的值验证效果。训练模式调用自动微分训练场景下接口通过 autograd 包装自动衔接反向计算。注意由于 inplace 特性需先对叶子张量x做一次拷贝import torch import torch_npu from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B, S, N, D 2, 32, 8, 128 slice_start, slice_end 0, 64 x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16, requires_gradTrue) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) y x * 1.0 y.retain_grad() # 正向自动追踪计算图y被inplace修改无需接收返回值 inplace_partial_rotary_mul( y, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], ) # 继续前向计算 loss y.sum() loss.backward() # 自动调用inplace_partial_rotary_mul_backward print(y.grad.shape) print(x.grad.shape) # r1.grad, r2.grad始终为Nonecos/sin不计算梯度此处的关键点在于y x * 1.0构造出非叶子的中间张量作为 inplace 修改对象规避x不能是requires_gradTrue的叶子张量的约束loss.backward()会经由InplacePartialRotaryMulFn.backward自动调用反向算子inplace_partial_rotary_mul_backward源码见 inplace_partial_rotary_mul.py并在反向中对r1、r2返回None梯度。图模式调用图模式torch.compiletorchairNPU 后端下将算子封装进nn.Module后编译执行import torch import torch_npu import torchair from cann_ops_transformer.ops import inplace_partial_rotary_mul torch_npu.npu.set_device(0) B 2 S 32 N 8 D 128 slice_start 0 slice_end 64 class InplacePartialRotaryMulModel(torch.nn.Module): def forward(self, x, r1, r2): inplace_partial_rotary_mul( x, r1, r2, rotary_modeinterleave, partial_slice[slice_start, slice_end], ) return x model InplacePartialRotaryMulModel().npu() npu_backend torchair.get_npu_backend() model torch.compile(model, backendnpu_backend, dynamicFalse) x torch.randn(B, S, N, D, devicenpu, dtypetorch.float16) r1 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) r2 torch.randn(B, S, 1, slice_end - slice_start, devicenpu, dtypetorch.float16) output model(x, r1, r2)从实现角度看图模式之所以可行得益于 torch 扩展侧对算子做了两层准备一是注册了Meta实现inplace_partial_rotary_mul.py对x.dim() ! 4、rotary_mode ! interleave、partial_slice长度不等于 2 等非法输入在编译期即抛出明确错误二是通过_ensure_initialized()与load()预加载算子模块避免 dynamo 跟踪期才触发torch.utils.cpp_extension.load()见 inplace_partial_rotary_mul.py。同时 graph_convert_inplace_partial_rotary_mul.py 负责在构图阶段完成算子到图 IR 的转换。底层实现与调用链Python 到 AICore 的完整调用链该算子的完整调用链可以概括为四层Python 接口层inplace_partial_rotary_mul.py 负责默认参数归一化partial_sliceNone - [0, 0]、autograd 分派最终下发到torch.ops.cann_ops_transformer.inplace_partial_rotary_mulC torch 扩展层csrc/inplace_partial_rotary_mul.cpp 校验x.dim() 4与rotary_mode将字符串模式映射为整数mode_map中interleave - 1再调用aclnnInplacePartialRotaryMulACLNN 计算接口层op_api/aclnn_inplace_partial_rotary_mul.h 提供标准的GetWorkspaceSizeExecute两段式接口rotary_mode以int64_t传入当前仅支持 interleave即 rotary_mode1Host/Kernel 层inplace_partial_rotary_mul_def.cpp 完成算子注册与 AICore 配置inplace_partial_rotary_mul_tiling.cpp 根据输入 shape、partial_slice、dtype 计算分核与分块策略最终由 op_kernel 下的 Kernel 代码在 NPU AICore 上执行旋转计算。Tiling 与 Kernel 的工程细节从 tiling 源码可以看出该算子在工程实现上的几个关键决策多种旋转布局支持Kernel 头文件 inplace_rotate_half.h 中定义了LAYOUT_BNSD / LAYOUT_BSND / LAYOUT_SBND / LAYOUT_R_B1SD / LAYOUT_BND / LAYOUT_NO_BROADCAST等多种布局路径分别对应r1/r2广播形态不同时的访存策略如 B1ND 广播需要按 batch 复用旋转系数对应RB1sdProcess混合精度计算非 FP32 输入会在 Kernel 内先Cast到 FP32 参与乘加最终以CAST_RINT模式写回原精度见 inplace_partial_rotary_mul.h对应 tiling 中的TILING_KEY_BFLOAT16_FLOAT32_MIXED等混合精度 tiling 分支no-op 短路切片长度为零时 tiling 直接标记TILING_KEY_NOOP正向与反向均跳过实际计算多核并行分片tiling 依据 batch、seq、head 维度做 Split 分核TilingSplitN / TilingSplitB / TilingSplitS每个核处理独立的 batch/head 分片平衡负载。示例与测试验证仓库为该算子提供了完整的示例与单测C 单算子示例examples/test_aclnn_inplace_partial_rotary_mul.cpp 与 examples/test_geir_inplace_partial_rotary_mul.cpp分别演示 ACLNN 接口与 GEIR 图接口的直接调用Host 侧单测tests/ut/op_host/test_inplace_partial_rotary_mul_a3_tiling.cpp 与 test_inplace_partial_rotary_mul_a5_tiling.cpp覆盖 Atlas A3 与 A5 平台的 tiling 策略Kernel 侧单测tests/ut/op_kernel/test_inplace_partial_rotary_mul_kernel.cpp配套 tests/ut/op_kernel/inplace_partial_rotary_mul_data 下的gen_data.py构造输入、gen_golden.py生成 golden 结果、gen_tiling.py生成 tiling 参数三个脚本构成完整的数据生成—计算—校验闭环。总结inplace_partial_rotary_mul通过输入输出同地址的 inplace 设计配合partial_slice局部切片能力为 transformer 类大模型在 NPU 上的旋转位置编码提供了一种低显存开销、支持自动微分的实现方案。其核心价值可以归纳为三点省显存、免拷贝直接改写输入张量x不产生额外输出适合长序列、大 batch 场景局部旋转partial_slice[start, end)精确控制旋转范围支持仅旋转头维度前 D/2等常见 RoPE 变体start end时自动 no-op全链路支持单算子、训练autograd 自动衔接反向、图模式torch.compile torchair三种调用方式覆盖推理与训练全场景且r1/r2作为常量参与计算、不产生梯度符合 RoPE 的标准用法。在接入该算子时务必重点核对x的 BSND 布局与 D 维约束、r1/r2最后一维等于切片长度且满足目标产品的广播排布A2/A3 系列为 BS1D、B11D950 系列为 BSND、B1ND、B11D、111D、以及训练场景下对叶子张量的拷贝处理即可稳定获得正确的部分旋转位置编码结果。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 局部旋转位置编码反向算子 aclnnInplacePartialRotaryMulGrad 原理与调用指南CANN ops transformer 局部旋转位置编码反向算子 aclnnInplacePartialRotaryMulGrad 原理与调用指南 本文以 C算子库人工智能深度学习Ascend使用Next-on-Netlify实现服务器端渲染的10个实用技巧使用Next on Netlify实现服务器端渲染的10个实用技巧 Next on Netlify是一个强大的工具它允许开发者在Netlify平台上构建和部署算子库人工智能深度学习AscendCANN ops-transformer 反向旋转位置编码算子 aclnnRotaryPositionEmbeddingGrad 使用与实现解析CANN ops transformer 反向旋转位置编码算子 aclnnRotaryPositionEmbeddingGrad 使用与实现解析 本文围绕 CA算子库人工智能深度学习Ascend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考