Source-linked AI summary
SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference
Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, Subhabrata Mukherjee
TL;DR
Token-level early exit can reduce LLM inference computation, but existing methods are not readily compatible with batching and KV caching. SkipDecode uses shared, decreasing exit points and layer skipping; the paper reports reduced computation across datasets and models, with 2× speedup and almost no degradation in all cases.
Problem
Existing token-level early exit methods are difficult to use with batching and KV caching, limiting their practical application.
Source: §1 · p. 1, §2.1 · p. 3, §2.2 · p. 3Method
SkipDecode assigns a shared exit point to each batch position, decreases exit points over the sequence, and skips lower layers so later tokens can benefit from earlier computation.
Source: §1 · p. 2, §2.4 · p. 5Results
Across datasets and model sizes, SkipDecode reduces computational demands and reaches 2× speedup with almost no degradation in all cases.
Source: §3.3 · p. 7Takeaways & Limitations
SkipDecode combines reduced computational requirements with batching, KV caching, and a controlled, predictable computational budget.
Source: §3.3 · p. 7, §6 · p. 9The decaying policy does not naturally support infinite-loop inference, because new samples can join a batch only when their position matches the remaining samples’ positions.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Autoregressive large language models (LLMs) have made remarkable progress in various natural language generation tasks. However, they incur high computation cost and latency resulting from the autoregressive token-by-token generation. To address this issue, several approaches have been proposed to reduce computational cost using early-exit strategies. These strategies enable faster text generation using reduced computation without applying the full computation graph to each token. While existing token-level early exit methods show promising results for online inference, they cannot be readily applied for batch inferencing and Key-Value caching. This is because they have to wait until the last token in a batch exits before they can stop computing. This severely limits the practical application of such techniques. In this paper, we propose a simple and effective token-level early exit method, SkipDecode, designed to work seamlessly with batch inferencing and KV caching. It overcomes prior constraints by setting up a singular exit point for every token in a batch at each sequence position. It also guarantees a monotonic decrease in exit points, thereby eliminating the need to recompute KV Caches for preceding tokens. Rather than terminating computation prematurely as in prior works, our approach bypasses lower to middle layers, devoting most of the computational resources to upper layers, allowing later tokens to benefit from the compute expenditure by earlier tokens. Our experimental results show that SkipDecode can obtain 2x to 5x inference speedups with negligible regression across a variety of tasks. This is achieved using OPT models of 1.3 billion and 6.7 billion parameters, all the while being directly compatible with batching and KV caching optimization techniques.
1 Introduction
SkipDecode addresses token-level early exit’s difficulty with batching and KV caching by assigning shared, monotonically decreasing exit points. The paper reports speedups on OPT models while retaining support for those inference optimizations.
1 Introduction
- Existing token-level early exit methods face practical barriers because batched computation waits for the last token, while later exits can require recomputing prior tokens’ KV caches.Dynamic exit indicators also provide no guarantee that worst-case computation will be lower than the full network’s cost.Source: §1 · pp. 1–2
- SkipDecode assigns all tokens at a sequence position one shared exit point and makes exit points decrease monotonically as generation proceeds.The policy is motivated by the observation that later tokens are generally more predictable as context grows.Source: §1 · p. 2
- Instead of ending computation early, SkipDecode bypasses lower layers and spends most computation on upper layers so later tokens can benefit from earlier tokens’ computation.The paper presents this as improving the speedup-task performance trade-off.Source: §1 · p. 2
- Up to 5x speedup was reported on OPT models with 1.3 billion and 6.7 billion parameters across three benchmark generation datasets.The method also supports batching and KV caching with a controlled computational budget.Source: §1 · p. 2
2 Addressing Practical Challenges in Token-Level Early Exit Strategies
This section describes how SkipDecode addresses practical early-exit challenges through synchronized batch exits, decreasing exit depths, bounded computation, and layer skipping.
2.1 Optimizing Batched Inference with Unified Exit Points
- At each sequence position, all tokens in a batch share one exit layer, avoiding the need to keep computing until the last token in the batch finishes.Tokens are processed column-wise, so each position’s computation can use computation from the preceding position.Source: §2.1 · p. 3
2.2 KV Caching and Monotonically Decreasing Exit Points
- Exit layers decrease across sequence positions, so prior tokens have computed at least as much as the current token and their KV caches need no recomputation.The policy reflects the assumption that earlier tokens are harder to predict with less context, while later tokens become more predictive.Source: §2.2 · pp. 3–4
2.3 Controlling Computational Budget with Batched Exit Points
- Predefined minimum and maximum exit layers bound computation, and a predefined function assigns the exit points across the sequence.Evaluation uses linear decay; prompt tokens run through all decoder layers, while generation decays active layers within the preset bounds.Source: §2.3 · pp. 4–5
2.4 Skipping vs. Early Termination
- Skipping lower layers rather than terminating early lets tokens attend to upper-layer computation from previous tokens; one warmup layer worked best across experimental settings.The warmup computation bridges the representation gap between initial embeddings and top-layer hidden states.Source: §2.4 · p. 5
3 Evaluation
Evaluation tests SkipDecode on two OPT model sizes and three generation datasets, then compares it with adapted early-exit baselines. The paper reports almost no degradation at 2× speedup across evaluated cases and less performance degradation than the compared approaches as speedup increases.
3.2 Datasets
- The evaluation uses OPT models with 1.3B and 6.7B parameters on E2E, Reddit-TLDR, and CNN Daily Mail generation tasks.Generation uses beam size 1, top-sampling 0.7, and temperature 0.3.Source: §3 · p. 5, §3.2 · p. 6
3.1 Experimental Design
- SkipDecode is evaluated at 2×–5× speedups against a 1× OPT baseline that supports batching and KV caching, with its computational budget bounded by the selected minimum layer.Configurations are tuned on E2E validation perplexity; the method requires no Transformer-architecture changes beyond its token-skipping policy.Source: §3.1 · p. 5
3.3 Key Results
- E2E scores remain relatively steady through 5×, whereas Reddit and CNN-DM show more pronounced task-metric declines as speedup increases.The paper attributes the larger reductions to the relative difficulty of Reddit and CNN-DM.Source: §3.3 · p. 7
- 2× speedup brings almost no task-performance degradation across evaluated cases, while computational demands decrease across datasets and model sizes.The reported performance trade-off varies by dataset.Source: §3.3 · p. 7, Table 4 · p. 8
3.4 Comparison to other Methods
- Table 5 reports less task-performance degradation for SkipDecode than for Multi-layer and CALM-DEC as speedup increases.Multi-layer supports batching and KV caching, while CALM-DEC is limited to batch size one and incurs KV backfill overhead.Source: §3.4 · pp. 7–8, Table 5 · p. 8
4 Related Work
Related work distinguishes static model-compression approaches from adaptive early-exit methods. For generation, prior token-level early exit faces practical batching and KV-caching constraints that SkipDecode is designed to address.
Model compression
Early-exit
- Adaptive early-exit methods vary computation across input parts, but most prior work targets encoder-only NLU models [24] [22] [11] [9] [2].Generation is more complex because it proceeds autoregressively, token by token.Source: §4 · p. 9
- CALM studies token-level early exit for generation [17], but prior early-exit methods face batching and KV-caching blockers and can incur full-network worst-case cost.SkipDecode is presented as supporting non-trivial batching and KV caching with predictable computational load.Source: §4 · p. 9
5 Limitations and Future Directions
SkipDecode’s decaying policy constrains when new samples can join a batch, so it does not naturally support infinite-loop inference. The paper also identifies open directions in decay-function choice and applying the policy to prompt processing.
5 Limitations and Future Directions
- New samples can join a batch only when their current position matches the remaining samples, so SkipDecode does not naturally support infinite-loop inference.This constraint arises as samples finish computation during generation under the decaying policy.Source: §5 · p. 9
- Preliminary experiments found no improvement from power-law decay over the linear decay used in this study [17].Exploring alternative decay functions remains a future research direction.Source: §5 · p. 9
- The decaying policy’s impact on prompts remains a research direction because the study uses the full network for prompt processing.The paper suggests that extending the policy to prompts and using more aggressive decay functions may yield additional speedup.Source: §5 · p. 9
6 Conclusions
SkipDecode is presented as bridging token-level early exit with practical inference requirements. The conclusion emphasizes controlled computation and identifies dynamic batching and decay functions as future improvement areas.
6 Conclusions
- SkipDecode bridges token-level early-exit benefits with real-world application requirements, including batch processing and KV caching.The conclusion frames this as addressing practical inference challenges.Source: §6 · p. 9
- SkipDecode reduces computational requirements by identifying hidden-state saturation points under a controlled computational budget.The conclusion describes this capability as consistent.Source: §6 · p. 9
- Future work can focus on improving dynamic batching and investigating decay-function behavior.Source: §6 · p. 9