cs.LGApr 3, 2026

Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing

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

Organizations: Apple

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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 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.

Figures & tables

Figure 1: Language models are often described as extremely efficient databases, able to store training data in their parameters. The KV cache however doesn’t reflect that, and a single book may require as much space as the model itself. Memory footprint per token can be obtained as follows for transformer decoders: 2×2\times #layers ×\times #kv_heads ×\times #head_dim ×2\times 2 , where the leftmost factor of 22 is due to storing both KK and VV representations, while the rightmost factor of 22 is due to each floating point value requiring two bytes, assuming FP16 precision.
Figure 2: Performance comparison under varying cache retention rates for base models and their counterparts trained with random cross-layer attention. We report F1 scores for Question-Answer evaluations under different levels of cache retention for standard and R-CLA fine-tuned models. R-CLA avoids much of the degradation base models incur upon cache sharing, and preserves or improves performance at full cache.
Figure 3: Cache size and prediction performance trade-offs for different models.
Figure 4: Behavior of a model implementing R-CLA. During training , each layer decides whether to attend to KV states obtained from the past layer’s hidden states, or from a randomly chosen layer in the past, so each batch observes a different cache sharing strategy. During testing , a fixed deterministic cache sharing scheme is employed, and the same group of layers always share the same KV cache.
Figure 5: Cache size vs. evaluation loss for Qwen3-1.7B. R-CLA models (red) use the full 2828 -layer architecture with depth-wise cache sharing, while baselines (blue) are shallower Transformers with the same cache footprint. Deeper models with shared cache consistently outperform shallower models, showing that cache sharing preserves the benefit of depth.
pp0.000.250.500.600.75
Eval Loss2.4242.4282.4412.4462.461
Table 1: Pre-training evaluation loss for Qwen3-1.7B trained with R-CLA at varying cross-layer attention probability pp . Training remains stable across all values, with loss increasing by less than 2% even at p=0.75p=0.75 .
  Llama-3.1-8BMistral-7BQwen3-8B
DatasetRetentionBaseR-CLAΔ%\Delta\%BaseR-CLAΔ%\Delta\%BaseR-CLAΔ%\Delta\%
HotpotQA100%0.2030.306+51.10.2150.242+12.40.2330.357+53.1
50%0.1710.305+78.70.1620.211+30.00.0850.314+267.8
25%0.0800.237+196.20.0310.162+421.30.0110.098+826.6
MSMarco100%0.3010.318+5.60.3100.316+1.70.2910.329+13.0
50%0.2260.324+43.20.2360.300+26.90.1120.227+102.4
Table 2: Effect of R-CLA on F1 Score for models fine-tuned for Question-Answer. We compare the baseline standard self-attention with R-CLA trained under p=0.6p=0.6 . Evaluation is performed across three depth-wise cache retention levels.
DatasetRetentionMethod
R-CLACLA@2CLA@4RD-CLA@2RD-CLA@4
HotpotQA100%0.3060.1920.3310.2230.167
50%0.3050.1490.2840.1790.134
25%0.2370.0890.1720.1000.071
MSMarco100%0.3180.2820.3030.3230.300
50%0.3240.1840.1590.2470.211
Table 3: Ablations on top of Llama-3.1-8B. We compare R-CLA ( p=0.6p=0.6 ) against deterministic CLA fixed at every 2nd2^{nd} or 4th4^{th} layer, and their randomly applied counterparts (RD-CLA, p=0.6p=0.6 ).
InputGroupPeak Mem.KV CacheTTFTThroughput
LengthSize(MB)(MB)(ms)(tok/s)
2,048116,5733066632.9
416,357776343.0
8,192119,3191,17029734.0
418,45529328641.6
16,384122,9812,32271334.2
Table 4: Inference efficiency across input lengths. Architecture matches Qwen3-8B (36 layers, 8 KV heads, bfloat16) on a single 80GB GPU. Group size g=1g{=}1 is the standard baseline; g=4g{=}4 shares every 4 consecutive layers. Time to first token (TTFT) measures prefill latency.
BatchGroupPeak Mem.KV CacheTTFTThroughput
SizeSize(MB)(MB)(ms)(tok/s)
2122,9732,32258833.3
421,24558057340.3
4130,2814,6431,20122.7
426,8251,1611,15925.9
8144,8979,2872,39512.8
Table 5: Batch size scaling at 8,192-token context. Same architecture as Table 4 . At batch size 16, the baseline ( g=1g{=}1 ) exceeds GPU memory, while cache sharing ( g=4g{=}4 ) completes successfully.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Training loss as a function of tokens processed for Qwen3-1.7B with varying KV cache sharing ratios.
Figure 7: Train loss (left) and eval loss (right) over fine-tuning steps for Base (blue) vs. R-CLA (red) on Qwen3-8B (top) and Llama-3.1-8B (bottom). R-CLA shows slower learning throughout, consistent with a regularization-like effect from stochastic KV routing.
Figure 8: F1 score vs. cache retention (%) for Llama-3.1-8B across five QA tasks. R-CLA (red, solid) maintains performance across retention levels, while CLA@ kk variants (dashed) degrade outside their trained retention level. Data from Table 3 .
DatasetModelRetentionF1Exact Match (EM)ROUGE-L
BaseR-CLAΔ%\Delta\%BaseR-CLAΔ%\Delta\%BaseR-CLAΔ%\Delta\%
HotpotQALlama-3.1-8B100%0.2030.306+51.1%0.1280.221+72.3%0.2060.307+49.0%
50%0.1710.305+78.7%0.1020.216+112.6%0.1730.305+75.8%
25%0.0800.237+196.2%0.0350.153+339.6%0.0880.240+171.6%
Mistral-7B100%0.2150.242+12.4%0.1220.141+15.8%0.2160.243+12.6%
50%0.1620.211+30.0%0.0810.106+30.5%0.1650.212+28.6%
Appendix
Table 6: We report Base, R-CLA ( p=0.6p=0.6 ), and relative improvement ( Δ%\Delta\% ) for F1, Exact Match (EM), and ROUGE-L across three cache retention levels.
More in this topic

KV-Cache Management

Key-value (KV) cache management reduces and coordinates Transformer attention state to make long-context inference faster and more memory-efficient.

See all 199 papers

Explore similar work

CardsList
Jul 7, 2026

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression

Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches, yet existing compression methods often apply uniform budgets across layers or tokens and degrade retrieval when lexical cues and semantic states require different preservation. We introduce DepthWeave-KV, a token-adaptive cache compression method that factorizes key and value states across neighboring transformer layers using shared low-rank channel bases while retaining lightweight token-specific residuals where attention behavior is sensitive. DepthWeave-KV combines cross-depth residual factorization with a token-conditional depth router that allocates higher reconstruction rank to instruction-bearing and retrieval-critical tokens, and uses calibration-free online error tracking from attention-output probes to adapt compression during generation without retraining the base model. A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic. Across LongBench, Needle-in-a-Haystack, L-Eval, and long-form QA and summarization benchmarks, DepthWeave-KV achieves near-full-cache task quality with substantially lower memory use, improving average score and retrieval accuracy over prior compressed caches while reaching 8.3x KV memory reduction and 72.8 tokens per second at 64K context.

Anna Cordoba, Adam Puente Tercero, Nerea Angulo Hijo +4

Instituto de Investigación en Visión Artificial

KV CachingLong-Context Language Model InferenceMemory-Efficient OptimizationLLM Inference AccelerationLow-Rank CompressionKV-Cache Compression
AskPDF
Aug 9, 2026

RippleKV: Cross-Layer KV Cache Allocation via Perturbation Propagation

Long-context LLM inference is bottlenecked by KV cache memory, yet distributing a limited cache budget across layers remains challenging. Existing methods rely on proxies such as layer depth, attention statistics, or representation change. These proxies do not measure how perturbations at each layer propagate to the output and may therefore cause sensitive layers to be underallocated while tolerant layers are overallocated. To address this issue, we propose RippleKV, which allocates cache across layers by estimating how perturbations to each layer's value cache affect the final predictive distribution. RippleKV independently injects norm-adaptive perturbations into each layer's value cache and measures the induced KL divergence at the model output over a small calibration set. Averaging these responses yields a sensitivity profile specific to the model that need not vary monotonically with depth. RippleKV then converts the sensitivity profile into layer budget multipliers by normalizing the sensitivity scores and applying an exponential mapping. A ratio parameter controls the allocation disparity between sensitive and tolerant layers, while a final normalization preserves the KV cache budget. Experiments on LongBench demonstrate that RippleKV achieves the highest average performance among the evaluated KV cache compression methods under matched cache budgets.

Dongjie Xu, Kai Qian, Julius +6

1Soochow University · 2Tencent Inc. · 3The Hong Kong University of Science and Technology +2

KV CachingKV-Cache ManagementKV-Cache Compression
AskPDF
Jul 7, 2026

FreqDepthKV: Frequency-Guided Depth Sharing for Robust KV Cache Compression in Long-Context LLM Inference

Long-context LLM inference is increasingly limited by the memory and bandwidth cost of KV caches, yet aggressive compression can remove the layer-specific evidence needed for retrieval and multi-step reasoning. We introduce FreqDepthKV, an inference-time cache compression method that factorizes adjacent-layer KV states into shared low-frequency depth components and sparse high-frequency residuals. A lightweight online probe assigns attention heads to shared-depth, residual-depth, or exact cache modes according to their contribution to reconstruction-sensitive attention logits, allowing the compression policy to adapt to prompt structure without retraining. Across long-context question answering, needle retrieval, summarization, and code generation benchmarks, FreqDepthKV preserves task accuracy under substantially smaller cache budgets. With a 32k-token prefill window, FreqDepthKV reaches 58.3 Exact Match, 63.0 F1, 32.5 ROUGE-L, and 48.1 pass@1, closely matching full KV while outperforming prior compressed-cache methods. It also improves decoding throughput to 70.4 tokens/s, reduces TTFT to 2.06 seconds, and lowers peak KV memory to 6.2 GB, achieving a 3.9x effective compression ratio.

Anna Córdoba, Adam Puente Tercero, Nerea Angulo Hijo +4

Instituto de Investigación en Visión Artificial

KV CachingLong-Context Language Model InferenceMemory-Efficient OptimizationLLM Inference AccelerationKV-Cache Compression
AskPDF