[object Object]

Note: This API will be deprecated in later versions. Use the latest API aclnnMlaPrologV3WeightNz instead.

[object Object][object Object]undefined[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 to obtain $q^N$, while the second path multiplies the result with $W^{QR}$ and applies rotary position encoding (ROPE) to obtain $q^R$.
  • The third path multiplies the input $x$ with $W^{DKV}$ for downsampling and RmsNorm, and then passes the result into the cache to obtain $k^C$.
  • 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 calculation formula of the query includes downsampling, RmsNorm, and two upsamplings.

  $$c^Q = 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 = 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} = RmsNorm(x \cdot W^{DKV})$$ $$k^C = Cache(c^{KV})$$

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

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

  Dequant Scale Query Nope calculation formula:

  $$dequantScaleQNope = {RowMax(abs(q^{N})) / 127}$$ $$q^{N} = {round(q^{N} / 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 performs input parameter validation. The following errors may be returned:

  [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 deterministic implementation. Currently, A5 does not fully support the generalization scope of the input parameters and supported scenarios of the MlaProlog API on A2 and A3. The details are as follows:

• Shape field description

  [object Object]undefined

• When transposing is not performed, the dimensions of [object Object], [object Object], and [object Object] are represented as (k, n).

• Shape constraints:

  • When BS fusion is used for [object Object], that is , (T, He)
    • The shape of ropeSin and ropeCos is (T, Dr).
    • The shape of cacheIndex is (T).
    • The shape of dequantScaleXOptional is (T, 1).
    • The shape of queryOut is (T, N, Hckv).
    • The shape of queryRopeOut is (T, N, Dr).
    • In the int8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (T, N, 1). In other scenarios, the shape is (1).
  • When BS fusion is not used for [object Object], that is , (B, S, He)
    • The shape of ropeSin and ropeCos is (B, S, Dr).
    • The shape of cacheIndex is (B, S).
    • The shape of dequantScaleXOptional is (B*S, 1).
    • The shape of queryOut is (B, S, N, Hckv).
    • The shape of queryRopeOut is (B, S, N, Dr).
    • In the int8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (B*S, N, 1). In other scenarios, the shape is (1).
  • 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.

• The aclnnMlaPrologV2WeightNz API supports the following scenarios: A2 and A3 support all the following scenarios, and A5 supports only non-quantization scenarios.

  [object Object]

• In different quantization scenarios, the combination of dtype and shape 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 following is the sample code for A5, which is for reference only.