[object Object]

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

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

• Description: Compatible with the functions of , this API supports positional encoding, paged attention, and KV cache dequantization features.

  For an auto-regressive language model, the length of the inference input increases as new words are generated. Incremental inference is implemented based on full inference. The S axis of [object Object] is fixed at [object Object]. The [object Object] and [object Object] are their respective content in the KV cache plus the state information generated by inference. The actual length of the S axis corresponding to each batch may be different, and the input data is the fixed-length data after padding.

  Compared with the FlashAttention operator () in the full inference scenario, this operator (in the incremental inference scenario) has a different inference process but similar accuracy.

  Note:

  The KV cache is a common technology for optimizing the inference performance of foundation models. During sampling, the transformer model uses the given prompt/context as the initial input for inference (parallel processing supported), and then generates additional tokens one by one to improve the generated sequence (reflecting the auto-regressive property of the model). The transformer performs the self-attention operation during sampling. Therefore, KV vectors need to be extracted for each item (regardless of the prompt/context or generated token) in the current sequence. These vectors are stored in a matrix called KV cache.

• Formula:

  Self-attention constructs an attention model by leveraging the relationships within input samples. The principle assumes there is an input sample sequence $x$ of length $n$, where each element of $x$ is a $d$-dimensional vector. Each $d$-dimensional vector can be regarded as a token embedding. Such a sequence is transformed by three weight matrices to obtain three $n$d matrices.

  The calculation formula for self-attention is generally defined as follows, where $Q$, $K$, and $V$ are key attribute elements of the input sample, obtained through spatial transformation and unified into a single feature space.

  $$Attention(Q,K,V)=Score(Q,K)V$$

  The [object Object] function in this operator employs the [object Object] function. The self-attention calculation formula is as follows:

  $$Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{d}})V$$

  The product of $Q$ and $K^T$ represents the attention to the input $x$. To prevent this value from becoming too large, the product is usually divided by the square root of $d$ for scaling, normalized by softmax for each row, and multiplied by $V$ to result in an $n*d$ matrix.

[object Object]

Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation 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 deterministic implementation.

• [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT:

  • The B axis can be less than or equal to [object Object], the N axis can be less than or equal to [object Object], and the D axis can be less than or equal to [object Object].
  • The data type of [object Object] can be FLOAT16 or BFLOAT16. The data types of [object Object], [object Object], and [object Object] can be FLOAT16, INT8, or BFLOAT16.
  • The data types of [object Object] and [object Object] can be UINT64 or FLOAT32.
  • The data types of [object Object], [object Object], and [object Object] can be FLOAT32.
  • The data type of [object Object] can be INT64.

• [object Object]Atlas inference accelerator cards[object Object]:

  • The B axis can be less than or equal to [object Object], the N axis can be less than or equal to [object Object], and the D axis can be less than or equal to [object Object].
  • The S axis of [object Object] and [object Object] is less than or equal to [object Object].
  • The data types of [object Object], [object Object], [object Object], and [object Object] can only be FLOAT16.
  • The data types of [object Object], [object Object], [object Object], [object Object], and [object Object] can only be [object Object].
  • [object Object] supports only the value [object Object].

• In non-contiguous scenarios, the number of tensors in the tensor list of parameters [object Object] and [object Object] is equal to [object Object] in [object Object] (due to tensor list restrictions, [object Object] must be less than or equal to [object Object] in non-contiguous scenarios). The shapes (except for that of [object Object]) must be identical, and the batch size can only be [object Object].

• The values of [object Object] and [object Object] are the same in [object Object]. The values of [object Object] and [object Object] are the same in [object Object] and [object Object]. The value of [object Object] is a multiple of that of [object Object].

• The S axis of [object Object] can only be [object Object].

• When the data type of [object Object] is INT8 or UINT8, the value in the tensor must be [object Object] or [object Object].

• The scenario where the input types of [object Object], [object Object], and [object Object] are INT8 is not supported.

• When the input type of [object Object], [object Object], and [object Object] is FLOAT16 and the output type is INT8, the input parameter [object Object] is required, and [object Object] is optional. Parameters [object Object], [object Object], and [object Object] (with their values being [object Object]) cannot be passed.

• Restrictions on parameters [object Object] and [object Object]:

  • Per-channel mode: The shapes of the two parameters can be (2, N, 1, D), (2, N, D), (2, H), where [object Object]·is the value of [object Object]. The data type is the same as that of [object Object].
  • Per-tensor mode: The shapes of the two parameters are (2), and the data type is the same as that of [object Object].

• The input parameters [object Object] and [object Object] support both the per-tensor and per-channel modes and the FLOAT32 and BFLOAT16 data types. If [object Object] is passed, ensure that its type and shape are consistent with those of [object Object]. When the input type is BFLOAT16, both FLOAT32 and BFLOAT16 are supported. Otherwise, only FLOAT32 is supported. In per-channel format, when the output layout is BSH, the product of all dimensions of [object Object] must be equal to [object Object]. For other layouts, the product must be equal to [object Object] * [object Object]. (When the output layout is BSH, it is recommended that the shape of [object Object] be set to [1, 1, H] or [H]. When the output layout is BNSD, it is recommended that the shape of [object Object] be set to [1, N, 1, D] or [N, D]. When the output layout is BSND, it is recommended that the shape of [object Object] be set to [1, 1, N, D] or [N, D].)

• Paged attention scenario:

  • The prerequisite for enabling paged attention is that [object Object] exists and is valid. In addition, [object Object] and [object Object] are arranged in a contiguous memory based on the index in the [object Object]. The data types of [object Object] and [object Object] can be FLOAT16, BFLOAT16, or INT8. In this scenario, the [object Object] parameter of [object Object] and [object Object] is invalid.
  • [object Object] is a user-defined parameter. The value of this parameter affects the paged attention performance. When paged attention is enabled, a non-zero value must be passed for [object Object], and the maximum value is [object Object]. If the input type is FLOAT16 or BFLOAT16, [object Object] and [object Object] must be 16-byte aligned. If the input type is INT8, [object Object] and [object Object] must be 32-bytes aligned. The value [object Object] is recommended. Generally, paged attention can improve the throughput but deteriorate the performance.
  • In the paged attention scenario, when [object Object] of [object Object] is BNSD, the KV cache layout can be (blocknum, blocksize, H) or (blocknum, KV_N, blocksize, D). When [object Object] of [object Object] is BSH or BSND, the KV cache layout can only be (blocknum, blocksize, H). The value of [object Object] cannot be less than the total number of blocks of each batch computed based on [object Object] and [object Object]. The shape of [object Object] must be the same as that of [object Object].
  • In the paged attention scenario, the performance is generally better when the KV cache layout is (blocknum, KV_N, blocksize, D) than when it is (blocknum, blocksize, H). Therefore, (blocknum, KV_N, blocksize, D) is recommended.
  • When paged attention is enabled, if the input KV cache layout format is (blocknum, blocksize, H) and the product of [object Object] multiplied by [object Object] exceeds 64K, an error is reported due to hardware instruction restrictions. This problem can be solved by enabling GQA (decreasing [object Object]) or adjusting the KV cache format to (blocknum, numKeyValueHeads, blocksize, D).
  • In the paged attention scenario, [object Object] must be passed.
  • In the paged attention scenario, [object Object] must be two-dimensional. The length of the first dimension must be equal to [object Object], and the length of the second dimension cannot be less than the value of [object Object] (the maximum number of blocks corresponding to [object Object] in each batch).
  • When paged attention is enabled, the input [object Object] must be greater than or equal to the value of [object Object] * [object Object] in the following scenarios:
  • When attention mask is enabled, for example, the mask shape is (B, 1, 1, S).
  • When [object Object] is enabled, for example, the pseShift shape is (B, N, 1, S).

[object Object]

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