RIFT: KV Cache Compression for Draft Models with Complementary Compressibility Training

概要

Speculative decoding is a popular decode acceleration method. As such, it is particularly relevant to long-output generation, which requires an increasingly large KV cache. KV cache compression offers a way to constrain costs in long-context drafting. However, compression methods that focus on prefill information can lose relevance as the output length grows. We term this phenomenon decode staleness and show that it degrades draft model accuracy under selection-based KV cache compression. We introduce RIFT, a family of selection-based KV cache compression strategies that balance prefill information and recent generated tokens. We show that refreshing the KV cache selection during token generation significantly improves the drafter's quality at the cost of additional decode latency. To alleviate the refresh costs, we propose an adaptive KV refresh strategy that reduces the number of score calculations while retaining the drafter's quality. We further demonstrate how a compressibility-driven drafter training phase, based on sparse attention masks, can improve the drafter's quality under a compressed KV cache. Combined with compressibility training, RIFT enables Qwen3.5-4B, drafting for Qwen3.5-122B, to retain over 98% of its full-cache acceptance length on 32K-token inputs, 4K-token outputs, and an 8K-token KV cache budget.

参考リンク

• Paper ↗

リンク先のライセンス・再現性・データ条件をご確認ください。自動分類や実装の目安は保証ではありません。

実装を始めるプロンプト

参考リンクを確認したうえで、次の内容をコピーして利用できます。

# 論文・記事
RIFT: KV Cache Compression for Draft Models with Complementary Compressibility Training

# 概要
Speculative decoding is a popular decode acceleration method. As such, it is particularly relevant to long-output generation, which requires an increasingly large KV cache.

# 原典
https://arxiv.org/abs/2610.31098

上記の手法を実装してください。実行前に依存関係とライセンスを確認してください。