NormRopeReshapeOperation（代码开放）
产品支持情况
硬件型号 |
是否支持 |
|---|---|
√ |
|
√ |
|
x |
|
x |
|
x |
功能
对输入tensor x进行RmsNorm操作，再对keyrope进行Rope操作，将两个操作的输出合并后进行ReshapeAndCache操作 。
定义
1 2 3 4 5 6 | struct NormRopeReshapeParam { uint32_t precisionMode = 0; uint32_t rotaryCoeff = 2; float epsilon = 1e-5; uint8_t rsv[16] = {0}; }; |
参数列表
成员名称 |
类型 |
默认值 |
取值范围 |
是否必选 |
描述 |
|---|---|---|---|---|---|
precisionMode |
uint32_t |
0 |
- |
否 |
精度模式。目前只支持0。 |
epsilon |
float |
1e-5 |
- |
否 |
归一化时加在分母上防止除0。 |
rotaryCoeff |
uint32_t |
2 |
- |
否 |
算子内Rope部分计算的旋转系数。 |
rsv |
uint8_t |
uint8_t |
- |
否 |
预留参数。 |
输入
参数 |
维度 |
数据类型 |
格式 |
描述 |
|---|---|---|---|---|
X |
[ntokens, 1, head_size_x] |
float16 |
ND |
维度目前只支持[ntokens, 1, 512]。 |
gamma |
[head_size_x] |
float16 |
ND |
维度目前只支持[512]。 |
keyRope |
[ntokens, hiddenSizeK] |
float16 |
ND |
[ntokens, hiddenSizeK] hiddenSizeK = headNum * head_size（这里headNum=1）。维度目前只支持[ntokens, 64]。 |
cos |
[ntokens, head_size] |
float16 |
ND |
[ntokens, head_size]，大小与hiddenSizeK一致（headNum=1）。维度目前只支持[ntokens, 64]。 |
sin |
[ntokens, head_size] |
float16 |
ND |
同cos。 |
slotMapping |
[dnslot] |
int32 |
ND |
元素值大小不能超过blockNum * blockSize。维度目前只支持[64]。 |
keycachein |
[blockNum, blockSize, 1,dnrac] |
float16 |
ND |
dnrac=head_size_x+hiddenSizeK。维度目前只支持[blockNum, blockSize, 1, 576]。blockNum * blockSize ≥ ntokens。 |
输出
参数 |
维度 |
数据类型 |
格式 |
描述 |
|---|---|---|---|---|
keycacheout |
[blockNum, blockSize, 1, dnrac] |
与keycachein一致 |
ND |
同keycachein。 |
规格约束
- x、gamma的最后一维大小相等且32字节对齐。
- keyRope、cos、sin的数据类型需保持一致，hiddenSizeK和head_size大小需相等。
- 11*dn*sizeof(float16) + dnrac * sizeof(float16) < maxUbSize(196352)。
- ntokens小于等于65536。