特性简介

在NPU芯片上存在多种硬件资源，不同硬件资源之间可以并发执行来最大化芯片的资源利用率，从而提升模型端到端性能。多流并发主要面向以下两种场景：

计算/通信并行：通信算子主要利用SDMA资源进行数据搬运，与计算算子使用的aicore资源独立，可以将计算/通信任务排布到不同流，实现overlap。
Cube/Vector并行： NPU的A3芯片是CV 分离架构，因此在没有数据依赖时，可以将cube/vector任务排布到不同流，实现overlap。

如大模型推理中，moe模块的共享专家和路由专家无数据依赖，可以用路有专家的gating计算和dispatch通信去掩盖共享专家的cube计算，缩短推理的总耗时。

特性说明

Torchair 提供了对外的 API 方便用户在脚本中，直接在网络模型中将串行任务拆解，基于多流完成并发的任务编排。

对外接口说明：

def npu_stream_switch(stream_tag: str, stream_priority: int = 0):
    """
    创建一个在指定stream上执行的context manager

    Args:
        stream_tag: string类型入参，用于指定当前context manager中所使用stream的tag
        stream_priority: 预留参数，int类型入参，默认值0，用于指定当前context manager中所使用stream的优先级
        
    Returns:
        返回基于用户入参创建的stream的context manager
    """

接口调用说明：

可以在多次调用 npu_stream_switch 时传入相同的 stream_tag，此时表示创建的多个 context manager 会在相同 stream_tag 表示的流上执行。
调用 npu_stream_switch 创建 context manager 时支持嵌套使用。
请在确认有资源可以并发使用时再开启多流并发。如果aicore资源已全部被使用，此时再开启多流并发，则无法带来计算耗时缩减，还会导致额外的调度耗时，影响端到端性能。

并行时序控制（可选）：

def npu_wait_tensor(self: torch.Tensor, dependency: torch.Tensor) 
    """
    控制图内两个节点的执行时序。一般用于多stream场景，两条并行stream上无数据依赖的两个task的执行时序。
    控制之后的目标时序是 dependency 执行完毕之后，所有用到self的节点再开始执行。

    Args:
        self: Tensor类型入参，指定后执行的算子的任意一个输入Tensor。
        dependency: Tensor类型入参，指定被等待的算子的任意一个输出Tensor。
    """

使用示例

使用前面的多流 API，构造一个示例模型，包含两组并发及时序控制。

import torch, torch_npu, torchair

# 定义模型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 torchair.scope.npu_stream_switch('1'): 
            # torch.mm算子(mm_result)等待torch.add算子(add_result)执行完再执行
            torchair.scope.npu_wait_tensor(in4, add_result)
            mm_result = torch.mm(in3, in4)
        mm1 = torch.mm(in3, in4)
        with torchair.scope.npu_stream_switch('2'):
            # torch.add算子(add2)等待torch.mm算子(mm_result)执行完再执行
            torchair.scope.npu_wait_tensor(in4, mm_result)
            add2 = torch.add(in3, in4)
        return add_result, mm_result, mm1, add2

model = Model()
npu_backend = torchair.get_npu_backend()
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")