Function: The aclnnBlockSparseAttention API performs backward sparse attention computation. It supports flexible block-level sparse modes and uses BlockSparseMask to specify the KV block selected by each Q block, implementing efficient sparse attention computation.
The calculation formula is as follows: Sparse block size: . BlockSparseMask specifies the sparse mode.
The known forward computation formula is as follows:
For convenience, the formula can be represented using variables and :
The backward computation formula is as follows:
The data layout formats of the BlockSparseAttentionGrad inputs dout, query, key, value, and attentionOut can be interpreted from multiple dimensions. The formats can be passed through qInputLayout and kvInputLayout. To facilitate understanding of the supported layout formats (such as BNSD and TND), the meaning of each dimension represented by the abbreviations in the layout formats is described as follows:
[object Object]([object Object]): input batch size- T: total token length with combined B and S dimensions
[object Object]([object Object]): sequence length of input samples[object Object]([object Object]): hidden-layer size[object Object]([object Object]): number of heads[object Object]([object Object]): minimum unit size of the hidden layer ([object Object]=[object Object]/[object Object])
The following layouts are supported:
- qInputLayout: "TND" "BNSD"
- kvInputLayout: "TND" "BNSD"
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.
- Parameters
Returns:
[object Object]: status code. For details, see .The first-phase API implements input parameter validation. The following error codes may be returned.
[object Object]
- When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
- actualSeqLengthsOptional is mandatory when qInputLayout is set to TND. actualSeqLengthsKvOptional is mandatory when kvInputLayout is set to TND.
- The size of the head dimension in the shape of the query tensor is denoted as N1, and the size of the head dimension in the shape of the key and value tensors is denoted as N2. N1 must be greater than or equal to N2, and N1 % N2 must be equal to 0. (For example, in the BNSD layout, N1 corresponds to the second dimension of the query, and N2 corresponds to the second dimension of the key/value.)
- headdim=128.
- Currently, only BNSD and MHA (N1 == N2) are supported.
The following example is for reference only. For details, see .