想象你是一家餐厅的主厨。最忙的午市,你不可能每做一道菜就开车去批发市场买一次食材——那样蒜蓉粉丝虾还没下锅,这道菜的成本已经翻了三倍。真正高效的后厨,会在清晨一次性备好一整天要用的食材,分门别类放进分格收纳柜;炒菜时伸手即取,用完洗净放回原格。AI 编译器里的内存池,做的正是这件事。深度学习推理时,每个算子都会临时“购买”一块内存来存放中间张量,用完再“扔掉”。如果每次采购都走一遍系统调用(malloc/free),几百个算子叠加起来,推理延迟就悄悄膨胀了。内存池的思路是:初始化时向系统一次性“批发”一大块内存,之后所有张量都在池内“取用—归还—复用”,把昂贵的系统调用从“每菜一次”降为“全天一次”。本文就用这条“后厨备菜”的线索,带你理解内存池的预分配、分块复用、碎片管理,以及它如何与计算图协作,让推理又快又省内存。
掌握内存池设计,是 AI 推理引擎与系统优化工程师优化延迟与内存占用的核心技能,直接决定模型上线后的吞吐能力与运行成本。
目录
可视化演示
展示哈希表链地址法与开放寻址法的冲突解决过程,拖动滑块调节负载因子观察冲突率变化。
本节配有交互动效(进入视口后自动加载,约 3KB)。
链地址法 相同哈希值的元素以链表形式存储 h(k)=k mod m 查询时遍历链表,冲突越多链表越长
分配瓶颈:malloc 慢在哪
如果主厨每做一道菜就开车去一趟批发市场,午市两小时最多出十道菜。更糟的是,他还得反复登记“买了什么、剩了多少”。深度学习推理里的 malloc/free 正是这种“低效采购”:每个算子计算完都会产生一到多个中间张量,这些张量的内存申请与释放,动辄在全球内存分配器里排几百上千次队。
malloc 并不是表面看起来的“在内存里画个圈”那么简单。它要维护空闲链表、查找合适大小的块、在必要时发起 brk 或 mmap 系统调用切换用户态与内核态。想象你买一根葱也要过一遍“出库单、过磅、开票”的完整流程,单次开销当然不便宜。而推理过程中张量大小各异、生灭频繁,分配器还要反复处理大小匹配和锁竞争,缓存局部性也被打乱。
把一次推理里所有分配次数记为 N,总时间大致是:
正因如此,后厨必须建立自己的食材柜——这就是内存池登场的原因。主厨不必每次炒菜都出门采购,他只需要在小推车和食材柜之间完成“取、放”。
- N 表示分配总次数推理中每个算子都会产生一个或多个中间张量,N 就是这类分配次数的总和。
- t\_分配 表示单次分配时间一次 malloc 涉及用户态/内核态切换与元数据维护,远贵于普通的内存赋值。
- t\_释放 表示单次释放时间free 同样是一次系统级操作,需要更新空闲链表、可能触发合并。
- 整体相乘N 很大时就形成了“系统调用风暴”,这是推理延迟膨胀的隐蔽来源。
池化思想:提前批发,后厨自取
中央厨房的做法是:清晨派一辆采购车,把全天要用的食材一次性拉回后厨,按类别放进分格柜。此后主厨的每一次取料都只发生在自己熟悉的几平方米内,快、可预期、不堵车。内存池完全复制了这套逻辑。
内存池在初始化阶段调用一次大块分配(常称为 arena 或 superblock),向操作系统批发几十甚至几百 MB 的连续内存。这块内存随后被切成小块,装入空闲链表。此后每次“内存分配”不再是系统调用,而是从链表头部摘下一个空闲块——常数时间完成,像从抽屉里拿一个调料罐。释放时也只需把块插回链表,等待下次复用。
关键是复用的链路:张量 A 释放出来的块,立刻可以被张量 B 使用。系统调用的次数从“每个张量一次”降为“整个池生命周期一次”。总时间变成:
后厨从此再也不用在午市高峰期和批发市场的堵车较劲了。主厨只需记住:所有食材都已在分格柜里,伸手、取用、放回,循环往复。
- T\_预分配 表示一次性批发成本进程启动或编译期规划时,向系统申请大块内存,仅此一次进入内核。
- N 表示后续分配次数所有中间张量的池内分配仍然发生 N 次,但每次都走本地链表。
- O(1) 表示单次取块开销从空闲链表头部弹出或插入,时间与池内已有多少块无关。
- 整体对比当 N 很大时,池化的总耗时几乎被压平成一条水平线。
池化设计:分格收纳与多级桶池
食材柜不能只有一种格子:冬瓜要占大格,一勺盐只需小格。如果所有食材都挤进同一个规格的格子,要么大材小用,要么装不下。内存池同样面临“块的大小规格”问题。
经典的解法是多级桶池:把内存按大小分成若干等级(size class),例如 16 字节、64 字节、256 字节、4KB、1MB……每个等级对应一条独立的空闲链表。分配时先按需求取整到最近的等级,再从对应链表中取块。取整规则最常用的是 2 的幂:
池化设计还要回答“谁来管理”的问题。多线程推理时,所有线程都抢同一个池会产生锁竞争,于是派生出线程私有池——每个线程一个小池,免锁访问;跨线程的大张量则交给共享池协调。这就像后厨给每位厨师配了随身调料盒,共用的蔬菜柜再配一个总管理员。多级桶池 + 线程私有池的组合,正是 tcmalloc、jemalloc 这些工业级分配器的看家本领,AI 编译器直接继承了这套智慧。
- log2 表示需求大小的数量级33 字节取对数后约为 5.04,说明它落在 2^5 与 2^6 之间。
- ceil 表示向上取整取 6,而不是四舍五入到 5,因为 32 字节装不下 33 字节。
- 2 的幂表示实际桶大小最终归入 64 字节的等级。
- 为什么要统一取整把形状各异的张量归到固定等级后,空闲链表可以做到 O(1) 精确匹配,不必每次从头搜索。
碎片管理:格柜里的空隙
食材柜用久了,你会发现总有些格子里塞着半袋蒜、一片姜,恰好占住一个格子的 20%,剩下的空间什么也放不下。这种“大格装小鱼”的浪费叫内部碎片;而“格子之间散布着零碎空隙、连一条鱼都塞不下”的情况叫外部碎片。内存池绕不开这两种碎片。
内部碎片来自取整:33 字节的需求占了 64 字节的块,31 字节就白扔了。外部碎片来自释放顺序的不均匀:块 A、C 被使用,中间的 B 被释放,B 的空洞太小,后续大张量无法入住。张量大小各异、释放顺序受计算图决定,空洞会越积越多。可以用利用率来衡量碎片严重程度:
控制碎片的手段主要有两种:一是合并策略,当相邻块都空闲时把它们拼成一个大块,伙伴系统(buddy system)就是为此设计的;二是对齐规则,把桶大小设计成与张量形状、缓存行对齐,减少取整浪费。后厨的收拾动作也一样:每隔一段时间把半袋蒜并进同一格,把零散空隙腾出来变成可用的大格。
- S\_有效 表示张量真正需要的字节数例如一个矩阵实际只占 33 字节。
- S\_已占用 表示池中实际占用的字节数因取整可能高达 64 字节。
- 对全部块求和把所有张量的有效大小与占用大小分别累加。
- 1 减去比值差值就是白白浪费的碎片比例。利用率 80% 意味着每 5 字节有 1 字节被浪费。
生命周期:按菜单规划备菜
真正的中央厨房不会等客人点菜才思考食材量——宴席菜单一出,主厨就知道每道菜要多少料、哪些菜可以共用同一种料。AI 编译器有一个巨大的优势:整个计算图在执行前就已经完整可见,它天生就握着一份“菜单”。
这份菜单上的每一笔“菜”,就是一个张量的生存期(liveness)。编译期的静态分析可以算出每个张量从哪个算子开始存活、到哪个算子之后死亡。于是编译器发现:张量 A 在算子 3 之后就没用了,而张量 B 在算子 4 才出生,两者从未同时活着——它们可以共用同一块内存,像两班倒的厨师共用同一张操作台。
内存池的规模因此可以精确到一个理论下界:任意时刻所有存活张量的总字节数中的最大值。这个值写出来就是:
TVM、MLIR 等编译器框架正是这样工作的:先做内存规划,再生成内存池的分配与释放指令,把“动态 malloc”替换成“静态位移”。内存安全也由此受益——因为物理地址在编译期就已确定,越界访问比运行期临时拼地址更可控。这就像宴席还没开,备菜表已写好:哪道菜用哪个格子,心中有数。
- 活跃(t) 表示 t 时刻存活张量的集合由计算图的 liveness 分析得出,张量创建与销毁之间记为存活。
- size(v) 表示张量 v 的字节数由张量形状与数据类型确定。
- 对每个时刻求和把所有同时存活的张量大小加起来,得到那一刻的内存水位。
- 取最大所有时刻水位线的最高点,就是安全池容量的最低要求。低于它会溢出,高于它是浪费。
实战评估:效果如何衡量
食材柜建得好不好,不能靠主厨“感觉快了”,而是要看午市高峰期同一时间段出了多少桌菜。内存池的收益同样需要量化,衡量的指标大致有四类:峰值内存、分配耗时、系统吞吐(每秒推理次数)和尾部延迟(p99)。
分配耗时最容易观察——池化后单次分配从“系统调用”变成“链表取块”,耗时往往下降一到两个数量级。峰值内存则是双刃剑:预分配会一次性占住大块内存,但如果配合计算图的存活期分析,总容量反而能压到比无池更紧凑的水平。综合收益常用加速比表达:
需要清醒的是,内存池不是万能药。若推理中只有少数几个巨型张量、生命周期很长且几乎没有释放复用,池化收益就很微薄;而频繁的小张量分配、短生命周期、多请求并发的推理服务,收益最大。落地时还要区分编译期规划(静态固定池)与运行期管理(动态可增长池)两种形态。后厨的经验是:先看菜单再建柜——菜单越可预测,食材柜越高效。
- T\_无池 表示传统 malloc/free 的总耗时每次分配都承受系统调用。
- T\_池 表示池化后的总耗时大部分分配退化为 O(1) 链表操作。
- 两者相除得到加速比,通常能达到 2 到 10 倍以上,取决于分配频率。
AI 应用场景
在线推理服务的高并发优化
阶段:推理部署与性能优化
解决的问题:在线推理服务每请求会产生大量中间张量,传统 malloc/free 在并发高时频繁进入内核,导致 p99 延迟抖动。内存池预分配 + 复用后,分配路径缩短到用户态的常数时间,延迟曲线趋稳。
具体案例:某推荐系统使用 TVM 部署 CTR 模型,请求峰值 2000 QPS 时,每个请求约有 300 次张量分配。改用内存池后,单请求平均分配耗时从 180 微秒降到 12 微秒,p99 延迟下降 38%,CPU 占用降低 15%。
训练迭代中的显存调度
阶段:训练图编译与显存复用
解决的问题:训练每一步迭代都会创建大量临时激活张量,并在反向传播后释放。若每次都向显存分配器申请,会拖慢迭代速度并产生碎片。PyTorch 的 caching allocator 正是按“池化”思路:预申请大块显存、按大小分级缓存、释放后不归还给驱动而是留池复用。
具体案例:在英伟达 A100 上训练 GPT-2 时,开启 caching allocator 后,训练吞吐从每秒 420 样本提升到 610 样本,显存峰值反而下降约 20%,因为释放的块被后续层立即复用。
端侧编译期的静态内存规划
阶段:模型编译与端侧部署
解决的问题:手机、MCU 等端侧设备内存紧张且不支持动态分配,运行时 malloc 既慢又不可控。AI 编译器在编译期对计算图做 liveness 分析,规划一个固定大小的全局内存池,把所有中间张量的地址静态编码进指令。
具体案例:某视觉模型部署到 2MB RAM 的嵌入式芯片,原先运行时动态分配导致反复内存不足;编译器用内存池静态规划后,内存占用固定在 1.2MB,推理首帧延迟从 2.8 秒降到 0.9 秒。
Python 演示
import random def system_cost(size_bytes): # 每次系统级内存请求的固定开销 + 与大小相关的开销 return 0.32 + size_bytes * 1e-6 class NaiveAllocator: # 无池版本: 每次分配/释放都触发一次系统级开销 def __init__(self): self.requests = 0 self.total_cost = 0.0 def alloc(self, size): self.requests += 1 self.total_cost += system_cost(size) def free(self, size): self.requests += 1 self.total_cost += system_cost(size) class PoolAllocator: # 池化版本: 只有大块不足时才向系统申请一次 def __init__(self, chunk_size=1 << 20): self.chunk_size = chunk_size self.requests = 0 self.total_cost = 0.0 self.avail = 0 def alloc(self, size): if self.avail < size: self.requests += 1 self.total_cost += system_cost(self.chunk_size) self.avail += self.chunk_size self.avail -= size # 池内取出是 O(1) 廉价操作 def free(self, size): self.avail += size # 归还池中, 等待复用 def run_workload(allocator, times=5000): random.seed(42) active = [] for _ in range(times): if not active or random.random() < 0.55: size = random.randint(256, 16384) allocator.alloc(size) active.append(size) else: victim = active.pop() allocator.free(victim) if __name__ == "__main__": naive = NaiveAllocator() pool = PoolAllocator() run_workload(naive) run_workload(pool) print("naive requests:", naive.requests, "cost:", round(naive.total_cost, 2)) print("pool requests:", pool.requests, "cost:", round(pool.total_cost, 2)) print("speedup:", round(naive.total_cost / pool.total_cost, 2), "x")
def system_cost(size_bytes):模拟一次系统级内存请求的成本:固定开销超大块,大数据量只占零头。self.requests += 1记录真正进入系统的次数,池化后这个数字会急剧下降。self.chunk_size = chunk_size池子每次批发 1MB,相当于厨房采购车的装载量。if self.avail < size:只有池内余量不足时才去“批发”,这是池化省钱的开关。self.avail -= size从池内划出空间,O(1) 完成,不经过系统。self.avail += size张量释放后空间归还池中,等待下一个张量复用。pool = PoolAllocator()初始化时并不立即批发,真正需要时才触发预分配。print("speedup:", ...)输出加速比,直观看到池化把开销压低了几个数量级。
常见误区
误区我以为内存池跟 Redis、CPU 缓存差不多,都是把常用数据存起来,下次直接读,减少重复计算。
正解两者管的根本不是一回事。缓存管的是“数据内容”的复用——热点数据二次访问时命中;内存池管的是“内存块”的复用——张量死掉后它的地址空间被新张量继承。内存池不关心块里存了什么,只关心“这块地皮能不能快速过户”。
为何常见:两者都叫“预先准备、避免重复获取”,外形相似,但一个是数据层的概念,一个是资源层的概念,初学者很容易混淆。
误区既然内存池能减少系统调用,那所有模型推理都应该无脑套上内存池,收益必然为正。
正解如果推理过程中只有几个大张量、生命周期贯穿始终、几乎没有“分配—释放—再分配”的交替,池化几乎没有复用机会,反而多了一次预分配的搬家和池管理的开销。池化的收益前提是“频繁分配 + 短生命周期 + 可复用”。
为何常见:教程里展示的加速数据都来自分配密集场景,读者容易把“场景优势”误当成“普适结论”。
误区我觉得内存池既然自己管理内存,分配释放都有记录,碎片问题就自然而然消失了。
正解池化只是把碎片问题从“系统分配器”搬到了“池内部”,内部碎片(取整浪费)和外部碎片(空洞)依然存在,甚至因为块大小固定而更明显。碎片必须靠分级桶、合并策略、对齐规则主动治理。
为何常见:看到“复用”就以为所有空间都被利用,忽略了取整和释放顺序这两大碎片来源。
误区我觉得池里的块都在池子掌控中,就算忘了 free,内存也还在池里,不算泄漏,不用管。
正解池内确实不会把内存还给系统,但若张量地址被错误持有、块被重复归还、或者全局池在进程结束时未销毁,照样会造成“池内逻辑泄漏”。它比普通泄漏更难察觉,因为从系统看内存占用没有变化,实际可用块却在不断减少。
为何常见:“系统内存没涨”让人误以为安全,可池内的可用块被白白占住,最终表现为推理中途分配失败,排查时非常隐蔽。
误区我想一个池子管所有张量最简单,分成十几个桶反而增加代码复杂度,不如一个大池。
正解单一的大池会让所有分配都在同一把锁下排队,多线程并发时锁竞争严重,且大小悬殊的块交替分配会迅速制造碎片。多级桶池 + 线程私有池的复杂度,换来的是免锁分配和碎片可控,这正是工业级分配器的共同选择。
为何常见:初学阶段追求“代码简单”,容易把工程上的取舍理解为越简单越好,忽略了并发与碎片这两个隐藏成本。
面试题
Q1. 内存池降低分配开销的核心做法是什么?
- A.提高模型层数精度
- B.预分配并复用内存块← 正确答案
- C.压缩输入数据大小
- D.增加并发线程数量
答案:B。内存池的本质是“一次批发、反复使用”:初始化时向系统预分配大块内存,之后所有张量在池内取块、归还、复用,避免每次都触发系统调用。A、C、D 分别涉及模型结构、数据压缩和并行度,都不是内存池要解决的问题。
Q2. 内部碎片产生的主要原因是什么?
- A.分配的块向上取整← 正确答案
- B.空闲区域过于零散
- C.释放顺序产生混乱
- D.请求系统调用过多
答案:A。内部碎片是“大格装小鱼”——需求 33 字节却被归入 64 字节桶,多出的 31 字节被浪费。B 描述的是外部碎片(空洞),C 是产生空洞的诱因而非内部碎片本身,D 是性能问题而非碎片问题。
Q3. 需求 33 字节的内存,2 的幂桶池最可能分配多少?
- A.按需求原样分配
- B.向上取整到 48 字节
- C.向上取整到 64 字节← 正确答案
- D.向上取整到 128 字节
答案:C。2 的幂取整规则是 2 的 ceil(log2(S)) 次方。log2(33) 约 5.04,向上取整得 6,2^6=64,故选 C。A 违背了“分级取整”的池化设计,B 不是 2 的幂,D 超过了必要等级。
Q4. 内存池为什么能降低推理延迟?
- A.减少计算图节点总数
- B.减少系统调用与切换开销← 正确答案
- C.提高模型参数量大小
- D.压缩激活张量的数值
答案:B。池化后分配从系统调用退化为常数时间的链表操作,省去用户态/内核态切换。A、C、D 分别涉及图结构、模型大小和数值精度,与内存分配路径无关。
Q5. 多线程推理场景下,最宜采用哪种池策略?
- A.线程私有的小型池← 正确答案
- B.单一大共享池更优
- C.直接关闭内存池机制
- D.只分配内存不释放
答案:A。线程私有池让每个线程访问自己的池,免去锁竞争,是并发场景的首选;共享池只用于跨线程的大张量协调。B 会在高并发下发生严重锁竞争,C 回到系统调用,D 会造成逻辑泄漏。
Q6. 编译期内存池规划主要依据什么?
- A.训练数据的规模大小
- B.显卡品牌与驱动型号
- C.张量的生存期分析← 正确答案
- D.模型参数量的多少
答案:C。编译器对计算图做 liveness 分析,得知每个张量从哪个算子存活到哪个算子,才能让不重叠的张量共享内存、算出峰值水位。训练数据、显卡型号和参数量都不决定中间张量的生命周期。
Q7. 外部碎片的主要危害是什么?
- A.空闲块过小无法利用← 正确答案
- B.已分配块内部空间浪费
- C.导致模型推理精度下降
- D.造成计算图规模膨胀
答案:A。外部碎片是分散的小空洞,单独看都不够大,拼接起来又不能被连续分配利用,导致池内“有空间却用不上”。B 是内部碎片的危害,C、D 与碎片无关。
Q8. 哪些场景最适合使用内存池?
- A.频繁的小张量分配场景← 正确答案
- B.张量生存期极短的场景← 正确答案
- C.推理多请求并发场景← 正确答案
- D.一次性巨型张量场景
答案:ABC。池化的收益来自“分配—释放—复用”的高频循环:频繁小张量、短生命周期、多并发都能放大复用收益。D 是一次性大块、几乎不复用,池化收益微薄。
Q9. 控制内存池碎片的手段包括哪些?
- A.空闲块的合并策略← 正确答案
- B.按大小分级的桶池← 正确答案
- C.统一对齐取整规则← 正确答案
- D.尽可能扩大取整粒度
答案:ABC。伙伴系统等合并策略可把相邻空闲块拼成大块;多级桶池让大小相近的块共用链表,减少交错分配;统一对齐规则让块尺寸规整。D 反向操作——扩大取整粒度会加剧内部碎片。
Q10. 使用内存池可能带来哪些隐患?
- A.池内泄漏难定位排查← 正确答案
- B.并发访问需要同步← 正确答案
- C.指针悬挂或误用风险← 正确答案
- D.分配开销反而增加
答案:ABC。池内块被误持有会造成逻辑泄漏;共享池多线程访问需要加锁;块归还后原指针再被使用就是悬挂引用。D 是误解——分配开销会大幅下降,否则内存池就失去了意义。