Routage KV Stochastique : Permettre le Partage Adaptatif du Cache en Profondeur

Résumé

L'optimisation du débit des modèles de langage de type transformer nécessite la mise en cache des paires Clé-Valeur (KVs) pour éviter les calculs redondants lors de la génération auto-régressive. L'empreinte mémoire de la mise en cache des KVs est considérable et impacte fortement les coûts de service. Ce travail propose de réduire ces besoins en mémoire. Alors que les travaux récents ont largement abordé la réduction du cache KV via la compression et l'éviction le long de l'axe temporel, nous soutenons que la dimension en profondeur offre une voie d'optimisation orthogonale et robuste. Bien que des recherches antérieures suggèrent qu'un cache complet pour chaque couche est redondant, la mise en œuvre du partage de cache inter-couches reste un défi pratique ; les méthodes existantes souffrent généralement d'une réduction du débit ou d'une augmentation du temps jusqu'au premier token. Dans cet article, nous démontrons que supprimer le cache d'une couche offre une optimisation efficace sans perte d'information. Nous proposons une approche d'entraînement simple : l'attention inter-couches aléatoire. Pendant l'entraînement, les couches choisissent aléatoirement de porter leur attention soit sur leurs propres états KV, soit sur ceux d'une couche précédente. Ce processus stochastique adapte le modèle pour qu'il soit robuste à diverses stratégies de partage de cache en profondeur, garantissant une flexibilité face à des contraintes matérielles inconnues au moment du déploiement. Nos évaluations montrent que l'application de ce schéma lors du pré-entraînement ou du fine-tuning permet le partage de cache en profondeur pour diverses familles de modèles. De plus, pour les modèles plus grands dans des contextes où les données sont limitées, cette approche suggère un effet semblable à une régularisation, préservant ou améliorant fréquemment les performances tout en réduisant significativement l'empreinte mémoire du 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.