Roteamento KV Estocástico: Permitindo o Compartilhamento Adaptativo de Cache em Profundidade

Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing

April 3, 2026

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

cs.AI

Resumo

A otimização do throughput no serviço de modelos de linguagem baseados em transformers requer o *caching* de pares Chave-Valor (KVs) para evitar computação redundante durante a geração autoregressiva. A pegada de memória do cache de KVs é significativa e impacta fortemente os custos de serviço. Este trabalho propõe a redução desses requisitos de memória. Enquanto pesquisas recentes abordaram a redução do cache de KV principalmente por meio de compressão e evição ao longo do eixo temporal, argumentamos que a dimensão de profundidade oferece uma via de otimização ortogonal e robusta. Embora estudos anteriores sugiram que um cache completo para cada camada é redundante, a implementação do compartilhamento de cache entre camadas permanece um desafio prático; os métodos existentes geralmente sofrem com throughput reduzido ou aumento do tempo para o primeiro *token*. Neste artigo, demonstramos que descartar o cache de uma camada oferece uma otimização eficiente sem perda de informação. Propomos uma abordagem de treinamento simples: atenção cross-layer aleatória. Durante o treinamento, as camadas escolhem aleatoriamente atentar aos seus próprios estados de KV ou aos de uma camada precedente. Este processo estocástico adapta o modelo a ser robusto para várias estratégias de compartilhamento de cache na dimensão de profundidade, garantindo flexibilidade para restrições de hardware desconhecidas no momento da implantação. Nossas avaliações mostram que aplicar este esquema durante o pré-treinamento ou *fine-tuning* permite o compartilhamento de cache em profundidade para várias famílias de modelos. Além disso, para modelos maiores em contextos com restrição de dados, esta abordagem sugere um efeito similar à regularização, frequentemente preservando ou melhorando o desempenho enquanto reduz significativamente a pegada de memória do 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.