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 步:准备好你的厨师团队(初始化后端)

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

c 复制代码
#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 步:聘请厨房总管(创建调度器)

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

c 复制代码
// 创建一个调度器来管理这两个后端
// 最后一个参数 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 告诉调度器某些张量(通常是模型的权重)应该驻留在哪个后端。这是调度器分配任务的最重要依据。

c 复制代码
// ... 创建上下文 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 步:开火!(执行计算)

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

c 复制代码
// 计算图,让调度器去处理所有事情
ggml_backend_sched_graph_compute(sched, graph);

调用这个函数时,调度器内部会:

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

第 5 步:清理

任务完成后,记得释放调度器(它会释放其管理的所有资源)。

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

ggml_backend_sched 结构体

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

c 复制代码
// 来自 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.cppwisper.cpp 等项目核心代码的能力,也为你将来利用 ggml 构建自己的高效推理应用打下了坚实的基础。恭喜你!

相关推荐
拉姆哥的小屋5 小时前
VAE-NPN跨域室内定位的实战与思考
人工智能·毕设
IT_陈寒6 小时前
JavaScript性能优化:这7个V8引擎技巧让我的应用速度提升了50%
前端·人工智能·后端
奔跑吧邓邓子6 小时前
【C++实战(64)】C++ 邂逅SQLite3:数据库编程实战之旅
数据库·c++·sqlite·实战·sqlite3·数据库编程
拉姆哥的小屋6 小时前
突破传统!基于SAM架构的双模态图像分割:让AI“看见“红外与可见光的完美融合
人工智能·架构
会开花的二叉树6 小时前
RabbitMQ C++ 客户端封装与实战
c++·rabbitmq·ruby
AI数据皮皮侠8 小时前
中国上市公司数据(2000-2023年)
大数据·人工智能·python·深度学习·机器学习
我爱计算机视觉8 小时前
ICCV 2025 (Highlight) Being-VL:师夷长技,用NLP的BPE算法统一视觉语言模型
人工智能·算法·语言模型·自然语言处理
FunTester8 小时前
人工智能:技术分类、核心领域与应用全景
人工智能·语言模型·分类
xwz小王子9 小时前
首个零样本跨本体泛化开源具身模型:智源RoboBrain-X0 技术细节全解析
人工智能·团队开发
Vect__9 小时前
从直线到环形:解锁栈、队列背后的空间与效率平衡术
数据结构·c++