Papers

BeNice!:让循环架构共享 KV 缓存,推理内存大幅下降

BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing

Why we picked it

循环模型每次循环都要单独存一份 KV 缓存,内存爆炸。这篇论文的两阶段训练方法让缓存能跨层共享,首 token 时间砍了 83.6%,做长推理的可以看看。

循环架构通过反复使用同一层块来保持参数量不变,但每次循环都会生成独立 KV 缓存,内存随深度线性增长。论文提出两阶段方法:预训练阶段用随机深度诱导收敛的残差动态,后训练阶段用 BeNice! 把末次循环的 KV 缓存广播给首次循环。该方法在匹配标准微调验证困惑度的同时实现 KV 共享,并与多隐注意力和 fp8 等 KV 量化方法兼容。模拟低内存长上下文推理中,缓存命中率提升最高 31.4%,吞吐量提升最高 23.6%,p90 首 token 时间缩短最高 83.6%。

Original · arXiv cs.LG

BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing

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%.