CANN-NPU 显存回收策略：内存碎片整理与显存池化机制实战

一、显存碎片从哪来1.1 碎片的两种形态外部碎片——总空闲内存够用但不连续。比如有 4 块 128MB 空闲但需要一块 512MB 的连续内存分配失败。内部碎片——分配器按固定大小的块分配实际使用的比分配的小。比如分配 400KB分配器给 512KB浪费了 112KB。1.2 为什么 NPU 显存碎片更严重批量推理不同 batch size 需要不同大小的 buffer反复分配释放动态 Shape每次推理的中间 tensor 大小不同无法复用多模型多个模型交替执行显存需求波动大以一个典型推理服务为例时间线: 0s: 分配模型A权重 200MB → 空闲 600MB (连续) 1s: 分配激活缓冲 100MB → 空闲 500MB (连续) 2s: 释放缓冲 → 空闲 500MB (连续) 3s: 分配模型B权重 300MB → 空闲 200MB (连续) 4s: 分配激活缓冲 150MB → 空闲 50MB 5s: 释放模型B权重 → 空闲 350MB (碎片: 150MB 200MB) 6s: 分配 250MB → 失败! (碎片无法满足)二、池化分配器设计2.1 核心思想预分配一大块显存自己管理分配和释放避免频繁调用系统 API。同时维护不同大小的空闲块列表减少碎片。importthreadingfromcollectionsimportdefaultdictclassNPUAllocator:NPU 显存池化分配器 设计要点: 1. 预分配: 启动时一次性申请整块显存 2. 大小类: 按 2 的幂分级管理减少内部碎片 3. 空闲链表: 每个大小类维护空闲块链表O(1) 分配 4. 线程安全: 推理服务多线程调用需要加锁 为什么用 2 的幂分级? - 分配 400KB → 给 512KB 块内部碎片率约 22% - 如果用 1KB 粒度 → 碎片率 0.5%但管理开销大 - 2 的幂是折中碎片率可控管理高效 def__init__(self,total_size8*1024*1024*1024):# 默认 8GBself.total_sizetotal_size self.used_size0self.lockthreading.Lock()# 按大小类管理空闲块# key 块大小 (2 的幂), value 空闲块列表self.free_blocksdefaultdict(list)# 地址 → 大小 映射用于释放时查找self.block_map{}# 预分配的整块内存模拟self._poolbytearray(total_size)self._base_addrid(self._pool)# 初始化: 把整块内存加入空闲列表self._add_free_block(self._base_addr,total_size)def_next_power_of_2(self,size):将大小向上取到 2 的幂power1whilepowersize:power1returnpowerdef_add_free_block(self,addr,size):self.free_blocks[size].append(addr)defallocate(self,size):分配内存 策略: 1. 找到 size 的最小 2 的幂大小类 2. 从该类的空闲链表取一块 3. 如果没有尝试拆分更大的块 4. 如果都没有触发碎片整理 withself.lock:aligned_sizeself._next_power_of_2(size)# 尝试从对应大小类分配blockself._try_allocate(aligned_size)ifblockisnotNone:addr,block_sizeblock self.block_map[addr]block_size self.used_sizeblock_sizereturnaddr# 尝试从更大的块拆分blockself._split_allocate(aligned_size)ifblockisnotNone:addr,block_sizeblock self.block_map[addr]block_size self.used_sizeblock_sizereturnaddr# 内存不足raiseMemoryError(f显存分配失败: 需要{size}bytes, f已用{self.used_size}/{self.total_size})def_try_allocate(self,aligned_size):从指定大小类尝试分配ifself.free_blocks[aligned_size]:addrself.free_blocks[aligned_size].pop()returnaddr,aligned_sizereturnNonedef_split_allocate(self,aligned_size):从更大的空闲块拆分forblock_sizeinsorted(self.free_blocks.keys()):ifblock_sizealigned_sizeandself.free_blocks[block_size]:# 拆分: 一大变两小addrself.free_blocks[block_size].pop()remaining_sizeblock_size-aligned_size# 剩余部分放回空闲列表self._add_free_block(addraligned_size,remaining_size)returnaddr,aligned_sizereturnNonedeffree(self,addr):释放内存 关键操作: 尝试合并相邻的空闲块减少外部碎片。 withself.lock:ifaddrnotinself.block_map:returnblock_sizeself.block_map.pop(addr)self.used_size-block_size# 尝试合并相邻空闲块merged_addr,merged_sizeself._try_merge(addr,block_size)self._add_free_block(merged_addr,merged_size)def_try_merge(self,addr,size):尝试与相邻空闲块合并 这是减少外部碎片的核心。 每次释放时检查左右两侧是否有空闲块有则合并。 merged_addraddr merged_sizesize# 检查右侧相邻块right_addraddrsizeforblock_size,addrsinlist(self.free_blocks.items()):forainaddrs:ifaright_addr:merged_sizeblock_size addrs.remove(a)break# 检查左侧相邻块left_addraddr-merged_sizeforblock_size,addrsinlist(self.free_blocks.items()):forainaddrs:ifablock_sizeleft_addr:merged_addra merged_sizeblock_size addrs.remove(a)breakreturnmerged_addr,merged_sizedefstats(self):返回显存使用统计withself.lock:free_sizeself.total_size-self.used_sizereturn{total:self.total_size,used:self.used_size,free:free_size,usage_rate:f{self.used_size/self.total_size*100:.1f}%,fragment_count:len(self.block_map),}2.2 使用示例# 初始化 8GB 显存池allocatorNPUAllocator(total_size8*1024*1024*1024)# 分配模型权重weight_bufallocator.allocate(200*1024*1024)# 200MBinput_bufallocator.allocate(100*1024*1024)# 100MB# 使用后释放allocator.free(input_buf)# 查看统计print(allocator.stats())# {total: 8589934592, used: 209715200, free: 8380219392, ...}三、碎片整理策略3.1 紧凑式碎片整理把所有已分配的块移到一端空闲空间连成一片。需要暂停所有推理任务。defcompact(self):紧凑式碎片整理 步骤: 1. 暂停推理阻止新分配 2. 计算所有已分配块的新地址连续排列 3. 搬移数据到新地址 4. 更新地址映射 5. 恢复推理 代价: - 暂停时间与已分配块数量成正比 - 数据搬移消耗带宽 - 适合在低峰期执行 withself.lock:# 收集所有已分配块allocatedsorted(self.block_map.items(),keylambdax:x[0])# 计算新地址new_addrself._base_addr addr_mapping{}# 旧地址 → 新地址forold_addr,block_sizeinallocated:addr_mapping[old_addr]new_addr new_addrblock_size# 搬移数据模拟forold_addr,new_addrinaddr_mapping.items():# 实际中需要用 memcpy 或 DMA 搬移pass# 更新映射self.block_map.clear()forold_addr,new_addrinaddr_mapping.items():self.block_map[new_addr]self.block_map.get(old_addr,0)# 重置空闲列表self.free_blocks.clear()free_startnew_addr free_sizeself._base_addrself.total_size-free_startiffree_size0:self._add_free_block(free_start,free_size)3.2 增量式碎片整理不需要暂停每次释放时主动合并。这是日常策略。defincremental_compact(self):增量式碎片整理 在每次 free() 时自动执行不需要暂停。 通过合并相邻空闲块逐步减少碎片。 这是日常策略大部分场景够用。 # 已经在 free() 的 _try_merge 中实现了pass3.3 策略对比策略适用场景暂停时间碎片消除效果增量合并日常运行无中等紧凑整理低峰期高最佳重新分配模型切换中等好四、针对推理场景的优化4.1 预分配策略推理服务启动时就预估好需要多少显存一次性分配到位。defpreallocate_for_inference(model_configs):根据模型配置预分配显存 推理服务的显存需求相对可预测: - 模型权重: 固定大小启动时确定 - KV Cache: 最大序列长度 × batch size × 层数 - 激活缓冲: 最大输入 shape × 中间层最大宽度 allocatorNPUAllocator()forconfiginmodel_configs:# 预分配权重缓冲weight_bufallocator.allocate(config[weight_size])# 预分配 KV Cache按最大序列长度kv_cache_size(config[max_seq_len]*config[max_batch]*config[num_layers]*config[hidden_dim]*2# K 和 V)kv_bufallocator.allocate(kv_cache_size)returnallocator4.2 多级缓存把显存分成几个区域优先用小的、快的区域。classMultiLevelCache:多级显存缓存 L1: 寄存器文件 — 最快容量最小几十 KB L2: SRAM 缓存 — 快容量中等几百 KB L3: HBM 显存 — 慢容量大GB 级 策略: 频繁访问的数据放 L1/L2不频繁的放 L3。 def__init__(self):self.l1{}# 寄存器级缓存self.l2{}# SRAM 级缓存self.l3{}# HBM 级缓存defget(self,key):ifkeyinself.l1:returnself.l1[key]ifkeyinself.l2:valself.l2.pop(key)self.l1[key]valreturnvalifkeyinself.l3:valself.l3.pop(key)self.l2[key]valreturnvalreturnNonedefput(self,key,value,levell3):iflevell1:self.l1[key]valueeliflevell2:self.l2[key]valueelse:self.l3[key]value五、监控与诊断classMemoryMonitor:显存使用监控 持续跟踪显存使用情况发现异常及时告警。 def__init__(self,allocator,threshold0.9):self.allocatorallocator self.thresholdthreshold self.history[]defcheck(self):statsself.allocator.stats()self.history.append(stats)usage_ratestats[used]/stats[total]ifusage_rateself.threshold:print(f[WARN] 显存使用率{usage_rate:.1%}超过阈值{self.threshold:.0%})returnstatsdefreport(self):生成显存使用报告ifnotself.history:return无数据latestself.history[-1]peak_usedmax(h[used]forhinself.history)reportf 显存使用报告: 总量:{latest[total]/1024**3:.1f}GB 当前使用:{latest[used]/1024**3:.1f}GB ({latest[usage_rate]}) 峰值使用:{peak_used/1024**3:.1f}GB ({peak_used/latest[total]:.1%}) 碎片数:{latest[fragment_count]}returnreport六、常见问题问题原因解决方案OOM 但总空闲够用外部碎片严重执行紧凑式碎片整理分配速度变慢空闲链表太长调整大小类粒度减少链表长度显存泄漏忘记释放 buffer用 MemoryMonitor 监控定期检查推理延迟抖动碎片整理暂停了推理用增量式整理避免暂停相关仓库CANN- 昇腾计算架构 https://gitee.com/ascend/cannjemalloc- 高性能内存分配器 https://github.com/jemalloc/jemalloctcmalloc- Google 线程级内存分配器 https://github.com/google/tcmallocPoolAllocator- 虚拟显存池化分配器 https://github.com/vllm-project/vllm