TorchAir npu_stream_switch 接口详解：Ascend IR 图内多流并行调度实战指南

人工智能模型编译Ascend【免费下载链接】torchairTorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。项目地址https://gitcode.com/Ascend/torchair点击查看免费下载npu_stream_switch是 TorchAir 提供的一个 Scope 类上下文接口用于在 Ascend IRGE 图场景下将with语句块内的算子显式分发到指定 stream 上执行从而把原本串行的多条计算分支拆到不同流上并行实现多流表达功能。本文基于仓库文档与源码完整讲解该接口的函数原型、参数语义、实现原理、配套时序控制npu_wait_tensor以及约束条件并给出可直接运行的调用示例帮助读者在昇腾 NPU 上通过图模式提升大模型推理场景的资源利用率。功能说明为什么需要图内多流并行大模型推理场景中模型内往往存在多条没有数据依赖的计算分支默认情况下这些算子会被串行调度执行造成 Cube矩阵计算等计算资源空闲。多流表达功能通过为算子指定执行 stream将原本串行的多个算子分发到不同 stream 进行并行计算多个 stream 上的计算形成 overlap从而降低整体计算耗时。并行主要分为两种形态计算与计算并行基于数据依赖关系分析出可并行的多条计算分支分别指定 stream 并行执行。计算与通信并行针对没有数据依赖的通信操作提前占用通信资源执行通信任务与计算 overlap。npu_stream_switch正是多流表达功能的核心入口用于处理Ascend IR 图内资源并发max-autotune 模式尤其适合 Cube 计算资源未完全使用的场景。若 Cube 计算资源已完全使用不建议开启本功能否则可能引入额外调度导致原计算性能劣化。函数原型与参数说明接口定义位于 python/torchair/scope/init.py函数原型如下npu_stream_switch(stream_tag: str, stream_priority: int 0, enable_inner_parallel: bool True)| 参数 | 输入/输出 | 说明 | |--|--|--| | stream_tag | 输入 | 字符串类型指定算子执行的目标 stream 标签。相同的标签代表相同的流标签到流的映射由用户控制。 | | stream_priority | 输入 | int 类型表示切换到 stream_tag 流的优先级即 Runtime 在并发时优先给高优先级的流分配核资源。当前版本为预留参数建议取默认值 0。 | | enable_inner_parallel | 输入 | bool 类型表示是否启用 stream_tag 流内的算子按照 GE 原有的并发策略分流默认开启。 |返回值无该接口返回的是一个_Scope上下文管理器对象配合with语句使用。关于enable_inner_parallel的底层行为文档与源码一致地指出GE 图中的算子将自动添加_enable_inner_parallel属性属性值与enable_inner_parallel参数值一致。_enable_inner_parallel属性的详细介绍和约束可参见 CANN GE 图引擎 API 中的属性名列表章节且配套 CANN 版本大于 9.0.0 时才生效。接口实现原理源码级解析1. Scope 上下文对象与属性封装调用npu_stream_switch时TorchAir 并不会立即切换流而是构造一个_Scope对象将三个参数封装成键值对列表见 python/torchair/scope/init.pyclass _Scope: def __init__(self, attrs): self.attrs attrs def __enter__(self): return _npu_scope_enter(self.attrs) def __exit__(self, *args): return _npu_scope_exit() def npu_stream_switch(stream_tag: str, stream_priority: int 0, enable_inner_parallel: bool True): return _Scope( [ (_user_stream_label, stream_tag), (_user_stream_priority, str(stream_priority)), (_enable_inner_parallel, str(enable_inner_parallel)), ] )可以看出三个参数最终映射为三个图节点属性键_user_stream_label对应stream_tag标识算子所属流。_user_stream_priority对应stream_priority以字符串形式透传。_enable_inner_parallel对应enable_inner_parallel控制流内是否按 GE 原有并发策略再分流。2. scope_enter / scope_exit 哨兵节点与属性栈__enter__/__exit__最终调用 python/torchair/scope/_scope.py 中注册的torch.ops.air.scope_enter/torch.ops.air.scope_exit自定义算子并在 FX 图中插入对应的哨兵节点带has_side_effect标记。属性以栈结构维护在 python/torchair/scope/_scope_attr.py 的ScopeAttrs中__enter__时push当前属性字典__exit__时pop。在 GE 图转换阶段apply_scope_attr会把栈内所有属性逐条写到新生成的算子节点上其中if key _op_exec_never_timeout or key _enable_inner_parallel: op.attr[key].b value True else: op.attr[key].s compat_as_bytes(str(value))即_enable_inner_parallel以bool 型属性.b写入_user_stream_label、_user_stream_priority以字符串型属性.s写入。这正是文档中GE 图中的算子将自动添加_enable_inner_parallel属性的实现来源。3. 实际 stream 切换在 Eager/AclGraph 等执行路径下guard_with_user_stream_scope见 python/torchair/scope/_scope_attr.py会在算子执行时检查属性栈栈顶是否含有_user_stream_label若存在则通过_npu_get_or_create_stream_with_tag按标签缓存并获取对应的torch_npu.npu.Stream再以torch.npu.stream(target_stream)上下文包裹算子执行实现真正意义上的流切换。相同stream_tag会命中_GLOBAL_TAG_TO_STREAM缓存从而保证相同的标签代表相同的流。4. 跨流控制边的构建在 GE 图转换器 python/torchair/_ge_concrete_graph/fx2ge_converter.py 中_get_current_stream从属性栈中读取当前_user_stream_label_build_wait_control_edge、_build_record_control_edge依据当前流上最后一个算子构造显式跨流依赖控制边配合_handle_wait_control_edge把控制边挂接到后续算子输入上从而在 GE 图中形成图 1 所示的流间控制关系。使用方法与完整调用示例npu_stream_switch的使用分为三步分析可并行算子用户自行分析模型中可进行并行计算的算子找出无数据依赖的计算分支。开启多流表达使用如下with语句块语句块内下发的算子切换至stream_tag流同时自动添加_enable_inner_parallel属性属性值与enable_inner_parallel参数值一致语句块外的算子使用默认 stream 计算。with torchair.scope.npu_stream_switch(stream_tag: str, stream_priority: int 0, enable_inner_parallel: bool True):可选控制并行计算的时序通过 npu_wait_tensor 接口实现时序控制指定算子 a 等待算子 b 执行完后再执行。完整示例代码以下示例改编自多流表达功能使用示例import 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在流1add2在流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)示例中的执行时序为add_result、mm1在默认 streammm_result在流1等待add_result完成后执行add2在流2等待mm_result完成后执行三条流之间通过显式控制边形成 overlap。如上图所示来自 pbtxt 文件样例图中展示了流间的控制关系虚线控制边表示显式的跨流数据依赖同步关系用于保证并行流之间的计算顺序符合用户预期。时序控制npu_wait_tensor 配套接口多流并行后若下游算子与上游算子存在数据依赖需要通过 npu_wait_tensor 显式建立跨流同步。其函数原型为npu_wait_tensor(self: torch.Tensor, dependency: torch.Tensor)| 参数 | 输入/输出 | 说明 | |--|--|--| | self | 输入 | Tensor 类型表示等待算子的入参即算子 a后执行的任意一个输入 Tensor。 | | dependency | 输入 | Tensor 类型表示被等待算子的出参即算子 b先执行的任意一个输出 Tensor。 |返回值返回self本身。从源码实现看python/torchair/scope/_wait_tensor.pynpu_wait_tensor在 GE 图转换时被转换为ge.Identity算子并把被依赖算子的ControlTensor作为控制输入挂接到 Identity 上register_fx_node_ge_converter(torch.ops.air.wait_tensor.default) def convert_wait_tensor(self: Tensor, dependency: Tensor, meta_outputs: TensorSpec None): control ControlTensor(dependency.node).controller identity ge.Identity(self) identity.node.input.append(control) return identity这样就在 GE 图中形成了等待算子 ← 控制边 ← 被等待算子的显式依赖关系配合npu_stream_switch完成图内多流计算配置。该接口一般与npu_stream_switch配套使用。约束与注意事项npu_stream_switch 自身约束with 语句块内不支持断图即不能在with npu_stream_switch(...)语句块内部进行图的截断/分段操作。enable_inner_parallel参数仅支持 max-autotune 模式若未开启 max-autotune该参数不生效。多流表达功能的使用约束本功能仅适用于 GE 图模式场景Ascend IR 图内资源并发。对于纯 Vector 场景其计算耗时一般在可接受范围内开启收益有限对于含 Cube 计算的场景开启本功能后的效益往往优于纯 Vector 计算场景。静态 Shape 场景下本功能与单流执行功能enable_single_stream冲突不支持同时开启。本功能不推荐在 SuperKernel 内设置算子多 stream 并行如有需要请使用图内标定 SuperKernel 范围中的 stream-fusion 编译选项配置。动态 Shape 场景下默认单流模式用户可通过如下 CANN 环境变量开启多流一旦开启多流其功能优先级低于本功能export ENABLE_DYNAMIC_SHAPE_MULTI_STREAM1测试用例验证仓库中的单测用例 tests/st/torchair_scope_st.py 直接验证了npu_stream_switch的图级行为可作为接口语义的可执行参考test_npu_stream_switch在with torchair.scope.npu_stream_switch(1, 3)内调用npu_wait_tensor与torch.mm断言生成的 MatMul 算子带有_user_stream_label b1、_user_stream_priority b3属性并断言 Identity 算子输入中出现了 Add 算子的控制边即跨流依赖生效。test_npu_stream_switch_parallel在多个嵌套with块中分别使用npu_stream_switch(1, 2)与npu_stream_switch(2, 3, False)断言流1内的 MatMul 带_user_stream_label b1、_user_stream_priority b2且_enable_inner_parallel True流2内的 Add 带_user_stream_label b2、_user_stream_priority b3且_enable_inner_parallel False而默认 stream 上的算子不包含上述任何属性。这两组用例印证了文档中的参数语义stream_tag决定_user_stream_labelstream_priority决定_user_stream_priorityenable_inner_parallel决定_enable_inner_parallel且该属性仅作用于with块内的算子。总结npu_stream_switch是 TorchAir 多流表达功能的核心 Scope 接口通过with上下文把算子按stream_tag分组分发到不同 stream实现计算与计算、计算与通信的并行 overlap。使用时应结合npu_wait_tensor显式控制跨流时序并遵守GE 图模式、max-autotune、静态 Shape 不与单流/SuperKernel 冲突等约束。对于含 Cube 计算且资源未饱和的推理模型该接口是提升昇腾 NPU 执行效率的有效手段。赞分享人工智能模型编译Ascend【免费下载链接】torchairTorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。项目地址https://gitcode.com/Ascend/torchair点击查看免费下载相关推荐CANN Ascend Transformer Boost 多流Multi-StreamDemo 实战图内多流并行与图间同步CANN Ascend Transformer Boost 多流Multi StreamDemo 实战图内多流并行与图间同步 导读 本文围绕 exampl人工智能算子库大模型深度学习CANNAscendascend-transformer-boost 多流Multi-Stream编程实战图内多流并行与图间同步 Demo 全解析ascend transformer boost 多流Multi Stream编程实战图内多流并行与图间同步 Demo 全解析 导读 多流Multi S人工智能算子库大模型深度学习CANNAscend如何在 KernelSU 上运行 LSPosedZygiskNext 三步配置与避坑指南如何在 KernelSU 上运行 LSPosedZygiskNext 三步配置与避坑指南 KernelSU 没有内置 ZygiskLSPosed 这类 Xp操作系统驱动开发上一篇微信红包助手2024自动抢红包插件全攻略无需ROOT也能快人一步下一篇Rust By Practice 类型转换实战as、From/Into、TryFrom/TryInto 与 transmute 全解析创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考