CANN ops-transformer 算子解析：SparseFlashMlaSoftmaxL1Norm 的算法原理、参数体系与双端调用实践-平芜编程栈

算子库
人工智能
深度学习
Ascend

【免费下载链接】ops-transformer

本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。

项目地址：https://gitcode.com/cann/ops-transformer

SparseFlashMlaSoftmaxL1Norm 是 CANN ops-transformer 算子库中服务于 Sparse Flash MLA（Multi-head Latent Attention）训练反向链路的正向配套算子，用于计算注意力概率矩阵 P 在 query 头组（group）维度上的 Softmax L1Norm 结果，为反向梯度计算提供中间量。本文以算子 README 为主体，结合仓库内 aclnn 接口文档、PyTorch 接口文档、算子定义、tiling、kernel 与调用示例源码，系统讲解该算子的适用平台、三阶段计算公式、全部输入输出参数与约束、metadata 负载均衡机制，以及 aclnn API 与 PyTorch API 两种调用方式，读者可按文直接完成单算子模式与 TorchAir 图模式的接入与调试。

产品支持情况

该算子仅面向最新一代昇腾 AI 处理器架构，当前仓库中对其余产品线均明确标注为不支持，具体如下表所示：

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

从源码看，算子定义文件中仅注册了AICore().AddConfig("ascend950", ...)一种硬件配置，与上表"仅 Ascend 950 支持"的结论相互印证。因此在使用本算子前，请先确认目标设备型号为 Ascend 950PR 或 Ascend 950DT。

功能说明：一次 Softmax L1Norm 的完整计算流程

算子定位与配套关系

SparseFlashMlaSoftmaxL1Norm的功能是计算 Sparse Flash MLA 注意力的 Softmax L1Norm 结果，支持 Compressed Attention（压缩注意力）以及 Sparse Compressed Attention（稀疏压缩注意力）两类场景。在训练链路中，它是aclnnDenseLightningIndexerKLLossGrad（dense_lightning_indexer_kl_loss_grad）反向算子的配套正向接口，其输出可用于反向梯度计算。

调用过程遵循"先负载均衡、后主计算"的两段式模式，由两个算子协作完成：

SparseFlashMlaSoftmaxL1NormMetadata（metadata 前置算子）：根据主算子的 shape、layout、mask 等信息，采用 strided 方式将任务均衡切分到可用 AIC 核上，输出 metadata 供主算子使用。
SparseFlashMlaSoftmaxL1Norm（主算子）：根据 metadata 中的分核信息，对 Q 和 K 计算 Softmax L1Norm。

metadata 的字段布局在 torch 接口文档与 kernel 侧头文件中均有明确定义，固定为 64 个 INT32 元素（SMLA_METADATA_SIZE = 64），核心字段如下：

字段	index	说明
totalNum	0	参与负载均衡的 seq 总数。
formerCoreProcessNum	1	常规核处理的 seq 数，即`ceil(totalNum / totalCoreNum)`。
remainCoreProcessNum	2	尾核处理的 seq 数，即`floor(totalNum / totalCoreNum)`。
remainCoreNum	3	尾核数目。
totalCoreNum	4	实际使用的 AIC 核数，取`min(totalNum, aicCoreNum, 36)`。

其中"最多 36 核"的限制在 kernel 侧由常量SMLA_METADATA_MAX_CORE_NUM = 36约束；metadata 结构体定义还包含SMLA_METADATA_SIZE = 64，并附带static_assert保证 64 个 INT32 足以容纳全部字段。

三阶段计算公式

算子内部按如下三个阶段完成从原始输入到 Softmax L1Norm 的推导：

阶段一：根据是否为 sparse 场景对输入 K 进行选择

当为 sparse 场景时，根据稀疏索引对 K 做 Gather，仅取权重较高的 block：

$$ selectedKv = Gather(K, sparseIndices[i]),\ 0 \le i < selectBlockCount $$

否则直接使用完整 K：

$$ selectedKv = K $$

阶段二：计算 P（SimpleSoftmax）

对 Q 与 selectedKv 的转置做矩阵乘并乘以缩放系数 scale，经 mask 处理后，利用正向 softmax 输出的 softmaxLse 做数值稳定的 SimpleSoftmax 归一化：

$$ P = SimpleSoftmax(Mask(Q @ selectedKv^{T} \cdot scale), softmaxLse) $$

阶段三：计算 Softmax L1Norm

在 G 维度（q head group 维度）上对 softmax 概率 P 求和后取平均：

$$ softmaxL1Norm = \frac{ReduceSum(P, dim=G)}{G} $$

其中 $G$ 为 group 数，$G = N1 / N2$（N1 为 q 的 head 数，N2 为 k 的 head 数）。以 kernel 源码 sparse_flash_mla_softmax_l1_norm.py 中的向量函数_muls_sel_vf_pse_type1_inner为例，其计算路径为：对 QK 结果先乘 scale（vf.muls），再与 lse 做exp_sub得到逐元素 exp 值，按 8 路展开累加求和，最后除以g_scalar（即 1/G）后写入输出，与上述公式逐项对应。

参数说明

主算子 sparse_flash_mla_softmax_l1_norm

参数名	输入/输出/属性	描述	数据类型	数据格式	维度（shape）
q	输入	attention 结构的输入 Q。	FLOAT16、BFLOAT16	ND	(B,S1,N1,D) 或 (T1,N1,D)
k	输入	attention 结构的输入 K(V)。	FLOAT16、BFLOAT16	ND	(B,S2,N2,D) 或 (T2,N2,D)
softmax_lse	输入	注意力正向计算的输出 softmaxLse，计算公式详见 sparse_flash_mla 文档。	FLOAT32	ND	(B,N2,S1,G) 或 (N2,T1,G)
sparse_indices	可选输入	稀疏场景下选择的 k 中权重较高的注意力索引。	INT32	ND	(B,S1,N2,K) 或 (T1,N2,K)
cu_seqlens_q	可选输入	每个 Batch 中 Query 的有效 token 数的累加和形式，layout 为 TND 时该参数必传。	INT32	ND	(B+1,)
cu_seqlens_k	可选输入	每个 Batch 中 Key 的有效 token 数的累加和形式，layout 为 TND 时该参数必传。	INT32	ND	(B+1,)
seqused_q	可选输入	表示不同 batch 中 q 实际参与运算的 token 数。	INT32	ND	(B,)
seqused_k	可选输入	表示不同 batch 中 k 实际参与运算的 token 数。	INT32	ND	(B,)
cmp_residual_k	可选输入	表示每个 batch`S2 // cmpRatio`后的余数，当 k 不为空且 mask_mode=3 时必须传入。	INT32	ND	(B,)
topk_length	可选输入	表示每行 q 对应的 k 实际可选的 topk 长度，mask_mode=0 且存在稀疏索引时需要传。	INT32	ND	(B,S1,N2) 或 (T1,N2)
metadata	可选输入	表示 tiling 下沉的 aicpu 算子输出结果，由`SparseFlashMlaSoftmaxL1NormMetadata`算子生成。	INT32	ND	(64,)
softmax_scale	可选属性	缩放系数，默认值 1.0，推荐值取 sqrt(head_dim) 的倒数。	FLOAT32	-	-
max_seqlen_k	可选属性	k 的最大序列长度，TND dense 场景用于输出 shape 推导，默认值 0。	INT64	-	-
cmp_ratio	可选属性	对 k 的压缩率，取值范围 1~128，默认值 1。	INT64	-	-
mask_mode	可选属性	q 和 k 计算的 mask 模式，0：No mask；3：rightDownCausal 模式，默认值 0。	INT64	-	-
layout_q	可选属性	q 的数据排布格式，支持 "BSND"、"TND"，默认值 "BSND"。	STRING	-	-
layout_k	可选属性	k 的数据排布格式，支持 "BSND"、"TND"，默认值 "BSND"。	STRING	-	-
softmax_l1_norm	输出	q 与 k 计算得出的 softmax L1Norm 结果，公式为 reduceG(softmax)/G。若存在 sparse_indices 则该输出不为空，其他场景下输出为空。	FLOAT32	ND	(B,S1,N2,S2) 或 (T1,N2,T2)

关于 shape 的几点补充（可在 aclnn 接口文档中查到完整说明）：

q、k 的 N 轴对应关系需满足 GQA 约束（N1 = N2 × G）；k 的 B 与 q 的 B 保持一致；k 的 D 与 q 的 D 保持一致。
softmax_lse 的 B 与 q 的 B 保持一致，S1 与 q 的 S1 保持一致，G = N1/N2。
cu_seqlens_q/cu_seqlens_k 为可选项：仅当 layout 为 TND 时存在；长度与 B+1 保持一致；累加和分别与 T1、T2 保持一致。
cmp_residual_k 的强制传入条件是maskMode=3 且 cmpRatio!=1。
topk_length 在 mask_mode=0 且存在稀疏索引时必传，且必须为准确值。
metadata 为必传项，形状为 (64,)。

metadata 前置算子 sparse_flash_mla_softmax_l1_norm_metadata

前置算子不接收张量形式的 q/k，而是接收描述注意力结构的标量参数与可选张量：

参数名	参数类型	可选/必选	描述	数据类型
num_heads_q	int	必选	公式中 Q 的头数（即 N1），当前支持 1~128。	int
num_heads_k	int	必选	公式中 key 的头数（即 N2），当前仅支持 1。	int
head_dim	int	必选	头的维度（即 D），当前仅支持 512。	int
cu_seqlens_q	Tensor	可选	每个 Batch 中 q 的有效 token 数的累加和形式，layout_q 为 TND 时必传，累加和与 T1 一致。	int32
cu_seqlens_k	Tensor	可选	每个 Batch 中 k 的有效 token 数的累加和形式，layout_k 为 TND 时必传，累加和与 T2 一致。	int32
seqused_q	Tensor	可选	不同 batch 中 q 实际参与运算的 token 数。	int32
seqused_k	Tensor	可选	预留接口参数，当前 kernel 路径暂不使用。	int32
cmp_residual_k	Tensor	可选	预留接口参数，当前 kernel 路径暂不使用。	int32
topk_length	Tensor	可选	每行 q 对应的 k 实际可选的 topk 长度。	int32
batch_size	int	可选	输入样本批量大小（即 B），默认 None（BSND 场景需传正数，TND 场景可为 None 自动推导）。	int
max_seqlen_q	int	可选	q 的最大序列长度，默认 None，BSND 场景必须为正数。	int
max_seqlen_k	int	可选	k 的最大序列长度，默认 None，BSND 场景必须为正数。	int
topk	int	可选	从 k 中筛选出的关键 token 个数，0 表示无稀疏，默认 None。	int
layout_q	str	可选	q 的数据排布格式，支持 "BSND"、"TND"，默认 None（内部转为 "BSND"）。	string
layout_k	str	可选	k 的数据排布格式，支持 "BSND"、"TND"，默认 None（内部转为 "BSND"）。	string
mask_mode	int	可选	0：No mask；3：rightDownCausal 模式，默认 None（内部转为 0）。	int
cmp_ratio	int	可选	对 k 的压缩率，取值范围 1~128，默认 None（内部转为 1）。	int

前置算子返回形状为 (64,) 的 INT32 负载均衡结果，即上文所述的 metadata。

约束说明

使用本算子必须同时满足以下约束，否则可能无法编译或产生非法结果：

确定性：aclnnSparseFlashMlaSoftmaxL1Norm默认为确定性实现，每次运行结果可复现。
layout：仅支持 BSND 或 TND 两种排布，且layout_q与layout_k必须保持一致。
数据 shape 约束：
- B：泛化支持；
- S1、S2：泛化支持，且支持 S1、S2 不等长；
- N1：支持 1~128，且 num_heads_q 必须能被 num_heads_k 整除；
- N2：仅支持 1；
- D：仅支持 512（q、k 最后一维保持一致）。
mask_mode 支持：

模式	含义	备注
0	不做 mask 操作	支持
3	rightDownCausal 模式的 mask，对应以右顶点为划分的下三角场景	支持

cmp_ratio取值范围：1~128。
TND 场景：必传 cu_seqlens_q 和 cu_seqlens_k；batch_size 可为 None（通过 cu_seqlens_q 推导）。
BSND 场景：必传 batch_size、max_seqlen_q 和 max_seqlen_k。
metadata 必须传入，且由SparseFlashMlaSoftmaxL1NormMetadata算子生成。
入参为空处理：q 为空 Tensor 时直接返回。
数据类型：q、k 的数据类型必须保持一致（FLOAT16/BFLOAT16）。
适用场景：该接口支持训练场景使用，支持单算子模式和 TorchAir 图模式调用。

调用方式一：aclnn API（C++ 两段式接口）

aclnn 接口采用 CANN 标准的两段式调用模式：先调用aclnnSparseFlashMlaSoftmaxL1NormGetWorkspaceSize获取计算所需 workspace 大小与执行器，再调用aclnnSparseFlashMlaSoftmaxL1Norm执行计算。

函数原型

aclnnStatus aclnnSparseFlashMlaSoftmaxL1NormGetWorkspaceSize( const aclTensor *q, const aclTensor *k, const aclTensor *softmaxLse, const aclTensor *sparseIndicesOptional, const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKOptional, const aclTensor *sequsedQOptional, const aclTensor *sequsedKOptional, const aclTensor *cmpResidualKOptional, const aclTensor *topkLengthOptional, const aclTensor *metadataOptional, double softmaxScale, int64_t cmpRatio, int64_t maskMode, char *layoutQOptional, char *layoutKOptional, const aclTensor *softmaxL1Norm, uint64_t *workspaceSize, aclOpExecutor **executor);

aclnnStatus aclnnSparseFlashMlaSoftmaxL1Norm( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);

第二段接口中：workspace为 Device 侧申请的 workspace 内存地址；workspaceSize为第一段接口返回的 workspace 大小；executor为包含算子计算流程的 op 执行器；stream指定执行任务的 Stream 流。

第一段接口完成入参校验，主要错误码如下：

返回值	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	参数中存在非法的 nullptr。
ACLNN_ERR_PARAM_INVALID	161002	输入的数据类型不满足支持类型。
ACLNN_ERR_PARAM_INVALID	161002	q、k、softmaxLse、softmaxL1Norm 必选输入/输出未传。

完整调用示例（TND 场景）

仓库在 examples/test_aclnn_sparse_flash_mla_softmax_l1_norm.cpp 提供了可直接编译运行的示例，核心流程如下：

#include "acl/acl.h" #include "aclnnop/aclnn_sparse_flash_mla_softmax_l1_norm.h" #include "aclnnop/aclnn_sparse_flash_mla_softmax_l1_norm_metadata.h" // TND layout: qShape=(T1,N1,D), kShape=(T2,N2,D), softmaxLseShape=(N2,T1,G) std::vector<int64_t> qShape = {16, 128, 512}; std::vector<int64_t> kShape = {2048, 1, 512}; std::vector<int64_t> softmaxLseShape = {1, 16, 128}; std::vector<int64_t> cuSeqQLenshape = {2}; std::vector<int64_t> cuSeqKLenshape = {2}; std::vector<int64_t> cmpResidualKShape = {1}; std::vector<int64_t> softmaxL1NormShape = {16, 1, 2048}; std::vector<int64_t> metadataShape = {64}; double softmaxScale = 0.088388; // 1 / sqrt(512) int64_t maxSeqlenK = 2048; int64_t cmpRatio = 128; int64_t maskMode = 3; char layoutQ[4] = {'T', 'N', 'D', 0}; char layoutK[4] = {'T', 'N', 'D', 0}; // 1. 调用 metadata 前置算子，完成负载均衡计算 uint64_t metadataWorkspaceSize = 0; aclOpExecutor* metadataExecutor = nullptr; ret = aclnnSparseFlashMlaSoftmaxL1NormMetadataGetWorkspaceSize( cuSeqQLen, cuSeqKLen, nullptr, nullptr, cmpResidualK, nullptr, 0, 16, 2048, 128, 1, 512, 0, cmpRatio, maskMode, layoutQ, layoutK, metadata, &metadataWorkspaceSize, &metadataExecutor); // ... 申请 metadataWorkspaceAddr 并调用 ret = aclnnSparseFlashMlaSoftmaxL1NormMetadata( metadataWorkspaceAddr, metadataWorkspaceSize, metadataExecutor, stream); aclrtSynchronizeStream(stream); // 2. 调用主算子 uint64_t workspaceSize = 0; aclOpExecutor* executor = nullptr; ret = aclnnSparseFlashMlaSoftmaxL1NormGetWorkspaceSize( q, k, softmaxLse, nullptr, cuSeqQLen, cuSeqKLen, nullptr, nullptr, cmpResidualK, nullptr, metadata, softmaxScale, maxSeqlenK, cmpRatio, maskMode, layoutQ, layoutK, softmaxL1Norm, &workspaceSize, &executor); // ... 申请 workspaceAddr 后执行 ret = aclnnSparseFlashMlaSoftmaxL1Norm( workspaceAddr, workspaceSize, executor, stream); aclrtSynchronizeStream(stream);

注意示例中 metadata 前置算子与主算子的入参顺序并不完全一致，两段接口的参数一一对应即可。示例还演示了标准的aclInit → aclrtSetDevice → aclrtCreateContext → aclrtCreateStream → 创建 aclTensor（ND 格式 + 连续 strides）→ 执行 → 释放资源 → aclFinalize完整生命周期；编译与运行方式可参考仓库的编译与运行样例。

调用方式二：PyTorch API

函数原型

cann_ops_transformer.sparse_flash_mla_softmax_l1_norm_metadata( num_heads_q, num_heads_k, head_dim, *, cu_seqlens_q=None, cu_seqlens_k=None, seqused_q=None, seqused_k=None, cmp_residual_k=None, topk_length=None, batch_size=None, max_seqlen_q=None, max_seqlen_k=None, topk=None, layout_q=None, layout_k=None, mask_mode=None, cmp_ratio=None ) -> Tensor

cann_ops_transformer.sparse_flash_mla_softmax_l1_norm( q, k, softmax_lse, *, sparse_indices=None, cu_seqlens_q=None, cu_seqlens_k=None, seqused_q=None, seqused_k=None, cmp_residual_k=None, topk_length=None, metadata=None, softmax_scale=1.0, max_seqlen_k=0, cmp_ratio=1, mask_mode=0, layout_q="BSND", layout_k="BSND" ) -> Tensor

Torch 扩展层通过 torch_extension/sparse_flash_mla_softmax_l1_norm.py 中的SparseFlashMlaSoftmaxL1NormOpBuilder注册算子 schema 与 Meta 实现，实际计算经由 csrc 的 C++ 扩展（csrc/sparse_flash_mla_softmax_l1_norm.cpp）落盘到 aclnn 接口。

单算子模式调用

在脚本中import cann_ops_transformer后即可直接调用，以 TND layout、Compressed Attention（cmp_ratio=128、mask_mode=3）场景为例：

import math import torch import torch_npu import cann_ops_transformer torch_npu.npu.set_device(0) S1 = 16 S2 = 2048 cmp_ratio = 128 actual_seq_q = [S1] actual_seq_k = [S2] T1 = sum(actual_seq_q) T2 = sum(actual_seq_k) B = 1 N1 = 128 N2 = 1 D = 512 scale_value = 1.0 / math.sqrt(D) dtype = torch.float16 input_layout = "TND" q_shape = (T1, N1, D) k_shape = (T2, N2, D) softmax_lse_shape = (N2, T1, N1 // N2) softmax_l1_norm_shape = (T1, N2, T2) cu_seq_qlen = [0] + [sum(actual_seq_q[:x+1]) for x in range(len(actual_seq_q))] cu_seq_klen = [0] + [sum(actual_seq_k[:x+1]) for x in range(len(actual_seq_k))] q = (torch.rand(q_shape).to(dtype)) * 2 k = (torch.rand(k_shape).to(dtype)) * 2 softmax_lse = (torch.rand(softmax_lse_shape).to(torch.float32)) cu_seq_qlen_tensor = torch.tensor(cu_seq_qlen).to(torch.int32).npu() cu_seq_klen_tensor = torch.tensor(cu_seq_klen).to(torch.int32).npu() cmp_residual_k = torch.zeros(B, dtype=torch.int32, device="npu") # 调用 sparse_flash_mla_softmax_l1_norm_metadata 完成负载均衡计算 metadata = cann_ops_transformer.sparse_flash_mla_softmax_l1_norm_metadata( N1, N2, D, cu_seqlens_q=cu_seq_qlen_tensor, cu_seqlens_k=cu_seq_klen_tensor, cmp_residual_k=cmp_residual_k, max_seqlen_q=S1, max_seqlen_k=S2, topk=0, cmp_ratio=cmp_ratio, mask_mode=3, layout_q=input_layout, layout_k=input_layout, ) # 调用 sparse_flash_mla_softmax_l1_norm 执行算子计算 softmax_l1_norm = cann_ops_transformer.sparse_flash_mla_softmax_l1_norm( q.npu(), k.npu(), softmax_lse.npu(), cu_seqlens_q=cu_seq_qlen_tensor, cu_seqlens_k=cu_seq_klen_tensor, cmp_residual_k=cmp_residual_k, metadata=metadata, softmax_scale=scale_value, max_seqlen_k=S2, cmp_ratio=cmp_ratio, mask_mode=3, layout_q=input_layout, layout_k=input_layout, ) torch_npu.npu.synchronize() assert softmax_l1_norm.shape == softmax_l1_norm_shape assert softmax_l1_norm.dtype == torch.float32 assert torch.isfinite(softmax_l1_norm.float()).all().item()

TorchAir 图模式调用

图模式将 metadata 前置算子的调用封装进torch.nn.Module，metadata 输出显式迁移到 NPU 后再作为主算子入参，配合torch.compile+torchair的 NPU 后端执行：

import math import torch import torch_npu import torchair import cann_ops_transformer from torchair.configs.compiler_config import CompilerConfig class SparseFlashMlaSoftmaxL1NormModel(torch.nn.Module): def __init__(self): super(SparseFlashMlaSoftmaxL1NormModel, self).__init__() def forward(self, m_inputs, npu_inputs): # 调用 sparse_flash_mla_softmax_l1_norm_metadata 完成负载均衡计算 metadata = torch.ops.cann_ops_transformer.sparse_flash_mla_softmax_l1_norm_metadata( **m_inputs ) # metadata 输出 Tensor 显式迁移到 NPU 后作为主算子入参 npu_inputs["metadata"] = metadata.npu() # 调用 sparse_flash_mla_softmax_l1_norm 执行算子计算 return torch.ops.cann_ops_transformer.sparse_flash_mla_softmax_l1_norm( **npu_inputs ) def sparse_flash_mla_softmax_l1_norm_acl_graph(m_inputs, npu_inputs): npu_mode = SparseFlashMlaSoftmaxL1NormModel().npu() config = CompilerConfig() config.mode = "reduce-overhead" npu_backend = torchair.get_npu_backend(compiler_config=config) torch._dynamo.reset() npu_mode = torch.compile(npu_mode, fullgraph=True, backend=npu_backend, dynamic=True) return npu_mode(m_inputs, npu_inputs) # m_inputs: 与单算子模式一致的 metadata 前置算子入参（dict） # npu_inputs: 主算子入参（dict），其中 softmax_scale 取 1.0 / math.sqrt(D) softmax_l1_norm = sparse_flash_mla_softmax_l1_norm_acl_graph(m_inputs, npu_inputs) torch_npu.npu.synchronize() assert softmax_l1_norm.shape == (T1, N2, T2) assert softmax_l1_norm.dtype == torch.float32 assert torch.isfinite(softmax_l1_norm.float()).all().item()

源码实现解读：从 Host 侧到 Kernel 侧

算子定义与默认属性

sparse_flash_mla_softmax_l1_norm_def.cpp 以OpDef方式注册算子：

输入：q、k（REQUIRED，DT_FLOAT16/DT_BF16，ND 格式）；softmax_lse（REQUIRED，DT_FLOAT）；其余 sparse_indices、cu_seqlens_q/k、seqused_q/k、cmp_residual_k、topk_length、metadata 均为 OPTIONAL，数据类型 INT32，格式 ND。所有张量输入均声明AutoContiguous()。
输出：softmax_l1_norm（REQUIRED，DT_FLOAT，ND）。
属性默认值：softmax_scale=1.0、max_seqlen_k=0、cmp_ratio=1、mask_mode=0、layout_q="BSND"、layout_k="BSND"，与 README 参数表完全一致。
AICore 配置仅注册ascend950，且开启了DynamicCompileStaticFlag(true)、DynamicFormatFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等动态能力。

shape 推导逻辑

sparse_flash_mla_softmax_l1_norm_infershape.cpp 负责输出 shape 推导：

首先校验layout_q必须为 "BSND" 或 "TND"，否则报错返回。
BSND：输出 4 维，第 0/1/2 维取 q 的 B、S1 与 k 的 N2；第 3 维在存在非空 sparse_indices 时取 K（sparseIndiceShape 第 3 维），否则取 k 的 S2。
TND：输出 3 维，第 0 维取 q 的 T1、第 1 维取 k 的 T2；第 2 维在存在非空 sparse_indices 时取 K，否则取max_seqlen_k——这正是max_seqlen_k属性"用于 TND dense 场景输出 shape 推导"的落点。
输出数据类型固定为 FLOAT32。

tiling 数据结构

sparse_flash_mla_softmax_l1_norm_tiling.h 定义了下沉到 kernel 的 tiling 结构SmlaSoftmaxL1NormTilingData，字段包括 b、sq、sk、g、d、t1、t2、max_seqlen_k、k_length、cmp_ratio、init_per_core_num、init_total_num、softmax_scale 以及 has_seqused_q/has_seqused_k/has_topk_length 三个布尔开关，同时携带 aivNum/aicNum/ubSize/l1Size/l0a/l0b/l0c 等硬件资源参数，供 tiling 决策分核与分块策略。

kernel 侧分核与分块策略

从 sparse_flash_mla_softmax_l1_norm.py（pypto-pro 编写）可以观察到以下关键实现细节：

tile 常量：TS=1、TG=128（G 最大分块）、TG_HALF=64、TKV=128、TD=128、D_TOTAL=512，与"G 最大为 128、D 固定 512"的规格呼应。
TilingKey：包含 SPARSE_MODE（2bit，取值 0~3）、IS_TND、IS_SPARSE 三个编译期开关，据此分支出 sparse/dense 与 BSND/TND 共 4 种 kernel 变体。
sparse 场景的 Gather：gather_k阶段按GATHER_ROW_NUM=32行一组，逐行读取 sparse_indices 索引并对 K 做 load，经 ND→NZ 转换后写入 L1，供 Cube 单元 QK matmul 使用；非 sparse 场景则直接按 ki 偏移加载 K 分块。
QK 计算：compute_qk/compute_qk_dense在 Cube 单元按 TD=128 的 D 维分块做matmul/matmul_acc累加，结果经AccToVecMode.DualModeSplitM从 Accumulator 搬到 Vector 单元。
Softmax L1Norm：compute_softmax_l1_norm阶段在 Vector 单元对 QK 结果与 lse 做exp_sub并跨 G 维累加，再乘以g_scale（= 1/G）完成归一化，最后以AtomicAdd原子写方式累加写入输出张量（因为同一 seq 的 G 个 head group 会被拆到多个核上并行处理）。
G 维拆分：单个核内部将 G 拆成前后两半（first_half_g / second_half_g），配合子核（sub_id）实现 2 路并行，再通过 metadata 的 strided 分核信息在各 AIC 核间分配 seq 任务。

该实现路径完整印证了 README 中"metadata 采用 strided 方式将任务均衡切分到可用 AIC 核"的表述。

总结

SparseFlashMlaSoftmaxL1Norm 是 CANN ops-transformer 面向 Ascend 950 平台提供的 MLA 训练配套算子，通过"metadata 前置算子做负载均衡 + 主算子做 Softmax L1Norm 计算"的两段式设计，兼顾了稀疏压缩注意力的索引 Gather、数值稳定的 SimpleSoftmax 归一化，以及 G 维跨核并行累加的性能要求。接入时需重点关注三点：其一，严格按照约束表控制 N1/N2/D/cmp_ratio/mask_mode 等规格；其二，TND 场景务必成对传入 cu_seqlens_q 与 cu_seqlens_k，sparse 场景务必传入准确的 topk_length；其三，metadata 必须由前置算子生成后传入主算子。开发者可分别参考 aclnn 接口文档、PyTorch 接口文档与 aclnn 调用示例完成两种调用方式的接入。