AIV MoE All-to-All

导言

MoE 的 All-to-All 不是一次整齐的等长集合通信，而是由路由结果决定长度的 pack、跨 rank 搬运、完成通知、expert 计算、反向搬运与 combine。固定源码表明：TileXR 主仓已经预留 sendCountMatrix、UDMA flag 与设备指针，却没有落地可追踪的 MoE dispatch/combine kernel；真正实现分别位于它固定的 SHMEM fork 和 cann/shmem。本文只描述源码能证明的部分，所有缺少完整实验条件的性能数字均标为未公开或不可归因。

先说结论

三条边界

1. 能力预留不等于算子实现：TileXR 的类型、注册项、sendCountMatrix 和 DeepEP 设计文档不能证明主仓已有 MoE kernel。
2. 子模块能力不等于主仓能力：融合 Dispatch-GMM-Combine 样例属于 LingquLab/shmem@b79bda38；classic/doubleplane/MegaMoE 属于 cann/shmem@382afa08。
3. 端到端提升不等于 AIV 收益：直驱、融合、量化、传输引擎与拓扑会同时改变，缺少消融实验时不能把总收益归给 AIV。

层次	固定 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

flowchart LR
    T["TileXR@46c58f3d<br/>CommArgs + InitUDMA<br/>无 MoE kernel 调用方"]
    F["LingquLab/shmem@b79bda38<br/>固定子模块<br/>融合 Dispatch-GMM-Combine"]
    C["cann/shmem@382afa08<br/>classic/doubleplane/MegaMoE<br/>公开 UDMA 原语"]
    X["私有 Host ABI<br/>aclshmemx_get_udma_info"]
    D["公开 Device ABI<br/>aclshmemx_udma_*"]
    T -->|gitlink| F
    T -.期望但固定源码未发现导出.-> X
    F -->|实际 MoE: aclshmem_ptr + GM 搬运| D
    C --> D
    T -.不能直接宣称兼容.-> C

图中断点来自 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：源码。

ret = aclshmemx_get_udma_info(&udmaInfoPtr, &udmaInfoSize);
if (ret != ACLSHMEM_SUCCESS || udmaInfoPtr == nullptr) {
    aclshmem_finalize();
    return TILEXR_SUCCESS;
}
commArgs_.udmaInfoPtr = reinterpret_cast<uint8_t*>(udmaInfoPtr);
commArgs_.extraFlag |= ExtraFlag::UDMA;

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 数。对一条路由：

global_expert_count = P * E_l
dst_rank             = expert_id / E_l
dst_local_expert     = expert_id % E_l
segment              = dst_local_expert * P + src_rank
data_block           = src_rank * E_l + dst_local_expert
dispatch_slot        = data_block * (B*K) + slot_in_segment
combine_slot         = token_id * K + topk_id

这是一种 等价 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。

自绘技术图：上半部按 route、量化、dispatch、两次 expert GEMM、写回和 combine 展开设备侧流水；下半部给出跨 PE 的 payload、完成信号和 expert 输出时序。该图以 cann/shmem MegaMoE 为实现锚点，不把融合结构本身换算成百分比收益。

调用锚点：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：

const int64_t route_dst_rank = expert_id / local_expert_num;
const int64_t dst_local_expert = expert_id % local_expert_num;
const int64_t data_block = my_rank * local_expert_num + dst_local_expert;
const int64_t global_slot = data_block * max_tokens_per_segment + slot;
aclshmemx_mte_put_nbi(payload_base + global_slot * payload_stride,
    x + token_id * h, tmp_buff, UB_DMA_MAX_SIZE, h, dst_rank, EVENT_ID0);
aclshmem_quiet();

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：

const int32_t src_rank = assist_info_for_combine[i * 3];
const int32_t token_id = assist_info_for_combine[i * 3 + 1];
const int32_t topk_id = assist_info_for_combine[i * 3 + 2];
const int64_t slot = token_id * k + topk_id;
aclshmemx_mte_put_nbi(data_base + slot * data_stride,
    expand_x + i * h, tmp_buff, UB_DMA_MAX_SIZE, h, src_rank, EVENT_ID0);

token owner 等待每个 token 的全部 K 个 status，执行：

x_out[token,j] = sum(data[token,topk,j] * expert_scales[token,topk])

之后清理 status，避免下一代复用旧标志。源码 L73-L136

组件时序

sequenceDiagram
    participant S as 源 Rank AIV
    participant D as Expert Rank AIV
    participant X as Expert AIC/AIV
    S->>D: Put payload(slot)
    S->>D: assist(src,token,topk)
    S->>D: ready(slot) + count(segment)
    Note over S,D: payload quiet 先于 ready/count 可见
    D->>D: wait count/ready
    D->>D: compact expand_x + assist
    D->>X: 按 local expert 提交有效行
    X-->>D: expert result
    D->>S: Put result(combine_slot)
    D->>S: status ready
    S->>S: wait all Top-K status
    S->>S: weighted sum -> x_out

张量数据流

flowchart LR
    X["x [B,H]"] --> R["route by expert_ids [B,K]"]
    I["expert_ids [B,K]"] --> R
    R --> P["peer window<br/>payload+assist+ready+count"]
    P --> C["compact"]
    C --> E["expand_x [<=P*B*K,H]"]
    C --> A["assist [<=P*B*K,3]"]
    E --> G["local expert"]
    A --> G
    G --> W["owner combine window [B*K,H]"]
    S["expert_scales [B,K]"] --> U["Top-K weighted sum"]
    W --> U
    U --> O["x_out [B,H]"]

完整教学伪代码

下面的伪代码保留 classic 协议的 producer、consumer、shape、完成顺序与清理动作；它不是可编译实现。

PROCEDURE MOE_ROUND(x[B,H], expert_ids[B,K], expert_scales[B,K], P, E_l):
    REQUIRE P > 0 AND B > 0 AND H > 0 AND K > 0 AND E_l > 0
    FOR token FROM 0 TO B-1:
        FOR topk FROM 0 TO K-1:
            REQUIRE 0 <= expert_ids[token,topk] < P*E_l

    ZERO dispatch_payload_window
    ZERO dispatch_assist_window
    ZERO dispatch_ready_window
    FILL dispatch_count_window WITH -1
    BARRIER_ALL_RANKS()

    FOR dst_rank FROM 0 TO P-1 IN PARALLEL_ON_AIV:
        segment_count[0:E_l] = 0
        segment_offset[0:E_l] = 0

        FOR flat FROM 0 TO B*K-1:
            expert = expert_ids[flat]
            IF expert / E_l == dst_rank:
                local_expert = expert MOD E_l
                segment_count[local_expert] += 1

        FOR flat FROM 0 TO B*K-1:
            expert = expert_ids[flat]
            IF expert / E_l == dst_rank:
                token = flat / K
                topk = flat MOD K
                local_expert = expert MOD E_l
                slot = segment_offset[local_expert]
                segment_offset[local_expert] += 1
                global_slot = (MY_RANK*E_l + local_expert)*(B*K) + slot
                PUT_NBI(dst_rank, payload[global_slot], x[token,0:H])
                QUIET(dst_rank)
                SIGNAL_SET(dst_rank, assist[global_slot,0], MY_RANK)
                SIGNAL_SET(dst_rank, assist[global_slot,1], token)
                SIGNAL_SET(dst_rank, assist[global_slot,2], topk)

        QUIET(dst_rank)
        FOR local_expert FROM 0 TO E_l-1:
            FOR slot FROM 0 TO segment_count[local_expert]-1:
                global_slot = (MY_RANK*E_l + local_expert)*(B*K) + slot
                PUT_NBI(dst_rank, ready[global_slot,0:8], READY_VECTOR)
            QUIET(dst_rank)
            segment = local_expert*P + MY_RANK
            PUT_NBI(dst_rank, count[segment,0:8], [1,segment_count[local_expert],0,0,0,0,0,0])
            QUIET(dst_rank)

    SYNC_LOCAL_AIV_CORES()
    running = 0
    FOR local_expert FROM 0 TO E_l-1:
        expert_total = 0
        FOR src_rank FROM 0 TO P-1:
            segment = local_expert*P + src_rank
            WAIT_EQ(count[segment,0], 1)
            received = count[segment,1]
            running += received
            expert_total += received
            ep_recv_count[segment] = running
        expert_token_nums[local_expert] = expert_total

    FOR local_expert FROM 0 TO E_l-1:
        FOR src_rank FROM 0 TO P-1:
            segment = local_expert*P + src_rank
            begin = 0 IF segment == 0 ELSE ep_recv_count[segment-1]
            received = count[segment,1]
            FOR i FROM 0 TO received-1:
                global_slot = (src_rank*E_l + local_expert)*(B*K) + i
                WAIT_ALL_EQ(ready[global_slot,0:8], 1)
                COPY expand_x[begin+i,0:H] = payload[global_slot,0:H]
                COPY assist_out[begin+i,0:3] = assist[global_slot,0:3]
                ZERO ready[global_slot,0:8]
            FILL count[segment,0:8] WITH -1

    expert_result = RUN_LOCAL_EXPERTS(expand_x, ep_recv_count, expert_token_nums)

    ZERO combine_data_window
    ZERO combine_status_window
    BARRIER_ALL_RANKS()

    FOR owner_rank FROM 0 TO P-1 IN PARALLEL_ON_AIV:
        FOR local_expert FROM 0 TO E_l-1:
            segment = local_expert*P + owner_rank
            begin = 0 IF segment == 0 ELSE ep_recv_count[segment-1]
            end = ep_recv_count[segment]
            FOR row FROM begin TO end-1:
                src_rank = assist_out[row,0]
                token = assist_out[row,1]
                topk = assist_out[row,2]
                slot = token*K + topk
                PUT_NBI(src_rank, combine_data[slot,0:H], expert_result[row,0:H])
                QUIET(src_rank)
                PUT_NBI(src_rank, combine_status[slot,0:8], READY_VECTOR)
                QUIET(src_rank)

    FOR token FROM 0 TO B-1 IN PARALLEL_ON_AIV:
        FOR topk FROM 0 TO K-1:
            slot = token*K + topk
            WAIT_ALL_EQ(combine_status[slot,0:8], 1)
        FOR column FROM 0 TO H-1:
            accumulator = 0.0
            FOR topk FROM 0 TO K-1:
                slot = token*K + topk
                accumulator += FLOAT(combine_data[slot,column]) * expert_scales[token,topk]
            x_out[token,column] = CAST_TO_T(accumulator)
        FOR topk FROM 0 TO K-1:
            slot = token*K + topk
            ZERO combine_status[slot,0:8]

    SYNC_LOCAL_AIV_CORES()
    BARRIER_ALL_RANKS()
    RETURN x_out

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_NOTIFY WQE，但返回仍只表示提交。

接口契约 Put+Notify 实现

flowchart LR
    B["瓶颈: Host 往返/小操作提交"] --> A["机制: AIV 设备侧 issue"]
    A --> M["MTE/SDMA/UDMA 数据面"]
    M --> O["可观察: comm_only"]
    A --> F["代价: issue/quiet/wait/状态协议"]
    Q["量化"] --> T["更少通信字节 + 不同 GEMM"]
    G["融合"] --> E["更少阶段边界 + 更复杂流水"]
    S["segment-aware 双平面"] --> D["大段 SDMA、小段 MTE"]
    O -.缺少消融时.-> N["不能把 full_op/e2e 总收益归因 AIV"]
    T -.缺少同结构 baseline.-> N
    E -.缺少同结构 baseline.-> N
    D -.只证明算法差异.-> N

两个融合变体

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

不是 UDMA 证据

该样例目录没有 aclshmemx_udma_* 调用，真实路径是 aclshmem_ptr + GM-to-GM/MTE。其默认示例是 EP=2, M=64, K=7168, N=4096, E_l=2, TopK=8, maxOutputSize=1024，Host 只循环一次且没有计时，因此只能证明功能路径，不能证明性能提升。

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。

异常与边界

1. 无效 expert id：classic dispatch 没有在 kernel 内显式拒绝 expert_id >= P*E_l。无效 route 可能被发送端丢弃，而 combine 仍等待其 status，生产接入必须前置校验。
2. 无限等待：classic signal wait、固定 fork barrier、MegaMoE generation wait 均无故障超时；peer 失败或状态缺失会挂住。
3. 容量截断：固定 fork 的 maxOutputSize 和 MegaMoE 的 maxReceivedTokens 会限制接收行数；这是有损容量语义，不是无损 All-to-AllV。
4. 状态复用：ready/count/status 必须在下一代前清理；generation wait 应比较 >= generation，否则快 rank 跨代会让慢 rank永远等不到相等值。
5. 对称地址：aclshmem_ptr、UDMA Put/Get/Notify 要求对称分配和一致布局；TileXR IPC peerMems[] 不能自动等价为 SHMEM symmetric memory。
6. TileXR 清理所有权：TileXRCommDestroy 和析构函数都可能 finalize SHMEM，前者还直接 free udmaInfoPtr，但注释称该指针由 SHMEM 管理；需要统一所有权和失败回滚。
7. 平台不可混用：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%+”这一受限事实。

可复现实验最小矩阵

同一硬件、节点、链路、版本、rank 映射、shape、dtype 和路由分布下，至少分别比较：Host/设备发起、classic/doubleplane、未融合/融合、BF16/FP8，以及 comm_only、kernel_event、full_op、e2e。只有单变量对照才能归因。

总结

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 已实现更重要。

固定来源

• TileXR 46c58f3d
• LingquLab/shmem b79bda38
• cann/shmem 382afa08

评论