【Runtime推理引擎②】内存池设计：Arena、Size-Class与显存碎片治理

本文是「Runtime推理引擎」系列的第二篇，带你深入推理引擎的"内存层"。如果说上一篇讲的是 Runtime 的"骨架"，那这一篇讲的就是 Runtime 的"血液系统"——显存怎么分配、怎么复用、怎么避免碎片化。

一、开篇引入：为什么内存管理是Runtime最容易被低估的部分？

很多团队做推理引擎，一开始精力全放在 kernel 优化上：这个算子融合一下，那个算法换个更高效的实现……折腾了俩月，端到端延迟只降了 5%。

然后有个老工程师花了一周把内存池重写了一下，端到端延迟直接降了 30%，并发能力翻了一倍。

这不是段子，这是很多推理团队真实发生过的事情。

内存管理是 Runtime 里最容易被低估、但对整体性能影响最大的模块之一。为什么？因为在大模型推理中：

• 显存是最稀缺的资源，往往是并发能力的天花板
• 内存分配/释放的开销，在高频调用场景下会被放大到惊人的程度
• 显存碎片会悄无声息地"吃掉"你的可用显存，让你以为显存不够用，其实是被浪费了

打个比方：kernel 优化像是给赛车换更好的发动机，而内存优化像是修一条更平整的赛道。发动机再强，坑坑洼洼的路也跑不快。

二、为什么不能热路径cudaMalloc/cudaFree？

刚接触 GPU 编程的同学最自然的想法是：需要显存就cudaMalloc，用完了就cudaFree，天经地义。为什么还要搞什么内存池？

我们从三个维度来看这个问题。

2.1 性能问题：慢，非常慢

cudaMalloc和cudaFree不是廉价操作。一次cudaMalloc的耗时通常在微秒到几十微秒级别，听起来好像还行？

但算一笔账：假设一个大模型有 80 层，每层推理要分配 10 个中间 Tensor，那一步推理就要 800 次分配。800 × 10μs = 8ms。而大模型一步推理的总耗时可能也就 20~30ms。内存分配就吃掉了 30% 以上的时间！

更要命的是，cudaMalloc是同步操作——它会阻塞 CPU 线程，甚至可能导致 GPU 流水线断掉。在 GPU 上，流水线断一次的代价可能比分配本身还大。

2.2 碎片问题：显存明明够，就是分配不出来

假设你有 10GB 空闲显存，但它们被分成了 1000 个 10MB 的"碎块"。这时候你想分配一个 100MB 的连续显存，会怎么样？

答案是：分配失败，尽管你的"总空闲显存"有 10GB。

这就是外部碎片（External Fragmentation）——空闲内存总量够，但没有一块足够大的连续空间。

频繁地cudaMalloc/cudaFree，而且每次分配的大小都不一样，显存会像被老鼠啃过的奶酪一样，千疮百孔。跑的时间越长，碎片越严重，到最后可能一半显存都被碎片"吃掉"了。

2.3 并发问题：锁竞争

cudaMalloc内部有全局锁。当你有多个线程同时分配显存时，它们会排队等锁。在高并发场景下，这个锁可能成为整个系统的瓶颈。

有人做过测试：8 个线程同时高频调用cudaMalloc，单个分配的平均延迟会比单线程时高 5~10 倍，因为大部分时间都在等锁。

2.4 一句话总结

在推理的热路径（也就是每一步推理都会执行的代码路径）上，绝对不能直接调用cudaMalloc/cudaFree。内存池就是为了解决这个问题而生的——提前申请一大块显存，然后自己管理分配和释放，把开销从"微秒级"降到"纳秒级"。

三、Arena Allocator：最简单的内存池

Arena（竞技场）是最朴素、也是最容易理解的内存分配器。它的思想可以用一句话概括：申请一大块，然后指针往后挪。

3.1 基本思想与代码

classArenaAllocator{public:ArenaAllocator(size_t total_size){cudaMalloc(&base_ptr_,total_size);current_=base_ptr_;total_size_=total_size;}void*alloc(size_t size){// 对齐到 256 字节（GPU 内存通常要求对齐）size_t aligned_size=(size+255)&~255;if(current_+aligned_size>base_ptr_+total_size_){returnnullptr;// 空间不够了}void*ptr=current_;current_+=aligned_size;returnptr;}// Arena 不支持单个释放，只能整体重置voidreset(){current_=base_ptr_;}~ArenaAllocator(){cudaFree(base_ptr_);}private:char*base_ptr_;char*current_;size_t total_size_;};

使用方式：

ArenaAllocatorarena(1024*1024*1024);// 预先申请 1GB// 推理过程中，所有小分配都走 arenavoid*a=arena.alloc(1024);// 1KBvoid*b=arena.alloc(2048);// 2KBvoid*c=arena.alloc(4096);// 4KB// 一轮推理结束，整体重置arena.reset();// 只移动一下指针，几乎零开销

3.2 优缺点

3.3 适用边界

Arena 适合什么场景？所有分配对象的生命周期相同，并且可以一次性全部释放的场景。

比如：

• 单轮推理的 activation 内存：一轮推理开始前 reset，推理结束后再 reset，中间所有 activation 都从 Arena 里分配
• 算子的 workspace 临时内存：算子执行前分配，执行完就可以整体释放

但如果你的分配对象生命周期参差不齐——有的活 1 秒，有的活 1 分钟——Arena 就完全不适用了，因为它没法单独释放某个对象。

四、Size-Class Free List：工业级内存池的标配

Arena 虽然快，但太"笨"了。实际场景中，我们需要既能快速分配，又能单独释放的内存池。这时候就轮到Size-Class Free List出场了。

4.1 基本原理

核心思想就两句话：

1. 按大小分级：把内存分配请求按大小分成若干个等级（size class），比如 64B、128B、256B、512B……每个等级维护一个空闲链表（free list）
2. 取了再用，还了再收：分配时从对应大小的空闲链表里取一块；释放时放回对应的空闲链表

为什么要分级？因为如果所有大小的释放块都扔到一个链表里，分配时要找一块"刚好够大"的，就得遍历整个链表，效率很低。分级之后，每个链表里的块大小都一样，拿第一个就行，O(1) 搞定。

4.2 简化版代码

classSizeClassAllocator{public:void*alloc(size_t size){// 1. 找到对应的 size classintcls=get_size_class(size);size_t block_size=class_to_size(cls);// 2. 如果这个 class 的 free list 不空，直接取一个if(!free_lists_[cls].empty()){void*ptr=free_lists_[cls].back();free_lists_[cls].pop_back();returnptr;}// 3. free list 空了，向操作系统（GPU驱动）申请一批returnrefill(cls,block_size);}voidfree(void*ptr,size_t size){intcls=get_size_class(size);free_lists_[cls].push_back(ptr);// 直接放回对应链表}private:// size class 表：64, 96, 128, 192, 256, 384, 512, 768, 1024, ...staticconstexprsize_t kClasses[]={64,96,128,192,256,384,512,768,1024,1536,2048,3072,4096,6144,8192,12288,16384,24576,32768,49152,65536,98304,131072};std::vector<void*>free_lists_[sizeof(kClasses)/sizeof(size_t)];void*refill(intcls,size_t block_size){// 一次申请多个块，减少系统调用次数constintnum_blocks=16;char*big_block;cudaMalloc(&big_block,block_size*num_blocks);// 把大快切成小块，放入 free listfor(inti=1;i<num_blocks;i++){free_lists_[cls].push_back(big_block+i*block_size);}returnbig_block;// 第一块直接返回}};

4.3 Size Class 怎么定？

这是一个非常关键的问题：size class 设哪些值？设密了，内部碎片多；设疏了，free list 命中率低。

工业界常用的几种策略：

• 2的幂：实现最简单（用位运算就能算），但碎片率最高
• 1.5倍：兼顾碎片率和实现复杂度，很多内存分配器（如 tcmalloc）采用
• 更密的分级：碎片率低，但 free list 数量多，管理开销大

对于大模型推理，推荐用1.5倍或更密的分级。因为大模型里 Tensor 大小通常比较大（几 MB 到几十 MB），50% 的内部碎片是绝对不可接受的——100MB 的 Tensor 浪费 50MB，80 层就是 4GB 显存没了。

小白提示：内部碎片和外部碎片是两个不同的概念。内部碎片是"分配的比实际需要的大，多出来的部分浪费了"；外部碎片是"空闲的内存都是碎块，凑不出一块大的"。内存池设计本质上就是在这两种碎片之间做权衡。

五、分池策略：四池分离，各管各的

有了 Size-Class Free List，是不是一个全局内存池就够了？对于大模型推理来说，还不够。

大模型推理中有几种完全不同"性格"的内存，把它们放在同一个池子里管理，效果往往不好。工业级推理引擎通常采用分池策略——不同类型的内存用不同的池来管理。

5.1 四池分离架构

为什么要分开？因为它们的生命周期模式和分配特征差异太大了：

• 把短命的 activation 和长命的 KV Cache 放一起：短命的频繁分配释放，会在长命的之间留下大量空洞，外部碎片会非常严重
• 把 workspace 和 activation 放一起：workspace 大小波动大（有的算子只要几 KB，有的要几百 MB），会加剧碎片
• 把 weight 和其他放一起：weight 占了显存的大头，而且从不释放，它本身就是最大的"钉子户"，会把空间分割成上下两块，其他分配只能在缝隙里折腾

5.2 分池的额外好处

除了减少碎片，分池还有几个额外的好处：

1. 精细化统计：你可以清楚地知道"权重占了多少、KV Cache 占了多少、激活值占了多少"，方便做容量规划和问题排查
2. 差异化策略：每个池可以用最适合它的分配算法，而不是"一把钥匙开所有锁"
3. 隔离性：activation 池就算碎片严重，也不会影响到 KV 池，更不会把权重池搞崩
4. 水位控制：你可以给每个池设上限——比如 KV 池最多用 10GB，超过了就拒绝新请求，避免把系统拖死

六、Zero Copy技术：能不拷贝就不拷贝

讲完分配和释放，我们再聊一个和内存相关但经常被忽略的优化方向：Zero Copy（零拷贝）。

6.1 什么是 Zero Copy？

字面意思：零次拷贝。就是说，数据从产生到消费，中间不经过额外的内存拷贝。

举个例子：A 算子的输出要作为 B 算子的输入。最朴素的做法是：

A 算子把结果写到 output_buf → 把 output_buf 拷贝到 B 的 input_buf → B 算子从 input_buf 读

而 Zero Copy 的做法是：

A 算子直接把结果写到 B 的输入位置 → B 算子直接从那里读

少了一次拷贝，省时间也省内存。

6.2 常见的 Zero Copy 场景

场景一：算子融合（Operator Fusion）

把多个小算子融合成一个大算子，中间结果直接放在寄存器或共享内存里，不需要写回全局显存。这是最常见、也是效果最显著的 Zero Copy。

// 不融合：两次写显存，两次读显存add(x,y,tmp);// tmp = x + y，写显存gelu(tmp,out);// 读 tmp，写 out// 融合：只写一次显存add_gelu(x,y,out);// 中间结果在寄存器里，直接算完写 out

场景二：In-place 操作

有些算子的输入和输出大小一样，而且计算模式允许"边读边写"，那就可以直接在输入内存上覆盖写入，省掉一份输出内存。

比如 ReLU、Element-wise Add（当其中一个输入可以被覆盖时）等。

// 非 in-place：需要额外的输出内存voidrelu(constfloat*input,float*output,intn);// in-place：直接在输入上修改，省一半内存voidrelu_inplace(float*data,intn);

场景三：指针转发（Pointer Forwarding）

在计算图层面，如果某个节点的输出就是下一个节点的输入，而且中间没有其他消费者，可以让下一个节点直接"继承"这个指针，不重新分配。

6.3 Zero Copy 的边界

Zero Copy 虽好，但不是万能的。以下情况要格外小心：

• 有多个消费者：如果一个 Tensor 被多个节点使用，就不能 in-place 改，否则第一个消费者改完，第二个消费者读到的就是错的
• 有数据依赖重叠：比如卷积操作，输出的每个元素依赖输入的一个窗口，in-place 写会覆盖还没读到的数据
• 需要保留中间结果：比如做调试、做 profiling 的时候，中间结果还是要留着的

七、显存碎片治理：指标与治理手段

显存碎片是 Runtime 运维中最让人头疼的问题之一。它不像 crash 那样立刻暴露，而是像慢性病一样，慢慢蚕食你的显存，直到某一天你发现"同样的模型，上周还能跑 8 个并发，这周怎么只能跑 5 个了"。

7.1 怎么衡量碎片？

首先要有量化指标，不然你都不知道碎片严不严重。常用的两个指标：

指标一：碎片率（Fragmentation Ratio）

碎片率 = 1 - (最大连续空闲块 / 总空闲显存)

• 碎片率 = 0：所有空闲显存都是一整块，完美
• 碎片率 = 0.5：一半空闲显存是"散的"，凑不成大块
• 碎片率越接近 1，情况越糟

指标二：分配失败率

分配失败率 = 分配失败的次数 / 总分配次数

这个指标更"实用"——如果你的总空闲显存明明够，但还是经常分配失败，那碎片就很严重了。

7.2 治理手段

碎片治理是一个"组合拳"，没有银弹，但以下手段都很有效：

对于大模型推理来说，分池 + Size-Class + Coalescing是最常用的组合。内存整理（Defrag）虽然效果好，但 GPU 上做数据搬运的成本很高，而且需要暂停推理，一般只在低峰期或者运维窗口做。

八、生产内存池接口建议与小结

8.1 生产级内存池应该长什么样？

最后，给大家一个生产级内存池的接口设计参考。一个好用的内存池，接口应该简洁而完整：

classMemoryPool{public:// ===== 核心接口 =====void*alloc(size_t size,MemType type=MemType::Activation);voidfree(void*ptr,size_t size,MemType type=MemType::Activation);// ===== 统计接口（运维必备） =====size_ttotal_allocated()const;// 已分配的总显存size_ttotal_free()const;// 空闲的总显存size_tlargest_free_block()const;// 最大连续空闲块doublefragmentation_ratio()const;// 碎片率size_talloc_count()const;// 累计分配次数size_tfail_count()const;// 累计失败次数// ===== 运维接口 =====voidrelease_free_memory();// 把空闲显存还给系统voiddefrag();// 碎片整理（低峰期调用）std::stringprint_stats()const;// 打印详细统计信息// ===== 配置接口 =====voidset_pool_limit(MemType type,size_t limit);// 给某个池设上限voidset_alignment(size_t alignment);// 设置对齐粒度};

很多人写内存池只写了alloc和free，上线之后才发现出了问题根本没法排查——你连碎片率是多少、哪个池在涨都不知道。统计和运维接口是生产级和玩具级的最大区别。

8.2 小结

今天我们系统地梳理了推理 Runtime 中内存池的设计：

• 为什么不能直接用 cudaMalloc：性能差、碎片多、并发有锁，热路径绝对不能碰
• Arena Allocator：最简单的内存池，分配极快但只能整体释放，适合生命周期一致的场景
• Size-Class Free List：工业级标配，按大小分级管理空闲块，分配释放都是 O(1)
• 分池策略：权重、KV Cache、Activation、Workspace 四池分离，各用各的分配器，从架构上减少碎片
• Zero Copy：算子融合、in-place 操作、指针转发，能省则省
• 显存碎片治理：先量化指标，再用分池、coalescing、defrag 等手段组合拳治理
• 生产级内存池：不仅要有 alloc/free，更要有统计和运维接口

8.3 思考问题

1. 如果你的推理服务需要支持动态 batch（batch size 随时变化），Activation Pool 用 Arena 还合适吗？为什么？如果不合适，你会怎么改？
2. KV Cache 的大小和序列长度正相关。如果用户的输入长度差异很大（有的 128 token，有的 4096 token），KV Pool 的碎片问题会特别严重。你有什么好的策略来缓解？
3. 内存池的设计中，"分配速度"和"内存利用率"往往是矛盾的。你觉得在大模型推理场景下，哪个更重要？为什么？

系列回顾：

• 第①篇：Runtime架构设计：从Session到Graph Executor的完整骨架
• 第②篇：本篇，内存池设计

下一篇我们会聊Kernel 调度与算子融合，敬请关注！