论文

RoLaQ:面向 2-bit KV 缓存的旋转与范围联合校准方法

RoLaQ: Joint Rotation and Range Calibration for 2-Bit KV Caches in Multi-Head Latent Attention

精选理由

论文提出 RoLaQ,让 MLA 模型的 KV 缓存压到 2-bit,在 DeepSeek-V2-Lite 上吞吐提升 2.51 倍,做长上下文推理的可以看看。

RoLaQ 是一种离线的旋转与范围联合校准方法,面向多头潜在注意力(MLA)模型的 2-bit KV 缓存量化,校准过程中模型权重保持冻结。该方法在由冻结的 key/value 投影定义的子空间内优化共享旋转,并联合学习与重构增益耦合的量化范围。在 3 个 MLA 模型和 6 个基准上,RoLaQ 在有无 token 保护的设置下均取得对比的 2-bit 方法中最高的平均准确率。在 DeepSeek-V2-Lite 上,基于 SGLang 的实现在相同显存预算下提供 3.96 倍于 BF16 的 KV 容量,在 256K 输入 token 下解码吞吐达到 2.51 倍。

原文 · arXiv: DeepSeek

RoLaQ: Joint Rotation and Range Calibration for 2-Bit KV Caches in Multi-Head Latent Attention

Multi-head latent attention (MLA) compresses the content of a key-value cache into a shared latent per token, making low-bit storage attractive for long-context inference. At two bits, errors in this representation perturb both attention matching and value readout, with effects that depend on error direction and reconstruction scale. We introduce RoLaQ (Rotation-calibrated Latent Quantization), an offline method for joint rotation and range calibration. RoLaQ refines a shared rotation within subspaces defined by the frozen key and value projections, and jointly learns quantization ranges whose thresholds and reconstruction gains are coupled to control the shared signal's response. Prediction matching and attention-output supervision guide calibration with model weights frozen. Across three MLA models and six benchmarks, RoLaQ achieves the highest average accuracy among the compared 2-bit methods both without and with token protection. On DeepSeek-V2-Lite, an SGLang implementation provides 3.96 times the BF16 KV capacity under the same memory budget and 2.51 times its decode throughput with 256K input tokens.