AIV UDMA Direct Drive

导言

本系列最后把视角移到 cann/shmem 的 UDMA Device 数据面：AIV 如何找到某个 PE/QP 的队列资源，为什么 WQE 常先在 UB 中组装，st_dev 到底做了什么，以及“指定 QP 就能并行”成立的条件。这里的代码用于展示机制，底层结构和接口必须与目标 CANN、ops 包、SoC 和 SHMEM revision 成套使用。

系列位置

1. URMA Mental Model：对象与数据路径
2. URMA Write and Read：最小读写程序
3. URMA Completion and Concurrency：完成、顺序与并发
4. SHMEM UDMA Programming：对称内存与 put 接口
5. AIV UDMA Direct Drive：st_dev、多 QP 与 Relay

源码截面

本文的 QP 接口、Relay 约束和 UDMA 提交路径核对到 cann/shmem@59f0313325ef3179306ce6fa3bfef8ba2ff4cbc5。底层 WQE 布局不属于跨版本稳定教程接口，实际工程必须与目标 CANN、ops 包和硬件配套。

Host 提交与 AIV 直驱

常规 URMA Host 路径是：

Host 应用构造 WR
    → 调用 urma_post_jetty_send_wr
    → UDMA Provider 编码并写入 WQE
    → CPU 写 Doorbell
    → UDMA 执行

AIV 直驱路径是：

Host 预先创建 Endpoint/Channel、注册内存并下发资源表
    ↓
AIV 根据 pe 和 qp_idx 查资源
    → 换算远端地址
    → 在 UB 中组装 WQE
    → 将 WQE 写入 HBM SQ
    → 更新 head
    → st_dev 写 Doorbell
    → UDMA 执行

二者的主要区别是谁位于每次通信操作的提交热路径。AIV 直驱并不意味着 Host 完全消失：资源创建、注册、建链和状态下发仍然需要 Host 控制面。

QP 是什么

QP，Queue Pair，是一条带独立队列状态的通信通道。 在本文关注的 UDMA 发送路径中，最重要的是它拥有或关联一组独立状态：

• SQ 与 WQE 槽位。
• SQ head/PI 和设备消费进度。
• CQ 与完成统计。
• Doorbell 地址。
• 目标端点、内存访问与通道信息。

PE 和 QP 解决不同问题：

参数	回答的问题	示例
pe	数据发给谁	发给 PE 3
qp_idx	通过到该 PE 的哪条队列通道提交	使用 PE 3 的 QP 1

一个 PE 可以对应多个 QP。它们都通向同一个逻辑参与者，但拥有各自的队列生产者状态。

AIV 如何提交一条 UDMA WQE

下面是按源码职责归一化后的伪代码，不是可独立编译的公共接口：

void post_udma_write(int pe, int qp_idx,
                     __gm__ uint8_t *dst,
                     __gm__ uint8_t *src,
                     uint32_t bytes,
                     __ubuf__ uint8_t *wqe_buf)
{
    UDMA_QP *qp = lookup_qp(pe, qp_idx);

    wait_until_sq_has_credit(qp);

    uint32_t slot = qp->head % qp->sq_depth;
    build_write_wqe(wqe_buf, dst, src, bytes, qp);

    copy_wqe_from_ub_to_sq(qp->sq_base, slot, wqe_buf);

    uint32_t new_head = qp->head + 1;
    make_wqe_visible_before_doorbell();
    st_dev(new_head, qp->doorbell_addr, 0);
    qp->head = new_head;
}

逐步解释：

1. lookup_qp(pe, qp_idx)：从初始化阶段下发的表中找到目标 PE 的指定 QP。
2. 检查 credit：确认 SQ 还有设备未占用的槽位。
3. 计算 slot：环形队列用 head 对深度取模选择物理槽位。
4. 构造 WQE：填 opcode、源地址、目标地址、长度、token 和完成控制等字段。
5. 写入 SQ：WQE 可先在 UB 临时区组装，再搬到 GM/HBM 中的 SQ。
6. 可见性屏障：确保设备先看到完整 WQE，再看到新的 Doorbell 值。
7. 敲 Doorbell：通知设备 SQ head 已推进。
8. 保存 head：为本 QP 下一次提交准备索引。

buf 参数之所以出现在显式 UDMA put 中，就是因为 AIV 需要一块 UB 空间暂存 WQE。它不是为了把全部 payload 先复制进 UB。

st_dev 到底是什么

st_dev 是 Ascend C/CCE Device 侧 store 指令，用于从 AIV 向特定 Device 地址写值。直驱 UDMA 时，常见用途是把新的 SQ head 写到 Doorbell：

make_wqe_visible_before_doorbell();
st_dev(new_head, doorbell_addr, 0);

三个参数可直观理解为：

• new_head：要写入的新队列生产者位置。
• doorbell_addr：设备可识别的 Doorbell 地址。
• 0：相对该地址的偏移。

最重要的边界有三条：

1. st_dev 不是 urma_post_jetty_send_wr() 的另一种写法。
2. st_dev 只发布 Doorbell，不会替调用者构造 WQE。
3. WQE 对设备可见必须发生在 Doorbell 更新之前，否则设备可能读到半成品。

不要手写未知布局

WQE、Doorbell、token 和队列状态是强版本相关的数据结构。除非正在开发或适配底层 Provider，不要脱离目标版本的 cann/shmem/HCOMM 实现自行猜字段布局。

指定 QP 为什么可以并行

假设一个 PE 配置两个 QP，并由两个 AIV 各自独占：

AIV block 0 → QP 0 → SQ 0 / head 0 / CQ 0 / Doorbell 0
AIV block 1 → QP 1 → SQ 1 / head 1 / CQ 1 / Doorbell 1

两个 AIV 不再同时读改写同一个 head，也不会争抢同一个 SQ slot，所以可以独立生产 WQE。这里的并行来自队列元数据和槽位被隔离，不是 qp_idx 这个整数本身具有魔法。

对比错误写法：

AIV block 0 ─┐
              ├→ 同一个 QP → 同一个 head 和 SQ slot → 数据竞争
AIV block 1 ─┘

多 QP Case 先在 Host 初始化前配置 QP 数。当前公共接口允许每个 peer 配置 1 到 ACLSHMEM_MAX_QP_NUM 个 QP，其中当前最大值为 32，而且所有 PE 必须传相同值：

// Host 侧，必须在 aclshmemx_init_attr() 之前调用。
int status = aclshmemx_set_qp_num(
    ACLSHMEM_DATA_OP_UDMA,
    qp_count);
if (status != ACLSHMEM_SUCCESS) {
    return status;
}

Device Case 可以按以下方式组织：

int block = GetBlockIdx();
int qp_idx = block;

if (block >= qp_count) {
    return;
}

uint64_t offset = block * slice_bytes;

aclshmemx_udma_qp_put_nbi<uint8_t>(
    remote_dst + offset,
    local_src + offset,
    wqe_ub_for_this_block,
    slice_bytes,
    peer_pe,
    qp_idx,
    sync_id_for_this_block);

aclshmemx_udma_qp_quiet(peer_pe, qp_idx);

Host 启动的 block 数应等于 qp_count，这样 GetBlockIdx() 可以直接作为 qp_idx，每个 QP 恰好只有一个 AIV 提交者。wqe_ub_for_this_block 也必须是当前 block 独占、容量不少于 ACLSHMEM_UDMA_MTE_STAGING_UB_SIZE 的 UB 临时区。

多 QP 成立的四个条件

1. 一核一 QP：同一个 QP 不再被多个无锁生产者共享。
2. 独立 UB 临时区：不同 AIV 不覆盖彼此正在组装的 WQE。
3. 独立数据切片：不同 QP 不无序写同一个远端 payload 区域。
4. 分别等待完成：每个 QP 的在途请求按其完成接口收敛，再进入跨核或跨 PE 同步。

如果只满足第一条而多个 QP 仍写同一远端地址，队列冲突消失了，业务数据竞争仍然存在。

多 QP 的完成与 signal

每个 QP 有独立进度，通常也应为不同生产者准备独立 signal：

QP 0 写 slice 0 → signal[0] = round
QP 1 写 slice 1 → signal[1] = round

接收端分别等待 signal[0] 和 signal[1]
两者都满足后才消费完整 buffer

共用一个 signal 会带来两个问题：

• 多个 writer 对同一值执行 SET 时，接收端无法知道哪些切片已到。
• 即使使用 ADD，也要处理原子性、轮次复用和计数溢出协议。

初学阶段优先使用“一 QP 一切片一 signal”，先确保正确，再考虑压缩同步状态。

Relay 不是 Replay

接口或注释中的 Relay 表示“中继路径”，不是“重放请求”的 Replay。

一条 Relay UDMA 请求涉及三个角色：

self PE ──经 relay PE 的端口/链路路径──→ target PE

• self PE：真正发起请求的当前 PE。
• relay PE：提供中继出口或指定 Fabric 路径的 PE。
• target PE：最终数据要写入或读取的 PE。

Relay 的主要用途是显式选择可达路径或利用多条链路。最终内存对象仍属于 target PE；它不是要求应用先把 payload 写进 relay PE 的对称缓冲区，再由 relay Kernel 做第二次 put。

使用时应满足：

• self、relay_pe 和目标 pe 的角色明确，符合接口限制。
• relay 所属端口确实能到达目标。
• 路径选择不会破坏上层要求的顺序和完成协议。
• 多路径写入仍使用不重叠切片或明确同步。

从入门到实际算子

建议按以下顺序增加复杂度：

1. 单 PE 内确认 Kernel、对称分配和 signal 地址正确。
2. 两 PE、单 AIV、单 QP 跑通一次 put-with-signal。
3. 改成 nbi，并加入 quiet，验证缓冲区生命周期。
4. 两 AIV、两 QP，各写一半不重叠数据，各用独立 signal。
5. 增加轮次序号，验证重复运行不会读到旧 signal。
6. 最后才引入 Relay、更多 PE、轮转调度和性能流水。

每一步至少检查：源数据、目标数据、signal、QP 完成状态和错误码。不要只看 Kernel 没有报错。

本篇结论

• AIV 直驱改变的是提交者，UDMA 仍是执行数据搬运的引擎。
• st_dev 负责 Doorbell store，不负责 payload 搬运或 WQE 构造。
• QP 是独立队列通道；PE 指目标参与者，QP 指到该参与者的提交车道。
• 指定不同 QP 能并行，是因为 SQ、head、CQ 和 Doorbell 状态被隔离。
• 多 QP 不自动提供跨 QP 顺序，也不消除远端地址冲突。
• Relay 是中继选路，不是 Replay，也不是应用层二次拷贝。

参考资料

• CANN SHMEM 源码仓
• UDMA Device 公共头文件
• UDMA Device 实现
• UDMA 多 QP 官方示例
• st_dev CCE Intrinsic API
• 进阶实现：AIV Direct Drive Programming
• 实际场景：Ascend DeepEP MoE Dispatch Optimization

评论