原题:LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers
循环 Transformer 反复调用同一组层,用较少参数就能追上大几倍的普通模型,但每个循环步都要存自己的 KV,缓存随循环步数线性增长,成了主要的内存瓶颈。作者观察到 KV 在循环轴上高度冗余,而不同循环步对量化的敏感度差别很大,于是提出 LoopVQ:把连续几个循环步的通道合成一个子向量、共用一个码本索引(循环轴向量量化),再微调 Fisher 加权 k-means 码本,让量化后模型的输出贴近全精度模型(训练时每个循环步都读量化缓存)。它和面向普通 LLM 的方法正交,可与 TaSQ 的键变换叠加。
点赞 8 · 代码星数 3
理由
- 作者称每元素 0.5 比特时,在 Ouro 和 Huginn 模型的通用、长思维链推理和长上下文检索任务上,优于现有面向普通 LLM 的 KV 向量量化方法;叠加 TaSQ 后到 0.25 比特仍保留大部分 BF16 精度。
- 系统收益落到了实处:作者在 SGLang 里实现后,32K 上下文下解码延迟几乎减半,Ouro 1.4B / 2.6B 的峰值服务吞吐分别提到 BF16 基线的 4.3 倍和 6.1 倍。
- 代码 Apache-2.0:校准流程、带打包索引存储和融合解码内核的 SGLang 分支、论文每张表的复现脚本都在仓库里。
用法
- ./setup_envs.sh 建三个环境(loopvq-serve 引擎、loopvq-eval 评测、loopvq-build 校准),export MODELS_DIR 指向 Ouro-1.4B、Ouro-2.6B、Ouro-*-Thinking 和 huginn-0125;复现通用表:MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.sh,再 MODEL=ouro_1p4b bash experiments/loopvq_general/run_all.sh,最后 python3 experiments/loopvq_general/make_table.py。
- 起服务:在 loopvq-serve 环境里 PREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <码本目录> --port 30000 --disable-piecewise-cuda-graph,得到 SGLang 的 OpenAI 兼容接口;开头和最近的 token 保持 bf16。
风险
- 只适用于循环 Transformer(论文用 Ouro 1.4B / 2.6B 和 Huginn-3.5B),普通 LLM 不能直接套用;校准只用了 16 条 GPQA 序列。
- 延迟和吞吐数字来自作者环境;抱脸点赞 8、代码星数 3(只代表关注度)。