arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2502.01068cs.LGcs.CL

FastKV: 解耦上下文缩减与KV缓存压缩以加速预填解码

FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration

Seoul National University(首尔国立大学);Sungkyunkwan University(庆州大学)

Dongwon Jo, Jiwon Song, Yulhwa Kim, Jae-Joon Kim

AI总结:

FastKV通过解耦上下文缩减与KV缓存压缩,提升预填和解码效率,实验显示在预填和解码阶段分别获得1.82倍和2.87倍加速,同时保持解码-only基线的准确性。

摘要:

尽管大型语言模型(LLMs)擅长处理长上下文序列,但它们需要大量的预填计算和键值(KV)缓存,这在预填和解码阶段都会严重影响计算效率和内存使用。最近一些压缩KV缓存并加速预填的工作虽然降低了成本,但无意中将预填计算减少与解码KV预算联系在一起。这种耦合源于忽视了关键上下文的层依赖性变化,通常导致精度下降。为了解决这个问题,我们引入了FastKV,一种KV缓存压缩框架,旨在通过利用后续层中令牌重要性的稳定来减少预填和解码的延迟。FastKV在Token-Selective Propagation(TSP)层进行全上下文计算,该层仅将最信息性的令牌传递给后续层。从这些传递的令牌中,FastKV独立选择显著的KV条目进行缓存,从而基于TSP决策解耦KV预算与预填计算减少。这种对TSP率和KV保留率的独立控制使效率和精度的优化更加灵活。实验结果表明,与全上下文基线相比,FastKV在预填和解码阶段分别实现了1.82倍和2.87倍的加速,同时保持解码-only基线的准确性。我们的代码可在https://github.com/dongwonjo/FastKV上获取。

英文摘要:

While large language models (LLMs) excel at handling long-context sequences, they require substantial prefill computation and key-value (KV) cache, which can heavily burden computational efficiency and memory usage in both prefill and decoding stages. Recent works that compress KV caches with prefill acceleration reduce this cost but inadvertently tie the prefill compute reduction to the decoding KV budget. This coupling arises from overlooking the layer-dependent variation of critical context, often leading to accuracy degradation. To address this issue, we introduce FastKV, a KV cache compression framework designed to reduce latency in both prefill and decoding by leveraging the stabilization of token importance in later layers. FastKV performs full-context computation until a Token-Selective Propagation (TSP) layer, which forwards only the most informative tokens to subsequent layers. From these propagated tokens, FastKV independently selects salient KV entries for caching, thereby decoupling KV budget from the prefill compute reduction based on the TSP decision. This independent control of the TSP rate and KV retention rate enables flexible optimization of efficiency and accuracy. Experimental results show that FastKV achieves speedups of up to 1.82$\times$ in prefill and 2.87$\times$ in decoding compared to the full-context baseline, while matching the accuracy of the decoding-only baselines. Our code is available at https://github.com/dongwonjo/FastKV.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑