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

• API function: Uses the FlashAttention algorithm to perform self-attention computation in training scenarios. Different from , this API supports the calculation of variable-length sequences (S), allowing multiple sequences of unequal lengths to be passed at once. When this API is used, [object Object], [object Object], and [object Object] data is passed in the TND format. T indicates the total number, which is the total length of all sequences. Use [object Object] and [object Object] to input the cumulative length of each sequence to distinguish different sequences. The attention result of each sequence is calculated separately.

• Formula:

  The forward computation formula for attention is as follows:

  • When psetype is set to 1, the calculation formula is the same as that of .

  • When [object Object] is set to other values, the formula is as follows:

    $$attention\_out=Dropout(Softmax(Mask(scale*(query*key^T) + pse),atten\_mask),keep\_prob)*value$$

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

• The constraints for input [object Object], [object Object], and [object Object] are as follows:

  • B: The batch sizes must be equal.
  • D: Head-Dim must satisfy (qD == kD && kD >= vD).
  • [object Object] must be consistent.

• Constraints on the data shape:

  • T: The value ranges from 1 to 1M.
  • N: The value ranges from 1 to 256.
  • D: The value ranges from 1 to 768.

• The data format of [object Object], [object Object], and [object Object] can only be TND. T indicates the data closely arranged on the B and S axes (SeqLenQ and SeqLenKV of each batch). B (Batch) indicates the batch size of the input sample, and 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. D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N).

• realShiftOptional: If Sq is greater than 1024, the length of Sq in each batch is the same as that of Skv, and the sparseMode is 0, 2, or 3, the alibi positional encoding compression can be enabled. In this case, only the last 1024 rows of the original PSE need to be input to optimize the memory. That is, alibi_compress = ori_pse[:, :, -1024:, :]. The details are as follows:

  • If the parameters of each batch are different, the shape is BNHSkv (H=1024).
  • When each batch is the same, the shape is 1NHSkv (H=1024).
  • If [object Object] is 2 or 3, the data type must be FLOAT32, and the supported shapes are [B,N] and [N].
  • If this parameter is not enabled, pass a null pointer to [object Object] and 1 to [object Object].

• Meanings of [object Object] values:

  [object Object]undefined

• When [object Object] is set to 2 or 3, Sq and Skv must be of the same length.

• [object Object]: 0 and 1 are reserved, and 2 indicates that invalid row calculation is enabled. This function is used to prevent precision loss caused by the mask of the entire row during calculation. However, this configuration deteriorates the performance. If the operator can determine that invalid rows exist, the invalid row computation is automatically enabled, such as in scenarios where [object Object] is set to 3 and Sq is greater than Skv.

• The constraints for [object Object] are as follows:

  • If the shape values of all [object Object] are the same and less than 2048, you are advised to use the default mode to reduce memory usage.
  • When this parameter is set to 1, 2, or 3, the preTokens and nextTokens configured by the user do not take effect.
  • When the value is set to 0 or 4, ensure that the ranges of [object Object], [object Object], and [object Object] are consistent.
  • If no specific value is required, you are advised to set it to 0.
  • For details about the sparse modes, see .
  • When the value is set to 3, computation on invalid rows is not supported, and Sq <= Skv must be satisfied for each batch.
  • When the value is set to 7, [object Object] is not supported.
  • When the value is set to 8, [object Object] is supported when the q and kv of each sequence have the same length. PSE generation is performed globally. The q direction can be used for external splitting. The q and kv of each sequence must have the same length before external splitting. Then, actualSeqQLenOptional[0] - actualSeqKvLenOptional[0] + qStartIdxOptional - kvStartIdxOptional == 0 (experimental function).

• In some scenarios, if the computation load is too large, the operator execution may time out (AI Core error, errorStr: timeout or trap error). In this case, you are advised to perform axis splitting. Note: The computation load is affected by parameters such as B, S, N, and D. Larger values indicate larger computation loads.

• In the band scenario, the values of [object Object] and [object Object] must overlap.

• The [object Object] sparse computing scenario is [object Object]. When Sq > Skv, the value range of N of [object Object] is [0, Skv]. When Sq ≤ Skv, the value range of N of [object Object] is [Skv – Sq, Skv].

• The actualSeqQLenOptional input supports the S length of 0 in a batch. In this case, the realShiftOptional input is not supported. The length of [object Object] ranges from 1 to 2K. When [object Object] is present, its maximum length is 1K.

• The attenMaskOptional input does not support padding. That is, attenMaskOptional cannot contain a row that contains all 1s.

• The length of S in a batch in actualSeqQLenOptional can be 0. If S is 0, the pse input is not supported. If the actual length of S is [2,2,0,2,2], the value of actualSeqQLenOptional is [2,4,4,6,8].

[object Object]

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