昇腾多卡之间搬数据太麻烦？CANN SHMEM 了解一下——个人开发者的一点体验分享

仓库地址：https://gitcode.com/cann/shmem

写在前面

最近在昇腾 NPU 上做多卡算子开发，遇到一个绕不开的问题：卡间数据交互。用集合通信库（HCCL）吧，粒度太粗、模式固定；自己手搓 RDMA / DMA 吧，链路管理、内存注册、同步原语一堆坑，对个人开发者实在不友好。

后来发现 CANN 官方在 GitCode 上开源了 SHMEM——面向昇腾平台的多机多卡分布式内存通信加速库，用了一段时间，感觉思路被打开了，分享给大家。

SHMEM 是什么

SHMEM 基于 OpenSHMEM 标准协议，核心思路是"对称内存 + 远程内存访问（RMA）"：

• 每张卡（每个 PE）开辟一块对称内存堆，大家地址布局一致；
• 通过 put / get 这类 RMA 接口，一张卡可以直接读写另一张卡的内存；
• 配合 barrier、锁、信号量等同步原语完成数据同步。

它把通信这件事从"调用集合通信 API"下沉成了"像访问普通指针一样访问远端内存"，编程模型一下子简洁了不少。

几个我觉得比较亮眼的特性：

• 双侧接口体系：Host 侧负责初始化、内存堆管理、通信域（Team）管理；Device 侧提供 RMA、设备级同步等接口，可以直接在算子 Kernel 里发起通信；
• AICore 直驱：支持 AICore 直驱 MTE、xDMA 引擎，实现 D2D / D2H / H2D / D2rH / rH2D 多种通路，"计算 + 通信"融合类算子（比如 matmul + allreduce）有非常大的优化空间；
• 兼容 MPI 生态：支持集合通信场景，也能接入 PyTorch 分布式训练流程替代传统 MPI；
• 默认 TLS 加密：跨设备传输默认加密，安全这块有兜底。

上手体验

SHMEM 提供三种安装方式，我最喜欢 pip 这种：

# 方式一：pip 安装（Python 3.10~3.12，依赖 torch-npu）
python3 -m pip install cann-shmem \
  --index-url https://ascend.devcloud.huaweicloud.com/cann/pypi/simple/
shmem-config --version   # 查版本
shmem-config --diagnose  # 环境自检 + 问题诊断

shmem-config 这个工具好评，装完跑一遍 --diagnose，后端、加载路径、环境问题都能查出来，省了很多排障时间。

习惯源码编译的也可以：

# 方式二：源码编译
git clone https://gitcode.com/cann/shmem.git
cd shmem
bash scripts/build.sh              # 910 A2/A3 平台
# bash scripts/build.sh -soc_type Ascend950   # 950 平台
source install/set_env.sh

仓库里的 examples 目录 提供了从基础通信（rdma_demo）到通算融合（配合 CATCCOS 算子库，矩阵乘法、ReduceScatter 等）的一整套样例，跟着跑一遍就能理解 RMA 编程模型。

一点个人看法

适合谁：

• 做多卡 / 多机分布式算子开发，被卡间通信折磨的人；
• 想做"通算融合"（计算与通信融合）类自定义算子的开发者；
• 搞分布式训练框架适配、想替代 MPI 降低通信开销的团队；
• 研究对称内存编程模型的高校学生（华南理工陆璐教授团队也参与了贡献）。

需要注意的：

• 仅适配昇腾平台，不支持 x86 通用服务器和 NVIDIA GPU；
• 需要真实的 NPU 硬件环境（无 NPU 只能做编译验证，跑不了用例）；
• RDMA 场景对网络环境有要求（网卡、无损网络配置、时钟同步等），FAQ 里排障条目不少，建议提前看快速入门文档；
• CANN 版本升级可能带来接口兼容问题，建议按文档锁定配套版本。

总体来说，SHMEM 给昇腾上的多卡编程补上了一块重要拼图——尤其 Device 侧能在 Kernel 里直接发起远端内存读写这一点，让"通信融进计算"变成了随手可做的事。项目 2025 年 12 月首次上线，目前迭代很活跃（v1.6.0 已发布），感兴趣的建议先 Star 一个：

https://gitcode.com/cann/shmem

更详细的接口文档可以看官方文档站：https://shmem-doc.pages.dev/，配合仓库里的 快速入门 和 FAQ，基本不会踩大坑。