torchtitan-npu NPU融合算子清单DSA、SMLA、GMM、MoE Dispatch四大Transformer加速秘籍【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-nputorchtitan-npu 是 torchtitan 的昇腾Ascend扩展插件通过NPU 融合算子将 Transformer 模型中的原始模块自动替换为 Ascend 亲和实现从而在 NPU 平台上大幅提升训练性能。本文将带你快速了解其中最重要的四大 Transformer 加速算子——DSA 稀疏注意力、SMLA 稀疏 MLA、GMM 分组矩阵乘、MoE Dispatch 专家分发以及它们的启用方式和源码位置帮你开箱即优地跑通昇腾大模型训练。上图展示了 torchtitan-npu 在整体技术栈中的位置它夹在 torchtitan模型定义、并行策略与 CANNAscendC Kernel、运行时之间其中的Module Convert能力正是 NPU 融合算子适配的核心入口。一步启用 NPU 融合算子转换器配置指南融合算子不需要你手写任何算子代码只需在训练配置中声明要启用的 converter 列表from torchtitan.protocols.model_converter import ModelConvertersContainer from torchtitan_npu.converters import get_model_converter_config model_converters ModelConvertersContainer.Config( converters[ get_model_converter_config(npu_dsa), get_model_converter_config(npu_rms_norm), get_model_converter_config(npu_moe_dispatch), get_model_converter_config(npu_gmm), ], )每个 converter 启动训练时会自动匹配模型中的目标模块并替换为融合实现。各模型的默认组合已预置在对应的注册表中例如DeepSeek-V3config_registry.pyDeepSeek-V3.2config_registry.pyDeepSeek-V4 / Qwen3config_registry.py秘籍一DSA 稀疏注意力——只算最值钱的 Token 适合场景DeepSeek-V3.2 的长上下文 Attention 加速DSADeepSeek Sparse Attention是 DeepSeek-V3.2 引入的稀疏注意力机制分两步走Lightning Indexer先筛选出少量高价值 token 的索引稀疏 Attention 只对选中索引做计算大幅减少无效注意力开销。npu_dsaconverter 将这两个核心组件替换为对应的 NPU 融合算子DeepSeek-V3.2 组件对应 NPU 融合算子Lightning Indexer 前向npu_lightning_indexerLightning Indexer 反向梯度 Lossnpu_sparse_lightning_indexer_grad_kl_loss稀疏注意力计算npu_sparse_flash_attention 源码dsa.py冒烟测试见 test_dsa.py。秘籍二SMLA 稀疏 MLA——DeepSeek-V4 的注意力快车道npu_smla面向 DeepSeek-V4 的稀疏 MLA 注意力路径会根据硬件自动选择最优实现A5 场景走cann_ops_transformer提供的sparse_flash_mla、lightning_indexer及对应的 metadata/反向融合算子其他场景沿用兼容 DeepSeek-V4 sparse attention / LI / LI loss 的既有 NPU 实现。 一个小技巧SMLA 路径的 LI loss 计算已整合进融合反向算子默认不再额外打印以避免性能开销。调试时可将torchtitan_npu.converters.kernels.npu_smla日志开到 DEBUG 级别查看。 源码npu_smla.py。秘籍三GMM 分组矩阵乘——一次调用算完所有专家 适合场景MoE 模型专家层矩阵运算加速MoE 模块中每个专家都要执行结构相同的 FFN 矩阵乘升维w1→ 激活 → 降维w2。逐个专家循环调用矩阵乘效率低下——npu_grouped_matmulGMM算子将所有专家的同一线性层合并为一次算子调用并行完成。它依赖 MoE Dispatch 输出的重排后 token 及专家索引。DSV3 / DSV3.2 / DSV4 / Qwen3 MoE 的标准 ExpertParallel 场景统一搭配npu_moe_dispatch使用。 源码gmm.py测试见 test_gmm.py。进阶npu_swiglu_group还能把 GMM 两次矩阵乘之间的限幅、SwiGLU 激活、routed-score 缩放合并为一次算子调用减少中间张量和算子下发开销源码见 swiglu_group.py。秘籍四MoE Dispatch——专家分发与重排一体化npu_moe_dispatch是 DS 系列和 Qwen3 MoE 标准ExpertParallel路径的总入口负责整条分发链路router 后使用npu_moe_token_permute完成第一组 token/expert 聚合expert 计算后用 unpermute 还原通过并行策略更新器把标准ExpertParallel自动替换为NpuExpertParallelEP all-to-all 之后用npu_moe_re_routing优化本卡局部重排替换旧的repeat_interleave npu_moe_token_permute削减热点冗余算子。⚠️ 注意router 后的第一次 permute按 top-k 索引生成专家槽位与 all-to-all 后的本卡重排按 experts 计数调整本地 buffer语义不同两者不可互相替代模型配置直接启用npu_moe_dispatch即可。 源码moe_dispatch.py、permutation.py。顺手了解其他常用融合算子速查表除了四大主角torchtitan-npu 还内置了一批高频小算子的融合适配按需启用即可Converter 名称加速目标源码位置npu_rms_normRMSNorm 融合rms_norm.pynpu_rope旋转位置编码融合rope.pynpu_rope_inplace_partialDeepSeek-V4 partial RoPE 原地写入inplace_partial_rope.pynpu_mhc_pre/npu_mhc_postDeepSeek-V4 MHC 前后处理mhc_prepost.pynpu_mhc_head_compute_mix_tilelangHcHead 多头混合计算mhc_head_compute_mix_tilelang.pynpu_moe_reduce_fused_tilelangMoE combine 阶段 TileLang 融合tilelang 算子目录更完整的介绍、配置示例和硬件路径说明请查阅官方文档npu_fused_ops.md。如何选择你的加速组合跑 DeepSeek-V3.2直接采用默认注册表npu_dsa npu_rms_norm npu_rope npu_moe_dispatch npu_gmm组合已就绪跑 DeepSeek-V4A5 上优先启用npu_smlanpu_mhc_post建议按需打开A3 上保持关闭避免性能退化跑 Qwen3 MoEnpu_moe_dispatch npu_gmm是标准搭配想验证算子正确性跑一遍 tests/smoke_tests/features/ 下的对应测试文件即可。注册机制本身很简洁所有 converter 通过 registry.py 中的register_model_converter登记名字get_model_converter_config取回配置对象——想新增或理解某个算子沿着名字找到 kernels 目录下的对应文件就够了。【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考