SHMEM 框架与 NDA (Network Device Abstraction) 层深度分析
代码仓库: 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）。