技术博客 | 使用 CUDA 图形优化 Llama.cpp AI 推理

开源 llama.cpp 代码库最初于 2023 年发布，是一种轻量级但高效的框架，用于在 Meta Llama 模型上执行推理。llama.cpp 基于去年发布的 GGML 库构建，由于专注于 C/C++ 而无需复杂的依赖项，因此很快就吸引了许多用户和开发者（尤其是在个人工作站上使用）。

自首次发布以来，Llama.cpp 已得到扩展，不仅支持各种模型、量化等，还支持多个后端，包括支持 NVIDIA CUDA 的 GPU。在撰写本文之时， Llama.cpp 在所有 GitHub 库中排名第 123 位 ，在所有 C++ GitHub 库中排名第 11 位。

在 NVIDIA GPU 上使用 Llama.cpp 执行 AI 推理已经带来了显著的优势，因为它们能够以极高的性能和能效执行基础 AI 推理的计算，同时在消费设备和数据中心中也很普遍。NVIDIA 和 Llama.cpp 开发者社区继续合作，以进一步提高性能。本文介绍了最近通过在 Llama.cpp 中引入 CUDA Graphs 功能而实现的改进。

GPU 会随着每一代产品的推出而不断加速，而且通常情况下，GPU 上的每个事件（例如内核或内存复制）都会很快完成。过去，每个事件都必须由 CPU 单独调度（启动），相关的开销可能会累积成为性能瓶颈。

CUDA Graphs 工具通过将多个 GPU 活动调度为单个计算图形来解决此问题。在上一篇文章“ Getting Started with CUDA Graphs ”中，我介绍了 CUDA Graphs 并演示了入门方法。在后续的博文“ A Guide to CUDA Graphs in GROMACS 2023 ”中，我将介绍如何将 CUDA Graphs 成功应用于 GROMACS 生物分子模拟科学软件包。

使用传统流模型时，每个 GPU 活动都是单独调度的，而 CUDA Graphs 支持对多个 GPU 活动进行一致调度。这减少了调度开销。调整现有基于流的代码以使用图形相对简单。该功能通过几次额外的 CUDA API 调用将流执行“捕获”到图形中。

本文将介绍如何利用此工具，使用图形（而非流）来执行预先存在的 llama.cpp 代码。

本节重点介绍现有代码中的开销，并描述如何引入 CUDA Graphs 来减少这些开销。

现有代码中的开销

图 1 显示了在引入 CUDA Graphs 之前，使用 Linux 在 NVIDIA A100 GPU 上执行 Llama 7B Q4 推理的现有代码的配置文件片段。它是通过 NVIDIA Nsight Systems 获得的。图中顶部配置文件中的每个 GPU 活动块都对应于对单个令牌的评估，其中缩放设置为显示正在评估的两个完整令牌。可以看到，每个令牌的评估（对应于与采样和计算图形准备相关的 CPU 活动）之间存在间隙。（我将在博文结束后返回到这一点）

图 1： 通过分析时间轴中 GPU 处于空闲状态的部分，观察到与 GPU 推理中涉及的 GPU 活动相关的开销

图 1 底部显示了相同的配置文件，但放大以显示令牌评估中的几个活动。令牌评估中的内核之间存在差距是由于启动开销引起的。正如本文所展示的，使用 CUDA Graphs 消除这些开销可显著提高性能。突出显示的事件是从 CPU 启动内核（左下角），并在 GPU 上执行相应的内核（右上角）：CPU 能够在 GPU 上执行之前很长时间成功启动。

因此，此处的 CPU 端启动开销并不在关键路径上，而是由于每次内核启动相关的 GPU 端活动而产生的开销。此行为可能因不同的模型和硬件而异，但 CUDA Graphs 适用于减少 CPU 和/或 GPU 启动开销。

Llama.cpp 已使用 GGML 格式的“图形”概念。每个令牌的生成涉及以下步骤：

• 根据所用模型准备 GGML 图形结构。

• 评估正在使用的后端结构（在本例中为 NVIDIA GPU），以获取“logits”，记录下一个令牌在词汇表中的对数概率分布。

• 在 CPU 上执行采样，使用 Logits 从词汇表中选择令牌。

通过截取 GPU 图形评估阶段引入了 CUDA Graphs。添加了代码以将现有流捕获到图形中，将捕获的图形实例化为可执行图形，并将其启动到 GPU 以执行单个令牌的评估。

为了获得适当的效率，必须多次重复使用相同的图形；否则，新引入的捕获和实例化所涉及的开销将大于收益。但是，图形会随着推理的进行而动态演变。面临的挑战是开发一种机制，以跨令牌对图形进行微调（低开销），从而实现整体收益。

随着推理的进行，操作长度会随着上下文大小的增加而增加，从而导致计算图形发生实质性的（但不频繁）更改。GGML 图形会接受检查，并仅在需要时重新捕获。 cudaGraphExecUpdate 用于更新先前实例化的可执行文件图形，其开销远远低于完全重新实例化。

计算图形也有频繁但非常微小的更改，其中每个令牌的某些节点（与 KV 缓存相关）的内核参数会发生变化。NVIDIA 开发了一种机制，仅在可重复使用的 CUDA Graphs 中更新这些参数。在启动每个图形之前，我们利用 CUDA Graphs API 功能来识别图形中需要更新的部分，并手动替换相关参数。

请注意，CUDA Graphs 目前仅限于批量大小为 1 的推理（Llama.cpp 的关键用例），并计划针对更大的批量大小开展进一步的工作。有关这些进展以及为解决问题和限制而正在进行的工作的更多信息，请参阅 GitHub 问题、 NVIDIA 为在 Llama.cpp 中使用 CUDA Graphs 而进行的新优化 ，以及此处链接的拉取请求。

在引入 CUDA Graphs 之前，由于 GPU 端启动开销，内核之间存在显著差距，如图 1 底部配置文件所示。

图 2 显示了与 CUDA Graphs 的等效关系。所有内核都作为同一计算图形的一部分（通过单个 CUDA API 启动调用）提交给 GPU。这极大地减少了开销：图形中每个内核之间的差距现在非常小。

图 3 显示了 llama.cpp 中新的 CUDA Graphs 功能的优势。测量到的加速因模型大小和 GPU 变体而异，随着模型大小的减少和 GPU 功能的增加，收益也在增加。这符合预期，因为使用 CUDA Graph 可减少与快速 GPU 上的小问题相关的开销。在速度最快的 NVIDIA H100 GPU 上，最小的 Llama 7B 模型实现的最高加速是 1.2 倍。所有结果都使用 Linux 系统。

在 Llama.cpp 主分支中，NVIDIA GPU 上的批量大小为 1 的推理现在默认启用 CUDA Graphs。

图 3： 使用 CUDA Graphs 相比传统流实现的加速，对于多个不同大小的 Llama 模型（均为批量大小 1）和多个 NVIDIA GPU 变体的结果。

图 1 中的顶部配置文件显示了时间轴中令牌评估之间的差距（GPU 处于空闲状态）。这些差距是与准备 GGML 图形和采样相关的 CPU 活动造成的。正如此 GitHub 问题 及其链接的拉取请求中所述，减少这些开销的工作处于高级阶段。这项工作预计将提供高达 10% 的进一步改进。

在本文中，我展示了在热门的 Llama.cpp 代码库中引入 CUDA Graph 如何显著提高了 NVIDIA GPU 上的 AI 推理性能，并且正在进行的工作有望实现进一步的增强。要将这项工作用于您自己的 AI 支持的工作流， 请按照使用说明进行操作 。