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

• Interface function: In inference scenarios, this operator performs the preprocessing computation for Multi-Head Latent Attention. The computation process is as follows:

  • First, the input $x$ is processed by RmsNormQuant, then multiplied by $W^{DQKV}$ for downsampling, and finally split into path 1 and path 2.
  • Path 1 performs RmsNormQuant and then multiplies the result by $W^{UQ}$, and then splits the result into paths 3 and 4.
  • Path 3 is multiplied by $W^{uk}$ and then $q^N$ is output.
  • Path 4 is rotated and position-encoded, and then $q^R$ is output.
  • Path 2 is split into path 5 and path 6.
  • Path 5 is passed to the cache after RmsNorm to obtain $k^N$.
  • Path 6 undergoes rotary position encoding (ROPE) and then is stored into another cache to obtain $k^R$.

• Computing flowchart


• Formula:

  RmsNormQuant formula

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

  Query calculation formula, including W^{DQKV} matrix multiplication, W^{UK} matrix multiplication, RmsNormQuant, and ROPE rotation position encoding.

  $$q^N = RmsNormQuant(x) \cdot W^{DQKV} \cdot W^{UK}$$ $$q^R = ROPE(x^Q)$$

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

  $$k^N = Cache({RmsNorm}(RmsNormQuant(x)))$$ $$k^R = Cache(ROPE(RmsNormQuant(x)))$$

[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]

• Returns:

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

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

  [object Object]

[object Object]

• Parameters

  [object Object]

• Returns:

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

[object Object]

• Deterministic computation:
  • [object Object] defaults to a deterministic implementation.
• Meanings and restrictions of the shape format fields
  • [object Object]: batch size of input samples. The value ranges from 0 to 256.
  • [object Object]: size of the hidden layer. The value ranges from 2048 to 10240 and must be a multiple of 256.
  • [object Object]: number of heads. Value range: 16, 32, 64, 128.
  • [object Object]: number of blocks in the PagedAttention scenario. The value range is 192.
  • [object Object]: block size in the PagedAttention scenario. The value range is 128.
  • When the data type of [object Object] and [object Object] is bfloat16, the input must also be bfloat16. In addition, [object Object] can only be 6144, and [object Object] can only be 0 or 1.

[object Object]

The following example is for reference only. For details, see .