Large language model (LLM) inference is often bottlenecked by both compute and memory, especially in offloading‑based deployments where model weights are shuttled across memory hierarchies during autoregressive decoding. Reducing the number of executed layers can cut per‑token latency and avoid costly weight movement, motivating a fresh look at early‑exit techniques.
FlexEE makes early exiting practical for LLM decoding through three key mechanisms:
- Layer‑wise exit supervision – each layer is equipped with a lightweight prediction head trained with reliable intermediate supervision, ensuring accurate exit decisions.
- Self‑speculative decoding – a low‑cost speculation over a local Top‑K vocabulary is performed; an exit is triggered only when confidence exceeds a threshold, eliminating expensive full‑vocab searches.
- Dynamic hidden‑state management – hidden states are allocated and reclaimed on‑the‑fly while preserving KV‑cache correctness, providing memory‑aware execution that adapts to different weight‑offloading ratios.
Across generative and downstream benchmarks, FlexEE achieves substantial speedups with negligible accuracy loss. End‑to‑end acceleration on Llama2‑7B and Llama3‑8B under 0% and 50% weight offloading reaches $1.27\times/3.16\times$ and $1.25\times/2.83\times$, respectively.
Review: FlexEE tightly integrates early exit with KV‑cache handling, offering an efficient and deployment‑friendly solution for resource‑constrained LLM inference.