aclnnMlaPrologV2WeightNz

产品支持情况

[object Object]undefined

功能说明

• 接口功能：推理场景，Multi-Head Latent Attention前处理的计算。主要计算过程分为五路；

  • 首先对输入$x$乘以$W^{DQ}$进行下采样和RmsNorm后分为两路，第一路乘以$W^{UQ}$和$W^{UK}$经过两次上采样后得到$q^N$；第二路乘以$W^{QR}$后经过旋转位置编码（ROPE）得到$q^R$。
  • 第三路是输入$x$乘以$W^{DKV}$进行下采样和RmsNorm后传入Cache中得到$k^C$；
  • 第四路是输入$x$乘以$W^{KR}$后经过旋转位置编码后传入另一个Cache中得到$k^R$；
  • 第五路是输出$q^N$经过DynamicQuant后得到的量化参数。
  • 权重参数WeightDq、WeightUqQr和WeightDkvKr需要以NZ格式传入

• 计算公式：

  RmsNorm公式

  $$\text{RmsNorm}(x) = \gamma \cdot \frac{x_i}{\text{RMS}(x)}$$ $$\text{RMS}(x) = \sqrt{\frac{1}{N} \sum_{i=1}^{N} x_i^2 + \epsilon}$$

  Query的计算公式，包括下采样、RmsNorm和两次上采样

  $$c^Q = RmsNorm(x \cdot W^{DQ})$$ $$q^C = c^Q \cdot W^{UQ}$$ $$q^N = q^C \cdot W^{UK}$$

  对Query进行ROPE旋转位置编码

  $$q^R = ROPE(c^Q \cdot W^{QR})$$

  Key的计算公式，包括下采样和RmsNorm，将计算结果存入cache

  $$c^{KV} = RmsNorm(x \cdot W^{DKV})$$ $$k^C = Cache(c^{KV})$$

  对Key进行ROPE旋转位置编码，并将结果存入cache

  $$k^R = Cache(ROPE(x \cdot W^{KR}))$$

  Dequant Scale Query Nope 计算公式：

  $$dequantScaleQNope = {RowMax(abs(q^{N})) / 127}$$ $$q^{N} = {round(q^{N} / dequantScaleQNope)}$$

函数原型

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

aclnnMlaPrologV2WeightNzGetWorkspaceSize

• 参数说明

  [object Object]undefined

• 返回值

  aclnnStatus：返回状态码，具体参见。[object Object] 第一段接口完成入参校验，出现以下场景时报错：

  [object Object]undefined

aclnnMlaPrologV2WeightNz

• 参数说明

  [object Object]undefined

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

约束说明

• 确定性计算：

  • aclnnMlaPrologV2WeightNz默认确定性实现。

• shape 格式字段含义说明

  [object Object]undefined

• weight_dq，weight_uq_qr，weight_dkv_kr在不转置的情况下各个维度的表示：（k，n）。

• shape约束

  • 若tokenX的维度采用BS合轴，即(T, He)
    • ropeSin和ropeCos的shape为(T, Dr)
    • cacheIndex的shape为(T)
    • dequantScaleXOptional的shape为(T, 1)
    • queryOut的shape为(T, N, Hckv)
    • queryRopeOut的shape为(T, N, Dr)
    • 全量化场景下，dequantScaleQNopeOutOptional的shape为(T, N, 1)，其他场景下为(1)
  • 若tokenX的维度不采用BS合轴，即(B, S, He)
    • ropeSin和ropeCos的shape为(B, S, Dr)
    • cacheIndex的shape为(B, S)
    • dequantScaleXOptional的shape为(B*S, 1)
    • queryOut的shape为(B, S, N, Hckv)
    • queryRopeOut的shape为(B, S, N, Dr)
    • 全量化场景下，dequantScaleQNopeOutOptional的shape为(B*S, N, 1)，其他场景下为(1)
  • B、S、T、Skv值允许一个或多个取0，即Shape与B、S、T、Skv值相关的入参允许传入空Tensor，其余入参不支持传入空Tensor。
    • 如果B、S、T取值为0，则queryOut、queryRopeOut输出空Tensor，kvCacheRef、krCacheRef不做更新。
    • 如果Skv取值为0，则queryOut、queryRopeOut、dequantScaleQNopeOutOptional正常计算，kvCacheRef、krCacheRef不做更新，即输出空Tensor。

• aclnnMlaPrologV2WeightNz接口支持场景：

  [object Object]

• 在不同量化场景下，参数的dtype和shape组合需要满足如下条件：

  [object Object]

调用示例

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