Batched Early-Exit LLM Inference Done Right with DREX
Xuting Liu, Daniel Alexander, Jiaheng Lu, Siva Kesava Reddy Kakarla, Behnaz Arzani, Vincent Liu
原文摘要
Early-exit (EE) LLMs let confident tokens skip a model's deep layers, but batched serving loses most of this gain. Current systems either force tokens to exit before their own ramp is confident, an error the ramp does not bound, or save almost nothing when decoding is weight-bound, because any continuing token makes a step read the deep weights. We present DREX, a batched EE serving system that never forces a token to exit. DREX introduces cross-iteration batching: tokens that continue past the exit layer wait there, and every few steps one deep pass over all waiting tokens reads the deep weights once for several steps. Its two policies differ in what such a token's request does: holding pauses it until the deep pass, and speculative holding lets it keep decoding from the exit layer's own guess, which the next deep pass verifies. Exited tokens are trusted and never verified. Using a calibrated cost model, a controller picks speculative holding where reading weights dominates a step and guesses are often right, and holding elsewhere. DREX keeps vLLM's asynchronous scheduling by planning each deep pass from exit counts of earlier steps. With synthetic exits where 90% of tokens exit at a quarter of the depth and 40% of guesses are correct, DREX achieves 2.22-2.24$\times$ decode speedup over non-EE across batch sizes on Llama-2-70B, while per-row skipping (vSkipper's design) gains at most 7.0%. On LayerSkip models, DREX adds no statistically resolved quality loss to the EE model's own. On CNN/DM at full memory, DREX with LayerSkip achieves 1.58-1.91$\times$ decode speedup while remaining within 1.4 ROUGE-2 points of the non-EE model. These speedups are 9-63% larger than EAGLE's across the evaluated model sizes, with each method measured against its own non-EE target.
cs.DC · cs.LG