[object Object]

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

• Function updates (compared with [object Object]):

  • The scale correction factors of the query and key are added, corresponding to qcQrScale ($\alpha_q$) and kcScale ($\alpha_{kv}$), respectively.
  • Optional input parameters (such as actualSeqLenOptional, kNopeClipAlphaOptional, queryNormFlag, weightQuantMode, kvCacheQuantMode, queryQuantMode, ckvkrRepoMode, quantScaleRepoMode, tileSize, [object Object], and [object Object]) are added, and [object Object] is changed from mandatory to optional.
  • The name and position of the [object Object] parameter are adjusted to match the current [object Object].

• Interface function: In inference scenarios, this operator performs the preprocessing computation for Multi-Head Latent Attention. The main computation process consists of five paths.

  • After multiplying the input $x$ with $W^{DQ}$ for downsampling and RmsNorm, the first path multiplies the result with $W^{UQ}$ and $W^{UK}$, followed by two upsampling operations, and then multiplies the result with the query scale correction factor $\alpha_q$ to obtain $q^N$. The second path multiplies the result with $W^{QR}$ and applies rotary position encoding (ROPE) to obtain $q^R$.
  • The third path is to multiply the input $x$ by $W^{DKV}$, perform downsampling and RmsNorm, and then multiply the result by the key scale correction factor $\alpha_{kv}$ to obtain $k^C$ and pass it to the cache.
  • The fourth path multiplies the input $x$ with $W^{KR}$, applies rotary position encoding (ROPE), and then stores the result into another cache to obtain $k^R$.
  • The fifth path processes the output $q^N$ through DynamicQuant to generate quantization parameters.
  • The weight parameters [object Object], [object Object], and [object Object] are required to be provided in NZ format.

• Formula:

  RmsNorm formula

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

  The computation formula of the query, including downsampling, RmsNorm, and two upsampling operations.

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

  Performs rotary position encoding (ROPE) to the query.

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

  The computation formula of the key, including downsampling and RmsNorm. The computation result is stored in the cache.

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

  Performs rotary position encoding (ROPE) to the key and stores the result in the cache.

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

  Dequant Scale Query Nope calculation formula:

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

[object Object]

Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters

  [object Object]undefined

• Returns

  [object Object]: status code. For details, see .

  The first-phase API implements input parameter verification. The following errors may be thrown.

  [object Object]undefined

[object Object]

• Parameters

  [object Object]undefined

• Returns

  [object Object]: status code. For details, see .

[object Object]

• Deterministic computation:

  • [object Object] defaults to a non-deterministic implementation. You can call [object Object] to enable deterministic computation.

• Shape field description

  [object Object]undefined

• Shape restrictions:

  • When BS fusion is used for [object Object], that is , (T, He)
    • The shape of ropeSin and ropeCos is (T, Dr).
    • When CacheMode is set to PA_BSND or PA_NZ, the shape of cacheIndex is (T).
    • When CacheMode is set to PA_BLK_BSND or PA_BLK_NZ, the shape of cacheIndex is (Sum(Ceil(S_i/BlockSize))), where S_i is the length of S in each batch.
    • When CacheMode is set to PA_BLK_BSND or PA_BLK_NZ, actualSeqLenOptional needs to be passed, and the dimension is (B).
    • In the int8/fp8/hif8 full quantization scenario, the shape of dequantScaleXOptional is (T, 1). In the mxfp8 full quantization scenario, the shape of dequantScaleXOptional is (T, He/32).
    • The shape of queryOut is (T, N, Hckv).
    • The shape of queryRopeOut is (T, N, Dr).
    • In the int8/mxfp8/fp8/hif8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (T, N, 1). In other scenarios, the shape is nullptr.
  • When BS fusion is not used for [object Object], that is , (B, S, He)
    • The shape of ropeSin and ropeCos is (B, S, Dr).
    • When CacheMode is set to PA_BSND or PA_NZ, the shape of cacheIndex is (B, S).
    • When CacheMode is set to PA_BLK_BSND or PA_BLK_NZ, the shape of cacheIndex is (B, Ceil(S/BlockSize)).
    • In the int8/fp8/hif8 full quantization scenario, the shape of dequantScaleXOptional is (B*S, 1). In the mxfp8 full quantization scenario, the shape of dequantScaleXOptional is (B*S, He/32).
    • The shape of queryOut is (B, S, N, Hckv).
    • The shape of queryRopeOut is (B, S, N, Dr).
    • In the int8/mxfp8/fp8/hif8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (B*S, N, 1). In other scenarios, the shape is nullptr.
  • One or more of the B, S, T, and Skv values can be 0. That is, input parameters related to the shape and B, S, T, and Skv values can be empty tensors. Other input parameters do not support empty tensors.
    • If the values of B, S, and T are 0, [object Object] and [object Object] output empty tensors, and [object Object] and [object Object] are not updated.
    • If the value of Skv is 0, [object Object], [object Object], and [object Object] are calculated normally, and [object Object] and [object Object] are not updated. That is, an empty tensor is output.
  • When CacheMode is BSND
    • BS fusion should not be used for [object Object], that is, the dimensions are (B, S, He).
    • The dimension of kvCache is (B, S, Nkv, Dr).
  • When CacheMode is TND
    • BS fusion should be used for [object Object], that is, the dimension is (T, He).
    • The dimension of kvCache is (T, Nkv, Dr).
  • When ckvkrRepoMode is 1
    • The dimension of krCache must contain 0, and the supported dimension is (0).

• Special constraints

  • When actualSeqLenOptional is passed, the last number of actualSeqLenOptional must be the same as T.
  • In per-tile quantization mode, both ckvkrRepoMode and quantScaleRepoMode must be 1. In other quantization modes and non-quantization scenarios, both ckvkrRepoMode and quantScaleRepoMode must be 0.
  • In per-tile quantization mode, [object Object] supports only [object Object], [object Object], and [object Object].
  • When the value of ckvkrRepoMode is 1, krCache must be an empty tensor (that is, the product of the shape is 0).
  • In kvcache per-tensor quantization mode, both kvCacheQuantMode and queryQuantMode must be 1.

• The aclnnMlaPrologV3WeightNz API supports the following scenarios: A2 and A3 do not support full quantization scenarios of fp8/hif8/mxfp8. A5 supports all quantization scenarios.

  [object Object]

• In different quantization scenarios, the dtype combinations of parameters must meet the following conditions:

  [object Object]

[object Object]

The following sample code for A2 and A3 is for reference only. For details, see .

The sample code for A5 is as follows: