QuantSparseFlashMlaMetadata

产品支持情况

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

功能说明

• 算子功能：QuantSparseFlashMlaMetadata是QuantSparseFlashMla算子的前置算子，用于后续Attention计算生成负载均衡的任务划分方案。本算子不执行实际的Attention计算，而是根据输入参数在AI CPU计算出每个AI Core应处理的Attention计算起止范围，从而最大化计算资源的利用率，避免各Core间负载不均衡的问题。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
cu_seqlens_q	可选输入	表示不同Batch中Query的有效Sequence Length，shape为(B+1, )，仅layout_q为TND场景需传入。	INT32	ND
cu_seqlens_ori_kv	可选输入	表示不同Batch中ori_kv的有效Sequence Length，shape为(B+1, )，仅layout_kv为TND场景需传入。	INT32	ND
cu_seqlens_cmp_kv	可选输入	表示不同Batch中cmp_kv的有效Sequence Length，shape为(B+1, )，仅layout_kv为TND场景需传入。	INT32	ND
seqused_q	可选输入	表示不同Batch中Query实际参与运算的Sequence Length，shape为(B, )。	INT32	ND
seqused_ori_kv	可选输入	表示不同Batch中ori_kv实际参与运算的Sequence Length，shape为(B, )。	INT32	ND
seqused_cmp_kv	可选输入	表示不同Batch中cmp_kv实际参与运算的Sequence Length，shape为(B, )。	INT32	ND
cmp_residual_kv	可选输入	表示不同Batch中cmp_kv压缩后Sequence Length的余数，配合cmp_ratio实现cmp_kv部分的mask和负载计算。cmp_mask_mode=3且cmp_ratio≠1时必须传入，shape为(B, )。	INT32	ND
ori_topk_length	可选输入	预留参数，当前不生效。	INT32	ND
cmp_topk_length	可选输入	预留参数，当前不生效。	INT32	ND
num_heads_q	属性	表示Query的head个数，支持1到128。	INT32	-
num_heads_kv	属性	表示Key和Value对应的多头数，当前仅支持1。	INT32	-
head_dim	属性	表示注意力头的维度，当前仅支持512。	INT32	-
quant_mode	属性	表示量化模式，1表示Q、K、V 为per-tensor量化，Q、K、V类型为HIFLOAT8。	INT32	-
batch_size	可选属性	表示Batch数量，默认值为0。	INT32	-
max_seqlen_q	可选属性	表示Query的最长Sequence Length，默认值为0。	INT32	-
max_seqlen_ori_kv	可选属性	表示ori_kv的最长Sequence Length，默认值为0。	INT32	-
max_seqlen_cmp_kv	可选属性	表示cmp_kv的最长Sequence Length，默认值为0。	INT32	-
ori_topk	可选属性	表示ori_kv中筛选出的关键稀疏token的个数，0表示非稀疏场景，默认值为0。	INT32	-
cmp_topk	可选属性	表示cmp_kv中筛选出的关键稀疏token的个数，0表示非稀疏场景，默认值为0。	INT32	-
cmp_ratio	可选属性	表示对cmp_kv的压缩率，默认值为1。	INT32	-
ori_mask_mode	可选属性	表示q和ori_kv计算的mask模式。 0: No Mask。 3: RightDownCausal模式。 4: Band模式。	INT32	-
cmp_mask_mode	可选属性	表示q和cmp_kv计算的mask模式。 0: No Mask。 3: RightDownCausal模式。	INT32	-
ori_win_left	可选属性	表示q和ori_kv计算中q对过去token计算的数量，支持-1或非负数，其中-1表示窗口不受限。	INT32	-
ori_win_right	可选属性	表示q和ori_kv计算中q对未来token计算的数量，支持-1或非负数，其中-1表示窗口不受限。	INT32	-
layout_q	可选属性	表示Query的排列格式，支持BSND、TND，默认值为BSND。	STRING	-
layout_kv	可选属性	表示Key的排列格式，支持BSND、TND、PA_BBND，默认值为BSND。	STRING	-
has_ori_kv	可选属性	用于标识是否含有ori_kv，默认值为true。	BOOL	-
has_cmp_kv	可选属性	用于标识是否含有cmp_kv，默认值为true。	BOOL	-
metadata	输出	表示负载均衡结果输出，shape固定为[1024]。	INT32	ND

约束说明

• QuantSparseFlashMlaMetadata算子需要与QuantSparseFlashMla算子配套使用。
• B（Batch）表示输入样本批量大小。
• 参数cu_seqlens_q、cu_seqlens_ori_kv及cu_seqlens_cmp_kv要求其值为当前Batch与前序Batch有效token数的累加值，后一个元素的值必须大于等于前一个元素的值。
• 参数seqused_q、seqused_ori_kv、seqused_cmp_kv要求其值表示每个Batch中的有效token数。
• 参数cmp_residual_kv需满足cmp_residual_kv[i] < cmp_ratio。
• ori_mask_mode及cmp_mask_mode所表示的mask模式的详细介绍见sparse_mode参数说明。

调用说明

调用方式	样例代码	说明
aclnn API	test_aclnn_quant_sparse_flash_mla_metadata	通过aclnnQuantSparseFlashMlaMetadata接口调用QuantSparseFlashMlaMetadata算子
PyTorch API	test_torch_quant_sparse_flash_mla_metadata	通过quant_sparse_flash_mla_metadata接口调用QuantSparseFlashMlaMetadata算子