要点:拆开 CPU 与 GPU 的工作负载,让两者并行执行,可以显著提高推理性能。
这是高效大语言模型推理系列的第二篇。第一篇从基本原理介绍了连续批处理,并引入本文继续使用的 KV 缓存、FlashAttention、注意力掩码等概念。
原文发表时,Inference Endpoints 上一张 H200 的价格约为每小时5美元。一小时看起来不贵,一天下来却已经是120美元,因此自然希望充分利用 GPU。连续批处理通过紧凑地组织批次,提高 GPU 利用率,避免把计算浪费在填充上。但还有一种浪费没有解决:它默认是同步的。CPU 与 GPU 轮流工作——GPU 计算时 CPU 等待,CPU 准备下一批时 GPU 等待。在每秒运行数百步的循环中,这些空闲间隙不断累积。后文将展示,它们可能接近总运行时间的四分之一。要让 GPU 始终忙于计算,就需要消除这些间隙。
为此可以采用异步批处理:把 CPU 的批次准备与 GPU 的批次计算拆开,让它们并行运行,持续让 GPU 做有用的工作。
同步批处理
最直接的同步批处理是这样工作的:
CPU 准备新批次时,会选择纳入的请求、更新 KV 缓存表、移除前几轮中已完成的请求,并接纳新请求填补腾出的空间。准备完成后,把输入传到 GPU。GPU 执行前向传播,为每个请求采样,也就是选出一个新 token。结果返回 CPU,让 CPU 知道每个请求刚生成哪个 token,然后整个循环重新开始。
注意图右侧的红色标注:GPU 完成计算后会进入空闲状态。CPU 必须完成更新——采样输出 token、更新请求状态、重新调度批次——下一批才能开始。
这正是同步批处理的低效核心:CPU 与 GPU 轮流工作。GPU 计算时 CPU 空闲,CPU 更新时 GPU 空闲,二者没有同时做有用工作的时刻。一次前向传播中的代价似乎很小,但每秒数百步的连续批处理循环会让这些间隙积累为真实的吞吐量损失。
为了展示这一点,作者用一个80亿参数(8B)模型、批大小32,生成8K token,并分析 CPU 与 GPU 的耗时:
要生成类似图表,可以给连续批处理代码加入埋点,输出 CPU 与 GPU 的活动区间,再使用这个脚本。
时间线在绿色(GPU 活跃、CPU 空闲)和红色(CPU 活跃、GPU 空闲)之间交替,二者不重叠。总生成时间为300.6秒,其中24.0%的时间 GPU 在等待 CPU 完成工作。从 GPU 的角度看,近四分之一的生成时间被浪费了。这是悲观的看法。
乐观地看,如果能完全消除 CPU 开销,生成时间就可以从约300秒降到228秒,耗时减少约24%。这不需要新内核,也不需要修改模型,只需要仔细协调硬件。
思路其实很简单:在批次 N 计算时,准备批次 N+1。但它隐藏了几个技术难点:
- 怎样在 GPU 上启动任务后,把控制权交还 CPU?
- 怎样确保 CPU 或 GPU 的任务启动时,所需数据已经就绪?
- 批次 N+1 依赖批次 N 的预测结果,怎样提前准备?
回答这些问题,就能从零构建异步批处理。作者正是按这些步骤,将它实现为 transformers 库的连续批处理的一部分。可以查看代码进行对照。
创建并发
最终目标是让 CPU 与 GPU 操作并发执行。需要一种操作分类方式,告诉机器哪些操作可以并发,这可以通过 CUDA 流实现。
什么是 CUDA 流?
要理解 CUDA 怎样排列操作,先认识CUDA 流(stream)。流是一个有序的 GPU 操作队列,包含内核启动、内存复制和同步屏障,按照提交顺序执行。每个 GPU 操作总是在某个流中调度。同一流内的操作是顺序的:上一项完成前,GPU 不会启动下一项。不同流之间的操作彼此独立,可以并发运行。如果在三个不同流里分别启动三项操作,执行过程可用下图表示:
图中三项操作同时开始。这是一种简化:每个 GPU 操作最终都由 CPU 发起,发起也需要少量时间,例如寻找合适的内核、发出调用、把命令从 CPU 传给 GPU。这称为CPU 启动开销。更贴近实际的图如下:
操作仍然并发,但开始时间因每次 CPU 启动的开销而错开。后文持续展示这些启动事件,因为它们确实耗时,也有助于跟踪异步工作流里“何时启动什么”。例如,我们经常检查一个流是否已经排空(flushed),即其中所有操作都已执行。
默认流与非默认流
如果你从未在 PyTorch 中显式使用 CUDA 流,可能会惊讶它们竟然存在。典型 PyTorch 脚本不会提到流,也不像在异步执行 GPU 操作:CPU 看起来总是等 GPU 完成才继续。原文将这种现象归因于默认流。
没有指定流的 PyTorch 操作会进入默认流。原文采用的简化模型是:默认流具有同步性质。默认流上的操作会等待其他流排空,即 GPU 上的其他工作完成后才开始;反过来,其他流上的操作也要等默认流排空后才启动。具体默认流行为及例外,见文末官方文档核对提示。
因此,在上述同步模型下,把默认流操作的结果传回 CPU 时,即使传输本来应当对 CPU 非阻塞,CPU 仍可能等到所有相关 GPU 操作完成,因为这些操作安排在默认流上。这会破坏构建并发的努力。
这就是要用非默认流的原因。把内核启动或非阻塞内存复制加入队列后,控制权立即返回 CPU。GPU 在后台执行操作,CPU 不必等待。这回答了第一个问题:用非默认流,在启动 GPU 工作后取回 CPU 控制权。
下文假设两个设备之间的所有内存传输都是非阻塞的,因此必须自行进行同步。
回到连续批处理
按照前面的设计,GPU 操作不应落到默认流。但问题仍在:应该使用哪些流?再看同步批处理图:
可以识别出三类不同的 GPU 操作:
- 把输入从 CPU 传到 GPU。
- 在 GPU 上执行计算。
- 把输出从 GPU 传回 CPU。
所以需要三个流:一个计算流,一个 CPU 到 GPU 的传输流,一个 GPU 到 CPU 的传输流。两种传输相互独立,没有必要把它们串行化,因此各自使用一个流。
术语约定:CUDA 文档把 CPU 称为主机(host),GPU 称为设备(device)。下文沿用此约定:CPU 到 GPU 是主机到设备(H2D)传输,GPU 到 CPU 是设备到主机(D2H)传输。因此三个流分别称为 H2D 流、计算流、D2H 流。
现在尝试用流异步启动 GPU 批次,并取回 CPU 控制权。CPU 依次执行:
- 在 CPU 上准备批次输入(没有流,只有 CPU 操作)。
- 用 H2D 流把输入传给 GPU。
- 用计算流在 GPU 上计算。
- 用 D2H 流取回批次输出。
- 查看结果(没有流)。
如果只使用 CUDA 流而不加入依赖,结果几乎立即可用,但它们是错误的。看看发生了什么:
由于各流独立,三个 GPU 操作几乎同时启动。计算流没有等 H2D 完成,所以前向传播读取的是 GPU 内存里原先已有的数据;D2H 流没有等计算完成,传回的是尚未算出的结果。第5步立即返回,因为没有东西阻塞 CPU,也没有默认流替它们同步。
各项操作单独看都正常。问题是,我们从未告诉流要相互等待。我们知道计算必须在 H2D 完成后开始,D2H 必须在计算完成后开始,但没有强制这种顺序。需要跨流边界表达“那项完成后,这项才能开始”的机制。
强制同步
使用CUDA 事件(event)来保证流之间的同步。
什么是 CUDA 事件?
CUDA 事件是可以记录进流里的标记。GPU 执行到该标记时,将事件设为完成;其他流可以被要求等事件完成后再开始下一项操作。原文用两个概念性操作表示:stream.record(event) 在流的当前位置插入标记,stream.wait(event) 阻止该流前进,直到事件完成。wait 阻塞的是流,而不是 CPU 或其他并行流:CPU 调用立即返回,只有等待的流被暂时拦住。具体 PyTorch API 名称见文末提示。
上图用一个事件同步两个流。CPU 快速连续发出三项操作(图中三个小块):在流1启动输入准备、在流1记录事件、通知流2等待事件。随后 CPU 立即继续。流1执行操作,完成后设置事件;流2一直停在等待标记处,事件完成才开始计算。CPU 没有参与这个等待过程,顺序完全在 GPU 端保证。
在连续批处理中使用事件
应用到本例,修复很直接。H2D 传输入队后,原文用 h2d_stream.record(h2d_done) 记录事件;只有传输完成,事件才完成。前向传播入队前,用 compute_stream.wait(h2d_done),让计算流等待输入就绪。计算与 D2H 之间也一样:通过 model.forward 启动前向传播后,记录 compute_stream.record(compute_done),再在输出传输入队前执行 d2h_stream.wait(compute_done)。这形成明确排序的流水线:
- H2D 传输在
h2d_stream上运行。 compute_stream等待h2d_done,然后执行前向传播。d2h_stream等待compute_done,然后把输出传回。
CPU 顺序提交所有操作,随后继续工作,提交期间不阻塞。GPU 通过事件保证顺序;每个流的依赖一满足,就可以开始执行。
上图展示了过程。CPU 准备批次后,迅速提交 H2D、前向传播、D2H,并在阶段间插入 record 和 wait。随后 CPU 自由了,GPU 按依赖事件依次执行各流。注意右侧绿色标注:D2H 完成后,CPU 回来读取结果。这次最终同步是整步中唯一阻塞 CPU 的位置。实现方式是:输出传输后在 D2H 流上记录第三个事件,再在 CPU 端调用 d2h_done_event.synchronize(),直到流抵达该标记,CPU 才继续。
与同步批处理的关键差别在这里:以前 CPU 每项操作后都阻塞;现在 GPU 工作时,CPU 可以做些事情。但还要明确做什么,因为目前 GPU 利用率还没有改变。
填满空闲窗口
CPU 可用的窗口位于向 GPU 提交批次 N 与提交批次 N+1 之间。最自然的用途,是准备 N+1 的输入,让它在 N 的计算结束时已经在 GPU 上就绪。下面看看如何做到。
准备 N+1 时,可以复用准备 N 的 CPU 侧对象,例如当前请求列表、缓存状态和主机端张量缓冲区。不过,需要注意两件事:
- 数据损坏:N+1 的设备端输入缓冲区不能与 N 相同,否则会覆盖 GPU 仍在读取的数据。
- 数据传递:某请求同时出现在 N 与 N+1 中时,N 生成的新 token 必须成为 N+1 的输入。
下面两节分别解决数据损坏和数据传递。
竞态条件
先处理潜在的数据损坏。如果 N 和 N+1 共用设备端输入缓冲区,N 仍在计算时 N+1 的 H2D 传输就开始,那么 GPU 读取 N 的同一内存时,可能读到被部分覆盖的数据,结果就会损坏。这是竞态条件。主机端也有同样风险:N 的 H2D 复制尚未完成时,复用其复制源也会破坏传输。
修复方法是使用两组张量,交替使用。GPU 从槽位 A 处理批次 N 时,CPU 用 N-1 的结果更新请求状态,再在槽位 B 准备 N+1 的输入。下一步交换槽位,如下图:
代价当然存在:存放输入和输出张量的 RAM、显存用量翻倍。不过这是可接受的折中,尤其在使用 FlashAttention 时:它不需要注意力掩码,而掩码通常是最大的输入张量。
两个槽位又带来另一个问题。推理通常使用CUDA Graph 降低延迟。简单来说,CUDA Graph 是预先记录的一系列 CUDA 操作,绑定具体内存地址。针对槽位 A 捕获的图不能在 B 的缓冲区上重放,因此需要两张图。如果每张图有自己的内存缓冲区,显存又会翻倍。
解决方案是内存池,两张图从同一个共享缓冲区分配内存。唯一约束是,同一池的两张图不能并发执行。N 必须完成后 N+1 才开始,因此满足条件。实践中,两张图合计的显存与一张接近,只多出初始化时两次捕获的成本。同一个池可建立任意数量的 CUDA Graph,总内存仍以各图需求的最大值为上限,如下图。
知道如何避免数据损坏后,再解决第二个问题:怎样把 N 的输出 token 送进 N+1 的输入。
结果承接(carry-over)
考虑一个同时出现在 N 和 N+1 中的请求。N 生成的新 token 是 N+1 的输入。但准备 N+1 输入缓冲区时,N 还在运行,我们尚未取得该 token。为此,构建 N+1 时先放一个占位 token,取值0,原因后文会解释。N 完成计算后、N+1 开始前向传播前,再替换占位值。这个步骤称为carry-over(承接),因为它把 N 生成的新 token 带到 N+1。原理如下:
承接只需要三样东西:N 的输出 token ID、N+1 的输入 token ID,以及说明如何承接的张量。这个张量称为承接掩码(carry-over mask):对需要承接的 token 记录目标位置,不需要承接则填 -1。示例如下:
承接本身包含四项操作:
- 从 N 的输出中选出需承接的 token,放进新张量 T。
- 把 T 中不需承接的 token 置零。
- 截断 T,使其长度与 N+1 的输入长度匹配。
- 把 T 加到 N+1 的输入 ID 上,这也是占位输入 ID 为0的原因。
这四项操作非常便宜,因此放在每个新批次的开头,并捕获进 CUDA Graph。如果承接掩码全为 -1(表示不承接该位置),最后一步就是加一个零张量。这样的情况不常见,因为跨多个批次的解码请求通常连续调度。
完整的异步循环
把这些部分组合起来,跟踪最初两步。
第0步是冷启动:没有上一批正在运行,CPU 在槽位 A 准备批次0,像同步批处理一样提交,此时还没有重叠。
第1步开始异步循环。GPU 正在槽位 A 执行批次0,CPU 已自由,立即在槽位 B 准备批次1:移除完成的请求、接纳新请求、更新 KV 缓存路由表、构建承接掩码。它们与 GPU 计算完全重叠。批次1输入就绪后,CPU 依次提交工作:启动槽位 B 的 H2D 传输,记录并等待计算流及 D2H 流的事件,然后继续。
GPU 上两件事并行发生。槽位 A 完成计算并设置 compute_done,释放批次0输出的 D2H 传输;槽位 B 则正在传入批次1输入。传输完成后设置 h2d_done,批次1开始计算。批次0到1的承接是这次计算的一部分,在正常前向传播前执行。槽位 A、B 独立,因此这些操作可以重叠。
与此同时,CPU 在 d2h_done_event.synchronize() 处阻塞,等批次0输出到达。随后处理输出、更新批次0中所有请求的状态,开始调度批次2。循环已经运行,后续每步都重复同样的模式。完整负载如下图:每个槽位的 CPU、GPU 操作及事件有各自颜色,事件也属于具体槽位。为清楚起见,图中省略 CPU 启动 GPU 计算或搬移数据的事件,但它们仍存在;相比图中操作,这些启动延迟很小。
只要 N 完成时 N+1 的输入已在 GPU 上就绪,批次间就不会出现 GPU 空闲。唯一问题是:CPU 能否在 GPU 计算完成前做完准备。通常可以:模型不断增长,批次调度仍较便宜,瓶颈是 GPU 计算,而不是 CPU。
它真的有效吗?
作者重新运行与之前相同的实验:8K token、批大小32、8B 模型。
时间线几乎全为深绿色,表示 CPU 与 GPU 同时运行。偶尔出现浅绿色细条,表示 GPU 活跃,但 CPU 已完成准备而在等待。几乎看不见的红色标记是批次间同步点:CPU 阻塞以处理批次 N 输出。GPU 活跃时间占比由76.0%升到99.4%,总生成时间从300.6秒降为234.5秒,耗时减少约22%。完全消除 CPU 开销时预期减少24%,剩下的小差距来自不可避免的同步点。没有新内核,也没有模型变更,只是让 CPU 与 GPU 同时工作。
结论
起点是 CPU 与 GPU 轮流工作的同步负载,两者都没有充分利用。通过把基于调度顺序的依赖改为基于数据的依赖,并细化同步点,作者成功拆开 CPU 与 GPU 工作负载,使两种硬件能够并行执行。这让 GPU 工作队列持续饱和,始终有活可做,最终在保持模型准确性的同时大幅提高生成速度。
完整实现位于 transformers 库。想对照实际代码,可从连续批处理入口 continuous_api.py 看起(原文链接对应此文件,正文将其称为 continuous_batching.py)。更聚焦异步的代码位于 ContinuousBatchingAsyncIOs 类。
对于强化学习等生成长度达到16K以上的任务,异步批处理让我们离最先进的长序列生成吞吐量更近一步。但实现这个目标还需要一些较小的改进。下一篇会讨论请求卸载、解码专用内核、细粒度编译等内容,敬请期待。
致谢:感谢 Pedro Cuenca 与 Aritra Roy Gosthipaty 的帮助和富有洞见的评审。
来源与版权
原文署名:Rémi Ouazan Reboul、Pedro Cuenca、Aritra Roy Gosthipaty。Hugging Face Blog,2026-05-14。依用户已声明的转载授权译载,保留原图与链接。所引用实现代码以各自仓库许可证为准;不把本文或图片自行改标为代码仓库许可证。 阅读原文。




























暂无评论内容