算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载SparseFlashMlaSoftmaxL1Norm 是 CANN ops-transformer 算子库中服务于 Sparse Flash MLAMulti-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稀疏压缩注意力两类场景。在训练链路中它是aclnnDenseLightningIndexerKLLossGraddense_lightning_indexer_kl_loss_grad反向算子的配套正向接口其输出可用于反向梯度计算。调用过程遵循先负载均衡、后主计算的两段式模式由两个算子协作完成SparseFlashMlaSoftmaxL1NormMetadatametadata 前置算子根据主算子的 shape、layout、mask 等信息采用 strided 方式将任务均衡切分到可用 AIC 核上输出 metadata 供主算子使用。SparseFlashMlaSoftmaxL1Norm主算子根据 metadata 中的分核信息对 Q 和 K 计算 Softmax L1Norm。metadata 的字段布局在 torch 接口文档 与 kernel 侧头文件 中均有明确定义固定为 64 个 INT32 元素SMLA_METADATA_SIZE 64核心字段如下字段index说明totalNum0参与负载均衡的 seq 总数。formerCoreProcessNum1常规核处理的 seq 数即ceil(totalNum / totalCoreNum)。remainCoreProcessNum2尾核处理的 seq 数即floor(totalNum / totalCoreNum)。remainCoreNum3尾核数目。totalCoreNum4实际使用的 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 $$阶段二计算 PSimpleSoftmax对 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, dimG)}{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 结果先乘 scalevf.muls再与 lse 做exp_sub得到逐元素 exp 值按 8 路展开累加求和最后除以g_scalar即 1/G后写入输出与上述公式逐项对应。参数说明主算子 sparse_flash_mla_softmax_l1_norm参数名输入/输出/属性描述数据类型数据格式维度shapeq输入attention 结构的输入 Q。FLOAT16、BFLOAT16ND(B,S1,N1,D) 或 (T1,N1,D)k输入attention 结构的输入 K(V)。FLOAT16、BFLOAT16ND(B,S2,N2,D) 或 (T2,N2,D)softmax_lse输入注意力正向计算的输出 softmaxLse计算公式详见 sparse_flash_mla 文档。FLOAT32ND(B,N2,S1,G) 或 (N2,T1,G)sparse_indices可选输入稀疏场景下选择的 k 中权重较高的注意力索引。INT32ND(B,S1,N2,K) 或 (T1,N2,K)cu_seqlens_q可选输入每个 Batch 中 Query 的有效 token 数的累加和形式layout 为 TND 时该参数必传。INT32ND(B1,)cu_seqlens_k可选输入每个 Batch 中 Key 的有效 token 数的累加和形式layout 为 TND 时该参数必传。INT32ND(B1,)seqused_q可选输入表示不同 batch 中 q 实际参与运算的 token 数。INT32ND(B,)seqused_k可选输入表示不同 batch 中 k 实际参与运算的 token 数。INT32ND(B,)cmp_residual_k可选输入表示每个 batchS2 // cmpRatio后的余数当 k 不为空且 mask_mode3 时必须传入。INT32ND(B,)topk_length可选输入表示每行 q 对应的 k 实际可选的 topk 长度mask_mode0 且存在稀疏索引时需要传。INT32ND(B,S1,N2) 或 (T1,N2)metadata可选输入表示 tiling 下沉的 aicpu 算子输出结果由SparseFlashMlaSoftmaxL1NormMetadata算子生成。INT32ND(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 模式0No mask3rightDownCausal 模式默认值 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 则该输出不为空其他场景下输出为空。FLOAT32ND(B,S1,N2,S2) 或 (T1,N2,T2)关于 shape 的几点补充可在 aclnn 接口文档 中查到完整说明q、k 的 N 轴对应关系需满足 GQA 约束N1 N2 × Gk 的 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 时存在长度与 B1 保持一致累加和分别与 T1、T2 保持一致。cmp_residual_k 的强制传入条件是maskMode3 且 cmpRatio!1。topk_length 在 mask_mode0 且存在稀疏索引时必传且必须为准确值。metadata 为必传项形状为 (64,)。metadata 前置算子 sparse_flash_mla_softmax_l1_norm_metadata前置算子不接收张量形式的 q/k而是接收描述注意力结构的标量参数与可选张量参数名参数类型可选/必选描述数据类型num_heads_qint必选公式中 Q 的头数即 N1当前支持 1~128。intnum_heads_kint必选公式中 key 的头数即 N2当前仅支持 1。inthead_dimint必选头的维度即 D当前仅支持 512。intcu_seqlens_qTensor可选每个 Batch 中 q 的有效 token 数的累加和形式layout_q 为 TND 时必传累加和与 T1 一致。int32cu_seqlens_kTensor可选每个 Batch 中 k 的有效 token 数的累加和形式layout_k 为 TND 时必传累加和与 T2 一致。int32seqused_qTensor可选不同 batch 中 q 实际参与运算的 token 数。int32seqused_kTensor可选预留接口参数当前 kernel 路径暂不使用。int32cmp_residual_kTensor可选预留接口参数当前 kernel 路径暂不使用。int32topk_lengthTensor可选每行 q 对应的 k 实际可选的 topk 长度。int32batch_sizeint可选输入样本批量大小即 B默认 NoneBSND 场景需传正数TND 场景可为 None 自动推导。intmax_seqlen_qint可选q 的最大序列长度默认 NoneBSND 场景必须为正数。intmax_seqlen_kint可选k 的最大序列长度默认 NoneBSND 场景必须为正数。inttopkint可选从 k 中筛选出的关键 token 个数0 表示无稀疏默认 None。intlayout_qstr可选q 的数据排布格式支持 BSND、TND默认 None内部转为 BSND。stringlayout_kstr可选k 的数据排布格式支持 BSND、TND默认 None内部转为 BSND。stringmask_modeint可选0No mask3rightDownCausal 模式默认 None内部转为 0。intcmp_ratioint可选对 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仅支持 1D仅支持 512q、k 最后一维保持一致。mask_mode 支持模式含义备注0不做 mask 操作支持3rightDownCausal 模式的 mask对应以右顶点为划分的下三角场景支持cmp_ratio取值范围1~128。TND 场景必传 cu_seqlens_q 和 cu_seqlens_kbatch_size 可为 None通过 cu_seqlens_q 推导。BSND 场景必传 batch_size、max_seqlen_q 和 max_seqlen_k。metadata 必须传入且由SparseFlashMlaSoftmaxL1NormMetadata算子生成。入参为空处理q 为空 Tensor 时直接返回。数据类型q、k 的数据类型必须保持一致FLOAT16/BFLOAT16。适用场景该接口支持训练场景使用支持单算子模式和 TorchAir 图模式调用。调用方式一aclnn APIC 两段式接口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_NULLPTR161001参数中存在非法的 nullptr。ACLNN_ERR_PARAM_INVALID161002输入的数据类型不满足支持类型。ACLNN_ERR_PARAM_INVALID161002q、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::vectorint64_t qShape {16, 128, 512}; std::vectorint64_t kShape {2048, 1, 512}; std::vectorint64_t softmaxLseShape {1, 16, 128}; std::vectorint64_t cuSeqQLenshape {2}; std::vectorint64_t cuSeqKLenshape {2}; std::vectorint64_t cmpResidualKShape {1}; std::vectorint64_t softmaxL1NormShape {16, 1, 2048}; std::vectorint64_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 → 创建 aclTensorND 格式 连续 strides→ 执行 → 释放资源 → aclFinalize完整生命周期编译与运行方式可参考仓库的编译与运行样例。调用方式二PyTorch API函数原型cann_ops_transformer.sparse_flash_mla_softmax_l1_norm_metadata( num_heads_q, num_heads_k, head_dim, *, cu_seqlens_qNone, cu_seqlens_kNone, seqused_qNone, seqused_kNone, cmp_residual_kNone, topk_lengthNone, batch_sizeNone, max_seqlen_qNone, max_seqlen_kNone, topkNone, layout_qNone, layout_kNone, mask_modeNone, cmp_ratioNone ) - Tensorcann_ops_transformer.sparse_flash_mla_softmax_l1_norm( q, k, softmax_lse, *, sparse_indicesNone, cu_seqlens_qNone, cu_seqlens_kNone, seqused_qNone, seqused_kNone, cmp_residual_kNone, topk_lengthNone, metadataNone, softmax_scale1.0, max_seqlen_k0, cmp_ratio1, mask_mode0, layout_qBSND, layout_kBSND ) - TensorTorch 扩展层通过 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 Attentioncmp_ratio128、mask_mode3场景为例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[:x1]) for x in range(len(actual_seq_q))] cu_seq_klen [0] [sum(actual_seq_k[:x1]) 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, dtypetorch.int32, devicenpu) # 调用 sparse_flash_mla_softmax_l1_norm_metadata 完成负载均衡计算 metadata cann_ops_transformer.sparse_flash_mla_softmax_l1_norm_metadata( N1, N2, D, cu_seqlens_qcu_seq_qlen_tensor, cu_seqlens_kcu_seq_klen_tensor, cmp_residual_kcmp_residual_k, max_seqlen_qS1, max_seqlen_kS2, topk0, cmp_ratiocmp_ratio, mask_mode3, layout_qinput_layout, layout_kinput_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_qcu_seq_qlen_tensor, cu_seqlens_kcu_seq_klen_tensor, cmp_residual_kcmp_residual_k, metadatametadata, softmax_scalescale_value, max_seqlen_kS2, cmp_ratiocmp_ratio, mask_mode3, layout_qinput_layout, layout_kinput_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.Modulemetadata 输出显式迁移到 NPU 后再作为主算子入参配合torch.compiletorchair的 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_configconfig) torch._dynamo.reset() npu_mode torch.compile(npu_mode, fullgraphTrue, backendnpu_backend, dynamicTrue) 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、kREQUIREDDT_FLOAT16/DT_BF16ND 格式softmax_lseREQUIREDDT_FLOAT其余 sparse_indices、cu_seqlens_q/k、seqused_q/k、cmp_residual_k、topk_length、metadata 均为 OPTIONAL数据类型 INT32格式 ND。所有张量输入均声明AutoContiguous()。输出softmax_l1_normREQUIREDDT_FLOATND。属性默认值softmax_scale1.0、max_seqlen_k0、cmp_ratio1、mask_mode0、layout_qBSND、layout_kBSND与 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 时取 KsparseIndiceShape 第 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.pypypto-pro 编写可以观察到以下关键实现细节tile 常量TS1、TG128G 最大分块、TG_HALF64、TKV128、TD128、D_TOTAL512与G 最大为 128、D 固定 512的规格呼应。TilingKey包含 SPARSE_MODE2bit取值 0~3、IS_TND、IS_SPARSE 三个编译期开关据此分支出 sparse/dense 与 BSND/TND 共 4 种 kernel 变体。sparse 场景的 Gathergather_k阶段按GATHER_ROW_NUM32行一组逐行读取 sparse_indices 索引并对 K 做 load经 ND→NZ 转换后写入 L1供 Cube 单元 QK matmul 使用非 sparse 场景则直接按 ki 偏移加载 K 分块。QK 计算compute_qk/compute_qk_dense在 Cube 单元按 TD128 的 D 维分块做matmul/matmul_acc累加结果经AccToVecMode.DualModeSplitM从 Accumulator 搬到 Vector 单元。Softmax L1Normcompute_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_ksparse 场景务必传入准确的 topk_length其三metadata 必须由前置算子生成后传入主算子。开发者可分别参考 aclnn 接口文档、PyTorch 接口文档 与 aclnn 调用示例 完成两种调用方式的接入。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 的 AllGatherMatmul 通算融合算子原理、参数与 aclnn 调用实战CANN ops transformer 的 AllGatherMatmul 通算融合算子原理、参数与 aclnn 调用实战 导读 AllGatherMatm算子库人工智能深度学习AscendCANN ops-transformer 算子解析DistributeBarrierExtend 全卡同步算子原理与 aclnn 调用实践CANN ops transformer 算子解析DistributeBarrierExtend 全卡同步算子原理与 aclnn 调用实践 本技术指南围绕 C算子库人工智能深度学习AscendCANN ops-transformer 的 FusedFloydAttention 算子深度解析算法原理、aclnn 调用与 Tiling 设计CANN ops transformer 的 FusedFloydAttention 算子深度解析算法原理、aclnn 调用与 Tiling 设计 导读 Fu算子库人工智能深度学习Ascend上一篇1Remote主题与语言定制终极指南打造个性化远程管理界面的10个技巧下一篇fastai AzureML集成云端训练与部署一体化终极指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考