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

What it claims

RIFT is an NVIDIA paper (arXiv, October 2026) on compressing the KV cache of speculative-decoding draft models to a fixed, input-length-independent budget. Its central finding is that prefill-based selection such as SnapKV degrades as output grows on hybrid-attention drafters. In that setting, a plain sliding window (SWA) beats SnapKV. At 32K input and 4x compression, Qwen3.5-4B (drafting for Qwen3.5-122B) retains up to 96% of full-cache acceptance length (AL). Nemotron3.5-Lightning (drafting for Nemotron3-Ultra) retains 97%. SnapKV retains 63-67% in both cases.

How it works

The paper defines three variants of increasing complexity. All use a sliding-window ratio ρ = 0.5.

• RIFT-stat splits budget B into ⌊ρB⌋ most-recent tokens and B − ⌊ρB⌋ positions chosen by a prefill scorer (SnapKV by default). The selection is fixed for the whole decode. A block-level version scores contiguous blocks by mean key score, to fit paged serving.
• RIFT-dyn rescores the whole stored draft cache every Δ decode steps, using a recent observation window. This requires the full draft KV to stay available in HBM or offloaded memory.
• RIFT-ada decides at each step whether to rescore. It checks every 8 steps. Signals are an exponentially weighted mean of recent AL, draft entropy, and query drift (1 minus the dot product of normalized query summaries over full-attention layers). Each signal is compared to a calibrated threshold.

flowchart TD
  A[Every 8 decode steps] --> B{age >= a_max?}
  B -- yes --> R[Rescore full draft KV]
  B -- no --> C{policy fires and age >= a_min?}
  C -- yes --> R
  C -- no --> K[Keep current selection]
  R --> A
  K --> A

Separately, compressibility training distills the drafter with sparse attention masks that mimic the compressed cache. This reduces the mismatch between training and inference. The fetched text does not detail the training recipe.

Evidence

The evaluation uses SPEED-Bench low-entropy splits, 8K budget, 4K output, draft length 11. AL values (T=0 / T=1):

Other findings:

• Staleness: Over generation, SnapKV's AL falls 36.7% while SWA's rises 41%. The share of full-cache attention mass that SnapKV's selection captures drops from 67.0% to 25.4%. SWA's share rises from 43.5% to 68.7%. Disabling reasoning shrinks the SnapKV-SWA gap from 2.11 to 0.60 AL.
• Architecture: On full-attention Llama3.1, SnapKV beats SWA throughout. On hybrid Qwen3.5, Nemotron3 and GPT-OSS, SWA wins.
• Refresh interval: At 100K input, RIFT-dyn with Δ = 128 retains 95.7% of full-KV AL, against 85.6% for RIFT-stat. Projected rescoring overhead ranges from about 0.1% to 6% of draft decode time across Δ = 1024 down to 16.

Limits

• Scope: Tail-augmented drafters (MTP, Eagle, DFlash) are excluded. Preliminary tests showed negligible benefit for them.
• Memory: RIFT-dyn and RIFT-ada keep the full draft cache in HBM or offload. The budget bounds the attended working set, not total storage.
• Latency: Rescoring overhead is projected, not measured end to end.
• Generality: Results outside speculative decoding (Appendix F) are labelled preliminary.
• Missing details: The fetched text is truncated mid-results. It gives no code, weights or licence, and does not describe the training data or cost for compressibility training.