跳到主要内容

数据截至 (上游 commit d7a2074112d2)

01 · ggml 张量库:静态图与后端抽象

这一章讲什么: 整套栈的地基 ggml——张量怎么表示、计算为什么拆成「建图」和「执行」两段、同一张图怎么跑到 CPU 和十几种 GPU 上。读懂这一章,后面所有层都是它的上层建筑。

1. 它要解决的小问题​

在纯 C 里做张量计算,会遇到三个挠头的问题:

  1. 算一个 op 就 malloc 一次,推理热循环每秒几千次分配,性能被内存管理吃掉。
  2. 算子在哪台设备上跑是运行时才知道的事(用户可能把一半层放 GPU、一半留 CPU)。
  3. 每种硬件的 kernel 完全不同(x86 的 AVX、ARM 的 NEON、NVIDIA 的 CUDA),但上层只想写 mul_mat(a, b)。

ggml 的回答是:把「算什么」和「在哪算、什么时候算」彻底分开。先建一张静态计算图,再让一个了解所有设备的调度器去执行它。

2. 思路:先记账,后结账​

直觉来自一个观察——推理时每张图的形状都是固定的。batch 大小定了、模型定了,算哪些 op、每个张量多大,全能提前算出来。

那就不必边算边分配:

  • 建图时只创建描述结构(张量的类型、形状、stride、用什么 op、输入是谁),数据区一个字节不碰。
  • 执行前一次性把整个图需要的内存预留好,执行时零分配。
  • 同一张图可以算一万次——每次只是换输入数据,图本身原样复用。

官方头文件里的注释把这套用法写得明明白白:ggml_mul(ctx, x, x) 这类调用「不涉及任何实际计算」,计算只发生在显式调用 ggml_graph_compute_with_ctx 时(ggml/include/ggml.h:54-69 的注释块)。

3. 图示:两段式的数据流​

建图阶段(只记账,不算)                 执行阶段(结账)
─────────────────────────              ─────────────────────────

ggml_init(mem_buffer)  ──► ctx         ggml_build_forward_expand(gf, out)
      │                                     │
ggml_mul_mat(ctx, W, x) ──► t1          调度器 split_graph()
      │  (t1 只有 op=MUL_MAT、             │   按张量所在设备切成 N 段
      │   src=[W,x]、形状,无数据)         ▼
      ▼                                backend[i].graph_compute(split_i)
…… 挂几百个节点 ……                        │   CPU: 线程池逐节点算
      ▼                                     │   CUDA: 提交 kernel
ggml_cgraph gf (节点列表)                  ▼
                                       输出张量的 data 被填好

怎么读这张图: 左边全程没有浮点运算,产物是一张「图」;右边拿到图才真正动数据。两边的桥梁是 ggml_cgraph。

4. 张量:一个纯描述结构​

ggml_tensor 的全部字段(ggml/include/ggml.h:673-705):

字段含义
type数据类型(F32/F16/Q4_0/Q4_K…,枚举在 ggml/include/ggml.h:389)
ne[4] / nb[4]形状(最多 4 维,GGML_MAX_DIMS=4,ggml/include/ggml.h:222)/ 各维字节 stride
op这个张量是哪个算子的产物(ADD、MUL_MAT、ROPE…)
src[10]输入张量指针(最多 10 个,GGML_MAX_SRC)
data / buffer数据指针 / 它躺在哪个后端 buffer 里
view_src / view_offs若是视图,指向底张量 + 偏移
name[64]名字(调试用,如 blk.12.attn_k.weight)

注意两点:

  • 结构里没有所有权语义——它既不 new 也不 free 数据,只是个「标签」。整个结构定长成 GGML_TENSOR_SIZE(ggml/include/ggml.h:707),可以成批放在数组里。
  • 维度上限写死为 4。这不是偷工减料:LLM 推理里的张量到 4 维足够,写死让 stride 计算、kernel 循环都变成编译期展开的定长逻辑。

创建张量时不分配数据——ggml_init 可以传 no_alloc = true(ggml/include/ggml.h:665-669 的 ggml_init_params),先把整图的内存需求「干跑」一遍算出来,再统一分配。头文件注释直白地说明了动机:「图定义一次、计算多次……避免运行时的内存分配开销」(ggml/include/ggml.h:87-90)。

5. 计算图:一个节点数组​

图本身朴素得惊人(ggml/src/ggml-impl.h:329-348):

struct ggml_cgraph {
    int size;    // 容量
    int n_nodes; // 当前节点数
    struct ggml_tensor ** nodes;  // 按依赖序排好的算子
    struct ggml_tensor ** leafs;  // 常量/输入
    struct ggml_hash_set visited_hash_set; // 建图去重
    // ……(grad 相关字段略)
};

所谓「建图」就是把输出张量传给 ggml_build_forward_expand,它顺着 src[] 指针做一次拓扑遍历,把途经的 op 节点按序填进 nodes[]。执行就是按下标从头走到尾,逐节点调用对应 kernel——没有图优化器、没有 JIT,刻意保持简单。优化发生在别处:上层(llama-graph)在建图时就把节点按「减少设备切换」的顺序挂好(注释见 src/llama-graph.cpp:2716-2717)。

6. 原理演示:一张图的生命周期​

下面这段示意代码演示「先记账后结账」的完整循环:

# 示意,非源码 —— 演示 ggml 的两段式用法(对应 ggml/include/ggml.h:35-69 的 C 示例)
ctx = ggml_init(mem_size=1 << 20)      # 一次性给一块内存,之后零 malloc

x = ggml_new_tensor_1d(ctx, F32, 1)    # 只创建描述结构
a = ggml_new_tensor_1d(ctx, F32, 1)
b = ggml_new_tensor_1d(ctx, F32, 1)

x2 = ggml_mul(ctx, x, x)               # 不算!只建一个 op=MUL 的节点
f  = ggml_add(ctx, ggml_mul(ctx, a, x2), b)   # f = a*x² + b,仍只是图

gf = ggml_new_graph(ctx)
ggml_build_forward_expand(gf, f)       # 拓扑展开:f 依赖谁,全部入图

for value in [2.0, 3.0, 4.0]:          # 同一张图复用多次
    ggml_set_f32(x, value)             # 换输入数据
    ggml_graph_compute(gf, n_threads)  # 这次才真的算

重点看: 循环体内没有建图、没有分配——这在 llama 库里对应「每张 batch 重建一次图,但图内存通过 llm_graph_result 复用」(src/llama-context.cpp:601-602)。

7. 真实实现:执行一张图(CPU 后端)​

CPU 侧的入口是 ggml_graph_compute(ggml/src/ggml-cpu/ggml-cpu.c:3356)。它做的事:

  1. 拿出(或临时创建)线程池 ggml_threadpool(结构在 ggml/src/ggml-cpu/ggml-cpu.c:480)。
  2. 主线程自己也是 worker:ggml_graph_compute_thread(&threadpool->workers[0])(ggml/src/ggml-cpu/ggml-cpu.c:3418 一带)。
  3. 每个 worker 沿 nodes[] 顺序认领节点,逐节点跑 ggml_compute_forward_*;节点间用 barrier 同步(ggml_barrier,ggml/src/ggml-cpu/ggml-cpu.c:575)。

以最重的算子 mul_mat 为例,ggml_compute_forward_mul_mat(ggml/src/ggml-cpu/ggml-cpu.c:1254)先按 type_traits_cpu[src0->type] 查出这个量化类型配对的 vec_dot 内核(ggml/src/ggml-cpu/ggml-cpu.c:1271-1273),再把工作切成 chunk 让线程用原子自增抢任务(atomic_fetch_add_explicit(&params->threadpool->current_chunk, ...),ggml/src/ggml-cpu/ggml-cpu.c:1450)——没有静态分工,谁闲谁拿,天然适应大小核。

8. 后端抽象:四层接口,十几种硬件​

「同一张图跑在任何设备上」靠的是一套函数指针接口,分四层(ggml/src/ggml-backend-impl.h):

层结构职责定义处
reg(注册项)ggml_backend_reg_i一类硬件(如 CUDA):枚举设备、查自定义函数ggml/src/ggml-backend-impl.h:215
device(设备)ggml_backend_device_i一块具体的卡:名字、显存、创建 backend/buffer 类型ggml/src/ggml-backend-impl.h:161
backend(执行流)ggml_backend_i一个执行上下文:graph_compute、异步拷贝、同步、事件ggml/src/ggml-backend-impl.h:106
buffer(存储)ggml_backend_buffer_i / _type_i一块内存/显存:set_tensor/get_tensor、对齐、分配ggml/src/ggml-backend-impl.h:41、:17

每层都是「一组函数指针 + 一个 void * context」。backend 的核心只有一个必实现函数:

// ggml/src/ggml-backend-impl.h:131,struct ggml_backend_i 内
enum ggml_status (*graph_compute)(ggml_backend_t backend, struct ggml_cgraph * cgraph);

谁来决定图的哪段在哪跑? 调度器 ggml_backend_sched。它按「每个张量的权重存在哪个设备的 buffer 里」把图切成若干 split(ggml_backend_sched_split_graph,ggml/src/ggml-backend.cpp:1057),执行时逐 split 调对应 backend 的 graph_compute(ggml_backend_sched_graph_compute_async,ggml/src/ggml-backend.cpp:1963),段与段之间用 set_tensor/cpy_tensor 搬运数据。llama 的 context 初始化时为所有设备建一个调度器(src/llama-context.cpp:604)。

后端从哪来? CPU 永远内置;GPU 后端编译成动态库,启动时 ggml_backend_load_all() 扫描目录、逐个 dlopen(ggml/src/ggml-backend-reg.cpp:562、ggml/src/ggml-backend-dl.cpp)。所以官方一个发布包就能装上所有后端,运行时按机器有什么硬件加载什么。

9. 关键细节/坑​

  • 图节点数有上限,要预留。 ggml_cgraph.size 在建图前定死;llama 侧用 graph_max_nodes() 按 token 数估算(src/llama-context.cpp:597)。低估会在建图时断言失败。
  • 视图张量不拷数据。 view_src + view_offs 意味着 reshape/切片只是改标签;但 kernel 普遍要求输入连续,mul_mat 里有 GGML_ASSERT(nb00 == ggml_type_size(src0->type)) 这类检查(ggml/src/ggml-cpu/ggml-cpu.c:1282)——布局不对要在建图时插 ggml_cont。
  • 执行顺序就是数组顺序。 没有自动重排,想减少跨设备切换得建图时自己安排好(见 §5 末)。
  • 调度器切图有代价。 每多一个 split 就多一次设备间拷贝;llama-graph 的注释明确说「把 Q/K/V 节点连续挂图,就是为了减少 split 数」(src/llama-graph.cpp:2716-2717)。

10. 代码地图​

主题文件路径符号名
张量结构ggml/include/ggml.hggml_tensor、ggml_type、ggml_op
图结构ggml/src/ggml-impl.hggml_cgraph、ggml_build_forward_expand
建图算子入口ggml/src/ggml.cggml_mul_mat、ggml_rope_ext、ggml_soft_max_ext
CPU 图执行ggml/src/ggml-cpu/ggml-cpu.cggml_graph_compute、ggml_threadpool、ggml_barrier
矩阵乘(CPU)ggml/src/ggml-cpu/ggml-cpu.cggml_compute_forward_mul_mat、type_traits_cpu
后端接口ggml/src/ggml-backend-impl.hggml_backend_i、ggml_backend_device_i、ggml_backend_buffer_i
调度器ggml/src/ggml-backend.cppggml_backend_sched_split_graph、ggml_backend_sched_graph_compute_async
后端动态加载ggml/src/ggml-backend-reg.cppggml_backend_load_all、ggml_backend_load
图内存分配器ggml/src/ggml-alloc.cggml_gallocr_new、ggml_gallocr_alloc_graph