4. 显存池核心设计：预分配策略、块管理算法
各位同学，今天我们来聊聊显存池里最核心的两个东西——预分配策略和块管理算法。说实话，这块内容我当年刚接触时也觉得挺枯燥的，直到自己在项目里被显存碎片折磨得死去活来，才真正体会到这些设计有多重要。
你想想看，GPU显存就那么几十个GB，训练大模型时动不动就占满。如果每次申请释放都走cudaMalloc，那性能简直没法看。更可怕的是，频繁分配释放会产生大量碎片，最后明明还有空闲显存，却申请不到一块连续的大块内存。嗯，这就是我们要搞显存池的原因。
4.1 预分配策略：先下手为强
预分配，说白了就是「先占坑」。在程序启动时，一次性向GPU申请一大块显存，然后自己管理这块内存的分配和回收。这样做的好处很明显——避免了频繁的cudaMalloc/cudaFree开销。
我个人习惯把预分配策略分成三种：
- 静态预分配：启动时就把所有显存都申请好。简单粗暴，适合显存需求确定的场景。我在做推理服务时常用这种方式，因为模型大小是固定的。
- 动态预分配：先申请一个初始大小，不够了再追加。灵活性更好，适合训练场景。我记得有一次训练一个超大模型，初始分配了80%的显存，结果中间有个大tensor死活放不下，后来改成动态增长才搞定。
- 按需预分配：根据历史分配记录预测未来需求。这个比较高级，我一般不用，因为预测不准反而容易浪费。
核心要点：预分配的大小很关键。分配太多，其他进程没显存用；分配太少，又得频繁扩容。我建议初始分配总显存的70%-80%，留点余量给系统和其他进程。
4.2 块管理算法：Buddy System
好，显存已经预分配好了，接下来怎么管理这些块？我最常用的就是Buddy System（伙伴系统）。
Buddy System的核心思想很简单：把内存块不断对半切，直到满足需求。释放时再合并相邻的「伙伴」块。这样做的好处是分配和释放都很快，而且能有效减少外部碎片。
举个例子，假设我们有1GB的显存池：
// Buddy System 分配示例
// 初始：1GB 空闲块
// 申请 128MB：切分 1GB -> 512MB + 512MB -> 256MB + 256MB + 512MB -> 128MB + 128MB + 256MB + 512MB
// 分配第一个 128MB 给用户
// 释放时：检查伙伴是否空闲，是则合并
// 释放 128MB -> 检查伙伴 128MB 是否空闲 -> 合并成 256MB -> 继续向上合并我在项目中遇到过一个问题：Buddy System虽然外部碎片少，但内部碎片可能很严重。比如你申请70MB，它只能给你128MB的块，那58MB就浪费了。这就是所谓的「内部碎片」。
避坑指南：我曾经在一个训练任务里用了纯Buddy System，结果显存利用率只有60%多。后来发现是大量小tensor分配导致的内部碎片。解决办法是结合Slab分配器来处理小对象。
4.3 块管理算法：Slab分配器
Slab分配器，说白了就是「按规格生产」。它专门用来处理固定大小的对象分配，比如4KB、8KB、16KB这些常见尺寸。
Slab的工作原理是这样的：
- 预先创建多个slab缓存，每个缓存管理一种固定大小的对象
- 每个slab内部维护一个空闲列表，分配时直接从列表取
- 释放时把对象放回空闲列表，不需要合并操作
我一般在显存池里这样搭配使用：
// 显存池混合策略
// 1. 大块分配（>1MB）：走 Buddy System
// 2. 小块分配（<=1MB）：走 Slab 分配器
// 3. 中间尺寸：走普通链表管理
struct MemoryPool {
BuddySystem* buddy; // 大块管理
SlabAllocator* slab; // 小块管理
FreeList* mid; // 中间尺寸
};你想想看，为什么需要这种混合策略？因为Buddy System处理小对象时内部碎片太多，而Slab处理大对象时又不够灵活。两者结合，才能发挥各自优势。
注意：Slab分配器的颜色标记（比如用不同颜色区分不同大小的slab）虽然方便调试，但会引入额外开销。我建议只在debug模式下启用，release模式关掉。
4.4 核心逻辑流程图
下面我用一张SVG图来展示显存池的核心工作流程。这张图我画了好几次才满意，希望能帮大家理清思路。
4.5 实际项目中的经验总结
做了这么多年显存池，我踩过不少坑，也积累了一些经验。这里分享几个关键点：
| 场景 | 推荐策略 | 原因 |
|---|---|---|
| 大模型训练 | 动态预分配 + Buddy System | 训练过程中tensor大小变化大，需要灵活扩容 |
| 推理服务 | 静态预分配 + Slab | 推理时tensor大小固定，Slab效率最高 |
| 混合负载 | 动态预分配 + 混合策略 | 兼顾灵活性和效率 |
个人经验：我建议在显存池里加一个「碎片整理」功能。当空闲块数量超过阈值时，触发一次整理。虽然整理过程会暂停分配，但能有效避免碎片积累。我曾经在一个训练任务里，每1000步整理一次，显存利用率从70%提升到了90%。
嗯，关于显存池的核心设计，今天就讲到这里。预分配策略和块管理算法是显存池的基石，理解透了，后面讲算子优化时你会觉得轻松很多。记住，没有银弹，实际项目中要根据场景灵活选择。
一句话总结：预分配解决性能问题，块管理解决碎片问题。两者配合，才能打造一个高效的显存池。
想要更多资料，联系可可