MixedQuantSparseFlashMla

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	×

Function Description

• Operator function:

  The MixedQuantSparseFlashMla operator performs MLA (Multi-head Latent Attention) computation in quantization and sparsity scenarios. It supports three types of attention computation scenarios: SWA (Sliding Window Attention), CSA (Compressed Sparse Attention), and HCA (Heavily Compressed Attention). Different from SparseFlashMla, this operator supports per-token-group quantized KV input. The three typical scenarios of this operator are as follows:

  • SWA (Sliding Window Attention): Only ori_kv is passed in, and sliding window attention is performed on the original KV.
  • CSA (Compressed Sparse Attention): ori_kv, cmp_kv, and cmp_sparse_indices are passed in, and attention is performed on both the original KV window and the compressed KV selected by topK.
  • HCA (Heavily Compressed Attention): ori_kv and cmp_kv are passed in, and attention is performed on both the original KV window and the continuous compressed KV segments.

  The invocation requires the task list metadata generated by MixedQuantSparseFlashMlaMetadata, which is generated before the main operator is executed. In the current version, this metadata must be passed to the main operator. The typical calling process is as follows:

  1. Prepare the corresponding inputs such as q, ori_kv, and cmp_kv based on the calling scenario.
  2. Call MixedQuantSparseFlashMlaMetadata to generate metadata, which is used as an input parameter of MixedQuantSparseFlashMla.
  3. Call MixedQuantSparseFlashMla and pass the metadata obtained in the previous step to the main operator to generate the computation result.

• Computation formula:

  MixedQuantSparseFlashMla uses MLA to compute sparse attention on inputs where Key and Value share the same data. The principle is to selectively compress and quantize the input KV, and then obtain the attention weights by applying Softmax to the computation result of the Query and the concatenated KV.

  The MLA computation formula is generally defined as follows, where $\tilde{K}=\tilde{V}$ is the KV that actually participates in the computation, controlled by the input parameters and composed of the sliding window part of ori_kv and the compressed part of cmp_kv. The range of the KV that actually participates in the computation is determined by cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, and cmp_sparse_indices.

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

Parameter Description

Parameter Name	Input/Output/Attribute	Description	Data Type	Data Format
q	Input	Represents Q in the formula.	FLOAT16, BFLOAT16	ND
ori_kv	Optional input	Represents part of K and V in the formula. It is the original uncompressed quantized KV, where Key and Value share the same data. It is concatenated from nope, rope, scale, and padding. For details, see quant_mode. The default value is None.	FLOAT8_E4M3FN, FLOAT16, BFLOAT16	ND
cmp_kv	Optional input	Represents part of K and V in the formula. It is the compressed quantized KV, where Key and Value share the same data. It is concatenated from nope, rope, scale, and padding. For details, see quant_mode. The default value is None.	FLOAT8_E4M3FN, UINT8	ND
ori_sparse_indices	Optional input	Represents the topK indexes of the original KV. Invalid positions are filled with -1. The default value is None.	INT32	ND
cmp_sparse_indices	Optional input	Represents the topK indexes of the compressed KV. Invalid positions are filled with -1. The default value is None.	INT32	ND
ori_block_table	Optional input	Represents the block mapping table used by ori_kv in the PageAttention scenario. The default value is None.	INT32	ND
cmp_block_table	Optional input	Represents the block mapping table used by cmp_kv in the PageAttention scenario. The default value is None.	INT32	ND
cu_seqlens_q	Optional input	Represents the cumulative sequence lengths of q in different batches under the TND layout. The default value is None.	INT32	ND
cu_seqlens_ori_kv	Optional input	Represents the cumulative sequence lengths of ori_kv in different batches under the TND layout. The default value is None.	INT32	ND
cu_seqlens_cmp_kv	Optional input	Represents the cumulative sequence lengths of cmp_kv in different batches under the TND layout. The default value is None.	INT32	ND
seqused_q	Optional input	Represents the number of tokens of q that actually participate in the computation in each batch. The default value is None.	INT32	ND
seqused_ori_kv	Optional input	Represents the number of tokens of ori_kv that actually participate in the computation in each batch. The default value is None.	INT32	ND
seqused_cmp_kv	Optional input	Represents the number of tokens of cmp_kv that actually participate in the computation in each batch. The default value is None.	INT32	ND
cmp_residual_kv	Optional input	Represents the remainder of the compressed KV, used to restore the pre-compression KV length used by the mask on the cmp side. The default value is None.	INT32	ND
ori_topk_length	Optional input	Represents the length of ori_sparse_indices that actually participates in the computation. The default value is None.	INT32	ND
cmp_topk_length	Optional input	Represents the length of cmp_sparse_indices that actually participates in the computation. The default value is None.	INT32	ND
sinks	Optional input	Represents the independent learnable virtual offset items set for each attention head, used to maintain stability during long-text inference. The default value is None.	FLOAT	ND
metadata	Optional input	Represents the core partitioning information generated by MixedQuantSparseFlashMlaMetadata. The default value is None.	INT32	ND
quant_mode	Attribute	Represents the quantization mode. Supports 1, 2, and 3. Q is not quantized. quant_mode=1: q uses BFLOAT16. The nope of ori_kv and cmp_kv uses per-token-group quantization with groupSize=64. Each token is concatenated in sequence from rope (64 BFLOAT16 values), nope (448 FLOAT8_E4M3FN values), scale (7 BFLOAT16 values), and pad (18 bytes). kv_d=608. quant_mode=2: q uses BFLOAT16. The nope of ori_kv and cmp_kv uses per-token-group quantization with groupSize=64. Each token contains nope (448 FLOAT8_E4M3FN values), rope (64 BFLOAT16 values), scale (7 FLOAT8_E8M0 values), and pad (1 byte). kv_d=584. Only layout_kv=PA_BBND is supported, and the data in a block is organized as block_size*(nope+rope)+block_size*(scale+pad). When quant_mode is 1 or 2, ori_kv and cmp_kv support UINT8 or FLOAT8_E4M3FN as the single-byte storage view, and the underlying byte content remains unchanged. The actual types of the fields are as described above. quant_mode=3: indicates the TurboQuant path that fuses TQ4 dequantization. Only the CSA scenario is supported. q uses FLOAT16 or BFLOAT16. ori_kv remains non-quantized, and its data type is the same as that of q, with kv_d=512. cmp_kv uses UINT8 to store TQ4 data. The 512-dimension data of each token is encoded into 256 bytes of 4-bit codebook indexes, and then concatenated with a 2-byte FLOAT16 per-token scale, totaling 258 bytes, with kv_d=258. Both KV paths use the PA_BBND layout, with shapes (block_num, block_size, 1, 512) and (block_num, block_size, 1, 258) respectively. The codebook is consistent with the agreement on the TurboQuant quantization side, and no external codebook or scale parameters are added.	INT	-
rope_head_dim	Optional attribute	Represents the dimension of the rope head. Only 64 is supported.	INT	-
softmax_scale	Optional attribute	Represents softmax_scale in the formula.	FLOAT	-
cmp_ratio	Optional attribute	Represents the compression ratio of cmp_kv relative to the pre-compression KV length, used to restore the pre-compression KV length used by the mask on the cmp side. When only ori_kv is passed in, this parameter does not participate in the compressed KV computation. Supports 1 to 128.	INT	-
ori_mask_mode	Optional attribute	Represents the mask mode for the computation between q and ori_kv. 0: No mask. 3: rightDownCausal mode. 4: sliding window mode. The default value is 0.	INT	-
cmp_mask_mode	Optional attribute	Represents the mask mode for the computation between q and cmp_kv. 0: No mask. 3: rightDownCausal mode. The default value is 0.	INT	-
ori_win_left	Optional attribute	Represents the number of historical tokens computed for each q in the computation between q and ori_kv. -1 indicates infinity, that is, all tokens participate in the computation. The default value is -1.	INT	-
ori_win_right	Optional attribute	Represents the number of future tokens computed for each q in the computation between q and ori_kv. -1 indicates infinity, that is, all tokens participate in the computation. The default value is -1.	INT	-
layout_q	Optional attribute	Represents the data layout format of the input q. Supports "BSND" and "TND". The default value is "BSND".	STRING	-
layout_kv	Optional attribute	Represents the data layout format of the input ori_kv and cmp_kv. Supports "BSND", "TND", and "PA_BBND". The default value is "BSND".	STRING	-
topk_value_mode	Optional attribute	Represents the value mode of the topK indexes. The default value is 1. When the value is 1, sparse indexes represent logical token indexes. When the value is 2, sparse indexes under the PA_BBND layout represent physical token offsets, that is, physical_block_id * block_size + offset_in_block, where physical_block_id is the actual physical block ID and offset_in_block is the offset within the actual physical block.	INT	-
return_softmax_lse	Optional attribute	Represents whether to return the lse result of softmax. The default value is False.	BOOL	-
attn_out	Output	Represents the output O in the formula.	FLOAT16, BFLOAT16	ND
softmax_lse	Optional output	Represents the result obtained as follows: first take the max of the result of query multiplied by key to obtain softmax_max, then subtract softmax_max from the result of query multiplied by key, take exp and then sum to obtain softmax_sum, and finally take the log of softmax_sum and add softmax_max. By default, no valid result is returned (return_softmax_lse=False). When the value is True, valid values are output.	FLOAT	ND

Constraints

• Ascend 950PR&950DT products: Only quant_mode=1/2 is supported. q and attn_out support only BFLOAT16, and ori_kv and cmp_kv support only FLOAT8_E4M3FN. When quant_mode is 1 or 2, kv_d is 608 or 584 respectively. When quant_mode is 2, layout_kv supports only PA_BBND, and topk_value_mode supports 1 and 2.

• Atlas A3 products: Only the TurboQuant CSA scenario of quant_mode=3 is supported. q, ori_kv, and attn_out support FLOAT16 and BFLOAT16, and their data types must be the same. The last dimension of ori_kv is 512. cmp_kv supports only UINT8 and its last dimension is 258. topk_value_mode supports only 1.

• Atlas A2 products: Only the TurboQuant CSA scenario of quant_mode=3 is supported. q, ori_kv, and attn_out support FLOAT16 and BFLOAT16, and their data types must be the same. The last dimension of ori_kv is 512. cmp_kv supports only UINT8 and its last dimension is 258. topk_value_mode supports only 1.

• This API is used in inference scenarios.
• This API supports the aclgraph mode.
• This API supports batch consistency.
• This API currently supports three computation scenarios: in the SWA (Sliding Window Attention) scenario, only ori_kv is passed in; in the CSA (Compressed Sparse Attention) scenario, ori_kv, cmp_kv, and cmp_sparse_indices are passed in; in the HCA (Heavily Compressed Attention) scenario, ori_kv and cmp_kv are passed in.

Common Field Definitions

Name	Description
b	Batch size of input samples
q_s	Sequence length of the input q
ori_kv_s	Sequence length of the input ori_kv
cmp_kv_s	Sequence length of the input cmp_kv
q_n	Number of heads of the input q
kv_n	Number of heads of the input ori_kv/cmp_kv
q_d	Dimension of the attention head of the input q
kv_d	Dimension of the attention head of the input ori_kv/cmp_kv
q_t	Cumulative sum of the sequence lengths of all batches of the input q
ori_kv_t	Cumulative sum of the sequence lengths of all batches of the input ori_kv
cmp_kv_t	Cumulative sum of the sequence lengths of all batches of the input cmp_kv
ori_kv_k	Number of tokens selected by topK in the input ori_sparse_indices
cmp_kv_k	Number of tokens selected by topK in the input cmp_sparse_indices
ori_kv_s_max	Maximum sequence length of the input ori_kv
cmp_kv_s_max	Maximum sequence length of the input cmp_kv
ori_kv_block_size	Block size of the input ori_kv in the PagedAttention scenario
cmp_kv_block_size	Block size of the input cmp_kv in the PagedAttention scenario
ori_kv_block_nums	Number of blocks of the input ori_kv in the PagedAttention scenario
cmp_kv_block_nums	Number of blocks of the input cmp_kv in the PagedAttention scenario

• The general specification constraints are as follows:
  • kv_n supports only 1, and q_d supports only 512. The kv_d of ori_kv and cmp_kv is concatenated from nope, rope, scale, and padding. For details, see quant_mode.
  • 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 1. Supports 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 values, where -1 indicates that the corresponding direction is unrestricted. Only when ori_mask_mode is 4, ori_win_left and ori_win_right can be >= 0.
  • rope_head_dim supports only 64.
  • The combinations of layout_q and layout_kv 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 PA_BBND scenarios, block_size supports 1 to 1024.
• When layout_q is TND, the function usage restrictions are as follows:
  • The shape of q must be [q_t, q_n, q_d].
  • The shape of ori_sparse_indices must be [q_t, kv_n, ori_kv_k].
  • The shape of cmp_sparse_indices must be [q_t, kv_n, cmp_kv_k].
  • cu_seqlens_q must be passed in. The input dimension is b+1, and the value of each element indicates the total number of tokens in the current batch and all previous batches, that is, the prefix sum. Therefore, the value of each subsequent element must be >= the value of the previous element, and the first element must be 0.
• When layout_q is BSND, the function usage restrictions are as follows:
  • The shape of q must be [b, q_s, q_n, q_d].
  • The shape of cmp_sparse_indices must be [b, q_s, kv_n, ori_kv_k].
  • The shape of cmp_sparse_indices must be [b, q_s, kv_n, cmp_kv_k].
• In the PageAttention scenario, the function usage restrictions are as follows:
  • The shapes of ori_kv and cmp_kv are [ori_kv_block_nums, ori_kv_block_size, kv_n, kv_d] and [cmp_kv_block_nums, cmp_kv_block_size, kv_n, kv_d] respectively, where ori_kv_block_nums and cmp_kv_block_nums are the numbers of blocks in the PagedAttention scenario, and ori_kv_block_size and cmp_kv_block_size are the numbers of tokens in a block, ranging from 1 to 1024.
  • The shapes of ori_block_table and cmp_block_table are 2D, where the length of the first dimension is b, and the length of the second dimension is not smaller than the number of blocks corresponding to the maximum ori_kv_s and cmp_kv_s across all batches, that is, ori_kv_s_max / ori_kv_block_size and cmp_kv_s_max / cmp_kv_block_size rounded up.
• metadata is the core partitioning result actually used by the operator. Currently, this parameter is mandatory, and its shape size is fixed to [2048].
• layout_kv supports "BSND", "TND", and "PA_BBND" as inputs, and must satisfy the preceding combination constraints of layout_q and layout_kv.
  • 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. When the input is TND, the valid length of ori_kv is expressed by cu_seqlens_ori_kv.
  • 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, kv_d], and the shape of cmp_kv is [b, cmp_kv_s, kv_n, kv_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.
• Reserved inputs such as ori_topk_length and cmp_topk_length may not be passed in or may be passed in as empty Tensors. Other passed-in Tensors cannot be 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 not passed in, it is derived from the cmp_kv shape, cu_seqlens_cmp_kv, or the semantics related to the PA block table.
• cmp_residual_kv is an optional input parameter of the operator. After it is passed in, it is used to restore the pre-compression KV length used by the mask on the cmp side according to cmp_len * cmp_ratio + residual, where cmp_len preferentially comes from the explicitly passed-in seqused_cmp_kv.

Calling Instructions

Calling Method	Sample Code	Description
aclnn API	test_aclnn_mixed_quant_sparse_flash_mla	Calls the MixedQuantSparseFlashMla operator through aclnnMixedQuantSparseFlashMla
PyTorch API	mixed_quant_sparse_flash_mla	Calls the MixedQuantSparseFlashMla operator through cann_ops_transformer.mixed_quant_sparse_flash_mla

TurboQuant Parameter Constraints

The following constraints apply to the quant_mode=3 scenario of Atlas A2 products and Atlas A3 products. The following fields are all existing parameters, and no new positional parameters are added. The newly added Tensor data types and layouts take effect only when quant_mode=3.

Parameter	quant_mode=3 Constraint
quant_mode	3
q, attn_out	FLOAT16 or BFLOAT16, with the same type and shape. TND layout, with shape (q_t, q_n, 512). q_n is a multiple of 4 from 4 to 128, and q_t can be 0.
ori_kv	Mandatory. For details about the data type, layout, and shape, see the description of quant_mode.
cmp_kv	Mandatory. For details about the data type, layout, and shape, see the description of quant_mode.
layout_q, layout_kv	TND and PA_BBND respectively. The block_size of both KV paths is a multiple of 16 from 16 to 1024, and the stride of axis 0 must cover a complete physical block.
cmp_sparse_indices	Mandatory INT32, with shape (q_t, 1, 512) or (q_t, 1, 1024). Invalid positions are filled with -1.
ori_block_table, cmp_block_table, cu_seqlens_q, seqused_ori_kv	All mandatory. The original INT32 type and parameter meanings are retained.
ori_sparse_indices, ori_topk_length, cmp_topk_length, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, seqused_q, seqused_cmp_kv, cmp_residual_kv	Passing in is not supported. The original optional parameter positions are retained.
ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right	4, 3, a non-negative value, and 0 respectively.
cmp_ratio, rope_head_dim, topk_value_mode	Support 4 or 128, 64, and 1 respectively. The supported range of the compression ratio of the operator does not mean that the framework quantizes all compressed KV.
metadata	Mandatory. It is generated by the preceding Metadata API with the same quant_mode=3 and the same input attributes. INT32, with shape (2048,).

sinks and Softmax LSE continue to use the original parameters and return forms.

The parameter validation conditions of quant_mode=1 and quant_mode=2 of Ascend 950PR&950DT products remain unchanged.