CAM (Communication Acceleration for Matrix) 架构深度分析报告
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 芯片的软件栈演进具有重要参考价值。