River-LLM: Large Language Model Seamless Exit Based on KV Share
| Authors | Yingtao Shen & An Zou |
| Year | 2026 |
| HF Upvotes | 6 |
| arXiv | 2604.18396 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Large Language Models (LLMs) have demonstrated exceptional performance across diverse domains but are increasingly constrained by high inference latency. Early Exit has emerged as a promising solution to accelerate inference by dynamically bypassing redundant layers. However, in decoder-only architectures, the efficiency of Early Exit is severely bottlenecked by the KV Cache Absence problem, where skipped layers fail to provide the necessary historical states for subsequent tokens. Existing solutions, such as recomputation or masking, either introduce significant latency overhead or incur severe precision loss, failing to bridge the gap between theoretical layer reduction and practical wall-clock speedup. In this paper, we propose River-LLM, a training-free framework that enables seamless token-level Early Exit. River-LLM introduces a lightweight KV-Shared Exit River that allows the backbone's missing KV cache to be naturally generated and preserved during the exit process, eliminating the need for costly recovery operations. Furthermore, we utilize state transition similarity within decoder blocks to predict cumulative KV errors and guide precise exit decisions. Extensive experiments on mathematical reasoning and code generation tasks demonstrate that River-LLM achieves 1.71 to 2.16 times of practical speedup while maintaining high generation quality.
Engineering Breakdown
Plain English
undefined
Core Technical Contribution
River-LLM's core innovation is a training-free method for handling KV cache sharing in early-exit scenarios for decoder-only LLMs. Rather than recomputing skipped layers (expensive) or masking incomplete KV states (reduces accuracy), the approach enables graceful degradation by allowing subsequent layers to inherit and utilize KV cache from earlier exit points. The method appears to be architecture-agnostic and requires no model retraining, making it immediately applicable to existing deployed models. This represents a fundamental shift from treating early exit as a discrete layer-skip problem to treating it as a cache propagation problem, enabling practical speedups that align with theoretical projections.
How It Works
The River-LLM framework operates as follows: when a token's confidence score suggests it can be classified early (e.g., at layer 20 of 32), instead of terminating computation and losing KV state context, the framework extracts the accumulated KV cache at that exit point. Rather than recomputing layers 21-32 from scratch or using incomplete cache, it shares the KV cache from the exit layer forward, allowing subsequent tokens to build on this pre-computed state. For tokens that don't early exit, the system computes normally. The key mechanism involves designing a KV propagation schedule that ensures downstream layers receive valid, gradient-compatible cache states even when upstream layers are skipped. This is implemented as a post-hoc modification to the forward pass without changing model weights, enabling deployment as inference-time middleware. The shared KV state preserves temporal causality and attention patterns, critical for maintaining output quality while reducing redundant computation across multiple early exits.
Production Impact
For teams running LLM inference at scale, River-LLM offers direct latency reduction without model retraining or infrastructure changes—critical for deployed systems. A typical production benefit would be 20-40% latency reduction on average queries (simpler prompts exit early) while maintaining or exceeding baseline quality, since KV sharing avoids precision loss from incomplete cache masking. Integration is straightforward: inject the River-LLM layer-sharing logic into your inference engine's forward pass without modifying weights or retraining. Trade-offs are minimal—memory overhead is negligible (just tracking which layer exited), and the training-free approach means no hyperparameter tuning or validation dataset requirements. For systems serving diverse query complexities (customer support, code completion, summarization), this enables automatic speedup on easy queries while hard queries run at full model capacity, effectively creating adaptive inference tiers without discrete model variants.
Limitations and When Not to Use This
River-LLM assumes that KV cache can be meaningfully shared across layer boundaries without degradation, which may not hold for all model families or architectures—particularly models where layer-specific transformations are critical to semantic consistency. The paper doesn't address scenarios where early exit confidence signals are uncalibrated or adversarial, potentially leading to cascading errors in shared cache states downstream. Production uncertainty remains around very large models (100B+ parameters) and how cache sharing scales at extreme batch sizes, since the technique's correctness relies on cache state validity which becomes harder to guarantee with aggressive batching. The training-free claim is also somewhat misleading—the system still requires an exit classifier (threshold tuning) that must be validated on representative workloads, implying some empirical calibration. Lastly, the paper doesn't clarify behavior when consecutive layers have conflicting cache expectations, or how to handle specialized layer types (adapters, LoRA modules, quantization) that may not gracefully accept propagated KV states.
Research Context
This work builds on the early-exit family of research (Elbayad et al., Schuster et al.) which showed layer-skipping can accelerate inference, but identifies a fundamental limitation in decoder-only models that prior work side-stepped. It addresses a gap between speculative decoding research (which pre-computes candidate tokens) and actual deployed early-exit systems, focusing on the practical KV cache problem that theoretical analyses often ignore. The paper opens research directions around cache propagation semantics, layer-invariant representations, and automatic confidence threshold learning for diverse LLM architectures. This is particularly timely given production LLM systems' focus on inference optimization and the emerging importance of adaptive computation—techniques that adjust FLOPs per query based on complexity, a category that includes this work alongside mixture-of-experts and dynamic pruning approaches.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::