InplacePartialRotaryMul

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：执行单路旋转位置编码的Inplace计算，直接修改输入张量，不产生新的输出张量。支持通过partial_slice参数对输入张量的最后一维进行部分旋转位置编码计算，仅对指定范围内的数据进行旋转位置编码。

• 计算公式：

  interleave模式（rotary_mode等于1）：

  $$x1=x[...,::2]$$

  $$x2=x[...,1::2]$$

  $$x\_rotate=torch.cat((-x2,x1),dim=-1)$$

  $$x=x\ast cos+x\_rotate\ast sin$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
xRef	输入	公式中的x，待执行旋转位置编码的张量。Inplace模式，xRef同时作为输出写入结果。	BFLOAT16、FLOAT16、FLOAT32	ND
cos	输入	公式中的cos，参与计算的位置编码张量。	BFLOAT16、FLOAT16、FLOAT32	ND
sin	输入	公式中的sin，参与计算的位置编码张量。	BFLOAT16、FLOAT16、FLOAT32	ND
rotary_mode	输入	公式中的旋转模式，当前仅支持interleave模式（rotary_mode=1）。	INT64	-
partial_slice	输入	部分旋转的切片范围[start, end)，作用于最后一维。默认值为[0, 0]。	INT64数组	-

约束说明

• 不支持非连续。

• 输入张量xRef最后一维（D）大小不超过1024。

• interleave模式（rotary_mode = 1）下，xRef最后一维（D）大小必须为2的倍数。

• 输入张量cos、sin最后一维大小必须相同，且必须等于partialSlice的切片长度（即partialSlice[1] - partialSlice[0]）。

• partialSlice约束：sliceStart ≥ 0，sliceEnd ≥ 0，sliceEnd ≤ xRef最后一维（D）大小，sliceLength = sliceEnd - sliceStart >= 0，当sliceEnd和sliceStart相同时，不做旋转位置编码，直接返回。

• 仅支持interleave模式（rotary_mode = 1）。

• Inplace执行：输入xRef和输出共享同一个Tensor，计算结果直接写回输入xRef。

• xRef的shape为BSND。

• cos/sin的shape必须与xRef满足广播关系，且存在如下约束：

  • Ascend 950PR/Ascend 950DT： cos/sin的shape当前只支持BSND、B1ND、B11D、111D排布。

  • Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品： cos/sin的shape当前只支持BS1D、B11D排布。

调用说明

调用方式	调用样例	说明
aclnn调用	test_aclnn_inplace_partial_rotary_mul	通过aclnnInplacePartialRotaryMul接口方式调用InplacePartialRotaryMul算子。
PyTorch API	inplace_partial_rotary_mul	通过cann_ops_transformer.inplace_partial_rotary_mul接口方式调用InplacePartialRotaryMul算子。
图模式调用	test_geir_inplace_partial_rotary_mul	通过算子IR构图方式调用InplacePartialRotaryMul算子。