PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing

Anastasiia Filippova, David Grangier, Marco Cuturi, João Monteiro

arXiv:2604.22782v1cs.LGcs.AI

TL;DR

KV caching reduces redundant generation computation but creates substantial memory and serving-cost pressure, while full layer-wise caching may be redundant. The paper introduces random cross-layer attention to train models for flexible depth-wise cache sharing, showing memory reductions from dropping 50–75% of layers while often preserving or improving performance.

  • Problem

    KV caches impose substantial memory and serving-cost pressure, while evidence suggests full caching across model layers is redundant and depth-wise sharing remains underused.

  • Method

    Random Cross-Layer Attention trains layers to attend stochastically to their own or preceding-layer KV states, supporting flexible cache-sharing strategies at deployment.

  • Results

    Dropping 50–75% of layers’ caches enables memory reduction while outperforming equivalent self-attention models, with larger models often preserving or improving full-retention performance.

  • Takeaways & Limitations

    R-CLA enables models to use substantially leaner caches and supports deployment across hardware environments with different cache capacities.

  • Takeaways & Limitations

    R-CLA requires training resources, and the study does not evaluate Mixture-of-Experts architectures or broader fine-tuning tasks beyond QA.

Abstract

from arXiv · show

Serving transformer language models with high throughput requires caching Key-Values (KVs) to avoid redundant computation during autoregressive generation. The memory footprint of KV caching is significant and heavily impacts serving costs. This work proposes to lessen these memory requirements. While recent work has largely addressed KV cache reduction via compression and eviction along the temporal axis, we argue that the \emph{depth} dimension offers an orthogonal and robust avenue for optimization. Although prior research suggests that a full cache for every layer is redundant, implementing cross-layer cache sharing remains a practical challenge; existing methods typically suffer from reduced throughput or increased time-to-first-token. In this paper, we demonstrate that dropping a layer's cache offers efficient optimization without information loss. We propose a simple training approach: random cross-layer attention. During training, layers randomly choose to attend either to their own KV states or those of a preceding layer. This stochastic process adapts the model to be robust to various depth-wise cache sharing strategies, ensuring flexibility for unknown hardware constraints at deployment time. Our evaluations show that applying this scheme during pre-training or fine-tuning enables depth-wise cache sharing for various model families. Furthermore, for larger models in data-constrained settings, this approach is suggestive of a regularization-like effect, frequently preserving or improving performance while significantly reducing the cache's memory footprint.

1 Introduction

KV caching imposes a large memory cost because it stores key and value states for every token at every layer, while inter-layer redundancy creates an opportunity for depth-wise sharing. R-CLA trains layers to tolerate randomized reuse of preceding layers’ KV states, enabling flexible cache reduction across deployment settings while often preserving or improving performance.

  • Motivation: KV cache memory scales with batch size, sequence length, and model depth, and can exceed the memory required for model parameters.Caching avoids redundant autoregressive computation but introduces a substantial serving-cost and memory-footprint trade-off.
  • Problem: Depth-wise cache sharing is orthogonal to temporal eviction, but prior approaches incur cache-update overhead or require estimating higher-layer outputs.The paper argues that high inter-layer redundancy makes full per-layer caching unnecessary, while existing sharing methods can reduce throughput or increase time-to-first-token.
  • Method: R-CLA randomly makes each layer attend either to its own KV states with probability p or to a randomly selected preceding layer with probability 1 − p.This training process breaks the rigid dependency between layers and their specific KV states by simulating structural cache faults.
  • Deployment: R-CLA supports flexible test-time policies, such as caching only every 4th layer and reusing the nearest cached states for intermediate layers.Training across randomized reuse patterns allows one model to accommodate different cache-sharing strategies and hardware environments.
  • Results: Dropping 50-75% of layers’ caches reduces memory while outperforming equivalent self-attention models.For larger models, the training randomness is also suggestive of regularization, with full-retention performance frequently preserved or improved over standard full-cache baselines across multiple tasks.

2 Related work

Prior KV-cache reduction work has focused mainly on temporal eviction/compression and architectural changes, while depth-wise sharing is less explored. Existing depth-wise methods support shared KV states but incur costly updates or increased time-to-first-token, motivating a randomized training scheme for robust cache sharing.

  • Temporal Eviction and Compression: Temporal methods reduce cache size by dropping or compressing tokens along the time axis, often retaining recent tokens and attention sinks.StreamLM uses a small sliding window of recent tokens alongside crucial attention sinks; H2O and FastGen focus on cache updates during generation.
  • Temporal Eviction and Compression: H2O and FastGen do not reduce peak memory during pre-filling for long contexts.
  • Architectural Improvements and Efficiency: Architectural approaches reduce stored K and V vectors through mechanisms such as Multi-Query Attention, Grouped-Query Attention, SSMs, hybrid attention-SSM designs, and local/global attention.SSM and hybrid architectures require no cache in non-attention blocks, while local layers attend only to a small neighborhood and therefore require a smaller cache.
  • Depth-Wise Cache Sharing: Depth-wise sharing exploits redundancy across layers, but XC-Cache requires expensive updates and Layer-Condensed KV Cache significantly increases time-to-first-token.XC-Cache uses calls to an external bi-directional encoder, while Layer-Condensed KV Cache requires sequential processing.
  • Depth-Wise Cache Sharing: Prior work shows Transformer decoders can function with shared KV states, motivating randomized training to keep models robust to varied cache-sharing strategies.The approach parallels GQA's sharing across heads while applying randomized sharing across layers.

3 Enabling Cross-Layer Attention

Cross-Layer Attention lets a layer reuse a preceding layer’s KV cache, eliminating the need for a unique cache at that layer. Random Cross-Layer Attention trains models to tolerate varied cache sources, enabling deterministic depth-wise sharing at inference while preserving performance.

  • Cross-Layer Attention: Cross-Layer Attention computes a fresh query from the current layer input while attending to keys and values from a preceding layer.Using the current layer’s own keys and values recovers standard self-attention; using an earlier layer’s cache bypasses unique cache storage and computation.
  • Cache Sharing Strategy: Depth-wise cache sharing maps uncached layers to the nearest preceding cached layer, so one loaded cache can serve multiple subsequent computations.Load and Update operations are needed only for layers authorized to maintain a cache.
  • Random Cross-Layer Attention: Random Cross-Layer Attention samples self-attention or attention to a uniformly selected preceding layer during each training forward pass.A Bernoulli decision variable d ~ Bernoulli(p) selects between the layer’s own KV states and a preceding layer’s KV states.
  • Random Cross-Layer Attention: Training across varied key/value sources makes models robust to deployment-time cache-sharing strategies without retraining.The method teaches the query projection to interact with diverse previous-layer key/value distributions and avoids dependence on a layer’s own cache.
  • Evaluation: R-CLA models preserve or improve full-cache performance and withstand cache sharing better than base models across different cache-retention levels.In F1-versus-cache-size trade-offs, R-CLA dominates standard self-attention in the Pareto sense.

4 Evaluation · 4.1 Pre-training

The evaluation tests whether R-CLA supports depth-wise cache sharing without destabilizing pre-training, using compute-constrained pre-training and comparisons against shallower models with equivalent cache footprints. In Qwen-style experiments, training remains stable across cache-sharing probabilities, with only a small evaluation-loss increase even at high sharing.

  • 4 Evaluation: The evaluation examines R-CLA under cache sharing through compute-constrained pre-training and task-specific QA fine-tuning.QA is used because performance depends on parsing and retaining information from the input prompt.
  • 4.1 Pre-training: Pre-training compares R-CLA decoder-only Transformers against shallower baselines matched to the effective cache size.For p = 0.5, an R-CLA model retaining x layers’ worth of cache is compared with an x-layer baseline Transformer.
  • 4.1.1 Experiment setup: The experiments use a Qwen-1.7B–style decoder-only Transformer pretrained from scratch on a subset of the OpenWeb corpus.All runs use a fixed context length of 2,048 tokens.
  • 4.1.1 Experiment setup: All models receive an identical 34B-token training budget, with R-CLA probabilities p ∈{0.25, 0.5, 0.6, 0.75} and baseline depths from 7 to 28 layers.Experiments are conducted on NVIDIA H100 GPUs.
  • 4.1.2 Results: 2.42 at p = 0 and 2.46 at p = 0.75: evaluation loss increases by less than 2% across R-CLA training configurations.The result is reported for Qwen3-1.7B and indicates stable training despite three-quarters of attention operations being redirected to preceding layers.
  • 4.1.2 Results: R-CLA does not destabilize pre-training dynamics, with stable training reported across all tested cache-sharing probabilities.The paper also reports this stability in Appendix A.
  • 4.1.2 Results: Figure 5 compares full-depth R-CLA models with depth-wise cache sharing against shallower fully cached baselines under equivalent cache budgets.The comparison evaluates the cache-size versus evaluation-loss trade-off.

4.2 Fine-tuning

Fine-tuning across multiple context-retrieval tasks and model families shows that R-CLA enables robust depth-wise cache sharing, with substantial benefits under low KV retention and no full-retention penalty. Ablations indicate that sharing creates the main regularization effect, while stochasticity improves robustness when KV states are missing.

  • Fine-tuning setup: Models were fine-tuned for 50,000 steps on context-retrieval tasks, using batch size 128 and maximum input length 8,192 tokens.The evaluation included models of different sizes trained with and without R-CLA.
  • Evaluation data: The evaluation covered HotpotQA, SQuAD v2, MSMarco, TriviaQA, and fictional-content RepLiQA to test information retrieval from context rather than parametric knowledge.HotpotQA specifically tests multi-hop reasoning across context passages.
  • Model-family results: Across Qwen3-8B, Mistral-7B, and Llama3.1-8B, cross-layer attention improved tolerance to depth-wise cache sharing under matched compute budgets and data mixtures.Comparisons differed strictly in whether R-CLA or standard self-attention was used.
  • Model-family results: At 25% retention, base models suffered catastrophic collapse while R-CLA retained substantial capabilities in some cases; at full retention, R-CLA incurred no penalty and improved HotpotQA by +50.7% for Llama-3.1.The authors associate these gains with a regularization effect from stochastic training in data-constrained settings.
  • Ablation: At 100% retention, fixed CLA schemes often matched or slightly exceeded R-CLA, whereas randomness became important at 50% and 25% retention when KV states were missing.The ablation separates the benefits of KV sharing from those of stochastic training.

4.3 Inference Efficiency

Depth-wise cache sharing reduces KV memory and improves decode throughput, with larger savings at longer contexts and batches. The measured gains are conservative because attention still reloads shared K,V from HBM.

  • Context-length scaling: At 8K context, g=4 reduces KV cache memory from 1170 MB to 293 MB, a 4× reduction, while throughput rises from 34.0 to 41.6 tok/s (+22%).The throughput gain comes from skipping K/V projections on non-leader layers.
  • Context-length scaling: At 32K tokens, g=4 saves 3.5 GB of peak GPU memory versus the baseline.Peak memory savings increase with context length.
  • Batch-size scaling: At batch size 16 and 8K context, the baseline runs out of GPU memory, whereas g=4 completes successfully.This demonstrates higher serving capacity on the same hardware.
  • Implementation limits: The reported gains are conservative because the implementation skips non-leader K/V projections and cache allocation but leaves attention’s HBM loading unchanged.The attention computation loads K,V from HBM regardless of whether layers read their own cache or a leader’s.
  • Implementation limits: Backend optimizations such as retaining shared K,V in SRAM across consecutive layers or fusing attention computations could further improve efficiency.These optimizations would reduce repeated memory loads.

5 Conclusion

The paper argues that depth-wise correlations enable substantially leaner KV caches and introduces random cross-layer attention (R-CLA) to make pretrained models robust to cache sharing. It also identifies training-resource requirements and several unexplored settings as limitations.

  • Full per-layer KV caches are redundant, and models can operate with significantly leaner caches by exploiting depth-wise correlations.
  • Random cross-layer attention (R-CLA) is a simple randomized training strategy that adapts pretrained models to various depth-wise cache sharing strategies.
  • R-CLA can be applied effectively early during pre-training.
  • Limitations: R-CLA requires training resources, motivating post-hoc methods or lightweight adapters that avoid full parameter updates.
  • Limitations: The experiments focus on compute-optimal training regimes, do not explore overtraining, and do not evaluate Mixture-of-Experts architectures.

A Training curves

The training-curve experiments pre-train Qwen-1.7B–style decoders with varying depth-wise KV cache sharing ratios under a fixed 34B-token budget. R-CLA remains stable even at aggressive sharing probabilities.

  • Training setup: The experiments pre-train Qwen-1.7B–style decoders from scratch on a subset of the OpenWeb corpus.Models use different levels of cache sharing or dropping ratios.
  • Training setup: R-CLA is evaluated with sharing probabilities p ∈{0.25, 0.5, 0.6, 0.75}.These settings vary the cache sharing or dropping ratio.
  • Training setup: 34B tokens is the identical training budget used for all models.The training curves compare loss across tokens processed under varying KV cache sharing ratios.
  • Training stability: R-CLA incurs no training instability even under aggressive p.This finding comes from the training curves for the reported experiments.

B Fine-tuning Training Dynamics

During fine-tuning, stochastic KV routing slows learning, while on Llama-3.1-8B it delays overfitting, indicating a regularization-like training effect. In data-constrained settings, this allows R-CLA models to train for more epochs before overfitting.

  • Fine-tuning Training Dynamics: Stochastic KV routing produces consistently slower learning and higher train loss than the base model during fine-tuning.This pattern appears on both Qwen3-8B and Llama-3.1-8B and is consistent with a regularization-like effect.
  • Fine-tuning Training Dynamics: On Llama-3.1-8B, the base model overfits while R-CLA delays the onset of overfitting.The delayed overfitting means R-CLA models can train for more epochs in data-constrained fine-tuning.
  • Fine-tuning Training Dynamics: On Qwen3-8B, R-CLA learns consistently more slowly, matching the slower-learning pattern reported during pre-training.The Qwen3-8B result is explicitly described as consistent with the pre-training results in Table 1.

C R-CLA vs. CLA@k

R-CLA provides robust cache-sharing performance across retention levels, whereas deterministic CLA@k variants are sharply sensitive to the retention level used during training.

  • C R-CLA vs. CLA@k: R-CLA maintains competitive F1 performance from 100% to 25% cache retention across five QA tasks on Llama-3.1-8B.Figure 8 compares R-CLA, CLA@2, CLA@4, RD-CLA@2, and RD-CLA@4.
  • C R-CLA vs. CLA@k: Deterministic CLA@k variants degrade sharply when cache retention moves away from their trained level.CLA@k can nevertheless match or exceed R-CLA at its specific trained retention level on some tasks, such as CLA@4 on TriviaQA at 100%.

D Additional fine-tuning results

Table 6 reports Base, R-CLA (p = 0.6), and relative improvement (∆%) for F1, Exact Match (EM), and ROUGE-L across three cache retention levels.

  • D Additional fine-tuning results: Table 6 compares Base and R-CLA (p = 0.6) using relative improvement (∆%) for F1, Exact Match (EM), and ROUGE-L at three cache retention levels.The reported metrics are F1, Exact Match (EM), and ROUGE-L.
Loading 2604.22782v1…