RIFT: KV Cache Compression for Draft Models with Complementary Compressibility Training
Speculative decoding draft models face severe KV cache bottlenecks during long-output generation, but standard prefill-based compression methods lose relevance as decoding progresses. This paper identifies that "decode staleness" phenomenon and proposes RIFT, a family of selection-based KV compression strategies that balance important prefill tokens with a recent-token sliding window, plus adaptive decode-time refresh and a complementary compression-aware distillation training phase. The core claim is that these methods allow small external drafters to retain over 96% of their full-cache acceptance length under tight memory budgets.
The paper makes a solid, empirically grounded contribution: it isolates decode staleness as a key failure mode for prefill-only KV compression in speculative decoding, and it addresses the problem with simple, effective mechanisms backed by systematic benchmarks. However, the work's immediate relevance is narrower than presented because it deliberately excludes tail-augmented drafters (MTP, EAGLE), which currently dominate GPU-only deployments, and because the strongest results rely on assumptions—full KV history retained elsewhere for rescoring and per-setting calibrated adaptive policies—that complicate real-world deployment.
The characterization of decode staleness is well-supported with both architecture-level and mechanistic evidence: the authors show that SnapKV retains a shrinking fraction of attention mass while sliding window retains more as generation progresses, and that this gap widens in hybrid-attention models and reasoning traces. The proposed RIFT methods are convincingly evaluated across multiple model families (Qwen, Llama, Nemotron, GPT-OSS) and cross-family settings, with RIFT-dyn and RIFT-ada consistently outperforming SnapKV, StreamingLLM, H2O, and TriAttention. The compression-aware distillation training provides complementary gains of up to 16% additional acceptance length at tight budgets.
First, the scope is explicitly limited to external drafts and self-speculation, so the gains do not apply to the tail-augmented drafters that currently dominate latency-sensitive GPU serving. Second, RIFT-dyn and RIFT-ada assume the full draft KV history remains available in HBM or offloaded memory for rescoring, meaning the active budget limits the drafting device's footprint but does not reduce the system's total stored KV memory. Third, on a naive GPU setup the rescoring latency overhead reaches up to 6.0% of draft decode time, and its magnitude in the paper's target disaggregated/SRAM-accelerator deployments is unmeasured. Fourth, the adaptive refresh controller is calibrated separately for each draft/target pair and temperature, requiring per-setting threshold tuning that limits out-of-the-box transferability.
The empirical evidence is extensive within its chosen benchmark: the authors report acceptance-length comparisons on SPEED-Bench across low- and high-entropy splits, 32K and 100K input lengths, multiple KV budgets, temperatures, draft lengths, and block sizes. The comparison to H2O, StreamingLLM, and TriAttention is fair and generally favors RIFT-dyn/ada, though H2O sometimes beats RIFT-stat. That said, the evaluation exclusively uses acceptance length rather than end-to-end wall-clock speedup, which is the metric that ultimately matters for speculative decoding. Claims such as retaining 96–98% of full-cache acceptance length are accurate to the reported tables, but without measured end-to-end speedups in realistic serving conditions, the operational value remains projected rather than demonstrated.
The paper provides substantial implementation detail: budget sizes, observation windows, block sizes, refresh intervals, batch size, training steps, and the distillation loss $\mathcal{L}=\lambda_{1}\operatorname{CE}(q_{\theta},y)+\lambda_{2}D_{\mathrm{KL}}(p\|q_{\theta})$ are all specified, with the full training recipe deferred to Appendix J.2. The SPEED-Bench benchmark is referenced as the evaluation substrate, and many appendices contain ablations. However, no code, data pipelines, or model checkpoints are mentioned as publicly released, which is a significant barrier to independent reproduction. In addition, reproducing RIFT-ada would require replicating the per-model, per-temperature calibration procedure described in Appendix E.3, including the choice of threshold percentiles and signal aggregations, for which no automation or general rule is provided.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.