通过连续深度批处理实现循环语言模型的深度自适应推理

TL;DR 本文提出连续深度批处理(CDB),让循环语言模型在 token 级动态调整循环深度并首次实现高效批处理,实测可达最高 99% 的理论加速。

问题

问题背景

循环语言模型（Looped LMs）通过共享层块循环可变次数实现深度自适应推理，根据 token 难度动态分配计算量，是降低 LLM 推理成本的重要方向。

现有方法局限

标准批处理系统（如 vLLM）要求所有 token 在一个统一前向传播中处理，而深度自适应推理中不同 token 的循环次数不同，无法直接纳入同一 batch。现有工作要么忽略批处理效率，要么为每个深度单独执行前向传播，导致 GPU 利用率低、调度开销大，深度自适应的吞吐量优势无法体现。此外，循环模型中的非循环层（如 token embedding、LM head、未共享 Transformer 块）若体积较大，会进一步拖慢调度并增加复杂度。

为什么难/重要

深度自适应推理的批处理需要动态形成新 batch、管理循环 KV cache、预测提前退出的 token 并异步准备，这些在传统静态计算图推理系统中缺乏支持。业界对推理成本高度敏感，若批处理效率不解决，循环模型省下的 FLOPs 会被低效调度抵消。论文提出的连续深度批处理（CDB）在 Ouro 1.4B 和 Huginn 3.5B 上实现了高达 99% 的估计最大加速，证明批处理是释放深度自适应潜力的关键瓶颈。

行业类比

类似推荐系统中动态序列长度导致批处理困难，需要动态 bucketing 和 padding；深度自适应推理则需要对循环步骤进行动态批次重组，否则服务吞吐量受限。

核心洞察

• 深度自适应推理的实用化瓶颈不在模型本身，而在批处理系统对动态深度的不兼容。CDB 通过连续深度批处理在循环步骤之间重组 batch，使每个 token 按其需要深度执行而不拖累整体吞吐。与标准 vLLM 等静态批处理相比，这是首个专门为 looped LM 设计的动态调度框架，将深度自适应从理论承诺转化为可部署的工程方案，最高实现 99% 估计加速比。
• 异步退出预测与提前 batch 准备是 CDB 实现高效率的核心机制。系统在 token 进入循环前预测其是否会提前退出，从而在循环步骤间异步准备下一批输入，隐藏调度延迟。这不同于传统 early-exit 模型在固定深度后统一判断，CDB 将预测与调度深度融合，使批量内 token 深度不同也不会产生空闲等待，显著提升在线服务吞吐，并揭示了退出行为对整体效率的决定性影响。
• 架构中非循环组件的开销决定了深度自适应推理的上限，完全 looped 架构（如 Ouro、Huginn）优于混合架构。实验表明 token embedding、LM head 和非共享 transformer block 等大计算量层会拖慢并复杂化调度，因此未来设计应尽量减少循环核心外的非共享层。这一观察将优化焦点从算法调度转移到模型架构选择，为 looped LM 的设计提供了明确的工程指引。

方法

方法：连续深度批处理（CDB）

CDB 的输入是待推理的 token 序列，模型为带有共享循环 block 的 looped LM。方法核心在于将推理过程分解为可动态重组的计算阶段，具体包含三个关键模块。

Stage-wise Scheduling 将模型分为非循环部分（token embedding、LM head 等）和循环部分（共享 transformer block）。非循环部分一次性执行，循环部分则按步迭代。在每一步循环之间，CDB 会重新组织 batch，使不同深度的 token 可以混合进入下一次循环，从而避免为浅层 token 等待深层 token 的同步开销。

Depth-aware KV Cache 为每个 token 维护与其当前循环深度对应的 KV 缓存条目。不同深度的 token 在拼接 batch 时，缓存按深度对齐存储，避免读写冲突，同时支持缓存复用与显存高效管理。

Asynchronous Scheduling 通过一个轻量预测器提前判断哪些 token 在当前循环后会退出（即不需继续循环），并在等待当前循环计算完成期间，异步准备下一批次的输入数据。这种预测与准备重叠执行，隐藏了调度和 IO 延迟，使得 GPU 利用率保持高位。

整体输出是各 token 按其所需深度完成前向计算后的 logits，但计算过程中 token 不再被强制统一深度，而是动态退出并持续重组。

与 vLLM 等标准批处理系统要求所有 token 共享同一前向路径不同，CDB 首次在循环语言模型推理中实现了深度自适应的连续批处理，使批量效率不再成为深度自适应推理的瓶颈。

实验

实验设计

实验在 Ouro 1.4B 与 Huginn 3.5B 两个 looped LM 上评估 Continuous Depth Batching (CDB)，关注离线吞吐与在线服务场景。通过模拟可变 loop 次数的 token 分布，对比标准 batching 不可用时的效率上限，以及 CDB 实际达到的加速比。

关键发现

• CDB 实现最高 99% 的估计最大加速比，表明调度开销几乎消除。
• 完全 looped 架构 更适合深度自适应推理：非 looped 的大层（如 token embedding、LM head、非共享 transformer block）会显著拖慢并复杂化调度。
• 异步调度与深度感知 KV cache 是保持高吞吐的关键。

与基线对比

标准系统 vLLM 无法处理不同 loop 次数的 token 混合 batch，因此 CDB 是首个可行的深度自适应 batching 方法。对比“无 batching”或“固定深度”基线，CDB 在保持自适应计算节省的同时，将 batching 效率恢复至接近理想状态。

行业影响

落地场景

深度自适应推理（Depth-adaptive Inference）最适配在线对话与批量生成 混合负载，如电商智能客服、内容平台摘要 / 标题生成、代码补全与文档处理。此类场景中 token 难度差异大：大量短句 / 常见词为 “easy token”，专业术语、逻辑推理为 “hard token”。CDB 让模型在循环块上按需分配计算，避免统一深度带来的浪费或欠拟合。

商业价值

• 降本：CDB 实现高达 99% 的估计最大加速比，对 Ouro 1.4B 与 Huginn 3.5B 验证有效；同等硬件下可服务更多请求，直接降低 GPU 时租 / 电费。
• 提升吞吐与服务容量：异步预退出预测减少空闲等待，批处理密度更高，对需大规模推理的内容平台或 RLHF 数据生成尤为有利。
• 体验增强：简单 token 低延迟返回，复杂 token 保留深层计算，改善长尾响应一致性。

与现有工作流接口

CDB 可作为 serving 引擎插件 接入 vLLM 类推理栈。模型需暴露循环块与 early-exit 门控，CDB 调度器负责循环步间动态组批、深度感知 KV cache 管理及异步批准备。工程上建议优先将 全循环架构（如 Universal Transformer 变体） 迁移至 CDB，因为大型非循环层（embedding、LM head、非共享 transformer block）会拖慢调度。

具体落地用例

1. 电商智能客服：高峰时段大量简单咨询（物流查询、退换货）与少量复杂投诉并存。CDB 让 easy tokens 浅层退出、hard tokens 深层推理，显著降低平均每 token 计算量，提升客服机器人并发量。
2. 医疗文档辅助摘要：临床记录中常规描述与罕见病症描述难度差异大。深度自适应确保罕见 token 获得充分循环计算，同时常规部分快速通过，使推理成本与质量达到较好均衡。

局限

• **架构依赖性**：CDB 的加速效果高度依赖模型 looped 核心的占比。论文实验显示全 looped 架构（Ouro 和 Huginn）能获得最高效率，但若存在较大的非 looped 层（如 token embedding、LM head、非共享 Transformer block），这些层会成为调度瓶颈，拖慢整体吞吐并增加调度的复杂度。这意味着 CDB 的实际收益有限于特定的模型设计，对于许多已有的混合架构或未来可能引入非 looped 组件的模型，可能无法达到接近理论最大加速比，需要模型设计阶段就考虑与 CDB 的兼容性。
• **评估范围有限**：实验主要基于 Ouro 1.4B 和 Huginn 3.5B 两个模型以及特定的 benchmark 工作负载。虽然离线吞吐和在线服务均有测试，但未充分覆盖长序列、极端并发、混合请求长度以及不同硬件配置等实际部署场景。异步调度依赖 exit prediction 的准确性，该准确性可能对数据分布漂移敏感，论文未深入分析预测误差对系统性能的影响，也未报告尾部延迟等关键服务指标。
• **工程整合难度高**：相比标准推理系统如 vLLM 的 PagedAttention 等成熟方案，CDB 需要深度定制调度器、深度感知的 KV cache 管理以及异步批准备流程，实现复杂度显著增加，难以直接融入现有推理框架。此外，CDB 依赖模型显式提供 early-exit head 或置信度信号，对于未针对 looped 推理专门训练的模型无法直接应用，限制了其向已有预训练模型迁移的便利性。