打破 GPU 气泡

# 戳破 GPU 气泡 | Moondream 来源: https://moondream.ai/blog/popping-the-gpu-bubble Moondream 工程团队 Photon，Moondream 的推理引擎，实现了接近实时的 VLM 推理（在 NVIDIA B200 上约 33 毫秒）。本文揭秘它如何通过优化 GPU 工作方式，将解码吞吐量提升高达 35%。 2026 年 6 月 4 日 如何让 AI 模型跑得尽可能快？这是 Moondream 总部团队痴迷的问题。GPU 负责模型推理中的所有数学运算，所以乍一看似乎没什么门道：只需告诉它做什么，然后等待答案即可。但如果你开始研究其底层工作原理，会发现 GPU 常常处于闲置状态——并非因为没有工作可做，而是因为 CPU 还没来得及告诉它下一步要做什么。这种现象被称为 **GPU 气泡**。 当典型的 AI 模型生成文本时，它会一次产生一个 **token**（token 是一段文本，大约几个字符）。每个 token 依赖于之前的 token，这种特性称为*自回归*，因此生成是顺序进行的。你无法在得到第二个 token 之前计算第三个 token。这个解码循环涉及 CPU 和 GPU 之间的往返。GPU 执行大部分繁重工作来运行实际模型，执行数十亿次算术运算以生成下一个 token。但令人惊讶的是，CPU 也做了大量工作。它选择下一步要运行哪些请求，设置 GPU 所需的元数据，从模型输出中挑选实际 token 并记录下来，等等。 挑战在于：一个 token 的 GPU 工作量是*微小*的，而 CPU 的日常维护工作是每次往返的固定开销。如果 GPU 必须等待这些维护工作完成后才能开始下一个 token，那么它在每个循环中都会有一部分时间处于闲置状态。这就是 GPU 气泡产生的原因。 在本文中，我们将深入探讨 **Photon** (https://moondream.ai/p/photon) 如何利用一种称为*流水线解码*的技术来隐藏这些气泡。其思路是重叠两种工作：在 CPU 仍在完成上一个 token 的处理时，我们就开始下一个 token 的 GPU 工作。 ## 气泡 问题如下所示。 阻塞式与流水线式解码时间线对比 在阻塞版本（上方）中，每一步都是一次接力棒传递。CPU 规划并启动一次前向传播，GPU 运行它，然后 CPU *同步*，等待结果落地，提交结果，然后才开始规划下一步。这是因为规划依赖于我们选择的 token。例如，如果模型表明它已完成回答，那么我们需要从队列中调度一个新的待处理请求。GPU 在等待 CPU 完成其提交-规划-启动的工作时处于闲置状态。 解决方案是**将循环流水线化**。在当前步骤的 token 仍在返回并被提交的过程中，启动下一次前向传播。这就是**流水线**版本（下方）：前向传播连续运行，而 CPU 的工作在其下方重叠进行。 我们之所以能做到这一点，是因为我们刚刚采样的 token 不需要离开 GPU。下一次前向传播直接从 GPU 内存中读取它作为输入。我们最终确实需要在 CPU 上保留一份副本，用于去 token 化、流式传输以及判断请求是否完成，但这些是可以在后台稍后处理的簿记工作，同时下一次前向传播已经在运行中。不等待那份副本，就是消除气泡的关键。 要使它安全可靠，需要满足三个条件，我们将在本文后续部分介绍：防止步骤缓冲区冲突（乒乓槽位）、为约束解码设置正确的采样顺序（先前向传播，后采样）、以及在请求完成后进行清理（僵尸）。 ## 机制一：乒乓槽位 为了运行一个解码步骤，GPU 需要一组工作缓冲区：一个用于放置输入的位置（最后生成的 token 及其在序列中的位置），一个供模型写入输出的位置（*logits*，词汇表中每个单词的得分），一个用于存放采样的 token 的位置，以及一些注意力内核所需的簿记信息，用于查找每个序列缓存的键和值（其 KV 缓存）。我们在两端都保留*固定*（页面锁定）的主机缓冲区，这样在 GPU 上进行的数据复制以后台 DMA（直接内存访问）传输的方式运行，而不是阻塞 CPU。 这些缓冲区一次性分配，并在每一步重复使用。我们努力避免在运行时进行 GPU 内存分配，因为这可能导致设备同步并引入气泡。固定的缓冲区地址也是将解码步骤捕获为 **CUDA 图** (https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/) 并重放所需的，从而减少内核启动开销。我们将这个包称为 `DecodeSlot` (https://github.com/m87-labs/kestrel/blob/bb530fad318ff82c1367af4629964938cff72eaa/kestrel/models/moondream/decode_slot.py)。 这很有效，但为流水线带来了一个障碍。缓冲区在步骤完成之前一直处于使用状态，因此我们无法在当前步骤完成之前启动下一步。为了重叠两个步骤，第二步需要自己的工作集，否则它可能会在 CPU 读取第一步的结果之前覆盖它们。因此，我们保留两个槽位，并以乒乓方式交替使用它们。 乒乓槽位 关于启动，有一点需要注意：我们不会在 CPU 发出启动命令的瞬间就执行内核。相反，我们将它们排队到一个*流*上——一个有序的队列，GPU 按顺序从中取出执行。同一流上的工作顺序执行，而不同流上的工作可以重叠。两个槽位都将它们的前向传播放到同一个计算流上。槽位并非用于 GPU 并行。它们的存在仅仅是为了让 CPU 可以在 GPU 运行另一个槽位的前向传播时，处理一个槽位的结果。 所有前向传播共享同一个计算流，但数据复制则不然。每个步骤的设备到主机复制（用于将采样的 token 传回进行簿记）放在一个*单独的*复制流上，这样它就可以在 GPU 忙于下一次前向传播时运行。这就是我们不需要等待它的原因。我们将复制锚定到在步骤输出被写入的那一刻记录的一个事件上，因此它只等待该步骤的工作，而不等待其后面排队的工作。 复制在后台运行 一个槽位只有在它的结果被读取后才变得可用，而不仅仅是在 GPU 完成它之后。它的固定主机缓冲区是一个可能仍在进行中的复制的着陆点，因此过早地将槽位交给新步骤会覆盖正在传输中的复制，导致难以调试的数据损坏错误。因此，槽位在读取它的提交操作完成之前一直保持保留状态，并且只有在提交完成后才释放。 ## 机制二：先前向传播，后采样 下一次前向传播可以提前运行，因为它不依赖于 CPU 对上一个 token 所做的任何事情。但是，*下一步*的两个方面确实依赖于上一步的提交结果。一个是批次中还有哪些序列：如果一个请求刚刚完成，它不应该出现在下一次前向传播中。这是下一节（僵尸）的内容。另一个是下一步允许采样哪些 token，这就是本节的内容。 这源于**约束解码**。Moondream 的空间技能返回结构化输出而不是自由文本：`point` 返回坐标，`detect` 返回边界框，`segment` 返回轮廓。我们通过在每一步限制模型可以生成的 token 来从同一个解码循环中获得这些结果：我们在采样之前将不允许的 token 的分数（*logits*）设置为负无穷。一个 `point` 步骤必须发出一个坐标，一个 `detect` 请求循环遍历 x、y、大小，等等。允许哪些 token，即*掩码*，取决于到目前为止已生成的内容，因此步骤 *t+1* 的掩码取决于我们在步骤 *t* 采样的 token。 依赖关系存在于*采样*中，而不是前向传播中。 前向传播不需要掩码；只有采样需要 每个调度器滴答经历三个阶段：**启动、提交和最终确定**： 1. **启动** 步骤 *t+1* 的前向传播。它不依赖于掩码，因此立即执行。 2. **提交** 步骤 *t*：等待正在进行的复制，并推进请求的解码状态。这是确定 *t+1* 掩码所必需的。 3. **最终确定** 步骤 *t+1* 的采样：在当前状态下，构建掩码并采样。 步骤 *t+1* 的采样在提交步骤 *t* 之后完成，因为提交使得步骤 *t+1* 的掩码正确。我们称之为“先提交后最终确定”的顺序。GPU 在步骤 2 和 3 期间运行步骤 *t+1* 的前向传播，因此提交从关键路径中消失。 对于纯文本，没有掩码，因此前向传播和采样都可以提前一步运行。对于约束序列，前向传播仍然提前运行，但采样等待前一次提交，这限制了在没有特殊处理的情况下我们能提前多少步。一个循环处理两者。 ## 机制三：僵尸：提前最终确定，延迟释放 回到*先前向传播，后采样*中，我们指出了下一步依赖于上一步提交结果的两种方式。采样掩码是其中之一。批次成员资格是另一个，它需要小心处理才能正确。 要启动步骤 *t+1*，我们首先决定其批次——哪些序列在其中——并且我们在提交步骤 *t* 之前就做了这件事。那么，当一个序列在步骤 *t* 达到其停止 token，但已经被包含在步骤 *t+1* 的前向传播中时，会发生什么？你不能取消已经启动的 GPU 工作。该序列已完成，但在物理上仍然存在于正在执行的批次中。 Photon 将这些称为**僵尸**，它没有采用取消逻辑，而是让行为从两个每序列字段中涌现： - `finalized`：在序列达到 EOS 或其长度上限后为 `True`。 - `inflight_refs`：引用该序列的正在执行中的步骤数（0、1 或 2）。 一个完成的序列作为僵尸搭乘步骤 t+1 当步骤 *t* 提交并检测到 EOS 时，该序列被标记为 `finalized`，并且其结果被发出——但它不会被拆除，因为 `inflight_refs` 仍然非零（步骤 *t+1* 引用它）。在步骤 *t+1* 提交时，该序列已经是 `finalized`，因此提交被**跳过**：不会追加 token，状态不变。这个僵尸无害地搭乘了一段——它占用了自己的槽位，并写入了一些没人会读取的 KV。只有当 `inflight_refs` 最终达到 0 时，它的 KV 页面和 LoRA 槽位才会被释放。 这种提前最终确定、延迟释放的操作是一种小型的引用计数，它取代了原本可能是“取消中间这个行”特殊情况的复杂逻辑。 ## 预填也走同一条流水线 到目前为止，这一切都关于解码步骤，但一个真正的服务循环不断执行两种*不同*类型的工作：**预填**（处理新请求的提示 + 图像，一次昂贵的一步前向传播，跨越许多 token）和**解码**（为每个正在运行的请求一次生成一个 token）。 Photon 不将它们分开。预填只是同一个双槽流水线中的另一个 `kind="prefill"` 启动。因为流水线只关心一个槽位是否空闲，而不关心上次使用它的工作类型，所以可以在一个槽位中启动预填前向传播，同时还在提交另一个槽位中的解码步骤，反之亦然。昂贵的预填前向传播在 GPU 上运行，同时 CPU 提交解码结果；下一个解码前向传播在 GPU 上运行，同时 CPU 完成接纳刚刚预填的请求。相同的提交顺序（以及相同的 `inflight_refs` 簿记）确保了两种工作在整个过程中的正确性，因此僵尸或约束解码逻辑都不需要为“如果预填正在执行中”而设置特殊情况。 这在输出较短时最为重要。一个只产生三个 token 的请求几乎将其整个生命周期都花费在预填和接纳上，而不是解码，因此许多短请求的工作负载本质上是一串预填流，偶尔夹杂一点解码。共享一个流水线可以让这个流重叠其自身的 CPU 簿记工作，而不是将预填与解码来回串行化。 ## 气泡的成本模型 流水线化实际上能带来多少收益？你可以从解码步骤的各个组成部分来预测，然后根据测量结果来验证预测。 一个解码步骤由三部分组成： - **前向传播**：繁重的 GPU 矩阵乘法。在解码时，这是受内存带宽限制的：每个 token 都将整个权重集流式穿过计算核心，因此其下限接近 `weight_bytes / memory_bandwidth`。随着内存变快或模型变小，这个时间会缩短。 - **采样**：将分数转换为提交的 token：约束解码掩码、argmax/采样、空间（接地）解码，以及结果从设备到主机的复制。全部是 GPU 工作。 - **簿记**：围绕它的 CPU 工作。选择下一个批次（`plan`），启动图（`launch`），提交上一步（`commit`）。 阻塞循环按顺序运行这三个部分，因此 GPU 在簿记期间处于闲置状态——这个空闲就是气泡。流水线化将一个步骤的簿记滑动到下一步*前向传播 + 采样*的下方，因此周期趋近于 `forward + sampling`，气泡消失。从每一步的测量结果来看，流水线化后确实如此——GPU 在整个周期内几乎始终忙碌（稳态中位数，moondream2，毫秒）： 前向传播 (ms) 采样 (ms) 周期 (ms) 3090 · 1 流 4.87 0.20 5.10 8 流 6.66 0.27 6.97 32 流 10.24 0.26 10.52 B200 · 1 流 2.45 0.14 2.63 8 流 3.12 0.14 3.30 32 流 3.80 0.14 3.98 `forward + sampling ≈ period`；剩余的 GPU 空闲时间低于 0.05 毫秒。那么，隐藏它到底值不值？这归结为两件事之间的拔河——你能隐藏多少步骤工作，对比提前运行所带来的微小代价： ``` 加速比 = T_block / T_pipe × (1 - z) └─ 隐藏气泡 ─┘ └─ 僵尸税 ─┘ ``` 两个符号，两个概念。第一项是收益，它讲述了 GPU 加速的全部故事：阻塞步骤耗时（`T_block`）除以流水线步骤耗时（`T_pipe`）——即，一旦簿记被隐藏在下面，步骤运行速度快了多少。 第二项 `z` 是提前运行的代价——机制 3 中的**僵尸税**。在提交步骤 *t* 之前启动步骤 *t+1*，一个刚刚完成的序列仍然有一个前向传播正在执行中：一个浪费的步骤。在单个流上，对于请求生成的每 `L` 个 token，就有一个浪费的前向传播，因此在 `L ≈ 110` 时大约为 1%。然而，一旦批次打包，它几乎消失——僵尸只是步骤中的额外一行，该步骤已经在为流式传输权重支付全价，因此它几乎免费搭乘。这个税在单个流上影响最大，并在吞吐量真正所在的地方减弱，这就是为什么预测它需要同时考虑 `L` 和批次大小。 下面是以两种方式测量的步骤——阻塞式每一步都空闲，等待 CPU 提交上一个 token 并重新启动；流水线化将这些工作（以及异步掩码上传）放在前向传播之下运行，因此前向传播永不停止： 在 B200 上测量的每步阻塞式与流水线式解码对比 现在代入实际数字。分别测量每个部分——两个步骤时间和 `L`——模型的预测应该与基准测试实际交付的结果相符（深度-1 阻塞与深度-2 流水线，其他不变）： 阻塞 (ms) 流水线 (ms) L 预测 观察 3090 · 1 流 5.44 5.10 104 +5.7% +6.5% 8 流 7.52 6.97 113 +7.6% +7.8% 32 流 11.74 10.52 113 +11.1% +11.6% B200 · 1 流 3.11 2.63 115 +17.2% +17.6% 8 流 4.04 3.30 115 +22.2% +21.9% 32 流 5.55 3.98 104 +39.1% +35.4% 从中可以读出三点： 1. **收益随 GPU 速度增长而增长**。相同工作负载，在 3090 上约 +12%，但在 32 流的 B200 上达到 +35%。簿记工作与 GPU 速度无关，因此随着前向传播时间缩短（更快的内存，或更小的模型），气泡在步骤中所占比例更大。流水线化是对 GPU 变快的一种保险，对我们来说，这与模型变小是同一回事。 2. **僵尸税是真实存在的，但很小，并且可以被摊销**。在单流上，僵尸是一个完整的前向传播浪费——在 L≈110 时大约为 1%。一个批