PaperlayerView exact revision on arXiv

Source-linked AI summary

River-LLM: Large Language Model Seamless Exit Based on KV Share

Yingtao Shen, An Zou

arXiv:2604.18396v3cs.CL

TL;DR

Early Exit promises lower LLM inference cost, but KV Cache Absence leaves subsequent tokens without historical states from skipped decoder layers, limiting practical speedup. River-LLM is a training-free framework that uses a KV-Shared Exit River and state transition similarity for seamless token-level exit. Across mathematical reasoning and code generation, it achieves practical speedups while maintaining high generation quality.

  • Problem

    KV Cache Absence prevents decoder-only Early Exit from converting theoretical layer reduction into practical speedup because skipped layers lack historical states for later tokens.

  • Method

    River-LLM uses a lightweight KV-Shared Exit River to preserve missing KV information and state transition similarity to guide exit decisions without training.

  • Results

    1.53× to 2.16× wall-clock speedup is achieved across mathematical reasoning and code generation without additional training or fine-tuning.

  • Takeaways & Limitations

    River-LLM enables seamless token-level Early Exit while maintaining high generation quality and avoiding costly KV-cache recovery operations.

  • Takeaways & Limitations

    Evaluation focuses on models up to 8B parameters, so larger-scale validation on 24B and 70B models remains necessary.

Abstract

from arXiv · show

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.53 to 2.16 times of practical speedup while maintaining high generation quality.

1 Introduction

Early Exit can reduce redundant LLM computation, but KV Cache Absence prevents decoder-only models from realizing its theoretical efficiency in practice. River-LLM addresses this bottleneck with a lightweight shared KV structure and reports substantial wall-clock speedups.

  • Over 50% of tokens can theoretically exit at early layers, yet autoregressive wall-clock speedup remains marginal because skipped layers lack historical KV states.Subsequent tokens therefore cannot access the necessary keys and values for self-attention in bypassed layers.
  • Existing remedies, including masking, recomputation, and state propagation, either add significant latency or cause severe accuracy degradation.These approaches do not fully reconcile token-level exiting with KV cache integrity.
  • River-LLM introduces a KV-Shared Exit River that naturally preserves or approximates missing backbone KV information during token-level exit.The framework is designed to eliminate costly recovery operations.
  • State transition similarity between decoder inputs and outputs predicts cumulative KV errors and guides precise exit decisions.
  • 1.53× to 2.16× wall-clock speedup is achieved across mathematical reasoning and code generation without additional training or fine-tuning.

2 Related Works

Related work develops Early Exit through dynamic skipping, specialized exit strategies, and applications such as speculative decoding. River-LLM is positioned within token-level, skip-based approaches for fixed decoder backbones.

  • Early Exit reduces computational redundancy by dynamically varying the depth used for inference.
  • CALM uses confidence metrics for LLM exits, while AdaInfer trains an SVM to predict layers matching final-layer outputs.
  • SkipDecode uses a predefined computational budget, whereas DiffSkip and AdaSkip use attention-vector similarity for on-the-fly skipping.
  • LayerSkip and SpecEE combine Early Exit with Speculative Decoding, while DEER truncates Chain-of-Thought reasoning based on confidence.
  • Early Exit bypasses computation within a fixed backbone, distinguishing it from methods that treat dynamic routing as an intrinsic architectural property.
  • Token-level Exit lets each autoregressive token terminate at its own optimal depth, unlike Sequence-level Exit, which uses one fixed exit position per sequence.

3 Motivation

Token-level Early Exit offers finer-grained acceleration, but decoder-only models face KV Cache Absence when skipped layers fail to preserve historical states. River-LLM addresses this challenge with seamless exit, preserving KV integrity while retaining token-level flexibility.

  • Early Exit Motivation: Theoretical Early Exit speedup can reach 3.3×, motivating finer-grained token-level decisions for LLM inference.The profile uses the shallowest exit layer that yields the same prediction as the final layer.
  • Sequence-level Exit: Sequence-level Exit fixes one stopping layer for every token, so the latest-exiting token determines the sequence’s depth and reduces early-termination opportunities.This simplicity limits flexibility during long-sequence generation.
  • Token-level Exit: Token-level Exit lets each autoregressive token exit at its own optimal depth, providing finer granularity than sequence-level strategies.Its flexibility is constrained by the need to maintain historical KV states for later tokens.
  • KV Cache Absence: KV Cache Absence occurs because an early-exiting token does not compute KV pairs for bypassed layers, leaving later tokens without required history.The KV cache creates a dependency between current-token computation and previous-token states.
  • Existing Remedies: Existing remedies trade off latency, accuracy, or exit flexibility: recomputation adds overhead, masking degrades precision, and constrained exits limit opportunities.The evaluated strategies include Batching Recompute, Mono-Decreasing Exit, State Propagation, and KV Mask.
  • River-LLM: River-LLM defines seamless exit as independently exiting tokens while automatically populating skipped-layer KV caches during the exit path.Its KV-Shared Exit River lets exit layers substitute for skipped decoder layers without additional operational overhead.

4 Seamless Exit: River-LLM

River-LLM enables token-level seamless exit by using a lightweight KV-Shared Exit River that preserves missing KV information and state-transition similarity to guide exit decisions.

  • KV-Shared Exit Layer: River-LLM uses a lightweight KV-Shared Exit Layer as a plug-in for decoder-only Transformers, with each exit layer mapped to a corresponding backbone decoder.The exit layers use the backbone’s KV addressing scheme, 4-bit weight-only quantization, and FP16 KV Cache.
  • Inference with Seamless Exit: The serially connected exit layers begin at a predefined entry layer and can receive hidden states before the next backbone decoder block.When exit criteria are met, remaining computation is offloaded to accelerated exit layers before the original LM Head produces logits.
  • Error-Aware Exit Decisions: State transition similarity predicts cumulative backbone-exit value discrepancy, with first-layer similarity correlating with final-layer value similarity at r = 0.5536.Similarity generally increases across layers, except at the last layer, supporting its use for exit decisions.
  • Error-Aware Exit Decisions: The exit criterion uses state transition similarity and an exit threshold τ to determine whether a token can leave the backbone at layer l.The criterion is evaluated during autoregressive generation for each token’s hidden state.
  • Deployment Efficiency: Backbone Offloading evicts sparsely activated later backbone blocks from primary VRAM while retaining the Exit River for tokens requiring deeper processing.Its total memory footprint can approach fully quantized-model efficiency while preserving selective high-precision inference.
  • Inference with Seamless Exit: River-LLM switches from sequence-level exit during prefill to token-level exit during generation, allowing individual tokens to terminate at different depths.This phase-adaptive design preserves parallelized attention efficiency during prefill while maximizing token-level speedup during generation.
  • Accuracy-Speed Trade-off: On GSM8K, River-LLM achieves a 98.03% next-token match rate and 0.0045 KL divergence at an average exit depth of 3.79 layers.The framework uses estimated quantization discrepancy to guide exits while preserving semantic and KV integrity.

5 Evaluation

River-LLM is evaluated across eight benchmarks and four decoder-only backbones, with results covering accuracy, practical generation speed, memory use, and decision overhead.

  • Evaluation Setup: River-LLM is evaluated on four backbones—Llama3.2 1B, Llama3.1 8B, Phi4-mini, and Ministral3 8B—across eight benchmarks.The benchmark suite includes common-sense reasoning, mathematical reasoning, and code generation tasks.
  • Evaluation Setup: The evaluation uses 5-shot GSM8K, 4-shot MATH, and 0-shot configurations for the remaining benchmarks on an NVIDIA A40 GPU.Long-sequence generation is measured on GSM8K, MATH, and HumanEval.
  • Accuracy and Exit Depth: Under τ = 0.5, River-LLM remains competitive with the baseline while executing only 3 to 4 backbone layers on average.The reported efficiency is attributed to KV integrity maintained by the KV-Shared Exit River.
  • Generation Speedup: River-LLM maintains near-lossless fidelity while delivering speedups approximately 10% below Full Quantization, which can achieve higher peak throughput but incurs accuracy degradation.Full Quantization is evaluated with HQQ and compilation optimization.
  • Trade-off Analysis: River-LLM defines a superior speedup-accuracy Pareto frontier relative to masking, recomputation, and state propagation strategies.Its threshold τ enables a flexible trade-off instead of the cliff-like accuracy drops reported for alternative methods.
  • Memory Efficiency: River-LLM’s memory footprint remains significantly below the backbone and existing early-exit baselines across sequence lengths from 4K to 64K tokens.The footprint approaches fully quantized-model efficiency through shared KV and backbone offloading.
  • Inference Overhead: The exit decision takes approximately 100 microseconds on Llama3.1 8B, accounting for 0.0688% of total per-token inference time.The state-transition similarity computation has time complexity O(d) for hidden dimension d.

6 Conclusion

River-LLM addresses KV Cache Absence in decoder-only LLM Early Exit with a training-free KV-Shared Exit River and state-transition-based decisions.

  • Conclusion: Existing recomputation and masking remedies fail to bridge theoretical layer reduction and practical wall-clock speedup because of latency overhead or precision degradation.The paper identifies KV Cache Absence as the bottleneck behind this gap.
  • Conclusion: River-LLM naturally generates and preserves missing backbone KV cache during seamless token-level exit without additional training or fine-tuning.State transition similarity guides exit decisions to maintain fidelity to the original backbone output.
  • Conclusion: 1.71× to 2.16× practical speedup is reported while maintaining near-lossless generation quality across diverse benchmarks.The framework also provides a flexible accuracy-speed trade-off relative to prior dynamic inference methods.

Limitations

River-LLM’s evaluation is bounded by model scale and workload type: it has been tested only up to 8B parameters, and speedups are less pronounced for prefill-dominant tasks.

  • Evaluation covers representative models up to 8B parameters, leaving 24B and 70B-scale performance and scalability for future validation.The framework is designed to be architecture-agnostic, but larger-scale confirmation remains necessary.
  • Speedup is most significant during token-level autoregressive generation and less pronounced for prefill-dominant workloads such as MMLU.A sequence-level exit strategy is currently used for prefill-dominant tasks.

A.1 Impact of Different Quantization Backends in Exit River

River-LLM supports multiple quantization backends in its Exit River path, and upgrading from HQQ to AWQ improves reasoning accuracy while preserving competitive throughput.

  • River-LLM integrates alternative quantization backends while preserving Intrinsic KV Integrity during exit.The evaluation compares default HQQ with AWQ on Llama-3.1-8B.
  • Upgrading the Exit River from HQQ to AWQ yields higher accuracy across reasoning tasks.The results indicate that improved static quantization transfers to River-LLM performance.
  • The River + AWQ configuration achieves a GSM8K score of 77.3, outperforming the static AWQ baseline at 76.2.Throughput remains competitive with the baseline.
  • River-LLM complements operator-level quantization by routing high-entropy tokens through the full-precision backbone with negligible acceleration impact.This routing mitigates cumulative quantization noise in static low-bit models.

A.2 Benchmark results of Phi4-mini and Ministral3 8B

Additional evaluations on Phi4-mini and Ministral3 8B show that River-LLM preserves generation quality while reducing computational depth across diverse benchmarks.

  • River-LLM preserves generation quality across diverse benchmarks on Phi4-mini and Ministral3 8B while reducing overall computational depth.Newer architectures show stronger early-layer semantic understanding, allowing tokens to converge at shallower layers.

A.3 Detailed Memory Usage of River-LLM and Other Baselines

River-LLM reduces memory overhead relative to dynamic-depth baselines, especially for long sequences, and can lower parameter memory through backbone offloading while remaining training-free.

  • River-LLM’s memory analysis partitions peak GPU usage into model parameters, KV cache, and temporary activations.The reported comparison targets Llama3.1 8B with batch_size = 1.
  • River-LLM maintains a single KV-cache set identical to the backbone baseline, unlike Balcony and EE-LLM, which require auxiliary KV-cache sets.This avoids the higher GPU memory consumption associated with exit-layer or recomputation states.
  • At sequence lengths of 64K, River-LLM avoids the memory inflation observed in competing dynamic-depth strategies.The KV-shared architecture scales favorably in long-sequence scenarios.
  • Backbone Offloading reduces parameter memory from 14.96 GB to 6.73 GB.Under-utilized backbone layers are offloaded while the Exit River is retained, bringing total memory closer to a static 4-bit baseline.
  • River-LLM is training-free and adds approximately 0.0688% decision overhead compared with mixed-precision approaches requiring calibration.Its objective prioritizes preserving backbone reasoning fidelity rather than aggressive low-bit hardware throughput.
Loading 2604.18396v3…