昇腾 NPU 的显存管理和 CUDA 不太一样。虽然从用户接口上看都是acl.rt.malloc/acl.rt.free这套东西但底层的分配策略、碎片管理、跨卡共享有一些独特的坑。搞清楚这些底层逻辑能帮你写出更高效的推理代码。显存布局概览昇腾 910 单卡有 32GB HBM 显存。这 32GB 不是全部给用户用的驱动会预留一部分给系统用固件、驱动日志、临时缓冲区。实际能给用户分配的大概在 30GB 左右。总显存: 32GB ├── 系统预留: ~1GB (固件、驱动元数据) ├── 用户可用: ~30GB │ ├── 权重区: 模型参数占用 │ ├── 激活区: 中间计算结果 │ └── 临时区: 算子临时缓冲驱动在分配显存的时候按 2MB 对齐这是硬件层面的对齐要求。所以即使用户申请了 1 字节驱动也会分配 2MB。这和 CUDA 的显存分配有区别CUDA 默认是 4KB 页对齐昇腾的粒度粗很多。分配接口昇腾的显存分配有两套接口高层次的 PyTorch 接口最常用importtorch# 分配一个 tensor设备是 npuatorch.randn(1024,1024,devicenpu)# 底层调的是 ACL 的分配接口# 显式分配固定大小的显存btorch.empty([8192,8192],dtypetorch.float16,devicenpu)# 8192*8192*2 bytes 128MB低层次的 ACL 接口当需要更精细控制时importacl acl.rt.set_device(0)# 分配 1GB 显存ptr,retacl.rt.malloc(1024*1024*1024,acl.CT_MEM_MALLOC_HUGE_FIRST)ifret!0:raiseRuntimeError(f分配失败:{ret})# 释放acl.rt.free(ptr)acl.CT_MEM_MALLOC_HUGE_FIRST是分配策略告诉驱动优先从大块空闲区分配。对于大 tensor这个策略能减少碎片。碎片化问题昇腾的显存分配器用的是伙伴系统的变体分配和释放的效率很高但长时间运行的推理服务容易遇到碎片问题。碎片化的典型场景# 假设我们跑一个动态 shape 的推理服务# 每次请求的 batch size 和 seq len 都可能不同batch_sizes[1,2,4,8,16,32]seq_lens[128,256,512,1024,2048]# 模拟一批请求分配不同大小的 tensorimporttorchforbsinbatch_sizes:forslinseq_lens:# 每个请求分配一个中间 tensor# shape: [batch, seq_len, hidden]hidden4096xtorch.randn(bs,sl,hidden,devicenpu)# 模拟推理计算...resultx.sum(dim1)# 用完释放delxdelresult# 运行一段时间后显存可能变成这样# [occupied: 1GB][free: 500MB][occupied: 2GB][free: 300MB][occupied: 1.5GB]# 总空闲: 800MB但最大连续块: 500MB# 申请 600MB 的 tensor 会失败但总空闲够解决这个问题有几个思路1. 预分配固定大小的内存池# 预分配一个内存池大小覆盖最常用场景MEMORY_POOL_SIZE4*1024*1024*1024# 4GBpool_ptr,retacl.rt.malloc(MEMORY_POOL_SIZE,acl.CT_MEM_MALLOC_HUGE_FIRST)# 自己管理这块内存避免碎片classMemoryPool:def__init__(self,ptr,size):self.ptrptr self.sizesize self.free_list[(0,size)]# (offset, size) 列表defallocate(self,size):# 第一次适配找第一个够大的块fori,(offset,avail)inenumerate(self.free_list):ifavailsize:# 分割self.free_list.pop(i)remainingavail-sizeifremaining0:self.free_list.insert(i,(offsetsize,remaining))returnptroffsetraiseRuntimeError(内存池空间不足)deffree(self,offset,size):# 归还后尝试合并相邻块self.free_list.append((offset,size))self.free_list.sort()# 合并逻辑省略遍历列表合并相邻块内存池的好处是把碎片控制在池子内部池子的大小固定后续分配释放都在这个范围内。2. 重置设备如果内存池方案太复杂一个取巧的办法是定期重置设备importacl# 当碎片化严重到影响分配时# 重新初始化设备清空所有显存分配acl.rt.reset_device(0)# 之前分配的指针全部失效# 重置之后需要重新初始化推理模型modelload_model()# 继续跑但这个方法有代价重置期间服务不可用而且之前分配的所有资源都要重建。显存的跨卡共享多卡训练的时候有时候需要在卡之间共享数据比如共享权重。这时候要用到昇腾的跨卡内存复制接口。importacl# 假设 rank 0 上有一个权重 tensorweighttorch.randn(1024,1024,devicenpu:0)# rank 1 要访问这个权重# 方式一HCCL broadcast通信方式需要网络传输importtorch.distributedasdist dist.broadcast(weight,src0)# 方式二P2P 直接访问内存映射不走通信# 把 rank 0 的显存地址传给 rank 1weight_ptrweight.data_ptr()# 虚拟地址# rank 1 通过 P2P 映射这个地址# 注意昇腾的 P2P 访问要求物理地址连续# 如果 rank 0 的显存碎片化严重可能无法做连续映射P2P 映射在某些场景下比 HCCL broadcast 快因为不走 NCCL/HCCL 的通信栈直接用硬件的 P2P 访问通路。但前提是两个卡都在同一个 PCIe switch 下并且驱动支持 P2P。内存泄漏检测和 CPU 内存一样GPU/NPU 显存也有泄漏问题。排查方法importgcdefcheck_memory_leak():importsubprocess# 记录初始显存使用beforesubprocess.check_output(npu-smi query -t memory -i 0 | grep Memory Used,shellTrue).decode()# 跑一段代码循环分配释放for_inrange(100):xtorch.randn(4096,4096,devicenpu)yx x.t()delxdely gc.collect()# 显式 GC# 记录最终显存使用aftersubprocess.check_output(npu-smi query -t memory -i 0 | grep Memory Used,shellTrue).decode()print(f前:{before})print(f后:{after})# 如果数字明显增加说明有泄漏check_memory_leak()泄漏的常见原因Python 对象引用没有释放del之后还有其他引用持有ACL 分配的内存没有调用对应的acl.rt.freePyTorch 分配走的是 torch不会有这个问题但自己用 ACL 接口写算子的时候要小心循环引用导致 GC 无法回收仓库在 https://atomgit.com/cann/runtimeCANN Runtime 的内存管理源码在里面。