Dynamic Rebatching for Efficient Early-Exit Inference with DREX
On the map: Efficient LLM inference and serving → Early exit and edge inference
Abstract (arXiv)
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.
Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving. Builds on 36 of its 47 references that are in the map. First posted on arXiv 17 Dec 2025.
Builds on
36 of 47 references.
Closest ancestors · 8
Same direction, newest first
- 2025HELIOS: Adaptive Model And Early-Exit Selection for Efficient LLM Inference Serving14 Apr
- 2024An Efficient Inference Framework for Early-exit Large Language Models25 Jul
- 2024LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding25 Apr
- 2024EE-Tuning: An Economical yet Scalable Solution for Tuning Early-Exit Large Language Models1 Feb
- 2023Apparate: Rethinking Early Exits to Tame Latency-Throughput Tensions in ML Serving8 Dec
- 2023EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism8 Dec
- 2023DEED: Dynamic Early Exit on Decoder for Accelerating Encoder-Decoder Transformer Models15 Nov
- 2022Confident Adaptive Language Modeling14 Jul
Other references · 28
- 2017Attention Is All You Needcited by 60,361
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
- 2019BERTScore: Evaluating Text Generation with BERTcited by 5,512
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
- 2022Fast Inference from Transformers via Speculative Decodingcited by 1,653
- 2023AWQ: Activation-aware Weight Quantization for LLM Compression and Accelerationcited by 1,298
- 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,141
- 2015Teaching Machines to Read and Comprehendcited by 1,049
- 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headscited by 843
- 2017BranchyNet: Fast Inference via Early Exiting from Deep Neural Networkscited by 596
- 2024Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Servecited by 526
16 more
- 2024ShortGPT: Layers in Large Language Models are More Redundant Than You Expectcited by 383
- 2020BERT Loses Patience: Fast and Robust Inference with Early Exitcited by 273
- 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 258
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 258
- 2019Depth-Adaptive Transformercited by 236
- 2023CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Servingcited by 203
- 2025Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computationcited by 133
- 2024vAttention: Dynamic Memory Management for Serving LLMs without PagedAttentioncited by 98
- 2024Not All Layers of LLMs Are Necessary During Inferencecited by 96
- 2023SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inferencecited by 89
- 2023Fast and Robust Early-Exiting Framework for Autoregressive Language Models with Synchronized Parallel Decodingcited by 80
- 2015Conditional Deep Learning for Energy-Efficient and Enhanced Pattern Recognitioncited by 67
- 2024Speculative Decoding via Early-exiting for Faster LLM Inference with Thompson Sampling Control Mechanismcited by 32
- 2025Adaptive Layer-skipping in Pre-trained LLMscited by 27
- 2025SpecEE: Accelerating Large Language Model Inference with Speculative Early Exitingcited by 23
- 2025Balcony: A Lightweight Approach to Dynamic Inference of Generative Language Modelscited by 5
Built on by
Citing papers per year: 2026: 2.
2 papers.
Similar, not cited
Close in wording, no citation either way
- 2025Accelerating Large Language Model Inference via Early-Exiting Algorithms
- 2025DEL: Context-Aware Dynamic Exit Layer for Efficient Self-Speculative Decoding
- 2026River-LLM: Large Language Model Seamless Exit Based on KV Share
- 2024RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference
- 2025Seesaw: High-throughput LLM Inference via Model Re-sharding
- 2024Dynamic Vocabulary Pruning in Early-Exit LLMs
- 2025PrefillOnly: An Inference Engine for Prefill-only Workloads in Large Language Model Applications
- 2024Multi-Bin Batching for Increasing LLM Inference Throughput