An Efficient Inference Framework for Early-exit Large Language Models

TL;DR

The paper tackles the inefficiency of inference in early-exit LLMs by introducing a unified framework that combines iteration-level batch inference with KV-cache management tailored for early-exit behavior. Implemented on vLLM with CALM-style encoder-decoder adaptations to T5, it evaluates three early-exit strategies (softmax response, hidden-states similarity, and a dedicated exit classifier) and demonstrates practical speedups. Training CALM variants on CNN/DM confirms the approach can reproduce performance-efficiency trade-offs. Empirically, the framework delivers up to 1.25x throughput improvement and up to 3.39x reduction in inner-token latency, highlighting the significance of accounting for early-exit dynamics in inference engines for transformer models.

Abstract

Building efficient inference framework has gained increasing interests for research community. Early-exit models, a variant of LLMs, improves the inference efficiency of LLMs by skipping rest layers and directly generate output tokens when they are confident enough. However, there is no work of LLM inference framework that takes early-exit models into consideration. This is non-trivial as prior art on LLM inference cannot be directly applied to early-exit models. In this work, we solves two key challenges in building efficient inference framework for early-exit models: (1) batch inference at iteration-level granularity; and (2) KV cache management. For the former, we propose to process the batch until all sequences surpass the early-exit confidence threshold. For the latter, we propose to fill the KV cache of rest layers before the iteration terminates. Our evaluation shows that, compared with the original vLLM operating at full layers, our solution achieves up to 1.25x speed up.

Figures (2)

• Figure 1: Token generation throughput of two models.
• Figure 2: Inner-token latency of two models.