Probabilistic KV Routing: Adaptive Layer-wise Cache Sharing

In Transformer language model serving, KV cache preserves inference speed but greatly increases memory cost. Instead of temporal-axis compression or eviction, this research targets depth-wise cache sharing that reduces redundancy along the layer depth direction.

Existing cross-layer cache sharing tended to cause throughput degradation or increased time-to-first-token. Instead, the model is trained with random cross-layer attention so that it pre-adapts to multiple cache sharing schemes.

• During training, each layer randomly references either its own KV state or the KV state of the immediately preceding layer.
• As a result, at deployment time, some layers' caches can be removed without loss of information.
• It can be applied to both pre-training and fine-tuning, and supports depth-wise cache sharing across multiple model families.
• In larger models and data-constrained settings, it acts like regularization, maintaining or improving performance while significantly reducing cache memory.