FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates

2026-09-24

AI Summary

Recurrent Transformers suffer from escalating computational and memory overheads as depth increases, severely limiting their practical deployment efficiency. To address this challenge, this work proposes FlashLoop, a training-free inference framework that exploits cross-loop redundancy to achieve efficient computation and storage optimization. Specifically, the framework identifies sparsity patterns in state transitions, attention discrepancies, and KV residuals, leveraging these observations to design token-sparse updates, sparse attention mechanisms, and low-bit KV residual quantization techniques. Experimental results demonstrate that FlashLoop achieves up to 1.64× end-to-end speedup and a 6× reduction in KV cache memory consumption while preserving lossless accuracy.

Abstract

Looped Transformers have attracted substantial attention as a parameter-efficient approach to increasing computational depth through repeated application of shared Transformer blocks. However, their practical advantages over conventional Transformers remain under debate: each additional loop incurs another Transformer pass and requires caching another set of KV states, causing inference FLOPs and KV-cache memory to grow continuously with loop depth. This overhead becomes particularly severe at large loop counts and long context, preventing the parameter efficiency of Looped Transformers from translating into practical inference efficiency. In this paper, we find that much of the additional computation and storage introduced by looping is redundant. As recurrence proceeds, state changes become increasingly concentrated on a small subset of tokens; attention-output differences are dominated by a sparse and stable subset of key columns; and KV residuals between adjacent loops become progressively more amenable to low-bit quantization. Building on these observations, we introduce FlashLoop, a training-free inference framework that reduces cross-loop redundancy through token-sparse updates, sparse attention, and KV-residual quantization. Across several Looped Transformers models, \textsc{FlashLoop} delivers lossless accuracy while achieving up to 1.64$\times$ end-to-end speedup and up to 6$\times$ KV-cache memory reduction, substantially improving the practicality of scaling Looped Transformers to greater computational depths and longer context.

Problem

• Looped Transformers
• inference efficiency
• KV-cache memory
• computational overhead
• redundancy

Innovation

• Looped Transformers
• Lazy Updates
• Token-Sparse Updates
• Sparse Attention
• KV-Residual Quantization