AIV MoE All-to-All
先说结论
| 层次 | 固定 revision 中实际存在的内容 | 证据边界 |
|---|---|---|
| 概念 | MoE route 产生不等长 peer 流量,dispatch 与 combine 构成往返数据交换 | 通用机制,由 classic/MegaMoE 源码具象化 |
| 可复用机制 | 对称内存、rank -> peer 地址映射、Put/Get、ready/count signal、quiet、barrier |
cann/shmem@382afa08 有公开实现 |
| TileXR 框架接入 | CommArgs::sendCountMatrix、extraFlag::UDMA、udmaInfoPtr、Host InitUDMA |
是预留与初始化草稿,不是 MoE 实现 |
| SHMEM 算子样例 | classic Dispatch/Combine、双平面、MegaMoE;固定 fork 的融合样例 | 必须按仓库与 commit 单独归属 |
归属与 ABI
TileXR 固定 revision 是 46c58f3d0c8704f67e37bf0af322efcf2880f42c。它的 gitlink 固定 3rdparty/shmem 为 b79bda38953d39e88b191e7805659298f0829d73,.gitmodules 指向 LingquLab/shmem 的 feature/CrossPOD_fix_2_ops_transformer 分支,而不是 cann/shmem@382afa08。.gitmodules
1 | flowchart LR |
图中断点来自 TileXR src/comm/tilexr_comm.cpp::InitUDMA L123-L190、src/include/tilexr_udma.h L32-L161,以及两个 SHMEM revision 的 include/、src/ 全局负检索。
TileXR 做了什么
Host 侧先交换 SHMEM UID,指定 UDMA 引擎并初始化 SHMEM,再调用私有 aclshmemx_get_udma_info,把设备指针和 UDMA bit 写入 CommArgs:源码。
1 | ret = aclshmemx_get_udma_info(&udmaInfoPtr, &udmaInfoSize); |
CommArgs 还保留一个 rankSize * rankSize 的 sendCountMatrix,注释明确写着 for all2allv;它表达每个源 rank 发往每个目标 rank 的元素数。源码
但这些字段没有连到 MoE dispatch/combine kernel。受控 src/ 中没有 tilexr_udma.h 调用方;DeepEP dispatch/combine 只出现在设计目标中。设计稿
为什么不能直接链接 cann/shmem main
| TileXR 薄封装期待 | 固定 SHMEM 实际接口 | 结果 |
|---|---|---|
Host aclshmemx_get_udma_info |
两个固定 SHMEM 的 include/src 均无该符号 |
私有 ABI 来源未固定 |
shmem/include/device/udma.h |
固定 fork 与 main 均无此文件 | include 不匹配 |
shmem::udma_put_nbi/get_nbi/quiet/atomic |
公开符号是 aclshmemx_udma_* |
namespace/signature 不匹配 |
args.udma_enabled、peer_mem_ptrs、peer_flag_ptrs |
实际 CommArgs 是 extraFlag、peerMems、udmaInfoPtr |
头文件与自身结构不匹配 |
因此,tilexr_udma.h 只能描述为 AICore 内联薄封装草稿。它没有消费传入的 udmaInfoPtr,也没有形成当前 revision 可证明的 MoE 调用链。薄封装源码
All-to-AllV 如何编码
用 P 表示 rank 数,B 表示每 rank token 数,H 表示 hidden size,K 表示 Top-K,E_l 表示每 rank 的 local expert 数。对一条路由:
1 | global_expert_count = P * E_l |
这是一种 等价 All-to-AllV 编码:代码没有调用名为 alltoallv 的集合 API,而是把不等长通信拆为 (src_rank,dst_rank,local_expert) 段,用 count 描述长度,用对称窗口描述接收地址,再以 one-sided Put 和 signal 完成交换。
对象账本
| 对象 | 生产者 -> 消费者 | shape/容量 | 生命周期 |
|---|---|---|---|
x |
上游 -> dispatch | [B,H], T |
本 rank,一次 MoE 层 |
expert_ids |
router -> dispatch/combine | [B,K], INT32 |
本 rank,一次 MoE 层 |
| dispatch payload window | 源 rank AIV -> 目标 rank compact | [P*E_l*B*K, align32(H*sizeof(T))] |
对称内存,迭代复用 |
| assist window | dispatch -> combine | 每 route 为 [src_rank,token_id,topk_id] |
ready 前写入 |
| ready/count | 源 rank -> 目标 rank wait | 每 slot 8 个 INT32 | compact 后清零或置 -1 |
expand_x |
compact -> local expert | 容量 [P*B*K,H] |
本 rank expert 输入 |
ep_recv_count |
dispatch -> expert/combine | [E_l*P] 累计段结束位置 |
一次 dispatch/combine 往返 |
| expert result | local expert -> combine | 与 expand_x 有效行对应 |
计算完成后回传 |
| combine window | expert rank -> token owner | [B*K,align32(H*sizeof(T))] |
status 后可读 |
expert_scales |
router -> combine | [B,K], FP32 |
加权归约 |
x_out |
combine -> 下游 | [B,H], T |
本 rank 最终输出 |
对象 shape 和分配来自 RunDispatchCase L264-L306 与 RunCombineCase L355-L370。Dispatch Host Combine Host
真实代码路径
cann/shmem@382afa08 的 classic 样例最适合解释协议,因为数据面与控制面都显式可见。它是 AIV 设备侧直驱 MTE,不是 UDMA benchmark。
调用锚点:dispatch_classic L72-L281、combine_classic_once L48-L137。Dispatch Combine
Dispatch:pack、Put、通知、compact
Kernel 启动 P 个 active AIV core,core_id == dst_rank 的 core 负责一个目标 rank。它扫描本地 [B,K] 路由,先计数,再将 token 写入目标 rank 的固定 slot:
1 | const int64_t route_dst_rank = expert_id / local_expert_num; |
payload 完成后,发送端发布三类控制信息:
- assist:
[src_rank, token_id, topk_id]; - ready:该 payload slot 可读;
- count:该
(src_rank, local_expert)段的实际长度。
目标 rank 先等待所有 count,构造 ep_recv_count 前缀和;随后逐 slot 等 ready,把固定窗口 compact 为按 (local_expert,src_rank) 排列的 expand_x 与 assist。源码 L155-L278
Combine:反向 Put 与加权归约
expert rank 从 assist 恢复 token owner 和原始 Top-K 槽位,把结果写回 combine_slot = token_id*K+topk_id:
1 | const int32_t src_rank = assist_info_for_combine[i * 3]; |
token owner 等待每个 token 的全部 K 个 status,执行:
1 | x_out[token,j] = sum(data[token,topk,j] * expert_scales[token,topk]) |
之后清理 status,避免下一代复用旧标志。源码 L73-L136
组件时序
1 | sequenceDiagram |
张量数据流
1 | flowchart LR |
完整教学伪代码
下面的伪代码保留 classic 协议的 producer、consumer、shape、完成顺序与清理动作;它不是可编译实现。
1 | PROCEDURE MOE_ROUND(x[B,H], expert_ids[B,K], expert_scales[B,K], P, E_l): |
UDMA 原语应怎样理解
cann/shmem@382afa08 公开 UDMA Put/Get/Quiet/Atomic/Put+Notify,但其 classic MoE 没有使用这些 UDMA 接口。公开契约要求:
- Put/Get 是非阻塞提交,Get 结果读取或 Put 源复用前必须
aclshmemx_udma_quiet(pe); - 单请求最多 256 MB,超出需分块;
- 同一 PE 不支持并发 RMA/AMO;
aclshmemx_udma_put_signal_nbi将数据写与 signal 更新编码为一个WRITE_WITH_NOTIFYWQE,但返回仍只表示提交。
1 | flowchart LR |
两个融合变体
TileXR 固定 fork
LingquLab/shmem@b79bda38 的 examples/dispatch_gmm_combine 确实实现融合路径:AIV 执行 Dispatch/Combine,AIC 执行两次 GMM,中间以 cross-core flag 协作。AIC/AIV 分工
Dispatch 先路由和量化,再用 aclshmem_ptr 映射 peer 的 token-count、A 与 per-token scale;Combine 写回 peer D,做跨 rank 同步和 token unpermute。Dispatch/Combine
cann/shmem MegaMoE
cann/shmem@382afa08 的 Ascend950 MegaMoE 把 route mask、FP8 量化、跨 rank dispatch、两次 Catlass GEMM、SwiGLU、写回和 Top-K combine 放进一个 Device 流水。主流程
它先把 inbound mask 写入目标 rank,对本地 BF16 token 做 FP8+E8M0 量化;目标 rank 通过 aclshmem_ptr 远端地址 pull 量化 token,建立 routing metadata,再运行两次 expert GEMM。第二次投影后按 metadata 写回源 rank 的 combinedTokens,最后按 routing weights 合并。远端地址
这个版本同时改变 融合、量化、GEMM 实现、AIC/AIV 流水和通信组织。没有同硬件同 shape 的逐项消融时,任何端到端数字都不可单独归因于 AIV。
异常与边界
- 无效 expert id:classic dispatch 没有在 kernel 内显式拒绝
expert_id >= P*E_l。无效 route 可能被发送端丢弃,而 combine 仍等待其 status,生产接入必须前置校验。 - 无限等待:classic signal wait、固定 fork barrier、MegaMoE generation wait 均无故障超时;peer 失败或状态缺失会挂住。
- 容量截断:固定 fork 的
maxOutputSize和 MegaMoE 的maxReceivedTokens会限制接收行数;这是有损容量语义,不是无损 All-to-AllV。 - 状态复用:ready/count/status 必须在下一代前清理;generation wait 应比较
>= generation,否则快 rank 跨代会让慢 rank永远等不到相等值。 - 对称地址:
aclshmem_ptr、UDMA Put/Get/Notify 要求对称分配和一致布局;TileXR IPCpeerMems[]不能自动等价为 SHMEM symmetric memory。 - TileXR 清理所有权:
TileXRCommDestroy和析构函数都可能 finalize SHMEM,前者还直接 freeudmaInfoPtr,但注释称该指针由 SHMEM 管理;需要统一所有权和失败回滚。 - 平台不可混用:doubleplane 暂不支持 Ascend950,MegaMoE 只支持 Ascend950,固定 fork 代码选择 AtlasA2。不同平台结果不能互作 baseline。
性能证据与归因
| benchmark/主张 | 硬件与卡数 | shape/dtype | baseline | 数值 | 口径与结论 |
|---|---|---|---|---|---|
| TileXR UDMA/MoE | README 仅给 Ascend950 构建要求;卡数未公开 | 未公开 | 未公开 | 未公开 | 主仓无落地 MoE kernel,不能建立提升主张 |
| 固定 fork Dispatch-GMM-Combine | 代码选择 AtlasA2;示例 2 PE,节点未公开 | M64,K7168,N4096,E_l2,TopK8;FP16/INT8 |
未公开 | 未公开 | 一次功能运行、无计时;不可归因 |
| shmem classic | 支持 A2/A3/Ascend950 构建;脚本扫描 2/4/8 PE | dispatch 可 sweep;combine 默认 H7168,TopK8;INT32/FP16/BF16 边界不同 |
无 Host collective baseline | 未公开 | comm_only 与 full_op,默认 warmup 5/loops 50;仓库未提交实测 CSV |
| doubleplane 对 classic | A2/A3;不支持 Ascend950 | 相同 shape/dtype | MTE-only classic | 未公开 | 只可归因“大段 SDMA、小段/控制 MTE”的算法差异,不能归因 AIV |
| MegaMoE | Ascend950;示例单机 2 NPU | tokens256,model4096,ffn1024,TopK6,E_l4;BF16+FP8 |
未公开 | 未公开/不可归因 | kernel_event_avg_ms 与含 stream sync/Host barrier 的 e2e_avg_ms;warmup 5/loop 20 |
| UDMA perftest | 实际硬件/节点未提交 | message size/batch 可配 | 接口级 put/get/notify | 未公开,不能外推 MoE | 不含 route、pack、GMM、wait、unpack/combine |
| CANN 官方 MoE 直驱案例 | 公开文章未给出足以复现的完整硬件与 topology | MoeDistributeDispatch / MoeDistributeCombine,完整 shape 未公开 | 前一阶段 AIV+AICPU 分层路径 | 平均 10%+ | 只属于该案例对照;不能归给 TileXR、SHMEM 仓库或外推为通用 AIV 收益 |
MegaMoE 的计时定义可见 README L364-L375;classic 的 full_op/comm_only 定义可见 Dispatch README L98-L124 与 Combine README L125-L168。MegaMoE 口径
上表唯一的公开加速比来自昇腾社区 MoE 计算通信融合算子文章。由于文章没有给出复现实验所需的全部版本、shape、topology 与原始测量,本文只保留“该案例相对上一阶段路径平均 10%+”这一受限事实。
总结
MoE All-to-All 的核心不是一个 API 名字,而是 路由元数据如何编码 peer 与变长段、payload 与 signal 如何保持顺序、接收端如何 compact、expert 结果如何按原 token/Top-K 槽位写回。cann/shmem classic 给出了最清楚的教学实现,固定 fork 与 MegaMoE 展示了进一步的融合和量化流水。
TileXR 当前 revision 则停在另一条边界:Host UDMA 初始化、CommArgs 能力字段和薄封装已经出现,但 ABI 不闭合、没有 MoE 调用方。把这一状态准确写成“能力预留/设计意图”,比把子模块样例或未来设计误报为 TileXR 已实现更重要。