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

• 人工智能
• 模型编译
• Ascend

npu_stream_switch是 TorchAir 提供的一个 Scope 类上下文接口，用于在 Ascend IR（GE 图）场景下将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.py）：

class _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的使用分为三步：

1. 分析可并行算子：用户自行分析模型中可进行并行计算的算子，找出无数据依赖的计算分支。

2. 开启多流表达：使用如下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):

3. （可选）控制并行计算的时序：通过 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在默认stream，mm_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_config=config) # 调用compile接口编译模型 model = torch.compile(model, backend=npu_backend, dynamic=False, fullgraph=True) 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(f"Result:\n{result}\n")

示例中的执行时序为：add_result、mm1在默认 stream；mm_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.py），npu_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_STREAM=1

测试用例验证

仓库中的单测用例 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 = b'1'、_user_stream_priority = b'3'属性，并断言 Identity 算子输入中出现了 Add 算子的控制边（即跨流依赖生效）。
• test_npu_stream_switch_parallel：在多个嵌套with块中分别使用npu_stream_switch('1', 2)与npu_stream_switch('2', 3, False)，断言流"1"内的 MatMul 带_user_stream_label = b'1'、_user_stream_priority = b'2'且_enable_inner_parallel = True；流"2"内的 Add 带_user_stream_label = b'2'、_user_stream_priority = b'3'且_enable_inner_parallel = False；而默认 stream 上的算子不包含上述任何属性。

这两组用例印证了文档中的参数语义：stream_tag决定_user_stream_label，stream_priority决定_user_stream_priority，enable_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