昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

RopeOperation

功能

旋转位置编码。

约束

  • hiddenSizeQ必须是hiddenSizeK的整数倍且满足hiddenSizeQ = headDim * headNum。
  • cos/sin tensor从cos/sin table中选取时，应选取与输入q/k的position id对应的cos/sin tensor。

定义

struct RopeParam {
    int32_t rotaryCoeff = 4;
    int32_t cosFormat = 0;
};

成员

成员名称

描述

rotaryCoeff

rope，旋转系数，对半旋转是2，支持配置2、4或headDim / 2。

cosFormat

训练用参数，支持配置0或1。

输入

参数

维度

数据类型

格式

描述

query

[ntokens, hiddenSizeQ]

float16/bf16

ND

当前step多个token的query。

仅Atlas 800I A2推理产品支持bfloat16。

key

[ntokens, hiddenSizeK]

float16/bf16

ND

当前step多个token的key。

仅Atlas 800I A2推理产品支持bfloat16。

cos

[ntokens, headDim] / [ntokens, headDim / 2]

float16/float/bf16

ND
  • 当cos的第二个维度与参数rotaryCoeff相等时，其值为headDim/2。
  • 当cos的第二个维度与参数rotaryCoeff不相等时，其值为headDim。
  • 仅Atlas 800I A2推理产品支持bfloat16。
  • ROPE高精度模式，需要输入cos的数据类型为float时生效。

sin

[ntokens, headDim] / [ntokens, headDim/ 2]

float16/float/bf16

ND
  • 当sin的第二个维度与参数rotaryCoeff相等时，其值为headDim/2。
  • 当sin的第二个维度与参数rotaryCoeff不相等时，其值为headDim。
  • 仅Atlas 800I A2推理产品支持bfloat16。
  • ROPE高精度模式，需要输入sin的数据类型为float时生效。

seqlen

[batch]

uint32/int32

ND

-

输出

参数

维度

数据类型

格式

描述

ropeQ

[ntokens, hiddenSizeQ]

float16/bf16

ND

旋转后的query。

仅Atlas 800I A2推理产品支持bfloat16。

ropeK

[ntokens, hiddenSizeK]

float16/bf16

ND

旋转后的key。

仅Atlas 800I A2推理产品支持bfloat16。
父主题： atb/infer_op_params.h