前言

单机8卡的推理集群里，模型参数需要在所有卡之间共享——每张卡运行相同的模型副本，推理时各自独立处理不同的请求。如果把模型参数在每个卡的HBM上各存一份，8张卡就需要8倍的显存——一个7B模型的FP16权重约14GB，8份就是112GB，而每张Ascend 910只有64GB HBM。shmem（Shared Memory）是昇腾CANN生态里的单机共享内存通信库，它允许多张NPU卡共享同一块Host Memory区域，通过PCIe DMA按需读取共享数据，避免在每张卡的HBM上都存一份完整模型。CANN社区在atomgit.com/cann上开源了shmem仓库，是单机多卡推理场景下节省显存的关键组件。

shmem的通信模型

shmem的通信模型基于PGAS（Partitioned 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启动次数。

import shmem
import numpy as np

# 初始化shmem运行时
# 必须在所有参与通信的进程中调用
shmem.init()

# 获取当前进程编号和总进程数
rank = shmem.my_pe()       # Process Element编号
size = shmem.n_pes()       # 总进程数

# 分配共享内存
# 每个进程分配自己负责的那部分，所有进程的部分拼接成完整的共享区域
# 为什么不是单个进程分配全部？因为PGAS模型的分配是对称的——
# 每个进程贡献相同大小的内存，拼接后形成全局地址空间
model_size = 14 * 1024 * 1024 * 1024  # 14GB模型权重
per_rank_size = model_size // size      # 每个进程负责的部分

# 在Host Memory中分配共享内存
# shmem_malloc分配的内存对所有进程可见
# 为什么用shmem_malloc而不是malloc？因为shmem_malloc分配的内存
# 会被映射到所有进程的地址空间，而malloc分配的内存只有当前进程可见
shared_weights = shmem.shmem_malloc(per_rank_size)

# 主控进程（rank=0）加载模型权重到共享内存
if rank == 0:
    # 加载完整模型权重
    weights = np.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搬运
for layer_idx in range(32):
    # 计算当前层权重在共享内存中的偏移
    offset = layer_idx * layer_size
    # 从共享内存读取到Device Memory
    # 为什么按层读取而不是一次性读取全部？因为NPU的HBM有限，
    # 不可能同时存放所有层的权重，按层读取可以配合推理流水线
    shmem.shmem_getmem(
        dest=device_buffer,      # Device Memory目标地址
        source=shared_weights + offset,  # 共享内存源地址
        size=layer_size,         # 当前层权重大小
        pe=0                     # 从rank 0的共享内存读取
    )
    # 在Device上执行当前层的推理
    run_layer(layer_idx, device_buffer)

权重共享的显存节省分析

以LLaMA-7B的8卡推理为例，分析shmem权重共享的显存节省：

不使用shmem（每卡独立加载模型）：

• 模型权重：14GB/卡 * 8卡 = 112GB总显存
• KV Cache：2GB/卡 * 8卡 = 16GB总显存
• 运行时开销：1GB/卡 * 8卡 = 8GB总显存
• 总计：136GB，每卡需要17GB（8卡可用64GB*8=512GB，利用率27%）

使用shmem（权重共享在Host Memory）：

• 模型权重：14GB共享在Host Memory（不占HBM）
• KV Cache：2GB/卡 * 8卡 = 16GB总显存
• 运行时开销：1GB/卡 * 8卡 = 8GB总显存
• 当前推理层的权重缓冲：0.5GB/卡（只缓存当前层的权重）
• 总计：24GB HBM + 14GB Host Memory，每卡只需3GB HBM

显存占用从每卡17GB降到3GB，节省82%。节省出的14GB/卡显存可以用来增加batch size或KV Cache长度——原来batch=4的配置，用shmem后可以跑batch=24，单卡吞吐提升6倍。

但shmem的代价是推理延迟增加——每层推理前需要通过PCIe DMA从Host Memory读取权重。PCIe Gen4 x16的带宽约32GB/s，一层权重约450MB，读取延迟约14ms。32层推理的权重读取总延迟约32 * 14 = 448ms——这对于延迟敏感的在线推理不可接受。

流水线预取优化

shmem通过流水线预取来掩盖权重读取延迟。思路是：在当前层推理的同时，异步预取下一层的权重。

# 流水线预取示例
import shmem

# 创建两个Stream：一个用于推理，一个用于预取
compute_stream = acl.rt.create_stream()
prefetch_stream = acl.rt.create_stream()

# 预取第一层权重
prefetch_layer(0, prefetch_stream)

for layer_idx in range(32):
    # 等待当前层权重预取完成
    acl.rt.synchronize_stream(prefetch_stream)

    # 启动当前层推理
    run_layer_async(layer_idx, compute_stream)

    # 在推理执行的同时，预取下一层权重
    # 为什么能和推理并行？因为推理和PCIe DMA走不同的硬件路径，
    # 推理用AI Core，DMA用PCIe控制器，两者互不干扰
    if layer_idx < 31:
        prefetch_layer(layer_idx + 1, prefetch_stream)

# 等待最后一层推理完成
acl.rt.synchronize_stream(compute_stream)

流水线预取的效果取决于权重读取时间和推理时间的比值。如果推理时间大于权重读取时间（计算密集型），预取可以完全掩盖读取延迟。LLaMA-7B的一层推理约5ms（batch=1），权重读取约14ms——推理时间小于读取时间，预取不能完全掩盖。但对于batch=8的推理，一层推理约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都支持多卡通信，但机制和适用场景完全不同：

对比维度	HCCL	shmem
通信模式	双边（Send/Recv）	单边（Put/Get）
底层机制	HCCS/NIC网络	Host共享内存 + PCIe DMA
通信范围	单机和多机	仅单机
典型场景	梯度同步（训练）	权重共享（推理）
带宽	HCCS: 392GB/s, RoCE: 200Gbps	PCIe: 32GB/s
延迟	约5微秒	约10微秒
对端参与	需要	不需要

HCCL的带宽远高于shmem（HCCS 392GB/s vs PCIe 32GB/s），因为HCCL走NPU之间的直连链路（HCCS），shmem走Host Memory中转（PCIe）。但shmem的优势是单边操作——不需要对端配合，适合推理场景中主控进程主动推送模型参数的模式。

使用前后效率对比

以LLaMA-7B 8卡推理为例，对比三种部署方案的资源利用和性能：

对比维度	每卡独立加载	shmem权重共享	shmem + 预取
每卡显存占用	17GB	3GB	3GB
最大batch size	4	24	24
单卡吞吐（tokens/s）	42	40（延迟增加）	42（预取掩盖）
总吞吐（8卡）	336	320	336
Host Memory占用	0	14GB	14GB
模型加载时间	14s * 8 = 112s	14s * 1 = 14s	14s * 1 = 14s

shmem的核心收益是显存节省（82%）和模型加载加速（8倍）。预取优化后吞吐量和独立加载方案持平，但每卡可以跑6倍大的batch——在延迟不敏感的离线推理场景下，总吞吐可以提升6倍。

结尾

shmem通过Host Memory共享和PCIe DMA单边通信，实现了单机多卡推理的权重共享，每卡显存节省82%。流水线预取可以部分掩盖权重读取延迟，在batch较大时实现零额外延迟。shmem最适合单机多卡、模型权重较大、需要节省显存来增大batch的推理场景。理解shmem的PGAS通信模型和预取优化机制，有助于在部署时权衡显存占用和推理延迟，选择最优的权重管理策略。

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