LoopKV: Loop-Aware KV Cache Compression for Looped Transformers

Looped Transformers improve accuracy by executing the same shared module multiple times without increasing the parameter count. But for models like Ouro, each loop maintains an independent key-value (KV) cache, and cache memory becomes the bottleneck limiting batch size. The recently proposed LoopKV method achieves more efficient KV cache compression by finely perceiving the quantization sensitivity of each loop and layer. When serving Ouro-1.4B on a single RTX PRO 6000, it reduces KV storage by 35% compared with uniform 4-bit quantization (KIVI-4) and achieves a 1.7–2.2× throughput improvement under a latency service-level objective.

The KV Cache Dilemma of Looped Transformers

The core idea of looped Transformers is to apply the same block to each token multiple times, thereby improving expressive power without increasing parameters. Ouro is one representative of this class of models, and it retains an independent KV cache at each loop. Although this approach is flexible, it incurs significant memory overhead — the cache size grows linearly with the number of loops, directly limiting the batch size that can be served. Batch size in turn determines the amortization efficiency of weight reads, so cache memory becomes a key constraint on end-to-end serving throughput.

Traditional KV cache compression methods, such as uniform bit quantization and fixed per-loop precision allocation rules, cannot allocate resources according to the actual sensitivity of each layer at each loop. In other words, they treat all (loop, layer) units equally, whereas in reality the quantization sensitivity of different units varies greatly.

KV cache is an intermediate state saved during autoregressive inference to avoid recomputation: for each generated token, the model needs the key (Key) and value (Value) vectors of all previous tokens to complete the attention computation. These vectors are cached, and their size grows linearly with sequence length. In a standard Transformer, each token goes through a single forward pass, so each layer has only one KV cache; in a looped Transformer, the same block is applied repeatedly, and if an independent KV cache is retained for each loop, the total cache must be multiplied by the number of loops. This is why Ouro-style models have significantly higher memory overhead than standard models with the same parameter count. Cache occupancy directly consumes the GPU memory available for storing batch samples, and the larger the batch, the more the gains from weight reads can be amortized across more samples. Therefore, KV cache size indirectly determines inference throughput through batch size.

Quantization refers to approximating floating-point values with low-bit-width integers; for example, 4 bits can represent only 16 discrete levels, and 2 bits only 4 levels. The lower the bit width, the smaller the storage and bandwidth overhead, but the greater the rounding error introduced, which may harm model output quality. Uniform quantization uses the same bit width for all parameters; it is simple to implement but ignores differences in how sensitive different parts are to error.

Fine-Grained Observations of Quantization Sensitivity

The research team conducted measurements on Ouro and found that quantization sensitivity varies significantly across (loop, layer) units. More importantly, the sum of the damage caused by each unit individually can approximate the overall output damage. This means that the sensitivity of each unit can be evaluated independently, and compression strategies can be formulated accordingly.

Another key finding is that sharing damage also varies by unit — later loops are generally more tolerant of sharing. This provides a basis for reusing the KV cache of the final loop. In addition, the practice of skipping low-attention values to save reads introduces errors, and these errors accumulate across loops, so it must be handled carefully.

Sensitivity (quantization sensitivity) is usually defined as the degree of deviation between the model output after applying a quantization perturbation to a unit and the full-precision output, and can be measured by KL divergence, mean squared error, or task metric degradation. High sensitivity means that the unit requires a higher bit width to maintain accuracy, while low sensitivity means that a lower bit width or even a reused approximation can be used. The article's statement that "damage can be approximately additive" refers to an additivity assumption: if the sum of output deviations caused by quantizing each (loop, layer) unit individually is close to the total deviation when all units are quantized simultaneously, then sensitivity can be measured independently unit by unit in an offline stage without exhaustively enumerating combinations. This reduces the planning problem from exponential complexity to a linearly solvable one.

Sharing damage refers to the additional error introduced when multiple units reuse the same cache (for example, all reusing the KV of the final loop). The representations of later loops often tend to converge and rely less on cache precision, so reusing the final loop's KV causes less damage at these positions.

LoopKV's Three-Level Storage Strategy

Based on the above observations, LoopKV designs a one-time offline analysis procedure that selects one of three storage modes for each (loop, layer) unit: 4-bit storage, 2-bit storage, or reuse of the final loop's KV. The planner minimizes the number of reads under the constraints of predicted damage and storage budget.

In the decoding stage, each loop first reads the high 2 bits of each 4-bit value, and fetches the low bits only when that loop's own attention weights require 4-bit precision. This on-demand reading mechanism further reduces memory bandwidth pressure. The entire scheme is executed through a single fused kernel in vLLM, ensuring an efficient implementation.

vLLM is a high-throughput serving framework for large-model inference, whose core mechanisms include paged KV cache management and fused operators. A fused kernel refers to combining multiple fine-grained operations into a single GPU kernel for execution, in order to reduce kernel launch overhead and GPU memory round-trips for intermediate results. LoopKV's decoding path needs to "first read the high 2 bits, then read the low 2 bits on demand"; splitting this into multiple independent operations would introduce extra latency, so completing the reads and dequantization within a single fused kernel is crucial to meeting the latency target. CUDA graphs, meanwhile, record a sequence of fixed GPU calls into a graph and replay it as a whole, further reducing scheduling overhead. This is also the infrastructure used by both sides in the comparison experiments, ensuring that the throughput difference mainly comes from cache compression itself.

Performance and Significance

When serving Ouro-1.4B on a single RTX PRO 6000, the memory-optimized LoopKV-M scheme reduces KV storage by 35% compared with uniform 4-bit quantization (KIVI-4), which also uses fused kernels and CUDA graphs. Under a maximum length of 16K tokens and a latency service-level objective, its effective throughput (goodput) reaches 1.7 to 2.2 times that of KIVI-4. On six public benchmarks, LoopKV's average performance is not significantly lower than the BF16 baseline.

This result shows that cache compression for looped Transformers requires more fine-grained awareness. By combining offline planning with online on-demand reading, LoopKV substantially improves serving efficiency while maintaining accuracy, providing strong support for the practical deployment of looped Transformers.

Implications for AI Application Development

For developers building AI applications with looped Transformers, KV cache management is a key factor affecting inference cost and response speed. LoopKV's approach — dynamically allocating precision based on sensitivity and exploiting differences in tolerance across loops — can inspire optimizations for other long-context or iterative reasoning scenarios. Although the method targets a specific model architecture, its core principle is general: fine-grained resource allocation is often more effective than a one-size-fits-all strategy.

As looped Transformers attract more attention in scenarios that require high precision and are parameter-constrained, similar system-level optimizations will become an important force driving their adoption. LoopKV demonstrates the potential of algorithm-system co-design and points the way for future research.

One-line Summary

LoopKV compresses KV caches by sensitivity across loops and layers, boosting Ouro inference throughput by 1.7–2.2×.

Paragraph Summary

Looped Transformers improve accuracy by reusing the same module multiple times without increasing parameters, but models such as Ouro maintain an independent KV cache at each loop, and cache memory grows linearly with the number of loops, becoming a bottleneck limiting batch size and throughput. LoopKV finds that quantization sensitivity varies significantly across (loop, layer) units and that unit-level damage is approximately additive, with later loops more tolerant of cache sharing. Based on this, it offline-plans each unit for 4-bit storage, 2-bit storage, or reuse of the final loop's KV, and during decoding reads the high 2 bits first and the low bits on demand, all executed by a single fused kernel in vLLM. Serving Ouro-1.4B on a single RTX PRO 6000, LoopKV-M reduces KV storage by 35% versus uniform 4-bit quantization and improves throughput by 1.7–2.2× under the latency target, with average performance on six benchmarks not significantly below BF16.