PaperlayerView exact revision on arXiv

Source-linked AI summary

Strong Drafts Need Compact Memories: Long-Context Speculative Decoding with Compressed KV Cache

Tong Yuan, Chengxi Liao, Zeyi Wen

arXiv:2608.30252v1cs.LG

TL;DR

Long-context speculative decoding must preserve draft acceptance without incurring the growing KV-access cost of strong full-KV drafts. The paper equips an independent draft with compressed KV memory that retains distant information and exact recent context while leaving target verification unchanged. At prefixes up to 32K tokens, the method reduces draft-side memory by over 70% and reaches speedups up to 2.08× on 8B and 3.33× on 70B targets over autoregressive decoding.

  • Problem

    Long-context speculative decoding must balance lightweight drafts' declining acceptance against strong independent drafts' growing full-KV access cost.

  • Method

    Memory-Augmented Sliding-Window drafting gives an independent draft compressed memory for distant history and exact recent context, while the target retains its full KV cache.

  • Results

    Over 70% draft-side memory reduction and speedups up to 2.08× on the 8B target and 3.33× on the 70B target are reported at prefixes up to 32K tokens.

  • Takeaways & Limitations

    The design preserves speculative decoding's lossless guarantee while reducing draft-side memory and latency across long-context settings.

  • Takeaways & Limitations

    The adaptor is trained on 2B tokens at 8K context, while local-window size, sink-token count, and slot interval are fixed.

Abstract

from arXiv · show

Long-context LLM applications such as document summarization and multi-turn agents require generation from prefixes spanning tens of thousands of tokens, making decoding latency a major bottleneck. Speculative decoding (SD) reduces latency without changing model outputs, but its speedup depends on both accepted draft tokens and draft-step latency: Lightweight drafts are fast but lack the capacity to capture long-range dependencies, whereas strong independent drafts recover acceptance but incur growing KV-access cost at long prefixes. We introduce memory-augmented drafting for long-context SD, equipping a strong independent draft with compressed draft-side KV memory: A lightweight adaptor constructs and incrementally updates this memory to retain distant information and exact recent context. The target verifier retains its full KV cache and applies the standard accept/reject rule, preserving SD's lossless guarantee. Experiments on Llama~3.1-8B and 70B targets at prefix lengths up to 32K show that our method reduces draft-side memory by over 70%. It achieves speedups of up to 2.08x and 3.33x , respectively, over autoregressive decoding.

1 Introduction

Long-context speculative decoding must balance draft capacity for acceptance against draft-side KV-access cost. The proposed memory-augmented draft uses compressed memory to retain distant information and exact recent context while preserving lossless verification.

  • Long-context applications face increasing decoding latency because prefixes span tens of thousands of tokens and each step accesses more KV entries.
  • Lightweight drafts are fast but lose acceptance on long contexts, whereas strong full-KV drafts recover acceptance at higher KV-access cost.
  • Memory-augmented drafting equips an independent draft with compressed KV memory for distant information and exact recent context.
  • The memory is initialized during prefill, incrementally updated during decoding, and produced by a lightweight adaptor while the draft backbone remains frozen.
  • The target verifier retains its unmodified full KV cache and standard accept/reject rule, preserving speculative decoding's lossless guarantee.
  • Over 70% draft-side memory reduction and speedups up to 2.08× on 8B and 3.33× on 70B targets are reported at prefixes up to 32K tokens.

2 Related Work

Prior long-context methods reduce KV storage or access through eviction, quantization, sparse attention, or context compression, with different hardware and conditioning trade-offs.

  • Eviction-based methods discard selected positions, which can create draft–verifier mismatch when applied asymmetrically.
  • Quantization-based methods reduce memory bytes but still scan the full prefix and typically require custom GPU kernels.
  • Context-compression methods replace raw tokens with learnable summary tokens that standard attention can consume without specialized kernels.
  • Training-aware context compression is described as hardware-friendly because it produces dense, fixed-shape KV entries compatible with existing serving infrastructure.

3 Background & Motivation

Speculative-decoding speedup depends on accepted draft length and relative draft cost. Long contexts expose a capacity–KV-access bottleneck: weak drafts lose acceptance, while strong drafts incur latency from full-history access.

  • 3.1 Speculative Decoding: Speculative decoding proposes candidate tokens with a draft model and verifies them in parallel with the target while preserving the target output distribution.
  • 3.1 Speculative Decoding: Speedup is governed by accepted draft length and relative draft cost, so effective drafting requires both faithful approximation and low latency.
  • 3.2 Model Capacity and Acceptance: EAGLE's accepted length drops monotonically with context length, while an independent Llama3.2-3B draft maintains high accepted length.
  • 3.2 Model Capacity and Acceptance: Long-context speculative decoding requires sufficient draft capacity because lightweight single-layer drafts cannot approximate the target over long histories.
  • 3.3 Context Length Dominates Latency: Full-KV draft latency scales nearly linearly with prefix length, with the slope dominated by historical KV access rather than model size.
  • 3.3 Context Length Dominates Latency: The resulting design target is to retain draft capacity while accessing a substantially shorter KV memory.

4 Method

MASW Drafting equips an independent draft model with compact memory that preserves distant context alongside exact recent context, reducing draft-side latency while leaving target verification unchanged and lossless.

  • 4.1 From Full-KV Drafting to Memory-Augmented Sliding Window: Full-KV drafting streams the entire historical cache at every step, whereas pure sliding windows discard distant history and reduce acceptance.MASW addresses this trade-off by augmenting the window with memory slots.
  • 4.1 From Full-KV Drafting to Memory-Augmented Sliding Window: MASW combines sink tokens, an exact local window, and compressed memory slots to retain both recent and distant context.The working memory is structured as Msink ∪ Mlocal,t ∪ Mslot,t.
  • 4.2 Memory Materialization: At every r-token boundary, the draft materializes one memory-slot token from completed history and retains it after raw KV eviction.Each slot is processed through the transformer blocks and appended to the slot memory during drafting.
  • 4.2 Memory Materialization: Successive memory slots form an incremental compression chain by reading prior slots while encoding the current window.This design makes each slot a compressed checkpoint rather than a disjoint segment summary.
  • 4.2 Memory Materialization: MASW uses dedicated mirrored projections for memory slots while keeping raw-token projections and draft backbone weights frozen.The adaptor’s trainable parameters are restricted to the mirrored memory-slot projection branch.
  • 4.3 Integration into Speculative Decoding: The target verifier retains the full KV cache and standard verification, while MASW compresses only the draft-side cache.Rollback discards speculative local KV and post-acceptance memory slots after rejection.

5 Evaluation

MASW is evaluated across long-context summarization tasks, target scales, compression ratios, and ablations. It reduces draft-side memory and latency while maintaining strong accepted-token performance and preserving lossless speculative decoding.

  • Experimental Setup: The evaluation covers Llama 3.1 8B and 70B targets, prefixes from 8K to 32K tokens, and mixed summarization tasks.Experiments use GovReport, QMSum, and MultiNews with greedy decoding on a single 8×H100 node.
  • Main Results: MASW avoids the long-context speedup decline seen in lightweight, sliding-window, and full-KV speculative baselines.EAGLE methods show little or inconsistent acceleration, SWA loses accepted length, and full-KV drafting incurs increasing KV traffic.
  • Main Results: 4× compression generally preserves more tokens per iteration, while 8× compression further reduces draft latency and memory.The best compression ratio depends on target scale and prefix length.
  • Effectiveness & Efficiency Analysis: The memory-augmented 8B draft prefills faster than the 70B target at every tested prefix length.Concurrent prefills make the additional draft-side work hidden behind target prefill, so TTFT does not increase in this setting.
  • Effectiveness & Efficiency Analysis: MASW reduces extra peak memory from 17.21–18.02 GB to 3.98–5.05 GB and draft latency from 41.34–54.92 ms to 12.39–15.37 ms at 32K context.Tok./Iter decreases from 4.05 to 3.43 for the 3B draft and from 4.82 to approximately 3.78 for the 8B draft.
  • Local-Window Size: Accepted length remains relatively stable for local-window sizes from 32 to 512 tokens and peaks at 3.77 with a 128-token window.The main configuration therefore sets W = S = 128.
  • Training Choices: Copy-weight initialization outperforms random initialization during adaptor training, while PT + SFT achieves the highest Tok./Iter across all tested settings.The two-stage recipe is especially advantageous under nominal 8× compression.

6 Conclusion

The conclusion presents MASW as a compact draft-memory design that reduces draft-side resource use while retaining much of full-KV acceptance quality. Full target verification preserves speculative decoding’s lossless guarantee, while the authors note potential misuse risks from lower generation cost.

  • 6 Conclusion: MASW combines sink tokens, an exact local window, and periodically materialized memory slots from trainable mirrored projections.These components replace the full draft KV cache with a compact working memory.
  • 6 Conclusion: MASW substantially reduces draft-side memory and latency while preserving much of the acceptance quality of full-KV drafting.The target verifier continues operating on the unmodified full KV cache.
  • 6 Conclusion: Using the unmodified full target KV cache preserves speculative decoding’s lossless guarantee.The draft-side compression does not alter the target verifier’s standard verification basis.
  • 6 Conclusion: Lower latency and serving cost may facilitate large-scale generation of misleading, abusive, or otherwise harmful content.The conclusion identifies this as a potential undesirable use of existing language models.

Limitations

The study’s limitations concern restricted adaptor training, fixed memory-allocation settings, and open extensions to broader compression and drafting configurations.

  • Adaptor training uses only 2B tokens and an 8K context length because training resources are limited.
  • The adaptor trains only mirrored K/V projections, which may limit how much information each memory slot can encode.
  • MASW fixes the local-window size, sink-token count, and slot interval to maintain regular masks, stable cache layouts, and bounded materialization overhead.
  • Longer-context training, larger corpora, broader draft fine-tuning, adaptive allocation, and self-speculative applications remain future directions.
  • The work relies on existing public research models and datasets, with usage governed by providers’ licenses or terms.

B Potential Risks

MASW improves inference efficiency but does not add content-generation capabilities and retains the underlying models’ safety, bias, and privacy risks.

  • MASW accelerates inference without adding new content-generation capabilities.
  • The lossless guarantee depends on verifying every draft proposal with the target model.
  • Deployments should retain full target verification and apply the underlying model’s access controls, monitoring, and content safeguards.

C Memory Prefill

During memory prefill, an asymmetric attention mask controls which tokens write to and read from the compact memory representation.

  • Memory prefill uses the same memory-write policy as decoding.
  • Sink tokens remain globally visible to stabilize attention.
  • Memory-slot tokens attend to earlier materialized slots, raw tokens in their local unit, and neighboring boundary tokens.
  • Ordinary raw tokens see only sink tokens, available memory slots, and a bounded local window.

D Detailed Experimental Setting

The experiments compare component-level draft-side KV reductions and report decoding under controlled proposal, prefix-length, dataset, and output-generation settings.

  • Component-level baselines: SWA and SnapKV are reproduced in the same Hugging Face Transformers framework, using identical draft and target backbones and weights.
  • Component-level baselines: SWA uses a sliding-window size of 1,024 tokens for draft-side KV reduction.
  • Component-level baselines: All speculative methods in Table 1 propose five draft tokens per iteration.
  • Component-level baselines: The Figure 3a motivation experiment uses a 10-token proposal horizon.
  • Decoding and timing: Evaluated 8K, 16K, 24K, and 32K lengths denote input prefixes before generation and exclude output tokens.
  • Decoding and timing: Runs use GovReport, QMSum, and MultiNews examples matched to evaluated prefix lengths and generate at most 1,024 output tokens.

E Notations

This section presents a table listing the notation used throughout the paper.

  • Table 6 lists the notation used in the paper.
Loading 2608.30252v1…