ggml_cgraph 的内存分配器:计算期零动态分配实现

封面信息图

在许多高性能 AI 运行时中,内存管理策略往往决定了推理底座在边缘端与多并发场景下的生死存亡。如果一个推理引擎在前向传播过程中,每执行一个算子(如一次矩阵乘或一次 LayerNorm)都要调用操作系统的 malloc 申请几兆字节的中间激活值(Activation Buffer),运行期的内存碎片、系统调用开销以及多线程争锁会瞬间将性能拖垮。

llama.cpp 底层的 GGML 库之所以能在各类极小算力设备上表现得如同一台精密的机械表,其核心秘密在于其为计算图量身定制的 静态内存规划器(ggml-alloc

它在模型真正跑第一步计算之前,就能完成整张计算图的内存生命周期分析,实现计算期零动态分配与内存最大化复用

+--------------------------------------------------------------------------+
|                     ggml-alloc 静态内存规划生命周期                         |
+--------------------------------------------------------------------------+
| 1. 计算图拓扑构建 (ggml_build_forward)                                      |
|    -> 确定 100+ 个节点的拓扑执行顺序序列                                     |
+--------------------------------------------------------------------------+
                                    |
                                    v
| 2. 离线生命周期分析 (Liveness Analysis)                                    |
|    -> 标记每个 Tensor 首次创建 (First-use) 与最终消费完毕 (Last-use) 时间点   |
+--------------------------------------------------------------------------+
                                    |
                                    v
| 3. 图着色与内存重叠复用 (Memory Arena Allocation)                          |
|    -> 生命期互不重叠的中间激活值共享同一块物理内存 Offset                      |
+--------------------------------------------------------------------------+
                                    |
                                    v
| 4. 固化物理指针,进入推理循环                                               |
|    -> 推理期间零 malloc, 零 free, 内存水位一条直线                           |
+--------------------------------------------------------------------------+

1. 激活值张量的生命周期(Liveness)分析

在深度学习前向推理中,大部分中间张量是“短命”的:例如第 3 层的注意力得分张量,在完成与 Value 矩阵相乘后就再也不会被后续算子使用。如果为每个中间张量都分配独立的显存空间,7B 模型的中间激活值可能需要消耗数 GB 显存。

ggml-alloc 在离线阶段,对 cgraph->nodes 数组从头到尾进行一次扫描:

  • 为每个 Tensor 记录其出生的拓扑下标 born_idx
  • 逆向扫描所有算子的输入依赖,记录每个 Tensor 被最后一次读取的拓扑下标 died_idx
  • 区间 [born_idx, died_idx] 即为该张量的活跃生命周期区间

2. 空间复用算法:最佳适应(Best-Fit)与内存着色

有了生命周期区间后,内存规划问题就转化为了经典的二维矩形装箱问题(时间轴 $\times$ 空间大小)。

ggml-alloc 维护一个虚拟的连续内存池(Virtual Arena):

  1. 按照拓扑顺序依次遍历每个算子节点;
  2. 当需要为一个新生成的 Tensor 分配空间时,在当前空闲的物理块列表中,寻找一个尺寸足够大且起始地址满足硬件对齐(如 32 字节对齐)的最小空闲块(Best-Fit)
  3. 将该物理块的偏移量(Offset)绑定到该 Tensor 的 data 指针上;
  4. 当某个 Tensor 的生命周期到达 died_idx 时,规划器将其占用的内存块标记为重新释放,立即归还给空闲池,供给后续的算子复用。
// 伪代码:生命周期复用规划
void ggml_alloc_graph(struct ggml_gallocr * galloc, struct ggml_cgraph * gf) {
    // 遍历所有前向节点
    for (int i = 0; i < gf->n_nodes; i++) {
        struct ggml_tensor * node = gf->nodes[i];

        // 1. 如果该节点需要存储空间且尚未分配,从池中获取最佳复用块
        if (node->data == NULL) {
            node->data = allocate_best_fit_block(galloc, ggml_nbytes(node));
        }

        // 2. 检查是否有父节点的生命周期在当前节点终结,若是则就地释放回池
        for (int j = 0; j < GGML_MAX_SRC; j++) {
            struct ggml_tensor * parent = node->src[j];
            if (parent && is_last_user(parent, gf, i)) {
                free_block_back_to_pool(galloc, parent->data);
            }
        }
    }
}

通过这种贪心复用算法,原本需要数 GB 的中间激活值空间,被压缩到了仅需 几百 MB(即网络中单层激活值峰值所需的最大内存),显存开销直接暴降 70% ~ 85%

3. 计算期零开销的工程红利

当这套静态规划完成后,所有的 Tensor 都已经被赋予了固定的物理内存绝对地址。

在实际推理循环中:

  • CPU 指令执行极其平滑:没有运行时的动态内存分配器介入,消除了一切系统调用(brk / mmap)开销;
  • 多线程无锁并发:Worker 线程在执行各个算子时,直接读写预定好的内存偏移,完全不需要在内存分配器上争抢互斥锁;
  • 内存水位绝对可预测:系统在启动后的第一秒,内存占用就定格为一条笔直的物理水平线,绝不会在运行数小时后因为内存碎片而发生意外 OOM。

这种将动态不确定性在编译期彻底消解的工程设计,正是系统级编程中最动人心魄的硬核之美。

Logo

openEuler 是由开放原子开源基金会孵化的全场景开源操作系统项目,面向数字基础设施四大核心场景(服务器、云计算、边缘计算、嵌入式),全面支持 ARM、x86、RISC-V、loongArch、PowerPC、SW-64 等多样性计算架构

更多推荐