PaperlayerView exact revision on arXiv

Source-linked AI summary

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache

Rishabh Tiwari, Haocheng Xi, Aditya Tomar, Coleman Hooper, Sehoon Kim, Maxwell Horton, Mahyar Najibi, Michael W. Mahoney, Kurt Keutzer, Amir Gholami

arXiv:2502.10424v1cs.LGcs.AI

TL;DR

Long-context decoding is constrained by KV-cache memory and low speculative-token acceptance. QuantSpec uses self-speculation with 4-bit weights and a hierarchical 4-bit KV cache; quantized KV improves target–draft acceptance and overall speedups.

  • Problem

    In long-context settings, the KV cache becomes a primary memory bottleneck, while conventional speculative decoding can have low acceptance and retain caches for both models.

    Source: §1 · p. 1

  • Method

    QuantSpec uses self-speculative decoding with 4-bit weights and a hierarchical 4-bit KV cache that shares bits between target and draft caches.

    Source: §1 · p. 2

  • Results

    Quantizing the KV cache improves acceptance rates between target and draft models and thus overall speedups.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Bit-sharing in the hierarchical cache eliminates the need for additional memory for a separate draft-model KV cache.

    Source: §1 · p. 2

    Combining QuantSpec with sparse-KV methods for additional speedup is left for future work.

    Source: §2.3 · p. 3

Abstract

from arXiv · show

Large Language Models (LLMs) are increasingly being deployed on edge devices for long-context settings, creating a growing need for fast and efficient long-context inference. In these scenarios, the Key-Value (KV) cache is the primary bottleneck in terms of both GPU memory and latency, as the full KV cache must be loaded for each decoding step. While speculative decoding is a widely accepted technique to accelerate autoregressive decoding, existing methods often struggle to achieve significant speedups due to inefficient KV cache optimization strategies and result in low acceptance rates. To address these challenges, we propose a novel self-speculative decoding framework, QuantSpec, where the draft model shares the architecture of the target model but employs a hierarchical 4-bit quantized KV cache and 4-bit quantized weights for acceleration. QuantSpec maintains high acceptance rates ($>$90%) and reliably provides consistent end-to-end speedups upto $\sim2.5\times$, outperforming other self-speculative decoding methods that use sparse KV cache for long-context LLM inference. QuantSpec also reduces the memory requirements by $\sim 1.3\times$ compared to these alternatives.

1. Introduction

Long-context inference faces KV-cache memory and latency bottlenecks; existing cache-efficiency methods can degrade generation quality, while conventional speculative decoding faces low acceptance and dual-cache memory costs. QuantSpec applies self-speculative decoding with 4-bit weights and a hierarchical 4-bit KV cache.

1. Introduction

  • QuantSpec achieves >1.78× throughput speedup over the autoregressive baseline across several context lengths on LWM-Text-Chat-128k.Figure 1 reports throughput in tokens/sec.Source: Figure 1 · p. 1
  • KV-cache quantization benefits long contexts, while weight quantization is more beneficial for short contexts.Source: §1 · p. 2
  • Hierarchical quantization shares bits between target and draft KV caches, eliminating additional memory for a separate draft cache.Source: §1 · p. 2
  • A double full-precision buffer stores recent KV entries to improve acceptance and avoid wasteful quantization and dequantization.Source: §1 · p. 2
  • Quantizing the KV cache improves target–draft acceptance rates and thereby overall speedups.Source: §1 · p. 2
  • Custom CUDA attention kernels with the hierarchical cache achieve up to ∼2.88× speedup at 4-bit precision relative to FP16 FlashAttention kernels.Source: §1 · p. 2

2. Related Work

Related work addresses long-context efficiency through sparse or selective KV computation, quantization, and speculative decoding. QuantSpec is positioned against self-speculative approaches that use sparse KV caches.

2.1. Efficient Long-Context Inference

2.2. Quantization

2.3. Speculative Decoding

3. LLM Inference Bottlenecks

The section uses arithmetic intensity to classify inference bottlenecks and guide latency optimizations. Its analysis finds that long-context decoding is memory-bound, while the dominant operations vary with context length and batch size.

3.1. Arithmetic Intensity

  • Arithmetic intensity measures floating-point operations per byte loaded and helps classify inference as compute-bound or memory-bound to guide latency optimizations.Memory-bound operations are limited by memory bandwidth; compute-bound operations are limited by peak FLOP/s.Source: §3.1 · p. 3

3.1.1. ASYMPTOTIC ANALYSIS OF ARITHMETIC INTENSITY FOR PREFILL AND DECODING

  • During prefill, linear projections dominate when sequence length is much smaller than hidden dimension, while attention dominates as sequence length becomes much larger.The aggregate arithmetic intensity reflects linear projections in the shorter-context regime and attention in the longer-context regime.Source: §3.1.1 · p. 3
  • Long-context decoding has extremely low arithmetic intensity regardless of batch size, as per-sequence attention cannot benefit from batching like linear layers.As context length grows, KV-cache load-store operations dominate decoding latency; aggregate decoding arithmetic intensity does not scale with context length.Source: §3.1.1 · p. 3
  • The analysis compares arithmetic intensity across linear, attention, and aggregate operations during prefill and decoding, under an assumption of FlashAttention.The table's variables include batch size B, sequence length SL, hidden dimension d, and generation length k.Source: §3.1.1 · p. 3, Table 1 · p. 4

3.1.2. COMPUTE VERSUS MEMORY-BOUND REGIMES

  • The roofline analysis classifies operations below the hardware ridge point as memory-bound and those above it as compute-bound.The study extrapolates the ridge point to a ridge plane using NVIDIA A6000 specifications for Llama-2-7B at 16-bit precision.Source: §3.1.2 · p. 4
  • For small-batch, short-context decoding, linear projections dominate latency; attention dominates in long-context or large-batch regimes, while intermediate regimes favor quantizing both weights and KV cache.When linear and attention contributions are approximately equivalent, both quantization strategies are identified as suitable.Source: §3.1.2 · p. 4

4. QuantSpec

QuantSpec is a self-speculative framework that quantizes weights and KV cache for inference. Its hierarchical cache lets the draft and target models share a representation while using different cache precisions.

4.1. Overview of QuantSpec

  • QuantSpec is a self-speculative framework that quantizes model weights and the KV cache to INT4 precision for short- and long-context generation.The overview describes dynamic switching between INT4 and INT8 KV-cache representations without on-the-fly quantization.Source: §4.1 · pp. 4–5
  • The draft model shares the target model's architecture, which the paper says improves acceptance rate and long-context capability.A naive sparse-KV self-speculative implementation would maintain a separate fully quantized draft cache, adding memory use and computational overhead.Source: §4.1 · p. 5

4.2. Hierarchical KV Cache

  • INT8 KV cache is reported to maintain competitive generation quality with the FP16 baseline while using half the KV-cache memory.The comparison is described for Llama-2-7B on WikiText-2 and C4.Source: §4.2 · p. 5
  • The hierarchical representation uses an INT4 cache plus its INT4 residual to represent the target model's INT8 cache, avoiding a separate draft-cache copy.The two 4-bit components correspond to the upper and lower segments of an INT8 value.Source: §4.2 · p. 5

L . CINT4

The cache description distinguishes upper and lower 4-bit components and describes their use in draft generation and target verification.

L . CINT4

  • The upper 4-bit component uses asymmetric round-to-nearest quantization over the range 0 to 15.The passage describes this component as CINT4 U.Source: L . CINT4 · p. 5
  • The lower 4-bit residual uses symmetric round-to-nearest quantization over the range -8 to 7.The passage says this choice better matches the error distribution.Source: L . CINT4 · p. 5
  • During decoding, draft-model candidate generation loads and dequantizes only the upper 4-bit representation; target-model verification uses both upper and lower 4-bit representations to reconstruct the KV cache in higher INT8 precision.Source: L . CINT4 · pp. 5–6

U + CINT4

The passage describes asymmetric quantization of the KV cache using an INT8 representation with a scale and zero point.

U + CINT4

  • Asymmetric quantization represents the FP32 cache as CINT8 multiplied by SINT8, then offset by ZINT8.Source: U + CINT4 · p. 6
  • SINT8 denotes the scaling factor and ZINT8 denotes the zero point in the cache representation.Source: U + CINT4 · p. 6
  • The passage places CINT8 in the integer range from 0 to 2^8 - 1.Source: U + CINT4 · p. 6

U + CINT4

QuantSpec addresses the overhead and acceptance-rate costs of quantizing KV caches for speculative decoding with a double full-precision buffer. Its hierarchical INT4 cache and buffer support flexible rejection handling and reduce how often quantization and cache movement occur.

4.3.1. CHALLENGES WITH KV CACHE QUANTIZATION AND SPECULATIVE DECODING

  • Quantizing keys by channel and values by token minimizes quantization error; the method uses asymmetric, per-group INT4 quantization with group size equal to the head dimension.These cache-specific strategies are described as supporting a higher acceptance rate.Source: §4.3.1 · p. 6
  • Naively quantizing new value tokens at every step is computationally expensive, while delaying key quantization until a group fills can reduce acceptance and trigger repeated quantization after rollbacks.The rollback overhead includes quantizing and dequantizing cache entries for rejected tokens.Source: §4.3.1 · p. 6

4.3.2. ADAPT TO SPECULATIVE DECODING USING FULL PRECISION BUFFER

  • A double full-precision buffer of size 2G keeps recent cache entries unquantized and stages generated tokens until target verification.CF1 and CF2 each hold G entries; rejected entries are removed before CF1 is quantized and appended to the cache.Source: §4.3.2 · p. 6
  • Keeping at least G recent tokens in CF1 supports acceptance, while quantization and cache movement occur only every G decoding steps.Rejected-token cache entries can be discarded by operating on CF2, without extra quantize and dequantize operations.Source: §4.3.2 · p. 6

4.3.3. SUMMARY

  • Drafting loads only the upper 4-bit cache, while verification loads both halves and dequantizes them to INT8; QuantSpec also supports FlashDecoding.The method uses the same architecture for draft and target models and clears and quantizes half the full-precision buffer after verification when it is saturated.Source: §4.3.3 · p. 7, §4.3.2 · p. 7

5. Evaluation

QuantSpec is evaluated across datasets and context lengths for acceptance, GPU memory, and serving speed. It delivers higher speedups than sparse-KV self-speculative baselines, with the strongest reported speedup at 128k context, while ablations show the most effective quantization choice varies with context length.

5. Evaluation

  • QuantSpec evaluation measures the draft–target acceptance ratio, GPU memory consumption, and end-to-end serving speedup, benchmarking these three dimensions across multiple datasets and context lengths.Experiments use long-context LLaMA-2 and LWM targets, PG-19, ∞BENCH Sum, and Multi-LexSum, comparing StreamingLLM and SnapKV. Baseline draft-KV budget is one-fourth of the context length; group size is 128, residual length R is 256, and outputs are limited to 90 tokens.Source: §5 · p. 7, §5.1 · p. 8

5.1. Setup

  • QuantSpec decoding-speedup benchmarks use long-context LLaMA-2 and LWM variants with PG-19, ∞BENCH Sum, and Multi-LexSum, comparing against StreamingLLM and SnapKV.The sparse-KV baselines’ draft KV budget is one-fourth of context length, matching QuantSpec’s 4-bit quantized KV cache.Source: §5.1 · p. 8

5.2. Speedup Evaluation

  • At 128k context, QuantSpec reaches ∼2.49× speedup over autoregressive decoding and achieves consistently better speedups across context lengths than sparse-KV baselines.The evaluation also reports lower GPU memory use than the baselines and higher acceptance rates on summarization tasks.Source: §5.2 · p. 8

5.2.1. KERNEL SPEEDUPS

  • For 128k context, the INT4 attention kernel is ∼2.88× faster than standard FP16 FlashAttention kernels.Table 4 reports kernel-level latency for attention with the hierarchical quantized KV cache.Source: Table 4 · p. 8, §5.2.1 · p. 8

5.3. Ablation Results

  • For short contexts, weight quantization contributes most of the speedup; both forms contribute at medium lengths, and KV-cache quantization is most effective at long lengths.Figure 4 compares weight-only, KV-cache-only, and combined quantization against the autoregressive baseline as context length increases.Source: Figure 4 · p. 8, §5.3 · p. 8

6. Conclusions

The conclusion presents quantized speculative decoding as a way to improve long-context LLM efficiency, while the accompanying analysis and design details explain inference bottlenecks, cache quantization, and Flash Decoding compatibility.

6. Conclusions

  • QuantSpec combines KV-cache optimization and a double full-precision buffer to address quantization/speculation conflicts; the authors report faster processing without compromising accuracy or performance.The conclusion frames this as an approach to efficiency and scalability in long-context settings.Source: §6 · pp. 8–9

A. Attention Module’s Inference Workflow

  • Prefill writes sequence-derived keys and values to caches for multi-head attention; decoding appends each token’s key/value before computing attention.The workflow distinguishes processing an input sequence during prefill from processing one current token per decoding step.Source: A. Attention Module’s Inference Workflow · p. 12

B. More Related Works

  • Related work describes high-bandwidth-memory loading and KV sharing to manage cache budgets, precision schemes to reduce costs, and QSpec/Sirius findings for speculative decoding.The cited work notes higher performance but higher memory consumption for some dynamic-loading approaches, while QSpec uses complementary quantization and Sirius reports contextual sparsity can hurt speculative-decoding performance.Source: B. More Related Works · p. 12

C.1. Prefill Arithmetic Intensity Analysis

  • All prefill regimes are compute-bound; for long contexts, quadratic attention dominates linear operations, and attention intensity does not scale with batch size.The figure places all prefill regimes above the ridge plane.Source: §C.1 · pp. 12–13, Figure 5 · p. 13

C.2. Modern GPU Hardware VRAM Size Constraints

  • At B = 16 and SL = 262k, KV cache memory is 160x model-weights memory; a long-context cache can exceed node VRAM, so batch scaling does not make decoding compute-bound.The cited hardware discussion concerns a Llama-2-7B model and a node with eight A100/H100 GPUs, each with 80 GB of memory.Source: Figure 6 · p. 13, §C.2 · p. 13

D. Quantization Strategies for KV cache

E. Compatibility with Flash Decoding

  • A recent-token full-precision buffer, with upper bound 2G, fits Flash Decoding as an independent chunk summed with quantized chunks.The quantized section divides into chunks that align with Flash Decoding’s structure.Source: E. Compatibility with Flash Decoding · p. 14

F. Details about the Datasets Used

The appendix describes the datasets, the procedure for selecting speculation length, and acceptance-rate behavior across speculation lengths, alongside QuantSpec’s cache-buffer algorithm.

F. Details about the Datasets Used

  • Evaluation datasets include WikiText-2 and PG-19 for language modeling, plus ∞BENCH Sum and Multi-LexSum for long-context summarization.The ∞BENCH Sum prompts average ∼171k tokens and summarize a fake book created by core-entity substitution; Multi-LexSum prompts average ∼90k tokens and concern civil-rights litigation lawsuits.Source: F. Details about the Datasets Used · pp. 14–15

G. Hyperparameter Search

  • The search selects γ for each dataset-method pair using prompt length 8192, reuses that value at other context lengths; sparse-based methods peak at γ = 1, while quantization-based methods usually perform best at γ = 4 or 6.Source: G. Hyperparameter Search · p. 15

H. Comparing Acceptance Rates

  • QuantSpec acceptance remains higher as speculation length increases, while sparse-KV methods’ acceptance degrades faster.Figure 9 compares acceptance rates by speculation length for LWM-Text-Chat-128k on Multi-LexSum.Source: Figure 9 · p. 15, H. Comparing Acceptance Rates · p. 16
  • The algorithm initializes upper/lower 4-bit caches and recent-token full-precision buffers, drafts and verifies tokens, then quantizes buffered cache as it fills.On rejection, it clears rejected draft KV entries from the full-precision buffer.Source: H. Comparing Acceptance Rates · p. 17
Loading 2502.10424v1…