SparseFlashMla

English | 简体中文

Product Support

Product	Supported
Ascend 950PR&950DT products	√
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Description

• Operator function:

  The SparseFlashMla operator performs the Attention computation described by the following formula and supports three types of Attention computation scenarios: SWA (Sliding Window Attention), CSA (Compressed Sparse Attention), and HCA (Heavily Compressed Attention). The invocation requires the task list metadata generated by SparseFlashMlaMetadata.

  The typical invocation flow is as follows:

  1. Prepare inputs such as q, ori_kv, cmp_kv, sequence lengths, block table, and sinks.
  2. Call SparseFlashMlaMetadata to generate metadata.
  3. Call SparseFlashMla and pass the metadata obtained in the previous step to the main operator.

• Formula:

  $$O=\text{softmax}(Q@{\tilde{K}}^T\cdot \text{softmax\_scale})@\tilde{V}$$

  Here, $\tilde{K}=\tilde{V}$ is the $KV$ that actually participates in the computation, controlled by input parameters such as ori_kv, cmp_kv, and cmp_ratio.

Parameters

Description:
Parameter dimension meanings: B denotes Batch Size. Q_S, ORI_KV_S, and CMP_KV_S denote the Sequence Length of q, ori_kv, and cmp_kv, respectively. Q_N and KV_N denote the Head Num of q and ori_kv/cmp_kv, respectively. Q_T, ORI_KV_T, and CMP_KV_T denote the Total Tokens of q, ori_kv, and cmp_kv, respectively.

Parameter Name	Input/Output/Attribute	Description	Data Type	Data Format
q	Input	Corresponds to Q in the formula.	BFLOAT16, FLOAT16	ND
ori_kv	Optional input	Part of K and V in the formula, indicating the original KV without compression. Default value: None.	BFLOAT16, FLOAT16	ND
cmp_kv	Optional input	Part of K and V in the formula, indicating the compressed KV. Default value: None.	BFLOAT16, FLOAT16	ND
ori_sparse_indices	Optional input	In the SWA sparse ori_kv scenario, indicates the logical indices for sparsely fetching data from ori_kv. -1 indicates an invalid or padded slot. The shape is (T1, N2, K) or (B, S1, N2, K). Default value: None.	INT32	ND
cmp_sparse_indices	Optional input	Indicates the indices for sparsely fetching data from cmp_kv. Default value: None.	INT32	ND
ori_block_table	Optional input	Indicates the block mapping table used by ori_kv in PageAttention. Default value: None.	INT32	ND
cmp_block_table	Optional input	Indicates the block mapping table used by cmp_kv in PageAttention. Default value: None.	INT32	ND
cu_seqlens_q	Optional input	Indicates the cumulative sequence lengths of q in different batches under the TND layout. Default value: None.	INT32	ND
cu_seqlens_ori_kv	Optional input	Indicates the cumulative sequence lengths of ori_kv in different batches under the TND layout. Default value: None.	INT32	ND
cu_seqlens_cmp_kv	Optional input	Indicates the cumulative sequence lengths of cmp_kv in different batches under the TND layout. Default value: None.	INT32	ND
seqused_q	Optional input	Indicates the number of tokens of q that actually participate in the computation in different batches. Default value: None.	INT32	ND
seqused_ori_kv	Optional input	Indicates the number of tokens of ori_kv that actually participate in the computation in different batches. Default value: None.	INT32	ND
seqused_cmp_kv	Optional input	Indicates the number of tokens of cmp_kv that actually participate in the computation in different batches. Default value: None.	INT32	ND
cmp_residual_kv	Optional input	Indicates the compression KV remainder, used to restore the pre-compression KV length used by the cmp-side mask. Default value: None.	INT32	ND
ori_topk_length	Optional input	In the SWA sparse ori_kv scenario, indicates the actual number of critical sparse ori_kv tokens corresponding to each q token. The shape is (T1, N2) or (B, S1, N2). Default value: None.	INT32	ND
cmp_topk_length	Optional input	In the CSA scenario of the Ascend 950PR/950DT series, indicates the number of valid entries in cmp_sparse_indices corresponding to each q token and KV head. This parameter is mandatory when cmp_mask_mode=0 and cmp_sparse_indices is passed in. The shape is (T1, 1) under TND and (B, S1, 1) under BSND, consistent with cmp_sparse_indices after removing the last dimension K2. Atlas A2/A3 do not support passing in a non-empty Tensor. Default value: None.	INT32	ND
sinks	Optional input	Indicates the attention sinks input. Default value: None.	FLOAT	ND
metadata	Optional input	Task splitting result generated by the accompanying metadata front-end API. Default value: None.	INT32	ND
softmax_scale	Optional attribute	Corresponds to softmax_scale in the formula. Default value: 1.0.	FLOAT	-
cmp_ratio	Optional attribute	Indicates the compression ratio of cmp_kv relative to the pre-compression KV length, used to restore the pre-compression KV length used by the cmp-side mask. When only ori_kv is passed in, it does not participate in the compressed KV computation. Values from 1 to 128 are supported. Default value: 1.	INT	-
ori_mask_mode	Optional attribute	Indicates the mask mode for the computation of q and ori_kv. 0: No Mask. 3: RightDownCausal mode. 4: Band mode. Default value: 0.	INT	-
cmp_mask_mode	Optional attribute	Indicates the mask mode for the computation of q and cmp_kv. 0: No Mask. 3: RightDownCausal mode. Default value: 0.	INT	-
ori_win_left	Optional attribute	Indicates the number of past tokens with which q is computed in the computation of q and ori_kv. Both -1 and non-negative numbers are supported, where -1 indicates that the window is unlimited. Default value: -1.	INT	-
ori_win_right	Optional attribute	Indicates the number of future tokens with which q is computed in the computation of q and ori_kv. Both -1 and non-negative numbers are supported, where -1 indicates that the window is unlimited. Default value: -1.	INT	-
layout_q	Optional attribute	Indicates the data layout format of the input q. "BSND" and "TND" are supported. Default value: "BSND".	STRING	-
layout_kv	Optional attribute	Indicates the data layout format of the input ori_kv and cmp_kv. "BSND", "TND", and "PA_BBND" are supported. Default value: "BSND".	STRING	-
topk_value_mode	Optional attribute	Indicates the TopK index value mode. Default value: 1. When the value is 1, sparse indices represent logical token indices. When the value is 2, sparse indices in the PA_BBND layout represent physical token offsets, that is, physical_block_id * block_size + offset_in_block.	INT	-
return_softmax_lse	Optional attribute	Indicates whether to return softmax_lse. Default value: False.	BOOL	-
attn_out	Output	Corresponds to the output O in the formula.	BFLOAT16, FLOAT16	ND
softmax_lse	Optional output	Returns the log-sum-exp result of softmax. By default, no valid result is returned (return_softmax_lse=False). When the value is True, valid values are output.	FLOAT	ND

Constraints

• This API can be used in training and inference scenarios.

• This API supports aclgraph mode.

• This API supports batch consistency.

• This API currently supports four computation scenarios: the SWA (Sliding Window Attention) scenario passes in only ori_kv; the SWA sparse ori_kv scenario passes in ori_kv, ori_sparse_indices, and ori_topk_length; the CSA (Compressed Sparse Attention) scenario passes in ori_kv, cmp_kv, and cmp_sparse_indices; the HCA (Heavily Compressed Attention) scenario passes in ori_kv and cmp_kv.

• The general specification constraints are as follows:

  • Only KV_N of 1 is supported, and only D of 512 is supported. The D_kv of ori_kv and cmp_kv is formed by concatenating nope(448) and rope(64).
  • cmp_ratio indicates the compression ratio of cmp_kv relative to the pre-compression KV length. When only ori_kv is passed in, cmp_ratio does not participate in the compressed KV computation and must remain the default value 1. Values from 1 to 128 are supported.
  • The values of ori_mask_mode, cmp_mask_mode, ori_win_left, and ori_win_right vary by product model. For details, see the "Product Model Constraints" section.
  • The block_size of PageAttention supports 1 to 1024.
  • The layout_q and layout_kv combinations support only "BSND"/"BSND", "TND"/"TND", "BSND"/"PA_BBND", and "TND"/"PA_BBND". In non-PA_BBND scenarios, layout_q and layout_kv must be the same.
  • In the SWA sparse ori_kv scenario, ori_topk_length must be passed in. The ori_topk of the accompanying Metadata API is the last dimension K of ori_sparse_indices, and the element values of ori_topk_length must be in the range [0, K]. In other scenarios, pass in nullptr or an empty Tensor for ori_topk_length.

• The product model constraints are as follows:

  • Atlas A2 products and Atlas A3 products: Q_N supports 1, 2, 4, 8, 16, 32, 64, and 128, and KV_N supports only 1. SWA does not pass in cmp_kv, and cmp_ratio does not participate in the computation. CSA supports passing in 1, 2, or 4, and HCA supports passing in 128. block_size must be a multiple of 16, and the maximum supported value is 1024. The SWA sparse ori_kv scenario supports ori_sparse_indices and ori_topk_length; layout_q and layout_kv support only the "TND"/"PA_BBND" combination; ori_mask_mode is 0; ori_win_left and ori_win_right are non-negative numbers. In non-SWA sparse ori_kv scenarios, ori_mask_mode is 4, ori_win_left is 127, and ori_win_right is 0. The last dimension K2 of cmp_sparse_indices supports any integer in [1, 8192]; cmp_mask_mode supports only 3, and topk_value_mode supports only 1.
  • Ascend 950PR&950DT products: Q_N supports 1-128, and KV_N supports only 1. When only ori_kv is passed in, cmp_ratio must be 1. The CSA/HCA scenarios support 1 to 128. ori_mask_mode supports 0, 3, and 4, and cmp_mask_mode supports 0 and 3. ori_win_left and ori_win_right support -1 or non-negative numbers, where -1 indicates that the corresponding direction is unlimited. Only when ori_mask_mode is 4, ori_win_left and ori_win_right can be >=0. topk_value_mode supports 1 and 2.

• When layout_q is TND, the function usage restrictions are as follows:

  • The shape of q must be [Q_T, Q_N, D].
  • In the SWA sparse ori_kv scenario, the shape of ori_sparse_indices is [Q_T, KV_N, K], and the shape of ori_topk_length is [Q_T, KV_N].
  • The shape of cmp_sparse_indices must be [Q_T, KV_N, K2], where K2 is the number of tokens selected in a single sparse selection from cmp_kv.
  • cu_seqlens_q must be passed in, with the shape [B+1,]. The first number is fixed to 0, that is, the prefix 0. Each of the following elements indicates the total number of tokens (prefix sum) of the current batch and all preceding batches. Therefore, the value of each element must be greater than or equal to that of the preceding element.

• When layout_q is BSND, the function usage restrictions are as follows:

  • The shape of q must be [B, Q_S, Q_N, D].
  • In the SWA sparse ori_kv scenario, the shape of ori_sparse_indices is [B, Q_S, KV_N, K], and the shape of ori_topk_length is [B, Q_S, KV_N].
  • The shape of cmp_sparse_indices must be [B, Q_S, KV_N, K2], where K2 is the number of tokens selected in a single sparse selection from cmp_kv.

• In the PageAttention scenario, the function usage restrictions are as follows:

  • The shapes of ori_kv and cmp_kv are [ori_block_num, ori_block_size, KV_N, D] and [cmp_block_num, cmp_block_size, KV_N, D], respectively, where ori_block_num and cmp_block_num are the total number of blocks in PageAttention, and ori_block_size and cmp_block_size are the number of tokens in a block. ori_block_size and cmp_block_size support values from 1 to 1024, and only KV_N of 1 is supported.
  • The shapes of ori_block_table and cmp_block_table are 2-dimensional. The length of the first dimension is B, and the length of the second dimension must not be less than the number of blocks corresponding to the maximum S2 and S3 among all batches, that is, the ceiling of ORI_KV_S_max / block_size and CMP_KV_S_max / block_size.

• metadata is the core splitting result that the operator actually needs to use. This parameter is currently mandatory, and the shape size is fixed to [2048].

• layout_kv supports the input "BSND", "TND", and "PA_BBND" values, which must satisfy the preceding layout_q and layout_kv combination constraints.

  • When the input is PA_BBND, seqused_ori_kv and ori_block_table must be passed in. When the input is BSND, seqused_ori_kv can be used to express the valid length of ori_kv in each batch, and the valid length must not exceed KV_S of key/value and must not be less than 0. When the input is TND, the total length of ori_kv is expressed by cu_seqlens_ori_kv. If seqused_ori_kv is passed in, the valid length of ori_kv is expressed by seqused_ori_kv; otherwise, the valid length is the same as the total length.
  • When the input is BSND, the layouts of both ori_kv and cmp_kv must be BSND. The shape of ori_kv is [B, ORI_KV_S, KV_N, D], and the shape of cmp_kv is [B, CMP_KV_S, KV_N, D].
  • When the input is TND, cu_seqlens_ori_kv must be passed in. If cmp_kv exists, cu_seqlens_cmp_kv must also be passed in.

• When return_softmax_lse is False, a placeholder Tensor is returned. When it is True, the log-sum-exp result of softmax is returned.

• The SWA sparse ori_kv scenario supports only the SWA template. Only ori_kv is passed in, and ori_sparse_indices and ori_topk_length must be passed in together. ori_win_left and ori_win_right take non-negative numbers only when ori_mask_mode is 4. The ori_topk of the accompanying Metadata API is K, and cmp_kv is not passed in this scenario.

• ori_topk_length indicates the actual number of valid index entries for each q token and KV head, and the values must be in the range [0, K]. In ori_sparse_indices, the range [0, ori_topk_length) contains left-aligned valid index entries, and the range [ori_topk_length, K) contains invalid or padded entries. It is recommended to fill in -1.

• For unused reserved inputs, they can be omitted or passed in as empty Tensors. Other passed-in Tensors do not support being empty.

• seqused_cmp_kv is an optional input for all layout_kv values. When explicitly passed in, it is used to overwrite the logical valid length on the cmp side. When it is not passed in, the length is derived from the cmp_kv shape, cu_seqlens_cmp_kv, or the PA block table semantics.

• cmp_residual_kv is an optional input parameter of both the main API and the metadata front-end API. After it is passed in, it is used to restore the pre-compression KV length used by the cmp-side mask as cmp_len * cmp_ratio + residual, where cmp_len preferentially comes from the explicitly passed-in seqused_cmp_kv.

• The data layout formats of q, ori_kv, and cmp_kv can be interpreted from multiple dimensions. B (Batch) indicates the batch size of the input samples. S (Seq-Length) indicates the sequence length of the input samples. H (Hidden-Size) indicates the size of the hidden layer. N (Head-Num) indicates the number of heads. D (Head-Dim) indicates the smallest unit size of the hidden layer, satisfying D=H/N. T indicates the cumulative sum of the input sample sequence lengths of all batches.

• Q_S indicates the S in the q shape. ORI_KV_S indicates the S in the ori_kv shape. CMP_KV_S indicates the S in the cmp_kv shape. Q_N indicates num_q_heads. KV_N indicates num_ori_kv_heads and num_cmp_kv_heads. Q_T indicates the cumulative sum of the input sample sequence lengths in the q shape.

Invocation Description

Invocation Method	Sample Code	Description
aclnn API	test_aclnnSparseFlashMla	Calls the SparseFlashMla operator through aclnnSparseFlashMla
PyTorch API	sparse_flash_mla	Calls the SparseFlashMla operator through cann_ops_transformer.sparse_flash_mla