[object Object]

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

• 接口功能：融合了MLA（Multi-head Latent Attention）结构中RMSNorm归一化计算（对应$rms\_size$）与RoPE（Rotary Position Embedding）位置编码（对应$rope\_size$)，以及更新KVCache的ScatterUpdate操作。本接口支持两种场景，向下兼容aclnnKvRmsNormRopeCache。

• 支持场景：

  [object Object]undefined
  • [object Object]Ascend 950PR/Ascend 950DT[object Object]：仅支持V1场景。
  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：支持V1和V2场景。

• 计算公式：

  定义输入张量kv的shape为$[Bkv, N, Skv, Dkv]$以及张量vOptional的shape为$[Bkv, N, Skv, Dv]$。

  (1) RmsNorm：

  $$x=kv[...,:rms\_size]$$ $$\operatorname{RmsNorm}(x_i)=\frac{1}{\operatorname{Rms}(\mathbf{x})} * x_i * gamma_i, \quad \text { where } \operatorname{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+epsilon}$$ $$y = \operatorname{RmsNorm}(x)$$

  (2) interleaveRope：

  $$x=\begin{cases} kv[...,Dv:], \quad vOptional = None\\ {\operatorname{RmsNorm}(\mathbf{x})}[...,:rope\_size], \quad vOptional != None \end{cases}$$ $$x1=x[...,::2]$$ $$x2=x[...,1::2]$$ $$x\_part1=torch.cat((x1,x2),dim=-1)$$ $$x\_part2=torch.cat((-x2,x1),dim=-1)$$ $$y\_rope=x\_part1*cos+x\_part2*sin$$ $$rope\_out=\begin{cases}y\_rope, \quad vOptional = None \\ concat(y\_rope, {\operatorname{RmsNorm}(\mathbf{x})}[...,rope\_size:]), \quad vOptional != None \end{cases}$$

  (3) 量化计算：

  x表示将要写入到kCacheRef和ckvCacheRef上的原始数据，作为量化过程的输入。

  $$x = x * scale,\ if\ scale\ !=\ None$$ $$x = x + offset,\ if\ offset\ !=\ None$$ $$y = \begin{cases}x, \quad scale == None \space and \space offset == None \\ round(x).clamp(-128,127), \quad others \end{cases}$$

  (4) Scatter写出：

  输入张量index对应输入kv缓存中各元素的索引映射表，取x中具体元素的索引$b \in Bkv$以及$s \in Skv$，$n$为注意力头索引，

  $$scatter\_idx = index(b, s) \\$$

  Quant表示前述量化计算过程，对原地更新参数k_cache和ckv_cache：

  $$k\_cache[scatter\_idx, ...] = Quant(x = rope\_out, scale = k\_scale, offset = k\_offset)[b, n, s]$$ $$ckv\_cache[scatter\_idx, ...] = \begin{cases} Quant(x = \operatorname{RmsNorm}(x), scale = v\_scale, offset = v\_offset)[b, n, s], \quad vOptional = None \\ Quant(x = vOptional, scale = v\_scale, offset = v\_offset)[b, n, s], \quad vOptional != None \end{cases}$$

  (5) 原始结果写出：

  当$is\_output\_kv=True$且有效时：

  $$k\_rope = rope\_out$$ $$c\_kv = \begin{cases}\operatorname{RmsNorm}(x), \quad vOptional = None \\ vOptional, \quad vOptional != None \end{cases}$$

[object Object]

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

[object Object]

• 参数说明

  [object Object]

• 返回值

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

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

  [object Object]

[object Object]

• 参数说明

  [object Object]

• 返回值

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

[object Object]

• 本算子默认确定性实现。

• 输入shape限制：

  • kv为四维张量，shape为[Bkv,N,Skv,D]，Bkv为输入kv的batch size，Skv为输入kv的sequence length，大小由用户输入场景决定，无明确限制。

  • N为输入kv的head number。V1场景与DeepSeekV3网络结构强相关，仅支持N=1的场景。V2场景支持N=1/2/4/8。

  • D为输入kv的head dim。根据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），block_size需32B对齐。

  • 关于上述32B对齐的情形，对齐值由cache的数据类型决定。以block_size为例，若cache的数据类型为int8，则需block_size%32=0；若cache的数据类型为float16，则需block_size%16=0；若kCacheRef与ckvCacheRef参数的dtype不一致，block_size需同时满足block_size%32=0和block_size%16=0。

  • block_num为写入cache的内存块数，大小由用户输入场景决定，无明确限制。

  • 旋转位置编码（RoPE）参数项(cos, sin)需要满足shape约束：

    • shape允许为4维[Bkv,N,Skv,Dk]或[Bkv,N,1,Dk]。
    • cos与sin的shape必须保持一致。

  • 量化参数项(kRopeScaleOptional, kRopeOffsetOptional, ckvScaleOptional, cKvOffsetOptional)需要满足shape约束：

    • 所有量化参数项的维度数量和N轴尺寸（如果存在）必须保持一致。

    • 不同场景下，量化参数的合法shape约束存在差异：

      [object Object]undefined
      • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：V1场景不支持[1,]，V2场景不支持量化参数项广播。

  • 输入张量均不支持空Tensor。

  • 所有输入均不支持无效值，包括且不限于：±inf，nan。

• cache的数据类型支持：

  • 非量化模式：cache类型必须与kv保持一致。
    • [object Object]Ascend 950PR/Ascend 950DT[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：可支持BFLOAT16、FLOAT16。
  • 量化模式：
    • [object Object]Ascend 950PR/Ascend 950DT[object Object]：可支持INT8、HIFLOAT8、FLOAT8E5M2、FLOAT8E4M3FN。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：仅支持INT8。

• 参数说明：

  • 输入参数中kv, gamma, cos, sin, vOptional的数据类型必须完全一致。

  • kCacheRef和ckvCacheRef是[object Object]原地更新参数[object Object]，它们的数据类型取决于相应的输入分量，以及相应的scale和offset。详情见下：

    [object Object]undefined
    • 非量化模式时，量化参数(kRopeScaleOptional, kRopeOffsetOptional, ckvScaleOptional, cKvOffsetOptional)必须设为None，且kCacheRef和ckvCacheRef的dtype必须与kv保持一致。
    • 量化模式时，kCacheRef和ckvCacheRef的dtype应为相应产品上支持的数据类型。

  • 输入分量关于量化因子scale与量化偏移scale的对应关系如下：

    [object Object]undefined
    • kCacheRef：量化系数为kRopeScaleOptional和kRopeOffsetOptional。
    • ckvCacheRef：对应量化系数为ckvScaleOptional和cKvOffsetOptional。

  • 输出参数中，k_rope和c_kv的类型必须与kv保持一致。

• 量化模式约束：

  [object Object]
  • 静态量化模式支持细节：
    • Ascend 950PR/Ascend 950DT产品：仅支持V1场景，支持[object Object]静态对称量化[object Object]和[object Object]静态非对称量化[object Object]。
    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：
      • V1场景：对除[object Object]以外的cachemode，仅支持[object Object]静态对称量化[object Object]。即使传入合法的offset，也不会被算子处理，仍视为[object Object]静态对称量化[object Object]。对[object Object]模式的cachemode，不支持任何静态量化。
      • V2场景：对所有cachemode，支持[object Object]静态对称量化[object Object]和[object Object]静态非对称量化[object Object]。

• cache与index相关约束：

  [object Object]undefined

  • Scache为输入cache的sequence length，大小由用户输入场景决定，无明确限制。

  • 当cacheModeOptional为Norm时，shape为2维[Bkv,Skv]，要求index的value值范围为[-1,Scache)。不同的Bkv下，value数值可以重复。

  • 当cacheModeOptional为PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ时，cache中的数据排布方式为：

    • 非量化模式下：kCacheRef 为 [block_num, Dk//16, block_size, 1, 16]；ckvCacheRef 为 [block_num, Dv//16, block_size, 1, 16]。
    • 静态量化模式下：kCacheRef 为 [block_num, Dk//32, block_size, 1, 32]；ckvCacheRef 为 [block_num, Dv//32, block_size, 1, 32]。

  • 当cacheModeOptional为PA_BNSD、PA_NZ时，shape为1维[Bkv Skv]，要求index的value值范围为[-1,block_num block_size)。value数值不能重复。

  • 当cacheModeOptional为PA_BLK_BNSD、PA_BLK_NZ时，shape为1维[Bkv ceil_div(Skv,block_size)]，要求index的value的数值范围为[-1,block_num block_size)。value/block_size的值不能重复。

• isOutputKv约束：

  • 作用是输出具体场景的中间处理结果，使能 kRopeOut 和 cKvOut 两项输出。
  • 在cacheModeOptional为PA, PA_BNSD, PA_NZ, PA_BLK_BNSD, PA_BLK_NZ模式时有效。
  • 在cacheModeOptional为Norm时，仅在V2场景中使能量化模式时有效。

• 输入组合约束：

  • 本约束内条目，仅适用于[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]。

  • 在所有cacheModeOptional下，必定支持[object Object]：即[object Object]三个维度与[object Object]严格一致的[object Object]和[object Object]。

  • 在各种cacheModeOptional下，支持的合法输入shape模式如下表：

    [object Object]

  • [object Object]未在上表覆盖范围内的输入组合，将导致算子的未定义行为[object Object]。

• vOptional：

  • 该参数仅限aclnnKvRmsNormRopeCacheV2接口，aclnnKvRmsNormRopeCache接口不支持该参数！

  • 该参数仅限[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]。

    • 该参数仅在[object Object]kv分离场景(V2)[object Object]中作为必须入参，在其他类型中会作为无效参数被忽略。
    • 当vOptional存在时，它的类型必须与kv一致，[object Object]维度也必须与kv一致。

  • Ascend 950PR/Ascend 950DT：不会拦截该参数，但实际功能不支持，也不会处理该参数。

[object Object]

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