ZB (Zero-Buffer) Dispatch 算子深度分析

项目: UMDK/CAM — Communication Acceleration for Matrix on Ascend NPU 分支: remotes/origin/cam_feature 分析日期: 2026-07-21 分析范围: 4 个 ZB (Zero-Buffer) Dispatch 系列算子 + 共享基础设施

ZB (Zero-Buffer) 设计哲学
共享基础设施分析
moe_dispatch_layout_zb — 布局算子
moe_dispatch_normal_zb — 标准分发算子
moe_dispatch_lowlatency_zb — 低延迟分发算子
moe_notify_dispatch_zb — 元数据通知算子
SHMEM 集成方式
性能优势分析
关键数据结构与常量汇总

1. ZB (Zero-Buffer) 设计哲学

1.1 什么是 Zero-Buffer？

ZB (Zero-Buffer) 是 CAM 库在 cam_feature 分支上引入的下一代 MoE Dispatch 通信范式。其核心思想是：

完全摒弃传统的 HCCL Window 通信模式，改用 CANN SHMEM（用户态 RDMA）API 直接在卡间对称内存（symmetric memory）上进行数据搬运和同步。

1.2 传统 HCCL Window 模式的问题

在传统的 MoE Dispatch 流程中（参见 07_hccl_window_deepdive.md），通信依赖 HCCL Window 机制：

维度	HCCL Window 模式	Zero-Buffer 模式
内存占用	每卡预分配 ~200MB HCCL Buffer 作为通信窗口	零额外内存，直接使用对称内存布局
通信路径	数据 → HCCL Window → RDMA → 对端 HCCL Window → 用户 Buffer	数据 → RDMA → 对端用户 Buffer（直达）
同步机制	HCCL Collective + 信号量	SHMEM flag 轮询 + Magic 机制
延迟	多次内存拷贝 + HCCL 软件栈开销	单次 RDMA + 用户态 flag 同步
灵活性	窗口大小固定，需预分配	按需使用，无固定窗口限制

1.3 Zero-Buffer 的核心优势

零额外内存: 不需要预分配 HCCL Window Buffer，节省 ~200MB/卡
单次 RDMA 直达: 数据直接从源卡写入目标卡的用户可见内存，无需中间拷贝
用户态同步: 通过 SHMEM flag + Magic 机制实现轻量级跨卡同步，避免重量级 HCCL Collective
对称内存布局: 所有卡上的 flag 区域布局完全一致，通过 shmem_ptr(base, rankId) 即可定位对端内存

1.4 代码组织

ZB 系列算子在 cam_feature 分支上的文件组织遵循 CANN 算子标准三层结构：

moe_dispatch_XXX_zb/
├── op_api/          # aclnn 对外接口层
│   ├── aclnn_XXX_zero_buffer.cpp
│   └── aclnn_XXX_zero_buffer.h
├── op_host/         # Host 端 tiling 与算子注册
│   ├── XXX.cpp          # OpDef 注册 + 算子配置
│   └── XXX_tiling.cc    # Tiling 逻辑 + 工作空间计算
└── op_kernel/       # Device 端 kernel 实现
    ├── XXX.cpp          # 入口函数
    ├── XXX.h            # Kernel 核心逻辑
    └── XXX_tiling.h     # TilingData 结构定义

共享工具位于 utils/:

utils/op_kernel/
├── zero_buffer_api.h          # SHMEM API 封装（#include "shmem.h"）
├── zero_buffer_sync_flag.h    # 跨卡 flag 同步机制（核心！）
├── comm_args_zb.h             # 通信参数（#include "comm_args.h"）
├── data_copy_zb.h             # 数据拷贝（#include "data_copy.h"）
├── moe_lowlatency_zero_buffer_base.h  # 低延迟基类工具
└── sync_collectives_zb.h      # 集合同步（旧版，部分算子仍使用）

2. 共享基础设施分析

2.1 zero_buffer_sync_flag.h — 核心同步机制

这是整个 ZB 体系最关键的共享组件。ZeroBufferSyncFlag 类实现了基于对称共享内存的跨卡 flag 同步。

内存布局

gva_gm + FLAG_AREA_BASE (100KB):
┌──────────────────────────────────────────────────────────┐
│ Magic Area (24 KB)                                       │
│   slot[coreIdx] at coreIdx × 512B                        │
│   每核单调递增，每次 IncrementMagic() +1                  │
├──────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 0                                │
│   slot[srcRank] (32B each), epWorldSize slots            │
├──────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 1 (同上)                          │
├──────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 0                                   │
│   flag[srcRank][eventID] (32B each)                      │
│   srcRank ∈ [0, epWorldSize)                             │
│   eventID ∈ [0, slotsPerRank)                            │
├──────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 1 (同上)                             │
└──────────────────────────────────────────────────────────┘

Total = Magic(24KB) + 2×BarrierBuf + 2×FlagBuf

Magic 机制

// 高 32 位存储 magic，低 32 位存储 phase
__aicore__ inline int64_t MakeFlagValue(int32_t magic, int32_t value)
{
    return (static_cast<int64_t>(static_cast<uint32_t>(magic)) << 32) |
           static_cast<int64_t>(static_cast<uint32_t>(value));
}

每个 kernel 入口调用 IncrementMagic() 一次
Magic 单调递增：kernel K 得到 M，kernel K+1 得到 M+1
Ping-pong buffer 选择：magic % 2
关键不变量: 单调 magic 保证任何旧 kernel 的 stale flag 都会被拒绝

同步协议

Writer (rank A, core i):
  1. 完成数据写入（DMA 到远端 shmem）
  2. SetFlag(destRank=B, eventID=i, phase)
     → MTE3 fence（flush 所有先前的写入）
     → 写 flag 到 B 的本地 shmem slot [A][i]

Reader (rank B, core j):
  1. WaitFlag(srcRank=A, eventID=i, phase)
     → 轮询本地 shmem 直到 flag 匹配（magic, phase）
  2. 返回后，rank A 的数据保证可见

同步粒度

模式	slotsPerRank	使用场景
Per-core	blockNum	每个 core 向所有 rank 发信号
Per-slot	N (任意)	细粒度事件同步
Per-rank	1	粗粒度 rank 间同步

2.2 zero_buffer_api.h — SHMEM API 封装

// zero_buffer_api.h
#include "shmem.h"  // CANN SHMEM 提供的用户态 RDMA API

shmem.h 是 CANN 工具链提供的外部头文件，提供： - aclshmem_ptr(baseAddr, rankId) — 获取对端 rank 的远端内存指针 - aclshmem_half_put_nbi() — half 类型非阻塞 RDMA put - aclshmem_bfloat16_put_nbi() — bf16 类型非阻塞 RDMA put - aclshmem_float_put_nbi() — float 类型非阻塞 RDMA put - shmem_get_int32_mem_nbi() — int32 类型非阻塞 RDMA get

2.3 sync_collectives_zb.h — 旧版集合同步

SyncCollectives 是一个更早期的同步实现，使用双内存区域设计： - Memory A（basicSyncAddr）: 卡内同步 flag - Memory B（blockOuterSyncAddr）: 卡间同步 flag

部分 ZB 算子（如 lowlatency_zb）仍混合使用 SyncCollectives 和 ZeroBufferSyncFlag。

2.4 moe_lowlatency_zero_buffer_base.h — 低延迟基类

提供 ZeroBufferInitWinState() 函数，用于初始化数据状态（dataState）并原子递增操作计数。

3. moe_dispatch_layout_zb — 布局算子

3.1 算子概述

moe_dispatch_layout_zb 是 MoE Dispatch 流程的第一阶段，负责： 1. 统计每个 rank 和每个 expert 接收的 token 数量 2. 计算 isTokenInRank 标记矩阵 3. 计算 sendTokenIdxSmall（token 在目标 expert 内的局部序号）

关键特征: 此算子不涉及跨卡通信，纯本地计算。它使用 SyncAll<true>() 进行卡内多核同步。

3.2 算子注册 (op_host/dispatch_layout.cpp)

// 输入
this->Input("topkIdx")         // DT_INT64, [numTokens, numTopk]
this->Attr("num_tokens").Int()
this->Attr("num_ranks").Int()
this->Attr("num_experts").Int()
this->Attr("num_topk").Int()
this->Attr("local_ranksize").Int()

// 输出
this->Output("numTokensPerRank")   // DT_INT32, [numRanks]
this->Output("numTokensPerExpert") // DT_INT32, [numExperts]
this->Output("isTokenInRank")      // DT_INT32, [numTokens, numRanks]
this->Output("notifySendData")     // DT_INT32, 复合缓冲区
this->Output("sendTokenIdxSmall")  // DT_INT32, [numTokens, numTopk]

// 平台配置
this->AICore().AddConfig("ascend910_93", a3_config);  // 910_93
this->AICore().AddConfig("ascend910b", a2_config);    // 910B (A2)

3.3 Tiling 策略 (op_host/dispatch_layout_tiling.cc)

// 约束检查
const int64_t MAX_COMM_WORLD_SIZE = 384;
const int64_t MAX_MOE_EXPERTS_NUM = 512;
const int64_t MAX_LOCAL_RANKSIZE = 8;
const uint32_t K_MAX = 16;

// 工作空间
constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024;  // 16MB
constexpr uint32_t KERNEL_USE_WORKSPACE = 1 * 1024 * 1024;    // 1MB
constexpr uint32_t KERNEL_A2_ARG_SIZE = 1 * 1024 * 1024;      // 1MB

// Tiling Key
constexpr static int TILING_KEY_INT = 23;       // A1 (910_93)
constexpr static int TILING_KEY_A2_TYPE = 100;  // A2 (910B) → key = 123

平台检测: 通过 CheckIfA2Machine() 检测 SoC 版本：

std::string socVersion;
platformInfo.GetPlatformResWithLock("version", "Short_SoC_version", socVersion);
if (socVersion == "Ascend910B") { return true; }

3.4 Kernel 实现 — A1 版本 (dispatch_layout.h)

核心算法：两阶段 token 计数

Phase 1: CountTokensInRound() — 分轮次统计

// 每轮处理 tokensPerRound 个 token
// UB 预算: 190KB max
uint32_t fixedSize = numTokensPerRank32AlignIntLen_ + numTokensPerExpert32AlignIntLen_ +
    Ceil(numRanks_ * sizeof(T), UB_32_ALIGN) * UB_32_ALIGN;

uint32_t perTokenSize = Ceil(numTopk_ * sizeof(int64_t), UB_32_ALIGN) * UB_32_ALIGN +
    Ceil(numRanks_ * sizeof(T), UB_32_ALIGN) * UB_32_ALIGN +
    Ceil(numTopk_ * sizeof(T), UB_32_ALIGN) * UB_32_ALIGN;

uint32_t tokensPerRound = (UB_MAX_SIZE - fixedSize) / perTokenSize;

每个 token 的处理逻辑：

for (uint32_t j = 0; j < numTopk_; ++j) {
    int64_t expertIdx = topkIdxLocal.GetValue(i * numTopk_ + j);
    expertLocal.SetValue(expertIdx, expertLocal.GetValue(expertIdx) + 1);
    int rankId = expertIdx / expertsPerRank;
    if (!seenLocal.GetValue(rankId)) {
        rankLocal.SetValue(rankId, rankLocal.GetValue(rankId) + 1);
        isTokenLocal.SetValue(i * numRanks_ + rankId, 1);
        seenLocal.SetValue(rankId, 1);
    }
}

Phase 2: CalcSendTokenIdxInRound() — 计算发送序号

for (uint32_t i = 0; i < roundTokens; ++i) {
    for (uint32_t j = 0; j < numTopk_; ++j) {
        int64_t expertIdx = topkIdxLocal.GetValue(i * numTopk_ + j);
        T val = expertLocal(expertIdx);
        sendIdxLocal(i * numTopk_ + j) = val;
        expertLocal(expertIdx) = val + 1;  // 递增，下一个 token 的序号
    }
}

跨核原子累加

Phase 1 中，每个 core 计算完自己的 numTokensPerExpert 后，通过 SetAtomicAdd 将结果写入后续 core 的临时区域，实现跨核累加：

AscendC::SetAtomicAdd<T>();
for (uint32_t i = coreIdx_ + 1; i < aivNum_; ++i) {
    DataCopyPad(tempExpertGM_[i * numExperts_], expertLocal, tempExpertParams);
}
AscendC::SetAtomicNone();

3.5 Kernel 实现 — A2 变体 (dispatch_layout_a2.h)

A2 变体（DispatchLayoutA2）是为 Ascend910B (A2) 芯片设计的增强版本，引入了服务器层级（server-level hierarchy）概念。

关键差异

维度	A1 (DispatchLayout)	A2 (DispatchLayoutA2)
处理模式	分轮次 (round-based)	单次全量 (single-pass)
层级	Rank 级别	Server + Rank 级别
新增输出	无	localTokenServer* 系列
同步	SyncAll (多轮)	单次 SyncAll

A2 新增数据结构

// Server 级别元数据
GlobalTensor<T> localTokenServerUniqCountGM_;    // 每个 server 的唯一 token 数
GlobalTensor<T> localTokenServerTotalCountGM_;    // 每个 token 在每个 server 的总数
GlobalTensor<T> localTokenServerNumGM_;           // 每个 token 需要发送的 server 数
GlobalTensor<T> localTokenServerOffsetGM_;        // server 内偏移
GlobalTensor<T> expertRankTokenIdxGM_;            // expert 内 token 索引
GlobalTensor<T> sendTokenIdxGM_;                  // 全局发送序号
GlobalTensor<T> tempServerGM_;                    // server 临时缓冲区

A2 核心计算

// 计算 server ID
int server_id = rank_id / localRankSize_;

// 统计每个 token 涉及的 server 数量
if (!seenServerTensor.GetValue(server_id)) {
    uint32_t uniqCount = localTokenServerUniqCountTensor.GetValue(server_id) + 1;
    localTokenServerUniqCountTensor.SetValue(server_id, uniqCount);
    localTokenServerOffsetTensor.SetValue(i * serverNum_ + server_id, 1);
    seenServerTensor.SetValue(server_id, 1);
    uint32_t sendServerNum = localTokenServerNumTensor.GetValue(i);
    localTokenServerNumTensor.SetValue(i, sendServerNum + 1);
}

A2 前缀和计算

A2 使用向量指令计算前缀和，为每个 token 在目标 expert 内的位置生成全局唯一偏移：

// 前缀和：计算每个 token 在 expert 内的全局偏移
DataCopy(prefixCountPerExpertTensor, numTokensPerExpertTensor, numExperts_);
for (int i = 0; i < tempTokens_; ++i) {
    for (int j = 0; j < serverNum_; ++j) {
        int value1 = localTokenServerOffsetTensor.GetValue(i * serverNum_ + j);
        int value2 = tempServerTensor.GetValue(j);
        localTokenServerOffsetTensor.SetValue(i * serverNum_ + j, value1 * value2 + value1);
        tempServerTensor.SetValue(j, value2 + value1);
    }
    // 向量乘法 + 加法计算 sendTokenIdx
    AscendC::Mul(sendTokenIdxTensor[i * numExperts_], prefixCountPerExpertTensor, intermediateExpertTensor, numExperts_);
    AscendC::Add(sendTokenIdxTensor[i * numExperts_], tempExpertTensor, intermediateExpertTensor, numExperts_);
    AscendC::Add(prefixCountPerExpertTensor, intermediateExpertTensor, tempExpertTensor, numExperts_);
}

3.6 入口函数 (op_kernel/dispatch_layout.cpp)

#define TILING_KEY_INT 23
#define TILING_KEY_A2_INT 123

extern "C" __global__ __aicore__ void dispatch_layout(...) {
    if (TILING_KEY_IS(TILING_KEY_INT)) {
        MoeDispatchLayout::DispatchLayout<int32_t> op;
        op.Init(...);
        op.Process();
    } else if (TILING_KEY_IS(TILING_KEY_A2_INT)) {
        MoeDispatchLayoutA2::DispatchLayoutA2<int32_t> op;
        op.Init(...);
        op.Process();
    }
}

3.7 aclnn 接口层

// 注意：设置 HCCL Server Type 为 MTE（而非 AICPU）
aclnnStatus aclnnDispatchLayoutZeroBuffer(void *workspace, uint64_t workspaceSize,
    aclOpExecutor *executor, aclrtStream stream)
{
    if (NnopbaseSetHcclServerType) {
        NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
    }
    return aclnnInnerDispatchLayout(workspace, workspaceSize, executor, stream);
}

关键设计: NnopbaseSetHcclServerType(executor, MTE) 将算子执行模式设置为 MTE（Memory Transfer Engine），而非默认的 AICPU。这确保算子通过硬件 MTE 通道执行，减少 CPU 干预。

4. moe_dispatch_normal_zb — 标准分发算子

4.1 算子概述

moe_dispatch_normal_zb 是 MoE Dispatch 流程的第二阶段（数据搬运阶段），负责： 1. 从 layout 算子获取 put_offset（目标偏移）和 send_tokenIdx（发送序号） 2. 通过 SHMEM RDMA 将 token 数据直接写入对端 rank 的 expandXOut 缓冲区 3. 支持可选的 DynamicQuant（INT8 动态量化）

关键特征: 使用 ZeroBufferSyncFlag 进行跨卡同步，使用 shmem_ptr() 获取远端指针。

4.2 算子注册 (op_host/moe_dispatch_normal_zero_buffer.cpp)

// 输入
this->Input("x")              // DT_BF16/DT_FLOAT16, [bs, h]
this->Input("topk_idx")       // DT_INT32, [bs, k]
this->Input("send_tokenIdx")  // DT_INT32, [bs, k]
this->Input("put_offset")     // DT_INT32, [moeExpertNum * epWorldSize]

// 输出
this->Output("recv_x")                   // DT_BF16/DT_INT8/DT_FLOAT16
this->Output("recv_x_scale")             // DT_FLOAT (量化时)
this->Output("assist_info_for_combine")  // DT_INT32
this->Output("waitRecvCostStats")        // DT_INT32 (诊断用)

// 属性
this->Attr("ep_world_size").Int()
this->Attr("ep_rank_id").Int()
this->Attr("tp_world_size").Int()
this->Attr("tp_rank_id").Int()
this->Attr("moe_expert_num").Int()
this->Attr("quant_mode").Int()
this->Attr("global_bs").Int()
this->Attr("zero_buffer_ptr").Int()  // ← ZB 核心：对称内存基地址

4.3 Tiling 数据

struct MoeDispatchNormalZeroBufferInfo {
    uint32_t epWorldSize;
    uint32_t tpWorldSize;
    uint32_t epRankId;
    uint32_t tpRankId;
    uint32_t moeExpertNum;
    uint32_t quantMode;
    uint32_t globalBs;
    uint32_t bs;
    uint32_t k;
    uint32_t h;
    uint32_t aivNum;
    bool isQuant;
    bool isEnableDiagnose;
    uint64_t totalUbSize;
    uint64_t totalWinSize;
};

struct MoeDispatchNormalZeroBufferTilingData {
    MoeDispatchNormalZeroBufferInfo moeDispatchNormalInfo;
    uint64_t zeroBufferPtr;  // ← ZB 对称内存基地址
};

4.4 Kernel 核心实现

模板参数

template <typename XType, typename ExpandXOutType, bool DynamicQuant,
          bool IsSmoothScaleExist, bool IsShareExpertRank>
class MoeDispatchNormalZeroBuffer {
    constexpr static int32_t PHASE_ENTRY = 1;
    constexpr static int32_t PHASE_DONE  = 2;
    // ...
};

内存布局常量

constexpr uint64_t NOTIFY_MAGIC_OFFSET = 50UL * 1024UL;          // 50KB
constexpr uint64_t WIN_MAGIC_OFFSET = 100UL * 1024UL;            // 100KB
constexpr uint64_t HALF_WIN_STATE_OFFSET = 8 * 1024UL * 1024UL;  // 8MB
constexpr uint64_t NOTIFY_WIN_STATE_OFFSET = 2 * 1024UL * 1024UL;// 2MB

远端地址计算

__aicore__ inline GM_ADDR GetWindStateAddrByRankId(const int32_t rankId)
{
    auto ptr = shmem_ptr(gva_gm, rankId);  // 获取对端 rank 的对称内存指针
    return (GM_ADDR)(ptr) + WIN_MAGIC_OFFSET + NOTIFY_WIN_STATE_OFFSET + dataState * HALF_WIN_STATE_OFFSET;
}

核心执行流程 (Process)

template <CamTypeClass>
__aicore__ inline void MoeDispatchNormalZeroBuffer<CamTypeFunc>::Process()
{
    if ASCEND_IS_AIV {
        // Step 1: IncrementMagic — 获取 magic M+1
        syncFlag_.IncrementMagic();

        // Step 2: BarrierAll — 确保所有 rank 的 Notify 结果可见
        //         (putOffset, recvData 等已被 Notify 计算完成)
        syncFlag_.BarrierAll();

        // Step 3: 通过 shmem 写 token 数据到远端 rank
        InputToDstOutput();

        // Step 4: BarrierAll — 确保所有 rank 完成写入
        //         后续 GMM kernel 可以安全消费 expandXOut
        syncFlag_.BarrierAll();
    }
}

数据搬运核心 (InputToDstOutput)

for (int32_t tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
    uint32_t dstExpertId = expertIdsTensor(tokenIndex - startTokenId);
    uint32_t dstRankId = dstExpertId / moeExpertNumPerRank;
    int32_t curExpertIdx = sendTokenIdxTensor(tokenIndex - startTokenId);
    int32_t dstExpertOffset = putOffsetTensor(dstExpertId * epWorldSize + epRankId);

    // ★ 关键：通过 shmem_ptr 获取对端 rank 的远端内存指针
    auto ptr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(expandXOutGM, dstRankId));
    uint64_t expandXByteOffset = static_cast<uint64_t>(hUBAlignSize) *
        static_cast<uint64_t>(dstExpertOffset + curExpertIdx);
    dstGT.SetGlobalBuffer((__gm__ ExpandXOutType *)(ptr + expandXByteOffset));

    // 写入数据（支持量化路径）
    if constexpr (DynamicQuant) {
        // INT8 量化路径：量化后写入 + 写入 scale
        QuantProcess();
        DataCopyPad(dstGT, xOutTensor, xOutCopyParams);
        DataCopyPad(dstScaleOutGT, xOutFp32Tensor[hUBAlignSize / sizeof(float)], scaleCopyParams);
    } else {
        // BF16/FP16 直接写入
        DataCopyPad(dstGT, xTmpTensor, xOutCopyParams);
    }
}

量化处理 (QuantProcess)

// DynamicQuant: BF16/FP16 → INT8 动态量化
Cast(floatLocalTemp, xInTensor, RoundMode::CAST_NONE, h);  // 转 float
Abs(floatLocalAbsTemp, floatLocalTemp, h);                  // 取绝对值
ReduceMaxInplace(floatLocalAbsTemp, h);                     // 求 max
dynamicScale = float(127.0) / (floatLocalAbsTemp.GetValue(0) + 1e-12f);
Muls(floatLocalTemp, floatLocalTemp, dynamicScale, h);      // 缩放
Cast(int32LocalTemp, floatLocalTemp, RoundMode::CAST_RINT, h);  // 转 int32
Cast(halfLocalTemp, int32LocalTemp, RoundMode::CAST_ROUND, h);  // 转 half
Cast(xOutTensor, halfLocalTemp, RoundMode::CAST_TRUNC, h);      // 转 int8

4.5 ZeroBufferSyncFlag 初始化

// 每个 core 的 syncFlagBuf 仅 32B（一个 flag slot 大小）
tpipe_->InitBuffer(syncFlagBuf_, ZeroBufferSyncFlagImpl::FLAG_SLOT_SIZE);

// per-core 粒度：slotsPerRank = blockNum
syncFlag_.Init(gva_gm, epRankId, epRankSize, blockNum, syncFlagBuf_);

4.6 与 master 版 moe_dispatch_normal 的差异

维度	master 版 (HCCL Window)	ZB 版 (Zero-Buffer)
远端访问	通过 HCCL Window 间接访问	`shmem_ptr()` 直接访问对端内存
同步	HCCL Collective + 信号量	`ZeroBufferSyncFlag` + Magic
内存开销	~200MB HCCL Buffer	零额外内存
数据路径	源 → Window → RDMA → 对端 Window → 目标	源 → RDMA → 对端目标（直达）
量化支持	无	DynamicQuant (INT8)
TP 支持	无	有 (tpWorldSize/tpRankId)

5. moe_dispatch_lowlatency_zb — 低延迟分发算子

5.1 算子概述

moe_dispatch_lowlatency_zb 是 ZB 系列中最复杂的算子（869 行 header），专为低延迟 MoE Dispatch 设计。它整合了 Notify + Dispatch 的功能，支持： 1. Shared Expert + MoE Expert 混合路由 2. 多种量化模式（StaticQuant / DynamicQuant / SmoothScale） 3. AllGather 模式（IsNeedAllgather） 4. 弹性专家（zero/copy/const expert）

5.2 算子注册 (op_host/moe_dispatch_lowlatency_zero_buffer.cpp)

// 输入
this->Input("x")              // DT_BF16/DT_FLOAT16, [bs, h]
this->Input("expert_ids")     // DT_INT32, [bs, k]
this->Input("scales")         // DT_FLOAT, optional (静态量化)
this->Input("x_active_mask")  // DT_BOOL, optional
this->Input("elastic_info")   // DT_INT32, optional

// 输出
this->Output("expand_x")                // DT_BF16/DT_INT8/DT_FLOAT16
this->Output("dynamic_scales")          // DT_FLOAT
this->Output("assist_info_for_combine") // DT_INT32 (expandIdx)
this->Output("expert_token_nums")       // DT_INT64
this->Output("ep_send_counts")          // DT_INT32
this->Output("tp_send_counts")          // DT_INT32

// 属性
this->Attr("ep_world_size").Int()
this->Attr("ep_rank_id").Int()
this->Attr("moe_expert_num").Int()
this->Attr("expert_shard_type").Int()
this->Attr("shared_expert_num").Int()
this->Attr("shared_expert_rank_num").Int()
this->Attr("quant_mode").Int()
this->Attr("expert_token_nums_type").Int()
this->Attr("zero_expert_num").Int()
this->Attr("copy_expert_num").Int()
this->Attr("const_expert_num").Int()

5.3 Tiling Key 矩阵

低延迟算子支持多种 Tiling Key，覆盖不同量化/AllGather 组合：

Tiling Key	StaticQuant	DynamicQuant	SmoothScale	AllGather
10000	false	false	false	false
10100	false	false	false	true
10011	true	false	false	false
10002	false	true	false	false
10012	false	true	true	false
10111	true	false	false	true
10102	false	true	false	true
10112	false	true	true	true

5.4 Kernel 核心实现

模板参数

template <typename XType, typename ExpandXOutType, bool StaticQuant,
          bool DynamicQuant, bool IsSmoothScaleExist, bool IsNeedAllgather>
class MoeDispatchLowlatencyZeroBuffer {
    // ...
};

内存布局常量

// Dispatch 同步 flag 区域从 FLAG_AREA_BASE + 64KB 开始
// 原因：某些 rank 可能已进入 combine 阶段，而其他 rank 仍在 dispatch
// 使用独立 flag slot 避免与 cleanup 竞争
constexpr uint64_t DISPATCH_SYNC_FLAG_BASE_OFFSET = ZeroBufferSyncFlagImpl::FLAG_AREA_BASE + 64UL * 1024UL;
constexpr uint32_t DISPATCH_SYNC_STATE_OFFSET = 32U;

元数据布局

metaDataGvaGM_ 布局:
┌──────────────────────────────────────────────────────────┐
│ Per-core magic value (aivNum × 32B)                      │
├──────────────────────────────────────────────────────────┤
│ Per-rank layout flags (coreNum × 32B)                    │
├──────────────────────────────────────────────────────────┤
│ Per-core send-count flags (coreNum × 64B)                │
├──────────────────────────────────────────────────────────┤
│ Send-count data (LocalNotifyDataTensor_)                 │
└──────────────────────────────────────────────────────────┘

核心执行流程 (Process)

template <TemplateMC2TypeClass>
__aicore__ inline void MoeDispatchLowlatencyZeroBuffer<TemplateMC2TypeFunc>::Process()
{
    if ASCEND_IS_AIV {
        // Phase 1: 本地统计每个 expert 的 token 发送数量
        LocalLayout();
        SyncAll<true>();

        // Phase 2: 通知所有 rank 本地统计完成
        SetLayoutStatus();
        SendCountNotify();  // 等待所有 rank 的统计结果
        SyncAll<true>();

        // Phase 3: 处理接收到的全局统计信息
        WaitNotify();
        SyncAll<true>();

        // Phase 4: 实际数据搬运
        InputToDstOutput();

        // Phase 5: 同步确保数据写入完成
        SetSyncFlag();
        WaitSyncFlag();
        SyncAll<true>();

        // Phase 6: 清理
        CleanUp();
    }
}

Phase 1: LocalLayout — 本地统计

__aicore__ inline void LocalLayout()
{
    // 1. 加载 expert_ids
    DataCopyPad(expertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, expertIdsCntCopyPadParams);

    // 2. 按 core 分片统计每个 expert 的 token 数
    for (int32_t i = startTokenId; i < endTokenId; ++i) {
        int32_t expertIdx = static_cast<int32_t>(expertIdsTensor_(i));
        int32_t curCnt = sendCountLt.GetValue(expertIdx) + 1;
        sendCountLt.SetValue(expertIdx, curCnt);
    }

    // 3. 原子写入共享元数据区域
    AscendC::SetAtomicAdd<int32_t>();
    DataCopyPad(LocalNotifyDataTensor_, sendCountLt, sendCountDataCopyParams);
    AscendC::SetAtomicNone();
}

Phase 2: SetLayoutStatus — 通知统计完成

__aicore__ inline void SetLayoutStatus()
{
    // 每个 AIV 负责通知一部分 rank
    for (uint32_t targetRankId = startRankId; targetRankId < endRankId; targetRankId++) {
        // ★ 使用 aclshmem_float_put_nbi 写入对端 rank 的 layoutStatus
        aclshmem_float_put_nbi(layoutStatusTensor_[epRankId_], statusTensor_, 1, targetRankId);
    }
}

Phase 3: SendCountNotify — 等待并获取对端统计

__aicore__ inline void SendCountNotify()
{
    for (uint32_t targetRankId = startRankId; targetRankId < endRankId; targetRankId++) {
        // 1. 轮询等待对端 rank 的 layout flag 匹配当前 magic
        while ((curVal < minFlagVal) || (curVal > maxFlagVal)) {
            DataCopy(layoutStatusFp32Tensor_, layoutStatusTensor_[targetRankId], intriParams);
            curVal = layoutStatusFp32Tensor_.GetValue(0);
        }

        // 2. ★ 使用 shmem_get_int32_mem_nbi 从对端 rank 读取 send count
        shmem_get_int32_mem_nbi(allExpertTokenNumsGMTensor_[targetRankId * moeExpertNum_],
            LocalNotifyDataTensor_, moeExpertNum_, targetRankId);
    }
}

Phase 4: InputToDstOutput — 数据搬运

__aicore__ inline void SendToMoeExpert()
{
    for (int32_t index = startTokenId; index < endTokenId; ++index) {
        int32_t tokenIndex = index / axisK_;
        uint32_t dstExpertId = expertIdsTensor_(index);
        int32_t dstRankId = dstExpertId / moeExpertNumPerRank_ + sharedExpertRankNum_;

        // 计算当前 token 在目标 expert 内的序号
        CalTokenSendExpertCnt(dstExpertId, index, curExpertIdx);

        // 计算目标偏移
        int32_t dstExpertOffset = sendCountsGlobal(offsetIdx - 1);

        // ★ 通过 shmem_ptr 获取对端 rank 的远端内存指针
        auto dstPtr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(expandXOutGM_, dstRankId));
        dstWinGMTensor.SetGlobalBuffer(
            (__gm__ ExpandXOutType *)(dstPtr + hOutSizeAlign_ * (dstExpertOffset + curExpertIdx)));

        // 写入数据
        if constexpr (DynamicQuant || StaticQuant) {
            // 量化路径
            QuantProcess();
            DataCopyPad(dstWinGMTensor, xOutTensor_, expandXCopyParams_);
            DataCopyPad(dstScaleGMTensor, xOutFp32Tensor[hOutSizeAlign_ / sizeof(float)], scaleCopyParams_);
        } else {
            // 直接写入
            DataCopyPad(dstWinGMTensor, xTmpTensor_, expandXCopyParams_);
        }
    }
}

Phase 5: SetSyncFlag / WaitSyncFlag — 数据写入同步

__aicore__ inline void SetSyncFlag()
{
    // 每个 core 向负责的 rank 写入 flag
    for (uint32_t targetRankId = curBlockStartRankId_; targetRankId < curBlockEndRankId_; ++targetRankId) {
        auto ptr = GetSyncFlagAddrByRankId(targetRankId) + epRankId_ * DISPATCH_SYNC_STATE_OFFSET;
        gmRemoteStatusGT.SetGlobalBuffer((__gm__ int32_t *)(ptr));
        DataCopy<int32_t>(gmRemoteStatusGT, statusTensor[...], 8UL);
    }
}

__aicore__ inline void WaitSyncFlag()
{
    // 轮询等待所有对端 rank 的 flag
    while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
        DataCopy(statusFp32Tensor, statusFp32TensorGT[...], intriParams);
        ReduceSum(statusSumOutTensor, statusFp32Tensor, gatherMaskOutTensor, mask, rankNumPerBlock_, 1);
        sumOfFlag = statusSumOutTensor.GetValue(0);
    }
    // 清理 flag
    Duplicate<int32_t>(cleanStateTensor, 0, ...);
    DataCopy(statusFp32TensorGT[...], cleanStateTensor.ReinterpretCast<float>(), intriOutParams);
}

5.5 AIV 分配策略

低延迟算子将 AIV 分为两组：

// Shared Expert 使用的 AIV 数量
if (sharedExpertRankNum_ != 0U) {
    sharedUsedAivNum_ = (aivNum_ * sharedExpertNum) / (axisK_ + sharedExpertNum);
    if (sharedUsedAivNum_ == 0) { sharedUsedAivNum_ = 1; }
}
// MoE Expert 使用的 AIV 数量
moeUsedAivNum_ = aivNum_ - sharedUsedAivNum_;

aivId < moeUsedAivNum_: 处理 MoE Expert 数据搬运
aivId >= moeUsedAivNum_: 处理 Shared Expert 数据搬运

5.6 与 normal 版的差异

维度	normal_zb	lowlatency_zb
Notify 阶段	独立算子 (moe_notify_dispatch_zb)	内联在算子中
同步机制	ZeroBufferSyncFlag	SyncCollectives + 自定义 flag
Shared Expert	不支持	支持
量化	DynamicQuant	StaticQuant + DynamicQuant + SmoothScale
AllGather	不支持	支持 (IsNeedAllgather)
AIV 分配	全部用于数据搬运	分 shared/moe 两组
执行阶段	2 步 (BarrierAll → 写数据 → BarrierAll)	6 步

6. moe_notify_dispatch_zb — 元数据通知算子

6.1 算子概述

moe_notify_dispatch_zb 是 MoE Dispatch 流程的元数据交换阶段，负责： 1. AllGather 每个 rank 的 tokenPerExpertData（每个 expert 的 token 数） 2. 计算 recvData（全局 rank × expert 矩阵） 3. 计算 totalRecvTokens（最大接收 token 总数） 4. 计算 maxBs（最大 batch size） 5. 计算 recvTokensPerExpert（每个 expert 接收的 token 数） 6. 计算 putOffset（前缀和偏移表）

6.2 算子注册 (op_host/notify_dispatch_zero_buffer.cpp)

// 输入
this->Input("tokenPerExpertData")  // DT_FLOAT16/DT_FLOAT/DT_INT32

// 输出
this->Output("recvData")             // 全局 rank × expert 矩阵
this->Output("totalRecvTokens")      // 最大接收 token 总数
this->Output("maxBs")                // 最大 batch size
this->Output("recvTokensPerExpert")  // DT_INT64, 每个 expert 接收的 token 数
this->Output("putOffset")            // DT_INT32, 前缀和偏移表

// 属性
this->Attr("send_count").Int()
this->Attr("rank_size").Int()
this->Attr("rank_id").Int()
this->Attr("local_rank_size").Int()
this->Attr("local_rank_id").Int()
this->Attr("topk_num").Int()
this->Attr("zero_buffer_ptr").Int()  // ← ZB 对称内存基地址

6.3 Tiling 数据

struct NotifyDispatchZeroBufferInfo {
    uint32_t rankSize;
    uint32_t rankId;
    uint32_t localRankSize;
    uint32_t localRankId;
    uint32_t sendCount;    // = numExperts
    uint32_t topkNum;
    uint32_t aivNum;
    uint64_t totalUbSize;
};

struct NotifyDispatchZeroBufferTilingData {
    NotifyDispatchZeroBufferInfo notifyDispatchInfo;
    uint64_t zeroBufferPtr;  // ← ZB 对称内存基地址
};

6.4 Kernel 核心实现

核心执行流程 (Process)

template <typename T>
__aicore__ inline void NotifyDispatchZeroBuffer<T>::Process()
{
    // Step 1: IncrementMagic — 获取 magic M
    syncFlag_.IncrementMagic();

    // Step 2: BarrierAll — 确保所有 rank 输入就绪
    syncFlag_.BarrierAll();

    // Step 3: AllGather 每个 rank 的 tokenPerExpertData
    AllGatherSendData();
    SyncAll<true>();

    // Step 4: 后处理（本地计算，无跨卡依赖）
    ReloadRecvData();
    BuildTotalRecvCount();
    BuildMaxBs();
    BuildTotalRecvTokens();
    BuildRecvTokenPerExp();
    SyncAll<true>();
}

AllGatherSendData — 跨卡数据收集

__aicore__ inline void AllGatherSendData()
{
    // 按 core 分片负责一部分 rank
    SplitCoreCal(epWorldSize_, rankNumPerBlock, startRankId, endRankId);

    for (uint32_t targetRankId = startRankId; targetRankId < endRankId; targetRankId++) {
        // ★ 通过 shmem_ptr 获取对端 rank 的 tokenPerExpertData
        auto ptr = shmem_ptr(tokenPerExpertData_, targetRankId);
        gmRemoteDataGt.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(ptr));

        // GM→GM 拷贝（支持类型转换）
        CpGM2GMPingPong<int32_t>(numExperts * sizeof(int32_t), gmRemoteDataGt,
            recvDataGt_[targetRankId * numExperts], COPYONLY);
    }
}

CpGM2GMPingPong — Ping-Pong GM→GM 拷贝

template <typename K, typename U>
__aicore__ inline void NotifyDispatchZeroBuffer<T>::CpGM2GMPingPong(
    int64_t dataSizeRemain, const GlobalTensor<U> &sendDataInputGt,
    const GlobalTensor<K> &recvDataOutputGT, int op)
{
    // 双 buffer ping-pong: UB_HEAD_OFFSET 和 UB_MID_OFFSET
    __ubuf__ U *inputUB[2] = {UB_HEAD_OFFSET, UB_MID_OFFSET};
    __ubuf__ K *outputUB[2] = {inputUB[0], inputUB[1]};

    SetAtomic<K>(op);
    for (int64_t i = 0; dataSizeRemain > 0; i++) {
        event_t eventId = (i & 1) ? EVENT_ID0 : EVENT_ID1;
        CpGM2UB(inputUB[i&1], input + inputOffsetNum, size);  // GM→UB
        // 类型转换（如需要）
        CastImpl(outputUB[i&1], inputUB[i&1], RoundMode::CAST_NONE, size / sizeof(K));
        CpUB2GM(output + outputOffsetNum, outputUB[i&1], size);  // UB→GM
        dataSizeRemain -= size;
    }
    UnsetAtomic(op);
}

BuildMaxBs — 计算最大 batch size

__aicore__ inline void BuildMaxBs()
{
    int32_t maxBsNum = 0;
    for (uint32_t srcRankId = 0; srcRankId < epWorldSize_; srcRankId++) {
        DataCopy(tokenPerExpertDataLt, recvDataTensor_[numExperts * srcRankId], numExperts);
        Cast(floatExpTokenCntLt, tokenPerExpertDataLt, RoundMode::CAST_NONE, numExperts);
        ReduceSum(floatExpTokenSumCntLt, floatExpTokenCntLt, sharedTmpBuffer, numExperts);
        int32_t curRankBsNum = static_cast<int32_t>(floatExpTokenSumCntLt(0));
        maxBsNum = curRankBsNum > maxBsNum ? curRankBsNum : maxBsNum;
    }
    maxBsGt.SetValue(0, maxBsNum / topkNum_);
}

BuildTotalRecvCount — 计算前缀和偏移表

__aicore__ inline void BuildTotalRecvCount()
{
    for (uint32_t rank = startRankId; rank < endRankId; ++rank) {
        // 转置并计算前缀和
        ReorderRecvDataOutput(rank, recvTokenLt, true);  // isCumSum = true
        DataCopyPad(recvCntGt[rank * numExperts], recvTokenLt, copyParams);
    }
}

ReorderRecvDataOutput — 矩阵转置 + 前缀和

__aicore__ inline void ReorderRecvDataOutput(int32_t rankId, LocalTensor<int32_t> &transLt, bool isCumSum)
{
    uint32_t moeExpertPerRankNum = numExperts / epWorldSize_;
    uint32_t startExpId = rankId * moeExpertPerRankNum;
    uint32_t endExpId = rankId * moeExpertPerRankNum + moeExpertPerRankNum;

    int32_t prefixSum = 0;
    for (uint32_t expId = startExpId; expId < endExpId; ++expId) {
        for (uint32_t srcRank = 0; srcRank < epWorldSize_; ++srcRank) {
            uint32_t index = (expId - startExpId) * epWorldSize_ + srcRank;
            uint32_t pairIdx = srcRank * numExperts + expId;
            int32_t curRecvCount = recvDataTensor_(pairIdx);
            transLt(index) = isCumSum ? prefixSum : curRecvCount;
            prefixSum += curRecvCount;
        }
    }
}

6.5 远端地址计算

template <typename T>
__aicore__ inline GM_ADDR NotifyDispatchZeroBuffer<T>::GetWindStateAddrByRankId(const int32_t rankId)
{
    auto ptr = shmem_ptr((__gm__ T *)gva_gm, rankId);
    return (GM_ADDR)(ptr) + WIN_MAGIC_OFFSET + bufferId_ * HALF_WIN_STATE_OFFSET;
}

6.6 ZeroBufferSyncFlag 初始化

// 使用 8KB 的 flag 缓冲区
constexpr static uint32_t UB_FLAG_SIZE = 8U * 1024U;
pipe_.InitBuffer(tBuf, UB_FLAG_SIZE);

// per-core 粒度：slotsPerRank = blockNum
syncFlag_.Init(gva_gm, epRankId_, epWorldSize_, blockNum_, tBuf);

7. SHMEM 集成方式

7.1 SHMEM API 调用汇总

ZB 系列算子使用以下 SHMEM API（来自 CANN shmem.h）：

API	功能	使用算子
`aclshmem_ptr(base, rankId)`	获取对端 rank 的远端内存指针	全部
`aclshmem_half_put_nbi(dst, src, count, dstRank)`	half 类型非阻塞 RDMA put	lowlatency_zb
`aclshmem_bfloat16_put_nbi(dst, src, count, dstRank)`	bf16 类型非阻塞 RDMA put	lowlatency_zb
`aclshmem_float_put_nbi(dst, src, count, dstRank)`	float 类型非阻塞 RDMA put	lowlatency_zb
`shmem_get_int32_mem_nbi(dst, src, count, srcRank)`	int32 类型非阻塞 RDMA get	lowlatency_zb

7.2 类型分发宏 (lowlatency_zb)

#define ZERO_BUFFER_PUT_BY_DTYPE(dtype, ...)                            \
    do {                                                          \
        if constexpr (std::is_same_v<dtype, half>) {              \
            aclshmem_half_put_nbi(__VA_ARGS__);                  \
        } else if constexpr (std::is_same_v<dtype, bfloat16_t>) { \
            aclshmem_bfloat16_put_nbi(__VA_ARGS__);              \
        } else if constexpr (std::is_same_v<dtype, float>) {      \
            aclshmem_float_put_nbi(__VA_ARGS__);                 \
        }                                                         \
    } while (0)

7.3 远端指针获取模式

所有 ZB 算子使用相同的模式获取远端指针：

// 1. 获取对端 rank 的对称内存基地址
auto ptr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(baseAddr, dstRankId));

// 2. 计算偏移量
uint64_t offset = ...;  // 根据 expertId、tokenIdx 等计算

// 3. 设置 GlobalTensor 指向远端地址
dstGT.SetGlobalBuffer((__gm__ DataType *)(ptr + offset));

// 4. 使用 DataCopyPad 进行 RDMA 写入
DataCopyPad(dstGT, srcTensor, copyParams);

7.4 同步原语

ZB 算子使用两种同步原语：

卡内同步: SyncAll<true>() — 确保所有 AIV core 完成当前阶段
跨卡同步: ZeroBufferSyncFlag — 基于共享内存 flag 的轻量级跨卡同步

8. 性能优势分析

8.1 内存占用对比

组件	HCCL Window 模式	Zero-Buffer 模式
HCCL Buffer	~200MB/卡	0
Flag 区域	N/A	~100KB/卡 (Magic + Barrier + Flag)
通信窗口	预分配固定大小	按需使用，无固定窗口
总计	~200MB/卡	~100KB/卡

内存节省: 每卡节省约 200MB，对于大规模 EP（如 256 卡）集群，总计节省约 50GB 内存。

8.2 通信延迟对比

阶段	HCCL Window 模式	Zero-Buffer 模式
数据写入	用户 Buffer → HCCL Window	用户 Buffer → RDMA → 对端用户 Buffer
同步信号	HCCL Collective	SHMEM flag 轮询
数据读取	对端 HCCL Window → 对端用户 Buffer	直接读取（已在用户 Buffer）

延迟降低原因: 1. 消除中间拷贝: 数据直接从源写入目标，无需经过 Window 中转 2. 用户态同步: SHMEM flag 轮询避免 HCCL Collective 的软件栈开销 3. 硬件直通: RDMA 引擎直接访问用户内存，无需 CPU 干预

8.3 同步机制优势

Magic + Ping-Pong 机制: - 每个 kernel 调用获取唯一 magic 值 - magic % 2 选择 ping-pong buffer，避免与上一次调用的 flag 冲突 - 单调递增 magic 保证 stale flag 被自动拒绝 - 无需显式清理 flag 区域（ping-pong 自动覆盖）

BarrierAll 协议:

SyncAll (卡内) → Write flags to all ranks → Wait all ranks' flags → SyncAll (卡内)

卡内同步确保所有 core 完成计算
分 core 策略：每个 core 负责一部分 rank，避免热点
无环依赖保证：DAG 结构，不会死锁

8.4 量化加速

normal_zb 和 lowlatency_zb 支持 INT8 动态量化： - 将 BF16/FP16 数据量化为 INT8 传输，带宽需求减半 - 在目标卡上反量化恢复精度 - 量化/反量化使用向量单元并行处理

8.5 AIV 并行策略

ZB 算子充分利用 Ascend 910 的 AIV（AI Vector）架构： - 最多 48 个 AIV core 并行工作 - 每个 core 负责一部分 token 或一部分 rank - 分片策略：SplitToCore() / SplitCoreCal() 均匀分配负载

9. 关键数据结构与常量汇总

9.1 共享常量

// zero_buffer_sync_flag.h
constexpr uint32_t FLAG_SLOT_SIZE = 32U;           // 每个 flag slot 32B
constexpr uint32_t FLAG_ELEM_NUM = 4;              // 4 × int64_t
constexpr uint64_t FLAG_AREA_BASE = 100UL * 1024UL; // Flag 区域起始偏移 100KB
constexpr uint32_t MAGIC_SLOT_ALIGN = 512U;        // 每个 core 的 magic slot 512B
constexpr uint32_t MAX_CORE_NUM_SYNC = 48U;        // 最大 48 个 core
constexpr uint64_t MAGIC_AREA_SIZE = 48UL * 512UL; // Magic 区域 24KB

// comm_args.h
constexpr int CAM_MAX_RANK_SIZE = 384;             // 最大 384 卡
constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024;  // UB 单次 DMA 最大 190KB
constexpr int UB_ALIGN_SIZE = 32;                  // UB 32B 对齐

9.2 算子特定常量

算子	关键常量	值
layout_zb	UB_MAX_SIZE	190KB
layout_zb	TILING_KEY_INT / TILING_KEY_A2_INT	23 / 123
normal_zb	NOTIFY_MAGIC_OFFSET	50KB
normal_zb	WIN_MAGIC_OFFSET	100KB
normal_zb	HALF_WIN_STATE_OFFSET	8MB
normal_zb	NOTIFY_WIN_STATE_OFFSET	2MB
lowlatency_zb	DISPATCH_SYNC_FLAG_BASE_OFFSET	FLAG_AREA_BASE + 64KB
lowlatency_zb	DISPATCH_SYNC_STATE_OFFSET	32B
notify_dispatch_zb	WIN_MAGIC_OFFSET	100KB
notify_dispatch_zb	HALF_WIN_STATE_OFFSET	8MB
notify_dispatch_zb	UB_FLAG_SIZE	8KB

9.3 工作空间需求

算子	SYSTEM	KERNEL	A2_ARGS	总计
layout_zb	16MB	1MB	1MB	18MB
normal_zb	-	-	-	由 tiling 计算
lowlatency_zb	-	-	-	由 tiling 计算
notify_dispatch_zb	16MB	1MB	1MB	18MB

9.4 文件路径汇总

src/cam/comm_operator/ascend_kernels/
├── moe_dispatch_layout_zb/
│   ├── op_api/aclnn_dispatch_layout_zero_buffer.{cpp,h}
│   ├── op_host/dispatch_layout.{cpp, tiling.cc}
│   └── op_kernel/dispatch_layout.{cpp,h, tiling.h, a2.h}
├── moe_dispatch_normal_zb/
│   ├── op_api/aclnn_moe_dispatch_normal_zero_buffer.{cpp,h}
│   ├── op_host/moe_dispatch_normal_zero_buffer.{cpp, tiling.cc}
│   └── op_kernel/moe_dispatch_normal_zero_buffer.{cpp,h, tiling.h}
├── moe_dispatch_lowlatency_zb/
│   ├── op_api/aclnn_moe_dispatch_lowlatency_zero_buffer.{cpp,h}
│   ├── op_host/moe_dispatch_lowlatency_zero_buffer.{cpp, tiling.cpp}
│   └── op_kernel/moe_dispatch_lowlatency_zero_buffer.{cpp,h, tiling.h}
├── moe_notify_dispatch_zb/
│   ├── op_api/aclnn_notify_dispatch_zero_buffer.{cpp,h}
│   ├── op_host/notify_dispatch_zero_buffer.{cpp, tiling.cc}
│   └── op_kernel/notify_dispatch_zero_buffer.{cpp,h, tiling.h}
└── utils/op_kernel/
    ├── zero_buffer_api.h
    ├── zero_buffer_sync_flag.h
    ├── comm_args_zb.h
    ├── data_copy_zb.h
    ├── moe_lowlatency_zero_buffer_base.h
    └── sync_collectives_zb.h

附录：ZB Dispatch 完整执行流程

┌─────────────────────────────────────────────────────────────────────────┐
│                    ZB MoE Dispatch Pipeline                              │
│                                                                          │
│  ┌──────────────────┐    ┌──────────────────┐    ┌──────────────────┐   │
│  │  layout_zb       │    │  notify_dispatch │    │  normal_zb /     │   │
│  │  (本地计算)       │    │  _zb (元数据交换) │    │  lowlatency_zb   │   │
│  │                  │    │                  │    │  (数据搬运)       │   │
│  │  1. 统计token数   │───▶│  1. AllGather    │───▶│  1. RDMA写入     │   │
│  │  2. 计算isToken   │    │  2. 计算偏移表    │    │  2. 跨卡同步     │   │
│  │  3. 计算sendIdx   │    │  3. 计算maxBs    │    │  3. 清理        │   │
│  └──────────────────┘    └──────────────────┘    └──────────────────┘   │
│           │                       │                       │              │
│           ▼                       ▼                       ▼              │
│  numTokensPerRank        recvData                 expandXOut            │
│  numTokensPerExpert      totalRecvTokens          dynamicScalesOut      │
│  isTokenInRank           maxBs                    assistInfoForCombine   │
│  notifySendData          recvTokensPerExpert      expertTokenNums        │
│  sendTokenIdxSmall       putOffset                epSendCounts           │
│                                                                          │
│  ★ 所有跨卡通信通过 SHMEM RDMA 完成                                       │
│  ★ 所有跨卡同步通过 ZeroBufferSyncFlag 完成                                │
│  ★ 零额外内存开销（无需 HCCL Window）                                      │
└─────────────────────────────────────────────────────────────────────────┘

分析说明: 本报告基于 remotes/origin/cam_feature 分支的实际代码分析，所有函数名、常量值、数据结构均来自源码。SHMEM API (shmem.h) 为 CANN 工具链外部头文件，其具体实现未在 CAM 仓库中。

← FusedDeepMoE SHMEM 路径 📑 索引 ZB Combine 算子族 →

ZB Dispatch 算子族 — UMDK/CAM

ZB (Zero-Buffer) Dispatch 算子深度分析

目录

1. ZB (Zero-Buffer) 设计哲学

1.1 什么是 Zero-Buffer？

1.2 传统 HCCL Window 模式的问题

1.3 Zero-Buffer 的核心优势

1.4 代码组织

2. 共享基础设施分析

2.1 zero_buffer_sync_flag.h — 核心同步机制

内存布局

Magic 机制

同步协议

同步粒度

2.2 zero_buffer_api.h — SHMEM API 封装

2.3 sync_collectives_zb.h — 旧版集合同步

2.4 moe_lowlatency_zero_buffer_base.h — 低延迟基类

3. moe_dispatch_layout_zb — 布局算子

3.1 算子概述

3.2 算子注册 (op_host/dispatch_layout.cpp)

3.3 Tiling 策略 (op_host/dispatch_layout_tiling.cc)

3.4 Kernel 实现 — A1 版本 (dispatch_layout.h)

核心算法：两阶段 token 计数

跨核原子累加

3.5 Kernel 实现 — A2 变体 (dispatch_layout_a2.h)

关键差异

A2 新增数据结构

A2 核心计算

A2 前缀和计算

3.6 入口函数 (op_kernel/dispatch_layout.cpp)

3.7 aclnn 接口层

4. moe_dispatch_normal_zb — 标准分发算子

4.1 算子概述

4.2 算子注册 (op_host/moe_dispatch_normal_zero_buffer.cpp)

4.3 Tiling 数据

4.4 Kernel 核心实现

模板参数

内存布局常量

远端地址计算

核心执行流程 (Process)

数据搬运核心 (InputToDstOutput)

量化处理 (QuantProcess)

4.5 ZeroBufferSyncFlag 初始化

4.6 与 master 版 moe_dispatch_normal 的差异

5. moe_dispatch_lowlatency_zb — 低延迟分发算子

5.1 算子概述

5.2 算子注册 (op_host/moe_dispatch_lowlatency_zero_buffer.cpp)

5.3 Tiling Key 矩阵

5.4 Kernel 核心实现

模板参数

内存布局常量

元数据布局

核心执行流程 (Process)

Phase 1: LocalLayout — 本地统计

Phase 2: SetLayoutStatus — 通知统计完成

Phase 3: SendCountNotify — 等待并获取对端统计

Phase 4: InputToDstOutput — 数据搬运

Phase 5: SetSyncFlag / WaitSyncFlag — 数据写入同步

5.5 AIV 分配策略

5.6 与 normal 版的差异

6. moe_notify_dispatch_zb — 元数据通知算子

6.1 算子概述

6.2 算子注册 (op_host/notify_dispatch_zero_buffer.cpp)

6.3 Tiling 数据

6.4 Kernel 核心实现

核心执行流程 (Process)

AllGatherSendData — 跨卡数据收集

CpGM2GMPingPong — Ping-Pong GM→GM 拷贝

BuildMaxBs — 计算最大 batch size

BuildTotalRecvCount — 计算前缀和偏移表

ReorderRecvDataOutput — 矩阵转置 + 前缀和

6.5 远端地址计算

6.6 ZeroBufferSyncFlag 初始化

7. SHMEM 集成方式

7.1 SHMEM API 调用汇总

7.2 类型分发宏 (lowlatency_zb)

7.3 远端指针获取模式

7.4 同步原语

8. 性能优势分析

8.1 内存占用对比