[object Object]

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

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

• API function: FlashAttention operator in the full inference scenario. Compared with , this API supports the following new functions: sparse optimization, [object Object] optimization, and INT8 quantization.

• 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" 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]

The 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 computing:

  • [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. When attentionOut is an empty tensor, the AscendCLNN 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.

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

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

    • The B axis must be less than or equal to 65536 (64K). If the input type is INT8 and the D axis is not 32-byte aligned, or the input type is FLOAT16 or BFLOAT16 and the D axis is not 16-byte aligned, the B axis can be up to 128.

    • The N axis must be less than or equal to 256.

    • 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 (AI Core error, errorStr: timeout or trap error). 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]

    • The D axis must be less than or equal to 512. If [object Object] is BSH or BSND, N x D must be less than 65535.

  • Atlas Inference Accelerator Card

    • When inputLayout is BSH, the B axis must be less than or equal to 300. In other cases, the B axis must be less than or equal to 128.
    • The N axis must be less than or equal to 256.
    • The S axis must be less than or equal to 65535 (64K). Q_S or KV_S is not 128-byte aligned. [object Object] cannot be configured when Q_S and KV_S have different lengths.
    • The D axis must be less than or equal to 512.

  • Input data type restrictions:

    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be FLOAT16, BFLOAT16, or INT8.
    • Atlas inference accelerator cards: Only the FLOAT16 data type is supported.

• Restrictions on [object Object]:

  • This parameter is reserved. It is an aclTensor on the device. Its data type must be compatible with that of query according to the type deduction rules. Currently, this parameter is fixed as a null pointer.
  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be FLOAT16 or BFLOAT16.
    • Atlas inference accelerator cards: Only nullptr is supported.

• Restrictions on [object Object]:

  • Input shape restriction: If this parameter is not used, pass [object Object]. Recommended 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. Q_S is S in the shape of [object Object], and KV_S is S in the shape of [object Object] and [object Object]. If KV_S of [object Object] is not 32-byte aligned, it is recommended that it be padded to 32 bytes to improve the performance, filling excess positions with ones.
  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be BOOL, INT8, or UINT8.
    • Atlas inference accelerator cards: Only BOOL is supported.
  • When the data type of [object Object] is INT8 or UINT8, the value in the tensor must be 0 or 1.

• Restrictions on the input of [object Object] and [object Object]:

  • Input value range restrictions:
    • For [object Object], if the sequence length is not specified, [object Object] can be passed, indicating that the valid sequence length is the same as S in the shape of the [object Object]. Note that the valid sequence length of each batch in this parameter cannot exceed the sequence length of the corresponding batch in [object Object].
    • For [object Object], if the sequence length is not specified, [object Object] can be passed, indicating that the valid sequence length is the same as S in the shape of the [object Object] and [object Object]. Note that the valid sequence length of each batch in this parameter cannot exceed the sequence length of the corresponding batch in [object Object] and [object Object].
  • The rules for the input length of [object Object] are as follows: 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] is used. Other lengths are not supported.
  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be INT64.
    • Atlas inference accelerator cards: The data type can be INT64.

• Restrictions on the input of [object Object] and [object Object]:

  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be UINT64 or FLOAT32.
    • Atlas inference accelerator cards: Only nullptr is supported.

• Restrictions on the input of [object Object]:

  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be FLOAT32.
    • Atlas inference accelerator cards: Only nullptr is supported.

• Restrictions on the input of [object Object] and [object Object]:

  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be FLOAT32 or BFLOAT16.
    • Atlas inference accelerator cards: Only nullptr is supported.

• Restrictions on the input of [object Object]:

  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be INT64.
    • Atlas inference accelerator cards: Only the value 2147483647 is supported.

• Restrictions on the input of [object Object]:

  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be INT64.
    • Atlas inference accelerator cards: Only the values 0 and 2147483647 are supported.

• Restrictions on the input of [object Object]:

  • Input data type restrictions:
    • The input format can be BSH, BSND, BNSD, or BNSD_BSND. When the input format is BNSD, the output format is BSND. If no specific format is required, you are advised to set it to BSH.
    • The data format of [object Object], [object Object], and [object Object] can be interpreted from multiple dimensions. To be specific, B (Batch) indicates the size of an input sample batch, S (Seq-Length) indicates the length of the input sample sequence, 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). T indicates the total length of all input sample sequences.

• Restrictions on the input of [object Object]:

  • [object Object] is a host-side integer representing the number of heads in [object Object] and [object Object], supporting grouped-query attention (GQA) scenarios. If no specific value is required, you are advised to set it to 0, indicating that [object Object], [object Object] and [object Object] have the same number of heads. Restrictions: [object Object] must be divisible by [object Object], and in BSND, BNSD, BNSD_BSND scenarios, it must match the N axis shape value of [object Object] and [object Object] in [object Object]. Otherwise an error is reported.
  • Input data type restrictions:
    • [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT: The data type can be INT64.
    • Atlas inference accelerator cards: Only the value 0 is supported.

• Restrictions on the input of [object Object]:

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

    • When [object Object] is 0, the defaultMask mode is used. If [object Object] is not passed, the mask operation is not performed, and [object Object] and [object Object] are ignored (internally set to INT_MAX). If [object Object] is passed, a complete [object Object] matrix (S1 × S2) needs to be passed, indicating that the portion between [object Object] and [object Object] needs to be calculated.
    • When [object Object] is 1, the allMask mode is used. A complete [object Object] matrix (S1 × S2) must be passed.
    • When [object Object] is 2, the leftUpCausal mode is used. An optimized [object Object] matrix (2048 × 2048) needs to be passed.
    • When [object Object] is 3, the rightDownCausal mode for lower-triangle scenarios with right vertex as the dividing line. An optimized [object Object] matrix (2048 × 2048) needs to be passed.
    • When [object Object] is 4, the band mode is used. An optimized [object Object] matrix (2048 × 2048) needs to be passed.
    • When [object Object] is 5, 6, 7, or 8, the prefix, global, dilated, and block_local modes are used respectively, which are not supported currently. If no specific value is required, you are advised to set it to 0.

  • Atlas inference accelerator cards: Only the value 0 is supported.

• Restrictions on the input of [object Object]:

  • Shape restrictions: When [object Object] is set to BNSD_BSND, the shape of the input query is BNSD and the output shape is BSND. In other cases, the shape of this input parameter must be the same as that of the input parameter [object Object].
  • Data type restrictions:
    • For the [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and Ascend 950PR/Ascend 950DT, the data type can be FLOAT16, BFLOAT16, or INT8.
    • Atlas inference accelerator cards: Only FLOAT16 is supported.

• Other constraints:

  • 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. If this parameter is not specified, the default value 0 is used.
    • 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 (0 is used if no value is 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].)
    • When the output is INT8 and [object Object] and [object Object] are set to per-channel mode, the left padding, ring attention, or non-32-byte alignment of the D axis are not supported.
    • If the output is INT8, the sparse mode cannot be band, and [object Object] and [object Object] cannot be negative numbers.
  • When the output is INT8, the input parameter [object Object] is a non-null pointer and a non-empty tensor, and [object Object], [object Object], and [object Object] meet the following conditions, some rows of the matrix are not involved in computation. As a result, the computation result is inaccurate. The following scenarios trigger interception (solution: perform post-quantization outside the PFA interface):
    • If sparseMode is set to 0 and attenMask is a non-null pointer, the interception condition is met when actualSeqLengths - actualSeqLengthsKV - preTokens > 0 or nextTokens < 0 for each batch.
    • When [object Object] is 1 or 2, interception does not occur.
    • When [object Object] is 3, interception occurs if [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] < 0 for any batch.

[object Object]

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