[object Object][object Object][object Object]undefined
[object Object]
  • Uses the FlashAttention algorithm to perform self-attention computation in training scenarios. Compared with , this API supports multiple query/key inputs, that is, [object Object], [object Object], [object Object], and [object Object] are used as inputs. In non-multi-input scenarios, use or other APIs.

  • Formula

    The forward computation formula for attention is as follows:

    attention_out=Dropout(Softmax(Mask(scale(querykeyT+queryRopekeyRopeT)+pse),atten_mask),keep_prob)valueattention\_out=Dropout(Softmax(Mask(scale*(query*key^T + queryRope*keyRope^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]
[object Object]
[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.
  • B (batch size) of the input [object Object], [object Object], [object Object], [object Object], and [object Object] must be the same.
  • D (Head-Dim) of the input [object Object], [object Object], and [object Object] must satisfy (qD == kD && kD >= vD). D must be an integer multiple of 8.
  • D (Head-Dim) of the input [object Object] and [object Object] must satisfy (qRopeD == kRopeD). D must be an integer multiple of 8 and less than or equal to the D of [object Object], [object Object], and [object Object].
  • The [object Object] of the input [object Object], [object Object], and [object Object] must be TND.
  • 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 shape must be TND.
    • [object Object] must be 1.
  • 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).
  • Restrictions on [object Object]
    • When the shapes of all [object Object] are less than 2048 and are the same, the default mode is recommended to reduce memory usage.
    • When the value is set to 1, 2, 3, or 5, the user-configured [object Object] and [object Object] 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, 0 is recommended.
    • 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. Outer splitting in the q direction is supported. q and kv of each sequence must have the same length before outer splitting, and [object Object] is passed after outer splitting.
  • In some scenarios, if the computation load is too large, the operator execution may time out (an AI Core error is reported, and [object Object] is [object Object]). 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]. [0] - actualSeqKvLenOptional[0] + qStartIdxOptional - kvStartIdxOptional == 0 (experimental feature)
  • The [object Object] input supports the S length of 0 in a batch. In this case, the [object Object] 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 [object Object] input does not support padding. That is, [object Object] cannot contain a row of all 1s.
  • The S length of a batch in [object Object] can be 0. If the S length is 0, the [object Object] input is not supported. If the actual S length is [2,2,0,2,2], the value of [object Object] is [2,4,4,6,8].
  • [object Object] can only be 0 or 1.
  • [object Object] must be null.
  • [object Object] must be null.
  • [object Object] cannot be null.
[object Object]

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

[object Object]