PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.31098v1cs.CL

TL;DR

Long-output speculative decoding increases KV-cache pressure, while fixed prefill selections can lose relevance as generation proceeds. RIFT balances recent-token allocation with refreshed importance selection and compression-aware drafter training. Together, these methods retain most of draft acceptance length while using a fraction of the KV cache.

  • Problem

    Under a fixed KV budget, prompt-selected KVs compete with accumulated generated context, and prefill-based selections can become stale as generation evolves.

    Source: How can we best compress the drafter’s… · p. 2

  • Method

    RIFT balances recent-token allocation with importance-selected context, refreshes KV selection during decode, and pairs these policies with compression-aware drafter training.

    Source: §4.1 · p. 4, §4.2 · p. 4, §7 · p. 8

  • Results

    Together, the methods retain most of the draft’s acceptance length while using a fraction of the total KV cache.

    Source: §7 · p. 8

  • Takeaways & Limitations

    The results demonstrate the value of adapting both KV-cache selection policies and draft models under fixed cache budgets.

    Source: §7 · p. 8

    The work covers external autoregressive drafts and KV-based self-speculation, but leaves tail-augmented drafters out of scope after preliminary studies found negligible compression benefit there.

    Source: §2 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

Long-output speculative decoding increases KV-cache pressure, while fixed prefill selections can lose relevance as generation proceeds. The paper introduces RIFT’s selection strategies and compression-aware training to preserve draft acceptance length under a fixed cache budget.

How can we best compress the drafter’s KV cache

  • RIFT balances recency allocation with importance-selected context through a family of draft KV compression methods.The methods balance implementation complexity and drafting quality.Source: How can we best compress the drafter’s… · p. 2
  • Prefill selection loses its initial advantage over SWA as generation progresses, with the gap increasing in hybrid-attention models.The paper terms this changing relevance decode staleness.Source: How can we best compress the drafter’s… · p. 2
  • Compression-driven attention masks in knowledge distillation are used to improve draft alignment under a restricted KV budget.This training phase is proposed to reduce the training–inference mismatch under compression.Source: How can we best compress the drafter’s… · p. 2
  • At 32K ISL and 4× compression, Qwen3.5-4B retains up to 96% of full-cache AL and Nemotron3.5-Lightning retains 97%, while SnapKV retains 63−67% in both cases.The targets are Qwen3.5-122B and Nemotron3-Ultra, respectively.Source: How can we best compress the drafter’s… · p. 2

2 Background and Setting

The study considers draft-side KV compression under memory constraints, measuring drafting accuracy with average acceptance length. It compares sliding-window retention with prefill-based importance selection.

Scope

  • The paper studies external autoregressive draft models and KV-based self-speculative decoding, leaving tail-augmented drafters outside its scope.Preliminary studies found negligible KV-compression benefit for tail-augmented drafters.Source: §2 · p. 2

Draft-Side Compression

  • Average acceptance length (AL) measures expected accepted tokens plus the extra target-generated token, while rejection sampling preserves the target distribution under draft-cache compression.The target cache is left uncompressed because compressing it can degrade output quality.Source: §2 · p. 2
  • Sliding-window attention retains the B most recent tokens at each decode step, whereas SnapKV selects high-scoring positions using prompt-query attention scores and an observation-window-sized sliding window.The SnapKV observation window used here is 64 tokens.Source: §2 · pp. 2–3

3 Decode Staleness In KV Compression

The section examines why prefill-based KV selection becomes less effective during generation and how this varies across model architectures. It attributes the decline to growing relevance of generated context that prefill selection omits.

Prefill-based KV compression

  • SnapKV’s advantage over SWA on a 32K-ISL low-entropy evaluation reverses at the same KV budget, where SWA outperforms it.The broader evaluation challenge is that long-input-centered benchmarks often overlook long output sequences and decode-time tradeoffs.Source: §3 · p. 3

Output-position dependence

  • SnapKV AL decreases by 36.7% while SW AL improves by 41% as generation evolves under an 8K-token budget.The paper calls the evolving SnapKV–SW difference the staleness gap and reports that disabling reasoning reduces it from 2.11 to 0.60 AL.Source: §3 · p. 3

Architecture dependence

  • Decode staleness is architecture-dependent: SnapKV consistently outperforms SWA in full-attention Llama3.1, while SWA outperforms SnapKV in Nemotron3 and GPT-OSS.The latter models use hybrid attention architectures.Source: §3 · p. 3

Mechanistic evidence

  • Retained full-cache attention mass shifts from 67.0% to 25.4% for SnapKV and from 43.5% to 68.7% for SW as generation progresses.This crossover is consistent with attention shifting toward generated context that the prefill-selected cache largely omits.Source: §3 · p. 3, Figure 2 · p. 4

4 RIFT: Staleness-Aware KV cache Compression

RIFT balances prefill-selected context with recent tokens and refreshes selections during decoding to address stale KV choices. Across evaluated settings, its variants retain more acceptance length than static and decode-time baselines.

4.1 Effect of Sliding Window Ratio on Prefill-Based Methods

  • With an 8K budget, RIFT-stat at ρ = 0.5 reaches AL 5.614, 37.8% above SnapKV’s 4.073 and above pure SWA’s 5.230.RIFT-stat allocates part of the budget to recent tokens and the remainder to attention-based selection.Source: §4.1 · p. 4, Figure 3 · p. 4

4.2 RIFT-dyn: KV Cache Refresh During Decode

  • At 100K ISL, RIFT-dyn with a 128-step refresh interval retains 95.7% of full-KV AL, compared with 85.6% for RIFT-stat.It rescores the working cache using a recent observation window while retaining a sliding-window component.Source: §4.2 · p. 4

4.3 Adaptive KV cache Refresh

  • RIFT-ada chooses refresh times from draft-side signals, including recent acceptance, entropy, and query drift, rather than following a fixed periodic schedule.The policy evaluates calibrated signals at decision intervals and can combine them with gates or a fitted weighted score.Source: §4.3 · p. 5

4.4 Results

  • At 100K ISL, Qwen3.5 RIFT-stat, RIFT-dyn, and RIFT-ada retain 85.6%, 95.5%, and 96.0% of full-cache AL, respectively, versus 59.7% for SnapKV.The comparison uses an 8K budget at T = 0.Source: Table 1 · p. 5, §4.4 · p. 6
  • Adding RIFT-stat improves CriticalKV, Expected Attention, and CurDKV at both evaluated budgets.The results indicate the recency allocation can be combined with multiple prefill-based scoring methods.Source: §4.4 · p. 6
  • RIFT extends to cross-family drafting after tokenizer matching and target-model KD, and RIFT-stat’s advantage also appears beyond speculative decoding.The latter result is from a long-input, long-output reasoning benchmark.Source: §4.4 · pp. 5–6
  • RIFT-ada achieves 2.03× fewer refreshes, or up to +0.122 AL at the same refresh count, than periodic RIFT-dyn.The comparison sweeps the query-drift threshold.Source: §4.4 · p. 6, Figure 6 · p. 7
  • Across two context lengths and temperatures, RIFT-dyn and RIFT-ada achieve higher AL than StreamingLLM and H2O, while H2O outperforms RIFT-stat.StreamingLLM is outperformed by all compared compression methods.Source: Table 4 · p. 7, §4.4 · p. 7

5 Improving Draft KV Compressibility with Training

Compression-aware training narrows the mismatch between training and inference KV caches, using restricted-cache attention masks to improve target-aligned drafting. Under RIFT, training improves acceptance length, with gains especially notable at lower budgets.

5 Improving Draft KV Compressibility with Training

  • Training uses a dense frozen-draft scoring pass to select prompt positions, then a trainable sparse pass combines those positions with a sliding-window component.The effective cache budget is B = ⌊pL⌋, with r = ⌊ρB⌋ positions reserved for the sliding window; the loss combines cross-entropy and KL distillation using teacher rollouts and logits.Source: §5 · p. 7

5.1 Results

  • 16.19% gains for Qwen3.5-4B (RIFT-stat, 2K budget, T = 1) and 16.42% for Nemotron3.5-Lightning (RIFT-dyn, 2K, T = 1) over untrained drafters.Both comparisons use the same compression policy and budget; gains are smaller at larger budgets, where baseline methods often retain over 95% of full-KV AL.Source: §5.1 · p. 8, Table 5 · p. 8
  • Compression-aware training significantly improves over the same post-training pipeline without attention masking, especially at lower KV budgets.Table 6 reports relative AL improvement over a similar training recipe with no compression-driven masks.Source: §5.1 · p. 8, Table 6 · p. 8

6 Related Work

Related work addresses KV selection, decode-time cache adaptation, and compression-aware training through distinct approaches. The paper positions RIFT around decode staleness and draft-side adaptation.

6 Related Work

  • SnapKV estimates prompt importance from observation-window attention, while PyramidKV and AdaKV refine KV allocation across layers and heads.These methods optimize which KVs survive compression; the paper instead studies how prefill decisions become stale during long generation.Source: §6 · p. 8
  • Decode-time alternatives include StreamingLLM and H2O cache policies, TriForce's periodically refreshed or acceptance-gated retrieved target cache, and methods using verification for KV compression.MagicDec studies sparse-KV self-speculation at high concurrency; SpecExtend, SparseSpec, and Vegas study verification-based KV compression.Source: §6 · p. 8
  • CAT proposes continual pretraining with a learned KV scoring model for compressibility training.Source: §6 · p. 8

7 Conclusion

The paper identifies decode staleness as a source of reduced draft-target alignment and presents RIFT strategies that combine recency with importance-based KV selection. It also proposes compression-aware training and identifies combining draft-side staleness signals with target-verification attention as future work.

7 Conclusion

  • Prefill-selected context loses relevance as output length grows, and this decode staleness degrades draft-target alignment, especially in hybrid-attention architectures.Source: §7 · p. 8
  • RIFT-stat combines recency with important prompt positions, RIFT-dyn rescores during decode, and RIFT-ada adaptively rescores to reduce latency overhead.The paper also proposes compression-aware training; a future direction is combining draft-side staleness signals with target-verification attention.Source: §7 · p. 8

Limitations

The evaluation centers on acceptance-length trade-offs under a fixed active draft KV budget, while the full history remains stored in GPU HBM or offloaded memory. The appendix results also describe KV-footprint conditions and a tail-augmented MTP setting where simple sliding-window attention suffices.

Limitations

  • Decode-time refresh can add up to 6.0% of total draft decode latency in naive GPU measurements, and its overhead may be greater on SRAM-based drafting devices.Its magnitude in disaggregated deployments remains unmeasured; the active cache budget constrains drafting-device memory, not total stored KV memory.Source: Limitations · p. 9

A KV cache Costs in Hybrid Draft Models

  • For Qwen3.5-4B, four requests at 64K tokens use 51.1% of occupied memory for KV storage; for Nemotron-3-Nano, KV is 62.0% at 256K ISL and 64 concurrent requests.The memory accounting assumes bf16 weights and KV entries, fp32 recurrent matrices, and bf16 convolution buffers.Source: A KV cache Costs in Hybrid Draft Models · p. 12

B KV Cache Compression in Tail-Augmented Drafters

  • For Nemotron3-Super MTP, simple sliding-window attention retains 94% of full-KV AL at 32K ISL and 95% at 100K ISL, making it a sufficient, simpler option in the tested settings.The reported system TPS reduction from MTP KV compression was negligible.Source: B KV Cache Compression in… · p. 12

C Additional Staleness Analysis

This appendix adds staleness evidence and examines the sliding-window ratio, rescoring overhead, and block-size effects.

D.1 RIFT-stat Sliding-Window Ratio

  • Across five additional output-cap and KV-budget settings, the best intermediate sliding-window ratios outperform both fixed-64 SnapKV and pure SWA.The best observed ratios range from 0.3 to 0.7, with modest gains over the default ρ = 0.5.Source: §D.1 · p. 13

D.2 Rescoring Overhead Estimation

  • The rescoring-overhead estimate combines RIFT-dyn acceptance length, forced-AL vLLM latency, and estimated rescoring time.Added time is reported as a percentage of measured decode latency.Source: §D.2 · p. 13

D.3 Block-Size Ablation

  • Moderate block sizes have negligible effects when small relative to the cache, but overly coarse blocks reduce selection resolution and drafting performance.The experiments sweep block sizes and two temperatures for 32K and 100K throughput splits.Source: §D.3 · p. 13

D.4 Budget, Context, and Cross-Family Comparisons

The appendix tests RIFT across models, budgets, and output lengths, and examines adaptive-refresh tradeoffs. Dynamic rescoring generally improves drafting performance, while selected adaptive policies reduce refreshes at matched acceptance length.

D.4 Budget, Context, and Cross-Family Comparisons

  • Across considered settings, RIFT performance is maintained, and dynamic rescoring dramatically improves drafting performance; some setups obtain similar performance with SWA and RIFT-stat.Nemotron3-Nano drafting for Nemotron3-Ultra or Super is one example where SWA can match RIFT-stat.Source: §D.4 · pp. 13, 15

D.5 Additional Prefill Selectors

  • Table 14 reports the absolute acceptance lengths underlying the relative gains in Table 2.Source: §D.5 · p. 15

D.6. Effect of draft length

  • As draft length increases, RIFT-dyn’s relative AL gain over SnapKV rises from 17.7% to 60.0% at T = 0 and from 17.9% to 54.8% at T = 1.RIFT-stat and RIFT-dyn consistently outperform SnapKV across the tested draft lengths; RIFT-dyn is closest to Full-KV.Source: §D.6 · p. 15, Figure 13 · p. 15

D.7. AL vs. OSL

  • Between 2K and 64K forced OSL, the RIFT-stat–SWA mean-AL gap narrows from 0.42 to 0.09 at T = 0 and from 0.29 to 0.03 at T = 1.The passage cautions that forced high OSLs produce highly repetitive text, so these trends describe forced continuations rather than naturally long, meaningful responses.Source: §D.7 · p. 15, Figure 14 · p. 16

E RIFT-Ada: Additional Results and Analysis

E.1 Additional Results

  • Figure 16 compares Qwen3.5-4B AL–refresh tradeoffs across temperatures, KV budgets, and input lengths using setting-specific controllers.Controller parameters are held fixed across budgets.Source: §E.1 · p. 15, Figure 16 · p. 17

E.2 Adaptive Refresh for Nemotron Lightning Self-Speculation

  • For Nemotron3.5-Lightning self-speculation, RIFT-ada achieves similar ALs with a fraction of RIFT-dyn’s refreshes.The controller uses calibrated acceptance thresholds at T = 0, and acceptance or entropy triggers at T = 1.Source: §E.2 · pp. 15, 18, Figure 17 · p. 20

E.3 Adaptive-Refresh Policies and Calibration

  • At matched AL, selected Qwen policies reduce refreshes by 53.5% at T = 0 and 51.1% at T = 1, but results include development subsets rather than an untouched held-out evaluation.Source: §E.3 · p. 20

F Beyond Speculative Decoding: Long-Context Quality Evaluation

RIFT-stat also improves long-context standalone-generation scores over SnapKV in preliminary LongBench evaluations. The appendix further reports task breakdowns, cross-family tokenizer adaptation, and comparisons with other compression policies.

F Beyond Speculative Decoding: Long-Context Quality Evaluation

  • At an 8K cache budget, RIFT-stat retains 96.4% and 89.9% of full-KV scores on LongBench-v1 and LongBench-v2, respectively.RIFT-stat outperforms SnapKV at every tested budget; the authors describe these standalone-generation results as preliminary.Source: F Beyond Speculative Decoding:… · p. 21, Table 15 · p. 21

G RIFT Performance: Task Breakdown

H Tokenizer Adaptation for Cross-Family Drafting

  • Cross-family draft adaptation resizes embeddings, maps matching tokens, initializes unmatched tokens from draft-tokenizer embeddings, and aligns output heads and special-token IDs.After vocabulary alignment, the authors distill on the Nemotron post-training V2 dataset for 10K steps.Source: H Tokenizer Adaptation for… · p. 21

I Comparison With Additional KV Compression Policies

J.1 Mask Construction

  • SnapKV mask construction captures post-RoPE queries and keys in a frozen dense forward pass, then uses attention weights to assess key-position importance.Source: §J.1 · p. 22

J.2 Training Configuration

  • Compressibility training runs 10K steps on Nemotron V2 teacher rollouts with batch size 64; evaluation uses 4K OSL and draft length 11 on SPEED-Bench 32K low-entropy data.Source: §J.2 · pp. 22–23

J.3 Ablation Studies

The ablations favor attention-based training masks and show tradeoffs among mask composition, fine-tuning cost, and model size. Additional results report benefits across longer contexts, other drafters, and RIFT-ada.

J.4 Additional Draft Compressibility Training Results

  • Compressibility-training gains also appear on 100K-input evaluations, with higher gains in non-greedy settings.Source: §J.4 · p. 23, Table 22 · p. 25
  • SnapKV benefits from compressibility training, but the authors attribute the magnitude of its gain to its lower starting AL relative to RIFT-stat.Source: §J.4 · p. 24, Table 23 · p. 25
  • For Lightning, training improves RIFT-dyn AL by approximately 16.7% (3.72 to 4.34) at 2K and T = 1; at 8K and T = 0, trained AL 5.27 is 97.8% of original full-KV AL 5.39.Qwen3.5-0.8B gains reach up to 10% at T = 1; the Lightning 8K result is 91.7% of the trained draft’s full-KV AL of 5.75.Source: §J.4 · p. 24, Table 25 · p. 26, Table 24 · p. 26

J.5 Additional RIFT-ada Results on Draft Compressibility

J.3 Ablation Studies

  • Attention-based masks outperform random and sliding-window controls at both temperatures and all three KV budgets, but require an additional dense forward pass.Table 19 reports the trained draft’s AL under the masking strategies.Source: §J.3 · p. 23, Table 19 · p. 25
  • The high-biased sparsity distribution performs best, while the training-time sliding-window ratio has no significant effect on performance.Table 20 compares mask-distribution and ratio settings.Source: §J.3 · p. 23, Table 20 · p. 25
  • Full fine-tuning yields slightly higher AL than LoRA, but has significant computational overhead; LoRA is preferable for a lightweight modification.Source: §J.3 · p. 23, Table 21 · p. 25

K Extended Related Work

  • RIFT complements prior approaches by targeting recency allocation and draft-side refresh signals under a fixed active-cache budget, for external or self-speculative drafts.Source: K Extended Related Work · pp. 24, 27

Related papers

Loading 2610.31098v1…