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

a training-free inference framework that reduces cross-loop redundancy through token-sparse updates, sparse attention, and KV-residual quantization

Wanqi Yang^1,2,3 Shiwei Liu^1,2,3

¹ Max Planck Institute for Intelligent Systems · ² ELLIS Institute Tübingen · ³ Tübingen AI Center

Code BibTeX

FlashLoop MMLU comparison demo looping at 2× playback speed

Quick Start

Run FlashLoop with Ouro

Install

pip install flashloop
flashloop-build-kernels

Python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from flashloop import FlashLoopEngine

model_path = "/path/to/Ouro-1.4B"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
input_ids = tokenizer.apply_chat_template(
    [{"role": "user", "content": "What is the capital of France?"}],
    tokenize=True, add_generation_prompt=True, return_tensors="pt",
).cuda()

# Optimized engine.
engine = FlashLoopEngine.from_pretrained(model_path)
output_ids = engine.generate(input_ids, max_new_tokens=64)
print(tokenizer.decode(output_ids[0, input_ids.shape[1]:], skip_special_tokens=True))
del engine
torch.cuda.empty_cache()

Three Complementary Components

Cross-Loop Token-Sparse Updates

FlashLoop selectively updates only the tokens that remain active across loops.

Loop-Aware Sparse Attention

FlashLoop uses the preceding loop to identify a small subset of columns for recomputation.

Cross-Loop KV Residual Quantization

FlashLoop therefore represents the KV cache as a quantized base state followed by quantized residual updates across loops.

Performance-matched Looped and Non-looped Models

FlashLoop substantially reduces both computation and KV-cache memory for looped Transformers.

Cross-Loop Redundancy in Looped Transformers

Recurrent refinement becomes structurally redundant along three complementary axes.

As recurrence progresses, changes between adjacent loops become increasingly sparse, predictable, and compressible.

01 / TOKENS

Token-update redundancy

Updates to most token states become increasingly small, with only a small fraction of token rows accounting for the majority of changes in hidden states and KV representations.

02 / ATTENTION

Attention-computation redundancy

Cross-loop differences in attention outputs are dominated by a sparse subset of columns whose indices can be reliably predicted from the preceding loop.

03 / MEMORY

KV-storage redundancy

KV residuals between adjacent loops are substantially more friendly to quantization than full KV states, exhibiting lower reconstruction error under identical quantization settings.

**Recurrent refinement becomes structurally redundant along three complementary axes.**

**Additional evidence for cross-loop redundancy across different looped models.**

The Method

FlashLoop

Three Coupled Compression Strategies

FlashLoop includes three coupled compression strategies: token-sparse updates, sparse attention, and KV residual quantization.

STEP 01

Cross-Loop Token-Sparse Updates

Active tokens generate new K/V states for the current loop, whereas inactive tokens continue to reuse their KV cache from preceding loop.

STEP 02

Loop-Aware Sparse Attention

FlashLoop uses the preceding loop to identify a small subset of columns for recomputation. Thus, FlashLoop caches their global probability mass, i.e., the total probability assigned to the selected columns under the full attention distribution of the preceding loop.

STEP 03

Cross-Loop KV Residual Quantization

The K/V states from the first loop are quantized as the base, while each subsequent loop stores only the residual relative to the reconstructed state of the preceding loop.

Dynamic token and column selection can introduce irregular memory accesses and fragmented computation, while separate cache-update and quantization steps may add additional data movement and kernel-launch overhead. FlashLoop addresses these issues by gathering selected tokens and keys into compact, compute-friendly buffers to preserve regular GEMM execution, and by fusing cross-loop KV updates and quantization into the attention kernel.

Experiments

Main Results

We evaluate FlashLoop on five representative Looped Transformers: four models from the Ouro family and Huginn-3.5B. We evaluate FlashLoop on five benchmarks: MATH-500, GSM8K, ARC-Challenge, HellaSwag and WinoGrande, using the EleutherAI LM Harness.

up to 1.64× end-to-end speedup

up to 6× KV-cache memory reduction

Model	Δ Avg.	Speedup ↑	KV Memory ↑
Ouro-1.4B	+0.37	1.59×	5.85×
Ouro-1.4B-Thinking	−0.88	1.59×	5.85×
Ouro-2.6B	−0.13	1.64×	6.06×
Ouro-2.6B-Thinking	−0.56	1.64×	6.06×
Huginn-3.5B	+0.17	1.52×	5.18×

Δ Avg. denotes the average-score difference from the original model in percentage points. KV Memory means the reduction of peak KV cache memory relative to original models.

**Efficiency analysis of each component.** We cumulatively enable cross-loop token-sparse updates, loop-aware sparse attention, and 4-bit KV residual quantization, reporting KV-cache memory, prefill latency, and decode latency.

**Accuracy–efficiency trade-offs of FlashLoop hyperparameters on Ouro-1.4B.** We vary the retained attention-key ratio (left), Loop-3/4 token-update ratios (middle), and KV residual quantization precision (right), while keeping the remaining components fixed.

**Scaling with context length.** FlashLoop delivers increasingly substantial memory savings and speedups as the context length grows.

**Scaling with loops.** As the number of loops increases, FlashLoop progressively reduces the accumulated prefill FLOPs and KV-cache memory of both Ouro-1.4B and Huginn-3.5B.

Conclusion

In this paper, we find that much of the additional computation and storage introduced by looping is redundant.

FlashLoop includes three coupled compression strategies: token-sparse updates, sparse attention, and KV residual quantization.

Across several Looped Transformers models, FlashLoop largely preserves model accuracy while achieving up to 1.64× end-to-end speedup and up to 6× KV-cache memory reduction.

Cite This Work

@misc{yang2026flashloop,
  title={FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates},
  author={Wanqi Yang and Shiwei Liu},
  year={2026}
}