# TPUT\_ASYNC ## 简介 `TPUT_ASYNC` 是异步远程写原语。它启动一次从本地GM到远端GM的传输，并立即返回 `AsyncEvent`。数据流： `srcGlobalData（本地 GM）` → DMA引擎 → `dstGlobalData（远端 GM）` ## 模板参数 - `engine`： - `DmaEngine::SDMA`（默认） - `DmaEngine::URMA`（Ascend 950PR/Ascend 950DT，仅NPU\_ARCH 3510） - `DmaEngine::RDMA`（Ascend 950PR/Ascend 950DT，仅NPU\_ARCH 3510；当前网卡平台仅支持 HNS1825） > **注意（SDMA路径）** > `TPUT_ASYNC` 配合 `DmaEngine::SDMA` 目前**仅支持扁平连续的逻辑一维tensor**。 > 当前SDMA异步实现不支持非一维或非连续布局。 ## C++内建接口声明于 `include/pto/comm/pto_comm_inst.hpp`： ```cpp template PTO_INST AsyncEvent TPUT_ASYNC(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData, const AsyncSession &session, WaitEvents &... events); ``` `AsyncSession` 是引擎无关的会话对象。使用 `BuildAsyncSession()` 构建一次后，传递给所有异步调用和事件等待。模板参数 `engine` 在编译期选择DMA后端，使代码对未来引擎（CCU等）保持前向兼容。 ## AsyncSession构建使用 `include/pto/comm/async_common/async_event_impl.hpp` 中的 `BuildAsyncSession`。该函数有两个重载——分别用于SDMA和URMA，参数列表不同。 ### SDMA构建（默认） ```cpp template PTO_INTERNAL bool BuildAsyncSession(ScratchTile &scratchTile, __gm__ uint8_t *workspace, AsyncSession &session, uint32_t syncId = 0, const sdma::SdmaBaseConfig &baseConfig = {sdma::kDefaultSdmaBlockBytes, 0, 1}, uint32_t channelGroupIdx = sdma::kAutoChannelGroupIdx); ``` | 参数 | 默认值 | 说明 | |---|---|---| | `scratchTile` | — | 用于SDMA控制元数据的UB scratch tile（参见 [scratchTile的作用](#scratchtile的作用)）。| | `workspace` | — | 由主机侧 `SdmaWorkspaceManager` 分配的GM指针。| | `session` | — | 输出的 `AsyncSession` 对象。| | `syncId` | `0` | MTE3/MTE2管道同步事件ID（0-7）。若kernel在相同ID上使用了其他管道屏障，则需覆盖此值。| | `baseConfig` | `{kDefaultSdmaBlockBytes, 0, 1}` | `{block_bytes, comm_block_offset, queue_num}`。适用于大多数单队列传输场景。| | `channelGroupIdx` | `kAutoChannelGroupIdx` | SDMA通道组索引。默认内部使用 `get_block_idx()` 映射到当前AI Core。多block、并发或自定义通道映射场景下需覆盖此值。| ### URMA构建（仅NPU\_ARCH 3510） > URMA（User-level RDMA Memory Access）是Ascend 950PR/Ascend 950DT（NPU\_ARCH 3510）上的硬件加速RDMA传输引擎。 > URMA要求CANN Toolkit **\>= 9.1.0**。 ```cpp #ifdef PTO_URMA_SUPPORTED template PTO_INTERNAL bool BuildAsyncSession(__gm__ uint8_t *workspace, uint32_t destRankId, AsyncSession &session); #endif ``` | 参数 | 说明 | |---|---| | `workspace` | 由主机侧 `UrmaWorkspaceManager` 分配的GM指针。| | `destRankId` | 此会话通信的远端PE rank id。对于 `TPUT_ASYNC`，这是数据写入的目标rank。| | `session` | 输出的 `AsyncSession` 对象。| URMA不需要 `scratchTile`——轮询通过 `ld_dev`/`st_dev` 硬件原语直接操作。 ### RDMA构建（仅NPU\_ARCH 3510） RDMA面向经典跨机场景组网，当前网卡平台仅支持 HNS1825。 ```cpp #ifdef PTO_RDMA_SUPPORTED template PTO_INTERNAL bool BuildAsyncSession(ScratchTile &scratchTile, __gm__ uint8_t *workspace, uint32_t myPe, AsyncSession &session, uint32_t syncId = 0); #endif ``` | 参数 | 说明 | |---|---| | `scratchTile` | 用于 WQE/CQE 控制数据的 UB/Vec tile，至少 64 字节。| | `workspace` | Host侧 RDMA 初始化流程返回的 GM 指针。| | `myPe` | 本地 rank id，用于选择已注册的本地内存区域。| | `session` | 输出的 `AsyncSession` 对象。| | `syncId` | MTE/Scalar 同步事件 ID，取值范围为 0-7。| 构建一次不绑定 peer 的Session，再通过显式 `peer` 参数选择目标rank： ```cpp comm::AsyncSession session; if (comm::BuildAsyncSession(scratchTile, rdmaWorkspace, myPe, session, syncId)) { auto event = comm::TPUT_ASYNC(dstG, srcG, session, peer); (void)event.Wait(session); } ``` ## 约束 - `GlobalSrcData::RawDType == GlobalDstData::RawDType` - `GlobalSrcData::layout == GlobalDstData::layout` - SDMA和URMA路径均要求源tensor为**扁平连续的逻辑一维** - SDMA workspace必须是由主机侧 `SdmaWorkspaceManager` 分配的有效GM指针 - URMA workspace必须是由主机侧 `UrmaWorkspaceManager` 分配的有效GM指针 - Session和workspace的生命周期必须覆盖相关Event的完成阶段 - URMA仅在NPU\_ARCH 3510（Ascend 950PR/Ascend 950DT）上可用 - URMA要求CANN Toolkit **\>= 9.1.0** - 传给 `UrmaWorkspaceManager::Init()` 的对称数据缓冲区必须由大页内存支撑（使用 `ACL_MEM_MALLOC_HUGE_ONLY` 分配）。底层MR注册要求大页背景；`ACL_MEM_MALLOC_HUGE_FIRST` 在小尺寸分配时可能静默回退到4KB小页，导致注册失败 - RDMA仅在Ascend 950PR/Ascend 950DT（NPU\_ARCH 3510）上可用，当前网卡平台仅支持 HNS1825 - RDMA源、目的tensor必须是扁平连续的逻辑一维，且本地和远端完整传输范围必须位于Host初始化阶段注册的内存区域内 - 单次RDMA传输不能超过 `0x7fffffff` 字节若不满足一维连续要求，当前实现返回无效async event（`handle == 0`）。 ## scratchTile的作用 `scratchTile` **不是**用于存放用户数据负载的暂存缓冲区。它被转换为 `TmpBuffer`，用作临时UB工作区，用于： - 写入/读取SDMA控制字（flag、sq\_tail、channel\_info） - 轮询事件完成标志 - 完成时提交队列尾部实际数据负载直接在GM缓冲区之间传输；`scratchTile` 仅用于控制和同步元数据。 ## scratchTile类型与大小约束 - 必须是 `pto::Tile` 类型 - 必须是UB/Vec tile（`ScratchTile::Loc == TileType::Vec`） - 可用字节数至少为 `sizeof(uint64_t)`（8字节）推荐使用：`Tile`（256Byte）。 ## 完成语义（Quiet语义）不同引擎的底层完成机制不同，但用户侧的quiet语义行为一致： - **SDMA**：每次`TPUT_ASYNC`都会提交数据传输SQE和用于标记本次操作完成的flag SQE。对其返回的Event调用`Wait`或`Test`时，通过轮询对应flag判断该次`TPUT_ASYNC`是否完成；完成后也能保证同一Session中此前提交的所有SDMA操作均已完成。 - **URMA**：`TPUT_ASYNC` 立即提交RDMA WRITE WQE并敲门铃。`Wait` 通过轮询Completion Queue（CQ）等待所有预期的CQE被消费。 - **RDMA**：`TPUT_ASYNC` 向 `peer` 选择的队列提交RDMA WRITE WQE；不同peer/queue分别跟踪完成状态。 - `event.Wait(session)` —阻塞，直到**自上次Wait以来所有已发出的异步操作**全部完成这意味着多次 `TPUT_ASYNC` 调用后，只需对最后一个返回的 `AsyncEvent` 调用一次 `Wait`，即可等待所有pending操作完成（类似shmem的quiet语义）。 RDMA操作涉及不同peer时，必须分别等待每个peer的最后一个Event。同一Session最多可有64个未完成操作，超过后提交可能产生背压。 wait成功后，所有已发出的 `dstGlobalData` 写入均已全部完成。 ## SDMA并发与Session所有权 - 同一个Session不能被多个执行流并发使用。 - 共用同一Channel Group的操作必须共享同一个Session。 - 并发Kernel或Kernel内多个独立Session必须使用隔离的Channel Group。 - 重新构建Session或复用Channel Group前，必须先完成此前所有Event。 ## 示例 ### 单次传输 ```cpp #include #include using namespace pto; template __global__ AICORE void SimplePut(__gm__ T *remoteDst, __gm__ T *localSrc, __gm__ uint8_t *sdmaWorkspace) { using ShapeDyn = Shape; using StrideDyn = Stride; using GT = GlobalTensor; using ScratchTile = Tile; ShapeDyn shape(1, 1, 1, 1, 1024); StrideDyn stride(1024, 1024, 1024, 1024, 1); GT dstG(remoteDst, shape, stride); GT srcG(localSrc, shape, stride); ScratchTile scratchTile; TASSIGN(scratchTile, 0x0); comm::AsyncSession session; if (!comm::BuildAsyncSession(scratchTile, sdmaWorkspace, session)) { return; } auto event = comm::TPUT_ASYNC(dstG, srcG, session); (void)event.Wait(session); } ``` ### 批量传输（Quiet语义） ```cpp template __global__ AICORE void BatchPut(__gm__ T *remoteDstBase, __gm__ T *localSrc, __gm__ uint8_t *sdmaWorkspace, int nranks) { using ShapeDyn = Shape; using StrideDyn = Stride; using GT = GlobalTensor; using ScratchTile = Tile; ShapeDyn shape(1, 1, 1, 1, 1024); StrideDyn stride(1024, 1024, 1024, 1024, 1); GT srcG(localSrc, shape, stride); ScratchTile scratchTile; TASSIGN(scratchTile, 0x0); comm::AsyncSession session; if (!comm::BuildAsyncSession(scratchTile, sdmaWorkspace, session)) { return; } comm::AsyncEvent lastEvent; for (int rank = 0; rank < nranks; ++rank) { GT dstG(remoteDstBase + rank * 1024, shape, stride); lastEvent = comm::TPUT_ASYNC(dstG, srcG, session); } (void)lastEvent.Wait(session); // 一次 Wait 等待所有 pending 操作 } ``` ### URMA示例（NPU\_ARCH 3510） ```cpp #include #include using namespace pto; template __global__ AICORE void SimplePutUrma(__gm__ T *remoteDst, __gm__ T *localSrc, __gm__ uint8_t *urmaWorkspace, uint32_t destRankId) { using ShapeDyn = Shape; using StrideDyn = Stride; using GT = GlobalTensor; ShapeDyn shape(1, 1, 1, 1, 1024); StrideDyn stride(1024, 1024, 1024, 1024, 1); GT dstG(remoteDst, shape, stride); GT srcG(localSrc, shape, stride); comm::AsyncSession session; if (!comm::BuildAsyncSession(urmaWorkspace, destRankId, session)) { return; } auto event = comm::TPUT_ASYNC(dstG, srcG, session); (void)event.Wait(session); } ```