通过自适应秩控制，实现高达75%的KV缓存压缩。
STAR-KV: Low-Rank KV Cache Compression via Soft Thresholding for Adaptive Rank Control
- 引入可微分阈值机制，动态调节每层注意力头与块的压缩秩。
- 在多个LLM上实现最高75%缓存压缩，结合量化达20倍缩减。
- 适合需要高效推理的部署场景，尤其对显存敏感的应用。
低秩投影已成为利用隐藏维度冗余压缩KV缓存的有前景方法。然而，先前方法依赖固定或启发式秩选择，在实现激进压缩时难以保持精度。本文提出STAR-KV，一种具备细粒度秩控制的自适应低秩缓存压缩框架。该框架包含：1）可微分阈值机制，支持在注意力头和块级别实现最优秩选择；2）混合分解策略，根据键值投影的敏感性采用不同低秩因子分解方式；3）面向低秩的混合精度量化，利用数据统计实现近乎无损的低位量化。在多个LLM和基准测试上评估，STAR-KV实现最高75%的KV缓存压缩，结合量化后整体缓存减少达20倍。基于定制Triton GPU内核，其使注意力模块加速6.9倍，端到端生成吞吐提升3.1倍。代码已开源：https://github.com/PriyanshBhatnagar/STAR-KV。
原文摘要 · Abstract (English)
Low-rank projection has emerged as a promising approach for compressing the KV cache by exploiting hidden-dimension redundancy. However, prior methods rely on fixed or heuristic rank selection and struggle to achieve aggressive compression with minimal accuracy degradation. We propose STAR-KV, an adaptive low-rank KV cache compression framework with fine-grained rank control. STAR-KV encompasses 1) a differentiable thresholding mechanism that enables optimal rank selection at both attention-head and block levels, 2) a hybrid decomposition strategy that applies different low-rank factorizations according to the sensitivity of key and value projections, and 3) a low-rank-aware mixed precision quantization that leverages data statistics for near lossless low-bit quantization. Evaluated across multiple LLMs and benchmarks, STAR-KV achieves up to 75% KV cache compression and up to 20x overall KV cache reduction when combined with quantization. Enabled by custom Triton-based GPU kernels, STAR-KV delivers up to 6.9x speedup for the attention module and 3.1x end-to-end generation throughput. Our code is publicly available at: https://github.com/PriyanshBhatnagar/STAR-KV.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。