[object Object][object Object][object Object]undefined
[object Object]
  • Function: BlockSparseAttention performs sparse attention computation and supports flexible block-level sparsity modes. It uses BlockSparseMask to specify the KV blocks selected for each Q block, achieving efficient sparse attention computation.

  • Formula: Sparse block size of blockShapeX×blockShapeYblockShapeX \times blockShapeY, with sparsity pattern defined by [object Object]

    attentionOut=Softmax(scalequerykeysparseT+atten_mask)valuesparseattentionOut = Softmax(scale \cdot query \cdot key_{sparse}^T + atten\_mask) \cdot value_{sparse}

    The input query, key, and value of BlockSparseAttention can be arranged in multiple dimensions, which can be specified by qInputLayout and kvInputLayout.

    • [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"
[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 validation. The following error codes may be returned.

    [object Object]
[object Object]
  • Parameters

    [object Object]
  • Returns

[object Object]: status code. For details, see .

[object Object]
  • Deterministic computation:
    • The default deterministic implementation of aclnnBlockSparseAttention is used.
  • When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
  • Currently, [object Object] can only be TND or BNSD.
  • Currently, [object Object] can only be TND or BNSD.
  • The input layouts of the query, key, and value must be the same.
  • The input query, key, and value must be of the same data type, which can be either FLOAT16 or BFLOAT16.
  • The dimension D of the query, key, and value can only be set to 64 or 128.
  • If blockShapeOptional is passed, it must contain at least two elements [blockShapeX, blockShapeY], and the values must be greater than 0. blockShapeY must be a multiple of 128.
  • blockSparseMaskOptional must be passed, and its shape must be [batch, headNum, ceilDiv(maxQS, blockShapeX), ceilDiv(maxKVS, blockShapeY)].
  • Currently, only nullptr can be passed to attentionMaskOptional.
  • actualSeqLengthsOptional is required when qInputLayout is set to "TND". actualSeqLengthsKvOptional is required when kvInputLayout is set to "TND".
  • actualSeqLengthsOptional and actualSeqLengthsKvOptional must be both configured or both not configured. If only one of them is configured, the operator will intercept the configuration.
  • Currently, only nullptr can be passed to blockTableOptional, indicating that the PagedAttention feature is disabled.
  • [object Object] must be 0 (for float32 Softmax) or 1 (for fp16 Softmax). If the [object Object] input is BFLOAT16, [object Object] can only be set to 0.
  • [object Object] and [object Object] do not need to be exactly divided by [object Object]. When they are not divisible, the actual number of blocks is determined by ceiling division.
  • If headNum of input [object Object] is N1 and headNum of input [object Object] and [object Object] is N2, then [object Object].
  • Currently, maskType can only be set to 0, indicating that no mask is added.
  • Currently, blockSize can only be set to 0, indicating that the paged cache is not supported.
  • Currently, preTokens and nextTokens can only be set to 2147483647, indicating that all tokens before and after the current token are involved in the attention operation. That is, sliding window attention is not supported.
[object Object]

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

[object Object]