2. 显存池数据结构设计

聊到显存池，很多人第一反应就是「搞个内存池改改不就行了？」嗯，我刚开始也这么想。但真上手做分布式训练优化时才发现——显存这玩意儿，比普通内存娇气多了。它没有虚拟内存机制，碎片多了就是真的多了，回收不回来。

所以今天咱们就聊聊显存池里的三个核心数据结构：空闲链表、哈希表和伙伴系统。说白了，它们就是显存池的「仓库管理员」，负责记录哪块显存空着、哪块被占着、怎么分配最不浪费。

1. 空闲链表：最朴素的管理方式

空闲链表，顾名思义，就是把所有空闲的显存块串成一个链表。每次分配时从头遍历，找到一块足够大的就切出去。

我个人习惯把空闲链表分成两种：

• 单链表：每个节点只存大小和地址，适合小规模场景
• 双链表：多了前后指针，合并相邻空闲块时更方便

我在项目中遇到过一个问题：用单链表管理显存，跑了一个大模型训练，结果分配速度越来越慢。为什么？因为每次都要从头遍历，空闲块多了以后，O(n)的复杂度就扛不住了。

// 一个简单的空闲链表节点
struct FreeBlock {
    void* addr;        // 显存起始地址
    size_t size;       // 块大小
    FreeBlock* next;   // 下一个空闲块
};

2. 哈希表：快速查找的利器

空闲链表慢在查找，那哈希表就是来治这个病的。我们把显存块的大小作为key，地址作为value，存到哈希表里。分配时直接根据请求大小查表，O(1)复杂度搞定。

你想想看，如果每次分配都要遍历链表，那多卡通信时频繁的显存申请和释放就会成为瓶颈。哈希表正好解决了这个问题。

我建议这样设计：

• 哈希函数：用块大小取模，或者用位运算快速映射
• 冲突处理：链地址法，每个桶里挂一个空闲链表
• 动态扩容：当负载因子超过0.75时，重新哈希

数据结构	分配复杂度	释放复杂度	碎片控制
空闲链表	O(n)	O(1)	差
哈希表	O(1) 平均	O(1) 平均	中
伙伴系统	O(log n)	O(log n)	好

3. 伙伴系统：对抗碎片的终极方案

说到碎片问题，伙伴系统是我最推荐的方式。它的核心思想是：把显存按2的幂次分成不同大小的块，分配时找最合适的那块，释放时尝试合并相邻的「伙伴」。

举个例子：你申请3MB显存，伙伴系统会给你一个4MB的块（2的2次方）。释放时，如果相邻的4MB块也是空闲的，就合并成8MB。这样层层合并，碎片自然就少了。

我记得有一次做多卡通信优化，用伙伴系统管理显存池，训练了三天三夜，显存碎片率始终控制在5%以内。换成空闲链表，第二天就涨到了30%。

// 伙伴系统的核心数据结构
struct BuddySystem {
    void* base_addr;           // 显存池基地址
    size_t total_size;         // 总大小（2的幂次）
    FreeList free_lists[32];   // 每个级别一个空闲链表
    // free_lists[i] 管理大小为 2^i 的块
};

4. 三种结构的组合使用

在实际的显存池设计中，我很少只用一种结构。通常是这样组合的：

• 小块分配（< 1MB）：用哈希表，追求速度
• 中块分配（1MB ~ 256MB）：用伙伴系统，平衡速度和碎片
• 大块分配（> 256MB）：用空闲链表，直接管理

说白了，没有银弹。不同的场景用不同的策略，这才是工程实践的真谛。

5. 核心逻辑流程图

下面这张图展示了显存池中三种数据结构的协作关系：

6. 实际项目中的选择建议

最后，我根据实际项目经验，给几个选择建议：

• 单卡小模型：空闲链表就够了，简单可靠
• 多卡大模型训练：必须上伙伴系统，碎片控制是关键
• 动态显存分配频繁：哈希表+伙伴系统混合，兼顾速度和碎片
• 显存池大小固定：伙伴系统是最优解，没有之一

好了，关于显存池的数据结构设计就聊到这儿。这三种方式各有千秋，关键是根据你的场景选对工具。嗯，希望这些经验对你有帮助。