大模型显存碎片消除实战：基于段式虚拟内存的 Buddy System 伙伴内存分配器手写

在大语言模型（LLM）长文本推理与多模态生成任务中，显存（VRAM / HBM）的物理容量往往直接决定了系统并发服务能力的生死线。在长达数十小时的在线服务运行后，一个令无数算法和运维团队抓狂的致命现象频繁发生：物理显存监控明明显示还有 20GB 以上的空闲空间，但在为某一个新到达的并发请求分配一个仅需 2GB 的连续张量或 KV-Cache 缓冲区时，系统却突然抛出惨烈的 Out of Memory (OOM) 异常而崩溃！

罪魁祸首正是被称为“内存黑洞”的外部显存碎片（External Memory Fragmentation）。

在自回归解码中，不同用户的 Prompt 序列长度千差万别，Token 逐字生成的生命周期动态交错。如果依赖标准的显存分配器（如 cudaMalloc 或通用的首次适配贪心池），频繁申请和释放不同尺寸的内存块，会迅速将原本平坦连续的大显存切割成数以万计不可复用的“离散孤岛”。当需要申请一个稍大的大页连续内存时，没有任何一个碎片空隙能够容纳，导致系统直接瘫痪。

解决这一物理顽疾的经典终极利器，是计算机体系结构中历经半个世纪检验的经典机制——伙伴内存分配系统（Buddy Memory Allocation System）。

本文将从显存碎片的微架构机理切入，深入剖析伙伴系统递归二分拆解与地址异或极速合并（XOR Coalescing）的数学法则，并手写一个基于 C++23 的工业级大模型显存段式伙伴分配器。

一、显存碎片的微架构灾难：内部碎片 vs 外部碎片

在动态内存管理中，碎片分为两类互相对立的形态：

显存碎片的两种形态对比:
1. 内部碎片 (Internal Fragmentation):
申请 180MB —> 分配器对齐给 256MB 块 —> 浪费 76MB 内部死空间
(特点: 空间在块内被浪费，但系统始终保持物理对齐规整)

2. 外部碎片 (External Fragmentation – 致命杀手!):
空闲显存总和 = 20GB，但物理分布为:
[ 已用 1GB ] [ 空闲 200MB ] [ 已用 2GB ] [ 空闲 500MB ] … [ 空闲 100MB ]
此时尝试申请一个 2GB 的连续特征图张量:
===> 没有任何一段物理地址连续长度达到 2GB! 触发惨烈 OOM 假死!

对于超大张量与高性能算子而言，外部碎片是绝对不可容忍的毒瘤。因为现代 GPU 和高速 DMA 引擎在执行宽位向量加载（如 512 位向量指令或 TMA 张量描述符搬运）时，必须要求张量在物理地址上具备严格的连续性与大页对齐。一旦地址被撕裂，算子根本无法启动。

二、伙伴系统（Buddy System）的数学之美：地址异或的 O(1) 合并

伙伴分配系统的核心哲学，是用可控的微小内部碎片，彻底根除失控的外部碎片。

设整个显存池的初始总容量为 $2^M$ 字节（如 16GB），最小可分配粒度为 $2^K$ 字节（例如 2MB 大页）。整个系统被划分为一个阶梯状的空闲链表数组（Free Lists: Order $0$ 到 Order $M-K$）：

这种自底向上的级联大块自动回缩能力，确保了系统能够随着请求的释放，100% 自动抚平所有内存裂痕，重新凝聚成数吉字节的规整大页！

三、工业级 C++23 伙伴内存分配器实现

以下代码演示了一个支持大页对齐、阶梯链表管理与位运算极速合并的工业级 Buddy Allocator 实现：

#include <iostream>
#include <vector>
#include <cstdint>
#include <cstddef>
#include <bit>
#include <concepts>
#include <algorithm>
#include <cassert>

// 伙伴系统空闲块节点
struct BuddyBlock {
BuddyBlock* next{nullptr};
BuddyBlock* prev{nullptr};
};

class BuddyMemoryAllocator {
private:
uint8_t* pool_base_{nullptr}; // 显存物理底座指针
size_t total_capacity_{0}; // 必须是 2 的幂
size_t min_block_size_{0}; // 最小分配粒度 (如 2MB)
size_t max_order_{0};

// 每个阶 (Order) 对应的空闲双向链表表头
std::vector<BuddyBlock*> free_lists_;

public:
// 初始化伙伴系统: total_bytes 与 min_block_bytes 必须为 2 的幂
BuddyMemoryAllocator(size_t total_bytes, size_t min_block_bytes = 2 * 1024 * 1024)
: total_capacity_(total_bytes), min_block_size_(min_block_bytes)
{
assert(std::has_single_bit(total_bytes));
assert(std::has_single_bit(min_block_bytes));
assert(total_bytes >= min_block_bytes);

// 计算最大阶数
max_order_ = std::countr_zero(total_bytes / min_block_bytes);
free_lists_.resize(max_order_ + 1, nullptr);

// 模拟申请整块连续的大显存空间 (物理 64 字节对齐)
pool_base_ = static_cast<uint8_t*>(::operator new(total_capacity_, std::align_val_t{64}));

// 初始状态: 将整块最大的内存作为一个整块挂入最高阶链表
BuddyBlock* root_block = reinterpret_cast<BuddyBlock*>(pool_base_);
root_block->next = nullptr;
root_block->prev = nullptr;
free_lists_[max_order_] = root_block;
}

~BuddyMemoryAllocator() {
if (pool_base_) {
::operator delete(pool_base_, std::align_val_t{64});
}
}

BuddyMemoryAllocator(const BuddyMemoryAllocator&) = delete;
BuddyMemoryAllocator& operator=(const BuddyMemoryAllocator&) = delete;

// 分配显存块
[[nodiscard]] void* allocate(size_t bytes) {
if (bytes == 0) return nullptr;

// 计算需要的实际阶数
size_t needed_size = std::max(min_block_size_, std::bit_ceil(bytes));
if (needed_size > total_capacity_) return nullptr;

size_t target_order = std::countr_zero(needed_size / min_block_size_);

// 寻找满足要求的最小阶空闲块
size_t current_order = target_order;
while (current_order <= max_order_ && free_lists_[current_order] == nullptr) {
current_order++;
}

if (current_order > max_order_) {
return nullptr; // 内存彻底耗尽，无法分配
}

// 从当前阶链表中弹出一个大块
BuddyBlock* block = pop_free_block(current_order);

// 自顶向下递归二分拆解，直到拆解到 target_order 为止
while (current_order > target_order) {
current_order–;
size_t half_bytes = (min_block_size_ << current_order);

// 孪生伙伴块的指针通过简单偏移行偏移计算
BuddyBlock* buddy = reinterpret_cast<BuddyBlock*>(
reinterpret_cast<uint8_t*>(block) + half_bytes
);

// 将伙伴块推入低一阶的空闲链表
push_free_block(current_order, buddy);
}

return static_cast<void*>(block);
}

// 释放显存块并级联自动合并
void deallocate(void* ptr, size_t bytes) {
if (!ptr) return;

size_t needed_size = std::max(min_block_size_, std::bit_ceil(bytes));
size_t order = std::countr_zero(needed_size / min_block_size_);

uint8_t* block_ptr = static_cast<uint8_t*>(ptr);
size_t offset = block_ptr – pool_base_;

// 递归向上尝试伙伴合并
while (order < max_order_) {
size_t block_bytes = (min_block_size_ << order);
// 核心: 利用单次按位异或运算，单周期瞬间算出唯一的孪生伙伴偏移量！
size_t buddy_offset = offset ^ block_bytes;
BuddyBlock* buddy = reinterpret_cast<BuddyBlock*>(pool_base_ + buddy_offset);

// 检查伙伴块是否在当前阶空闲链表中
if (!remove_if_free(order, buddy)) {
break; // 伙伴正在被占用，合并终止
}

// 伙伴块空闲，成功抓取并合并！
// 合并后大块的首地址为两者中较小的那一个
offset = std::min(offset, buddy_offset);
order++;
}

// 将最终合并得到的大块挂入对应阶的空闲链表
BuddyBlock* final_block = reinterpret_cast<BuddyBlock*>(pool_base_ + offset);
push_free_block(order, final_block);
}

private:
BuddyBlock* pop_free_block(size_t order) {
BuddyBlock* b = free_lists_[order];
if (b) {
free_lists_[order] = b->next;
if (b->next) b->next->prev = nullptr;
b->next = b->prev = nullptr;
}
return b;
}

void push_free_block(size_t order, BuddyBlock* b) {
b->next = free_lists_[order];
b->prev = nullptr;
if (free_lists_[order]) {
free_lists_[order]->prev = b;
}
free_lists_[order] = b;
}

bool remove_if_free(size_t order, BuddyBlock* target) {
BuddyBlock* curr = free_lists_[order];
while (curr) {
if (curr == target) {
// 从链表安全摘除
if (curr->prev) curr->prev->next = curr->next;
else free_lists_[order] = curr->next;

if (curr->next) curr->next->prev = curr->prev;
curr->next = curr->prev = nullptr;
return true;
}
curr = curr->next;
}
return false;
}
};

四、显存碎片压力压测与连续性保持对比

在模拟 64 线程并发、变长上下文（16MB 到 2GB 动态交错）的大模型推理压力测试中，连续运行 1,000,000 次分配与随机释放，对比朴素首次适配链表分配器与本文伙伴内存分配器的关键指标：

测试数据展现了数学原理带来的降维打击：

五、总结

大模型时代的底层架构设计，是一场与物理存储极限的残酷拉锯战。每一吉字节显存的浪费，背后都是数以万计服务器成本与延迟的沉重代价。伙伴内存分配系统以极其典雅的二叉树代数结构与位运算几何学，向我们展示了经典算法在现代 AI 异构世界里的永恒生命力。通过用 C++23 的极致严谨手写这套显存防线，我们让庞大而狂暴的动态数据流在芯片的物理地址空间上重新归于整洁与安宁。