[论文] 动态 Rebatching 实现高效 Early-Exit 推理（DREX）

Source: arXiv - 2512.15705v1

Overview

Early‑Exit (EE) 技术允许大型语言模型（LLMs）在处理“容易”的 token 时跳过不必要的层，从而降低推理延迟。论文 Dynamic Rebatching for Efficient Early‑Exit Inference with DREX 表明，传统的批处理流水线浪费了这一机会，因为它们将批次视为整体——要么所有请求一起退出，要么全部不退出。作者提出了 Dynamic Rebatching，一种在运行时动态重塑批次的机制，使每个请求能够在最优层退出，而不牺牲吞吐量或质量。

关键贡献

• Dynamic Rebatching concept – 在每个早期退出检查点重新组织批次，立即完成满足退出条件的 token，同时将其余 token 缓冲以进行更深层处理。
• DREX system – 一个面向生产的推理引擎，实现了动态重批，并采用两项提升性能的技巧：
  1. Copy‑free rebatching buffer – 在重新排列请求时避免昂贵的内存拷贝。
  2. EE‑and SLA‑aware scheduler – 通过分析预测重批步骤是否能提升整体延迟/吞吐量，防止有害的重组。
• Efficient KV‑cache handling – 提供一种内存占用轻量的方法，重建因层被跳过而缺失的键值缓存条目，保持 Transformer 缓存的加速优势。
• Guarantee of no involuntary exits – DREX 确保请求的退出时间不会早于 EE 模型自身的置信阈值，从而保护输出质量。
• Empirical validation – 表明相较于之前的 EE 批处理基线，吞吐量提升了 2–12%，且生成质量保持不变。

方法论

1. 提前退出检查点 – 模型在选定的 Transformer 层后装配了多个“退出头”。每个头会产生一个置信度分数；如果该分数超过预先调好的阈值，令牌即可提前输出。
2. 动态重批循环：
   • 在退出点，运行时会扫描当前批次。
   • 符合置信度标准的令牌被 最终确定 并从后续处理移除。
   • 剩余的令牌被放入 重批缓冲区，该缓冲区记录它们的原始位置，但 不 复制底层张量数据。
   • 缓冲区将待处理的令牌分组为一个新批次（大小可能不同），并将其送入下一层更深的网络。
3. 无拷贝缓冲区设计 – 使用基于索引的间接方式（例如指针/偏移量向量），使相同的激活内存可以在重批步骤之间复用，消除 O(N) 的数据移动。
4. 调度器分析 – 对于每个可能的重批，DREX 使用轻量模型估算成本/收益权衡（额外计算 vs. 节省的延迟），该模型考虑层延迟、批次规模的伸缩以及 SLA 约束（如每令牌最大延迟）。如果预测的收益为负，调度器会推迟重批并保持当前批次不变。
5. KV‑缓存重建 – 当令牌跳过某层时，DREX 通过复制最近的缓存状态并施加一个廉价的线性投影来合成缺失的键值条目，从而保持缓存大小受限。

Results & Findings

• Throughput gains 随着更大的批量大小和更深的退出点而提升，因为 DREX 能在保持批次充分填充的同时，让早期退出的 token 及时离开，从而让 GPU 持续忙碌。
• Zero involuntary exits 消除了早期 EE 系统中出现的微妙质量缺陷——即为了保持批次统一而强制提前退出的 token。
• Scheduler accuracy：分析利润模型能够在 >95 % 的情况下正确预测有利的重新批次，避免了本会增加开销的不必要重排。

实际影响

• LLM API 部署者 可以集成 DREX，以降低云 GPU 成本，同时保持早退出模型所承诺的质量保证。
• 对延迟敏感的应用（例如实时代码补全、对话代理）受益于更低的尾部延迟，因为容易的 token 可以立即完成，而不必等待批次中最慢的请求。
• 框架开发者（PyTorch、TensorFlow、Triton）可以采用无拷贝的重新批处理缓冲区模式，以支持动态批大小而不产生内存拷贝开销。
• 感知 SLA 的调度 为混合工作负载打开了可能性，其中一些请求有严格的延迟上限，而其他请求则优先考虑吞吐量；DREX 可以自动在两者之间平衡。
• KV 缓存重建 展示了一种实用方法，使得即使跳过层也能保持 Transformer 缓存的有效性，这一模式可用于其他条件执行模型（例如 Mixture‑of‑Experts）。

限制与未来工作

• 模型特定调优 – 退出阈值和调度器成本模型需要针对每种模型架构和硬件进行校准；没有提供“一刀切”的配置。
• GPU 内存碎片化 – 虽然免拷贝，但间接缓冲区可能导致分散的内存访问，这可能会影响缓存有限的 GPU 的性能。
• 大批量规模的可扩展性 – 报告的收益在约 256 token 批次以上趋于平稳；需要进一步研究以在极大规模下保持优势。
• 多节点推理的扩展 – DREX 目前面向单节点 GPU 部署；分布式场景需要在节点之间协调重新分批。
• 更广泛的 EE 策略 – 本文聚焦于基于置信度的退出；未来工作可以探索混合标准（例如 token 级别的难度估计器），并与其他条件执行技术（如自适应深度或专家混合）结合。

作者

• Xuting Liu
• Daniel Alexander
• Siva Kesava Reddy Kakarla
• Behnaz Arzani
• Vincent Liu

论文信息

• arXiv ID: 2512.15705v1
• 分类: cs.DC, cs.LG
• 出版时间: 2025年12月17日
• PDF: 下载 PDF