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 的抽象层级过高,无法满足:

  1. 延迟敏感:需要直接 post RDMA WQE 到硬件 SQ,避免 HCCL 的软件调度开销
  2. 计算-通信重叠:需要在 kernel 内部发起 RDMA 操作后立即进行 GEMM 计算
  3. 原子操作: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 端管理两类内存注册:

  1. 用户 HBM 缓冲区(RegisterMemoryRegion):通过 HcommMemReg 注册,区分 COMM_MEM_TYPE_DEVICE(HBM)和 COMM_MEM_TYPE_HOST(DDR)。
  2. 原子操作共享内存(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 通信的垂直优化软件栈,其架构设计体现了以下核心洞察:

  1. MoE EP 通信不适合传统集合通信抽象 — 需要用户态 RDMA 获得极致延迟
  2. 计算-通信融合是性能关键 — FusedDeepMoE 将 5 阶段流水线融合为单 kernel
  3. 硬件多样性需要编译时多态 — traits 系统零开销适配 IN_DIE/XSCALE 后端
  4. 生态兼容性需要分层设计 — 控制面复用 HCCP,数据面自研,三层编程模型符合 CANN 规范

NDA 层是 CAM 的"皇冠上的明珠",它实现了一个轻量级用户态 RDMA 子系统,在 Ascend 芯片上直接操作 RoCE 硬件,同时保持与 HCCP 控制面的兼容性。这种设计模式对理解国产 AI 芯片的软件栈演进具有重要参考价值。

📑 索引 编译系统与部署 →