QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Authors: Rishabh Tiwari
Haocheng Xi
Aditya Tomar
Coleman Richard Charles Hooper
Sehoon Kim
Maxwell Horton
Mahyar Najibi
Michael W. Mahoney
Kurt Keutzer
Amir Gholami
Venue: ICML 2025 poster (2025)
Keywords: speculative decoding
quantization
long-context inference
Primary area: deep_learning->large_language_models
TL;DR: We propose a novel self-speculative decoding framework to accelerate long-context inference using KV cache and weight quantization.

Abstract: 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.

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Rishabh Tiwari Haocheng Xi! Aditya Tomar'.
Coleman Hooper's
Sehoon Kim Maxwell Horton
Mahyar Najibi? Michael W. Mahoney 134
Kurt Keutzer
Amir Gholami 13
Abstract
AutoRegressive
StreamingLLM
QuantSpec
40
Large Language Models (LLMs) are increasingly being deployed on edge devices for long-context settings, creating a growing need for fast and ef- ficient long-context inference. In these scenar- ios, the Key-Value (KV) cache 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 decod- ing is a widely accepted technique to accelerate autoregressive decoding, existing methods often struggle to achieve significant speedups due inefficient KV cache optimization strategies and result in low acceptance rates. To address these challenges, propose novel self-speeu 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 accelera- tion. QuantSpec maintains high acceptance rates (>90%) and reliably provides consistent end- to-end speedups upto 2.5x, outperforming other self-speculative decoding methods that use sparse KV cache for long-context LLM inference. QuantSpec also reduces the memory requirements by ~ 1.3x compared to these alternatives.
1.78%
2.23*
Token
1,99x
20
10
16k, GPU
32k, GPU
64k, 2 GPU
Context Length
Figure I. Throughput in tokens/sec of various decoding methods. QuantSpec achieves >> 1.78x speedup over the autoregressive baseline across several context lengths. Benchmarked LWM- Text-Chat-12255
tails across extended sequences. However, long-context inference presents significant challenges in terms of effi- ciency and scalability. For example, token eviction (Zhang et al., 2024d; Ge et al., 2023; Liu et al., 2024b) and KV cache quantization (Liu et al., 2024c; Kang et al., 2024; Hooper et al., 2024) have been proposed to improve the efficiency for long-context inference. However, they often entail noticeable degradation in generation quality.
One promising alternative to enhance the efficiency of LLMs while preserving generation quality is speculative decod- ing (Leviathan et al., 2023; Chen et al., 2023; Kim et al., 2024). This method accelerates inference by using a smaller (draft) model rapidly generate candidate tokens, and uses the original (target) model to verify these tokens to ensure generation quality. However, the efficient application of speculative decoding in long-context settings has not been thoroughly explored.
1. Introduction
Large Language Models (LLMs) have been widely used in recent years, revolutionizing natural language processing (NLP) and artificial intelligence (AI) applications. As their applications expand, there is growing demand to deploy LLMs in long-context settings handling extended text inputs such as document summarization, lengthy conver- sations, or comprehensive instructions. The model must maintain coherence in such contexts and track intricate de-
Traditional speculative decoding approaches often rely on using smaller models as the draft model in order mini- mize the weemory-baddwee weights of the larger target model. In long-context scenarios, however, the primary bottleneck shifts from model weights to the KV cache, which grows linearly with the context length. Additionally, since small models do not usually pos- sess good long context understanding ability, the acceptance rates of the candidate tokens by the target model drop signifi-
overhead of loading the model
"quual contribution 'UC Berkeley AAppee PICSI LLBNL. Cor- respondence to: Amir Gholami <amirgh beereelcy.rk.....
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
cantly, leading suboptimal speedup. Moreover, traditional speculative decoding methods maintain the KV cache for both the target model and the draft model, causing a large memory footprint. Therefore, finding solution that both optimizes the KV cache's memory efficiency and improves the acceptance rate within speculative decoding is essential for performant LLMs in long-context applications.
for tokens relevant for next token prediction. KV Predic- tion (Horton et al., 2024) improves prompt processing time by predicting the KV cache needed for autoregressive gener- ation. Retrieval-augmented generation (Tan et al., 2024; Liu et al., 2024a) enhances the accuracy of language model out- puts by combining generative models with external retrieval mechanisms whose context length is very long.
To mitigate these issues and to enable efficient and accu- rate long-context inference, we propose QuantSpec, self- speculative decoding method that utilizes 4-bit weights and a 4-bit hicrarchical KV cache to speedup long-contex infer- ence. In particular we make the following contributions:
2.2. Quantization
Quantization has emerged as a powerful technique to re- duce the memory footprint and computational complexity in large-scale neural networks. Weight-only quantization (Lin et al., 2024; Kim et al., 2023a; Shao et al., 2023; Chee et al., 2024) focuses on reducing the precision of model weights to reduce the memory requirements of the model. As models grow larger, the memory footprint of KV caches can become substantial, especially for long input sequences. KV cache quantization (Liu et al., 2024c; Hooper et al., 2024; Kang et al., 2024) addresses this issue by quantizing the key and value caches to enable longer sequence inference.
We perform a comprehensive analysis of LLM infer- ence to identify bottlenecks across various context lengths, demonstrating that quantizing the KV cache improves efficiency for long contexts, while quantizing model weights is more beneficial for short contexts (see Section 3.1).
We introduce a novel hierarchical quantization tech- nique that enables bit-sharing between the target and draft models' KV caches, eliminating the need for ad- ditional memory for the draft model (see Section 4.2).
2.3. Speculative Decoding
Speculative decoding has become an important technique for improving the inference efficiency of L.LMs (Leviathan et al., 2023; Chen et al., 2023; Kim et al., 2024). It uses smaller draft model to rapidly generate candidate tokens, which are then verified by a larger target model to ensure correctness. Parallelization in speculative decoding has also been studied to enhance the efficiency by predicting multiple tokens at one time (Cai et al., 2024; Bhendawade et al., 2024; Li et al., 2024b; Chen et al., 2024b). We include additional related works in Appendix B.
We propose a double full call-precision cache buffer used for storing the most recent KV cache in full precision to improve acceptance rates and also eliminate waste- ful quantization and dequantization operations (see Section 4.3).
We show that using quantized KV cache leads to better acceptance rates between the target and the draft model, and thus leads to better overall speedups (see Section 5.2).
Self-speculative decoding is the class of speculative de- coding methods in which the draft model shares the same architecture as that of target model for better alignment. Recent works like Magicdec (Sadhukhan et al., 2024) and TriForce (Sun al., 2024) have shown that self-speeulation with sparse KV can effectively speedup the draft model in long-context settings, where KV is the main bottle- neck. While this design avoids loading the entire KV cache throughout the autoregressive generation process, KV cache sparsification can lead to noticeable performance degrada- tion as evidenced in previous works (Zhang ct al., 2024d; Liu al., 2024b; Ge et al., 2023; Zhou et al., 2024). This can potentially yield a mismatch between the draft and tar- get model's 's predictions (i.c., lower acceptance rate), which is critical factor in overall speedup. QuantSpec addresses this limitation by proposing a draft model with a novel hi- erarchical quantized KV cache, which maintains a higher acceptance rate between the draft and target models, there- fore leading to better speedup. Note that our method can be combined with sparse KV methods (Sun et al., 2024;
We implement custom CUDA kernels for attention with our hierarchical quantized KV cache achieving up to ~ 2.88x speedups at 4-bit precision relative to FP16 FlashAtrention kernels. (see Section 5.2.1)
2. Related Work
2.1. Efficient Long-Context Inference
An important challenge in optimizing long-cowe Long-context infer- ence lies in reducing memory and computation requirements while retaining high performance on tasks that involve long sequences. Sparse attention mechanisms (Liu et al., 2021: Xiao al., 2023b; Yao et al., 2024; Tang et al., 2024; Yang et al., 2024; Liu et al., 2024b; Ge et al., 2023; Jiang al., 2024) have been widely adopted to manage the quadratic complexity of traditional full attention in long contexts. These techniques typically maintain efficiency by dropping non-essential Key-Value (KV) pairs from the cache. To- ken pruning (Fu al., 2024) selectively computes the KV
2

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Sadhukhan et al., 2024) for additional speedup, which we leave for future work.
sand to hundreds of thousands of tokens, this phase consists of large matrix-matrix multiplications (matmuls) with high arithmetic intensities. Table I shows asymptotic analysis of arithmetic intensity for prefill and decoding broken up into linear, attention, and aggregate operations for batch size B, sequence length St, hidden dimension d, and a generation length of tokens. During prefill, the aggregate arithmetic intensity is similar to the arithmetic intensity of the linear projections when St < d because self-attention is relatively inexpensive for short contexts. Thus the linear projections dominate latency in this regime. However, as the context length increases and St >> d, the aggregate arithmetic intensity reflects the arithmetic intensity of atten- tion, which begins to dominate latency since self-attention incurs additional cost with longer context lengths. Note that our analysis assumes the use of Flash (Dao et al.. 2022), such that the attention scores matrix which grows on the order of O(B Sy." ?) is never fully materialized, and thus the memory operations for this matrix are limited to O(B.Sp
3. LLM Inference Bottlenecks
3.1. Arithmetic Intensity
To understand the primary bottlenecks in LLM inference and to motivate our method, we perform a thorough analysis of inference under several different regimes. These regimes include combination of small versus large batch sizes and short versus long context lengths during both the prefill and decoding stages. We use arithmetic intensity as the central metric in our analysis, where arithmetic intensity is defined as the number of floating point operations (FI.Ps) that can be performed per byte loaded from memory, or memory operations (MOPs) (Williams et al., 2009):
# FLOPs
Arithmetic Intensity = # MOPs
Arithmetic intensity allows us to classify which regimes of LLM inference are compute-bound or memory-bound and determine appropriate optimizations to improve latency. Compute-bourd bound operations are limited by the hardware's 's peak F1.OP/s (FIOPs per second) performance and benefit from algorithmic improvements that reduce computational complexity (e.g., subquadratic attention). On the other hand, memory-bound operations are limited by the hardware's ' memory bandwidth (GB/s) and benefit from techniques that optimize memory load-store operations, such as quantizing the weights of model even if they are later scaled up to a higher precision during computation preserve accuracy.
On the other hand, in the decoding stage, generating tokens requires loading and storing the weights and KV cache k times. Since the input each iteration is a single token per sequence in the batch (r E R BxIx), these oper- ations mainly consist of small matmuls with low arithmetic intensity. For short context lengths where St. & d, the ag- gregate arithmetic intensity for decoding again reflects the arithmetic intensity of the linear projections as loading and storing a small KV cache is relatively inexpensive compared to loading and storing the model weights. However, as the context length grows (St >> d), the load-store operations for the large KV cache exacerbate and dominate latency, and the aggregate arithmetic intensity reflects the arithmetic intensity of attention. Ultimately, the aggregate arithmetic intensity for decoding is much lower than that of prefill:
For finer- fineregrained analysis, break down the major operations in the Transformer into two categories: lin- ear. which consists of the weight-to-activation matrix multiplications (i.e.. Wo, WK, Wy, Wout, mlp mwp_up.ppppppmmtmmmmmmmm.ppppppp ml p.lpddon.prroj.cm and the linear classification layer), and attention, which consists of the activation-to-activan activation ma- trix multiplications (i.e., query key and attention weights x values). Note that the aggregate of all Transformer oper- ations includes the above operations as well as non-linear operations like activation functions in the feed-forward net- work, softmax in the attention mechanism, and layer normal- ization. Because we are interested in studying the linear and attention operations, we do not explicitly focus on the non- linear operations and classification layer in our asymptotic analysis, although we include them in our final results.
wpp.pppp. j..proj.
SO(B-SD) SL), O(Sr), St. > d
S2 < d
OOB), St <<< d 00), S, > d'
>
prefill
decode
While the aggregate arithmetic intensity for prefill scales proportionally to the context length which can be in the hundreds of thousands, the aggregate arithmetic intensity for decoding does not scale with the context length at all. Moreover, using larger batch sizes only seems increase the arithmetic intensity for decoding in the short-conter setting. For long contexts, decoding has an extremely low arithmetic intensity irrespective of the batch size since every sequence in the batch undergoes self-attention separately and therefore cannot benefit from batching in the same way linear layers do.
3. ASYMPTOTIC ANALYSIS OF ARITUMETIC INTENSITY FOR PREFILI AND DECODING
During prefill, the model weights are only loaded once to process all tokens in the input and generate the first token.
Because the context length can range from couple thou-
3

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Table Asymptotic analysis of arithmetic intensity for linear, attention, and aggregate operations under prefill and decoding for batch size B. sequence length St., hidden dimension d. and generation length of & tokens.
Prefill
L Linear
Attention
Aggregate
FLOPs
O(B SL d²)
O SL
SL d)) O OB
MOPs
O(B S2 dd)
Old
O(B -
O O -d
activations
weights
scores activations (Q)
O(B-SE)
S < d
(O(S)). S <d ((S))) S
O(B- O(St),
Sy < d Sr >
Arithmetic Intensity
O(d),
St >< d
Decode
Linear
Attention
Aggregate
FLOPs
O B.d²)
O SL SL.d)
O(k. (.k.B.d') + O B-
MOPs
O(k B dd
Okk .F)
O B St) + O(k- Okk-B.S B..
O(k d²) O B St,
activations
weights
attention scores
activations (Cx,CV)
O(B),
SL <d
0(1), SL <d ((1), ST S
O(B), SL <
Arithmetic Intensity
O(B), SL >
(O(), St > d
3.1.2. COMPUTE VERSUS SEMORY-BOUND REGIMES
generating 1k tokens at different context lengths and batch sizes for the Linear/Attention components as well as the ag- gregate arithmetic intensity. To decide the ideal quantization strategy for different regimes, we consider the aggregate arithmetic intensity, which colored by the percentage of the total latency taken up by attention and provides plete view of decoding in all regimes. Based these results, we can clearly see that in the small batch + short context regime, the memory operations for the linear projections dominate latency, so weight quantization could provide considerable speedup in this regime. In the small batch + long context, large batch short context, and large batch + long context regimes, attention dominates latency due the expensive load-store operations for the large KV cache. KV cache quantization could help provide performance im- provements in these regimes. In the small batch + medium context and short context + medium batch regimes, the lin- ear and attention operations are approximately equivalent in their contributions to total latency. Thus, both weight and KV cache quantization are ideal here.
The asymptotic analysis suggests that in general, decoding suffers from low arithmetic intensities compared to prefill in all regimes. However, to decide which optimizations will most effectively improve latency, all regimes must be classi- fied either compute-bound or memory-boune Whether an operation compute or memory-bound depends on the hardware it is being run well the magnitude of the arithmetic intensity achieved by the operation.
We utilize an analytical roofline model (Williams et al., 2009; Kim et al., 2023a;b) to help determine which regimes are compute or memory-boune in practical inference set- ting. The roofline model defines ridge point which is calculated
peak compute performance (FLOP/s)
peak memory- (GB/s)
Note that the ridge point has the same units as arithmetic intensity (FL.OPs/byte). In the roofline model, any operation with an arithmetic intensity smaller than the ridge point is memory-boune -boudd, and any operation with an arithmetic inten- sity greater than the ridge point is compute-bound. For our analysis, extrapolate this to ridge plane and use hard- ware specifications for an NVIDIA A6000 GPU to study inference for the Llama-2-7B model in 16 bit precision.
4. QuantSpec
4.1. Overview of QuantSpec
In this section, we introduce QuantSpec, a self-speculative decoding framework designed to accelerate both short- and long-context generation by quantizing the model weights and KV cache into INT4 precision. We begin by noting that self-speculative decoding is particularly well-suited
For optimizing speculative decoding, we specifically focus on the decoding phase, although we include results for prefill in Appendix C. Figure 2 shows the arithmetic intensity for

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Linear
Attention
Aggregate
200
Compute Bound
Ridge Plane
Attemton of Roal Latency
Arthhmettc intensity
Arthimettc Interniity
Artthmmete Intersity
Merorry Boun
800
230
230
200
200
233
359
120
100
8. tgfe Batch Size
50
Batch Size
Batch Size
36
18 26
Context Length
Context Length
Context Length
200k
200k
Figure 2. Breakdown of how arithmetic intensity changes during decoding as the context length and batch size scaled logarithmically for linear, attention, and aggregate operations. All regimes lie below the ridge plane and thus are mnemory-bund The ridge plane is calculated for an NVIDIA A6000 GPU. The colors for the linear and attention surface plots simply represent the magnitude of the arithmetic intensity. The aggregate plot is colored by attention's 's runtime as percentage of the total latency. Prefill results in Appendix C.I.
for long-context generation, as the draft model shares the same architecture as the target model. This architectural alignment improves both the acceptance rate and the model's ability to handle long contexts effectively. However, a naive implementation of self speeulleulttee decoding (e.g. based on sparse KV) would require maintaining a separate, fully quantized copy of the KV cache, leading to inefficiencies in memory usage and computational overhead.
petitve generation quality with respect to the FP16 baseline while using half the KV cache's 's memory.
Datasets
KV Cache
WikiTexx2
C4
FP16 (Baseline)
6.4595
7.2617
INT8 (QuantSpec target)
6.4696
7.2620
To address this limitation, QuantSpec introduces a novel hierarchical KV cache design, which we discuss in detail in Section 4.2. This design enables dynamic switching between INT4 and INT8 representations of the KV cache without the overhead of on-the-fly quantization. By elimi- nating redundancy between the draft and target models' KV caches, our method significantly reduces the total memory footprint while preserving efficiency. We also address the inefficient combination of conventional quantization strate- gies with the snjjet-and-rever--bae back mechanism specific to speculative decoding methods by proposing pellipreiiioon KV cache buffer in QuantSpec. As we further explain in Section 4.3, this helps achieve high acceptance rates for the draft model and thus results in greater end-to-end speedup.
Table 2. Perplexity evaluations of Llama-2-70 with FP16, INT8 with group size = 128, residual length = 256 on different datasets.
Having observed this, we further note that an INT8 KV cache can be represented as an INT4 KV cache plus its INT4 residual. This works by decomposing an INT8 value into two INT4 components corresponding to its first and second 4-bit segments, which we call the upper and lower bitss. This effectively allows us to use hierarchical design to represent the KV cache of the draft model in INT4 and the target model in INT8 at the same time, removing the need to store a separate INT4 copy.
Our method is visualized in Figure 3. During prefill, QuantSpec quantizes the FP16 KV cache to form the up- per and lower INT4 representations. To obtain the upper 4-bits CENTS and lower 4-bits CINTA values, we first calcu- late CANT4 then quantize the quantization error ENTT to get CINTA C E [0, 15) uses asymmetric and round- to-nearest quantization. Since the distribution of E is symmetric and has an expectation close to zero, for C E [-8, 7] we use symmetric and round-to-nearest quantization to better match the distribution of errors.
4.2. Hierarchical KV Cache
We propose 4-bit hierarchical KV cache wherein strate- gically structure each tensor's 's representation such that the draft and target models are able to dynamically reconstruct their KV cache without any on-the-fly quantization over- head. Firstly, we observe that using an INT8 KV cache for the target model is comparable in terms of accuracy and performance with the same target model using an FP16 KV cache. To demonstrate this, we conduct a perplexity analysis for Llama-2-7B on the WikiText-2 (Merity et al., 2016) and C4 (Raffel et al., 2020) datasets in Table 2, which shows that the target model with an INT8 KV cache maintains com-
Then during decoding, when using the draft model to gener- ate candidate tokens, we only load the upper 4-bit representa- tion in our kernel and dequantize it for inference. When ver- ifying the drafted tokens using the target model, we utilize both the upper and lower 4-bit representations to reconstruct
5

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Input Tokens
Candidate Tokeas
Target Model
IP16 KV Cache
Draß Model (INT4 Weight)
Generre
Varey
Target Model
(EP16 Weight
(FP1 Weight)
Accepted Tokens
Rejected Tokens
2
Quantize the Quantization Error
Maintain the recent cache in Full Precision
Load Upper INT4
Load All 8-bit
① Quantize the KV Cache
INT4 Cacho
Upper INT4
Residual Cache
Upper INT4 Coche
Lower INT4 Cache
Cacbe
(a) During prefill, our hierarchical KV cache quantizes the FP16 KV cache to get the upper and lower INT4 representations.
(b) During decoding, our hierarchical KV Cache loads the Upper INT4 KV Cache for the draft model and loads all 8-bits for the target model.
Figure 3. How Hierarchical KV Cache works in the speculative decoding setting.
the KV cache in the higher INT8 precision. To represent the INT8 KV cache CANTS as the upper INT4 KV cache CINT4 and the lower INT4 cache C the INT8 KV cache can be expressed as C 2 + CINTA where we multiply by 2' to align their represented values. The asym- metric quantization for the KV cache can be represented as CP C şINTS + ZNNT,, where SINTS is the scaling factor, ZINT8 is the zero point, and CANTS E [0, 2 - 1]. In this scenario, its 4-bit representation can be viewed
Key Cache
Value Cache
token-wise
channel-wise
token-wise 6,587
6.507
channel-wise 7.041
6.911
Table 3. Perplexity of Llama-2-7B on WikiText-2 dataset with different quantization strategies. Group size G = 128. Channel- wise quantization for key cache and token-wist quantization for value cache gives the best performance.
FPP32_ = (2' + C + INS
CIITT GINTH + INTT S -CU
+CINT
24 + INT
in full precision until they equal the quantization group size, and then quantize them. However, since the KV cache for the most recent tokens is no longer preserved in full preci- sion after quantization, this strategy adversally affects the acceptance rate, thus reducing the effectiveness of specula- tive decoding. Moreover, since speculative decoding may result in frequent rollbacks due to the target model rejecting the draft tokens, the quantized KV cache for the rejected tokens need to be discarded and replaced with new tokens in the quantization group. This leads to repeated quantization and dequantization, slowing down the decoding process.
where ZINT4_ ZINTS
SINTA
S
4.3. KV Cache with Double Full Precision Buffer
4.3.1. CHALLENGES WITH KV CACHE QUANTIZATION AND SPECULATIVE DECODING
The key and value caches have each been found to exhibit unique characteristics indicating that they should be quan- tized with different strategies (Liu et al., 2024c). Specifi- cally, quantizing the key cache along the channel axis and quantizing the value cache along the token axis minimizes quantization error (as shown in Table 3), and therefore leads to higher acceptance rate in speculative decoding. We ap- ply asymmetric quantization and per-group quantization to both the key and value caches in INT4 precision, and we set the group size G to be equal to the head dimension to reduce overhead. These quantization techniques are illustrated in Appendix D Figure& 8.
4.3.2. ADAPT TO SPECULATIVE DECODING USING FULL PRECISION BUFFER
To enhance efficiency and ensure compatibility with spec- ulative decoding, we propose maintaining double full- precision buffer of size 2G, where G is the quantization group size. This buffer is divided into two equal parts: Op, and CF., each of size G. During prefill, we quantize the input tokens in batches of G while ensuring that at least G but no more than 2G of the most recent tokens remain in full precision. This ensures that CF, is always filled, In the decoding stage, newly generated tokens are stored in full precision in the second buffer, Cy,- Once the full- precision buffer reaches its maximum capacity of 2G, we wait for the target model to verify the generated tokens. If any tokens are rejected, we first remove the corresponding full full-precision KV cache entries. Then, we quantize CF and append it to the quantized KV cache. We then move
However, these quantization strategies for the key and value caches pose efficiency challenges when combined with spec- ulative decoding, Regarding the value cache wherein the values are quantized along the token axis, the naive strategy of directly quantizing newly generated tokens at each decod- ing step is expensive, as it introduces high computational overhead that occurs very frequently. Moreoever, regarding the key cache for which we apply quantization along the channel axis, the naive approach is to store multiple tokens
6

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
INT4 Quantized KV Cache
① Quantize the first half the full precision cache buffer
2 Move the sccond half to the first
half when the second half is full
Allocated Static Cache (Currently Empty)
head dim
First Half of the Full precision cache buffer
Quantized Cache (INT4)
Verified Tokens
Full Precision
Second Half of the Full precision cache buffer
sequence length
cache buffer
Move rest second half'
head dim
of the cache buffer
Tokens verified by the Target Model
Quantized Cache (INT4)
Figure 4. Our KV cache with 2 full precision cache buffers for recent KV cache.
Cry to Cr., which fully occupies Cr, while leaving Crs, empty and ready for tokens generated in future decoding steps. This whole process is visualized in Figure 4.
5. Evaluation
In this section, evaluate the performance of QuantSpec across multiple datasets and context lengths. Our evalua- tion focuses on three key dimensions: (1) the acceptance ratio between the draft and target models, (2) GPU mem- ory consumption, and (3) end-to-end serving speedup. We begin by presenting detailed benchmarking acceptance rate, memory usage, and end-to-end speedup across dif- ferent datasets. Then, we highlight the performance gains achieved by our custom kernels for quantized KV cache. Fi- nally, We also present an extensive ablation study focusing on the contribution of weight versus KV cache quantization to the final speed-up.
Using this design, we ensure that (1) at every step Cr, is always filled, so there are at least recent G tokens kept in full precision, which is beneficial for the acceptance rate. (2) Quantization and KV cache movement will only happen every G' decoding steps, which significantly reduces the overhead. (3) The design is compatible with speculative decoding since we can discard the KV cache for rejected tokens very flexibly by only operating on the second full- precision buffer CF2 and without needing extra quantize and dequantize operations. We also show that our method is fully compatible with FlashDecoding in Appendix E.
5.1. Setup
4.3.3. SUMMARY
All experiments are performed on a node equipped with 8 NVIDIA RTX A6000 GPUs. We evaluate QuantSpec using long-context variants of LLaMA-2 and LWM mod- els as target models. For benchmarking decoding speedup, we PG-19 (Rae et al., 2019), (an open-vocabulary lan- guage modeling benchmark derived from books) and two long context summarization datasets, namely BoBeNcu Sum (Zhang et al.. 2024c; Yen et al., 2024) and Multi-LexSum (Shen et al., 2022; Yen et al., 2024). More details about the datasets are provided Appendix F. Following Sadhukhan et al. (2024), we compare against two recent sparse KV- based self-speculative decoding baselines: StreamingI.L.M
In summary, QuantSpec allows the draft and target models to share the same architecture in a self-speculative decoding manner, ensuring greater consistency between drafting and verification opposed to traditional big-little speculative decoding methods. Our approach is mainly designed for long-context scenarios, where efficient KV cache manage- ment is critical, but it also supports short contexts where using weight quantization becomes more critical. We quan- tize the KV cache using our hierarchical INT4 design and use double full-precision cache buffer for higher accep- tance rates and flexibility with speculative decoding. Then in the decoding stage, when generating draft tokens, we only load the upper b-bit of the KV cache and achieve speedup by significantly reducing the memory load/store operations. When verifying these draft tokens, we load both the upper and lower 4-bit KV cache representations and dequantize them into their INTS representation to achieve performance that is comparable with an FP16 KV cache. If the full- quantize
(Sadhukhan al., 2024; Xiao al., 2023a) and SnapKV (Sadhukhan et al., 2024; Li et al., 2024a). To ensure fair comparison, the draft KV budget for the basclines is set to one-fourth of the context length, matching our 4-bit quan- tized KV cache. We fix the quantization group size at 128, the residual length R for the KV cache 256, and limit the number of output tokens to 90. The optimal speculation length y for each dataset is determined through hyperpa- rameter search for each dataset-mode pair. Details of the hyperparameter search are provided in Appendix G.
precision buffer is saturated, after verification and clear one-half of the full-precision buffer to prepare for the next round of generation. The whole algorithm can be visualized in Figure 3 (and Algorithm in Appendix).

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Table Efficiency result of QuantSpec on Llama-2-7B-32K, 2-1-78---2K-nnnstrutt and LWW-Texx--ha-- 128k. We benchmark on multiple context length settings, ranging from 4k to 128k for batch size 1, and take their average across 10 different examples. Speedup ratio is compared with autoregressive generation the target model (AR), compare with Sparsity-basee -based self-speculative baselines that use Lt and SnapKV to quantize the KV cache. Acceptance rate of each method is shown for the respective optimal We outperform these baselines and achieve a maximum of 2.49x speedup compared with autoregressive generation.
Llama-2-7B- 2-t-----------tttteeee
Datasct
Context Length
# GPUs
Method
Acceptance Rate (%)
Peak GPU Memory (GB) 4
Speedup (x AR)
StreamingL.IMM
88.87
15.20
1.13
4k
SnapKV
93.59
15.39
1.17
QuantSpec
92.46
16.26
1.35
PG19
StrcamingLLMM
90.31
17.75
1.27
8k
SnapKV
94.39
18.10
QuantSpec
89.88
17.66
1.44
StreamingLLMM
90.78
17.75
1.27
8k
SnapKV
55.55
18.10
1.01
QuantSpec
91.23
17.66
1.61
Multi-LexSum
Streaming
91.19
32.61
1.84
32k
SnapKV
72.54
QuantSpee
91.16
25.84
2.08
LWM-Text- SWWM-TextCCht-11225
Datasct
Context Length
GPUs
Method
Acceptance Rate (%)
Peak GPU Memory (GB)
Specdup AR)
StreamingL.LM
87.41
22.63
1.49
16k
SnapKV
88.61
23.33
1.50
QuantSpec
91.64
20.36
1.78
BoBENC Sum
StreamingLLMM
85,38
32,63
1.15
32k
SnapKV
84.55
33,98
1.78
QuantSpec
90.99
25.
1.99
StreamingLLM
78.95
52.96
2.16
64k
2
SnapKV
82.79
55.66
QuantSpec
92.18
38.18
2.23
Multi-LexSum
StreamingI.I
OOM
128k
2
SnapKV
0OM
QuantSpec
94.31
61.22
2.49
5.2. Speedup Evaluation
5.2.1. KERNEL SPEEDUPS
Table 4 shows the acceptance rate, GPU memory required, and speedup achieved compared to autoregressive decod- ing. We observe that QuantSpec provides consistently better speedups for all context lengths. For short and medium context lengths (e.g. 8k and 32k prompt length), QuantSpec achieves ~1.61 x to ~2.08 speedups respectively on the Multi-LexSum dataset. For longer context lengths (e.g. 128k), our speedups while using lower GPU memory than the baselines. We also see that acceptance rates of QuantSpec are considerably higher than the baselines for summarization tasks (refer to Appendix for detailed comparison); this shows that for such tasks where the whole context is important, sparse KV cache methods are much more lossy, whereas quantization preserves most of the information in the context. Conse- quently, QuantSpec proves to be more reliable choice, de- livering consistent speedups across varying context lengths and query complexities.
In Table 5 show the speedup achieved using custom attention kernel that makes of quantized KV cache ver- sus the standard FP16 FlashAttention kemnlss. For context length of 128k, our INT4 attention kernel is ~ 2.88 faster than the standard FlashAttention kernel.
even greater, up ~2.99 all
Table 5. Latency benchmark of custom attention kernels for calculating attention with quantized hicrarchical KV cache. QuantSpec INT4 refers to only loading upper-4-bit, QuantSpec INTS refers to loading both the upper and lower 4-bit. Bench- marked on kernel level.
Context Length
Kernels
64k
256k
FlashAttention (FP16)
3.07 ms
6.16 ms
QuantSpec INT 8
1.08 ms (1.44x)
4.06 ms (1.51x)
QuantSpec INT
0.54 ms (2.88x)
ms (2.86x)
8

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Autoregressive
KV
INT4 Weight
QuaniSpec
consequences of our work, none which we feel must be specifically highlighted here.
Spocdup Ratio
1.6
1.4
Acknowledgmentss
1.2
We acknowledge gracious support from Apple team, as well as Nvidia for providing GPU hardware. We also appreciate the support from Microsoft through their Accelerating Foun- dation Model Research. Furthermore, we appreciate support from Google Cloud, the Google TRC team, and specifically Jonathan Caton, Divy Thakkar, and Prof. David Patterson. Prof. Keutzer's lab is sponsored by the Intel corporation, Intel One-API, Intel VLAB team, the Intel One-API center of excellence, as well as funding through BDD, BAIR, and Furiosa. We appreciate great feedback and support from Ellick Chan, Saurabh Tangri, Andres Rodriguez, and Kit- tur Ganesh. Sehoon Kim and Suhong Moon would like to acknowledge the support from the Korea Foundation for Advanced Studies (KFAS). MWM would like to acknowl- edge DARPA, DOE, NSF, and ONR for providing partial support of this work. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
0.8
2048
4096
8192
16384
32768
Context Length
Figure 5. Speedup ratio of QuantSpec compared to autoregressive baseline as we scale the context length. We report QuantSpec with KV cache-only quantization, weight-only quantization, and both. Benchmarked on Llama-2-7B- 2-n-a-----322 22k-Instruct using PG-19.
5.3. Ablation Results
We present an extensive ablation study of QuantSpec focus- ing on the contribution of weight versus KV cache quantiza- tion the final speed-up.
Weight versus KV Quantization: Figure 5 illustrates the speedup ratio of QuantSpec compared to autoregressive baseline as context length increases. The figure benchmarks QuantSpec with KV cache-only quantization, weight-only quantization, and both. The results are aligned with the analysis done in Section 3.1, showing that for short contexts most of the speedup comes from quantizing weights, for medium length prompts both weight and KV cache quanti- zation contribute to the final speedup, and KV cache quanti- zation is most effective for long contexts.
References
Bhendawade, N., Belousova, I., Fu, Q., Mason, H., Raste- gari, M., and Najibi, M. Speculative streaming: Fast 11m inference without auxiliary models. arXiv preprint arXiv:2402.1113 2024.
6. Conclusions
Brandon, W., Mishra, M., Nrusimha, A., Panda, R., and Kelly, J. R. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint SKGV:2200.1298.1 2024.
In this paper, we have introduced a novel approach en- hance the efficiency and scalability of Large Language Models (LLMs) in long-context settings through quantized speculative decoding. Our method addresses the increas- ing memory and computational demands by optimizing the Key-Value (KV) cache operations, which become a sig- nificant bottleneck as the context length grows. We pro- pose a double full- full-precision cache buffer resolve conflicts between per-group quantization and speculative decoding. Our comprehensive approach shows that by integrating ad- vanced quantization techniques with speculative decoding, it is possible to significantly improve processing speed with- out compromising the accuracy and performance of L.Ms.. This work paves the way for more scalable and effective deployment of LLMs in applications that require extensive contextual understanding, offering a robust solution to the challenges posed by long-context settings.
Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint arXi::2401.10774 2024.
Chee, J., Cai, Y., Kuleshov, V., and De Sa, C. M. Quip: 2-bit quantization of large language models with guarantees. Advances in Neural Information Processing Systems, 36, 2024.
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint urXiv:2302.013188 2023.
Impact Statement
Chen, S., Liu, Z. Wu, Z., Zheng, C., Cong, P., Jiang, Z., Su, L.,, and Yang, T. Int-flashattention: Enabling flash attention for int8 quantization. arXiv preprint arXiv:2409. 16997, 2024a.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Chen, Z.,, May, A., Svirschevski, R., Huang. Y., Ryabinin, M., Jia, Z., and Chen, B. Sequoia: Scalable, robust, and hardware-aware speculative decoding. arXiv preprint arXiv:2402.12377, 2024b.
M. W., Shao, Y. S., and Gholami, A. Full stack opti- mization of transformer inference: survey, 2023b. URL
Sttps://arxier wttpaxiwww.rrg/b////eee... 14017.
Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M. W., Gholami, A., and Keutzer, K. Speculative decod- ing with big little decoder. Advances in Neural Informa- tion Processing Systems, 36, 2024.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io- awareness, 2022. URL https: //wpii/aavvvv.//aa org/ abs/ 2205.14133
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Inter- national Conference Machine Learning, pp. 19274 19286. PMLR, 2023.
Dao. T.. Haziza, D., Massa, F., and Sisov, G. Flash-decoding for long-context inference: /ttdd//222///////
https: http://cff...ss..f...osss 12/ 22/flashdecoding, htt.eeeoodinnn.m.......... 2023.
Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404. 14469, 2024a.
Fishman, M., Chmiel, B., Banner, R., and Soudry, D. Seal- ing fp8 training trillion-token 1lms. arXiv preprint arXiv::409 12.409...77.... 2024.
Fu, Q. Cho, M. Merth, T., Mehta, S. Rastegari, M., and Najibi, M. Lazyllm: Dynamic token pruning for efficient long context IIm inference. arXiv preprint arXiv:2407. 14057, 2024.
Li, Y., Wei, F,, Zhang, C., and Zhang, H. Eagle: Speculative sampling requires rethinking feature uncertainty, arXiv preprint arXiv:2401. 15077, 2024b.
Lin, J., Tang, J., Tang, H., Yang, S., Chen, W.-M., Wang, W.-C., Xiao, G., Dang, X., Gan, C., and Han, S. Awq: Activation-aware weight quantization for on-device I1m compression and acceleration. Proceedings of Machine Learning and Systems, 6:87-100, 2024.
Ge, S., Zhang. Y., Liu, L., Zhang, M., Han, J., and Gao, J. Model tells you what discard: Adaptive kv cache compression for Ilms. arXiv preprint arXlv:2310.0180, 2023.
Hooper, C., Kim, S., Mohammadzadeh, H., Mahoney, M. W., Shao, Y. S., Keutzer, K., and Gholami, A. Kvquant: Towards 10 million context length IIm infer- ence with kv cache quantization. Advances in Neural Information Processing Systems, 2024.
Liu, D., Chen, M., Lu, B., Jiang, H., Han, Z., Zhang, Q.. Chen, Q. Zhang, C., Ding, B., Zhang, K., et al. Re- trievalattention: Accelerating long-context Ilm inference via vector retrieval. arXiv preprint arXiv:2409.100 10516. 2024a.
Horton, M., Cao, Q., Sun, C., Jin, Y., Mehta, S., Rastegari, M.. and Nabi, M. Kv prediction for improved time to first token, 2024. URL 2410. 08391.
Liu, L., Qu, Z., Chen, Z., Ding, Y., and Xie, Y. Trans- former acceleration with dynamic sparse attention. arXiv preprint arXiv:21 200:::2100.110 12999, 2021.
Liu, Z., Desai, A., Liao, F,, Wang, W., Xie, V., Xu, Z., Kyril- lidis, A., and Shrivastava, A. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024b.
Rang. H.. Li, Y.. Zhang, C., Wu. Q., Luo, X., Ahn, S., Han, 7., Abdi, A. H., Li, D., Lin, C.-Y., et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprini
2024.
Kang, H., Zhang, Q. Kundu, S., Jeong, G., Liu, Z., Krishna, T., and Zhao, T. Gear; An efficient kv cache compression recipefor near-lossless generative inference of Ilm. arXiv preprint arXiv:2403.05522 2024.
Liu, Z., Yuan, J., Jin, H., Zhong, S., Xu, Z., Braverman, V., Chen, B., and Hu, X. Kivi: A tuning-free asym- metric 2bit quantization for kv cache. arXiv preprint arXiv.240.002550, 2024c.
Kim, S.. Hooper, C., Gholami, A., Dong, Z., Li, X.. Shen, S., Mahoney, M. W., and Keutzer, K. Squeezellm: Dense-and-spe quantization. arXiv preprint arXiv:2306.07623 2023a.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models, 2016.
Nawrot, P., Laúcucki, A., Chochowski, M., Tarjan, D., and Ponti, E. M. Dynamic memory compression: Retrofitting llms for accelerated inference, arXiv preprint arXiv:2403.09mm 096666 2024.
Kim, S., Hooper, C., Wattanawong, T., Kang, M., Yan, R., Gene, H., Dinh, G., Huang, Q., Keutzer, K., Mahoney,
10

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Park, Y. Hyun, J., Cho, S., Sim, B., and Lee, J. W. Any-precision IIm: Low-cost deployment of multiple, different-sized Ilms. arXiv preprint arXiv:2402.10517, 2024.
accelerators. In Proceedings of the ACM/IEEE 47th An- nual International Symposium on Computer Architecture, ISCA '20, 2020.
Williams, S., Waterman, A., and Patterson, D. Roofline: insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65-76, 2009.
Peng, H., Wu. K., Wei, Y., Zhao, G., Yang, Y., Liu, Z., Xiong, Y., Yang, Z., Ni, B., Hu, J., et al. Fp8-Im: Training fp8 large language models. arXiv preprint AXiv:2310.1831133 2023.
Xi, H.,, Cai, H., Zhu, L., Lu, Y., Keutzer, K., Chen, J., and Han, S. Coat: Compressing optimizer states and activation for memory-efficient fp8 training, 2024a. URL /ttps://arxiv..er wttp://gg.bb.....a...
Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long- range sequence modelling. arXiv preprint, 2019. URL ht https: /tapiii//xxiv...///////// owgg/bgg//s... 9111. 05507.
Xi, H., Chen, Y., Zhao, K., Zheng, K., Chen, J., and Zhu, J. Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization. arXiv preprint arXiv:2403. 12.3.124:5 12422, 2024b.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., Liu, P. J., et al. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., 21 21(140):166) 2020.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv, 2023a.
Sadhukhan, R., Chen, J., Chen, Z.,, Tiwari, V., Lai, R., Shi, J.,, Yen, I. E.-H., May, A., Chen, T., and Chen, Magicdec: Breaking the latency -hrrouguhtu tradeoff for long context generation with speculative decoding. arXiv preprinl arXiv:2408. / 2024.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. arXiv preprint arXiv:2300.. 1..3099.7443. /7453, 2023b.
Shah, J., Bikshandi, G., Zhang. Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate atten- tion with asynchrony and low-precision. arXiv preprini arXiv:2407.08600, 2024.
Yang, S., Sheng, Y., Gonzalez, J. E., Stoica, L., and Zheng, L. Post-training sparse attention with double sparsity. arXiv preprint arXiv:2408.07092. 2024.
Yao, Y., Li, Z., and Zhao, H. Sirllm: Streaming infinite retentive 1lm. arXiv preprint 245. 2024.
Shao, W., Chen, M., Zhang, Z., Xu, P., Zhao, L., Li, Z., Zhang, K., Gao, P., Qiao, Y., and Luo, P. Omniquant: Omnidirectionally calibrated quantization for large lan- guage models. arXiv preprint arXiv:2308. 13137, 2023.
Yen, H., Gao, T., Hou, M., Ding, K., Fleischer, D., Izsak, P., Wasserblat, M., and Chen, D. Helmet: How to evaluate long-context language models effectively and thoroughly. arXiv preprint WxXi:2241000264 2024.
Shen, Z., Lo, K., Yu, L., Dahlberg, N., Schlanger, M., and Downey, D. Multi-lexsum: Real-world summaries of civil rights lawsuits at multiple granularities. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. URL https: // //openrevies net. teee/fformet forum? d= 12l222288 s8cr.
Zhang, J., Huang. H., Zhang, P., Wei, J., Zhu, J., and Chen, J. Sageattention2 technical report: Accurate bit attention for plug-and-play inference acceleration. arXiv preprint Xrx: 2et::2211..999.. 10958, 2024a.
Sun, H., Chen, 7.,, Yang, X., Tian, Y., and Chen, B. Tri- force: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint 2r0iv:2404. 19/2, 2024.
Zhang, J., Zhang, P., Zhu, J., Chen, J., al. Sageatten- tion: Accurate 8-bit attention for plug-and-play inference acceleration. arXiv preprint arXiv:2410.02367. 2024b.
Zhang, X., Chen, Y., Hu, S., Xu, Z., Chen, J., Hao, M. K., Han, X., Thai, Z. L., Wang. S., Liu, Z., and Sun, M. oobench: Extending long context evaluation beyond 100k tokens, 2024c. URL https: http::///xxxvvvo. /bs/ 2402.
Tan, S., Li, X., Patil, S., Wu, Z., Zhang, T., Keutzer, K., Gonzalez, J. E., and Popa, R. A. Lloco: Learning long contexts offline. arXiv preprint 2rX1v:2404.07979, 2024.
Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikci, B., and Han, S. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406.100 107744 2024.
Zhang. Z., Sheng, Y., Zhou, T., Chen, T., Zheng, I.., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024d.
Weng, J., Liu, S., Dadu, V., Wang, Z., Shah, P., and Nowatzki, Dsagen: Synthesizing programmable spatial

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Zhao, J., Lu, W. Wang, S., Kong, L.., and Wu, C. Qspec: Speculative decoding with complementary quantization schemes. arXiv preprint 2rX.v:2410.11305, 2024.
Zhou, Y., Chen, Z., Xu, Z., Lin, V., and Chen, B. Sir- jus: Contextual sparsity with correction for efficient Ilms. arXiv preprint 2rXiv:2409003556 2024.
12

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Appendix
A. Attention Module's Inference Workflow
The inference of LLMs can be divided into 2 parts: the prefill stage and the decoding stage. In the prefill stage, for the input sequence E R Bx5xxx xd the KV cache update rule can be calculated as
Q XWo: CK = XWK, Cv=XWi
where we denote the query, key, and value weight matrices as Wq, WK, Wy E Rd and denote the key and value caches as CK and Cv respectively. B refers to batch size, Sz refers to sequence length, and d refers to hidden size. We then calculate the multi-head attention (MHA)
O MultiHHeadAttu(Q. Ck: Ov).
In the decode stage, for input token E RP Rx xxd, we first calculate the query, key, and value of the current token:
zWQ, Cx = xWK, = xWv,
then concatenate the KV cache with the current token's 's key and value to update the KV cache:
Ck = concat(Ck, Cv = concat(Civ,..)
Then, the multi-head attention output is calculated:
O MMlttlleadAttnlg Ck₁ Cv).
B. More Related Works
We list some related works that we find interesting, but can not elaborate in the related works section due to space limitations.
Efficient Long Context Inference Some research maintains the full key-value pairs but dynamically loads them from high- high-bandwidth memory (Yang et al., 2024; Tang et al., 2024), and usually achieves higher performance at the cost of higher memory consumption. Shared KV cache across tokens (Nawrot et al., 2024) and layers (Brandon et al., 2024) provides a new way to reduce the KV cache budget through sharing.
Quantization Any Precision representation (Park al., 2024) incorporates multiple precision levels (e.g., INT2, INT4, and INT8) within single representation, eliminating the need to store separate KV caches for each precision and allowing the framework to dynamically select the optimal precision based on the complexity of the task. Training quantization (Peng et al., 2023; Xi et al., 2024b; Fishman et al., 2024; Xi et al., 2024a) reduces the bit precision of various model parameters, gradients, and activations to accelerate training. Attention quantization (Chen et al., 2024a; Zhang et al., 2024b;a; Weng et al., 2020; Shah et al., 2024) reduces the computational overhead associated with attention computations, which becomes dominant in the prefill stage of the long-context inference setting.
Speculative Decoding Zhao et al., (Zhao et al., 2024) explored complementary quantization schemes speculative decoding with QSpec, enhancing efficiency without significant performance degradation. Sirius (Zhou et al., 2024) finds that contextual sparsity will lead poor performance under the speculative decoding setting since the model performance is degraded, and thus it cannot accelerate LLM inference.
C. Additional LLM Inference Bottlenecks Analysis
C.1. Prefill Arithmetic Intensity Analysis
Keeping in line with the asymptotic analysis in Table 1, the arithmetic intensity of attention during prefill does not scale with batch size at all, as attention unable to benefit from batching in the same way that linear operations do. Moreover,
13

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
for long contexts, attention entirely dominates the linear operations due the quadratic nature of self-attention. For short contexts however, this quadratic cost relatively inexpensive when compared the linear operations. As shown in Figure 6, the arithmetic intensity for all prefill operations in all regimes is above the ridge plane, which means that prefill is entirely compute-bound
Linear
Attention
Aggregate
00
2000
Arithieeic Intensity
Arithmelic Intensity
0000
#0000
Attertion of
Arthmmttc
1390
Ridge Plane
69333
3000
tife
20000
300
Compute Bound
Batch Size
1-b Batch Size
Batch Size
Contest Length
12R 33
Context Length
Contest Length
2014 anc Memory Bound
salk 2924
zack
Figure 6. During prefill, all regimes lie above the ridge plane and thus are compute-boure bound.
C.2. Modern GPU Hardware VRAM Size Constraints
Cache Memory Footprint
140
120
100
80
2000
60
Rutio of KV Cache to Modal Weights
1500
Memory (60)
1000
900
250
Batch Size
Contast Length
xx NVIDIA T4 (16 GB)
GeForce RTX 4090 (24 GB) NVIDIA 8x NVIDIA A100/1100 (80 GBI
P100 GB)
GB)
Geforce RTX 3090 (24 GB)
single node (8 GPUs) context length and batch size scaled logarithmically. The surface plot's 's color represents the ratio of KV cache memory to the model weights memory. The dotted- lines represent GPU DRAM capacities for several different GPUs. At (B- 16, Sr 262k), the KV cache takes up 160x more memory than the model weights.
Figure 7. KV cache memory usage by Llama-2-7B
The relatively higher linear arithmetic intensities observed in decoding for batch sizes greater than 8 in Figure 2 are misleading due to the limited VRAM sizes in modern GPUs. As shown in Figure 7, the size of the KV cache for Llama-2-7B model exceeds the total VRAM capacities of a single node equipped with 8 A100/H100 GPUs with 80 GB of memory each. This means that simply scaling the batch size for decoding will not translate the memory- bound nature of generation to being compute-bound
14

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
D. Quantization Strategies for KV cache
Here we provide a visualization our quantization scheme Figure 8. We apply asymmetric quantization for both keys and values cache, and upply channel-wise quantization the key cache and apply token-wise quantization to the value cache.
sequence length
group-wise quantization along channel axis
Quantized Keys
Scale Factors (FP16)
Zero Points (FP16)
(INT)
Zero Points (FP16)
Scale Factors (FP16)
sequence length
Quantized Values (INT)
quantize along
Figure 8. We apply asymmetric and per-group quantization for both the key cache and value cache, along the channel axis and token axis, respectively. This figure describes how it works when only the upper-4 bit cache is applied.
E. Compatibility with Flash Decoding
Our full-precision buffer design, as shown in Figure 4, is fully compatible with Flash Decoding (Dao al., 2023), fast attention implementation available for the decoding stage. In this setup, the quamtized section divides naturally into several chunks, aligning perfectly with the structure of Flash Decoding. For the full-precision buffer, given its upper bound length of 2G, can be processed independently with minimal overhead. This full-precision segment can be treated as an additional chunk, which can then be summed with the quantized segments and seamlessly integrated into the Flash Decoding algorithm.
F. Details about the Datasets Used
We provide overview of the datasets used in experiments, highlighting their key characteristics.
WikiTexxt22 (Merity et al., 2016): WikiText-2 is widely used dataset for language modeling. It is a subset of the larger WikiText dataset and consists of high-quality, clean, and well-structured English text extracted from Wikipedia articles.
C4 (Raffel al., 2020): C4 is large scale web-crawled language modelling dataset mostly used for pretraining LLMs.
PG-19 (Rae et al., 2019): It is a dataset of books from Project Gutenberg, designed for long-context language modeling.
BeBEN Sum (Zhang et al., 2024c): InfiniteBench benchmark tailored for evaluating the capabilities of language models to process, understand, and reason super long contexts. We used of its summarization datasets where the task is to summarize fake book created by core entity substitution. The average length of input prompt is 177k.
Multi-LexSum (Shen et al., 2022): Multi-LexSum is a multi-doc summarization dataset for civil rights litigation lawsuits. The average length of prompt in this dataset is 90kk.
15

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
G. Hyperparameter Search
Here we present details about the hyperparameter search done to select optimal 7y for each experiment. We search 7 for each dataset and method pair using prompt length of 8192 and use the same value for all other context length experiments. Table 6 shows the results of the search. We find that sparse-based methods achieves maximum performance when ? equals to 1, while our quantization-bart based method usually achieves the best performance with larger y such as 4 or 6.
Table 6. Hyperparameter Search for Llama-2-7B-32K and LWM-Text-Caa-1284 models on PG19, Multi-LexSum, and ccBENCI Sum datasets. Context length is kept as 8k.
(a) Llama-2-7B-32k on PG19
(b) Llama-2-7B-32X on Multi-LexSum
Method
Acceptance Rate 1
Speedup f
Method
Acceptance Rate
Speedup
1
90.78
39.1
1
90.78
39.17
StreamingLLM
2
89.42
38.5
StreamingLLM
2
86.82
37.84
3
90.21
38.66
3
83.29
36
1
94.39
40.34
1
55.55
31.05
SnapKV
2
91.03
39.38
SnapKV
2
43.96
24.39
3
91.84
39.28
3
36.61
19.78
1
91.88
41.52
96.58
42.83
89.88
44.51
2
96.61
47.51
QuantSpec
QuantSpec
83.17
43.84
95.59
49.47
77.07
41.88
91.23
49.62
(c) I.W.-Text-Chat 128k on BoBENNH Sum
(d) LWM Tex--CCa---haa on Multi-LexSuu
Method
7
Acceptance Rate t
Speedup f
Method
Y
Acceptance Rate
Speedup
1
81.79
37.17
83.96
37,80
StreamingLLMM
2
74.86
34.33
Streaming
2
77.41
35.13
3
64.48
30.14
3
71.28
32.37
1
85.55
38.27
1
89.25
39.04
SnapKV
2
82.92
36.97
SnapKV
2
83.53
37.22
3
77.13
34.40
3
80.04
35.48
93.83
42.01
95.94
42.79
2
94.38
46.74
2 95.06
47.10
QuantSpec
QuantSpec
4
90.33
47.30
4
92.55
48.15
6
82.13
45.51
6
87.73
48.20
H. Comparing Acceptance Rates
Although the acceptance rates Table 4 for the sparse KV baselines and QuantSpec do seem exceedingly different, this misleading because the acceptance rates shown are for the optimal y values observed from our hyperparameter search in Table 6. Table effectively compares the acceptance rates of the baselines very low (e.g. with those of QuantSpec at much higher 7 (e.g. 6). Here, we compare the acceptance rates of different self-speculative decoding algorithms. For fair comparison, Figure 9 illustrates the acceptance rate between the draft and target models as function of speculation length. We observe that QuantSpec consistently outperforms sparse KV approaches terms of acceptance rate. Notably, as speculation length increases, the acceptance rate sparse KV methods degrades much faster, whereas our method maintains high acceptance rates.
16

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
90
Acceptance Rate
80
70
60
QuantSpec
SnapKV
StreamingLLM
50
2
4
5
Speculation Length (v)
Figure Acceptance rate self-speculative decoding methods different speculation length measured model LWM-Text Bexxt-hat--iee 128k on Multi-LexSuu dataset.
Additional Results
Here we present some additional results newer model architectures. Table 7 shows acceptance rate and speedup evaluations on Mistral and Llama 3.1 models. The results follow the same trends and conclusions as the results in the main paper.
Table Efficiency result of QuantSpec on Mistral-v0.3 and Llama models on the Multi-LexSum dataset different context lengths on single RTX A6000 GPU. We adopt the best speculative length for each method.
Mistral-7B-v0.33
Context Length
Method
Acceptance Rate (optimal 7)
Peak GPU Memory (GB) |
Speedup (x AR) t
StreamingLLM
0.86 (1)
16.25
0.94
16k
SnapKV
0.86 (1)
16.70
0.93
QuantSpec
0.94 (6)
17.42
1.55
StreamingLLM
0.89 (1)
18.67
1.07
32k
SnapKV
0.85 (1)
19.79
0.98
QuantSpec
0.93 (6)
18.34
1.61
Llama-3.1-8B
Context Length
Method
Acceptance Rate (optimal 7)
Peak GPU Memory (GB) |
Speedup (x AR)
StrcamingLLM
0.63 (1)
17.73
0.82
16k
SnapKV
0.66 (1)
18.18
0.85
QuantSpec
0.90 (6)
18.79
1.48
StreamingLLM
0.81 (1)
20.15
0.93
32k
SnapKV
0.73 (1)
21.27
0.90
QuantSpec
0.92 (6)
19.82
1.54
StreamingILLM
0.89 (1)
34.91
1.06
128k
SnapKV V
0.80 (1)
39.95
1.06
QuantSpec
0.91 (6)
26.05
1.63
17

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
J. Kernel Implementation Detail
In our approach, we implement the algorithm using Flash Decoding framework. In the initial stage 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 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 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, separate computation 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.
18

---

QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache
Algorithm 1 QuantSpec Algorithm
Input: Model M, Upper 4-bit Cache Cu, Lower 4-bit Cache CL. Full Precision Cache Buffer Cp.]
Input: Prefill Length Sp, Target Decode Length Sp. Prefill Context X (zo,-" _T5p-11, Speculate Length
Input: Number of Layers L, Sensitive Layer Number Ls. Quantization Group Size G
Function: PREFILL, DRAFT, TARGET, VERITY, QUANTIZE, REJECTCACHE
Notation: Verified tokens 2i, generated draft tokens Ji, logits of draft model dis logits of target model pi, number of tokens already been generated total N, number of tokens already been generated this speculate step n, number of tokens accepted in this speculate step
Prefill Stage
1: Xsp: ORV < PREFILL(M, X xsPP)
D KV Cache is written together for simplicity Prepare the hierarchical quantized KV Cache b Prepare the full-precision cache buffer for recent tokens
2: Cv. + QUANTIZINECCCV.S....
Ls)
3: CA, CP. + -KVsp-o: None
Decode Stage
4: while SD do
5: n 0 and
6: while y do
7: Qn+1, Cr. + DRAFT(M, Cu. Cri Cr,, Ls, XSSp+N + Ien)
8: Sample gn+1 4n+1 and n + n +
9: end while
PI. Py, CF2 TARGET(M, CU, CR,, OF.s XES++N + 9<g)
for to to do
if VERIFY (gi, Pi, ②)) then
13:
2N++ 9: and + ++1
14:
else
15:
2N+1 + CORRECT(p: 9:) and u ++1
16:
Cr.- - REJECTCACHE(CC..
D Clear the rejected KV cache from the full precision cache buffer
Break
end if
18:
if = then
IN+++1 + P++1 and N t- N +1
20:
end if
21:
end for +- +
22:
if Cr, is full then
23:
Concatenate QUANTIZE(CH, with Cy and Cr.
> Quantize the first half of the full precision cache buffer
OF, CF. :-G and CF, CF. -G:
Move the second part to the first part of full precision buffer
25:
end if
26: end while
19


Reviews:
[{"review_id": "sirDmnMuMN", "reviewer": "Reviewer_ikhW", "summary": "This submission introduces QuantSpec, a novel self-speculative decoding method that employs KV-Cache quantisation during token drafting, to optimise the inference efficiency on long-context LLMs. Based on the insight that in long-context sequences much of the pressure in memory bandwidth is attributed to loading KV Cache entries during decoding, QuantSpec introduces a hierachical 4-8 bit quantized KV cache, used for drafting and verification respectively. A custom CUDA kernel implementation tailored to the proposed KV cache structure allows the proposed approach to realise inference speed-up in real-world GPU deployment.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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: \n\n> 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. \n\n**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.\n\nKV cache dominates when\n$\\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.\n\n---\n\n> R1-2: Are there any representative works of 7B models adopting >100k token context length?\n\n**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.\n\n---\n\n> 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?\n\n**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.\n\n---\n\n> R1-4: How does your quantization scheme for KV cache different from other nested quantization approaches like AnyPrecision LLM.\n\n**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. \n\n| Method                    | Quantization Error (MSE) ↓ | Acceptance Rate ↑ |\n|:-------------------:|:--------------------:|:------------------:|\n| Anyprecision 4 bit  |        0.761                                  |       0.86                     |\n| QuantSpec  4 bit    |        0.013                                  |       0.92                     |\n\nWe will add this analysis in Appendix and will also add a discussion about nested quantization in related works.\n\n---\n\n> R1-5:  Can your method be applied with traditional self-speculative decoding approaches like SWIFT or LayerSkip?\n\n**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. \nWe will add a discussion about compatibility with other self-speculative decoding algorithms in the Appendix. \n\n---\n\n\n[1] Any-Precision LLM: Low-Cost Deployment of Multiple, Different-Sized LLMs, ICML 2024\n\n[2] LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding, ACL 2024\n\n[3] SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration, ICLR 2025", "claims_and_evidence": "The claims made in this submission are supported by clear evidence for the examined cases. However, a broader analysis is required to solidify the findings and demonstrate the generality of the proposed approach. Please see below for more details suggestions.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria are meaningful and suit well the examined problem and usecases.", "theoretical_claims": "Not Applicable.", "experimental_designs_or_analyses": "The presented experiments are convincing and demonstrate the effectiveness of the proposed approach via real-world deployments measurements, in the examined use-cases.", "relation_to_broader_scientific_literature": "The key contributions of the manuscript focuses on improving the inference efficiency of self-speculative decoding in long-context LLM inference. The proposed approach is based on KV Cache quantisation during drafting, with minimal overheads and respecting the assumptions of the original inference scheme.", "essential_references_not_discussed": "The manuscript adequately cites relevant work in general. Particularly for the adopted quantization scheme, a comparative discussion with the approach of AnyPrecisionLLM [ICML'24] (which solely focuses on weights, however is also hierarchical and thus relevant to the proposed method) is required for completeness.", "other_strengths_and_weaknesses": "Strengths:\n- Overall the manuscript studies a very interesting and timely problem.\n- The provided discussion and analysis offers numerous insights for the efficient deployment of LLMs.\n- The proposed approach is simple and effective and suits well the examined scenario, while minimizing any deployment overheads.\n- The proposed double cache buffer and custom CUDA kernel allows the proposed methodology to yield speed-ups, realisable during deployment on commodity GPUs.\n\nComments:\n- It is unclear whether the findings of the analysis in Fig.2 and corresponding results in Fig.4 are representative of real world use-cases. Although the provided results study a wide range of realistic context lengths, the size of the model (Llama2-7B) is disproportional to the typical size of LLMs used for long context inference, which usually approach or surpass 1T params. Please see questions below too.\n- Some of the parameters of the analysis remain vague and need to be discussed in more detail. For example what is considered \"small\" or \"large\" batch in the discussion of Sec. 3.1.2.\n- It is unclear whether the proposed methodology can also be applied on the more general self-speculative decoding setting, where the draft model is a subset of the original (verification LLM), extracted e.g. through pruning (SWIFT [ICLR'25]) or early-exiting (LayerSkip [ACL'24]).", "other_comments_or_suggestions": "POST REBUTTAL EDIT: Provisionally increasing my score from 2 to 3, having read the thorough replies of the authors to all raised comments.", "questions_for_authors": "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. Are there any representative works of 7B models adopting >100k token context length?\n2. How do the finding of Sec.3 change with batch size of 1? What is defined as a \"small\" and \"large\" batch in the discussion? \n3. How is the proposed quantisation scheme positioned to other nested quantisation approaches, typically applied in weights, as in AnyPrecision LLM? What is the main benefit of the proposed scheme compared to these works?\n4. Can the proposed be applied in more traditional self-speculative decoding, where the draft model is a subset of the verification one e.g. SWIFT or LayerSkip ?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hmDYmAUtHF", "reviewer": "Reviewer_srdL", "summary": "The authors introduce a novel speculative decoding based technique for speeding up LLM inference. The basic idea is to use a hierarchical quantized KV cache and quantize model weights instead of storing a separate KV cache for both the target model and the draft model and storing separate model weights.The basic technical Insight is that you can decompose an int8 value into two int4 components which motivates a hierarchical KV cache structure and also allows for the shared architecture between the draft and target model.\n\n## post rebuttal\nThanks for your thorough response! I'll update my score; your response was very thorough.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> R2-1: Do you have intuition on why key and value caches have distinct quantization strategies?\n\n**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.\n\n[1] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, ICML 2024", "claims_and_evidence": "The claims are all supported by two sets of experiments. QuantSpec achieves speedups against streamingllm and snapKV, supporting the design choices of using quantize weights and cache rather than a sparse KV cache. The author is also provide insights into regimes in which quantizing weights versus the KV cache perform better. I also appreciated the analysis on using customized Cuda kernels.", "methods_and_evaluation_criteria": "I think the evaluations and baselines make sense. SnapKV and streaming LM weren't necessarily designed for the purposes of being used in a speculative decoding framework but I think that makes them reasonable baselines.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "I did check the soundness and validity of experimental designs. I don't have any issues to point out.", "relation_to_broader_scientific_literature": "The paper addresses a well-known technical problem (e.g,. speeding up inference in llms). I think the paper motivates why naive approaches to speculative decoding will fail in this regime and also provides interesting insights on how the draft model should be instantiated", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Do you have intuition on why key and value caches have distinct quantization strategies?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "czGtnu0anU", "reviewer": "Reviewer_5aMm", "summary": "In long context senarios, loading KV cache is a major bottleneck in both memory and latency. This paper introduces QuantSpec, which is a self-speculative decoding framework designed to accelerate long-context inference. Unlike existing speculative decoding methods that using a smaller model as the draft model, QuantSpec uses the same model with 4-bit quantized weights as the draft model, along with the hierarchical 4-bit quantized KV cache. This approach enables end-to-end speedups of up to 2.5×.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments. We address the questions below:\n\n> R3-1: Can the author elaborate more details on the CUDA kernel design?\n\n**A**: We will add a section in Appendix outlining our kernel design. We include a short summary here for the reviewer:\n\nIn 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.\n\nFor 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.\n\nFinally, 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.", "claims_and_evidence": "As listed near the end of introduction, the main claims made in the submission are (I skipped the roofline analysis here):\n\n(1) A new hierarchical quantization technique that enables bit-sharing between the target and draft models’ KV caches\n\n(2) A double full-precision cache buffer used for storing the most recent KV cache in full precision to improve acceptance rates and also reduce the number of quant. and dequant. operations\n\n(3) Provide custom CUDA kernels for attention with hierarchical quantized KV cache achieving up to 2.88× speedups at 4-bit precision relative to FP16 FlashAttention kernels\n\nAll three claims are verified empirically in the experiment section", "methods_and_evaluation_criteria": "Yes, the applied benchmark dataset makes sense for the problem and application.", "theoretical_claims": "There is no theoretical claims in this paper", "experimental_designs_or_analyses": "Yes, I checked the validity of the experimental  analyses. It makes sense to me", "relation_to_broader_scientific_literature": "Speed up long context generation is a very important problem. As more and more people rely on chatbot to process long documents.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "Can the author elaborate more details on the CUDA kernel design?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oIff0ddfKG", "reviewer": "Reviewer_ce5e", "summary": "This paper introduces QuantSpec, a self-speculative decoding framework designed specifically for long-context LLM inference. The framework employs a draft model that shares the architecture of the original model but implements hierarchical 4-bit quantized KV cache and 4-bit quantized weights for acceleration. QuantSpec achieves up to 2.5x end-to-end speedup while reducing memory usage by 1.3x compared to other self-speculative decoding methods for long-context scenarios.\n\n## update after rebuttal \nI raised my score to 3 since the authors addressed most of my questions.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> R4-1: Evaluate your approach on more recent models that better represent the current state of the field.\n\n**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.\n\n\n### Mistral-7B-v0.3\n\n|Context Length|Method|Acceptance Rate ↑ (optimal γ)|Peak GPU Memory (GB) ↓|Speedup (× AR) ↑|\n|:-------------:|:----:|:---------------------------:|:--------------------:|:--------------:|\n|16k|StreamingLLM|0.86 (1)|16.25|0.94|\n||SnapKV|0.86 (1)|16.70|0.93|\n||**QuantSpec**|0.94 (6)|17.42|**1.55**|\n|32k|StreamingLLM|0.89 (1)|18.67|1.07|\n||SnapKV|0.85 (1)|19.79|0.98|\n||**QuantSpec**|0.93 (6)|18.34|**1.61**|\n\n### Llama-3.1-8B\n\n|Context Length|Method|Acceptance Rate ↑ (optimal γ)|Peak GPU Memory (GB) ↓|Speedup (× AR) ↑|\n|:-------------:|:----:|:---------------------------:|:--------------------:|:--------------:|\n|16k|StreamingLLM|0.63 (1)|17.73|0.82|\n||SnapKV|0.66 (1)|18.18|0.85|\n||**QuantSpec**|**0.90 (6)**|18.79|**1.48**|\n|32k|StreamingLLM|0.81 (1)|20.15|0.93|\n||SnapKV|0.73 (1)|21.27|0.90|\n||**QuantSpec**|0.92 (6)|19.82|**1.54**|\n|128k|StreamingLLM|0.89 (1)|34.91|1.06|\n||SnapKV|0.80 (1)|39.95|1.06|\n||**QuantSpec**|0.91 (6)|26.05|**1.63**|\n\n\n---\n\n> R4-2: Your paper lacks a comparison with MagicDec, which employs a similar self-speculative approach.\n\n**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.\n\n---\n\n> R4-3: What advantages does your method offer over MagicDec specifically?\n\n**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.\n\n---\n\n> 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.\n\n**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.\n\n[1] MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding, ICLR 2025\n\n[2] https://www.nvidia.com/content/dam/en-zz/Solutions/gtcf21/jetson-orin/nvidia-jetson-agx-orin-technical-brief.pdf\n\n[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", "claims_and_evidence": "Yes, the evidence presented adequately supports the claims proposed in the paper.", "methods_and_evaluation_criteria": "The method is technically sound but lacks significant novelty in its approach.", "theoretical_claims": "No issues were found in the theoretical claims presented.", "experimental_designs_or_analyses": "The experimental analysis is generally well-executed. However, the experimental design has several notable issues:\n\n1. 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 A5000 GPU, which is not an edge device but rather a professional-grade GPU.\n2. The tested models are outdated. Llama-2-7b was released two years ago, and demonstrating that the method works on this model has limited relevance to current applications. Contemporary models have different characteristics that may affect the method's efficacy.\n3. The paper lacks comparison with MagicDec, which is a significant omission given its relevance to the self-speculative decoding approach.", "relation_to_broader_scientific_literature": "The idea of using the model itself for speculative decoding is not novel, as similar approaches have been implemented in prior works such as MagicDec and TriForce. Additionally, the application of 4-bit hierarchical KV cache represents an incremental improvement rather than a substantial innovation in the field.", "essential_references_not_discussed": "No essential references appear to be missing from the discussion.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-written with clear presentation that enhances readability and comprehension.\n2. The authors demonstrate significant technical effort through the implementation of high-performance custom CUDA kernels, which is commendable.", "other_comments_or_suggestions": "No.", "questions_for_authors": "I have two primary questions regarding the evaluation settings:\n\n1. Llama-2-7b is relatively outdated at this point. Why didn't you evaluate your approach on more recent models that better represent the current state of the field? Including contemporary models would significantly strengthen your claims about the method's broad applicability.\n2. Your paper lacks comparison with MagicDec, which employs a similar self-speculative approach. What advantages does your method offer over MagicDec specifically? This comparison seems essential for properly positioning your contribution in the literature.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Summary:

QuantSpec is a novel self-speculative decoding framework designed for efficient long-context LLM inference. Recognizing that memory bandwidth bottlenecks primarily arise from loading large KV caches, QuantSpec employs a hierarchical KV-cache quantization strategy, using 4-bit quantization for token drafting and higher precision (8-bit) for verification. Unlike conventional speculative decoding methods relying on smaller auxiliary models, QuantSpec leverages the same model architecture with 4-bit quantized weights for drafting, which enables sharing between draft and target models. A custom CUDA kernel implementation further optimizes real-world GPU deployments.

Strengths:

1. The proposed method is simple and effective. The authors develop custom CUDA kernels for the proposed method, which makes the experiments much more convincing.

2. All reviewers agree that the problem of study is important and timely.

3. The provided discussion and analysis offers numerous insights for the efficient deployment of LLMs.

Weaknesses:

1. The authors are encouraged to discuss the case of bs=1 in more detail, as bs=1 and bs>1 are two different scenarios for applications.

2. As the motivation of this work is on edge device, the authors are encouaged to run experiments on real edge devices instead of professional-grade A5000 GPUs.

3. Some missing experiments are suggested by Reviewer ce5e.

All reviewers vote for (weak) acceptance. AC agrees with reviewers' recommendation and suggest weak accept for this paper.
Decision: Accept (poster)