跳到主要内容
返回时间线
arXiv来源发表:

CDB 让循环语言模型按 token 自适应深度推理,在 Ouro 1.4B 上实现约 99% 的早退可用加速

核心概要

该工作提出并首次端到端实现连续深度批处理(CDB),通过把循环语言模型的推理拆成 prefill、prelude、循环核心与 coda 四个阶段队列、在循环步之间重组批次、管理深度索引 KV 缓存并用前瞻门提前一步预测退出,从而让不同循环深度的 token 共享前向过程;在 Ouro 1.4B 与 Huginn 3.5B 上,CDB 分别实现约 99% 与 83%–96% 的估计最大加速,并显示完全循环架构最适合深度自适应推理。

AI-generated editorial illustration: Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching

深度剖析

CDB 首次给出深度自适应循环语言模型的端到端高效推理实现,把解码拆成 prefill、prelude、循环核心、coda 四个可分别批处理的阶段队列,并在循环步之间重组批次。 此前 CDB 只在理论上被提出并用模拟估计吞吐加速,后续消融实验排除了 prefill、调度与 KV 缓存更新;该工作把这些成本全部纳入并给出首个完整实现与评测。 在单张 80 GB H100 上,基于 Hugging Face Transformers、vLLM 与 SGLang 的同一代码库、内核与准入策略,对 Ouro 1.4B 与 Huginn 3.5B 做离线吞吐与在线服务基准,并给出调度轨迹与延迟拟合。

该工作用 FLOP 上界与硬件感知的循环步延迟模型解释早退何时真正省时间:早退减少循环核心 FLOPs,但在内存受限区间缩小批次几乎不降低延迟,因此 refill 模式关键。 把仅按解码 FLOPs 估计的加速,修正为包含 prefill 占比的端到端估计,并给出饱和批大小作为内存受限与计算受限的分界。 推导出解码加速上界与 prefill 修正后的端到端上界;实测 prefill 占 GPU 时间的 3.7%–36%,在长上下文 ArXiv 上占比最大,使端到端加速明显低于解码上界。

架构决定深度自适应推理的收益:边界阶段(token 嵌入、LM head、非共享 transformer 层)越贵,refill 越难,完全循环模型越占优。 给出面向推理的循环架构设计指引,指出非循环层虽可能提升精度,却会拖慢并复杂化调度。 Ouro 1.4B 为完全循环(0-4-0),refill 在三个批大小扫描中始终优于 no-refill,并在饱和批大小附近或以下优势最大;Huginn 3.5B 为 2-4-2,50% 的 transformer 层在核心之外,no-refill 反而更好,两者随批大小增大而趋同。

异步调度与前瞻门把 GPU 空闲降到接近零,且不损失精度。 退出信号原本只在循环步结束后才可得,导致 GPU 在每步后等待调度;前瞻门让退出决策提前一步,使 CPU 有完整一个循环步准备下一批。 门蒸馏训练的 lookahead gate 与教师门在退出行为、精度与阈值曲线上接近;Huginn 的延迟退出只需重新校准阈值,也无精度代价;代价是最小退出深度受约束。

启示与展望

该结果面向循环语言模型的服务系统:在单张 80 GB H100 上,对完全循环的 Ouro 1.4B 与带边界 transformer 层的 Huginn 3.5B,用离线吞吐与在线服务两类基准验证 CDB 能把早退的算力节省转成墙钟加速。它适用于希望按 token 分配循环深度、并在内存受限区间保持大批次的推理场景;对完全循环、边界阶段仅含嵌入与 LM head 的模型收益最大。方法本身不改变模型预测与退出决策,因此可与既有早退门、分页 KV 缓存与 FlashAttention 内核配合使用。

KV 缓存共享的精度代价仍取决于模型:在固定深度评测中,Huginn 对共享布局不敏感甚至在单槽共享下更好,而 Ouro 在单槽共享下精度近乎崩溃,作者也未能复现原论文报告的近乎无损共享。前瞻门把最小退出深度约束为 2,对只在一次循环后就产生有用输出的模型会限制可用加速。作者只评测两个循环架构与单张 H100,未包含多 GPU、抢占、前缀缓存与投机解码,这些特性可能改变 prefill 与解码的平衡。此外,本文为全文阅读,但附录 D 的完整在线服务结果与部分图表未在文本中展开,跨工作负载的在线服务细节仍属开放问题。

来源