Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing - Apple Machine Learning Research
Machine Learning Research
Open MenuClose Menu
Overview
Research Highlights
Publications
Events
Work with us
research area Methods and Algorithms, research area Speech and Natural Language Processing
content type paperpublished May 2026
Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing
AuthorsAnastasiia Filippova, David Grangier, Marco Cuturi, João Monteiro
View publication
Copy Bibtex
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.
Related readings and updates.
EpiCache: Episodic KV Cache Management for Long Conversational Question Answering
September 23, 2025research area Methods and Algorithms, research area Speech and Natural Language Processingconference ICML
Recent advances in large language models (LLMs) have extended context lengths, enabling assistants to sustain long histories for coherent, personalized responses. This ability, however, hinges on Key-Value (KV) caching, whose memory grows linearly with dialogue length and quickly dominates under strict resource constraints. An active line of research for reducing this overhead is KV cache compression, which seeks to limit cache size while…
Read more
KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
May 14, 2024research area Methods and Algorithms, research area Speech and Natural Language Processingconference ICML
Large Language Model or LLM inference has two phases, the prompt (or prefill) phase to output the first token and the extension (or decoding) phase to the generate subsequent tokens. In this work, we propose an efficient parallelization scheme, KV-Runahead to accelerate the prompt phase. The key observation is that the extension phase generates tokens faster than the prompt phase because of key-value cache (KV-cache). Hence, KV-Runahead…
Read more
Discover opportunities in Machine Learning.
Our research in machine learning breaks new ground every day.
Work with us
Machine Learning Research
Publications
Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing
Privacy Policy
Terms of Use
Legal
Copyright © 2026 Apple Inc. All rights reserved.