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

hccl 的通信模型是消息传递——发送方调 send接收方调 recv两边同步。hixl 的模型是单边推送——发送方调 put接收方不用参与。shmem 是第三种模型PGASPartitioned Global Address Space每张 NPU 的 HBM 都是全局地址空间的一部分任何一张 NPU 可以直接读写其他 NPU 的 HBM 地址像访问本地内存一样。这意味着8 张 Ascend 910 的 HBM8×32GB 256GB在编程模型上是一块 256GB 的全局内存。不用消息传递不用同步配对直接用地址访问。PGAS 和消息传递的区别用一段代码感受差异。8 张 NPU 各持有一个大矩阵的一行需要更新第 3 张 NPU 上的某个元素// 消息传递模型hccl 风格// 卡0想更新卡3上的数据必须让卡3配合// 卡0发一个更新请求消息hcclSend(update_value,1,FLOAT,dst3,comm);// 卡3必须调 recv 来收这个值hcclRecv(local_buffer,1,FLOAT,src0,comm);local_matrix[row][col]local_buffer;// 卡3手动写入// PGAS 模型shmem 风格// 卡0直接写卡3的 HBM 地址// 初始化时注册了远程地址映射shmem_float_p(remote_matrix[3][row][col],update_value,pe3);// 一行代码卡3不参与shmem_float_p的语义是把update_value写到 PEProcessing Element3 的remote_matrix[3][row][col]地址上。PE 3 不需要调任何函数——数据直接出现在它的 HBM 里。这种编程模型在什么场景下最有价值不规则数据访问。消息传递擅长规则的全局通信AllReduce但对于「卡 0 需要更新卡 3 上的第 7 行第 12 列」这种随机地址写入消息传递的开销在同步上PGAS 的开销只在网络延迟上。shmem 的核心操作Put / Get单边读写#includeshmem.h// 初始化 shmemshmem_init();intmy_peshmem_my_pe();// 当前 NPU 编号 (0-7)intn_peshmem_n_pes();// 总 NPU 数量 (8)// 每张 NPU 分配本地数组也是全局可见的double*local_datashmem_malloc(1024*sizeof(double));// Put把本地的值写到远程 NPU 的地址// 语义local_data[100] 42.0但写的是 PE 3 上的地址doubleval42.0;shmem_double_p(local_data[100],val,3);// 写到 PE 3// Get从远程 NPU 的地址读值// 语义val remote_data[200]但读的是 PE 5 上的地址doubleremote_valshmem_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(inti0;i256;i)send_buf[i]my_pe*1000i;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(pe3);// 和 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获取锁intoldshmem_int_atomic_compare_swap(lock,0,1,5);if(old0){// 成功获取锁可以安全操作 PE 5 上的共享数据}// 远程原子取值并加// 读取 PE 2 上的 ticket 的当前值然后加 1intmy_ticketshmem_int_atomic_fetch_add(ticket,1,2);原子操作在分布式锁、全局计数器、任务队列等场景中直接可用不需要消息传递那套 request-response 往返。图计算场景为什么 shmem 比 hccl 合适图神经网络GNN的分布式训练是不规则访问的典型场景。一张大图的节点分布在 8 张 NPU 上每个节点需要聚合邻居节点的特征——但邻居可能在其他 NPU 上且每个节点的邻居数量和位置完全不规则。# GNN 的邻居聚合用 shmem 实现importshmem4pyasshm# 初始化shm.init()my_peshm.my_pe()# 本地节点特征local_featuresshm.alloc(NODES_PER_PE,dtypenp.float32)# 邻居聚合对于每个本地节点从远程 PE 聚合邻居特征fornodeinlocal_nodes:forneighborinnode.neighbors:# 邻居在远程 PE 上——直接读它的特征不需要对方配合remote_peneighbor.owner_pe remote_offsetneighbor.local_idx neighbor_featshm.float_get(remote_features[remote_offset],peremote_pe)aggregated[node]neighbor_feat# 用 fence 确保所有远程读取完成shm.fence()同样的逻辑用 hccl 实现需要先通过 AllGather 把所有节点的特征聚到每张卡上再做本地聚合——对于稀疏图来说AllGather 传了大量不需要的冗余数据。shmem 的按需读取避免了这种冗余。和 hixl 的区别两者都做单边通信但编程模型不同维度hixlshmem编程模型显式连接 内存窗口全局地址空间内存管理手动注册 远程地址获取shmem_malloc自动全局可见同步hixl 内部处理fence/quiet/barrier 显式控制适用粒度大块传输MB 级任意粒度元素级到 MB 级原子操作不支持支持典型场景PD 分离 KV Cache 传输GNN / 不规则计算 / 全局数据结构hixl 更适合大规模、单向、低频的传输——设计目标就是 PD 分离。shmem 更适合细粒度、双向、高频的随机访问——设计目标是通用分布式计算。三种通信模型的本质区别在于谁拥有数据的主动权。消息传递hccl双方都要参与。单边推送hixl发送方单方面推数据。全局地址空间shmem任何一方随时读写任何位置。主动权越分散编程灵活性越高但一致性管理的复杂度也越高。shmem 的 fence/quiet/barrier 体系就是为这个一致性管理而设计的——用显式同步换来了完全自由的访问模式。