[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 main computation process consists of five paths.

    • After multiplying the input xx with WDQW^{DQ} for downsampling and RmsNorm, the first path multiplies the result with WUQW^{UQ} and WUKW^{UK}, followed by two upsampling operations to obtain qNq^N, while the second path multiplies the result with WQRW^{QR} and applies rotary position encoding (ROPE) to obtain qRq^R.
    • The third path multiplies the input xx with WDKVW^{DKV} for downsampling and RmsNorm, and then passes the result into the cache to obtain kCk^C.
    • The fourth path multiplies the input xx with WKRW^{KR}, applies rotary position encoding (ROPE), and then stores the result into another cache to obtain kRk^R.
    • The fifth path processes the output qNq^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

    RmsNorm(x)=γxiRMS(x)\text{RmsNorm}(x) = \gamma \cdot \frac{x_i}{\text{RMS}(x)} RMS(x)=1Ni=1Nxi2+ϵ\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.

    cQ=RmsNorm(xWDQ)c^Q = RmsNorm(x \cdot W^{DQ}) qC=cQWUQq^C = c^Q \cdot W^{UQ} qN=qCWUKq^N = q^C \cdot W^{UK}

    Performs rotary position encoding (ROPE) to the query.

    qR=ROPE(cQWQR)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.

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

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

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

    Dequant Scale Query Nope calculation formula:

    dequantScaleQNope=RowMax(abs(qN))/127dequantScaleQNope = {RowMax(abs(q^{N})) / 127} qN=round(qN/dequantScaleQNope)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]
[object Object]
[object Object]
  • Parameters

    [object Object]undefined
  • Returns

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

    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.
  • Shape field description

    [object Object]undefined
  • When transposing is not performed, the dimensions of [object Object], [object Object], and [object Object] are represented as [object Object].

  • Shape restrictions:

    • When BS fusion is used for [object Object], that is , [object Object]
      • The shape of [object Object] and [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • In the full quantization scenario, the shape of [object Object] is [object Object]. In other scenarios, the shape is [object Object].
    • When BS fusion is not used for [object Object], that is , [object Object]
      • The shape of [object Object] and [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • The shape of [object Object] is [object Object].
      • In the full quantization scenario, the shape of [object Object] is [object Object]. In other scenarios, the shape is [object Object].
    • 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 [object Object] API supports the following 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 example is for reference only. For details, see .

[object Object]