自定义 Tensor 内存分配器：避免 glibc 碎片化

邵宇然

16人浏览 · 2026-09-09 22:59:43

邵宇然 · 2026-09-09 22:59:43 发布

自定义 Tensor 内存分配器：避免 glibc 碎片化

封面信息图

在长周期、高并发运行的 AI 推理服务与高性能计算底座中，内存碎片化（Memory Fragmentation） 往往是一个随着时间缓慢积累、但最终会导致服务物理 OOM 崩溃的“隐形杀手”。

在很多 C++/Rust 系统中，开发者习惯于直接依赖操作系统默认的通用分配器（如 Linux 标准的 glibc ptmalloc）。

然而，大模型推理场景具有极具破坏性的内存分配特征：

极高频地申请和释放尺寸跨度极大的临时张量（从几百字节的 Token 数组，到数十兆的中间激活值，再到几吉字节的 KV Cache 块）；
这些不同生命周期、不同尺寸的对象在堆上犬牙交错。

运行数小时后，监控大盘就会暴露出经典故障：系统实际使用的有效张量数据只有 4GB，但操作系统物理内存（RSS）却被撑到了 16GB 且居高不下！由于内存空洞遍地，分配器再也无法找到一块连续的 64MB 空间来容纳新 Batch，最终触发直接崩溃。

手写一个专为张量（Tensor）生命周期量身定制的 Size-Class 分箱与 Arena 内存分配器，是根治内存碎片的必由之路。

+--------------------------------------------------------------------------+
|                       通用 glibc vs 自定义 Tensor Arena 分配器架构            |
+--------------------------------------------------------------------------+
| [glibc 默认分配器 (产生碎片与内存空洞)]:                                    |
| [4KB 对象] -> [空闲空洞 128KB] -> [8MB 对象] -> [空洞 64KB] -> [4KB 对象]  |
| -> 🚨 无法向操作系统内核归还物理页 (madvise 失效)，总内存无限膨胀           |
+--------------------------------------------------------------------------+
                                    | 架构升级
                                    v
| [自定义 Tensor Arena 分箱分配器 (Zero Fragmentation)]:                     |
| 1. Small Arena (Size-Class 规整分箱: 4KB, 16KB, 64KB 紧凑数组，绝对无空洞)  |
| 2. Large Arena (针对 > 1MB 的超大激活值，直接通过 mmap 按需向内核申请/释放) |
| 3. Ring Arena (单次前向传播专享，推理结束一键 O(1) 重置游标!)               |
+--------------------------------------------------------------------------+

1. glibc 碎片化的底层物理成因

glibc malloc 采用 Chunk 链表与 Bin 机制管理堆内存。

当它通过 brk 扩展堆顶时：

堆内存是从低地址向高地址线性延伸的连续虚拟空间；
只要高地址处还有哪怕 1 个 4KB 的小对象 处于活跃状态，即使它下方有 10GB 的空间已经被 free 释放，brk 也无法向下回缩！
操作系统内核无法回收这些物理页面，导致内存监控上的常驻物理内存（RSS）永远定格在历史最高峰。

2. 基于 Size-Class 与 Arena 的张量分配器设计

针对张量生命周期的确定性，我们设计了分层分配策略：

瞬态激活值（Activation Arena）：
专为单次前向推理服务。采用 Bump Allocator（游标分配器），每次申请仅原子递增偏移量。前向传播完成后，直接调用 reset() 将游标归零，彻底不走任何细粒度的 free 流程！
中等定长张量（Size-Class Bins）：
将内存划分为离散的阶梯（如 16KB、64KB、256KB、1MB）。每个阶梯维护一个独立的连续 Block 数组。相同尺寸的对象永远挤在同一个物理页面中，完全杜绝了大小对象交错产生的空洞；
超大权重与 KV Cache（Direct mmap）：
对于超过 4MB 的巨型张量，直接绕过分配器，调用 mmap(MAP_ANONYMOUS) 独立向内核申请。用完后立即调用 munmap 或 madvise(MADV_DONTNEED) 瞬间将物理 RAM 释放给操作系统。

3. Rust 最小可用 Tensor Arena 分配器实现

use std::alloc::{alloc, dealloc, Layout};
use std::sync::atomic::{AtomicUsize, Ordering};

pub struct TensorArena {
    raw_ptr: *mut u8,
    capacity: usize,
    offset: AtomicUsize,
}

impl TensorArena {
    pub fn new(capacity_bytes: usize) -> Self {
        // 强制按 64 字节（Cache Line 对齐）申请大连续物理内存
        let layout = Layout::from_size_align(capacity_bytes, 64).unwrap();
        let raw_ptr = unsafe { alloc(layout) };
        if raw_ptr.is_null() {
            panic!("Failed to preallocate Tensor Arena memory");
        }

        Self {
            raw_ptr,
            capacity: capacity_bytes,
            offset: AtomicUsize::new(0),
        }
    }

    /// 纳秒级极速分配张量缓冲区（零锁、零系统调用）
    pub fn allocate_tensor(&self, num_elements: usize, elem_size: usize) -> Option<*mut u8> {
        let total_bytes = num_elements * elem_size;
        // 严格按 64 字节向上对齐
        let aligned_bytes = (total_bytes + 63) & !63;

        let current_offset = self.offset.fetch_add(aligned_bytes, Ordering::SeqCst);
        if current_offset + aligned_bytes <= self.capacity {
            unsafe { Some(self.raw_ptr.add(current_offset)) }
        } else {
            None // 显存/内存池溢出
        }
    }

    /// 一键 O(1) 重置全量内存，零析构遍历开销
    #[inline(always)]
    pub fn reset(&self) {
        self.offset.store(0, Ordering::Release);
    }
}

impl Drop for TensorArena {
    fn drop(&mut self) {
        let layout = Layout::from_size_align(self.capacity, 64).unwrap();
        unsafe { dealloc(self.raw_ptr, layout) };
    }
}