🤖 machine learning

Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

本文介绍了连续深度批处理（Continuous Depth Batching, CDB），这是一种新颖的调度框架，通过将循环迭代与边界阶段解耦，并将调度与 GPU 计算进行重叠，实现了针对循环语言模型的高效、端到端深度自适应推理，从而与标准方法相比，实现了高达 1.9 倍的吞吐量提升并显著降低了延迟。

原作者： Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

发布于 2026-08-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在经营着一家规模宏大、速度极快的图书馆，一位图书管理员就能同时回答成千上万个问题。在人工智能的世界里，这位图书管理员就是“大语言模型”（LLM），而问题则是用户输入的文本提示词。通常，为了节省时间和金钱，图书管理员会按照严格的队列处理所有问题，无论问题简单还是复杂，都给予完全相同的关注度。这种方式很高效，但却很浪费：回答“2+2等于几？”所消耗的脑力与解决一个复杂的物理问题竟然一样多，尽管前者易如反掌。

最近，科学家们发现了一种被称为“循环”（looped）模型的聪明技巧。这些模型不再是直线运行，而是拥有一个“思考环路”。它们可以检查自己的工作，将答案送入一个特殊的思考室进行处理；如果答案看起来不错，它们就停止；如果看起来不太稳妥，它们就会再次通过思考室进行循环。这意味着简单的问题可以快速通过，而难题则需要经过更多的循环进行深度思考。这就像一个学生，面对简单的数学题一眼就能看出答案并立即写下，但面对一道难题时，却会在白板上涂涂画画写上十分钟。这听起来很完美，但对于试图管理队列的图书管理员来说，这简直是一场噩ty（噩梦）。如果队列中的每个人都需要不同的思考时间，整个队列就会卡住。图书管理员必须等待当前这一组人全部结束，才能处理下一组，如果有些人还在思考而其他人已经完成了，系统就会陷入停滞。这篇论文正是针对这个交通拥堵问题展开研究的。

作者引入了一种名为**连续深度批处理（Continuous Depth Batching, CDB）**的新系统来解决由这些“思考环路”引起的交通拥堵。他们意识到，旧的管理方式（称为标准批处理）过于僵化。它将每个问题都视为需要相同次数的循环，这浪费了新模型的速度优势。为了解决这个问题，他们构建了一个动态调度器，充当一名超级组织有序的交通控制器。该控制器不再等待整个小组完成，而是让完成的问题立即离开队列，并瞬间用等待开始的新问题替换它们。

把这想象成一个繁忙的餐厅厨房。在旧系统中，如果厨师同时在做一份简单的沙拉（简单的 token）和一份复杂的舒芙蕾（复杂的 token），整个厨房必须等到舒芙蕾做好后，才能端出沙拉或开始下一道订单。新的 CDB 系统则不同。一旦沙拉准备好了，它就会被装盘送出，厨师会立即拿起新订单开始切菜，保持厨房始终处于忙碌状态。论文表明，通过这样做，系统可以让“思考环路”以全速运行，而绝不会让厨师的手闲下来。

研究人员在两个不同的 AI 模型——Ouro 和 Huginn 上对该系统进行了测试。他们发现 CDB 的效果极其出色。对于 Ouro 模型，新系统达到了理论最大加速比的 99%。用通俗的话说，这意味着该系统的速度几乎可以达到通过消除所有浪费时间后所能达到的极限。在现实世界的测试中，这意味着每秒处理的请求量比旧方法增加了 1.5 到 1.9 倍。当系统处于高负载状态（例如拥有数千名用户的繁忙网站）时，等待答案的时间降低了 45% 到 90%。

然而，论文也指出，这种魔力并非自动发生的。只有当“思考环路”成为工作的主要部分时，系统效果才最好。如果循环前后的步骤（例如阅读问题或编写最终答案）过于缓慢或沉重，它们会拖慢整个过程。作者发现，对于这些额外步骤较重的模型，虽然加速效果仍然很好，但不会像这样显著。他们还发现了一个处理问题的“甜点区”（最佳平衡点）：如果你尝试同时处理太少的问题，系统会因为为每一个步骤重新加载模型的“大脑”（权重）而浪费时间；如果你处理太多，系统则会变得不堪重负。CDB 通过在任何位置空出时不断用新问题填补队列，帮助找到了这种完美的平衡。

作者使用的最巧妙的技巧之一是被称为“前瞻门控”（lookahead gate）的技术。通常情况下，系统必须等待一个问题完成其思考环路，才能知道答案是否足够好以至于可以离开。这会导致计算机在等待时产生一个微小的停顿。新系统会提前预测一步答案，从而在当前问题完成思考环路的同时，准备好下一批次的请求。这就像一位指挥家在当前乐章结束之前，就预先提示下一个乐段，从而确保音符之间没有沉默。这一微小的改变将计算机的“空闲时间”降到了几乎为零，使整个过程变得异常流畅。

简而言之，这篇论文证明了我们可以让这些聪明的、具有自我纠错能力的 AI 模型在现实世界中真正跑得快。通过将“思考环路”视为一种灵活的调度而非僵化的直线，并通过不断用新任务替换已完成的任务，作者构建了一个不仅更聪明、而且显著更快、更高效的系统。他们不仅仅是建议这可能奏效，他们还构建了它，在真实硬件上进行了测试，并展示了它能带来巨大的速度提升，将一个理论上的想法转化为了面向 AI 未来的实用工具。

技术摘要：循环语言模型的连续深度批处理

问题陈述

循环语言模型（Looped LMs）为深度自适应推理提供了一种极具前景的架构，其中共享的循环核心根据每个 token 被迭代的可变次数进行迭代。这使得模型能够为“难”token 分配更多计算量，为“易”token 分配较少计算量，从而优化计算与准确度之间的权衡。然而，这种自适应性从根本上破坏了现代推理框架（如 vLLM）中使用的标准批处理机制。

在标准连续批处理（Continuous Batching）中，前向传播是一个统一的操作，其中批次中的所有 token 都经历相同数量的层。在深度自适应的循环语言模型中，token 会在不同的深度退出循环核心。因此，单个批次包含需要不同操作的 token：一些 token 需要进行下一次循环步骤，而另一些已经退出的 token 则需要运行非循环的“尾部阶段”（Coda，即 LM head）以生成下一个 token。现有的 token 级调度器无法处理这种异构性，因为它们无法将不同的 token 路由到单次前向传播中的不同阶段。此外，循环架构引入了“边界阶段”（Prelude 和 Coda），这些阶段针对每个 token 运行一次，但其频率与循环不同，增加了调度复杂度。目前尚无现有的推理系统能够支持这些具有动态深度的模型的端到端服务。

方法论：连续深度批处理 (CDB)

作者引入了 连续深度批处理 (CDB)，这是一种在单个循环迭代而非 token 边界粒度上进行调度的推理方法。CDB 解决了三个主要挑战：阶段异构性、关键路径退出以及不规则的键值（KV）缓存。

1. 基于队列的解码调度器

CDB 将生成过程分解为四个不同的队列：

预填充 (Prefill)： 等待中的请求。
序幕 (Prelude)： 等待初始嵌入/预处理的 token。
循环核心 (Recurrent Core)： 准备好进行下一次循环步骤的 token。
尾部 (Coda)： 已从核心退出并等待最终预测的 token。

调度器将这些阶段视为独立的运算，允许灵活的执行顺序。论文实现了两种模式：

无填充 (No-Refill)： 固定顺序（序幕 $\to$ 核心 $\to$ 尾部）。当 token 退出时，批次规模缩小，且在下一个解码轮次之前不会有新 token 进入。
填充 (Refill)： 阶段根据优先级以任意顺序运行。当一个 token 退出时，其槽位会立即由预填充队列中的新 token 填补。这最大化了批次大小和吞吐量，利用了循环核心无论循环迭代次数多少都重用相同权重的特性。

2. 深度感知 KV 缓存

标准的 KV 缓存按深度索引，当 token 提前退出时会产生空隙。CDB 提出了两种布局：

最后退出缓存 (Last-Exited Cache)： 通过复制已退出 token 的最终状态来填充跳过的步骤中的空槽位。这保留了注意力语义，但浪费了内存和计算量进行复制。
共享缓存 (Shared Cache)： 每个物理层仅维护一个 KV 槽位，该槽位在每次循环步骤中都会被覆盖。这使内存占用减少了 $r_{max}$ 倍，但也改变了注意力语义（token 会关注前序 token 的最新状态，无论其退出深度如何）。作者指出，现代注意力算子（如 PagedAttention, FlashAttention-3）与两者均兼容。

3. 异步执行与前瞻

为了最小化 CPU-GPU 同步带来的开销，CDB 采用了异步执行：

重叠 (Overlapping)： 当 CPU 根据第 $N-1$ 步的退出决策准备下一个批次时，GPU 正在执行当前步骤（ $N$ ）。
前瞻门控 (Lookahead Gate)： 为了消除因等待退出决策而导致的 GPU 空闲时间，调度器使用了一个“前瞻门控”。步骤 $r$ 的退出信号决定的是参与步骤 $r+2$ 而非 $r+1$ 的情况。这给了 CPU 一个完整的循环步骤来准备下一个批次，而不会浪费计算资源，从而将 GPU 空闲时间降至接近零（实验中为 0.67%）。

核心贡献

端到端实现： 首个针对深度自适应循环语言模型的完整 CDB 实现，包括服务调度器、退出门控集成和 KV 缓存处理。
理论界限： 推导了基于 FLOPs 的自适应深度加速比上限，以及一个硬件感知的屋顶模型（roofline model），用于识别“填充”模式在哪些批次规模区间中最具优势（特别是在内存受限区间）。
实证评估： 在两种架构（Ouro 1.4B 和 Huginn 3.5B）上进行的全面评估，展示了模型结构、循环次数和边界阶段成本对推理效率的影响。

实验结果

作者在 Ouro 1.4B（全循环，低成本边界阶段）和 Huginn 3.5B（2-4-2 分割，高成本边界阶段）上使用 Alpaca 和 ShareGPT 工作负载对 CDB 进行了评估。

吞吐量： CDB 达到了理论 FLOP 加速上限的 99%。
- Ouro 1.4B： 与标准连续批处理 (CB) 相比，CDB (refill) 实现了 1.5–1.58× 的离线吞吐量提升。
- Huginn 3.5B： 尽管边界阶段成本较高，CDB 仍实现了高达 1.9× 的吞吐量提升。
延迟： 在动态在线服务负载下，CDB 相比标准 CB 降低了 45–90% 的归一化延迟。
填充有效性： “填充”模式在内存受限区间（小批次规模）最为有效，因为在这些区间，提前退出会导致 GPU 计算资源闲置。在计算受限区间（大批次规模）中，由于总工作量随时间缩放，其优势会减弱。
准确度： 作者指出，前瞻门控（提前一步决定退出）与标准退出门控相比，不会显著降低准确度。然而，他们观察到 KV 缓存共享（特别是单槽位共享布局）可能会导致 Ouro 出现准确度崩溃，除非模型针对此进行了训练，而 Huginn 则表现得较为鲁棒。

意义与主张

论文声称，使用 CDB 服务的循环语言模型不仅是参数高效的，也是推理高效的。通过将可变深度的前向传播分解为调度问题，CDB 使现有硬件能够处理深度自适应，而无需承受此前认为会使此类模型在服务实践中变得不切实际的巨大开销。

作者强调，效率增益高度依赖于模型架构与推理系统的协同设计。具体而言：

边界阶段： 具有高成本边界阶段的模型（如 Huginn）加速效果会降低，因为 token 在重新进入填充队列之前必须等待这些阶段。
内存 vs. 计算： “填充”策略对于克服小批次推理的内存受限特性至关重要，它有效地隐藏了循环核心的权重重载成本。
权衡： 这项工作突出了准确度与效率之间的张力；虽然深度自适应解码节省了计算量，但退出阈值和 KV 缓存策略的选择直接影响模型性能。

论文总结道，CDB 为服务深度自适应模型提供了一条可行路径，实现了在实际高吞吐场景下实现可变深度推理的理论收益。