An Efficient Inference Framework for Early-exit Large Language Models

Ruijie Miao,Yihan Yan,Xinshuo Yao,Tong Yang

2024-07-25

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.

Computation and Language,Artificial Intelligence,Machine Learning

What problem does this paper attempt to address?

This paper aims to solve the problem of constructing an efficient inference framework for early - exit large language models (LLMs). Specifically, the paper focuses on two main challenges:

1. **Granularity of batching at the iteration level**: The paper proposes a method to handle batching in each iteration until all sequences exceed the confidence threshold for early - exit. This means that at each decoding step, the model will check whether the current output is accurate enough. If the condition is met, the remaining layer calculations can be ended in advance, thereby improving efficiency.

2. **KV cache management**: Since the early - exit model may skip the calculations of certain layers, this leads to the challenge of KV cache management. The paper proposes to fill the KV caches of the remaining layers before each iteration termination to ensure that when these layers are required for subsequent calculations, the KV caches are already prepared.

By addressing these two challenges, the framework proposed in the paper can significantly improve the inference speed while maintaining the model performance. Experimental results show that compared with the complete vLLM, this framework can achieve a speed - up of up to 1.25 times. In addition, the early - exit technique using hidden - state similarity performs best in terms of inference efficiency.