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

【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(size255)~255;if(current_aligned_sizebase_ptr_total_size_){returnnullptr;// 空间不够了}void*ptrcurrent_;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*aarena.alloc(1024);// 1KBvoid*barena.alloc(2048);// 2KBvoid*carena.alloc(4096);// 4KB// 一轮推理结束整体重置arena.reset();// 只移动一下指针几乎零开销3.2 优缺点优点缺点分配极快O(1)只是指针加法不支持单个释放只能整体重置实现简单几乎不会有 bug生命周期必须严格一致同生共死完全没有碎片问题空间利用率取决于峰值使用量不需要加锁单线程使用时多线程并发需要加锁3.3 适用边界Arena 适合什么场景所有分配对象的生命周期相同并且可以一次性全部释放的场景。比如单轮推理的 activation 内存一轮推理开始前 reset推理结束后再 reset中间所有 activation 都从 Arena 里分配算子的 workspace 临时内存算子执行前分配执行完就可以整体释放但如果你的分配对象生命周期参差不齐——有的活 1 秒有的活 1 分钟——Arena 就完全不适用了因为它没法单独释放某个对象。四、Size-Class Free List工业级内存池的标配Arena 虽然快但太笨了。实际场景中我们需要既能快速分配又能单独释放的内存池。这时候就轮到Size-Class Free List出场了。4.1 基本原理核心思想就两句话按大小分级把内存分配请求按大小分成若干个等级size class比如 64B、128B、256B、512B……每个等级维护一个空闲链表free list取了再用还了再收分配时从对应大小的空闲链表里取一块释放时放回对应的空闲链表为什么要分级因为如果所有大小的释放块都扔到一个链表里分配时要找一块刚好够大的就得遍历整个链表效率很低。分级之后每个链表里的块大小都一样拿第一个就行O(1) 搞定。4.2 简化版代码classSizeClassAllocator{public:void*alloc(size_t size){// 1. 找到对应的 size classintclsget_size_class(size);size_t block_sizeclass_to_size(cls);// 2. 如果这个 class 的 free list 不空直接取一个if(!free_lists_[cls].empty()){void*ptrfree_lists_[cls].back();free_lists_[cls].pop_back();returnptr;}// 3. free list 空了向操作系统GPU驱动申请一批returnrefill(cls,block_size);}voidfree(void*ptr,size_t size){intclsget_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::vectorvoid*free_lists_[sizeof(kClasses)/sizeof(size_t)];void*refill(intcls,size_t block_size){// 一次申请多个块减少系统调用次数constintnum_blocks16;char*big_block;cudaMalloc(big_block,block_size*num_blocks);// 把大快切成小块放入 free listfor(inti1;inum_blocks;i){free_lists_[cls].push_back(big_blocki*block_size);}returnbig_block;// 第一块直接返回}};4.3 Size Class 怎么定这是一个非常关键的问题size class 设哪些值设密了内部碎片多设疏了free list 命中率低。工业界常用的几种策略策略说明典型比例内部碎片率2的幂64, 128, 256, 512…2x最高 50%1.5倍增长64, 96, 144, 216…1.5x最高 33%1.25倍增长64, 80, 100, 125…1.25x最高 20%2的幂实现最简单用位运算就能算但碎片率最高1.5倍兼顾碎片率和实现复杂度很多内存分配器如 tcmalloc采用更密的分级碎片率低但 free list 数量多管理开销大对于大模型推理推荐用1.5倍或更密的分级。因为大模型里 Tensor 大小通常比较大几 MB 到几十 MB50% 的内部碎片是绝对不可接受的——100MB 的 Tensor 浪费 50MB80 层就是 4GB 显存没了。小白提示内部碎片和外部碎片是两个不同的概念。内部碎片是分配的比实际需要的大多出来的部分浪费了外部碎片是空闲的内存都是碎块凑不出一块大的。内存池设计本质上就是在这两种碎片之间做权衡。五、分池策略四池分离各管各的有了 Size-Class Free List是不是一个全局内存池就够了对于大模型推理来说还不够。大模型推理中有几种完全不同性格的内存把它们放在同一个池子里管理效果往往不好。工业级推理引擎通常采用分池策略——不同类型的内存用不同的池来管理。5.1 四池分离架构池子存放内容生命周期分配模式推荐分配器Weight Pool权重池模型权重weights极长进程级启动时一次性分配基本不释放裸 cudaMalloc不需要池KV PoolKV缓存池自回归推理的 KV Cache长Session级几秒到几分钟Session 创建时分配销毁时释放Size-Class 内存池Activation Pool激活池推理中间激活值极短单步推理级毫秒级逐层分配、逐层释放Arena 或 Bump AllocatorWorkspace Pool工作池算子临时工作空间极短算子级微秒级算子开始时分配结束时释放固定大小的 Ring Buffer 或 Arena为什么要分开因为它们的生命周期模式和分配特征差异太大了把短命的 activation 和长命的 KV Cache 放一起短命的频繁分配释放会在长命的之间留下大量空洞外部碎片会非常严重把 workspace 和 activation 放一起workspace 大小波动大有的算子只要几 KB有的要几百 MB会加剧碎片把 weight 和其他放一起weight 占了显存的大头而且从不释放它本身就是最大的钉子户会把空间分割成上下两块其他分配只能在缝隙里折腾5.2 分池的额外好处除了减少碎片分池还有几个额外的好处精细化统计你可以清楚地知道权重占了多少、KV Cache 占了多少、激活值占了多少方便做容量规划和问题排查差异化策略每个池可以用最适合它的分配算法而不是一把钥匙开所有锁隔离性activation 池就算碎片严重也不会影响到 KV 池更不会把权重池搞崩水位控制你可以给每个池设上限——比如 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 对齐分配大小向上取整到 size class释放后能精确匹配通用场景内存整理Defrag把数据挪一挪把碎块拼成大块低峰期、离线场景Coalescing合并释放时检查相邻空闲块能合并就合并成大块所有需要 free list 的场景Slab 分配按固定大小的块来分配从根本上消除外部碎片分配大小比较一致的场景定期重启终极手段碎片太多了干脆重启碎片累积无法避免的场景对于大模型推理来说分池 Size-Class Coalescing是最常用的组合。内存整理Defrag虽然效果好但 GPU 上做数据搬运的成本很高而且需要暂停推理一般只在低峰期或者运维窗口做。八、生产内存池接口建议与小结8.1 生产级内存池应该长什么样最后给大家一个生产级内存池的接口设计参考。一个好用的内存池接口应该简洁而完整classMemoryPool{public:// 核心接口 void*alloc(size_t size,MemType typeMemType::Activation);voidfree(void*ptr,size_t size,MemType typeMemType::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 思考问题如果你的推理服务需要支持动态 batchbatch size 随时变化Activation Pool 用 Arena 还合适吗为什么如果不合适你会怎么改KV Cache 的大小和序列长度正相关。如果用户的输入长度差异很大有的 128 token有的 4096 tokenKV Pool 的碎片问题会特别严重。你有什么好的策略来缓解内存池的设计中分配速度和内存利用率往往是矛盾的。你觉得在大模型推理场景下哪个更重要为什么系列回顾第①篇Runtime架构设计从Session到Graph Executor的完整骨架第②篇本篇内存池设计下一篇我们会聊Kernel 调度与算子融合敬请关注