昇腾CANN shmem：把多张 NPU 的 HBM 变成一块全局内存

hccl 的通信模型是消息传递——发送方调 send，接收方调 recv，两边同步。hixl 的模型是单边推送——发送方调 put，接收方不用参与。shmem 是第三种模型：PGAS（Partitioned Global Address Space），每张 NPU 的 HBM 都是全局地址空间的一部分，任何一张 NPU 可以直接读写其他 NPU 的 HBM 地址，像访问本地内存一样。

这意味着：8 张 Ascend 910 的 HBM（8×32GB = 256GB）在编程模型上是一块 256GB 的全局内存。不用消息传递，不用同步配对，直接用地址访问。

PGAS 和消息传递的区别

用一段代码感受差异。8 张 NPU 各持有一个大矩阵的一行，需要更新第 3 张 NPU 上的某个元素：

// 消息传递模型（hccl 风格）// 卡0想更新卡3上的数据：必须让卡3配合// 卡0：发一个"更新请求"消息hcclSend(&update_value,1,FLOAT,dst=3,comm);// 卡3：必须调 recv 来收这个值hcclRecv(&local_buffer,1,FLOAT,src=0,comm);local_matrix[row][col]=local_buffer;// 卡3手动写入

// PGAS 模型（shmem 风格）// 卡0直接写卡3的 HBM 地址// 初始化时注册了远程地址映射shmem_float_p(&remote_matrix[3][row][col],update_value,pe=3);// 一行代码，卡3不参与

shmem_float_p的语义是：把update_value写到 PE（Processing Element）3 的remote_matrix[3][row][col]地址上。PE 3 不需要调任何函数——数据直接出现在它的 HBM 里。

这种编程模型在什么场景下最有价值：不规则数据访问。消息传递擅长规则的全局通信（AllReduce），但对于「卡 0 需要更新卡 3 上的第 7 行第 12 列」这种随机地址写入，消息传递的开销在同步上，PGAS 的开销只在网络延迟上。

shmem 的核心操作

Put / Get：单边读写

#include<shmem.h>// 初始化 shmemshmem_init();intmy_pe=shmem_my_pe();// 当前 NPU 编号 (0-7)intn_pe=shmem_n_pes();// 总 NPU 数量 (8)// 每张 NPU 分配本地数组（也是全局可见的）double*local_data=shmem_malloc(1024*sizeof(double));// Put：把本地的值写到远程 NPU 的地址// 语义：local_data[100] = 42.0，但写的是 PE 3 上的地址doubleval=42.0;shmem_double_p(&local_data[100],val,3);// 写到 PE 3// Get：从远程 NPU 的地址读值// 语义：val = remote_data[200]，但读的是 PE 5 上的地址doubleremote_val=shmem_double_g(&local_data[200],5);// 从 PE 5 读

shmem_malloc和aclrtMalloc的区别：shmem_malloc分配的内存自动注册到全局地址空间——其他 NPU 可以用shmem_*_p/shmem_*_g直接访问。aclrtMalloc分配的内存只有本地可见。

批量传输

单元素 Put/Get 的粒度太细，网络延迟吃不掉。shmem 提供了批量操作：

// 批量 Put：把本地的 256 个 double 写到 PE 3 的地址doublesend_buf[256];for(inti=0;i<256;i++)send_buf[i]=my_pe*1000+i;shmem_double_put(local_data,send_buf,256,3);// 参数：目标地址（远程PE上的）、源地址（本地的）、元素数、目标PE编号// 批量 Get：从 PE 5 的地址读 256 个 doubledoublerecv_buf[256];shmem_double_get(recv_buf,local_data,256,5);// 参数：目标地址（本地的）、源地址（远程PE上的）、元素数、源PE编号

批量操作走的是 RDMA 传输——数据从一张 NPU 的 HBM 经 RoCE 网卡直接写到另一张 NPU 的 HBM，不经过 CPU。批量越大，RDMA 的启动开销（约 5μs）摊得越薄，有效带宽越接近 RoCE 的峰值（100 Gbps 单链路）。

同步原语

PGAS 模型下，Put 完成不代表远程 NPU 能看到数据——RDMA Write 完成只代表数据到了对端 HBM，对端 CPU/NPU 的缓存可能还没更新。shmem 提供了几种同步机制：

// fence：保证之前所有的 Put 操作在全局可见shmem_fence();// quiet：阻塞等待，直到之前所有的 Put/Get 操作全部完成shmem_quiet();// barrier：全局同步——所有 PE 到齐后才继续shmem_barrier_all();// sync：两个 PE 之间的点对点同步shmem_sync(pe=3);// 和 PE 3 同步

实际开发中最常用的模式是 Put + quiet：

// 写一批数据到远程，等写入完成后再通知对方shmem_double_put(remote_buf,local_buf,count,target_pe);shmem_quiet();// 确保数据已经到达远程 HBMshmem_int_atomic_inc(&flag,target_pe);// 原子递增通知标志

原子操作

shmem 支持在远程 HBM 上执行原子操作——不需要加锁，不需要对方配合：

// 远程原子加：把 PE 3 上的 counter 原子地加 1shmem_int_atomic_add(&counter,1,3);// 远程原子比较交换：// 如果 PE 5 上的 lock == 0，把它设为 1（获取锁）intold=shmem_int_atomic_compare_swap(&lock,0,1,5);if(old==0){// 成功获取锁，可以安全操作 PE 5 上的共享数据}// 远程原子取值并加：// 读取 PE 2 上的 ticket 的当前值，然后加 1intmy_ticket=shmem_int_atomic_fetch_add(&ticket,1,2);

原子操作在分布式锁、全局计数器、任务队列等场景中直接可用，不需要消息传递那套 request-response 往返。

图计算场景：为什么 shmem 比 hccl 合适

图神经网络（GNN）的分布式训练是不规则访问的典型场景。一张大图的节点分布在 8 张 NPU 上，每个节点需要聚合邻居节点的特征——但邻居可能在其他 NPU 上，且每个节点的邻居数量和位置完全不规则。

# GNN 的邻居聚合用 shmem 实现importshmem4pyasshm# 初始化shm.init()my_pe=shm.my_pe()# 本地节点特征local_features=shm.alloc(NODES_PER_PE,dtype=np.float32)# 邻居聚合：对于每个本地节点，从远程 PE 聚合邻居特征fornodeinlocal_nodes:forneighborinnode.neighbors:# 邻居在远程 PE 上——直接读它的特征，不需要对方配合remote_pe=neighbor.owner_pe remote_offset=neighbor.local_idx neighbor_feat=shm.float_get(remote_features[remote_offset],pe=remote_pe)aggregated[node]+=neighbor_feat# 用 fence 确保所有远程读取完成shm.fence()

同样的逻辑用 hccl 实现，需要先通过 AllGather 把所有节点的特征聚到每张卡上，再做本地聚合——对于稀疏图来说，AllGather 传了大量不需要的冗余数据。shmem 的按需读取避免了这种冗余。

和 hixl 的区别

两者都做单边通信，但编程模型不同：

hixl 更适合大规模、单向、低频的传输——设计目标就是 PD 分离。shmem 更适合细粒度、双向、高频的随机访问——设计目标是通用分布式计算。

三种通信模型的本质区别在于谁拥有数据的主动权。消息传递（hccl）：双方都要参与。单边推送（hixl）：发送方单方面推数据。全局地址空间（shmem）：任何一方随时读写任何位置。主动权越分散，编程灵活性越高，但一致性管理的复杂度也越高。shmem 的 fence/quiet/barrier 体系就是为这个一致性管理而设计的——用显式同步换来了完全自由的访问模式。