Papers

LoopVQ:面向 Looped Transformer 的亚 1-bit KV Cache 量化方法

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

Why we picked it

.looped 模型参数省但 KV cache 爆炸,这篇提出 0.5bit 量化方案,SGLang 实测吞吐最高涨 6.1 倍,做部署的可以看看。

Looped transformer 通过跨多个 loop 步骤复用同一组层,参数量小但 KV cache 随 loop 步数线性增长。LoopVQ 利用 KV cache 在 loop 轴上的冗余性,联合编码多个 loop 步骤的通道并微调 codebook。在 0.5 bit 每元素下,LoopVQ 在通用、长链推理和长上下文检索任务上超过了面向标准 LLM 的现有 KV cache 向量量化方法,测试对象为 Ouro 和 Huginn 模型。与 TaSQ 结合后在 0.25 bit 每元素下仍保留大部分 BF16 精度。作者在 SGLang 中实现该方法,32K 上下文下解码延迟近乎减半,1.4B 和 2.6B Ouro 模型的峰值吞吐分别提升 4.3 倍和 6.1 倍。

Original · arXiv cs.LG

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

Looped transformers are a parameter-efficient alternative to standard transformers for language modeling. By reapplying the same layers across multiple loop steps, recent looped language models match the performance of non-looped models several times their size. While this architecture effectively reduces parameter memory, each loop step still stores its own keys and values, so the key-value (KV) cache grows linearly with the number of loop steps and becomes the main memory bottleneck. We observe that the KV cache of looped models is highly redundant along the loop axis and that the loop steps differ substantially in their sensitivity to quantization. Based on these observations, we propose $\textbf{LoopVQ}$, a KV cache vector quantization method for looped transformers that codes channels of multiple loop steps jointly, and fine-tunes the codebooks to match the output of the full-precision model. At 0.5 bits per element, LoopVQ outperforms existing KV cache vector quantization methods for standard LLMs on general, long chain-of-thought reasoning, and long-context retrieval tasks with Ouro and Huginn models. Its design is orthogonal to these methods, and combined with TaSQ, it retains much of the BF16 accuracy even at 0.25 bits per element. We also implement LoopVQ in SGLang, where it nearly halves the decode latency at a 32K-token context and raises the peak serving throughput by up to $4.3\times$ and $6.1\times$ on the 1.4B and 2.6B Ouro models compared to the BF16 baseline.