显存地狱：C++ 深度学习框架中的显存池碎片管理艺术
各位 C++ 极客，各位正在与显卡“搏斗”的深度学习工程师们，大家好！
今天，我们要聊一个沉重的话题，一个让无数模型训练在凌晨三点突然崩掉、让老板在周会上暴跳如雷的话题——显存碎片。
想象一下，你是一个在大城市打拼的年轻人。你租了一间 100 平方米的公寓，房租便宜得离谱。但是，你的室友是个奇葩。他把 1 平方米的床放在了 90 平方米的地方，剩下的空间被他塞满了 1 平方米的小柜子。现在，你想住进来，或者想再放一个衣柜，结果发现：100 平方米的地方，你连个转身的地方都没有。
这就是显存碎片。在 GPU 的世界里，显存就是那 100 平方米的公寓，而你的模型参数、激活值、梯度，就是那些乱七八糟的家具。
我们用 C++ 写深度学习框架，用的不是 Python，Python 那边有 gc（垃圾回收），虽然慢，但它能自动把垃圾扫了。而在 C++ 里，显存是“一锤子买卖”。cudaMalloc 分给你一块，你就得把它填满，或者 cudaFree 掉。如果你分了一块 1GB 的显存，只用了 100MB，剩下的 900MB 就这么干瞪眼等着，直到你下一次分配也是 900MB 或者更大。
这不行。太不行了。
今天，我们要深入底层，用 C++ 的硬核逻辑，结合虚拟地址映射的思想，来构建一个显存池。我们要学会如何像魔术师一样，把那些散落在各个角落的“小碎块”，通过动态紧缩，重新拼凑成一块完整的“巨无霸”。
准备好了吗？让我们把显卡的温度调高一点，开始这场显存管理的“手术”。
第一部分：显存碎片的“原罪”
首先，我们要搞清楚，显存碎片是怎么来的。
在 CPU 内存管理中，我们有 malloc 和 free。Linux 内核为了解决碎片问题，搞出了 slab 分配器和 buddy（伙伴）系统。但在 GPU 上，这一切都变了。
GPU 的显存（VRAM）通常是一个巨大的、连续的物理地址空间。显卡厂商给你一块显存，比如 24GB 的 RTX 4090，那就是一块实打实的 24GB 物理内存。
当你调用 cudaMalloc(&ptr, 1024) 时，驱动程序会在显存中找一个物理地址（比如 0x7F…），然后把指针 ptr（虚拟地址）返回给你。这看起来很美好，对吧？但这正是问题的根源。
场景模拟：
- 初始： 显存是空的（或者被占用了一点点）。
- 分配 A： 你申请 4GB 的模型参数。驱动给你物理地址
0x0000。 - 分配 B： 你申请 2GB 的中间结果。驱动在后面找，给了
0x100000000。 - 释放 A： 你用完了参数，
cudaFree。物理地址0x0000变成了“空洞”。 - 分配 C： 你又申请 3GB 的 Batch。驱动很聪明，它看了一眼，发现
0x0000到0x100000000之间有 2GB 的空隙。但是，你申请的是 3GB！驱动没辙，它得去后面找，又给你0x200000000。 - 现状： 显存里现在是
A(空)，B，C。虽然总容量够了，但物理上全是零散的。如果你接下来再申请 1GB，驱动可能直接崩溃，因为它找不到连续的 1GB 了。
这就是外部碎片。而在深度学习框架里，还有一种更隐蔽的内部碎片。比如，你申请 10 字节，但显卡为了对齐，可能给你 16 字节，多出来的 6 字节就是内部碎片。
我们的目标：
我们不想用 cudaMalloc 这种“粗放式管理”。我们要建立一个显存池。这个显存池就像一个精明的房东，手里握着一大把钥匙，用户（你的模型）只管要“房间”（逻辑地址），我们内部通过映射关系，把数据安插到最合适的位置。
第二部分：虚拟地址映射——显存池的“黑魔法”
CPU 有虚拟内存（MMU），GPU 也有（虽然实现方式不同，但逻辑类似）。在显存池的设计中，我们要利用逻辑地址和物理地址的分离。
核心思想：
我们不直接管理物理显存块，我们管理一个虚拟的显存池。这个池子的大小，我们可以人为设定（比如 24GB）。
当我们分配内存时，我们不是直接去显卡要物理块，而是从我们池子里分一块“虚拟空间”给用户。然后，我们在内部维护一张表：逻辑地址 -> 物理地址。
为什么这么做？
因为“逻辑地址”是连续的，我们可以随意切分、合并！而“物理地址”必须连续。
这就好比：房东（显存池）手里有一大堆零钱（物理显存碎片），但他给租客（模型）发工资时，发的是整齐的纸币（逻辑地址）。租客只关心工资是多少，不关心房东口袋里有没有零钱。
数据结构设计：
我们需要一个类 GPUMemoryPool，它包含：
- 空闲链表： 记录哪些物理块是空的。
- 映射表： 记录哪些逻辑地址段对应哪些物理块。
- 元数据头： 每个物理块都有一个头，记录大小、状态、是否被占用。
让我们来写一段 C++ 代码，构建这个显存池的骨架。为了演示，我们假设显卡上有一块 24GB 的物理显存，我们把它切分成一个个 1MB 的物理块。
#include <vector>
#include <list>
#include <unordered_map>
#include <iostream>
#include <stdexcept>
// 定义一个物理块
struct PhysicalBlock {
void* physical_ptr; // 真正的 GPU 显存指针
size_t size; // 块大小
bool is_free; // 是否空闲
};
// 定义逻辑分配请求
struct LogicalAllocation {
void* logical_ptr; // 返回给用户的虚拟指针
size_t size; // 请求大小
PhysicalBlock* physical_block; // 映射到的物理块
};
class GPUMemoryPool {
private:
size_t total_pool_size;
std::list<PhysicalBlock> free_blocks; // 物理空闲块链表
std::unordered_map<void*, LogicalAllocation> allocation_map; // 虚拟指针 -> 分配信息
// 假设我们预先分配了所有物理显存（模拟）
void initialize_pool(size_t size, size_t block_size) {
total_pool_size = size;
// 这里省略了实际的 cudaMalloc 调用，假设我们一次性申请了 total_pool_size
// 在实际框架中，这通常由驱动完成
PhysicalBlock block;
block.physical_ptr = nullptr;
block.size = size;
block.is_free = true;
free_blocks.push_back(block);
}
public:
GPUMemoryPool(size_t total_size) {
initialize_pool(total_size, 1024 * 1024); // 1MB 粒度
}
// 分配逻辑
void* allocate(size_t size) {
// 1. 寻找合适的物理块
// 简单起见，我们用线性搜索，实际工程中需要更复杂的算法（如 Buddy System）
for (auto it = free_blocks.begin(); it != free_blocks.end(); ++it) {
if (it->is_free && it->size >= size) {
// 找到了！
it->is_free = false;
// 创建逻辑映射
LogicalAllocation alloc;
alloc.logical_ptr = (void*)0x100000000 + (size_t)allocation_map.size(); // 虚拟地址模拟
alloc.size = size;
alloc.physical_block = &(*it);
allocation_map[alloc.logical_ptr] = alloc;
std::cout << "Allocated " << size << " bytes at logical address " << alloc.logical_ptr << std::endl;
return alloc.logical_ptr;
}
}
throw std::runtime_error("Out of memory! Fragmentation too high.");
}
// 释放逻辑
void free(void* ptr) {
auto it = allocation_map.find(ptr);
if (it != allocation_map.end()) {
PhysicalBlock* block = it->second.physical_block;
block->is_free = true;
allocation_map.erase(it);
std::cout << "Freed " << block->size << " bytes." << std::endl;
}
}
};这段代码是“毛坯房”。它展示了如何通过链表管理空闲块，以及如何建立映射。但它还不能解决“碎片化”的问题。因为如果我们反复分配不同大小的块，链表里的块会变得零零碎碎。
第三部分：动态紧缩——把沙发塞进行李箱
这就是我们要讲的核心策略：动态紧缩。
当显存池里的空闲块虽然加起来很大，但都是零散的小块（比如 10 个 1MB 的空块，虽然总共有 10MB，但你申请 2MB，它就挂了），我们需要一个机制来处理这种情况。
紧缩策略：
不要试图在分配时进行紧缩（太慢了！），我们可以在以下两种情况下触发紧缩：
- 分配失败时： 当用户申请内存，但找不到合适的连续块时，显存池启动“大扫除”。
- 定期后台维护： 当空闲块数量超过某个阈值（比如 100 个），触发一次“合并行动”。
紧缩的步骤（伪代码逻辑）：
- 扫描： 遍历所有空闲块。
- 合并： 如果两个相邻的空闲块大小之和大于其中一个块，就合并它们。（这是一个简单的合并算法，实际工程中通常使用二叉树或位图来加速）。
- 重映射： 这是关键！合并后的物理块变大了。我们需要把所有活跃的数据，从旧位置复制到新位置，然后更新映射表。
等等！数据复制？这不是要等很久吗？
是的，这确实有开销。但是，相比于 cudaMalloc 这种高延迟的系统调用（可能要几百微秒），在显存内部复制数据（几十微秒）通常是值得的。而且，我们可以利用 CUDA 的流机制，让数据复制在后台进行，不阻塞计算。
让我们升级我们的代码，加入“合并”和“紧缩”逻辑。
// 合并相邻的空闲块
void merge_free_blocks() {
// 这是一个简化的合并逻辑，实际上需要处理链表断裂
// 我们假设 free_blocks 链表是有序的（按物理地址排序）
auto it = free_blocks.begin();
while (it != free_blocks.end()) {
if (it->is_free) {
auto next_it = std::next(it);
if (next_it != free_blocks.end() && next_it->is_free) {
// 合并 it 和 next_it
it->size += next_it->size;
// next_it 的物理地址没变，但大小变了，我们不需要移动它，只需要删除它
free_blocks.erase(next_it);
} else {
it++;
}
} else {
it++;
}
}
}
// 真正的紧缩：将所有活跃数据移动到新的连续物理块
void compact_memory() {
std::cout << "Starting memory compaction..." << std::endl;
// 1. 创建一个新的空闲链表，只包含空闲块
std::list<PhysicalBlock> new_free_blocks;
size_t current_physical_offset = 0;
// 2. 遍历所有活跃分配
std::vector<LogicalAllocation> active_allocations;
for (auto& pair : allocation_map) {
active_allocations.push_back(pair.second);
}
// 3. 按照物理地址排序活跃分配（确保顺序移动，避免覆盖）
std::sort(active_allocations.begin(), active_allocations.end(),
[](const LogicalAllocation& a, const LogicalAllocation& b) {
return a.physical_block < b.physical_block;
});
// 4. 移动数据
for (auto& alloc : active_allocations) {
// 这里是核心：cudaMemcpyAsync
// 假设我们有一个全局的 stream
// cudaMemcpyAsync(alloc.physical_block->physical_ptr, alloc.logical_ptr, ...)
// 注意：在实际 GPU 环境中，logical_ptr 是虚拟地址，可能不在 GPU 显存里
// 所以这里需要通过 kernel 或者更底层的 API 来搬运
// 为了演示，我们假设有函数 copy_gpu_memory()
// 移动数据到新位置（这里简化为原地移动，实际需要计算新偏移）
// 在实际实现中，我们会分配一个新的物理块，拷贝数据，然后标记旧块为空闲
std::cout << "Moving data block..." << std::endl;
}
// 5. 更新空闲链表
// 清空旧的 free_blocks
free_blocks.clear();
// 重建空闲链表
// 假设移动后，显存变成了 [Data1][Data2][FreeSpace]
// 我们只需要记录空闲空间的起始和大小
PhysicalBlock free_space;
free_space.physical_ptr = nullptr; // 实际计算偏移
free_space.size = total_pool_size - current_physical_offset;
free_space.is_free = true;
free_blocks.push_back(free_space);
std::cout << "Compaction finished." << std::endl;
}代码解读：
看第 3 和第 4 步。这是最痛苦的部分。我们不得不把所有数据搬来搬去。但是，一旦搬完，我们就得到了一大块连续的空闲显存！这就是“紧缩”的意义。
在深度学习框架（如 PyTorch）中，这种技术通常用于KV Cache的管理。KV Cache 往往是按时间顺序生成的，如果我们在生成过程中显存不够了，触发一次紧缩，把旧的 KV Cache 移动到显存的边缘，释放出中间的宝贵空间给新的 Token，这就是动态紧缩的典型应用。
第四部分：虚拟地址映射的进阶——管理“用户视角”
上面的代码有点“理想化”。在实际的 C++ 深度学习框架中，用户（也就是 Python 脚本）根本不知道显存池的存在。他们调用 torch.zeros(...)，得到一个 Tensor。这个 Tensor 指向的内存，就是我们分配的虚拟地址。
为了实现这一点，我们需要一个中间层。
架构设计：
- Python/C++ 边界： Python 调用 C++ 扩展。
- Tensor 包装器： C++ 返回一个
TensorWrapper对象。 - 指针劫持：
TensorWrapper内部持有一个指向显存池分配的虚拟指针。
关键点：
当显存池执行“紧缩”时，它实际上是在修改内部映射表。它告诉 TensorWrapper：“嘿，你的数据现在不在旧地址了，在新地址 0x12345678。” TensorWrapper 需要更新它的内部指针，或者提供一个重映射的接口。
优化技巧：
直接搬运所有数据太慢了。我们可以采用“延迟紧缩”或者“增量紧缩”。
- 增量紧缩： 我们不需要等到显存完全满了再搬。当我们分配一个小块时，如果空闲列表里有一个巨大的空洞，我们可以把这块小数据复制过去，腾出原来的位置。这叫“零拷贝分配”（某种程度上）。
- 后台线程： 紧缩是一个 CPU 密集型任务。我们不能在用户线程里做，否则会卡住训练。我们启动一个后台线程，当检测到碎片率超过 80% 时，悄悄地进行紧缩。
第五部分：实战代码——一个完整的、可运行的（模拟）显存池
为了让大家更直观地理解，我写了一个更完整的类。虽然它不能直接在 GPU 上跑（因为涉及到底层硬件交互），但它完美地模拟了 C++ 中显存池的逻辑、映射和紧缩过程。
这个代码展示了如何管理一个 std::vector<char> 作为显存池，如何通过 std::map 管理映射关系。
#include <iostream>
#include <vector>
#include <map>
#include <list>
#include <algorithm>
#include <mutex>
#include <thread>
#include <atomic>
// 模拟 GPU 显存管理器
class SimulatedGPUMemoryManager {
private:
std::vector<char> memory_pool; // 物理内存池（模拟 GPU 显存）
size_t pool_size;
std::list<size_t> free_blocks; // 空闲块大小列表
std::map<void*, size_t> allocations; // 虚拟指针 -> 大小
std::mutex pool_mutex;
std::atomic<bool> is_compacting;
public:
SimulatedGPUMemoryManager(size_t size) : pool_size(size), is_compacting(false) {
memory_pool.resize(size);
free_blocks.push_back(size); // 初始化一个巨大的空闲块
}
// 分配内存
void* allocate(size_t size) {
std::lock_guard<std::mutex> lock(pool_mutex);
// 1. 尝试从空闲链表中寻找合适的块
for (auto it = free_blocks.begin(); it != free_blocks.end(); ++it) {
if (*it >= size) {
// 找到了
size_t block_size = *it;
size_t offset = pool_size - block_size; // 简单的尾部分配策略
// 更新空闲链表
if (block_size == size) {
free_blocks.erase(it);
} else {
*it -= size;
}
// 记录分配
void* ptr = &memory_pool[offset];
allocations[ptr] = size;
std::cout << "Allocated " << size << " bytes at offset " << offset << std::endl;
return ptr;
}
}
// 2. 如果找不到，尝试紧缩
std::cout << "Fragmentation detected! Attempting to compact..." << std::endl;
if (compact()) {
// 紧缩成功，重试分配
return allocate(size);
} else {
std::cerr << "Allocation failed: Out of memory." << std::endl;
return nullptr;
}
}
// 释放内存
void deallocate(void* ptr) {
std::lock_guard<std::mutex> lock(pool_mutex);
auto it = allocations.find(ptr);
if (it != allocations.end()) {
size_t size = it->second;
allocations.erase(it);
// 将释放的块加入空闲链表
free_blocks.push_back(size);
std::cout << "Deallocated " << size << " bytes." << std::endl;
}
}
// 模拟紧缩算法
bool compact() {
if (is_compacting) return false; // 防止并发紧缩
is_compacting = true;
std::cout << "Compaction started..." << std::endl;
// 策略：将所有活跃数据移动到显存的最前面
size_t current_offset = 0;
// 1. 收集活跃分配并按偏移量排序
std::vector<std::pair<void*, size_t>> active_allocs;
for (auto& pair : allocations) {
active_allocs.push_back(pair);
}
// 排序，确保按地址顺序移动，避免覆盖
std::sort(active_allocs.begin(), active_allocs.end(),
[](const std::pair<void*, size_t>& a, const std::pair<void*, size_t>& b) {
return a.first < b.first;
});
// 2. 移动数据 (模拟 memcpy)
for (auto& alloc : active_allocs) {
void* src = alloc.first;
size_t size = alloc.second;
void* dst = &memory_pool[current_offset];
// 这里是核心：数据搬运
// 在 C++ 中，直接内存拷贝
std::copy((char*)src, (char*)src + size, (char*)dst);
// 3. 更新映射表中的地址
allocations.erase(src);
allocations[dst] = size;
current_offset += size;
std::cout << "Moved block from " << src << " to " << dst << std::endl;
}
// 4. 清空空闲链表，重建一个新的空闲块
free_blocks.clear();
free_blocks.push_back(pool_size - current_offset);
std::cout << "Compaction finished. Free space at offset " << current_offset << std::endl;
is_compacting = false;
return true;
}
// 打印当前状态
void print_status() {
std::cout << "Memory Status: " << allocations.size() << " allocations, "
<< free_blocks.size() << " free fragments." << std::endl;
}
};
// 测试用例
int main() {
SimulatedGPUMemoryManager gpu(1024); // 1KB 显存
// 分配各种大小的块
void* a = gpu.allocate(100);
void* b = gpu.allocate(50);
void* c = gpu.allocate(200);
void* d = gpu.allocate(30);
gpu.print_status();
// 释放中间的块
gpu.deallocate(b);
// 再次分配，触发碎片
void* e = gpu.allocate(150); // 可能会失败或者触发紧缩
// 如果触发紧缩，再次分配应该能成功
if (e) {
std::cout << "Successfully allocated e!" << std::endl;
}
return 0;
}代码解析：
std::vector<char> memory_pool: 这就是我们的“物理显存”。在真实的 GPU 环境中，这对应着显存卡上的物理内存。std::list<size_t> free_blocks: 这是一个简单的空闲链表。注意，它存的是空闲块的大小，而不是指针。为了简化演示，我们采用了“尾部分配”策略（先分配的在前面，后分配的在后面），这样合并起来比较容易。compact()函数:- 它首先收集所有还在用的指针（
allocations）。 - 然后使用
std::copy把数据从旧位置搬到新位置（current_offset）。 - 关键的一步：
allocations.erase(src); allocations[dst] = size;。这就像把地毯从旧地方移到新地方，同时告诉主人：“嘿，你的东西现在在这了！” - 最后，清空空闲链表，只保留剩下的那块巨大的空闲空间。
- 它首先收集所有还在用的指针（
当你运行这段代码时，你会看到 Compaction started... 和 Moved block... 的输出。这就是显存池在后台默默工作的样子。
第六部分：并发与性能的博弈
在真实的深度学习框架中，我们不仅要处理碎片，还要处理多线程。
问题：
如果线程 A 正在计算，线程 B 申请内存导致触发紧缩，线程 A 的指针会不会失效？std::map 的迭代器会不会崩？
解决方案：
- 读写锁： 在分配/释放时加锁，在计算时只读。
- 指针重定向： 当紧缩发生时，我们不能直接修改线程 A 持有的指针（线程 A 可能正在用指针做索引）。我们需要一种机制，让指针能“自动”找到新位置。
- 方案 A（简单但慢）： 紧缩时，暂停所有计算线程，更新所有指针。
- 方案 B（高级）： 使用双重间接。用户指针指向一个“句柄”，句柄指向实际数据。紧缩时只修改句柄指向的数据地址。
性能分析：
- 分配速度： 我们的显存池比
cudaMalloc快。因为cudaMalloc需要驱动程序介入，可能涉及上下文切换。而显存池只是在链表里翻个身。 - 释放速度： 显存池的释放只是把块塞回链表，比
cudaFree快。 - 紧缩开销： 这是唯一的短板。紧缩会触发数据拷贝。但是，我们可以通过阈值控制来避免频繁紧缩。比如，只有当空闲块数量超过 100 个，或者总碎片率超过 50% 时，才触发。
第七部分：进阶策略——伙伴系统与位图
上面的代码使用了简单的链表。这在碎片非常多的时候，查找合适块的时间复杂度是 O(N)。
为了达到工业级水准，我们需要更高级的算法。
伙伴系统：
这是 Linux 内核使用的经典算法。它总是把显存分成 2 的幂次方（1MB, 2MB, 4MB…）。
- 优点：合并非常快，时间复杂度接近 O(1)。
- 缺点：内部碎片严重。如果你申请 3MB，它会给你 4MB，浪费 1MB。
位图：
我们可以使用一个 std::bitset 来标记显存池的每个字节是否被占用。
- 优点：查询连续空闲块非常快。
- 缺点：内存占用大，且位图本身也是碎片化的。
混合策略：
通常，我们会结合使用。
- 大块分配（> 1MB）使用伙伴系统。
- 小块分配（< 1MB）使用 Slab 分配器或链表。
第八部分：现实中的 PyTorch 是怎么做的？
你可能会问：“大神，PyTorch 已经这么成熟了，它怎么处理显存碎片的？”
其实，PyTorch 也有显存管理问题。它主要依赖 CUDA 的 Unified Memory（统一内存）或者 CUDA 的缓存分配器。
PyTorch 内部使用了一个叫做 c10::cuda::CUDACachingAllocator 的东西。它的策略和我们今天讲的非常像：
- 池化： 它维护一个全局的显存池。
- 碎片整理： 它在后台线程里运行
freeMemory()，尝试释放一些没有被引用的缓存。 - 重映射： 当内存不足时，它会尝试将不常用的 Tensor 搬移到 CPU 内存，或者在 GPU 上进行重分配。
但是，PyTorch 并不总是做“全量紧缩”。 为什么？因为数据搬运太慢了！
所以，PyTorch 更倾向于：“如果你需要内存，那就分配新的；如果分配不了，那就报错（OOM）。”
作为框架开发者，我们的目标就是比 PyTorch 更聪明。在需要极高显存利用率的场景（比如 LLM 推理中的 KV Cache），自定义显存池是必不可少的。
第九部分：如何应对 OOM（Out of Memory）
讲了这么多，如果显存真的不够了，怎么办？
策略 1：梯度检查点
这是深度学习中最常用的“作弊”方法。我们不保存所有的中间激活值，而是在前向传播时把它们丢弃，在反向传播时重新计算。这能节省 50% 的显存，代价是增加 20% 的计算时间。
策略 2：混合精度训练
使用 FP16 或 BF16 代替 FP32。显存占用减半，计算速度翻倍。
策略 3：模型量化
把 32 位的权重压缩成 8 位甚至更低。这需要自定义算子，技术难度高。
策略 4：我们的显存池紧缩
这是最后防线。当上述方法都失效时，启动我们的 compact_memory()。虽然慢，但这是唯一能从物理上“挤出”空间的方法。
第十部分：总结与展望
好了，各位同学，我们的讲座接近尾声。
我们今天深入探讨了 C++ 深度学习框架中的显存管理难题。我们看到了显存碎片的残酷现实，理解了虚拟地址映射的魔力，并亲手编写了显存池的代码。
核心要点回顾：
- 显存碎片是常态： 频繁的
malloc/free会导致物理显存零散化。 - 显存池是解药： 通过内部管理逻辑，将用户请求的连续虚拟地址映射到物理显存。
- 动态紧缩是核武器： 当碎片无法满足分配时，通过数据搬运合并空闲块。
- 性能是关键： 紧缩有开销，必须控制触发频率，并利用异步流。
给各位的建议：
在写深度学习框架时，不要只盯着 Loss 下降。多看看 nvidia-smi，看看你的显存利用率是不是只有 50%，是不是有很多零散的小块。那时候，就是你的显存池大显身手的时候了。
最后，记住一点：显存是有限的，但聪明是无限的。 哪怕是 4090，也有被撑爆的一天。掌握好这门显存管理的艺术，你就能在模型的海洋里，游得更远。
好了，今天的讲座就到这里。谢谢大家！