前言单机8卡的推理集群里模型参数需要在所有卡之间共享——每张卡运行相同的模型副本推理时各自独立处理不同的请求。如果把模型参数在每个卡的HBM上各存一份8张卡就需要8倍的显存——一个7B模型的FP16权重约14GB8份就是112GB而每张Ascend 910只有64GB HBM。shmemShared Memory是昇腾CANN生态里的单机共享内存通信库它允许多张NPU卡共享同一块Host Memory区域通过PCIe DMA按需读取共享数据避免在每张卡的HBM上都存一份完整模型。CANN社区在atomgit.com/cann上开源了shmem仓库是单机多卡推理场景下节省显存的关键组件。shmem的通信模型shmem的通信模型基于PGASPartitioned Global Address Space——所有参与通信的进程共享一个全局地址空间每个进程可以独立访问全局地址空间中的任意位置不需要对方的参与。PGAS模型和传统的消息传递模型MPI Send/Recv有本质区别消息传递模型中数据交换需要双方配合——发送方调用Send接收方调用Recv数据从发送方的内存拷贝到接收方的内存。如果接收方没有及时调用Recv发送方会阻塞等待。PGAS模型中数据交换是单边操作——发起方直接读写远端内存远端不需要任何配合。shmem提供的Put和Get操作就是单边的shmem_put——把本地数据写入远端的共享内存区域。远端进程不需要调用任何接收函数数据直接出现在远端内存中。shmem_get——从远端的共享内存区域读取数据到本地。远端进程不需要调用任何发送函数数据直接从远端内存拷贝到本地。这种单边通信模型非常适合推理场景——主控进程把模型参数写入共享内存每张NPU卡按需读取自己需要的部分不需要和其他卡同步。如果模型参数更新了比如动态加载新模型主控进程直接覆盖共享内存中的数据所有卡下次读取时自动获得最新值。shmem在昇腾NPU上的实现shmem在昇腾NPU上的实现依赖两个底层机制Host侧的共享内存和Device侧的PCIe DMA。Host侧共享内存。shmem在Host Memory中分配一块共享内存区域映射到所有参与通信的进程的虚拟地址空间。这块内存使用大页Huge Page分配——2MB或1GB的大页减少TLB Miss。共享内存的创建和映射通过POSIX shm_open mmap实现跨进程可见。Device侧PCIe DMA。每张NPU卡通过PCIe DMA直接读取Host Memory中的共享数据。昇腾NPU的PCIe控制器支持Scatter-Gather DMA——可以把共享内存中不连续的数据段比如模型的不同层权重聚合成一次DMA传输减少DMA启动次数。importshmemimportnumpyasnp# 初始化shmem运行时# 必须在所有参与通信的进程中调用shmem.init()# 获取当前进程编号和总进程数rankshmem.my_pe()# Process Element编号sizeshmem.n_pes()# 总进程数# 分配共享内存# 每个进程分配自己负责的那部分所有进程的部分拼接成完整的共享区域# 为什么不是单个进程分配全部因为PGAS模型的分配是对称的——# 每个进程贡献相同大小的内存拼接后形成全局地址空间model_size14*1024*1024*1024# 14GB模型权重per_rank_sizemodel_size//size# 每个进程负责的部分# 在Host Memory中分配共享内存# shmem_malloc分配的内存对所有进程可见# 为什么用shmem_malloc而不是malloc因为shmem_malloc分配的内存# 会被映射到所有进程的地址空间而malloc分配的内存只有当前进程可见shared_weightsshmem.shmem_malloc(per_rank_size)# 主控进程rank0加载模型权重到共享内存ifrank0:# 加载完整模型权重weightsnp.load(model_weights.npy)# 14GB# 把权重写入共享内存# 为什么rank 0写入全部因为只有rank 0加载了完整的权重文件# 其他rank只负责自己那部分共享内存shmem.shmem_putmem(shared_weights,weights.ctypes.data,model_size,0)# 同步等待所有数据写入完成# 为什么需要同步因为shmem_put是异步的写入操作可能还没完成# shmem_barrier确保所有进程都到达同步点后才继续shmem.shmem_barrier_all()# 每个NPU卡从共享内存读取自己需要的权重层# 使用PCIe DMA从Host Memory读取到Device Memory# 为什么用Get而不是直接访问因为共享内存在Host Memory# NPU不能直接访问Host Memory需要通过PCIe DMA搬运forlayer_idxinrange(32):# 计算当前层权重在共享内存中的偏移offsetlayer_idx*layer_size# 从共享内存读取到Device Memory# 为什么按层读取而不是一次性读取全部因为NPU的HBM有限# 不可能同时存放所有层的权重按层读取可以配合推理流水线shmem.shmem_getmem(destdevice_buffer,# Device Memory目标地址sourceshared_weightsoffset,# 共享内存源地址sizelayer_size,# 当前层权重大小pe0# 从rank 0的共享内存读取)# 在Device上执行当前层的推理run_layer(layer_idx,device_buffer)权重共享的显存节省分析以LLaMA-7B的8卡推理为例分析shmem权重共享的显存节省不使用shmem每卡独立加载模型模型权重14GB/卡 * 8卡 112GB总显存KV Cache2GB/卡 * 8卡 16GB总显存运行时开销1GB/卡 * 8卡 8GB总显存总计136GB每卡需要17GB8卡可用64GB*8512GB利用率27%使用shmem权重共享在Host Memory模型权重14GB共享在Host Memory不占HBMKV Cache2GB/卡 * 8卡 16GB总显存运行时开销1GB/卡 * 8卡 8GB总显存当前推理层的权重缓冲0.5GB/卡只缓存当前层的权重总计24GB HBM 14GB Host Memory每卡只需3GB HBM显存占用从每卡17GB降到3GB节省82%。节省出的14GB/卡显存可以用来增加batch size或KV Cache长度——原来batch4的配置用shmem后可以跑batch24单卡吞吐提升6倍。但shmem的代价是推理延迟增加——每层推理前需要通过PCIe DMA从Host Memory读取权重。PCIe Gen4 x16的带宽约32GB/s一层权重约450MB读取延迟约14ms。32层推理的权重读取总延迟约32 * 14 448ms——这对于延迟敏感的在线推理不可接受。流水线预取优化shmem通过流水线预取来掩盖权重读取延迟。思路是在当前层推理的同时异步预取下一层的权重。# 流水线预取示例importshmem# 创建两个Stream一个用于推理一个用于预取compute_streamacl.rt.create_stream()prefetch_streamacl.rt.create_stream()# 预取第一层权重prefetch_layer(0,prefetch_stream)forlayer_idxinrange(32):# 等待当前层权重预取完成acl.rt.synchronize_stream(prefetch_stream)# 启动当前层推理run_layer_async(layer_idx,compute_stream)# 在推理执行的同时预取下一层权重# 为什么能和推理并行因为推理和PCIe DMA走不同的硬件路径# 推理用AI CoreDMA用PCIe控制器两者互不干扰iflayer_idx31:prefetch_layer(layer_idx1,prefetch_stream)# 等待最后一层推理完成acl.rt.synchronize_stream(compute_stream)流水线预取的效果取决于权重读取时间和推理时间的比值。如果推理时间大于权重读取时间计算密集型预取可以完全掩盖读取延迟。LLaMA-7B的一层推理约5msbatch1权重读取约14ms——推理时间小于读取时间预取不能完全掩盖。但对于batch8的推理一层推理约20ms已经大于14ms的读取时间预取可以完全掩盖。实际效果不预取时32层推理延迟约32 * (14 5) 608ms预取后约32 * 20 640ms计算时间主导…不对预取后延迟应该是max(计算, 读取) * 32 20 * 32 640ms受最慢步骤约束而非608ms。但对比每卡独立加载模型的延迟32 * 5 160ms权重已在HBM中shmem方案即使有预取也慢4倍——这是节省显存的代价。shmem和HCCL的对比shmem和HCCL都支持多卡通信但机制和适用场景完全不同对比维度HCCLshmem通信模式双边Send/Recv单边Put/Get底层机制HCCS/NIC网络Host共享内存 PCIe DMA通信范围单机和多机仅单机典型场景梯度同步训练权重共享推理带宽HCCS: 392GB/s, RoCE: 200GbpsPCIe: 32GB/s延迟约5微秒约10微秒对端参与需要不需要HCCL的带宽远高于shmemHCCS 392GB/s vs PCIe 32GB/s因为HCCL走NPU之间的直连链路HCCSshmem走Host Memory中转PCIe。但shmem的优势是单边操作——不需要对端配合适合推理场景中主控进程主动推送模型参数的模式。使用前后效率对比以LLaMA-7B 8卡推理为例对比三种部署方案的资源利用和性能对比维度每卡独立加载shmem权重共享shmem 预取每卡显存占用17GB3GB3GB最大batch size42424单卡吞吐tokens/s4240延迟增加42预取掩盖总吞吐8卡336320336Host Memory占用014GB14GB模型加载时间14s * 8 112s14s * 1 14s14s * 1 14sshmem的核心收益是显存节省82%和模型加载加速8倍。预取优化后吞吐量和独立加载方案持平但每卡可以跑6倍大的batch——在延迟不敏感的离线推理场景下总吞吐可以提升6倍。结尾shmem通过Host Memory共享和PCIe DMA单边通信实现了单机多卡推理的权重共享每卡显存节省82%。流水线预取可以部分掩盖权重读取延迟在batch较大时实现零额外延迟。shmem最适合单机多卡、模型权重较大、需要节省显存来增大batch的推理场景。理解shmem的PGAS通信模型和预取优化机制有助于在部署时权衡显存占用和推理延迟选择最优的权重管理策略。仓库地址https://atomgit.com/cann/shmem
昇腾CANN共享内存通信库shmem：单机多卡数据共享的深度优化实践
发布时间：2026/6/8 13:34:17
前言单机8卡的推理集群里模型参数需要在所有卡之间共享——每张卡运行相同的模型副本推理时各自独立处理不同的请求。如果把模型参数在每个卡的HBM上各存一份8张卡就需要8倍的显存——一个7B模型的FP16权重约14GB8份就是112GB而每张Ascend 910只有64GB HBM。shmemShared Memory是昇腾CANN生态里的单机共享内存通信库它允许多张NPU卡共享同一块Host Memory区域通过PCIe DMA按需读取共享数据避免在每张卡的HBM上都存一份完整模型。CANN社区在atomgit.com/cann上开源了shmem仓库是单机多卡推理场景下节省显存的关键组件。shmem的通信模型shmem的通信模型基于PGASPartitioned Global Address Space——所有参与通信的进程共享一个全局地址空间每个进程可以独立访问全局地址空间中的任意位置不需要对方的参与。PGAS模型和传统的消息传递模型MPI Send/Recv有本质区别消息传递模型中数据交换需要双方配合——发送方调用Send接收方调用Recv数据从发送方的内存拷贝到接收方的内存。如果接收方没有及时调用Recv发送方会阻塞等待。PGAS模型中数据交换是单边操作——发起方直接读写远端内存远端不需要任何配合。shmem提供的Put和Get操作就是单边的shmem_put——把本地数据写入远端的共享内存区域。远端进程不需要调用任何接收函数数据直接出现在远端内存中。shmem_get——从远端的共享内存区域读取数据到本地。远端进程不需要调用任何发送函数数据直接从远端内存拷贝到本地。这种单边通信模型非常适合推理场景——主控进程把模型参数写入共享内存每张NPU卡按需读取自己需要的部分不需要和其他卡同步。如果模型参数更新了比如动态加载新模型主控进程直接覆盖共享内存中的数据所有卡下次读取时自动获得最新值。shmem在昇腾NPU上的实现shmem在昇腾NPU上的实现依赖两个底层机制Host侧的共享内存和Device侧的PCIe DMA。Host侧共享内存。shmem在Host Memory中分配一块共享内存区域映射到所有参与通信的进程的虚拟地址空间。这块内存使用大页Huge Page分配——2MB或1GB的大页减少TLB Miss。共享内存的创建和映射通过POSIX shm_open mmap实现跨进程可见。Device侧PCIe DMA。每张NPU卡通过PCIe DMA直接读取Host Memory中的共享数据。昇腾NPU的PCIe控制器支持Scatter-Gather DMA——可以把共享内存中不连续的数据段比如模型的不同层权重聚合成一次DMA传输减少DMA启动次数。importshmemimportnumpyasnp# 初始化shmem运行时# 必须在所有参与通信的进程中调用shmem.init()# 获取当前进程编号和总进程数rankshmem.my_pe()# Process Element编号sizeshmem.n_pes()# 总进程数# 分配共享内存# 每个进程分配自己负责的那部分所有进程的部分拼接成完整的共享区域# 为什么不是单个进程分配全部因为PGAS模型的分配是对称的——# 每个进程贡献相同大小的内存拼接后形成全局地址空间model_size14*1024*1024*1024# 14GB模型权重per_rank_sizemodel_size//size# 每个进程负责的部分# 在Host Memory中分配共享内存# shmem_malloc分配的内存对所有进程可见# 为什么用shmem_malloc而不是malloc因为shmem_malloc分配的内存# 会被映射到所有进程的地址空间而malloc分配的内存只有当前进程可见shared_weightsshmem.shmem_malloc(per_rank_size)# 主控进程rank0加载模型权重到共享内存ifrank0:# 加载完整模型权重weightsnp.load(model_weights.npy)# 14GB# 把权重写入共享内存# 为什么rank 0写入全部因为只有rank 0加载了完整的权重文件# 其他rank只负责自己那部分共享内存shmem.shmem_putmem(shared_weights,weights.ctypes.data,model_size,0)# 同步等待所有数据写入完成# 为什么需要同步因为shmem_put是异步的写入操作可能还没完成# shmem_barrier确保所有进程都到达同步点后才继续shmem.shmem_barrier_all()# 每个NPU卡从共享内存读取自己需要的权重层# 使用PCIe DMA从Host Memory读取到Device Memory# 为什么用Get而不是直接访问因为共享内存在Host Memory# NPU不能直接访问Host Memory需要通过PCIe DMA搬运forlayer_idxinrange(32):# 计算当前层权重在共享内存中的偏移offsetlayer_idx*layer_size# 从共享内存读取到Device Memory# 为什么按层读取而不是一次性读取全部因为NPU的HBM有限# 不可能同时存放所有层的权重按层读取可以配合推理流水线shmem.shmem_getmem(destdevice_buffer,# Device Memory目标地址sourceshared_weightsoffset,# 共享内存源地址sizelayer_size,# 当前层权重大小pe0# 从rank 0的共享内存读取)# 在Device上执行当前层的推理run_layer(layer_idx,device_buffer)权重共享的显存节省分析以LLaMA-7B的8卡推理为例分析shmem权重共享的显存节省不使用shmem每卡独立加载模型模型权重14GB/卡 * 8卡 112GB总显存KV Cache2GB/卡 * 8卡 16GB总显存运行时开销1GB/卡 * 8卡 8GB总显存总计136GB每卡需要17GB8卡可用64GB*8512GB利用率27%使用shmem权重共享在Host Memory模型权重14GB共享在Host Memory不占HBMKV Cache2GB/卡 * 8卡 16GB总显存运行时开销1GB/卡 * 8卡 8GB总显存当前推理层的权重缓冲0.5GB/卡只缓存当前层的权重总计24GB HBM 14GB Host Memory每卡只需3GB HBM显存占用从每卡17GB降到3GB节省82%。节省出的14GB/卡显存可以用来增加batch size或KV Cache长度——原来batch4的配置用shmem后可以跑batch24单卡吞吐提升6倍。但shmem的代价是推理延迟增加——每层推理前需要通过PCIe DMA从Host Memory读取权重。PCIe Gen4 x16的带宽约32GB/s一层权重约450MB读取延迟约14ms。32层推理的权重读取总延迟约32 * 14 448ms——这对于延迟敏感的在线推理不可接受。流水线预取优化shmem通过流水线预取来掩盖权重读取延迟。思路是在当前层推理的同时异步预取下一层的权重。# 流水线预取示例importshmem# 创建两个Stream一个用于推理一个用于预取compute_streamacl.rt.create_stream()prefetch_streamacl.rt.create_stream()# 预取第一层权重prefetch_layer(0,prefetch_stream)forlayer_idxinrange(32):# 等待当前层权重预取完成acl.rt.synchronize_stream(prefetch_stream)# 启动当前层推理run_layer_async(layer_idx,compute_stream)# 在推理执行的同时预取下一层权重# 为什么能和推理并行因为推理和PCIe DMA走不同的硬件路径# 推理用AI CoreDMA用PCIe控制器两者互不干扰iflayer_idx31:prefetch_layer(layer_idx1,prefetch_stream)# 等待最后一层推理完成acl.rt.synchronize_stream(compute_stream)流水线预取的效果取决于权重读取时间和推理时间的比值。如果推理时间大于权重读取时间计算密集型预取可以完全掩盖读取延迟。LLaMA-7B的一层推理约5msbatch1权重读取约14ms——推理时间小于读取时间预取不能完全掩盖。但对于batch8的推理一层推理约20ms已经大于14ms的读取时间预取可以完全掩盖。实际效果不预取时32层推理延迟约32 * (14 5) 608ms预取后约32 * 20 640ms计算时间主导…不对预取后延迟应该是max(计算, 读取) * 32 20 * 32 640ms受最慢步骤约束而非608ms。但对比每卡独立加载模型的延迟32 * 5 160ms权重已在HBM中shmem方案即使有预取也慢4倍——这是节省显存的代价。shmem和HCCL的对比shmem和HCCL都支持多卡通信但机制和适用场景完全不同对比维度HCCLshmem通信模式双边Send/Recv单边Put/Get底层机制HCCS/NIC网络Host共享内存 PCIe DMA通信范围单机和多机仅单机典型场景梯度同步训练权重共享推理带宽HCCS: 392GB/s, RoCE: 200GbpsPCIe: 32GB/s延迟约5微秒约10微秒对端参与需要不需要HCCL的带宽远高于shmemHCCS 392GB/s vs PCIe 32GB/s因为HCCL走NPU之间的直连链路HCCSshmem走Host Memory中转PCIe。但shmem的优势是单边操作——不需要对端配合适合推理场景中主控进程主动推送模型参数的模式。使用前后效率对比以LLaMA-7B 8卡推理为例对比三种部署方案的资源利用和性能对比维度每卡独立加载shmem权重共享shmem 预取每卡显存占用17GB3GB3GB最大batch size42424单卡吞吐tokens/s4240延迟增加42预取掩盖总吞吐8卡336320336Host Memory占用014GB14GB模型加载时间14s * 8 112s14s * 1 14s14s * 1 14sshmem的核心收益是显存节省82%和模型加载加速8倍。预取优化后吞吐量和独立加载方案持平但每卡可以跑6倍大的batch——在延迟不敏感的离线推理场景下总吞吐可以提升6倍。结尾shmem通过Host Memory共享和PCIe DMA单边通信实现了单机多卡推理的权重共享每卡显存节省82%。流水线预取可以部分掩盖权重读取延迟在batch较大时实现零额外延迟。shmem最适合单机多卡、模型权重较大、需要节省显存来增大batch的推理场景。理解shmem的PGAS通信模型和预取优化机制有助于在部署时权衡显存占用和推理延迟选择最优的权重管理策略。仓库地址https://atomgit.com/cann/shmem
