MixedQuantSparseFlashMla

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	×
Atlas A2 训练系列产品/Atlas A2 推理系列产品	×
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：

  MixedQuantSparseFlashMla算子旨在完成以下公式描述的Attention计算，支持SWA（Sliding Window Attention）、CSA（Compressed Sparse Attention）、HCA（Heavily Compressed Attention）三类Attention计算场景。与SparseFlashMla的区别在于，本算子支持KV的per-token-group量化输入。调用时需要使用MixedQuantSparseFlashMlaMetadata生成的任务列表metadata。

  典型调用流程如下：

  1. 准备q、ori_kv、cmp_kv、序列长度、block table、sinks等输入。
  2. 调用MixedQuantSparseFlashMlaMetadata生成metadata。
  3. 调用MixedQuantSparseFlashMla，将上一步得到的metadata传入主算子。

• 计算公式：

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

  其中$\tilde{K}=\tilde{V}$为基于ori_kv、cmp_kv以及cmp_ratio等入参控制的实际参与计算的$KV$。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
q	输入	对应公式中的Q。	BFLOAT16	ND
ori_kv	可选输入	对应公式中K和V的一部分，表示原始不经压缩的量化KV。由nope、rope、scale、padding拼接而成，详见quant_mode。	详见quant_mode	ND
cmp_kv	可选输入	对应公式中K和V的一部分，表示经过压缩的量化KV。由nope、rope、scale、padding拼接而成，详见quant_mode。	详见quant_mode	ND
ori_sparse_indices	可选输入	表示从ori_kv中离散取数的索引，当前版本不支持传入。	INT32	ND
cmp_sparse_indices	可选输入	表示从cmp_kv中离散取数的索引。	INT32	ND
ori_block_table	可选输入	表示PageAttention中ori_kv使用的block映射表。	INT32	ND
cmp_block_table	可选输入	表示PageAttention中cmp_kv使用的block映射表。	INT32	ND
cu_seqlens_q	可选输入	表示TND布局下不同batch中q的累积序列长度。	INT32	ND
cu_seqlens_ori_kv	可选输入	表示TND布局下不同batch中ori_kv的累积序列长度。	INT32	ND
cu_seqlens_cmp_kv	可选输入	表示TND布局下不同batch中cmp_kv的累积序列长度。	INT32	ND
seqused_q	可选输入	表示不同batch中q实际参与计算的token数。	INT32	ND
seqused_ori_kv	可选输入	表示不同batch中ori_kv实际参与计算的token数。	INT32	ND
seqused_cmp_kv	可选输入	表示不同batch中cmp_kv实际参与计算的token数。	INT32	ND
cmp_residual_kv	可选输入	表示压缩KV余数，用于恢复cmp侧mask使用的压缩前KV长度。	INT32	ND
ori_topk_length	可选输入	预留输入，当前版本不支持传入非空Tensor。	INT32	ND
cmp_topk_length	可选输入	预留输入，当前版本不支持传入非空Tensor。	INT32	ND
sinks	可选输入	表示attention sinks输入。	FLOAT	ND
metadata	可选输入	MixedQuantSparseFlashMlaMetadata生成的任务切分结果。	INT32	ND
quant_mode	属性	表示量化模式。量化模式1表示K、V nope为per-token-group量化，K、V依次由rope（64，bfloat16）、nope（448，FLOAT8_E4M3FN）、scale（7，bfloat16）、pad（18B）拼接而成；量化模式2表示K、V nope为per-token-group量化，K、V依次由nope（448，FLOAT8_E4M3FN）、rope（64，bfloat16）、scale（7，FLOAT8_E8M0）、pad（1B）拼接而成。当前仅支持1和2，量化模式2仅支持layout_kv为PA_BBND。	INT	-
rope_head_dim	可选属性	表示rope头的维度，仅支持64。	INT	-
softmax_scale	可选属性	对应公式中的softmax_scale。	FLOAT	-
cmp_ratio	可选属性	表示cmp_kv相对于压缩前KV长度的压缩倍率，用于恢复cmp侧mask使用的压缩前KV长度；仅传入ori_kv时不参与压缩KV计算。支持1到128。	INT	-
ori_mask_mode	可选属性	表示q和ori_kv计算的mask模式。 0: No Mask。 3: RightDownCausal模式。 4: Band模式。	INT	-
cmp_mask_mode	可选属性	表示q和cmp_kv计算的mask模式。 0: No Mask。 3: RightDownCausal模式。	INT	-
ori_win_left	可选属性	表示q和ori_kv计算中q对过去token计算的数量，支持-1或非负数，其中-1表示窗口不受限。	INT	-
ori_win_right	可选属性	表示q和ori_kv计算中q对未来token计算的数量，支持-1或非负数，其中-1表示窗口不受限。	INT	-
layout_q	可选属性	表示输入q的数据排布格式，支持"BSND"和"TND"。	STRING	-
layout_kv	可选属性	表示输入ori_kv和cmp_kv的数据排布格式，支持"BSND"、"TND"和"PA_BBND"。	STRING	-
topk_value_mode	可选属性	表示TopK索引取值模式。	INT	-
return_softmax_lse	可选属性	表示是否返回softmax_lse。	BOOL	-
attn_out	输出	对应公式中的输出O。	BFLOAT16	ND
softmax_lse	可选输出	返回softmax的log-sum-exp结果。	FLOAT	ND

约束说明

• 该接口支持推理场景下使用。
• 该接口支持aclgraph模式。
• 该接口当前支持三种计算场景：SWA（Sliding Window Attention）场景仅传入ori_kv；CSA（Compressed Sparse Attention）场景传入ori_kv、cmp_kv及cmp_sparse_indices；HCA（Heavily Compressed Attention）场景传入ori_kv及cmp_kv。
• 通用规格约束如下：
  • KV_N仅支持1，D仅支持512。其中，ori_kv和cmp_kv的D_kv由nope、rope、scale、padding拼接而成，详见quant_mode。
  • cmp_ratio表示cmp_kv相对于压缩前KV长度的压缩倍率；仅传入ori_kv时，cmp_ratio不参与压缩KV计算，需保持默认值1；支持1到128。
  • ori_mask_mode支持0、3和4，cmp_mask_mode支持0和3，ori_win_left和ori_win_right支持-1或非负数，-1表示对应方向不受限。
  • rope_head_dim仅支持64。
  • layout_q和layout_kv组合仅支持"BSND"/"BSND"、"TND"/"TND"、"BSND"/"PA_BBND"、"TND"/"PA_BBND"；非PA_BBND场景下layout_q和layout_kv必须一致；PA_BBND场景下block_size支持1到1024。
  • ori_topk_length和cmp_topk_length为预留输入，全平台均不支持传入非空Tensor。
• 当layout_q为TND时，功能使用限制如下：
  • q的shape需要为[Q_T, Q_N, D]。
  • ori_sparse_indices当前暂不支持。
  • cmp_sparse_indices的shape需要为[Q_T, KV_N, K2]，其中K2为对cmp_kv一次离散选取的token数。
  • cu_seqlens_q必须传入，输入维度为B+1，大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和，即前缀和，因此后一个元素的值必须>=前一个元素的值且首元素必须为0。
• 当layout_q为BSND时，功能使用限制如下：
  • q的shape需要为[B, Q_S, Q_N, D]。
  • ori_sparse_indices当前暂不支持。
  • cmp_sparse_indices的shape需要为[B, Q_S, KV_N, K2]，其中K2为对cmp_kv一次离散选取的token数。
• PageAttention场景下，功能使用限制如下：
  • ori_kv和cmp_kv的shape分别为[ori_block_num, ori_block_size, KV_N, D]和[cmp_block_num, cmp_block_size, KV_N, D]，其中ori_block_num和cmp_block_num为PageAttention时block总数，ori_block_size和cmp_block_size为一个block的token数，ori_block_size和cmp_block_size取值为1到1024，KV_N仅支持1。
  • ori_block_table和cmp_block_table的shape为2维，其中第一维长度为B，第二维长度不小于所有batch中最大的S2和S3对应的block数量，即S2_max / block_size和S3_max / block_size向上取整。
• metadata为算子实际需要使用的分核结果，目前该参数必传，shape大小固定为[1024]。
• layout_kv支持输入"BSND"、"TND"和"PA_BBND"，需满足上述layout_q和layout_kv组合约束。
  • 当输入为PA_BBND时，seqused_ori_kv和ori_block_table必须传入；当输入为BSND时，seqused_ori_kv可用于表达每个batch的ori_kv有效长度；当输入为TND时，ori_kv有效长度由cu_seqlens_ori_kv表达。
  • 当输入为BSND时，ori_kv和cmp_kv的layout都必须为BSND，ori_kv的shape为[B, S2, KV_N, D]，cmp_kv的shape为[B, S3, KV_N, D]。
  • 当输入为TND时，cu_seqlens_ori_kv必须传入；若存在cmp_kv，cu_seqlens_cmp_kv也必须传入。
• return_softmax_lse为False时返回占位Tensor；为True时返回softmax的log-sum-exp结果。
• 除ori_topk_length和cmp_topk_length等预留输入可不传或传入空Tensor外，其余已传入Tensor不支持为空。
• seqused_cmp_kv为所有layout_kv下的可选输入，显式传入时用于覆盖cmp侧逻辑有效长度；未传时由cmp_kv shape、cu_seqlens_cmp_kv或PA block table相关语义推导。
• cmp_residual_kv为算子的可选入参；传入后用于按cmp_len * cmp_ratio + residual恢复cmp侧mask使用的压缩前KV长度，其中cmp_len优先来自显式传入的seqused_cmp_kv。
• q、ori_kv、cmp_kv数据排布格式支持从多种维度解读，B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Hidden-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示hidden层最小的单元尺寸，且满足D=H/N、T表示所有Batch输入样本序列长度的累加和。
• Q_S表示q shape中的S，S2表示ori_kv shape中的S，S3表示cmp_kv shape中的S；Q_N表示num_q_heads，KV_N表示num_ori_kv_heads和num_cmp_kv_heads；Q_T表示q shape中的输入样本序列长度的累加和。

调用说明

调用方式	样例代码	说明
aclnn API	test_aclnn_mixed_quant_sparse_flash_mla	通过aclnnMixedQuantSparseFlashMla调用MixedQuantSparseFlashMla算子
PyTorch API	mixed_quant_sparse_flash_mla	通过cann_ops_transformer.mixed_quant_sparse_flash_mla调用MixedQuantSparseFlashMla算子