Looped models share memory caches to reduce inference costs
BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing
Machine Learning
Summary
Looped models run the same layers repeatedly to improve quality without adding parameters, but this increases memory use because each repetition stores its own cache. The authors study how to let looped models share these caches safely to save memory. They first train models with a technique called stochastic depth to make outputs similar across repeats, then fine-tune them with a method named BeNice! which helps the model use a shared cache without losing quality. This reduces memory growth during use, makes the models faster, and cuts down waiting times when generating text.
What this means in practice
- •For machine learning engineers: Reduce memory usage and increase throughput during inference of looped transformer models by enabling KV-cache sharing with BeNice!.
- •For data center operators: Lower resource demands for hosting large language models by integrating KV-sharing to decrease memory footprint and speed up responses.
Authors
Hayden Prairie, Jacob M. Ryan, Taylor Berg-Kirkpatrick, Daniel Y. Fu
Abstract
Looped architectures scale quality while keeping parameter count fixed by repeatedly passing activations through the same block of layers. During inference, each recurrence generates independent KV-caches, causing memory to grow linearly. To address this challenge, prior work has proposed KV-sharing (i.e., sharing KV-caches across depth), finding that, in some cases, the KV-cache of the final recurrence can be shared with the earlier recurrences without quality loss, a property we denote as KV-shareability. However, the mechanisms underlying KV-shareability remain poorly understood, with even identically trained looped architectures not consistently exhibiting it. In this paper, we study how KV-shareability can be enabled in looped models. We propose a two-stage method: in the first stage, we use stochastic depth during pretraining, which induces contractive residual dynamics that drive representations toward a common limit across recurrences. In the second stage, we propose BeNice!, a post-training method that broadcasts the final recurrence's KV-cache into the first recurrence, training it to produce compatible residuals while preserving the original model's quality. As a result, BeNice! enables KV-shareability while matching standard finetuning's unshared validation perplexity. Using BeNice!, we show that KV-sharing is complementary to alternative KV-compression methods, composing with memory-efficient sequence mixers (e.g., multi-latent attention, Kimi delta attention) and low-precision KV-cache quantization (e.g., fp8, mxfp8, nvfp4). Finally, in simulated low-memory, long-context inference, KV-sharing improves cache hit rates by up to 31.4%, increases throughput by up to 23.6%, and reduces p90 time-to-first-token by up to 83.6%.