KV Routing Stocastico: Abilitazione della Condivisione Adattiva della Cache in Profondità

Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing

April 3, 2026

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

cs.AI

Abstract

La gestione di modelli linguistici transformer ad alta produttività richiede la memorizzazione nella cache di chiavi e valori (KV) per evitare calcoli ridondanti durante la generazione autoregressiva. L'ingombro di memoria della cache KV è significativo e influisce pesantemente sui costi di servizio. Questo lavoro si propone di ridurre questi requisiti di memoria. Mentre i recenti studi hanno affrontato la riduzione della cache KV principalmente attraverso la compressione e l'evizione lungo l'asse temporale, noi sosteniamo che la dimensione della profondità offre un percorso di ottimizzazione ortogonale e robusto. Sebbene ricerche precedenti suggeriscano che una cache completa per ogni livello sia ridondante, l'implementazione della condivisione della cache tra livelli rimane una sfida pratica; i metodi esistenti soffrono tipicamente di una ridotta produttività o di un aumento del tempo per il primo token. In questo articolo, dimostriamo che l'eliminazione della cache di un livello offre un'ottimizzazione efficiente senza perdita di informazioni. Proponiamo un semplice approccio di addestramento: l'attenzione casuale tra livelli. Durante l'addestramento, i livelli scelgono casualmente di attendere sia ai propri stati KV che a quelli di un livello precedente. Questo processo stocastico adatta il modello per essere robusto a varie strategie di condivisione della cache lungo la profondità, garantendo flessibilità per vincoli hardware sconosciuti al momento del deployment. Le nostre valutazioni mostrano che l'applicazione di questo schema durante il pre-addestramento o il fine-tuning abilita la condivisione della cache lungo la profondità per varie famiglie di modelli. Inoltre, per modelli più grandi in contesti con dati limitati, questo approccio suggerisce un effetto simile alla regolarizzazione, preservando o migliorando frequentemente le prestazioni mentre riduce significativamente l'ingombro di memoria della cache.

English

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.