[object Object]

[object Object][object Object]undefined[object Object]

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

• 计算公式：

  (1) interleaveRope:

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

  (2) rmsNorm:

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

[object Object]

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

[object Object]

• 参数说明

  [object Object]

• 返回值

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

  第一段接口完成入参校验，出现以下场景时报错：

  [object Object]

[object Object]

• 参数说明

  [object Object]

• 返回值

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

[object Object]

• 参数说明里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、HIFLOAT8、FLOAT8E5M2、FLOAT8E4M3FN，则需BlockSize%32=0；若cache的数据类型为float16或bfloat16，则需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_BNSD、PA_BLK_NZ时，shape为1维[Bkv * ceil_div(Skv,BlockSize)]，要求index的value的数值范围为[-1,BlockNum * BlockSize)。value/BlockSize的值不能重复。
• 量化场景的相关约束：
  • 量化场景支持的情况1：kCacheRef的数据类型为FLOAT16或BFLOAT16，ckvCacheRef的数据类型为INT8、HIFLOAT8、FLOAT8E5M2、FLOAT8E4M3FN。
  • 量化场景支持的情况2：ckvCacheRef的数据类型为FLOAT16或BFLOAT16，kCacheRef的数据类型为INT8、HIFLOAT8、FLOAT8E5M2、FLOAT8E4M3FN。
  • 量化场景支持的情况3：kCacheRef与ckvCacheRef的数据类型一致，为INT8、HIFLOAT8、FLOAT8E5M2、FLOAT8E4M3FN。

[object Object]

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