解锁连续批处理中的异步性：实现 CPU 与 GPU 并行计算

背景

在大型语言模型（LLM）推理领域，效率至关重要。Hugging Face 此前发布的第一篇文章从基本原理出发，介绍了**连续批处理（Continuous Batching）**技术，解释了 KV Cache、FlashAttention、注意力掩码等概念如何提升 GPU 利用率。连续批处理通过紧密打包的请求批次，消除了填充（padding）带来的计算浪费，从而显著提高了 GPU 的使用效率。

然而，尽管连续批处理解决了计算资源的浪费问题，它默认采用**同步（Synchronous）**模式，这引入了第二个效率瓶颈。在同步模式下，CPU 和 GPU 轮流工作：当 GPU 进行计算时，CPU 处于空闲状态；而当 CPU 准备下一个批次时，GPU 又处于等待状态。

这种低效在高频循环中尤为明显。以 Inference Endpoints 上的 H200 GPU 为例，虽然每小时成本约为 5 美元，但若全天候运行，日成本可达 120 美元。为了最大化硬件投资回报，必须确保 GPU 100% 的时间都在进行有效计算。研究表明，在同步连续批处理中，由于 CPU 和 GPU 之间的空闲间隙，可能导致近四分之一的总运行时间被浪费。

核心内容

为了解决上述同步瓶颈，本文介绍了**异步批处理（Asynchronous Batching）**机制，旨在将 CPU 的批次准备工作与 GPU 的计算工作解耦，使两者能够并行运行。

同步批处理的低效本质

在 naive 的同步批处理流程中，执行顺序如下：

1. CPU 准备：选择请求、更新 KV Cache 表、驱逐已完成的请求、接纳新请求以填充空间。
2. 数据传输：将准备好的输入数据传输至 GPU。
3. GPU 计算：执行前向传播（Forward Pass）并为每个请求采样（选择）一个新的 token。
4. 结果回传：结果传回 CPU，CPU 记录生成的 token，然后循环重复。

关键问题在于：GPU 完成计算后会进入空闲状态，直到 CPU 完成下一批次的更新（包括采样输出 token、更新请求状态、重新调度批次）。这种“你等我，我等你”的交替模式，导致在任何时刻 CPU 和 GPU 都无法同时执行有用工作。

性能剖析数据

通过对使用 8B 模型、批次大小为 32、生成 8K token 的场景进行性能分析，结果显示：

• 总生成时间：300.6 秒。
• GPU 空闲时间占比：24.0%。
• 乐观估计：如果能消除 CPU 开销，生成时间可缩短至 228 秒，实现 24% 的性能提升。

这一优化无需修改任何内核或模型架构，仅需对硬件协调进行精细控制。

异步批处理的技术挑战

实现异步批处理的核心思想是：在批次 N 进行计算的同时，准备批次 N+1。但这面临三个技术难题：

1. 如何在 GPU 上启动任务并立即将控制权交还给 CPU？
2. 如何确保在任务启动时，CPU 或 GPU 所需的数据已经就绪？
3. 如果批次 N+1 依赖于批次 N 的预测结果，该如何准备？

实现并发：CUDA Streams

为了回答上述问题并构建异步批处理，我们需要利用 CUDA Streams。

什么是 CUDA Stream？ CUDA Stream 是一个有序的操作队列（包括内核启动、内存复制、同步屏障），操作按提交顺序执行。

• 同一 Stream 内：操作是串行的，GPU 必须等待前一个操作完成才能开始下一个。
• 不同 Stream 间：操作是独立的，可以并行执行。

默认 Stream 与非默认 Stream

• 默认 Stream（Default Stream）：在 PyTorch 中，如果不显式指定 stream，操作会落入默认 stream。默认 stream 具有同步特性：任何操作都必须等待所有其他 stream 刷新（完成）后才能开始，反之亦然。这意味着即使使用非阻塞内存传输，如果操作落在默认 stream，CPU 仍会被阻塞直到 GPU 所有工作完成。这破坏了并发的可能性。
• 非默认 Stream（Non-default Stream）：为了获得并发能力，必须使用非默认 stream。将内核启动或非阻塞内存复制排入非默认 stream 后，GPU 会在后台运行操作，而 CPU 立即返回控制权，无需等待。

构建异步流水线

基于上述原理，我们将 GPU 操作划分为三个独立的 Stream，以消除默认 stream 的同步阻塞：

1. 输入传输 Stream：负责将数据从 CPU 传输到 GPU。
2. 计算 Stream：负责 GPU 上的前向传播计算。
3. 输出传输 Stream：负责将结果从 GPU 传输回 CPU。

由于这些传输操作是独立的，我们可以让它们在非默认 stream 中并行或重叠执行，从而在批次 N 计算的同时，由 CPU 在另一个 stream 中准备批次 N+1 的数据，彻底消除空闲间隙。

关键要点

• 同步模式的浪费：默认的连续批处理是同步的，导致 CPU 和 GPU 轮流空闲，造成约 24% 的运行时间浪费。
• 异步解耦：通过异步批处理，将 CPU 的批次准备（调度、KV Cache 更新）与 GPU 的计算解耦，实现并行执行。
• CUDA Streams 的作用：
  • Stream 是 GPU 操作的有序队列。
  • 不同 Stream 间的操作可以并发执行。
  • 必须避免使用默认 Stream，因为它具有全局同步特性，会阻塞并发。
  • 使用非默认 Stream 可实现非阻塞操作，让 CPU 在 GPU 计算时立即返回控制权。
• 三流架构：为了实现高效的异步流水线，需要至少三个独立的 Stream：一个用于 CPU 到 GPU 的输入传输，一个用于 GPU 计算，一个用于 GPU 到 CPU 的输出传输。
• 零代码修改模型：这种优化不涉及模型架构或内核算子的更改，仅通过改进硬件调度和数据流管理即可实现显著的性能提升（约 24% 的速度提升）。

意义与影响

异步连续批处理技术的引入，对于降低 LLM 推理成本和提升吞吐量具有重大意义：

1. 显著降低成本：对于按小时计费的云推理服务（如 Inference Endpoints），消除 24% 的空闲时间意味着在相同硬件预算下，可以处理更多的请求，或者在相同请求量下大幅降低 GPU 租赁成本。
2. 最大化硬件利用率：现代 GPU（如 H200）算力强大但昂贵。异步机制确保 GPU 始终处于忙碌状态，避免了昂贵的计算资源闲置。
3. 工程实践价值：该方案无需修改模型本身，仅通过优化推理引擎的调度逻辑（如在 transformers 库中实现）即可落地。这为其他推理框架提供了可参考的架构模式，即通过细粒度的流管理来挖掘硬件并行潜力。
4. 推动实时交互体验：更高的吞吐量意味着更低的延迟和更高的并发处理能力，这对于需要高并发、低延迟响应的实时对话应用至关重要。