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

Paper

Dor Tsur, Maor Ashkenazi, Maximilian Golub, Peter Hicks, Yaniv Galron, Talor Abramovich, Yoav Miron, Tiyasa Mitra, Bita Darvish Rouhani, Ran Zilberstein

NVIDIA

TL;DR: We introduce RIFT, a KV-cache selection method for external drafters in speculative decoding. RIFT combines importance-selected positions with recent generated tokens and supports both static selection and dynamic refreshes during generation. We also introduce compressibility training to improve drafter performance under KV-cache compression. Together, RIFT and compressibility training retain over 98% of the original drafter’s full-cache acceptance length for Qwen3.5-4B → Qwen3.5-122B and over 97% for Nemotron3.5-Lightning → Nemotron3-Ultra, using 32K-token inputs, 4K-token outputs, and an 8K-token active KV-cache budget.

Why draft KV compression matters

Speculative decoding uses an inexpensive draft model to propose several tokens that a larger target model verifies in a single pass. Its effectiveness depends on a draft that is both cheap to run and closely aligned with the target. Long reasoning traces and agentic interactions put pressure on both requirements: the draft repeatedly accesses a growing key-value (KV) cache, while its next prediction must account for an increasingly long generated history.

This becomes especially relevant in a new tier of heterogeneous inference hardware. NVIDIA Groq 3 LPX paired with Vera Rubin supports an external-drafter configuration in which the draft model runs on LPUs while the target model runs on GPUs. The LPU’s fast SRAM must accommodate both model weights and the draft’s growing KV cache. Keeping that cache compact is therefore crucial to fitting long-context drafting within the available memory.

When the initial selection becomes stale

Compressing the draft’s cache can constrain its memory requirements, but deciding what to retain becomes particularly important during long-output generation. Keeping all generated KVs can exceed a fixed cache budget during long generation. On the other hand, retaining selected prompt KVs alongside a small recent window keeps the cache bounded but restricts access to the full context. Rescoring the history at every step updates the selection, but adds repeated computation.

Prefill-based methods select which KVs to retain using information collected during prefill, when the model processes the supplied context before generation. Using the SPEED-Bench benchmark to study draft behavior in long contexts, we identify decode staleness: as generation evolves, prefill-based selection loses its relevance, while the importance of recent tokens increases. In our Qwen3.5 analysis, a draft with a SnapKV-based cache starts with high acceptance length (AL), but loses its advantage over a simple sliding window as generation progresses.

Left: acceptance length across generation positions. Right: acceptance length across cache budgets; diagonal hatching denotes compressibility training. Qwen3.5-4B drafts Qwen3.5-122B on low-entropy SPEED-Bench with 32K-token inputs, up to 4K-token outputs, temperature 0, and draft length 11. The position-wise comparison uses an 8K budget.

RIFT: allocate, refresh, and adapt

The growing mismatch between the context selected at prefill and the context needed during generation, which we term the staleness gap, motivates the three RIFT variants. They address two decisions: how to divide the active cache between importance-selected and recent positions, and when to update the importance selection. We instantiate the methods with SnapKV-style scoring, which ranks keys using attention from a short observation window of queries, although RIFT can also use other prefill-based scoring methods.

• RIFT-stat: make room for recent context. We score the prompt once after prefill and divide the budget between selected prompt positions and a sliding window. In the results below, each receives half of the budget. Recent context is particularly important because it contains the latest target-verified tokens that the draft must condition on to remain aligned. The prompt selection stays fixed while the sliding window advances.
• RIFT-dyn: refresh the selection periodically. We periodically rescore the available history using recent queries. The importance-selected portion can now retain new prompt positions and generated tokens that have moved outside the sliding window. More frequent rescoring improves AL at the cost of additional computation during decoding.
• RIFT-ada: choose when to refresh. Different stages of generation require different refresh rates. We use various signals, including previous rounds’ acceptance, draft entropy, and query drift, to adaptively decide when to rescore. This makes refresh responsive to the evolving generation rather than tied to a fixed schedule.

The same prompt and generated history under three refresh rules. Bright green marks importance-selected KVs; dark green marks the sliding window. Black diamonds record decode-time rescoring events. The selection and timing are illustrative, not measured results.

The gains persist as input length increases while the active budget stays fixed:

Average AL for Qwen3.5-4B drafting Qwen3.5-122B on low-entropy SPEED-Bench, with an 8K active cache budget, up to 4K output tokens, and draft length 11. These results use the untrained draft. The paper also evaluates Llama and Nemotron drafts, self-speculation, and cross-family drafting.

RIFT can also be applied to self-speculative decoding: the same model drafts with a compressed cache and verifies with its full cache. Avoiding a separate draft model’s weights makes this an appealing option for memory-constrained devices.

Adaptive refresh improves the tradeoff between AL and rescoring frequency. At the highlighted operating point below, RIFT-ada uses 51.1% fewer refreshes at matched AL, or achieves 0.122 higher AL at the same refresh count, relative to RIFT-dyn’s periodic-refresh curve.

Qwen3.5-4B drafting Qwen3.5-122B on SPEED-Bench 32K low-entropy data, with an 8K budget, 4K output limit, and temperature 1.

Compression-aware draft training

The above RIFT variants improve cache management of a fixed draft. To improve drafting quality, we can also train the draft to operate under the same kind of restricted attention it will encounter at inference.

We propose a compression-aware knowledge distillation procedure. It starts with a frozen, dense draft pass that scores keys. Those scores define a sparse attention mask combining importance-selected positions with a sliding window. A second, masked draft pass then learns from precomputed target-model outputs and logits. We train LoRA adapters and vary compression ratios across samples so the draft sees different cache budgets during training.

Draft compressibility training scheme. The snowflake marks the frozen scoring pass; the flame marks the trainable masked pass. Gradients update the draft adapters through the distillation loss.

Training provides additional gains on top of both allocation and periodic refresh:

Acceptance length after compressibility training, with relative improvement over the untrained draft under the same RIFT variant and budget in parentheses. SPEED-Bench 32K low-entropy data, 2K active cache budget, 4K output limit, temperature 1, and draft length 11.

Together, RIFT’s inference variants and compressibility training address a shared objective: keeping a memory-constrained draft aligned throughout long generation.

Citation

@article{tsur2026rift,
  title={RIFT: KV Cache Compression for Draft Models with Complementary Compressibility Training},
  author={Tsur, Dor and Ashkenazi, Maor and Golub, Maximilian and Hicks, Peter and Galron, Yaniv and Abramovich, Talor and Miron, Yoav and Mitra, Tiyasa and Darvish Rouhani, Bita and Zilberstein, Ran},
  journal={arXiv preprint arXiv:2610.31098},
  year={2026},
  url={https://arxiv.org/abs/2610.31098}
}