ggml介绍 (9) 后端调度器 (ggml_backend_sched)

在上一章 图分配器 (ggml_gallocr) 中，我们掌握了如何聘请一位"仓管员"，为单个后端的计算任务高效地规划和复用内存。这解决了在一个厨房（如 CPU）里高效烹饪的问题。

但是，现代厨房往往设备齐全，既有普通的炉灶（CPU），也有高功率的专业烤箱（GPU）。对于一道超级复杂的盛宴（一个大型语言模型），有些步骤在普通炉灶上处理更灵活，而有些大块的烘烤任务则必须交给专业烤箱才能又快又好。这时，我们就需要一位"厨房总管"来统筹全局，决定哪个步骤由哪位厨师在哪台设备上完成，并安排好食材在不同工位间的流转。

这位厨房总管，就是本章的主角------ggml_backend_sched，后端调度器。

什么是后端调度器？为什么需要它？

核心思想 ：把它想象成一个厨房总管。当厨房里有多个厨师（后端）时，总管会根据菜谱的复杂程度和厨师的专长，决定哪个步骤由哪个厨师来完成。ggml_backend_sched 负责管理多个后端，它分析整个计算图，将不同的计算任务分配给最合适的硬件（例如，GPU处理大型矩阵乘法，CPU处理其他操作），并管理好食材在不同厨师之间的传递（数据拷贝），以最高效率完成整道大餐。

在实际应用中，我们经常遇到这样的场景：

• 模型太大：一个大模型无法完全装入 GPU 的显存中。
• 混合计算更优：某些操作在 CPU 上执行反而比在 GPU 上更快（比如一些小的、串行的操作），或者某些操作 GPU 不支持。

ggml_backend_sched 就是为了解决这个问题而生的。它允许你同时使用多个后端（比如一个 CUDA 后端和一个 CPU 后端），并自动地将一个大的计算图"分裂"成多个可以在特定后端上运行的子图。

后端调度器如何工作？

调度器的工作核心是图分裂 (Graph Splitting)。当它拿到一个完整的计算图时，会像一个经验丰富的总管一样进行分析和规划：

1. 分析菜谱 (分析计算图)：它会检查图中的每一个计算节点（操作）。
2. 分配任务 (分配后端)：根据预设的规则（比如，用户可以指定"这个权重张量在 GPU 上"，那么使用这个权重的计算就优先在 GPU 上进行），它会为图中的每个节点初步确定一个最适合的后端。
3. 划分工序 (分裂图)：当它发现相邻的两个计算节点被分配给了不同的后端时（比如，一个在 GPU，下一个在 CPU），它就在这里"切一刀"，形成一个"分裂点"。
4. 安排流转 (处理数据拷贝)：在每个分裂点，调度器都明白，上一个子图的结果（比如 GPU 计算完的中间张量）需要被拷贝到下一个子图所需的后端缓冲区中（比如从 VRAM 拷贝到 RAM）。它会自动处理这些数据传输。

最终，一个大的计算图被分解成了一系列按顺序执行的子任务，每个子任务都在最合适的硬件上运行。

graph TD subgraph "原始计算图 (整个菜谱)" A[输入] --> B(CPU 操作1) --> C(GPU 操作2) --> D(GPU 操作3) --> E(CPU 操作4) --> F[输出] end subgraph "调度器的工作流程" G["分析图，分配后端"] end subgraph "分裂后的执行计划" direction LR subgraph "子图 1 (在 CPU 上执行)" B1("CPU 操作1") end subgraph "子图 2 (在 GPU 上执行)" C1("GPU 操作2") --> D1("GPU 操作3") end subgraph "子图 3 (在 CPU 上执行)" E1("CPU 操作4") end B1 --> |数据从 RAM 拷贝到 VRAM| C1 D1 --> |数据从 VRAM 拷贝到 RAM| E1 end A --> G G --> B1

如图所示，调度器将一个混合了 CPU 和 GPU 操作的图，巧妙地分成了三个子图，并自动管理了中间的数据拷贝。

动手实践：让 CPU 和 GPU 协同工作

让我们来看一下如何使用 ggml_backend_sched 来执行一个部分层在 GPU、部分在 CPU 的计算图。

第 1 步：准备好你的厨师团队（初始化后端）

首先，我们需要初始化所有要参与工作的后端。

#include "ggml-backend.h"
#include "ggml-cuda.h" // 假设使用 CUDA

// 初始化 GPU 后端
ggml_backend_t backend_gpu = ggml_backend_cuda_init(0);
// 初始化 CPU 后端
ggml_backend_t backend_cpu = ggml_backend_cpu_init();

if (!backend_gpu || !backend_cpu) { /* 错误处理 */ }

// 将后端实例放入一个数组中，优先级从高到低 (GPU > CPU)
ggml_backend_t backends[] = { backend_gpu, backend_cpu };

注意：后端的顺序很重要！调度器会优先尝试将任务分配给列表中靠前的后端（这里是 GPU）。CPU 通常作为最后的"接盘侠"，放在数组末尾。

第 2 步：聘请厨房总管（创建调度器）

有了后端团队，我们就可以创建调度器了。

// 创建一个调度器来管理这两个后端
// 最后一个参数 true 表示开启自动权重卸载
ggml_backend_sched_t sched = ggml_backend_sched_new(
    backends,  // 后端数组
    NULL,      // 缓冲区类型 (NULL 表示使用各后端默认值)
    2,         // 后端数量
    GGML_DEFAULT_GRAPH_SIZE, // 预估的图大小
    true,      // 是否并行 (高级用法，暂设为 true)
    true       // 是否允许自动卸载操作
);

sched 现在就是我们的厨房总管，准备好接收计算任务了。

第 3 步：定义菜谱并指定关键食材位置

在构建计算图时，我们可以通过 ggml_backend_sched_set_tensor_backend 告诉调度器某些张量（通常是模型的权重）应该驻留在哪个后端。这是调度器分配任务的最重要依据。

// ... 创建上下文 ctx 和输入张量 ...

// 假设 weight_gpu 是一个很大的权重，我们希望它在 GPU 上
struct ggml_tensor * weight_gpu = ggml_new_tensor_...;
ggml_backend_sched_set_tensor_backend(sched, weight_gpu, backend_gpu);

// res1 = input * weight_gpu
struct ggml_tensor * res1 = ggml_mul_mat(ctx, weight_gpu, input);

// 假设 weight_cpu 是一个小权重，或者是一个 CPU 更擅长处理的操作
struct ggml_tensor * weight_cpu = ggml_new_tensor_...;
ggml_backend_sched_set_tensor_backend(sched, weight_cpu, backend_cpu);

// res2 = res1 + weight_cpu
struct ggml_tensor * res2 = ggml_add(ctx, res1, weight_cpu);

// ... 构建完整的计算图 graph ...
struct ggml_cgraph * graph = ...;

通过这两次 set_tensor_backend 调用，我们给了调度器明确的指令："凡是用到 weight_gpu 的计算，请尽量交给 GPU 厨师；凡是用到 weight_cpu 的，请交给 CPU 厨师。"

第 4 步：开火！(执行计算)

现在，我们只需要一个简单的调用，剩下的所有复杂工作都由调度器完成。

// 计算图，让调度器去处理所有事情
ggml_backend_sched_graph_compute(sched, graph);

调用这个函数时，调度器内部会：

1. 自动调用其内部的图分配器来为 GPU 和 CPU 分别规划和分配内存。
2. 执行图分裂。
3. 按顺序执行每个子图，并在子图之间拷贝数据。

第 5 步：清理

任务完成后，记得释放调度器（它会释放其管理的所有资源）。

ggml_backend_sched_free(sched);
ggml_backend_free(backend_gpu);
ggml_backend_free(backend_cpu);
// ... 其他清理 ...

深入幕后：调度器的决策过程

调度器的核心逻辑在 ggml_backend_sched_split_graph 函数中（位于 src/ggml-backend.cpp）。这是一个多趟（multi-pass）处理过程，我们可以简化理解为：

sequenceDiagram participant User as 用户代码 participant Sched as ggml_backend_sched participant Graph as 计算图 User->>Sched: 调用 ggml_backend_sched_graph_compute(sched, graph) Note over Sched: (如果需要) 内部调用 alloc_graph Sched->>Sched: **Pass 1: 初步分配**<br>根据用户指定的张量后端(set_tensor_backend)<br>和权重位置，为节点分配初始后端。 Sched->>Sched: **Pass 2: 扩展分配**<br>从GPU节点开始，向上下游<br>扩展分配，把兼容的操作也拉到GPU上。 Sched->>Sched: **Pass 3: 最终分配**<br>处理剩余未分配的节点，<br>选择最优的后端。 Sched->>Sched: **Pass 4: 分裂图**<br>遍历图，在后端切换处<br>插入"分裂点"并记录需要拷贝的张量。 Sched->>Sched: **执行子图**<br>循环遍历分裂出的子图... loop 每个子图 Sched->>Sched: 拷贝输入数据到当前子图的后端 Sched->>Graph: 在对应后端上执行子图 end Sched-->>User: 计算完成

ggml_backend_sched 结构体

这个结构体是调度器的"大脑"，存储了所有管理信息：

// 来自 src/ggml-backend.cpp 的简化版结构
struct ggml_backend_sched {
    int n_backends;
    ggml_backend_t backends[GGML_SCHED_MAX_BACKENDS]; // 后端列表
    ggml_gallocr_t galloc; // 内置的图分配器

    // 用于记录图中每个张量被分配到哪个后端的哈希表
    struct ggml_hash_set  hash_set;
    int                 * hv_tensor_backend_ids;

    // 分裂后的子图信息
    struct ggml_backend_sched_split * splits;
    int n_splits;

    // ... 其他用于并行和调试的字段 ...
};

• backends: 存储了我们传入的后端列表。
• galloc: 调度器拥有自己的图分配器实例，用于为所有后端统一管理内存。
• hv_tensor_backend_ids: 一个哈希表，是决策的核心，记录了 (ggml_tensor*) -> backend_id 的映射。
• splits: 一个数组，存储了所有分裂出的子图的信息，包括每个子图的后端、起止节点，以及需要从外部拷贝的输入。

ggml_backend_sched 真正地将我们之前学到的所有核心概念------张量、计算图、后端、后端缓冲区和图分配器------完美地粘合在一起，形成了一个强大而灵活的计算执行引擎。

总结

在本章中，我们认识了 ggml 的"厨房总管"------后端调度器 (ggml_backend_sched)。

• 它负责管理和协调多个后端（如 CPU 和 GPU）共同完成一个计算任务。
• 它的核心能力是图分裂，即将一个大计算图分解为可以在不同硬件上执行的多个子图。
• 它会自动处理子图之间的数据依赖和拷贝。
• 我们通过向调度器注册后端列表，并使用 ggml_backend_sched_set_tensor_backend 提供提示，来引导它的决策过程。
• ggml_backend_sched_graph_compute 是一个高度封装的接口，它为我们处理了所有底层的复杂性，包括内存分配、图分裂和执行。

至此，我们已经完成了 ggml 核心概念的探险之旅！从最基础的张量到最顶层的调度器，你已经掌握了 ggml 工作的基本原理。现在，你已经具备了阅读和理解 llama.cpp、wisper.cpp 等项目核心代码的能力，也为你将来利用 ggml 构建自己的高效推理应用打下了坚实的基础。恭喜你！