notify_dispatch & dispatch_layout 代码深度剖析

一、概述

本文档对 CAM (Collective Arithmetic Module) 算子库中的两个核心算子进行完整的代码深度分析：

算子	类型	核心职责
`notify_dispatch`	通信算子	元数据 All-to-All 交换：在各 rank 之间交换 token-per-expert 计数信息，为后续数据 dispatch 做准备
`dispatch_layout`	纯计算算子	按 expert 粒度重排 token，计算每个 token 在目标 expert 接收缓冲区中的索引位置

关键设计特征：notify_dispatch 使用 自管理 IPC（基于 HCCL Window 的单边 RDMA 写），而非 HCCL AllToAll 集合通信原语。

二、notify_dispatch 完整执行流程

2.1 类定义与常量

template <typename T> class NotifyDispatch {
    constexpr static int64_t MAX_RANK_PER_CORE = 8;
    constexpr static int64_t MULTI_RANK_SIZE = 48;
    constexpr static int64_t MAX_BUFFER_NUMBER = 10;
    constexpr static uint32_t UB_FLAG_SIZE = 8U * 1024U;  // 8KB flag buffer
    constexpr static int64_t FLAG_UNIT_INT_NUM = 4;        // 每个 flag 占 4 个 int32 = 16 Bytes
    constexpr static int64_t MAGIC_MASK = ~((1LL << 32) - 1);
    constexpr static uint32_t FIRST = 1, SECOND = 2, THIRD = 3, FOURTH = 4;

关键常量（定义于 comm_args.h）：

常量	值	含义
`IPC_DATA_OFFSET`	2 MB	共享内存前 2MB 为 flag 区，其后 100MB 为数据区
`IPC_BUFF_MAX_SIZE`	100 MB	IPC 数据区大小
`PING_PONG_SIZE`	2	Ping-Pong 双缓冲数量
`UB_HEAD_OFFSET`	96	第一块 UB 偏移
`UB_MID_OFFSET`	~97 KB	第二块 UB 偏移
`STATE_WIN_OFFSET`	900 KB	Magic 值在 Window 中的偏移
`sendPerGroup`	3	每个 expert 发送 (count, offset, numTokens)

2.2 内存布局

共享内存布局 (每个 rank 的 Window):
┌─────────────────────┐ ← shareAddrs[rank]
│   Flag 区 (2MB)     │  ← IPC_DATA_OFFSET 之前
│   - 同步 flag       │     每个 flag = 4 × int32 = 16 Bytes
│   - 按 blockIdx 索引 │
├─────────────────────┤ ← IPC_DATA_OFFSET (2MB)
│   Data 区 (100MB)   │  ← IPC_BUFF_MAX_SIZE
│   - sendData 数据   │     每个 rank 的 perRankDataNum 数据
│   - Ping-Pong 交替  │     magic % 2 选择缓冲
└─────────────────────┘

2.3 Process() 主流程

__aicore__ inline void Process()
{
    if (blockIdx < 1) {
        AssembleSendData();          // Phase 0: 仅 core 0 组装发送数据
    }
    SyncAll<true>();                 // 全局 barrier
    if (blockIdx < coreNumPerStageX) {
        InputToShareSlice();         // Phase 1: GM → 远端 Window
    }
    if (blockIdx < coreNumPerStageY) {
        ShareToShareSlice();         // Phase 2: 远端 Window → 本地 recvData
    }
    SyncAll<true>();                 // 全局 barrier
    ReorderOutput();                 // Phase 3: 重排输出
    BuildTotalRecvTokens();          // 仅 core 0: 计算总接收 token 数
    BuildRecvCount();                // 仅 core 1: 计算 recvCount 前缀和
    BuildRecvOffset();               // 仅 core 2: 计算 recvOffset
    BuildMaxBs();                    // 仅 core 3: 计算最大 batch size
    BuildRecvTokenPerExp();          // 仅 core 4: 计算每个 expert 接收的 token 数
}

2.4 Phase 0: AssembleSendData()

仅 blockIdx < 1（即 core 0）执行。

__aicore__ inline void AssembleSendData()
{
    // 1. 从 GM 读取 tokenPerExpertData (numExperts 个 int32)
    CpGM2UB(tokenPerExpertUB, tokenPerExpertDataInputGt, tokenPerExpertDataAlignLen);

    // 2. 构建 sendDataUB 和 sendDataOffsetUB
    int prefixSum = 0;
    for (int i = 0; i < numExperts; ++i) {
        int numTokensExpert = tokenPerExpertUB[i];
        sendDataUB[i * sendPerGroup]     = numTokensExpert;  // count
        sendDataUB[i * sendPerGroup + 1] = prefixSum;        // offset
        sendDataUB[i * sendPerGroup + 2] = numTokens;        // total tokens
        sendDataOffsetUB[i] = prefixSum;
        prefixSum += numTokensExpert;
    }

    // 3. 写回 GM
    CpUB2GM(sendDataInputGt, sendDataUB, sendDataAlignLen);
    CpUB2GM(sendDataOffsetOutputGt, sendDataOffsetUB, sendDataOffsetAlignLen);
}

数据结构： - sendDataUB[i * 3 + 0] = 发送给 expert i 的 token 数量 - sendDataUB[i * 3 + 1] = 在本地缓冲区中的偏移（前缀和） - sendDataUB[i * 3 + 2] = 本 rank 的总 token 数

2.5 Phase 1: InputToShareSlice()

每个 core 负责 rankNumPerCore 个 rank 的数据搬运。

__aicore__ inline void InputToShareSlice()
{
    int64_t copyOffset = blockIdx * rankNumPerCore;
    copyLen = min(rankSize - copyOffset, rankNumPerCore);

    if (copyLen > 0) {
        // 1. 从本地 GM 拷贝 sendData 到各 rank 的 Window
        readGt = sendDataInputGt[copyOffset * perRankDataNum];
        CpGM2GMPingPong<T>(copyLen * perRankDataNum * sizeof(T), 
                           readGt, writeGt[copyOffset * perRankDataNum], COPYONLY);

        // 2. 写入 flag 通知远端数据就绪
        uint64_t v = MergeMagicWithValue(magic, 1);
        *inputUB = v;
        for (int i = copyOffset; i < copyOffset + copyLen; ++i) {
            CpUB2GM(shareAddrs[i] + rank * FLAG_UNIT_INT_NUM, inputUB, sizeof(uint64_t));
        }
    }
}

关键行为： - 使用 CpGM2GMPingPong 将本地 sendData 广播到所有 rank 的共享内存 - 通过单边 RDMA 写 flag（magic * 2 + 1）通知远端"数据已就绪" - 这是 自管理 IPC 而非 HCCL 集合通信的核心体现

2.6 Phase 2: ShareToShareSlice()

__aicore__ inline void ShareToShareSlice()
{
    // 1. 计算需要读取的 rank（轮转避免冲突）
    for (int i = copyOffset; i < copyOffset + copyLen; ++i) {
        checkRank[i - copyOffset] = i + rank % copyLen;
        if (checkRank[i - copyOffset] >= copyOffset + copyLen)
            checkRank[i - copyOffset] -= copyLen;
    }

    // 2. 等待本 rank 的 flag 被远端写入（数据就绪信号）
    WaitSyncFlag(magic, 1, copyOffset, rank, copyLen);

    // 3. 从各远端 rank 的 Window 读取数据到本地 recvDataOutput
    for (int i = 0; i < copyLen; i++) {
        CpGM2GMPingPong<T>(perRankDataNum * sizeof(T), 
                           readGt1[i][rank * perRankDataNum],    // 源：远端 Window
                           recvDataOutputGt[checkRank[i] * perRankDataNum],  // 目的：本地
                           COPYONLY);
    }
}

轮转模式 (checkRank 计算)： - rank r 从 rank copyOffset + (i + r) % copyLen 读取数据 - 避免多个 rank 同时读取同一远端造成的竞争

2.7 Phase 3: 输出构建

各 core 分工明确（按 blockIdx 分发）：

blockIdx	函数	输出
0	`BuildTotalRecvTokens()`	`totalRecvTokens` — 所有 expert 接收 token 总数
1	`BuildRecvCount()`	`recvCount` — 每个 expert 的累计接收计数（前缀和）
2	`BuildRecvOffset()`	`recvOffset` — 每个 expert 的接收偏移
3	`BuildMaxBs()`	`maxBs` — 所有 expert 中最大 batch size
4	`BuildRecvTokenPerExp()`	`recvTokensPerExpert` — 本 rank 每个 expert 接收的 token 数

Reorder 逻辑（以 ReorderSendCountOutput 为例）：

__aicore__ inline void ReorderSendCountOutput()
{
    for (uint32_t expId = 0; expId < numExperts / rankSize; ++expId) {
        for (uint32_t srcRank = 0; srcRank < rankSize; ++srcRank) {
            uint32_t index = expId * rankSize + srcRank;
            uint32_t pair_idx = sendPerGroup * (srcRank * numExperts / rankSize + expId);
            sendCountTensor(index) = recvDataTensor(pair_idx);
        }
    }
}

将 [numExperts/sendPerGroup, rankSize] 的二维布局转置为 [numExperts/rankSize, rankSize] 的按 expert 分组布局。

三、notify_dispatch Ping-Pong 机制

3.1 CpGM2GMPingPong 实现

template <typename T>
template <typename K, typename U>
__aicore__ inline void NotifyDispatch<T>::CpGM2GMPingPong(
    int64_t dataSizeRemain,
    const GlobalTensor<U> &sendDataInputGt,
    const GlobalTensor<K> &recvDataOutputGT, int op)
{
    constexpr int32_t ubBlockSize = UB_SINGLE_PING_PONG_ADD_SIZE_MAX;  // 95KB
    constexpr int32_t inputUbBlockSize = ubBlockSize;  // K == U 时
    constexpr int32_t outputUbBlockSize = ubBlockSize;

    // 双缓冲 UB 分配
    __ubuf__ U *inputUB[2] = {(__ubuf__ U *)(UB_HEAD_OFFSET), (__ubuf__ U *)(UB_MID_OFFSET)};
    __ubuf__ K *outputUB[2] = {(__ubuf__ K *)inputUB[0], (__ubuf__ K *)inputUB[1]};

3.2 双缓冲流水线

UB 内存布局:
┌──────────────────┐ 0
│   偏移 96 Bytes  │ ← UB_HEAD_OFFSET (跳过头部)
├──────────────────┤
│                  │
│  inputUB[0]     │ ← 95KB (Ping)
│  outputUB[0]    │  (K==U 时共用)
│                  │
├──────────────────┤ ~97KB
│  inputUB[1]     │ ← 95KB (Pong)
│  outputUB[1]    │
│                  │
└──────────────────┘

流水线执行：

SetFlag<MTE3_MTE2>(EVENT_ID0);  // MTE2 等待 MTE3 完成
SetFlag<MTE3_MTE2>(EVENT_ID1);

for (int64_t i = 0; dataSizeRemain > 0; i++) {
    uint32_t size = min(dataSizeRemain, outputUbBlockSize);
    event_t eventId = (i & 1) ? EVENT_ID0 : EVENT_ID1;

    // 1. 等待上一轮 MTE3 完成（UB→GM 写回完毕）
    WaitFlag<MTE3_MTE2>(eventId);

    // 2. GM → UB 加载当前块
    CpGM2UB((i & 1) ? inputUB[0] : inputUB[1], input + inputOffsetNum, size);

    // 3. [可选] 类型转换
    if constexpr (!std::is_same_v<K, U>) {
        SetWaitEvent<MTE2_V>(eventId);   // 等 MTE2 完成
        CastImpl(outputUB[...], inputUB[...]);
        SetWaitEvent<V_MTE3>(eventId);   // 等 Cast 完成
    }

    // 4. UB → GM 写回
    SetFlag<MTE2_MTE3>(eventId);
    WaitFlag<MTE2_MTE3>(eventId);
    CpUB2GM(output + outputOffsetNum, outputUB[...], size);

    // 5. 通知 MTE2 可开始下一轮
    SetFlag<MTE3_MTE2>(eventId);

    dataSizeRemain -= size;
    inputOffsetNum += size / sizeof(K);
    outputOffsetNum += size / sizeof(K);
}

流水线时序：

时间 →
MTE2:  load[0]          load[1]          load[2]
MTE3:       store[0]          store[1]          store[2]
           ←── 块 0 ──→←── 块 1 ──→

EVENT_ID0 对应偶数块（i=0,2,4...）
EVENT_ID1 对应奇数块（i=1,3,5...）
MTE2（GM→UB）和 MTE3（UB→GM）形成三级流水线

3.3 跨类型拷贝

当 K != U 时（如 float16 → int32），UB 被分为两部分：

constexpr int32_t ubAlignNum = ubBlockSize / (sizeof(K) + sizeof(U)) / UB_ALIGN_SIZE * UB_ALIGN_SIZE;
constexpr int32_t inputUbBlockSize = ubAlignNum * sizeof(U);
constexpr int32_t outputUbBlockSize = ubAlignNum * sizeof(K);

outputUB[0] = inputUB[0] + inputUbBlockSize / sizeof(U);
outputUB[1] = inputUB[1] + inputUbBlockSize / sizeof(U);

四、notify_dispatch 同步协议

4.1 Magic 值机制

template <typename T> __aicore__ inline uint64_t NotifyDispatch<T>::GetMagicValue(void)
{
    GM_ADDR statusDataSpaceGm = (GM_ADDR)(winContext_[COMM_EP_IDX]->localWindowsExp);
    selfDataStatusTensor.SetGlobalBuffer(statusDataSpaceGm + STATE_WIN_OFFSET);

    // 读取当前 magic
    magic = selfDataStatusTensor(blockIdx * UB_ALIGN_SIZE);
    if (magic <= 0) magic = 1;

    // 自增并写回
    selfDataStatusTensor(blockIdx * UB_ALIGN_SIZE) = magic + 1;
    return magic;
}

作用：区分不同批次的算子调用，防止旧调用的 flag 被新调用误判。

4.2 Flag 编码

// NotifyDispatch 中的编码
__aicore__ inline uint64_t MergeMagicWithValue(uint64_t magic, uint64_t value)
{
    return (magic * 2ULL + value);  // magic 在高位，value 在低位
}

// SyncCollectives 中的编码（更标准）
__aicore__ inline int64_t MergeMagicWithValue(int32_t magic, int32_t value)
{
    return (static_cast<int64_t>(magic) << 32) | static_cast<int64_t>(value);
}

4.3 WaitOneRankPartFlag 轮询

__aicore__ inline void WaitOneRankPartFlag(__gm__ uint64_t *waitAddr, int64_t flagNum, uint64_t checkValue)
{
    do {
        // 1. GM → UB 拷贝 flag
        DataCopy(localWait, globalWait[checkedFlagNum * FLAG_UNIT_INT_NUM],
                 (flagNum - checkedFlagNum) * FLAG_UNIT_INT_NUM);

        // 2. 逐 flag 检查
        isSync = true;
        for (auto i = 0; i < remainToCheck; ++i) {
            uint64_t v = localWait.GetValue(i * FLAG_UNIT_INT_NUM);
            if ((v & MAGIC_MASK) != (checkValue & MAGIC_MASK) || v < checkValue) {
                isSync = false;
                checkedFlagNum += i;  // 跳过已确认的 flag
                break;
            }
        }
    } while (!isSync);
}

优化：checkedFlagNum 记录已确认的 flag 数量，后续轮询跳过已完成的。

4.4 同步协议总结

Rank A 写入数据到 Rank B 的 Window:
┌─────────────────────────────────────────────────────────┐
│ 1. Rank A: CpGM2GMPingPong → 写数据到 B 的 Data 区     │
│ 2. Rank A: CpUB2GM → 写 flag (magic*2+1) 到 B 的 Flag 区│
│ 3. Rank B: WaitSyncFlag → 轮询等待 flag == magic*2+1    │
│ 4. Rank B: CpGM2GMPingPong → 从 A 的 Window 读数据     │
└─────────────────────────────────────────────────────────┘

五、dispatch_layout 算法细节

5.1 类定义与输入输出

template <typename T> class DispatchLayout {
    // 输入
    GM_ADDR topkIdx;           // [numTokens, numTopk] int64 — 每个 token 选中的 expert
    GM_ADDR numTokensPerRank;  // [numRanks] int32 — 输出：每个 rank 的 token 数
    GM_ADDR numTokensPerExpert;// [numExperts] int32 — 输出：每个 expert 的 token 数
    GM_ADDR isTokenInRank;     // [numTokens, numRanks] int32 — 输出：token 是否发往某 rank
    GM_ADDR notifySendData;    // 可变大小 int32 — 输出：给 notify_dispatch 的临时数据
    GM_ADDR sendTokenIdxSmall; // [numTokens, numTopk] int32 — 输出：token 在 expert 中的索引
};

5.2 核心算法：Expert-level Token Counting

int experts_per_rank = numExperts_ / numRanks_;
for (int i = 0; i < tempTokens_; ++i) {
    Duplicate<T>(seenRankTensor, 0, numRanks_);  // 每 token 重置
    for (int j = 0; j < numTopk_; ++j) {
        int64_t expert_idx = topkIdxTensor.GetValue(i * numTopk_ + j);

        // 1. 统计 expert 维度计数
        uint32_t per_expert_num = numTokensPerExpertTensor.GetValue(expert_idx) + 1;
        numTokensPerExpertTensor.SetValue(expert_idx, per_expert_num);

        // 2. 统计 rank 维度计数（去重）
        int rank_id = expert_idx / experts_per_rank;
        if (!seenRankTensor.GetValue(rank_id)) {
            uint32_t per_rank_num = numTokensPerRankTensor.GetValue(rank_id) + 1;
            isTokenInRankTensor.SetValue(i * numRanks_ + rank_id, 1);
            seenRankTensor.SetValue(rank_id, 1);
            numTokensPerRankTensor.SetValue(rank_id, per_rank_num);
        }
    }
}

算法语义： - numTokensPerExpert[e] = 发往 expert e 的 token 总数（含 topk 重复计数） - numTokensPerRank[r] = 发往 rank r 的 token 数（去重：同一 token 发往同一 rank 多个 expert 只计一次） - isTokenInRank[i][r] = token i 是否发往 rank r（0/1 标记）

5.3 sendTokenIdxSmall 计算

// 第二次遍历（聚合后）
for (int i = 0; i < tempTokens_; ++i) {
    for (int j = 0; j < numTopk_; ++j) {
        int64_t expert_idx = topkIdxTensor.GetValue(i * numTopk_ + j);
        T valT = numTokensPerExpertTensor(expert_idx);  // 读取当前计数
        sendTokenIdxSmallTensor(i * numTopk_ + j) = valT;  // 作为索引
        numTokensPerExpertTensor(expert_idx) = valT + 1;   // 递增
    }
}

语义：sendTokenIdxSmall[i][j] = token i 在其目标 expert 接收缓冲区中的位置索引。

示例：

token 0: topk = [expert_3, expert_5]
  → sendTokenIdxSmall[0][0] = 0  (expert_3 的第 0 个位置)
  → sendTokenIdxSmall[0][1] = 0  (expert_5 的第 0 个位置)

token 1: topk = [expert_3, expert_7]
  → sendTokenIdxSmall[1][0] = 1  (expert_3 的第 1 个位置)
  → sendTokenIdxSmall[1][1] = 0  (expert_7 的第 0 个位置)

六、dispatch_layout 原子操作

6.1 跨 Core 聚合策略

// 1. 写 isTokenInRank 到 GM
DataCopyPad(isTokenInRankGM_, isTokenInRankTensor, isTokenInRankDataCopyParams);

// 2. 开启原子加
AscendC::SetAtomicAdd<T>();

// 3. 每个 core 将自己的 numTokensPerExpert 原子加到后续 core 的临时区
for (int i = coreIdx_ + 1; i < aivNum_; ++i) {
    DataCopyPad(tempExpertGM_[i * numExperts_], numTokensPerExpertTensor, 
                tempExpertDataCopyParams);
}

// 4. 写 numTokensPerRank 和 numTokensPerExpert 到 GM
DataCopyPad(numTokensPerRankGM_, numTokensPerRankTensor, numTokensPerRankDataCopyParams);
DataCopyPad(numTokensPerExpertGM_, numTokensPerExpertTensor, numTokensPerExpertDataCopyParams);

// 5. 关闭原子操作
AscendC::SetAtomicNone();
PipeBarrier<PIPE_MTE3>();
SyncAll<true>();

6.2 前缀和构建原理

tempExpertGM_ 布局 (每个 core 写到的位置):
┌─────────────────────────────────────────────────────────┐
│ core 0: 写到 tempExpertGM_[1*N], tempExpertGM_[2*N], ...│
│ core 1: 写到 tempExpertGM_[2*N], tempExpertGM_[3*N], ...│
│ core 2: 写到 tempExpertGM_[3*N], tempExpertGM_[4*N], ...│
│ ...                                                     │
└─────────────────────────────────────────────────────────┘

原子加后:
tempExpertGM_[i * N] = sum(core_0..core_{i-1} 的 numTokensPerExpert)

效果：tempExpertGM_[i * N] 存储的是 core 0 到 core i-1 的累计计数，即 core i 的前缀和。

6.3 读取聚合结果

// SyncAll 后，每个 core 读取自己的前缀和
DataCopyPad(numTokensPerExpertTensor, tempExpertGM_[coreIdx_ * numExperts_], 
            tempExpertDataCopyParams);

此时 numTokensPerExpertTensor 包含了所有前面 core 的累计计数，第二次遍历时： - 第一次读取 numTokensPerExpertTensor[expert_idx] 得到全局索引 - 递增后为下一个 token 准备

七、pybind 层实现

7.1 get_dispatch_layout

std::tuple<at::Tensor, at::Tensor> GetDispatchLayoutImplNpu(
    const at::Tensor &topkIdx, int64_t numExperts, int64_t numRanks)
{
    const int numTokens = topkIdxInt64.size(0);
    const int numTopk = topkIdxInt64.size(1);
    const int localRanksize = LOCAL_RANK_SIZE;  // 8
    auto serverNum = numRanks / localRanksize;

    // 分配输出 tensor
    auto numTokensPerExpert = at::zeros({numExperts}, kInt);
    auto numTokensPerRank = at::zeros({numRanks}, kInt);
    auto isTokenInRank = at::zeros({numTokens, numRanks}, kInt);
    auto sendTokenIdxSmall = at::zeros({numTokens, numTopk}, kInt);

    // notifySendData 大小计算
    const int notifySendDataSize = numExperts * EXPERT_DATA_SIZE + serverNum 
                                 + MAX_BATCH_SIZE * (1 + 2 * serverNum + numExperts);
    auto notifySendData = at::zeros({notifySendDataSize}, kInt);

    // 调用 aclnnDispatchLayout
    EXEC_NPU_CMD(aclnnDispatchLayout, topkIdxInt64, numTokens, numRanks, 
                 numExperts, numTopk, localRanksize,
                 numTokensPerRank, numTokensPerExpert, isTokenInRank, 
                 notifySendData, sendTokenIdxSmall);

    return {numTokensPerExpert, sendTokenIdxSmall};
}

7.2 moe_dispatch_prefill

std::tuple<...> MoeDispatchPrefillImplNpu(
    const at::Tensor &x, const at::Tensor &topkIdx, ...)
{
    // 1. 准备 notify_dispatch 的输入
    auto sendData = torch::empty({numExperts * sendPerGroup}, kInt);
    int64_t sendCount = sendPerGroup * numLocalExperts * numRanks;

    // 2. 调用 aclnnNotifyDispatch — 元数据 All-to-All
    EXEC_NPU_CMD(aclnnNotifyDispatch, sendData, numTokensPerExpert, sendCount, numTokens,
                 groupEpPtr, numRanks, rank, localRankSize, localRankId,
                 sendDataOffset, recvData, totalRecvToken, recvCount, recvOffset, 
                 maxBs, recvTokensPerExpert);

    // 3. 调用 aclnnMoeDispatchNormal — 实际数据 dispatch
    EXEC_NPU_CMD(aclnnMoeDispatchNormal, x, expertIds, sendDataOffset, sendTokenIdxSmall,
                 recvOffset, recvCount, groupEpPtr, numRanks, rank, ...);

    return {expandxOut, dynamicScalesOut, expandIdxOut, recvCount, recvTokensPerExpert};
}

7.3 moe_combine_prefill

at::Tensor MoeCombinePrefillImplNpu(
    const at::Tensor &x, const at::Tensor &topkIdx, const at::Tensor &topkWeights,
    const at::Tensor &srcIdx, const at::Tensor &sendHead, ...)
{
    // 调用 aclnnMoeCombineNormal — 将 expert 输出聚合回原 token
    EXEC_NPU_CMD(aclnnMoeCombineNormal, x, tokenSrcInfo, epSendCounts, expertScales,
                 tpSendCounts, groupEpPtr, numRanks, rank, ...);
    return combinedX;
}

八、两个算子的协作关系

8.1 完整数据流

                    ┌─────────────────────┐
                    │   dispatch_layout   │
                    │   (纯计算，无通信)   │
                    └──────────┬──────────┘
                               │
        ┌──────────────────────┼──────────────────────┐
        ▼                      ▼                      ▼
  numTokensPerExpert    sendTokenIdxSmall      isTokenInRank
        │                      │                      │
        ▼                      │                      │
  ┌─────────────────────────────┴──────────────────────┐
  │              notify_dispatch                        │
  │  (元数据 All-to-All，自管理 IPC)                    │
  │                                                      │
  │  输入: numTokensPerExpert (本 rank 各 expert 发送数) │
  │  输出:                                                │
  │    - totalRecvTokens: 本 rank 总接收 token 数         │
  │    - recvCount: 各 expert 累计接收计数                │
  │    - recvOffset: 各 expert 接收偏移                   │
  │    - maxBs: 最大 batch size                          │
  │    - recvTokensPerExpert: 本 rank 各 expert 接收数    │
  └──────────────────────────────┬───────────────────────┘
                                 │
        ┌────────────────────────┼────────────────────────┐
        ▼                        ▼                        ▼
   recvOffset/recvCount    totalRecvTokens/maxBs    recvTokensPerExpert
        │                        │                        │
        ▼                        ▼                        ▼
  ┌─────────────────────────────────────────────────────────────┐
  │                    aclnnMoeDispatchNormal                    │
  │                   (实际数据 All-to-All)                      │
  │                                                              │
  │  输入: x, expertIds, sendDataOffset, sendTokenIdxSmall       │
  │        recvOffset, recvCount                                 │
  │  输出: expandxOut, dynamicScalesOut, expandIdxOut            │
  └─────────────────────────────────────────────────────────────┘

8.2 数据依赖关系

生产者	输出	消费者
`dispatch_layout`	`numTokensPerExpert`	`notify_dispatch` 输入
`dispatch_layout`	`sendTokenIdxSmall`	`MoeDispatchNormal` 输入
`notify_dispatch`	`recvOffset`, `recvCount`	`MoeDispatchNormal` 输入
`notify_dispatch`	`totalRecvTokens`, `maxBs`	确定接收缓冲区大小
`notify_dispatch`	`recvTokensPerExpert`	后续 combine 阶段

8.3 notifySendData 的作用

dispatch_layout 输出的 notifySendData 是一个临时缓冲区，其大小计算：

const int notifySendDataSize = numExperts * EXPERT_DATA_SIZE    // expert 数据
                             + serverNum                          // server 数量
                             + MAX_BATCH_SIZE * (1 + 2 * serverNum + numExperts);  // batch 数据

这个缓冲区在 notify_dispatch 中被用作 tempExpertGM_，用于跨 core 聚合时的原子操作目标区。

九、SHMEM vs 非 SHMEM 差异分析

9.1 结论

算子	是否涉及 SHMEM	原因
`notify_dispatch`	是（间接）	使用 HCCL Window 机制实现单边 RDMA 写，这是华为平台上的 SHMEM 等价物
`dispatch_layout`	否	纯计算算子，无任何跨 rank 通信

9.2 notify_dispatch 的通信机制

notify_dispatch 不使用 HCCL AllToAll，而是使用 自管理 IPC：

// 获取远端 rank 的 Window 地址
shareAddrs[i] = GetWindAddrByRankId(i, ctxIdx) + (magic % PING_PONG_SIZE) * (IPC_BUFF_MAX_SIZE + IPC_DATA_OFFSET);

// 直接 RDMA 写到远端 Window
CpGM2GMPingPong<T>(..., readGt, writeGt[...], COPYONLY);

// 直接 RDMA 写 flag 通知远端
CpUB2GM(shareAddrs[i] + rank * FLAG_UNIT_INT_NUM, inputUB, sizeof(uint64_t));

HCCL Window 机制： - winContext_->localWindowsIn — 本 rank 的输入 Window（本地内存映射） - winContext_->remoteRes[rankId].windowsIn — 远端 rank 的输入 Window（RDMA 远端映射） - 这是华为 HCCL 提供的 单边通信（one-sided communication）能力

9.3 与 SHMEM 的关系

传统 SHMEM 模型:
  shmem_putmem(dest, src, len, pe)  →  远端内存写
  shmem_getmem(dest, src, len, pe)  →  远端内存读

notify_dispatch 模型:
  CpGM2GMPingPong(..., localSrc, remoteDest, COPYONLY)  →  远端内存写（通过 HCCL Window）
  CpGM2GMPingPong(..., remoteSrc, localDest, COPYONLY)  →  远端内存读（通过 HCCL Window）

本质相同：都是通过单边 RDMA 操作实现跨 rank 内存访问，无需远端 CPU 参与。

9.4 为什么不用 HCCL AllToAll？

对比维度	HCCL AllToAll	自管理 IPC
同步方式	集合同步（需所有 rank 同时参与）	异步单边写 + flag 轮询
延迟	高（需建立连接、同步）	低（直接 RDMA 写）
适用场景	大数据量数据交换	小尺寸元数据交换
灵活性	低（固定 API）	高（可自定义协议）

notify_dispatch 交换的是元数据（每个 expert 的 token 计数，约 numExperts × 3 × sizeof(T) 字节），数据量小但延迟敏感。使用自管理 IPC 可以： 1. 与数据 dispatch 流水线化（overlap） 2. 避免 AllToAll 的建立开销 3. 更精细地控制同步时机

9.5 dispatch_layout 为什么无通信

dispatch_layout 的所有输入（topkIdx）和输出（numTokensPerExpert, sendTokenIdxSmall 等）都是 rank-local 的：

topkIdx 是本地 MoE routing 的输出
numTokensPerExpert 统计的是 本 rank 发往各 expert 的 token 数
sendTokenIdxSmall 是 本 rank 内部的索引

跨 rank 的信息交换完全由 notify_dispatch 完成。

十、关键代码路径索引

notify_dispatch

文件	关键函数/类	行号
`notify_dispatch.h`	`NotifyDispatch<T>` 类定义	30-513
`notify_dispatch.h`	`Process()` 主流程	91-110
`notify_dispatch.h`	`AssembleSendData()`	129-160
`notify_dispatch.h`	`InputToShareSlice()`	163-181
`notify_dispatch.h`	`ShareToShareSlice()`	230-254
`notify_dispatch.h`	`CpGM2GMPingPong()`	627-686
`notify_dispatch.h`	`WaitOneRankPartFlag()`	190-217
`notify_dispatch.h`	`GetMagicValue()`	541-555
`notify_dispatch.h`	`InitSmallFullMesh()`	557-610
`notify_dispatch_tiling.cpp`	`NotifyDispatchTilingFuncImpl()`	237-294
`notify_dispatch.cpp` (op_host)	`NotifyDispatch` OpDef	13-86
`aclnn_notify_dispatch.cpp`	`aclnnNotifyDispatch()`	41-47

dispatch_layout

文件	关键函数/类	行号
`dispatch_layout.h`	`DispatchLayout<T>` 类定义	35-197
`dispatch_layout.h`	`Init()`	39-81
`dispatch_layout.h`	`Process()`	83-166
`dispatch_layout.h`	Expert counting 循环	113-129
`dispatch_layout.h`	Atomic 聚合	134-147
`dispatch_layout.h`	sendTokenIdxSmall 计算	154-161
`dispatch_layout_tiling.cpp`	`DispatchLayoutTilingFuncImpl()`	199-230
`dispatch_layout.cpp` (op_host)	`DispatchLayout` OpDef	13-69
`aclnn_dispatch_layout.cpp`	`aclnnDispatchLayout()`	38-44

pybind 层

文件	关键函数	行号
`get_dispatch_layout.cpp`	`GetDispatchLayoutImplNpu()`	33-60
`moe_dispatch_prefill.cpp`	`MoeDispatchPrefillImplNpu()`	38-127
`moe_combine_prefill.cpp`	`MoeCombinePrefillImplNpu()`	28-64

公共工具

文件	关键定义
`comm_args.h`	IPC 常量、UB 偏移、Op 枚举
`moe_distribute_base.h`	`HcclOpResParam`、`HcclRankRelationResV2`
`sync_collectives.h`	`SyncCollectives` 类、flag 操作

十一、性能优化要点

11.1 notify_dispatch 优化

Ping-Pong 双缓冲：MTE2 和 MTE3 流水线化，隐藏 GM 访问延迟
轮转读取：checkRank 计算避免多 rank 同时读取同一远端
Flag 轮询优化：checkedFlagNum 跳过已确认的 flag
按 core 分工：5 个输出由 5 个 core 并行计算

11.2 dispatch_layout 优化

Token 分片：tempTokens_ = numTokens / aivNum，余数分配给前几个 core
原子操作聚合：SetAtomicAdd 实现无锁跨 core 累加
前缀和策略：通过 tempExpertGM_ 不同偏移实现隐式前缀和
seenRank 去重：避免同一 token 对同一 rank 重复计数

十二、总结

notify_dispatch

核心功能：元数据 All-to-All 交换
通信方式：基于 HCCL Window 的自管理 IPC（单边 RDMA 写）
同步协议：Magic 值 + Flag 轮询
流水线：Ping-Pong 双缓冲 GM↔GM 拷贝
输出：recvCount、recvOffset、totalRecvTokens、maxBs、recvTokensPerExpert

dispatch_layout

核心功能：按 expert 粒度重排 token
算法：Expert-level counting + 前缀和索引
同步方式：SetAtomicAdd 跨 core 聚合 + SyncAll
输出：numTokensPerExpert、sendTokenIdxSmall、isTokenInRank

协作关系

dispatch_layout (本地计算)
    ↓
notify_dispatch (元数据交换)
    ↓
MoeDispatchNormal (数据交换)

dispatch_layout 和 notify_dispatch" 共同构成了 MoE dispatch 的 **控制平面**，为数据平面的MoeDispatchNormal` 提供完整的路由信息。

← 非 SHMEM — HCCL Window 📑 索引 FusedDeepMoE SHMEM 路径 →

notify_dispatch + dispatch_layout — UMDK/CAM

notify_dispatch & dispatch_layout 代码深度剖析

一、概述

二、notify_dispatch 完整执行流程

2.1 类定义与常量

2.2 内存布局

2.3 Process() 主流程

2.4 Phase 0: AssembleSendData()

2.5 Phase 1: InputToShareSlice()

2.6 Phase 2: ShareToShareSlice()

2.7 Phase 3: 输出构建

三、notify_dispatch Ping-Pong 机制

3.1 CpGM2GMPingPong 实现

3.2 双缓冲流水线

3.3 跨类型拷贝

四、notify_dispatch 同步协议

4.1 Magic 值机制

4.2 Flag 编码

4.3 WaitOneRankPartFlag 轮询

4.4 同步协议总结

五、dispatch_layout 算法细节

5.1 类定义与输入输出

5.2 核心算法：Expert-level Token Counting

5.3 sendTokenIdxSmall 计算

六、dispatch_layout 原子操作

6.1 跨 Core 聚合策略

6.2 前缀和构建原理

6.3 读取聚合结果

七、pybind 层实现

7.1 get_dispatch_layout

7.2 moe_dispatch_prefill

7.3 moe_combine_prefill

八、两个算子的协作关系

8.1 完整数据流

8.2 数据依赖关系

8.3 notifySendData 的作用

九、SHMEM vs 非 SHMEM 差异分析

9.1 结论

9.2 notify_dispatch 的通信机制

9.3 与 SHMEM 的关系

9.4 为什么不用 HCCL AllToAll？

9.5 dispatch_layout 为什么无通信

十、关键代码路径索引

notify_dispatch

dispatch_layout

pybind 层

公共工具

十一、性能优化要点

11.1 notify_dispatch 优化

11.2 dispatch_layout 优化

十二、总结

notify_dispatch

dispatch_layout

协作关系