分离预填充计算与缓存压缩,显著提升长文本生成速度
FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration
- 通过选择性传播关键词元,减少预填充计算量
- 预填充加速1.82倍,解码加速2.87倍,精度不降
- 支持灵活调节效率与精度,适合大模型推理部署
大型语言模型在处理长序列时面临预填充计算和键值缓存占用过大的问题。现有方法虽压缩缓存并加速预填充,但二者耦合导致精度下降。本文提出FastKV,利用后期层中词元重要性趋于稳定的特点,在某一层(TSP层)前进行全上下文计算,仅将最具信息量的词元传递至后续层,并独立选择用于缓存的关键条目。该设计使预填充计算量与缓存预算可独立控制,实现灵活优化。实验表明,相比全上下文基线,FastKV在预填充阶段提速1.82倍,解码阶段提速2.87倍,同时保持与仅解码基线相当的精度。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。