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

aclnnKvRmsNormRopeCache

产品支持情况

产品 是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object] ×
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object] ×
[object Object]Atlas 200I/500 A2 推理产品[object Object] ×
[object Object]Atlas 推理系列产品 [object Object] ×
[object Object]Atlas 训练系列产品[object Object] ×

功能说明

  • 算子功能：对输入张量(kv)的尾轴，拆分出左半边用于rms_norm计算，右半边用于rope计算，再将计算结果分别scatter到两块cache中。

  • 计算公式：

    (1) interleaveRope:

    x=kv[...,Dv:]x=kv[...,Dv:] x1=x[...,::2]x1=x[...,::2] x2=x[...,1::2]x2=x[...,1::2] x_part1=torch.cat((x1,x2),dim=−1)x\_part1=torch.cat((x1,x2),dim=-1) x_part2=torch.cat((−x2,x1),dim=−1)x\_part2=torch.cat((-x2,x1),dim=-1) y=x_part1∗cos+x_part2∗siny=x\_part1*cos+x\_part2*sin

    (2) rmsNorm:

    x=kv[...,:Dv]x=kv[...,:Dv] square_x=x∗xsquare\_x=x*x mean_square_x=square_x.mean(dim=−1,keepdim=True)mean\_square\_x=square\_x.mean(dim=-1,keepdim=True) rms=torch.sqrt(mean_square_x+epsilon)rms=torch.sqrt(mean\_square\_x+epsilon) y=(x/rms)∗gammay=(x/rms)*gamma

函数原型

每个算子分为undefined，必须先调用“aclnnKvRmsNormRopeCacheGetWorkspaceSize”接口获得入参并根据流程计算所需workspace大小，再调用“aclnnKvRmsNormRopeCache”接口执行计算。

  • aclnnStatus aclnnKvRmsNormRopeCacheGetWorkspaceSize(const aclTensor* kv, const aclTensor* gamma, const aclTensor* cos, const aclTensor* sin, const aclTensor* index, aclTensor* kCacheRef, aclTensor* ckvCacheRef, const aclTensor* kRopeScaleOptional, const aclTensor* ckvScaleOptional, const aclTensor* kRopeOffsetOptional, const aclTensor* cKvOffsetOptional, double epsilon, char* cacheModeOptional, bool isOutputKv, aclTensor* kRopeOut, aclTensor* cKvOut,uint64_t workspaceSize, aclOpExecutor* executor)
  • aclnnStatus aclnnKvRmsNormRopeCache(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnKvRmsNormRopeCacheGetWorkspaceSize

  • 参数说明：

    • kv(aclTensor*，计算输入)：必选参数，公式中用于切分出rms_norm计算所需数据Dv和rope计算所需数据Dk的输入数据。Device侧的aclTensor，支持undefined。shape仅支持4维[Bkv,N,Skv,D]。undefined支持ND，数据类型支持FLOAT6、BFLOAT16。
    • gamma(aclTensor*，计算输入)：必选参数，公式中用于rms_norm计算的输入数据。Device侧的aclTensor，支持undefined。shape为1维[Dv,]。undefined支持ND。数据类型与输入kv一致。
    • cos(aclTensor*，计算输入)：必选参数，公式中用于rope计算的输入数据，对输入张量Dk进行余弦变换，Device侧的aclTensor，支持undefined。shape为4维[Bkv,1,Skv,Dk]或[Bkv,1,1,Dk]，undefined支持ND。数据类型与输入kv一致。
    • sin(aclTensor*，计算输入)：必选参数，公式中用于rope计算的输入数据，对输入张量Dk进行正弦变换。Device侧的aclTensor，支持undefined。shape为4维[Bkv,1,Skv,Dk]或[Bkv,1,1,Dk]，与cos的shape保持一致。undefined支持ND。数据类型与输入kv一致。
    • index(aclTensor*，计算输入)：必选参数，用于指定写入cache的具体索引位置，当index的value数值为-1时，代表跳过更新。Device侧的aclTensor，支持undefined。当cacheModeOptional为Norm时，shape为2维[Bkv,Skv]；当cacheModeOptional为PA_BNSD、PA_NZ时，shape为1维[Bkv * Skv]；当cacheModeOptional为PA_BLK_BSND、PA_BLK_NZ时，shape为1维[Bkv*ceil_div(Skv,BlockSize)]。undefined支持ND。数据类型支持INT64。
    • kCacheRef(aclTensor*，计算输入/输出)：必选参数，提前申请的cache，输入输出同地址复用。Device侧的aclTensor，支持undefined。当cacheModeOptional为PA场景（cacheModeOptional为PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ）时，shape为4维[BlockNum,BlockSize,N,Dk]；当cacheModeOptional为Norm场景时，shape为4维[Bcache,N,Scache,Dk]。undefined支持ND。数据类型与输入kv一致或者INT8。
    • ckvCacheRef(aclTensor*，计算输入/输出)：必选参数，提前申请的cache，输入输出同地址复用。Device侧的aclTensor，支持undefined。当cacheModeOptional为PA场景（cacheModeOptional为PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ）时，shape为4维[BlockNum,BlockSize,N,Dv]；当cacheModeOptional为Norm场景时，shape为4维[Bcache,N,Scache,Dv]。undefined支持ND。数据类型与输入kv一致或者INT8。
    • kRopeScaleOptional(aclTensor*，计算输入)：可选参数，当kCacheRef数据类型为INT8时需要此输入参数。Device侧的aclTensor，支持undefined。shape为2维[N,Dk]；或者shape为1维[Dk,]；或者shape为1维[1,]。undefined支持ND。数据类型支持FLOAT32。
    • ckvScaleOptional(aclTensor*，计算输入)：可选参数，当ckvCacheRef数据类型为INT8时需要此输入参数。Device侧的aclTensor，支持undefined。shape为2维[N,Dv]；或者shape为1维[Dv,]；或者shape为1维[1,]。undefined支持ND。数据类型支持FLOAT32。
    • kRopeOffsetOptional(aclTensor*，计算输入)：可选参数。当kCacheRef数据类型为INT8且对应的kRopeScaleOptional输入存在并量化场景为非对称量化时，需要此参数输入。Device侧的aclTensor，支持undefined。shape为2维[N,Dk]；或者shape为1维[Dk,]；或者shape为1维[1,]。undefined支持ND。数据类型支持FLOAT32。
    • cKvOffsetOptional(aclTensor*，计算输入)：可选参数，当ckvCacheRef数据类型为INT8且对应的ckvScaleOptional输入存在并量化场景为非对称量化时，需要此参数输入。Device侧的aclTensor，支持undefined。shape为2维[N,Dv]；或者shape为1维[Dv,]；或者shape为1维[1,]。undefined支持ND。数据类型支持FLOAT32。
    • epsilon(double，输入)：必选参数，rms_norm计算防止除0。float类型浮点数。建议设为1e-5。
    • cacheModeOptional(char*，输入)：必选参数，cache格式的选择标记。char*类型。类型有Norm、PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ，建议设为Norm。
    • isOutputKv(bool，输入)：必选参数，kRopeOut和cKvOut输出控制标记。bool类型。当isOutputKv为true时，表示需输出kRopeOut和cKvOut。建议设为false。
    • kRopeOut(aclTensor*，计算输出)：由isOutputKv控制，当isOutputKv为true时，需输出。shape为4维[Bkv,N,Skv,Dk]。undefined支持ND。数据类型与输入kv一致。
    • cKvOut(aclTensor*，计算输出)：由isOutputKv控制，当isOutputKv为true时，需输出。shape为4维[Bkv,N,Skv,Dv]。undefined支持ND。数据类型与输入kv一致。
    • workspaceSize(uint64_t*，出参)：返回需要在Device侧申请的workspace大小。
    • executor(aclOpExecutor**，出参)：返回op执行器，包含了算子计算流程。

  • 返回值： aclnnStatus：返回状态码，具体参见undefined。

    [object Object]

aclnnKvRmsNormRopeCache

  • 参数说明：

    • workspace(void*，入参)：在Device侧申请workspace内存地址。
    • workspaceSize(uint64_t，入参)：在Device侧申请的workspace大小，由第一段接口clnnKvRmsNormRopeCacheGetWorkspaceSize获取。
    • executor(aclOpExecutor*，入参)：op执行器，包含了算子计算流程。
    • stream(aclrtStream，入参)：指定执行任务的Stream。

  • 返回值：

    aclnnStatus：返回状态码，具体参见undefined。

约束说明

  • 参数说明里shape格式说明：
    • Bkv为输入kv的batch size，Skv为输入kv的sequence length，大小由用户输入场景决定，无明确限制。
    • N为输入kv的head number。此算子与DeepSeekV3网络结构强相关，仅支持N=1的场景，不存在N非1的场景。
    • D为输入kv的head dim。rms_norm计算所需数据Dv和rope计算所需数据Dk由输入kv的D切分而来。故Dk、Dv大小需满足Dk+Dv=D。同时，Dk需满足rope规则。根据rope规则，Dk为偶数。若cacheModeOptional为NZ场景（cacheModeOptional为PA_NZ、PA_BLK_NZ），Dk、Dv需32B对齐。
    • 若cacheModeOptional为PA场景（cacheModeOptional为PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ），BlockSize需32B对齐。
    • 关于上述32B对齐的情形，对齐值由cache的数据类型决定。以BlockSize为例，若cache的数据类型为int8，则需BlockSize%32=0；若cache的数据类型为float16，则需BlockSize%16=0；若kCacheRef与ckvCacheRef参数的dtype不一致，BlockSize需同时满足BlockSize%32=0和BlockSize%16=0。
    • Bcache为输入cache的batch size，Scache为输入cache的sequence length，大小由用户输入场景决定，无明确限制。
    • BlockNum为写入cache的内存块数，大小由用户输入场景决定，无明确限制。
  • index相关约束：
    • 当cacheModeOptional为Norm时，shape为2维[Bkv,Skv]，要求index的value值范围为[-1,Scache)。不同的Bkv下，value数值可以重复。
    • 当cacheModeOptional为PA_BNSD、PA_NZ时，shape为1维[Bkv * Skv]，要求index的value值范围为[-1,BlockNum * BlockSize)。value数值不能重复。
    • 当cacheModeOptional为PA_BLK_BSND、PA_BLK_NZ时，shape为1维[Bkv * ceil_div(Skv,BlockSize)]，要求index的value的数值范围为[-1,BlockNum * BlockSize)。value/BlockSize的值不能重复。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。

[object Object]