QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache

Thank you for your valuable comments. We appreciate that you find the paper technically sound and well-written. We address your questions and comments in detail below:

R4-1: Evaluate your approach on more recent models that better represent the current state of the field.

A: We appreciate the reviewer’s feedback on including newer models. Below we show results on Mistral-v0.3 and Llama 3.1 models on the Multi-LexSum dataset for different context lengths on a single RTX A6000 GPU. We adopt the best speculative length γ for each method. These results follow the same trends and conclusions as the results in the main paper. We will add these results and more in the Appendix.

Mistral-7B-v0.3

Llama-3.1-8B

R4-2: Your paper lacks a comparison with MagicDec, which employs a similar self-speculative approach.

A: Thank you for pointing this out. We would like to clarify that the baselines we adopt in our paper, that we referred as StreamingLLM and SnapKV, are in fact the two variants introduced by MagicDec [1]. As such, our experiments already provide a direct empirical comparison with the methods proposed in MagicDec. We show that our approach outperforms these baselines in efficiency, demonstrating the effectiveness of our method. We will clarify this more explicitly in the revised version.

R4-3: What advantages does your method offer over MagicDec specifically?

A: For long context settings, MagicDec employs self-speculative decoding with sparse KV cache for draft model which leads to poor acceptance rates especially in tasks where full context is important e.g. summarization. In contrast, our method employs KV cache quantization to accelerate draft models which leads to better acceptance rate and thus much better speedups (please refer to Results section and Appendix H for exact numbers). Additionally, due to the hierarchical nature of quantization, we enable bit sharing between the target and draft KV cache, which saves GPU memory. However for the sparse KV methods used by MagicDec, draft model and target model must use a separate copy of KV Cache, leading to higher memory consumption.

R4-4: The paper states in its abstract that KV-cache has been the main bottleneck in long-context LLM inference for edge devices. However, evaluations were conducted on an A6000 GPU, which is not an edge device but rather a professional-grade GPU.

A: Yes that is definitely right that A6000 GPU is not an edge device. However, we kindly note that edge hardware often have worse memory bandwidth to compute ratio which actually makes KV Cache bottleneck an even bigger problem. For instance, the Nvidia Jetson AGX Orin 64GB, a typical edge device, provides up to 204 GB/s of memory bandwidth and 275 TOPS [2] of compute performance. In comparison, the RTX A6000 offers a higher memory bandwidth of 768 GB/s and 309.7 TFLOPS [3] . Notably, the memory-to-compute ratio on the Jetson AGX Orin is lower than that of the A6000, making it more susceptible to memory bandwidth limitations. As a result, the speedup benefits of QuantSpec by reducing KV cache size are expected to be even more pronounced on edge devices, where memory bandwidth is a more significant bottleneck. In the final version of the paper we will add an experiment with an edge hardware to showcase this along with a discussion on this.

[1] MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding, ICLR 2025

[2] https://www.nvidia.com/content/dam/en-zz/Solutions/gtcf21/jetson-orin/nvidia-jetson-agx-orin-technical-brief.pdf

[3] https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/quadro-product-literature/proviz-print-nvidia-rtx-a6000-datasheet-us-nvidia-1454980-r9-web%20(1).pdf

We thank the reviewer for the valuable comments. We address the questions below:

R3-1: Can the author elaborate more details on the CUDA kernel design?

A: We will add a section in Appendix outlining our kernel design. We include a short summary here for the reviewer:

In our approach, we implement the algorithm using a Flash Decoding framework. In the initial stage of the Flash Decoding process, we compute the log-sum-exp (LSE) values over the INT4-quantized key-value (KV) cache. To facilitate parallelism, the keys and values are partitioned into smaller chunks, with each chunk length set as a multiple of the quantization group size. This segmentation enables parallel computation of attention between the query and each chunk. During this process, the LSE values for individual chunks are recorded.

For the draft model, we begin by loading only the upper 4 bits of the INT4-quantized KV cache within each chunk, along with the corresponding scaling factors and zero points. These are then dequantized in the kernel to reconstruct the KV cache, and the LSE is computed following the standard Flash Decoding procedure. During the verification phase, both the upper and lower bits of the INT4-quantized KV cache are loaded and dequantized. Simultaneously, a separate computation of the LSE is performed using the full-precision KV cache retained in BF16 format. Since the residual cache length exceeds the speculative length, the attention computation over the quantized region is inherently non-causal. Consequently, attention masking is applied only to the full-precision segment.

Finally, in the second stage of the Flash Decoding algorithm, the LSE values obtained from both the quantized and BF16 segments are merged to form an integrated representation that captures information from the complete KV cache.

Thank you for your valuable comments. We are happy that you found the insights from the paper interesting. We address your questions in detail below:

R2-1: Do you have intuition on why key and value caches have distinct quantization strategies?

A: The KV cache in transformer models serve distinct purposes, which inform their optimal quantization strategies. Key vectors are reused across tokens and accessed along the channel dimension, making per-channel quantization better for preserving directional consistency. In contrast, value vectors are consumed per token and contribute to weighted outputs, making per-token quantization more effective. We follow a previous work KIVI [1] which also leverages this asymmetry design to apply a 2-bit quantization that reduces memory and achieves a relatively good performance.

[1] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, ICML 2024

Thank you for your valuable comments. We are happy that you find discussion and analysis in the paper insightful. We address your questions and comments in detail below:

R1-1: How do the findings of Fig.2 and 4 scale with the parameter count of the backbone LLM ? Long-context LLMs typically adopt larger model backbones, which may again be bound by weight transfers from memory.

A: This is a great question. The findings still hold even for larger model backbones for long context lengths. To be specific, the minimum context length for which KV cache starts becoming the main bottleneck increases linearly with active parameter count (in case of MoE, active parameters can be much less than the total parameter count). However, with increasing batch size KV cache again starts to become bottleneck even for shorter context lengths. Below, we have included the exact formula to check when KV cache becomes the main bottleneck. Also it is worth noting that our method uses weight quantization for draft model to provide reliable speedups even for short contexts lengths where weight transfers are the main bottleneck.

KV cache dominates when $\text{KV cache elements} \approx 2 \cdot L \cdot B \cdot S \cdot d_{\text{model}} > \text{Model Weights}$, where $L$ is the number of layers, $B$ is the batch size, $S$ is the sequence length and $d_{model}$ is the hidden dim.

R1-2: Are there any representative works of 7B models adopting >100k token context length?

A: Yes, Qwen offers a 7B model (Qwen2.5-7B-Instruct-1M) trained with a context length of 1M tokens. Other examples include Meta’s Llama-3.1-8B model and Gemma 3 4B which were both trained with a context length of 128K. Moreover, given recent inference-time scaling research, the demand for small models with long context lengths is further increasing. This is because the long reasoning traces outputted by the model becomes part of its context, so model providers are increasingly extending the context lengths of small reasoning models to outperform larger models. Thus, long context use cases are important even for smaller (<7B) language models.

R1-3: How do findings of Section 3 change with batch size of 1, what is defined as a "small" and "large" batch in the discussion?

A: We define small batch sizes as 1-8, and large batch sizes as greater than 32. The findings of Section 3 show that for small batch sizes (including batch size of 1), short contexts (1-8k) benefit from weight quantization, medium contexts (8k-64k) benefit from both weight and KV cache quantization, and long contexts (>64k) benefit from KV cache-only quantization. We will clarify this in the final version.

R1-4: How does your quantization scheme for KV cache different from other nested quantization approaches like AnyPrecision LLM.

A: AnyPrecision LLM [1] truncates the 8-bit representation to get the corresponding upper 4-bit version. However, truncation introduces bias in the 4-bit quantization case, leading to higher quantization error of the upper 4-bit cache and lower acceptance rates. Our approach reduces this bias by first quantizing the upper 4-bit, then quantizing the residual part to get the lower 4-bit. This results in a more accurate quantization and a better acceptance rate. Below we show the quantization error on a random fp16 tensor of length 100. We also report acceptance rate we get when using Anyprecision vs our quantization scheme. For this analysis, we used Llama-2-7B-32K-Instruct model on Multi-LexSum dataset with a prefill length of 32k.

We will add this analysis in Appendix and will also add a discussion about nested quantization in related works.

R1-5: Can your method be applied with traditional self-speculative decoding approaches like SWIFT or LayerSkip?

A: This is a very interesting point. QuantSpec can be applied on top of other traditional self-speculative decoding methods like LayerSkip [2] and SWIFT [3]. For example, the algorithm proposed in SWIFT could be used to identify important and unimportant layers, and QuantSpec could then be used to retain important layers in higher precision. We will add a discussion about compatibility with other self-speculative decoding algorithms in the Appendix.

[1] Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs, ICML 2024

[2] LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding, ACL 2024

[3] SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration, ICLR 2025