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

背景

在大型语言模型（LLM）推理领域，效率至关重要。Hugging Face 此前发布过一篇关于“连续批处理（Continuous Batching）”的基础文章，介绍了 KV Cache、FlashAttention、Attention Masks 等核心概念。本文是该系列的第二篇，旨在解决连续批处理中一个常被忽视的性能瓶颈。

以 Inference Endpoints 上的 H200 GPU 为例，其每小时成本约为 5 美元。虽然单次使用看似便宜，但若全天候运行，日成本即达 120 美元。因此，最大化 GPU 利用率是降低推理成本的关键。

虽然连续批处理通过紧密打包批次（tightly packed batches）消除了填充（padding）带来的计算浪费，提高了 GPU 利用率，但它默认是**同步（synchronous）**的。这意味着 CPU 和 GPU 的工作是交替进行的：当 GPU 进行计算时，CPU 处于空闲状态；而当 CPU 准备下一个批次时，GPU 又处于空闲状态。在一个每秒运行数百步的循环中，这些空闲间隙会累积，导致显著的性能损失。

核心内容

同步批处理的低效本质

在传统的同步批处理流程中，工作流如下：

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

这种模式的核心低效在于CPU 和 GPU 无法同时执行有用工作。同步批处理存在明显的“空闲间隙”：GPU 完成计算后必须等待 CPU 完成采样、状态更新和重新调度才能开始下一轮。

性能剖析：近 25% 的时间被浪费

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

• 总生成时间为 300.6 秒。
• 其中 24.0% 的时间 GPU 处于空闲状态，等待 CPU 完成更新。
• 从 GPU 的角度看，近四分之一的生成时间被浪费。

如果完全消除 CPU 开销，生成时间可从 300 秒降至 228 秒，实现 24% 的速度提升。这不需要任何新的内核或模型更改，仅需对硬件协调进行优化。

异步批处理：解耦 CPU 与 GPU 工作

为了实现 GPU 100% 时间都在计算，我们需要引入异步批处理（Asynchronous Batching）。其核心思想是：将批次 N+1 的准备工作与批次 N 的计算并行执行。

这带来了三个技术挑战：

1. 如何在 GPU 上启动任务并将控制权交还给 CPU？
2. 如何确保在任务启动时，数据（无论是 CPU 还是 GPU 任务）已准备就绪？
3. 如果批次 N+1 依赖于批次 N 的预测结果，如何准备批次 N+1？

技术实现：CUDA Streams

为了解决上述问题，我们需要利用 CUDA Streams 来管理并发。

什么是 CUDA Stream？

• 定义：Stream 是 GPU 操作（内核启动、内存拷贝、同步屏障）的有序队列。
• 顺序性：同一 Stream 内的操作按提交顺序执行，前一个完成前，后一个不会开始。
• 并发性：不同 Stream 之间的操作相互独立，可以并发执行。

默认 Stream 与非默认 Stream

• 默认 Stream（Default Stream）：PyTorch 中未指定 Stream 的操作默认落入此 Stream。它具有同步属性：
  • 默认 Stream 上的操作会等待所有其他 Stream 完成。
  • 任何其他 Stream 上的操作也会等待默认 Stream 完成。
  • 这导致即使使用非阻塞内存拷贝，CPU 仍会阻塞，直到 GPU 所有操作结束，从而破坏了并发性。
• 非默认 Stream（Non-default Stream）：
  • 内核启动或非阻塞内存拷贝返回控制权给 CPU 立即执行。
  • GPU 在后台运行操作，CPU 不等待。
  • 结论：为了实现异步，必须使用非默认 Stream。

构建异步连续批处理架构

既然不能使用默认 Stream，我们需要为不同的 GPU 操作分配独立的 Stream。在同步批处理中，主要涉及三类 GPU 操作，因此需要三个 Stream：

1. Compute Stream：用于 GPU 上的计算（前向传播）。
2. CPU-to-GPU Transfer Stream：用于从 CPU 到 GPU 的数据传输。
3. GPU-to-CPU Transfer Stream：用于从 GPU 到 CPU 的结果回传。

通过为这些操作分配不同的 Stream，我们可以实现 CPU 和 GPU 的并行工作。例如，当 GPU 正在 Compute Stream 上处理批次 N 时，CPU 可以在另一个 Stream 上准备批次 N+1 的数据，并通过 CPU-to-GPU Stream 传输，从而消除空闲间隙。

关键要点

• 性能瓶颈：默认同步的连续批处理导致 CPU 和 GPU 交替工作，产生大量空闲时间。性能分析显示，近 24% 的生成时间浪费在 GPU 等待 CPU 上。
• 解决方案：引入异步批处理，解耦 CPU 的批次准备工作和 GPU 的计算工作，实现并行执行。
• 核心技术：利用 CUDA Streams 管理并发。
• 避免默认 Stream：PyTorch 中的默认 Stream 是同步的，会阻塞并发。必须使用非默认 Stream 来实现非阻塞操作和真正的并行。
• Stream 分配策略：至少需要三个独立的 Stream 分别处理：GPU 计算、CPU 到 GPU 的数据传输、GPU 到 CPU 的数据传输。
• 零代码改动模型：这种优化不需要修改模型架构或内核，仅通过优化硬件调度和数据流即可实现显著加速（约 24%）。

意义与影响

1. 显著降低推理成本：通过消除 CPU 和 GPU 之间的空闲等待，GPU 利用率接近 100%。对于高成本的 GPU 实例（如 H200），这意味着在相同时间内可以处理更多请求，或在不增加硬件成本的情况下提升吞吐量。
2. 提升系统吞吐量：24% 的性能提升直接转化为更高的每秒请求处理量（QPS），对于高并发推理服务至关重要。
3. 优化资源协调：展示了在深度学习推理中，细粒度的硬件资源管理（如 CUDA Streams）对整体性能的巨大影响。这要求工程师不仅关注模型算法，还要深入理解底层硬件执行机制。
4. 可复现性与开源贡献：Hugging Face Transformers 库已按照此逻辑实现了异步连续批处理，开发者可以直接利用这些优化，无需从零开始实现复杂的并发逻辑。