Flexee speeds up large language model inference with smart early exit
FlexEE: Self-Speculative and KV-Compatible Early Exiting for Offloading-Aware LLM Inference
Artificial Intelligence
Summary
Large language models can be slow and use a lot of memory, especially when they must move parts of their model in and out of memory while generating text. The authors developed FlexEE, a method that lets the model stop early in its computation to save time and memory without losing much accuracy. FlexEE makes smart decisions about when to stop by checking partial results and managing memory better during text generation. This approach can make using large language models faster and more efficient on limited hardware.
What this means in practice
- •For model deployment engineers: Reduce latency and memory use when deploying large language models on hardware with limited memory by using dynamic early exit methods.
- •For cloud service providers: Enable faster inference in multi-tenant environments by minimizing weight transfers during large language model serving with KV-cache aware early exit.
Authors
Qihu Xie, Ziwei Li, Yi Kang
Abstract
Large language model (LLM) inference is often constrained by both computation and memory, especially in offloading-based deployments where model weights are transferred across memory hierarchies during autoregressive decoding. In this setting, reducing the number of executed layers can lower per-token latency while also avoiding costly weight movement. Motivated by this observation, we present FlexEE, an early exiting framework for resource-constrained and offloading-based LLM inference. FlexEE makes early exiting practical for LLM decoding through layer-wise exit supervision for reliable intermediate-layer prediction, self-speculative decoding over a Top-K local vocabulary for low-cost exit decisions, and dynamic hidden state management for KV-cache-correct and memory-aware execution. Across generative and downstream tasks, FlexEE enables efficient early exit with minimal accuracy degradation, delivering up to 1.27$\times$/3.16$\times$ and 1.25$\times$/2.83$\times$ end-to-end speedups on Llama2-7B and Llama3-8B under 0\%/50\% weight offloading, respectively.