[KongchangAI]
· 2 min read· 1,112 words

LoopVQ: Sub-1-Bit KV Cache Quantization for Looped Transformers

LoopVQ: Sub-1-Bit KV Cache Quantization for Looped Transformers

LoopVQ compresses looped Transformer KV caches with joint vector quantization, preserving performance and greatly accelerating at 0.5 bits.

Looped Transformers reduce parameters by reusing the same layers, but each loop step requires an independent KV cache, causing linear cache growth and blocking long-context deployment. LoopVQ exploits redundancy along the loop axis by jointly encoding channels across loop steps into vectors and fine-tuning the codebook to align with full-precision outputs, outperforming standard LLM KV quantization methods at 0.5 bits/element. It is orthogonal to methods like TaSQ, retaining most BF16 precision at 0.25 bits when combined. In SGLang, 32K-context decoding latency is nearly halved, and peak throughput improves 4.3x and 6.1x for 1.4B and 2.6B Ouro models.

Looped Transformers reduce parameter memory by repeatedly applying the same layers, but they introduce a new bottleneck: each loop step must store an independent key-value (KV) cache, causing the KV cache to grow linearly with the number of loop steps. LoopVQ proposes a vector quantization method for KV caches tailored to looped architectures, maintaining model performance even at the extremely low precision of 0.5 bits/element, and substantially improving serving throughput in a SGLang implementation.

LoopVQ paper title screenshot

The KV Cache Dilemma of Looped Transformers

Looped language models reuse the same set of layers multiple times, achieving performance comparable to non-looped models with fewer parameters. This parameter efficiency advantage is significant, but the cost is that during inference, each loop step must save its own KV cache. As the number of loop steps increases, the memory occupied by the KV cache grows linearly, becoming a major obstacle to deploying long-context, long-chain reasoning tasks.

The research team observed that the KV cache of looped models is highly redundant along the loop axis, and that different loop steps vary greatly in their sensitivity to quantization. This offers two angles for compression: exploiting cross-loop correlations and allocating bit resources based on sensitivity.

KV cache is a core mechanism in autoregressive Transformer inference: when generating each new token, the model needs to recompute attention, but the key (Key) and value (Value) vectors of historical tokens were already computed in previous steps and can be cached to avoid redundant computation. A standard Transformer has only one KV cache per layer, with size proportional to context length; a looped Transformer executes the same set of layers repeatedly, and each loop is equivalent to the attention computation of a "new layer," so a separate KV cache must be stored for each loop step, making the total cache roughly the number of loop steps times that of a standard model. This is precisely the root cause of the looped architecture "saving parameters but costing cache," and the main source of GPU memory usage in long-context scenarios.

Vector quantization differs from scalar quantization in that scalar quantization rounds each value independently, whereas vector quantization groups multiple values into a vector and maps it as a whole to the nearest codeword in a pre-learned codebook. Because the codebook can capture inter-dimensional correlations, vector quantization typically preserves more information at the same bit rate, but requires additional codebook storage and nearest-neighbor search overhead. LoopVQ leverages this property to jointly encode the values of the same channel across multiple loop steps as a vector, thereby maintaining accuracy at extremely low bit widths.

LoopVQ's Core Design: Joint Encoding and Codebook Fine-Tuning

LoopVQ is a vector quantization method whose key innovation is jointly encoding channels across multiple loop steps rather than processing each step independently. By capturing redundancy along the loop axis, it preserves more information at extremely low bit widths. In addition, LoopVQ fine-tunes the codebook so that the quantized output aligns with the full-precision model's output, thereby compensating for quantization error.

This design enables LoopVQ, at 0.5 bits/element, to outperform existing KV cache vector quantization methods designed for standard LLMs on general tasks, long-chain-of-thought reasoning, and long-context retrieval tasks. Experiments were conducted based on the Ouro and Huginn models.

Codebook fine-tuning is a common post-processing technique in vector quantization. The initial codebook is typically obtained via clustering algorithms such as k-means on a small amount of calibration data, but the clustering objective (minimizing reconstruction error) is not aligned with the downstream task objective (preserving model output). LoopVQ additionally optimizes the codebook after quantization so that the attention output of the quantized model approximates that of the full-precision model as closely as possible, effectively correcting quantization error with the task loss. This approach is consistent with the idea of post-training quantization methods such as GPTQ, but the object of optimization is the codebook rather than per-weight scaling factors.

Orthogonality and Synergy with Existing Methods

LoopVQ's design is orthogonal to KV cache quantization methods for standard LLMs, meaning they can be stacked. When combined with TaSQ, LoopVQ still retains most of BF16 precision even at the extremely low bit rate of 0.25 bits/element. This compatibility provides flexible precision-efficiency trade-offs for practical deployment.

Inference Acceleration: SGLang Implementation and Throughput Gains

The research team implemented LoopVQ in SGLang. At a 32K token context, decoding latency is nearly halved. Compared with the BF16 baseline, peak serving throughput improved by 4.3x and 6.1x for the 1.4B and 2.6B Ouro models, respectively. These numbers show that LoopVQ not only compresses memory but also translates into real latency and throughput benefits.

SGLang is a high-performance serving framework for large model inference, centered on RadixAttention—using a radix tree to manage prefix KV caches shared across multiple requests, thereby reusing common prefixes and reducing redundant computation in batch serving. Integrating LoopVQ into SGLang means the quantized KV cache must be compatible with this prefix-reuse and paged memory management mechanism: the cache is stored in a more compact format and decompressed during decoding to participate in attention computation. The throughput gains come from two aspects: lower per-request GPU memory usage enables larger batches, while reduced bandwidth demand for reading KV during decoding directly shortens per-step latency.

Implications for Deploying Looped Architectures

The parameter efficiency of looped Transformers makes them a strong candidate for resource-constrained scenarios, but KV cache bloat has been an obstacle to deployment. LoopVQ approaches this from a quantization angle, maintaining model capability at sub-1-bit precision and working in concert with existing quantization methods, providing a feasible memory solution for long-context serving of looped models. As looped architectures gain attention in language modeling, compression techniques tailored to their characteristics will become increasingly important.

One-line Summary

LoopVQ uses joint vector quantization to compress the KV cache of looped Transformers, preserving performance and greatly accelerating at 0.5 bits.

Paragraph Summary

Looped Transformers reduce parameters by reusing the same set of layers, but each loop step must store an independent KV cache, causing the cache to grow linearly with the number of loop steps and becoming a major obstacle to long-context deployment. LoopVQ exploits high redundancy along the loop axis by jointly encoding channels across multiple loop steps into vectors for quantization, and fine-tunes the codebook to align with full-precision outputs, outperforming KV quantization methods for standard LLMs on general tasks, long-chain reasoning, and long-context retrieval at 0.5 bits/element. The method is orthogonal to standard quantization methods such as TaSQ, and when stacked, 0.25 bits still retains most of BF16 precision. In a SGLang implementation, decoding latency at 32K context is nearly halved, and peak throughput improves by 4.3x and 6.1x for the 1.4B and 2.6B Ouro models, respectively.

#LoopVQ#KV cache quantization#looped Transformer#vector quantization#SGLang#sub-1-bit quantization
Share:

Related articles