CANN 多流表达功能实战指南：基于 TorchAir npu_stream_switch 的 GE 图模式并行优化

CANN 多流表达功能实战指南基于 TorchAir npu_stream_switch 的 GE 图模式并行优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读大模型推理中模型的执行关键路径上常常存在 Cube 与 Vector 利用率不均衡、HCCL 通信与本地计算串行等硬件资源空闲现象。本指南以 torchair 的Ascend IR 图内资源并发多流表达功能为核心讲解在 GE 图模式下通过npu_stream_switch把算子分发到不同 stream 实现并行、用npu_wait_tensor精确控制跨流时序的完整方法。读完本文你将掌握多流的适用场景判定、API 使用约束、可运行的完整示例以及如何与控核limit_core_num联动并学会从当前仓库的多流优化案例中复用编排思路。术语说明上游 torchair 文档中的「Ascend IR」与「GE 图模式」指同一条执行路径分别从 IR 侧和执行侧描述——通过torchair.CompilerConfig()把 PyTorch FX 图转为 Ascend IR再由 GEGraph Engine编译执行因此「多流表达Ascend IR」文档中写「仅适用于 GE 图模式」并不矛盾。本文基于仓库内ascend_ir_multi_stream.md展开并辅以仓库源码佐证。一、功能定位Ascend IR 图内资源并发大模型推理中某些场景天然存在可并行的计算分支。多流表达功能允许用户在脚本中为每个算子显式指定执行 stream把原本需要串行执行的多个算子分发到不同 stream 上并行计算使多条 stream 上的计算形成 overlap从而降低整体计算耗时。并行包含两种典型形态计算与计算并行基于数据依赖关系分析出可以并行的多条计算分支分别指定 stream 并行执行。计算与通信并行针对没有数据依赖的通信操作提前占用通信资源执行通信任务与主计算流上的计算形成 overlap。本功能主要面向Ascend IR 图内资源并发max-autotune 模式尤其针对Cube 计算资源未完全使用的场景。这一点是选型的第一判断依据若 Cube 计算资源已完全使用不建议开启本功能否则可能引入额外的调度开销导致原计算性能劣化。从仓库源码可以印证这一资源未用满才值得并行的工程判断在docs/cann/zh/multi_stream_principles.md中将算子按算力使用情况划分为纯 Cube 算子如非量化 Matmul、纯 Vector 算子如 RMSNorm和混合算子如量化 Matmul、FlashAttention 系列指出一般算子不能将全部算力用满这正是多流加速的优化空间所在同时指出 Cube 与 Vector 可独立调度的分离式平台上将 Cube 算子与 Vector 算子交替编排在两条 Stream 上形成互补是收益最稳定的并行形态之一。二、使用约束多流表达功能不是任意场景都能开启使用前必须确认以下约束仅适用于 GE 图模式场景即torchair.CompilerConfig()torchair.get_npu_backend()的图编译路径。纯 Vector 场景与含 Cube 场景纯 Vector 场景的计算耗时一般在可接受范围内收益有限含 Cube 计算的场景开启本功能后的效益往往优于纯 Vector 计算场景这与上一节的定位面向 Cube 资源未用满一致。静态 Shape 场景本功能与单流执行功能enable_single_stream冲突不支持同时开启。本功能不推荐在 SuperKernel 内设置算子多 stream 并行如有需要应使用图内标定 SuperKernel 范围时的stream-fusion编译选项配置。动态 Shape 场景默认单流模式用户可通过如下 CANN 环境变量开启多流。一旦开启多流其功能优先级低于脚本内的显式多流表达即通过npu_stream_switch显式指定的多流优先。export ENABLE_DYNAMIC_SHAPE_MULTI_STREAM1这些约束在仓库的技能规则中进一步收敛为两条可操作的执行路径定界原则见 SKILL.mdtorchair.CompilerConfig()torchair.get_npu_backend()走 Ascend IR / GE 图模式路径torch.compile(..., backendnpugraph_ex)走 npugraph_ex / aclgraph 路径两条路径的 API 与约束不能混用也不要把 Ascend IR 的图内 API 和 aclgraph 的显式 stream API 混写进同一套实现。三、使用方法多流表达的使用分为三步前两步为必选第三步为可选的时序控制。3.1 分析可并行算子用户自行分析模型中可进行并行计算的算子判断依据是多个算子的输入都依赖于同一个前驱算子的输出、彼此之间没有数据依赖这些算子在拓扑上即可并发执行。典型结构是前置算子 A 的输出同时被 B、C 两条分支消费B、C 之间无数据依赖可分别承载于两条 stream 并行执行最终在汇合算子 D 合并。3.2 开启多流表达使用如下 with 语句块npu_stream_switch语句块内下发的算子切换至stream_tag流语句块外的算子使用默认 stream 计算with torchair.scope.npu_stream_switch(stream_tag: str, stream_priority: int 0, enable_inner_parallel: bool True):三个参数的语义如下参数类型/默认值说明stream_tagstr必填需要切换到的流的标签相同的标签代表相同的流由用户控制。tag 在全局范围内唯一标识一条 stream跨模块复用同 tag 会采用相同 stream。stream_priorityint默认0切换到stream_tag流的优先级即 Runtime 在并发时优先给高优先级的流分配核资源当前版本使用默认值 0 即可。enable_inner_parallelbool默认True是否使能stream_tag流内的算子按照 GE 原有的并发策略分流默认开启。3.3 控制并行计算的时序可选通过npu_wait_tensor接口实现时序控制指定算子 a 等待算子 b 执行完后执行。该接口用于两条流之间需要精确控制流排序的场景——当数据流存在跨 stream 依赖一个 stream 的输入来自另一个 stream或需要控制算子的执行顺序以错开资源利用时使用。结合仓库源码看npu_wait_tensor在实际模型中的典型用法是把来自副流tag 流的中间结果作为等待锚点确保主流消费该结果前其计算已完成。例如在models/deepseek_v3_2_exp/models/modeling_deepseek.py中进入npu_stream_switch(True, 22)副流后先执行hidden_states npu_wait_tensor(True, hidden_states, topk_idx)即让副流中的计算等待topk_idx就绪后再继续保证跨流数据依赖的时序正确。四、完整使用示例原文档给出了一个可直接运行的完整示例它演示了三个不同 stream 上的算子如何通过npu_wait_tensor建立跨流时序add_result、mm1在默认 streammm_result在流 1add2在流 2import torch, os import torch_npu import torchair as tng from torchair.configs.compiler_config import CompilerConfig from torchair.core.utils import logger import logging logger.setLevel(logging.DEBUG) # 定义模型model # add_result、mm1在默认streammm_result在流“1”add2在流“2” class Model(torch.nn.Module): def __init__(self): super().__init__() def forward(self, in1, in2, in3, in4): add_result torch.add(in1, in2) with tng.scope.npu_stream_switch(1): # torch.mm算子(mm_result)等待torch.add算子(add_result)执行完再执行 tng.scope.npu_wait_tensor(in4, add_result) mm_result torch.mm(in3, in4) mm1 torch.mm(in3, in4) with tng.scope.npu_stream_switch(2): # torch.add算子(add2)等待torch.mm算子(mm_result)执行完再执行 tng.scope.npu_wait_tensor(in4, mm_result) add2 torch.add(in3, in4) return add_result, mm_result, mm1, add2 model Model() config CompilerConfig() config.debug.graph_dump.type pbtxt npu_backend tng.get_npu_backend(compiler_configconfig) # 调用compile接口编译模型 model torch.compile(model, backendnpu_backend, dynamicFalse, fullgraphTrue) in1 torch.randn(1000, 1000, dtype torch.float16).npu() in2 torch.randn(1000, 1000, dtype torch.float16).npu() in3 torch.randn(1000, 1000, dtype torch.float16).npu() in4 torch.randn(1000, 1000, dtype torch.float16).npu() result model(in1, in2, in3, in4) print(fResult:\n{result}\n)对示例的要点解读示例同时使用了config.debug.graph_dump.type pbtxt开启图结构 dump。编译后可在 dump 出的 pbtxt 文件中查看流间的控制关系虚线控制边表示显式的跨流数据依赖同步关系即npu_wait_tensor在图中生成的跨流等待边这是验证多流是否按预期落图的重要手段。dynamicFalse, fullgraphTrue表示静态 shape 全图编译与本文第二部分静态 shape 场景的约束对应。整个示例用torch.mm/torch.add构造了三条逻辑流默认流上同时存在add_result被流 1 等待与mm1流 1 上的mm_result被流 2 等待最终四条结果统一返回。可见npu_wait_tensor的关键语义是当前 scope 内的算子等待指定的前驱 tensor 就绪后再执行。五、多流与控核联动limit_core_num多流并行时可能出现所有核Core都被一个流占用的情况导致另一个流上的算子并行度降低这时需要把核分给不同的流使用。GE 图模式提供了两种限核方法算子级核数与全局核数session配置其中算子级优先级高于全局核数配置。5.1 算子级核数with torchair.scope.limit_core_num(op_aicore_num: int, op_vectorcore_num: int):op_aicore_num该算子运行时的最大 AI Core 数取值范围[1, max_aicore]。op_vectorcore_num该算子运行时的最大 Vector Core 数取值范围[1, max_vectorcore]当 AI 处理器上仅存在 AI Core 不存在 Vector Core 时仅支持取值为 0。最大核数可通过 CANN 软件安装目录下的platform_config/soc_version.ini文件查看例如[SoCInfo] ai_core_cnt24 cube_core_cnt24 vector_core_cnt48表示该 AI 处理器上存在 24 个 Cube Core、48 个 Vector Core。配置的核数不能超过处理器允许的最大值实际运行核数可能少于配置的最大核数。配置结果可通过图结构 dump 查看设置config.debug.graph_dump.type txt在算子attr属性中查看 key 为_op_aicore_num和_op_vectorcore_num的取值。5.2 全局核数通过torchair.get_npu_backend的compiler_config配置import torch_npu, torchair config torchair.CompilerConfig() # 全局核数配置项 config.ge_config.aicore_num 24|100 npu_backend torchair.get_npu_backend(compiler_configconfig) opt_model torch.compile(model, backendnpu_backend)aicore_num为字符串类型形如${aicore_num}|${vectorcore_num}必须用|分隔。${aicore_num}表示全局 AI Core 数取值范围[1, max_aicore]${vectorcore_num}表示全局 Vector Core 数取值范围[1, max_vectorcore]。对于仅存在 AI Core 无 Vector Core 的处理器可配置为24|或24。5.3 多流场景下的控核要点在两条并行 stream 上两条流通常都要限核并行分支所有核数加起来约等于总核数若超过总核数可能导致无法并行甚至资源互锁卡死。仓库中 LongCat-Flash 的多流实现是这一原则的典型落地——在models/longcat_flash/models/modeling_longcat_flash.py中副流以with npu_stream_switch_gegraph(True, 1): with limit_core_num(True, self.aic_num1, self.aiv_num1):承载 shortcut MoE 路径主流以with limit_core_num(not self.enable_afd, self.aic_num2, self.aiv_num2):承载 dense/attention 主路径两套核数配额独立配置对应案例文档 longcat-flash-multi-stream-limit-core.md 中aic_num1/aiv_num1与aic_num2/aiv_num2的设计。该案例明确指出多流但不控核收益可能被拖尾抵消控核参数不是通用值要结合 profile 调整。六、仓库中的多流实战案例当前仓库沉淀了多个已落地的多流优化案例均可在 examples 目录 下找到可作为不同模型结构的编排参考案例优化方法代表模型对应源码MoE 共享专家双流并行共享专家放副流与路由专家路径重叠DeepSeek-V3.2-Exp / DeepSeek-R1modeling_deepseek.py、modeling_glm.pyIndexer Prolog 多流并行Attention 前处理 Q 路径与权重投影路径拆流DeepSeek-V3.2-Exp / GLM-5indexer.pyKVCache Offload 异步搬运流独立流完成 KVCache 搬运降低主计算流阻塞DeepSeek-V3.2-Exp / GLM-5offload_cache.pyPrefill Micro-Batch 双流流水两个 micro-batch 分跑两条流event 编排 dispatch/combineDeepSeek-R1modeling_deepseek.pyLongCat-Flash 多流与控核联动多流 limit_core_num分核减少拖尾LongCat-Flashmodeling_longcat_flash.pyHunyuanImage-3.0 MoE 多流变体初始化阶段注入share_mlp_stream承载共享 MLPHunyuanImage-3.0hunyuan.py以 MoE 共享专家双流为例详细见 moe-shared-expert-dual-stream.md其核心是decode 阶段共享专家与路由专家之间无数据依赖把共享专家放进副流即可与 gating、dispatch、路由专家计算重叠执行。最小切流形态if self.n_shared_experts 0: enable_multi_streams self.enable_multi_streams and not is_prefill with npu_stream_switch(enable_multi_streams, 11): hidden_states_share self.shared_experts( hidden_states.view(-1, hidden_states.shape[-1]) )叠加 ACL graph event 时则在切流前后记录和等待 tagged event常见写法为tng.ops.npu_record_tagged_stream(hidden_states, 11)、tng.ops.npu_tagged_event_record(moe_npu_events[0])与副流内的tng.ops.npu_tagged_event_wait(moe_npu_events[0])注意切流/时序的npu_stream_switch、npu_wait_tensor在torchair.scope命名空间而 record/wait/tagged event 在torchair.ops命名空间是两个不同的命名空间。若同时开启图模式 SuperKernel可通过stream-fusion1编译选项在 SuperKernel 内实现多流融合。七、常见误区与调试建议结合 SKILL.md 与 api-routing.md 中沉淀的实践规则使用多流表达功能时需注意Stream ID 切对 ≠ 物理并行成立npu_stream_switch给副流贴的是逻辑标签GE 编译期会重排算子——尤其会把消费副流输出的轻量 precompute如 Cast/Swish/Sigmoid拉到主流导致主流被 barrier 卡住、副流运行时主流 idle物理执行仍是串行。任何已落到副流的结论都必须用 profiling 的 Stream ID Start Time Duration 验证真实 overlap不能凭逻辑切流就认定方案生效。不要默认开启控核limit_core_num只解决资源分配问题不解决依赖错误只有 overlap 已成立但出现资源争抢或明显拖尾时才引入。保留开关与回退路径多流路径要保留 enable 开关和原始回退路径如上述案例中的enable_multi_streams开关方便调试与淘汰时干净隔离。先证明正确再追性能先验证依赖、功能和精度再看 overlap 与时延收益跨流汇合点不要遗漏等待避免读到未完成结果。执行路径不要混用eager 路径不要照搬 TorchAir 的图内 tag 风格npugraph_ex 路径的切流torch.npu.Stream()with torch.npu.stream(s)与同步torch.npu.Event()record()/wait()与 GE 图模式的npu_stream_switch不能混写进同一套实现。八、总结多流表达是 CANN 平台上利用 GE 图模式挖掘模型并行性的核心手段通过npu_stream_switch将无数据依赖的算子分支分发到不同 stream 形成 overlap通过npu_wait_tensor精确控制跨流时序再视资源争抢情况用limit_core_num进行算子级或全局级控核。其适用前提是存在真实可并行的分支且 Cube 资源未用满使用前务必核对静态/动态 shape 约束、单流与 SuperKernel 冲突约束。仓库中的 DeepSeek 系列、GLM 系列、LongCat-Flash 等模型已提供多套可复用的编排模板建议按先整网拆模块判并行性 → 每个并行点派生多种编排候选 → 切流同步 → profile 验证真实 overlap → 按需控核的顺序落地优化。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考