戳破GPU泡沫：Photon推理引擎的流水线解码技术揭秘
1. 问题：GPU泡沫与解码循环
如何让AI模型运行得尽可能快？这是Moondream总部一直痴迷的问题。GPU负责模型推理中的所有数学运算，因此乍一看似乎没什么大不了的：只需告诉它做什么，然后等待答案即可。但如果你开始研究它在底层实际是如何工作的，你会发现GPU经常处于空闲状态，不是因为没活干，而是因为CPU还没有告诉它下一步该做什么。这种现象被称为GPU泡沫（GPU bubble）。
当典型的AI模型生成文本时，它一次生成一个token（token是文本块，大约几个字符）。每个token都依赖于它之前的token，这种属性称为自回归（autoregressive），因此生成是顺序的。在得到第二个token之前，你无法计算第三个token。这个解码循环涉及CPU和GPU之间的往返。GPU承担了运行实际模型的大部分繁重工作，执行数十亿次算术运算以生成下一个token。但CPU也完成了数量惊人的工作。它选择接下来运行哪些请求，设置GPU所需的元数据，从模型输出中挑选实际的token并记录下来，等等。
挑战在于，一个token的GPU工作量很小，而CPU的日常开销是每次往返都要支付的固定成本。如果GPU必须等待这些日常开销才能开始下一个token，那么它在每个循环中都会有一部分时间处于空闲状态。这就是我们得到GPU泡沫的原因。在本文中，我们将深入探讨Photon如何使用一种称为流水线解码（pipelined decoding）的技术来隐藏这些泡沫。其思想是重叠两种工作：我们在CPU仍在完成上一个token时就开始下一个token的GPU工作。
2. 泡沫的形状与流水线修复
以下是问题的形状。在阻塞版本（顶部）中，每一步都是一次接力棒传递。CPU规划并启动一次前向传播（forward），GPU运行它，然后CPU同步（synchronize），等待结果落地，提交它们，然后才开始规划下一步。这是因为规划取决于我们选择的token。例如，如果模型指示它已完成回答，那么我们需要从队列中调度一个新的待处理请求。GPU在等待CPU完成其提交-规划-启动工作时处于空闲状态。
修复方法是流水线化循环。在当前步骤的token仍在返回和提交时，启动下一次前向传播。这就是流水线版本（底部）：前向传播背靠背运行，CPU工作重叠在它们之下。我们之所以能做到这一点，是因为我们刚刚采样的token不必离开GPU。下一次前向传播直接从GPU内存读取它作为输入。我们最终仍然希望在CPU上有一个副本，用于去token化、流式传输以及判断请求是否完成，但这是可以在稍后时刻、在后台进行的簿记工作，而下一次前向传播已经在运行。不等待那个副本就是消除泡沫的关键举措。
使其安全需要三件事，我们将在本文的其余部分介绍：防止步骤缓冲区冲突（乒乓槽位，ping-pong slots）、为约束解码（constrained decoding）正确设置采样顺序（先前向传播，后采样，forward now, sample later），以及在请求完成后进行清理（僵尸，zombies）。
3. 机制一：乒乓槽位
为了运行一个解码步骤，GPU需要一组工作缓冲区：一个用于暂存输入（最后生成的token及其在序列中的位置）的地方，一个供模型写入其输出（logits，词汇表中每个单词一个分数）的地方，一个用于存放采样token的地方，以及一些注意力内核（attention kernel）需要的簿记信息，用于找到每个序列缓存的键和值（即KV缓存，KV cache）。我们在两端都保留固定（页面锁定，page-locked）的主机缓冲区，因此进出GPU的复制作为后台DMA（直接内存访问，direct memory access）传输运行，而不是阻塞CPU。这些缓冲区一次性分配，并在每一步重复使用。我们努力避免在运行时执行GPU内存分配，因为它们可能导致设备同步并引入泡沫。固定缓冲区地址也是将解码步骤一次性捕获为CUDA图（CUDA graph）并重放所必需的，从而减少内核启动开销。
我们将这个捆绑包称为DecodeSlot。这可行，但为流水线化引入了一个障碍。缓冲区在步骤完成之前一直处于使用状态，因此我们无法在当前步骤完成之前开始下一步。为了重叠两个步骤，第二步需要自己的工作集，否则它可能会在CPU读取第一步的结果之前覆盖它们。因此，我们保留两个槽位并在它们之间交替，采用乒乓方式。
关于启动需要注意的一点：我们从CPU发出启动命令时，并不会立即执行内核。相反，我们将它们排入一个流（stream）——一个有序队列，GPU按顺序从中取出执行。同一流上的工作顺序执行，而不同流上的工作可以重叠。两个槽位都将它们的前向传播放入同一个计算流。这些槽位不是为了GPU并行性。它们的存在仅仅是为了让CPU可以在GPU运行另一个槽位的前向传播时处理一个槽位的结果。所有前向传播共享那个计算流，但复制操作不共享。每一步的设备到主机复制（将采样的token带回进行簿记）放在一个单独的复制流上，因此它可以在GPU忙于下一次前向传播时运行。这就是我们不必等待它的原因。我们将复制锚定到一个事件（event），该事件在步骤的输出被写入的瞬间记录，因此它只等待该步骤的工作，而不等待其后排队的任何内容。
一个槽位只有在它的结果被读取后才变为空闲，而不仅仅是在GPU处理完它之后。它的固定主机缓冲区是可能仍在传输中的复制的着陆点，因此过早地将槽位交给新步骤会覆盖正在传输中的复制，导致难以调试的损坏错误。因此，该槽位在读取它的提交（commit）过程中保持保留状态，并且只有在提交完成后才被释放。
4. 机制二：先前向传播，后采样
下一次前向传播可以提前运行，因为它不依赖于CPU对上一个token所做的任何事情。但是，下一步的两件事确实依赖于上一步的提交结果。一个是哪些序列仍在批次中：如果一个请求刚刚完成，它不应该出现在下一次前向传播中。这是下一节（僵尸）的内容。另一个是下一步甚至允许采样哪些token，这就是本节的内容。它来自约束解码（constrained decoding）。
Moondream的空间技能返回结构化输出而不是自由文本：point返回坐标，detect返回边界框，segment返回轮廓。我们通过限制模型每一步可能产生的token，从同一个解码循环中获得这些输出：我们在采样之前将不允许的token的分数（logits）强制设为负无穷。一个point步骤必须发出一个坐标，一个detect请求经历x、y、size循环，等等。哪些token被允许（即掩码，mask）取决于到目前为止已经产生的内容，因此步骤t+1的掩码取决于我们在t处采样的token。依赖关系在于采样，而不在于前向传播。
每个调度器滴答（scheduler tick）经历三个阶段：启动（launch）、提交（commit）和最终确定（finalize）：
- 启动t+1的前向传播。它不依赖于掩码，因此立即进行。
- 提交步骤t：等待正在传输中的复制，并推进请求的解码状态。这是决定t+1掩码所必需的。
- 最终确定t+1的采样：在状态当前的情况下，构建掩码并采样。
t+1的采样在提交t之后进行，因为提交使得t+1的掩码正确。我们称这种顺序为提交-前于-最终确定（commit-before-finalize）。GPU在步骤2和3期间运行t+1的前向传播，因此提交从关键路径中消失。对于纯文本，没有掩码，因此前向传播和采样都可以提前一步运行。对于约束序列，前向传播仍然提前运行，但采样等待前一次提交，这限制了我们在没有特殊处理的情况下能提前多少。一个循环处理两者。
5. 机制三：僵尸：提前最终确定，延迟释放
回到“先前向传播，后采样”部分，我们指出了下一步依赖于上一步提交结果的两种方式。采样掩码是其中之一。批次成员资格是另一个，需要小心处理才能正确。
为了启动步骤t+1，我们首先决定它的批次，即哪些序列在其中，并且我们在提交步骤t之前就这样做。那么，当一个序列在t处遇到其停止token，但已经被烘焙到t+1的前向传播中时，会发生什么？你无法取消已启动的GPU工作。该序列已完成，但仍然物理存在于正在执行的批次中。Photon称这些为僵尸（zombies），并且不是强行加入取消逻辑，而是让行为从两个每序列字段中显现出来：
- finalized：在序列遇到EOS（序列结束标记，end-of-sequence）或其长度上限后为True。
- inflight_refs：仍在飞行中且仍引用此序列的步骤数（0、1或2）。
当步骤t提交并检测到EOS时，该序列被标记为finalized，并且其结果被发出——但它不会被拆除，因为inflight_refs仍然非零（步骤t+1引用它）。在步骤t+1提交时，该序列已经finalized，因此提交被跳过：不追加token，状态不变。僵尸无害地随行——它占据了它的槽位并写入了一些没人会读取的KV。只有当inflight_refs最终达到0时，它的KV页面和LoRA槽位才会被释放。这种“提前最终确定，延迟释放”的舞蹈是一种少量的引用计数，它取代了原本可能是一堆“在飞行中取消这一行”的特殊情况。
6. 预填充共享同一流水线
到目前为止，这都关于解码步骤，但一个真正的服务循环不断执行两种不同类型的工作：预填充（prefill，处理新请求的提示+图像，一次昂贵的多token前向传播）和解码（decode，为每个已在运行的请求一次生成一个token）。Photon不将它们分开。预填充只是同一个双槽流水线中的另一种launch（类型=prefill）。因为流水线只关心槽位是否空闲，而不关心上次使用它的是什么类型的工作，所以一个预填充前向传播可以被启动到一个槽位中，而另一个槽位的解码步骤仍在提交中，反之亦然。昂贵的预填充前向传播在GPU上运行，而CPU提交解码结果；下一个解码前向传播运行，而CPU完成接纳刚刚预填充的请求。相同的提交顺序（以及相同的inflight_refs簿记）使两种类型之间的所有内容保持正确，因此僵尸或约束解码逻辑都不需要针对“如果预填充正在飞行中”的特殊情况。
这在输出较短时最为重要。一个发出三个token的请求几乎将其整个生命周期都花在预填充和接纳上，而不是解码，因此许多短请求的工作负载实际上是一系列预填充，其中夹杂着少量解码。共享一个流水线正是让该流能够重叠其自身的CPU簿记，而不是将预填充序列化在解码之后，然后再回来。
7. 泡沫的成本模型与实测验证
流水线化实际上能带来多少收益？你可以从解码步骤的组成部分来预测它，然后对照测量结果检查预测。一个解码步骤由三部分组成：
- 前向传播（forward）：繁重的GPU矩阵乘法。在解码时，这受内存带宽限制：每个token将整个权重集流经核心，因此它有一个接近weight_bytes / memory_bandwidth的下限。随着内存变快或模型变小，它会缩小。
- 采样（sampling）：将分数转换为已提交的token：约束解码掩码、argmax/采样、空间（接地，grounding）解码以及结果的设备→主机复制。全部是GPU工作。
- 簿记（bookkeeping）：围绕它的CPU工作。选择下一个批次（plan）、启动图（launch）、提交上一步（commit）。
阻塞循环串行运行这三部分，因此GPU在簿记期间处于空闲状态——这个空闲就是泡沫。流水线化将一个步骤的簿记滑入下一个步骤的前向传播+采样之下，因此周期趋近于前向传播+采样，泡沫消失。
按步骤测量，流水线化后，这正是我们看到的——GPU在几乎整个周期内都处于忙碌状态（稳态中位数，moondream2，毫秒）：
- 3090 · 1流：前向传播4.87ms，采样0.20ms，周期5.10ms
- 3090 · 8流：前向传播6.66ms，采样0.27ms，周期6.97ms
- 3090 · 32流：前向传播10.24ms，采样0.26ms，周期10.52ms
- B200 · 1流：前向传播2.45ms，采样0.14ms，周期2.63ms
- B200 · 8流：前向传播3.12ms，采样0.14ms，周期3.30ms
- B200 · 32流：前向传播3.80ms，采样0.14ms，周期3.98ms
前向传播+采样≈周期；剩余的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并重新启动；流水线化在前向传播之下运行该工作（以及异步掩码上传），因此前向传播永不停止：
现在放入实际数字。单独测量每一部分——两个步骤时间和L——模型的预测应该落在基准测试实际交付的结果上（深度1阻塞 vs 深度2流水线化，其他不变）：
- 3090 · 1流：阻塞5.44ms，流水线化5.10ms，L=104，预测+5.7%，观测+6.5%
- 3090 · 8流：阻塞7.52ms，流水线化6.97ms，L=113，预测+7.6%，观测+7.8%
- 3090 · 32流：阻塞11.74ms，流水线化10.52ms，L=113，预测+11.1%，观测+11.6%
- B200 · 1流：阻塞3.11ms，流水线化2.63ms，L=115，预测+17.2%，观测+17.6%
- B200 · 8流：阻塞4.04ms，流水线化3.30ms，L=115，预测+22.2%，观测+21.9%
- B200 · 32流：阻塞5.55ms，流水线化3.98ms，L=104，预测+39.1%，观测+35.4%
从中可以读出三件事：
- 收益随着GPU速度的提高而增长。相同工作负载，在3090上为+12%，但在B200上32流时为+35%。簿记与GPU速度无关，因此随着前向传播缩小（更快的内存或更小的模型），泡沫在步骤中占据更大份额。流水线化是对GPU变快的一种保险，对我们来说，这与模型变小是同一回事。
- 僵尸税是真实存在的，但很小，并且可以摊销。在单个流上，僵尸是一个完全浪费的前向传播——在L≈110时约为1%。在批次中，它是受内存带宽限制（而非行数）的步骤中的额外一行，因此几乎不花代价：在32流时，3090观测到的+11.6%正好落在无僵尸的每步骤比率上。该税在单个流上影响最大，并在吞吐量所在之处消失。（B200的32流行比预测低几个百分点，原因更平淡——在约4毫秒/步时，整个运行不到半秒，因此预填充和运行结束时的批次缩减是墙上时间的一个可见部分。）
- 它只在泡沫实际可隐藏时才付出代价。（事实上，这就是我们如何发现一个bug：流水线化数字以阻塞速度出现，追溯到在构建约束解码掩码时的一次意外同步复制。将其移动到复制流在3090上价值+11%，在B200上价值+34%。）
8. 总结：从来不是单一因素
这就是整个技术：乒乓槽位使两个步骤不冲突，前向传播/采样分离使即使是约束解码也能提前运行，以及一点僵尸引用计数使完成的请求干净地拆除。GPU不再等待CPU，你就能获得从几个百分点到三分之一的收益；加速器/模型越快，收益越大。
但Photon之所以快，并非因为这一种技术，或任何一种单一技术。它之所以快，是因为数十个这样的细节在整个服务栈中叠加：我们如何调整和分块输入图像、运行模型的内核、这里的调度器排序，以及我们从热路径中移除的同步点。没有哪一个部分是全部故事；当足够多的部分排列一致时，整个栈才会变快。
我们将继续撰写这些文章，一次一个栈的角落。在Twitter上关注我们，这样你就不会错过下一篇。并请关注即将推出的Photon 2.0：我们还不能分享细节，但它是一个大版本。