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

• API function: Inputs the qkv fusion tensor, splits the q, k, and v tensors using SplitVD, performs the RmsNorm, ApplyRotaryPosEmb, Quant, and Scatter fusion operations, and outputs qOut, kCache, vCache, qBeforeQuant (optional), kBeforeQuant (optional), and vBeforeQuant (optional).

• The following table lists the scenarios supported by this API.

  [object Object]undefined

• Formulas:

  (1) SplitVD:

  In the following formula, $N_q$, $N_k$, and $N_v$ indicate the number of attention heads of the q, k, and v components, respectively. The following conditions must be met:

  $$\begin{cases} N_k = N_v \\ N_{qkv} = N_k + N_v + N_q \\ D_{qkv} = D_q = D_k = D_v \end{cases}$$ $$\begin{aligned} q &= qkv[..., [:N_q] * D_{qkv}] \\ k &= qkv[..., [N_q:-N_v] * D_{qkv}] \\ v &= qkv[..., [-N_v:] * D_{qkv}] \end{aligned}$$

  (2) RmsNorm:

  Here, x and y indicate the input and output tensors of RmsNorm, respectively. The normalization is performed along the last dimension (feature dimension). This calculation rule is also applicable to the q and k components.

  $$squareX = x * x$$ $$meanSquareX = squareX.mean(dim = -1, keepdim = True)$$ $$rms = \sqrt{meanSquareX + epsilon}$$ $$y = (x / rms) * gamma$$

  (3) RoPE (Half-and-Half):

  y here refers to the output result of the RmsNorm calculation.

  $$y1 = y[\ldots, :d/2]$$ $$y2 = y[\ldots, d/2:]$$ $$y\_RoPE = torch.cat((-y2, y1), dim = -1)$$ $$y\_embed = (y * cos) + y\_RoPE * sin$$

  (4) Quant:

  Without quantization:

  $$kQuant = kRoPE \\ vQuant = v$$

  Symmetric quantization:

  $$kQuant = kRoPE / kScale \\ vQuant = v / vScale$$

  Asymmetric quantization:

  $$kQuant = kRoPE / kScale + kOffset \\ vQuant = v / vScale + vOffset$$

[object Object]

Each operator is divided into . You must call the aclnnQkvRmsNormRopeCacheGetWorkspaceSize API to obtain the input parameters, calculate the required workspace size based on the process, and then call the aclnnQkvRmsNormRopeCache API to perform computation.

[object Object]

• Parameters:

  [object Object]

• Returns

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

  The first-phase API implements input parameter validation. The following error codes may be returned.

  [object Object]

[object Object]

• Parameters:

  [object Object]

• Returns

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

[object Object]

• Deterministic computation:

  • The default deterministic implementation of aclnnQkvRmsNormRopeCache is used.

• Input shape restrictions:

  • B[object Object]qkv[object Object] indicates the batch size of the input qkv, and S[object Object]qkv[object Object] indicates the sequence length of the input qkv. The size is determined by qkvSize.
  • N[object Object]qkv[object Object] indicates the number of input qkv heads. D[object Object]qkv[object Object] indicates the head dimension of the input qkv. Currently, only 128 is supported. D[object Object]q[object Object], D[object Object]k[object Object], and D[object Object]k[object Object] are the head dimensions of q, k, and v, respectively. It is required that D[object Object]qkv[object Object] = D[object Object]q[object Object] = D[object Object]k[object Object] = D[object Object]v[object Object] and D[object Object]qkv[object Object] be divisible by 32 (D[object Object]qkv[object Object]*qkv data type occupies bytes).
  • According to the rope rule, D[object Object]k[object Object] and D[object Object]q[object Object] must be even numbers. If cacheMode is set to PA_NZ, D[object Object]k[object Object] and D[object Object]q[object Object] must be 32-byte aligned, and BlockSize must be 32-byte aligned.
  • The alignment value in the preceding 32-byte alignment scenario is determined by the data type of the cache. Take BlockSize as an example. If the data type of the cache is int8, BlockSize must be a multiple of 32. If the data type of the cache is float16, BlockSize must be a multiple of 16. If the data types of kCache and vCache are different, BlockSize must be a multiple of both 32 and 16.
  • BlockNum indicates the number of memory blocks written to the cache. The size is determined by the user input scenario and must be BlockNum >= Ceil(S[object Object]qkv[object Object] / BlockSize) * B[object Object]qkv[object Object].
  • The requireMemory parameter indicates the size of the space required for storing data. The following condition must be met: requireMemory >= (B[object Object]qkv[object Object] S[object Object]qkv[object Object] N[object Object]qkv[object Object] D[object Object]qkv[object Object] + 2 D[object Object]qkv[object Object] + 2 B[object Object]qkv[object Object] S[object Object]qkv[object Object] D[object Object]qkv[object Object] + B[object Object]qkv[object Object] S[object Object]qkv[object Object] N[object Object]q[object Object] D[object Object]qkv[object Object] + BlockNum BlockSize N[object Object]v[object Object] D[object Object]qkv[object Object] + BlockNum BlockSize N[object Object]k[object Object] D[object Object]qkv[object Object]) sizeof(FLOAT16) + B[object Object]qkv[object Object] S[object Object]qkv[object Object] sizeof(INT64) + (2 N[object Object]k[object Object] D[object Object]qkv[object Object] + 2 N[object Object]v[object Object]) * sizeof(FLOAT). If the calculated value of requireMemory exceeds the total GM space of the current AI processor, this API cannot be used.

• Other restrictions:

  • For index, the value range of index must be [-1, BlockNum * BlockSize). The value cannot be repeated. When index is -1, the update is skipped.
  • kScaleOptional and vScaleOptional indicate the scaling factors for symmetric quantization. Therefore, if the parameters are passed, the values cannot be 0.

[object Object]

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