动态重批处理让大模型提前退出更高效,不丢质量还提速。
Dynamic Rebatching for Efficient Early-Exit Inference with DREX
- 按每层出口情况动态重组请求批次,避免无效等待
- 吞吐提升2-12%,完全杜绝非自愿提前退出
- 适合需要稳定输出质量的实时大模型服务场景
早期退出(EE)是一种大语言模型加速推理架构,通过仅使用部分模型层生成简单输入的结果。但传统批处理框架难以适配EE模型,因批量中各请求并非同时满足退出条件。现有方法要么强制统一决策,错失退出机会;要么强制过早退出,降低输出质量。本文提出动态重批处理机制,在每个早期退出点动态重组请求:满足退出条件的立即处理,未完成的保留在缓冲区,重新分组后送入更深层。我们实现DREX系统,包含两项优化:1)无拷贝重批缓冲,避免数据移动开销;2)结合早期退出与服务等级协议(SLA)的智能调度器,可预测重批操作是否有益。DREX还通过高效状态复制解决跳过层带来的缺失键值缓存问题。评估显示,相比基线方法,DREX提升吞吐2%-12%且保持输出质量,关键在于彻底消除非自愿退出,保障模型设计初衷。
原文摘要 · Abstract (English)
Early-Exit (EE) is a Large Language Model (LLM) architecture that accelerates inference by allowing easier tokens to be generated using only a subset of the model's layers. However, traditional batching frameworks are ill-suited for EE LLMs, as not all requests in a batch may be ready to exit at the same time. Existing solutions either force a uniform decision on the batch, which overlooks EE opportunities, or degrade output quality by forcing premature exits. We propose Dynamic Rebatching, a solution where we dynamically reorganize the batch at each early-exit point. Requests that meet the exit criteria are immediately processed, while those that continue are held in a buffer, re-grouped into a new batch, and forwarded to deeper layers. We introduce DREX, an early-exit inference system that implements Dynamic Rebatching with two key optimizations: 1) a copy-free rebatching buffer that avoids physical data movement, and 2) an EE and SLA-aware scheduler that analytically predicts whether a given rebatching operation will be profitable. DREX also efficiently handles the missing KV cache from skipped layers using memory-efficient state-copying. Our evaluation shows that DREX improves throughput by 2-12% compared to baseline approaches while maintaining output quality. Crucially, DREX completely eliminates involuntary exits, providing a key guarantee for preserving the output quality intended by the EE model.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。