aclnnMlaProlog

产品支持情况

[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$。

• 计算公式：

  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}))$$

函数原型

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

aclnnMlaPrologGetWorkspaceSize

• 参数说明

  [object Object]undefined

• 返回值

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

  [object Object]undefined

aclnnMlaProlog

• 参数说明

  [object Object]undefined

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

约束说明

• 确定性计算：

  • aclnnMlaProlog默认确定性实现。

• shape 格式字段含义说明

  [object Object]undefined

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

• aclnnMlaProlog接口支持场景：

  [object Object]

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

  [object Object][object Object]

调用示例

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