BatchLLM：离线大 Batch 场景 LLM 推理优化

最新推荐文章于 2025-04-18 20:15:21 发布

程序猿李巡天

最新推荐文章于 2025-04-18 20:15:21 发布

阅读量1.5k

点赞数 19

文章标签： batch 开发语言金融语言模型人工智能 ui 音视频

本文链接：https://blog.csdn.net/m0_59235945/article/details/144384923

版权

一、背景

我们之前已经介绍过很多针对 LLM 在线推理的文章，而最近针对离线 LLM 推理的应用也越来越多，比如我们前段时间介绍过的 Best-of-N 优化方案，本文中我们继续介绍一个针对 Batch Inference 场景的新工作：BatchLLM。

二、摘要

LLM 在各类信息处理等任务中的应用场景越来越多。许多此类任务以大规模批处理或离线方式执行，相应更关注的性能指标为吞吐量。这些任务通常存在前缀共享（Prefix Sharing）的特性，即不同 Prompt 输入间存在部分共同前缀（Token）。然而，现有 LLM 推理引擎主要针对在线推理，更多聚焦在优化流式请求，对具备前缀共享特性的大批量任务支持有限。

当前的解决方案主要采用基于 LRU 的缓存机制，以复用请求间的共同前缀的 KV Cache。然而，隐式 Cache 管理可能导致即将复用的 KV Cache 过早被驱逐。即便未被驱逐，因为共享同一上下文的 Request 并没有同步调度，导致共享 KV Cache 的生命周期被延长，进而导致内存占用增加。
这些面向流式的系统通常按先到先服务或类似顺序方式调度 Request，致使 Decoding 步骤占比较高的 Request 可能因调度过晚，无法与 Prefill Block 混合，从而影响硬件利用率。
此外，基于 Token 和 Request 数量的 Batching 机制限制了 Token Batch 规模，使得 GPU 在以 Decoding Token 为主的迭代中无法达到饱和状态。

针对以上问题，作者提出了 BatchLLM。BatchLLM 通过全局显式识别共同前缀，将共享同一前缀的 Request 同步调度，以最佳方式复用 KV Cache，同时缩短共同 KV 内存的生命周期。此外，BatchLLM 重新排序 Request，优先调度 Decoding 步骤占比较高的 Request，以更好地将 Decoding Token 与后续 Prefill Block 混合，并采用以内存为中心的 Token Batchiing 策略，扩大 Token Batch 规模，从而提升 GPU 利用率。最后，BatchLLM 还通过融合优化前缀共享 Attention Kernel，减少长尾效应与 Kernel 启动开销。

三、引言

3.1 LLM 推理过程

当前主流 LLM 基本都是 Decoder Only 的 Transformer 模型，其推理过程可以分为两个关键阶段，如下图 Fig.2 所示（图片来自 [2404.14294] A Survey on Efficient Inference for Large Language Models [2]）：

Prefill：根据输入的多个 Token（I，like，natural，language）生成第一个输出 Token（Processing），通过一次 Forward 就可以完成。在 Forward 中，输入 Token 间可以并行执行（类似 Bert 这些 Encoder 模型），执行效率很高。
Decoding：从生成第一个 Token（Processing）之后开始，采用自回归方式一次生成一个 Token，直到满足终止条件。假设输出总共有 N 个 Token，则 Decoding 阶段需要执行 N-1 次 Forward，这 N-1 次 Forward 只能串行执行，效率很低。

如下图 Fig.3 所示为生成过程中 Latency 和 Memory 的关系，显存的增量与序列的长度成正比：

在红色的 Prefill 阶段，可以并行计算，Latency 相对比较小，会一次生成所有输入 Token 的 KV Cache。
在绿色的 Decoding 阶段，不能并行计算，Latency 比较高，显存增加比较慢，但如果生成长度较长，依然会增加很多。678
黄色为 Decoding 阶段生成一个 Token 的 Latency，需要说明的是，随着序列的变长，这个 Latency 会逐渐增加。
当一个序列处理完之后，会释放相应的 KV Cahce；此外，模型参数在不同序列间是共享的，而 KV Cache 则不是（如果没有公共 Prefix）。

3.2 Prefix Sharing

Prefix Sharing 是一种非常常见的 LLM 推理优化手段，得益于 LLM 推理的 Causal Mask 特性，每一个 Token 只能关注到该 Token 之前的 Token。对于不同的 Request，如果它们有一个公共的 Prefix，那么 Prefix 部分的 Token 只用计算一次即可以被其他 Request 用于共享 KV Cache，如下图 Figure 9 所示为 [2312.07104] SGLang: Efficient Execution of Structured Language Model Programs [3] 中的一些示例：

3.3 Chunked Prefill

微软团队在 [2308.16369] SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills [4] 中尝试将 Prefill 阶段的 Prompt 切分为不同的 Chunks，并将不同请求间的 Prompt Chunk 与 Decoding 结合，其主要是为了减少流水线并行（Pipeline Parallelism）推理中的 Bubble，如下图所示：

在 [2401.08671] DeepSpeed-FastGen: High-throughput Text Generation for LLMs via MII and DeepSpeed-Inference [5] 中，作者更进一步在每个 Prompt Chunk 中都与其他请求的 Decoding 结合，进一步提升吞吐，如下图 Figure 3 所示：

3.4 新兴场景

除了常见的聊天和在线问答场景外，LLM 也越来越多的被应用到大规模的信息处理等任务中，鉴于这些任务的流量巨大，而 LLM 推理成本又很高，很多任务的处理方式都趋向于大规模批处理（离线处理），其主要的性能指标是吞吐量而不是延迟。

以搜索场景的问答和 Snippet 生成任务为例，该任务根据文档和用户 Query，在搜索结果页面生成网页文档的片段，如下图所示为 [2310.14587] Large Search Model: Redefining Search Stack in the Era of LLMs [6] 中的示例：

由于在线执行大规模 LLM 推理时难以满足 SLO，常见的做法是对高频网页文档和 Query 进行离线 Snippet 生成，并在在线服务期间检索相应的离线结果。

许多基于 LLM 的信息处理任务中，不同 Request 往往是在相同的内容（比如网页）上执行，因此不同 Request 间表现出共享 Prefix 的特征。还是以搜索引擎为例，往往会出现针对同一网页文档中提取不同信息，使文档成为 Prompt 中的共享前缀。

共享前缀也可以在多个层级进行管理，比如，第一级可以是全局的 System Prompt，第二级可以是共享的文档。然而，随着 LLM 的发展及任务特定模型微调的应用，System Prompt 的前缀变的越来越短，最终前缀长度主要由共享的文档决定。

3.5 新的优化需求

当前常见的 LLM 推理框架已经提供了各种优化方案，然而，其针对日益重要的大批量场景却有一定的局限性，主要是 3 个方案：

3.5.1 隐式前缀缓存无法实现大批量推理中的最优 KV 复用

现有的 LLM 服务系统往往采用前缀树来维护 KV Block。相同的前缀可映射至相同的 KV Block，以避免重复计算。系统通常采用 LRU 策略从 Block 表中淘汰 Block。对于待处理的大批量输入，系统并未在全局视角下优化，而仅依据历史输入。因此，可共享的 KV Block 容易被过早淘汰。作者评估了 vLLM 的前缀共享性能，采用典型行业工作负载。通过分析输入数据集，前缀复用带来的最优 Token 节省率为 56%，而基于 vLLM 的隐式前缀缓存仅实现了 44.2% 的效率。需注意，节省率计算公式如下图所示：

鉴于批处理任务中所有 Prefill 特征均可提前获知，因此也就有机会提升大规模批处理任务中的前缀共享效率。

如下图所示为 vLLM 中 KV Cache Block 管理的一个示例，一行是一个 Block，Seq A 和 Seq B 可以共享 “The future of artificial” 这个 Block：

3.5.2. 针对在线的调度与 Token Batching 对于前缀共享和高吞吐任务可能不是最优

当前的 LLM 推理系统以 Request 为粒度进行任务调度，并未分析整个大批次中的前缀共享特性，也未将共享相同前缀的 Request 一并调度。这不仅可能导致可共享的 KV Cache 过早被驱逐，还延长了可共享 KV Cache 的生命周期，加剧了大规模 KV 内存消耗的问题。此外，这些系统设计倾向于支持在线服务，它们在每次迭代中根据 Request 的到达顺序形成 Token Batch，这可能导致 Token Batch 不是最优性能。

以下图 Figure 1 为例，按 Request 到达顺序进行的简单调度无法将 Request 3 的 Decoding 与其它 Request 1 的 Prefill Chunk 混合。相反，通过提前调度 Request 3，其 Decoding 可以与其它 Request 的 Prefill Chunk 混合。

另一个问题是，当前系统将 Token Batch 中的 Token 数量和 Request 数量作为限制 Batching 的阈值，在以 Decoding 为主的 Token Batch 中限制了 Batching 更多 Token 的机会，即使在内存充足的情况下也依然无法充分利用 GPU 算力。如下图 Figure 2 展示了一个典型任务在 vLLM 最佳配置下执行时每次迭代的 Token 数量。图中所示，在许多迭代中存在“低谷”，此时 Token 数量可能不足以使 GPU 达到饱和状态。

因此，可以重新安排任务，使得共享相同前缀的 Request 更为接近，并促进 Decoding Token 与 Prefill Chunk 的更好混合。此外，还有空间扩大由 Decoding Token 主导的 Token Batch 的规模。

3.5.3. 共享前缀的 Attention 优化仍有改进空间

现有工作中前缀共享 Attention 的实现方式通常是：通过计算不同 KV Block 上的 Attention，对公共前缀执行矩阵乘法，并对非共享块进行基本的矩阵向量乘法。然而，不同块的计算在不同的 Kernel 中进行。一方面，分离的 Kernel 增加了 GPU Kernel 的长尾效应；另一方面，多个 Kernel 的存在也带来了挑战。

四、方案

4.1 概览

如下图 Figure 3 展示了 BatchLLM 的概览，基于 vLLM 实现。

首先，对大批量输入 Prompt 进行分析，并识别出共同的前缀部分，这一步骤在 Request 调度之前完成。通过将多层次前缀简化为单一层次前缀，BatchLLM 洞察到信息处理等任务的前缀通常由长篇上下文（如网页文档）主导。
随后，将 Request 分组，每组对应于共享相同前缀的 Request。这些组将成为 Request 调度的基本单元。在调度各组之前，BatchLLM 还会根据前缀 Prefill 长度与估计的 Decoding 长度的比例对组进行重新排序，并推迟那些前缀 Prefill 较长而 Decoding 较短的组。
接着，充分考虑 KV 内存使用情况来实现 Batching，旨在更好地混合 Decoding 步骤与前缀 Prefill Chunk，以增大整体 Token Batch 规模。
此外，BatchLLM 还实现了融合 Attention Kernel，以减少长尾效应及 Kernel 启动开销。

4.2 显式全局 KV 复用标识

如下图 Figure 4 所示，BatchLLM 利用前缀树来识别和表示 Prompt 之间的共同前缀。然而，直接使用前缀树的第一层前缀作为共同前缀可能导致前缀复用效果不佳。以下图 Figure 4d 的 Prompt 为例，6 个 Prompt 按 Figure 4a 所示组成前缀树，其中每个节点代表若干连续 Token（节点中数字表示连续 Token 数），节点中的 Token 由子节点共享。

根据 Figure 4a，根结点的右子节点（‘b’）总共被 5 个 Prompt 共享（Prompt 2-6），因此复用一层前缀可以节约 4 个 Token 的 KV 计算。
然而，很容易推断出，在 Prompt 4 和 5 之间共享了前 8 个 Token，可以节约 8 个 Token 的 KV 计算。
因此，需要重构前缀树以最大化 Token 复用。

如上图 Figure 4b 和 4c 展示了上述的过程，作者设计了一种基于前缀树的动态规划算法，以最大化一级 Token 的复用，如下图 Algorithm 1 所示。其基本思路为：对于一个节点 D，为了最大化其一级 Token 的复用：

首先，解决所有子节点的子问题，以最大化每一个子节点的复用（第 6 行）。
然后，尝试将 D 的一级前缀与其子节点的前缀合并，以观察是否能扩大 B 的 Token 复用（第 7-11 行）。通过调用 MaximizeReuse(root)，递归执行上述过程，直到到达根结点。

如图 Figure 4b 中，树的第 3 层左侧节点被分叉并与其右侧子节点合并，从而扩大第二层右侧节点的一级前缀复用。类似的，第二层右侧节点被分叉并与其第二个子节点合并，以扩大根结点的第一级前缀复用。在最大化一级前缀后，其他级别的前缀将在 Token Batching 中扩展，不在被视为共享上下文。

作者对比原始多级前缀与扩展单级前缀的 Token 节省比例。对于 Snippet 生成任务数据集，原始多级前缀的节省比例为 56%，而扩展后的一级前缀的节省比例为 55%。这也就意味着，扩展后的一级前缀在 Token 复用仅损失 1% 的情况下，可以同时减少 Token Batching 和 Attention Kernel 优化的复杂性和开销。

4.3 面向吞吐量的 Token 批处理

4.3.1 基于前缀分组的调度策略

为了简化 KV 复用并缩短公共前缀在内存中的生命周期，BatchLLM 将共享相同前缀的请求集中调度。具体而言，一个前缀组对应一组共享相同前缀的 Request 的集合，然后以组的粒度进行调度，一旦分组完成，公共前缀的内存即可立即释放，避免内存被无谓占用。

为了实现上述调度，BatchLLM 维护了 3 个待 Batching 的 Token 队列：公共前缀队列，专有 Prompt（Request 中除了公共前缀之外的 Prompt）队列和 Decoding 队列。每次迭代中，BatchLLM 将按照 Decoding 队列、专有 Prompt 队列和公共前缀队列的顺序（一个约束条件为：公共前缀应在专有 Prompt 之前被处理）从 3 个队列中提取 Token 以形成 Token Batch。

一方面，先提取 Decoding Token 再提取 Prefill Token 有助于在 Token Batch 中更好地混合这两类 Token；
另一方面，先提取 Decoding Token 和专有 Prompt Token 再提取前缀 Token，可以确保活跃请求（已经被部分处理）在非活跃请求之前被调度，从而有助于提前完成活跃请求并尽可能释放其内存。

4.3.2 Request 重排序

BatchLLM 根据 Prompt 长度和 Decoding 长度（R）的比例对 Request 进行重新排序以调度请求，R 较大的请求将提前调度。这有助于更早地发出具有长 Decoding 步骤的请求，从而可以更好地与前面的 Decoding Token 混合以扩大 Token Batch 大小。然而，在完成请求之前，并不知道输出 Token 的确切数量。作者观察到一个任务的输出长度分布相对集中，因此假设输出长度是恒定的，可以按照下述公式进行归一化，然后 Rgroup 较大的 Request 组将被更优先调度：

4.3.3 以内存为核心的 Token Batching

上述重排可以减少 Figure 2 中的“低谷”。另一个限制 Token Batch 大小的因素是：当前很多工作采用固定 Request 数量来限制 Token Batch，即请求数量不能超过固定的阈值。然而，如果请求都是以 Decoding Token 为主，则会导致总的 Token 数量可能达不到较大规模，从而影响 GPU 利用率。比如，vLLM 默认的最大请求数为 256，若一个 Token Batch 已经超过 256 个 Decoding Token，即便还有更多内存容纳 Prefill Chunk 的 KV 内存，也无法添加更多 Prefill Chunk。然而直接扩大上述阈值也可能导致因 GPU 内存不足以容纳新的 KV Cache 而频繁进行内存交换和重计算。

4.4 前缀共享 Attention 优化

Token Batching 的目前是在 GPU 内存大小的限制下，增加 Batch 中 Token 数量，这个过程主要受到两个因素影响：

Batch 中的 Token 数量是否足够大，以使 GPU 达到包含状态。
剩余内存是否可以容纳更多的 Prefill Chunk，这决定了是否可以进一步优化。

基于此，BatchLLM 采用 KV 内存本身和预设的 Chunk 大小 Schunk 作为限制条件，而非使用 Request 数量间接限制。具体而言，BatchLLM 预设了 KV 内存的上限，记为 Mthreshold。在每次迭代组合 Batch 时，仅检查添加一个 Prefill Chunk 是否会超过 Mthreshold，而不考虑 Request 的数量。只要剩余的 GPU 内存容量允许，Prefill Chunk 就会被加入到 Token Batch 中。如下图 Algorithm 2 详细描述了这一过程。值得注意的是，作者同样还使用一个固定数值来限制每个 Batch 的 Token 数，与 vLLM 的做法一致，有助于在不同的迭代中更均匀的分配 Token 数量。

五、实验 & 评估

5.1 实验配置

相关代码基于 vLLM v0.5.4 实现，使用 FlashAttention v2.6.1 作为 vLLM 的 Attention 后端，并且使用 vLLM 的 Prefix-Caching 和 Chunked-Prefill 作为基线。对于 Chunked-Prefill，使用 2048 作为 Token Batch 大小。

其他设置如下所示：

自定义的 GPU Kernel 使用 OpenAI Triton 3.0.0 实现。
测试的硬件为 NVIDIA A100 80GB GPU 和 AMD MI200 GPU。
测试模型为 Qwen2-7B、Mistral-7B 和 LLaMA3-8B。
测试数据集包含 3 类，对应的公共前缀/专有 Prompt 分别为 2000/200，1000/1000，200/2000。
测试任务，主要是网页搜索中的两种典型场景：Snippet 生成任务和离线排序任务。

5.2 端到端比较

5.2.1 整体评估

如下图所示，其中 “+p” 表示打开 Prefix Caching，“+c” 表示打开 Chunked-Prefill，可以看出，BatchLLM 在几乎所有配置下都取得最优吞吐，尤其是也优于 vLLM + Prefix Caching + Chunked-Prefill：

5.2.2 案例 1：长公共前缀+短专有 Prompt

这里的行业负载测试类似于 Snippet 生成任务，原始输入包括两级前缀：一是全局共享指令，二是不同 Query 共享的文档。一级前缀往往很短。多级前缀和扩展后的一级前缀节省的 Token 比例接近，分别为 56% 和 55%。

根据这个任务，作者合成了评估数据集，平均共享相同前缀的不同 Prompt 个数为 7，平均前缀长度为 1100 个 Token，平均专有 Prompt Token 为 400 Token。作者抽取了 3000 个 Query 在 A100 上使用 Mistral-7B 模型进行测试。如下图 Table 1 所示，BatchLLM 相较于 vLLM 的最佳配置，吞吐提升约 1.3x。

5.2.3 案例 2：短公共前缀+长专有 Prompt

另一个行业工作负载为搜索场景的离线排序任务，根据相关性对给定用户 Query 的一组 Web 文档进行排序。这种情况下，公共前缀的长度比专有 Prompt 短。作者选择了 1000 个样本，平均共享相同前缀的不同 Prompt 个数为 11.36，公共前缀的长度为 770，而专有 Prompt 的长度为 978。

作者使用 Qwen2-7B 模型在 AMD MI200 GPU 上进行实验。与 Table 1 中的 vLLM 最优配置相比，BatchLLM 可以获得 1.47x 的吞吐提升。

5.3 实验分析

5.3.1 Token 复用效率对比

在 5.2.2 部分的 Snippet 生成任务中，vLLM 的 Prefix Caching 节约的 Token 比例为 44.2%，而 BatchLLM 节约了 54.9% 的 Token，证明 BatchLLM 可以有效节约 Token 计算。

5.3.2 Token-Batching

如下图 Figure 6 展示了每次迭代的 Token 数量，采用与 Figure 2 相同的工作负载。显而易见，BatchLLM 中的 Token Batching 优化可以有效减少迭代中的“低谷”，从而提升整体 GPU 利用率。

为了探究 Token Batching 优化的可迁移性，作者针对以内存为中心的 Token Batching 和基于组的重新排序/调度进行了消融实验，结果如下图 Figure 7 所示。同时，作者提出了一种适用于大规模 Batch 处理的度量标准，即排除冷却时间的端到端服务时间（TECD）。该度量排查了服务的冷却阶段，如上图 Figure 6 的最右侧部分，因为随着数据规模的增加，冷却阶段的时间占比会逐渐减小。

如下图 Figure 7a 所示，展示了启用组重新排序、以内存为中心的 Token Batching 以及两者结合时的吞吐量提升情况。可以看出，不管是在基线上，还是在全局前缀共享的情况下，它们都能获得更高吞吐，并且两者结合时最优的。如下图 Figure 7b 所示，使用 BatchLLM 最多可以将 TECD 时间缩短 15%（1600 Step -> 1350 Step）：

5.3.3 Kernel 性能对比

作者同样对比了 BatchLLM 的前缀共享 Attention 实现与多个基线在 NVIDIA A100 GPU 和 AMD MI200 GPU 上的表现。如下图 Figure 8 所示，作者在两种场景下比较了不同 Kernel 函数的性能：纯 Decoding 场景（请求数 256）和分块 Prefill 场景，其中 Token Batch 包含 Prefill Chunk 和 Decoding Token（7 个 Prefill Request 和 256 个 Decoding Request）。

可以看出，基于 Triton 的 FlashAttention 与基于 CUDA 的 FlashAttention 相比，性能仍有差距。表明对于复杂计算任务，Triton 在性能上仍有提升空间。然而，作者基于 Triton 的 BatchLLM 实现相比 Cascade Inference Kernel 仍有一定的优势。显示出 BatchLLM 内核优化的潜力。总而言之，BatchLLM Kernel 在有些情况下表现不是最优，但整体仍然具有竞争力：

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述