RopeWithSinCosCache
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | √ |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
| Kirin X90 处理器系列产品 | √ |
| Kirin 9030 处理器系列产品 | √ |
功能说明
-
算子功能：推理网络为了提升性能，将sin和cos输入通过cache传入，执行旋转位置编码计算。
-
计算公式：
1、mrope模式：positions的shape输入是[3, numTokens]：
cosSin[i]=cosSinCache[positions[i]]cosSin[i] = cosSinCache[positions[i]]
cos,sin=cosSin.chunk(2,dim=−1)cos, sin = cosSin.chunk(2, dim=-1)
cos0=cos[0,:,:mropeSection[0]]cos0 = cos[0, :, :mropeSection[0]]
cos1=cos[1,:,mropeSection[0]:(mropeSection[0]+mropeSection[1])]cos1 = cos[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])]
cos2=cos[2,:,(mropeSection[0]+mropeSection[1]):(mropeSection[0]+mropeSection[1]+mropeSection[2])]cos2 = cos[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])]
cos=torch.cat((cos0,cos1,cos2),dim=−1)cos = torch.cat((cos0, cos1, cos2), dim=-1)
sin0=sin[0,:,:mropeSection[0]]sin0 = sin[0, :, :mropeSection[0]]
sin1=sin[1,:,mropeSection[0]:(mropeSection[0]+mropeSection[1])]sin1 = sin[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])]
sin2=sin[2,:,(mropeSection[0]+mropeSection[1]):(mropeSection[0]+mropeSection[1]+mropeSection[2])]sin2 = sin[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])]
sin=torch.cat((sin0,sin1,sin2),dim=−1)sin= torch.cat((sin0, sin1, sin2), dim=-1)
queryRot=query[...,:rotaryDim]queryRot = query[..., :rotaryDim]
queryPass=query[...,rotaryDim:]queryPass = query[..., rotaryDim:]
（1）rotate_half（GPT-NeoX style）计算模式：
x1,x2=torch.chunk(queryRot,2,dim=−1)x1, x2 = torch.chunk(queryRot, 2, dim=-1)
o1[i]=x1[i]∗cos[i]−x2[i]∗sin[i]o1[i] = x1[i] * cos[i] - x2[i] * sin[i]
o2[i]=x2[i]∗cos[i]+x1[i]∗sin[i]o2[i] = x2[i] * cos[i] + x1[i] * sin[i]
queryRot=torch.cat((o1,o2),dim=−1)queryRot = torch.cat((o1, o2), dim=-1)
query=torch.cat((queryRot,queryPass),dim=−1)query = torch.cat((queryRot, queryPass), dim=-1)
（2）rotate_interleaved（GPT-J style）计算模式：
x1=queryRot[...,::2]x1 = queryRot[..., ::2]
x2=queryRot[...,1::2]x2 = queryRot[..., 1::2]
o1[i]=x1[i]∗cos[i]−x2[i]∗sin[i]o1[i] = x1[i] * cos[i] - x2[i] * sin[i]
o2[i]=x2[i]∗cos[i]+x1[i]∗sin[i]o2[i] = x2[i] * cos[i] + x1[i] * sin[i]
queryRot=torch.stack((o1,o2),dim=−1)queryRot = torch.stack((o1, o2), dim=-1)
query=torch.cat((queryRot,queryPass),dim=−1)query = torch.cat((queryRot, queryPass), dim=-1)
2、rope模式：positions的shape输入是[numTokens]：
cosSin[i]=cosSinCache[positions[i]]cosSin[i] = cosSinCache[positions[i]]
cos,sin=cosSin.chunk(2,dim=−1)cos, sin = cosSin.chunk(2, dim=-1)
queryRot=query[...,:rotaryDim]queryRot = query[..., :rotaryDim]
queryPass=query[...,rotaryDim:]queryPass = query[..., rotaryDim:]
（1）rotate_half（GPT-NeoX style）计算模式：
x1,x2=torch.chunk(queryRot,2,dim=−1)x1, x2 = torch.chunk(queryRot, 2, dim=-1)
o1[i]=x1[i]∗cos[i]−x2[i]∗sin[i]o1[i] = x1[i] * cos[i] - x2[i] * sin[i]
o2[i]=x2[i]∗cos[i]+x1[i]∗sin[i]o2[i] = x2[i] * cos[i] + x1[i] * sin[i]
queryRot=torch.cat((o1,o2),dim=−1)queryRot = torch.cat((o1, o2), dim=-1)
query=torch.cat((queryRot,queryPass),dim=−1)query = torch.cat((queryRot, queryPass), dim=-1)
（2）rotate_interleaved（GPT-J style）计算模式：
x1=queryRot[...,::2]x1 = queryRot[..., ::2]
x2=queryRot[...,1::2]x2 = queryRot[..., 1::2]
o1[i]=x1[i]∗cos[i]−x2[i]∗sin[i]o1[i] = x1[i] * cos[i] - x2[i] * sin[i]
o2[i]=x2[i]∗cos[i]+x1[i]∗sin[i]o2[i] = x2[i] * cos[i] + x1[i] * sin[i]
queryRot=torch.stack((o1,o2),dim=−1)queryRot = torch.stack((o1, o2), dim=-1)
query=torch.cat((queryRot,queryPass),dim=−1)query = torch.cat((queryRot, queryPass), dim=-1)
参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| positions | 输入 | Device侧的aclTensor，输入索引。 | INT32、INT64 | ND |
| queryIn | 输入 | Device侧的aclTensor，表示要执行旋转位置编码的第一个张量，公式中的`query`。 | BFLOAT16、FLOAT16、FLOAT32 | ND |
| keyIn | 输入 | Device侧的aclTensor，表示要执行旋转位置编码的第二个张量。 | BFLOAT16、FLOAT16、FLOAT32 | ND |
| cosSinCache | 输入 | Device侧的aclTensor，表示参与计算的位置编码张量。 | BFLOAT16、FLOAT16、FLOAT32 | ND |
| mropeSection | 输入 | mrope模式下用于整合输入的位置编码张量信息，公式中的`mropeSection`。 | INT64 | - |
| headSize | 输入 | 表示每个注意力头维度大小。 | INT64 | - |
| isNeoxStyle | 输入 | true表示rotate\_half（GPT-NeoX style）计算模式，false表示rotate\_interleaved（GPT-J style）计算模式。 | BOOL | - |
| queryOut | 输出 | 输出query执行旋转位置编码后的结果。 | FLOAT、FLOAT16、BFLOAT16 | ND |
| keyOut | 输出 | 输出key执行旋转位置编码后的结果。 | FLOAT、FLOAT16、BFLOAT16 | ND |
- Kirin X90/Kirin 9030 处理器系列产品: 不支持 BFLOAT16。
约束说明
- queryIn、keyIn、cosSinCache只支持2维shape输入。
- headSize: 数据类型为BFLOAT16或FLOAT16时为32的倍数，数据类型为FLOAT32时为16的倍数。
- rotaryDim: 始终小于等于headSize；数据类型为BFLOAT16或FLOAT16时为32的倍数，数据类型为FLOAT32时为16的倍数;mrope模式下应满足rotaryDim = mropeSection[0] + mropeSection[1] + mropeSection[2]。
- 输入tensor positions的取值应小于cosSinCache的0维maxSeqLen。
- aclnnRopeWithSinCosCache默认确定性实现。
- mropeSection:取值限制为[16, 24, 24]。
调用说明
| 调用方式 | 样例代码 | 说明 |
|---|---|---|
| aclnn接口 | test_aclnn_rope_with_sin_cos_cache | 通过aclnnRopeWithSinCosCache接口方式调用RopeWithSinCosCache算子。 |