1. 整体架构设计哲学
1.1 定位与使命
CAM 是华为 Ascend NPU 上面向 MoE Expert Parallelism (EP) 通信加速的软件栈。其核心使命是：
在 Ascend 910_93 (A3) 芯片上，为 MoE 模型提供高性能的 EP 通信算子（Dispatch/Combine），并进一步将通信与计算（GEMM/SwiGLU）融合为单一算子（FusedDeepMoE），以最大化 HCCS/RoCE 互联带宽利用率。
1.2 设计哲学总结
CAM 的设计体现了以下核心哲学：
| 哲学维度 | 具体表现 |
|---|---|
| 硬件原生 | 直接操作 RDMA QP/SQ/CQ 硬件队列，绕过 HCCL 抽象层以获得极致延迟 |
| 垂直融合 | FusedDeepMoE 将 Dispatch + GMM1 + SwiGLU + GMM2 + Combine 融合为单算子，消除多次 kernel launch 和 HBM 读写 |
| 编译时多态 | 通过 C++ template traits 系统（aclshmemi_backend_traits）在编译期选择 IN_DIE / XSCALE 后端，零运行时开销 |
| SOC 世代解耦 | operator_registry.json + select_ops.py 实现 SOC 世代与算子版本的声明式注册 |
| 分层传输抽象 | NDA (Network Device Abstraction) 层分离 host 端资源管理（RdmaTransportManagerV2）与 device 端 RDMA 原语 |
| Ascend C 编程范式 | 严格遵循 op_host（tiling 计算）/ op_kernel（计算 kernel）/ op_api（aclnn 接口）三层模型 |
2. 模块分解与职责
src/cam/
├── comm_operator/ # 算子实现主目录
│ ├── ascend_kernels/ # Ascend C 原生算子源码
│ │ ├── operator_registry.json # SOC → 算子映射注册表
│ │ ├── dispatch_layout/ # EP layout 计算（纯计算，无通信）
│ │ ├── moe_dispatch_normal/# MoE Dispatch（prefill）
│ │ ├── moe_combine_normal/ # MoE Combine（prefill）
│ │ ├── notify_dispatch/ # Dispatch 通知同步算子
│ │ ├── fused_deep_moe/ # 融合算子（decode: dispatch+gemm+combine）
│ │ ├── utils/ # 共享工具（tiling_args, comm_args, sync, data_copy）
│ │ └── cmake_files/ # 构建模板（op_host/CMakeLists.txt, op_kernel/CMakeLists.txt）
│ └── pybind/ # Python 绑定层
│ ├── pybind.cpp # PYBIND11_MODULE + TORCH_LIBRARY 注册
│ ├── get_dispatch_layout.cpp # aclnnDispatchLayout 的 NPU 执行逻辑
│ ├── fused_deep_moe.cpp # FusedDeepMoe 的 pybind 实现
│ ├── moe_dispatch_prefill.cpp
│ ├── moe_combine_prefill.cpp
│ └── pytorch_extension/ # PyTorch 扩展（NPUStorageImpl, NPUBridge）
├── framework/ # 框架层
│ └── nda/ # Network Device Abstraction（核心传输层）
│ ├── include/device/gm2gm/engine/
│ │ └── shmem_device_rdma.h # 用户态 RDMA API 声明
│ └── src/
│ ├── device/gm2gm/engine/ # Device 端 RDMA 原语
│ │ ├── shmem_device_rdma.hpp # RDMA 操作实现（get/put/atomic）
│ │ ├── shmemi_device_rdma.h # 底层数据结构定义
│ │ └── rdma_backends/ # 后端特化
│ │ ├── rdma_device_backend_base.h # 后端 traits 系统接口
│ │ ├── rdma_device_backend_base.hpp # 通用 write/read/atomic 实现
│ │ └── rdma_device_backend_xscale.hpp # XSCALE 后端完整特化
│ └── host/
│ ├── transport/device_rdma/
│ │ ├── device_rdma_transport_manager_v2.h # Host 端传输管理器
│ │ └── device_rdma_transport_manager_v2.cpp # HCCP/Hcomm 初始化与 QP 信息提取
│ └── utils/under_api/
│ ├── dl_hccp_def.h # HCCP 协议结构体定义（QP/CQ/MR）
│ └── dl_hcomm_def.h # Hcomm 上层结构体（ChannelEntity, ProtectionInfo）
├── examples/ # 使用示例
│ ├── moe_dispatch_combine_prefill_sample.py # Dispatch+Combine 性能 benchmark
│ ├── fused_deep_moe_sample.py # FusedDeepMoE vs 分步算子对比
│ └── utils.py # 分布式初始化、benchmark 工具
└── third_party/
└── catlass/ # CATLASS 模板库（git submodule）
2.1 关键设计决策：为何需要 NDA 层
传统 HCCL 通信库提供高阶集合通信原语（AllReduce/AllToAll 等）。但 MoE EP 通信的模式是 细粒度的、动态的、不规则的 P2P 数据传输（每个 token 发往不同 expert 所在的 rank），HCCL 的抽象层级过高，无法满足：
- 延迟敏感：需要直接 post RDMA WQE 到硬件 SQ，避免 HCCL 的软件调度开销
- 计算-通信重叠：需要在 kernel 内部发起 RDMA 操作后立即进行 GEMM 计算
- 原子操作：Combine 阶段需要 remote atomic 操作更新 token 计数器
因此 NDA 层本质上是一个 用户态 RDMA 子系统，在 Ascend 芯片上直接操作 RoCE 硬件。
3. framework/nda 传输抽象层深度分析
3.1 架构概览
NDA 层是 CAM 的"操作系统网络栈"，分为 host 端 和 device 端 两部分：
┌─────────────────────────────────────────────────────────────┐
│ Host Side │
│ RdmaTransportManagerV2 │
│ ├── OpenDevice() → 获取 phyId, deviceIp │
│ ├── CreateEndpoint() → HcommEndpointCreate (RoCE) │
│ ├── RegisterMemoryRegion → HcommMemReg (注册 HBM 为 MR) │
│ ├── Connect() → HcommChannelCreate + FillRdmaInfo│
│ └── GetQpInfo() → 返回 device-visible 的 QP 信息 │
│ │
│ 底层依赖：dl_hccp_def.h (HCCP 协议) + dl_hcomm_def.h │
└─────────────────────────────────────────────────────────────┘
│
│ qpInfo_ (AiQpRMAQueueInfo)
│ 通过 AclrtMemcpy Host→Device 拷贝
▼
┌─────────────────────────────────────────────────────────────┐
│ Device Side │
│ shmem_device_rdma.hpp │
│ ├── aclshmemx_roce_get_nbi() ← RDMA READ (远端→本地) │
│ ├── aclshmemx_roce_put_nbi() ← RDMA WRITE (本地→远端) │
│ ├── aclshmemx_roce_quiet() ← 等待所有 WQE 完成 │
│ └── aclshmemx_roce_atomic_*() ← RDMA 原子操作 (CAS/FA) │
│ │
│ 后端选择 (compile-time): │
│ ├── IN_DIE — 片内互联（HCCS） │
│ └── XSCALE — 跨机 RoCE（DDRM Diamond 引擎） │
└─────────────────────────────────────────────────────────────┘
3.2 Host 端：RdmaTransportManagerV2
文件位置：src/cam/framework/nda/src/host/transport/device_rdma/
RdmaTransportManagerV2 继承自 TransportManager 接口，是 host 端 RDMA 资源管理的核心类。其生命周期：
3.2.1 初始化流程
// 1. OpenDevice: 获取设备物理 ID 和 RDMA 网络地址
Result OpenDevice(const TransportOptions& options) {
DlAclApi::AclrtGetDevice(&userId);
DlAclApi::AclrtGetPhyDevIdByLogicDevId(userId, &phyId_);
TopoReader::ParseRdmaNetAddr(phyId_, deviceIp_); // 从拓扑表获取 RDMA IP
CreateEndpoint(); // HcommEndpointCreate (COMM_PROTOCOL_ROCE)
ReserveRdmaInfoSpace(); // 在 HBM 上预留 qpInfo + atomic 共享内存
}
// 2. Connect: 创建通信通道并填充 QP 信息
Result Connect() {
ValidateRanksPerNic(); // 验证单 NIC 不超过 16 ranks
PrepareChannelDescs(); // 为每个 remote rank 准备 HcommChannelDesc
CreateChannelsAndFillInfo();
// ├── RegisterAtomicMemory() → HcommMemReg("AtomicBuffer")
// ├── HcommChannelCreate() → COMM_ENGINE_AIV (AIV 直接驱动 RDMA)
// └── FillRdmaInfo() → 从 ChannelEntity 提取 SQ/CQ/MR 信息，拷贝到 device
}
3.2.2 关键数据结构
AiQpRMAQueueInfo（host 侧镜像，最终拷贝到 device）：
struct AiQpRMAQueueInfo {
uint32_t count;
AiQpRMAWQ* sq; // Send Queue 信息 [rankCount]
AiQpRMAWQ* rq; // Receive Queue 信息 [rankCount]
AiQpRMACQ* scq; // Send Completion Queue [rankCount]
AiQpRMACQ* rcq; // Receive Completion Queue [rankCount]
RdmaMemRegionInfo* mr; // Memory Region keys [rankCount]
};
每个 AiQpRMAWQ 包含：
- wqn — QP 编号
- bufAddr — SQ ring buffer 的 HBM 地址
- wqeSize / depth — WQE 大小和队列深度
- headAddr / tailAddr — producer/consumer index 的 HBM 地址
- dbAddr — doorbell register 地址
- atomicAddr / atomicLkey — 原子操作的本地缓冲区
3.2.3 内存注册策略
Host 端管理两类内存注册：
- 用户 HBM 缓冲区（
RegisterMemoryRegion）：通过HcommMemReg注册，区分COMM_MEM_TYPE_DEVICE（HBM）和COMM_MEM_TYPE_HOST（DDR）。 - 原子操作共享内存（
RegisterAtomicMemory）：大小为ATOMIC_MAX_NUM(128) × sizeof(uint64_t) × rankCount，512 字节对齐，用于存储原子操作的 fetch-before-write 结果。
3.3 Device 端：RDMA 原语层
文件位置：src/cam/framework/nda/src/device/gm2gm/engine/
3.3.1 用户 API 层（shmem_device_rdma.h）
这一层提供类似 OpenSHMEM 的 RMA 接口，但底层是原生 RDMA 操作：
// 对称地址转换：将本地对称地址转为远端可访问地址
ACLSHMEM_DEVICE void* aclshmem_roce_ptr(void* ptr, int pe);
// 异步 RDMA READ：从远端对称内存读取到本地
template <typename T>
ACLSHMEM_DEVICE void aclshmemx_roce_get_nbi(T* dst, T* src, T* buf, uint32_t elem_size, int pe);
// 异步 RDMA WRITE：从本地写入远端对称内存
template <typename T>
ACLSHMEM_DEVICE void aclshmemx_roce_put_nbi(T* dst, T* src, T* buf, uint32_t elem_size, int pe);
// Quiet：等待指定 PE 的所有先前 WQE 完成
template <typename T>
ACLSHMEM_DEVICE void aclshmemx_roce_quiet(uint32_t pe, T* buf, uint32_t sync_id);
// 原子操作（仅 XSCALE 后端 / Ascend950 支持）
template <typename T> T aclshmemx_roce_atomic_fetch(T* src, int32_t pe);
template <typename T> void aclshmemx_roce_atomic_set(T* dst, T value, int32_t pe);
template <typename T> T aclshmemx_roce_atomic_compare_swap(T* dst, T cond, T value, int32_t pe);
template <typename T> void aclshmemx_roce_atomic_add(T* dst, T value, int32_t pe);
// ... 还有 atomic_inc, atomic_fetch_add, atomic_and/or/xor 等
关键约束（WARNING 注释明确指出）：
"When using RDMA as the underlying transport, concurrent RMA/AMO operations to the same PE are not supported. Use sync_id in device_state.rdma_config for pipeline synchronization."
这意味着对同一远端的并发 RMA 操作需要通过 sync_id（MTE3 event）进行流水线同步。
3.3.2 后端抽象层（rdma_backends/）
这是 NDA 最精妙的设计 — 编译时多态 traits 系统：
// base.h: 后端 traits 接口
template <aclshmemi_rdma_backend_t B>
struct aclshmemi_backend_traits {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...); // 构造 WQE
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...); // 提交并 ring doorbell
};
};
// xscale.hpp: XSCALE 后端特化
template <>
struct aclshmemi_backend_traits<aclshmemi_rdma_backend_t::XSCALE> {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
// 使用 XSCALE Diamond 引擎特定的 WQE 格式
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...) {
return aclshmemi_roce_xscale_fill_wqe_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
};
};
后端选择在编译期通过宏确定：
#if defined(ACLSHMEMI_RDMA_K_BACKEND_XSCALE)
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::XSCALE)
#else
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::IN_DIE)
#endif
3.3.3 XSCALE 后端详解（rdma_device_backend_xscale.hpp）
XSCALE 后端直接操作 DDRM Diamond RDMA 引擎，这是 Ascend 910_93 跨机通信的硬件核心。
WQE 格式（XSCALE v2）：
struct aclshmemi_xscdv_wqe_ctrl_seg_v2 {
uint8_t msg_opcode; // 操作码（RDMA_WRITE=1, RDMA_READ=2, ATOMIC_CAS=26, ATOMIC_FA=27...）
uint8_t with_imm : 1;
uint8_t csum_en : 2;
uint8_t ds_data_num : 5; // data segment 数量
uint16_t rsv1;
uint32_t msg_len;
uint32_t opcode_data;
uint32_t se : 1; // solicited event
uint32_t ce : 1; // 是否产生 CQE
uint32_t in_line : 1; // inline 模式
uint32_t fence_mode : 2;
uint32_t mask : 2;
uint32_t rsv : 5;
uint32_t wqe_id : 20; // WQE ID
}; // 共 16 bytes
CQE 格式（Completion Queue Entry）：
struct aclshmemi_xscdv_diamond_cqe_v2 {
uint32_t error_code : 8;
uint32_t qp_id : 15;
uint32_t se : 1;
uint32_t has_pph : 1;
uint32_t type : 1;
uint32_t with_imm : 1;
uint32_t csum_err : 4;
uint32_t imm_data;
uint32_t msg_len;
uint32_t vni;
uint32_t ts_l; // 时间戳低 32 位
uint32_t ts_h; // 时间戳高 16 位（共 48 位）
uint32_t msg_opcode : 8;
uint32_t wqe_id : 20;
uint32_t owner : 1; // ownership bit
}; // 共 32 bytes，cqe64 补齐到 64 bytes
关键操作流程（以 RDMA WRITE 为例）：
// 1. 检查 SQ 剩余容量，不足则 poll_cq 等待 WQE 完成
if ((cur_head + XSCALE_POLL_CQ_THRESHOLD) % depth == (depth + cur_tail) % depth) {
aclshmemi_roce_poll_cq<XSCALE>(pe, qp_idx, cur_head, ...);
}
// 2. 获取 remote/local memory keys
wr.rkey = remote_mem_info->rkey;
wr.lkey = local_mem_info->lkey;
// 3. 填充 WQE（ctrl seg + data seg）
aclshmemi_roce_fill_wqe<XSCALE, OP_RDMA_WRITE>(wr, sq_context, wqe_addr, cur_head);
// 4. 刷 cacheline 到 HBM
dcci_cachelines(wqe_addr, wqe_total_size);
// 5. Ring SQ Doorbell（更新 head pointer + 写 doorbell register）
aclshmemi_roce_ring_sq_doorbell<XSCALE>(sq_context, cur_head, ...);
Poll CQ 流程：
while (cur_tail != target_idx) {
// 检查 owner bit 判断 CQE 是否有效
while (!aclshmemi_roce_xscale_check_cqe_owner(cqe_addr, cur_tail, depth) &&
run_cycles < ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_CYCLES) {
run_cycles++;
dcci_cachelines((__gm__ uint8_t *)cqe_addr, sizeof(aclshmemi_xscdv_cqe64_t));
}
// 检查 error_code
status = cqe_addr->cqe.error_code;
cur_tail++;
}
// 更新 CQ tail 并 ring CQ doorbell
aclshmemi_roce_ring_cq_doorbell<XSCALE>(...);
超时机制：5 分钟超时（ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_DURATION = 5ULL * 60 * 1000000 微秒），超过后返回 ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_ERROR = 0x10000。
3.4 NDA 如何抽象 HCCL/HCCP
NDA 并非完全绕过 HCCL/HCCP，而是 分层使用：
| 层级 | 组件 | 作用 |
|---|---|---|
| 初始化 | HcommEndpointCreate / HcommChannelCreate | 使用 HCCP 协议栈建立 QP 连接，获取 SQ/CQ/MR 句柄 |
| 内存注册 | HcommMemReg / HcommMemUnreg | 通过 HCCP 注册内存区域，获取 lkey/rkey |
| QP 信息提取 | ChannelEntity → AiQpRMAQueueInfo | 从 HCCP 内部结构体提取硬件寄存器地址 |
| 运行时数据面 | 直接操作 WQE/doorbell | 绕过 HCCL/HCCP 软件栈，AIV 核直接写硬件 SQ 队列 |
这种设计的好处： 1. 控制面复用 HCCP：利用其成熟的拓扑发现、连接建立、QP 状态机管理 2. 数据面自研：针对 MoE EP 的 P2Gather 模式做极致优化（inline write、masked atomic、pipeline sync）
4. Operator Registry 与 SOC 世代机制
4.1 operator_registry.json 结构
{
"soc_versions": {
"ascend910_93": [
"dispatch_layout",
"fused_deep_moe",
"moe_combine_normal",
"moe_dispatch_normal",
"notify_dispatch"
]
},
"operator_meta": {
"dispatch_layout": { "requires_shmem": false },
"fused_deep_moe": {
"requires_shmem": false,
"family": "fused_deep_moe",
"variant": "base"
},
...
}
}
设计要点：
- soc_versions 定义每个 SOC 世代支持的算子列表
- operator_meta 提供算子级元数据（如 requires_shmem、family、variant）
- 算子身份 = 源码目录名（ascend_kernels/ 下的子目录）
4.2 select_ops.py 解析逻辑
文件位置：build/cam/comm_operator/select_ops.py
def resolve(registry, soc_list, shmem, quant, user_ops):
# 1. 确定候选集：全量 SOC 列表 or 用户指定的 -a 列表
candidates = user_ops if user_ops else soc_list
# 2. fused_deep_moe base/w4a8 互斥规则
if quant:
# -q 选择 w4a8 变体，移除 base
candidates = [o for o in candidates if o != BASE]
else:
# 无 -q 时 base 优先，移除 w4a8
candidates = [o for o in candidates if o != W4A8]
# 3. SHMEM 过滤：未安装 SHMEM 时跳过 requires_shmem 的算子
if not shmem:
candidates = [o for o in candidates if not requires_shmem(o)]
# 4. 保持 registry 顺序（确定性构建）
return [o for o in soc_list if o in candidates]
互斥规则的本质：fused_deep_moe（base）和 fused_deep_moe_w4a8 共享完全相同的源文件名（fused_deep_moe.cpp），区别仅在于编译时宏定义（-DW4A8_QUANTIZATION），因此不能同时编译进一个 run package。
4.3 构建系统集成
build/cam/build.sh 通过以下流程调用 select_ops.py：
# 解析出最终要编译的算子列表
OPS=$(python3 select_ops.py \
--registry operator_registry.json \
--soc ascend910_93 \
--shmem ${SHMEM_AVAILABLE} \
--quant ${QUANT_FLAG} \
--ops "${USER_OPS}")
# 遍历算子，逐个编译
for op in $OPS; do
# 复制 cmake_files/ 模板到算子目录
cp -r cmake_files/op_host ${op}/op_host/cmake
cp -r cmake_files/op_kernel ${op}/op_kernel/cmake
# 调用 ascend_c_impl_build 构建
done
5. op_host / op_kernel / op_api 三层编程模型
5.1 模型概述
Ascend C 编程模型严格遵循三段式结构：
┌──────────────────────────────────────────────────────────────┐
│ op_host (CPU 侧) │
│ ├── xxx.cpp → OpDef 注册（输入/输出/属性声明） │
│ └── xxx_tiling.cpp → Tiling 函数（计算数据切分、workspace） │
│ 运行时调用，将 tilingData 传给 device │
├──────────────────────────────────────────────────────────────┤
│ op_kernel (Device AIV/Vector 侧) │
│ ├── xxx.h → Kernel 类定义（Init/Process 方法） │
│ ├── xxx.cpp → Kernel 实现 │
│ └── xxx_tiling.h → TilingData 结构体定义 │
│ 在 NPU 上并行执行的真正计算逻辑 │
├──────────────────────────────────────────────────────────────┤
│ op_api (aclnn 接口层) │
│ ├── aclnn_xxx.h → C 接口声明（GetWorkspaceSize + Execute）│
│ └── aclnn_xxx.cpp → 封装 aclnnInner 调用 │
│ 对外暴露的统一算子接口 │
└──────────────────────────────────────────────────────────────┘
5.2 典型示例：dispatch_layout
5.2.1 op_host 层
dispatch_layout.cpp — OpDef 注册：
class DispatchLayout : public OpDef {
public:
explicit DispatchLayout(const char *name) : OpDef(name) {
this->Input("topkIdx")
.ParamType(REQUIRED)
.DataType({ge::DT_INT64})
.Format({ge::FORMAT_ND});
this->Attr("num_tokens").Int();
this->Attr("num_ranks").Int();
// ...
this->Output("numTokensPerRank")...
this->Output("numTokensPerExpert")...
// ...
// 指定支持的 SOC 世代和配置
OpAICoreConfig a3_config;
a3_config.DynamicCompileStaticFlag(true)
.DynamicFormatFlag(true)
.PrecisionReduceFlag(true)
.ExtendCfgInfo("jitCompile.flag", "static_true");
this->AICore().AddConfig("ascend910_93", a3_config);
}
};
OP_ADD(DispatchLayout);
dispatch_layout_tiling.cpp — Tiling 计算：
// Tiling 函数签名
ge::graphStatus DispatchLayoutTilingFunc(gert::TilingContext *context) {
// 1. 从 context 获取输入 shape 和 attr
auto attrs = context.GetAttrs();
// 2. 参数校验（shape 维度、数据类型、值域）
CheckTensorDataType(context, nodeName);
CheckTensorShape(context, nodeName);
// 3. 填充 TilingData
GetAttrAndSetTilingData(context, nodeName, *tilingData);
// 4. 设置 workspace 大小
SetWorkSpace(context, nodeName); // SYSTEM_NEED_WORKSPACE + KERNEL_USE_WORKSPACE
// 5. 设置 tiling key（选择 kernel 分支）
context.SetTilingKey(TILING_KEY_INT); // =23
// 6. 设置 blockDim（AIV 核数）
context.SetBlockDim(aivNum);
}
// 注册 tiling 函数
IMPL_OP_OPTILING(DispatchLayout)
.Tiling(DispatchLayoutTilingFunc)
.TilingParse<DispatchLayoutCompileInfo>(TilingParseForDispatchLayout);
5.2.2 op_kernel 层
dispatch_layout.h — Kernel 类定义：
template <typename T>
class DispatchLayout {
public:
__aicore__ inline void Init(...) {
// 从 tilingData 读取配置
numTokens_ = tilingData->dispatchLayoutInfo.numTokens;
// 计算每个 AIV core 负责处理的 token 数
aivNum_ = numTokens_ <= maxAivNum ? numTokens_ : maxAivNum;
// 分配 UB buffer
}
__aicore__ inline void Process() {
// 1. 从 HBM 加载 topk_idx 到 UB
DataCopyPad(topkIdxTensor, topkIdxGM_, ...);
// 2. 统计每个 expert/rank 的 token 数
for (int i = 0; i < tempTokens_; ++i) {
for (int j = 0; j < numTopk_; ++j) {
int64_t expert_idx = topkIdxTensor.GetValue(i * numTopk_ + j);
numTokensPerExpertTensor.SetValue(expert_idx, per_expert_num);
}
}
// 3. 写回 HBM
DataCopyPad(numTokensPerExpertGM_, numTokensPerExpertTensor, ...);
// 4. 跨 core 同步
SyncAll<true>();
}
};
dispatch_layout.cpp — Kernel 入口：
extern "C" __global__ __aicore__ void dispatch_layout(
GM_ADDR topkIdx, GM_ADDR numTokensPerRank, ..., GM_ADDR workspace, GM_ADDR tiling)
{
REGISTER_TILING_DEFAULT(DispatchLayoutTilingData);
GET_TILING_DATA_WITH_STRUCT(DispatchLayoutTilingData, tilingData, tiling);
TPipe pipe;
if (TILING_KEY_IS(TILING_KEY_INT)) {
MoeDispatchLayout::DispatchLayout<int32_t> op;
op.Init(topkIdx, ..., workspace, &pipe, &tilingData);
op.Process();
}
}
5.2.3 op_api 层
aclnn_dispatch_layout.h — C 接口声明：
__attribute__((visibility("default")))
aclnnStatus aclnnDispatchLayoutGetWorkspaceSize(
const aclTensor *topkIdx, int64_t numTokens, ...,
uint64_t *workspaceSize, aclOpExecutor **executor);
__attribute__((visibility("default")))
aclnnStatus aclnnDispatchLayout(
void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream);
aclnn_dispatch_layout.cpp — 实现：
aclnnStatus aclnnDispatchLayout(void *workspace, uint64_t workspaceSize,
aclOpExecutor *executor, aclrtStream stream) {
// 设置 HCCL server 类型为 MTE（避免走 AICPU 路径）
if (NnopbaseSetHcclServerType) {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
}
return aclnnInnerDispatchLayout(workspace, workspaceSize, executor, stream);
}
5.3 三层交互时序
Python (torch.ops.umdk_cam_op_lib.get_dispatch_layout)
│
▼
TORCH_LIBRARY 分发 → GetDispatchLayoutImpl (pybind)
│
▼
EXEC_NPU_CMD(aclnnDispatchLayout, ...) → aclrtLaunchKernel
│
▼
┌─ Host ─────────────────────────────────────────────┐
│ Tiling 阶段 (CPU): │
│ DispatchLayoutTilingFunc(context) │
│ ├── 计算 blockDim, tilingKey │
│ ├── 填充 DispatchLayoutTilingData │
│ └── 设置 workspace 大小 │
└────────────────────────────────────────────────────┘
│
▼ tilingData (HBM)
┌─ Device ───────────────────────────────────────────┐
│ Kernel 阶段 (AIV cores): │
│ dispatch_layout() │
│ ├── GET_TILING_DATA(tilingData) │
│ ├── op.Init() → 分配 UB, 计算偏移 │
│ └── op.Process() → 执行计算 + 同步 │
└────────────────────────────────────────────────────┘
6. Examples 使用模式
6.1 moe_dispatch_combine_prefill_sample.py
这是一个 prefill 阶段的 EP 通信 benchmark，展示完整的 Dispatch-Combine 流程：
# 1. 初始化分布式环境
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
# 2. 计算 layout（纯计算，无通信）
num_tokens_per_expert, send_token_idx_small = torch.ops.umdk_cam_op_lib.get_dispatch_layout(
topk_idx=topk_idx, num_experts=num_experts, num_ranks=num_ranks
)
# 3. Dispatch：将 token 按 expert 分布发送到目标 rank
recv_x, dynamic_scales, expand_idx, recv_count, recv_tokens_per_expert = \
torch.ops.umdk_cam_op_lib.moe_dispatch_prefill(
x=x, topk_idx=topk_idx, topk_weights=topk_weights,
num_tokens_per_expert=num_tokens_per_expert,
send_token_idx_small=send_token_idx_small,
group_ep=ep_hcomm_info, rank=rank, num_ranks=num_ranks, use_quant=False
)
# 4. Combine：将 expert 输出加权聚合回原 rank
combined_x = torch.ops.umdk_cam_op_lib.moe_combine_prefill(
x=recv_x, topk_idx=topk_idx, topk_weights=topk_weights,
src_idx=expand_idx, send_head=recv_count,
group_ep=ep_hcomm_info, rank=rank, num_ranks=num_ranks
)
性能测量：
# 使用 torch.npu.Event 精确计时
t = bench(lambda: torch.ops.umdk_cam_op_lib.moe_dispatch_prefill(**dispatch_args))
# 计算有效带宽
bandwidth = recv_bytes / t # GB/s (HCCS)
6.2 fused_deep_moe_sample.py
这是一个 decode 阶段的融合算子验证，对比 FusedDeepMoE 与分步执行（SmallOps）的结果一致性：
# SmallOps: 分步执行 (使用华为 CANN 内置算子)
expand_x, ..., expert_token_nums, ... = torch_npu.npu_moe_distribute_dispatch_v2(...)
y1 = torch_npu.npu_grouped_matmul(x=[expand_x], weight=[gmm1_weight], ...)
y1 = torch_npu.npu_dequant_swiglu_quant(x=y1, ...)
y2 = torch_npu.npu_grouped_matmul(x=[y1], weight=[gmm2_weight], ...)
combine_output = torch_npu.npu_moe_distribute_combine_v2(expand_x=y2, ...)
# FusionOp: CAM 融合算子（单 kernel）
output, share_output, expert_token_nums = torch.ops.umdk_cam_op_lib.fused_deep_moe(
x=x, expert_ids=expert_ids,
gmm1_weight=gmm1_weight, gmm1_weight_scale=gmm1_weight_scale,
gmm2_weight=gmm2_weight, gmm2_weight_scale=gmm2_weight_scale,
expert_scales=expert_scales, share_gmm1_weight=..., x_active_mask=...,
group_ep=ep_hcomm_info, ep_rank_size=ep_world_size, ep_rank_id=global_rank_id,
moe_expert_num=moe_expert_num, quant_mode=0, global_bs=global_bs
)
# 精度验证
torch.testing.assert_close(small_op_token_output, fused_op_token_output, atol=2.0, rtol=0.02)
关键参数模式：
- quant_mode=0：当前仅支持 W8A8（int8 权重 + bf16 激活）
- global_bs：全局 batch size，用于 alltoall 通信的 buffer 分配
- x_active_mask：支持 dynamic batch（部分 token 不参与 dispatch）
- dynamic_eplb：支持 expert-level load balancing（权重按 expert 维度拆分）
6.3 使用模式总结
| 场景 | 调用模式 | 关键约束 |
|---|---|---|
| Prefill EP | get_dispatch_layout → moe_dispatch_prefill → moe_combine_prefill | 三步分离，中间可插入自定义计算 |
| Decode 融合 | fused_deep_moe (单步) | 通信+计算融合，延迟最低 |
| 精度验证 | 分步执行 vs 融合执行对比 | atol=2.0, rtol=0.02 (W8A8 精度) |
7. 关键设计决策与 Trade-offs
7.1 Trade-off 1: 用户态 RDMA vs HCCL 集合通信
| 方案 | 优势 | 劣势 |
|---|---|---|
| 用户态 RDMA (CAM 选择) | 极致低延迟、可流水线化、支持 inline 操作 | 需要管理 QP/CQ 生命周期、代码复杂度高 |
| HCCL AllToAll | 抽象简单、自动处理拓扑 | 无法做到 kernel 内融合、调度开销大 |
决策理由：MoE EP 的通信模式是高度动态的（每个 token 的目的地不同），HCCL 的静态规划 AllToAll 无法高效处理。用户态 RDMA 允许在 kernel 内部按需发起单次 RDMA 操作，实现真正的计算-通信重叠。
7.2 Trade-off 2: 编译时后端选择 vs 运行时调度
| 方案 | 优势 | 劣势 |
|---|---|---|
| 编译时 traits (CAM 选择) | 零运行时开销、编译器可优化 | 每个后端需要单独编译 |
| 运行时 switch | 单 binary 适配多后端 | 分支预测开销、无法内联 |
决策理由：RDMA 操作在关键路径上，每次操作节省几个时钟周期在百万次调用下累积显著。且 Ascend 芯片的后端（IN_DIE vs XSCALE）在部署时即确定，无需运行时适配。
7.3 Trade-off 3: 三层编程模型 vs 单层 kernel
| 方案 | 优势 | 劣势 |
|---|---|---|
| op_host/op_kernel/op_api (CAM 选择) | 符合 CANN 生态规范、支持 GE 图编译、自动并行 | 代码分散在多个文件 |
| 单层 kernel | 代码集中 | 无法利用 CANN 的调度优化、不兼容 ACLGraph |
决策理由：Ascend C 三层模型是华为 CANN 生态的标准约定。op_host 的 tiling 计算允许运行时根据输入 shape 动态调整数据切分策略，op_kernel 的模板化支持多种数据类型，op_api 层则保证了与 torchair 图编译模式的兼容。
7.4 Trade-off 4: 注册式 SOC 管理 vs 硬编码
| 方案 | 优势 | 劣势 |
|---|---|---|
| JSON registry + Python resolver | 新增 SOC 世代只需编辑 JSON、算子与 SOC 解耦 | 引入外部依赖 |
| 硬编码 CMake 条件 | 无额外依赖 | 每加 SOC 需修改构建脚本 |
决策理由：CAM 未来需要支持 A2、A3 等多代芯片，注册式设计使得算子开发者无需了解构建系统细节。select_ops.py 的互斥规则（base/w4a8）和 SHMEM 过滤逻辑也更容易维护。
7.5 Trade-off 5: 控制面 HCCP + 数据面自研
| 方案 | 优势 | 劣势 |
|---|---|---|
| HCCP 控制面 + 自研数据面 (CAM 选择) | 复用成熟的连接管理、同时获得极致数据面性能 | 需要从 HCCP 内部结构体提取信息 |
| 完全自研 | 无外部依赖 | QP 状态机、错误处理、重传机制需要大量工作量 |
| 纯 HCCL | 最简单 | 无法满足 kernel 内 RDMA 需求 |
决策理由：这是 NDA 层最精妙的设计。控制面（连接建立、QP 状态机、内存注册）使用 HCCP 的 HcommEndpointCreate / HcommChannelCreate API，数据面则直接从 ChannelEntity 结构体提取硬件寄存器地址（SQ head/tail pointer, doorbell address），在 kernel 中直接操作硬件。
7.6 Trade-off 6: FusedDeepMoE 的融合粒度
| 方案 | 优势 | 劣势 |
|---|---|---|
| Dispatch+GMM1+SwiGLU+GMM2+Combine (CAM 选择) | 最大程度减少 HBM 读写和 kernel launch | 编译复杂度高、需要定制 epilogue |
| 仅 Dispatch+Combine 融合 | 实现简单 | GEMM 仍需独立 kernel |
决策理由：MoE decode 阶段的 FFN 计算（GMM1 + SwiGLU + GMM2）受限于 HBM 带宽。融合后，dispatch 的输出直接在 UB/Cache 中进入 GMM1，GMM2 的输出也直接在 UB 中进入 combine 阶段，避免多次 HBM round-trip。从 fused_deep_moe/gemm/ 和 fused_deep_moe/epilogue/ 目录结构看，华为为融合算子定制了：
- block_mmad — 矩阵乘累加 tile
- block_mmad_preload_async_with_callback_resident_a — 异步预取 + 回调
- block_epilogue_per_token_dequant_swiglu — SwiGLU + 反量化融合 epilogue
8. 总结
CAM 是一个 面向 MoE EP 通信的垂直优化软件栈，其架构设计体现了以下核心洞察：
- MoE EP 通信不适合传统集合通信抽象 — 需要用户态 RDMA 获得极致延迟
- 计算-通信融合是性能关键 — FusedDeepMoE 将 5 阶段流水线融合为单 kernel
- 硬件多样性需要编译时多态 — traits 系统零开销适配 IN_DIE/XSCALE 后端
- 生态兼容性需要分层设计 — 控制面复用 HCCP，数据面自研，三层编程模型符合 CANN 规范
NDA 层是 CAM 的"皇冠上的明珠"，它实现了一个轻量级用户态 RDMA 子系统，在 Ascend 芯片上直接操作 RoCE 硬件，同时保持与 HCCP 控制面的兼容性。这种设计模式对理解国产 AI 芯片的软件栈演进具有重要参考价值。
分析范围：
build/cam/、src/cam/comm_operator/、docker/cam/、umdk.spec、src/cam/third_party/、RELEASE-NOTES.md分析日期：2026-07-21 项目根：/Users/backyes/work/claude_workspace/umdk_research/umdk
1. 编译系统完整链路分析
CAM 的编译系统采用 三层 shell 脚本 + CMake + 华为 CANN npu_op_* 构建框架 的架构，从顶层入口到内核编译的完整调用链如下：
build/cam/build.sh ← 顶层入口，解析 module 名 + 转发 -c/-a/-q/-p/-r
└─ build/cam/comm_operator/build.sh ← 模块级构建，管理编译阶段开关
├─ build/cam/comm_operator/compile_ascend_proj.sh ← 单 SOC 编译主流程
│ ├─ select_ops.py ← 算子选择（registry 解析 + 互斥规则）
│ ├─ set_conf.py ← 修改 CMakePresets.json（build_type/vendor）
│ ├─ msopgen gen ← CANN 工具：生成 SOC 特定算子工程骨架
│ ├─ cmake -S . -B build_out --preset=default
│ └─ cmake --build build_out --target binary/package
└─ build/cam/comm_operator/build_pybind.sh ← Python whl 包构建
└─ src/cam/comm_operator/pybind/setup.py ← NpuExtension + bdist_wheel
└─ pytorch_extension/bdist_wheel_build.py ← auditwheel repair 封装
1.1 顶层入口 build/cam/build.sh
关键设计：
- 通过 getopts ":dhc:" 仅解析 -d（Debug）、-c（SOC 类型）、-h（help），其余选项（-a/-q/-p/-r）原样转发给子模块脚本。
- 导出全局路径变量：ROOT_PATH、SRC_PATH、BUILD_OUT_PATH、SCRIPTS_PATH、TEST_PATH、BUILD_PATH、CAM_THIRD_PARTY_PATH。
- 通过 CPATH 注入 catlass 头文件路径：export CPATH=${CAM_THIRD_PARTY_PATH}:${CAM_THIRD_PARTY_PATH}/catlass/include:${CPATH}。
- prepare_cam_third_party() 负责初始化 catlass 子模块（git submodule update --init --recursive）。
- 模块分发逻辑：all/comm_operator 走 comm_operator 构建；all/nda 且 -c ascend950 走 nda 构建。
关键代码（build.sh:137-156）：
if [[ "$MODULE_NAME" == "all" || "$MODULE_NAME" == "comm_operator" ]]; then
if [[ -z "${SOC_TYPE}" || "${SOC_TYPE}" == "ascend910_93" ]]; then
IS_MODULE_EXIST=1
if ! prepare_cam_third_party "${CAM_THIRD_PARTY_PATH}"; then
exit 1
fi
${SCRIPTS_PATH}/comm_operator/build.sh $@
fi
fi
1.2 模块级构建 build/cam/comm_operator/build.sh
编译阶段开关（互斥设计）：
| 变量 | 默认值 | 含义 |
|---|---|---|
ENABLE_SRC_BUILD |
1 | 编译 .run 算子包 |
ENABLE_PYBIND_BUILD |
1 | 编译 .whl Python 包 |
ENABLE_UT_BUILD |
0 | 编译单元测试 |
ENABLE_RUN_ONLY |
0 | 仅编译 run 包（-r） |
USE_W4A8 |
0 | 选择 w4a8 量化变体（-q） |
OP_SELECT |
"" | 算子选择列表（-a） |
互斥规则：
- -r（run only）和 -p（pybind only）互斥，同时指定会 exit 1。
- -a 要求先指定 -c（SOC 版本），因为算子选择需要 SOC 上下文。
- -q 也依赖 -c，因为 w4a8 变体是 SOC 相关的。
关键转发（build.sh:107-108）：
export CAM_OP_SELECT="$OP_SELECT"
export CAM_USE_W4A8="$USE_W4A8"
1.3 单 SOC 编译主流程 compile_ascend_proj.sh
这是最核心的编译脚本，完成从算子选择到 .run 包产出的全过程：
核心函数链：
list_soc_versions()— 从operator_registry.json读取所有注册的 SOC 版本。resolve_ops()— 调用select_ops.py解析最终算子列表。copy_ops_include()— include-list 模式：仅拷贝选中算子的op_host/op_kernel/op_api到目标工程目录（区别于旧的 full-copy + exclude-list 模式，避免同名算子文件覆盖）。build_ascend_proj()— 单 SOC 编译主函数。
build_ascend_proj() 关键步骤：
# 1. 调用 msopgen 生成 SOC 特定算子工程
msopgen gen -i ./ascend_kernels/AddCustom.json -c ai_core-${soc_version} -f pytorch -lan cpp -out ${MODULE_BUILD_PATH}/${proj_name}
# 2. 删除模板自带的 AddCustom 示例算子
rm -rf ${MODULE_BUILD_PATH}/${proj_name}/op_host/add_custom*
rm -rf ${MODULE_BUILD_PATH}/${proj_name}/op_kernel/add_custom*
# 3. 拷贝 CMakeLists.txt 和 cmake_files/
cp ./ascend_kernels/CMakeLists.txt ${MODULE_BUILD_PATH}/${proj_name}/
cp ./ascend_kernels/cmake_files/op_host/CMakeLists.txt ${MODULE_BUILD_PATH}/${proj_name}/op_host/
cp ./ascend_kernels/cmake_files/op_kernel/CMakeLists.txt ${MODULE_BUILD_PATH}/${proj_name}/op_kernel/
cp -rf ./ascend_kernels/cmake_files/cmake ${MODULE_BUILD_PATH}/${proj_name}/
# 4. 修改 CMakePresets.json（build_type、ENABLE_SOURCE_PACKAGE、vendor_name）
python3 $SCRIPTS_PATH/comm_operator/set_conf.py ${MODULE_BUILD_PATH}/${proj_name}/CMakePresets.json $build_type True CAM
# 5. 配置 CANN 包路径
export ASCEND_CANN_PACKAGE_PATH="${ASCEND_TOOLKIT_HOME}" # 默认 /usr/local/Ascend/ascend-toolkit/latest
# 6. 推导 CATLASS_HOME_PATH（从 CPATH 解析）
export CATLASS_HOME_PATH=$(echo "${CPATH}" | tr ':' '\n' | grep "catlass" | sed 's|/include||')
# 7. CMake configure + build
cmake -S . -B build_out --preset=default -DCMAKE_BUILD_TYPE=$build_type
cmake --build build_out --target binary -j${BUILD_JOBS}
cmake --build build_out --target package -j${BUILD_JOBS}
# 8. 产出 .run 包
cp build_out/*.run ${BUILD_OUT_PATH}/comm_operator/run/CAM_${soc_version}_${os_id}_${arch}.run
关键 patch 机制（compile_ascend_proj.sh:177-192）：
由于 npu_op_* 构建系统的 simple_kernel_compile 不读取 custom_compile_options.ini，脚本通过 sed 直接 patch .make 文件来注入 SHMEM/CATLASS 的 include 路径：
KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
if [ -n "${CATLASS_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="${KERNEL_COMPILE_OPTS} -I${CATLASS_HOME_PATH}/include"
fi
find build_out/op_kernel/CMakeFiles -name "build.make" -exec sed -i "s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
1.4 Python whl 包构建 build_pybind.sh
build_pybind() {
local build_tmp_dir="${MODULE_BUILD_PATH}/pybind"
local wheel_out_dir="${MODULE_BUILD_OUT_PATH}/dist"
rm -rf "${build_tmp_dir}" "${wheel_out_dir}"
mkdir -p "${wheel_out_dir}" "$(dirname "${build_tmp_dir}")"
cp -r $MODULE_SRC_PATH/pybind "${build_tmp_dir}"
export MAX_JOBS=${CAM_BUILD_JOBS:-8}
cd "${build_tmp_dir}"
python3 setup.py bdist_wheel --dist-dir="${wheel_out_dir}"
}
注意： cp -r $MODULE_SRC_PATH/pybind "${build_tmp_dir}" 仅在 build_tmp_dir 不存在时才会将 pybind 目录内容直接拷贝到 build_tmp_dir（即 build_tmp_dir/setup.py）。如果 build_tmp_dir 已存在，会嵌套一层 pybind/pybind/。脚本通过 rm -rf "${build_tmp_dir}" 确保每次都是干净构建。
2. 算子选择机制
2.1 算子注册表 operator_registry.json
{
"soc_versions": {
"ascend910_93": [
"dispatch_layout",
"fused_deep_moe",
"moe_combine_normal",
"moe_dispatch_normal",
"notify_dispatch"
]
},
"operator_meta": {
"fused_deep_moe": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "base" },
"dispatch_layout": { "requires_shmem": false },
"moe_combine_normal": { "requires_shmem": false },
"moe_dispatch_normal": { "requires_shmem": false },
"notify_dispatch": { "requires_shmem": false }
}
}
设计要点：
- 算子身份 = 源码目录名（ascend_kernels/<op_name>/）。
- soc_versions 定义每个 SOC 支持的算子列表（当前仅注册 ascend910_93）。
- operator_meta 存储构建元数据（requires_shmem、family、variant）。
- 注释说明：要支持新 SOC 只需在 soc_versions 添加条目；要添加算子只需在对应 SOC 列表追加目录名。
2.2 算子选择逻辑 select_ops.py
解析流程：
- 加载注册表 →
load_registry() - 获取 SOC 支持列表 →
get_soc_list() - 解析用户 -a 列表 →
split_ops()（分号分隔，禁止空条目） - 验证用户选择 →
validate_user_ops()（每个 -a 条目必须在 SOC 支持列表中） - resolve() 核心规则：
互斥规则（fused_deep_moe 家族）：
- fused_deep_moe（base）和 fused_deep_moe_w4a8 共享相同的源码文件名（fused_deep_moe.*），因此互斥——同一 run 包只能编译其中一个。
- --quant（-q）选择 w4a8 变体，此时 base 被禁止。
- fused_deep_moe_fwk 是独立算子（文件名为 fused_deep_moe_fwk.*），不受互斥约束。
- 全量构建（无 -a）时默认保留 base，丢弃 w4a8（带 NOTE 提示）。
SHMEM 过滤：
- 当 SHMEM_HOME_PATH 未设置时，自动丢弃 requires_shmem=true 的算子。
确定性排序： - 最终列表按注册表中的顺序排列，确保构建可复现。
输出格式：
- 每行一个算子目录名，由 compile_ascend_proj.sh 通过 tr '\n' ' ' 转为空格分隔列表。
3. 编译依赖分析
3.1 核心依赖矩阵
| 依赖项 | 版本 | 用途 | 来源 |
|---|---|---|---|
| CANN | 8.5 / 9.0 | 华为 Ascend 算子开发工具包（msopgen、npu_op_*、ACL） | 官方 .run 安装包 |
| torch | 2.8.0 | PyTorch 框架（pybind 编译需要） | pip |
| torch-npu | 2.8.0.post2 (CANN 8.5) / 2.8.0.post4 (CANN 9.0) | 华为 PyTorch NPU 适配（NpuExtension、NpuBridge） | 华为 gitcode 发布 |
| gtest | 1.16.0 | 单元测试框架 | GitHub 源码编译 |
| OpenMPI | 5.0.7 | 多线程测试（MPI 接口） | 官方源码编译 |
| nlohmann/json | 3.12.0 | JSON 解析（测试/配置） | GitHub 源码编译 |
| catlass | v1-stable（pin） | 华为 CANN 模板库（GEMM/Epilogue 模板） | git submodule |
| Python | 3.11 | 构建脚本 + pybind | Miniconda |
| cmake | ≥3.16 | 构建系统 | 系统包 |
| gcc/g++ | - | C++17 编译 | 系统包 |
3.2 CANN 版本与 torch-npu 版本匹配
| CANN 版本 | torch-npu wheel | 来源 |
|---|---|---|
| 8.5.T63 | torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl |
gitcode Ascend/pytorch v7.3.0-pytorch2.8.0 |
| 9.0.0 | torch_npu-2.8.0.post4-cp311-cp311-manylinux_2_28_aarch64.whl |
gitcode Ascend/pytorch v26.0.0-pytorch2.8.0 |
3.3 catlass 子模块
.gitmodules 配置：
[submodule "src/cam/third_party/catlass"]
path = src/cam/third_party/catlass
url = https://gitcode.com/cann/catlass.git
branch = catlass-v1-stable
catlass 在编译中的作用：
- 提供 fused_deep_moe 算子的 GEMM/Epilogue 模板（grouped_matmul_slice_m_per_token_dequant_*.h、block_epilogue_*.hpp 等）。
- 通过 CPATH 环境变量注入（${CAM_THIRD_PARTY_PATH}/catlass/include）。
- 在 compile_ascend_proj.sh 中自动推导 CATLASS_HOME_PATH 并 patch 到 kernel compile options。
3.4 pybind 编译依赖
setup.py 中的 include 路径：
include_dirs=[
os.path.join(torch_npu_path, "include"),
os.path.join(torch_npu_path, "include/third_party/acl/inc/acl/"),
os.path.join(torch_npu_path, "include/third_party/acl/inc"),
os.path.join(os.environ["ASCEND_HOME_PATH"], f"{arch}-linux", "include"),
os.path.join(os.environ["ASCEND_HOME_PATH"], f"{arch}-linux", "include", "hccl"),
os.path.join(os.environ["ASCEND_HOME_PATH"], f"{arch}-linux", "include", "experiment", "runtime"),
os.path.join(os.environ["ASCEND_HOME_PATH"], f"{arch}-linux", "include", "experiment", "msprof"),
os.path.join(torch_path, "include"),
os.path.join(os.path.dirname(__file__), "./", "pytorch_extension")
]
链接库：
libraries=["torch_npu", "gcov", "runtime", "torch", "ascendcl", "profapi"]
FP8/FP4 探测机制：
- setup.py 中的 _acl_has_fp8_fp4() 函数通过编译探测 acl_base.h 是否声明了 ACL_FLOAT8_E5M2/ACL_FLOAT8_E8M0/ACL_FLOAT4_E2M1 枚举值。
- 如果支持，定义 -DCAM_ACL_HAS_FP8_FP4，使 pytorch_npu_helper.hpp 填充低精度 dtype 表。
- 这是为了兼容不同 CANN 版本（CANN 8.5 的 acl_base.h 枚举到 BF16 为止，CANN 9.0+ 才包含 FP8/FP4）。
4. Docker 部署环境完整分析
4.1 Dockerfile 多阶段构建
基础阶段（base）：
- 镜像：ubuntu:24.04
- 系统依赖：vim git cmake net-tools wget gcc g++ make iputils-ping iproute2 libnuma-dev + 开发库（zlib1g-dev libbz2-dev libssl-dev 等）
- Miniconda 安装到 /opt/conda
- pip 镜像：https://pypi.tuna.tsinghua.edu.cn/simple
- Python 环境：conda create -n cam_py311_pt28 python=3.11
- Python 包：attrs cython numpy==1.24.0 decorator sympy cffi pyyaml pathlib2 psutil protobuf==3.20 scipy requests absl-py pybind11 pytest setuptools_scm ninja
- gtest 1.16.0：从 GitHub 源码编译安装到 /usr/local
- OpenMPI 5.0.7：源码编译，prefix /usr/local/openmpi-5.0.7，启用 --enable-orterun-prefix-by-default
- nlohmann/json 3.12.0：源码编译安装到 /usr
CANN 阶段（二选一）：
| 阶段名 | CANN 版本 | 安装内容 |
|---|---|---|
config_a3_8.5_open |
8.5.T63 | Ascend-cann_8.5.0_linux-aarch64.run + Ascend-cann-A3-ops_8.5.0_linux-aarch64.run |
config_a3_9.0_open |
9.0.0 | Ascend-cann-toolkit_9.0.0_linux-aarch64.run + Ascend-cann-A3-ops_9.0.0_linux-aarch64.run |
最终阶段（final）：
- SHMEM 1.3.0：从 gitcode.com/cann/shmem.git 源码构建 whl + 安装到 /usr/local/Ascend/shmem/latest/
- 环境变量注入 /etc/bashrc：
- GTEST_ROOT=/usr/local
- PATH 添加 OpenMPI bin
- LD_LIBRARY_PATH 包含 OpenMPI、CANN lib64、CANN devlib
- 自动 source CANN set_env.sh + SHMEM set_env.sh
- 默认激活 cam_py311_pt28 conda 环境
4.2 docker-compose.yml
services:
cam-dev:
image: ${IMAGE_NAME}
container_name: ${CONTAINER_NAME}
privileged: true # 必须特权模式（NPU 设备访问）
tty: true
stdin_open: true
command: /bin/bash
volumes:
- /home:/home
- /ssd_0:/ssd_0
- /mnt:/mnt
- /sys:/sys
- /usr/local/Ascend/driver:/usr/local/Ascend/driver # NPU 驱动挂载
- /usr/local/bin/npu-smi:/usr/local/bin/npu-smi # NPU 管理工具
- /etc/ascend_install.info:/etc/ascend_install.info # NPU 安装信息
- /etc/localtime:/etc/localtime
devices:
- /dev/vfio/vfio # SR-IOV VF 设备
- /dev/uio0 ~ /dev/uio7 # UIO 设备（8 卡 NPU）
关键部署约束：
- privileged: true 是 NPU 设备访问的硬性要求。
- 挂载 /usr/local/Ascend/driver 使容器内可见宿主机的 NPU 驱动。
- devices 列表暴露 VFIO 和 UIO 设备给容器（用于直通 NPU 设备）。
- 默认网络为 bridge（cam-net）。
4.3 .env 配置
COMPOSE_PROJECT_NAME=cam90
IMAGE_NAME=cam-dev:a3_9.0_open
CONTAINER_NAME=my_cam_dev
注意： 注释提醒 COMPOSE_PROJECT_NAME 和 CONTAINER_NAME 在多容器同机部署时必须全局唯一。
5. .run 安装包与 .whl 打包机制
5.1 .run 包（算子运行时包）
产出路径：
output/cam/comm_operator/run/CAM_${soc_version}_${os_id}_${arch}.run
# 例如：CAM_ascend910_93_ubuntu24.04_aarch64.run
安装方式：
./CAM_ascend910_93_ubuntu24.04_aarch64.run --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp
source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/CAM/bin/set_env.bash
内部结构（npu_op_* 构建系统生成）：
- packages/vendors/CAM/op_impl/ai_core/tbe/ — 算子实现目录
- kernel/${compute_unit}/ — 编译好的内核二进制
- config/${compute_unit}/ — 算子配置 JSON
- ${vendor_name}_impl/dynamic/ — 动态 shape 实现（如果启用）
- framework/pytorch/ — npu_supported_ops.json（支持的算子列表）
5.2 .whl 包（Python 接口包）
产出路径：
output/cam/comm_operator/dist/umdk_cam_op_lib-${version}-cp311-cp311-linux_aarch64.whl
版本号格式： V_version.R_version.C_version.B_version（默认 209.0.0.B001，通过 CAM_WHL_VERSION 环境变量覆盖）。
安装方式：
pip install --force-reinstall ./output/cam/comm_operator/dist/umdk_cam_op_lib-*.whl
whl 构建关键机制：
- 使用 torch_npu.utils.cpp_extension.NpuExtension（华为扩展的 CMake 构建）。
- bdist_wheel_build.py 在构建后调用 auditwheel repair 修复 manylinux 标签，并排除 PyTorch 内部 so（libc10.so、libtorch.so 等）。
- 链接的 so：torch_npu、runtime、torch、ascendcl、profapi、gcov。
6. RPM spec 分析
6.1 总体结构
umdk.spec 是 UMDK 项目的顶层 RPM 打包规范，CAM 模块本身不单独打包 RPM，而是作为 UMDK 源码的一部分存在。
关键宏定义：
| 宏 | 默认值 | 含义 |
|---|---|---|
with_64kb |
aarch64 下可通过 -with_64kb 启用 |
64KB 页大小内核支持 |
build_all |
1 | 全量构建开关 |
with_urma/with_urpc/with_dlock/with_ums |
0 | 各子模块开关 |
kernel_version |
自动检测 | 内核版本（用于 kmod 构建） |
6.2 构建流程
%build
cmake ./src/ -DCMAKE_INSTALL_PREFIX=/usr \
-DBUILD_DATE=%{rpm_build_date} \
-DASAN="enable" # 仅 with asan
-DTSAN="enable" # 仅 with tsan
-DCODE_COVERAGE="enable" # 仅 with gcov
-DRELEASE_ENABLE="enable" # 仅 with release_enable
-DURMA_TEST="enable" # 仅 with test
-DBUILD_ALL="enable" # build_all
-DBUILD_URMA="enable" # build_all 或 with urma/urpc
-DBUILD_URPC="enable" # build_all 或 with urpc
-DBUILD_UMS="enable" # with ums
-DBUILD_DLOCK="enable" # with dlock
make %{?_smp_mflags}
6.3 子包划分
| 子包 | 内容 |
|---|---|
umdk-urma-lib |
liburma.so.*、liburma_common.so.*、liburma_ubagg.so.* |
umdk-urma-devel |
liburma.so、头文件 urma_*.h |
umdk-urma-tools |
urma_admin、urma_ping、urma_perftest |
umdk-urma-bin |
libtpsa.so、TPSA 配置 |
umdk-urpc-framework |
liburpc_framework.so.* |
umdk-urpc-umq |
libumq.so.*、libumq_buf.so.*、libumq_ub.so.* |
umdk-dlock-lib |
libdlockm.so.*、libdlocks.so.*、libdlockc.so.* |
umdk-ums |
ums.ko（内核模块） |
umdk-ums-tools |
libums-preload.so、ums_run |
umdk-ums-agent |
ums_agent 守护进程 + systemd service |
6.4 UMS 内核模块特殊处理
- 支持 64KB 页大小内核（
umdk-ums-64kb子包）。 %pre阶段检查运行时页大小是否匹配。%post阶段处理 weak-modules 兼容（内核版本不一致时）。- 自动排除 dracut（
omit_drivers+=" ums "）。
7. 第三方依赖（catlass submodule）
7.1 catlass 的角色
catlass（CANN Template Library for Ascend）是华为提供的 AscendC 内核模板库，为 fused_deep_moe 算子提供高性能 GEMM/Epilogue 实现。
7.2 目录结构（pin 在 v1-stable 分支）
src/cam/third_party/catlass/
└── include/ ← 头文件（通过 CPATH 注入）
7.3 在 fused_deep_moe 中的使用
fused_deep_moe 的 op_kernel/fused_deep_moe/ 目录包含大量基于 catlass 的模板：
| 文件 | 功能 |
|---|---|
grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h（24KB） |
基础 grouped GEMM（per-token dequant） |
grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h（78KB） |
SwiGLU + 量化融合 GEMM |
post_swiglu_dynamic_quant.h（12KB） |
SwiGLU 后动态量化 |
block_mmad_preload_async_with_callback_resident_a.h（20KB） |
MMAD 预取 + resident A 优化 |
block_epilogue_per_token_dequant.hpp（21KB） |
per-token dequant epilogue |
block_epilogue_per_token_dequant_swiglu.h（17KB） |
SwiGLU dequant epilogue |
tile_stride_binary.h（5KB） |
tile stride 二元操作 |
tile_stride_muls.h（3KB） |
tile stride 乘法 |
7.4 编译注入路径
- 顶层 build.sh：
export CPATH=${CAM_THIRD_PARTY_PATH}/catlass/include:${CPATH} - compile_ascend_proj.sh：从 CPATH 解析
CATLASS_HOME_PATH - op_kernel/CMakeLists.txt：
npu_op_kernel_options("" ALL OPTIONS -I${CATLASS_HOME_PATH}/include) - op_host/CMakeLists.txt：
npu_op_code_gen(COMPILE_OPTIONS -I${CATLASS_HOME_PATH}/include) - build.make patch：
sed -i "s|--compile-options=\"\"|--compile-options=\"-I${CATLASS_HOME_PATH}/include\"|g"
8. 构建产物与安装路径
8.1 构建产物清单
| 产物 | 路径 | 说明 |
|---|---|---|
.run 算子包 |
output/cam/comm_operator/run/CAM_<soc>_<os>_<arch>.run |
单 SOC 完整算子包 |
.whl Python 包 |
output/cam/comm_operator/dist/umdk_cam_op_lib-*.whl |
pybind 接口 |
| 解压的 run 包 | output/cam/comm_operator/extract/<soc>/ |
通过 -x 触发 |
| 中间工程 | build/cam/build_master/comm_operator/ascend_kernels_<soc>_proj/ |
msopgen 生成的工程 |
8.2 安装路径
| 组件 | 安装路径 |
|---|---|
| 算子包 | /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/CAM/ |
| Python 包 | site-packages/umdk_cam_op_lib/ |
| SHMEM | /usr/local/Ascend/shmem/latest/ |
| CANN | /usr/local/Ascend/ascend-toolkit/latest/ |
| OpenMPI | /usr/local/openmpi-5.0.7/ |
| gtest | /usr/local/ |
8.3 环境变量
# CANN
export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest
export ASCEND_CANN_PACKAGE_PATH=/usr/local/Ascend/ascend-toolkit/latest
source /usr/local/Ascend/ascend-toolkit/set_env.sh
# SHMEM
export SHMEM_HOME_PATH=/usr/local/Ascend/shmem/latest
source /usr/local/Ascend/shmem/latest/set_env.sh
# OpenMPI
export PATH=/usr/local/openmpi-5.0.7/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/openmpi-5.0.7/lib:...:${LD_LIBRARY_PATH}
# catlass（构建时）
export CPATH=${CAM_THIRD_PARTY_PATH}/catlass/include:${CPATH}
9. 关键设计洞察
9.1 include-list vs exclude-list 演进
compile_ascend_proj.sh 的注释明确说明了从 full-copy + exclude-list 到 include-list 的迁移原因：
Unlike the old full-copy + exclude_list approach, only the operators resolved by select_ops.py are copied here, which fundamentally avoids same-name operators (the fused_deep_moe family) overwriting each other's files.
这是 fused_deep_moe / fused_deep_moe_w4a8 共享文件名的直接后果——全量拷贝会导致文件覆盖。
9.2 build.make patch 的必要性
npu_op_ 构建系统的 simple_kernel_compile 不读取 custom_compile_options.ini，因此 --compile-options="" 默认为空。华为选择通过 sed patch .make 文件来注入 SHMEM/CATLASS 路径，这是一种非侵入式的 workaround*，避免了修改 CANN 构建系统源码。
9.3 算子身份 = 源码目录名
整个算子选择机制的核心设计是算子身份与源码目录名一一对应：
- operator_registry.json 中的列表项 = ascend_kernels/ 下的目录名
- select_ops.py 输出 = 目录名列表
- copy_ops_include() 按目录名拷贝
这种设计使得添加新算子只需：创建目录 → 在 registry 中添加条目 → 完成。
9.4 版本匹配的脆弱性
CANN ↔ torch-npu 的版本匹配是硬编码在 Dockerfile 中的： - CANN 8.5 → torch_npu 2.8.0.post2 - CANN 9.0 → torch_npu 2.8.0.post4
升级 CANN 必须同步修改 Dockerfile 中的 wheel URL。setup.py 的 FP8/FP4 探测机制部分缓解了运行时版本差异，但构建时依赖仍需人工维护。
9.5 CAM 不单独打包 RPM
CAM 的 .run 和 .whl 包是独立于 UMDK RPM 体系的。umdk.spec 中的 RPM 打包仅覆盖 URMA/URPC/DLock/UMS 等子模块，CAM 通过独立的安装脚本（run --install-path + pip install）部署。
10. 相关文件索引
| 文件 | 作用 |
|---|---|
build/cam/build.sh |
顶层构建入口 |
build/cam/comm_operator/build.sh |
comm_operator 模块构建 |
build/cam/comm_operator/compile_ascend_proj.sh |
单 SOC 编译主流程 |
build/cam/comm_operator/select_ops.py |
算子选择逻辑 |
build/cam/comm_operator/set_conf.py |
CMakePresets.json 修改 |
build/cam/comm_operator/build_pybind.sh |
whl 包构建 |
src/cam/comm_operator/ascend_kernels/operator_registry.json |
算子注册表 |
src/cam/comm_operator/ascend_kernels/CMakeLists.txt |
顶层 CMake |
src/cam/comm_operator/ascend_kernels/cmake_files/op_host/CMakeLists.txt |
host 库 CMake |
src/cam/comm_operator/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt |
kernel 库 CMake |
src/cam/comm_operator/ascend_kernels/cmake_files/cmake/func.cmake |
自定义 CMake 函数 |
src/cam/comm_operator/ascend_kernels/cmake_files/cmake/util/ascendc_impl_build.py |
AscendC 内核编译生成器 |
src/cam/comm_operator/pybind/setup.py |
pybind whl 构建配置 |
src/cam/comm_operator/pybind/pytorch_extension/bdist_wheel_build.py |
auditwheel 修复 |
docker/cam/Dockerfile |
多阶段构建（CANN 8.5/9.0） |
docker/cam/docker-compose.yml |
容器编排 |
docker/cam/.env |
容器配置 |
umdk.spec |
UMDK RPM 打包规范 |
.gitmodules |
catlass 子模块配置 |
src/cam/third_party/catlass/ |
CANN 模板库（submodule） |
RELEASE-NOTES.md / RELEASE-NOTES.ch.md |
版本发布说明 |
分析范围:
src/cam/comm_operator/ascend_kernels/下的 4 个核心 MoE EP 通信算子 目标硬件: Ascend 910_93 (A3) 分析日期: 2026-07-21 分析层级: op_api → op_host → op_kernel → pybind 全栈
目录
- 架构总览与设计哲学
- 共享基础设施分析
- dispatch_layout 算子
- notify_dispatch 算子
- moe_dispatch_normal 算子
- moe_combine_normal 算子
- Pybind 绑定层
- 通信模式全局分析
- 量化机制深度分析
- 关键性能优化点
- HCCL Window 内存布局
1. 架构总览与设计哲学
1.1 CAM 的定位
CAM (Communication Acceleration for Matrix) 是华为 Ascend NPU 面向 MoE (Mixture-of-Experts) Expert Parallelism 场景的通信加速库。其核心设计目标是：
- 消除 All-to-All 瓶颈: MoE 中 token 需要跨卡路由到对应 expert 所在的 rank，传统 All-to-All 通信开销巨大
- 零拷贝 (IPC): 利用 HCCL Window 机制实现卡间直接内存访问，避免显式收发
- 计算通信重叠: 通过状态位同步 + 双缓冲 (ping-pong) 实现流水线隐藏延迟
1.2 算子调用链
Python 前端
│
▼
Pybind 层 (TORCH_LIBRARY + EXEC_NPU_CMD)
│
▼
aclnn API 层 (op_api/aclnn_*.cpp) ─── 调用 aclnnInner* 内部实现
│
▼
Tiling 层 (op_host/*_tiling.cpp) ─── 计算分片参数、校验合法性、设置 HCCL 配置
│
▼
Kernel 层 (op_kernel/*.cpp) ─── 设备端执行，操作 HCCL Window + IPC
1.3 核心数据流
┌─────────────────────────────────────────────────────────┐
│ MoE EP Pipeline │
│ │
TopkIdx │ Notify MoeDispatch Expert MoeCombine │
(tokens×k) │ Dispatch Normal Compute Normal │
│ │ │ │ │ │ │
▼ │ ▼ ▼ ▼ ▼ │
┌────────┐ │ ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐ │
│Dispatch│────────▶│ │SendData│ │ X+Idx │ │Expert │ │Weighted│ │
│Layout │ │ │Offset │ │→Window │ │Forward │ │ Sum │ │
└────────┘ │ └────────┘ └────────┘ └────────┘ └────────┘ │
│ │
└─────────────────────────────────────────────────────────┘
2. 共享基础设施分析
2.1 CommArgs (utils/op_kernel/comm_args.h)
这是整个通信层的核心参数结构体：
struct CommArgs {
int rank; // 全局 rank ID
int localRank; // 本地 rank (fullmesh 内部编号)
int rankSize; // 全局 rank 数量 (最大 384)
int localRankSize; // 本地 fullmesh 互联 rank 数
GM_ADDR peerMems[CAM_MAX_RANK_SIZE]; // 初始化时获取的各卡 Buffer 地址
int64_t sendCountMatrix[CAM_MAX_RANK_SIZE * CAM_MAX_RANK_SIZE]; // All-to-All 计数矩阵
int64_t sendCounts[CAM_MAX_RANK_SIZE];
int64_t recvCounts[CAM_MAX_RANK_SIZE];
// ...
};
关键常量：
| 常量 | 值 | 含义 |
|---|---|---|
CAM_MAX_RANK_SIZE |
384 | 最大支持 384 张 NPU 卡 |
IPC_BUFF_MAX_SIZE |
100 MB | IPC 数据区最大 100MB |
IPC_DATA_OFFSET |
2 MB | IPC 前 2MB 作为 flag 区 |
UB_SINGLE_DMA_SIZE_MAX |
190 KB | 单次 DMA 最大数据量 |
PING_PONG_SIZE |
2 | 双缓冲深度 |
2.2 SyncCollectives (utils/op_kernel/sync_collectives.h)
这是卡间同步原语的核心实现，使用 HCCL Window 上的 flag 实现跨卡同步：
同步标志结构：
constexpr int64_t FLAG_UNIT_INT_NUM = 4; // 每个 flag 占 4 个 int64
constexpr int64_t SYNC_UNIT_SIZE = 4 * sizeof(int64_t); // 32 Bytes
constexpr int64_t MAGIC_OFFSET = 32; // magic 占高 32 位
核心机制 — Magic + Value 编码：
__aicore__ inline int64_t MergeMagicWithValue(int32_t magic, int32_t value) {
return (static_cast<int64_t>(static_cast<uint32_t>(magic)) << MAGIC_OFFSET)
| static_cast<int64_t>(value);
}
设计洞察: 高 32 位 magic 用于区分不同轮次（避免 ABA 问题），低 32 位 value 表示具体事件。这种无锁设计允许跨卡原子操作。
内存布局 (每个 rank 的 shareAddr)：
┌──────────────────────────────────────────────────────┐
│ 段 A: 卡内同步 flags (blockNum × FLAG_UNIT_INT_NUM) │
├──────────────────────────────────────────────────────┤
│ 段 B: 卡间同步 flags (blockNum × FLAG_UNIT_INT_NUM) │
├──────────────────────────────────────────────────────┤
│ Chunk Flags: IPC 传输完成标记 (per destRank) │
└──────────────────────────────────────────────────────┘
WaitOneRankPartFlag 实现 — 核心等待循环：
do {
DataCopy(localWait, globalWait[...], ...); // GM→UB 拷贝
WaitFlag<MTE2_S>(); // 等待 DMA 完成
for (i = 0; i < remainToCheck; ++i) {
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，减少 UB 读取开销。
2.3 DataCopy 工具 (utils/op_kernel/data_copy.h)
提供 GM↔UB、UB↔UB 的模板化拷贝函数：
- CpUB2GM: UB → Global Memory (含原子操作支持)
- CpGM2UB: Global Memory → UB
- CopyUB2UB: UB 内部拷贝
- SetAtomicOpType<T>(op): 设置原子操作类型（ADD/MUL/MAX/MIN）
2.4 Mc2TilingUtils (utils/op_host/mc2_tiling_utils.h)
关键功能：读取 HCCL_BUFFSIZE 环境变量，计算 Window 可用大小：
uint64_t maxWindowSize = defaultWindowSize * 1024UL * 1024UL; // 默认 200MB
2.5 TilingArgs (utils/op_host/tiling_args.h)
定义关键偏移常量：
constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 4UL * 1024UL * 1024UL; // Combine 状态区偏移
constexpr uint64_t NOTIFY_DISPATCH_WIN_OFFSET = 204UL * 1024UL * 1024UL; // NotifyDispatch 偏移
3. dispatch_layout 算子
3.1 功能定位
dispatch_layout (aclnnDispatchLayout) 是 MoE 路由的前置分析算子，纯计算、无通信。它根据 topkIdx 计算：
- 每个 rank 需要接收的 token 数 (
numTokensPerRank) - 每个 expert 需要处理的 token 数 (
numTokensPerExpert) - 每个 token 是否在某个 rank 内 (
isTokenInRank) - 用于 notify 的发送数据 (
notifySendData) - 小 sendTokenIdx (
sendTokenIdxSmall)
3.2 op_api 层
aclnnStatus aclnnDispatchLayout(void *workspace, uint64_t workspaceSize,
aclOpExecutor *executor, aclrtStream stream) {
if (NnopbaseSetHcclServerType) {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE); // 关键: 使用 MTE 模式
}
return aclnnInnerDispatchLayout(workspace, workspaceSize, executor, stream);
}
关键设计: 强制设置 HCCL Server Type 为 MTE (Memory Transfer Engine)，绕过 AICPU 中转。
3.3 op_host 层 (Tiling 逻辑)
输入/输出定义：
- Input: topkIdx (int64, [numTokens, numTopk])
- Attrs: num_tokens, num_ranks, num_experts, num_topk, local_ranksize
- Outputs: numTokensPerRank, numTokensPerExpert, isTokenInRank, notifySendData, sendTokenIdxSmall
核心参数校验：
- numRanks ∈ (0, 384]
- numExperts ∈ (0, 512]
- numTopk ∈ (0, 16]
- numTokens ∈ (0, 8000]
- 约束: numExperts % numRanks == 0（expert 必须均匀分布）
Workspace 计算：
workSpaces[0] = SYSTEM_NEED_WORKSPACE + KERNEL_USE_WORKSPACE; // 16MB + 1MB
Tiling Key 固定为 23。
3.4 op_kernel 层 (核心算法)
并行策略: 按 token 维度切分到多个 AIV core：
aivNum_ = numTokens_ <= maxAivNum ? numTokens_ : maxAivNum;
// 每个 core 处理 tempTokens_ 个 token
核心算法 — Process()：
Phase 1: 统计 (每个 core 独立)
for each token i in [startToken, endToken):
for each topk j:
expert_idx = topkIdx[i][j]
numTokensPerExpert[expert_idx]++
rank_id = expert_idx / experts_per_rank
if !seenRank[rank_id]:
numTokensPerRank[rank_id]++
isTokenInRank[i][rank_id] = 1
seenRank[rank_id] = 1
Phase 2: 全局聚合 (原子加)
SetAtomicAdd<T>()
for each other core i:
DataCopyPad(notifySendData[i * numExperts], numTokensPerExpert)
SetAtomicNone()
Phase 3: SyncAll 全卡同步
SyncAll<true>()
Phase 4: 计算 sendTokenIdxSmall
for each token i:
for each topk j:
expert_idx = topkIdx[i][j]
sendTokenIdxSmall[i][j] = aggregatedNumTokensPerExpert[expert_idx]++
算法洞察: Phase 2 使用
SetAtomicAdd而非显式 reduce，利用硬件原子操作实现跨 core 聚合，避免额外的同步开销。
UB 内存规划：
topkIdx32AlignIntLen_ = Ceil(tempTokens_ * numTopk_ * sizeof(int64_t), 32) * 32;
numTokensPerRank32AlignIntLen_ = Ceil(numRanks_ * sizeof(T), 32) * 32;
numTokensPerExpert32AlignIntLen_ = Ceil(numExperts_ * sizeof(T), 32) * 32;
isTokenInRank32AlignIntLen_ = Ceil(tempTokens_ * numRanks_ * sizeof(T), 32) * 32;
4. notify_dispatch 算子
4.1 功能定位
notify_dispatch (aclnnNotifyDispatch) 是元数据 All-to-All 通信算子。它在实际 token 数据传输前，先让每张卡知道：
- 其他卡要发送给自己多少 token (
recvCount) - 从哪个 offset 开始存放 (
recvOffset) - 本卡要给其他卡发送的 offset (
sendDataOffset) - 全局最大 batch size (
maxBs) - 每个本地 expert 接收的 token 总数 (
recvTokensPerExpert)
4.2 op_api 层
参数列表 (极复杂，反映通信元数据交换的本质)：
aclnnNotifyDispatchGetWorkspaceSize(
sendData, tokenPerExpertData, sendCount, numTokens,
commGroup, rankSize, rankId, localRankSize, localRankId,
sendDataOffset, recvData, totalRecvTokens, recvCount, recvOffset,
maxBs, recvTokensPerExpert, ...
)
4.3 op_host 层
HCCL 通信配置：
uint32_t opType1 = OP_TYPE_ALL_TO_ALL; // 8 = AlltoAll
std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise";
通信拓扑: level0 使用 fullmesh（适合机内全互联），level1 使用 pairwise（适合跨机逐步传输）。这是经典的两级 All-to-All 策略。
Tiling Key 选择（基于数据类型）: - Float16 → 20 - BFloat16 → 21 - Float → 22 - Int → 23 - Ascend910B 额外 +100
4.4 op_kernel 层 (核心算法)
NotifyDispatch 是 CAM 中最复杂的算子，它实现了自管理的 All-to-All（不依赖 HCCL 集合通信）：
四阶段流水线：
Process():
if (blockIdx < 1):
AssembleSendData() // Phase 0: Core 0 组装发送数据
SyncAll<true>()
if (blockIdx < coreNumPerStageX):
InputToShareSlice() // Phase 1: 将自己的数据写入其他卡的 Window
if (blockIdx < coreNumPerStageY):
ShareToShareSlice() // Phase 2: 从其他卡的 Window 读取数据
SyncAll<true>()
ReorderOutput() // Phase 3a: 重排接收到的数据
BuildTotalRecvTokens() // Phase 3b: 构建 totalRecvTokens (blockIdx==0)
BuildRecvCount() // Phase 3c: 构建 recvCount (blockIdx==1)
BuildRecvOffset() // Phase 3d: 构建 recvOffset (blockIdx==2)
BuildMaxBs() // Phase 3e: 构建 maxBs (blockIdx==3)
BuildRecvTokenPerExp() // Phase 3f: 构建 recvTokensPerExpert (blockIdx==4)
AssembleSendData 详解：
// 每个 expert 的 sendPerGroup = 3 个字段:
// [numTokensExpert, prefixSum, numTokens]
for (int i = 0; i < numExperts; ++i) {
int numTokensExpert = tokenPerExpertUB[i];
sendDataUB[i * sendPerGroup] = numTokensExpert; // 该 expert 有多少 token
sendDataUB[i * sendPerGroup + 1] = prefixSum; // 偏移量
sendDataUB[i * sendPerGroup + 2] = numTokens; // 总 token 数
sendDataOffsetUB[i] = prefixSum;
prefixSum += numTokensExpert;
}
InputToShareSlice — 核心 IPC 写入：
// 每个 core 负责对 copyLen 个 rank 进行写入
CpGM2GMPingPong<T>(copyLen * perRankDataNum * sizeof(T),
readGt, writeGt[copyOffset * perRankDataNum], COPYONLY);
// 写入完成后设置 flag 通知对端
uint64_t v = MergeMagicWithValue(magic, 1);
CpUB2GM(shareAddrs[i] + rank * FLAG_UNIT_INT_NUM, inputUB, sizeof(uint64_t));
设计亮点: 使用 Ping-Pong 双缓冲 GM→GM 拷贝，
UB_HEAD_OFFSET(96B) 和UB_MID_OFFSET(96B + 95KB + 32B) 交替使用，实现流水隐藏。
ShareToShareSlice — 读取对端数据：
// 等待本卡数据已被写入对端
WaitSyncFlag(magic, 1, copyOffset, rank, copyLen);
// 从对端 Window 读取属于本卡的数据
for (int i = 0; i < copyLen; i++) {
CpGM2GMPingPong<T>(perRankDataNum * sizeof(T),
readGt1[i][rank * perRankDataNum],
recvDataOutputGt[checkRank[i] * perRankDataNum], COPYONLY);
}
CpGM2GMPingPong — 核心 DMA 引擎：
// 双缓冲交替:
// 偶数轮: 用 inputUB[0] + outputUB[0]
// 奇数轮: 用 inputUB[1] + outputUB[1]
for (int64_t i = 0; dataSizeRemain > 0; i++) {
event_t eventId = (i & 1) ? EVENT_ID0 : EVENT_ID1;
WaitFlag<MTE3_MTE2>(eventId); // 等待上一轮写完成
CpGM2UB(inputUB[i&1], input, size); // GM→UB
CpUB2GM(output, outputUB[i&1], size); // UB→GM
SetFlag<MTE3_MTE2>(eventId); // 通知下一轮可开始
dataSizeRemain -= size;
}
4.5 同步机制
Magic 值获取：
__aicore__ inline uint64_t NotifyDispatch<T>::GetMagicValue(void) {
// 从 localWindowsExp + STATE_WIN_OFFSET + blockIdx * UB_ALIGN_SIZE 读取
magic = selfDataStatusTensor(blockIdx * UB_ALIGN_SIZE);
if (magic <= 0) magic = 1;
selfDataStatusTensor(blockIdx * UB_ALIGN_SIZE) = magic + 1; // 自增防 ABA
return magic;
}
Window 地址计算：
shareAddrs[rank] = GetWindAddrByRankId(rank, ctxIdx)
+ (magic % PING_PONG_SIZE) * (IPC_BUFF_MAX_SIZE + IPC_DATA_OFFSET);
Ping-Pong 双 Window: 交替使用两个 102MB+2MB 的 buffer，避免读写冲突。
5. moe_dispatch_normal 算子
5.1 功能定位
moe_dispatch_normal (aclnnMoeDispatchNormal) 是 token 数据实际分发算子。它将 hidden states 通过 HCCL Window 路由到对应 expert 所在的 rank。
5.2 op_api 层
aclnnMoeDispatchNormalGetWorkspaceSize(
x, topkIdx, sendOffset, sendTokenIdx, // 输入
recvOffset, recvCount, // 从 notify 获取的接收信息
groupEp, epWorldSize, epRankId, // EP 通信域
groupTpOptional, tpWorldSize, tpRankId, // 可选 TP 通信域
moeExpertNum, quantMode, globalBs, // 配置参数
recvX, recvXScales, assistInfoForCombine, // 输出
waitRecvCostStats, // 可选诊断
...
)
5.3 op_host 层 (Tiling 逻辑)
输入/输出多态（根据 quantMode 变化）：
| Index | Input/Output | quantMode=0 | quantMode=2 |
|---|---|---|---|
| 0 | x | BF16 | BF16 |
| 1 | topk_idx | INT32 | INT32 |
| 0 | recv_x | BF16 | INT8 |
| 1 | x_scales | FLOAT | FLOAT |
HCCL 通信配置：
// EP 域: All-to-All (fullmesh + pairwise)
std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise";
// TP 域: AllGather (ring)
std::string algConfigAllGatherStr = "AllGather=level0:ring";
Window 大小校验：
tokenActualLen = Ceil(h * 2, 32) * 32 + 44; // h_align_32B + scale(32B) + expandIdx(3*4B)
tokenNeedSizeDispatch = Ceil(tokenActualLen, 512) * 512;
tokenNeedSizeCombine = Ceil(h * 2, 512) * 512;
actualSize = (maxBs * k * (tokenNeedSizeCombine + tokenNeedSizeDispatch)
+ 4MB + 204MB) * 2;
// 校验 actualSize <= HCCL_BUFFSIZE
Tiling Key 计算：
tilingKey = 10000 + quantMode + (tpWorldSize == 2 ? 100 : 0);
// 10000: 无量化, 10002: 动态量化, 10100: TP=2 无量化
Workspace：
workSpaces[0] = 16MB + 512 * aivNum * aivNum; // 用于诊断统计
5.4 op_kernel 层 (核心算法)
模板参数化：
template <typename XType, typename ExpandXOutType,
bool dynamicQuant, bool IsSmoothScaleExist, bool IsShareExpertRank>
class MoeDispatchNormal { ... };
四阶段流水线 (Process)：
InputToShare() → 将 token 数据写入目标 rank 的 Window
SetStatus() → 设置目标 rank 的状态位 (通知数据已到达)
WaitStatus() → 等待所有源 rank 的数据到达
ShareToOutput() → 从 Window 读取数据到本地输出
InputToShare 详解：
// 每个 core 负责 [startTokenId, endTokenId) 的 token
for (tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
dstExpertId = expertIdsTensor(tokenIndex - startTokenId);
curExpertCnt = sendTokenIdxTensor(tokenIndex - startTokenId);
dstExpertOffset = sendOffsetTensor(dstExpertId);
// 计算目标 Window 地址
rankGM = shareGM + hOutGMAlignSize * (dstExpertOffset + curExpertCnt);
if (DynamicQuant) {
// 量化路径: BF16 → FP32 → Abs → ReduceMax → Scale → INT8
QuantProcess();
FillTriple(xOutTensor, tokenIndex / topK, tokenIndex % topK);
DataCopyPad(dstGT, xOutTensor, hCommuCopyOutParams);
} else {
// 非量化路径: 直接拷贝 + 填充三元组
FillTriple(xTmpTensor, tokenIndex / topK, tokenIndex % topK);
DataCopyPad(dstGT, xTmpTensor, hCommuCopyOutParams);
}
}
FillTriple — 元数据嵌入：
__aicore__ inline void FillTriple(LocalTensor<ExpandXOutType> &xOutTensor,
uint32_t tokenIndex, uint32_t k) {
LocalTensor<int32_t> xOutTint32 = xOutTensor.ReinterpretCast<int32_t>();
xOutTint32(expandIdxStartIdx) = epRankId; // 源 rank ID
xOutTint32(expandIdxStartIdx + 1) = tokenIndex; // 源 token index
xOutTint32(expandIdxStartIdx + 2) = k; // topk index
}
设计洞察: 在每个 token 的 hidden state 后面附加三元组 (srcRank, srcTokenIdx, topkIdx)，这样 combine 阶段无需额外通信即可知道数据来源。
SetStatus — 状态位同步：
// 将 status 写入目标 rank 的 Window 状态区
for (uint32_t i = startExpId; i < endExpId; ++i) {
targetRankId = i / moeExpertNumPerRank;
offset = stateOffset * (epRankId + i % moeExpertNumPerRank * epRankSize);
rankGM = GetWindStateAddrByRankId(COMM_EP_IDX, targetRankId) + offset;
DataCopy(dstStatusGT, statusTensor[(i - startExpId) * 8], 8UL);
}
WaitStatus — 等待数据到达：
while (sumOfFlag != compareTarget) {
DataCopy(statusFp32Tensor, windowInstatusFp32Tensor[startStatusId * stateOffset / 4], ...);
ReduceSum(statusSumOutTensor, statusFp32Tensor, gatherMaskOutTensor, ...);
sumOfFlag = statusSumOutTensor.GetValue(0);
}
// 清理状态
Duplicate<int32_t>(cleanStateTensor, 0, ...);
DataCopy(windowInstatusFp32Tensor[...], cleanStateTensor.ReinterpretCast<float>(), ...);
性能洞察: 使用
ReduceSum在 Vector 单元上快速求和所有状态位，当 sum == statusNumPerCore 时说明所有数据已到达。
ShareToOutput — 从 Window 读取：
for (uint32_t i = startStatusId; i < endStatusId; ++i) {
fromRank = i % epRankSize;
count = recvCountTensor(i) - preCount;
for (uint32_t j = 0; j < count; ++j++) {
// 从 Window 读取 token + expandIdx
DataCopyPad(xTmpTensor, srcTokenGT, hCommuCopyOutParams);
// 提取 expandIdx (三元组)
DataCopyPad(expandIdxOutGT[(targetOffset + j) * 3], xTmpTensorInt[expandIdxStartIdx], ...);
// 如果是量化模式, 提取 scale
if (DynamicQuant) {
DataCopyPad(dynamicScalesOutGT[targetOffset + j], xOutFp32Tensor[hUBAlignSize / 4], ...);
}
// 拷贝 token 到输出
DataCopyPad(dstTokenGT, xTmpTensor, expandXCopyParams);
}
}
5.5 量化机制 (QuantProcess)
__aicore__ inline void QuantProcess() {
Cast(floatLocalTemp, xInTensor, CAST_NONE, h); // BF16 → FP32
Abs(floatLocalAbsTemp, floatLocalTemp, h); // |x|
ReduceMaxInplace(floatLocalAbsTemp, h); // max(|x|)
dynamicScale = 127.0 / (floatLocalAbsTemp.GetValue(0) + 1e-12f);
Muls(floatLocalTemp, floatLocalTemp, dynamicScale, h); // x * scale
Cast(int32LocalTemp, floatLocalTemp, CAST_RINT, h); // FP32 → INT32 (四舍五入)
Cast(halfLocalTemp, int32LocalTemp, CAST_ROUND, h); // INT32 → FP16 (dequant)
Cast(xOutTensor, halfLocalTemp, CAST_TRUNC, h); // FP16 → INT8
// 存储 scale 供 combine 反量化
floatLocalTemp.SetValue(hUBAlignSize / sizeof(float), 1.0 / dynamicScale);
}
ReduceMaxInplace 优化：
// 两级 reduce: 先两两比较 (stride=8), 再 WholeReduceMax
if (repsFp32 > 1) {
Max(srcLocal, srcLocal[64], srcLocal, 64, repsFp32 - 1, {1,1,1,0,8,0});
}
WholeReduceMax(srcLocal, srcLocal, mask, 1, 8, 1, 8);
6. moe_combine_normal 算子
6.1 功能定位
moe_combine_normal (aclnnMoeCombineNormal) 是 token 加权聚合算子。Expert 计算完成后，将结果通过 Window 发回源 rank，并按 topk_weight 加权求和。
6.2 op_api 层
aclnnMoeCombineNormalGetWorkspaceSize(
recvX, tokenSrcInfo, epRecvCounts, recvTopkWeights, tpRecvCountsOptional,
epGroupName, epWorldSize, epRankId,
tpGroupNameOptional, tpWorldSize, tpRankId,
moeExpertNum, globalBs, out, sendCostStats, ...
)
6.3 op_host 层 (Tiling 逻辑)
输入/输出： - recv_x: [totalRecvTokens, h] (BF16/FP16) - token_src_info: [totalRecvTokens * 3] (INT32) — expandIdx - ep_recv_counts: [moeExpertNum] (INT32) - recv_topk_weights: [bs, k] (FLOAT) - x: [bs, h] (BF16/FP16)
HCCL 通信配置：
// EP: All-to-All (fullmesh + pairwise)
// TP: ReduceScatter (ring) ← 注意这里是 ReduceScatter 不是 AllGather
std::string algConfigReduceScatterStr = "ReduceScatter=level0:ring";
Tiling 约束：
// combine 不支持多 expert + TP 同时启用
OPS_ERR_IF((localMoeExpertNum > 1) && (tpWorldSize > 1),
"Cannot support multi-moeExpert in a rank when tpWorldSize > 1");
6.4 op_kernel 层 (核心算法)
两阶段流水线：
CopyBufferToShareAndSetStatus() → 将 expert 输出写入源 rank 的 Window，设置状态
ReadBufferFromRemote() → 等待数据到达，读取并按 weight 加权求和
CopyBufferToShareAndSetStatus 详解：
// 每个 core 负责 [startTokenId, endTokenId) 的 token
for (tokenIndex = startTokenId; tokenIndex < endTokenId; tokenIndex++) {
srcRankId = srcInfoLocal(index + 0); // 源 rank
srcTokenId = srcInfoLocal(index + 1); // 源 token
srcTopkId = srcInfoLocal(index + 2); // topk index
// 拷贝数据到源 rank 的 Window
CopyBufferToShare(srcRankId, srcTokenId, srcTopkId, tokenIndex);
// 设置状态位
SetStatusBySrcInfo(srcRankId, srcTokenId, srcTopkId);
}
CopyBufferToShare — Window 写入：
GM_ADDR dstGM = GetBufferAddrByRankId(srcRankId)
+ (srcTokenId * axisK_ + srcTopkId) * h512AlignRecvXLen_;
// 512 字节对齐存储，每个 slot 存一个 topk 结果
DataCopyPad(dstWindow, localCopyTensor, xOutCopyParams);
SetStatusBySrcInfo — 状态更新：
GM_ADDR stateGM = GetStateAddrByRankId(srcRankId)
+ (srcTokenId * axisK_ + srcTopkId) * UB_32_ALIGN;
// 写入 1.0 (0x3F800000) 表示该 slot 数据已就绪
DataCopy<uint32_t>(stateGMTensor, statusTensor, FLOAT_NUM_PER_ALIGN);
WaitBuffCopy — 等待数据到达：
target = axisK_ * FLOAT_NUM_PER_ALIGN; // 所有 topk 数据都到达
while (current != target) {
DataCopy<float>(stateTensorLocal, stateGMTensor, calCount);
Sum(tempStateTensorLocal, stateTensorLocal, sumPerKParams);
current = tempStateTensorLocal(0);
}
// 清零状态
Duplicate<float>(tempStateTensorLocal, 0.0, calCount);
DataCopy<float>(stateGMTensor, tempStateTensorLocal, calCount);
ReadBufferAndWeightedSum — 加权求和：
Duplicate(sumFloatBufLocal, 0, axisH_); // 初始化累加器
for (topkId = 0; topkId < axisK_; topkId++) {
scale = topkWeightsLocal.GetValue((tokenIndex - startTokenIndex) * axisK_ + topkId);
// 从本卡 Window 读取该 topk 对应的结果
localTokenAddr = localRankGM_ + (tokenIndex * axisK_ + topkId) * h512AlignRecvXLen_;
DataCopyPad(tmpToken, localTokenTensor, xOutCopyParams);
// 类型转换 + 乘权重 + 累加
Cast(tokenFloatLocal, tmpToken, CAST_NONE, axisH_); // BF16 → FP32
Muls(weightedMulBufLocal, tokenFloatLocal, scale, axisH_); // * weight
Add(sumFloatBufLocal, sumFloatBufLocal, weightedMulBufLocal, axisH_); // sum
}
// 最终结果转回 BF16
Cast(xOutLocal, sumFloatBufLocal, CAST_RINT, axisH_);
DataCopyPad(xOutGlobal_[tokenIndex * axisH_], xOutLocal, xOutCopyParams);
数值精度洞察: 累加使用 FP32 精度（避免 BF16 累加误差），最终输出时转回 BF16/FP16。这是混合精度计算的标准做法。
6.5 Magic 双缓冲机制
__aicore__ inline void MoeCombineNormal::InitMagic() {
selfMagicTensor.SetGlobalBuffer(
(__gm__ int32_t *)((GM_ADDR)epWinContext_->localWindowsExp
+ MAGIC_WIN_OFFSET + coreIdx_ * WIN_512_ALIGN));
DataCacheCleanAndInvalid<int32_t, ...>(selfMagicTensor);
magic_ = selfMagicTensor(0);
selfMagicTensor(0) = ((magic_ == 0) ? 1 : 0); // 0/1 翻转
}
// 使用 magic 选择 Window 的哪一半
winDataSizeOffset_ = magic_ * (totalWinSize / 2UL);
7. Pybind 绑定层
7.1 注册机制
使用 PyTorch 的 Dispatcher + TORCH_LIBRARY 双重注册：
// 1. pybind11 模块注册 (Python 直接调用)
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("get_dispatch_layout", &GetDispatchLayoutImplAutograd, ...);
m.def("moe_dispatch_prefill", &MoeDispatchPrefillImplAutograd, ...);
m.def("moe_combine_prefill", &MoeCombinePrefillImplAutograd, ...);
}
// 2. TORCH_LIBRARY Schema 注册 (供 GE 图引擎使用)
TORCH_LIBRARY(umdk_cam_op_lib, m) {
m.def("get_dispatch_layout(Tensor topk_idx, int num_experts, int num_ranks) -> (Tensor, Tensor)");
m.def("moe_dispatch_prefill(Tensor x, Tensor topk_idx, ...) -> (Tensor, Tensor, Tensor, Tensor, Tensor)");
m.def("moe_combine_prefill(Tensor x, Tensor topk_idx, ...) -> Tensor");
}
7.2 三层调用架构
Python 前端
│
▼
AutogradPrivateUse1 (支持反向传播)
│ MoeDispatchPrefill::apply() → forward() / backward()
▼
PrivateUse1 (NPU 实际执行)
│ MoeDispatchPrefillImplNpu() → EXEC_NPU_CMD
▼
aclnn API (CANN 运行时)
7.3 get_dispatch_layout 绑定
std::tuple<at::Tensor, at::Tensor> GetDispatchLayoutImplNpu(
const at::Tensor &topkIdx, int64_t numExperts, int64_t numRanks) {
// 类型转换: topkIdx → int64
at::Tensor topkIdxInt64 = topkIdx.scalar_type() == at::kLong ? topkIdx : topkIdx.to(at::kLong);
// 分配输出 tensor
auto numTokensPerExpert = at::zeros({numExperts}, at::dtype(at::kInt).device(device));
auto sendTokenIdxSmall = at::zeros({numTokens, numTopk}, at::dtype(at::kInt).device(device));
auto notifySendData = at::zeros({notifySendDataSize}, ...);
// 调用 aclnn API
EXEC_NPU_CMD(aclnnDispatchLayout, topkIdxInt64, numTokens, numRanks,
numExperts, numTopk, localRanksize,
numTokensPerRank, numTokensPerExpert, isTokenInRank,
notifySendData, sendTokenIdxSmall);
return std::make_tuple(numTokensPerExpert, sendTokenIdxSmall);
}
7.4 moe_dispatch_prefill 绑定
这是最复杂的 pybind 函数，它组合了 notify_dispatch + moe_dispatch_normal 两个算子：
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
MoeDispatchPrefillImplNpu(...) {
// Step 1: 准备 notify 的输入
auto sendData = torch::empty({numExperts * 3}, ...); // sendPerGroup=3
auto sendDataOffset = torch::empty({numExperts}, ...);
// Step 2: 调用 notify_dispatch (All-to-All 交换元数据)
EXEC_NPU_CMD(aclnnNotifyDispatch, sendData, numTokensPerExpert, sendCount, numTokens,
groupEpPtr, numRanks, rank, localRankSize, localRankId,
sendDataOffset, recvData, totalRecvToken, recvCount, recvOffset,
maxBs, recvTokensPerExpert);
// Step 3: 计算 globalBs 和 numRecvTokens
int64_t gBs = maxBs.item<int>() * numRanks;
int64_t trt = totalRecvToken.item<int>();
int numRecvTokens = (trt == 0) ? 1 : trt;
// Step 4: 分配 dispatch 输出
auto expandxOut = useQuant ? torch::empty({numRecvTokens, hidden}, at::kChar)
: torch::empty({numRecvTokens, hidden}, x.options());
auto dynamicScalesOut = torch::empty({numRecvTokens}, at::kFloat);
auto expandIdxOut = torch::empty({numRecvTokens * 3}, at::kInt);
// Step 5: 调用 moe_dispatch_normal (实际 token 传输)
EXEC_NPU_CMD(aclnnMoeDispatchNormal, x, expertIds, sendDataOffset, sendTokenIdxSmall,
recvOffset, recvCount, groupEpPtr, numRanks, rank,
groupEpPtr, tpSize, tpRank, numExperts, quantMode, gBs,
expandxOut, dynamicScalesOut, expandIdxOut, dispatchWaitRecvCostStatsOut);
return {expandxOut, dynamicScalesOut, expandIdxOut, recvCount, recvTokensPerExpert};
}
7.5 moe_combine_prefill 绑定
at::Tensor MoeCombinePrefillImplNpu(...) {
// 类型转换
at::Tensor topkIdxInt32 = topkIdx.scalar_type() == at::kInt ? topkIdx : topkIdx.to(at::kInt);
at::Tensor expertScales = topkWeights.scalar_type() == at::kFloat ? topkWeights : topkWeights.to(at::kFloat);
// 分配输出
auto combinedX = torch::empty({expertScales.size(0), hidden}, x.options());
// 调用 combine
EXEC_NPU_CMD(aclnnMoeCombineNormal, x, tokenSrcInfo, epSendCounts, expertScales,
tpSendCounts, groupEpPtr, numRanks, rank, groupEpPtr, tpWorldSize, tpRankId,
moeExpertNumber, globalBs, combinedX, combineSendCostStatsOut);
return combinedX;
}
7.6 异常处理 (torch_bind_exception.h)
class TorchBindException : public std::exception {
// 捕获文件名、行号、断言条件
message = "Failed: " + name + " error " + file + ":" + line + " error message: " + error;
};
#define TORCH_BIND_ASSERT(cond) \
if (not(cond)) { throw TorchBindException("Assertion", __FILE__, __LINE__, #cond); }
8. 通信模式全局分析
8.1 通信模式总结
| 算子 | 通信类型 | HCCL 操作 | 数据量 |
|---|---|---|---|
| dispatch_layout | 无 | 无 | 纯本地计算 |
| notify_dispatch | All-to-All | 自管理 IPC | 极小 (元数据) |
| moe_dispatch_normal | 间接写入 | Window IPC | 大 (hidden states) |
| moe_combine_normal | 间接写入 | Window IPC | 大 (expert 输出) |
8.2 自管理 All-to-All 的优势
CAM 的 notify_dispatch 不直接使用 HCCL 的 All-to-All API，而是自己管理 IPC：
- 精确控制: 直接操作 Window 地址，避免 HCCL 额外开销
- 流水线: Ping-Pong 双缓冲实现计算通信重叠
- 同步优化: Magic+Value 编码实现无锁同步
- 内存效率: 原地更新，无需额外 buffer
8.3 通信拓扑
Level 0 (机内): FullMesh
┌───┐ ┌───┐ ┌───┐ ┌───┐
│ 0 │──│ 1 │──│ 2 │──│ 3 │
└───┘ └───┘ └───┘ └───┘
│ │ │ │
┌───┐ ┌───┐ ┌───┐ ┌───┐
│ 4 │──│ 5 │──│ 6 │──│ 7 │
└───┘ └───┘ └───┘ └───┘
Level 1 (跨机): Pairwise
Machine 0 ←──→ Machine 1 ←──→ Machine 2 ...
9. 量化机制深度分析
9.1 量化模式
enum { NO_SCALES = 0, DYNAMIC_SCALES = 2 };
- quantMode=0: 无量化，BF16/FP16 直通
- quantMode=2: 动态 INT8 量化（per-token 动态 scale）
9.2 量化流程 (Dispatch 阶段)
BF16 Token → FP32 → |x| → max(|x|) → scale = 127/max → x*scale → INT8
↓
存储 1/scale 供反量化
9.3 反量化流程 (Combine 阶段)
// 如果 combine 需要反量化 (当前代码中 combine 直接操作 recvX)
Cast(tokenFloatLocal, tmpToken, CAST_NONE, axisH_); // INT8 → FP32 (如果量化)
Muls(weightedMulBufLocal, tokenFloatLocal, scale, axisH_);
9.4 量化对内存的影响
非量化模式 (per token):
[h * 2B (BF16)] + [32B (scale)] + [12B (expandIdx)] = h*2 + 44 Bytes
量化模式 (per token):
[h * 1B (INT8)] + [32B (scale)] + [12B (expandIdx)] = h + 44 Bytes
带宽节省: 当 h=7168 时，非量化 14380B → 量化 7212B，节省约 50% Window 带宽。
10. 关键性能优化点
10.1 双缓冲 (Ping-Pong)
NotifyDispatch 的 CpGM2GMPingPong：
__ubuf__ U *inputUB[2] = {UB_HEAD_OFFSET, UB_MID_OFFSET};
for (int64_t i = 0; dataSizeRemain > 0; i++) {
event_t eventId = (i & 1) ? EVENT_ID0 : EVENT_ID1;
WaitFlag<MTE3_MTE2>(eventId); // 等待写完成
CpGM2UB(inputUB[i&1], input, size);
CpUB2GM(output, outputUB[i&1], size);
SetFlag<MTE3_MTE2>(eventId); // 通知可开始下一轮
}
Combine 的 Magic 双 Window：
winDataSizeOffset_ = magic_ * (totalWinSize / 2UL); // 交替使用两半
10.2 地址对齐
所有 Window 访问都严格对齐到 512B (WIN_ADDR_ALIGN)：
hOutGMAlignSize = Ceil(hScaleIdxSize, WIN_ADDR_ALIGN) * WIN_ADDR_ALIGN;
h512AlignRecvXLen_ = Ceil(hRecvXTypeLen_, WIN_512_ALIGN) * WIN_512_ALIGN;
原因: HCCL Window 的跨卡访问需要 512B 对齐以达到最佳 DMA 性能。
10.3 增量 Flag 检查
checkedFlagNum += i; // 跳过已确认的 flag
避免重复检查已完成的 flag，减少 GM→UB 读取。
10.4 状态位批量 ReduceSum
// WaitStatus: 使用 ReduceSum 快速判断所有数据是否到达
ReduceSum(statusSumOutTensor, statusFp32Tensor, gatherMaskOutTensor, ...);
sumOfFlag = statusSumOutTensor.GetValue(0);
10.5 原子操作聚合
// dispatch_layout 中使用原子加实现跨 core 聚合
AscendC::SetAtomicAdd<T>();
for (int i = coreIdx_ + 1; i < aivNum_; ++i) {
DataCopyPad(tempExpertGM_[i * numExperts], numTokensPerExpertTensor, ...);
}
AscendC::SetAtomicNone();
10.6 诊断模式 (可选)
if (isEnableDiagnose) {
int32_t durationTime = (GetSystemCycle() - systemCycleStart) / CYCLE_TO_TIME; // us
waitRecvCostStatsTensor.SetValue(srcRank, preTime + durationTime);
}
设计亮点: 诊断输出是可选的 (OPTIONAL tensor)，不影响正常路径性能。
11. HCCL Window 内存布局
11.1 整体布局
┌──────────────────────────────────────────────────────────────────────────┐
│ HCCL Window (默认 200MB) │
├──────────────────────────────────────────────────────────────────────────┤
│ 0x000000 │ NotifyDispatch 数据区 (IPC Data, 100MB) │
│ │ ├── Flag 区 (2MB) │
│ │ └── Data 区 (100MB) │
├──────────┼───────────────────────────────────────────────────────────────┤
│ 0xCC0000 │ Combine 状态区 (4MB) — COMBINE_STATE_WIN_OFFSET │
├──────────┼───────────────────────────────────────────────────────────────┤
│ 0x1900000│ Dispatch 数据区 — NOTIFY_DISPATCH_WIN_OFFSET (204MB) │
│ │ ├── Token Data (ping) │
│ │ └── Token Data (pong) │
├──────────┼───────────────────────────────────────────────────────────────┤
│ │ Combine 数据区 │
│ │ ├── Token Data (ping) │
│ │ └── Token Data (pong) │
├──────────┼───────────────────────────────────────────────────────────────┤
│ │ 状态/Flag 区 │
│ │ ├── DataState (ping/pong 标志) │
│ │ ├── Magic (combine 双缓冲标志) │
│ │ └── Per-token Status (数据到达标志) │
└──────────┴───────────────────────────────────────────────────────────────┘
11.2 Token 存储格式 (Dispatch)
每个 token slot (512B 对齐):
┌──────────────────────────────────────────────────────┐
│ [0, h*2) │ BF16 hidden states (或 INT8 量化) │
│ [h*2, h*2+32) │ Scale (FP32, 量化模式) │
│ [h*2+32, h*2+44)│ ExpandIdx: [srcRank, srcToken, topk] │
│ [h*2+44, 512n) │ Padding │
└──────────────────────────────────────────────────────┘
11.3 Token 存储格式 (Combine)
每个 token slot (512B 对齐):
┌──────────────────────────────────────────────────────┐
│ [0, h*2) │ Expert 输出 (BF16) │
│ [h*2, 512n) │ Padding │
└──────────────────────────────────────────────────────┘
每个 slot 对应一个 status flag (32B 对齐):
┌──────────────────────────────────────────────────────┐
│ [0, 32) │ 1.0 (数据就绪) / 0.0 (未就绪) │
└──────────────────────────────────────────────────────┘
附录 A: 关键常量速查表
| 常量 | 值 | 来源文件 | 含义 |
|---|---|---|---|
| CAM_MAX_RANK_SIZE | 384 | comm_args.h | 最大 rank 数 |
| IPC_BUFF_MAX_SIZE | 100 MB | comm_args.h | IPC 数据区大小 |
| IPC_DATA_OFFSET | 2 MB | comm_args.h | IPC flag 区大小 |
| UB_SINGLE_DMA_SIZE_MAX | 190 KB | comm_args.h | 单次 DMA 上限 |
| UB_SINGLE_PING_PONG_ADD_SIZE_MAX | 95 KB | comm_args.h | Ping-Pong 单缓冲 |
| COMBINE_STATE_WIN_OFFSET | 4 MB | tiling_args.h | Combine 状态偏移 |
| NOTIFY_DISPATCH_WIN_OFFSET | 204 MB | tiling_args.h | Dispatch 数据偏移 |
| WIN_STATE_OFFSET | 500 KB/950 KB | moe_dispatch_normal.h | 状态区偏移 |
| MAGIC_WIN_OFFSET | 975 KB | moe_combine_normal.h | Magic 标志偏移 |
| WIN_ADDR_ALIGN | 512 B | 多处 | Window 地址对齐 |
| FLAG_UNIT_INT_NUM | 4 | sync_collectives.h | 每个 flag 占 4 个 int64 |
| STATE_OFFSET | 32 | moe_dispatch_normal.h | 状态位间隔 |
附录 B: 文件索引
src/cam/comm_operator/
├── ascend_kernels/
│ ├── AddCustom.json # 算子定义 (AddCustom 占位)
│ ├── operator_registry.json # 算子注册表
│ ├── dispatch_layout/
│ │ ├── op_api/aclnn_dispatch_layout.{h,cpp}
│ │ ├── op_host/dispatch_layout.cpp # 算子定义
│ │ ├── op_host/dispatch_layout_tiling.cpp # Tiling 逻辑
│ │ └── op_kernel/
│ │ ├── dispatch_layout.{h,cpp} # Kernel 实现
│ │ └── dispatch_layout_tiling.h # Tiling 数据结构
│ ├── notify_dispatch/
│ │ ├── op_api/aclnn_notify_dispatch.{h,cpp}
│ │ ├── op_host/notify_dispatch.cpp
│ │ ├── op_host/notify_dispatch_tiling.cpp
│ │ └── op_kernel/notify_dispatch.{h,cpp}
│ ├── moe_dispatch_normal/
│ │ ├── op_api/aclnn_moe_dispatch_normal.{h,cpp}
│ │ ├── op_host/moe_dispatch_normal.cpp
│ │ ├── op_host/moe_dispatch_normal_tiling.cpp
│ │ └── op_kernel/moe_dispatch_normal.{h,cpp}
│ ├── moe_combine_normal/
│ │ ├── op_api/aclnn_moe_combine_normal.{h,cpp}
│ │ ├── op_host/moe_combine_normal.cpp
│ │ ├── op_host/moe_combine_normal_tiling.cpp
│ │ └── op_kernel/moe_combine_normal.{h,cpp}
│ └── utils/
│ ├── op_host/
│ │ ├── tiling_args.h # 关键偏移常量
│ │ ├── mc2_tiling_utils.h # Window 大小计算
│ │ ├── dfx_base.h # 日志基础设施
│ │ ├── ops_log.h / ops_error.h / error_log.h
│ └── op_kernel/
│ ├── comm_args.h # CommArgs 结构体
│ ├── sync_collectives.h # 跨卡同步原语
│ ├── data_copy.h # GM↔UB 拷贝
│ └── moe_distribute_base.h # HCCL 拓扑结构
└── pybind/
├── pybind.cpp # pybind11 模块 + TORCH_LIBRARY
├── functions.h # 函数声明
├── get_dispatch_layout.cpp # dispatch_layout 绑定
├── moe_dispatch_prefill.cpp # dispatch 绑定 (notify + dispatch)
├── moe_combine_prefill.cpp # combine 绑定
└── torch_bind_exception.h # 异常处理
附录 C: 关键洞察总结
C.1 设计范式
- Window-based IPC: 所有跨卡数据交换通过 HCCL Window 直接内存访问，避免显式 Send/Recv
- Flag-based Synchronization: 使用 Magic+Value 编码的 flag 实现无锁同步
- Two-phase Pipeline: 先交换元数据 (notify)，再传输实际数据 (dispatch)
- Self-managed All-to-All: notify_dispatch 自己管理 IPC，不依赖 HCCL 集合通信
C.2 性能关键路径
Token 路由延迟 = notify_dispatch (All-to-All 元数据)
+ dispatch_normal (Window 写入 + 状态同步)
+ expert_compute
+ combine_normal (Window 读取 + 加权求和)
C.3 与竞品对比洞察
| 特性 | CAM (华为) | NCCL (NVIDIA) | 差异 |
|---|---|---|---|
| All-to-All | 自管理 IPC | NCCL 原生 | CAM 更精细控制 |
| 同步 | Magic+Flag | NCCL 内部 | CAM 显式可见 |
| 量化 | 内置 INT8 | 需外部实现 | CAM 端到端 |
| 双缓冲 | Magic 翻转 | 无 | CAM 避免冲突 |
C.4 潜在优化方向
- 量化 Combine: 当前 combine 阶段未利用 INT8 反量化，可在 Window 读取时直接反量化
- TP 支持: 当前 TP 仅支持 ring AllGather/ReduceScatter，可探索 2D-HMesh
- 动态 Topk: 当前 topk 固定，可探索动态 topk 路由
- Sparse Expert: 当前 expert 均匀分布，可探索负载均衡路由
CAM (Communication Acceleration for Matrix) 皇冠上的明珠 — 将 Dispatch + GMM1 + Swiglu + GMM2 + Combine 融合为单一算子的终极优化
目录
- 整体架构与融合策略
- op_host Tiling 策略
- op_kernel 执行流程
- raw_distributed 子模块
- gemm 子模块深度分析
- epilogue 子模块
- Swiglu + Quantization 融合机制
- 性能优化关键点
- 与标准 MoE (DeepSeek-style) 的对应关系
1. 整体架构与融合策略
1.1 算子定义
FusedDeepMoE 是 CAM 中最复杂的融合算子，将 MoE Decode 阶段的完整计算流水线融合进单一 kernel：
标准小算子组合 (5次kernel launch + 多次GM读写):
Dispatch → GMM1 → Swiglu → GMM2 → Combine
FusedDeepMoE (1次kernel launch, 全在片上完成):
Dispatch → GMM1 → Swiglu → GMM2 → Combine (全部融合)
1.2 输入/输出签名
输入 (14个): | 索引 | 名称 | 类型 | 说明 | |------|------|------|------| | 0 | x | bf16/fp16 | 输入 token [bs, h] | | 1 | expert_ids | int32 | 专家索引 [bs, topk] | | 2 | gmm1_weight | int8 (TensorList) | 路由专家 GMM1 权重 | | 3 | gmm1_weight_scale | fp32/bf16/fp16 | GMM1 权重 scale | | 4 | gmm2_weight | int8 (TensorList) | 路由专家 GMM2 权重 | | 5 | gmm2_weight_scale | fp32/bf16/fp16 | GMM2 权重 scale | | 6 | expert_scales | fp32 | token级专家权重 [bs, topk] | | 7-8 | share_gmm1_weight/scale | int8/fp32 | 共享专家 GMM1 权重 (可选) | | 9-10 | share_gmm2_weight/scale | int8/fp32 | 共享专家 GMM2 权重 (可选) | | 11 | expert_smooth_scales | fp32 | 平滑 scale (可选) | | 12 | share_smooth_scales | fp32 | 共享专家平滑 scale (可选) | | 13 | x_active_mask | bool | token 激活掩码 (可选) |
输出 (3个):
- output: [bs, h] — MoE 路由输出
- share_output: [bs, h] — 共享专家输出
- expert_token_nums: [moeExpertNumPerRank] — 每个专家接收的 token 数
属性 (6个):
- group_ep, ep_rank_size, ep_rank_id, moe_expert_num, quant_mode, global_bs
1.3 为什么要融合？
核心动机：消除中间数据的 Global Memory 往返
在标准小算子组合中： 1. Dispatch 写出 expand_x, dynamic_scales 到 GM → GMM1 读回 (读写 2次) 2. GMM1 写出 int32 累加结果到 GM → Swiglu 读回 (读写 2次) 3. Swiglu 写出 int8 + scale 到 GM → GMM2 读回 (读写 2次) 4. GMM2 写出 结果到 GM → Combine 读回 (读写 2次)
总计至少 8 次 GM 访问，对于 Decode 阶段（batch小、延迟敏感），GM 带宽是瓶颈。
FusedDeepMoE 的融合策略： - X1 (GMM1输入)：AIV 在 UB 中完成 Dispatch 量化后，直接通过 workspace 偏移量传递给 AIC 做 GMM1 - Swiglu 输出：GMM1 的 int32 结果在 L0C/UB 中直接做 Swiglu + 动态量化，写出 float 到 workspace - X2 (GMM2输入)：Swiglu 后立即做动态量化为 int8，通过 workspace 传递给 GMM2 - GMM2 输出：直接在 UB 中完成 Combine 的 AlltoAll 发送
关键收益：中间数据 X1、SwigluOut、X2 全部在片上 workspace 中流转，不访问 GM。
2. op_host Tiling 策略
2.1 核心 Tiling 数据结构
// fused_deep_moe_tiling.h
struct FusedDeepMoeInfo {
uint32_t epRankSize; // EP 并行度
uint32_t epRankId; // 当前 rank
uint32_t moeExpertNum; // 总专家数
uint32_t moeExpertNumPerRank; // 每 rank 专家数
uint32_t quantMode; // 量化模式
uint32_t globalBs; // 全局 batch size
uint32_t bs; // 本地 batch size
uint32_t k; // topK
uint32_t h; // hidden size
uint32_t aicNum; // AIC 核数
uint32_t aivNum; // AIV 核数
uint64_t gmm1HLen; // GMM1 隐藏维度 (intermediate*2)
uint64_t shareGmm1HLen; // 共享专家 GMM1 隐藏维度
bool isTensorList; // 权重是否为 TensorList
};
2.2 Workspace 布局计算
Workspace 是 FusedDeepMoE 的核心设计，所有中间数据在此共享：
// 各段偏移量计算 (启用 ENABLE_REUSE_MEMORY 时复用内存)
WorkSpaceOffset:
├── shareX1TokenOffset / x1TokenOffset // GMM1 输入 (int8)
├── shareX1ScaleOffset / x1ScaleOffset // GMM1 输入 per-token scale
├── swapSpaceOffset / shareSwigluOffset // GMM1 输出 (int32) / Swiglu 输出 (fp32)
├── shareX2TokenOffset / x2TokenOffset // GMM2 输入 (int8)
├── shareX2ScaleOffset / x2ScaleOffset // GMM2 输入 per-token scale
├── y2TokenOffset // GMM2 输出 (浅融合使用)
├── groupListOffset // 专家 token 数前缀和
├── expandIdxOffset // Dispatch 索引
├── epSendCountOffset // AlltoAll 发送计数
└── reservedOffset // 保留空间 (256KB)
内存复用策略 (ENABLE_REUSE_MEMORY):
- X1 和 X2 共享同一块 workspace 空间（因为 X1 和 X2 生命周期不重叠）
- Scale 空间同理复用
- SwapSpace 和 Swiglu 输出复用
2.3 Tile Shape 常量
// GMM1 Tile Shape
GMM1_L1M = 256, GMM1_L1N = 128, GMM1_L1K = 512
GMM1_L0K = 128, GMM1_EPIM = 64
GMM1_SWIZZLE_OFFSET = 3, GMM1_SWIZZLE_DIRECTION = 0
// GMM2 Tile Shape
GMM2_L1M = 128, GMM2_L1N = 256, GMM2_L1K = 512
GMM2_L0K = 128, GMM2_EPIM = 32
GMM2_L1A_STAGES = 4, GMM2_L1B_STAGES = 2 // GMM2 使用 ResidentA 优化
// 通用
WORKSPACE_STAGES = 4
CUSTOM_PRELOAD_STAGES = 1
CUSTOM_L1_STAGES = 2, CUSTOM_L0A_STAGES = 2, CUSTOM_L0B_STAGES = 2
2.4 Tiling Key 路由机制
使用华为 ASCENDC_TPL 机制，通过 5 个 bool 标志 + 1 个 ARCH 参数组合出多种执行路径：
// 关键执行标志
EXEC_FLAG_DEEP_FUSE = 1 << 0 // 是否深度融合 (moeExpertNumPerRank != 1)
EXEC_FLAG_TENSOR_LIST = 1 << 1 // 权重是否为 TensorList
EXEC_FLAG_X_ACTIVE_MASK = 1 << 2 // 是否启用 x_active_mask
EXEC_FLAG_SHARED_EXPERT = 1 << 3 // 是否有共享专家
EXEC_FLAG_SMOOTH_QUANT = 1 << 4 // 是否启用 smooth quant
Tiling Key 决定： - 是否跳过独立 Dispatch（深度融合时 Dispatch 在 kernel 内完成） - 是否走共享专家路径 - 是否启用 smooth scale 处理
3. op_kernel 执行流程
3.1 Kernel 入口
// fused_deep_moe.cpp
template <bool TPL_IS_DEEP_FUSE, bool TPL_IS_TENSOR_LIST, ...>
__global__ __aicore__ void fused_deep_moe(...) {
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); // 1C2V 模式
// 重组 EXEC_FLAG 位图
constexpr uint32_t EXEC_FLAG = (TPL_IS_DEEP_FUSE ? EXEC_FLAG_DEEP_FUSE : 0) | ...;
// 根据 TPL_IS_SHARED_EXPERT 选择 tiling struct
// 实例化 FusedDeepMoe<...> op; op.Init(...); op.Process();
}
1C2V 模式：每个 AI Core 组包含 1 个 AIC (AI Core Compute) + 2 个 AIV (AI Core Vector)。AIC 负责 GMM 计算，AIV 负责通信调度。
3.2 FusedDeepMoe::Process() 主流程
void Process() {
// 1. 非深度融合时，先执行独立 Dispatch
if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
if (g_coreType == AIV) {
CamMoeDistributeDispatch dispatcher;
dispatcher.Init(...);
dispatcher.Process(); // AlltoAll Dispatch
}
PipeBarrier<PIPE_ALL>();
CrossCoreBarrier/Wait(); // AIV→AIC 同步
}
// 2. GMM1 + Swiglu + 动态量化 (融合执行)
GmmDeqSwigluQuant<Gmm1L1TileShape, ...>(...);
PipeBarrier<PIPE_ALL>();
CrossCoreBarrier/Wait(); // AIC→AIV 同步
// 3. 初始化 Combine (AIV 端)
CamMoeDistributeCombine combiner;
combiner.Init(...);
// 4. GMM2 + Combine (融合执行)
GmmDeq<Gmm2L1TileShape, ...>(..., &combiner);
}
3.3 GmmDeqSwigluQuant 函数
这是 GMM1 的模板化入口，定义了完整的 GMM1 计算流水线：
template <..., class L1TileShape_, class L0TileShape_, class EpilogueTileShape_, class BlockScheduler_>
CATLASS_DEVICE void GmmDeqSwigluQuant(...) {
// 数据类型定义
using AType = Gemm::GemmType<int8_t, layout::RowMajor>; // 输入 X1
using BType = Gemm::GemmType<int8_t, layout::zN>; // 权重 (NZ格式)
using CType = Gemm::GemmType<int32_t, layout::RowMajor>; // 累加器
// BlockMmad: L1→L0 多级流水
using BlockMmad = Gemm::Block::BlockMmad<DispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType>;
// Epilogue: PerTokenDequantSwiglu — 核心融合点!
using EpilogueDispatchPolicy = Epilogue::EpilogueAtlasA2PerTokenDequantSwiglu<ubStages, 0>;
using BlockEpilogue = Epilogue::Block::BlockEpilogue<...>;
// Kernel: GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
using GemmKernel = Gemm::Kernel::GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace<...>;
GemmKernel gemm;
gemm(params);
}
3.4 GmmDeq 函数
GMM2 的模板化入口，结构类似但使用不同的 Epilogue：
// GMM2 使用 PerTokenDequantCombine epilogue
using EpilogueDispatchPolicy = Epilogue::EpilogueAtlasA2PerTokenDequantCombine<ubStages, EXEC_FLAG>;
// 注意 GMM2 使用 ResidentA Dispatch Policy (A矩阵常驻 L1)
using Gmm2DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncWithCallbackResidentA<...>;
4. raw_distributed 子模块
4.1 CamMoeDistributeDispatch
文件: raw_distributed/cam_moe_distribute_dispatch.h
负责将 token 按 expert_id 路由到对应 rank 的窗口内存中。
核心流程 (Process):
void Process() {
if ASCEND_IS_AIV {
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT)
QuantForShareExpert(); // 共享专家量化
AlltoAllDispatch(); // 路由专家 AlltoAll 分发
SetStatus(); // 通知对端数据就绪
WaitDispatch(); // 等待所有 rank 数据到达
LocalWindowCopy(); // 从窗口拷贝到本地 tensor
if constexpr (IsNeedAllgater) {
AllGatherSetStatusAndWait();
AllgatherProcessOut();
}
UpdataTokenNumsOut(); // 更新 expert_token_nums
}
}
关键优化 — AIV 循环优化 (enableAivOpt):
// 当 batch_size <= 64 且 expert_num 对齐后 <= 256 时启用
// 使用 table 法替代逐 token 计数，减少 UB 访问
if (axisBS_ <= LOOP_OPT_MAX_BS && moeExpertRankNumAligned_ <= LOOP_OPT_MAX_MOE_RANK) {
enableAivOpt_ = true;
// 构建 (bs+1) × expert_num 的 table，用前缀和计算 expand_idx
}
量化融合：Dispatch 阶段同时完成动态量化（bf16→int8），将 scale 附加在 token 后面传输：
// hCommuSize = hOutSize + scaleParamPad (128B for scale)
hCommuSize_ = axisH_ * sizeof(ExpandXOutType) + scaleParamPad_;
4.2 CamMoeDistributeCombine
文件: raw_distributed/cam_moe_distribute_combine.h
负责将 GMM2 计算结果路由回原始 rank，并做加权求和。
核心流程:
void Process() {
SyncAll<true>();
if constexpr (IsNeedReduceScatter) ReduceScatterTrans();
if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
BuffInit();
SetWaitTpStatusAndDisPatch(); // 非融合模式
}
AlltoAllBuffInit();
SetStatus();
WaitDispatch();
LocalWindowCopy(); // 从窗口读取 + 加权求和 + 写出
}
深度融合时的特殊路径:
// 融合模式下，Combine 被拆分为两部分：
void AllToAllSend() {
// AIV subblock 0: AlltoAll 发送
// AIV subblock 1: ReducePermute
if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
}
}
void ReducePermute() {
// 在 GMM2 的 AIV 端执行
AlltoAllBuffInit();
WaitDispatch();
LocalWindowCopy();
}
加权求和 (LocalWindowCopy):
// 对每个 token，遍历 topK 个专家的结果，加权求和
for (uint32_t tokenIndex = beginIndex; tokenIndex < endIndex; ++tokenIndex) {
Duplicate(sumFloatBufLocal, 0, axisH_);
for (uint32_t i = 0; i < axisK_; i++) {
float scaleVal = expandScalesLocal.GetValue(index);
// 从对应专家的窗口位置读取
rowTmpGlobal_.SetGlobalBuffer(wAddr + expertPerSizeOnWin_ * moeExpert + ...);
Cast(rowTmpFloatLocal, tmpUb, CAST_NONE, processLen);
Muls(mulBufLocal, rowTmpFloatLocal, scaleVal, processLen);
Add(sumFloatBufLocal, sumFloatBufLocal, mulBufLocal, processLen);
}
Cast(sumBufLocal, sumFloatBufLocal, CAST_RINT, processLen);
DataCopy(expandOutGlobal_[tokenIndex * axisH_ + tokenOffset], sumBufLocal, processLen);
}
5. gemm 子模块深度分析
5.1 Catlass 集成
FusedDeepMoE 使用 Catlass (CUDA Template Library analogue for Ascend) 作为 GMM 计算框架。Catlass 提供：
- 类型化的矩阵类型:
GemmType<Element, Layout> - Block 级矩阵乘法:
BlockMmad<DispatchPolicy, L1Tile, L0Tile, AType, BType, CType> - Epilogue 融合:
BlockEpilogue<DispatchPolicy, CType, ScaleType, PerTokenScaleType, DType, ...> - Kernel 级调度:
GroupedMatmul<...>
5.2 GroupedMatmul 变体
5.2.1 GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
用途: GMM1 — 完成 int8 GMM + PerTokenDequant + Swiglu + 动态量化
文件: gemm/kernel/grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h
核心特点:
- 1C2V 架构: AIC 做 GMM 计算，AIV 做通信 + Swiglu + 量化
- 三角色分离:
- isSendCore (aivIdx % 2 == 1): 负责将量化后的 token 发送给目标 rank
- isRecvCore (aivIdx % 2 == 0): 负责接收其他 rank 的 token
- isCompCore: 负责 Swiglu + 动态量化计算
AIC 端 (operator()
// 1. 先处理共享专家 (如果启用)
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
// 等待 AIV 完成共享专家量化
WaitGroupTokenNumReady(shareQuantTokenStateTensor, expected);
// 执行共享专家 GMM1
blockMmad(gmA[gmOffsetA], layoutA, gmB[gmOffsetB], layoutB, ...);
}
// 2. 处理路由专家 (按 group 循环)
for (uint32_t groupIdx = 0; groupIdx < localExpertNum; ++groupIdx) {
// 等待 AIV 完成该专家的 token 接收
WaitGroupTokenNumReady(groupTokenNumStateTensor, expected);
currentM = groupTokenNumStateTensor.GetValue(GROUP_TOKEN_COUNT);
// 执行该专家的 GMM1
blockMmad(gmA[gmGroupOffsetA + gmOffsetA], ...);
}
AIV 端 (operator()
// 1. 共享专家量化 (shareQuantCoreFunc)
if (isShareQuantCore) {
// 对每个 token 做动态量化: bf16 → int8 + per-token scale
shareQuantCoreFunc(gmX, gmShareSmoothScales, gmShareX1, gmShareX1Scale);
// 通知 AIC 量化完成
}
// 2. 路由专家 Dispatch (SendCoreFunc)
if (isSendCore) {
TokenActiveMaskCal(); // 计算有效 token 数
CalAndSendTokenCount(); // 计算每个专家的 token 数并通知对端
SendToMoeExprt(); // 量化 + 发送到目标 rank 窗口
}
// 3. 路由专家接收 (RecvCoreFunc)
if (isRecvCore) {
RecvCount(); // 等待所有 rank 的 token 计数到达
for (each expert) {
GetCumSum(); // 计算前缀和索引
RecvToken(); // 从各 rank 窗口接收 token
}
// 通知 AIC 接收完成 (通过 GM flag)
}
// 4. Swiglu + 动态量化 (CompCoreFunc)
if (isCompCore) {
CompCoreFunc(); // 调用 BlockEpilogue 做 Swiglu
}
// 5. 后处理: Swiglu 输出的动态量化
// 共享专家 Swiglu 输出量化
BlockQuant<ArchTag> blockQuant(resource, quantParams);
// 路由专家 Swiglu 输出量化
BlockQuant<ArchTag> blockQuant(resource, quantParams);
5.2.2 GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace
用途: GMM2 — 完成 int8 GMM + PerTokenDequant + Combine 发送
文件: gemm/kernel/grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h
核心特点:
- 使用 ResidentA 策略（GMM2 的 A 矩阵是 Swiglu 输出，可常驻 L1）
- Epilogue 中融合 Combine 的 AlltoAll 发送逻辑
- AIV 端直接调用 combiner->ProcessCombine() 或 combiner->AllToAllSend()
AIV 端特殊处理:
// 共享专家路径
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
if (AscendC::GetSubBlockIdx() == 1) {
combiner->TPipeSet(&resource.pipe);
combiner->ProcessCombine(); // 执行 Combine
}
}
// 深度融合路径
else if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
if (AscendC::GetSubBlockIdx() == 0) {
combiner->AllToAllSend(); // subblock 0: 发送
} else {
combiner->ReducePermute(); // subblock 1: 接收+permute
}
}
5.3 BlockMmad — L1/L0 多级流水
文件: gemm/block/block_mmad_preload_async_with_callback_resident_a.h
// 使用 ResidentA 策略的 BlockMmad
// 关键优化: 当 A 矩阵不变时，跳过 GM→L1 的加载
bool useResidentA = (kTileCount == L1A_STAGES) && (!isFirstLoad)
&& (gmBlockA.GetPhyAddr() == lastGmBlockA.GetPhyAddr());
if (!useResidentA) {
copyGmToL1A(l1ATensorList[l1AListId], gmTileA, L1A_LAYOUT, layoutTileA);
}
多级流水阶段: 1. PRELOAD_STAGES: 预加载指令提前发射 2. L1A_STAGES / L1B_STAGES: L1 上 A/B 矩阵的 double buffer 3. L0A_STAGES / L0B_STAGES: L0 上 A/B 矩阵的多级流水 4. L0C_STAGES: L0C 累加器 5. ENABLE_SHUFFLE_K: K 维度循环打乱，避免 bank conflict
5.4 同步机制
AIC→AIV 同步 (Soft Sync):
// AIC 完成一个 L1 tile 的计算后，设置 flag
aicSetFunc = {syncGmAddr + SOFT_SYNC_OFFSET, static_cast<uint8_t>(AscendC::GetBlockIdx())};
// callbackAfterFixpipe = MakeCallback(&aicSetFunc);
// AIV 等待 flag 满足条件
CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, static_cast<uint8_t>(coreIdx), target);
C/V 数据空间交替 (cvDataState):
// 使用 dataState 实现双缓冲，避免读写冲突
cvDataState = FlushAndSpinValue<int32_t>(selfDataStatusTensor, aivStateGlobalCoreIdx * UB_ALIGN);
vToCFlag = (cvDataState == 0) ? V_TO_C_FLAG_1 : V_TO_C_FLAG_2;
6. epilogue 子模块
6.1 EpilogueAtlasA2PerTokenDequantSwiglu
用途: GMM1 的 Epilogue — int32→float 反量化 + Swiglu 激活
文件: epilogue/block/block_epilogue_per_token_dequant_swiglu.h
计算流程:
// 1. 从 GM 加载 int32 累加结果到 UB
copyGmToUbC(ubC, gmTileC, layoutUbC, layoutGmTileC);
// 2. int32 → float 转换
Cast(ubTmpMxN, ubC, CAST_RINT, TileShape::COUNT);
// 3. 乘以 weight scale (行广播: 每列共享一个 scale)
tileRowBroadcastMul(ubTmpMxN, ubTmpMxN, ubFp32Scale);
// 4. 乘以 per-token scale (行广播: 每行共享一个 scale)
tileBroadcastOneBlk(ubTmpMx32B, ubPerTokenScale);
tileOneBlkColumnBroadcastMul(ubTmpMxN, ubTmpMxN, ubTmpMx32B);
// 5. Swiglu 激活 (左右两半)
// 左半: x / (1 + exp(-x)) — sigmoid 门控
// 右半: 保持不变
if (isLeft) {
Muls(ubDenominatorMxN, ubTmpMxN, -1.0f, TileShape::COUNT);
Exp(ubDenominatorMxN, ubDenominatorMxN, TileShape::COUNT);
Adds(ubDenominatorMxN, ubDenominatorMxN, 1.0f, TileShape::COUNT);
Div(ubD, ubTmpMxN, ubDenominatorMxN, TileShape::COUNT);
} else {
tileOneBlkColumnBroadcastMul(ubD, ubTmpMxN, ubTmpMx32B);
}
// 6. 写出 float 结果到 workspace
copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD);
Swiglu 数学公式:
Swiglu(x) = [x_left / (1 + exp(-x_left)), x_right]
= [x_left * sigmoid(x_left), x_right]
6.2 EpilogueAtlasA2PerTokenDequantCombine
用途: GMM2 的 Epilogue — int32→bf16/fp16 反量化 + Combine 发送
文件: epilogue/block/block_epilogue_per_token_dequant.hpp
计算流程:
// 1-4: 与 Swiglu Epilogue 相同 (int32 → float → scale → per-token scale)
Cast(ubCFp32, ubC, CAST_RINT, TileShape::COUNT);
tileRowBroadcastMul(ubMul, ubCFp32, ubFp32Scale);
tileBroadcastOneBlk(ubPerTokenScaleBrcb, ubPerTokenScale);
tileOneBlkColumnBroadcastMul(ubPerTokenMul, ubMul, ubPerTokenScaleBrcb);
// 5. float → bf16/fp16 转换
Cast(ubD, ubPerTokenMul, CAST_RINT, TileShape::COUNT);
// 6. 写出或 Combine 发送
if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
if (expertIdx == UINT32_MAX) {
copyUbToGmD(gmTileD, ubD, layoutGmTileD, layoutUbD); // 共享专家
} else {
DoCombineSend(ubD, layoutGmTileD, layoutUbD, ...); // 路由专家: 直接发送到对端窗口
}
}
DoCombineSend — 融合 Combine 的关键:
// 根据 epSendCount 计算每个 rank 的 token 范围
// 直接将 UB 中的结果 DataCopyPad 到对端窗口
for (uint32_t epRank = epRankStart; epRank < calcInfo.epWorldSize_; ++epRank) {
GM_ADDR rankGM = GetWinAddrByRankId(remoteEpRank, expertIdx) + ...;
rankWindow.SetGlobalBuffer((__gm__ ElementD *)rankGM);
DataCopyPad(rankWindow[(itToken - prevSendCount) * axisH_ + tokenOffset],
ubD[(itToken - startToken) * layoutUbD.stride(0)], dataCopyParams);
}
6.3 Tile 级操作
TileStrideBinary (tile_stride_binary.h):
- 支持 Mul/Div 操作，带自定义 stride
- 使用 BinaryRepeatParams 实现跨行跨列的灵活访问
TileStrideMuls (tile_stride_muls.h):
- 标量乘法，支持不同 DstTileShape 和 SrcTileShape
- 用于 per-token scale 广播
7. Swiglu + Quantization 融合机制
7.1 完整数据流
GMM1 输入: int8 → 累加: int32
↓
PerTokenDequant: int32 × weight_scale × per_token_scale → float32
↓
Swiglu: [x_left * sigmoid(x_left), x_right] → float32
↓
动态量化: float32 → int8 + per_token_scale → workspace
↓
GMM2 输入: int8 → 累加: int32
↓
PerTokenDequant: int32 × weight_scale × per_token_scale → float32
↓
类型转换: float32 → bf16/fp16
↓
Combine 发送: 直接到对端窗口
7.2 BlockQuant — 动态量化实现
文件: gemm/kernel/post_swiglu_dynamic_quant.h
// BlockQuant 完成 Swiglu 输出的动态量化
// 输入: float32 Swiglu 输出 [token_count, n]
// 输出: int8 [token_count, n/2] + per_token_scale [token_count]
// 关键步骤:
// 1. 计算 Swiglu 右半部分与左半部分相乘 (完成 Swiglu)
Mul(ubInput, ubInput, ubInputRightHalf, tileCount);
// 2. 取绝对值
Abs(ubAbs, ubInput, tileCount);
// 3. 对每行取 max (分半处理)
for (uint32_t rowIdx = 0; rowIdx < tileRow; ++rowIdx) {
Max(ubMax[rowIdx * halfTileColumn], ubAbs[rowIdx * tileColumn],
ubAbs[rowIdx * tileColumn + halfTileColumn], halfTileColumn);
}
// 4. 计算量化 scale = max / 127
Muls(ubDequantScale, ubReduceMax, 1.0f / 127.0f, tileRow);
// 5. 量化: x / scale → int8
Muls(ubInputTmp, ubInput, 127.f, tileCount);
for (uint32_t rowIdx = 0; rowIdx < tileRow; ++rowIdx) {
Muls(ubQuantF32[rowIdx * tileColumn], ubInputTmp[rowIdx * tileColumn],
1.f / ubReduceMax.GetValue(rowIdx), tileColumn);
}
Cast(ubQuantS32, ubQuantF32, CAST_RINT, tileCount);
Cast(ubQuantF16, ubQuantS32, CAST_RINT, tileCount);
Cast(ubOutput, ubQuantF16, CAST_RINT, tileCount);
7.3 量化参数传递
Per-token scale 的存储:
// 每个 token 附带 512B 的额外空间存储 scale
constexpr uint32_t TOKEN_EXTRA_SPACE = 512;
scaleParamPad = TOKEN_EXTRA_SPACE;
hCommuSize = hOutSize + scaleParamPad;
通信时的 scale 传递:
// 发送端: 将 scale 附加在 token 后面
yFp32Tensor.SetValue(tokenLength / sizeof(float), float(1.0) / dynamicQuantScale);
DataCopy(dstWinGMTensor, yInt8Tensor[index], tokenLength); // token
DataCopy(dstWinGMTensor[tokenLength], yInt8Tensor[index][tokenLength], scaleParamPad); // scale
// 接收端: 从 token 后面读取 scale
DataCopyPad(dynamicScalesOutGMTensor_[beginIdx + j],
xOutFp32Tensor_[tokenLength / sizeof(float)], dataCopyParamsFloat);
8. 性能优化关键点
8.1 内存带宽优化
| 优化点 | 收益 | 实现方式 |
|---|---|---|
| 中间数据片上消除 | 减少 8 次 GM 访问 | Workspace 复用 + 流水线传递 |
| X1/X2 Scale 压缩 | 减少 1/3 通信量 | Per-token scale 替代 per-channel |
| 权重 NZ 格式 | 提升 Cube 效率 | layout::zN 布局 |
| L1 常驻 A 矩阵 | 减少 GM→L1 加载 | ResidentA Dispatch Policy |
8.2 计算通信重叠
1C2V 架构的天然优势: - AIC 在做 GMM 计算时，AIV 可以并行做： - 下一组专家的 token 接收 - 当前组 Swiglu 输出的量化 - Combine 的 AlltoAll 发送
流水线示意:
时间 →
AIC: [GMM1_Expert0] [GMM1_Expert1] [GMM1_Expert2] [GMM2_Expert0] ...
AIV: [Dispatch_Send] [Recv_Expert0] [Swiglu_Quant_0] [Combine_Send_0] ...
8.3 同步优化
Soft Sync 替代 Hard Sync:
// 使用 GM 上的 flag 做轻量级同步
// 避免 PipeBarrier<PIPE_ALL>() 的全局阻塞
CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, coreIdx, target);
CV 双缓冲:
// dataState 在 0/1 之间交替
// C 端写 state0 时 V 端读 state1，反之亦然
winDataSizeOffset = dataState_ * epWorldSize_ * expertPerSizeOnWin_ * moeExpertNumPerRank_;
8.4 循环优化
AIV Dispatch 优化 (enableAivOpt): - 小 batch (≤64) 时，用 table 法替代逐 token 计数 - 将 expert_id 映射到 (bs+1) × expert_num 的 table - 用前缀和指令快速计算 expand_idx
K 维度 Shuffle:
// 打乱 K 维度的循环顺序，避免 bank conflict
if constexpr (ENABLE_SHUFFLE_K) {
startTileIdx = AscendC::GetBlockIdx() % kTileCount;
}
8.5 Workspace 内存复用
#ifdef ENABLE_REUSE_MEMORY
// X1 和 X2 共享空间 (生命周期不重叠)
shareX1TokenOffset == shareX2TokenOffset
// Scale 空间复用
shareX1ScaleOffset == shareX2ScaleOffset
// SwapSpace 和 Swiglu 输出复用
swapSpaceOffset == shareSwigluOffset
#endif
9. 与标准 MoE (DeepSeek-style) 的对应关系
9.1 数学对应
FusedDeepMoe 实现了 DeepSeek-V3 风格的 MoE Decode:
标准 MoE 计算:
1. Routing: expert_ids = TopK(router(x), k)
2. Dispatch: 将 token 发送到对应 expert 的 rank
3. GMM1: y = x @ W1 (int8 量化)
4. Swiglu: y = Swiglu(y) (GLU 变体)
5. GMM2: z = y @ W2 (int8 量化)
6. Combine: output = Σ(expert_scale_i * z_i)
FusedDeepMoE 对应:
1. expert_ids 由外部 router 计算后传入
2. Dispatch → CamMoeDistributeDispatch (融合在 kernel 内)
3. GMM1 → GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
4. Swiglu → BlockEpilogue 中的 Swiglu 激活
5. GMM2 → GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace
6. Combine → CamMoeDistributeCombine (融合在 kernel 内)
9.2 量化策略对应
| 标准 MoE | FusedDeepMoE | 说明 |
|---|---|---|
| BF16 输入 | Dynamic Quant → int8 | Dispatch 阶段完成 |
| W1 int8 per-channel | W1 int8 + per-channel scale | GMM1 权重 |
| Swiglu 输出 BF16 | Dynamic Quant → int8 | Swiglu 后立即量化 |
| W2 int8 per-channel | W2 int8 + per-channel scale | GMM2 权重 |
| 输出 BF16 | PerTokenDequant → BF16 | GMM2 Epilogue |
9.3 与 torch_npu 标准算子对比
标准小算子组合 (SmallOps):
# 5 次 kernel launch
expand_x, scales, assist_info, token_nums, ep_send_counts, _, expand_scales = \
torch_npu.npu_moe_distribute_dispatch_v2(...)
y1_int32 = torch_npu.npu_grouped_matmul(x=[expand_x], weight=[gmm1_weight], ...)
y1, y1_scale = torch_npu.npu_dequant_swiglu_quant(x=y1_int32, ...)
y2 = torch_npu.npu_grouped_matmul(x=[y1], weight=[gmm2_weight], ...)
output = torch_npu.npu_moe_distribute_combine_v2(expand_x=y2, ...)
FusedDeepMoE (FusionOp):
# 1 次 kernel launch
output, share_output, expert_token_nums = torch.ops.umdk_cam_op_lib.fused_deep_moe(
x, expert_ids, gmm1_weight, gmm1_weight_scale, gmm2_weight, gmm2_weight_scale,
expert_scales, share_gmm1_weight, ..., group_ep=..., ep_rank_size=..., ...)
9.4 支持的配置范围
| 参数 | 范围 | 说明 |
|---|---|---|
| batchSize | 0-256 | 受 AIV 循环优化限制 |
| tokenLength (h) | 512-7168 | 常见 7168 (DeepSeek-V3) |
| gmm1HLen | 1048-6144, 对齐 1024 | 中间维度 |
| moeExpertNum | 1-512, 整除 epRankSize | 总专家数 |
| topK | 1-12 | 路由 topK |
| epRankSize | 任意 | EP 并行度 |
9.5 执行路径总结
┌─────────────────────────────────────┐
│ FusedDeepMoe Entry │
└──────────────┬──────────────────────┘
│
┌──────────────────┼──────────────────┐
│ │ │
moeExpertNumPerRank==1 │ moeExpertNumPerRank>1
(非深度融合) │ (深度融合)
│ │ │
▼ │ ▼
独立 Dispatch kernel │ GMM1 内完成 Dispatch
│ │ │
└──────────────────┼──────────────────┘
│
▼
┌──────────────────────────────┐
│ GmmDeqSwigluQuant (GMM1) │
│ AIC: int8 GMM → int32 │
│ AIV: Swiglu + Dynamic Quant │
│ → workspace (int8 + scale) │
└──────────────┬───────────────┘
│
▼
┌──────────────────────────────┐
│ GmmDeq (GMM2 + Combine) │
│ AIC: int8 GMM → int32 │
│ AIV: PerTokenDequant + │
│ Combine AlltoAll Send │
└──────────────────────────────┘
附录：关键文件索引
| 文件 | 功能 |
|---|---|
op_kernel/fused_deep_moe.h |
主 kernel 类定义，Process() 入口 |
op_kernel/fused_deep_moe.cpp |
kernel 入口，模板实例化 |
op_kernel/fused_deep_moe_tiling.h |
Tiling 数据结构，Tile Shape 常量 |
op_kernel/fused_deep_moe_tiling_key.h |
Tiling Key 路由机制 |
op_host/fused_deep_moe_tiling.cpp |
Host 端 tiling 计算 |
op_host/fused_deep_moe.cpp |
算子定义 (输入/输出/属性) |
op_host/fused_deep_moe_infer.cpp |
Shape/DataType 推断 |
raw_distributed/cam_moe_distribute_dispatch.h |
Dispatch 实现 |
raw_distributed/cam_moe_distribute_combine.h |
Combine 实现 |
gemm/dispatch_policy.h |
GMM Dispatch Policy 定义 |
gemm/kernel/grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h |
GMM1 完整实现 |
gemm/kernel/grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h |
GMM2 完整实现 |
gemm/kernel/post_swiglu_dynamic_quant.h |
Swiglu 后动态量化 |
gemm/block/block_mmad.h |
BlockMmad 入口 |
gemm/block/block_mmad_preload_async_with_callback_resident_a.h |
ResidentA BlockMmad |
epilogue/dispatch_policy.h |
Epilogue Dispatch Policy |
epilogue/block/block_epilogue.h |
BlockEpilogue 入口 |
epilogue/block/block_epilogue_per_token_dequant_swiglu.h |
Swiglu Epilogue |
epilogue/block/block_epilogue_per_token_dequant.hpp |
Combine Epilogue |
epilogue/tile/tile_stride_binary.h |
Tile 级 Mul/Div |
epilogue/tile/tile_stride_muls.h |
Tile 级 Muls |
fused_deep_moe_utils.h |
同步工具函数 |
op_api/aclnn_fused_deep_moe.cpp |
aclnn API 入口 |
pybind/fused_deep_moe.cpp |
PyTorch 绑定 |
examples/fused_deep_moe_sample.py |
使用示例 |
总结
FusedDeepMoE 是 CAM 库的皇冠级算子，其核心价值在于：
- 极致融合: 将 5 个独立算子 (Dispatch + GMM1 + Swiglu + GMM2 + Combine) 融合为 1 个 kernel
- 中间数据全片上化: X1、SwigluOut、X2 等中间数据在 workspace 中流转，零 GM 访问
- 1C2V 异构并行: AIC 做 GMM 计算，AIV 做通信调度，天然计算通信重叠
- 全链路量化: 从输入动态量化到 Swiglu 后动态量化，全程 int8 计算
- 精细同步: Soft Sync + CV 双缓冲 + 多级 Flag 同步，最小化同步开销
这是华为昇腾 Ascend 910B 平台上 MoE Decode 阶段的终极优化实现，代表了 NPU 上 MoE 算子融合的最高水平。
分析范围：test/cam/ut/、PyTorch NPU Extension Bridge、集成测试框架、构建系统
分析日期：2026-07-21
项目路径：/Users/backyes/work/claude_workspace/umdk_research/umdk
一、单元测试体系分析
1.1 测试目录结构
test/cam/ut/
├── __init__.py
├── conftest.py # pytest fixtures：通信域初始化 + 用例同步
├── pytest.ini # marker 注册
├── run_shell.py # 覆盖率包装入口
├── util/
│ ├── __init__.py
│ ├── marker.py # 自定义 marker + 条件跳过
│ └── tool.py # 分布式工具函数 + 精度校验
├── fused_deep_moe/
│ ├── __init__.py
│ └── test_fused_deep_moe.py
└── moe_dispatch_combine_prefill/
├── __init__.py
└── test_moe_dispatch_combine_prefill.py
关键观察：当前 CAM 的 UT 体系是 纯 Python/pytest 风格，与 URMA 模块的 C++/gtest 体系形成鲜明对比。CAM 没有 C++ 层的 gtest UT，所有测试通过 Python 调用 torch.ops.umdk_cam_op_lib.* 自定义算子完成。
1.2 pytest.ini Marker 体系
[pytest]
markers =
author : "marks test as author"
multi_process_test : "marks test as multi_process_test"
single_process_test : "marks test as single_process_test"
a2_test : "marks test as A2/910B test"
a3_test : "marks test as A3/910C test"
| Marker | 含义 | 使用场景 |
|---|---|---|
MPTest (multi_process_test) |
多进程测试 | 所有 CAM 算子测试，通过 torchrun 拉起 |
SPTest (single_process_test) |
单进程测试 | 当前未使用 |
A2Test (a2_test) |
华为 910B 芯片测试 | 当前未使用 |
A3Test (a3_test) |
华为 910C 芯片测试 | fused_deep_moe、moe_dispatch_combine_prefill |
Author |
用例作者标记 | 当前未使用 |
1.3 conftest.py — Fixture 设计
@pytest.fixture(scope="function", autouse=True)
def test_cases_sync():
if tool.is_torchrun_process():
dist.barrier()
yield
if tool.is_torchrun_process():
dist.barrier()
@pytest.fixture(scope="session", autouse=True)
def mpt_global_fixture():
if tool.is_torchrun_process():
device = torch.device("npu", tool.get_local_rank())
torch_npu.npu.set_device(device)
dist.init_process_group(backend="hccl", rank=tool.get_rank(), world_size=tool.get_world_size())
yield
设计要点：
- session 级：mpt_global_fixture 在测试会话开始时通过 HCCL 后端初始化 ProcessGroup，使用 torchrun 注入的环境变量（RANK/WORLD_SIZE/LOCAL_RANK/LOCAL_WORLD_SIZE）
- function 级：test_cases_sync 在每个用例前后加 dist.barrier()，保证多进程用例执行同步
- autouse=True：两个 fixture 均自动生效，测试用例无需显式声明
1.4 util/marker.py — 条件跳过机制
def SKIP_ENV_RANKSIZE_UNEQUAL(expect_ranksize, reason='skip if ranksize is not equal to '):
skip = False
if expect_ranksize != tool.get_world_size():
skip = True
reason += str(expect_ranksize)
return pytest.mark.skipif(condition=skip, reason=reason)
这是一个运行时条件跳过装饰器：当实际 WORLD_SIZE 不等于期望值时自动 skip。例如 @SKIP_ENV_RANKSIZE_UNEQUAL(4) 表示仅在 4 卡环境执行。
1.5 util/tool.py — 分布式工具函数
| 函数 | 功能 |
|---|---|
is_torchrun_process() |
判断是否 torchrun 拉起的进程（检查 RANK/WORLD_SIZE/LOCAL_RANK/LOCAL_WORLD_SIZE 四个环境变量） |
get_rank() / get_local_rank() / get_world_size() / get_local_world_size() |
获取分布式信息 |
is_run_for_cov() |
检查是否覆盖率模式（ENABLE_COV 环境变量） |
allclose_nparray() |
带容错的 numpy 数组比较：先 np.allclose，失败则统计不满足 rtol/atol 的元素比例 |
精度校验策略（allclose_nparray）：
error = np.abs(data_expect - data_check)
greater = np.greater(error, atol + np.abs(data_check) * rtol)
loss_count = np.count_nonzero(greater)
assert (loss_count / total_count) < rtol
1.6 run_shell.py — 覆盖率包装
relative_path = os.path.join("..", "..", "..", "scripts", "comm_operator", "tool", "run_with_cov.sh")
script_path = os.path.abspath(os.path.join(script_dir, relative_path))
command = ["bash", script_path] + args
subprocess.run(command)
注意：该文件引用了 scripts/comm_operator/tool/run_with_cov.sh，但当前该文件不存在。这是一个预留的覆盖率收集入口，说明 CAM 规划了基于 shell 的覆盖率收集流程，但尚未落地。
二、fused_deep_moe 测试用例深度分析
2.1 测试策略：参考实现 vs 融合算子对比
测试采用 "Golden Reference vs Fused Op" 的对比策略：
| 组件 | 角色 | 实现方式 |
|---|---|---|
SmallOps |
参考实现（Golden Reference） | 手动拆解为 npu_moe_distribute_dispatch_v2 → npu_grouped_matmul → npu_dequant_swiglu_quant → npu_grouped_matmul → npu_moe_distribute_combine_v2 |
FusionOp |
被测融合算子 | 调用 torch.ops.umdk_cam_op_lib.fused_deep_moe 单一融合接口 |
核心断言：
tool.allclose_nparray(small_token_np, fused_token_np, rtol=0.02, atol=2.0) # token 输出
tool.allclose_nparray(small_share_token_np, fused_share_token_np, rtol=0.02, atol=2.0) # share 输出
np.testing.assert_array_equal(small_count_np, fused_count_np) # count 输出精确相等
2.2 测试场景参数化
CASE_4RANK = {
"totalExpertNum": 16,
"topk": 8,
"batchSize": 16,
"hiddenSize": 7168,
"intermediateHiddenSize": 2048,
"dynamicEPLB": False,
"with_mc2_mask": False,
}
CASE_8RANK = {
"totalExpertNum": 16,
"topk": 8,
"batchSize": 32,
"hiddenSize": 7168,
"intermediateHiddenSize": 2048,
"dynamicEPLB": True,
"with_mc2_mask": False,
}
注意：CASE_8RANK 已定义但当前测试仅使用 CASE_4RANK。
2.3 测试参数矩阵
| 维度 | 覆盖值 | 说明 |
|---|---|---|
| 执行模式 | GE（图模式）、Eager |
通过 @pytest.mark.parametrize("mode", [("GE"), ('Eager')]) |
| Rank 数 | 4 | @SKIP_ENV_RANKSIZE_UNEQUAL(4) |
| 数据类型 | bfloat16 | test_bfloat16 = True |
| 动态 batch | 关闭 | dynamicBS = False |
| Share Expert | 关闭 | with_share = False |
| Smooth Quant | 关闭 | with_smooth = False |
| MC2 Mask | 关闭 | with_mc2_mask = False |
2.4 覆盖率规避策略
if mode == "GE" and tool.is_run_for_cov():
return
图模式（GE）无法获得算子覆盖率，因此在覆盖率收集模式下自动跳过 GE 测试。
2.5 测试覆盖的算子链路
SmallOps._apply_ops 覆盖的完整 MoE 计算链路：
输入 x
↓
npu_moe_distribute_dispatch_v2 (EP 分发)
↓
npu_grouped_matmul (GMM1, int8 量化矩阵乘)
↓
npu_dequant_swiglu_quant (SwiGLU 激活 + 反量化)
↓
npu_grouped_matmul (GMM2, 带 pertoken scale)
↓
npu_moe_distribute_combine_v2 (EP 聚合)
↓
输出 + share_output + expert_token_nums
2.6 数据生成策略
generate_datas() 函数生成随机测试数据：
- 权重：int8 随机值 [-16, 16]，shape 为 [local_expert_num, input_dim, output_dim]
- Scale：float32 随机值 [0.0015, 0.0045]
- Expert IDs：顺序排列后取模 expert_ids % moe_expert_num
- Expert Scales：[0, 1) 均匀分布
- 输入 x：[-5, 5) 均匀分布
三、moe_dispatch_combine_prefill 测试用例深度分析
3.1 测试策略：自包含正确性验证
与 fused_deep_moe 的"参考实现对比"策略不同，dispatch_combine_prefill 采用数学定义验证：
# 参考输出计算
ref_weights = np.where(topk_idx_np == -1, np.float32(0), topk_weights_np).sum(axis=1).reshape(-1, 1)
ref_x = x_np * ref_weights
# 相似度计算（余弦相似度变体）
def calc_diff(x, y):
x, y = x.astype(np.float64) + 1, y.astype(np.float64) + 1
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
assert(diff < 5e-5)
3.2 测试场景参数
CASE_16RANK = {
"num_ranks": 16,
"num_experts": 16,
"num_topk": 8,
"num_tokens": 32,
"hidden_size": 7168,
}
| 维度 | 覆盖值 |
|---|---|
| Rank 数 | 16 |
| 执行模式 | Eager（GE 在覆盖率模式下跳过） |
| 量化模式 | use_quant=False 和 use_quant=True |
3.3 测试覆盖的算子
# 1. 获取 dispatch layout
output_layout = torch.ops.umdk_cam_op_lib.get_dispatch_layout(topk_idx, num_experts, num_ranks)
# 2. dispatch 阶段
output_dispatch = torch.ops.umdk_cam_op_lib.moe_dispatch_prefill(
x, topk_idx, topk_weights, num_tokens_per_expert, send_token_idx_small,
group_ep, rank, num_ranks, use_quant)
# 3. 反量化（如果 use_quant=True）
recv_x = per_token_cast_back(recv_x, dynamic_scales_out) if use_quant else recv_x
# 4. combine 阶段
combine_x = torch.ops.umdk_cam_op_lib.moe_combine_prefill(
recv_x, topk_idx, topk_weights, expand_idx_out, recv_count,
group_ep, rank, num_ranks)
3.4 FP8 反量化实现
def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor):
if x_scales.dtype == torch.int:
x_scales = x_scales.view(dtype=torch.int8).to(torch.int) << 23
x_scales = x_scales.view(dtype=torch.float)
x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128)
x_scales = x_scales.view(x_fp8.size(0), -1, 1)
return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16)
这是 128-element per-tile FP8 反量化：将 int8 的 scale 通过位操作还原为 float，然后按 128 元素粒度做 per-token 缩放。
3.5 数据来源声明
文件头部注明：
Part of codes modified from: MagiAttention
Original Source: https://github.com/SandAI-org/MagiAttention.git
说明该测试用例参考了 SandAI 的 MagiAttention 实现。
四、PyTorch NPU Extension Bridge 实现
4.1 架构概览
┌─────────────────────────────────────────────────────────┐
│ PyTorch 前端 │
│ torch.ops.umdk_cam_op_lib.* │
└─────────────────────┬───────────────────────────────────┘
│
┌─────────────────────▼───────────────────────────────────┐
│ pybind11 绑定层 │
│ fused_deep_moe.cpp / moe_dispatch_prefill.cpp / ... │
└─────────────────────┬───────────────────────────────────┘
│
┌─────────────────────▼───────────────────────────────────┐
│ pytorch_npu_helper.hpp (EXEC_NPU_CMD) │
│ at::Tensor → aclTensor 类型转换 + OpApi 调用 │
└─────────────────────┬───────────────────────────────────┘
│
┌─────────────────────▼───────────────────────────────────┐
│ ACL Runtime │
│ libopapi.so / libcust_opapi.so │
└─────────────────────────────────────────────────────────┘
4.2 NPUBridge — Tensor 到 StorageImpl 的桥接
class NPUBridge {
public:
static NPUStorageImpl *GetNpuStorageImpl(const at::Tensor &tensor);
static NPUStorageImpl *GetNpuStorageImpl(c10::StorageImpl *storageImpl);
static NPUStorageImpl *GetNpuStorageImpl(c10::Storage &&storage);
static NPUStorageDesc &GetNpuStorageImplDesc(const at::Tensor &tensor);
};
实现方式：通过 static_cast 将 PyTorch 的 c10::StorageImpl 强制转换为 CAM 自定义的 NPUStorageImpl。这是侵入式扩展——CAM 在创建 tensor 时用 NPUStorageImpl 替换默认 StorageImpl。
4.3 NPUStorageDesc — NPU 存储元数据
struct NPUStorageDesc {
c10::SmallVector<int64_t, 5> base_sizes_;
c10::SmallVector<int64_t, 5> base_strides_;
c10::SmallVector<int64_t, 5> storage_sizes_;
int64_t base_offset_ = 0;
use_byte_size_t base_dtype_ = {};
aclFormat origin_format_ = ACL_FORMAT_UNDEFINED;
aclFormat npu_format_ = ACL_FORMAT_ND;
caffe2::TypeMeta data_type_ = caffe2::TypeMeta::Make<uint8_t>();
};
关键字段：
- base_sizes_ / base_strides_：原始 tensor 的 shape/stride
- storage_sizes_：实际存储的 storage 维度（考虑 NZ 格式转换）
- origin_format_ / npu_format_：ACL 格式（ND、FRACTAL_NZ 等）
- data_type_：用于从 StorageImpl 重建 GE tensor
4.4 NPUStorageImpl — 自定义 Storage
struct NPUStorageImpl : public c10::StorageImpl {
explicit NPUStorageImpl(use_byte_size_t, size_t size_bytes, at::DataPtr data_ptr,
at::Allocator *allocator, bool resizable);
void release_resources() override;
NPUStorageDesc npu_desc_;
};
工厂函数 make_npu_storage_impl 创建带 NPU 描述的 Storage：
c10::intrusive_ptr<c10::StorageImpl> make_npu_storage_impl(...) {
if (data_ptr == nullptr) {
data_ptr = allocator->allocate(size_bytes.as_int_unchecked());
}
return c10::make_intrusive<NPUStorageImpl>(...);
}
4.5 pytorch_npu_helper.hpp — 核心转换层
4.5.1 类型映射表
#define AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(_) \
_(at::ScalarType::Byte, ACL_UINT8) \
_(at::ScalarType::Char, ACL_INT8) \
_(at::ScalarType::BFloat16, ACL_BF16) \
...
完整的 ATen ScalarType → ACL DataType 映射，覆盖 20+ 种数据类型。
4.5.2 FP8/FP4 低精度支持
#ifdef CAM_ACL_HAS_FP8_FP4
static struct LowPrecisionDtypeMapInitializer {
LowPrecisionDtypeMapInitializer() {
kATenScalarTypeToAclDataTypeTable[at::ScalarType::Float8_e5m2] = ACL_FLOAT8_E5M2;
kATenScalarTypeToAclDataTypeTable[at::ScalarType::Float8_e4m3fn] = ACL_FLOAT8_E4M3FN;
kATenScalarTypeToAclDataTypeTable[at::ScalarType::Float8_e8m0fnu] = ACL_FLOAT8_E8M0;
kATenScalarTypeToAclDataTypeTable[at::ScalarType::Float4_e2m1fn_x2] = ACL_FLOAT4_E2M1;
}
} low_precision_dtype_map_initializer;
#endif
关键设计：通过 CAM_ACL_HAS_FP8_FP4 宏在编译期探测 CANN 版本是否支持 FP8/FP4。旧版 CANN 8.5（aclDataType 到 ACL_BF16 截止）也能编译，低精度回退到 view(int8) 方案。
4.5.3 OpApi 动态加载
inline void *GetOpApiFuncAddr(const char *apiName)
{
static auto custOpApiHandler = GetOpApiLibHandler(GetCustOpApiLibName()); // libcust_opapi.so
if (custOpApiHandler != nullptr) {
auto funcAddr = GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName);
if (funcAddr != nullptr) return funcAddr;
}
static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); // libopapi.so
return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName);
}
双库回退：优先查找 libcust_opapi.so（自定义算子 API），找不到再回退到 libopapi.so（标准 ACLNN API）。
4.5.4 EXEC_NPU_CMD 宏 — 算子调用统一入口
#define EXEC_NPU_CMD(aclnn_api, ...)
这是整个 helper 的核心宏，封装了完整的 ACLNN 算子调用流程：
- 获取函数地址：
GetOpApiFuncAddr动态加载aclnn_api和aclnn_apiGetWorkspaceSize - 初始化大内存：
InitHugeMemThreadLocal - 类型转换：
ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr)将所有参数转为 ACL 类型 - 获取 Workspace：调用
GetWorkspaceSize获取所需 workspace 大小 - 分配 Workspace：通过
at::empty在 NPU 上分配 - 执行算子：调用 OpApiFunc，传入 workspace 和 acl_stream
- 释放资源：
ReleaseConvertTypes+UnInitHugeMemThreadLocal+ReleaseHugeMem
4.5.5 ConvertType — Tensor 到 aclTensor 的转换
inline aclTensor *ConvertType(const at::Tensor &at_tensor) {
// 1. 获取 ACL 数据类型
aclDataType acl_data_type = kATenScalarTypeToAclDataTypeTable[...];
// 2. 判断格式（NZ vs ND）
if (!IsOpInputBaseFormat(at_tensor)) {
format = NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.npu_format_;
storageDims = NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.storage_sizes_;
} else {
// 根据维度数推断格式：3D→NCL, 4D→NCHW, 5D→NCDHW
}
// 3. 创建 aclTensor
aclCreateTensor(sizes, dimNum, acl_data_type, strides, offset, format, storageDims, ...);
}
4.6 pybind.cpp — 算子注册
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fused_deep_moe", &FusedDeepMoeImplAutograd, "fused_deep_moe");
m.def("get_dispatch_layout", &GetDispatchLayoutImplAutograd, "get_dispatch_layout");
m.def("moe_dispatch_prefill", &MoeDispatchPrefillImplAutograd, "moe_dispatch_prefill");
m.def("moe_combine_prefill", &MoeCombinePrefillImplAutograd, "moe_combine_prefill");
}
TORCH_LIBRARY(umdk_cam_op_lib, m) {
m.def("fused_deep_moe(Tensor x, Tensor expert_ids, Tensor[] gmm1_weight, ...) -> Tensor[]");
m.def("get_dispatch_layout(Tensor topk_idx, int num_experts, int num_ranks) -> (Tensor, Tensor)");
m.def("moe_dispatch_prefill(...) -> (Tensor, Tensor, Tensor, Tensor, Tensor)");
m.def("moe_combine_prefill(...) -> Tensor");
}
双层注册：
1. pybind11 层：提供 Python 直接调用入口（torch.ops.umdk_cam_op_lib.fused_deep_moe）
2. TORCH_LIBRARY 层：注册到 PyTorch 的 dispatcher，支持 torch.compile 图模式捕获
五、集成测试框架
5.1 UMDK 集成测试整体架构
test/intergration_test/
├── conftest.py # 全局 pytest 配置（随机种子、路径注册）
├── test_framework/
│ ├── app/
│ │ ├── common/ # C 语言公共库（socket、thread pool、log）
│ │ ├── urma/ # URMA 测试应用层
│ │ ├── dlock/ umq/ ums/ urpc/
│ │ └── cam/ # ← 当前不存在，CAM 无集成测试
│ ├── common/
│ │ ├── constants.py # 全局常量（从 /etc/ubus_ci/test_env.yaml 加载）
│ │ ├── host.py # SSH 远程主机管理
│ │ ├── check_resource.py
│ │ └── utils.py
│ └── ubus_test/
│ └── base_test.py # 基类：setup_class/teardown、core dump 检测
└── test_suites/
├── URMA/ # URMA 功能测试
├── URMA_UT/ # URMA 单元集成测试
├── DLOCK/ UMQ/ UMS/ URPC_LIB/
└── （无 CAM 测试套件）
5.2 CAM 在集成测试中的位置
关键发现：test_framework/app/cam/ 目录不存在，test_suites/ 下也没有 CAM 相关测试套件。
这意味着： - CAM 当前没有参与 UMDK 的集成测试体系 - CAM 的测试完全依赖单元测试（pytest + 真实 NPU 多卡环境） - 与 URMA 模块（有完整的集成测试套件）形成对比
5.3 URMA 集成测试框架参考
URMA 的集成测试框架设计值得 CAM 借鉴：
BaseTest 基类（ubus_test/base_test.py）：
- setup_class：解析测试环境 YAML、记录 core dump 状态
- teardown_class：进程清理、core dump 检测
- parse_host：通过 SSH 连接多台测试主机
- 支持多主机（host_list）、多设备（test_nic1_dev）
测试环境配置（/etc/ubus_ci/test_env.yaml）：
host_info:
host1:
user: root
manage_nic: { name: xxx, ip: x.x.x.x }
test_nic1: { name: xxx, ip: x.x.x.x, eid: xxx }
host2: ...
conftest.py 全局配置：
def pytest_configure(config):
random_seed = config.getoption("--random_seed")
random.seed(int(random_seed))
const.CASE_PATH = ...
const.TMOUT = timeout
六、回归测试机制
6.1 当前回归测试现状
| 机制 | CAM 状态 | 说明 |
|---|---|---|
| CI Pipeline | ❌ 未发现 | 无 Jenkinsfile / .github/workflows / .gitlab-ci.yml |
| 自动化触发 | ❌ | 无 pre-commit / post-merge 自动触发证据 |
| 覆盖率门禁 | ❌ | run_with_cov.sh 不存在 |
| 多芯片覆盖 | ⚠️ 部分 | A3（910C）有 marker，A2（910B）有 marker 但无测试使用 |
| 多 Rank 覆盖 | ⚠️ 部分 | 4-Rank 和 16-Rank 有测试，但 8-Rank 场景已定义未使用 |
6.2 条件跳过作为回归防护
@SKIP_ENV_RANKSIZE_UNEQUAL(4) # 仅 4 卡环境执行
@SKIP_ENV_RANKSIZE_UNEQUAL(16) # 仅 16 卡环境执行
这种机制确保： - 测试不会在错误的硬件配置下执行（避免误报） - 但也不会在错误的配置下失败（skip 而非 fail）
6.3 与 URMA 回归测试的对比
URMA 有完善的回归测试基础设施（test/urma/script/urma_ut_common.sh）：
| 能力 | URMA | CAM |
|---|---|---|
| 分阶段构建 | ✅ common/core/cmd_tlv/bond/uvs/fuzz |
❌ |
| 代码覆盖率 | ✅ lcov + genhtml + 90% 门禁 | ❌ |
| ASAN 内存检测 | ✅ -fsanitize=address |
❌ |
| Fuzz 测试 | ✅ urma_api_fuzz.c + API 覆盖率检查 |
❌ |
| 增量覆盖率 | ✅ check_diff_coverage.py --threshold 90 |
❌ |
| JUnit 报告 | ✅ --output-junit |
❌ |
七、测试基础设施
7.1 构建系统
7.1.1 顶层构建流程
build/cam/build.sh (入口)
└── build/cam/comm_operator/build.sh
├── compile_ascend_proj.sh → 编译 .run 包（op_host + op_kernel）
├── build_pybind.sh → 编译 .whl 包（pybind11 扩展）
└── BuildTest (未实现) → UT 构建入口
7.1.2 compile_ascend_proj.sh — 算子编译
BuildTest 函数：在 build.sh 中被调用（if [ $ENABLE_UT_BUILD -eq 1 ]; then BuildTest; fi），但在整个代码库中未找到定义。这是一个未实现的预留接口。
编译流程：
1. select_ops.py 解析 operator_registry.json，根据 SOC/SHMEM/quant 筛选算子
2. msopgen gen 生成 CMakePresets.json 和项目骨架
3. copy_ops_include 仅复制选中算子的源码（include-list 模式）
4. cmake -S . -B build_out --preset=default 配置
5. cmake --build build_out --target binary -j${CAM_BUILD_JOBS:-8} 编译
6. 生成 .run 包
7.1.3 operator_registry.json — 算子注册表
{
"soc_versions": {
"ascend910_93": [
"dispatch_layout", "fused_deep_moe", "moe_combine_normal",
"moe_dispatch_normal", "notify_dispatch"
]
},
"operator_meta": {
"fused_deep_moe": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "base" }
}
}
设计亮点：
- SOC 版本 → 算子列表的映射
- requires_shmem 标记（SHMEM 未安装时自动过滤）
- family + variant 处理互斥变体（base vs w4a8 共享文件名，不可同时编译）
7.1.4 select_ops.py — 算子选择逻辑
def resolve(registry, soc_list, shmem, quant, user_ops):
# 1. 确定候选集（全量 or 用户指定）
# 2. fused_deep_moe base/w4a8 互斥规则
# 3. SHMEM 过滤
# 4. 保持 registry 顺序（确定性）
7.2 CMake 构建配置
7.2.1 op_kernel/CMakeLists.txt
npu_op_kernel_options("" ALL OPTIONS ${KERNEL_COMPILE_OPTS})
npu_op_kernel_sources(ascendc_kernels KERNEL_DIR ./)
npu_op_kernel_library(ascendc_kernels SRC_BASE ${CMAKE_CURRENT_SOURCE_DIR}/ TILING_LIBRARY cust_optiling)
npu_op_package_add(${package_name} LIBRARY ascendc_kernels)
编译选项：
- SHMEM 包含路径：-I$ENV{SHMEM_HOME_PATH}/shmem/include
- CATLASS 包含路径：-I$ENV{CATLASS_HOME_PATH}/include
- Debug 模式：-g -O0
- 生产模式：无额外 flag
7.2.2 op_host/CMakeLists.txt
CANN 版本探测：
set(CANN83_PATH "${CANN_ROOT}/include/platform/platform_infos_def.h")
set(CANN82_PATH "${CANN_ROOT}/include/experiment/platform/platform/platform_infos_def.h")
5 步构建流程：
1. npu_op_code_gen：生成 ACLNN stubs、op_proto.cc、fallback
2. npu_op_library(cust_opapi ACLNN ...)：ACLLN 单算子调用库
3. npu_op_library(cust_op_proto GRAPH ...)：算子原型库（图编译）
4. npu_op_library(cust_optiling TILING ...)：Tiling 库
5. npu_op_package_add：注册到 run 包
7.3 测试链接方式
CAM 测试不链接 C++ 测试库（gtest/gmock），而是：
1. 通过 import umdk_cam_op_lib 加载 pybind 编译的 .so
2. 通过 torch.ops.umdk_cam_op_lib.* 调用算子
3. 通过 torch.distributed + HCCL 实现多进程通信
八、测试覆盖率与质量保证策略
8.1 当前覆盖率机制
| 层次 | 机制 | 状态 |
|---|---|---|
| Python 算子覆盖率 | tool.is_run_for_cov() 跳过 GE 模式 |
✅ 已实现 |
| C++ 代码覆盖率 | run_with_cov.sh（不存在） |
❌ 未实现 |
| 算子级覆盖率 | 无 | ❌ |
| 分支覆盖率 | 无 | ❌ |
8.2 精度质量保证
8.2.1 容差设计
| 测试 | rtol | atol | 说明 |
|---|---|---|---|
| fused_deep_moe token output | 0.02 | 2.0 | 宽松容差（int8 量化累积误差） |
| fused_deep_moe share output | 0.02 | 2.0 | 同上 |
| fused_deep_moe count output | 精确相等 | - | assert_array_equal |
| dispatch_combine_prefill | 5e-5（余弦相似度） | - | 更严格（FP8 反量化后） |
8.2.2 容差选择的合理性
- rtol=0.02, atol=2.0：对于 int8 量化 MoE（hidden_size=7168），2.0 的绝对误差容忍度对应约 0.03% 的相对误差（假设值域 [-5,5]），是合理的量化精度验证
- 5e-5 余弦相似度：对于 dispatch/combine（纯通信 + 加权），要求几乎无损
8.3 测试质量评估
8.3.1 优势
- 真实硬件验证：测试在真实 NPU 多卡环境运行，非 mock/simulator
- 端到端对比：fused_deep_moe 有完整的 Golden Reference 实现
- 多模式覆盖：Eager + GE 双模式
- 条件跳过：通过 marker 实现环境自适应
8.3.2 不足
| 不足 | 影响 | 建议 |
|---|---|---|
| 无 C++ 单元测试 | 算子内部逻辑（kernel、tiling）无独立验证 | 参考 URMA 引入 gtest |
| 无 Fuzz 测试 | 边界条件覆盖不足 | 引入 fuzz 生成随机 shape/config |
| 无 ASAN 检测 | 内存问题无法在测试阶段发现 | 测试构建启用 -fsanitize=address |
| 无代码覆盖率门禁 | 无法衡量测试完整性 | 接入 lcov + diff coverage |
| 无集成测试 | 算子间交互、端到端训练无验证 | 添加 CAM 集成测试套件 |
| 无 CI 配置 | 回归测试无法自动触发 | 添加 Jenkinsfile / GitHub Actions |
| CASE_8RANK 已定义未使用 | 8 卡场景无覆盖 | 启用或删除 |
| 无动态 batch / share / smooth 测试 | 关键功能分支未覆盖 | 扩展参数化矩阵 |
| 无异常输入测试 | 鲁棒性未知 | 添加边界值、空输入、超范围值测试 |
8.4 与 URMA 测试体系的对比总结
| 维度 | URMA | CAM |
|---|---|---|
| 测试语言 | C++ (gtest) + Python | Python (pytest) |
| 测试层级 | UT + Fuzz + 集成 | 仅 UT |
| 代码覆盖率 | lcov + 90% 门禁 | 无 |
| 内存安全 | ASAN | 无 |
| 构建系统 | CMake + ctest | shell + pytest |
| 报告格式 | JUnit XML + HTML | 控制台输出 |
| 分阶段执行 | 6 个 phase | 无 |
| 多芯片支持 | 无芯片限制 | A2/A3 marker（部分使用） |
| 多 Rank 支持 | 单进程为主 | 多进程（torchrun） |
九、关键发现与建议
9.1 架构级发现
-
CAM 采用"Python 测试 + C++ 算子"的混合模式：测试逻辑完全在 Python 层，通过 pybind 调用 C++ 算子。这种模式开发效率高，但无法覆盖 C++ 内部逻辑。
-
NPUBridge 是侵入式设计：通过
static_cast将c10::StorageImpl转为NPUStorageImpl，依赖 CAM 自己创建 tensor。如果 PyTorch 内部实现变化（如 StorageImpl 布局调整），可能导致隐蔽 bug。 -
pytorch_npu_helper.hpp 是完整的 ACL 抽象层：
EXEC_NPU_CMD宏封装了完整的算子调用生命周期（workspace 管理、大内存初始化、类型转换、错误处理），设计成熟。 -
FP8/FP4 支持有版本兼容设计：通过
CAM_ACL_HAS_FP8_FP4宏实现新旧 CANN 的兼容编译。
9.2 测试策略建议
-
短期： - 启用
CASE_8RANK测试 - 添加with_share=True、with_smooth=True、with_mc2_mask=True的参数化组合 - 实现run_with_cov.sh覆盖率收集 -
中期： - 为 C++ 核心逻辑（NPUBridge、StorageImpl、类型转换）添加 gtest - 添加 Fuzz 测试覆盖随机 shape/config - 接入 CI Pipeline（Jenkinsfile）
-
长期： - 构建 CAM 集成测试套件（参考 URMA 框架） - 添加端到端训练验证（小模型 + 真实数据） - 建立性能回归测试（吞吐量/延迟基线）
附录：关键文件路径索引
| 文件 | 路径 |
|---|---|
| conftest.py | test/cam/ut/conftest.py |
| pytest.ini | test/cam/ut/pytest.ini |
| marker.py | test/cam/ut/util/marker.py |
| tool.py | test/cam/ut/util/tool.py |
| fused_deep_moe 测试 | test/cam/ut/fused_deep_moe/test_fused_deep_moe.py |
| dispatch_combine 测试 | test/cam/ut/moe_dispatch_combine_prefill/test_moe_dispatch_combine_prefill.py |
| NPUBridge.h | src/cam/comm_operator/pybind/pytorch_extension/NPUBridge.h |
| NPUBridge.cpp | src/cam/comm_operator/pybind/pytorch_extension/NPUBridge.cpp |
| NPUStorageImpl.h | src/cam/comm_operator/pybind/pytorch_extension/NPUStorageImpl.h |
| NPUStorageImpl.cpp | src/cam/comm_operator/pybind/pytorch_extension/NPUStorageImpl.cpp |
| pytorch_npu_helper.hpp | src/cam/comm_operator/pybind/pytorch_extension/pytorch_npu_helper.hpp |
| pybind.cpp | src/cam/comm_operator/pybind/pybind.cpp |
| compile_ascend_proj.sh | build/cam/comm_operator/compile_ascend_proj.sh |
| build.sh | build/cam/comm_operator/build.sh |
| operator_registry.json | src/cam/comm_operator/ascend_kernels/operator_registry.json |
| select_ops.py | build/cam/comm_operator/select_ops.py |
| op_kernel/CMakeLists.txt | src/cam/comm_operator/ascend_kernels/cmake_files/op_kernel/CMakeLists.txt |
| op_host/CMakeLists.txt | src/cam/comm_operator/ascend_kernels/cmake_files/op_host/CMakeLists.txt |
| URMA UT CMakeLists | test/urma/CMakeLists.txt |
| URMA UT 脚本 | test/urma/script/urma_ut_common.sh |
| 集成测试框架 | test/intergration_test/test_framework/ |
1. SHMEM 模式整体架构
1.1 定位与使命
SHMEM 模式是 CAM 的 用户态 RDMA 直通路径。它完全绕过 HCCL 集合通信库，在 AIV（Ascend Intelligent Vector）核心上直接操作 RoCE RDMA 硬件的 WQE/CQE/SQ/CQ 寄存器，实现 kernel 内部的细粒度 P2P 数据传输。
┌──────────────────────────────────────────────────────────────────────┐
│ Host Side (CPU) │
│ │
│ RdmaTransportManagerV2 │
│ ├── OpenDevice() → phyId 获取 + RDMA IP 解析 │
│ ├── CreateEndpoint() → HcommEndpointCreate(COMM_PROTOCOL_ROCE) │
│ ├── RegisterMemoryRegion() → HcommMemReg (HBM 注册为 MR) │
│ ├── Connect() → HcommChannelCreate(COMM_ENGINE_AIV) │
│ └── FillRdmaInfo() → AiQpRMAQueueInfo 拷贝到 HBM │
│ │
│ 关键常量: MAX_RANKS_PER_NIC=16, ATOMIC_MAX_NUM=128 │
│ 底层协议: dl_hccp_def.h (HCCP) + dl_hcomm_def.h (Hcomm) │
└──────────────────────────────────────────────────────────────────────┘
│
│ qpInfo_ → AclrtMemcpy(Host→Device)
▼
┌──────────────────────────────────────────────────────────────────────┐
│ Device Side (AIV Core, Ascend C) │
│ │
│ 粗粒度 API 层 (shmem_device_rdma.hpp) │
│ ├── aclshmemx_roce_get_nbi<T>() ← RDMA READ (远端→本地) │
│ ├── aclshmemx_roce_put_nbi<T>() ← RDMA WRITE (本地→远端) │
│ ├── aclshmemx_roce_quiet<T>() ← 同步等待所有 WQE 完成 │
│ └── aclshmemx_roce_atomic_*() ← 16 种原子操作变体 │
│ │
│ 细粒度原语层 (rdma_device_backend_base.h) │
│ ├── aclshmemi_roce_fill_wqe<B, OP>() ← 构造 WQE │
│ ├── aclshmemi_roce_ring_sq_doorbell<B>() ← 振铃 SQ Doorbell │
│ ├── aclshmemi_roce_post_send<B, OP>() ← post 发送 (fill+bell) │
│ ├── aclshmemi_roce_ring_cq_doorbell<B>() ← 振铃 CQ Doorbell │
│ └── aclshmemi_roce_poll_cq<B>() ← 轮询 CQE 完成 │
│ │
│ 后端特化层 (compile-time 选择) │
│ ├── IN_DIE — 片内互联 (HCCS, 当前未完整实现) │
│ └── XSCALE — 跨机 RoCE (DDRM Diamond 引擎, 完整实现) │
│ │
│ 硬件寄存器: SQ Doorbell / CQ Doorbell / WQE(128B) / CQE(64B) │
└──────────────────────────────────────────────────────────────────────┘
1.2 数据流全景
一次 aclshmemx_roce_put_nbi (RDMA WRITE) 的完整路径：
用户调用 aclshmemx_roce_put_nbi(dst, src, buf, elem_size, pe)
│
├─ 1. aclshmem_ptr(dst, pe) ← 本地对称地址 → 远端对称地址
│ 计算: remote = p2p_device_heap_base[pe] + (dst - heap_base)
│
├─ 2. aclshmemi_roce_write<B>() ← 填充 send_wr 结构体
│ wr.remote_addr = ptr
│ wr.local_addr = src
│ wr.message_len = elem_size * sizeof(T)
│
├─ 3. aclshmemi_roce_post_send<B, OP_RDMA_WRITE>()
│ │
│ ├─ 3a. 读取 SQ head/tail (dcci_cachelines 刷新缓存行)
│ │ 若 SQ 即将满 (head+16 == tail)，先 poll_cq 等待空闲
│ │
│ ├─ 3b. aclshmemi_roce_fill_wqe<B, OP_RDMA_WRITE>()
│ │ → 在 HBM 上构造 128B WQE (ctrl_seg + remote_seg + local_seg)
│ │
│ ├─ 3c. dcci_cachelines(wqe_addr, 128) ← 刷回 HBM 确保硬件可见
│ │
│ └─ 3d. aclshmemi_roce_ring_sq_doorbell<B>()
│ → 更新 head_addr (producer index)
│ → 写 doorbell register: {next_pid, qp_id}
│
└─ 4. 硬件自动执行 RDMA WRITE，数据通过 RoCE 网络直达远端 HBM
2. 核心数据结构深度分析
2.1 aclshmemi_rdma_info — 全局 RDMA 信息表
文件: shmemi_device_rdma.h:38-45
struct aclshmemi_rdma_info {
uint32_t qp_num; // 每个连接的 QP 数量 (支持多 QP 并行)
uint64_t sq_ptr; // SQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t rq_ptr; // RQ 上下文数组地址 [PE_NUM][qp_num]
uint62_t scq_ptr; // Send CQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t rcq_ptr; // Recv CQ 上下文数组地址 [PE_NUM][qp_num]
uint64_t mem_ptr; // Memory Region 信息数组地址 [MAX_PE_NUM]
};
设计要点:
- 这是 device 端唯一的全局入口，通过 aclshmemi_qp_info_fetch() 从 HBM 固定地址读取
- 数组维度 [PE_NUM][qp_num] 支持多 rank 多 QP 场景
- mem_ptr 存储每个 rank 的 MR (Memory Region) 信息，包含 lkey/rkey 用于 RDMA 访问权限控制
2.2 aclshmemi_rdma_sq_ctx — Send Queue 上下文
文件: shmemi_device_rdma.h:54-66
struct aclshmemi_rdma_sq_ctx {
uint32_t wqn; // Work Queue Number (硬件 QP 编号)
uint64_t buf_addr; // SQ ring buffer 起始地址 (HBM)
uint32_t wqe_size; // 每个 WQE 的大小 (XSCALE=128B)
uint32_t depth; // ring buffer 深度 (WQE 数量)
uint64_t head_addr; // Producer Index 地址 (HBM, 软件写)
uint64_t tail_addr; // Consumer Index 地址 (HBM, 硬件写)
aclshmemi_rdma_db_mode_t db_mode; // HW_DB(0) / SW_DB(1)
uint64_t db_addr; // Doorbell register 地址
uint32_t sl; // Service Level (RDMA 优先级)
uint64_t amo_addr; // 原子操作的本地结果缓冲区地址
uint32_t amo_lkey; // 原子操作缓冲区的 local key
};
关键字段解读:
- head_addr / tail_addr: 分离的 producer/consumer index，支持硬件异步消费。head_addr 由软件通过 aclshmemi_roce_ring_sq_doorbell 更新，tail_addr 由硬件通过 CQE 完成事件推进。
- db_mode: 区分硬件 doorbell (HW_DB) 和软件 doorbell (SW_DB)，适应不同硬件版本。
- amo_addr / amo_lkey: 原子操作的专用缓冲区。RDMA atomic fetch 操作的结果会写回此地址，软件通过 dcci_cachelines 刷新后读取。
2.3 aclshmemi_rdma_cq_ctx — Completion Queue 上下文
文件: shmemi_device_rdma.h:68-77
struct aclshmemi_rdma_cq_ctx {
uint32_t cqn; // Completion Queue Number
uint64_t buf_addr; // CQ ring buffer 起始地址
uint32_t cqe_size; // 每个 CQE 的大小 (XSCALE=64B, 含 32B padding)
uint32_t depth; // ring buffer 深度
uint64_t head_addr; // Producer Index (硬件写)
uint64_t tail_addr; // Consumer Index (软件写)
aclshmemi_rdma_db_mode_t db_mode;
uint64_t db_addr; // CQ Doorbell register 地址
};
与 SQ 的对称性: CQ 的 head_addr 由硬件更新（新 CQE 写入时），tail_addr 由软件更新（aclshmemi_roce_ring_cq_doorbell 在 poll 完成后更新）。
2.4 aclshmemi_rdma_send_wr — 发送工作请求
文件: shmemi_device_rdma.h:121-141
struct aclshmemi_rdma_send_wr {
// ---- Operation ----
uint32_t send_flags; // 发送标志位掩码
uint32_t imm_data; // immediate data (带内通知数据)
// ---- Remote side ----
__gm__ uint8_t *remote_addr; // 远端 HBM 目标地址
uint32_t rkey; // 远端 Memory Region key
// ---- Local side ----
__gm__ uint8_t *local_addr; // 本地 HBM 源地址
uint64_t message_len; // 消息长度 (Bytes)
uint32_t lkey; // 本地 Memory Region key
// ---- Multi-SGE ----
uint32_t num_sge; // SGE 数量 (scatter-gather)
aclshmemi_rdma_sge *sg_list; // SGE 列表
// ---- Atomic ----
aclshmemi_rdma_atomic_params atomic; // 原子操作参数
};
设计要点:
- __gm__ 修饰符标记全局内存指针（HBM 地址空间）
- num_sge / sg_list 支持 scatter-gather 传输（当前 XSCALE 后端主要使用单 SGE）
- atomic 联合体支持多种原子操作参数格式
2.5 aclshmemi_rdma_atomic_params — 原子操作参数联合体
文件: shmemi_device_rdma.h:85-119
struct aclshmemi_rdma_atomic_params {
union {
// 标准 8B FA/CAS (unmasked)
struct { uint64_t compare_add; uint64_t swap; };
// 通用 masked 格式 (XSCALE 使用)
struct {
uint64_t swap_add_data;
uint64_t compare_data;
uint64_t swap_add_mask;
uint64_t compare_mask;
} masked_common;
// 64-bit masked CAS (双 64-bit 段: swap+compare, mask_swap+mask_compare)
struct { uint64_t swap_data; uint64_t compare_data;
uint64_t swap_mask; uint64_t compare_mask; } masked_cas64;
// 64-bit masked FA (add_data + field_boundary)
struct { uint64_t add_data; uint64_t field_boundary; } masked_fa64;
// 32-bit masked CAS
struct { uint32_t swap_data; uint32_t compare_data;
uint32_t swap_mask; uint32_t compare_mask; } masked_cas32;
// 32-bit masked FA
struct { uint32_t add_data; uint32_t field_boundary;
uint64_t _reserved; } masked_fa32;
};
};
设计哲学: 使用 union 实现参数格式的零开销复用。masked_common 是 XSCALE 后端实际使用的格式，其 4 个 64-bit 字段直接映射到 WQE 的 atomic segment。
2.6 aclshmemi_rdma_mem_info — Memory Region 信息
文件: shmemi_device_rdma.h:47-52
struct aclshmemi_rdma_mem_info {
uint64_t size; // MR 大小
uint64_t addr; // MR 起始地址
uint32_t lkey; // Local Key (本地访问权限)
uint32_t rkey; // Remote Key (远端 RDMA 访问权限)
};
3. API 调用链
3.1 粗粒度 → 细粒度调用层次
┌─────────────────────────────────────────────────────────────┐
│ 用户代码 / 算子 kernel │
│ aclshmemx_roce_put_nbi(dst, src, buf, len, pe) │
│ aclshmemx_roce_get_nbi(dst, src, buf, len, pe) │
│ aclshmemx_roce_quiet(pe, buf, sync_id) │
│ aclshmemx_roce_atomic_fetch_add(dst, val, pe) │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 后端无关转发层 (shmem_device_rdma.hpp) │
│ aclshmemi_roce_write<T, ACLSHMEMI_K_RDMA_BACKEND>(...) │
│ → 根据 ACLSHMEMI_K_RDMA_BACKEND 宏选择 XSCALE 或 IN_DIE │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 细粒度原语层 (rdma_device_backend_base.h/.hpp) │
│ aclshmemi_roce_post_send<B, OP_RDMA_WRITE>(wr, pe, ...) │
│ → aclshmemi_roce_fill_wqe<B, OP_CODE>(wr, sq, addr, hd) │
│ → aclshmemi_roce_ring_sq_doorbell<B>(sq, head, ...) │
│ aclshmemi_roce_poll_cq<B>(pe, qp_idx, target, ...) │
│ aclshmemi_roce_ring_cq_doorbell<B>(pe, qp, tail, ...) │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 后端特化层 (rdma_device_backend_xscale.hpp) │
│ aclshmemi_xscale_post_send_read_write<OP_CODE>(...) │
│ → aclshmemi_roce_xscale_fill_wqe_ctrl_seg(...) │
│ → aclshmemi_roce_xscale_fill_wqe_data_seg(...) │
│ → aclshmemi_roce_write_ub_to_gm_with_sync(...) │
│ aclshmemi_rdma_xscdv_fill_wqe_fa<T, IS_MASKED>(...) │
│ aclshmemi_rdma_xscdv_fill_wqe_cas<T, IS_MASKED>(...) │
└─────────────────────────────────────────────────────────────┘
3.2 RDMA WRITE 完整调用链
// Step 1: 用户调用
aclshmemx_roce_put_nbi<uint32_t>(dst, src, buf, 256, pe);
// Step 2: 地址转换 + 转发
// shmem_device_rdma.hpp:183-197
aclshmem_ptr(dst, pe); // 本地对称地址 → 远端对称地址
aclshmemi_roce_write<uint32_t, XSCALE>(ptr, src, pe, 0, 1024, ...);
// Step 3: 填充 wr + post_send
// rdma_device_backend_base.h:137-149
wr.remote_addr = ptr;
wr.local_addr = src;
wr.message_len = 1024;
aclshmemi_roce_post_send<XSCALE, OP_RDMA_WRITE>(wr, pe, 0, ...);
// Step 4: XSCALE 后端 post_send
// rdma_device_backend_xscale.hpp:609-667
aclshmemi_xscale_post_send_read_write<OP_CODE>(wr, pe, qp_idx, ...) {
// 4a. 读取 SQ head/tail
dcci_cachelines(sq_pi_addr, 8);
cur_head = *(uint32_t*)sq_pi_addr;
// 4b. SQ 满检查 (阈值=16)
if ((cur_head + 16) % depth == (depth + cur_tail) % depth)
aclshmemi_roce_poll_cq<XSCALE>(...);
// 4c. 获取 rkey/lkey
wr.rkey = remote_mem_info->rkey;
wr.lkey = local_mem_info->lkey;
// 4d. 填充 WQE
aclshmemi_roce_fill_wqe<XSCALE, OP_RDMA_WRITE>(wr, sq, wqe_addr, cur_head);
// 4e. 刷回 + doorbell
dcci_cachelines(wqe_addr, 128);
aclshmemi_roce_ring_sq_doorbell<XSCALE>(sq, cur_head+1, ...);
}
3.3 aclshmemx_roce_quiet 同步机制
// shmem_device_rdma.hpp:253-276
template <typename T>
void aclshmemx_roce_quiet(uint32_t pe, __ubuf__ T *buf, uint32_t sync_id) {
for (qp_idx = 0; qp_idx < qp_num; qp_idx++) {
sq_context = ...;
dcci_cachelines(sq_pi_addr, 8); // 刷新 head 缓存行
cur_head = *(uint32_t*)sq_pi_addr; // 读取当前 producer index
aclshmemi_roce_poll_cq<XSCALE>(pe, qp_idx, cur_head, ...);
// poll_cq 会阻塞直到 cur_tail == cur_head (所有 WQE 完成)
}
}
关键设计: quiet 遍历所有 QP，对每个 QP 读取当前 SQ head（即已 post 的 WQE 总数），然后 poll_cq 等待 CQ tail 追上该值。
4. 编译时多态 traits 系统
4.1 设计理念
CAM 使用 C++ 模板特化 实现零运行时开销的后端选择。核心思想：
在编译期通过
ACLSHMEMI_K_RDMA_BACKEND宏确定后端，所有aclshmemi_roce_*<B>()调用在编译期绑定到具体后端实现，运行时无 switch/if 分支。
4.2 后端选择宏
// shmem_device_rdma.hpp:25-29
#if defined(ACLSHMEMI_RDMA_K_BACKEND_XSCALE)
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::XSCALE)
#else
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::IN_DIE)
#endif
4.3 traits 系统架构
// rdma_device_backend_base.h:263-315 — 主模板 (未定义, 强制特化)
template <aclshmemi_rdma_backend_t B>
struct aclshmemi_backend_traits {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...);
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...);
};
};
// rdma_device_backend_xscale.hpp:883-905 — XSCALE 特化
template <>
struct aclshmemi_backend_traits<aclshmemi_rdma_backend_t::XSCALE> {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...) {
return aclshmemi_roce_xscale_fill_wqe_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...) {
aclshmemi_rdma_xscdv_post_send_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
};
};
4.4 编译期错误防护
// rdma_device_backend_base.h:318-323
template <aclshmemi_rdma_atomic_op_t>
struct aclshmemi_atomic_op_dependent_false : std::false_type {};
template <aclshmemi_rdma_backend_t>
struct aclshmemi_rdma_backend_dependent_false : std::false_type {};
当调用未特化的后端或未实现的原子操作时，static_assert 触发编译错误而非运行时崩溃。
4.5 零运行时开销验证
所有关键路径函数标记为 ACLSHMEM_DEVICE（展开为 __forceinline__ __device__），配合 if constexpr 实现编译期分支消除：
// shmem_device_rdma.hpp:75-81
if constexpr (ACLSHMEMI_K_RDMA_BACKEND == aclshmemi_rdma_backend_t::XSCALE) {
return aclshmemi_roce_atomic_fetch_and_add<T, IS_MASKED, XSCALE>(...);
} else {
ACLSHMEM_DEBUG_FUNC(aclshmemi_kernel_abort, "ROCE atomic add is only supported on XSCALE backend.\n");
return T(0);
}
编译后，XSCALE 后端下 else 分支被完全消除，无运行时判断。
5. XSCALE 后端 WQE/CQE 格式
5.1 WQE 整体结构 (XSCALE, 128 Bytes)
XSCALE 后端的 WQE 为 128 Bytes (ACLSHMEMI_XSCALE_SND_WQE_SIZE = 1 << 7 = 128)，由以下部分组成：
┌──────────────────────────────────────────────────────┐
│ Control Segment (16 Bytes) │
│ ├── msg_opcode (8 bits) — 操作码 │
│ ├── with_imm (1 bit) — 是否带 immediate │
│ ├── csum_en (2 bits) — 校验和使能 │
│ ├── ds_data_num (5 bits) — Data Segment 数量 │
│ ├── wqe_id (16/20 bits) — WQE 标识 │
│ ├── msg_len (32 bits) — 消息长度 │
│ ├── opcode_data (32 bits) — 操作码扩展数据 │
│ ├── se (1 bit) — Solicitation Event │
│ ├── ce (1 bit) — Completion Enable │
│ ├── in_line (1 bit) — Inline 模式 │
│ ├── fence_mode (2 bits) — Fence 模式 │
│ └── mask (2 bits) — 掩码 │
├──────────────────────────────────────────────────────┤
│ Remote Data Segment (16 Bytes) │
│ ├── seg_len (31 bits) — 远端数据长度 │
│ ├── rkey (32 bits) — 远端 Memory Region Key │
│ └── va (64 bits) — 远端虚拟地址 │
├──────────────────────────────────────────────────────┤
│ Local Data Segment (16 Bytes) │
│ ├── seg_len (31 bits) — 本地数据长度 │
│ ├── lkey (32 bits) — 本地 Memory Region Key │
│ └── va (64 bits) — 本地虚拟地址 │
├──────────────────────────────────────────────────────┤
│ Atomic Segment (16 Bytes, 仅原子操作) │
│ ├── swap_add (64 bits) — 交换/加数据 │
│ └── compare (64 bits) — 比较数据 │
├──────────────────────────────────────────────────────┤
│ Padding (64 Bytes, 补齐到 128B) │
└──────────────────────────────────────────────────────┘
5.2 Control Segment 版本差异
XSCALE 支持两个 API 版本，通过 ACLSHMEMI_XSCALE_API_VERSION_VAR 宏选择：
// V1: wqe_id 占 16 bits (bits [0:15] of the last 32-bit field)
struct aclshmemi_xscdv_wqe_ctrl_seg_v1 {
uint8_t msg_opcode; // [0:7]
uint8_t with_imm : 1; // [8]
uint8_t csum_en : 2; // [9:10]
uint8_t ds_data_num : 5; // [11:15]
uint16_t wqe_id; // [16:31] — 16 bits
uint32_t msg_len; // [32:63]
uint32_t opcode_data; // [64:95]
uint32_t se : 1; // [96]
uint32_t ce : 1; // [97]
uint32_t in_line : 1; // [98]
uint32_t fence_mode : 2; // [99:100]
uint32_t mask : 2; // [101:102]
uint32_t rsv : 25; // [103:127]
};
// V2: wqe_id 扩展到 20 bits, 重新排列字段
struct aclshmemi_xscdv_wqe_ctrl_seg_v2 {
uint8_t msg_opcode;
uint8_t with_imm : 1;
uint8_t csum_en : 2;
uint8_t ds_data_num : 5;
uint16_t rsv1; // 保留
uint32_t msg_len;
uint32_t opcode_data;
uint32_t se : 1;
uint32_t ce : 1;
uint32_t in_line : 1;
uint32_t fence_mode : 2;
uint32_t mask : 2;
uint32_t rsv : 5;
uint32_t wqe_id : 20; // [27:47] — 20 bits (支持更多 WQE)
};
版本选择机制:
using aclshmemi_xscdv_wqe_ctrl_seg_t =
std::conditional_t<ACLSHMEMI_XSCALE_API_VERSION_VAR == 1,
aclshmemi_xscdv_wqe_ctrl_seg_v1,
aclshmemi_xscdv_wqe_ctrl_seg_v2>;
5.3 CQE 格式 (XSCALE Diamond, 64 Bytes)
struct aclshmemi_xscdv_cqe64_t {
aclshmemi_xscdv_diamond_cqe_t cqe; // 32 Bytes
uint8_t padding[32]; // 32 Bytes padding
};
// V2 CQE (32 Bytes):
struct aclshmemi_xscdv_diamond_cqe_v2 {
union {
struct {
uint32_t error_code : 8; // [0:7] 错误码
uint32_t qp_id : 15; // [8:22] 完成的 QP 编号
uint32_t rsv : 1; // [23]
uint32_t se : 1; // [24] Solicitation Event
uint32_t has_pph : 1; // [25] 是否有 PPH
uint32_t type : 1; // [26]
uint32_t with_imm : 1; // [27] 是否有 immediate
uint32_t csum_err : 4; // [28:31] 校验和错误
};
uint32_t flags_qp_id_err_code;
};
uint32_t imm_data; // immediate value
uint32_t msg_len; // 消息长度
uint32_t vni; // Virtual Network Identifier
uint32_t ts_l; // 时间戳低 32 位
uint32_t ts_h; // 时间戳高 16 位 (共 48 位)
union {
struct {
uint32_t msg_opcode : 8; // 对应 WQE 的 opcode
uint32_t rsv1 : 4;
uint32_t wqe_id : 20; // 对应 WQE 的 wqe_id
};
uint32_t wqe_id_rsv_opcode;
};
union {
struct {
uint32_t rsv2 : 31;
uint32_t owner : 1; // Owner bit: 软件解析权限标志
};
uint32_t owner_rsv;
};
};
5.4 Owner Bit 机制
XSCALE 使用 owner bit 实现硬件/软件同步，避免昂贵的 atomic 操作：
// rdma_device_backend_xscale.hpp:389-397
bool aclshmemi_roce_xscale_check_cqe_owner(__gm__ aclshmemi_xscdv_cqe64_t *cqe64,
uint32_t cur_tail, uint32_t depth) {
bool owner_bit = (cqe64->cqe.owner & ACLSHMEMI_XSC_CQE_OWNER_MASK);
bool expect_bit = ((cur_tail / depth) & 1); // 每翻转一轮 depth，owner bit 翻转
return owner_bit == expect_bit;
}
原理: 硬件每写完一轮 CQ（depth 个 CQE），自动翻转 owner bit。软件通过比较 owner_bit == (cur_tail / depth) & 1 判断 CQE 是否已写入。
5.5 Doorbell 格式
SQ Doorbell (8 Bytes):
union aclshmemi_xscdv_diamond_send_doorbell_v2 {
struct {
uint64_t next_pid : 21; // 下一个 WQE 的 producer index (已左移)
uint64_t qp_id : 16; // QP 编号
};
uint64_t raw;
};
CQ Doorbell (8 Bytes):
union aclshmemi_xscdv_diamond_cq_doorbell_t {
struct {
uint64_t cq_next_cid : 23; // 下一个待处理的 CQE ID
uint64_t cq_id : 16; // CQ 编号
uint64_t cq_sta : 2; // CQ 状态
};
uint64_t raw;
};
5.6 操作码映射
enum class aclshmemi_xscdv_msg_type_t : uint32_t {
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_WRITE = 1,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_READ = 2,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_CMP_AND_SWAP = 26,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_FETCH_AND_ADD = 27,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_CMP_AND_SWAP = 31,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_FETCH_AND_ADD = 32,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_CMP_AND_SWAP = 33,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_FETCH_AND_ADD = 34,
};
6. 原子操作机制
6.1 原子操作 API 全景
CAM 提供 16 种 原子操作 API，覆盖 fetch/set/compare-swap/inc/add/and/or/xor 及其组合：
| 类别 | API | 同步性 | 实现方式 |
|---|---|---|---|
| Fetch | atomic_fetch(src, pe) |
同步 | FA(add=0) + quiet + read |
| Set | atomic_set(dst, val, pe) |
异步 | CAS(swap=val, comp=0, mask_swap=MAX) |
| CAS | atomic_compare_swap(dst, cond, val, pe) |
同步 | CAS + quiet + read |
| Swap | atomic_swap(dst, val, pe) |
同步 | CAS(swap=val, comp=0) + quiet + read |
| Inc | atomic_inc(dst, pe) |
异步 | FA(add=1) |
| Add | atomic_add(dst, val, pe) |
异步 | FA(add=val) |
| And | atomic_and(dst, val, pe) |
异步 | CAS(swap=val, mask_swap=~val) |
| Or | atomic_or(dst, val, pe) |
异步 | CAS(swap=val, mask_swap=val) |
| Xor | atomic_xor(dst, val, pe) |
异步 | FA(add=val, boundary=MAX) |
| FetchInc | atomic_fetch_inc(dst, pe) |
同步 | FA(add=1) + quiet + read |
| FetchAdd | atomic_fetch_add(dst, val, pe) |
同步 | FA(add=val) + quiet + read |
| FetchAnd | atomic_fetch_and(dst, val, pe) |
同步 | CAS + quiet + read |
| FetchOr | atomic_fetch_or(dst, val, pe) |
同步 | CAS + quiet + read |
| FetchXor | atomic_fetch_xor(dst, val, pe) |
同步 | FA + quiet + read |
6.2 Masked vs Unmasked 变体
XSCALE 后端对原子操作区分 masked 和 unmasked 两种模式：
// 4B 操作强制使用 masked 模式
static_assert(!(TYPE_BYTES <= BYTES_32 && !IS_MASKED),
"XSCALE backend doesn't support fetch and add operation with type size of 4B without mask\n");
Masked 模式 允许对 64-bit 字中的特定位字段进行原子操作，通过 swap_add_mask / compare_mask 控制哪些位参与操作。
6.3 原子操作 WQE 构造
以 aclshmemi_rdma_xscdv_fill_wqe_fa 为例：
template <typename T, bool IS_MASKED>
uint32_t aclshmemi_rdma_xscdv_fill_wqe_fa(aclshmemi_rdma_send_wr &wr, ...) {
constexpr uint32_t DS_DATA_NUM = 3; // ctrl + data_seg + atomic_seg
constexpr uint32_t MSG_LEN = (TYPE_BYTES == BYTES_32 && IS_MASKED) ? 8 : 16;
constexpr uint32_t DATA_LEN = (TYPE_BYTES <= BYTES_32) ? BYTES_32 : BYTES_64;
// 根据类型选择操作码
constexpr aclshmemi_xscdv_msg_type_t OPCODE =
(TYPE_BYTES == BYTES_64 && IS_MASKED) ? ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_FETCH_AND_ADD
: (TYPE_BYTES == BYTES_64 && !IS_MASKED) ? ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_FETCH_AND_ADD
: ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_FETCH_AND_ADD;
// 填充 ctrl_seg (inline=1, 因为数据在 WQE 内)
cur_wqe_addr = aclshmemi_roce_xscale_fill_wqe_ctrl_seg(wqe_addr, DS_DATA_NUM, cur_head, OPCODE, MSG_LEN, 1);
// 填充 data_seg (remote + local)
cur_wqe_addr = aclshmemi_roce_xscale_fill_wqe_data_seg(cur_wqe_addr, wr.rkey, wr.remote_addr, wr.lkey, wr.local_addr, DATA_LEN);
// 填充 atomic_seg
return aclshmemi_rdma_xscdv_fill_fa_atomic_seg<T, IS_MASKED>(cur_wqe_addr, wr);
}
6.4 字节序转换
XSCALE 硬件要求 Big-Endian 格式的原子操作数据：
// rdma_device_backend_xscale.hpp:333-343
ACLSHMEM_DEVICE uint64_t aclshmemi_htobe64(uint64_t host_val) {
return ((host_val >> 0) & 0xFF) << 56 |
((host_val >> 8) & 0xFF) << 48 |
((host_val >> 16) & 0xFF) << 40 |
((host_val >> 24) & 0xFF) << 32 |
((host_val >> 32) & 0xFF) << 24 |
((host_val >> 40) & 0xFF) << 16 |
((host_val >> 48) & 0xFF) << 8 |
((host_val >> 56) & 0xFF) << 0;
}
读取 fetch 结果时需反向转换：
// shmem_device_rdma.hpp:293-297
if constexpr (sizeof(T) == 4 && ACLSHMEMI_K_RDMA_BACKEND == aclshmemi_rdma_backend_t::XSCALE) {
uint32_t fetch_bytes = *fetch_addr;
return (T)aclshmemi_htobe32(fetch_bytes); // 4B 结果需要字节序转换
}
6.5 原子操作结果读取
// shmem_device_rdma.hpp:278-301
uint64_t aclshmemi_roce_get_atomic_fetch_addr(uint32_t pe, uint32_t qp_idx) {
// 从 sq_ctx->amo_addr 获取原子操作结果缓冲区地址
return sq_context->amo_addr;
}
template <typename T>
T aclshmemi_roce_get_atomic_fetch_data(uint32_t pe, uint32_t qp_idx) {
auto amo_addr = aclshmemi_roce_get_atomic_fetch_addr(pe, qp_idx);
dcci_cachelines((__gm__ uint8_t *)amo_addr, sizeof(T)); // 刷新缓存行
__gm__ T *fetch_addr = reinterpret_cast<__gm__ T *>(amo_addr);
return *fetch_addr; // 读取 fetch 前的原始值
}
7. Host 端传输管理
7.1 RdmaTransportManagerV2 生命周期
OpenDevice() → RegisterMemoryRegion() × N → Prepare() → Connect() → [运行期] → CloseDevice()
7.2 OpenDevice 流程
// device_rdma_transport_manager_v2.cpp:44-101
Result OpenDevice(const TransportOptions& options) {
// 1. 获取设备 ID
DlAclApi::AclrtGetDevice(&userId);
DlAclApi::AclrtGetPhyDevIdByLogicDevId(userId, &phyId_);
// 2. 解析 RDMA 网络地址
TopoReader::ParseRdmaNetAddr(phyId_, deviceIp_);
devicePort_ = RDMA_PORT_PREFIX(60032) + (rankId_ % MAX_RANKS_PER_NIC(16));
nicInfo_ = GenerateDeviceNic(deviceIp_, devicePort_);
// 3. 创建 RoCE 端点
CreateEndpoint(); // → HcommEndpointCreate(&endpointDesc, &endpointHandle_)
// 4. 预留 device 内存 (qpInfo + atomic buffer)
ReserveRdmaInfoSpace();
}
7.3 Connect 流程
// device_rdma_transport_manager_v2.cpp:280-307
Result Connect() {
// 1. 验证单 NIC rank 数不超过 16
ValidateRanksPerNic();
// 2. 准备 channel 描述符
uint32_t channelNum = rankCount_ - 1;
PrepareChannelDescs(channelDescs, roceTc, roceSl);
// 对每个 remote rank:
// - 设置 protocol = COMM_PROTOCOL_ROCE
// - 设置 IP 地址 (IPv4/IPv6)
// - 设置 tc/sl (Traffic Class / Service Level)
// - 设置 port = 60032 + serverRank * 16 + clientRank
// - 设置 role = SERVER (rankId < remoteRank) / CLIENT
// 3. 创建 channel 并填充 QP 信息
CreateChannelsAndFillInfo(channelDescs, channelNum);
}
7.4 FillRdmaInfo — 核心数据提取
// device_rdma_transport_manager_v2.cpp:507-531
Result FillRdmaInfo() {
// 1. 分配临时 buffer
std::vector<uint8_t> qpInfoBuffer(qpInfoSize_);
auto copyInfo = (AiQpRMAQueueInfo*)qpInfoBuffer.data();
// 2. 设置指针布局 (sq/rq/scq/rcq/mr 紧跟在 AiQpRMAQueueInfo 之后)
FillQpPreSettingCopyInfo(copyInfo);
// 3. 从 ChannelEntity 读取 SQ/CQ/MR 信息
GetRdmaInfoFromChannelEntity(copyInfo, channelPtrs_);
// ├── ReadLocalBufferInfo() → 读取本地 MR 的 lkey/rkey
// └── ReadRemoteBufferInfo() → 遍历每个 remote rank:
// ├── ReadSingleRemoteRank() → D2H 拷贝 ChannelEntity
// │ ├── 读取 remote MR 的 lkey/rkey
// │ └── FillSqCqAtomicInfo() → 读取 SQ/CQ 上下文
// │ ├── CopyAiWQInfo() → 提取 wqn/bufAddr/wqeSize/depth/head/tail/dbAddr
// │ ├── CopyAiCQInfo() → 提取 cqn/bufAddr/cqeSize/depth/head/tail/dbAddr
// │ └── 设置 atomicAddr/atomicLkey
// └── 设置 remote MR 的 size/addr
// 4. 拷贝到 device HBM
DlAclApi::AclrtMemcpy(qpInfo_, qpInfoSize_, copyInfo, qpInfoSize_, ACL_MEMCPY_HOST_TO_DEVICE);
}
7.5 内存注册策略
// device_rdma_transport_manager_v2.cpp:191-224
Result RegisterMemoryRegion(const TransportMemoryRegion& mr) {
CommMem commMem{};
// 区分 HBM 和 Host DDR
if ((mr.flags & REG_MR_FLAG_HBM) || IsVirtualAddressNpu(mr.addr))
commMem.type = COMM_MEM_TYPE_DEVICE; // HBM
else
commMem.type = COMM_MEM_TYPE_HOST; // DDR
// 注册到 HCCP 协议栈
HcommMemReg(endpointHandle_, "HcclBuffer", &commMem, &memHandle);
registeredMRs_.emplace(mr.addr, entry);
}
7.6 原子操作共享内存
// device_rdma_transport_manager_v2.cpp:703-735
bool ReserveRdmaInfoSpace() noexcept {
// qpInfo 大小: sizeof(AiQpRMAQueueInfo) + rankCount * [2*(sizeof(AiQpRMAWQ) + sizeof(AiQpRMACQ)) + sizeof(RdmaMemRegionInfo)]
qpInfoSize_ = sizeof(AiQpRMAQueueInfo) + oneQpSize * rankCount_;
DlAclApi::AclrtMalloc(&ptr, qpInfoSize_, 0);
qpInfo_ = (AiQpRMAQueueInfo*)ptr;
// 原子操作共享内存: 128 * 8B * rankCount, 512B 对齐
uint32_t atomicSize = ATOMIC_MAX_NUM(128) * sizeof(uint64_t) * rankCount;
atomicSize = ALIGN_UP(atomicSize, MEMORY_ALIGNMENT(512));
DlAclApi::AclrtMalloc(&atomicSharedMemory_, atomicSize, 0);
}
设计要点: 原子操作共享内存用于存储 fetch-before-write 结果。每个 rank 占用 128 × 8B = 1KB，支持最大 128 个并发原子操作。
7.7 底层协议结构
HCCP 协议 (dl_hccp_def.h):
- HccpMrInfo: Memory Region 信息 (addr/size/access/lkey/rkey)
- HccpQpExtAttrs: QP 扩展属性 (qpMode/cqAttr/qp_attr/version/mem_align)
- ai_data_plane_wq / ai_data_plane_cq: 硬件队列描述 (wqn/cqn/buf_addr/wqebb_size/depth/head/tail/db_reg)
- AiQpRMAWQ / AiQpRMACQ: 运行时 SQ/CQ 信息 (含 atomicAddr/atomicLkey)
Hcomm 协议 (dl_hcomm_def.h):
- ChannelEntity: 通道实体 (engine/protocol/bufferNum/notifyNum/sqNum/cqNum)
- SqContext / CqContext: 队列上下文 (roceSq/roceCq 包含 qpn/cqn/sqVa/cqVa/dbVa 等)
- RegedBufferEntity: 注册缓冲区 (addr/size/protectionInfo)
- ProtectionInfo: 保护信息 (type=PROTECTION_TYPE_ROCE, lkey/rkey)
8. 编译注入机制
8.1 SHMEM_HOME_PATH 环境变量
# ascendc_impl_build.py:167-178
SHMEM_HOME_PATH = os.environ.get('SHMEM_HOME_PATH')
if SHMEM_HOME_PATH is None:
SHMEM_HOME_PATH = os.path.realpath("/usr/local/Ascend/shmem/latest")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/host_device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/internal")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src/device")
options.append("-I" + SHMEM_HOME_PATH + "/memfabric_hybrid/include/smem/device")
8.2 编译期后端选择
# compile_ascend_proj.sh:179-192
KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
# 通过 sed 注入到 build.make 的 --compile-options 中
find build_out/op_kernel/CMakeFiles -name "build.make" -exec sed -i \
"s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
8.3 算子选择过滤
# select_ops.py:55-60, 116-121
def requires_shmem(registry, op):
return bool(get_meta(registry, op).get("requires_shmem", False))
# 在 resolve() 中:
if not shmem:
dropped = [o for o in candidates if requires_shmem(registry, o)]
for o in dropped:
sys.stderr.write(f"NOTE: dropping {o} (requires SHMEM, which is not installed)\n")
candidates = [o for o in candidates if not requires_shmem(registry, o)]
8.4 当前算子注册表状态
// operator_registry.json
"operator_meta": {
"dispatch_layout": { "requires_shmem": false },
"fused_deep_moe": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "base" },
"moe_combine_normal": { "requires_shmem": false },
"moe_dispatch_normal":{ "requires_shmem": false },
"notify_dispatch": { "requires_shmem": false }
}
当前所有算子 requires_shmem: false, 即默认使用 HCCL Window 模式。SHMEM 模式基础设施已完整就绪，但尚未在算子层启用。
9. 与 HCCL Window 模式的本质区别
9.1 架构层面对比
| 维度 | SHMEM 模式 | HCCL Window 模式 |
|---|---|---|
| 通信路径 | AIV 核心直接操作 RoCE 硬件 WQE/CQE | 通过 HCCL Window 机制 (winContext_->localWindowsIn) |
| 运行时开销 | 单次 kernel 内完成 RDMA 操作，无 HCCL 调度 | 需经过 HCCL 调度层、Window 资源管理 |
| 延迟 | 微秒级 (直接 doorbell 振铃) | 数十微秒级 (HCCL 软件栈开销) |
| 灵活性 | 支持任意地址、任意长度的 P2P 传输 | 受限于 Window 预分配区域 |
| 原子操作 | 16 种 RDMA 原子操作 (CAS/FA/AND/OR/XOR) | 依赖 HCCL 集合通信原语 |
| 计算-通信重叠 | 同一 kernel 内交替执行 GEMM 和 RDMA | 需多 kernel 协作，HBM 读写开销 |
| 后端选择 | 编译期 traits 系统 (IN_DIE/XSCALE) | HCCL 内部自动选择 |
| 内存注册 | 显式 MR 注册 (HcommMemReg) | HCCL 隐式管理 |
| QP 管理 | 显式 SQ/CQ 上下文管理 | HCCL 抽象 |
9.2 代码层面对比
SHMEM 模式 (直接 RDMA):
// 算子 kernel 内部直接发起 RDMA
aclshmemx_roce_put_nbi(remote_addr, local_addr, ub_buf, len, remote_pe);
// ↑ 直接构造 WQE → doorbell → 硬件执行
HCCL Window 模式 (当前默认):
// 通过 HCCL Window 机制
winContext_->localWindowsIn[peer_rank] // 预分配的 Window 区域
HcclRankRelationResV2 // HCCL 管理的 rank 关系
// ↑ 需经过 HCCL 调度、Window 资源同步
9.3 性能影响分析
SHMEM 模式的核心优势在于 消除 HCCL 软件栈开销:
- 零调度延迟: 直接 doorbell 振铃 (~100ns) vs HCCL 调度 (~10-50μs)
- Kernel 内融合: FusedDeepMoE 可在同一 kernel 内交替执行 GEMM 和 RDMA，避免多次 kernel launch 和 HBM 读写
- 细粒度 P2P: 支持 token 级别的动态路由，无需预定义通信模式
- 硬件原子操作: 利用 RDMA 硬件原子操作实现无锁同步
9.4 启用路径
要将算子从 HCCL Window 模式切换到 SHMEM 模式:
- 修改 operator_registry.json: 将目标算子的
requires_shmem设为true - 确保 SHMEM 安装: 设置
SHMEM_HOME_PATH环境变量指向 SHMEM 安装目录 - 算子代码改造: 在 kernel 中调用
aclshmemx_roce_*API 替代 HCCL Window 操作 - 编译:
select_ops.py会自动包含该算子，编译注入 SHMEM include 路径
10. 关键设计洞察
10.1 用户态 RDMA 的极致优化
CAM 的 SHMEM 模式是 用户态 RDMA 在 AI 芯片上的极致实现:
- AIV 核心直接驱动: 不同于传统 RDMA 由 CPU 或专用 NIC 驱动，CAM 让 AIV 向量核心直接构造 WQE 和振铃 doorbell
- Cache-line 精确刷新: 使用
dcci_cachelines()精确刷新 64B 缓存行，确保 HBM 数据对硬件可见 - Inline 模式: 原子操作使用 inline 模式 (数据在 WQE 内)，避免额外的 DMA 读取
10.2 编译时多态的工程价值
traits 系统实现了 零成本抽象:
- 无虚函数表 (vtable) 开销
- 无运行时 switch/if 分支
- 未使用的后端代码在编译期完全消除
if constexpr确保编译期分支消除
10.3 版本兼容性的优雅处理
XSCALE 后端的 V1/V2 API 版本通过 std::conditional_t 实现编译期选择:
using aclshmemi_xscdv_wqe_ctrl_seg_t =
std::conditional_t<ACLSHMEMI_XSCALE_API_VERSION_VAR == 1,
aclshmemi_xscdv_wqe_ctrl_seg_v1,
aclshmemi_xscdv_wqe_ctrl_seg_v2>;
这允许同一份代码支持不同硬件世代的 WQE 格式差异。
10.4 同步机制的精心设计
- SQ head/tail 分离: 软件写 head (producer), 硬件写 tail (consumer via CQE)
- CQ owner bit: 硬件/软件同步无 atomic 操作，每轮 depth 翻转一次
- SQ 满检查: 阈值 16 的预检查避免死锁，同时最大化 SQ 利用率
- 超时保护: 5 分钟超时防止硬件异常导致无限等待
10.5 原子操作的统一抽象
通过 aclshmemi_rdma_atomic_params union + IS_MASKED 模板参数，实现了:
- 8B/4B 数据类型统一处理
- masked/unmasked 模式编译期选择
- 多种原子操作 (FA/CAS/AND/OR/XOR) 的统一 WQE 构造路径
11. 总结
CAM 的 SHMEM 模式是一个 完整的用户态 RDMA 子系统，其核心优势在于:
- 硬件直通: AIV 核心直接操作 RoCE 硬件，消除 HCCL 软件栈开销
- 编译时多态: traits 系统实现零运行时开销的后端选择
- 原子操作丰富: 16 种原子操作变体，支持 8B/4B masked/unmasked 模式
- 版本兼容: V1/V2 API 版本编译期选择，支持多硬件世代
- 工程完备: 超时保护、错误处理、日志调试机制齐全
当前所有算子默认使用 HCCL Window 模式 (requires_shmem: false)，但 SHMEM 模式基础设施已完整就绪，是 CAM 未来性能突破的关键路径。
日期: 2026-07-21 范围: CAM 库中
moe_dispatch_normal+moe_combine_normal的完整 3 层栈 (op_api → op_host → op_kernel) 模式: 生产默认模式 (requires_shmem: false)，所有算子当前均使用此路径
目录
- 整体架构概览
- Window 内存布局
- moe_dispatch_normal 完整执行流程
- Window 地址计算
- Ping-Pong 双缓冲机制
- Magic+Value 同步协议
- moe_combine_normal 完整执行流程
- Tiling 策略
- 常量速查表
- 与 SHMEM 模式的代码级对比
- 关键执行流程时序图
1. 整体架构概览
HCCL Window 模式是 CAM 库的生产默认通信路径。与 SHMEM 模式（基于 IPC 共享内存 + 自管理 flag 区）不同，Window 模式直接利用 HCCL 框架分配的 Window 内存窗口（localWindowsIn / localWindowsExp + 远端 windowsIn / windowsExp）进行跨片数据通信。
核心数据结构关系
HcclOpResParam (per-rank, 来自 HCCL 上下文)
├── localWindowsIn → 本卡 Window 输入区 (combine 读取)
├── localWindowsOut → 本卡 Window 输出区 (预留)
├── localWindowsExp → 本卡 Window 扩展区 (dispatch 状态同步)
├── remoteRes[] → 远端资源指针数组
│ └── HcclRankRelationResV2
│ ├── windowsIn → 远端 Window 输入区 (本卡 DMA 写入)
│ ├── windowsOut → 远端 Window 输出区
│ └── windowsExp → 远端 Window 扩展区 (flag 同步)
└── winSize → Window 总大小 (由 HCCL_BUFFSIZE 环境变量控制)
通信模式对比
| 维度 | HCCL Window 模式 (Non-SHMEM) | SHMEM 模式 |
|---|---|---|
| 内存来源 | HCCL 框架预分配 Window | IPC 共享内存 + 自管理 flag 区 |
| 同步机制 | Magic+Value flag (高32位 magic + 低32位 value) | 自管理 flag + chunk 协议 |
| 双缓冲 | Ping-Pong (dataState 0/1) | 同样 Ping-Pong |
| 数据区 | localWindowsIn / remoteWindowsIn |
IPC data 区 |
| Flag 区 | localWindowsExp / remoteWindowsExp |
IPC flag 区 (前 2MB) |
| 当前状态 | 所有算子默认 | 预留/实验性 |
2. Window 内存布局
Window 内存由 HCCL 框架通过 HCCL_BUFFSIZE 环境变量分配（默认 200MB）。CAM 将其划分为多个功能区域：
2.1 整体布局 (单卡视角)
Window Base (localWindowsIn / remoteWindowsIn 指向)
│
├─ [0, 204MB) NOTIFY_DISPATCH_BUFF_OFFSET
│ └── notify_dispatch 算子使用的缓冲区
│
├─ [204MB, 204MB + dataArea) COMBINE_STATE_WIN_OFFSET = 4MB 之后
│ └── combine 数据区 (token 数据)
│ 每个 token 占 h512AlignRecvXLen = Ceil(h*2, 512) * 512
│
├─ ... (其他区域)
│
└─ 总计由 totalWinSize 控制
Window Exp Base (localWindowsExp / remoteWindowsExp 指向)
│
├─ [0, 500KB * dataState) WIN_STATE_OFFSET = 500KB
│ └── 预留区域
│
├─ [500KB, 950KB) STATE_WIN_OFFSET = 950KB
│ └── 本卡 dataState 状态区 (每 core 512B 对齐)
│
├─ [950KB, 950KB + N*stateOffset) dispatch flag 区
│ └── 每个 expert 占 stateOffset=32B
│ └── 写入目标: 远端 windowsExp 对应位置
│
├─ [975KB, 975KB + N*512B) MAGIC_WIN_OFFSET = 975KB
│ └── combine 用 Magic 双缓冲标志 (每 core 512B)
│
└─ 总计由 winExpSize 控制
2.2 关键偏移量定义
// comm_args.h (Moe namespace)
constexpr uint64_t NOTIFY_DISPATCH_BUFF_OFFSET = 204UL * 1024UL * 1024UL; // 204MB
constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; // 100MB
constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // 2MB flag + 100MB data
// moe_dispatch_normal.h
constexpr uint64_t WIN_STATE_OFFSET = 500UL * 1024UL; // 500KB
constexpr uint64_t STATE_WIN_OFFSET = 950UL * 1024UL; // 950KB
constexpr uint64_t WIN_ADDR_ALIGN = 512UL; // 512B 对齐
constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 4UL * 1024UL * 1024UL; // 4MB
// moe_combine_normal.h
constexpr uint64_t MAGIC_WIN_OFFSET = 975UL * 1024UL; // 975KB
constexpr uint32_t STATE_OFFSET = 32U; // 32B per expert
2.3 数据区 Token 格式 (Dispatch → Combine)
每个 token 在 Window 数据区中的布局：
hOutGMAlignSize = Ceil(h*2 + 44, 512) * 512
┌─────────────────────────────────────────────────────────┐
│ token data (h * sizeof(ExpandXOutType)) │ h*2 bytes (BF16)
│ e.g., h=7168 → 14336 bytes │
├─────────────────────────────────────────────────────────┤
│ padding to 32B alignment │
├─────────────────────────────────────────────────────────┤
│ dynamicScale (float32) - 仅 DynamicQuant 模式 │ 4 bytes
├─────────────────────────────────────────────────────────┤
│ expandIdx[0] = epRankId (int32) │ 4 bytes
│ expandIdx[1] = tokenIndex (int32) │ 4 bytes
│ expandIdx[2] = topkIdx (int32) │ 4 bytes
├─────────────────────────────────────────────────────────┤
│ padding to 512B alignment │
└─────────────────────────────────────────────────────────┘
总计 hOutGMAlignSize (通常 14848 → 15360 for h=7168)
3. moe_dispatch_normal 完整执行流程
3.1 文件结构
| 文件 | 职责 |
|---|---|
op_api/aclnn_moe_dispatch_normal.cpp |
aclnn 接口层，调用 aclnnInner |
op_host/moe_dispatch_normal.cpp |
算子定义 (OpDef)，注册输入/输出/属性 |
op_host/moe_dispatch_normal_tiling.cpp |
Tiling 计算，校验参数，计算 Window 大小 |
op_kernel/moe_dispatch_normal.cpp |
内核入口，根据 tilingKey 选择模板 |
op_kernel/moe_dispatch_normal.h |
核心：完整 kernel 实现 |
3.2 模板参数
template <typename XType, typename ExpandXOutType, bool DynamicQuant,
bool IsSmoothScaleExist, bool IsShareExpertRank>
class MoeDispatchNormal;
XType: 输入数据类型 (BF16/FP16)ExpandXOutType: 输出数据类型 (BF16/FP16/INT8)DynamicQuant: 是否动态量化 (INT8 输出)IsSmoothScaleExist: 平滑 scale (未使用)IsShareExpertRank: 共享 expert rank (未使用)
3.3 TilingKey 路由
// moe_dispatch_normal.cpp (kernel entry)
#define TILINGKEY_NO_QUANT 10000
#define TILINGKEY_QUANT 10002
#if (ORIG_DTYPE_RECV_X == DT_BF16 || ORIG_DTYPE_RECV_X == DT_FLOAT16)
if (TILING_KEY_IS(TILINGKEY_NO_QUANT)) {
MoeDispatchNormal<DTYPE_X, DTYPE_RECV_X, false, false, false> op;
...
}
#elif (ORIG_DTYPE_RECV_X == DT_INT8)
if (TILING_KEY_IS(TILINGKEY_QUANT)) {
MoeDispatchNormal<DTYPE_X, DTYPE_RECV_X, true, false, false> op;
...
}
#endif
3.4 Init 阶段详解
__aicore__ inline void MoeDispatchNormal<TemplateTypeFunc>::Init(...)
{
// 1. 获取 HCCL 上下文 (两个通信域: EP + TP)
winContext_[COMM_EP_IDX] = GetHcclContext<HCCL_GROUP_ID_0>(); // EP 域
winContext_[COMM_TP_IDX] = GetHcclContext<1>(); // TP 域
// 2. 读取本卡 dataState (Ping-Pong 状态)
GM_ADDR statusDataSpaceGm = winContext_[COMM_EP_IDX]->localWindowsExp;
selfDataStatusTensor.SetGlobalBuffer(
statusDataSpaceGm + STATE_WIN_OFFSET + blockIdx * WIN_ADDR_ALIGN);
// 3. 获取当前 dataState 并翻转 (0→1, 1→0)
DataCacheCleanAndInvalid(selfDataStatusTensor); // 缓存失效，从内存读取
dataState = selfDataStatusTensor(0);
selfDataStatusTensor(0) = (dataState == 0) ? 1 : 0;
DataCacheCleanAndInvalid(selfDataStatusTensor); // 写回
PipeBarrier<PIPE_ALL>();
// 4. 计算 winDataSizeOffset (Ping-Pong 偏移)
winDataSizeOffset = dataState * (totalWinSize / 2)
+ globalBatchSize / epRankSize * topK * hSizeAlignCombine;
// 5. 计算 shareGM 地址 (本卡 Window 数据区基址)
shareGM = GetWindAddrByRankId(COMM_EP_IDX, epRankId);
// 6. 初始化 Buffer
hOutUBAlignSize = Ceil(hScaleIdxSize, UB_ALIGN) * UB_ALIGN;
tpipe_->InitBuffer(xQueue, BUFFER_NUM, hOutUBAlignSize); // 2 * 14K = 28K
}
3.5 Process 主流程
__aicore__ inline void Process()
{
if ASCEND_IS_AIV {
InputToShare(); // 阶段1: 数据从 GM → Window
SetStatus(); // 阶段2: 通知远端 "数据已到"
WaitStatus(); // 阶段3: 等待远端 "数据已到"
ShareToOutput(); // 阶段4: 数据从 Window → GM 输出
}
}
3.6 InputToShare — 数据写入 Window
对每个 token (batchSize * topK 个):
1. 读取 expertId 和 tokenIdx
2. 计算目标地址: shareGM + hOutGMAlignSize * (dstExpertOffset + curExpertCnt)
3. 从 GM 读取 token 到 UB (DataCopyPad)
4. [可选] DynamicQuant: BF16 → FP32 → Abs → ReduceMax → Scale → INT8
5. FillTriple: 在 token 后附加 (epRankId, tokenIndex, topkIdx)
6. DataCopyPad: UB → Window (GM)
关键代码路径 (非量化):
for (int32_t tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
uint32_t dstExpertId = expertIdsTensor(tokenIndex - startTokenId);
int32_t curExpertCnt = sendTokenIdxTensor(tokenIndex - startTokenId);
int32_t dstExpertOffset = sendOffsetTensor(dstExpertId);
// 计算 Window 目标地址
GM_ADDR rankGM = shareGM + hOutGMAlignSize * (dstExpertOffset + curExpertCnt);
dstGT.SetGlobalBuffer((__gm__ ExpandXOutType *)rankGM);
// GM → UB → Window
xTmpTensor = xQueue.AllocTensor<ExpandXOutType>();
DataCopyPad(xTmpTensor, xGT[tokenIndex / topK * h], xCopyParams, tokenCopyPadExtParams);
xQueue.EnQue(xTmpTensor);
xTmpTensor = xQueue.DeQue<ExpandXOutType>();
// 附加三元组 (srcRank, srcTokenIdx, topkIdx)
FillTriple(xTmpTensor, tokenIndex / topK, tokenIndex % topK);
// UB → Window (DMA 写入远端或本卡)
DataCopyPad(dstGT, xTmpTensor, hCommuCopyOutParams);
xQueue.FreeTensor(xTmpTensor);
}
3.7 SetStatus — 通知远端数据已到
__aicore__ inline void SetStatus()
{
// 1. 准备 flag 数据 (8 个 float = 1.0f, 即 0x3F800000)
Duplicate<int32_t>(statusTensor, 0, expNumPerCore * 8);
Duplicate<int32_t>(statusTensor, 0x3F800000, mask, statusCntAlign / 8, 1, 8);
// 每个 expert 占 8 个 int32 (32B), 全部填充为 1.0f 的 IEEE754 表示
SyncAll<true>(); // 全卡同步
// 2. 将 flag 写入远端 windowsExp 的对应位置
for (uint32_t i = startExpId; i < endExpId; ++i) {
uint32_t targetRankId = i / moeExpertNumPerRank;
uint32_t offset = stateOffset * (epRankId + i % moeExpertNumPerRank * epRankSize);
GM_ADDR rankGM = GetWindStateAddrByRankId(COMM_EP_IDX, targetRankId) + offset;
dstStatusGT.SetGlobalBuffer(rankGM);
DataCopy(dstStatusGT, statusTensor[(i - startExpId) * 8], 8UL);
}
}
Flag 区布局 (远端 windowsExp):
远端 windowsExp + dataState * WIN_STATE_OFFSET
│
├─ expert 0: [rank0_flag(32B), rank1_flag(32B), ..., rankN_flag(32B)]
├─ expert 1: [rank0_flag(32B), rank1_flag(32B), ..., rankN_flag(32B)]
├─ ...
└─ expert M-1: [...]
每个 flag 占 32B (8 个 float32, 全部为 1.0f)
3.8 WaitStatus — 等待远端数据到达
__aicore__ inline void WaitStatus()
{
// 1. 准备: 从本卡 windowsExp 读取 flag 区
GlobalTensor<float> windowInstatusFp32Tensor;
windowInstatusFp32Tensor.SetGlobalBuffer(
GetWindStateAddrByRankId(COMM_EP_IDX, epRankId));
float compareTarget = 1.0f * statusNumPerCore; // 期望的 flag 总和
// 2. 轮询直到所有 flag 就绪
while (sumOfFlag != compareTarget) {
// 从 GM 读取 flag 到 UB
DataCopy(statusFp32Tensor,
windowInstatusFp32Tensor[startStatusId * stateOffset / sizeof(float)],
intriParams);
SyncFunc<MTE2_V>();
// ReduceSum: 对所有 flag 求和
ReduceSum(statusSumOutTensor, statusFp32Tensor, gatherMaskOutTensor,
mask, statusNumPerCore, 1);
SyncFunc<V_S>();
sumOfFlag = statusSumOutTensor.GetValue(0);
}
// 3. 清除 flag (写回 0)
Duplicate<int32_t>(cleanStateTensor, 0, duplicateMask, Ceil(statusNumPerCore, 8), 1, 8);
DataCopy(windowInstatusFp32Tensor[...], cleanStateTensor.ReinterpretCast<float>(), intriOutParams);
SyncAll<true>();
}
轮询优化: 使用 ReduceSum 一次性对所有 flag 求和，与 compareTarget 比较。只有当所有 source rank 都写入了 flag 时才退出。
3.9 ShareToOutput — 从 Window 读取到输出
__aicore__ inline void ShareToOutput()
{
for (uint32_t i = startStatusId; i < endStatusId; ++i) {
fromRank = i % epRankSize;
count = recvCountTensor(i) - preCount;
recvOffset = recvOffsetTensor(i);
// 计算源 Window 地址
GM_ADDR recvStart = GetWindAddrByRankId(COMM_EP_IDX, fromRank)
+ recvOffset * hOutGMAlignSize;
for (uint32_t j = 0; j < count; ++j) {
srcTokenGT.SetGlobalBuffer(recvStart + j * hOutGMAlignSize);
// Window → UB
xTmpTensor = xQueue.AllocTensor<ExpandXOutType>();
DataCopyPad(xTmpTensor, srcTokenGT, hCommuCopyOutParams, copyPadExtParams);
xQueue.EnQue(xTmpTensor);
xTmpTensor = xQueue.DeQue<ExpandXOutType>();
// 提取三元组 → expandIdxOut
DataCopyPad(expandIdxOutGT[(targetOffset + j) * EXPAND_IDX_INFO],
xTmpTensorInt[expandIdxStartIdx], dataCopyExandIdxParams);
// 提取 token data → expandXOut
dstTokenGT.SetGlobalBuffer(expandXOutGM + (targetOffset + j) * h);
DataCopyPad(dstTokenGT, xTmpTensor, expandXCopyParams);
xQueue.FreeTensor(xTmpTensor);
}
}
}
4. Window 地址计算
4.1 GetWindAddrByRankId — 数据区地址
__aicore__ inline GM_ADDR GetWindAddrByRankId(uint8_t ctxIdx, const int32_t rankId)
{
uint32_t curRankId = (ctxIdx == COMM_EP_IDX) ? epRankId : tpRankId;
if (curRankId == rankId) {
// 本卡: 使用 localWindowsIn
return (GM_ADDR)(winContext_[ctxIdx]->localWindowsIn)
+ winDataSizeOffset + COMBINE_STATE_WIN_OFFSET
+ NOTIFY_DISPATCH_BUFF_OFFSET;
}
// 远端: 使用 remoteRes 中的 windowsIn
return (GM_ADDR)(((HcclRankRelationResV2 *)
(winContext_[ctxIdx]->remoteRes[rankId].nextDevicePtr))->windowsIn)
+ winDataSizeOffset + COMBINE_STATE_WIN_OFFSET
+ NOTIFY_DISPATCH_BUFF_OFFSET;
}
地址公式:
本卡: localWindowsIn + winDataSizeOffset + 4MB + 204MB
远端: remoteRes[rankId].windowsIn + winDataSizeOffset + 4MB + 204MB
4.2 GetWindStateAddrByRankId — 状态区地址
__aicore__ inline GM_ADDR GetWindStateAddrByRankId(uint8_t ctxIdx, const int32_t rankId)
{
uint32_t curRankId = ctxIdx == COMM_EP_IDX ? epRankId : tpRankId;
if (curRankId == rankId) {
// 本卡: 使用 localWindowsExp
return (GM_ADDR)(winContext_[ctxIdx]->localWindowsExp)
+ dataState * WIN_STATE_OFFSET;
}
// 远端: 使用 remoteRes 中的 windowsExp
return (GM_ADDR)(((HcclRankRelationResV2 *)
(winContext_[ctxIdx]->remoteRes[rankId].nextDevicePtr))->windowsExp)
+ dataState * WIN_STATE_OFFSET;
}
地址公式:
本卡: localWindowsExp + dataState * 500KB
远端: remoteRes[rankId].windowsExp + dataState * 500KB
4.3 winDataSizeOffset — Ping-Pong 偏移计算
uint64_t hSizeAlignCombine = Ceil(h * sizeof(XType), WIN_ADDR_ALIGN) * WIN_ADDR_ALIGN;
winDataSizeOffset = dataState * (totalWinSize / 2)
+ globalBatchSize / epRankSize * topK * hSizeAlignCombine;
含义:
- dataState * (totalWinSize / 2): Ping-Pong 双缓冲，dataState=0 用前半，dataState=1 用后半
- globalBatchSize / epRankSize * topK * hSizeAlignCombine: 跳过 combine 数据区 (combine 从前往后写，dispatch 从后往前写)
4.4 地址计算全景图
┌──────────────────────────────────────────────┐
│ totalWinSize │
│ (由 HCCL_BUFFSIZE 控制) │
├──────────────────────────────────────────────┤
│ [0, totalWinSize/2) ← dataState=0 │
│ [totalWinSize/2, end) ← dataState=1 │
└──────────────────────────────────────────────┘
dataState=0 时:
winDataSizeOffset = 0 + bs/epRankSize * topK * hAlign
Dispatch 数据区起始:
localWindowsIn + 0 + 4MB + 204MB + bs/epRankSize*topK*hAlign
Combine 数据区起始:
localWindowsIn + 0 + 4MB + 204MB
5. Ping-Pong 双缓冲机制
5.1 设计思想
Ping-Pong 双缓冲允许 dispatch 写下一批数据 与 combine 读上一批数据 并行执行，实现通信-计算重叠。
5.2 dataState 翻转逻辑
// Init 阶段
DataCacheCleanAndInvalid(selfDataStatusTensor); // 缓存失效
dataState = selfDataStatusTensor(0); // 读取当前状态
selfDataStatusTensor(0) = (dataState == 0) ? 1 : 0; // 翻转
DataCacheCleanAndInvalid(selfDataStatusTensor); // 写回
注意: DataCacheCleanAndInvalid 在这里用于确保跨 AIV core 的可见性。所有 core 都读取同一个 selfDataStatusTensor，但只有 blockIdx=0 的 core 实际参与翻转（因为偏移是 blockIdx * WIN_ADDR_ALIGN）。
5.3 数据区划分
Window 数据区 (localWindowsIn + 204MB + 4MB 起始)
│
├─ [0, combineAreaSize) Combine 数据区 (combine 写，combine 读)
│ 每个 token: h512AlignRecvXLen
│
├─ [combineAreaSize, end) Dispatch 数据区 (dispatch 写，dispatch 读)
│ 每个 token: hOutGMAlignSize
│
└─ 总计: totalWinSize / 2 (Ping-Pong 各占一半)
5.4 时序图
时间 ──────────────────────────────────────────────────────────→
Dispatch N-1:
InputToShare → SetStatus
↓ (flag 写入远端 windowsExp)
Combine N-1: ↓
WaitBuffCopy (轮询 flag) ←──────────┘
ReadBufferAndWeightedSum
Dispatch N:
InputToShare (写 dataState=1 区域) ← 与 Combine N-1 并行!
SetStatus
↓
Combine N:
WaitBuffCopy (轮询 flag) ←──────────┘
ReadBufferAndWeightedSum (读 dataState=0 区域)
6. Magic+Value 同步协议
6.1 编码方式
// sync_collectives.h
constexpr int64_t MAGIC_OFFSET = 32;
constexpr int64_t MAGIC_MASK = ~((1LL << MAGIC_OFFSET) - 1);
__aicore__ inline int64_t MergeMagicWithValue(int32_t magic, int32_t value)
{
return (static_cast<int64_t>(static_cast<uint32_t>(magic)) << MAGIC_OFFSET)
| static_cast<int64_t>(value);
}
64位编码:
┌─────────────────────────────┬─────────────────────────────┐
│ High 32 bits │ Low 32 bits │
│ magic (uint32) │ value (uint32) │
│ (防 ABA，批次标识) │ (事件/状态码) │
└─────────────────────────────┴─────────────────────────────┘
6.2 Magic 翻转机制 (Combine 阶段)
// moe_combine_normal.h — InitMagic()
__aicore__ inline void InitMagic()
{
// 读取本卡 Magic 标志
selfMagicTensor.SetGlobalBuffer(
epWinContext_->localWindowsExp + MAGIC_WIN_OFFSET + coreIdx_ * WIN_512_ALIGN);
DataCacheCleanAndInvalid(selfMagicTensor);
magic_ = selfMagicTensor(0); // 读取当前 magic
selfMagicTensor(0) = (magic_ == 0) ? 1 : 0; // 翻转
DataCacheCleanAndInvalid(selfMagicTensor);
}
Magic 标志布局 (localWindowsExp + 975KB):
localWindowsExp + MAGIC_WIN_OFFSET
│
├─ core 0: [512B] magic flag (int32, 0 或 1)
├─ core 1: [512B] magic flag
├─ ...
└─ core N: [512B] magic flag
6.3 Flag 等待逻辑
// sync_collectives.h — WaitOneRankPartFlag
__aicore__ inline void WaitOneRankPartFlag(__gm__ int64_t* waitAddr,
int64_t flagNum, int64_t checkValue)
{
do {
// GM → UB: 读取 flag
DataCopy(localWait, globalWait[checkedFlagNum * FLAG_UNIT_INT_NUM],
(flagNum - checkedFlagNum) * FLAG_UNIT_INT_NUM);
SetFlag<MTE2_S>(EVENT_ID0);
WaitFlag<MTE2_S>(EVENT_ID0);
// 检查每个 flag
isSync = true;
for (auto i = 0; i < remainToCheck; ++i) {
int64_t v = localWait.GetValue(i * FLAG_UNIT_INT_NUM);
// 检查: magic 位相等 且 value >= checkValue
if ((v & MAGIC_MASK) != (checkValue & MAGIC_MASK) || v < checkValue) {
isSync = false;
checkedFlagNum += i; // 跳过已确认的 flag
break;
}
}
} while (!isSync);
}
关键判断条件:
(v & MAGIC_MASK) != (checkValue & MAGIC_MASK) || v < checkValue
MAGIC_MASK = 0xFFFFFFFF00000000(高32位)- 先检查 magic 位是否匹配（防止读到旧一批的 flag）
- 再检查 value 是否 >= 目标值（支持乱序到达）
6.4 Dispatch 阶段的 Flag 编码
Dispatch 阶段不使用 SyncCollectives 类，而是直接用 float32 = 1.0f 作为 flag:
// SetStatus: 写入 1.0f (IEEE754 = 0x3F800000)
Duplicate<int32_t>(statusTensor, 0x3F800000, mask, statusCntAlign / 8, 1, 8);
// WaitStatus: 轮询求和
float compareTarget = 1.0f * statusNumPerCore;
while (sumOfFlag != compareTarget) {
DataCopy(statusFp32Tensor, windowInstatusFp32Tensor[...]);
ReduceSum(statusSumOutTensor, statusFp32Tensor, ...);
sumOfFlag = statusSumOutTensor.GetValue(0);
}
注意: Dispatch 阶段的 flag 不使用 Magic+Value 编码，而是直接用 float 求和。这是因为 dispatch 的 flag 区在 localWindowsExp + dataState * 500KB，已经通过 dataState 实现了 Ping-Pong 隔离。
6.5 Combine 阶段的 Flag 编码
Combine 阶段使用 float32 = 1.0f 作为 flag (与 dispatch 类似):
// CopyBufferToShareAndSetStatus
Duplicate<uint32_t>(statusTensor, 0x3F800000, FLOAT_NUM_PER_ALIGN);
// 写入远端 windowsIn 对应位置
// WaitBuffCopy
float target = 1.0f * axisK_ * FLOAT_NUM_PER_ALIGN;
while (current != target) {
DataCopy(stateTensorLocal, stateGMTensor, calCount);
Sum(tempStateTensorLocal, stateTensorLocal, sumPerKParams);
current = tempStateTensorLocal(0);
}
7. moe_combine_normal 完整执行流程
7.1 文件结构
| 文件 | 职责 |
|---|---|
op_api/aclnn_moe_combine_normal.cpp |
aclnn 接口层 |
op_host/moe_combine_normal.cpp |
算子定义 |
op_host/moe_combine_normal_tiling.cpp |
Tiling 计算 |
op_kernel/moe_combine_normal.cpp |
内核入口 |
op_kernel/moe_combine_normal.h |
核心：完整 kernel 实现 |
7.2 模板参数
template <typename RecvXType, typename XType, typename SrcInfoType>
class MoeCombineNormal;
RecvXType: 接收数据类型 (BF16/FP16)XType: 输出数据类型 (BF16/FP16)SrcInfoType: 源信息类型 (int32)
7.3 Init 阶段
__aicore__ inline void Init(...)
{
coreIdx_ = GetBlockIdx();
// 1. 初始化 Magic (Ping-Pong)
InitMagic();
// 2. 初始化 Tiling 数据
InitTilingData(tilingData);
// 3. 初始化 Global Buffer
InitGlobalBuffer(recvX, tokenSrcInfo, epRecvCount, topkWeights, XOut, sendCostStatsOut);
// 4. 计算 Buffer 长度
InitBuffLen();
// 5. 计算 winDataSizeOffset
PipeBarrier<PIPE_ALL>();
winDataSizeOffset_ = magic_ * (totalWinSize / 2UL);
localRankGM_ = GetBufferAddrByRankId(epRankId_);
// 6. 读取 selfSendCnt_ (本卡需要发送的 token 数)
DataCacheCleanAndInvalid(epRecvCountGM_[moeExpertNum_ - 1]);
selfSendCnt_ = epRecvCountGM_(moeExpertNum_ - 1);
}
7.4 Process 主流程
__aicore__ inline void Process()
{
if ASCEND_IS_AIV {
CopyBufferToShareAndSetStatus(); // 阶段1: 数据 GM → Window + 通知远端
ReadBufferFromRemote(); // 阶段2: 等待数据 + 加权求和
}
}
7.5 CopyBufferToShareAndSetStatus — 数据写入远端 Window
__aicore__ inline void CopyBufferToShareAndSetStatus()
{
// 1. 按 core 分配 token
SplitCoreCal(selfSendCnt_, perBlockSendNum, startTokenId, endTokenId);
// 2. 准备 flag 数据
Duplicate<uint32_t>(statusTensor, 0x3F800000, FLOAT_NUM_PER_ALIGN);
// 3. 读取 srcInfo (三元组: srcRankId, srcTokenId, srcTopkId)
DataCopyPad(srcInfoLocal, tokenSrcInfoGM_[startTokenId * TOKEN_SRC_INFO_LEN], ...);
// 4. 对每个 token
for (uint32_t tokenIndex = startTokenId; tokenIndex < endTokenId; tokenIndex++) {
uint32_t srcRankId = srcInfoLocal(index + RANK_ID_OFFSET_IN_SRC_INFO);
uint32_t srcTokenId = srcInfoLocal(index + TOKEN_IDX_OFFSET_IN_SRC_INFO);
uint32_t srcTopkId = srcInfoLocal(index + TOPK_IDX_OFFSET_IN_SRC_INFO);
// 4a. 复制数据到远端 Window
CopyBufferToShare(srcRankId, srcTokenId, srcTopkId, tokenIndex);
// 4b. 设置 flag
SetStatusBySrcInfo(srcRankId, srcTokenId, srcTopkId);
}
}
CopyBufferToShare 详解:
__aicore__ inline void CopyBufferToShare(uint32_t srcRankId, uint32_t srcTokenId,
uint32_t srcTopkId, uint32_t tkIndex)
{
uint32_t tokenOffset = tkIndex * axisH_;
// 计算远端 Window 目标地址
GM_ADDR dstGM = GetBufferAddrByRankId(srcRankId)
+ (srcTokenId * axisK_ + srcTopkId) * h512AlignRecvXLen_;
GlobalTensor<XType> dstWindow;
dstWindow.SetGlobalBuffer(dstGM);
// GM → UB → Window
LocalTensor<RecvXType> localCopyTensor;
localCopyTensor = localCopyQueue_.AllocTensor<RecvXType>();
DataCopyPad(localCopyTensor, recvXGM_[tokenOffset], xOutCopyParams, copyPadExtParams);
localCopyQueue_.EnQue(localCopyTensor);
localCopyTensor = localCopyQueue_.DeQue<RecvXType>();
DataCopyPad(dstWindow, localCopyTensor, xOutCopyParams);
localCopyQueue_.FreeTensor(localCopyTensor);
}
SetStatusBySrcInfo 详解:
__aicore__ inline void SetStatusBySrcInfo(uint32_t srcRankId, uint32_t srcTokenId,
uint32_t srcTopkId)
{
LocalTensor<uint32_t> statusTensor = stateBuf_.AllocTensor<uint32_t>();
// 计算远端 flag 地址
GM_ADDR stateGM = GetStateAddrByRankId(srcRankId)
+ (srcTokenId * axisK_ + srcTopkId) * UB_32_ALIGN;
GlobalTensor<uint32_t> stateGMTensor;
stateGMTensor.SetGlobalBuffer(stateGM);
// UB → GM (写入 1.0f)
DataCopy<uint32_t>(stateGMTensor, statusTensor, FLOAT_NUM_PER_ALIGN);
}
7.6 ReadBufferFromRemote — 等待数据 + 加权求和
__aicore__ inline void ReadBufferFromRemote()
{
// 1. 按 core 分配 token
SplitCoreCal(axisBS_, tokenPerBlock, startTokenIndex, endTokenIndex);
// 2. 初始化 Buffer
tpipe_->InitBuffer(xOutBuf_, h32AlignRecvXLen_);
tpipe_->InitBuffer(tokenFloatBuf_, h32AlignFloatLen_);
tpipe_->InitBuffer(weightedMulBuf_, h256AlignFloatLen_);
tpipe_->InitBuffer(sumFloatBuf_, h32AlignFloatLen_);
tpipe_->InitBuffer(weightedSumQueue_, DOUBLE_BUFFER, h32AlignRecvXLen_);
tpipe_->InitBuffer(stateBuf_, axisK_ * UB_32_ALIGN);
tpipe_->InitBuffer(tempStateBuf_, axisK_ * UB_32_ALIGN);
tpipe_->InitBuffer(topkWeightsBuf_, tokenPerBlock * axisK_ * sizeof(float));
// 3. 读取 topK weights
DataCopyPad(topkWeightsLocal, topkWeightsGM_[startTokenIndex * axisK_], ...);
// 4. 对每个 token
for (uint32_t tokenIndex = startTokenIndex; tokenIndex < endTokenIndex; tokenIndex++) {
WaitBuffCopy(tokenIndex); // 等待所有 topK 数据到达
ReadBufferAndWeightedSum(tokenIndex); // 加权求和
}
}
7.7 WaitBuffCopy — 等待 topK 数据全部到达
__aicore__ inline void WaitBuffCopy(uint32_t tokenIndex)
{
uint32_t calCount = axisK_ * FLOAT_NUM_PER_ALIGN;
// 本卡 Window 中该 token 对应的 flag 区
GM_ADDR stateGM = GetStateAddrByRankId(epRankId_) + tokenIndex * axisK_ * UB_32_ALIGN;
GlobalTensor<float> stateGMTensor;
stateGMTensor.SetGlobalBuffer(stateGM);
float target = 1.0f * axisK_ * FLOAT_NUM_PER_ALIGN;
// 轮询求和
while (current != target) {
SyncFunc<S_MTE2>();
DataCopy<float>(stateTensorLocal, stateGMTensor, calCount);
SyncFunc<MTE2_V>();
Sum(tempStateTensorLocal, stateTensorLocal, sumPerKParams);
SyncFunc<V_S>();
current = tempStateTensorLocal(0);
}
// 清除 flag
SyncFunc<S_V>();
Duplicate<float>(tempStateTensorLocal, 0.0f, calCount);
SyncFunc<V_MTE3>();
DataCopy<float>(stateGMTensor, tempStateTensorLocal, calCount);
}
7.8 ReadBufferAndWeightedSum — 加权求和
__aicore__ inline void ReadBufferAndWeightedSum(uint32_t tokenIndex, uint32_t startTokenIndex)
{
LocalTensor<float> tokenFloatLocal = tokenFloatBuf_.Get<float>();
LocalTensor<float> weightedMulBufLocal = weightedMulBuf_.Get<float>();
LocalTensor<float> sumFloatBufLocal = sumFloatBuf_.Get<float>();
LocalTensor<float> topkWeightsLocal = topkWeightsBuf_.Get<float>();
// 初始化累加器 (FP32)
Duplicate(sumFloatBufLocal, 0.0f, axisH_);
for (uint32_t topkId = 0; topkId < axisK_; topkId++) {
float scale = topkWeightsLocal.GetValue((tokenIndex - startTokenIndex) * axisK_ + topkId);
// 从本卡 Window 读取 token
GM_ADDR localTokenAddr = localRankGM_ + (tokenIndex * axisK_ + topkId) * h512AlignRecvXLen_;
GlobalTensor<XType> localTokenTensor;
localTokenTensor.SetGlobalBuffer(localTokenAddr);
LocalTensor<XType> tmpToken = weightedSumQueue_.AllocTensor<XType>();
DataCopyPad(tmpToken, localTokenTensor, xOutCopyParams, copyPadExtParams);
weightedSumQueue_.EnQue(tmpToken);
tmpToken = weightedSumQueue_.DeQue<XType>();
// BF16 → FP32
Cast(tokenFloatLocal, tmpToken, CAST_NONE, axisH_);
PipeBarrier<PIPE_V>();
// 乘 scale
Muls(weightedMulBufLocal, tokenFloatLocal, scale, axisH_);
PipeBarrier<PIPE_V>();
// 累加
Add(sumFloatBufLocal, sumFloatBufLocal, weightedMulBufLocal, axisH_);
weightedSumQueue_.FreeTensor(tmpToken);
}
PipeBarrier<PIPE_V>();
// FP32 → BF16 输出
LocalTensor<XType> xOutLocal = xOutBuf_.Get<XType>();
Cast(xOutLocal, sumFloatBufLocal, CAST_RINT, axisH_);
SyncFunc<V_MTE3>();
DataCopyPad(xOutGlobal_[tokenIndex * axisH_], xOutLocal, xOutCopyParams);
}
关键设计: 使用 FP32 累加 避免 BF16 累加误差。每个 token 先 Cast 到 FP32，乘 weight，累加，最后 Cast 回 BF16。
7.9 Combine 地址计算
// 状态区地址 (flag)
__aicore__ GM_ADDR GetStateAddrByRankId(const int32_t rankId)
{
GM_ADDR bufferAddr;
if (epRankId_ == rankId) {
bufferAddr = epWinContext_->localWindowsIn;
} else {
bufferAddr = ((HcclRankRelationResV2 *)
epWinContext_->remoteRes[rankId].nextDevicePtr)->windowsIn;
}
return bufferAddr + winDataSizeOffset_ + NOTIFY_DISPATCH_BUFF_OFFSET;
}
// 数据区地址 (token)
__aicore__ GM_ADDR GetBufferAddrByRankId(const int32_t rankId)
{
return GetStateAddrByRankId(rankId) + COMBINE_STATE_WIN_OFFSET;
}
地址公式:
状态区: windowsIn + winDataSizeOffset + 204MB
数据区: windowsIn + winDataSizeOffset + 204MB + 4MB
8. Tiling 策略
8.1 Tiling 数据结构
// moe_dispatch_normal_tiling.h
struct MoeDispatchNormalInfo {
uint32_t epWorldSize; // EP 通信域大小
uint32_t tpWorldSize; // TP 通信域大小
uint32_t epRankId; // 本卡 EP Rank ID
uint32_t tpRankId; // 本卡 TP Rank ID
uint32_t moeExpertNum; // MoE expert 总数
uint32_t quantMode; // 量化模式 (0=无, 2=动态)
uint32_t globalBs; // 全局 batch size
uint32_t bs; // 本地 batch size
uint32_t k; // topK
uint32_t h; // hidden size
uint32_t aivNum; // AIV core 数量
bool isQuant; // 是否量化
bool isEnableDiagnose; // 是否启用诊断
uint64_t totalUbSize; // UB 总大小
uint64_t totalWinSize; // Window 总大小
};
8.2 Window 大小校验
// moe_dispatch_normal_tiling.cpp
uint64_t maxWindowSize = Mc2TilingUtils::GetMaxWindowSize(); // 从 HCCL_BUFFSIZE 读取
// 计算每个 token 需要的空间
uint64_t tokenActualLen = Ceil(h * 2, 32) * 32 + 44; // h*2 + scale(32B) + triple(12B)
uint64_t tokenNeedSizeDispatch = Ceil(tokenActualLen, 512) * 512;
uint64_t tokenNeedSizeCombine = Ceil(h * 2, 512) * 512;
// 总需求 (双缓冲)
uint64_t actualSize = (maxBs * k * (tokenNeedSizeCombine + tokenNeedSizeDispatch)
+ COMBINE_STATE_WIN_OFFSET + NOTIFY_DISPATCH_WIN_OFFSET)
* DOUBLE_DATA_BUFFER;
// 校验
if (actualSize > maxWindowSize) {
// 报错: HCCL_BUFFSIZE too SMALL
}
8.3 分片策略
Dispatch 阶段:
// 按 core 分配 token
sendTokenNum = expertIdsCnt / blockNum; // expertIdsCnt = batchSize * topK
remainTokenNum = expertIdsCnt % blockNum;
startTokenId = sendTokenNum * blockIdx;
if (blockIdx < remainTokenNum) {
sendTokenNum += 1;
startTokenId += blockIdx;
} else {
startTokenId += remainTokenNum;
}
Combine 阶段:
// SplitCoreCal: 按 core 分配 token
perCoreNum = totalNum / aivNum_;
remainderRankNum = totalNum % aivNum_;
startIdx = perCoreNum * coreIdx_;
if (coreIdx_ < remainderRankNum) {
perCoreNum++;
startIdx += coreIdx_;
} else {
startIdx += remainderRankNum;
}
8.4 状态区分配
// 每个 core 负责的 expert 数量
statusNumPerCore = moeExpertNum / blockNum;
remainStatus = moeExpertNum % blockNum;
startStatusId = statusNumPerCore * blockIdx;
if (blockIdx < remainStatus) {
statusNumPerCore += 1;
startStatusId += blockIdx;
} else {
startStatusId += remainStatus;
}
endStatusId = startStatusId + statusNumPerCore;
8.5 Workspace 计算
// moe_dispatch_normal_tiling.cpp
workSpaces[0] = SYSTEM_NEED_WORKSPACE + WORKSPACE_ELEMENT_OFFSET * aivNum * aivNum;
// SYSTEM_NEED_WORKSPACE = 16MB
// WORKSPACE_ELEMENT_OFFSET = 512
// 例如 aivNum=48: 16MB + 512 * 48 * 48 = 16MB + 11.25MB = 27.25MB
// moe_combine_normal_tiling.cpp
workspace[0] = SYSTEM_NEED_WORKSPACE; // 16MB
8.6 HCCL 通信配置
// Dispatch: AlltoAll (EP) + AllGather (TP)
std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise";
std::string algConfigAllGatherStr = "AllGather=level0:ring";
// Combine: AlltoAll (EP) + ReduceScatter (TP)
std::string algConfigAllToAllStr = "AlltoAll=level0:fullmesh;level1:pairwise";
std::string algConfigReduceScatterStr = "ReduceScatter=level0:ring";
9. 常量速查表
9.1 Window 布局常量
| 常量 | 值 | 来源文件 | 用途 |
|---|---|---|---|
NOTIFY_DISPATCH_BUFF_OFFSET |
204MB (213909504) | comm_args.h | notify_dispatch 缓冲区偏移 |
IPC_BUFF_MAX_SIZE |
100MB | comm_args.h | IPC 缓冲区最大大小 |
IPC_DATA_OFFSET |
2MB | comm_args.h | IPC flag 区大小 |
WIN_STATE_OFFSET |
500KB | moe_dispatch_normal.h | dataState 间偏移 |
STATE_WIN_OFFSET |
950KB | moe_dispatch_normal.h | 本卡状态区偏移 |
WIN_ADDR_ALIGN |
512B | moe_dispatch_normal.h | Window 地址对齐 |
COMBINE_STATE_WIN_OFFSET |
4MB | moe_dispatch_normal.h | combine 数据区偏移 |
MAGIC_WIN_OFFSET |
975KB (998400) | moe_combine_normal.h | combine magic 标志偏移 |
STATE_OFFSET |
32B | moe_dispatch_normal.h | 每个 expert flag 占用 |
9.2 对齐常量
| 常量 | 值 | 用途 |
|---|---|---|
UB_ALIGN |
32B | UB 地址对齐 |
UB_32_ALIGN |
32B | UB 32字节对齐 |
MUL_256_ALIGN |
256B | 256B 对齐 (用于 FP32 累加) |
WIN_512_ALIGN |
512B | Window 地址对齐 |
FLOAT_NUM_PER_ALIGN |
8 | 每 32B 对齐的 float 数量 |
9.3 缓冲区常量
| 常量 | 值 | 用途 |
|---|---|---|
BUFFER_NUM |
2 | Ping-Pong 双缓冲 |
DOUBLE_BUFFER |
2 | 双缓冲 |
EXPAND_IDX_INFO |
3 | 三元组长度 (rank, token, topk) |
TOKEN_SRC_INFO_LEN |
3 | 源信息长度 |
9.4 通信常量
| 常量 | 值 | 用途 |
|---|---|---|
COMM_NUM |
2 | 通信域数量 (EP + TP) |
COMM_EP_IDX |
0 | EP 通信域索引 |
COMM_TP_IDX |
1 | TP 通信域索引 |
CAM_MAX_RANK_SIZE |
384 | 最大支持 rank 数 |
MAX_EP_WORLD_SIZE |
384 | 最大 EP world size |
MAX_TP_WORLD_SIZE |
2 | 最大 TP world size |
9.5 同步常量
| 常量 | 值 | 用途 |
|---|---|---|
MAGIC_OFFSET |
32 | magic 在高 32 位 |
MAGIC_MASK |
0xFFFFFFFF00000000 |
magic 位掩码 |
FLAG_UNIT_INT_NUM |
4 | 每个 flag 占 int64 数量 |
SYNC_UNIT_SIZE |
32B | 每个同步单元大小 |
CYCLE_TO_TIME |
50 | 系统周期→时间转换 |
9.6 限制常量
| 常量 | 值 | 用途 |
|---|---|---|
MOE_EXPERT_MAX_NUM |
512 | 最大 expert 数量 |
K_MAX |
16 | 最大 topK |
H_MIN |
1024 | 最小 hidden size |
H_MAX |
7168 | 最大 hidden size |
BS_UPPER_BOUND |
8000 | 最大 batch size |
DISPATCH_STATUS_MAX_SUPPORT_NUM |
1280 | 最大 dispatch 状态数 |
10. 与 SHMEM 模式的代码级对比
10.1 内存管理
| 维度 | HCCL Window 模式 | SHMEM 模式 |
|---|---|---|
| 内存获取 | GetHcclContext<HCCL_GROUP_ID_0>() 获取预分配 Window |
IPC 共享内存 + 自管理 flag 区 |
| 数据区 | localWindowsIn / remoteWindowsIn |
IPC data 区 (100MB) |
| Flag 区 | localWindowsExp / remoteWindowsExp |
IPC flag 区 (2MB) |
| 地址计算 | GetWindAddrByRankId + 固定偏移 |
IPC 基址 + chunk 偏移 |
10.2 同步机制
| 维度 | HCCL Window 模式 | SHMEM 模式 |
|---|---|---|
| Flag 编码 | float32 = 1.0f (dispatch/combine) | Magic+Value (int64) |
| 等待方式 | ReduceSum/Sum 求和比较 | WaitOneRankPartFlag 逐 flag 检查 |
| 清除方式 | Duplicate 0 + DataCopy | SetFlag(0) |
| 防 ABA | dataState Ping-Pong 隔离 | Magic 翻转 (高 32 位) |
10.3 双缓冲实现
| 维度 | HCCL Window 模式 | SHMEM 模式 |
|---|---|---|
| 缓冲切换 | dataState (0/1) 存储在 localWindowsExp + STATE_WIN_OFFSET |
magic 存储在 flag 区 |
| 偏移计算 | dataState * (totalWinSize / 2) |
magic * (totalSize / 2) |
| 翻转时机 | Init 阶段读取并翻转 | Init 阶段读取并翻转 |
10.4 地址计算差异
HCCL Window 模式:
// 数据区
localWindowsIn + winDataSizeOffset + COMBINE_STATE_WIN_OFFSET + NOTIFY_DISPATCH_BUFF_OFFSET
// 状态区
localWindowsExp + dataState * WIN_STATE_OFFSET
SHMEM 模式 (推测):
// 数据区
ipcDataAddr + magic * (totalSize / 2) + tokenOffset
// Flag 区
ipcFlagAddr + magic * (flagSize / 2) + flagOffset
10.5 代码路径差异
| 代码路径 | HCCL Window 模式 | SHMEM 模式 |
|---|---|---|
| 上下文获取 | GetHcclContext<HCCL_GROUP_ID_0>() |
IPC 共享内存初始化 |
| 远端访问 | remoteRes[rankId].nextDevicePtr |
IPC 映射 |
| Flag 写入 | DataCopy 到远端 windowsExp | SetFlag 到 IPC flag 区 |
| Flag 等待 | ReduceSum 求和 | WaitOneRankPartFlag 逐 flag 检查 |
| 同步原语 | SyncFunc (HardEvent) | SyncCollectives 类 |
10.6 HCCL Window 独有的代码路径
- 双通信域管理:
winContext_[COMM_EP_IDX]+winContext_[COMM_TP_IDX] - MC2 Tiling:
Mc2InitTiling+Mc2CcTiling配置 - HCCL 算法配置:
AlltoAll=level0:fullmesh;level1:pairwise - Window 大小校验:
Mc2TilingUtils::GetMaxWindowSize()从环境变量读取 - dataState 状态机:
selfDataStatusTensor存储在localWindowsExp + STATE_WIN_OFFSET - Magic 标志:
MAGIC_WIN_OFFSET用于 combine 阶段
11. 关键执行流程时序图
11.1 Dispatch → Combine 全流程
Rank 0 Rank 1
────── ──────
Dispatch:
InputToShare()
├─ 读取 token[0..N]
├─ 量化 (可选)
├─ FillTriple (srcRank=0, tokenIdx, topkIdx)
└─ DataCopyPad → Rank1.windowsIn
│
│ DMA 跨片写入
▼
SetStatus() SetStatus()
├─ 准备 flag (1.0f) ├─ 准备 flag (1.0f)
└─ DataCopy → Rank1.windowsExp └─ DataCopy → Rank0.windowsExp
│ │
▼ ▼
WaitStatus() WaitStatus()
├─ 轮询 windowsExp ├─ 轮询 windowsExp
│ (ReduceSum == target?) │ (ReduceSum == target?)
│ ← Rank1 已写入 flag │ ← Rank0 已写入 flag
├─ 清除 flag ├─ 清除 flag
└─ SyncAll<true>() └─ SyncAll<true>()
│ │
▼ ▼
ShareToOutput() ShareToOutput()
├─ 从 Rank1.windowsIn 读取 token ├─ 从 Rank0.windowsIn 读取 token
├─ 提取三元组 → expandIdxOut ├─ 提取三元组 → expandIdxOut
└─ DataCopyPad → expandXOut └─ DataCopyPad → expandXOut
│ │
▼ ▼
Combine: Combine:
CopyBufferToShareAndSetStatus() CopyBufferToShareAndSetStatus()
├─ 读取 recvX[tokenIdx] ├─ 读取 recvX[tokenIdx]
├─ DataCopyPad → Rank1.windowsIn ├─ DataCopyPad → Rank0.windowsIn
└─ SetStatusBySrcInfo → Rank1.windowsIn └─ SetStatusBySrcInfo → Rank0.windowsIn
│ │
▼ ▼
ReadBufferFromRemote() ReadBufferFromRemote()
├─ WaitBuffCopy(tokenIdx) ├─ WaitBuffCopy(tokenIdx)
│ (Sum == target?) │ (Sum == target?)
│ ← Rank1 已写入 flag │ ← Rank0 已写入 flag
├─ 清除 flag ├─ 清除 flag
├─ ReadBufferAndWeightedSum ├─ ReadBufferAndWeightedSum
│ ├─ for topk in 0..K: │ ├─ for topk in 0..K:
│ │ ├─ 从 localWindowsIn 读取 token │ │ ├─ 从 localWindowsIn 读取 token
│ │ ├─ Cast BF16→FP32 │ │ ├─ Cast BF16→FP32
│ │ ├─ Muls (scale) │ │ ├─ Muls (scale)
│ │ └─ Add (累加) │ │ └─ Add (累加)
│ ├─ Cast FP32→BF16 │ ├─ Cast FP32→BF16
│ └─ DataCopyPad → xOut │ └─ DataCopyPad → xOut
└─ done └─ done
11.2 Ping-Pong 时序
时间 ──────────────────────────────────────────────────────────────────→
┌─── Iteration N ─────────────────────┐
│ │
Dispatch: │ InputToShare → SetStatus │
│ (写 dataState=0 区域) │
│ ↓ │
Combine: │ │ WaitBuffCopy │
│ │ (轮询 flag) │
│ ↓ │
│ ReadBufferAndWeightedSum │
│ (读 dataState=0 区域) │
│ │
────────────────────────────────────────────────────────────────────
│ ┌─── Iteration N+1 ─────────────────┐
│ │ │
│ │ InputToShare → SetStatus │
│ │ (写 dataState=1 区域) │
│ │ ↓ │
│ │ │ WaitBuffCopy │
│ │ │ (轮询 flag) │
│ │ ↓ │
│ │ ReadBufferAndWeightedSum │
│ │ (读 dataState=1 区域) │
│ │ │
────────────────────────────────────────────────────────────────────
附录 A: 关键函数索引
moe_dispatch_normal.h
| 函数 | 行号 | 职责 |
|---|---|---|
Init |
173-261 | 初始化上下文、dataState、Buffer |
InputToShare |
346-410 | GM → Window 数据写入 |
SetStatus |
412-439 | 写入 flag 到远端 windowsExp |
WaitStatus |
441-543 | 轮询 flag 直到全部就绪 |
ShareToOutput |
545-591 | Window → GM 输出 |
Process |
593-601 | 主流程编排 |
GetWindAddrByRankId |
67-76 | 数据区地址计算 |
GetWindStateAddrByRankId |
78-87 | 状态区地址计算 |
FillTriple |
334-344 | 附加三元组 (rank, token, topk) |
QuantProcess |
295-332 | 动态量化 (BF16 → INT8) |
moe_combine_normal.h
| 函数 | 行号 | 职责 |
|---|---|---|
Init |
202-221 | 初始化 Magic、Tiling、Buffer |
InitMagic |
144-156 | 读取并翻转 Magic 标志 |
CopyBufferToShareAndSetStatus |
224-283 | GM → Window + flag 写入 |
CopyBufferToShare |
286-303 | 数据 DMA 到远端 Window |
SetStatusBySrcInfo |
307-314 | 写入 flag 到远端 |
WaitBuffCopy |
317-340 | 轮询 flag 直到 topK 全部到达 |
ReadBufferAndWeightedSum |
343-377 | 加权求和 (FP32 累加) |
ReadBufferFromRemote |
379-412 | 主循环: 等待 + 求和 |
Process |
414-420 | 主流程编排 |
GetStateAddrByRankId |
67-76 | 状态区地址计算 |
GetBufferAddrByRankId |
78-81 | 数据区地址计算 |
sync_collectives.h
| 函数 | 行号 | 职责 |
|---|---|---|
MergeMagicWithValue |
303-307 | Magic+Value 编码 |
SetSyncFlag |
46-67 | 设置同步 flag |
WaitSyncFlag |
87-108 | 等待同步 flag |
WaitOneRankPartFlag |
320-347 | 轮询部分 flag |
CheckOneRankPartFlag |
356-376 | 检查部分 flag (一次性) |
附录 B: 调试与诊断
B.1 Diagnose 模式
两个算子都支持 isEnableDiagnose 标志，用于收集通信耗时统计:
// Dispatch: waitRecvCostStats
if (isEnableDiagnose) {
int32_t durationTime = (GetSystemCycle() - systemCycleStart) / CYCLE_TO_TIME; // us
// 按 srcRank 累加耗时
waitRecvCostStatsTensor.SetValue(srcRank, preTime + durationTime);
}
// Combine: sendCostStats
if (isEnableDiagnose) {
int32_t durationTime = (GetSystemCycle() - sendStartCycle) / CYCLE_TO_TIME; // us
sendCostStatsTensor.SetValue(srcRank, preTime + durationTime);
}
B.2 关键日志
// Tiling 阶段
OPS_LOG_D(nodeName, "windowSize = %lu", maxWindowSize);
OPS_LOG_D(nodeName, "blockDim=%u, aivNum=%u, ubSize=%lu", blockDim, aivNum, ubSize);
OPS_LOG_E(nodeName, "HCCL_BUFFSIZE is too SMALL, ... NEEDED_HCCL_BUFFSIZE=%luMB, HCCL_BUFFSIZE=%luMB.",
actualSize / MB_SIZE + 1UL, maxWindowSize / MB_SIZE);
总结: HCCL Window 模式通过 HCCL 框架预分配的 Window 内存窗口，实现了高效的跨片通信。其核心设计包括: (1) Ping-Pong 双缓冲实现通信-计算重叠; (2) Magic+Value 编码实现无锁同步; (3) 三元组嵌入 (srcRank, srcTokenIdx, topkIdx) 避免额外通信; (4) FP32 累加避免精度损失。该模式是 CAM 库的生产默认路径，所有算子当前均使用此模式。
一、概述
本文档对 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` 提供完整的路由信息。
代码深读报告 | 聚焦 FusedDeepMoE 的 SHMEM vs 非 SHMEM 通信路径、raw_distributed 通信原语、三条执行路径差异、Workspace 内存布局，以及编译系统的 SHMEM 注入机制。
1. 总体架构概览
FusedDeepMoE 是 CAM 中最复杂的融合算子，将 Dispatch → GMM1 → SwiGLU → GMM2 → Combine 五个阶段融合进单一 kernel。其核心挑战在于：
- 通信与计算的融合：dispatch/combine 需要跨 EP rank 搬运 token，同时 GMM 需要密集 Cube 计算
- 三条执行路径：非深度融合 / 深度融合 / 共享专家，各自通信模式不同
- SHMEM vs HCCL Window：当前默认使用 HCCL Window 通信，SHMEM 路径为未来预留
┌─────────────────────────────────────────────────────────────────────┐
│ FusedDeepMoE 融合流程 │
│ │
│ x ──→ Dispatch ──→ GMM1 ──→ SwiGLU ──→ GMM2 ──→ Combine ──→ out │
│ (HCCL Window) (Cube) (Vector) (Cube) (HCCL Window) │
│ │
│ shared_x ──→ ShareGMM1 ──→ SwiGLU ──→ ShareGMM2 ──→ share_out │
│ (Cube) (Vector) (Cube) │
└─────────────────────────────────────────────────────────────────────┘
2. FusedDeepMoE 通信模式实现
2.1 核心入口：fused_deep_moe.h
文件：op_kernel/fused_deep_moe.h
通信完全通过 raw_distributed 子模块实现，主 kernel 类 FusedDeepMoe 本身不包含任何通信逻辑：
// fused_deep_moe.h 第30行
#include "fused_deep_moe/raw_distributed/cam_moe_distribute_dispatch.h"
// 注意：combine 的 include 在 GMM2 kernel 内部（通过 combiner 指针）
Process() 函数揭示的三阶段执行模型（第354-448行）：
template <TemplateMC2TypeClass>
__aicore__ inline void FusedDeepMoe<TemplateMC2TypeFunc>::Process()
{
// 阶段1: Dispatch（仅非深度融合路径）
if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
MoeDistributeDispatchImpl::CamMoeDistributeDispatch<...> dispatcher;
dispatcher.Init(...);
dispatcher.Process(); // ← HCCL Window AlltoAll
}
// 阶段2: GMM1 + SwiGLU（Cube 计算，无通信）
GmmDeqSwigluQuant<...>(...);
// 阶段3: GMM2 + Combine（融合通信计算）
MoeDistributeCombineImpl::CamMoeDistributeCombine<...> combiner;
GmmDeq<...>(..., &combiner); // ← combiner 在 GMM2 的 AIV 侧执行
}
2.2 通信上下文获取：HCCL Window 模式
关键发现：当前代码中没有任何 SHMEM 分支。 所有通信通过 HcclOpResParam 结构体（HCCL Window 描述符）实现：
// cam_moe_distribute_dispatch.h 第236-237行
winContext_[COMM_EP_IDX] = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
winContext_[COMM_TP_IDX] = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<1>();
// cam_moe_distribute_combine.h 第245-246行
auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
epWinContext_ = (__gm__ HcclOpResParam *)contextGM0;
HCCL Window 地址解析（dispatch 第83-91行）：
__aicore__ inline GM_ADDR GetWindAddrByRankId(uint8_t ctxIdx, const int32_t rankId)
{
uint32_t curRankId = ctxIdx == COMM_EP_IDX ? epRankId_ : tpRankId_;
if (curRankId == rankId) {
// 本地 Window：直接访问 localWindowsIn
return (GM_ADDR)(winContext_[ctxIdx]->localWindowsIn) + winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
}
// 远端 Window：通过 remoteRes 指针访问对端 windowsIn
return (GM_ADDR)(((HcclRankRelationResV2 *)(winContext_[ctxIdx]->remoteRes[rankId].nextDevicePtr))->windowsIn)
+ winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
}
通信模式总结：
- 数据通道：windowsIn — 存放实际 token 数据（int8/half/bf16）
- 状态通道：windowsExp — 存放同步状态标志（flag/sumTarget）
- 同步机制：通过 DataCopy 在 Window 上写入 flag，轮询对端 flag 实现握手
3. SHMEM vs 非 SHMEM 代码路径
3.1 当前状态：无 SHMEM 分支
对整个 fused_deep_moe/ 目录进行全量 grep：
grep -rn "aclshmem\|SHMEM\|shmem_" src/cam/comm_operator/ascend_kernels/fused_deep_moe/
# 结果：0 命中
结论：FusedDeepMoE kernel 代码中不存在任何 SHMEM 相关代码路径。 所有通信均通过 HCCL Window（HcclOpResParam）实现。
3.2 编译系统的 SHMEM 预留
虽然 kernel 代码没有 SHMEM 分支，但编译系统已完整预留 SHMEM 注入点：
| 层级 | 文件 | SHMEM 机制 |
|---|---|---|
| 算子选择 | select_ops.py |
requires_shmem() 过滤 |
| 编译脚本 | compile_ascend_proj.sh |
SHMEM_HOME_PATH 检测 + include 注入 |
| CMake 工具 | ascendc_impl_build.py |
SHMEM include 路径硬编码注入 |
3.3 抽象层设计
当前通信原语通过 HcclOpResParam 抽象，没有引入统一的 CommunicationInterface。dispatch/combine 直接依赖 HCCL 的 Window 结构体：
// 通信上下文是 HCCL 强绑定的
__gm__ HcclOpResParam *winContext_[COMM_NUM]; // dispatch
__gm__ HcclOpResParam *epWinContext_; // combine
__gm__ HcclOpResParam *tpWinContext_; // combine (TP domain)
如果要支持 SHMEM，需要：
1. 引入 ICommContext 抽象基类
2. 派生 HcclWindowContext 和 ShmemContext
3. dispatch/combine 通过基类指针访问通信原语
4. raw_distributed 子模块：通信原语详解
4.1 cam_moe_distribute_dispatch.h
文件：op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_dispatch.h
类模板参数（第19-24行 fused_deep_moe_base.h）：
#define TemplateDispatchTypeClass
typename XType, typename ExpandXOutType, bool StaticQuant,
bool DynamicQuant, bool IsSmoothScaleExist, bool IsNeedAllgater, uint32_t EXEC_FLAG
核心通信流程（第1023-1039行）：
void Process()
{
if ASCEND_IS_AIV {
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
QuantForShareExpert(); // 共享专家量化（本地操作）
}
AlltoAllDispatch(); // ← 核心：AlltoAll 数据搬运
SetStatus(); // 写入本 rank 状态到各远端 Window
WaitDispatch(); // 轮询等待所有远端状态就绪
LocalWindowCopy(); // 从本地 Window 读取数据到 workspace
if constexpr (IsNeedAllgater) {
AllGatherSetStatusAndWait();
AllgatherProcessOut();
}
UpdataTokenNumsOut(); // 更新 expertTokenNums 计数
}
}
AlltoAllDispatch 关键路径（第575-602行）：
void AlltoAllDispatch()
{
// 1. 加载 expertIds（决定每个 token 去哪个 expert/rank）
DataCopyPad(expertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, copyPadParams);
// 2. 计算每个 expert 的 token 计数 + 状态更新
if (enableAivOpt_) {
CountAndPrepareStatusWithAivOpt(expertIdsCnt); // 优化路径
} else {
for (uint32_t tokenIndex = 0; tokenIndex < expertIdsCnt; ++tokenIndex) {
expertCountTensor_(tokenIndex) = statusTensor_(expertId * INT32_NUM_PER_BLOCK + 1);
statusTensor_(expertId * INT32_NUM_PER_BLOCK + 1)++; // 原子计数
}
}
// 3. 将 token 数据写入目标 rank 的 Window
SendToMoeExpert();
}
SendToMoeExpert 数据搬运（第433-487行）：
void SendToMoeExpert()
{
for (uint32_t tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
int32_t dstExpertId = expertIdsTensor_(tokenIndex);
uint32_t tempRankId = dstExpertId / moeExpertNumPerRank_;
// 计算目标 Window 地址
GM_ADDR rankGM = GetWindAddrByRankId(COMM_EP_IDX, tempRankId)
+ expertPerSizeOnWin_ * (epRankId_ * moeExpertNumPerRank_ + dstExpertId % moeExpertNumPerRank_)
+ hCommuSize_ * expertCountTensor_(tokenIndex);
// 量化 + 写入目标 Window
if constexpr (DynamicQuant || StaticQuant) {
QuantProcess(expertIndex, scalesGMTensor_);
DataCopy(dstWinGMTensor, xOutTensor_, axisHCommu_);
} else {
DataCopy(dstWinGMTensor, xTmpTensor_, axisHCommu_);
}
}
}
WaitDispatch 同步机制（第684-726行）：
void WaitDispatch()
{
// 轮询 windowInstatusFp32Tensor_ 中的 flag
while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
DataCopy(statusFp32Tensor_, windowInstatusFp32Tensor_[startStatusIndex * stateOffset_ / sizeof(float)], intriParams);
GatherMask(gatherMaskOutTensor, statusFp32Tensor_, gatherTmpTensor, true, mask, ...);
Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams);
sumOfFlag = statusSumOutTensor.GetValue(0);
}
SyncAll<true>; // 全核同步
}
4.2 cam_moe_distribute_combine.h
文件：op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_combine.h
类模板参数（与主 kernel 共享 TemplateMC2TypeClass）：
#define TemplateMC2TypeClass
typename ExpandXType, typename W1ScaleType, typename W2ScaleType,
typename ExpandIdxType, bool IsNeedReduceScatter, uint32_t EXEC_FLAG
CombineCalcInfo 结构体（第46-56行）— 传递给 GMM2 AIV 侧的关键信息：
struct CombineCalcInfo {
uint64_t expertPerSizeOnWin_; // 每个 expert 在 Window 上占用的字节数
uint32_t epRankId_;
uint32_t epWorldSize_;
uint32_t moeExpertPerRankNum_;
uint32_t axisH_;
uint32_t moeSendNum_;
GM_ADDR epSendCount_;
__gm__ HcclOpResParam *epWinContext_; // ← HCCL 强绑定
uint64_t winDataSizeOffset_;
};
Process 流程（第719-734行）：
void Process()
{
SyncAll<true>;
if constexpr (IsNeedReduceScatter) {
ReduceScatterTrans(); // TP 域 ReduceScatter
}
if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
BuffInit();
SetWaitTpStatusAndDisPatch(); // 非融合路径的 AlltoAll
}
AlltoAllBuffInit();
SetStatus(); // 写入 EP 状态
WaitDispatch(); // 等待 EP 域数据就绪
LocalWindowCopy(); // 从 Window 读取 + 加权求和
}
LocalWindowCopy 核心逻辑（第627-716行）：
void LocalWindowCopy()
{
for (uint32_t tokenIndex = beginIndex; tokenIndex < endIndex; tokenIndex++) {
Duplicate(sumFloatBufLocal, (float)0, axisH_); // 清零累加器
for (uint32_t i = 0; i < axisK_; i++) {
int32_t moeExpert = expertIdsLocal.GetValue(index);
float scaleVal = expandScalesLocal.GetValue(index);
// 从 Window 读取该 expert 的 token 数据
GM_ADDR wAddr = epWindowGM_ + expertPerSizeOnWin_ * moeExpert
+ indexCountsLocal.GetValue(index) * axisHExpandXTypeSize_ + tokenOffset;
rowTmpGlobal_.SetGlobalBuffer((__gm__ ExpandXType *)wAddr);
// 量化反量化 + 缩放 + 累加
Cast(rowTmpFloatLocal, tmpUb, CAST_NONE, processLen);
Muls(mulBufLocal, rowTmpFloatLocal, scaleVal, processLen);
Add(sumFloatBufLocal, sumFloatBufLocal, mulBufLocal, processLen);
}
Cast(sumBufLocal, sumFloatBufLocal, CAST_RINT, processLen);
DataCopy(expandOutGlobal_[tokenIndex * axisH_ + tokenOffset], sumBufLocal, processLen);
}
}
5. 三条执行路径的通信差异
5.1 路径判定机制
EXEC_FLAG 位图定义（fused_deep_moe_tiling.h 第104-108行）：
constexpr uint32_t EXEC_FLAG_DEEP_FUSE = (1U << 0); // 深度融合
constexpr uint32_t EXEC_FLAG_TENSOR_LIST = (1U << 1); // 权重 tensor list
constexpr uint32_t EXEC_FLAG_X_ACTIVE_MASK = (1U << 2); // 激活 mask
constexpr uint32_t EXEC_FLAG_SHARED_EXPERT = (1U << 3); // 共享专家
constexpr uint32_t EXEC_FLAG_SMOOTH_QUANT = (1U << 4); // 平滑量化
Tiling 阶段计算 EXEC_FLAG（fused_deep_moe_tiling.cpp 第35-40行）：
constexpr uint32_t EXEC_FLAG =
(TPL_IS_DEEP_FUSE ? EXEC_FLAG_DEEP_FUSE : 0) |
(TPL_IS_TENSOR_LIST ? EXEC_FLAG_TENSOR_LIST : 0) |
(TPL_IS_X_ACTIVE_MASK ? EXEC_FLAG_X_ACTIVE_MASK : 0) |
(TPL_IS_SHARED_EXPERT ? EXEC_FLAG_SHARED_EXPERT : 0) |
(TPL_IS_SMOOTH_QUANT ? EXEC_FLAG_SMOOTH_QUANT : 0);
5.2 路径对比
| 特性 | 非深度融合 | 深度融合 (DEEP_FUSE) | 共享专家 (SHARED_EXPERT) |
|---|---|---|---|
| 触发条件 | moeExpertNumPerRank == 1 |
moeExpertNumPerRank > 1 |
share_gmm1_weight != nullptr |
| Dispatch | 独立执行（AIV 上） | 跳过（由 GMM2 AIV 侧融合） | 独立执行 |
| Combine | 独立执行（AIV 上） | 融合进 GMM2 AIV 侧 | 融合进 GMM2 AIV 侧 |
| GMM1 后同步 | PipeBarrier + CrossCoreFlag |
同左 | 同左 |
| GMM2 AIV 角色 | 仅做 epilogue | 执行 AllToAllSend + ReducePermute |
执行 ProcessCombine |
| 共享专家 GMM | 无 | 无 | 有（额外 GMM1+GMM2） |
| 核分配 | AIC + AIV 1:1 | AIC + AIV 1:2 (1C2V) | AIC + AIV 1:2 |
5.3 深度融合路径的通信融合
GMM2 AIV 侧执行流程（grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h 第291-458行）：
template <>
CATLASS_DEVICE void operator()<AscendC::AIV>(Params const ¶ms)
{
auto *combiner = (CamMoeDistributeCombine *)params.combiner;
// 1. 先执行 GMM2 的 epilogue（反量化 + SwiGLU 后处理）
for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) {
for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) {
CheckSyncFlag(syncGmAddr + SOFT_SYNC_OFFSET, coreIdx, target); // 等待 AIC 完成
blockEpilogue(gmGroupOffsetD, groupIdx, ...); // epilogue 计算
}
}
// 2. 共享专家 GMM2（如有）
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
// ... 执行共享专家 GMM2 epilogue
}
// 3. 融合 Combine 通信
if constexpr (EXEC_FLAG & EXEC_FLAG_SHARED_EXPERT) {
if (GetSubBlockIdx() == 1) {
combiner->ProcessCombine(); // ← 在 GMM2 AIV 上执行 Combine
}
} else if constexpr (EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) {
if (GetSubBlockIdx() == 0) {
combiner->AllToAllSend(); // ← 发送阶段
} else {
combiner->ReducePermute(); // ← 接收+累加阶段
}
} else {
combiner->Process(); // ← 非融合：完整 Combine
}
}
关键同步机制：
- SEND_SYNC_EVENT_ID = 6 — GMM2 epilogue 完成 → Combine 开始
- RECV_SYNC_EVENT_ID = 7 — Combine 接收完成 → 主 kernel 继续
- CheckSyncFlag — AIV 等待 AIC 的 Cube 计算完成（soft sync）
6. Workspace 内存布局
6.1 WorkSpaceOffset 结构体
文件：fused_deep_moe_tiling.h 第19-40行
struct WorkSpaceOffset {
// MM1/GMM1-Swiglu input
int64_t shareX1TokenOffset; // 共享专家 GMM1 输入
int64_t x1TokenOffset; // MoE GMM1 输入（dispatch 后）
int64_t shareX1ScaleOffset; // 共享专家量化 scale
int64_t x1ScaleOffset; // MoE 量化 scale
// MM1/GMM1-Swiglu output
int64_t shareSwigluOffset; // 共享专家 SwiGLU 输出
int64_t swigluOffset; // MoE SwiGLU 输出
// MM2/GMM2 input
int64_t shareX2TokenOffset; // 共享专家 GMM2 输入
int64_t x2TokenOffset; // MoE GMM2 输入
int64_t shareX2ScaleOffset; // 共享专家 GMM2 scale
int64_t x2ScaleOffset; // MoE GMM2 scale
int64_t swapSpaceOffset; // C↔V 数据交换区（GMM1 输出/GMM2 输入）
int64_t y2TokenOffset; // 浅融合路径 GMM2 输出
int64_t groupListOffset; // expert token 计数前缀和
int64_t expandIdxOffset; // dispatch 远端 token 索引
int64_t epSendCountOffset; // 每个 expert 从每个 rank 接收的 token 数
int64_t reservedOffset; // 保留区
};
6.2 内存复用策略（ENABLE_REUSE_MEMORY）
文件：fused_deep_moe_tiling.cpp 第671-711行
#ifdef ENABLE_REUSE_MEMORY
// shareX1 和 shareX2 复用同一块内存（因为 SwiGLU 后不再需要 X1）
tilingData.workSpaceOffset.shareX1TokenOffset = offset;
tilingData.workSpaceOffset.shareX2TokenOffset = offset;
tilingData.workSpaceOffset.x1TokenOffset = offset + shareExpertTokenNum * h * sizeof(int8_t);
tilingData.workSpaceOffset.x2TokenOffset = offset + shareExpertTokenNum * shareGmm2HLen * sizeof(int8_t);
offset += maxTokenSize; // 取 x1/x2 中较大者
// scale 同样复用
tilingData.workSpaceOffset.shareX1ScaleOffset = offset;
tilingData.workSpaceOffset.shareX2ScaleOffset = offset;
tilingData.workSpaceOffset.x1ScaleOffset = offset + shareExpertTokenNum * sizeof(float);
tilingData.workSpaceOffset.x2ScaleOffset = offset + shareExpertTokenNum * sizeof(float);
offset += tokenScaleSize;
// swapSpace 和 shareSwiglu 复用
tilingData.workSpaceOffset.swapSpaceOffset = offset;
tilingData.workSpaceOffset.shareSwigluOffset = offset;
tilingData.workSpaceOffset.swigluOffset = offset + shareExpertTokenNum * shareGmm1HLen * sizeof(float);
offset += maxSwapSwigluSize;
tilingData.workSpaceOffset.y2TokenOffset = offset;
#else
// 非复用模式：各区域独立分配
// ... 每个区域单独 CeilUp 对齐
#endif
6.3 Workspace 总大小计算
// 第716-717行
size_t usrSize = tilingData.workSpaceOffset.reservedOffset + reservedSize;
workSpaces[0] = SYSTEM_NEED_WORKSPACE + usrSize; // 16MB 系统预留 + 用户区
各区域大小公式：
| 区域 | 大小公式 | 说明 |
|---|---|---|
| x1/x2 token | max(shareBs*h, maxTokenNum*h) |
int8 存储 |
| scale | (shareBs + maxTokenNum) * sizeof(float) |
per-token scale |
| swap/swiglu | max(maxTokenNum*gmm1HLen, maxTokenNum*h) * sizeof(int32/float) |
GMM1 输出 |
| y2 | maxTokenNum * h * 2 (if moeExpertNumPerRank==1) |
GMM2 输出 |
| groupList | moeExpertNumPerRank * sizeof(int64_t) |
前缀和 |
| expandIdx | bs * topK * sizeof(int32_t) |
远端索引 |
| epSendCount | epRankSize * moeExpertNumPerRank * sizeof(int32_t) |
接收计数 |
| reserved | 256KB | 保留 |
6.4 数据流转图
HBM (GM)
│
├── x [bs, h] ─────────────────────────────────────────────────────────────┐
├── expert_ids [bs, k] ─────────────────────────────────────────────────────┤
├── gmm1_weight [moeExpertNumPerRank, h, gmm1HLen] ─────────────────────────┤
├── gmm2_weight [moeExpertNumPerRank, gmm1HLen/2, h] ───────────────────────┤
└── workspace ──────────────────────────────────────────────────────────────┘
│
├── shareX1Token (Window Copy ← Dispatch)
│ ↓ GMM1 (Cube) + SwiGLU (Vector)
├── shareSwigluOut
│ ↓
├── shareX2Token ──────────────────────────→ ShareGMM2 → share_output
│
├── x1Token (Window Copy ← Dispatch)
│ ↓ GMM1 (Cube) + SwiGLU (Vector)
├── swigluOut
│ ↓
├── x2Token
│ ↓ GMM2 (Cube, 输出到 swapSpace)
│ ↓ epilogue (Vector, 反量化 + 写回 Window)
│ ↓ Combine (AllToAllSend → Window → ReducePermute → output)
│
├── swapSpace (GMM2 AIC 输出 → AIV epilogue 输入)
├── y2Token (浅融合路径 GMM2 最终输出)
├── groupList (expert token 前缀和)
├── expandIdx (远端 token 索引)
├── epSendCount (各 rank 接收计数)
└── reserved
7. 编译系统 SHMEM 机制
7.1 select_ops.py — SHMEM 算子过滤
文件：build/cam/comm_operator/select_ops.py
requires_shmem() 函数（第55-56行）：
def requires_shmem(registry, op):
return bool(get_meta(registry, op).get("requires_shmem", False))
SHMEM 过滤逻辑（第116-121行）：
# 3. SHMEM filter: drop requires_shmem operators when SHMEM is not installed.
if not shmem:
dropped = [o for o in candidates if requires_shmem(registry, o)]
for o in dropped:
sys.stderr.write(f"NOTE: dropping {o} (requires SHMEM, which is not installed)\n")
candidates = [o for o in candidates if not requires_shmem(registry, o)]
完整选择管线：
1. 确定候选集（全 SOC 列表 或 用户 -a 列表）
↓
2. fused_deep_moe base/w4a8 互斥规则
-q → 选 w4a8，剔除 base
非 -q → 选 base，剔除 w4a8
fwk 不受影响
↓
3. SHMEM 过滤
shmem=0 → 剔除 requires_shmem=true 的算子
↓
4. 按 registry 顺序排列输出
7.2 compile_ascend_proj.sh — SHMEM 路径注入
文件：build/cam/comm_operator/compile_ascend_proj.sh
SHMEM_HOME_PATH 检测（第19-22行）：
# Whether SHMEM is installed (determined by the SHMEM_HOME_PATH env variable)
shmem_installed=1
if [ -z "${SHMEM_HOME_PATH}" ]; then
shmem_installed=0
fi
SHMEM include 注入（第179-192行）：
KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
if [ -n "${CATLASS_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="${KERNEL_COMPILE_OPTS} -I${CATLASS_HOME_PATH}/include"
fi
if [ -n "${KERNEL_COMPILE_OPTS}" ]; then
echo "Patching kernel compile options: ${KERNEL_COMPILE_OPTS}"
find build_out/op_kernel/CMakeFiles -name "build.make" -exec sed -i \
"s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
fi
关键机制：通过 sed 替换 .make 文件中的 --compile-options="" 为实际的 include 路径。这是因为 npu_op_* 构建系统的 simple_kernel_compile 不读取 custom_compile_options.ini。
7.3 ascendc_impl_build.py — SHMEM 编译选项生成
文件：src/cam/comm_operator/ascend_kernels/cmake_files/cmake/util/ascendc_impl_build.py
SHMEM include 路径硬编码（第167-179行）：
SHMEM_HOME_PATH = os.environ.get('SHMEM_HOME_PATH')
if SHMEM_HOME_PATH is None:
SHMEM_HOME_PATH = os.path.realpath("/usr/local/Ascend/shmem/latest")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include")
options.append("-I" + SHMEM_HOME_PATH + "shmem/include/device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/host_device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/internal")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src/device")
options.append("-I" + SHMEM_HOME_PATH + "/memfabric_hybrid/include/smem/device")
注意：这里有一个 bug — 第172行缺少 / 分隔符：
options.append("-I" + SHMEM_HOME_PATH + "shmem/include/device") # 缺少 /
# 应该是：
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device")
7.4 operator_registry.json — SHMEM 标记
文件：src/cam/comm_operator/ascend_kernels/operator_registry.json
{
"operator_meta": {
"dispatch_layout": { "requires_shmem": false },
"fused_deep_moe": {
"requires_shmem": false,
"family": "fused_deep_moe",
"variant": "base"
},
"moe_combine_normal": { "requires_shmem": false },
"moe_dispatch_normal": { "requires_shmem": false },
"notify_dispatch": { "requires_shmem": false }
}
}
当前所有算子的 requires_shmem 均为 false，即 SHMEM 过滤当前不影响任何算子。
8. SHMEM_HOME_PATH 环境变量
8.1 作用域
| 阶段 | 使用方式 | 默认值 |
|---|---|---|
| 编译时 | compile_ascend_proj.sh 检测是否设置 |
未设置 → shmem_installed=0 |
| 编译时 | ascendc_impl_build.py 读取路径 |
/usr/local/Ascend/shmem/latest |
| 运行时 | kernel 代码中 未使用 | N/A |
8.2 编译时行为
SHMEM_HOME_PATH 未设置:
├── select_ops.py: shmem=0 → 剔除 requires_shmem=true 的算子
├── compile_ascend_proj.sh: 不注入 SHMEM include
└── ascendc_impl_build.py: 使用默认路径 /usr/local/Ascend/shmem/latest
SHMEM_HOME_PATH 已设置:
├── select_ops.py: shmem=1 → 保留所有算子
├── compile_ascend_proj.sh: 注入 -I${SHMEM_HOME_PATH}/shmem/include 等
└── ascendc_impl_build.py: 使用 SHMEM_HOME_PATH 的值
8.3 运行时行为
当前 kernel 代码不读取 SHMEM_HOME_PATH。SHMEM 的运行时行为（如 aclshmem_* API 的链接）由 CANN 的 SHMEM 库在算子加载时动态解析。
9. 未来 SHMEM 集成路径
9.1 需要修改的代码
要让 FusedDeepMoE 支持 SHMEM，需要以下改动：
9.1.1 通信抽象层（新增）
新增文件：
op_kernel/fused_deep_moe/comm/icomm_context.h # 抽象基类
op_kernel/fused_deep_moe/comm/hccl_window_context.h # HCCL 实现
op_kernel/fused_deep_moe/comm/shmem_context.h # SHMEM 实现
// icomm_context.h (概念设计)
class ICommContext {
public:
virtual GM_ADDR GetDataAddr(int rankId) = 0;
virtual GM_ADDR GetStateAddr(int rankId) = 0;
virtual void SetFlag(int rankId, int value) = 0;
virtual int WaitFlag(int rankId, int target) = 0;
virtual void DataCopy(void* dst, void* src, size_t len) = 0;
};
9.1.2 dispatch/combine 改造
// cam_moe_distribute_dispatch.h 改造方向
template <typename CommContext>
class CamMoeDistributeDispatch {
CommContext* commCtx_; // 替代 winContext_[COMM_NUM]
// ...
};
9.1.3 编译系统改动
- operator_registry.json — 新增 SHMEM 变体：
{
"fused_deep_moe_shmem": {
"requires_shmem": true,
"family": "fused_deep_moe",
"variant": "shmem"
}
}
- select_ops.py — 新增 SHMEM 变体互斥规则：
SHMEM = "fused_deep_moe_shmem"
# base/w4a8/shmem 三者互斥
- compile_ascend_proj.sh — 已支持，无需改动（通过 SHMEM_HOME_PATH 自动注入）
9.1.4 运行时 SHMEM API 调用
SHMEM 路径的核心变化：
// HCCL Window 模式（当前）
GM_ADDR dstAddr = GetWindAddrByRankId(COMM_EP_IDX, rankId);
DataCopy(dstWinGMTensor, xOutTensor_, axisHCommu_);
// SHMEM 模式（未来）
// 使用 aclshmem_* API 直接 RDMA
aclshmem_put(dstAddr, srcAddr, len, dstRankId);
aclshmem_fence();
aclshmem_wait_flag(flagAddr, target);
9.2 SHMEM 集成挑战
| 挑战 | 说明 |
|---|---|
| 同步语义差异 | HCCL Window 是双向对称的（localWindowsIn + remoteRes），SHMEM 是单边 RDMA |
| 状态通道设计 | 当前 windowsExp 用于 flag 同步，SHMEM 需要额外设计 fence/notify 机制 |
| 多域通信 | dispatch 使用 COMM_EP_IDX + COMM_TP_IDX 两个域，SHMEM 需要对应两个通信组 |
| 地址计算 | expertPerSizeOnWin_、winDataSizeOffset_ 等偏移量需要重新映射到 SHMEM 内存池 |
| 量化融合 | dispatch 中量化+Window 拷贝是融合的，SHMEM 路径需要保持同样的融合度 |
9.3 推荐集成策略
- Phase 1：引入
ICommContext抽象，HCCL 路径保持不变 - Phase 2：实现
ShmemContext，在notify_dispatch等简单算子上验证 - Phase 3：将 SHMEM 路径扩展到 FusedDeepMoE 的 dispatch/combine
- Phase 4：性能调优，利用 SHMEM 的单边 RDMA 降低同步开销
10. 关键代码索引
| 文件 | 关键函数/结构 | 行号 |
|---|---|---|
fused_deep_moe.h |
FusedDeepMoe::Process() |
354-448 |
fused_deep_moe.h |
GmmDeqSwigluQuant() |
56-152 |
fused_deep_moe.h |
GmmDeq() |
154-238 |
fused_deep_moe_tiling.h |
WorkSpaceOffset |
19-40 |
fused_deep_moe_tiling.h |
EXEC_FLAG_* 定义 |
104-108 |
fused_deep_moe_tiling.cpp |
SetWorkSpace() |
636-719 |
fused_deep_moe_tiling.cpp |
SetHcommCfg() |
605-616 |
fused_deep_moe_tiling_key.h |
ASCENDC_TPL_SEL |
34-55 |
cam_moe_distribute_dispatch.h |
GetWindAddrByRankId() |
83-91 |
cam_moe_distribute_dispatch.h |
AlltoAllDispatch() |
575-602 |
cam_moe_distribute_dispatch.h |
SendToMoeExpert() |
433-487 |
cam_moe_distribute_dispatch.h |
WaitDispatch() |
684-726 |
cam_moe_distribute_combine.h |
CombineCalcInfo |
46-56 |
cam_moe_distribute_combine.h |
GetWinAddrByRankId() |
96-111 |
cam_moe_distribute_combine.h |
LocalWindowCopy() |
627-716 |
cam_moe_distribute_combine.h |
Process() |
719-734 |
grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h |
operator()<AIV>() |
291-458 |
select_ops.py |
requires_shmem() |
55-56 |
select_ops.py |
SHMEM 过滤逻辑 | 116-121 |
compile_ascend_proj.sh |
SHMEM_HOME_PATH 检测 | 19-22 |
compile_ascend_proj.sh |
SHMEM include 注入 | 179-192 |
ascendc_impl_build.py |
SHMEM include 路径 | 167-179 |
operator_registry.json |
operator_meta | 12-30 |
11. 结论
-
FusedDeepMoE 当前完全基于 HCCL Window 通信，kernel 代码中不存在任何 SHMEM 分支。
-
通信原语通过
HcclOpResParam强绑定，dispatch/combine 直接操作windowsIn（数据）和windowsExp（状态）。 -
三条执行路径的通信差异： - 非深度融合：dispatch 和 combine 独立执行 - 深度融合：combine 融合进 GMM2 AIV 侧，通过
AllToAllSend+ReducePermute实现 - 共享专家：额外执行共享专家 GMM1+GMM2，combine 通过ProcessCombine执行 -
Workspace 内存布局采用复用策略（
ENABLE_REUSE_MEMORY），shareX1/X2 共用、shareX1Scale/X2Scale 共用、swapSpace/shareSwiglu 共用。 -
编译系统已完整预留 SHMEM 注入点： -
select_ops.py通过requires_shmem过滤 -compile_ascend_proj.sh通过SHMEM_HOME_PATH注入 include -ascendc_impl_build.py硬编码 SHMEM include 路径（含一个路径 bug） -
SHMEM 集成需要引入通信抽象层，将
HcclOpResParam替换为ICommContext接口，并在 dispatch/combine 中通过模板参数化通信后端。 -
ascendc_impl_build.py第172行存在路径 bug：SHMEM_HOME_PATH + "shmem/include/device"缺少/分隔符。
项目: 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 仓库中。
项目: UMDK/CAM — Communication Acceleration for Matrix on Ascend NPU
分支:cam_feature(remotes/origin/cam_feature)
分析日期: 2026-07-21
分析范围:moe_combine_normal_zb+moe_combine_lowlatency_zb两个 Zero-Buffer Combine 算子
目录
- ZB Combine 设计哲学
- moe_combine_normal_zb 完整分析
- moe_combine_lowlatency_zb 完整分析
- Combine 阶段的 SHMEM 通信
- 与 master 版 moe_combine_normal 的对比
- Tiling 策略
- 关键数据结构与内存布局
1. ZB Combine 设计哲学
1.1 Zero-Buffer 的核心思想
Zero-Buffer 是 CAM 库在 cam_feature 分支上引入的全新 MoE 通信范式。其核心洞察是：
传统 MoE Combine 需要两次数据搬运：
1. 每个 rank 将本地 GMM 输出 CopyBufferToShare 写入远端 rank 的 HCCL Window
2. 通过 polling WaitBuffCopy 等待所有源 rank 的数据到位
3. 最后执行加权求和
Zero-Buffer 将其简化为一次直接读取：
- 数据在 GMM 完成后直接写入 SHMEM window（由 dispatch 阶段完成）
- Combine 阶段通过 shmem_ptr() 直接读取远端 SHMEM 中的数据
- 通过 ZeroBufferSyncFlag 的 magic + phase 协议保证数据可见性
1.2 设计原则
| 原则 | 说明 |
|---|---|
| 零拷贝 | Combine 不需要额外的 buffer 中转，直接从 SHMEM window 读取 |
| Magic 单调递增 | 每个 kernel 入口调用 IncrementMagic()，保证 flag 唯一性 |
| Ping-Pong 双缓冲 | magic % 2 选择 buffer，避免连续 kernel 调用间的 flag 冲突 |
| MTE3 Fence 保序 | SetFlag 前强制 MTE3 fence，确保数据写入先于 flag 可见 |
| 全 AIV 并行 | 所有 AIV 核同时参与计算，SplitCoreCal 均匀分配 token |
1.3 文件结构
moe_combine_normal_zb/
├── op_kernel/ # 设备端 kernel 代码
│ ├── moe_combine_normal_zero_buffer.cpp # kernel 入口 + 实现
│ ├── moe_combine_normal_zero_buffer.h # 类定义
│ └── moe_combine_normal_zero_buffer_tiling.h # tiling 数据结构
├── op_host/ # 主机端 tiling 代码
│ ├── moe_combine_normal_zero_buffer.cpp # tiling 入口
│ └── moe_combine_normal_zero_buffer_tiling.cc # tiling 实现
└── op_api/ # aclnn API 注册
├── aclnn_moe_combine_normal_zero_buffer.cpp
└── aclnn_moe_combine_normal_zero_buffer.h
moe_combine_lowlatency_zb/
├── op_kernel/
│ ├── moe_combine_lowlatency_zero_buffer.cpp
│ ├── moe_combine_lowlatency_zero_buffer.h
│ └── moe_combine_lowlatency_zero_buffer_tiling.h
├── op_host/
│ ├── moe_combine_lowlatency_zero_buffer.cpp
│ └── moe_combine_lowlatency_zero_buffer_tiling.cpp
└── op_api/
├── aclnn_moe_combine_lowlatency_zero_buffer.cpp
└── aclnn_moe_combine_lowlatency_zero_buffer.h
utils/op_kernel/ # 公共工具
├── zero_buffer_api.h # SHMEM API 封装 (include shmem.h)
├── zero_buffer_sync_flag.h # 跨卡 flag 同步机制
├── comm_args_zb.h # ZB 版通信参数
├── sync_collectives_zb.h # 同步原语
├── moe_lowlatency_zero_buffer_base.h # lowlatency 共享工具
└── data_copy_zb.h # 数据拷贝辅助
2. moe_combine_normal_zb 完整分析
2.1 算子签名与 I/O
Kernel 入口 (moe_combine_normal_zero_buffer.cpp):
extern "C" __global__ __aicore__ void moe_combine_normal_zero_buffer(
GM_ADDR recvX, // [bs*k, h] 本 rank 收到的所有 token（各 expert 输出）
GM_ADDR epRecvCount, // [moeExpertNum, epWorldSize] 每个 expert 从各 rank 收到的 token 数
GM_ADDR topkWeights, // [bs, k] top-k 权重
GM_ADDR topkIdx, // [bs, k] top-k expert ID
GM_ADDR sendTokenIdx, // [bs, k] 在远端 recvX 中的偏移索引
GM_ADDR probGrad, // [bs, h] 概率梯度输入（可选）
GM_ADDR XOut, // [bs, h] 输出
GM_ADDR sendCostStatsOut,// [epWorldSize] 诊断统计（可选）
GM_ADDR gradOut, // [bs, k] 梯度输出（可选）
GM_ADDR workspaceGM,
GM_ADDR tilingGM)
关键输入语义:
- recvX: 本 rank 上所有 expert 的 GMM 输出，按 [token_idx, topk_idx] 排列
- epRecvCount: 全局接收计数，epRecvCount[expertId][rankId] 表示该 expert 从 rankId 收到的 token 数
- sendTokenIdx: 每个 (token, topk) 对在其目标 rank 的 recvX 中的位置索引
2.2 执行流程
Process():
│
├── syncFlag_.IncrementMagic() // Step 1: 获取 magic M，递增
├── syncFlag_.BarrierAll() // Step 2: 全 rank 屏障，确保所有 GMM 完成
├── ReadBufferFromRemote() // Step 3: 从远端 SHMEM 读取并加权求和
└── SyncAll<true>() // Step 4: 全核同步
Step 1: IncrementMagic
__aicore__ inline void IncrementMagic() {
GM_ADDR magicAddr = flagBaseGM_ + coreIdx_ * MAGIC_SLOT_ALIGN; // 512B 对齐
// DataCacheCleanAndInvalid → 读 magic
magic_ = static_cast<int32_t>(magicGT(0));
magicGT(0) = magic_ + 1; // 递增
// DataCacheCleanAndInvalid → 写回
barrierSeq_ = 0;
}
每个 core 有独立的 magic slot（512B 对齐），避免 false sharing。
Step 2: BarrierAll
__aicore__ inline void BarrierAll() {
SyncAll<true>(); // 核内同步
SplitCoreCal(epWorldSize_, ...); // 分配远端 rank
// Write: 向所有分配的远端 rank 写 barrier flag
for (r = startRank; r < endRank; r++) {
DataCopy(remoteSlot, ubFlag, ...); // 写 flag 到远端 SHMEM
}
// Wait: 轮询本地 slot 等待所有远端 rank 的 flag
for (r = startRank; r < endRank; r++) {
do { observed = localSlot; } while (observed != barrierFlag);
}
SyncAll<true>();
}
关键保证: BarrierAll 返回时，所有 rank 的 GMM 输出已完整写入 SHMEM。
Step 3: ReadBufferFromRemote
__aicore__ inline void ReadBufferFromRemote() {
SplitCoreCal(axisBS_, ...); // 将 bs 个 token 均匀分配到各 AIV core
// 初始化所有 buffer
tpipe_->InitBuffer(xOutBuf_, h32AlignRecvXLen_); // 14KB (h=7168, bf16)
tpipe_->InitBuffer(tokenFloatBuf_, h32AlignFloatLen_); // 28KB (fp32)
tpipe_->InitBuffer(weightedMulBuf_, h256AlignFloatLen_); // 28KB
tpipe_->InitBuffer(sumFloatBuf_, h32AlignFloatLen_); // 28KB
tpipe_->InitBuffer(weightedSumQueue_, DOUBLE_BUFFER, h32AlignRecvXLen_); // 28KB
// ... topkWeightsBuf_, sendTokenIdxBuf_, topkIdxBuf_, allRecvCountBuf_
// 读取全局 recvCount 到本地
DataCopyPad(allRecvCountLocal, epRecvCountGT_, ...);
for (tokenIndex = startTokenIndex; tokenIndex < endTokenIndex; tokenIndex++) {
// 读取当前 token 的 topkWeights, topkIdx, sendTokenIdx
DataCopyPad(topkWeightsLocal, topkWeightsGT_[tokenIndex * axisK_], ...);
DataCopyPad(topkIdxLocal, topkIdxGT_[tokenIndex * axisK_], ...);
DataCopyPad(sendTokenIdxLocal, sendTokenIdxGT_[tokenIndex * axisK_], ...);
ReadBufferAndWeightedSum(tokenIndex, startTokenIndex);
}
}
2.3 加权求和核心逻辑
__aicore__ inline void ReadBufferAndWeightedSum(uint32_t tokenIndex, uint32_t startTokenIndex) {
Duplicate(sumFloatBufLocal, 0.0f, axisH_); // 清零累加器
for (uint32_t topkId = 0U; topkId < axisK_; topkId++) {
float scale = topkWeightsLocal.GetValue(topkId);
int32_t expertId = topkIdxLocal.GetValue(topkId);
int32_t remoteReadOffset = sendTokenIdxLocal(topkId);
int32_t remoteReadBase = allRecvCountLocal(expertId * epWorldSize_ + epRankId_);
uint64_t remoteReadAddr = (remoteReadBase + remoteReadOffset) * hRecvXTypeLen_;
int32_t dstRankId = expertId / moeExpertPerRankNum_;
auto ptr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(recvXGM_, dstRankId));
dstGT.SetGlobalBuffer((__gm__ XType *)(ptr + remoteReadAddr));
// 从远端 SHMEM 读取 token
LocalTensor<XType> tmpToken = weightedSumQueue_.AllocTensor<XType>();
DataCopyPad(tmpToken, dstGT, xOutCopyParams, ...);
weightedSumQueue_.EnQue(tmpToken);
tmpToken = weightedSumQueue_.DeQue<XType>();
// bf16 → fp32 → 乘权重 → 累加
Cast(tokenFloatLocal, tmpToken, CAST_NONE, axisH_); // bf16 → fp32
Muls(weightedMulBufLocal, tokenFloatLocal, scale, axisH_); // × weight
Add(sumFloatBufLocal, sumFloatBufLocal, weightedMulBufLocal, axisH_); // 累加
weightedSumQueue_.FreeTensor<XType>(tmpToken);
}
// fp32 → bf16 → 写输出
Cast(xOutLocal, sumFloatBufLocal, CAST_RINT, axisH_);
SyncFunc<V_MTE3>();
DataCopyPad(xOutGlobal_[tokenIndex * axisH_], xOutLocal, ...);
}
加权求和公式:
$$X_{out}[i] = \text{Cast}{fp32 \to bf16}\left(\sum{j=0}^{K-1} w_j \cdot \text{Cast}{bf16 \to fp32}(X{recv}[expert_j, token_j])\right)$$
2.4 地址计算详解
对于 token i 的第 j 个 top-k 选择：
expertId = topkIdx[i][j] // 目标 expert
dstRankId = expertId / moeExpertPerRankNum_ // expert 所在的 rank
remoteReadBase = epRecvCount[expertId][epRankId_] // 本 rank 在该 expert 上的基址偏移
remoteReadOffset = sendTokenIdx[i][j] // 在该 expert 接收列表中的位置
remoteReadAddr = (remoteReadBase + remoteReadOffset) * hRecvXTypeLen_
// 最终地址 = 远端 rank 的 SHMEM window + remoteReadAddr
finalAddr = shmem_ptr(recvXGM_, dstRankId) + remoteReadAddr
2.5 UB 内存预算
| Buffer | 大小 | 说明 |
|---|---|---|
xOutBuf_ |
14KB | 输出 token (h=7168 × bf16) |
tokenFloatBuf_ |
28KB | fp32 中间结果 |
weightedMulBuf_ |
28KB | 256 对齐的乘法结果 |
sumFloatBuf_ |
28KB | 累加器 |
weightedSumQueue_ |
28KB | 双缓冲 DMA queue |
topkWeightsBuf_ |
32B | k 个 float 权重 |
sendTokenIdxBuf_ |
32B | k 个 int32 索引 |
topkIdxBuf_ |
32B | k 个 int32 expert ID |
allRecvCountBuf_ |
128KB | epWorldSize × moeExpertNum × 4B (max 512×64×4=128KB) |
| 总计 | ~254KB |
3. moe_combine_lowlatency_zb 完整分析
3.1 算子签名与 I/O
Kernel 入口 (moe_combine_lowlatency_zero_buffer.cpp):
extern "C" __global__ __aicore__ void moe_combine_lowlatency_zero_buffer(
GM_ADDR expandX, // [A*tpWorldSize, h] 展开的 expert 输出
GM_ADDR expertIds, // [bs, k] top-k expert ID
GM_ADDR assistInfoForCombine, // [A*128] 辅助信息
GM_ADDR epSendCount, // [moeExpertNum*epWorldSize] EP 发送计数
GM_ADDR scales, // [bs, k] top-k 权重
GM_ADDR tpSendCount, // [tpWorldSize] TP 发送计数
GM_ADDR xActiveMask, // [bs] 或 [bs, k] 活跃 mask（可选）
GM_ADDR activationScale, // 预留
GM_ADDR weightScale, // 预留
GM_ADDR groupList, // 预留
GM_ADDR expandScales, // 预留
GM_ADDR sharedExpertX, // [bs, h] 共享专家输出（可选）
GM_ADDR elasticInfo, // 弹性信息（可选）
GM_ADDR oriX, // [bs, h] 原始输入（copy expert 用）
GM_ADDR constExpertAlpha1, // [constExpertNum] 常量专家参数
GM_ADDR constExpertAlpha2, // [constExpertNum]
GM_ADDR constExpertV, // [constExpertNum, h]
GM_ADDR XOut, // [bs, h] 输出
GM_ADDR workspaceGM, GM_ADDR tilingGM)
TilingKey 编码:
10000: tp=1, IsInt8Quant=0
10100: tp=2, IsInt8Quant=0
10020: tp=1, IsInt8Quant=1
10120: tp=2, IsInt8Quant=1
3.2 执行流程
Process():
│
├── CopyValidExpandXToRemote() // Step 1: 拷贝本地 expert 数据到 SHMEM
├── SyncAll<true>() // Step 2: 全核同步
├── tpipe_->Reset() // 重置 pipe
├── SetSyncFlag() // Step 3: 向远端 rank 写 flag
├── WaitSyncFlag() // Step 4: 轮询等待远端 flag
├── SyncAll<true>() // Step 5: 全核同步
└── InputToDstOutput() // Step 6: 从 SHMEM 读取并加权求和
Step 1: CopyValidExpandXToRemote
__aicore__ inline void CopyValidExpandXToRemote() {
if (!needCopyExpandX_) return; // oriX==0 或 validTokenNum==0 时跳过
SplitToCore(validTokenNum_, aivNum_, ...); // 分配 token
tpipe_->InitBuffer(copyInBuf_, COPY_BUFFER_NUM * GM2GM_COPY_BLOCK_SIZE); // 2×64KB
// 双缓冲 GM→GM 拷贝：oriX → expandX (SHMEM)
for (i = 0; dataSizeRemain > 0; ++i) {
copySize = min(dataSizeRemain, 64KB);
WaitFlag<MTE3_MTE2>(eventId); // 等待 DMA 通道
DataCopyPad(copyLocal[bufferOffset], oriXByteGT[srcOffset], ...); // GM→UB
SetFlag<MTE2_MTE3>(eventId);
WaitFlag<MTE2_MTE3>(eventId);
DataCopyPad(expandXByteGT[dstOffset], copyLocal[bufferOffset], ...); // UB→GM(SHMEM)
SetFlag<MTE3_MTE2>(eventId);
}
}
注意: 这是 lowlatency 特有的步骤。因为 lowlatency 模式下，本地 expert 的输出需要从 oriX（原始位置）拷贝到 expandX（SHMEM window）以供远端 rank 读取。normal_zb 不需要此步骤，因为数据已在 dispatch 阶段直接写入 SHMEM。
Step 3-4: SetSyncFlag / WaitSyncFlag
__aicore__ inline void SetSyncFlag() {
// 构造 flag 值：0x3F800000 = 1.0f (IEEE 754)
Duplicate(statusTensor, 0, rankNumPerBlock_ * 8);
Duplicate(statusTensor, 0x3F800000, mask, ...); // 按 mask 设置 1.0
for (targetRankId = curBlockStartRankId_; targetRankId < curBlockEndRankId_; ++targetRankId) {
auto ptr = GetSyncFlagAddrByRankId(targetRankId) + epRankId_ * COMBINE_SYNC_STATE_OFFSET;
gmRemoteStatusGT.SetGlobalBuffer(ptr);
DataCopy(gmRemoteStatusGT, statusTensor[...], 8UL);
}
}
__aicore__ inline void WaitSyncFlag() {
// 轮询本地 SHMEM 中的 flag 区域
compareTarget = 1.0f * rankNumPerBlock_;
while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
DataCopy(statusFp32Tensor, statusFp32TensorGT[curBlockStartRankId_ * ...], ...);
ReduceSum(statusSumOutTensor, statusFp32Tensor, ...);
sumOfFlag = statusSumOutTensor.GetValue(0);
}
// 清理 flag
Duplicate(cleanStateTensor, 0, ...);
DataCopy(statusFp32TensorGT[...], cleanStateTensor.ReinterpretCast<float>(), ...);
}
同步协议: 每个 rank 向远端 rank 的 SHMEM flag 区域写入 1.0f。当本 rank 收到所有远端 rank 的 flag 后（sum == rankNumPerBlock），表示所有远端数据已就绪。
Step 6: InputToDstOutput
__aicore__ inline void InputToDstOutput() {
SplitToCore(axisBS_, aivNum_, ...); // 分配 token
for (tokenIndex = startTokenId; tokenIndex < endTokenId; tokenIndex++) {
DataCopyPad(expertIdsTensor, expertIdsGM_[tokenIndex * axisK_], ...);
Duplicate(sumTokenTensor_, 0.0f, axisH_);
for (topkId = 0U; topkId < axisK_; topkId++) {
int32_t dstExpertId = expertIdsTensor.GetValue(topkId);
int32_t srcRankId = dstExpertId / moeExpertPerRankNum_;
uint32_t offsetIdx = dstExpertId * epWorldSize_ + epRankId_;
uint32_t col = offsetIdx % moeExpertNum_;
int32_t remoteBase = (col == 0) ? 0 : epSendCountGM_(offsetIdx - 1);
int32_t remoteOffset = expandIdxLocal(localTokenId * axisK_ + topkId);
float scale = expertScalesLocal_.GetValue(localTokenId * axisK_ + topkId);
// 从远端 SHMEM 读取
auto srcPtr = shmem_ptr(expandXGM_, srcRankId);
srcWinGMTensor.SetGlobalBuffer(srcPtr + hExpandXTypeSize_ * (remoteBase + remoteOffset));
// 加权求和（同 normal_zb）
DataCopyPad(tmpToken, srcWinGMTensor, ...);
Cast(tokenReduceTensor_, tmpToken, CAST_NONE, axisH_);
Muls(weightedSumTensor_, tokenReduceTensor_, scale, axisH_);
Add(sumTokenTensor_, sumTokenTensor_, weightedSumTensor_, axisH_);
}
Cast(xOutLocal, sumTokenTensor_, CAST_RINT, axisH_);
DataCopyPad(XOutGT_[tokenIndex * axisH_], xOutLocal, ...);
}
}
3.3 与 normal_zb 的关键差异
| 维度 | normal_zb | lowlatency_zb |
|---|---|---|
| 同步机制 | ZeroBufferSyncFlag (通用框架) |
自实现 SetSyncFlag/WaitSyncFlag |
| Magic 管理 | IncrementMagic() 自动管理 |
无 magic，使用固定 flag 值 1.0f |
| 本地数据拷贝 | 不需要（dispatch 已写入 SHMEM） | 需要 CopyValidExpandXToRemote |
| 地址计算 | epRecvCount 全局表 + sendTokenIdx |
epSendCount 前缀和 + expandIdx |
| TP 支持 | 仅 EP | EP + TP (tpWorldSize=1 or 2) |
| 共享专家 | 不支持 | 支持 (sharedExpertNum, sharedExpertRankNum) |
| 弹性缩容 | 不支持 | 支持 (elasticInfo) |
| 常量专家 | 不支持 | 支持 (constExpertAlpha1/2/V) |
| Copy Expert | 不支持 | 支持 (oriX → expandX 拷贝) |
| 量化通信 | 不支持 | INT8 量化 (IsInt8Quant) |
| TilingKey | 无 | 4 种组合 (tp × quant) |
| 输入复杂度 | 6 个输入 | 17 个输入 |
3.4 lowlatency 的地址计算
// 前缀和计算远端基址
offsetIdx = dstExpertId * epWorldSize_ + epRankId_
col = offsetIdx % moeExpertNum_
remoteBase = (col == 0) ? 0 : epSendCountGM_(offsetIdx - 1)
remoteOffset = expandIdxLocal(localTokenId * axisK_ + topkId)
// 最终地址
srcAddr = shmem_ptr(expandXGM_, srcRankId) + hExpandXTypeSize_ * (remoteBase + remoteOffset)
epSendCount 存储的是前缀和（prefix sum），epSendCount[i-1] 即为第 i 个 expert 在本 rank 上的数据基址。
4. Combine 阶段的 SHMEM 通信
4.1 SHMEM API 层次
shmem.h (CANN 提供)
├── shmem_ptr(baseGM, rankId) → 获取远端 rank 的 SHMEM window 基址
├── aclshmem_half_put_nbi(...) → 非阻塞 half 写入远端
├── aclshmem_bfloat16_put_nbi(...) → 非阻塞 bf16 写入远端
├── aclshmem_float_put_nbi(...) → 非阻塞 float 写入远端
├── aclshmem_half_get_nbi(...) → 非阻塞 half 读取远端
└── aclshmem_bfloat16_get_nbi(...) → 非阻塞 bf16 读取远端
zero_buffer_api.h (CAM 封装)
└── #include "shmem.h"
zero_buffer_sync_flag.h (同步框架)
└── ZeroBufferSyncFlag 类
4.2 shmem_ptr 的使用模式
两个 combine 算子都使用 shmem_ptr 进行远端 SHMEM 读取：
// normal_zb
auto ptr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(recvXGM_, dstRankId));
dstGT.SetGlobalBuffer((__gm__ XType *)(ptr + remoteReadAddr));
// lowlatency_zb
auto srcPtr = reinterpret_cast<__gm__ uint8_t *>(shmem_ptr(expandXGM_, srcRankId));
srcWinGMTensor.SetGlobalBuffer((__gm__ ExpandXType *)(srcPtr + hExpandXTypeSize_ * (remoteBase + remoteOffset)));
shmem_ptr(baseGM, rankId) 语义: 返回 rankId 的 SHMEM window 虚拟地址，该地址在本 rank 的地址空间中映射到远端 rank 的 SHMEM 物理内存（通过 HCCL 的 window 机制）。
4.3 ZeroBufferSyncFlag 内存布局
gva_gm + FLAG_AREA_BASE (100KB):
├── Magic Area (24KB = 48 cores × 512B)
│ ├── slot[0] @ 0×512B ← core 0 的 magic
│ ├── slot[1] @ 1×512B ← core 1 的 magic
│ └── ...
├── Barrier PingPong Buffer 0 (epWorldSize × 32B)
│ ├── slot[srcRank=0] ← rank 0 写入的 barrier flag
│ ├── slot[srcRank=1]
│ └── ...
├── Barrier PingPong Buffer 1 (同上)
├── Flag PingPong Buffer 0 (epWorldSize × slotsPerRank × 32B)
│ ├── slot[srcRank=0][eventID=0]
│ ├── slot[srcRank=0][eventID=1]
│ └── ...
└── Flag PingPong Buffer 1 (同上)
4.4 同步协议详解
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 的数据保证可见
DAG 保证（无环依赖）:
DataWrite(A→B) ── MTE3 fence ──► SetFlag(A→B)
│
▼
WaitFlag(B←A) ──► DataRead(B←A)
4.5 lowlatency 的自定义同步
lowlatency_zb 没有使用 ZeroBufferSyncFlag 框架，而是实现了自己的同步：
// 同步区域: metaDataGvaGM_ + COMBINE_SYNC_FLAG_BASE_OFFSET (100KB)
// 每个 rank 占 COMBINE_SYNC_STATE_OFFSET (32B)
// SetSyncFlag: 向远端 rank 的 flag 区域写入 1.0f
GetSyncFlagAddrByRankId(rankId) = shmem_ptr(metaDataGvaGM_, rankId) + 100KB + epRankId_ * 32B
// WaitSyncFlag: 轮询本 rank 的 flag 区域，直到所有远端 rank 都写入
// 使用 ReduceSum 求和，期望值 = rankNumPerBlock_
与 normal_zb 的区别:
- normal_zb 使用 magic + phase 的 64-bit flag，支持跨 kernel 同步
- lowlatency_zb 使用固定值 1.0f，仅支持单次 kernel 内同步
- lowlatency_zb 在等待完成后会清理 flag（写回 0），为下次同步做准备
5. 与 master 版 moe_combine_normal 的对比
5.1 架构差异总览
| 维度 | master (HCCL Window) | cam_feature (Zero-Buffer) |
|---|---|---|
| 数据搬运 | CopyBufferToShare → WaitBuffCopy | 直接 SHMEM 读取 |
| 同步机制 | polling WaitBuffCopy (per-token per-topk) |
ZeroBufferSyncFlag (per-rank) |
| Magic 管理 | GetHcclContext + localWindowsExp |
IncrementMagic() + per-core slot |
| Window 访问 | epWinContext_->localWindowsIn / remoteRes |
shmem_ptr(baseGM, rankId) |
| 状态管理 | per-token per-topk 的 32B status slot | per-rank 的 64B flag slot |
| 数据源 | tokenSrcInfo (3-tuple: rankId, tokenIdx, topkIdx) |
sendTokenIdx + epRecvCount |
| Copy 方向 | 本地 → 远端 (CopyBufferToShare) | 远端 → 本地 (DataCopyPad from SHMEM) |
5.2 master 版执行流程
// master: moe_combine_normal.cpp
Process():
├── CopyBufferToShareAndSetStatus() // 将本地 token 拷贝到远端 window
│ └── for each token:
│ ├── CopyBufferToShare() // DataCopyPad 到远端 SHMEM
│ └── SetStatusBySrcInfo() // 写 status flag
└── ReadBufferFromRemote() // 读取远端数据
└── for each token:
├── WaitBuffCopy() // 轮询等待所有源 rank 数据到位
└── ReadBufferAndWeightedSum() // 加权求和
5.3 ZB 版执行流程
// cam_feature: moe_combine_normal_zero_buffer.cpp
Process():
├── syncFlag_.IncrementMagic() // 获取 magic
├── syncFlag_.BarrierAll() // 全 rank 屏障
├── ReadBufferFromRemote() // 直接读取远端 SHMEM
│ └── for each token:
│ └── ReadBufferAndWeightedSum() // 直接从远端 SHMEM 读取 + 加权
└── SyncAll<true>()
5.4 关键差异分析
5.4.1 数据流方向
master: 推模式 (Push)
Rank A: CopyBufferToShare → 写入 Rank B 的 window
Rank B: WaitBuffCopy → 轮询等待 → 从本地 window 读取
ZB: 拉模式 (Pull)
Rank A: (dispatch 阶段已写入 SHMEM)
Rank B: shmem_ptr(recvXGM_, rankA) → 直接读取 Rank A 的 SHMEM
优势: ZB 消除了 CopyBufferToShare 阶段，Combine 不需要先将本地数据推送到远端。
5.4.2 同步粒度
master: per-token per-topk 粒度
// 每个 (token, topk) 对有一个 32B status slot
GM_ADDR stateGM = GetStateAddrByRankId(srcRankId) + (srcTokenId * axisK_ + srcTopkId) * UB_32_ALIGN;
// 轮询直到所有 K 个 source 都到位
while (current != target) {
DataCopy(stateTensorLocal, stateGMTensor, calCount);
Sum(tempStateTensorLocal, stateTensorLocal, sumPerKParams);
current = tempStateTensorLocal(0);
}
ZB: per-rank 粒度
// 每个 rank 有一个 64B flag slot
// BarrierAll 一次等待所有 rank
syncFlag_.BarrierAll();
优势: ZB 的同步粒度更粗，减少了 polling 次数和 flag 写入次数。
5.4.3 Magic 管理
master: 使用 HCCL context 中的 magic
auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
epWinContext_ = (__gm__ HcclOpResParam *)contextGM0;
selfMagicTensor.SetGlobalBuffer(epWinContext_->localWindowsExp + MAGIC_WIN_OFFSET + coreIdx_ * WIN_512_ALIGN);
magic_ = selfMagicTensor(0); // 0 或 1 翻转
selfMagicTensor(0) = ((magic_ == 0) ? 1 : 0);
ZB: 使用 ZeroBufferSyncFlag 的递增 magic
magic_ = static_cast<int32_t>(magicGT(0));
magicGT(0) = magic_ + 1; // 单调递增
优势: ZB 的 magic 单调递增，可以区分任意历史 kernel 的 flag，不会产生 ABA 问题。
5.4.4 内存占用
master:
- 需要 tokenSrcInfo 数组 (bs × k × 12B)
- 需要 per-token per-topk 的 status slot (bs × k × 32B)
- 需要 ping-pong 双缓冲
ZB:
- 需要 sendTokenIdx 数组 (bs × k × 4B)
- 需要 epRecvCount 表 (moeExpertNum × epWorldSize × 4B)
- 只需要 per-rank 的 flag slot
5.5 性能提升预估
| 优化点 | 预估收益 |
|---|---|
| 消除 CopyBufferToShare | 减少一次 GM→GM DMA，约节省 bs × k × h × sizeof(bf16) 的写入延迟 |
| 粗粒度同步 | 从 O(bs × K) 次 polling 降至 O(1) 次 BarrierAll |
| 消除 WaitBuffCopy | 消除 per-token 的 polling 循环（可能涉及数百次循环迭代） |
| 减少 flag 写入 | 从 bs×K 次 flag 写入降至 epWorldSize 次 |
6. Tiling 策略
6.1 normal_zb Tiling
核心流程 (MoeCombineNormalZeroBufferA3TilingFuncImpl):
// 1. 获取属性
GetAttrAndSetTilingData(context, *tilingData, nodeName);
// → 读取 meta_data_ptr, ep_world_size, ep_rank_id, moe_expert_num, global_bs
// 2. 校验输入输出
TilingCheckMoeCombineNormalZeroBuffer(context, nodeName, isEnableDiagnose);
// → 检查 dim/format/dataType
// 3. 校验属性
CheckAttrs(context, *tilingData, nodeName, localMoeExpertNum);
// → moeExpertNum % epWorldSize == 0
// → tp=2 时 localMoeExpertNum == 1
// → globalBs == 0 || globalBs >= bs * epWorldSize
// 4. 校验 shape
CheckTensorShape(context, *tilingData, nodeName, localMoeExpertNum);
// → k ∈ (0, min(16, moeExpertNum)]
// → h ∈ [1024, 7168]
// 5. 校验 win 区大小
tokenNeedSizeCombine = Ceil(h * 2 + 511, 512) * 512; // 512 对齐
actualSize = (maxBs * k * tokenNeedSizeCombine + COMBINE_STATE_WIN_OFFSET + NOTIFY_DISPATCH_WIN_OFFSET) * 2;
// 6. 设置 blockDim 和 schedule mode
blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, 0, aivNum);
context->SetScheduleMode(1); // batch mode，所有核同时启动
约束条件:
- epWorldSize ∈ [2, 384]
- tpWorldSize ∈ [0, 2]
- moeExpertNum ∈ (0, 512]
- k ∈ (0, min(16, moeExpertNum)]
- h ∈ [1024, 7168]
- (moeExpertNum / epWorldSize) × epWorldSize ≤ 1280
6.2 lowlatency_zb Tiling
核心流程 (MoeCombineLowlatencyA3TilingFuncImpl):
// 1. 获取属性（更复杂）
GetAttrAndSetTilingData(context, *tilingData, nodeName, groupEp, groupTp, commQuantMode);
// → 额外读取 sharedExpertNum, sharedExpertRankNum, zeroExpertNum, copyExpertNum, constExpertNum
// 2. 计算 A（expandX 的 dim0）
if (isShared) {
A = maxBs * maxSharedGroupNum; // 共享专家卡
} else {
A = globalBs * min(localMoeExpertNum, expertIdsDim1); // MoE 专家卡
}
// 3. 校验 win 区大小
tokenNeedSizeCombine = Ceil(h * 2 + 511, 512) * 512;
tokenActualLen = Ceil(h * 2 + 31, 32) * 32 + 44; // + scale(32B) + expandIdx(3*4B)
tokenNeedSizeDispatch = Ceil(tokenActualLen + 511, 512) * 512;
actualSize = ((maxBs * tokenNeedSizeDispatch * epWorldSize * localMoeExpertNum) +
(maxBs * tokenNeedSizeCombine * (k + sharedExpertNum))) * 2;
// 4. 计算 TilingKey
tilingKey = 10000;
if (tpWorldSize == 2) tilingKey += 100;
if (commQuantMode == 2) tilingKey += 20;
// 5. 设置 Mc2CcTiling（HCCL 通信配置）
SetHCommCfg(context, tilingData, groupEp, groupTp);
// → AlltoAll: level0:fullmesh;level1:pairwise
// → ReduceScatter: level0:ring
额外约束:
- epWorldSize ∈ [2, 768]（支持更大规模）
- moeExpertNum ∈ (0, 1024]
- h ∈ [1024, 8192]
- bs ∈ [1, 512]
- sharedExpertNum ∈ [0, 4]
- moeExpertNum % (epWorldSize - sharedExpertRankNum) == 0
- zeroBufferPtr != 0
6.3 Win 区大小计算对比
normal_zb:
actualSize = (maxBs × k × tokenNeedSizeCombine + COMBINE_STATE_WIN_OFFSET + NOTIFY_DISPATCH_WIN_OFFSET) × 2
= (maxBs × k × Ceil(h×2, 512) + 5MB + 204KB) × 2
lowlatency_zb:
actualSize = ((maxBs × tokenNeedSizeDispatch × epWorldSize × localMoeExpertNum) +
(maxBs × tokenNeedSizeCombine × (k + sharedExpertNum))) × 2
差异: lowlatency 需要同时容纳 dispatch 数据区和 combine 数据区，而 normal 只需要 combine 数据区（dispatch 数据区在 notify 算子中管理）。
6.4 Schedule Mode
两个算子都使用 SetScheduleMode(1) (batch mode):
- 所有 AIV 核同时启动
- 核间通过 SyncAll<true>() 同步
- 不需要核间通信，完全并行
7. 关键数据结构与内存布局
7.1 Tiling 数据结构
normal_zb:
struct MoeCombineNormalZeroBufferInfo {
uint32_t epWorldSize; // EP 通信域大小
uint32_t tpWorldSize; // TP 通信域大小 (0 or 1 or 2)
uint32_t epRankId; // 本 rank ID
uint32_t tpRankId; // 本 TP rank ID
uint32_t expertShardType; // 专家分片类型
uint32_t moeExpertNum; // MoE 专家总数
uint32_t moeExpertPerRankNum; // 每 rank 专家数
uint32_t globalBs; // 全局 batch size
uint32_t bs; // 本 rank batch size
uint32_t k; // top-K
uint32_t h; // hidden size
uint32_t aivNum; // AIV 核数
uint64_t totalUbSize; // UB 总大小
uint64_t totalWinSize; // Window 总大小
float armAvgFactor; // ARM 平均因子
float epsilon; // 数值精度
bool isEnableDiagnose; // 是否启用诊断
bool isGetProb; // 是否获取概率梯度
};
struct MoeCombineNormalZeroBufferTilingData {
MoeCombineNormalZeroBufferInfo moeCombineNormalInfo;
uint64_t zeroBufferPtr; // Zero-Buffer 对称地址
};
lowlatency_zb:
struct MoeCombineLowlatencyZeroBufferInfo {
// ... (同上，额外字段)
uint32_t sharedExpertNum; // 共享专家数
uint32_t sharedExpertRankNum; // 共享专家所在 rank 数
uint32_t zeroExpertNum; // 零专家数
uint32_t copyExpertNum; // 拷贝专家数
uint32_t constExpertNum; // 常量专家数
bool isTokenMask; // token 级 mask
bool isExpertMask; // expert 级 mask
bool hasSharedExpertX; // 是否有共享专家输入
bool hasElasticInfo; // 是否有弹性信息
uint64_t zeroBufferPtr; // Zero-Buffer 地址
};
struct MoeCombineLowlatencyZeroBufferTilingData {
Mc2InitTiling mc2InitTiling; // MC2 初始化 tiling
Mc2CcTiling mc2CcTiling1; // MC2 CC tiling 1 (AlltoAll)
Mc2CcTiling mc2CcTiling2; // MC2 CC tiling 2 (ReduceScatter)
MoeCombineLowlatencyZeroBufferInfo moeCombineLowlatencyZeroBufferInfo;
};
7.2 常量定义
// normal_zb 内存布局常量
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 = 8UL * 1024UL * 1024UL; // 8MB
constexpr uint64_t COMBINE_WIN_STATE_OFFSET = 5UL * 1024UL * 1024UL; // 5MB
constexpr uint64_t WIN_ADDR_ALIGN = 512UL;
// lowlatency_zb 同步常量
constexpr uint64_t COMBINE_SYNC_FLAG_BASE_OFFSET = 100UL * 1024UL; // 100KB
constexpr uint64_t COMBINE_SYNC_STATE_OFFSET = 32U; // 32B per rank
// 通用常量
constexpr uint32_t UB_32_ALIGN = 32U;
constexpr uint32_t MUL_256_ALIGN = 256U;
constexpr uint8_t DOUBLE_BUFFER = 2;
7.3 类继承/依赖关系
MoeCombineNormalZeroBuffer<RecvXType, XType, SrcInfoType>
├── 使用 ZeroBufferSyncFlag (通过 TBuf)
├── 使用 shmem_ptr() 访问远端 SHMEM
└── 使用 TQue/TBuf 管理 UB
MoeCombineLowlatencyZeroBuffer<ExpandXType, XType, ExpandIdxType, IsNeedReduceScatter, IsInt8Quant>
├── 自实现 SetSyncFlag/WaitSyncFlag
├── 使用 shmem_ptr() 访问远端 SHMEM
├── 使用 moe_lowlatency_zero_buffer_base.h 工具
└── 支持 TP=2, 共享专家, 弹性缩容, 常量专家
附录：关键函数索引
normal_zb
| 函数 | 文件 | 行号 | 说明 |
|---|---|---|---|
moe_combine_normal_zero_buffer |
cpp:18 | kernel 入口 | |
Init() |
h:108 | 初始化所有 buffer 和地址 | |
Process() |
h:175 | 主执行流程 | |
IncrementMagic() |
h:117 | 获取并递增 magic | |
ReadBufferFromRemote() |
h:143 | 从远端 SHMEM 读取数据 | |
ReadBufferAndWeightedSum() |
h:124 | 加权求和核心逻辑 | |
SplitCoreCal() |
h:82 | 均匀分配 token 到各 core | |
InitMagic() |
h:113 | 读取并翻转 magic |
lowlatency_zb
| 函数 | 文件 | 行号 | 说明 |
|---|---|---|---|
moe_combine_lowlatency_zero_buffer |
cpp:37 | kernel 入口 | |
Init() |
h:135 | 初始化 | |
Process() |
h:286 | 主执行流程 | |
CopyValidExpandXToRemote() |
h:214 | 拷贝本地数据到 SHMEM | |
SetSyncFlag() |
h:172 | 写同步 flag | |
WaitSyncFlag() |
h:193 | 轮询等待 flag | |
InputToDstOutput() |
h:241 | 从 SHMEM 读取并加权求和 | |
SplitToCore() |
h:155 | 分配 token 到各 core |
工具类
| 类/函数 | 文件 | 说明 |
|---|---|---|
ZeroBufferSyncFlag |
zero_buffer_sync_flag.h | 通用跨卡 flag 同步 |
IncrementMagic() |
zero_buffer_sync_flag.h:103 | magic 管理 |
SetFlag() |
zero_buffer_sync_flag.h:127 | 写 flag |
WaitFlag() |
zero_buffer_sync_flag.h:153 | 轮询等待 flag |
BarrierAll() |
zero_buffer_sync_flag.h:238 | 全 rank 屏障 |
SyncCollectives |
sync_collectives_zb.h | 底层同步原语 |
shmem_ptr() |
shmem.h (CANN) | 获取远端 SHMEM 地址 |
总结
CAM 库的 ZB (Zero-Buffer) Combine 算子代表了 MoE 通信的一次范式转变：
- 从推模式到拉模式: 消除了 Combine 阶段的
CopyBufferToShare，改为直接从远端 SHMEM 读取 - 从细粒度到粗粒度同步: 从 per-token per-topk 的 polling 升级为 per-rank 的 BarrierAll
- 从固定 magic 到递增 magic: 解决了 ABA 问题，支持跨 kernel 同步
- 从单一场景到多样化: lowlatency 版支持 TP=2、共享专家、弹性缩容、常量专家等高级特性
这些优化使得 Combine 阶段的延迟显著降低，尤其在大规模 EP 场景下（epWorldSize 达到 384/768），粗粒度同步的优势更加明显。
CAM 皇冠上的明珠 — 三变体架构全景解剖 分支:
remotes/origin/cam_feature| 分析日期: 2026-07-21
目录
- 家族概览与设计哲学
- 三变体架构对比总览
- Base 变体 (
fused_deep_moe) — 完整流程深度分析 - W4A8 变体 (
fused_deep_moe_w4a8) — INT4 量化融合 - FWK 变体 (
fused_deep_moe_fwk) — 框架原生范式 - 三条执行路径的 cam_feature 演进
- raw_distributed 子模块 — 通信原语解剖
- GEMM/Epilogue 子模块 — 计算内核
- SHMEM 集成机制
- 关键数据结构与内存布局
1. 家族概览与设计哲学
FusedDeepMoE 是 CAM (Communication Acceleration for Matrix) 中最复杂的算子家族，将 MoE Decode 阶段的 Dispatch → GMM1 → SwiGLU → GMM2 → Combine 五阶段完整流水线融合为单一 kernel launch。
核心动机：消除中间数据的 Global Memory 往返。 标准小算子组合中，中间数据 X1、SwigluOut、X2 需要至少 8 次 GM 访问；融合后全部在片上 workspace 中流转。
三个变体服务不同的部署场景：
| 维度 | Base (fused_deep_moe) |
W4A8 (fused_deep_moe_w4a8) |
FWK (fused_deep_moe_fwk) |
|---|---|---|---|
| 定位 | 通用 EP 融合 MoE | INT4/INT8 混合量化 | 框架原生 + OPX 流水线 |
| Weight 类型 | INT8 | INT4 (int4b_t) | INT8 |
| requires_shmem | true | false | false |
| 创建时间 | 2025-07-19 | 2026-06-29 | 2026-01-20 |
| 共享专家 | 可选 (share_gmm1/2_weight) | 可选 + Bias 支持 | 原生 sharedExpertNum/RankNum |
| 模板参数 | W1ScaleType, W2ScaleType |
同 Base | 仅 ExpandXType |
| GMM1_L1K/L0K | 512 / 128 | 1024 / 256 | 512 / 128 |
| Bias 输入 | 无 | gmm1/2_bias + share_bias | 无 |
| 通信机制 | SHMEM + HCCL Window | HCCL Window (MC2) | HCCL Window + OPX LibOps |
2. 三变体架构对比总览
2.1 统一执行模型
三个变体共享相同的五阶段融合架构，但通信与量化路径不同：
┌─────────────────────────────────────────────────────────────────────┐
│ FusedDeepMoE 统一执行流程 │
│ │
│ ┌──────────┐ ┌──────────┐ ┌─────────┐ ┌────────┐ ┌──────┐ │
│ │ Dispatch │──▶│ GMM1 │──▶│ SwiGLU │──▶│ GMM2 │──▶│Combine│ │
│ │ (AIV) │ │ (AIC+AIV)│ │(Epilogue│ │(AIC+ │ │(AIV) │ │
│ │ │ │ │ │ fused) │ │ AIV) │ │ │ │
│ └──────────┘ └──────────┘ └─────────┘ └────────┘ └──────┘ │
│ │ │ │ │ │ │
│ AlltoAll PerTokenDeq PerTokenDeq PerTokenDeq AlltoAll │
│ Send + Swiglu + Quant + Combine Recv │
│ + Quant + Quant │
│ │
│ ─────── 中间数据全部在片上 workspace 流转，不访问 GM ─────── │
└─────────────────────────────────────────────────────────────────────┘
2.2 模板体系对比
Base 变体 (fused_deep_moe.h):
// 完整的 6 参数模板，支持灵活的 Scale 类型
template <typename ExpandXType, typename W1ScaleType, typename W2ScaleType,
typename ExpandIdxType, bool IsNeedReduceScatter, uint32_t EXEC_FLAG>
class FusedDeepMoe { ... };
// 实例化: FusedDeepMoe<DTYPE_X, DTYPE_GMM1_WEIGHT_SCALE, DTYPE_GMM2_WEIGHT_SCALE, int32_t, false, TILING_KEY_VAR>
FWK 变体 (fused_deep_moe_fwk.h):
// 简化模板：去掉 Scale 类型参数，EXEC_FLAG 提升为第一参数
template <typename ExpandXType, typename ExpandIdxType,
bool IsNeedReduceScatter, uint32_t EXEC_FLAG>
class FusedDeepMoeFwk { ... };
// 实例化: FusedDeepMoeFwk<DTYPE_X, int32_t, false, TILING_KEY_VAR>
W4A8 变体: 同 Base 的 6 参数模板，但内部类型使用 AscendC::int4b_t。
2.3 EXEC_FLAG 位定义
三个变体共享大部分标志位，但存在关键差异：
| 标志位 | Base | W4A8 | FWK | 含义 |
|---|---|---|---|---|
EXEC_FLAG_DEEP_FUSE (bit 0) |
✅ | ✅ | ✅ | 深度融合（Dispatch 在 GMM Kernel 内） |
EXEC_FLAG_TENSOR_LIST (bit 1) |
✅ | ❌ | ✅ | 权重为 TensorList |
EXEC_FLAG_X_ACTIVE_MASK (bit 2) |
✅ | ❌ | ❌ | Token 激活掩码 |
EXEC_FLAG_SHARED_EXPERT (bit 3) |
✅ | ✅ | ❌ | 共享专家 |
EXEC_FLAG_SMOOTH_QUANT (bit 4) |
✅ | ❌ | ❌ | 平滑量化 |
EXEC_FLAG_ZERO_BUFFER (bit 5) |
✅ | ❌ | ❌ | 零缓冲区（Base 独有） |
关键差异: Base 有
EXEC_FLAG_ZERO_BUFFER（SHMEM 双缓冲清零），W4A8/FWK 无此标志。FWK 无SHARED_EXPERT（共享专家通过sharedExpertNum原生集成）。
3. Base 变体 (fused_deep_moe) — 完整流程深度分析
源文件位置: src/cam/comm_operator/ascend_kernels/fused_deep_moe/
3.1 文件结构
fused_deep_moe/
├── op_kernel/
│ ├── fused_deep_moe.{h,cpp} ← 主 kernel 入口 + FusedDeepMoe 类
│ ├── fused_deep_moe_base.h ← 模板宏定义 + 通信组结构
│ ├── fused_deep_moe_tiling.{h,cpp} ← TilingData 定义 + 编译期常量
│ ├── fused_deep_moe/ ← 子模块目录
│ │ ├── gemm/ ← GEMM block_mmad, kernel
│ │ ├── epilogue/ ← Epilogue block, tile, policy
│ │ └── raw_distributed/ ← dispatch/combine 通信原语
│ └── cam_moe_distribute_*.h ← 兼容性别名（实际指向子模块）
├── op_host/
│ └── fused_deep_moe_infer.cpp ← 仅 InferShape/InferDataType
└── op_api/
└── aclnn_fused_deep_moe.{h,cpp} ← aclnn 接口封装
注意: Base 变体的 host 侧只有
infer.cpp（无 tiling.cpp），说明其 tiling 参数全部在 kernel 侧通过 constexpr 计算，或依赖 MC2 框架的默认 tiling。
3.2 主 Kernel 入口
// fused_deep_moe.cpp
extern "C" __global__ __aicore__ void fused_deep_moe(...) {
icache_preload(8);
REGISTER_TILING_DEFAULT(FusedDeepMoeTilingData);
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); // 1C2V 混合调度
GET_TILING_DATA(tiling_data, tiling);
// 64 个 tiling key 分支（编译期展开）
CALL_FUSED_DEEP_MOE_IF_TILINGKEY(0)
CALL_FUSED_DEEP_MOE_ELIF_TILINGKEY(1)
// ... 直至 key 63
}
关键设计: 64 个 tiling key 通过宏展开为编译期 if constexpr 链，每个 key 对应不同的 EXEC_FLAG 组合。这允许编译器为每种配置生成特化代码。
3.3 FusedDeepMoe::Init — 参数初始化
Init 方法（fused_deep_moe.h L352-412）完成：
- GM 地址绑定: 所有输入/输出/权重的 Global Memory 指针
- 维度推导:
-
maxTokenNum_ = maxBs * epRankSize * min(topK, moeExpertNumPerRank)— 单 rank 需处理的最大 token 数 -gmm2InputDim_ = gmm1OutputDim_ / 2— SwiGLU 将隐藏维度减半 -shareGmm2InputDim_ = shareGmm1OutputDim_ / 2 - SHMEM workspace 指针:
shmemWorkspaceGM_从 tiling 数据的shmemWorkspacePtr获取 - statusDataSpace 偏移计算: 用于 cross-rank 状态同步的多缓冲区布局
3.4 FusedDeepMoe::Process — 五阶段融合主流程
Process() 是整个算子的核心（fused_deep_moe.h L414-558），分两大阶段：
阶段 A: Problem Shape 与 Workspace 布局
GemmCoord gmm1ProblemShape{maxTokenNum_, gmm1OutputDim_, tokenHiddenSize_}; // [M, N, K]
GemmCoord gmm2ProblemShape{maxTokenNum_, gmm2OutputDim_, gmm2InputDim_};
Workspace 内存布局（全部在 GM 上，但中间数据通过 SHMEM 在片上流转）：
shmemWorkspaceGM_ (Base 独有，SHMEM 双缓冲):
├── gmShareX1 / gmX1 — GMM1 输入 token (int8)
├── gmShareX2 / gmX2 — GMM2 输入 token (int8)
├── gmTokenFlag — token 到达标志
├── gmAllExpertTokenNums — 全局专家 token 计数
└── gmCombineSend — Combine 发送缓冲
workspaceGM_ (普通 GM workspace):
├── gmCVSwap — Cube 计算 workspace (blockDim * L1M*L1N * STAGES * int32)
├── gmShareSwigluOut/gmSwigluOut — SwiGLU 输出 (float)
├── gmGroupList — 专家分组列表
├── gmExpandIdx — 展开索引
├── gmEpSendCount — EP 发送计数
├── gmResvered — 保留空间 (256KB)
└── gmAllEpRecvCount — 全局 EP 接收计数
阶段 B: 非深度融合路径 (Shallow Dispatch)
当 EXEC_FLAG & EXEC_FLAG_DEEP_FUSE == 0 时，执行独立的 Dispatch 阶段：
if constexpr ((EXEC_FLAG & EXEC_FLAG_DEEP_FUSE) == 0) {
if constexpr (g_coreType == AscendC::AIV) {
// AIV 上执行 Dispatch：量化 + AlltoAll 发送
CamMoeDistributeDispatch<...> dispatcher;
dispatcher.Init(gmX_, gmexpertIds_, ...);
dispatcher.Process();
icache_preload(8); // 预取 GMM 指令到 icache
}
// AIV 完成 → CrossCoreFlag 通知 AIC
PipeBarrier<PIPE_ALL>();
Arch::CrossCoreFlag gmm1AivFinished{0};
if constexpr (g_coreType == AscendC::AIV) {
Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(gmm1AivFinished);
} else {
Arch::CrossCoreWaitFlag(gmm1AivFinished);
}
}
阶段 C: GMM1 + SwiGLU + Quant 融合
GmmDeqSwigluQuant<TemplateMC2TypeFunc, Gmm1L1TileShape, Gmm1L0TileShape, ...>(
gmm1ProblemShape, groupCount_, gmGroupList, gmX1, layoutX1, ...);
这是第一个 GMM kernel，在 AIC+AIV 上执行，内部完成：
- GMM1: INT8 × INT8 → INT32 (GroupedMatmul)
- PerTokenDequant: INT32 × weightScale × perTokenScale → float
- SwiGLU: x * sigmoid(x) * x_right (half-glu)
- BlockQuant: float → INT8 + dequantScale (为 GMM2 准备输入)
阶段 D: GMM2 + Combine 融合
// 先初始化 Combiner（AIV 上）
CamMoeDistributeCombine<...> combiner;
combiner.Init(gmGmm2DepOut, gmexpertIds_, gmExpandIdx, ...);
// 再执行 GMM2 kernel（AIC+AIV），Combiner 作为回调注入
GmmDeq<TemplateMC2TypeFunc, Gmm2L1TileShape, ...>(
gmm2ProblemShape, ..., &combiner);
GMM2 kernel 的 Epilogue 在完成 L0C→GM 写入后，直接调用 combiner.AllToAllSend() 将结果通过 HCCL Window 发送回源 rank。
3.5 GmmDeqSwigluQuant — 类型系统
// GMM1 的 GEMM 类型定义
using AType = Gemm::GemmType<int8_t, layout::RowMajor>; // 激活 (量化后)
using BType = Gemm::GemmType<int8_t, layout::zN>; // 权重 (zN 布局)
using CType = Gemm::GemmType<int32_t, layout::RowMajor>; // 累加器
// Epilogue 类型
using ScaleType = Gemm::GemmType<W1ScaleType, layout::VectorLayout>; // 权重 scale
using PerTokenScaleType = Gemm::GemmType<float, layout::VectorLayout>; // per-token scale
using DType = Gemm::GemmType<float, layout::RowMajor>; // SwiGLU 输出
Kernel 选择逻辑（fused_deep_moe.h L113-118）:
using GemmKernel = typename std::conditional<
(EXEC_FLAG & EXEC_FLAG_DEEP_FUSE),
// 深度融合：Dispatch 在 GMM kernel 内部完成（AIV 先做量化+发送，AIC 同时计算）
Gemm::Kernel::GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace<...>,
// 非深度融合：独立的 Dispatch + GMM
Gemm::Kernel::GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspaceWithShallowDispatch<...>
>::type;
3.6 Tiling 常量
// GMM1 Tile Shape
constexpr uint32_t GMM1_L1M = 256; // L1 tile M 维
constexpr uint32_t GMM1_L1N = 128; // L1 tile N 维
constexpr uint32_t GMM1_L1K = 512; // L1 tile K 维
constexpr uint32_t GMM1_L0K = 128; // L0 tile K 维 (L1→L0 切分)
constexpr uint32_t GMM1_EPIM = 64; // Epilogue interleaving M
// GMM2 Tile Shape
constexpr uint32_t GMM2_L1M = 128;
constexpr uint32_t GMM2_L1N = 256;
constexpr uint32_t GMM2_L1K = 512;
constexpr uint32_t GMM2_L0K = 128;
constexpr uint32_t GMM2_EPIM = 32;
// Pipeline 深度
constexpr uint32_t WORKSPACE_STAGES = 4; // workspace 4 级流水线
constexpr uint32_t CUSTOM_PRELOAD_STAGES = 1; // preload 1 级
constexpr uint32_t CUSTOM_L1_STAGES = 2; // L1 双缓冲
constexpr uint32_t CUSTOM_L0A_STAGES = 2; // L0A 双缓冲
constexpr uint32_t CUSTOM_L0B_STAGES = 2; // L0B 双缓冲
constexpr uint32_t GMM2_L1A_STAGES = 4; // GMM2 L1A 4 级 (ResidentA)
constexpr uint32_t GMM2_L0A_STAGES = 4; // GMM2 L0A 4 级
4. W4A8 变体 (fused_deep_moe_w4a8) — INT4 量化融合
源文件位置: src/cam/comm_operator/ascend_kernels/fused_deep_moe_w4a8/
4.1 设计定位
W4A8 是面向 W4A8 (Weight-INT4, Activation-INT8) 量化 MoE 模型的专用变体。与 Base 的核心差异：
- 权重类型从 INT8 降为 INT4 (
AscendC::int4b_t)，权重存储量减半 - 新增 Bias 输入 (gmm1_bias, gmm2_bias, share_gmm1_bias, share_gmm2_bias)
- 更大的 L1/L0 tile (L1K=1024, L0K=256) — 补偿 INT4 计算密度的降低
- 独立的 host tiling.cpp — 更复杂的参数校验和 workspace 计算
- 无 SHMEM 依赖 (
requires_shmem: false)
4.2 输入签名差异
W4A8 的 kernel 入口比 Base 多 4 个 Bias 输入：
extern "C" __global__ __aicore__ void fused_deep_moe(
// ... 同 Base 的 14 个输入 ...
GM_ADDR gmm1_bias, GM_ADDR gmm2_bias, // ← 新增
GM_ADDR share_gmm1_bias, GM_ADDR share_gmm2_bias, // ← 新增
// ... 输出 + workspace + tiling ...
)
对应的 FusedDeepMoe::Init 也扩展了 Bias 地址绑定：
gmBias1_ = gmm1_bias;
gmBias2_ = gmm2_bias;
gmShareBias1_ = share_gmm1_bias;
gmShareBias2_ = share_gmm2_bias;
4.3 类型系统差异
// W4A8: 权重和激活均为 INT4
using AType = Gemm::GemmType<AscendC::int4b_t, layout::RowMajor>;
using BType = Gemm::GemmType<AscendC::int4b_t, layout::zN>;
using CType = Gemm::GemmType<int32_t, layout::RowMajor>;
// 布局使用 int4b_t
layout::zN layoutWeight1 = layout::zN::template MakeLayout<AscendC::int4b_t>(tokenHiddenSize_, gmm1OutputDim_);
4.4 block_mmad_preload_async_with_callback_w4a8.hpp — 核心优化
这是 W4A8 最独特的文件，实现了 INT4 的 2:1 压缩存储与计算补偿。
4.4.1 INT4 打包比
constexpr int CONSTANT_INT4_PACKING_RATIO = 2; // 2 个 INT4 打包进 1 个 INT8
// L1/L0 tile 大小计算（注意除以打包比）
static constexpr uint32_t L1A_TILE_SIZE =
L1TileShape::M * L1TileShape::K * sizeof(ElementA) / CONSTANT_INT4_PACKING_RATIO;
static constexpr uint32_t L1B_TILE_SIZE =
L1TileShape::N * L1TileShape::K * sizeof(ElementB) / CONSTANT_INT4_PACKING_RATIO;
含义: INT4 权重在内存中以 INT8 存储（2 个 INT4 打包），实际 L1 存储需求减半。这使得 W4A8 可以使用更大的 L1 tile (L1K=1024)。
4.4.2 与 Base 的 block_mmad 关键差异
| 特性 | Base (INT8) | W4A8 (INT4) |
|---|---|---|
| DispatchPolicy | MmadAtlasA2PreloadAsyncWithCallback |
MmadAtlasA2PreloadAsyncWithCallbackW4a8 |
| L1A/L1B 大小 | M*K*sizeof(int8) |
M*K*sizeof(int4)/2 |
| L0A/L0B 大小 | M*K*sizeof(int8) |
M*K*sizeof(int4)/2 |
| Bias 支持 | 无 | 有 (BiasType_ 模板参数) |
| 对齐检查 | 256b 对齐 | 256b 对齐（相同） |
4.4.3 计算流程
W4A8 的 BlockMmad::operator() 流程与 Base 一致：
1. GM→L1 预取: K 轴分 tile，L1 双缓冲
2. L1→L0 搬运: K 轴进一步切分到 L0 tile
3. TileMmad: 调用硬件 tileMmad 指令完成 INT4×INT4→INT32 乘累加
4. Callback: K 轴最后一轮时触发 callbackBeforeFixpipe / callbackAfterFixpipe（用于触发 Epilogue 的 SwiGLU+Quant）
4.4.4 更大的 Tile Shape
// W4A8 tiling.h
constexpr uint32_t GMM1_L1K = 1024; // Base 是 512
constexpr uint32_t GMM1_L0K = 256; // Base 是 128
constexpr uint32_t GMM2_L1K = 1024;
constexpr uint32_t GMM2_L0K = 256;
原因: INT4 权重存储减半 → L1 空间可容纳更大的 K 维 tile → 提高 Cube 引擎利用率。
4.5 Host Tiling 差异
W4A8 有独立的 fused_deep_moe_tiling.cpp（Base 没有），包含：
- INT4 权重维度校验:
gmm1HLen = shape * INT4_WEIGHT_HIDDEN_SCALE(×8 展开) - Bias 形状校验:
CheckGmm2BiasShape() - 更严格的参数范围:
MAX_BATCH_SIZE=256,MAX_MOE_EXERT_NUM=512,SUPPORT_TOP_K=12 - HCCL 缓冲区检查:
CheckHcclBufferSize()验证窗口大小是否足够 - Workspace 计算: 包含
CVSwapBufferSize(Cube workspace) 的精确计算
// W4A8 workspace 计算
size_t CVSwapBufferSize = CeilUp(aicNum * L1_TILE_BYTE_SIZE * CUBE_WORKSPACE_STAGE * sizeof(int32_t), GM_ALIGN_SIZE);
workSpaces[0] = SYSTEM_NEED_WORKSPACE + usrSize; // SYSTEM_NEED_WORKSPACE = 16MB
4.6 Tiling Key 生成
// W4A8 tiling.cpp
uint64_t tilingKey = 0;
if (xActiveMaskEnable) tilingKey |= EXEC_FLAG_X_ACTIVE_MASK;
if (calShareExpert) tilingKey |= EXEC_FLAG_SHARED_EXPERT;
if (moeExpertNumPerRank != 1) tilingKey |= EXEC_FLAG_DEEP_FUSE;
if (isTensorList) tilingKey |= EXEC_FLAG_TENSOR_LIST;
if (expertSmoothScalesExist) tilingKey |= EXEC_FLAG_SMOOTH_QUANT;
context.SetTilingKey(tilingKey);
context.SetBlockDim(aicNum);
5. FWK 变体 (fused_deep_moe_fwk) — 框架原生范式
源文件位置: src/cam/comm_operator/ascend_kernels/fused_deep_moe_fwk/
5.1 设计定位
FWK 是面向框架原生集成的变体，核心特征：
- OPX 框架集成: 通过
lib_ops.h提供 OPX 兼容的 Op 类 - 原生共享专家:
sharedExpertNum/sharedExpertRankNum作为一等公民 - 无 SHMEM 依赖: 纯 HCCL Window 通信
- 简化模板: 去掉 Scale 类型参数
- Permuted Weight: 输入权重名为
gmm1_permuted_weight（预置换布局） - HCCL Context 直连:
GetHcclContext<AscendC::HCCL_GROUP_ID_0>()
5.2 输入签名差异
// FWK kernel 入口 — 无 share_gmm1/2_weight，无 x_active_mask
extern "C" __global__ __aicore__ void fused_deep_moe_fwk(
GM_ADDR x, GM_ADDR expert_ids,
GM_ADDR gmm1_permuted_weight, GM_ADDR gmm1_permuted_weight_scale, // ← permuted
GM_ADDR gmm2_weight, GM_ADDR gmm2_weight_scale,
GM_ADDR expert_smooth_scales, GM_ADDR expert_scales, // ← 无 share_smooth
GM_ADDR output, GM_ADDR expertTokenNums, // ← 无 share_output
GM_ADDR workspace, GM_ADDR tiling)
关键差异:
- 无 share_output 输出（共享专家输出直接写回源 rank）
- 无 x_active_mask 输入
- 权重名为 permuted_weight（暗示框架侧已做权重置换）
5.3 共享专家的原生支持
FWK 的 FusedDeepMoeFwk::Init 中：
sharedExpertNum_ = tilingData->...sharedExpertNum;
sharedExpertRankNum_ = tilingData->...sharedExpertRankNum;
bool isShareExpert = (epRankId_ < sharedExpertRankNum_);
if (isShareExpert) {
m_ = maxBs_ * epRankSize_ / sharedExpertRankNum_; // 共享专家处理全部 token
} else {
m_ = maxBs_ * epRankSize_ * min(topK_, moeExpertNumPerRank_); // 路由专家
}
groupCount_ = isShareExpert ? 1 : moeExpertNumPerRank_; // 共享专家只有 1 组
含义: 共享专家 rank 处理所有 token 的 GMM1/GMM2（groupCount=1），路由专家 rank 处理分配给它的 token。
5.4 lib_ops.h — OPX 框架操作库
lib_ops.h 是 FWK 独有的核心文件，提供 OPX (Operator Extension) 框架兼容的操作类。
5.4.1 文件结构
lib_ops.h
├── 工具函数
│ ├── PipeSync<EVENT>() — 管道同步
│ └── GetBufferByByte<T>() — UB 缓冲区分配
├── LibOps namespace
│ ├── MoEContext — MoE 全局上下文（窗口地址、状态管理）
│ ├── OpDispatchSendToken — Dispatch 发送端（量化 + 窗口写入）
│ ├── OpDispatchRecvToken — Dispatch 接收端（窗口读取 + 反量化）
│ ├── OpDequant — 反量化操作（INT32 → float）
│ ├── OpSwiglu — SwiGLU 激活
│ ├── OpSyncAllCleanAndUpdate — 全局同步 + 状态清零
│ ├── OpBlockQuant — 分块量化（float → INT8）
│ ├── OpCombineSend — Combine 发送端（窗口写入）
│ ├── OpCombineSetFlag — Combine 标志设置
│ ├── OpCombineRecv — Combine 接收端（ReducePermute）
│ ├── GmScaleReader — GM 标量读取器
│ ├── GmPerTokenScaleReader — GM per-token scale 读取器
│ └── GmSwigluOutputWriter — SwiGLU 输出写入器
5.4.2 MoEContext — 全局上下文
struct MoEContext {
__gm__ HcclOpResParam *winContext; // HCCL 窗口上下文
GM_ADDR gmStatusDataSpace; // 状态数据空间
uint32_t epRankSize, epRankId;
uint32_t aicNum, aivNum;
uint32_t moeExpertNumPerRank;
int32_t dataState; // 双缓冲状态 (0/1)
int32_t tokenFlag; // token 到达标志 (0x55555555 / 0x33333333)
int32_t vToCFlag; // AIV→AIC 通信标志
float sumTarget; // ReduceSum 目标值（用于 flag 检测）
uint64_t expertPerSizeOnWin; // 单个专家在窗口中的空间大小
uint64_t winDataSizeOffset; // 窗口数据偏移（双缓冲）
};
窗口地址计算:
GM_ADDR GetWindowStateAddrByRankId(uint32_t rankId) {
if (epRankId == rankId) return (GM_ADDR)winContext->localWindowsExp;
return (GM_ADDR)((HcclRankRelationResV2 *)winContext->remoteRes[rankId].nextDevicePtr)->windowsExp
+ dataState * WIN_STATE_OFFSET; // 双缓冲偏移
}
GM_ADDR GetWindowDataAddrByRankId(uint32_t rankId) {
// ... 类似，但使用 localWindowsIn + winDataSizeOffset
}
5.4.3 OpDispatchSendToken — 发送端
template <typename XType, uint32_t BUFFER_NUM, bool ENABLE_X_ACTIVE_MASK>
class OpDispatchSendToken {
// 1. TokenActiveMaskCal: 计算有效 token 数（可选）
// 2. LoadExpertIds: 加载 expert_ids 到 UB
// 3. CalAndSendTokenCount: 计算每个专家的 token 数并写入远端状态空间
// 4. SendToMoeExpert: 量化 token 并写入远端窗口
// - DataCopy: GM → UB (读取 token)
// - QuantToken: Cast to float → Abs → ReduceMax → scale → Cast to int8
// - DataCopy: UB → GM (写入远端窗口)
};
QuantToken 实现（lib_ops.h L524-558）:
void QuantToken(...) {
Cast(xFp32TmpTensor, xInTensor, CAST_NONE, tokenLength); // bf16/fp16 → float
Abs(xFp32AbsTensor, xFp32TmpTensor, tokenLength);
ReduceMax(xRowMaxTensor, xFp32AbsTensor, ..., tokenLength, false);
float dynamicQuantScale = 127.0 / xRowMaxTensor.GetValue(0);
// 写入 scale 到 token 尾部的 extra space
yFp32Tensor.SetValue(tokenLength / sizeof(float), 1.0 / dynamicQuantScale);
yInt32Tensor.SetValue(tokenLength / sizeof(int32_t) + 1, ctx.tokenFlag); // 写入 flag
Muls(xFp32TmpTensor, xFp32TmpTensor, dynamicQuantScale, tokenLength);
Cast(ytmpInt32Tensor, xFp32TmpTensor, CAST_RINT, tokenLength); // float → int32
Cast(yHalfTensor, ytmpInt32Tensor, CAST_ROUND, tokenLength); // int32 → fp16
Cast(yInt8Tensor, yHalfTensor, CAST_TRUNC, tokenLength); // fp16 → int8
}
5.4.4 OpDispatchRecvToken — 接收端
class OpDispatchRecvToken {
// 1. RecvTokenCountInfo: 轮询状态空间直到所有源 rank 的 flag 到达
// - GatherMask: 从 512B 状态块中提取 flag 和 count
// - ReduceSum: 检测 flag 总和是否达到目标值
// 2. RecvToken: 从窗口读取 token 并写入本地 GM
// - 轮询 token flag 直到到达
// - DataCopy: GM → UB → GM (窗口 → 本地)
// 3. NotifyNextStage: AtomicAdd 通知 AIC 该组就绪
};
状态检测机制（lib_ops.h L641-695）:
// 使用 GatherMask 从 256B 状态块中提取 flag
// 每个 rank 的状态: [flag(0), count(1), pad(2), pad(3)] × 8 blocks = 256B
// src1Pattern=3: 每 4 个 int32 取第 1 个（flag）
// src1Pattern=4: 每 4 个 int32 取第 2 个（count）
do {
DataCopy(statusFp32Tensor, windowInStatusFp32Tensor, dataCopyParams);
GatherMask(gatherMaskOutFp32Tensor, statusFp32Tensor, src1Pattern=3, ...);
ReduceSum(gatherMaskOutFp32Tensor, ..., recStatusNumPerCore);
} while (!(minTarget < sumOfFlag && sumOfFlag < maxTarget));
5.4.5 OpDequant / OpSwiglu / OpBlockQuant — 计算 Op
这三个类是 OPX 框架的标准 Op 接口实现：
// OpDequant: INT32 → float (per-token dequant)
// 输入: cQue(int32), scaleQue(float), perTokenScaleQue(float)
// 输出: outQue(float)
// 计算: float = int32 * weightScale * perTokenScale
// OpSwiglu: SwiGLU 激活
// 输入: inQue(float)
// 输出: outQue(float)
// 计算: swiglu(x) = x[:, :N/2] * sigmoid(x[:, :N/2]) * x[:, N/2:]
// = x_left * (1 / (1 + exp(-x))) * x_right
// OpBlockQuant: float → INT8 (per-token block quant)
// 输入: inQue(float)
// 输出: dequantScaleQue(float), outQue(int8)
// 计算: abs → Max(左右两半) → scale = Max/127 → quant
5.4.6 OpCombineSend / OpCombineRecv — Combine 端
// OpCombineSend: 将 GMM2 输出写入远端窗口
// 使用 epSendCount 前缀和确定每个 rank 的 token 范围
// DataCopyPad: UB → GM (远端窗口)
// OpCombineRecv: 调用 combiner->ReducePermute()
// ReducePermute 完成 AlltoAll 接收 + 专家输出累加 + 写回输出
5.5 FWK 的 GmmDeqSwigluQuant 差异
FWK 的 GmmDeqSwigluQuant 函数签名与 Base 显著不同：
template <uint32_t EXEC_FLAG, typename XType_, ...>
CATLASS_DEVICE void GmmDeqSwigluQuant(
GemmCoord problemShape, uint32_t groupCount, GM_ADDR gmGroupList,
GM_ADDR gmA, layout::RowMajor layoutA,
GM_ADDR gmB, layout::zN layoutB,
GM_ADDR gmScale, layout::VectorLayout layoutScale,
GM_ADDR gmPerTokenScale, layout::VectorLayout layoutPerTokenScale,
GM_ADDR gmD, layout::RowMajor layoutD,
GM_ADDR gmDequantScale, layout::VectorLayout layoutDequantScale,
GM_ADDR gmWorkspace,
GM_ADDR gmX, GM_ADDR debugGm, // ← debugGm 替代 smoothScales
GM_ADDR gmexpertIds, GM_ADDR gmExpandIdx,
GM_ADDR gmEpSendCount, GM_ADDR gmResvered,
GM_ADDR gmExpertTokenNums,
uint32_t epRankSize, uint32_t epRankId, // ← 标量参数替代 tiling 引用
uint32_t moeExpertNum, uint32_t moeExpertNumPerRank,
uint32_t sharedExpertNum, uint32_t sharedExpertRankNum, // ← 新增
uint32_t quantMode, uint32_t globalBs, uint32_t bs,
uint32_t topK, uint32_t tokenLen)
关键差异:
- 无 gmShareX1/gmShareX1Scale/gmShareSwigluOut/gmShareX2/gmShareX2Scale（无独立共享专家 workspace）
- 无 gmCombineSend（Combine 由 OPX Op 处理）
- 新增 sharedExpertNum/sharedExpertRankNum 标量参数
- 新增 debugGm 参数
5.6 FWK 的 GmmDeq 差异
FWK 的 GmmDeq 使用 layout::nZ（而非 layout::zN）作为权重布局：
using BType = Gemm::GemmType<int8_t, layout::nZ>; // FWK: nZ
// vs Base:
using BType = Gemm::GemmType<int8_t, layout::zN>; // Base: zN
zN vs nZ: 这是 CATLASS 中两种不同的矩阵布局。zN 是"Z 形 N 维优先"，nZ 是"n 维 Z 形"。FWK 使用 nZ 暗示权重已经过不同的置换处理（与
gmm1_permuted_weight命名一致）。
6. 三条执行路径的 cam_feature 演进
6.1 路径 1: 非深度融合 (Shallow Dispatch)
AIV: Dispatch (量化 + AlltoAll 发送) ──────────────────────▶
AIC: 等待 ◀── CrossCoreFlag ── 等待结束 ──▶ GMM1+Swiglu+Quant ──▶ GMM2+Combine
触发条件: EXEC_FLAG & EXEC_FLAG_DEEP_FUSE == 0（即 moeExpertNumPerRank == 1）
特点:
- Dispatch 在 AIV 上独立执行
- AIV→AIC 通过 CrossCoreFlag 同步
- 适用于单专家 per rank 场景（通信量小）
6.2 路径 2: 深度融合 (Deep Fuse)
AIV: [Dispatch 量化+发送] ─┐
├─▶ 在 GMM Kernel 内部流水线化
AIC: [GMM1 计算] ──────────┘ [SwiGLU+Quant] ──▶ [GMM2+Combine]
触发条件: EXEC_FLAG & EXEC_FLAG_DEEP_FUSE != 0（即 moeExpertNumPerRank > 1）
特点:
- Dispatch 嵌入 GMM kernel 的 GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
- AIV 在 GMM1 计算前完成量化+发送，AIC 同时计算
- 通过 workspace 流水线（WORKSPACE_STAGES=4）实现计算-通信重叠
- 适用于多专家 per rank 场景（通信量大，需要隐藏延迟）
6.3 路径 3: 共享专家 (Shared Expert)
路由专家 Rank:
Dispatch → GMM1_route → SwiGLU → GMM2_route → Combine
共享专家 Rank (epRankId < sharedExpertRankNum):
GMM1_share (全部 token) → SwiGLU → GMM2_share → 直接写回
FWK 的实现:
bool isShareExpert = (epRankId_ < sharedExpertRankNum_);
if (isShareExpert) {
m_ = maxBs_ * epRankSize_ / sharedExpertRankNum_; // 处理全部 token
} else {
m_ = maxBs_ * epRankSize_ * min(topK_, moeExpertNumPerRank_);
}
groupCount_ = isShareExpert ? 1 : moeExpertNumPerRank_;
Base 的实现: 通过 EXEC_FLAG_SHARED_EXPERT 标志，在 workspace 中为共享专家分配独立区域（gmShareX1, gmShareX2 等），与路由专家并行计算。
6.4 三条路径的协同
在 cam_feature 上，三条路径可以组合： - 非深度融合 + 无共享专家: tilingKey bit0=0, bit3=0 - 深度融合 + 无共享专家: tilingKey bit0=1, bit3=0 - 深度融合 + 共享专家: tilingKey bit0=1, bit3=1
7. raw_distributed 子模块 — 通信原语解剖
7.1 CamMoeDistributeDispatch — 发送端
文件: fused_deep_moe/op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_dispatch.h
7.1.1 核心常量
constexpr uint32_t STATE_OFFSET = 512; // 每个 rank 的状态空间偏移
constexpr uint32_t STATE_SIZE = 1024 * 1024; // 1MB 状态空间
constexpr uint64_t WIN_STATE_OFFSET = 512 * 1024; // 窗口状态双缓冲偏移
constexpr uint64_t STATE_WIN_OFFSET = 900 * 1024; // 状态在窗口中的偏移
constexpr uint32_t SELF_STATE_OFFSET = 256 * 1024; // 自状态偏移
7.1.2 状态空间布局
窗口状态空间 (每个 rank 1MB):
├── [0, 512KB): 数据状态空间 (dataState=0)
│ ├── rank0: [0, 512B) — 每个 rank 512B 状态块
│ │ ├── [0, 256B): flag + count 数组 (64 个 int32)
│ │ └── [256B, 512B): 保留
│ ├── rank1: [512B, 1024B)
│ └── ...
├── [512KB, 1024KB): 数据状态空间 (dataState=1) — 双缓冲
├── [900KB, 932KB): 窗口状态空间 (STATE_WIN_OFFSET)
├── [932KB, 964KB): 组 token 计数 (GROUP_TOKEN_NUM_OFFSET)
└── [964KB, 988KB): 软同步空间 (SOFT_SYNC_OFFSET)
7.1.3 窗口地址计算
GM_ADDR GetWindAddrByRankId(uint8_t ctxIdx, const int32_t rankId) {
if (curRankId == rankId)
return (GM_ADDR)(winContext_[ctxIdx]->localWindowsIn) + winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
return (GM_ADDR)(((HcclRankRelationResV2 *)(winContext_[ctxIdx]->remoteRes[rankId].nextDevicePtr))->windowsIn)
+ winDataSizeOffset_ + rankId * OPT_RANK_OFFSET;
}
含义: 通过 winContext 获取远端 rank 的窗口地址，winDataSizeOffset_ 实现双缓冲（dataState * WIN_STATE_OFFSET）。
7.1.4 数据空间布局
窗口数据空间 (每个 rank):
├── expert0: [0, expertPerSizeOnWin)
│ ├── token0: [0, hCommuSize) — token 数据 + extra space (scale + flag)
│ ├── token1: [hCommuSize, 2*hCommuSize)
│ └── ...
├── expert1: [expertPerSizeOnWin, 2*expertPerSizeOnWin)
└── ...
hCommuSize = tokenLength * sizeof(int8_t) + TOKEN_EXTRA_SPACE
TOKEN_EXTRA_SPACE = 512B (存储 scale + flag)
7.2 CamMoeDistributeCombine — 接收端
文件: fused_deep_moe/op_kernel/fused_deep_moe/raw_distributed/cam_moe_distribute_combine.h
7.2.1 CombineCalcInfo 结构
struct CombineCalcInfo {
uint64_t expertPerSizeOnWin_; // 单个专家在窗口中的空间
uint32_t epRankId_; // 本 rank ID
uint32_t epWorldSize_; // EP 世界大小
uint32_t moeExpertPerRankNum_; // 每 rank 专家数
uint32_t axisH_; // 隐藏维度
uint32_t moeSendNum_; // MoE 发送数
GM_ADDR epSendCount_; // EP 发送计数
__gm__ HcclOpResParam *epWinContext_; // EP 窗口上下文
uint64_t winDataSizeOffset_; // 窗口数据偏移
uint32_t aivNum_; // AIV 数量
uint32_t moeExpertNum_; // 总专家数
GM_ADDR metaInfoGm_; // 元信息 GM 地址
GM_ADDR allExpertTokenNums_; // 全局专家 token 数
GM_ADDR allEpRecvCount_; // 全局 EP 接收计数
GM_ADDR combineSend_; // Combine 发送缓冲
};
7.2.2 SHMEM 集成
Base 变体的 Combine 使用 SHMEM 进行跨 rank 同步：
__aicore__ GM_ADDR GetShmemAddrByRankId(GM_ADDR baseAddr, int64_t rankId) {
if (epRankId_ == rankId) return baseAddr;
return (GM_ADDR)aclshmem_ptr(baseAddr, rankId); // SHMEM 指针转换
}
__aicore__ GM_ADDR GetShmemStateAddrByRankId(int64_t rankId) {
return GetShmemAddrByRankId(metaInfoGm_ + statusDataSpaceOffset_ + dataState_ * WIN_STATE_OFFSET, rankId);
}
SHMEM 清理（ShmemCleanUp）:
void ShmemCleanUp() {
if (coreIdx_ != 0) return; // 仅 core 0 执行
// 清零本地 notify 数据
Duplicate(cleanTempLt, 0, moeExpertNum_);
DataCopyPad(localNotifyDataTensor, cleanTempLt, ...);
}
7.2.3 Combine 执行流程
// 1. AllToAllSend: 将 GMM2 输出发送到远端 rank
void AllToAllSend() {
// 读取 epSendCount 确定每个 rank 的 token 范围
// DataCopyPad: UB → GM (远端窗口)
}
// 2. ReducePermute: 接收远端 rank 的结果并累加
void ReducePermute() {
// 从窗口读取 token
// 按 expert 累加 (CustomAdd)
// 写回输出 GM
}
// 3. ProcessCombine: 主流程
void ProcessCombine() {
AllToAllSend();
ReducePermute();
}
7.3 通信原语总结
| 操作 | 方向 | 机制 | 触发方 |
|---|---|---|---|
| Dispatch Send | 本 rank → 远端 rank | DataCopy 到远端窗口 | AIV |
| Dispatch Recv | 远端 rank → 本 rank | 轮询窗口状态 + DataCopy | AIV |
| Count Send | 本 rank → 远端 rank | DataCopy 到远端状态空间 | AIV |
| Count Recv | 远端 rank → 本 rank | GatherMask + ReduceSum 检测 | AIV |
| Combine Send | 本 rank → 远端 rank | DataCopyPad 到远端窗口 | AIV (GMM2 Epilogue 回调) |
| Combine Recv | 远端 rank → 本 rank | 轮询窗口 + 累加 | AIV |
| SHMEM 同步 | 跨 rank | aclshmem_ptr + DataCacheClean | AIC+AIV |
8. GEMM/Epilogue 子模块 — 计算内核
8.1 子模块目录结构
fused_deep_moe/op_kernel/fused_deep_moe/
├── gemm/
│ ├── block/
│ │ ├── block_mmad.h ← BlockMmad 通用实现
│ │ ├── block_mmad_preload_async_with_callback_resident_a.h ← GMM2 ResidentA
│ │ └── (W4A8) block_mmad_preload_async_with_callback_w4a8.hpp
│ ├── kernel/
│ │ ├── grouped_matmul_slice_m_per_token_dequant_multistage_workspace.h ← GMM2 kernel
│ │ └── grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h ← GMM1 kernel
│ └── dispatch_policy.h
├── epilogue/
│ ├── block/
│ │ ├── block_epilogue.h ← BlockEpilogue 通用实现
│ │ └── (W4A8) block_epilogue_per_token_dequant.hpp / _swiglu.h
│ ├── tile/
│ │ ├── tile_broadcast_mul.hpp ← 行广播乘
│ │ ├── tile_broadcast_one_blk.hpp ← 单块广播
│ │ ├── tile_swizzle.hpp ← Tile 调度
│ │ └── (FWK) tile_stride_muls.h / tile_stride_binary.h
│ └── dispatch_policy.h
└── raw_distributed/
├── cam_moe_distribute_dispatch.h
└── cam_moe_distribute_combine.h
8.2 BlockMmad — 矩阵乘累加块
文件: block_mmad.h (Base), block_mmad_preload_async_with_callback_w4a8.hpp (W4A8)
8.2.1 通用结构
template <class DispatchPolicy, class L1TileShape, class L0TileShape,
class AType, class BType, class CType, class BiasType,
class TileCopy, class TileMmad>
struct BlockMmad {
// 1. InitL1: 初始化 L1 双缓冲
// 2. InitL0A/InitL0B/InitL0C: 初始化 L0 缓冲
// 3. operator(): 主计算循环
// - K 轴分 tile → GM→L1 预取 → L1→L0 搬运 → TileMmad
// 4. SynchronizeBlock: 排空流水线
};
8.2.2 流水线机制
K 轴循环:
for kTileIdx in [0, kTileCount):
1. GM→L1: copyGmToL1A/L1B (MTE2)
2. 如果 preloadCount == PRELOAD_STAGES:
L1TileMmad(preload 最早的 tile) ← 计算与预取重叠
3. 存储当前参数到 l1TileMmadParamsList
4. 更新 preloadCount / l1ListId
L1TileMmad (内部):
for mPartIdx, kPartIdx, nPartIdx:
L1→L0A: copyL1ToL0A (MTE1)
L1→L0B: copyL1ToL0B (MTE1)
TileMmad: 硬件乘累加 (FIX)
if isKLoopLast:
callbackBeforeFixpipe()
copyL0CToGm (MTE3)
callbackAfterFixpipe() ← 触发 Epilogue
8.3 BlockEpilogue — 后处理块
文件: block_epilogue.h
template <class EpilogueDispatchPolicy, class CType, class ScaleType,
class PerTokenScaleType, class DType, class TileRowBroadcastMul,
class TileBroadcastOneBlk, class TileOneBlkColumnBroadcastMul,
class TileCopy, class TileScheduler>
struct BlockEpilogue {
// 在 L0C→GM 写入前/后执行:
// 1. PerTokenDequant: INT32 * weightScale * perTokenScale → float
// 2. SwiGLU: x * sigmoid(x) * x_right (half-glu)
// 3. BlockQuant: float → INT8 + dequantScale
};
Epilogue 计算流程（以 GMM1 为例）:
L0C (int32) → Cast → float
→ RowBroadcastMul(weightScale) // 乘权重 scale
→ BroadcastOneBlk(perTokenScale) // 乘 per-token scale
→ SwiGLU: // 激活函数
left = x[:, :N/2]
right = x[:, N/2:]
silu = left * sigmoid(left)
out = silu * right
→ BlockQuant: // 量化为 INT8
abs → rowMax → scale = rowMax/127
quant = round(x * 127/rowMax)
→ GM (int8) + dequantScale (float)
8.4 GroupedMatmul Kernel
文件: grouped_matmul_slice_m_per_token_dequant_swiglu_quant_multistage_workspace.h
这是 GMM1 的 kernel 级别实现，核心特征：
- Grouped Matmul: 多个专家权重组成一个 group，每个 token 选择对应专家的权重做 GMM
- Slice M: M 维按 token 切片，每个 AIV 处理一部分 token
- PerTokenDequant: 每个 token 有独立的 dequant scale
- SwigluQuant: SwiGLU 激活后立即量化
- MultiStageWorkspace: 4 级 workspace 流水线
8.4.1 深度融合变体
// 深度融合: Dispatch 在 GMM kernel 内部
template <TemplateMC2TypeClass, class BlockMmad, class BlockEpilogue,
class BlockScheduler, uint32_t WORKSPACE_STAGES, class ElementGroupList>
struct GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace {
// AIV 在 GMM1 计算前完成 Dispatch (量化 + AlltoAll 发送)
// AIC 同时计算 GMM1
// 通过 workspace 流水线实现计算-通信重叠
};
// 非深度融合: 独立的 Dispatch + GMM
template <TemplateMC2TypeClass, class BlockMmad, class BlockEpilogue,
class BlockScheduler, uint32_t WORKSPACE_STAGES, class ElementGroupList>
struct GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspaceWithShallowDispatch {
// Dispatch 已在外部完成，直接执行 GMM1
};
8.4.2 BlockQuant 内嵌
GMM1 kernel 内部嵌入了 BlockQuant 类，用于 SwiGLU 输出后的量化：
template <class ArchTag>
class BlockQuant {
// 输入: SwiGLU 输出 (float)
// 输出: INT8 token + dequantScale
// 计算:
// 1. Abs + ReduceMax → rowMax
// 2. scale = rowMax / 127
// 3. quant = Cast(Cast(Cast(x * 127/scale, RINT), ROUND), TRUNC)
};
8.5 GMM2 的 ResidentA 策略
GMM2 使用 MmadAtlasA2PreloadAsyncWithCallbackResidentA 策略：
using Gmm2DispatchPolicy =
Gemm::MmadAtlasA2PreloadAsyncWithCallbackResidentA<
CUSTOM_PRELOAD_STAGES, GMM2_L1A_STAGES, GMM2_L1B_STAGES,
GMM2_L0A_STAGES, GMM2_L0B_STAGES, CUSTOM_L0C_STAGES,
CUSTOM_ENABLE_UNIT_FLAG, CUSTOM_ENABLE_SHUFFLE_K>;
ResidentA 含义: 激活矩阵 A (X2) 在 L1 中常驻（4 级缓冲），因为 GMM2 的输入 X2 来自 SwiGLU 的输出（已在 workspace 中），不需要频繁从 GM 读取。权重 B 仍然需要预取。
9. SHMEM 集成机制
9.1 Base 变体的 SHMEM 使用
Base 变体是唯一 requires_shmem: true 的变体，通过 SHMEM 实现跨 rank 的元数据同步。
9.1.1 SHMEM Workspace 分配
// fused_deep_moe.h Init()
if constexpr (EXEC_FLAG & EXEC_FLAG_ZERO_BUFFER) {
shmemWorkspaceGM_ = (GM_ADDR)(tilingData->disGmmDeqSwigluQuantGmmDeqComInfo.shmemWorkspacePtr);
}
shmemWorkspacePtr 来自 tiling 数据，指向 SHMEM 分配的窗口空间。
9.1.2 SHMEM Workspace 布局
shmemWorkspaceGM_:
├── gmShareX1 / gmX1 — GMM1 输入 token (int8)
├── gmShareX2 / gmX2 — GMM2 输入 token (int8)
├── gmShareX1Scale / gmX1Scale — GMM1 输入 scale
├── gmShareX2Scale / gmX2Scale — GMM2 输入 scale
├── gmTokenFlag — token 到达标志
├── gmAllExpertTokenNums — 全局专家 token 计数
└── gmCombineSend — Combine 发送缓冲
9.1.3 SHMEM 指针转换
// combine 中的 SHMEM 地址计算
__aicore__ GM_ADDR GetShmemAddrByRankId(GM_ADDR baseAddr, int64_t rankId) {
if (epRankId_ == rankId) return baseAddr;
return (GM_ADDR)aclshmem_ptr(baseAddr, rankId); // SHMEM 指针转换
}
aclshmem_ptr 是 CAM 提供的 SHMEM 指针转换函数，将本地 SHMEM 地址转换为远端 rank 可访问的指针。
9.1.4 双缓冲机制
// dataState 在 {0, 1} 之间交替
winDataSizeOffset_ = dataState * epRankSize * expertPerSizeOnWin * moeExpertNumPerRank;
// 窗口地址 = base + winDataSizeOffset_ + rankId * OPT_RANK_OFFSET
GM_ADDR GetWindowDataAddrByRankId(uint32_t rankId) {
return baseAddr + winDataSizeOffset_ + rankId * WIN_OPT_RANK_OFFSET;
}
双缓冲作用: 避免上一次迭代的通信数据被当前迭代覆盖。dataState 在每次迭代后翻转。
9.1.5 SHMEM 清理
void ShmemCleanUp() {
if (coreIdx_ != 0) return; // 仅 core 0 执行
// 清零本地 notify 数据
AscendC::LocalTensor<int32_t> cleanTempLt = cleanUpBuf_.Get<int32_t>();
AscendC::Duplicate(cleanTempLt, (int32_t)0, moeExpertNum_);
// DataCopyPad 到远端窗口的 notify 区域
DataCopyPad(localNotifyDataTensor, cleanTempLt, ...);
}
9.2 W4A8/FWK 的非 SHMEM 通信
W4A8 和 FWK 使用纯 HCCL Window 通信（通过 MC2 框架），不需要 SHMEM：
- W4A8: 依赖
Mc2InitTiling/Mc2CcTiling进行 AlltoAll 通信配置 - FWK: 通过
winContext->localWindowsIn/Exp直接访问窗口
10. 关键数据结构与内存布局
10.1 FusedDeepMoeInfo (Base/W4A8)
struct FusedDeepMoeInfo {
uint32_t epRankSize; // EP rank 数量
uint32_t epRankId; // 本 rank ID
uint32_t moeExpertNum; // 总 MoE 专家数
uint32_t moeExpertNumPerRank; // 每 rank 专家数
uint32_t quantMode; // 量化模式
uint32_t globalBs; // 全局 batch size
uint32_t bs; // 本 rank batch size
uint32_t k; // topK
uint32_t h; // 隐藏维度
uint32_t aicNum; // AIC 数量
uint32_t aivNum; // AIV 数量
uint64_t totalUbSize; // UB 总大小
uint64_t totalWinSize; // 窗口总大小
uint64_t gmm1HLen; // GMM1 隐藏维度
uint64_t shareGmm1HLen; // 共享专家 GMM1 隐藏维度
uint64_t shmemWorkspacePtr; // SHMEM workspace 地址 (Base 独有)
uint64_t metaInfoPtr; // 元信息地址 (Base 独有)
bool isTensorList; // 权重是否为 TensorList
};
10.2 FusedDeepMoeFwkInfo (FWK)
struct FusedDeepMoeFwkInfo {
uint32_t epRankSize;
uint32_t epRankId;
uint32_t moeExpertNum;
uint32_t moeExpertNumPerRank;
uint32_t sharedExpertNum; // ← FWK 独有：共享专家数
uint32_t sharedExpertRankNum; // ← FWK 独有：共享专家 rank 数
uint32_t quantMode;
uint32_t globalBs;
uint32_t bs;
uint32_t k;
uint32_t h;
uint32_t aicNum;
uint32_t aivNum;
uint64_t totalUbSize;
uint64_t totalWinSize;
uint64_t gmm1HLen;
bool isTensorList;
// 注意: 无 shmemWorkspacePtr / metaInfoPtr
};
10.3 TilingData 结构
// Base/W4A8
struct FusedDeepMoeTilingData {
Mc2InitTiling mc2InitTiling; // MC2 初始化 tiling
Mc2CcTiling mc2CcTiling; // MC2 通信 tiling
FusedDeepMoeInfo disGmmDeqSwigluQuantGmmDeqComInfo; // 算子信息
};
// FWK
struct FusedDeepMoeFwkTilingData {
Mc2InitTiling mc2InitTiling;
Mc2CcTiling mc2CcTiling;
FusedDeepMoeFwkInfo disGmmDeqSwigluQuantGmmDeqComInfo;
};
10.4 窗口状态空间布局 (全局)
HCCL Window (每个 rank):
├── DataSpace (数据空间)
│ ├── rank0: [0, expertPerSizeOnWin * moeExpertNumPerRank)
│ │ ├── expert0: [0, expertPerSizeOnWin)
│ │ │ ├── token0: [0, hCommuSize) — int8 token + 512B extra
│ │ │ └── ...
│ │ └── ...
│ └── ...
├── ExpSpace (状态空间, 1MB)
│ ├── [0, 512KB): dataState=0 状态
│ │ ├── rank0: [0, 512B)
│ │ │ ├── [0, 256B): [flag0, count0, pad, pad, flag1, count1, ...]
│ │ │ └── [256B, 512B): 保留
│ │ └── ...
│ ├── [512KB, 1024KB): dataState=1 状态 (双缓冲)
│ ├── [900KB, 932KB): 窗口状态 (STATE_WIN_OFFSET)
│ ├── [932KB, 964KB): 组 token 计数 (GROUP_TOKEN_NUM_OFFSET)
│ └── [964KB, 988KB): 软同步空间 (SOFT_SYNC_OFFSET)
附录 A: 三变体文件清单
Base (fused_deep_moe)
| 文件 | 作用 |
|---|---|
op_kernel/fused_deep_moe.h |
主 kernel 类 + GmmDeqSwigluQuant/GmmDeq 函数 |
op_kernel/fused_deep_moe.cpp |
kernel 入口 + 64 tiling key 展开 |
op_kernel/fused_deep_moe_base.h |
模板宏 + 通信组结构 |
op_kernel/fused_deep_moe_tiling.h |
TilingData + 编译期常量 |
op_kernel/fused_deep_moe_tiling.cpp |
(不存在，依赖 MC2 默认 tiling) |
op_host/fused_deep_moe_infer.cpp |
InferShape/InferDataType |
op_api/aclnn_fused_deep_moe.{h,cpp} |
aclnn 接口 |
W4A8 (fused_deep_moe_w4a8)
| 文件 | 作用 |
|---|---|
op_kernel/fused_deep_moe.h |
主 kernel 类 (INT4 类型) |
op_kernel/fused_deep_moe.cpp |
kernel 入口 + Bias 输入 |
op_kernel/fused_deep_moe_base.h |
模板宏 (含 CONSTANT_TWO) |
op_kernel/fused_deep_moe_tiling.h |
TilingData (L1K=1024) |
op_host/fused_deep_moe_tiling.cpp |
完整 tiling 逻辑 + Bias 校验 |
op_host/fused_deep_moe_infer.cpp |
InferShape/InferDataType |
op_kernel/fused_deep_moe/gemm/block/block_mmad_preload_async_with_callback_w4a8.hpp |
INT4 BlockMmad |
FWK (fused_deep_moe_fwk)
| 文件 | 作用 |
|---|---|
op_kernel/fused_deep_moe_fwk.h |
主 kernel 类 (简化模板) |
op_kernel/fused_deep_moe_fwk.cpp |
kernel 入口 (4 tiling key) |
op_kernel/fused_deep_moe_fwk_base.h |
模板宏 (OPT_RANK_OFFSET) |
op_kernel/fused_deep_moe_fwk_tiling.h |
TilingData (FWKInfo) |
op_host/fused_deep_moe_fwk.cpp |
tiling 函数 |
op_host/fused_deep_moe_fwk_infer.cpp |
InferShape (共享专家) |
op_kernel/fused_deep_moe_fwk/lib_ops.h |
OPX 框架操作库 |
附录 B: 关键常量对比表
| 常量 | Base | W4A8 | FWK |
|---|---|---|---|
| GMM1_L1M | 256 | 256 | 256 |
| GMM1_L1N | 128 | 128 | 128 |
| GMM1_L1K | 512 | 1024 | 512 |
| GMM1_L0K | 128 | 256 | 128 |
| GMM1_EPIM | 64 | 64 | 64 |
| GMM2_L1M | 128 | 128 | 128 |
| GMM2_L1N | 256 | 256 | 256 |
| GMM2_L1K | 512 | 1024 | 512 |
| GMM2_L0K | 128 | 256 | 128 |
| GMM2_EPIM | 32 | 32 | 32 |
| WORKSPACE_STAGES | 4 | 4 | 4 |
| CUSTOM_PRELOAD_STAGES | 1 | 1 | 1 |
| CUSTOM_L1_STAGES | 2 | 2 | 2 |
| GMM2_L1A_STAGES | 4 | 4 | 4 |
| GMM2_L0A_STAGES | 4 | 4 | 4 |
| GM_ALIGN_BYTE | 512 | 512 | 512 |
| STATE_OFFSET | 512 | — | — |
| WIN_STATE_OFFSET | 512KB | — | — |
| TOKEN_EXTRA_SPACE | 512B | — | 512B |
附录 C: 核心函数调用链
Base 变体 (深度融合路径)
fused_deep_moe()
└── FusedDeepMoe::Process()
├── [if !DEEP_FUSE] CamMoeDistributeDispatch::Init/Process() [AIV]
│ ├── SendToMoeExpert() — 量化 + AlltoAll 发送
│ ├── AlltoAllDispatch() — 窗口数据搬运
│ └── ResetStatus() — 状态清零
├── PipeBarrier<PIPE_ALL>()
├── CrossCoreFlag: AIV→AIC 同步
├── GmmDeqSwigluQuant() — GMM1 + SwiGLU + Quant
│ └── GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
│ ├── BlockMmad::operator() — GM→L1→L0→TileMmad
│ ├── BlockEpilogue::Process() — Dequant + SwiGLU
│ └── BlockQuant::operator() — float→INT8
├── PipeBarrier<PIPE_ALL>()
├── CrossCoreFlag: AIV→AIC 同步
├── CamMoeDistributeCombine::Init() [AIV]
└── GmmDeq() — GMM2 + Combine
└── GroupedMatmulSliceMPerTokenDequantMultiStageWorkspace
├── BlockMmad::operator() — ResidentA 策略
└── callbackAfterFixpipe → combiner.AllToAllSend()
FWK 变体
fused_deep_moe_fwk()
└── FusedDeepMoeFwk::Process()
├── [if !DEEP_FUSE] CamMoeDistributeDispatch::Init/Process() [AIV]
├── PipeBarrier<PIPE_ALL>()
├── CrossCoreFlag: AIV→AIC 同步
├── GmmDeqSwigluQuant<EXEC_FLAG, ExpandXType, ...>()
│ └── GroupedMatmulSliceMPerTokenDequantSwigluQuantMultiStageWorkspace
│ └── (内部使用 LibOps 的 Op 类)
├── PipeBarrier<PIPE_ALL>()
├── CrossCoreFlag: AIV→AIC 同步
├── CamMoeDistributeCombine::Init() [AIV]
└── GmmDeq() — GMM2 + Combine
└── callbackAfterFixpipe → combiner.AllToAllSend()
└── OpCombineSend/OpCombineRecv (LibOps)
总结: FusedDeepMoE 三变体共享五阶段融合架构，但在通信机制（SHMEM vs HCCL Window）、量化精度（INT8 vs INT4）、框架集成（原生 vs OPX）三个维度上分化。Base 是通用旗舰，W4A8 服务量化部署，FWK 面向框架原生集成。三者的演进体现了 CAM 从专用算子到通用框架的架构升级路径。
分支:
remotes/origin/cam_feature| 分析日期: 2026-07-21 目标芯片: Ascend 910_93 (a2e/e2a) / Ascend 910_93 + 910B (detour 系列)
目录
- 新算子总体定位
- a2e (All-to-Expert)
- e2a (Expert-to-All)
- all_to_all_detour
- reduce_scatter_detour
- Detour 机制深度分析
- SHMEM 使用
- 与现有算子链的关系
1. 新算子总体定位
1.1 背景问题
在 MoE（Mixture-of-Experts）+ EP（Expert Parallelism）架构中，Attention 侧与 Expert 侧通常运行在不同的 NPU 卡簇上。这两侧之间存在大规模的数据搬运需求：
- 前向: Attention 输出的 token hidden states 需要按 expert 路由结果分发到对应 Expert 卡 → All-to-Expert (a2e)
- 反向/回传: Expert 计算完成的 token 结果需要汇总回 Attention 卡 → Expert-to-All (e2a)
此外，当 MoE 集群规模超出单网络拓扑的 fullmesh 互连范围时（例如跨 pod、跨交换机），部分 rank 对之间没有直连链路，通信必须通过中间 rank 绕行 → Detour 机制。
1.2 四个算子的角色
| 算子 | 方向 | 核心功能 | 场景 |
|---|---|---|---|
| a2e | Attention → Expert | 将 token hidden states + expert IDs + scales 从 Attention 卡发送到 Expert 卡 | MoE 前向 dispatch 的跨簇通信 |
| e2a | Expert → Attention | 将 Expert 计算完成的 token 结果发回 Attention 卡 | MoE 前向 combine 的跨簇通信 |
| all_to_all_detour | 全对全 + 绕行 | 在 non-fullmesh 拓扑下，通过中间 rank 绕行完成 All-to-All | 跨 pod MoE EP 全交换 |
| reduce_scatter_detour | 归约散射 + 绕行 | 在 non-fullmesh 拓扑下，通过绕行完成 Reduce-Scatter | 跨 pod 梯度同步 / 归约 |
1.3 与 master 分支的差异
operator_registry.json 明确显示这四个算子仅存在于 cam_feature 分支的 ascend910_93 列表中。master 分支的 MoE 通信依赖 HCCL 原生 collective（AlltoAll、ReduceScatter）或已存在的 moe_dispatch_* / moe_combine_* 算子，不包含跨 Attention-Expert 簇的 SHMEM 直传算子和 Detour 绕行算子。
2. a2e (All-to-Expert)
源文件: src/cam/comm_operator/ascend_kernels/a2e/
2.1 算子签名
a2e(x, expert_ids, scales, batch_size, hidden_size, topk,
expert_rank_size, attention_rank_size, rank, group_ep, aiv_num, compute_gate)
→ (expand_x, simulate_expert_ids, simulate_expert_scales, atten_batch_size, x_active_mask)
2.2 核心数据结构
A2ETilingData (a2e_tiling.h):
struct A2ETilingData {
Mc2InitTiling mc2InitTiling; // MC2 初始化 tiling
Mc2CcTiling mc2CcTiling1; // MC2 通信 tiling
uint32_t batchSize;
uint32_t hiddenSize;
uint32_t topk;
uint32_t expertRankSize; // Expert 侧 rank 数
uint32_t attentionRankSize; // Attention 侧 rank 数
uint32_t rank; // 当前 rank
uint32_t computeGate; // 计算门控（决定发送内容）
};
2.3 执行流程
入口 (op_kernel/a2e.cpp):
- 仅 TILING_KEY_IS(21)（BF16）被实际编译
- 实例化 A2e<bfloat16_t, bfloat16_t, false>，调用 init() → process()
init() 阶段 (a2e.h):
-
获取 HCCL IPC 上下文:
cpp epWinContext_ = (__gm__ HcclOpResParam *)AscendC::GetHcclContext<HCCL_GROUP_ID_0>();通过localWindowsIn和remoteRes[i].nextDevicePtr->windowsIn获取每张卡的 IPC 共享内存窗口基址。 -
Magic 计数器初始化（跨 block 同步令牌）: ```cpp magicTensor_.SetGlobalBuffer((gm int32_t*)((epWinContext_->localWindowsIn)
- IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK));
`` 每个 block 通过AtomicAdd在共享内存中顺序写入递增的 blockIdx 值，读取到的值即为该 block 的唯一magic` 序号。这是无锁分布式 barrier 令牌。
- IPC_DATA_OFFSET - blockNum * sizeof(int32_t) * INT32_COUNT_PER_BLOCK));
-
共享内存地址映射 (
shareAddrs[]): - Attention 卡 (rank >= expertRankSize): 只需映射自己和对应 Expert 卡的地址 - Expert 卡 (rank < expertRankSize): 映射所有 rank 的地址 - 地址公式:localWindowsIn + (magic % 2) * 800MB + rank * 512-DOUBLE_BUFFER_COUNT = 2实现 ping-pong 避免上一轮/本轮数据冲突 -OPT_RANK_OFFSET = 512保证每张卡的 flag 区域不互相覆盖 -
attnToMoeRatio 计算:
cpp attnToMoeRatio = attentionRankSize > expertRankSize ? (attentionRankSize + expertRankSize - 1) / expertRankSize : 1;表示每个 Expert 卡需要接收来自多少张 Attention 卡的数据。当attentionRankSize不能整除expertRankSize时，部分 Expert 卡少收一份。
process() 阶段 — 根据 computeGate 分两模式:
模式 A: computeGate == 0（纯数据搬运模式）
Attention 卡 (rank >= expertRankSize):
sendWithMte() → 仅发送 token X
Expert 卡 (rank < expertRankSize):
for i in 0..attnToMoeRatio:
recvWithMte(i) → 从 rank + (i+1)*expertRankSize 接收 X
- 发送端 (
sendX): 多核协作copyGmToGmWithBlocks将batchSize * hiddenSize的 token 数据从本地 GM 拷贝到目标 rank 的 IPC 共享内存，然后写 flag 通知接收端 - 接收端 (
recvWithMte):waitFlagWithScalar轮询 flag 到位后，从源 rank 的共享内存拷贝到本地expandX
模式 B: computeGate == 1（分核分工模式）
Expert 卡:
blockIdx == blockNum-1 → sendExpertIds() // 发送 expert IDs
blockIdx == blockNum-2 → sendExpertScales() // 发送 expert scales
其他 block → sendX() // 发送 token X
每个 Expert 卡将 expert_ids（路由结果）和 scales（权重）也一并发送给对应的 Attention 卡，使得 Attention 卡可以本地模拟 Expert 分布，用于后续 simulate_expert_ids / simulate_expert_scales 输出。
2.4 数据布局（IPC 共享内存）
shareAddrs[rank] 指向的内存布局:
┌──────────────────────┐ offset 0
│ batchSize (int32) │ 4 bytes
├──────────────────────┤ offset 4
│ expertIds │ batchSize * topk * sizeof(int32)
├──────────────────────┤ offset 4 + expertIdsSize
│ expertScales │ batchSize * topk * sizeof(float)
├──────────────────────┤ offset 4 + expertIdsSize + scalesSize
│ X (token data) │ batchSize * hiddenSize * sizeof(bf16)
└──────────────────────┘
Flag 区域位于 IPC_DATA_OFFSET 之前（前 2MB），使用 mergeMagicWithValue(magic, 0) 编码:
uint64_t mergeMagicWithValue(uint32_t magic, uint32_t count) {
return (static_cast<uint64_t>(magic) << 32) | static_cast<uint64_t>(count);
}
Flag 偏移量按 (flagIndex * attnToMoeRatio + sendOffset) 索引，每个 Expert 卡可能接收多份 Attention 卡的数据，通过 attnToMoeRatio 区分。
2.5 同步机制
发送端写 flag (通知对端数据就绪):
GlobalTensor<uint64_t> shareFlagGt;
shareFlagGt.SetGlobalBuffer(shareAddrs[rank % expertRankSize]
+ (WAIT_FLAG_OFFSET_3 * attnToMoeRatio + sendOffset) * flagUnitInt64Num);
camCpUB2GM(shareFlagGt, flagLt, 1, sizeof(uint64_t));
接收端轮询 flag:
__aicore__ inline void waitFlagWithScalar(int addr, uint32_t magic) {
GlobalTensor<uint32_t> flagGt;
flagGt.SetGlobalBuffer((__gm__ uint32_t *)(shareAddrs[rank] + addr));
while(1) {
DATA_FULSH(flagGt, uint32_t); // cache flush + invalidate
if (flagGt.GetValue(0) == magic) return;
}
}
DATA_FULSH 宏确保每次轮询都重新从内存读取（绕过 cache），通过 DataCacheCleanAndInvalid + dsb(DSB_ALL) 实现。
2.6 参数约束
| 参数 | 约束 |
|---|---|
| batchSize | 1 ~ 1024 |
| hiddenSize | 1024 ~ 7168, 256 对齐 |
| topk | 1 ~ 8 |
| expertRankSize | > 0 且 ≤ attentionRankSize |
| aivAlgNum | 4 ~ 48 |
| computeGate | 0 或 1 |
2.7 与 dispatch 的关系
a2e 是 MoE dispatch 阶段的跨簇通信原语。传统 dispatch（moe_dispatch_* 算子）在同一 EP group 内部做 token-to-expert 路由。而 a2e 解决的是Attention 卡簇与 Expert 卡簇之间的物理分离场景 — Attention 卡持有 token，Expert 卡持有 expert 权重，中间通过非 HCCL 路径（可能是 PCIe/SIO 互连而非 RDMA）传输。
从 sample 代码可见，Attention 侧发送真实数据，Expert 侧用空 tensor 占位接收 — a2e 是角色感知的，根据 rank < expertRankSize 自动判断当前是 Expert 侧还是 Attention 侧。
3. e2a (Expert-to-All)
源文件: src/cam/comm_operator/ascend_kernels/e2a/
3.1 算子签名
e2a(expand_x, atten_batch_size, batch_size, hidden_size, topk,
expert_rank_size, attention_rank_size, rank, group_ep, aiv_num)
→ x
3.2 核心数据结构
E2ATilingData (e2a_tiling.h):
struct E2ATilingData {
Mc2InitTiling mc2InitTiling;
Mc2CcTiling mc2CcTiling1;
uint32_t batchSize;
uint32_t hiddenSize;
uint32_t topk;
uint32_t expertRankSize;
uint32_t attentionRankSize;
uint32_t rank;
};
注意: e2a 的 tiling 没有 computeGate，因为 e2a 只负责回传 token 数据。
3.3 执行流程
入口: 仅 TILING_KEY_IS(21) (BF16) 编译，实例化 E2a<bfloat16_t>。
process() 阶段 — 极为简洁:
if (rank < expertRankSize) {
sendWithMte(); // Expert 卡: 发送结果回 Attention 卡
} else {
recvWithMte(); // Attention 卡: 接收 Expert 结果
}
Expert 端 (sendWithMte):
int sendRank = rank + expertRankSize; // 目标 Attention 卡
for (int i = 0; i < attnToMoeRatio; i++) {
// 将 expandX 中对应分片拷贝到目标 Attention 卡的 IPC 内存
copyGmToGmWithBlocks(shareXGt, expandXGt[sendOffset],
sendBatchSize * hiddenSize, blockNum, blockIdx);
sendOffset += sendBatchSize * hiddenSize;
sendRank += expertRankSize;
}
SyncAll();
// 写完所有分片后，一次性通知所有目标 Attention 卡
for (int i = 0; i < attnToMoeRatio; i++) {
GlobalTensor<uint64_t> shareFlagGt;
shareFlagGt.SetGlobalBuffer((__gm__ uint64_t *)(shareAddrs[sendRank]));
camCpUB2GM(shareFlagGt, flagLt, 1, sizeof(uint64_t));
}
Attention 端 (recvWithMte):
// 轮询等待 flag
while (1) {
camCpGM2UB(flagLt, flagGt, sizeof(uint64_t));
if (flagLt.GetValue(0) == mergeMagicWithValue(magic, 0, 0)) break;
}
// flag 到位后，从 IPC 内存拷贝结果到本地
GlobalTensor<T> shareXGt;
shareXGt.SetGlobalBuffer((__gm__ T *)(shareAddrs[rank] + IPC_DATA_OFFSET));
copyGmToGmWithBlocks(xGt, shareXGt, batchSize * hiddenSize, blockNum, blockIdx);
3.4 a2e/e2a 对称性分析
| 维度 | a2e | e2a |
|---|---|---|
| 方向 | Attention → Expert | Expert → Attention |
| 发送方 | Attention 卡 (rank ≥ expertRankSize) | Expert 卡 (rank < expertRankSize) |
| 接收方 | Expert 卡 | Attention 卡 |
| 传输内容 | X + expertIDs + scales (computeGate=1) | 仅 X |
| 发送端分核 | computeGate=1 时分核发送 IDs/scales/X | 不分核，所有核协作发 X |
| 接收端行为 | 多核协作接收，按 attnToMoeRatio 拼接 | 多核协作接收完整 batch |
| Flag 机制 | 按消息类型分 flag 偏移 | 单一 flag 通知 |
| computeGate | 有 (0/1 两模式) | 无 |
本质区别: a2e 是多对多分发（一个 Expert 卡接收多个 Attention 卡的 token），e2a 是多对多回传（一个 Attention 卡接收多个 Expert 卡的结果）。a2e 需要更复杂的 flag 系统来区分不同消息类型（batchSize、expertIds、scales、X），e2a 只需一个 flag 通知"数据就绪"。
3.5 与 combine 的关系
e2a 是 MoE combine 阶段的跨簇通信原语。传统 combine（moe_combine_* 算子）在 EP group 内部聚合 expert 输出。e2a 解决 Expert 计算完成后，将结果从 Expert 卡簇回传到 Attention 卡簇的问题。
4. all_to_all_detour
源文件: src/cam/comm_operator/ascend_kernels/all_to_all_detour/
4.1 算子签名
all_to_all_detour(sendData, commRankIds, commArgs, magic) → recvData
sendData: 发送缓冲区，shape[sendCount]commRankIds: 参与通信的 rank ID 列表，shape[commRankCount]commArgs: 通信参数结构体（含peerMems和memLen）magic: 同步令牌
4.2 核心数据结构
All2AllDetourTilingData (all2_all_detour_tiling.h):
BEGIN_TILING_DATA_DEF(All2AllDetourTilingData)
TILING_DATA_FIELD_DEF(uint32_t, magic);
TILING_DATA_FIELD_DEF(uint32_t, sendCount);
TILING_DATA_FIELD_DEF(uint32_t, commRankCount);
END_TILING_DATA_DEF;
4.3 核心算法 (all2all_small_fullmesh_with_detour.h)
All2AllSmallFullmeshWithDetour<T> 继承自 Collectives 基类，实现两阶段 detour 通信。
4.3.1 Init 阶段
void Init(KERNELS_ARGS_FUN_ALL2ALL, GM_ADDR commRankIds, uint32_t commRankSize) {
// 1. 读取 commRankIds，判断当前 rank 是否参与通信
for (int i = 0; i < commRankSize; ++i)
this->commRankIds[i] = commRankIdsGM.GetValue(i);
if (!IsValidCommRank()) return; // 当前 rank 不在通信列表中则退出
// 2. 初始化 Collectives 基类（使用 peerMems 而非 HCCL window）
Collectives::InitByCamCommDetour(KERNELS_ARGS_CALL_MIX);
// 3. 计算每 rank 数据量和 detour 分片量
perRankDataNum = GetDataCount(len, commRankSize); // len / commRankSize
perRankDataNumForDetour = (perRankDataNum + rankSize - commRankSize)
/ (rankSize - commRankSize + 1);
// 4. 初始化核心分组
InitCoreGroup();
}
关键常量:
- MAX_CORE_NUM = 48（Ascend 910_93 的 AIV 核数）
- rankNumPerCore = (rankSize + 46) / 46 — 每核负责的绕行 rank 数
4.3.2 Process 阶段 — 两阶段流水线
Stage 1: InputToShareSlice() — 从 input 读取数据，经 detour rank 写入目标 rank 的共享内存
Stage 2: ShareToShareSlice() — 从共享内存读取数据，写入 output
Stage 1: InputToShareSlice() — 按 blockIdx 分三种情况:
-
blockIdx == 0: 本地直拷（自己发给自己的数据）
cpp CpGM2GMPingPong<T>(perRankDataNum * sizeof(T), inputGt[commRankIdx * perRankDataNum], outputGt[commRankIdx * perRankDataNum], COPYONLY); -
blockIdx == 1: 第一块 detour 数据 — 每个目标 comm rank 发送
perRankDataNumForDetour大小的首块数据，写入对应 detour rank 的共享内存，并写 flag 通知 ```cpp writeGt = shareAddrs[detourRankId] + GetShareMemOffset(rank, dstCommRank, detourRankId) + IPC_DATA_OFFSET; CpGM2GMPingPong(perRankDataNumForDetour * sizeof(T), readGt, writeGt, COPYONLY); // 写 flag CpUB2GM(shareAddrs[detourRankId] + GetShareMemOffset(rank, dstCommRank, detourRankId) - rank * FLAG_UNIT_INT_NUM, inputUB, sizeof(int64_t)); ```
-
blockIdx ≥ 2: 后续块 — 按
rankNumPerCore切分剩余数据，每核处理一部分 detour rank 的拷贝
Stage 2: ShareToShareSlice() — 从共享内存收数据写入 output:
- blockIdx == 0: 无操作（直拷已在 Stage 1 完成）
- blockIdx == 1: 接收每个源 rank 的首块数据
cpp sync.WaitSyncFlag(magic, 1, srcCommRank, rank, 1, GetShareMemOffset(srcCommRank, rank, rank)); CpGM2GMPingPong<T>(perRankDataNumForDetour * sizeof(T), readGt, outputGt[...], COPYONLY); - blockIdx ≥ 2: 接收后续块，通过 detour rank 中转的数据
4.3.3 共享内存地址计算
int64_t GetShareMemOffset(int srcRank, int dstRank, int memRank) {
return (rankSize * srcRank + dstRank) * memLen
+ memRank * OPT_RANK_OFFSET
+ (magic % PING_PONG_SIZE) * (memLen / 2);
}
(rankSize * srcRank + dstRank) * memLen: 二维矩阵寻址，每对 rank 有独立的内存条带memRank * OPT_RANK_OFFSET (512): 在条带内按 memRank 偏移，避免多核写同一地址(magic % 2) * (memLen / 2): ping-pong 双缓冲
4.4 Detour 路由逻辑
GetDstRankId(order) — 获取第 order 个非通信 rank（即 detour rank）的 ID:
int GetDstRankId(int order) {
int count = 0;
for (int i = 0; i < rankSize; ++i) {
if (!IsCommRank(i)) { // 遍历所有非 commRank 的 rank
if (count == order) return i;
++count;
}
}
}
Detour 的含义: 当 commRankSize < rankSize 时，部分 rank 不在通信列表中（例如它们是另一 pod 的 rank，或纯转发节点）。这些 rank 不直接参与 All-to-All 的数据生产/消费，但作为中转站帮助数据绕行到目标。
4.5 数据类型支持
| Tiling Key | 数据类型 |
|---|---|
| 20 | float16 |
| 21 | bfloat16 |
| 22 | float32 |
| 23 | int32 |
支持 5 种 input format 组合（含 FORMAT_NC1HWC0 用于特殊场景）。
4.6 平台适配
// aclnn_all2_all_detour.cpp
if (op::GetCurrentPlatformInfo().GetSocVersion() == op::SocVersion::ASCEND910B) {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_AICPU);
} else {
NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
}
910B 使用 AICPU 调度，910_93 使用 MTE（Memory Transfer Engine）调度 — 反映两种芯片的 DMA 能力差异。
5. reduce_scatter_detour
源文件: src/cam/comm_operator/ascend_kernels/reduce_scatter_detour/
5.1 算子签名
reduce_scatter_detour(input, commRankIds, commArgs, magic, rank_size, op) → output
input: 输入数据，shape[sendCount]commRankIds: 参与通信的 rank ID 列表commArgs: 通信参数（其中commArgs[0]=rank,commArgs[2]=rankSize）magic: 同步令牌rank_size: 通信域大小op: 归约操作类型（ADD=0, MAX=2 等）
5.2 核心数据结构
ReduceScatterDetourTilingData (reduce_scatter_detour_tiling.h):
struct ReduceScatterDetourTilingData {
uint32_t sendCount;
uint32_t op;
int64_t magic;
uint32_t commRankCount;
};
5.3 核心算法 (reduce_scatter_detour.h)
ReduceScatterWithDetour<T, U> 继承自 Collectives，模板参数:
- T: 输入数据类型（如 float16）
- U: 累加器类型（如 float），精度更高
5.3.1 Init 阶段
void Init(KERNELS_ARGS_FUN_MIX, GM_ADDR commRankIds, uint32_t commRankSize) {
// 1. 读取 commRankIds
// 2. 初始化 Collectives 基类
detourRankSize = rankSize - commRankSize; // 绕行 rank 数
// 3. 计算各阶段核心数
in2IpcCoreNum = (detourRankSize + 1) * (commRankSize - 1) + 1;
reduceCoreNum = detourRankSize + 2;
blockNum = max(in2IpcCoreNum, reduceCoreNum);
// 4. 计算 detour 分片量
perDataNumForDetourRank = (len + detourRankSize) / (detourRankSize + 1);
piplineNum = commRankSize - 1; // 归约流水线级数
}
核心分配策略:
- Input2Ipc: 需要 (detourRankSize + 1) * (commRankSize - 1) + 1 个核
- Reduce: 需要 detourRankSize + 2 个核
- 最终 blockNum = max(两者)，确保两阶段都有足够核
5.3.2 Process 阶段 — 两阶段流水线
Stage 1: Input2Ipc() — 将输入数据分片写入各 rank 的 IPC 共享内存
Stage 2: Reduce() — 从共享内存读取各分片，归约后写回输出
Stage 1: Input2Ipc() — 数据分发
-
blockIdx == 0: 本地数据直拷到自己的共享内存
cpp in2IpcDstGtU.SetGlobalBuffer(shareAddrs[rank] + shareMemOffsetInRank + IPC_DATA_OFFSET, len); CpGM2GMPingPong(len * sizeof(U), in2IpcSrcGt, in2IpcDstGtU, COPYONLY); // 写 flag 通知归约阶段 CpUB2GM(shareAddrs[rank] + shareMemOffsetInRank + rank * FLAG_UNIT_INT_NUM, inputUB, sizeof(int64_t)); -
blockIdx ≥ 1: 分发到其他 rank 的共享内存 ```cpp // 计算目标 comm rank 和 detour rank int targetCommRankIdx = (blockIdx - 1) / (detourRankSize + 1); // 跳过自己 if (targetCommRankIdx >= commRankIdx) targetCommRankIdx++; in2IpcTargetCommRank = commRankIds[targetCommRankIdx];
int offsetIdx = (blockIdx - 1) % (detourRankSize + 1); in2IpcTargetRank = (offsetIdx == 0) ? in2IpcTargetCommRank : GetDetourRankId(offsetIdx - 1);
// 拷贝数据分片到目标 detour rank 的共享内存 in2IpcDstGt.SetGlobalBuffer(shareAddrs[in2IpcTargetRank] + shareMemOffsetInDstRank + IPC_DATA_OFFSET, curDataNumForDetourRank); CpGM2GMPingPong(curDataNumForDetourRank * sizeof(T), in2IpcSrcGt, in2IpcDstGt, COPYONLY); ```
Detour 分发逻辑: 每个 comm rank 的数据被切成 detourRankSize + 1 份，其中一份发给目标 rank 本身（offsetIdx == 0），其余通过 detour rank 转发。
Stage 2: Reduce() — 归约计算
void Reduce() {
if (blockIdx == 0) return; // block 0 不参与归约
int offsetIdx = (blockIdx - 1) % (detourRankSize + 1);
int64_t curReduceDataNum = min(len - offsetIdx * perDataNumForDetourRank, perDataNumForDetourRank);
// 目标 buffer
reduceDstGt = shareAddrs[rank] + shareMemOffsetInRank + IPC_DATA_OFFSET
+ (blockIdx - 1) * perDataNumForDetourRank;
for (int i = 0; i < piplineNum; i++) {
int srcCommRankIdx = (commRankIdx + 1 + i) % commRankSize;
int srcCommRank = commRankIds[srcCommRankIdx];
int detourRank = (offsetIdx == 0) ? rank : GetDetourRankId(offsetIdx - 1);
// 等待源 rank 数据就绪
if (i == 0) {
sync.WaitSyncFlag(magic, 1, rank, rank, 1, shareMemOffsetInRank);
} else {
sync.WaitOuterFlagDetour(magic, i, rank, rankSize + rank, shareMemOffsetInRank);
}
sync.WaitSyncFlag(magic, 1, srcCommRank, detourRank, 1, shareMemOffsetInDstRank);
// 归约: 将源数据拷贝到目标（CpGM2GMPingPong 支持 ADD/MUL/MAX/MIN）
CpGM2GMPingPong(curReduceDataNum * sizeof(U), reduceSrcGt, reduceDstGt, atomOp);
// 设置 outer flag 通知下一级可以开始
sync.SetOuterFlagDetour(magic, i + 1, rank, rankSize + rank, shareMemOffsetInRank);
}
// 等待最后一级完成
sync.WaitOuterFlagDetour(magic, piplineNum, rank, rankSize + rank, shareMemOffsetInRank);
// 从共享内存写回输出
Ipc2OutDstGt = output + (blockIdx - 1) * perDataNumForDetourRank;
CpGM2GMPingPong(curReduceDataNum * sizeof(T), reduceDstGt, Ipc2OutDstGt, COPYONLY);
}
归约流水线: piplineNum = commRankSize - 1 级，每级从一个源 comm rank 读取数据并归约到本地 buffer。使用 WaitOuterFlagDetour / SetOuterFlagDetour 实现流水线级间同步 — 第 i 级完成后设置 flag，第 i+1 级等待该 flag。
5.4 与标准 Reduce-Scatter 的差异
| 维度 | 标准 Reduce-Scatter (HCCL) | reduce_scatter_detour |
|---|---|---|
| 拓扑要求 | 需要 fullmesh 或特定拓扑 | 支持 non-fullmesh，通过 detour 绕行 |
| 中间节点 | 无 | 使用非 commRank 作为转发节点 |
| 归约方式 | 硬件 tree/ring 归约 | 软件流水线逐级归约 |
| 精度 | 输入类型直接归约 | 支持高精度累加器 (T→U, 如 fp16→fp32) |
| 同步 | HCCL 内部 barrier | 显式 flag 同步 (SyncCollectives) |
| 内存 | HCCL 内部 buffer | IPC 共享内存 (peerMems) |
5.5 平台与类型
- 仅支持 Ascend 910B（
this->AICore().AddConfig("ascend910b")） - 仅支持 FP16 输入（
TILING_KEY_FLOAT16 = 20） - 累加器为 float（
ReduceScatterWithDetour<float16_t, float>） - 需要 16MB workspace（
SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024）
6. Detour 机制深度分析
6.1 什么是网络绕行 (Detour)
在超大规模 MoE 集群中，NPU 卡通常按网络拓扑分层组织:
Pod 0 (Fullmesh via HCCS) Pod 1 (Fullmesh via HCCS)
┌─────────────────────┐ ┌─────────────────────┐
│ Rank 0 ─ Rank 1 │ │ Rank 8 ─ Rank 9 │
│ │ │ │ RDMA/ │ │ │ │
│ Rank 2 ─ Rank 3 │◄──────────►│ Rank 10─ Rank 11 │
│ │ │ │ RoCE │ │ │ │
│ Rank 4 ─ Rank 5 │ (高延迟) │ Rank 12─ Rank 13 │
│ │ │ │ │ │ │ │
│ Rank 6 ─ Rank 7 │ │ Rank 14─ Rank 15 │
└─────────────────────┘ └─────────────────────┘
当 Rank 0 需要给 Rank 8 发送数据时，如果跨 pod 带宽不足或链路拥塞，可以通过中间 rank 绕行:
Rank 0 → Rank 3 (同 pod, 高带宽 HCCS) → Rank 11 (跨 pod) → Rank 8
Detour 机制的本质是: 将一次远距离通信拆分为多次短距离通信，利用拓扑局部性降低拥塞、提高有效带宽。
6.2 Detour 在 CAM 中的实现
CAM 的 Detour 机制通过 commRankIds 参数显式指定哪些 rank 参与实际通信:
// all2all_small_fullmesh_with_detour.h
bool IsValidCommRank() {
for (int i = 0; i < commRankSize; i++) {
if (commRankIds[i] == rank) { commRankIdx = i; return true; }
}
return false; // 当前 rank 是 detour rank，不直接参与通信
}
两种角色:
1. CommRank (通信 rank): 实际生产/消费数据的 rank，在 commRankIds 列表中
2. DetourRank (绕行 rank): 不在 commRankIds 中，仅做数据转发
6.3 Detour 数据流 (以 all_to_all_detour 为例)
Stage 1: InputToShareSlice Stage 2: ShareToShareSlice
───────────────────────── ──────────────────────────
CommRank A ──► DetourRank X ──► CommRank B CommRank B ◄── DetourRank X ◄── CommRank A
│ │ │ │ │ │
│ 1. A 读取本地 │ │ │ 4. B 从 X 的 │ │
│ input 数据 │ │ │ 共享内存读 │ │
│ 2. A 写入 X │ │ │ 取数据 │ │
│ 的共享内存 │ │ │ 5. B 写入 │ │
│ 3. A 写 flag │ │ │ output │ │
│ 通知 X │ │ │ │ │
6.4 为什么需要 Detour — 拥塞与拓扑不匹配
问题 1: 拓扑不匹配 - HCCL 原语（AlltoAll、ReduceScatter）假设 fullmesh 或特定分层拓扑 - 当实际拓扑是部分连接（如某些 rank 对之间无直连）时，HCCL 可能无法正确路由 - Detour 通过软件显式指定转发路径，绕过拓扑限制
问题 2: 跨 pod 拥塞 - 跨 pod 带宽远低于 pod 内带宽（HCCS vs RoCE） - 直接 All-to-All 会导致跨 pod 链路热点 - Detour 可以将跨 pod 流量分散到多个中间 rank，负载均衡
问题 3: 小规模通信优化 - 算子名中的 "small" 暗示这是针对小数据量场景的优化 - 小数据量下 HCCL 启动开销占比高，使用 SHMEM 直传 + detour 可降低延迟
6.5 Detour 与标准 All-to-All 的复杂度对比
| 维度 | 标准 All-to-All | Detour All-to-All |
|---|---|---|
| 跳数 | 1 跳 (直连) | 2 跳 (经 detour rank) |
| 总通信量 | N×(N-1) 条消息 | N×(N-1)×2 条消息 (翻倍) |
| 单跳距离 | 可能很远 (跨 pod) | 总是很短 (同 pod) |
| 带宽利用 | 跨 pod 瓶颈 | 利用 pod 内高带宽 |
| 延迟 | 单次长延迟 | 两次短延迟 (可能更优) |
适用场景: 当跨 pod 延迟 >> pod 内延迟时，2×短延迟 < 1×长延迟，detour 反而更快。
7. SHMEM 使用
7.1 SHMEM 在四个算子中的使用
| 算子 | SHMEM 类型 | 获取方式 | 用途 |
|---|---|---|---|
| a2e | HCCL IPC Window | epWinContext_->localWindowsIn + remoteRes[i].windowsIn |
Attention↔Expert 数据传输 |
| e2a | HCCL IPC Window | 同上 | Expert↔Attention 数据回传 |
| all_to_all_detour | peerMems (CommArgs) | commArgs->peerMems[rank] |
Detour rank 间数据中转 |
| reduce_scatter_detour | peerMems (CommArgs) | 同上 | 归约中间结果共享 |
7.2 两种 SHMEM 机制对比
HCCL IPC Window (a2e/e2a):
- 由 HCCL 初始化时分配，通过 GetHcclContext 获取
- 结构: HcclOpResParam.localWindowsIn (本地窗口) + remoteRes[i].nextDevicePtr->windowsIn (远端窗口)
- 大小: IPC_BUFF_MAX_SIZE_MUL_EXP = 800MB (a2e 中) / IPC_BUFF_MAX_SIZE = 100MB (comm_args.h)
- 布局: 前 2MB 为 flag 区 (IPC_DATA_OFFSET = 2MB)，后续为数据区
peerMems (detour 系列):
- 由 CommArgs 结构体携带，在算子初始化时由框架分配
- peerMems[rankId] 直接给出每张卡的可访问内存基址
- 通过 InitByCamCommDetour 初始化:
cpp
peerMemsAddrGm_.SetGlobalBuffer(&(commArgs)->peerMems[0], CAM_MAX_RANK_SIZE);
for (int i = 0; i < rankSize; i++)
shareAddrs[i] = (GM_ADDR)peerMemsAddrGm_.GetValue(i);
7.3 同步原语
a2e/e2a 的自定义 flag 同步:
// 发送端: 写 flag
LocalTensor<uint64_t> flagLt;
flagLt.SetValue(0, mergeMagicWithValue(magic, 0));
camCpUB2GM(shareFlagGt, flagLt, 1, sizeof(uint64_t));
// 接收端: 轮询 flag
while(1) {
DATA_FULSH(flagGt, uint32_t); // cache flush
if (flagGt.GetValue(0) == magic) return;
}
detour 系列的 SyncCollectives:
- SetSyncFlag(magic, value, eventID, rank): 在指定 rank 的共享内存写 flag
- WaitSyncFlag(magic, value, eventID, rank): 轮询等待 flag 到位
- SetInnerFlag / WaitInnerFlag: 卡内 block 间同步
- SetOuterFlag / WaitOuterFlag: 跨卡同步
- SetChunkFlag / GetChunkFlag: 分块传输同步
Magic 编码:
int64_t MergeMagicWithValue(int32_t magic, int32_t value) {
return (static_cast<int64_t>(static_cast<uint32_t>(magic)) << 32)
| static_cast<int64_t>(value);
}
高 32 位为 magic（防止不同轮次的 flag 误匹配），低 32 位为具体值。
7.4 Ping-Pong 双缓冲
a2e/e2a 和 detour 系列都使用 magic % PING_PONG_SIZE 实现双缓冲:
// a2e
shareAddrs[rank] = localWindowsIn + (magic % 2) * 800MB + rank * 512;
// detour
return (rankSize * srcRank + dstRank) * memLen + memRank * 512
+ (magic % 2) * (memLen / 2);
作用: 避免上一轮通信的数据被本轮覆盖。奇数轮使用内存前半，偶数轮使用后半。
8. 与现有算子链的关系
8.1 MoE EP 通信流水线全景
┌─────────────────────────────────────────────────────────────────────────┐
│ MoE Expert Parallelism 前向流水线 │
├─────────────────────────────────────────────────────────────────────────┤
│ │
│ Attention 卡簇 (rank ≥ expertRankSize) │
│ ┌──────────────────────────────────────────────┐ │
│ │ 1. Attention 计算 │ │
│ │ 2. Gate/Routing → expert_ids, scales │ │
│ │ 3. a2e() ─────────────发送 token + 路由 ─────┼────► Expert 卡簇 │
│ │ 4. e2a() ◄────────────接收 Expert 结果 ──────┼───── (rank < epSize) │
│ │ 5. 后续 Attention 层 │ │
│ └──────────────────────────────────────────────┘ │
│ │
│ Expert 卡簇 (rank < expertRankSize) │
│ ┌──────────────────────────────────────────────┐ │
│ │ 3. a2e() ◄─────────────接收 token + 路由 ───┼───── Attention 卡簇 │
│ │ 4. moe_dispatch_* (EP 内 AlltoAll) │ │
│ │ 5. Expert FFN 计算 │ │
│ │ 6. moe_combine_* (EP 内 AlltoAll) │ │
│ │ 7. e2a() ─────────────发送结果 ─────────────┼────► Attention 卡簇 │
│ └──────────────────────────────────────────────┘ │
│ │
│ 跨 Pod 场景 (rankSize > commRankSize) │
│ ┌──────────────────────────────────────────────┐ │
│ │ all_to_all_detour() — 经 detour rank 绕行 │ │
│ │ reduce_scatter_detour() — 经 detour rank 归约 │ │
│ └──────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────────────┘
8.2 新算子与现有算子的协作关系
| 现有算子 (master) | 新算子 (cam_feature) | 关系 |
|---|---|---|
moe_dispatch_normal_zb |
a2e | a2e 是 dispatch 的前置跨簇传输，将 token 从 Attention 卡送到 Expert 卡后，再由 moe_dispatch 做 EP 内分发 |
moe_combine_normal_zb |
e2a | e2a 是 combine 的后置跨簇传输，Expert 计算结果经 moe_combine 聚合后，由 e2a 送回 Attention 卡 |
moe_dispatch_layout_zb |
a2e | layout 算子管理 token-to-expert 映射，a2e 负责按映射传输 |
| HCCL AlltoAll | all_to_all_detour | detour 是 AlltoAll 在 non-fullmesh 拓扑下的替代实现 |
| HCCL ReduceScatter | reduce_scatter_detour | detour 是 ReduceScatter 在 non-fullmesh 拓扑下的替代实现 |
fused_deep_moe |
a2e/e2a | fused_deep_moe 融合 EP 内 dispatch+FFN+combine，a2e/e2a 处理跨簇部分 |
8.3 端到端数据流 (以 sample 为例)
# Attention 侧 (rank >= expert_rank_size)
# 1. 发送阶段
a2e_output = a2e(x, expert_ids, scales, ...)
# → expand_x (空), simulate_expert_ids, simulate_expert_scales, atten_batch_size, x_active_mask
# 2. 接收阶段
e2a_output = e2a(expand_x=x, atten_batch_size=atten_batch_size, ...)
# → x (Expert 计算结果)
# Expert 侧 (rank < expert_rank_size)
# 1. 接收阶段
a2e_output = a2e(dummy_x, dummy_ids, dummy_scales, ...)
# → expand_x (真实 token), simulate_expert_ids, simulate_expert_scales, atten_batch_size, x_active_mask
# 2. 发送阶段
e2a_output = e2a(expand_x=expand_x, atten_batch_size=atten_batch_size, ...)
# → x (空)
8.4 设计意图总结
- 解耦跨簇通信与 EP 内通信: a2e/e2a 处理 Attention↔Expert 跨簇传输，moe_dispatch/moe_combine 处理 EP 内传输，各司其职
- SHMEM 直传替代 HCCL collective: 对于小数据量或特定拓扑，SHMEM 直传延迟更低
- Detour 解决拓扑限制: 当 HCCL 原生 collective 无法适配非 fullmesh 拓扑时，detour 提供软件可控的绕行方案
- 精度提升: reduce_scatter_detour 支持 fp16 输入 + fp32 累加，避免低精度归约的精度损失
- 平台差异化: a2e/e2a 面向 910_93 (AI 集群)，detour 系列同时支持 910_93 和 910B (推理/混合场景)
附录 A: 关键常量汇总
| 常量 | 值 | 来源 | 含义 |
|---|---|---|---|
CAM_MAX_RANK_SIZE |
384 | comm_args.h | 最大支持 rank 数 |
IPC_BUFF_MAX_SIZE |
100MB | comm_args.h | IPC 数据区大小 |
IPC_BUFF_MAX_SIZE_MUL_EXP |
800MB | a2e.h | a2e 扩展 IPC 区 |
IPC_DATA_OFFSET |
2MB | comm_args.h | flag 区大小 |
PING_PONG_SIZE |
2 | comm_args.h | 双缓冲 |
OPT_RANK_OFFSET |
512 | a2e.h | rank 间内存偏移 |
MAX_CORE_NUM |
48 | all2all_small_fullmesh_with_detour.h | 最大 AIV 核数 |
FLAG_UNIT_INT_NUM |
4 | sync_collectives.h | flag 单元 int64 数 |
MAGIC_OFFSET |
32 | sync_collectives.h | magic 高位偏移 |
UB_SINGLE_DMA_SIZE_MAX |
190KB | comm_args.h | 单次 DMA 最大量 |
UB_SINGLE_TOTAL_SIZE_MAX |
192KB | comm_args.h | UB 总量 |
附录 B: 文件索引
src/cam/comm_operator/ascend_kernels/
├── a2e/
│ ├── op_kernel/a2e.cpp # 设备入口
│ ├── op_kernel/a2e.h # A2e 模板类（核心实现）
│ ├── op_kernel/a2e_tiling.h # Tiling 数据结构
│ ├── op_host/a2e.cpp # Tiling 函数 + 算子定义
│ └── op_api/aclnn_a2e.cpp/.h # aclnn API 封装
├── e2a/
│ ├── op_kernel/e2a.cpp # 设备入口
│ ├── op_kernel/e2a.h # E2a 模板类（核心实现）
│ ├── op_kernel/e2a_tiling.h # Tiling 数据结构
│ ├── op_host/e2a.cpp # Tiling 函数 + 算子定义
│ └── op_api/aclnn_e2a.cpp/.h # aclnn API 封装
├── all_to_all_detour/
│ ├── op_kernel/all2_all_detour.cpp # 设备入口
│ ├── op_kernel/all2all_small_fullmesh_with_detour.h # 核心算法
│ ├── op_host/all2_all_detour.cpp # Tiling 函数 + 算子定义
│ ├── op_host/all2_all_detour_tiling.h # Tiling 数据结构
│ └── op_api/aclnn_all2_all_detour.cpp/.h # aclnn API 封装
├── reduce_scatter_detour/
│ ├── op_kernel/reduce_scatter_detour.cpp # 设备入口
│ ├── op_kernel/reduce_scatter_detour.h # 核心算法
│ ├── op_kernel/reduce_scatter_detour_tiling.h # Tiling 数据结构
│ ├── op_host/reduce_scatter_detour.cpp # Tiling 函数 + 算子定义
│ ├── op_host/reduce_scatter_detour_tiling.cpp # Tiling 实现
│ └── op_api/aclnn_reduce_scatter_detour.cpp/.h # aclnn API 封装
├── utils/op_kernel/
│ ├── collectives.h # Detour 基类（CpGM2GMPingPong, InitByCamCommDetour）
│ ├── sync_collectives.h # Flag 同步原语
│ ├── comm_args.h # CommArgs 结构体 + 全局常量
│ ├── moe_distribute_base.h # HCCL 拓扑/资源结构体
│ └── data_copy.h # GM-GM/GM-UB 拷贝辅助
├── pybind/
│ ├── a2e.cpp # PyTorch 算子注册
│ └── e2a.cpp # PyTorch 算子注册
└── operator_registry.json # 算子注册表
附录 C: 关键函数调用链
a2e 调用链
a2e() [op_kernel/a2e.cpp]
└─ A2e<bf16,bf16,false>::init()
├─ GetHcclContext<HCCL_GROUP_ID_0>() // 获取 IPC 窗口
├─ AtomicAdd → magicTensor_ // 分布式 magic 分配
└─ shareAddrs[] 计算 // 每 rank IPC 地址映射
└─ A2e::process()
├─ [Attention 卡] sendWithMte()
│ ├─ sendX() → copyGmToGmWithBlocks + flag
│ ├─ sendExpertIds() → CpGM2GMPingPong + flag
│ └─ sendExpertScales() → CpGM2GMPingPong + flag
└─ [Expert 卡] recvWithMte()
├─ waitFlagWithScalar() → DATA_FULSH 轮询
├─ recvExpertIdsWithMte() → camCpGM2UB + camCpUB2GM
└─ copyGmToGmWithBlocks() → expandX
all_to_all_detour 调用链
all2_all_detour() [op_kernel/all2_all_detour.cpp]
└─ All2AllSmallFullmeshWithDetour<T>::Init()
├─ IsValidCommRank() // 检查是否参与通信
└─ Collectives::InitByCamCommDetour() // 初始化 peerMems 映射
└─ All2AllSmallFullmeshWithDetour<T>::Process()
├─ [blockIdx < coreNumPerStageX] InputToShareSlice()
│ ├─ blockIdx==0: 本地直拷
│ ├─ blockIdx==1: 首块 detour 数据 → 写 flag
│ └─ blockIdx≥2: 后续块 detour 数据 → 写 flag
└─ [blockIdx < coreNumPerStageY] ShareToShareSlice()
├─ blockIdx==0: 无操作
├─ blockIdx==1: WaitSyncFlag → 收首块 → output
└─ blockIdx≥2: WaitSyncFlag → 收后续块 → output
reduce_scatter_detour 调用链
reduce_scatter_detour() [op_kernel/reduce_scatter_detour.cpp]
└─ ReduceScatterWithDetour<fp16,fp32>::Init()
├─ IsValidCommRank()
├─ Collectives::InitByCamCommDetour()
└─ 计算 in2IpcCoreNum / reduceCoreNum / blockNum
└─ ReduceScatterWithDetour<T,U>::Process()
├─ [blockIdx < in2IpcCoreNum] Input2Ipc()
│ ├─ blockIdx==0: 本地数据 → 共享内存 + flag
│ └─ blockIdx≥1: 分片 → detour rank 共享内存 + flag
└─ [blockIdx < reduceCoreNum] Reduce()
├─ blockIdx==0: 无操作
└─ blockIdx≥1: 流水线归约
├─ WaitSyncFlag(本地/远端数据就绪)
├─ CpGM2GMPingPong(..., atomOp) // 归约
├─ SetOuterFlagDetour(通知下一级)
└─ WaitOuterFlagDetour(最后一级) → Ipc2Out
报告完
分析分支：
remotes/origin/cam_featurevsremotes/origin/master分析日期：2026-07-21 分析范围：共享工具（op_host / op_kernel utils）、构建系统、HFA 内存分配模块、pybind 层、operator_registry.json
目录
- ZB 共享工具变化
- 构建系统演进
- HFA (HBM Fast Alloc) 新模块
- pybind 层变化
- operator_registry.json 解读
- 从 master 到 cam_feature 的整体架构演进
1. ZB 共享工具变化
1.1 comm_args.h — 新增 memLen 字段与常量清理
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_kernel/comm_args.h
1.1.1 新增字段
// cam_feature 新增
int64_t memLen;
CommArgs 结构体新增 memLen 字段（位于 dfx[DFX_COUNT] 之后），用于传递内存长度参数。这是 ZB（Zero-Buffer）系列算子的关键参数——在 SHMEM 共享内存方案中，memLen 指定了预分配的共享内存窗口大小，算子通过该长度计算 flag 区域和数据区域的偏移。
1.1.2 删除的常量
| 删除的常量 | 原值 | 说明 |
|---|---|---|
NOTIFY_DISPATCH_BUFF_OFFSET |
204UL * 1024UL * 1024UL |
通知分发缓冲区偏移，已迁移到 tiling_args.h 的 Moe 命名空间 |
该常量在 master 中定义在 comm_args.h，cam_feature 中已删除（迁移到 tiling_args.h 的 NOTIFY_DISPATCH_WIN_OFFSET）。
1.1.3 关键常量保留
cam_feature 保留了以下核心常量：
constexpr int CAM_MAX_RANK_SIZE = 384; // 最大 NPU 卡数
constexpr int64_t IPC_BUFF_MAX_SIZE = 100 * 1024 * 1024; // IPC 缓冲区 100MB
constexpr int64_t IPC_DATA_OFFSET = 2 * 1024 * 1024; // 前 2MB 为 flag，后 100MB 为数据
constexpr int64_t UB_SINGLE_DMA_SIZE_MAX = 190 * 1024; // UB 单次 DMA 最大 190KB
constexpr int64_t UB_SINGLE_PING_PONG_ADD_SIZE_MAX = UB_SINGLE_DMA_SIZE_MAX / 2; // Ping-Pong 半区
constexpr int UB_ALIGN_SIZE = 32; // UB 对齐 32B
constexpr uint8_t COMM_NUM = 2; // 通信域数量（EP + TP）
constexpr uint8_t COMM_EP_IDX = 0;
constexpr uint8_t COMM_TP_IDX = 1;
constexpr int64_t STATE_WIN_OFFSET = 900 * 1024; // 状态窗口偏移
constexpr int64_t UB_SINGLE_TOTAL_SIZE_MAX = 192 * 1024; // UB 总大小 192KB
constexpr int64_t START_OFFSET_FOR_SHARE = 512; // 共享区域起始偏移
1.1.4 Op 枚举
enum Op : int {
COPYONLY = -1,
ADD = 0,
MUL = 1,
MAX = 2,
MIN = 3
};
该枚举定义了 CpGM2GMPingPong 的原子操作类型，用于 GM→GM 拷贝时的原子归约。
1.2 sync_collectives.h — Detour 原语与 Offset 重载
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_kernel/sync_collectives.h
这是 cam_feature 中变化最大的共享工具文件。SyncCollectives 类是 Ascend C kernel 间同步的核心原语库，管理两套 flag 区域（Memory A 卡内同步 / Memory B 卡间同步）。
1.2.1 新增 InitDetour — 共享内存偏移初始化
__aicore__ inline void InitDetour(int rank, int rankSize, GM_ADDR *shareAddrs,
int64_t shareMemOffset, TBuf<QuePosition::VECCALC> &tBuf)
{
// ...
localSyncAddr = (__gm__ int64_t*)(shareAddrs[rank] + shareMemOffset);
basicSyncAddr = (__gm__ int64_t*)(shareAddrs[rank] + shareMemOffset) + GetBlockIdx() * FLAG_UNIT_INT_NUM;
blockOuterSyncAddr = (__gm__ int64_t*)(shareAddrs[rank] + shareMemOffset) + segmentCount + GetBlockIdx() * FLAG_UNIT_INT_NUM;
// ...
}
设计意图： 在 ZB 系列中，共享内存窗口被划分为多个逻辑区域（combine 区域、notify-dispatch 区域等）。shareMemOffset 允许同一算子内的不同 kernel 使用同一块共享内存的不同偏移段，避免 flag 区域冲突。这是 Detour 算子的核心需求——detour 算子需要在同一共享内存窗口内复用同步原语。
1.2.2 新增 WaitSyncFlag 6 参数重载
__aicore__ inline void WaitSyncFlag(int32_t magic, int32_t value, int32_t eventID,
int32_t rank, int64_t flagNum, int64_t offset)
{
int64_t v = MergeMagicWithValue(magic, value);
WaitOneRankPartFlag((__gm__ int64_t*)(shareAddrs[rank] + offset) + eventID * FLAG_UNIT_INT_NUM, flagNum, v);
}
相比 5 参数版本，新增 offset 参数，支持在共享内存中指定偏移后等待 flag。
1.2.3 新增 SetOuterFlagDetour / WaitOuterFlagDetour
__aicore__ inline void SetOuterFlagDetour(int32_t magic, int32_t eventID,
int64_t setRank, int64_t setBlock, int64_t shareMemOffset)
{
__gm__ int64_t* flagAddr = GetOuterFlagAddrDetour(setRank, setBlock, shareMemOffset);
int64_t value = MergeMagicWithValue(magic, eventID);
SetFlag(flagAddr, value);
}
__aicore__ inline void WaitOuterFlagDetour(int32_t magic, int32_t eventID,
int64_t waitRank, int64_t waitBlock, int64_t shareMemOffset)
{
int64_t value = MergeMagicWithValue(magic, eventID);
__gm__ int64_t* flagAddr = GetOuterFlagAddrDetour(waitRank, waitBlock, shareMemOffset);
WaitOneRankPartFlag(flagAddr, 1, value);
}
1.2.4 新增 GetOuterFlagAddrDetour
__aicore__ inline __gm__ int64_t* GetOuterFlagAddrDetour(int64_t flagRank, int64_t flagBlock, int64_t shareMemOffset)
{
return (__gm__ int64_t*)(shareAddrs[flagRank] + shareMemOffset) + segmentCount + flagBlock * FLAG_UNIT_INT_NUM;
}
总结： 所有 Detour 变体都在原始地址计算基础上增加了 shareMemOffset，使得同步原语可以在共享内存窗口的任意偏移处工作。这是 ZB 系列算子（moe_combine_*_zb、moe_dispatch_*_zb、moe_notify_dispatch_zb）的核心基础设施。
1.2.5 同步原语内存布局
共享内存窗口 (peerMems[rank])
├── [0, segmentCount * FLAG_UNIT_INT_NUM * 8) ← Memory A: 卡内同步 flag
│ ├── block 0: [flag0, flag1, flag2, flag3] (FLAG_UNIT_INT_NUM = 4)
│ ├── block 1: [flag0, flag1, flag2, flag3]
│ └── ...
├── [segmentCount * 4 * 8, ...) ← Memory B: 卡间同步 flag
│ ├── block 0: [flag0, flag1, flag2, flag3]
│ └── ...
└── [IPC_CHUNK_FLAG, ...) ← Chunk Flag 区域 (IPC 用)
segmentCount = blockNum * FLAG_UNIT_INT_NUMFLAG_UNIT_INT_NUM = 4（每个同步单元 4 个 int64_t = 32 Bytes）MAGIC_OFFSET = 32（magic 占高 32 bits，value 占低 32 bits）
1.3 moe_distribute_base.h — 结构体无变化
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_kernel/moe_distribute_base.h
git diff 显示 master 与 cam_feature 之间无差异。该文件定义了 HCCL 通信库的核心资源描述结构：
| 结构体 | 用途 |
|---|---|
HcclSignalInfo |
信号资源（resId、addr、devId、tsId、rankId、flag） |
HcclStreamInfo |
流信息（streamId、sqId、cqId、logicCqId） |
LocalResInfoV2 |
本地资源 V2（64 个 signal、19 个 stream、2 个 AICPU notify） |
AlgoTopoInfo |
拓扑信息（rank、rankSize、superPod、RDMA、NIC 等） |
HcclOpConfig |
算子配置（deterministic、retry、highPerf、floatOverflowMode） |
HcclOpResParam |
算子资源参数（workspace、window、config、zero-copy ring buffer 等） |
HDCommunicateParams |
HD 通信参数（host/device addr、cache、flag） |
关键常量：
constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64; // 最大本地 notify 数
constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19; // 最大本地 stream 数
constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2; // AICPU notify 数
constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024; // AICPU 最大 rank 数
1.4 tiling_args.h — 窗口偏移常量
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_host/tiling_args.h
namespace Moe {
constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 4U * 1024UL * 1024UL; // 4MB
constexpr uint64_t NOTIFY_DISPATCH_WIN_OFFSET = 204U * 1024UL * 1024UL; // 204MB
} // namespace Moe
该文件定义了 ZB 系列算子的共享内存窗口偏移常量：
- COMBINE_STATE_WIN_OFFSET (4MB)：combine 算子的状态窗口偏移
- NOTIFY_DISPATCH_WIN_OFFSET (204MB)：notify-dispatch 算子的窗口偏移
这些常量在 master 中已存在，cam_feature 保持不变。
1.5 mc2_tiling_utils.h — HCCL_BUFFSIZE 解析
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_host/mc2_tiling_utils.h
class Mc2TilingUtils {
public:
static uint64_t GetMaxWindowSize() {
uint16_t defaultWindowSize = 200; // 默认 200MB
if (getenv("HCCL_BUFFSIZE") != nullptr) {
try {
std::string envStr(getenv("HCCL_BUFFSIZE"));
defaultWindowSize = std::stoi(envStr);
} catch (...) { /* 异常处理 */ }
}
return static_cast<uint64_t>(defaultWindowSize) * 1024UL * 1024UL;
}
};
通过环境变量 HCCL_BUFFSIZE 动态获取 HCCL 缓冲区大小（单位 MB），用于计算最大窗口尺寸。默认 200MB。
1.6 DFX 工具链（dfx_base.h / ops_log.h / ops_error.h / error_log.h）
1.6.1 dfx_base.h — 统一日志基础设施
class DfxBase {
public:
static constexpr const int MAX_LOG_LEN = 16000;
static constexpr const int MSG_HDR_LEN = 200;
static inline uint64_t GetTid(); // 获取线程 ID
static inline std::string GetOpInfo(const gert::TilingContext *context); // 获取算子信息
// ... 支持 TilingContext / TilingParseContext / InferShapeContext / InferDataTypeContext
};
核心宏：
#define OPS_LOG_STUB(MOD_ID, LOG_LEVEL, OPS_DESC, FMT, ...)
// 统一日志记录，包含文件名、行号、函数名、TID、算子名
#define OPS_INNER_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...)
// 内部错误报告 + 日志
#define OPS_CALL_ERR_STUB(ERR_CODE_STR, OPS_DESC, FMT, ...)
// 调用错误报告 + 日志
1.6.2 ops_log.h — 分级日志宏
#define OPS_LOG_D(OPS_DESC, ...) OPS_LOG_STUB_D(OPS_DESC, __VA_ARGS__) // Debug
#define OPS_LOG_I(OPS_DESC, ...) OPS_LOG_STUB_I(OPS_DESC, __VA_ARGS__) // Info
#define OPS_LOG_W(OPS_DESC, ...) OPS_LOG_STUB_W(OPS_DESC, __VA_ARGS__) // Warn
#define OPS_LOG_E(OPS_DESC, ...) OPS_INNER_ERR_STUB("EZ9999", OPS_DESC, __VA_ARGS__) // Error + Report
#define OPS_LOG_EVENT(OPS_DESC, ...) OPS_LOG_STUB_EVENT(OPS_DESC, __VA_ARGS__) // Event
条件日志宏：
#define OPS_LOG_D_IF(COND, OP_DESC, EXPR, ...)
#define OPS_LOG_E_IF_NULL(OPS_DESC, PTR, EXPR) // 空指针检查 + 错误报告
1.6.3 ops_error.h — 错误报告宏
#define OPS_REPORT_VECTOR_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E89999", OPS_DESC, __VA_ARGS__)
#define OPS_REPORT_CUBE_INNER_ERR(OPS_DESC, ...) OPS_INNER_ERR_STUB("E69999", OPS_DESC, __VA_ARGS__)
#define OPS_ERR_IF(COND, LOG_FUNC, EXPR) OPS_LOG_STUB_IF(COND, LOG_FUNC, EXPR)
错误码体系：
- EZ9999：通用错误
- E89999：Vector 内核内部错误
- E69999：Cube 内核内部错误
1.6.4 error_log.h — 轻量级 printf 日志
#define OP_LOGI(opname, ...) // Info 级别（空实现）
#define OP_LOGW(opname, ...) // Warn 级别（printf）
#define OP_LOGE_WITHOUT_REPORT(opname, ...) // Error 不报告
#define OP_LOGE(opname, ...) // Error 级别（printf）
#define OP_LOGD(opname, ...) // Debug 级别（空实现）
这是 tiling 侧的轻量级日志，使用 printf 而非 Alog。
1.7 data_copy.h — GM↔GM Ping-Pong 拷贝
文件路径： src/cam/comm_operator/ascend_kernels/utils/op_kernel/data_copy.h
核心函数 CpGM2GMPingPong：
template <typename K, typename U = K>
FORCE_INLINE_AICORE void CpGM2GMPingPong(int64_t dataSizeRemain,
const GlobalTensor<U> &sendDataInputGt,
const GlobalTensor<K> &recvDataOutputGT, int op)
设计要点：
1. Ping-Pong 双缓冲：UB 分为两个区域（UB_HEAD_OFFSET 和 UB_MID_OFFSET），MTE2（GM→UB）和 MTE3（UB→GM）流水线化
2. 类型转换支持：当 K != U 时（如 half→float），UB 按 sizeof(U):sizeof(K) 比例分割
3. 原子操作：通过 SetAtomic<K>(op) 支持 ADD/MUL/MAX/MIN 归约
4. 事件同步：使用 MTE3_MTE2、MTE2_MTE3、MTE2_V、V_MTE3 等硬事件
辅助函数：
template <typename T> void CpUB2GM(__gm__ T *gmAddr, __ubuf__ T *ubAddr, uint32_t size);
template <typename T> void CpGM2UB(__ubuf__ T *ubAddr, __gm__ T *gmAddr, uint32_t size);
template<typename T> void CopyUB2UB(__ubuf__ T *dst, __ubuf__ T *src, const uint32_t calCount);
template<typename T> void copyGmToGmWithBlocks(...); // 多核分块 GM→GM
2. 构建系统演进
2.1 operator_registry.json — 从 5 个算子到 13 个算子
2.1.1 master 版本（5 个算子）
{
"soc_versions": {
"ascend910_93": [
"dispatch_layout",
"fused_deep_moe",
"moe_combine_normal",
"moe_dispatch_normal",
"notify_dispatch"
]
},
"operator_meta": {
"dispatch_layout": { "requires_shmem": false },
"fused_deep_moe": { "requires_shmem": false, "family": "fused_deep_moe", "variant": "base" },
"moe_combine_normal": { "requires_shmem": false },
"moe_dispatch_normal": { "requires_shmem": false },
"notify_dispatch": { "requires_shmem": false }
}
}
2.1.2 cam_feature 版本（13 个算子）
{
"soc_versions": {
"ascend910_93": [
"a2e",
"all_to_all_detour",
"e2a",
"fused_deep_moe",
"fused_deep_moe_fwk",
"fused_deep_moe_w4a8",
"moe_combine_lowlatency_zb",
"moe_combine_normal_zb",
"moe_dispatch_layout_zb",
"moe_dispatch_lowlatency_zb",
"moe_dispatch_normal_zb",
"moe_notify_dispatch_zb",
"reduce_scatter_detour"
]
}
}
2.1.3 新增算子清单
| 算子名 | requires_shmem | family | variant | 说明 |
|---|---|---|---|---|
a2e |
- | - | - | Attention→Expert 数据搬运 |
e2a |
- | - | - | Expert→Attention 数据搬运 |
all_to_all_detour |
- | - | - | All-to-All Detour 通信 |
reduce_scatter_detour |
- | - | - | Reduce-Scatter Detour 通信 |
fused_deep_moe_fwk |
false | fused_deep_moe | fwk | 独立 FWK 变体 |
fused_deep_moe_w4a8 |
false | fused_deep_moe | w4a8 | W4A8 量化变体 |
moe_combine_lowlatency_zb |
true | - | - | ZB 低延迟 Combine |
moe_combine_normal_zb |
true | - | - | ZB 标准 Combine |
moe_dispatch_layout_zb |
true | - | - | ZB Dispatch 布局 |
moe_dispatch_lowlatency_zb |
true | - | - | ZB 低延迟 Dispatch |
moe_dispatch_normal_zb |
true | - | - | ZB 标准 Dispatch |
moe_notify_dispatch_zb |
true | - | - | ZB Notify-Dispatch |
2.1.4 SHMEM 依赖分析
requires_shmem = true（7 个）： 全部为 ZB 系列算子
- moe_combine_lowlatency_zb
- moe_combine_normal_zb
- moe_dispatch_layout_zb
- moe_dispatch_lowlatency_zb
- moe_dispatch_normal_zb
- moe_notify_dispatch_zb
requires_shmem = false（2 个）：
- fused_deep_moe_fwk
- fused_deep_moe_w4a8
未标记 requires_shmem（4 个）： 默认为 false
- a2e, e2a, all_to_all_detour, reduce_scatter_detour
2.1.5 fused_deep_moe 家族互斥规则
fused_deep_moe (base) ──┐
├── 互斥（共享 fused_deep_moe.* 文件名）
fused_deep_moe_w4a8 ────┘
fused_deep_moe_fwk ───── 独立（fused_deep_moe_fwk.* 文件名），可与 base/w4a8 共存
base和w4a8共享fused_deep_moe.h/fused_deep_moe.cpp，互斥fwk使用独立文件名fused_deep_moe_fwk.*，可与任意变体共存--quant标志选择w4a8变体
2.2 select_ops.py — 算子选择逻辑
文件路径： build/cam/comm_operator/select_ops.py
2.2.1 核心 resolve 函数流程
输入: registry, soc_list, shmem, quant, user_ops
│
├─ 1. 确定候选集: user_ops 或全量 soc_list
│
├─ 2. fused_deep_moe 互斥规则
│ ├─ quant=true: 强制 w4a8，移除 base
│ ├─ quant=false + 两者都在: 保留 base，移除 w4a8
│ └─ fwk 不受影响
│
├─ 3. SHMEM 过滤
│ └─ shmem=false: 移除所有 requires_shmem=true 的算子
│
└─ 4. 按注册表顺序排序输出
2.2.2 命令行参数
python3 select_ops.py \
--registry operator_registry.json \
--soc ascend910_93 \
--shmem 1 \
--quant \
--ops "a2e;e2a;fused_deep_moe"
2.2.3 master vs cam_feature 差异
select_ops.py 在 master 和 cam_feature 之间逻辑完全一致，仅注释语言不同（master 英文，cam_feature 中文）。变化的是 operator_registry.json 的内容。
2.3 compile_ascend_proj.sh — SHMEM 注入与中文注释
文件路径： build/cam/comm_operator/compile_ascend_proj.sh
2.3.1 SHMEM 检测逻辑
shmem_installed=1
if [ -z "${SHMEM_HOME_PATH}" ]; then
shmem_installed=0
fi
通过 SHMEM_HOME_PATH 环境变量判断 SHMEM 是否安装。
2.3.2 SHMEM 编译选项注入（关键 patch）
# Patch kernel compile .make files to pass custom include paths to OPC tool.
KERNEL_COMPILE_OPTS=""
if [ -n "${SHMEM_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="-I${SHMEM_HOME_PATH}/shmem/include -I${SHMEM_HOME_PATH}/shmem/src/device"
fi
if [ -n "${CATLASS_HOME_PATH}" ]; then
KERNEL_COMPILE_OPTS="${KERNEL_COMPILE_OPTS}-I${CATLASS_HOME_PATH}/include"
fi
if [ -n "${KERNEL_COMPILE_OPTS}" ]; then
find build_out/op_kernel/CMakeFiles -name "build.make" \
-exec sed -i "s|--compile-options=\"\"|--compile-options=\\\"${KERNEL_COMPILE_OPTS}\\\"|g" {} +
fi
设计意图： npu_op_* 构建系统的 simple_kernel_compile 不读取 custom_compile_options.ini，导致 --compile-options 为空。该 patch 通过 sed 直接修改 .make 文件，将 SHMEM 和 Catlass 的 include 路径注入到编译选项中。
2.3.3 Include-List 构建模式
copy_ops_include() {
# 始终复制 utils 公共头目录
cp -rf "$src_dir/utils/op_host/"* "$dst_dir/op_host/"
cp -rf "$src_dir/utils/op_kernel/"* "$dst_dir/op_kernel/"
# 仅复制选中算子的 op_host/op_kernel/op_api
for name in "${ops[@]}"; do
cp -rf "$subdir/op_host/"* "$dst_dir/op_host/"
cp -rf "$subdir/op_kernel/"* "$dst_dir/op_kernel/"
cp -rf "$subdir/op_api/"* "$dst_dir/pregen/build_out/autogen/"
done
}
相比旧版全量拷贝 + exclude-list，新版只复制 select_ops.py 求解出的算子，从根本上避免同名算子文件覆盖。
2.3.4 master vs cam_feature 差异
compile_ascend_proj.sh 在 master 和 cam_feature 之间逻辑完全一致，仅注释语言不同（master 英文，cam_feature 中文）。
2.4 ascendc_impl_build.py — SHMEM 路径修复
文件路径： src/cam/comm_operator/ascend_kernels/cmake_files/cmake/util/ascendc_impl_build.py
2.4.1 关键 Bug Fix
# master（有 bug）
options.append("-I" + SHMEM_HOME_PATH + "shmem/include/device") # 缺少 / 分隔符
# cam_feature（修复）
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device") # 正确路径
影响： 在 master 中，如果 SHMEM_HOME_PATH 以 / 结尾则正常，否则路径拼接错误。cam_feature 修复了此问题。
2.4.2 SHMEM Include 路径列表
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/host_device")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/include/internal")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src")
options.append("-I" + SHMEM_HOME_PATH + "/shmem/src/device")
options.append("-I" + SHMEM_HOME_PATH + "/memfabric_hybrid/include/smem/device")
2.5 CMakeLists.txt 演进
2.5.1 op_kernel/CMakeLists.txt — Catlass 架构宏
# cam_feature 新增
list(APPEND KERNEL_COMPILE_OPTS -I${CMAKE_CURRENT_SOURCE_DIR} -DASCENDC_DUMP=0 --cce-auto-sync=off -DCATLASS_ARCH=2201)
变化：
1. 新增 -I${CMAKE_CURRENT_SOURCE_DIR}：将当前源码目录加入全局头搜索路径
2. 新增 -DCATLASS_ARCH=2201：Catlass 架构版本宏（2201 对应 Ascend 910B）
2.5.2 op_host/CMakeLists.txt — CANN 版本检测
set(CANN83_PATH "${CANN_ROOT}/include/platform/platform_infos_def.h")
set(CANN82_PATH "${CANN_ROOT}/include/experiment/platform/platform/platform_infos_def.h")
if(EXISTS "${CANN83_PATH}")
add_definitions(-DUSE_CANN83_PATH)
elseif(EXISTS "${CANN82_PATH}")
add_definitions(-DUSE_CANN82_PATH)
endif()
自动检测 CANN 8.2 vs 8.3+ 版本，定义对应的宏。
2.5.3 顶层 CMakeLists.txt
cmake_minimum_required(VERSION 3.16.0)
project(opp)
find_package(ASC REQUIRED)
set(package_name ${vendor_name})
npu_op_package(${package_name} TYPE RUN CONFIG INSTALL_PATH ${CMAKE_BINARY_DIR}/)
add_subdirectory(op_host)
add_subdirectory(op_kernel)
使用 npu_op_package 宏定义算子包，类型为 RUN（run 包）。
3. HFA (HBM Fast Alloc) 新模块
文件路径： src/cam/utilities/hfa/
HFA 是 cam_feature 中全新引入的模块，master 中不存在。
3.1 模块概览
hfa/
├── .gitignore
├── CMakeLists.txt
├── README.md
├── doc/
│ └── HFA.pdf
├── include/
│ └── hbm_fast_alloc.h ← 对外头文件
├── src/
│ ├── CMakeLists.txt
│ ├── acl_free_chunks.cpp ← 空闲块管理
│ ├── acl_virtual_memory.cpp ← 虚拟内存管理
│ ├── acl_virtual_memory.h ← 虚拟内存头文件
│ ├── flex_alloc.cpp ← 弹性分配器
│ ├── flex_alloc.h ← 弹性分配器头文件
│ ├── hbm_fast_alloc.cpp ← HBM 快速分配器
│ └── hfa_log.h ← 日志工具
└── test/
3.2 FlexAlloc — 双索引弹性分配器
文件路径： src/cam/utilities/hfa/src/flex_alloc.h / flex_alloc.cpp
3.2.1 数据结构
class FlexAlloc {
private:
struct Range { length_t size; };
// 主索引：按起始偏移排序
using StartMap = std::map<offset_t, Range>;
// 辅助索引：按大小排序，值为 StartMap 的迭代器
using SizeMap = std::multimap<length_t, typename StartMap::iterator>;
StartMap m_free_ranges_by_offset; // 按偏移索引的空闲区间
SizeMap m_free_ranges_by_size; // 按大小索引的空闲区间
size_t m_mem_size { 0 };
size_t m_total_alloc_bytes { 0 };
size_t m_total_free_bytes { 0 };
size_t m_total_alloc_objs { 0 };
};
设计亮点： 双索引结构——StartMap 支持 O(log n) 的邻居合并，SizeMap 支持 O(log n) 的最佳适配查找。SizeMap 存储的是 StartMap 的迭代器，避免数据冗余。
3.2.2 分配算法
FlexAlloc::FlexRetCode FlexAlloc::alloc(length_t req_size, size_t &offset, size_t &alloc_size)
{
alloc_size = round_up_bits(req_size, OBJ_SIZE_BITS); // 64 字节对齐
// 1. 在 SizeMap 中查找第一个 >= alloc_size 的区间
auto sizeIt = m_free_ranges_by_size.lower_bound(alloc_size);
// 2. 如果没找到，区分"内存不足"和"碎片化"
if (sizeIt == m_free_ranges_by_size.end()) {
if (get_total_free_mem() >= alloc_size)
return FLEX_RC_ERR_FRAGMENTATION; // 总空闲够但无连续块
else
return FLEX_RC_ERR_NO_MEM; // 总空闲不够
}
// 3. 分配并分割剩余
offset = rangeIt->first;
blkSz = rangeIt->second.size;
// 从两个索引中删除
// 如果有剩余，将新区间插回两个索引
}
时间复杂度： O(log n)（红黑树操作）
3.2.3 释放与合并
bool FlexAlloc::free(offset_t offset, length_t size)
{
// 1. 查找第一个起始位置 >= offset 的区间
auto next = m_free_ranges_by_offset.lower_bound(offset);
// 2. 尝试与前一个区间合并
if (prevEnd == offset) { /* 合并 */ }
// 3. 尝试与后一个区间合并
while (newStart + newSize == rangeStart) { /* 合并 */ }
// 4. 插入合并后的区间
}
3.2.4 对齐策略
static size_t const OBJ_SIZE_BITS = 6; // 2^6 = 64 字节对齐
static size_t const MIN_OBJ_SIZE = (1UL << OBJ_SIZE_BITS); // 最小 64 字节
template <typename T>
T round_up_bits(T num, uint32_t nbits) {
T shift = ((static_cast<T>(1)) << nbits) - 1;
return (num + shift) & ~(shift);
}
3.2.5 返回码
enum FlexRetCode {
FLEX_RC_SUCCESS, // 成功
FLEX_RC_ERROR, // 通用错误
FLEX_RC_ERR_INVAL_ARG, // 无效参数
FLEX_RC_ERR_NO_MEM, // 内存不足
FLEX_RC_ERR_FRAGMENTATION // 碎片化（总空闲够但无连续块）
};
3.3 AclVirtualMemory — ACL 虚拟内存管理
文件路径： src/cam/utilities/hfa/src/acl_virtual_memory.h / acl_virtual_memory.cpp
3.3.1 核心能力
class AclVirtualMemory {
public:
// 分配 HBM 内存：虚拟内存最大 128GB，物理内存最大 32GB
std::pair<VmRet, void*> alloc(int32_t device_id, size_t vr_mem_size,
size_t ph_mem_size, size_t ph_page_size);
// 碎片整理：将分散的物理页重新映射为连续块
VmRet defrag(AclFreeChunks* p_free_chunks);
void* get_base_addr() { return mp_vr_base_addr; }
void free();
};
3.3.2 关键常量
const size_t MAX_VIRTUAL_MEM_SIZE = 128UL * 1024 * 1024 * 1024; // 128GB
const size_t PH_PAGE_ALIGN_SIZE = 2 * 1024 * 1024; // 2MB 页对齐
const size_t DEF_INITIAL_FREE_PAGES = 1024;
const size_t DEF_INITIAL_FREE_CHUNKS = 1024;
3.3.3 碎片整理算法（defrag）
输入: 空闲块列表 (AclFreeChunks)
│
├─ 1. 遍历每个空闲块，提取完整物理页
│ ├─ 对齐的首页 → 加入 free_pages
│ ├─ 对齐的尾页 → 加入 free_pages
│ └─ 中间完整页 → 加入 free_pages
│
├─ 2. 重新映射完整页到虚拟内存末尾
│ for each page in free_pages:
│ if 完整页: aclrtMapMem 到末尾
│ else: 作为不完整块保留
│
└─ 3. 输出: 一个连续大块 + 若干不完整边块
关键 API：
- aclrtMapMem：将物理页映射到指定虚拟地址
- aclrtUnmapMem：解除映射
- aclrtFreePhysical：释放物理内存
- aclrtReleaseMemAddress：释放虚拟地址空间
3.3.4 错误码体系
enum class VmRet {
SUCCESS = 0,
E_OUT_OF_MEMORY = 0x100,
E_INSTANCE_NOT_INITIALIZED = 0x101,
E_INSUFFICIENT_MEM = 0x102,
E_INVALID_PARAMETER = 0x103,
// ...
E_ACL_MAP_MEM_FAILED = ACL_BASE_ERROR + 5,
E_ACL_UNMAP_MEM_FAILED = ACL_BASE_ERROR + 6,
};
3.4 HbmFastAlloc — 统一分配器
文件路径： src/cam/utilities/hfa/include/hbm_fast_alloc.h / src/cam/utilities/hfa/src/hbm_fast_alloc.cpp
3.4.1 类定义
class HbmFastAlloc {
public:
bool init(int32_t device_id, size_t vr_mem_size, size_t ph_mem_size,
size_t ph_page_size, size_t initial_free_pages);
void deinit();
void *alloc(size_t req_size);
bool free(void *ptr);
void free_all();
size_t get_last_defrag_elem_size() { return m_last_defrag_elem_size; }
static void set_log_level(int lvl);
static int get_log_level();
private:
bool try_defrag();
bool handle_defrag(size_t size, size_t &alloc_offset, size_t &alloc_size);
AclVirtualMemory *mp_avm { nullptr };
FlexAlloc *mp_fa { nullptr };
std::mutex m_alloc_map_lock;
AllocMap m_alloc_map; // ptr → (offset, size)
size_t m_last_defrag_elem_size { 0 };
size_t m_defrag_ctr { 0 };
size_t m_defrag_failed { 0 };
};
3.4.2 分配流程
alloc(size)
│
├─ 1. mp_fa->alloc(size, offset, alloc_size)
│ ├─ 成功 → 返回对象地址
│ └─ 失败且为 FRAGMENTATION → 进入 defrag 流程
│
├─ 2. handle_defrag(size, offset, alloc_size)
│ ├─ try_defrag()
│ │ ├─ 获取所有空闲区间
│ │ ├─ 调用 mp_avm->defrag() 重新映射物理页
│ │ └─ 将新区间插回 FlexAlloc
│ └─ 再次尝试 mp_fa->alloc()
│
└─ 3. 记录到 m_alloc_map，返回地址
3.4.3 释放流程
bool HbmFastAlloc::free(void *ptr)
{
std::lock_guard<std::mutex> lock(m_alloc_map_lock);
auto iter = m_alloc_map.find(static_cast<uint8_t*>(ptr));
if (iter == m_alloc_map.end()) return false;
// 通过 FlexAlloc 释放
mp_fa->free(iter->second.first, iter->second.second);
m_alloc_map.erase(iter);
return true;
}
3.5 hfa_log.h — 分级日志
enum HfaLogLevel {
HFA_LOG_DEBUG = 0,
HFA_LOG_INFO = 1,
HFA_LOG_WARNING = 2,
HFA_LOG_ERROR = 3
};
#define DEBUG_PRINT(...) HFA_LOG_PRINT_IF(HFA_LOG_DEBUG, stdout, __VA_ARGS__)
#define INFO_PRINT(...) HFA_LOG_PRINT_IF(HFA_LOG_INFO, stdout, __VA_ARGS__)
#define WARNING_PRINT(...) HFA_LOG_PRINT_IF(HFA_LOG_WARNING, stdout, __VA_ARGS__)
#define ERROR_PRINT(...) HFA_LOG_PRINT_IF(HFA_LOG_ERROR, stderr, __VA_ARGS__)
默认日志级别：HFA_LOG_WARNING。
3.6 HFA 与 Zero-Buffer 的关系
HFA 虽然位于 src/cam/utilities/hfa/ 公共目录，但其设计目标是为 ZB 系列算子提供共享内存窗口的动态分配能力：
- ZB 算子的内存需求： ZB 系列（
moe_combine_*_zb、moe_dispatch_*_zb、moe_notify_dispatch_zb）需要在 HBM 上预分配大块共享内存窗口，用于卡间 flag 同步和数据交换 - HFA 的角色： 提供高效的 HBM 分配/释放/碎片整理，使 ZB 算子能动态管理共享内存窗口
- 碎片整理： ZB 算子的内存分配模式（频繁 alloc/free 不同大小块）容易产生碎片，HFA 的
defrag通过 ACL API 重新映射物理页解决
性能指标（README 描述）： - ~99.7% HBM 利用率（所有对象大小） - ~0.5 µs 分配/释放延迟（超低且稳定）
4. pybind 层变化
文件路径： src/cam/comm_operator/pybind/pybind.cpp / functions.h
4.1 master 版本（4 个算子）
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fused_deep_moe", &FusedDeepMoeImplAutograd);
m.def("get_dispatch_layout", &GetDispatchLayoutImplAutograd);
m.def("moe_dispatch_prefill", &MoeDispatchPrefillImplAutograd);
m.def("moe_combine_prefill", &MoeCombinePrefillImplAutograd);
}
4.2 cam_feature 版本（10 个算子 + 1 个类）
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fused_deep_moe", &FusedDeepMoeImplAutograd);
m.def("all2_all_detour", &all2_all_detour_impl_autograd);
m.def("reduce_scatter_detour", &reduce_scatter_detour_impl_autograd);
m.def("cam_get_comm", &cam_get_comm);
m.def("cam_free_comm", &cam_free_comm);
m.def("cam_get_rank_size", &cam_get_rank_size);
m.def("cam_get_and_increase_magic", &cam_get_and_increase_magic);
m.def("cam_get_magic", &cam_get_magic);
m.def("a2e", &A2eImplAutograd);
m.def("e2a", &E2aImplAutograd);
// ZB Fused Deep MOE Buffer 类
pybind11::class_<fused_deep_moe::Buffer>(m, "FusedDeepMoeBuffer")
.def(pybind11::init<>())
.def("init", &fused_deep_moe::Buffer::init)
.def("zb_fused_deep_moe", &fused_deep_moe::Buffer::zb_fused_deep_moe)
.def("is_initialized", &fused_deep_moe::Buffer::is_initialized)
.def("get_ext_info", &fused_deep_moe::Buffer::get_ext_info)
.def("get_shmem_workspace", &fused_deep_moe::Buffer::get_shmem_workspace);
}
4.3 新增算子详解
| 算子 | 功能 | 参数 |
|---|---|---|
all2_all_detour |
All-to-All Detour 通信 | sendData, commRankIds, commArgs, commId |
reduce_scatter_detour |
Reduce-Scatter Detour | sendData, commRankIds, commArgs, commId, op |
cam_get_comm |
获取通信资源 | comm_id, rank, group_size, server_ip_port |
cam_free_comm |
释放通信资源 | comm_id |
cam_get_rank_size |
获取 rank 数量 | comm_id |
cam_get_and_increase_magic |
原子获取并递增 magic | comm_id |
cam_get_magic |
获取当前 magic | comm_id |
a2e |
Attention→Expert 数据搬运 | x, expert_ids, scales, batch_size, hidden_size, topk, ... |
e2a |
Expert→Attention 数据搬运 | expand_x, atten_batch_size, batch_size, hidden_size, topk, ... |
4.4 FusedDeepMoeBuffer 类
namespace fused_deep_moe {
class Buffer {
void init(int rank, int num_ranks, uint64_t memsize, const std::string &ip_port);
std::vector<at::Tensor> zb_fused_deep_moe(
const at::Tensor &x,
const at::Tensor &expert_ids,
const TensorVector &gmm1_weight,
// ... 完整参数列表
);
bool is_initialized() const;
int64_t get_ext_info() const;
int64_t get_shmem_workspace() const;
};
}
设计意图： FusedDeepMoeBuffer 封装了 ZB Fused Deep MOE 算子的状态管理：
- init：初始化共享内存工作区（通过 HCCL 通信）
- zb_fused_deep_moe：执行 ZB 融合 MOE 前向计算
- get_shmem_workspace：返回共享内存工作区指针（供其他算子使用）
4.5 TORCH_LIBRARY 注册
TORCH_LIBRARY(umdk_cam_op_lib, m) {
m.def("fused_deep_moe(...) -> Tensor[]");
m.def("all2_all_detour(Tensor sendData, Tensor commRankIds, Tensor commArgs, int commId) -> Tensor");
m.def("reduce_scatter_detour(Tensor sendData, Tensor commRankIds, Tensor commArgs, int commId, int op) -> Tensor");
m.def("cam_get_comm(int comm_id, int rank, int group_size, str server_ip_port) -> Tensor");
m.def("cam_free_comm(int comm_id) -> Tensor");
m.def("a2e(...) -> Tensor[]");
m.def("e2a(...) -> Tensor");
}
通过 TORCH_LIBRARY 注册为 PyTorch 自定义算子，支持 torch.compile 和 AOTInductor。
4.6 functions.h 签名变化
| master | cam_feature |
|---|---|
GetDispatchLayoutImplAutograd |
移除 |
MoeDispatchPrefillImplAutograd |
移除 |
MoeCombinePrefillImplAutograd |
移除 |
| - | all2_all_detour_impl_autograd (新增) |
| - | reduce_scatter_detour_impl_autograd (新增) |
| - | A2eImplAutograd (新增) |
| - | E2aImplAutograd (新增) |
5. operator_registry.json 解读
5.1 架构设计
operator_registry.json
├── _comment: 元数据说明
├── soc_versions: 代际 → 算子目录名列表
│ └── ascend910_93: [13 个算子]
└── operator_meta: 算子元数据
├── requires_shmem: 是否需要 SHMEM
├── family: 家族（用于互斥）
└── variant: 变体（base/w4a8/fwk）
5.2 算子分类
5.2.1 按功能分类
| 类别 | 算子 | 说明 |
|---|---|---|
| ZB MOE 核心 | moe_combine_*_zb, moe_dispatch_*_zb, moe_notify_dispatch_zb |
7 个，全部 requires_shmem=true |
| Fused MOE | fused_deep_moe, fused_deep_moe_w4a8, fused_deep_moe_fwk |
3 个变体 |
| 数据搬运 | a2e, e2a |
Attention↔Expert |
| 通信 Detour | all_to_all_detour, reduce_scatter_detour |
通信加速 |
5.2.2 按 SHMEM 依赖分类
| requires_shmem | 算子数 | 算子 |
|---|---|---|
| true | 6 | moe_combine_lowlatency_zb, moe_combine_normal_zb, moe_dispatch_layout_zb, moe_dispatch_lowlatency_zb, moe_dispatch_normal_zb, moe_notify_dispatch_zb |
| false | 2 | fused_deep_moe_fwk, fused_deep_moe_w4a8 |
| 未指定（默认 false） | 5 | a2e, e2a, all_to_all_detour, reduce_scatter_detour, fused_deep_moe |
5.3 family/variant 互斥规则
fused_deep_moe (family=fused_deep_moe, variant=base)
fused_deep_moe_w4a8 (family=fused_deep_moe, variant=w4a8)
└── 互斥：共享 fused_deep_moe.* 文件名
fused_deep_moe_fwk (family=fused_deep_moe, variant=fwk)
└── 独立：fused_deep_moe_fwk.* 文件名，可与 base/w4a8 共存
互斥规则代码实现：
if quant:
# -q 选择 w4a8，移除 base
candidates = [o for o in candidates if o != BASE]
if W4A8 not in candidates:
candidates.append(W4A8)
else:
# 默认选择 base，移除 w4a8
if has_base and has_w4a8:
candidates = [o for o in candidates if o != W4A8]
5.4 构建流程中的算子选择
build.sh
│
├─ CAM_USE_W4A8=1 → --quant
├─ CAM_OP_SELECT="a2e;e2a" --ops
└─ SHMEM_HOME_PATH → --shmem 1/0
│
└─ select_ops.py
│
├─ 1. 确定候选集
├─ 2. 应用互斥规则
├─ 3. SHMEM 过滤
└─ 4. 输出最终列表
6. 从 master 到 cam_feature 的整体架构演进
6.1 算子体系演进
master (5 个算子) cam_feature (13 个算子)
───────────────── ──────────────────────
dispatch_layout → (移除，替换为 ZB 版本)
fused_deep_moe → fused_deep_moe (保留)
moe_combine_normal → moe_combine_normal_zb (SHMEM 化)
moe_dispatch_normal → moe_dispatch_normal_zb (SHMEM 化)
notify_dispatch → moe_notify_dispatch_zb (SHMEM 化)
+ moe_combine_lowlatency_zb (新增)
+ moe_dispatch_lowlatency_zb (新增)
+ moe_dispatch_layout_zb (新增)
+ fused_deep_moe_w4a8 (新增量化变体)
+ fused_deep_moe_fwk (新增 FWK 变体)
+ a2e (新增)
+ e2a (新增)
+ all_to_all_detour (新增)
+ reduce_scatter_detour (新增)
6.2 共享内存架构演进
master： 使用 HCCL 标准通信，无自定义共享内存
cam_feature： 引入 SHMEM 共享内存窗口
共享内存窗口布局 (HBM)
├── [0, 4MB) ← COMBINE_STATE_WIN_OFFSET
│ └── combine 算子 flag 区域
├── [204MB, ...) ← NOTIFY_DISPATCH_WIN_OFFSET
│ └── notify-dispatch 算子 flag 区域
├── [IPC_CHUNK_FLAG, ...) ← IPC Chunk 标志区域
└── [peerMems[rank], ...) ← 每个 rank 的 flag 区域
├── Memory A: 卡内同步 flag
└── Memory B: 卡间同步 flag
6.3 同步原语演进
| 特性 | master | cam_feature |
|---|---|---|
| 基础 Init | ✓ | ✓ |
| Detour Init（偏移初始化） | ✗ | ✓ |
| WaitSyncFlag 5 参数 | ✓ | ✓ |
| WaitSyncFlag 6 参数（+offset） | ✗ | ✓ |
| SetOuterFlagDetour | ✗ | ✓ |
| WaitOuterFlagDetour | ✗ | ✓ |
| GetOuterFlagAddrDetour | ✗ | ✓ |
6.4 构建系统演进
| 特性 | master | cam_feature |
|---|---|---|
| 构建系统 | npu_op_* | npu_op_* |
| 算子选择 | select_ops.py | select_ops.py |
| 构建模式 | include-list | include-list |
| SHMEM 检测 | SHMEM_HOME_PATH | SHMEM_HOME_PATH |
| SHMEM 编译注入 | sed patch .make | sed patch .make |
| SHMEM 路径修复 | 有 bug（缺 /） | 已修复 |
| Catlass 架构宏 | 无 | CATLASS_ARCH=2201 |
| 全局头搜索路径 | 无 | -I${CMAKE_CURRENT_SOURCE_DIR} |
6.5 新增模块
| 模块 | master | cam_feature |
|---|---|---|
| HFA (HBM Fast Alloc) | ✗ | ✓ |
| FusedDeepMoeBuffer (pybind) | ✗ | ✓ |
| a2e / e2a 算子 | ✗ | ✓ |
| all_to_all_detour / reduce_scatter_detour | ✗ | ✓ |
| cam_get_comm / cam_free_comm 等通信 API | ✗ | ✓ |
6.6 架构演进总结
cam_feature 分支代表了 CAM 库从标准 HCCL 通信向SHMEM 共享内存加速的架构转型：
- ZB 系列算子： 7 个新的 ZB 算子（全部 requires_shmem=true）使用自定义共享内存窗口替代部分 HCCL 通信，降低延迟
- Detour 同步原语：
SyncCollectives新增 Detour 变体，支持在共享内存窗口任意偏移处工作 - HFA 内存管理： 全新的 HBM 快速分配器，提供 O(log n) 分配/释放 + 碎片整理能力
- pybind 扩展： 新增 6 个通信管理 API + 1 个 ZB Fused MOE Buffer 类
- 构建系统： 修复 SHMEM 路径 bug，新增 Catlass 架构宏，支持更灵活的算子选择
附录：关键文件路径索引
| 文件 | 路径 |
|---|---|
| tiling_args.h | src/cam/comm_operator/ascend_kernels/utils/op_host/tiling_args.h |
| mc2_tiling_utils.h | src/cam/comm_operator/ascend_kernels/utils/op_host/mc2_tiling_utils.h |
| dfx_base.h | src/cam/comm_operator/ascend_kernels/utils/op_host/dfx_base.h |
| ops_log.h | src/cam/comm_operator/ascend_kernels/utils/op_host/ops_log.h |
| error_log.h | src/cam/comm_operator/ascend_kernels/utils/op_host/error_log.h |
| ops_error.h | src/cam/comm_operator/ascend_kernels/utils/op_host/ops_error.h |
| comm_args.h | src/cam/comm_operator/ascend_kernels/utils/op_kernel/comm_args.h |
| sync_collectives.h | src/cam/comm_operator/ascend_kernels/utils/op_kernel/sync_collectives.h |
| data_copy.h | src/cam/comm_operator/ascend_kernels/utils/op_kernel/data_copy.h |
| moe_distribute_base.h | src/cam/comm_operator/ascend_kernels/utils/op_kernel/moe_distribute_base.h |
| operator_registry.json | src/cam/comm_operator/ascend_kernels/operator_registry.json |
| select_ops.py | build/cam/comm_operator/select_ops.py |
| compile_ascend_proj.sh | build/cam/comm_operator/compile_ascend_proj.sh |
| ascendc_impl_build.py | src/cam/comm_operator/ascend_kernels/cmake_files/cmake/util/ascendc_impl_build.py |
| flex_alloc.h | src/cam/utilities/hfa/src/flex_alloc.h |
| flex_alloc.cpp | src/cam/utilities/hfa/src/flex_alloc.cpp |
| hbm_fast_alloc.h | src/cam/utilities/hfa/include/hbm_fast_alloc.h |
| hbm_fast_alloc.cpp | src/cam/utilities/hfa/src/hbm_fast_alloc.cpp |
| acl_virtual_memory.h | src/cam/utilities/hfa/src/acl_virtual_memory.h |
| acl_virtual_memory.cpp | src/cam/utilities/hfa/src/acl_virtual_memory.cpp |
| hfa_log.h | src/cam/utilities/hfa/src/hfa_log.h |
| pybind.cpp | src/cam/comm_operator/pybind/pybind.cpp |
| functions.h | src/cam/comm_operator/pybind/functions.h |
| buffer.h | src/cam/comm_operator/pybind/buffer.h |
| ext_utils.h | src/cam/comm_operator/pybind/ext_utils.h |
分析日期: 2026-07-21 分支:
remotes/origin/cam_feature分析范围: pybind 层完整代码、examples、docker 部署、release notes、测试体系
目录
- Pybind 层完整分析
- PyTorch NPU Extension Bridge
- Examples 变化
- Docker 部署与 SHMEM 配置
- Release Notes 解读
- 测试体系变化
- 编译系统变化
- 综合洞察与架构研判
1. Pybind 层完整分析
1.1 文件清单与架构概览
cam_feature 分支的 pybind 目录实际包含以下文件（注意：用户指定的 moe_dispatch_prefill.cpp、moe_combine_prefill.cpp、get_dispatch_layout.cpp、NPUBridge.h/.cpp、NPUStorageImpl.h/.cpp 在该分支上不存在，这些是其他分支或旧版的内容）：
| 文件 | 功能 | 创建日期 |
|---|---|---|
pybind.cpp |
模块入口，注册所有 Python 绑定 | 2025-12-10 |
functions.h |
所有 autograd 函数声明 | 2025-12-10 |
torch_bind_exception.h |
自定义异常类 | 2026-02-28 |
fused_deep_moe.cpp |
FusedDeepMoe 算子绑定 | 2025-12-10 |
a2e.cpp |
A2E (Attention-to-Expert) 算子绑定 | 2026-02-06 |
e2a.cpp |
E2A (Expert-to-Attention) 算子绑定 | 2026-02-06 |
all2_all_detour.cpp |
All2All Detour 算子绑定 | 2026-01-22 |
reduce_scatter_detour.cpp |
ReduceScatter Detour 算子绑定 | 2026-01-22 |
buffer.h/.cpp |
SHMEM ZeroBuffer 管理 | 2026-05-30 |
ext_utils.h/.cpp |
CamComm 通信管理工具 | - |
pytorch_npu_helper.hpp |
NPU 命令执行辅助宏 | 2025-11-29 |
setup.py |
whl 包构建脚本 | 2025-12-11 |
bdist_wheel_build.py |
wheel 打包辅助 | 2025-11-29 |
1.2 pybind.cpp — 模块注册核心
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 注册了以下 Python 可调用函数：
m.def("fused_deep_moe", &FusedDeepMoeImplAutograd, "fused_deep_moe");
m.def("all2_all_detour", &all2_all_detour_impl_autograd, "all2_all_detour");
m.def("reduce_scatter_detour", &reduce_scatter_detour_impl_autograd, "reduce_scatter_detour");
m.def("cam_get_comm", &cam_get_comm, "cam_get_comm");
m.def("cam_free_comm", &cam_free_comm, "cam_free_comm");
m.def("cam_get_rank_size", &cam_get_rank_size, "cam_get_rank_size");
m.def("cam_get_and_increase_magic", &cam_get_and_increase_magic, "cam_get_and_increase_magic");
m.def("cam_get_magic", &cam_get_magic, "cam_get_magic");
m.def("a2e", &A2eImplAutograd, "a2e");
m.def("e2a", &E2aImplAutograd, "e2a");
同时注册了 FusedDeepMoeBuffer 类（SHMEM ZeroBuffer 的 Python 接口），暴露 init、zb_fused_deep_moe、is_initialized、get_ext_info、get_shmem_workspace 方法。
1.3 TORCH_LIBRARY 注册 — 算子 Schema 定义
TORCH_LIBRARY(umdk_cam_op_lib, m) 定义了 6 个算子的完整 schema：
(1) fused_deep_moe — 最复杂的算子
m.def("fused_deep_moe(Tensor x, Tensor expert_ids, Tensor[] gmm1_weight, Tensor[] gmm1_weight_scale, \
Tensor[] gmm2_weight, Tensor[] gmm2_weight_scale, Tensor expert_scales, \
Tensor? share_gmm1_weight, Tensor? share_gmm1_weight_scale, \
Tensor? share_gmm2_weight, Tensor? share_gmm2_weight_scale, \
Tensor? expert_smooth_scales, Tensor? share_smooth_scales, Tensor? x_active_mask, \
Tensor[] gmm1_bias, Tensor[] gmm2_bias, \
Tensor? share_gmm1_bias, Tensor? share_gmm2_bias, \
str group_ep, int ep_rank_size, int ep_rank_id, int moe_expert_num, \
int quant_mode, int global_bs) -> Tensor[]");
关键设计点：
- 返回 Tensor[]（3 个输出：output, shareOutput, expertTokenNums）
- 支持 Shared Expert（share_gmm1/2_weight 可选）
- 支持 Smooth Quantization（expert_smooth_scales）
- 支持 Active Mask（x_active_mask，用于 MC2 动态 batch）
- 支持 Bias（gmm1/2_bias）
- quant_mode 参数控制量化模式（0=W8A8, 1=其他）
(2) all2_all_detour — All-to-All 通信 Detour
m.def("all2_all_detour(Tensor sendData, Tensor commRankIds, Tensor commArgs, int commId) -> Tensor");
(3) reduce_scatter_detour — ReduceScatter 通信 Detour
m.def("reduce_scatter_detour(Tensor sendData, Tensor commRankIds, Tensor commArgs, int commId, int op) -> Tensor");
(4) cam_get_comm / cam_free_comm — 通信域管理
m.def("cam_get_comm(int comm_id, int rank, int group_size, str server_ip_port) -> Tensor");
m.def("cam_free_comm(int comm_id) -> Tensor");
(5) a2e — Attention 到 Expert 的数据路由
m.def("a2e(Tensor x, Tensor? expert_ids, Tensor? scales, int batch_size, int hidden_size, \
int topk, int expert_rank_size, int atten_rank_size, int rank, str group_ep, \
int aiv_num, int compute_gate) -> Tensor[]");
(6) e2a — Expert 到 Attention 的数据回传
m.def("e2a(Tensor expand_x, Tensor atten_batch_size, int batch_size, int hidden_size, \
int topk, int expert_rank_size, int attention_rank_size, int rank, str group_ep, \
int aiv_num) -> Tensor");
1.4 统一的 Autograd 注册模式
所有算子遵循完全相同的四层注册架构（以 fused_deep_moe 为例）：
FusedDeepMoeImplAutograd (入口，调用 apply)
└─ FusedDeepMoe::apply (torch::autograd::Function)
└─ FusedDeepMoe::forward → FusedDeepMoeImpl
└─ Dispatcher::findSchemaOrThrow → 分发到具体实现
三层 TORCH_LIBRARY_IMPL 注册：
| 后端 | 注册函数 | 用途 |
|---|---|---|
PrivateUse1 |
FusedDeepMoeImplNpu |
NPU 实际执行 |
AutogradPrivateUse1 |
FusedDeepMoeImplAutograd |
Autograd 支持 |
Meta |
FusedDeepMoeImplMeta |
元设备（shape 推断） |
关键实现细节：
- FusedDeepMoeImplNpu 中调用 EXEC_NPU_CMD(aclnnFusedDeepMoe, ...) 执行实际 NPU 命令
- backward 返回全空 Tensor（所有算子反向传播均未实现，仅支持推理）
- Meta 实现仅做 shape 推断，不分配实际内存
1.5 ext_utils — CamComm 通信管理
constexpr int MAX_COMMS = 64;
extern CamCommPtr g_comms[MAX_COMMS];
- 全局通信域数组，最多支持 64 个并发通信域
cam_get_comm通过CamCreateComm创建通信域，返回commArgsTensorcam_get_and_increase_magic用于获取并递增 magic 号（防 ABA 问题）cam_get_rank_size查询通信域 rank 数量
1.6 torch_bind_exception.h — 异常处理
class TorchBindException : public std::exception {
// 格式: "Failed: <name> error <file>:<line> error message or error code is '<error>'"
};
#define TORCH_BIND_ASSERT(cond) // 不满足则抛出 TorchBindException
这是一个轻量级断言宏，用于在 C++ 层捕获错误并转换为 Python 异常。
2. PyTorch NPU Extension Bridge
2.1 pytorch_npu_helper.hpp — 核心执行引擎
这是整个 pybind 层最底层的基础设施，提供了 ATen → ACL 类型转换 和 NPU 命令执行 两大功能。
类型转换体系
at::Tensor → aclTensor* (ConvertType)
at::Scalar → aclScalar* (ConvertType)
at::IntArrayRef → aclIntArray* (ConvertType)
at::TensorList → aclTensorList* (ConvertType)
c10::optional<T> → T* or nullptr (ConvertType)
ScalarType → ACL DataType 映射表：
constexpr aclDataType kATenScalarTypeToAclDataTypeTable[] = {
at::ScalarType::Byte → ACL_UINT8,
at::ScalarType::Char → ACL_INT8,
at::ScalarType::Short → ACL_INT16,
at::ScalarType::Int → ACL_INT32,
at::ScalarType::Long → ACL_INT64,
at::ScalarType::Half → ACL_FLOAT16,
at::ScalarType::Float → ACL_FLOAT,
at::ScalarType::Double → ACL_DOUBLE,
at::ScalarType::BFloat16 → ACL_BF16,
// ... Complex 类型
};
Format 推断逻辑：
- dim=3 → ACL_FORMAT_NCL
- dim=4 → ACL_FORMAT_NCHW
- dim=5 → ACL_FORMAT_NC1HWC0
- INT8 + (dim=2 or 3) → ACL_FORMAT_FRACTAL_NZ
- 其他 → ACL_FORMAT_ND
EXEC_NPU_CMD 宏 — 核心执行宏
#define EXEC_NPU_CMD(aclnn_api, ...)
这是整个 CAM 算子执行的统一入口，执行流程：
- 动态加载 OpAPI 函数：通过
dlsym从libopapi.so/libcust_opapi.so加载 - 初始化大页内存：
InitHugeMemThreadLocal - 类型转换：
ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr) - 获取 Workspace 大小：调用
GetWorkspaceSize - 分配 Workspace：通过
at::empty分配 NPU Tensor - 执行算子：lambda 捕获所有参数，通过
OpCommand提交到 NPU 流 - 释放资源：
ReleaseConvertTypes+ReleaseHugeMem+UnInitHugeMemThreadLocal
关键设计点：
- 使用 static 局部变量缓存函数地址（首次调用时解析，后续复用）
- 线程本地存储 g_hashBuf[8192] 用于算子 hash 计算
- 支持 InitHugeMemThreadLocal / UnInitHugeMemThreadLocal 可选大页内存管理
- 通过 at_npu::native::OpCommand 提交命令，支持自定义 handler
2.2 关于 NPUBridge/NPUStorageImpl
注意： NPUBridge.h/.cpp 和 NPUStorageImpl.h/.cpp 在 cam_feature 分支上不存在。这些文件可能是：
- 其他分支（如 main 或 master）的内容
- 旧版 CAM 的 PyTorch 存储扩展实现
- 已被 pytorch_npu_helper.hpp 取代
当前 cam_feature 分支使用 pytorch_npu_helper.hpp 作为唯一的 NPU 交互层，不再需要单独的 Bridge/StorageImpl 层。
3. Examples 变化
3.1 文件清单
| 文件 | 功能 | 创建日期 |
|---|---|---|
fused_deep_moe_sample.py |
FusedDeepMoe 基础示例 | 2025-12-11 |
fused_deep_moe_sample_w4A8.py |
W4A8 量化变体示例 | 2026-06-29 |
fused_deep_moe_shmem_sample.py |
SHMEM ZeroBuffer 示例 | 2026-04-22 |
a2e_e2a_sample.py |
A2E/E2A 路由示例 | 2026-03-02 |
3.2 fused_deep_moe_sample.py — 基础精度验证
核心设计： 对比 FusionOp（融合算子）与 SmallOps（小算子组合）的输出精度。
# 调用方式
output, expert_token_nums = torch.ops.umdk_cam_op_lib.fused_deep_moe(
x=x,
expertIds=expert_ids,
gmm1PermutedWeight=self.gmm1_weight,
gmm1PermutedWeightScale=self.gmm1_weight_scale_fp32,
gmm2Weight=self.gmm2_weight,
gmm2WeightScale=self.gmm2_weight_scale_fp32,
expertSmoothScalesOptional=smooth_scales,
expertScalesOptional=expert_scales,
groupEp=self.ep_hcomm_info,
epRankSize=self.ep_world_size,
epRankId=self.global_rank_id,
moeExpertNum=self.moe_expert_num,
sharedExpertNum=1,
sharedExpertRankNum=self.shared_expert_rank_num,
quantMode=0,
globalBs=self.batch_size * self.ep_world_size
)
权重预处理（关键）：
# GMM1 权重需要特殊 permute
gmm1_weight = gmm1_weight.transpose(1,2).contiguous()\
.view(local_expert_num, 2, intermediate // 64, 64, hidden)\
.transpose(1,2).contiguous()\
.view(local_expert_num, intermediate * 2, hidden)\
.transpose(1,2).contiguous()
# 转为 FRACTAL_NZ 格式
gmm1_weight = torch_npu.npu_format_cast(gmm1_weight, Format.FRACTAL_NZ)
测试配置：
- batch_size=64, hidden_size=7168, intermediate_size=2048
- ep_world_size=16, moe_expert_num=64, top_k=8
- 精度验证：atol=2.0, rtol=0.02
3.3 fused_deep_moe_shmem_sample.py — SHMEM ZeroBuffer
核心创新： 使用 FusedDeepMoeBuffer 初始化 SHMEM 内存，实现零拷贝通信。
# SHMEM Buffer 初始化
self.buffer = umdk_cam_op_lib.FusedDeepMoeBuffer()
self.buffer.init(
rank=self.global_rank_id,
num_ranks=self.ep_world_size,
memsize=shmem_memsize, # 默认 3GB
ip_port="tcp://172.16.0.146:8766"
)
# ZeroBuffer 调用
output, share_output, expert_token_nums = self.buffer.zb_fused_deep_moe(
x, expert_ids,
self.gmm1_weight, self.gmm1_weight_scale,
self.gmm2_weight, self.gmm2_weight_scale,
expert_scales,
self.share_gmm1_weight, self.share_gmm1_weight_scale,
self.share_gmm2_weight, self.share_gmm2_weight_scale,
self.smooth_scales, self.share_smooth_scales_fp32,
x_active_mask,
"", group_ep, ep_rank_size, rank_id, expert_num, quant_mode, global_bs
)
测试场景覆盖：
- test_fused_deep_moe_shmem_base — 基础功能
- test_fused_deep_moe_shmem_with_mc2_mask — MC2 Mask 动态 batch
- test_fused_deep_moe_shmem_eplb — 动态 EPLB（Expert Load Balancing）
- test_fused_deep_moe_shmem_with_share — Shared Expert
- test_fused_deep_moe_shmem_with_smooth — Smooth Quantization
- test_graph — torchair 图模式编译
3.4 fused_deep_moe_sample_w4A8.py — W4A8 量化
核心特性：
- 支持 INT4 权重 × INT8 激活的 GEMM
- 包含 TorchSimulator 类用于 CPU 端模拟 NPU 量化行为
- 三段流水线模拟：Pre-process (AIV) → Mid-process (AIC) → Post-process
- 验证融合算子与 CPU 小算子组合的精度一致性
3.5 a2e_e2a_sample.py — MoE 与 Attention 分离架构
核心概念： 16 rank 分为 MOE 侧（rank 0-7）和 Attention 侧（rank 8-15）。
# Attention 侧：发送数据到 MOE
a2e_output = torch.ops.umdk_cam_op_lib.a2e(
x=x, expert_ids=expert_ids, scales=scales,
batch_size=16, hidden_size=512, topk=2,
expert_rank_size=8, atten_rank_size=8,
rank=rank, group_ep=ep_hcomm_info, aiv_num=4, compute_gate=1
)
# MOE 侧：接收数据并处理
e2a_output = torch.ops.umdk_cam_op_lib.e2a(
expand_x=expand_x, atten_batch_size=atten_batch_size,
batch_size=16, hidden_size=512, topk=2,
expert_rank_size=8, attention_rank_size=8,
rank=rank, group_ep=ep_hcomm_info, aiv_num=4
)
4. Docker 部署与 SHMEM 配置
4.1 Docker 文件缺失
注意： docker/cam/Dockerfile、docker/cam/docker-compose.yml、docker/cam/.env 在 cam_feature 分支上均不存在。docker 目录为空。
4.2 SHMEM 配置机制（从编译系统推断）
虽然 Docker 文件不存在，但从 setup.py 和编译脚本可以完整推断 SHMEM 部署要求：
环境变量依赖
| 变量 | 必需 | 说明 |
|---|---|---|
SHMEM_HOME_PATH |
否（但 SHMEM 算子需要） | Ascend-SHMEM 安装路径 |
CAM_COMM_PATH |
是 | cam_comm 库构建路径 |
ASCEND_HOME_PATH |
是 | CANN 工具包路径 |
SHMEM 编译时配置
# setup.py
shmem_home_path = os.environ.get("SHMEM_HOME_PATH")
if shmem_home_path:
include_dirs.append(os.path.join(shmem_home_path, 'shmem', 'include'))
library_dirs.append(os.path.join(shmem_home_path, 'shmem', 'lib'))
libraries.append("shmem")
compile_args.append("-DSHMEM_ENABLED")
SHMEM 运行时配置
从 buffer.cpp 可知：
- 使用 aclshmemx_init_attr 初始化 SHMEM 通信域
- 通过 ip_port（如 tcp://172.16.0.146:8766）建立跨节点连接
- aclshmem_malloc 分配共享内存（默认 3GB）
- 分为 ext_info（2MB）和 shmem_workspace（剩余空间）两部分
// buffer.cpp
status = aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, &attributes);
ext_info_ptr_ = aclshmem_malloc(EXT_INFO_SIZE); // 2MB
shmem_workspace_ptr_ = aclshmem_malloc(shmem_workspace_size_); // 剩余
SHMEM 算子自动剔除机制
# select_ops.py
def requires_shmem(registry, op):
return bool(get_meta(registry, op).get("requires_shmem", False))
当 SHMEM_HOME_PATH 未设置时，所有 requires_shmem=true 的算子自动跳过。
4.3 部署架构推断
┌─────────────────────────────────────────────┐
│ Ascend A2/A3 SuperPod │
│ CANN ≥ 8.3.RC1 │
│ Torch 2.8.0 + Torch-NPU 2.8.0-7.2.0 │
│ Ascend-SHMEM 1.0.0 (可选) │
│ │
│ Step 1: 安装 .run 包到 opp 目录 │
│ Step 2: source set_env.bash │
│ Step 3: pip install umdk_cam_op_lib.whl │
└─────────────────────────────────────────────┘
5. Release Notes 解读
5.1 版本信息
- 版本号： v25.12.0
- 发布日期： 2025/12/30
5.2 核心更新内容
Release Notes 主要描述 UMDK 整体（URMA/URPC/DLock/UMS），CAM 部分未单独列出。关键更新：
URMA — 聚合设备特性
- Jetty 多路径模式：CTP 通道，RM/RC 双模式，最大 4KB 消息
- Jetty 单路径模式：TP 通道，仅 RC 模式，最大 64KB 消息，支持本地回环
- JFS/JFR 多路径模式：CTP 通道，仅 RM 模式，最大 4KB 消息
URPC — 统一远程过程调用
- 超时机制、连接保活、Channel 多对一
- UMQ 消息队列（enqueue/dequeue + Post/Poll 两种模型）
- UMQ Token 安全校验
DLock — 分布式锁
- 原子锁/公平锁/读写锁的基本和批量操作
- 分布式对象管理（FAA/CAS 原子操作）
- 故障恢复（Primary Server 故障后重建）
- 安全传输加密
UMS — Socket 兼容
- 标准 Socket API 兼容
ums_run透明替换工具
5.3 CAM 相关更新（从代码推断）
虽然 Release Notes 未直接列出 CAM 更新，但从代码可推断 cam_feature 分支的关键更新：
| 特性 | 状态 | 说明 |
|---|---|---|
| FusedDeepMoe (W8A8) | ✅ 已完成 | A3 支持，低延迟模式 |
| FusedDeepMoe (W4A8) | ✅ 已完成 | 量化变体 |
| FusedDeepMoE_fwk | ✅ 已完成 | 独立框架算子 |
| A2E/E2A | ✅ 已完成 | MoE-Attention 分离架构 |
| All2All Detour | ✅ 已完成 | 自定义 All-to-All |
| ReduceScatter Detour | ✅ 已完成 | 自定义 ReduceScatter |
| SHMEM ZeroBuffer | ✅ 已完成 | 零拷贝通信 |
| MC2 Mask | ✅ 已完成 | 动态 batch 支持 |
| EPLB | ✅ 已完成 | Expert Load Balancing |
| Shared Expert | ✅ 已完成 | 共享专家支持 |
| torchair 图模式 | ✅ 已完成 | 编译优化 |
6. 测试体系变化
6.1 测试文件清单
| 文件 | 功能 | 创建日期 |
|---|---|---|
test/cam/ut/a2e_e2a/test_a2e_e2a.py |
A2E/E2A 单元测试 | 2026-03-02 |
6.2 测试覆盖分析
当前仅有 1 个测试文件，覆盖范围极其有限：
@MPTest # 多进程测试
@A3Test # A3 芯片测试
@SKIP_ENV_RANKSIZE_UNEQUAL(16) # 需要 16 rank
@pytest.mark.parametrize("mode", [('Eager')]) # 仅 Eager 模式
def test_base_test(mode):
# 测试配置
CASE_16RANK = {
"batchSize": 16,
"hiddenSize": 512,
"topk": 2,
"expertRankSize": 8,
"attentionRankSize": 8,
"aivNum": 4
}
验证内容：
- Attention 侧：E2A 输出 shape 与输入一致，数值 atol=1e-3
- MOE 侧：expand_x.shape[1] == hidden_size
6.3 测试缺口
| 缺口 | 严重程度 | 说明 |
|---|---|---|
| 无 FusedDeepMoe UT | 🔴 高 | 最复杂的算子没有单元测试 |
| 无 SHMEM 测试 | 🔴 高 | ZeroBuffer 功能无自动化测试 |
| 无 W4A8 测试 | 🟡 中 | 量化变体无独立测试 |
| 无 All2All/ReduceScatter UT | 🟡 中 | Detour 算子无测试 |
| 无图模式测试 | 🟡 中 | torchair 编译路径未覆盖 |
| 无精度回归测试 | 🟡 中 | 无数值精度自动化验证 |
| 无多配置参数化 | 🟡 中 | 仅单一配置 |
| 无 GE 模式测试 | 🟢 低 | 仅 Eager，但注释说明 GE 无法获取覆盖率 |
7. 编译系统变化
7.1 operator_registry.json — 算子注册表
{
"soc_versions": {
"ascend910_93": [
"a2e", "all_to_all_detour", "e2a",
"fused_deep_moe", "fused_deep_moe_fwk", "fused_deep_moe_w4a8",
"moe_combine_lowlatency_zb", "moe_combine_normal_zb",
"moe_dispatch_layout_zb", "moe_dispatch_lowlatency_zb",
"moe_dispatch_normal_zb", "moe_notify_dispatch_zb",
"reduce_scatter_detour"
]
}
}
13 个算子，其中 6 个需要 SHMEM。
7.2 构建选项
| 选项 | 功能 | cam_feature 特有 |
|---|---|---|
-c <soc> |
目标 SOC 代际 | |
-a <ops> |
选择算子列表 | |
-q |
W4A8 量化变体 | |
-d |
Debug 构建 | |
-t |
仅构建 UT | |
-p |
仅构建 whl | |
-r |
仅构建 run 包 | |
-m |
仅构建 cam_comm 库 | ✅ 仅 cam_feature |
7.3 关键编译规则
fused_deep_moe与fused_deep_moe_w4a8互斥（共享源文件名）fused_deep_moe_fwk独立，可与任意变体共存- SHMEM 算子在
SHMEM_HOME_PATH未设置时自动跳过 -m选项仅cam_feature有（master 无 cam_comm）
8. 综合洞察与架构研判
8.1 架构演进方向
从通信算子到融合算子的跃迁：
Phase 1 (master): dispatch + combine 分离通信算子
Phase 2 (cam_feature):
→ FusedDeepMoe: dispatch + GMM1 + SwiGLU + GMM2 + combine 五合一
→ A2E/E2A: MoE-Attention 分离通信
→ SHMEM ZeroBuffer: 零拷贝通信
→ All2All/ReduceScatter Detour: 自定义通信原语
8.2 关键设计决策
-
全推理导向：所有算子
backward返回空 Tensor，表明 CAM 定位为推理加速库，训练依赖 PyTorch 原生 autograd 或外部框架。 -
三层后端注册：
PrivateUse1（执行）+AutogradPrivateUse1（梯度）+Meta（shape 推断）是 PyTorch 自定义算子的最佳实践。 -
SHMEM 可选编译：通过
SHMEM_HOME_PATH环境变量控制，不依赖 SHMEM 时自动降级，保证兼容性。 -
Dispatcher 分发模式：所有算子通过
torch::Dispatcher分发，支持 lazy execution、symbolic tracing 和 torch.compile。 -
EXEC_NPU_CMD 统一宏：封装了 OpAPI 加载、workspace 管理、内存释放的完整生命周期，大幅降低算子开发复杂度。
8.3 与竞品对比
| 特性 | CAM (Huawei) | NCCL (NVIDIA) | MS-AMP (Microsoft) |
|---|---|---|---|
| 融合 MoE 算子 | ✅ FusedDeepMoe | ❌ 分离 | ❌ 分离 |
| 自定义 All-to-All | ✅ Detour | ❌ 标准 | ❌ |
| 零拷贝通信 | ✅ SHMEM ZB | ✅ GPUDirect | ❌ |
| MoE-Attn 分离 | ✅ A2E/E2A | ❌ | ❌ |
| 量化支持 | ✅ W8A8/W4A8 | ❌ | ✅ FP8 |
| 反向传播 | ❌ 仅推理 | ✅ | ✅ |
8.4 风险与建议
-
测试覆盖严重不足：13 个算子仅 1 个 UT，建议优先补充 FusedDeepMoe 和 SHMEM 测试。
-
反向传播缺失：如果未来需要训练支持，需实现完整的 backward 函数。
-
SHMEM 网络配置：
ip_port硬编码在示例中，生产环境需要动态服务发现。 -
W4A8 量化：示例文件较大（33KB），包含 CPU 模拟器，说明量化精度验证是难点。
-
Docker 部署缺失：无 Dockerfile/docker-compose，建议补充容器化部署方案。
附录：关键文件路径索引
src/cam/comm_operator/pybind/
├── pybind.cpp # 模块入口 + TORCH_LIBRARY 注册
├── functions.h # 所有 autograd 函数声明
├── torch_bind_exception.h # 自定义异常
├── fused_deep_moe.cpp # FusedDeepMoe 实现
├── a2e.cpp # A2E 实现
├── e2a.cpp # E2A 实现
├── all2_all_detour.cpp # All2All Detour 实现
├── reduce_scatter_detour.cpp # ReduceScatter Detour 实现
├── buffer.h / buffer.cpp # SHMEM ZeroBuffer
├── ext_utils.h / ext_utils.cpp # CamComm 通信管理
├── pytorch_npu_helper.hpp # NPU 执行辅助宏
├── setup.py # whl 构建
└── pytorch_extension/
├── bdist_wheel_build.py # wheel 打包
└── pytorch_npu_helper.hpp # (同上)
src/cam/examples/
├── fused_deep_moe_sample.py # 基础精度验证
├── fused_deep_moe_sample_w4A8.py # W4A8 量化
├── fused_deep_moe_shmem_sample.py # SHMEM ZeroBuffer
└── a2e_e2a_sample.py # MoE-Attn 分离
test/cam/ut/a2e_e2a/
└── test_a2e_e2a.py # A2E/E2A 单元测试
build/cam/comm_operator/
├── build.sh # 主编排
├── build_pybind.sh # pybind 构建
├── compile_ascend_proj.sh # 算子编译
└── select_ops.py # 算子选择
src/cam/comm_operator/ascend_kernels/
└── operator_registry.json # 算子注册表
RELEASE-NOTES.md # 英文版本
RELEASE-NOTES.ch.md # 中文版本
README.md # 项目 README
src/cam/README.md # CAM 模块 README
分析结论：
cam_feature分支代表了 CAM 从"通信加速库"向"融合算子库"的关键转型。FusedDeepMoe 五合一算子、SHMEM ZeroBuffer、A2E/E2A 分离架构是三大核心创新。但测试覆盖和反向传播是当前最大短板。
代码仓库: UMDK/CAM — Communication Acceleration for Matrix on Ascend NPU 分析分支:
master(SHMEM/NDA 框架) +cam_feature(ZB 算子消费端) 分析日期: 2026-07-21 分析范围: 从 host QP 建立到 device RDMA 操作的全链路基础设施
0. 关键发现：SHMEM 框架与 ZB 算子的分支分布
重要: SHMEM/NDA 框架代码位于 master 分支，而 ZB (Zero-Buffer) 算子位于 cam_feature 分支。cam_feature 分支本身不包含 NDA 框架源码，它通过外部依赖方式引用 shmem 包（从 gitcode.com/cann/shmem 克隆，commit 6a5f4b09）。
master 分支:
src/cam/framework/nda/
├── include/device/gm2gm/engine/shmem_device_rdma.h (公共设备 API)
└── src/device/gm2gm/engine/
├── shmemi_device_rdma.h (内部核心数据结构)
├── shmem_device_rdma.hpp (设备端实现/调度层)
└── rdma_backends/
├── rdma_device_backend_base.h (后端抽象基类)
├── rdma_device_backend_base.hpp (后端原子操作实现)
└── rdma_device_backend_xscale.hpp (XSCALE 硬件后端)
└── src/host/transport/device_rdma/
├── device_rdma_transport_manager_v2.h (Host 传输管理器)
└── device_rdma_transport_manager_v2.cpp (QP/MR 生命周期)
└── src/host/utils/under_api/
├── dl_hccp_def.h (HCCP 底层数据结构)
└── dl_hcomm_def.h (HCOMM 通道实体定义)
cam_feature 分支:
src/cam/comm_operator/ascend_kernels/
├── moe_dispatch_normal_zb/ (使用 aclshmem_* API)
├── moe_combine_normal_zb/
├── moe_dispatch_lowlatency_zb/
├── moe_combine_lowlatency_zb/
├── moe_dispatch_layout_zb/
├── moe_notify_dispatch_zb/
└── utils/op_kernel/
├── zero_buffer_api.h (→ #include "shmem.h")
├── zero_buffer_sync_flag.h (跨卡同步标志)
├── sync_collectives_zb.h (集合同步原语)
└── comm_args_zb.h / data_copy_zb.h
1. SHMEM 基础设施完整架构
1.1 全链路数据流
┌─────────────────────────────────────────────────────────────────────────────┐
│ Host 侧 (CPU) │
│ │
│ RdmaTransportManagerV2 │
│ ├── OpenDevice() → 获取 device/IP/port, 创建 Endpoint │
│ ├── RegisterMemoryRegion() → HcommMemReg 注册 HBM 为 MR │
│ ├── Prepare() → 解析对端 NIC 地址, 构建 rankInfo │
│ ├── Connect() → HcommChannelCreate 创建 QP 连接 │
│ │ ├── PrepareChannelDescs() → 填充 channel 描述符 (TC/SL/端口) │
│ │ └── CreateChannelsAndFillInfo() → 创建通道 + 填充 device 侧信息 │
│ └── FillRdmaInfo() → 读取 SQ/CQ/MR 信息, 拷贝到 device HBM │
│ │
│ 输出: aclshmemi_rdma_info 结构体 (驻留 device HBM) │
└────────────────────────────────┬────────────────────────────────────────────┘
│ AclrtMemcpy(Host→Device)
▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Device 侧 (NPU AIV Core) │
│ │
│ aclshmemi_qp_info_fetch() → 从 HBM 获取 rdma_info 基址 │
│ │
│ 用户调用: aclshmemx_roce_put_nbi / aclshmemx_roce_get_nbi / atomic_* │
│ │ │
│ ▼ │
│ aclshmem_device_rdma.hpp (调度层) │
│ │ aclshmemi_roce_write<T, ACLSHMEMI_K_RDMA_BACKEND>() │
│ │ aclshmemi_roce_read<T, ACLSHMEMI_K_RDMA_BACKEND>() │
│ │ aclshmemi_roce_amo_add / aclshmemi_roce_amo_cas │
│ ▼ │
│ rdma_device_backend_base.hpp (原子操作基类) │
│ │ aclshmemi_roce_atomic_fetch_and_add<T, IS_MASKED, B>() │
│ │ aclshmemi_roce_atomic_compare_and_swap<T, IS_MASKED, B>() │
│ ▼ │
│ aclshmemi_backend_traits<B>::atomic_op_traits<T, IS_MASKED> │
│ │ (编译时多态分发) │
│ ▼ │
│ rdma_device_backend_xscale.hpp (XSCALE 硬件后端) │
│ ├── fill_wqe() → 构造 WQE (控制段 + 数据段 + 原子段) │
│ ├── ring_sq_doorbell()→ 更新 SQ head + 写 doorbell 寄存器 │
│ ├── poll_cq() → 轮询 CQE owner bit 直到完成 │
│ └── ring_cq_doorbell()→ 更新 CQ tail + 写 CQ doorbell │
└─────────────────────────────────────────────────────────────────────────────┘
1.2 编译时后端选择
// shmem_device_rdma.hpp
#if defined(ACLSHMEMI_RDMA_K_BACKEND_XSCALE)
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::XSCALE)
#else
#define ACLSHMEMI_K_RDMA_BACKEND (aclshmemi_rdma_backend_t::IN_DIE)
#endif
所有设备端 API 通过 ACLSHMEMI_K_RDMA_BACKEND 宏在编译时绑定到具体后端，零运行时开销。当前仅 XSCALE 后端有完整实现。
2. 核心数据结构
2.1 aclshmemi_rdma_info — 设备端 RDMA 全局信息
// shmemi_device_rdma.h
struct aclshmemi_rdma_info {
uint32_t qp_num; // 每个连接的 QP 数量 (multi-QP 支持)
uint64_t sq_ptr; // Send Queue 上下文数组 [PE_NUM][qp_num]
uint64_t rq_ptr; // Receive Queue 上下文数组
uint64_t scq_ptr; // Send Completion Queue 上下文数组
uint64_t rcq_ptr; // Receive Completion Queue 上下文数组
uint64_t mem_ptr; // Memory Region 信息数组 [MAX_PE_NUM]
};
该结构体驻留于 device HBM，由 host 端 FillRdmaInfo() 通过 AclrtMemcpy 拷贝上去。设备端通过 aclshmemi_qp_info_fetch() 获取指针。
2.2 SQ/CQ 上下文
struct aclshmemi_rdma_sq_ctx {
uint32_t wqn; // Work Queue Number (硬件 QP 编号)
uint64_t buf_addr; // SQ ring buffer 基址
uint32_t wqe_size; // 每个 WQE 大小 (XSCALE: 128B)
uint32_t depth; // ring buffer depth
uint64_t head_addr; // Producer Index 地址 (host/device 共享)
uint64_t tail_addr; // Consumer Index 地址 (host/device 共享)
aclshmemi_rdma_db_mode_t db_mode; // HW_DB / SW_DB
uint64_t db_addr; // doorbell 寄存器地址
uint32_t sl; // Service Level
uint64_t amo_addr; // 原子操作结果回写地址
uint32_t amo_lkey; // 原子操作内存 lkey
};
struct aclshmemi_rdma_cq_ctx {
uint32_t cqn; // Completion Queue Number
uint64_t buf_addr; // CQ ring buffer 基址
uint32_t cqe_size; // 每个 CQE 大小 (XSCALE: 64B, 含 32B padding)
uint32_t depth; // ring buffer depth
uint64_t head_addr; // Producer Index 地址
uint64_t tail_addr; // Consumer Index 地址
aclshmemi_rdma_db_mode_t db_mode;
uint64_t db_addr; // CQ doorbell 寄存器地址
};
2.3 send_wr — 工作请求描述符
struct aclshmemi_rdma_send_wr {
uint32_t send_flags; // 发送标志位掩码
uint32_t imm_data; // immediate data (RDMA_WRITE_WITH_IMM)
__gm__ uint8_t *remote_addr; // 远端 HBM 地址
uint32_t rkey; // 远端 memory region key
__gm__ uint8_t *local_addr; // 本地 HBM 地址
uint64_t message_len; // 消息长度 (字节)
uint32_t lkey; // 本地 memory region key
uint32_t num_sge; // SGE 数量 (multi-sge)
aclshmemi_rdma_sge *sg_list; // scatter-gather 列表
aclshmemi_rdma_atomic_params atomic; // 原子操作参数
};
2.4 atomic_params — 16 种变体联合
struct aclshmemi_rdma_atomic_params {
union {
struct { uint64_t compare_add; uint64_t swap; }; // 基础 CAS/FA
struct { // 通用 masked
uint64_t swap_add_data; uint64_t compare_data;
uint64_t swap_add_mask; uint64_t compare_mask;
} masked_common;
struct { uint64_t swap_data; uint64_t compare_data; // masked CAS64
uint64_t swap_mask; uint64_t compare_mask; } masked_cas64;
struct { uint64_t swap_data; uint64_t compare_data; // masked CAS32
uint32_t swap_mask; uint32_t compare_mask; } masked_cas32;
struct { uint64_t add_data; uint64_t field_boundary; } // masked FA64
masked_fa64;
struct { uint32_t add_data; uint32_t field_boundary; // masked FA32
uint64_t _reserved; } masked_fa32;
};
};
设计洞察: 使用 union 覆盖所有原子操作的参数变体，通过 masked_common 统一接口传入 swap_add_data/compare_data/swap_add_mask/compare_mask，后端根据 IS_MASKED 和 sizeof(T) 选择正确的硬件 WQE 格式。
2.5 Host 端关键结构
// dl_hccp_def.h
struct AiQpRMAWQ {
uint32_t wqn; uint64_t bufAddr; uint32_t wqeSize; uint32_t depth;
uint64_t headAddr; uint64_t tailAddr; DBMode dbMode; uint64_t dbAddr;
uint32_t sl; uint64_t atomicAddr; uint32_t atomicLkey;
};
struct AiQpRMACQ {
uint32_t cqn; uint64_t bufAddr; uint32_t cqeSize; uint32_t depth;
uint64_t headAddr; uint64_t tailAddr; DBMode dbMode; uint64_t dbAddr;
};
struct AiQpRMAQueueInfo {
uint32_t count;
AiQpRMAWQ* sq; AiQpRMAWQ* rq;
AiQpRMACQ* scq; AiQpRMACQ* rcq;
RdmaMemRegionInfo* mr;
};
3. API 调用链
3.1 RDMA Write 完整路径
用户代码: aclshmemx_roce_put_nbi(dst, src, buf, elem_size, pe)
│
▼
aclshmemi_roce_write<T, B>(dst, src, pe, 0, len, ub64, ub32, sync_id)
│ 构建 wr = {remote_addr=dst, local_addr=src, message_len=len}
▼
aclshmemi_roce_post_send<B, OP_RDMA_WRITE>(wr, pe, qp_idx, ...)
│
▼ (XSCALE 特化)
aclshmemi_xscale_post_send_read_write<OP_RDMA_WRITE>(wr, pe, qp_idx, ...)
│
├── 1. 读取 SQ head/tail (dcci_cachelines 从 HBM)
├── 2. 检查 SQ 满: (cur_head + 16) % depth == cur_tail → poll_cq
├── 3. 获取 rkey/lkey (从 rdma_info->mem_ptr)
├── 4. 计算 WQE 地址: sq_base_addr + wqe_size * (cur_head % depth)
├── 5. fill_wqe → 构造 128B WQE
│ ├── fill_wqe_ctrl_seg: msg_opcode=RDMA_WRITE(1), ds_data_num=2
│ └── fill_wqe_data_seg: [rkey|remote_addr] + [lkey|local_addr]
├── 6. dcci_cachelines(wqe_addr, 128) // flush WQE 到 HBM
├── 7. cur_head++
└── 8. ring_sq_doorbell:
├── 写 cur_head 到 sq_context->head_addr (HBM)
└── 写 doorbell 寄存器: {next_pid=cur_head<<3, qp_id=wqn}
3.2 Atomic Fetch-and-Add 完整路径
用户代码: aclshmemx_roce_atomic_fetch_add(dst, value, pe)
│
▼
aclshmemi_roce_amo_add<T, IS_MASKED>(dst, src, pe, 0, add_val, 0, ...)
│ if constexpr (XSCALE) → aclshmemi_roce_atomic_fetch_and_add
▼
aclshmemi_roce_atomic_fetch_and_add<T, IS_MASKED, B>(dst, src, pe, ...)
│
├── 获取 sq_context, remote_mem_info
├── 构建 wr:
│ remote_addr = dst
│ local_addr = sq_context->amo_addr (原子结果回写地址)
│ atomic.masked_common.swap_add_data = add_val
│ atomic.masked_common.swap_add_mask = boundary
│ rkey = remote_mem_info->rkey
│ lkey = sq_context->amo_lkey
▼
aclshmemi_backend_traits<B>::atomic_op_traits<T, IS_MASKED>
::post_send<OP_ATOMIC_FA>(wr, ...)
│
▼ (XSCALE)
aclshmemi_rdma_xscdv_post_send_atomic<T, IS_MASKED, OP_ATOMIC_FA>(wr, ...)
│
├── 检查 SQ 满 → poll_cq
├── aclshmemi_roce_xscale_fill_wqe_atomic<T, IS_MASKED, OP_ATOMIC_FA>
│ └── aclshmemi_rdma_xscdv_fill_wqe_fa<T, IS_MASKED>
│ ├── fill_wqe_ctrl_seg: opcode=8B/4B_MSK_FETCH_AND_ADD
│ ├── fill_wqe_data_seg: [rkey|remote] + [lkey|amo_addr]
│ └── fill_fa_atomic_seg: [add_data|field_boundary] (big-endian)
├── dcci_cachelines(wqe_addr, 128)
└── ring_sq_doorbell
3.3 Quiet (同步等待完成)
aclshmemx_roce_quiet(pe, buf, sync_id)
│
▼
aclshmemi_roce_quiet(pe, qp_idx, ub64, ub32, sync_id)
│
├── 遍历所有 qp_idx ∈ [0, qp_num)
├── 读取 sq_context->head_addr → cur_head
▼
aclshmemi_roce_poll_cq<XSCALE>(pe, qp_idx, cur_head, ...)
│
├── 读取 CQ tail (HBM)
├── while (cur_tail != target_idx):
│ ├── 计算 CQE 地址: cq_base + cqe_size * (cur_tail % depth)
│ ├── 检查 owner_bit == (cur_tail / depth) & 1
│ ├── dcci_cachelines(cqe_addr, 64)
│ ├── 超时检测: 5 分钟 (按 cycle 计数)
│ └── 检查 error_code
├── cur_tail++
└── ring_cq_doorbell: 写 CQ tail + doorbell
4. 编译时多态 — aclshmemi_backend_traits 系统
4.1 设计理念
采用 CRTP 风格的特化模式（无需虚函数），通过模板特化实现零开销后端分发：
// 基类模板 (未定义, 强制特化)
template <aclshmemi_rdma_backend_t B>
struct aclshmemi_backend_traits {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t OP>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...);
template <aclshmemi_rdma_atomic_op_t OP>
static ACLSHMEM_DEVICE void post_send(...);
};
};
// XSCALE 特化 (完整实现)
template <>
struct aclshmemi_backend_traits<aclshmemi_rdma_backend_t::XSCALE> {
template <typename T, bool IS_MASKED>
struct atomic_op_traits {
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE uint32_t fill_wqe(...) {
return aclshmemi_roce_xscale_fill_wqe_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
template <aclshmemi_rdma_atomic_op_t ATOMIC_OP_CODE>
static ACLSHMEM_DEVICE void post_send(...) {
aclshmemi_rdma_xscdv_post_send_atomic<T, IS_MASKED, ATOMIC_OP_CODE>(...);
}
};
};
4.2 编译时错误保护
// 未实现的原子操作 → 编译错误
template <aclshmemi_rdma_atomic_op_t>
struct aclshmemi_atomic_op_dependent_false : std::false_type {};
// 未实现的后端 → 编译错误
template <aclshmemi_rdma_backend_t>
struct aclshmemi_rdma_backend_dependent_false : std::false_type {};
4.3 调用链中的编译时决议
// 粗粒度 API → 编译时选择后端
template <typename T, aclshmemi_rdma_backend_t B>
ACLSHMEM_DEVICE void aclshmemi_roce_write(...) {
aclshmemi_roce_post_send<B, OP_RDMA_WRITE>(...); // B = XSCALE
}
// 用户入口 (shmem_device_rdma.hpp)
aclshmemi_roce_write<T, ACLSHMEMI_K_RDMA_BACKEND>(...); // 编译时绑定
关键优势: 所有后端选择、原子操作类型判断均在编译时完成，生成的机器码中无任何 switch/if 分支。
5. XSCALE 后端 — WQE/CQE 硬件格式
5.1 WQE 控制段 (16 bytes)
struct aclshmemi_xscdv_wqe_ctrl_seg_v2 {
uint8_t msg_opcode; // 操作码 (见 aclshmemi_xscdv_msg_type_t)
uint8_t with_imm : 1;
uint8_t csum_en : 2;
uint8_t ds_data_num : 5; // 数据段数量
uint16_t rsv1;
uint32_t msg_len; // 消息长度
uint32_t opcode_data; // 操作码相关数据
uint32_t se : 1;
uint32_t ce : 1; // CQE 使能
uint32_t in_line : 1; // 内联模式 (原子操作使用)
uint32_t fence_mode : 2;
uint32_t mask : 2;
uint32_t rsv : 5;
uint32_t wqe_id : 20; // WQE 标识
};
5.2 WQE 数据段 (16 bytes each)
struct aclshmemi_xscdv_diamond_data_seg_t {
uint32_t length;
uint32_t key; // rkey 或 lkey
uint64_t addr; // VA 地址
};
5.3 CQE 格式 (32 bytes + 32 bytes padding = 64 bytes)
struct aclshmemi_xscdv_diamond_cqe_v2 {
union { uint32_t flags_qp_id_err_code; struct {
uint32_t error_code : 8;
uint32_t qp_id : 15;
uint32_t rsv : 1;
uint32_t se : 1;
uint32_t has_pph : 1;
uint32_t type : 1;
uint32_t with_imm : 1;
uint32_t csum_err : 4;
}; };
uint32_t imm_data;
uint32_t msg_len;
uint32_t vni;
uint32_t ts_l;
uint32_t ts_h;
union { uint32_t wqe_id_rsv_opcode; struct {
uint32_t msg_opcode : 8;
uint32_t rsv1 : 4;
uint32_t wqe_id : 20;
}; };
union { uint32_t owner_rsv; struct {
uint32_t rsv2 : 31;
uint32_t owner : 1; // 软件通过此 bit 判断 CQE 是否有效
}; };
};
5.4 Doorbell 机制
Send Doorbell (8 bytes):
union aclshmemi_xscdv_diamond_send_doorbell_v2 {
struct {
uint64_t next_pid : 21; // 下一个 WQE 索引 (cur_head << 3)
uint64_t qp_id : 16; // QP 编号
};
uint64_t raw;
};
CQ Doorbell (8 bytes):
union aclshmemi_xscdv_diamond_cq_doorbell_t {
struct {
uint64_t cq_next_cid : 23; // 下一个 CQE ID
uint64_t cq_id : 16; // CQ 编号
uint64_t cq_sta : 2;
};
uint64_t raw;
};
5.5 消息操作码
enum class aclshmemi_xscdv_msg_type_t : uint32_t {
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_WRITE = 1,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_READ = 2,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_CMP_AND_SWAP = 26,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_FETCH_AND_ADD = 27,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_CMP_AND_SWAP = 31,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_8B_MSK_FETCH_AND_ADD= 32,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_CMP_AND_SWAP = 33,
ACLSHMEMI_XSCALE_MSG_OPCODE_RDMA_ATOMIC_4B_MSK_FETCH_AND_ADD= 34,
};
5.6 Owner Bit 机制
XSCALE 使用 owner bit 实现 CQE 生产-消费同步：
bool aclshmemi_roce_xscale_check_cqe_owner(cqe64, cur_tail, depth) {
bool owner_bit = cqe64->cqe.owner & 1;
bool expect_bit = (cur_tail / depth) & 1; // 每轮翻转
return owner_bit == expect_bit;
}
原理: owner bit 在每轮 CQ 环遍历后翻转，硬件写入 CQE 时设置当前轮的 owner 值，软件通过比较判断 CQE 是否已更新。
5.7 字节序转换
XSCALE 硬件要求 big-endian 格式：
// 手动实现 htobe64 (无标准库依赖)
ACLSHMEM_DEVICE uint64_t aclshmemi_htobe64(uint64_t host_val) {
return ((host_val >> 0) & 0xFF) << 56 |
((host_val >> 8) & 0xFF) << 48 |
... // 逐字节反转
}
注意: 4B fetch/swap 操作在 XSCALE 上获取的数据为 little-endian，需要额外 aclshmemi_htobe32 转换（见 aclshmemi_roce_get_atomic_fetch_data）。
6. 原子操作 — 16 种变体的 IS_MASKED 设计
6.1 变体矩阵
| 操作 | 类型 | IS_MASKED | 硬件 OPCODE | WQE 结构 |
|---|---|---|---|---|
| FA | 64B | false | FETCH_AND_ADD(27) | ctrl + 2 data + atomic_seg(16B) |
| FA | 64B | true | 8B_MSK_FETCH_AND_ADD(32) | ctrl + 2 data + masked_fa_seg(16B) |
| FA | 32B | true | 4B_MSK_FETCH_AND_ADD(34) | ctrl + 2 data + masked_fa32_seg(16B) |
| FA | 32B | false | 不支持 | static_assert 报错 |
| CAS | 64B | false | CMP_AND_SWAP(26) | ctrl + 2 data + atomic_seg(16B) |
| CAS | 64B | true | 8B_MSK_CMP_AND_SWAP(31) | ctrl + 2 data + 2×masked_cas_seg(32B) |
| CAS | 32B | true | 4B_MSK_CMP_AND_SWAP(33) | ctrl + 2 data + masked_cas32_seg(16B) |
| CAS | 32B | false | 不支持 | static_assert 报错 |
6.2 高级 API 到硬件原语的映射
// atomic_set → CAS with mask=UINT64_MAX, comp_mask=0
aclshmemi_roce_amo_cas<T, true>(dst, nullptr, pe, 0, value, 0, UINT64_MAX, 0, ...);
// atomic_swap → CAS with comp=0, mask=UINT64_MAX, comp_mask=0
aclshmemi_roce_amo_cas<T, true>(dst, nullptr, pe, 0, value, 0, UINT64_MAX, 0, ...);
// atomic_fetch → FA with add=0, boundary=0
aclshmemi_roce_amo_add<T, true>(dst, nullptr, pe, 0, 0, 0, ...);
// atomic_and → CAS: swap_mask=~value, comp_mask=0
aclshmemi_roce_amo_cas<T, true>(dst, nullptr, pe, 0, value, 0, ~value, 0, ...);
// atomic_or → CAS: swap_mask=value, comp_mask=0
aclshmemi_roce_amo_cas<T, true>(dst, nullptr, pe, 0, value, 0, value, 0, ...);
// atomic_xor → FA: add=value, boundary=UINT64_MAX
aclshmemi_roce_amo_add<T, true>(dst, nullptr, pe, 0, value, UINT64_MAX, ...);
设计洞察: 所有逻辑原子操作 (and/or/xor) 都通过 CAS 或 FA 的 masked 变体实现，利用 mask 位控制实现条件更新。
6.3 原子操作结果获取
template <typename T>
ACLSHMEM_DEVICE T aclshmemi_roce_get_atomic_fetch_data(uint32_t pe, uint32_t qp_idx) {
auto amo_addr = sq_context->amo_addr; // 硬件回写地址
dcci_cachelines(amo_addr, sizeof(T));
T data = *(__gm__ T *)amo_addr;
if constexpr (sizeof(T) == 4 && XSCALE) {
return aclshmemi_htobe32(data); // 4B 需要字节序转换
}
return data;
}
7. Host 端传输管理 — RdmaTransportManagerV2
7.1 生命周期
OpenDevice() → RegisterMemoryRegion() → Prepare() → Connect() → [运行] → CloseDevice()
│ │ │ │
▼ ▼ ▼ ▼
创建 Endpoint HcommMemReg 解析 NIC HcommChannelCreate
分配 device IP (HBM → MR) 地址/角色 + FillRdmaInfo
7.2 关键常量
constexpr uint32_t RDMA_PORT_PREFIX = 60032; // 基础端口号
constexpr uint32_t MAX_RANKS_PER_NIC = 16; // 每 NIC 最大 rank 数
constexpr uint32_t ATOMIC_MAX_NUM = 128; // 最大原子操作并发数
constexpr uint32_t RDMA_NOTIFY_NUM = 3; // notify 数量
constexpr uint32_t RDMA_QUEUE_NUM = 1; // 每连接 QP 数
constexpr uint32_t MEMORY_ALIGNMENT = 512; // 内存对齐
7.3 端口分配策略
// Endpoint 端口
devicePort_ = RDMA_PORT_PREFIX + (rankId_ % MAX_RANKS_PER_NIC); // 60032 + rank%16
// Channel 端口 (server/client 模式)
port = RDMA_PORT_PREFIX + (serverRank % 16) * 16 + (clientRank % 16);
// 即: 60032 + server*16 + client
7.4 MR 注册与 key 管理
Result RegisterMemoryRegion(const TransportMemoryRegion& mr) {
CommMem commMem;
commMem.type = (mr.flags & REG_MR_FLAG_HBM) ? COMM_MEM_TYPE_DEVICE : COMM_MEM_TYPE_HOST;
commMem.addr = mr.addr;
commMem.size = mr.size;
HcommMemReg(endpointHandle_, "HcclBuffer", &commMem, &memHandle);
registeredMRs_.emplace(mr.addr, {memHandle, mr.addr, mr.size});
}
7.5 原子操作共享内存
bool RegisterAtomicMemory() {
uint32_t atomicSize = ATOMIC_MAX_NUM * sizeof(uint64_t) * rankCount_; // 128 * 8 * ranks
atomicSize = ALIGN_UP(atomicSize, 512);
HcommMemReg(endpointHandle_, "AtomicBuffer", &commMem, &memHandle);
}
每个 rank 拥有 128 × 8B = 1KB 的原子操作结果缓冲区，用于硬件回写 fetch/swap 结果。
7.6 FillRdmaInfo — 信息汇聚到 Device
Result FillRdmaInfo() {
// 1. 分配临时 buffer
// 2. FillQpPreSettingCopyInfo: 计算 sq/rq/scq/rcq/mr 指针布局
// 3. GetRdmaInfoFromChannelEntity:
// ├── ReadLocalBufferInfo: 从 channel entity 读取本地 MR lkey/rkey
// └── ReadRemoteBufferInfo:
// └── ReadSingleRemoteRank: 读取远端 SQ/CQ/MR 信息
// └── FillSqCqAtomicInfo: 填充 SQ/CQ/atomicAddr 信息
// 4. AclrtMemcpy(qpInfo_, qpInfoSize_, copyInfo, ..., HOST_TO_DEVICE)
}
关键: 所有 QP/CQ/MR 信息在 host 侧收集后，一次性拷贝到 device HBM，设备端无需再与 host 交互即可执行 RDMA 操作。
7.7 QP Info 内存布局 (Device HBM)
┌──────────────────────────────────────────────────────────────┐
│ AiQpRMAQueueInfo (count + 指针) │
├──────────────────────────────────────────────────────────────┤
│ sq[rankCount] : AiQpRMAWQ 数组 │
├──────────────────────────────────────────────────────────────┤
│ rq[rankCount] : AiQpRMAWQ 数组 │
├──────────────────────────────────────────────────────────────┤
│ scq[rankCount]: AiQpRMACQ 数组 │
├──────────────────────────────────────────────────────────────┤
│ rcq[rankCount]: AiQpRMACQ 数组 │
├──────────────────────────────────────────────────────────────┤
│ mr[rankCount] : RdmaMemRegionInfo 数组 │
└──────────────────────────────────────────────────────────────┘
8. ZB 算子如何调用 SHMEM
8.1 调用层次
moe_dispatch_normal_zero_buffer (kernel entry)
│
▼
MoeDispatchNormalZeroBuffer::Process()
│
├── syncFlag_.IncrementMagic() ← ZeroBufferSyncFlag (shmem_ptr)
├── syncFlag_.BarrierAll() ← 跨卡 barrier
├── InputToDstOutput() ← 数据搬运
│ └── shmem_ptr(expandXOutGM, dstRankId) ← 获取远端对称地址
│ └── DataCopyPad(dstGT, xOutTensor, ...) ← 写入远端 HBM
└── syncFlag_.BarrierAll() ← 确保所有写入完成
8.2 shmem_ptr — 对称地址转换
// shmem_device_rdma.hpp
ACLSHMEM_DEVICE __gm__ void *aclshmem_roce_ptr(__gm__ void *ptr, int pe) {
auto device_state = aclshmemi_get_state();
uint64_t offset = (uint64_t)ptr - (uint64_t)device_state->heap_base;
uint64_t remote_ptr = (uint64_t)device_state->p2p_device_heap_base[pe] + offset;
return (__gm__ void *)remote_ptr;
}
原理: 所有 rank 的对称内存在各自地址空间中具有相同偏移量，通过 远端 heap_base + 本地偏移 计算远端可访问地址。
8.3 ZeroBufferSyncFlag — 跨卡同步协议
内存布局 (per rank, 从 gva_gm + 100KB 开始):
┌─────────────────────────────────────────────────────────┐
│ Magic Area (24KB = 48 cores × 512B) │
│ slot[coreIdx] — 单调递增, 每 kernel +1 │
├─────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 0 (epWorldSize × 32B) │
├─────────────────────────────────────────────────────────┤
│ Barrier PingPong Buffer 1 (同上) │
├─────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 0 (epWorldSize × slotsPerRank × 32B)│
├─────────────────────────────────────────────────────────┤
│ Flag PingPong Buffer 1 (同上) │
└─────────────────────────────────────────────────────────┘
同步协议:
Writer (rank A, core i):
1. 完成数据写入 (DMA to remote shmem)
2. MTE3 fence (flush all prior writes)
3. SetFlag(destRank=B, eventID=i, phase)
→ 写 flag 到 B 的本地 shmem slot[A][i]
Reader (rank B, core j):
1. WaitFlag(srcRank=A, eventID=i, phase)
→ 轮询 B 的本地 shmem 直到 flag 匹配
2. 返回后, A 的数据保证可见
8.4 Dispatch 同步流程
void Process() {
// Step 1: IncrementMagic — 获取 magic M+1
syncFlag_.IncrementMagic();
// Step 2: BarrierAll — 确保 Notify 结果全局可见
// (putOffset, recvData 等已计算完成)
syncFlag_.BarrierAll();
// Step 3: 通过 shmem 写 token 数据到远端 rank
InputToDstOutput();
// └── shmem_ptr(expandXOutGM, dstRankId) → 远端地址
// └── DataCopyPad(dstGT, xOutTensor, ...) → 写入
// Step 4: BarrierAll — 确保所有 rank 写入完成
// 后续 GMM kernel 可安全消费
syncFlag_.BarrierAll();
}
8.5 BarrierAll 实现
void BarrierAll() {
// Step 1: 卡内同步
SyncAll<true>();
// Step 2: 分配 rank 到 core (SplitCoreCal)
// Step 3: 写 barrier flag 到所有远端 rank
for (r = startRank; r < endRank; r++) {
if (r == epRankId_) continue;
remotePtr = shmem_ptr(flagBaseGM_, r);
DataCopy(remoteSlot, ubFlag, FLAG_ELEM_NUM);
}
// Step 4: 轮询本地 slot 等待所有远端 flag
for (r = startRank; r < endRank; r++) {
if (r == epRankId_) continue;
do { DataCopy(ubFlag, localSlot, ...); }
while (ubFlag.GetValue(0) != barrierFlag);
}
// Step 5: 卡内同步
SyncAll<true>();
}
8.6 数据搬运核心逻辑
void InputToDstOutput() {
// 1. 读取 putOffset (全局 recv_count 前缀和)
DataCopyPad(putOffsetTensor, putOffsetGT, ...);
// 2. 读取 expertIds 和 sendTokenIdx
DataCopyPad(expertIdsTensor, expertIdsGT[startTokenId], ...);
DataCopyPad(sendTokenIdxTensor, sendTokenIdxGT[startTokenId], ...);
// 3. 逐 token 处理
for (tokenIndex = startTokenId; tokenIndex < endTokenId; ++tokenIndex) {
dstExpertId = expertIdsTensor(...);
dstRankId = dstExpertId / moeExpertNumPerRank;
curExpertIdx = sendTokenIdxTensor(...);
dstExpertOffset = putOffsetTensor(dstExpertId * epRankSize + epRankId);
// 计算远端输出地址
auto ptr = shmem_ptr(expandXOutGM, dstRankId);
uint64_t expandXByteOffset = hUBAlignSize * (dstExpertOffset + curExpertIdx);
dstGT.SetGlobalBuffer(ptr + expandXByteOffset);
// 写入数据
DataCopyPad(dstGT, xOutTensor, xOutCopyParams);
}
}
9. 与 master 版 NDA 的差异 — cam_feature 上的 SHMEM 框架演进
9.1 架构分层演进
| 维度 | master (原始 NDA) | cam_feature (ZB 算子) |
|---|---|---|
| API 层 | aclshmemx_roce_put_nbi/get_nbi/atomic_* |
shmem_ptr() + DataCopyPad |
| 传输模式 | RDMA one-sided (NIC 直接读写) | 设备端 AIV 直接写 HBM (无需 NIC) |
| 同步机制 | aclshmemx_roce_quiet (轮询 CQ) |
ZeroBufferSyncFlag (flag-based) |
| 内存模型 | 对称内存 + 远端指针转换 | 对称内存 + shmem_ptr 地址转换 |
| 后端 | XSCALE RDMA 硬件 | AIV core 直接 store (不走 RDMA) |
| 原子操作 | 16 种硬件原子 (FA/CAS masked) | 不需要 (单方面写入) |
9.2 关键差异分析
1. 数据传输路径不同
master 的 NDA 框架使用 RDMA one-sided 操作：
- 由 NIC 直接执行 RDMA_WRITE / RDMA_READ
- 需要 SQ doorbell / CQ poll 机制
- 数据路径: Local HBM → NIC → (wire) → NIC → Remote HBM
cam_feature 的 ZB 算子使用 AIV core 直接写入：
- 由 AI Vector core 执行 DataCopyPad (本质是 store 指令)
- 利用对称内存的地址直接可达性
- 数据路径: Local UB → AIV store → Remote HBM (通过内部互联)
2. 同步协议升级
master: aclshmemx_roce_quiet 需要轮询 CQ，等待硬件完成中断。
cam_feature: ZeroBufferSyncFlag 使用 magic + phase 的 flag 协议，支持：
- ping-pong buffer 避免 stale flag
- per-core / per-rank / per-expert 多级粒度
- 无需硬件中断，纯软件轮询
3. 内存布局精细化
cam_feature 引入了更精细的内存分区：
gva_gm + 0: Notify 区域 (2MB)
gva_gm + 2MB: Dispatch 区域 (3MB)
gva_gm + 5MB: Combine 区域 (3MB)
gva_gm + 100KB: Flag Area (同步标志区)
├── Magic Area (24KB)
├── Barrier PingPong (2 × epWorldSize × 32B)
└── Flag PingPong (2 × epWorldSize × slotsPerRank × 32B)
9.3 演进趋势
- 从 RDMA 到 AIV-direct: ZB 算子不直接使用 NDA 的 RDMA API，而是利用对称内存 + AIV 直接写入，避免了 NIC 交互开销。
- 从 quiet 到 flag-based sync: 同步机制从硬件 CQ 轮询演进为软件 flag 协议，更灵活且支持复杂 DAG 依赖。
- 从粗粒度到细粒度: 同步粒度从 per-QP 演进到 per-core/per-expert，支持更精细的流水线。
- SHMEM 作为基础设施: NDA 框架提供的
shmem_ptr地址转换和对称内存管理成为 ZB 算子的基础，但 RDMA 操作本身被绕过。
10. 编译与构建
10.1 NDA 框架构建 (master)
# build/cam/framework/nda/build.sh
# 1. 克隆 shmem 包 (外部依赖)
git clone https://gitcode.com/cann/shmem
cd shmem && git checkout 6a5f4b09 -b nda-base
# 2. 覆盖 rdma_backends (注入 UMDK 定制后端)
cp -f ${RDMA_BACKENDS_SRC}/* shmem/src/device/gm2gm/engine/rdma_backends/
# 3. 构建 (仅支持 Ascend950 + XSCALE)
bash scripts/build.sh -enable_rdma -rdma_backend XSCALE -soc_type Ascend950
10.2 ZB 算子构建 (cam_feature)
# CMakeLists.txt
# SHMEM 头文件路径 (外部包)
-I$ENV{SHMEM_HOME_PATH}/shmem/include
-I$ENV{SHMEM_HOME_PATH}/shmem/src/device
# 编译选项
-DASCENDC_DUMP=0 --cce-auto-sync=off -DCATLASS_ARCH=2201
关键: ZB 算子通过 SHMEM_HOME_PATH 环境变量引用外部 shmem 包的头文件，不直接包含 NDA 框架源码。
11. 关键设计洞察
11.1 编译时多态 vs 运行时多态
NDA 框架大量使用模板特化而非虚函数：
- 零运行时开销: 所有后端选择在编译时确定
- 类型安全: 未实现的操作在编译时报错 (static_assert)
- 代码生成: 每个后端生成独立的特化代码，无分支预测失败
11.2 内存一致性模型
XSCALE 后端使用 dcci_cachelines (cache line flush) 确保：
- WQE 写入对 NIC 可见后再 ring doorbell
- CQ polling 时读取最新 CQE
- 原子操作结果对 CPU 可见
11.3 超时与容错
// 5 分钟超时 (按 cycle 计数)
constexpr uint64_t ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_DURATION = 5ULL * 60 * 1000000;
constexpr uint64_t ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_CYCLES =
ACLSHMEMI_XSC_POLL_CQ_TIMEOUT_DURATION * ACLSHMEMI_XSC_CYCLE_TO_TIME_BASE;
// 950M 系列: ×50, 其他: ×1000
11.4 字节序陷阱
XSCALE 硬件使用 big-endian：
- 所有 WQE 中的 64-bit 数据需要 htobe64 转换
- 4B fetch/swap 结果为 little-endian，需要额外 htobe32 转换
- 这是设备端代码中容易出错的地方
11.5 SQ 流控
constexpr int XSCALE_POLL_CQ_THRESHOLD = 16;
// 当 SQ 即将满时 (head + 16 == tail), 先 poll_cq 等待空闲
if ((cur_head + XSCALE_POLL_CQ_THRESHOLD) % depth == (depth + cur_tail) % depth) {
aclshmemi_roce_poll_cq<XSCALE>(pe, qp_idx, cur_head, ...);
}
设计: 提前 16 个 WQE 触发反压，避免 SQ 完全填满导致死锁。
12. 总结
SHMEM/NDA 框架是 UMDK/CAM 中设备端 RDMA 通信的基础设施层，其核心设计特点：
- 分层架构: 公共 API → 调度层 → 后端抽象 → 硬件特化，每层职责清晰
- 编译时多态: 通过
aclshmemi_backend_traits实现零开销后端分发 - 硬件原生: XSCALE 后端直接操作 WQE/CQE/doorbell，无中间抽象层
- Host-Device 分离: Host 负责 QP 建立/MR 注册/信息收集，Device 独立执行 RDMA 操作
- ZB 算子演进: cam_feature 上的 ZB 算子利用对称内存绕过 RDMA，使用 AIV 直接写入 + flag 同步，实现了更低延迟的通信
该框架目前仅支持 Ascend950 + XSCALE 后端，IN_DIE 后端尚未实现（编译时 fallback 到 XSCALE）。
三种通信范式对比
| 维度 | 非 SHMEM (HCCL Window) | SHMEM (用户态 RDMA) | ZB (Zero-Buffer) |
|---|---|---|---|
| 通信原语 | winContext->localWindowsIn + flag 同步 | aclshmemx_roce_get_nbi / put_nbi / atomic_* | shmem_ptr() + DataCopyPad + ZeroBufferSyncFlag |
| 硬件操作 | 通过 HCCL 软件栈间接访问 | AIV 核直接操作 WQE/CQE 硬件队列 | AIV 核直接写入远端 HBM |
| 同步机制 | Magic+Value flag (远程状态位写入) | aclshmemx_roce_quiet (本地 CQ 轮询) | ZeroBufferSyncFlag (单调递增 magic + ping-pong) |
| 内存占用 | 每卡 ~200MB Window 内存 | QP/MR 资源，内存占用低 | 零中间拷贝，数据直达对端 User Buffer |
| 延迟特性 | 中等 (HCCL 调度开销) | 极低 (绕过 HCCL) | 极低 (零拷贝 + 直写远端) |
| 同步粒度 | per-token per-topk (O(bs×K) polling) | per-QP (quiet 操作) | per-rank (O(1) BarrierAll) |
| 后端选择 | 运行时 HCCL 拓扑感知 | 编译时 traits: IN_DIE / XSCALE | 外部 shmem 包 |
| master 状态 | 生产默认路径 | 基础设施就绪，算子 requires_shmem: false | 不存在 |
| cam_feature 状态 | 仍支持 (w4a8/fwk 变体) | Base 变体 requires_shmem: true | 生产默认路径 |
master → cam_feature 架构演进
| 方面 | master | cam_feature |
|---|---|---|
| 算子数 | 5 | 13 |
| SHMEM 算子 | 0 (全 false) | 9 (大部分 true) |
| 新增算子 | — | a2e, e2a, all_to_all_detour, reduce_scatter_detour |
| ZB 变体 | — | 6 个 (zero_buffer 系列) |
| FusedDeepMoE | 1 变体 | 3 变体 (base/w4a8/fwk) |
| 新模块 | — | HFA (HBM Fast Alloc) |
| 通信范式 | 纯 HCCL Window | ZB (Zero-Buffer) 为主 |
关键创新总结
Zero-Buffer 核心创新
- 用 CANN SHMEM（用户态 RDMA）完全替代 HCCL Window，每卡节省 ~200MB 内存
- 数据从源卡直达对端 User Buffer（零中间拷贝）
- 同步粒度从 O(bs×K) 降到 O(1)，彻底解决 ABA 问题
- Detour 机制利用拓扑局部性降低跨 pod 拥塞
- HFA 模块为 ZB 算子提供动态共享内存分配能力
项目内部文档
| 文档 | 路径 |
|---|---|
| CAM README | src/cam/README.md |
| CAM API Guide (EN) | doc/en/cam/CAM API Guide.md |
| CAM API Guide (中文) | doc/ch/cam/CAM API Guide.ch.md |
| UMDK 总 README | README.md |
| Release Notes | RELEASE-NOTES.md |
外部参考
| 资源 | 说明 |
|---|---|
| Huawei Ascend CANN | Ascend 计算架构与神经网络库 |
| googletest v1.16.0 | C++ 单元测试框架 |
| OpenMPI 5.0.7 | 多进程通信测试 |
| catlass (submodule) | AscendC 内核模板库，pin 在 v1-stable 分支 |
| CANN SHMEM (外部包) | gitcode.com/cann/shmem, commit 6a5f4b09 |
分析过程记录
分析方法本报告采用 16-Agent 并行深度分析架构：
- Agent-1~9：master 分支基础分析（架构、编译、EP 算子、FusedMoE、测试、SHMEM、HCCL Window、notify+layout、MoE SHMEM）
- Agent-10~16：cam_feature 分支深度分析（ZB Dispatch、ZB Combine、FusedMoE 三变体、新通信算子、ZB 工具/HFA、Pybind/部署、SHMEM 框架/NDA）
总计分析 15698 行，覆盖 98+ 个源文件 + 262 个测试文件。
分文件分析报告索引
| # | 报告 | 路径 |
|---|---|---|
| 1 | 整体架构与设计哲学 | analysis/cam/agents/01_architecture.md |
| 2 | 编译系统与部署 | analysis/cam/agents/02_build_deploy.md |
| 3 | EP 通信算子 (master) | analysis/cam/agents/03_ep_kernels.md |
| 4 | FusedDeepMoE (master) | analysis/cam/agents/04_fused_deep_moe.md |
| 5 | 测试体系与回归测试 | analysis/cam/agents/05_test_system.md |
| 6 | SHMEM 模式 (master) | analysis/cam/agents/06_shmem_deepdive.md |
| 7 | 非 SHMEM — HCCL Window (master) | analysis/cam/agents/07_hccl_window_deepdive.md |
| 8 | notify_dispatch + dispatch_layout (master) | analysis/cam/agents/08_notify_dispatch_layout_deepdive.md |
| 9 | FusedDeepMoE SHMEM 路径 (master) | analysis/cam/agents/09_fused_moe_shmem_deepdive.md |
| 10 | ZB Dispatch 算子族 (cam_feature) | analysis/cam/agents/10_zb_dispatch_operators.md |
| 11 | ZB Combine 算子族 (cam_feature) | analysis/cam/agents/11_zb_combine_operators.md |
| 12 | FusedDeepMoE 三变体 (cam_feature) | analysis/cam/agents/12_fused_moe_family.md |
| 13 | 新通信算子 — A2E/E2A/Detour (cam_feature) | analysis/cam/agents/13_new_comm_operators.md |
| 14 | ZB 工具 + 编译系统 + HFA (cam_feature) | analysis/cam/agents/14_zb_utils_build_hfa.md |
| 15 | Pybind + Examples + Docker (cam_feature) | analysis/cam/agents/15_pybind_examples_docker.md |
| 16 | SHMEM 框架 + NDA 层 (cam_feature) | analysis/cam/agents/16_shmem_framework_nda.md |