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

• A FlashAttention operator adapted for both full and incremental inference scenarios. It supports both full computation ([object Object]) and incremental computation ([object Object]). When [object Object] of the query matrix is 1, the IncreFlashAttention branch is executed. In other scenarios, the PromptFlashAttention branch is executed.

• Formula

  Self-attention constructs an attention model by leveraging the relationships within the input samples. The principle assumes 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 produce three matrices of $n*d$.

  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" in the formula and operator name is an abbreviation for "self-attention."

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

  The score function in this operator employs the softmax 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 excessively large, it is typically scaled by dividing by the square root of $d$, followed by row-wise softmax normalization. The result is then multiplied by $V$ to produce an $n*d$ matrix.

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

• When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.

• Processing logic for a null input parameter: The operator checks whether [object Object] is a null pointer. If so, an error is reported. If [object Object] is not an empty tensor but [object Object] and [object Object] are empty tensors (that is, S2 is 0), [object Object] is filled with all zeros. If [object Object] is an empty tensor, the framework will process it. For other input parameters which support the passing of null pointers as described in the preceding parameter description, no processing is performed when they are null pointers.

• The shapes of the tensors corresponding to the [object Object] and [object Object] parameters must be identical. In non-contiguous scenarios, the batch size in the tensor lists of [object Object] and [object Object] can only be 1, and the number of elements must be equal to the batch size (B) of [object Object]. The N and D dimensions must be the same. Due to the tensor list restrictions, B cannot be greater than 256 in non-contiguous scenarios.

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

  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The data type can be FLOAT16, BFLOAT16, INT8, or INT4 (INT32).

• Restrictions on [object Object]

  • If [object Object] is not 1, [object Object] must be of type FLOAT16 or INT8 when [object Object] is of type FLOAT16, and [object Object] must be of type BFLOAT16 when [object Object] is of type BFLOAT16. The input shape must be (B, N, Q_S, KV_S) or (1, N, Q_S, KV_S), where [object Object] is the S dimension in the shape of [object Object], and [object Object] is the S dimension in the shapes of [object Object] and [object Object]. For scenarios where [object Object] of [object Object] is not 32-byte aligned, it is recommended that the value be padded to 32 bytes to improve performance. There is no requirement on the padding value of the extra part.
  • If [object Object] is 1, [object Object] must be of type FLOAT16 when [object Object] is of type FLOAT16, and [object Object] must be of type BFLOAT16 when [object Object] is of type BFLOAT16. The input shape must be (B, N, 1, KV_S) or (1, N, 1, KV_S), where [object Object] is the S dimension in the shapes of [object Object] and [object Object]. For scenarios where [object Object] of [object Object] is not 32-byte aligned, it is recommended that the value be padded to 32 bytes to improve performance. There is no requirement on the padding value of the extra part.

• [object Object]: masks the QK product to define the attention visibility between tokens.

  • If [object Object] is not 1, the recommended input shapes are (Q_S, KV_S), (B, Q_S, KV_S), (1, Q_S, KV_S), (B, 1, Q_S, KV_S), and (1, 1, Q_S, KV_S).
  • If [object Object] is 1, the recommended input shapes are (B, KV_S), (B, 1, KV_S), and (B, 1, 1, KV_S).
  • [object Object] refers to the S dimension in the shape of [object Object], while [object Object] refers to the S dimension in the shapes of [object Object] and [object Object]. If [object Object] or [object Object] is not 16- or 32-byte aligned, it can be rounded up to the nearest aligned value.

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

• [object Object]: aclIntArray on the host, indicating the valid sequence lengths of [object Object] in different batches. The data type can be INT64. If you do not specify [object Object], [object Object] is passed, indicating that the length is the same as that of [object Object] in the shape of [object Object]. Restriction: The valid sequence length of each batch in the input parameter must be less than or equal to that of the corresponding batch in [object Object]. This parameter is invalid when [object Object] is [object Object]. If the input length of [object Object] is 1, all batches use the same [object Object]. If the input length is greater than or equal to the batch size, the first N elements (where N equals the batch size) of [object Object] are used. Other lengths are not supported.

• [object Object]: aclIntArray on the host. It can be set to [object Object], indicating the valid sequence lengths of [object Object]/[object Object] in different batches. The data type can be INT64. If you do not specify [object Object], [object Object] is passed, indicating that the length is the same as that of [object Object] in the shape of [object Object]/[object Object]. Restriction: The valid sequence length of each batch in the input parameter must be less than or equal to that of the corresponding batch in [object Object]/[object Object]. If the input length of [object Object] is 1, all batches use the same [object Object]. If the input length is greater than or equal to the batch size, the first N elements (where N equals the batch size) of [object Object] are used. Other lengths are not supported.

• Constraints on the number of input parameters and input and output data formats related to INT8 quantization:

  • If both the input and output are of the INT8 type, the input parameters [object Object], [object Object], [object Object], and [object Object] must exist at the same time. [object Object] is optional and defaults to 0 if not passed.
  • If the input is of the INT8 type and the output is of the FLOAT16 type, the input parameters [object Object], [object Object], and [object Object] must exist at the same time. If the input parameter [object Object] or [object Object] exists (not [object Object]), an error is reported and returned.
  • When the input is of the FLOAT16 or BFLOAT16 type and the output is of the INT8 type, the input parameter [object Object] must exist, and [object Object] is optional (defaults to 0 if not passed). If the input parameter [object Object], [object Object], or [object Object] exists (not [object Object]), an error is reported and returned.
  • The input parameters [object Object] and [object Object] support both the per-tensor and per-channel formats 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 is of the BFLOAT16 type, 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 H. For other layouts, the product must be equal to N × D. (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].)

• Constraints on the fake-quantization parameters [object Object] and [object Object]:

  • Only the fake-quantization scenario where [object Object] is INT8 is supported.
  • Per-channel mode: The shapes of the two parameters can be (2, N, 1, D), (2, N, D), (2, H), where N is [object Object]. The data type is the same as that of [object Object], and [object Object] is set to [object Object].
  • Per-tensor mode: The shapes of the two parameters are (2), the data type is the same as that of [object Object], and [object Object] is set to [object Object].
  • Per-token mode: The shapes of the two parameters are (2, B, S), the data type is fixed at FLOAT32, and [object Object] is set to [object Object].
  • In asymmetric quantization mode, both [object Object] and [object Object] must be present.
  • In symmetric quantization mode, [object Object] can be [object Object]. If [object Object] is [object Object], symmetric quantization is performed. Otherwise, asymmetric quantization is performed.
  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: When [object Object] is 1, the data type can be FLOAT16, BFLOAT16, or FLOAT32. If [object Object] is greater than or equal to 2, only FLOAT16 is supported.

• [object Object]: aclTensor on the device. The data type can be FLOAT16, BFLOAT16, or FLOAT32. The data format can be ND. When the KV fake-quantization parameters are separated, this parameter indicates the dequantization offset of [object Object]. If this function is used, the data type and shape must be the same as those of [object Object]. If this function is not used, pass [object Object].

  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: When [object Object] is greater than or equal to 2, only the per-token and per-channel modes are supported. When [object Object] is 1, only the per-tensor, per-channel, and per-token modes are supported. The per-tensor mode can be combined with the per-head mode, the per-token mode can be combined with the per-head mode, the per-token mode can be combined with the page attention mode to manage scale/offset, and the per-token mode can be combined with the per-head mode and the page attention mode to manage scale/offset.

• [object Object]: aclTensor on the device. The data format can be ND. When the KV fake-quantization parameters are separated, this parameter indicates the dequantization factor of [object Object]. When [object Object] is greater than or equal to 2, only the per-token and per-channel modes are supported. If this function is not used, pass [object Object].

  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The data type can be FLOAT16, BFLOAT16, or FLOAT32. The per-tensor, per-channel, and per-token modes are supported. The per-tensor mode can be combined with the per-head mode, the per-token mode can be combined with the per-head mode, the per-token mode can be combined with the page attention mode to manage scale/offset, and the per-token mode can be combined with the per-head mode and the page attention mode to manage scale/offset.

• [object Object]: aclIntArray on the host. [object Object] can be passed, indicating the valid sequence length of [object Object]/[object Object]. The data type can be INT64. If [object Object] is not specified, [object Object] can be passed, indicating that the length is the same as that of [object Object] in [object Object]/[object Object]. Restriction: The valid sequence length in this input parameter must be less than or equal to that in [object Object]/[object Object].

• [object Object]: data layout of the input [object Object], [object Object], and [object Object]. Currently, BSH, BSND, BNSD, and BNSD_BSND are supported. (If the input layout is BNSD, the output layout is BSND, and only [object Object] greater than 1 is supported.) If no specific layout is required, BSH is recommended.

  • Note: The data layout of [object Object], [object Object], and [object Object] can be interpreted from multiple dimensions. B (Batch) indicates the batch size of input samples, S (Seq-Length) indicates the sequence length of input samples, H (Head-Size) indicates the size of the hidden layer, N (Head-Num) indicates the number of heads, and D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N).

• Restrictions on [object Object]: [object Object] must be exactly divided by [object Object], and the ratio of [object Object] to [object Object] cannot be greater than 64. In the BSND, BNSD, and BNSD_BSND scenarios, the value must be the same as the N-axis value of [object Object]/[object Object] in the shape. Otherwise, an exception occurs.

• Restrictions on [object Object]

  [object Object]

• Restrictions on [object Object]

  • There are four modes (0, 1, 2, and 3) in total, represented by 2-bit combinations. Bit 0 indicates whether to use the high-precision or high-performance mode, and bit 1 indicates whether to perform invalid row correction.

    [object Object]

  • Note: The high-precision and high-performance modes are applicable to both BFLOAT16 and INT8. Invalid row correction takes effect for FLOAT16, BFLOAT16, and INT8. The values [object Object] and [object Object] are reserved. If the masks involved in the computation are all 1s, the precision may be affected. In this case, you can set this parameter to [object Object] or [object Object] to enable invalid row correction to improve the precision. However, this configuration deteriorates the performance. If the operator can determine that invalid rows exist, the invalid row correction is automatically enabled, such as in scenarios where [object Object] is set to [object Object] and Sq is greater than Skv.

• Restrictions on [object Object]

  [object Object]
  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: When [object Object] is 1, passing a value other than 0, 1, 2, 3, 4, or 5 will result in an execution error. If [object Object] is greater than or equal to 2, only the values 0 and 1 are supported. Other values will result in an execution error.

• Restrictions on [object Object]

  • Except for the scenario where [object Object] is 0 and [object Object] is 1, the value must be the same as that of [object Object].
  • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: When [object Object] is 1, passing a value other than 0, 1, 2, 3, 4, or 5 will result in an execution error. If [object Object] is greater than or equal to 2, only the values 0 and 1 are supported. Other values will result in an execution error.

• Restrictions on [object Object]

  • In the ring attention algorithm, the product of query and key is first processed to obtain softmax_max. This max value is subtracted from the product before calculating the exponential, which is then summed to yield softmax_sum. Finally, the log of softmax_sum is added back to softmax_max to obtain the final result.
  • When [object Object] is [object Object], the shape must be [B, N, Q_S, 1] in general. When [object Object] is [object Object] or [object Object], the shape must be [T, N, 1].
  • When [object Object] is [object Object], if the [object Object] tensor is not [object Object], the tensor data is returned directly. If [object Object] is [object Object], a tensor of shape {1} filled with zeros is returned.

• When Q_S is greater than 1

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

    • The B axis must be less than or equal to 65536.

    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: If the input type is INT8 and the D axis is not 32-byte aligned, the maximum value of the B axis is 128. If the input type is FLOAT16 or BFLOAT16 and the D axis is not 16-byte aligned, the maximum value of the B axis is also 128.

    • The N axis can be less than or equal to 256, and the D axis can be less than or equal to 512. If [object Object] is BSH or BSND, it is recommended that N × D be less than 65535.

    • The S axis must be less than or equal to 20971520 (20M). In some long sequence scenarios, if the computation load is too large, the PFA operator execution may time out (an AI Core error is reported, and [object Object] is [object Object]). In this case, S axis splitting is recommended. Note: The computation load is affected by parameters such as B, S, N, and D. Larger values indicate larger computation loads. The following lists some typical scenarios with long sequences (that is, the product of B, S, N, and D is large).

      [object Object]

    • Restrictions on the D axis: For the [object Object]Atlas A2 training products/Atlas A2 inference products[object Object], if [object Object], [object Object], [object Object], or [object Object] is INT8, the D axis must be 32-element aligned. If [object Object], [object Object], [object Object], or [object Object] is INT4, the D axis must be 64-element aligned. If they are all FLOAT16 or BFLOAT16, the D axis must be 16-element aligned.

  • Currently, [object Object] can only be set to 0, 1, 2, 3, or 4. An error will be reported if it is set to other values.

    • When [object Object] is set to 0, if [object Object] is a null pointer or is passed in the left padding scenario, the input parameters [object Object] and [object Object] are ignored.
    • When [object Object] is set to 2, 3, or 4, the shape of [object Object] must be (S, S), (1, S, S), or (1, 1, S, S), where S must be fixed at 2048. Additionally, the input[object Object] must be a lower triangular matrix. If [object Object] is not passed or the passed shape is incorrect, an error will be reported.
    • When [object Object] is set to 1, 2, or 3, the input parameters [object Object] and [object Object] are ignored, and their values are assigned based on related rules.

  • In the synthesis parameter scenario of KV cache dequantization, only when [object Object] is of the FLOAT16 type, [object Object] and [object Object] of the INT8 type can be dequantized to FLOAT16. If the product of the data ranges of the input [object Object] and [object Object] and the data range of the input [object Object] must be within the range of (–1, 1), the high-performance mode can ensure precision. Otherwise, the high-precision mode needs to be enabled to ensure precision.

  • Page attention scenario

    • The prerequisite for enabling page attention is that [object Object] exists and is valid, and [object Object] and [object Object] are arranged in a continuous memory based on the indexes in [object Object]. In this scenario, [object Object] of [object Object] and [object Object] is invalid. [object Object] is filled with block IDs. Currently, the validity of block IDs is not verified. You need to ensure the validity of block IDs.
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The data types of [object Object] and [object Object] can be FLOAT16 or BFLOAT16.
    • [object Object] is a user-defined parameter. Its value affects the page attention performance. When page attention is enabled, the value of [object Object] must be a multiple of 128, ranging from 128 to 512. Generally, page attention can improve the throughput but deteriorate the performance.
    • In the page attention scenario, if the input KV cache layout is BnBsH (blocknum, blocksize, H) and the product of [object Object] multiplied by [object Object] exceeds 65535, an error will be reported due to hardware instruction restrictions. This problem can be solved by enabling GQA (decreasing [object Object]) or adjusting the KV cache layout to BnNBsD (blocknum, KV_N, blocksize, D). When [object Object] of [object Object] is BNSD or TND, the KV cache layout can beBnBsH or BnNBsD. When [object Object] of [object Object] is BSH or BSND, the KV cache layout can only be BnBsH. The value of [object Object] cannot be less than the sum of blocks in each batch calculated based on [object Object] and [object Object]. The shapes of [object Object] and [object Object] must be the same.
    • Page attention fake-quantization scenario
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The data type of [object Object] can be FLOAT16 or BFLOAT16, and the data types of [object Object] and [object Object] can be INT8.
    • Page attention full-quantization scenario
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The data type of [object Object] cannot be INT8.
    • Page attention does not support the tensor list or left padding.
    • In the page attention scenario, [object Object] must be passed.
    • In the page attention scenario, [object Object] must be two-dimensional. The length of the first dimension must be equal to B, and the length of the second dimension must be greater than or equal to [object Object] (the maximum number of blocks corresponding to [object Object] in different batches).
    • When page attention is enabled, the input [object Object] must be greater than or equal to [object Object] × [object Object] in the following scenarios:
      • When [object Object] is passed, the mask shape is (B, 1, Q_S, KV_S).
      • When [object Object] is passed, the [object Object] shape is (B, N, Q_S, KV_S).

  • Left padding for [object Object]

    • The transfer start point of [object Object] is calculated as follows: [object Object] – [object Object] – [object Object]. The transfer end point of [object Object] is calculated as follows: [object Object] – [object Object]. The transfer start point of [object Object] cannot be less than 0, while the end point cannot be greater than [object Object]. Otherwise, the result will not meet the expectation.
    • If [object Object] is less than 0, it will be set to 0.
    • It must be enabled together with [object Object]. Otherwise, the default scenario is right padding for [object Object].
    • It does not support page attention and cannot be enabled together with [object Object].
    • It is not supported when Q is BF16/FP16 and KV is INT4.

  • Left padding for [object Object]:

    • The transfer start point of [object Object] and [object Object] is calculated as follows: [object Object] – [object Object] – [object Object]. The transfer end point of [object Object] and [object Object] is calculated as follows: [object Object] – [object Object]. The transfer start point of [object Object] and [object Object] cannot be less than 0, while the end point cannot be greater than [object Object]. Otherwise, the result will not meet the expectation.
    • If [object Object] is less than 0, it will be set to 0.
    • It must be enabled together with [object Object]. Otherwise, the default scenario is right padding for [object Object].
    • It does not support page attention and cannot be enabled together with [object Object].
    • It is not supported when Q is BF16/FP16 and KV is INT4.

  • When the output is of type INT8 and [object Object] and [object Object] are per-channel, left padding, ring attention, or non-32-byte alignment of D is not supported.

  • When the output is of type INT8, [object Object] cannot be [object Object] and [object Object] or [object Object] cannot be negative.

  • Restrictions on [object Object]

    • This function is supported when the data type of [object Object] is FLOAT16, BFLOAT16, or INT8.
    • When the data type of [object Object] is FLOAT16 and [object Object] exists, the high-precision mode is forcibly used. The corresponding restrictions are the same as those of the high-precision mode.
    • [object Object] must be greater than or equal to [object Object] of [object Object], and [object Object] must be greater than or equal to [object Object] of [object Object]. In the prefix scenario, [object Object] must be greater than or equal to the sum of [object Object] and [object Object] of [object Object].

  • When the output of type INT8, if the input parameter[object Object] is a non-null pointer and a non-null tensor, and [object Object], [object Object], and [object Object] meet the following conditions, certain rows of the matrix will not be involved in computation, resulting in a computation result error. In this scenario, the computation will be intercepted. (Solution: To prevent interception, perform post-quantization outside the FIA interface.)

    • When [object Object] is 0 and [object Object] is a non-null pointer, interception occurs if [object Object] – [object Object] – [object Object] – [object Object] > 0 or [object Object] < 0 for any batch.
    • When [object Object] is 1 or 2, interception does not occur.
    • When [object Object] is 3, interception occurs if [object Object] + [object Object] – [object Object] < 0 for any batch.
    • When [object Object] is 4, interception occurs if [object Object] < 0 or [object Object] + [object Object] + [object Object] – [object Object] < 0 for any batch.

  • Restrictions on prefix parameters

    • Both [object Object] and [object Object] must be either null or non-null.
    • If neither [object Object] nor [object Object] is null, the dimensions and data types of [object Object], [object Object], [object Object], and [object Object] must be the same.
    • If neither [object Object] nor [object Object] is null, the first dimension (batch) of the shape of [object Object] must be 1. When the layout is BNSD or BSND, the N and D axes must be the same as those of [object Object]. When the layout is BSH, the H axis must be the same as that of [object Object]. The same rules apply to [object Object]. [object Object] of [object Object] and [object Object] must be the same.
    • When [object Object] exists, its shape must be [1], and its value cannot be greater than [object Object] of [object Object] and [object Object].
    • The sum of [object Object] of the public prefix and [object Object] of [object Object] or [object Object] must meet the original restriction on [object Object] of [object Object] or [object Object].
    • The prefix does not support page attention, left padding, or tensor list.
    • In the prefix scenario, when [object Object] is 0 or 1 and [object Object] is passed, [object Object] must be greater than or equal to the sum of [object Object] and [object Object] of [object Object].
    • In the prefix scenario, the input [object Object] cannot be all INT8.

  • KV fake-quantization parameter separation

    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]
      • [object Object] and [object Object] must be the same.
      • Both [object Object] and [object Object] must be either null or non-null. Both [object Object] and [object Object] must be either null or non-null.
      • If neither [object Object] nor [object Object] is null, their shapes must be the same. If neither [object Object] nor [object Object] is null, their shapes must be the same.
      • Only the per-token and per-channel modes are supported. In per-token mode, the shapes of the two parameters must be both (B, S), and the data type is fixed at FLOAT32. In per-channel mode, the shapes of the two parameters must be (N, D), (N, 1, D), or (H), and the data type is fixed at BF16.
      • When both fake-quantization parameters and KV separation quantization parameters are passed, the KV separation quantization parameters take effect.
      • When [object Object] and [object Object] are non-null, [object Object] of [object Object] must be less than or equal to 16.
      • When [object Object] and [object Object] are non-null, the data type of [object Object] must be BFLOAT16, the data type of [object Object] and [object Object] must be INT8, and the data type of the output must be BFLOAT16.
      • When [object Object] and [object Object] are non-null, the tensor list, left padding, and page attention functions are not supported.

• When Q_S is equal to 1

  • Restrictions on [object Object], [object Object], and [object Object]:
    • The B axis can be less than or equal to 65536, the N axis can be less than or equal to 256, and the D axis can be less than or equal to 512.
    • The input types of [object Object], [object Object], and [object Object] cannot be all INT8.
    • In INT4 (INT32) fake-quantization scenarios, the aclnn single-operator call supports KV inputs in either INT4 format or INT4-packed INT32 format. Using [object Object] to generate INT4 data is recommended, which stores eight INT4 elements within one INT32.
    • In INT4 (INT32) fake-quantization scenarios, if KV INT4 values are packed into INT32 inputs, the N, D, or H dimensions of KV must be 1/8 of their actual values (the same applies to prefix).
    • Restrictions on the D axis for [object Object] and [object Object] in specific data types
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: When the input type of [object Object] and [object Object] is INT4 (INT32), the D axis must be 64-element aligned (or 8-element aligned for INT32).
  • Page attention scenario
    • The prerequisite for enabling page attention is that [object Object] exists and is valid, and [object Object] and [object Object] are arranged in a continuous memory based on the indexes in [object Object]. In this scenario, [object Object] of [object Object] and [object Object] is invalid.
    • [object Object] is a user-defined parameter. Its value affects the page attention performance. When page attention is enabled, [object Object] must be set to a non-zero value, and the maximum value of [object Object] cannot exceed 512. Generally, page attention can improve the throughput but deteriorate the performance.
    • In the page attention scenario, when [object Object] of [object Object] is BNSD or TND, the KV cache layout can be BnBsH (blocknum, blocksize, H) or BnNBsD (blocknum, KV_N, blocksize, D). When [object Object] of [object Object] is BSH or BSND, the KV cache layout can only be BnBsH. The value of [object Object] cannot be less than the sum of blocks in each batch calculated based on [object Object] and [object Object]. The shapes of [object Object] and [object Object] must be the same.
    • In the page attention scenario, the performance is generally better when the KV cache layout is BnNBsD than when it is BnBsH. Therefore, BnNBsD is recommended.
    • In the page attention scenario, if the input KV cache layout is BnBsH and numKvHeads × headDim exceeds 64 KB, an error will be reported due to hardware instruction constraints. This problem can be solved by enabling GQA (decreasing [object Object]) or adjusting the KV cache layout to BnNBsD.
    • Page attention does not support the tensor list, left padding, or the scenario where Q is BF16/FP16 and KV is INT4 (INT32).
    • In the page attention scenario, [object Object] must be passed.
    • In the page attention scenario, [object Object] must be two-dimensional. The length of the first dimension must be equal to B, and the length of the second dimension must be greater than or equal to [object Object] (the maximum number of blocks corresponding to [object Object] in each batch).
    • When page attention is enabled, the input [object Object] must be greater than or equal to [object Object] × [object Object] in the following scenarios:
      • When the attention mask is enabled, the mask shape is (B, 1, 1, S).
      • When [object Object] is enabled, the [object Object] shape is (B, N, 1, S).
      • When the fake-quantization per-token mode is enabled, the shape of the input parameters [object Object] and [object Object] is both (2, B, S).
    • Left padding for [object Object]:
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: Scenarios where Q is BF16/FP16 and KV is INT4 (INT32) are not supported.
      • The transfer start point of [object Object] is calculated as follows: [object Object] – [object Object] – [object Object] The transfer end point of [object Object] is calculated as follows: [object Object] – [object Object] If the transfer start point or end point is less than 0, the returned data is all 0s.
      • If [object Object] is less than 0, it will be set to 0.
      • It must be enabled together with [object Object]. Otherwise, the default scenario is right padding for [object Object].
      • Page attention and tensor list are not supported. Otherwise, the default scenario is right padding for [object Object].
      • When it is enabled together with [object Object], ensure that the meaning of [object Object] is correct, that is, invalid data can be correctly masked. Otherwise, precision issues may occur.
    • Restrictions on [object Object]
      • The data types of [object Object] and [object Object] must be consistent.
    • KV fake-quantization parameter separation
      • Except when [object Object] is 0 and [object Object] is 1, the values of [object Object] and [object Object] must be the same.
      • Both [object Object] and [object Object] must be either null or non-null. Both [object Object] and [object Object] must be either null or non-null.
      • If neither [object Object] nor [object Object] is null, their shapes must be the same, except when [object Object] is 0 and [object Object] is 1. If neither [object Object] nor [object Object] is null, their shapes must be the same, except when [object Object] is 0 and [object Object] is 1.
      • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object]: The following eight modes are supported: per-channel, per-tensor, per-token, per-tensor + per-head, per-token + per-head, per-token + page attention to manage scale/offset, per-token + per-head + page attention to manage scale/offset, and [object Object] in per-channel + [object Object] in per-token. In the following description, [object Object] indicates [object Object].
        • Per-channel mode: The shapes of the two parameters can be (1, N, 1, D), (1, N, D), or (1, H). Their data types are the same as that of [object Object]. This mode is supported when the data types of [object Object] and [object Object] are INT8 or INT4 (INT32).
        • Per-tensor mode: The shapes of the two parameters are both (1), and their data types are the same as that of [object Object]. This mode is supported when the data types of [object Object] and [object Object] are INT8.
        • Per-token mode: The shapes of the two parameters are both (1, B, S), and their data type is fixed at FLOAT32. This mode is supported when the data types of [object Object] and [object Object] are INT8 or INT4 (INT32).
        • Per-tensor + per-head mode: The shapes of the two parameters are both (N), and their data types are the same as that of [object Object]. This mode is supported when the data types of [object Object] and [object Object] are INT8.
        • Per-channel for [object Object] + per-token for [object Object] mode: For [object Object] in per-channel, the shapes of the two parameters can be (1, N, 1, D), (1, N, D), or (1, H), and their data types are the same as that of [object Object]. For [object Object] in per-token, the shapes of the two parameters are both (1, B, S), and their data types are fixed at FLOAT32. This mode is supported when the data types of [object Object] and [object Object] are INT8 or INT4 (INT32). When the data types of [object Object] and [object Object] are INT8, only the data types of [object Object] and [object Object] can be FLOAT16.
      • Per-token + per-head mode: The shapes of the two parameters are both (B, N, S), and their data types are fixed at FLOAT32. This mode is supported when the data types of [object Object] and [object Object] are INT8 or INT4 (INT32).
      • Per-token + page attention to manage scale/offset: The shapes of the two parameters are both (blocknum, blocksize), and their data types are fixed at FLOAT32. This mode is supported when the data types of [object Object] and [object Object] are INT8.
      • Per-token + per-head + page attention to manage scale/offset: The shapes of the two parameters are both (blocknum, N, blocksize), and their data types are fixed at FLOAT32. This mode is supported when the data types of [object Object] and [object Object] are INT8.
      • When both fake-quantization parameters and KV separation quantization parameters are passed, the KV separation quantization parameters take effect.
      • In the INT4 (INT32) fake-quantization scenario where only KV fake-quantization parameter separation is supported, the following modes are supported:
        • Per-channel mode
        • Per-token mode
        • Per-token + per-head mode
        • Per-channel for [object Object] + per-token for [object Object] mode
      • Post-quantization is not supported in INT4 (INT32) fake-quantization scenarios.
    • Restrictions on prefix parameters
      • Both [object Object] and [object Object] must be either null or non-null.
      • If neither [object Object] nor [object Object] is null, the dimensions and data types of [object Object], [object Object], [object Object], and [object Object] must be the same.
      • If neither [object Object] nor [object Object] is null, the first dimension (batch) of the shape of [object Object] must be 1. When the layout is BNSD or BSND, the N and D axes must be the same as those of [object Object]. When the layout is BSH, the H axis must be the same as that of [object Object]. The same rules apply to [object Object]. [object Object] of [object Object] and [object Object] must be the same.
      • When [object Object] exists, its shape must be [1], and its value cannot be greater than [object Object] of [object Object] and [object Object].
      • The sum of [object Object] of the public prefix and [object Object] of [object Object] or [object Object] must meet the original restriction on [object Object] of [object Object] or [object Object].

[object Object]

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