TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
Authors: Hanshi Sun
Zhuoming Chen
Xinyu Yang
Yuandong Tian
Beidi Chen
Venue: COLM (2024)
Keywords: Long-context model, Speculative decoding, LLM efficiency
TL;DR: Hierarchical Speculative Decoding for Lossless Acceleration of Long Sequence Generation

Abstract: With large language models (LLMs) widely deployed in long content generation recently, there has emerged an increasing demand for efficient long-sequence inference support. However, key-value (KV) cache, which is stored to avoid re-computation, has emerged as a critical bottleneck by growing linearly in size with the sequence length. Due to the auto-regressive nature of LLMs, the entire KV cache will be loaded for every generated token, resulting in low utilization of computational cores and high latency. While various compression methods for KV cache have been proposed to alleviate this issue, they suffer from degradation in generation quality. We introduce TriForce, a hierarchical speculative decoding system that is scalable for long sequence generation. This approach leverages the original model weights and dynamic sparse KV cache via retrieval as a draft model, which serves as an intermediate layer in the hierarchy and is further speculated by a smaller model to reduce its drafting latency. TriForce not only facilitates impressive speedups for Llama2-7B-128K, achieving up to 2.31$\times$ on an A100 GPU but also showcases scalability in handling even longer contexts. For the offloading setting on two RTX 4090 GPUs, TriForce achieves 0.108s/token—only half as slow as the auto-regressive baseline on an A100, which attains 7.78$\times$ on our optimized offloading system. Additionally, TriForce performs 4.86$\times$ than DeepSpeed-Zero-Inference on a single RTX 4090 GPU. TriForce's robustness is highlighted by its consistently outstanding performance across various temperatures. The code is available at https://github.com/Infini-AI-Lab/TriForce.

Published as a conference paper at COLM 2024
TRIFORCE: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
Hanshi Sun', Zhuoming Chen', Xinyu Yang', Yuandong Tian?, Beidi Chen 1.2
Carneeie Mellon Univeristy
Meta AI (FAIR)
(hanshis, zhuominc, xinyuya2, beidic)@andrew. cmu, edu yuandong@meta. com
Abstract
With large language models (LLMs) widely deployed in long content gen- eration recently, there has emerged an increasing demand for efficient long-sequence inference support. However, key-value (KV) cache, which is stored to avoid re-computation, has emerged by growing linearly in size with the sequence length. Due to the auto- regressive nature of LLMs, the entire KV cache will be loaded for every generated token, resulting in low utilization of computational cores and high latency. While various compression methods for KV cache have been proposed to alleviate this issue, they suffer from degradation in generation quality. We introduce TRIFORCE, a hierarchical speculative decoding sys- tem that is scalable for long sequence generation. This approach leverages the original model weights and dynamic sparse KV cache via retrieval as a draft model, which serves as an intermediate layer in the hierarchy and is further speculated by a smaller model to reduce its drafting latency. TRI- FORCE not only facilitates impressive speedups for Llama2-7B-1288 ing up to 2.31 x on an A100 GPU but also showcases scalability in handling even longer contexts. For the offloading setting on two RTX 4090 GPUs, TRIFORCE achieves 0.108s/ 11000/tokenn token-one half as slow as the auto-regressive baseline on A100, which attains 7.78x our optimized offloading system. Additionally, TRIFORCE performs 4.86x than DeepSpeed-Zent Inference on a single RTX 4090 GPU. TRIFORCE's robustness highlighted by its consistently outstanding performance across various temperatures. The code is available at https: h/gggtttuubb.u www..oownnnnnniiieeeeei.iiniieie att Al-aaarTT-----------r-oo
critical bottleneck
achiev-
1 Introduction
Large language models (LLMs)) with long-context capability, such as GPT-4 (Achiam et al., 2023), Gemini (Team et al., 2023), and LWM (Liu et al., 2024a) continue to emerge and gain proficient application in scenarios including chatbots, vision generation, and financial analysis (Touvron et al., 2023; Chowdhery et al., 2023; Zhao et al., 2023; Reddy et al., 2024). However, losslessly serving these LLMs efficiently is challenging. Because of the auto- regressive nature of LLMs, the entire key-value (KV) cache, which stores intermediate key-value states from previous contexts to avoid re-computation, together with model parameters will be loaded into GPU SRAM for every token generated, resulting in low utilization of computational cores. In addition to the large volume of model parameters, the memory footprint of KV cache, which grows linearly with sequence length (Pope et al., 2023), is emerging as a new bottleneck for long sequence generation.
Recent methodologies have proposed KV cache eviction strategies (Xiao et al., 2023b; Zhang et al., 2024b; Liu et al., 2024c; Jiang et al., 2023; Ge et al., 2023) to mitigate the substantial memory footprint of KV cache, which selectively discard KV pairs from the cache based designed eviction policy, allowing models to generate texts with a limited KV cache budget. However, considering that discarded KV pairs cannot be restored and the difficulty in

---

Published as a conference paper at COLM 2024
coppert ploce
Draft steas repestedly generiite 2 tolens [Fr, Ro. 2 Yi)
1.2
Uirra StramrnggLM Mf Cecher
150MG, B.Dma
Draft Model
duundlyy wihhuce duc sight antiinnns quality rode
Aouurey
0.2
Unmach 1280 Portid Cashe
Target Model
24
For the verily Duu TuuKKaa to. An
2808, 10rs
Ratrieual Cache
04
TriFoce
0.
StroningLLM
Cacte
Ulima--7 78-3286-11111 Tul Ceche
Target Model
78GB, 30me
Full KV Cache
10
20
32
40
50
Latency
Figure 1: Left: TRIFORCE employs retrieval-based drafting and hierarchical speculation to effectively address the dual bottlenecks. It integrates two models and three caches, comprising a draft model, a target model, a StreamingL.LM cache for the draft model, alongside retrieval cache and a full cache for the target model. The process initiates by repeatedly drafting for Y steps, assisting the target model with retrieved partial KV cache in generating over '2 tokens, which will be further verified by the target model using full KV cache. Right: Evaluating the Llama2-7B- B-8K on a needle retrieval task indicates that KV cache eviction-based methods, such as StreamingLLM, require trade-off between latency and accuracy. In contrast, our TRIFORCE maintains low latency without sacrificing accuracy.
precisely foreseeing which KV pairs will be crucial for future text generation, they struggle with potential information loss, including hallucination and contextual incoherency (Yang et al., 2024), particularly in long contexts. Such challenges prevent these approaches from boosting speed without sacrificing the performance of models, as illustrated in Figure 1.
Concurrently, speculative decoding, which leverages a lightweight draft model to sequen- tially predict the next few tokens and let the target model verify the predicted tokens in parallel, is introduced to accelerate LLM inference while provably precisely preserving model output (Leviathan et al., 2023; Chen et al., 2023a; Xia et al., 2024). Nonetheless, de- ploying it for long sequence generation faces several challenges. First, training draft models to match the context length of target LLMs requires massive computation and it remains questionable whether these small models can achieve the same accuracy with a context length around 1M (Beltagy et al., 2020; Peng et al., 2023; Yan et al., 2024). Second, we found that draft models with existing trining-free methods (e.g. KV cache eviction strategies) can result in poor speculating performance. A continuously increasing divergence (Leviathan et al., 2023) is witnessed as the sequence length increases, as shown in Figure 2a.
In pursuit of lossless acceleration, we utilize the lossless feature of speculative decoding as the foundation of our system. An ideal speculative decoding algorithm should (i) be training-free, (ii) maintain high acceptance rate (Leviathan et al., 2023) with long contexts, and (ii) have low-cost drafting. However, two technical challenges need to be addressed to achieve the goal. First, is not immediately apparent what we can use for low-latency drafting without training smaller draft model to match the long context length. Second, the key factors for attaining a high acceptance rate with long contexts remain unclear.
Fortunately, based on our preliminary exploration, three key observations pave the way for designing an applicable system for serving LLMs with long contexts.
Hierarchical Speculation for Dual Memory Bottlenecks: As illustrated in Figures 2b and 2c, we recognize two memory bottlenecks: model weights and KV cache, and the latter gradually becomes the dominant bottleneck as context length increases. This inspires us to apply hierarchical speculation to tackle the two bottlenecks sequentially by different draft models.
We identify considerable redundancy within KV cache, finding that a relatively small portion of it is sufficient to achieve a high acceptance rate by using partial KV cache as a draft cache for self-speculation.
Leveraging Attention Sparsity for Speculative Decoding:
Exploiting Contextual Locality for Drafting Efficiency: We discover that the information from long context tokens needed by adjacent tokens tends to be similar. This observation suggests
2

---

Published as a conference paper at COLM 2024
70
0.
Lamez size-1)
0.76
Lamad-112 (watch scmen))
0.8
0.50
0:95
32
0.000
(1144)
0.2
Sheee rightd pritne Liara 1an
10
0.0
Cor Lengtt
<<
Corgest
19K
30C
64K
128K
2X
LongetLonghh
8K
3DK
128K
(a)
(b)
(c)
Figure 2: (a) A continuously increasing natural divergence (Leviathan et al., 2023) between the draft model with StreamingL.LM or Sliding-window with Re-computation and Llama2- 7B-128K is witnessed as the sequence length increases, indicating a falling acceptance rate for speculative decoding with longer contexts. Additionally, temperature sensitivity signals lack of robustness. (b) Compared with model weights, KV cache gradually becomes another bottleneck with long contexts. (c) KV cache occupies most of the memory as the context length increases.
that specific segment of the cache can be effectively reused across multiple decoding steps, amortizing the overhead of constructing draft cache and enhancing drafting efficiency.
Building on these insights, we introduce a hierarchical speculation approach. For a long- context target model (e.g., Llama2-7B-128K (Peng et al., 2023)), we leverage the original model weights but only with small proportion (e.g., 3%) of KV cache as a draft to tackle the bottleneck of KV cache. Hierarchically, the draft model is further speculated by lightweight model (e.g,, Llama-68M) with StreamingLLM cache to address the bottleneck of model weights. We present TRIFORCE, depicted in Figure 1, a scalable and robust specula- tive decoding system that integrates sarrieval-based drafting and hierarchical speculation, optimized for both on-chip and offloading scenarios. Specifically,
In Section 4.1, by maintaining the full cache, gain the flexibility to select KV pairs, allowing us to devise a superior selection method, termed retrieval-baseee drafting. This strategy retrieves required context information for future needs, which is characterized as lossless, particularly in comparison to eviction-based methods such as StreamingLLM, and H2O. We further demonstrated its effectiveness and robustness on different datasets.
In Section 4.2, propose hierarchical system to address the dual memory bottlenecks. Using a lightweight model paired with a Streaming cache for initial speculations, we can reduce the drafting latency for the subsequent speculation stage, thereby accelerating end-to-end inference.
Empirically, in Section 5, we perform extensive experiments and ablation studies to demon- strate the effectiveness of TRIFORCE. We show that TRIFORCE achieves up to 2.31x speedup for Llama2-7B-128K on a single A100 GPU on top Hugging Face (Wolf et al., 2019) with CUDA graphs (NVIDIA & Fitzek, 2020). For the offloading setting, TRIFORCE attains an impressive 7.78 on two RTX 4090 GPUs, reaching 0.108 s s/tookneoonnt only half as slow as the auto-regressive baseline on A100. TRIFORCE can efficiently serve Llama2-13B with 128K contexts with 0.226s token, which is 7.94x faster on our optimized offloading sys- tem. On single RTX 4090 GPU, TRIFORCE is 4.86x faster than DeepSpeed-Zero- Zero-Inferenne
(Aminabadi et al., 2022) Further, we show that: (i) TRIFORCE has theoretical 13.1 x upper bound, demonstrating exceptional scalability when dealing with long contexts; (ii) TRIFORCE robust across various temperature settings, maintaining an acceptance rate above 0.9 even for temperature of 1.0; and (iii) TRIFORCE's abilitt to efficiently process large batches, achieving 1.9x speedup for batch size of six with 19K contexts per sample.
The official implementation of SeepSpee--ZRO-Infeeen offloading currently only supports single GPU, which computes attention on CPU. Our offloading system transfers KV cache from CPU to GPU, benefiting from Tensor Parallelism.
(Aminabadi et al., 2022) with KV cache
3

---

Published as a conference paper at COLM 2024
2 Background
2.1 Speculative Decoding
Speculative decoding (Stern et al., 2018; Leviathan et al., 2023; Chen et al., 2023a; Kim et al., 2024; Zhang et al., 2023; Santilli et al., 2023; Hooper et al., 2023) is featured by accelerating LLM decoding while precisely maintaining the model's output distribution. As the speed of the auto-regressive decoding process is mainly bound by the time for loading model weights and KV cache to GPU SRAM, speculative decoding leverages the observation that generating one token takes the same time as processing tens of tokens in parallel. Tree-based speculation methods are proposed to fully utilize the speculation budget (Fu et al., 2024; Li et al., 2024). Instead of making one prediction for the next token, tree-based methods leverage multiple candidates to boost the acceptance rate so that more tokens can get accepted (Miao et al., 2023; Sun et al., 2024; Chen et al., 2024). Staged speculation techniques (Spector & Re, 2023; Chen et al., 2023b) have been suggested to further accelerate inference by using a cascade of draft models, and hierarchical speculation shares similarities with these approaches. However, our method focuses on long sequence generation tasks, which presents unique challenges. We use different drafting methods for each speculation phase to address two distinct bottlenecks in long-context scenarios, instead of focusing on model weights at every stage for acceleration. Meanwhile, self-speculation approaches such as Medusa (Cai et al., 2024; Ankner et al., 2024), which are orthogonal to our method, require training efforts and can be integrated into our intermediate draft model.
2.2 KV Cache Eviction Strategies
StreamingLLM (Xiao et al., 2023b) addresses the limitations of window attention and slicing window with re-computation by presenting a straightforward yet effective method that allows LLMs to handle infinitely long text sequences without fine-tuning. StreamingL.LM stabilizes the performance by retaining critical attention sink tokens together with recent KV for attention computation. By prioritizing sink tokens, StreamingLLM ensures the attention score distribution remains stable, promoting consistent language modeling for long texts.
H2O (Zhang et al., 2024b) introduces a greedy but low-cost approach to processing infinite- length input streams, inspired by a simplified version of the heavy-hitters (H2) eviction policy. This method dynamically updates the KV cache based on the cumulative attention scores, systematically removing the least critical KV to maintain a fixed cache size. By leveraging greedy algorithm based on local statistics, H2O effectively selects which KV pairs to preserve in the cache, ensuring efficient inference without compromising quality.
However, it is important to recognize that these techniques do not increase the context window size (Zhang et al., 2024a; Jin et al., 2024; Ge et al., 2023; Jiang et al., 2023). They focus retaining only the most recent tokens along with either attention sinks or heavy-hitters, while discarding other tokens. These approaches limit the model to processing based on their designed eviction policies and recent tokens. Consequently, they might not be directly applicable to tasks that demand comprehensive, loriiiieeennnooooooooo long-context understanding.
2.3
KV Cache Quantization
Several approaches to KV cache quantization have been introduced to enhance the efficiency of inference for long sequence generation, aiming to maintain generation quality while reducing the memory consumption (Xiao et al., 2023a; Hooper et al., 2024; Sheng et al., 2023; Liu et al., 2023a; Zirui Liu et al., 2023; Yue et al., 2024). Quantization methods focus compressing the bit width of KV cache activations, which is orthogonal to our approach.
Observation
Our design of TRIFORCE is inspired by two critical empirical observations regarding LLMs when dealing with long contexts, detailed as follows.
4

---

Published as a conference paper at COLM 2024
Recovery Rate Tpp4200a59600%
1.0
1.0
0.8
0.6
Layer
24 16
Recouery
800
(1-0)
Layer
0.75
Snasniggal (1%)
0.2
(L074))
Layer
15 20 Head
24
25
122
192
256
KV Cache Budget
Tokee
(a)
(b)
(c)
Figure (a) The Llama2-7B-1284 model demonstrates significant attention sparsity with 120K context. Apart from the first two layers, the rest exhibit significant sparsity. (b) We can utilize partial KV cache and whole model weights to perform self-speculation. High acceptance rates are attainable using existing methods with a limited budget. (c) A notable degree of locality is observed in most layers, which gradually diminishes as context evolves.
3.1 Leveraging Attention Sparsity for Speculative Decoding
Observation The phenomenon of attention sparsity in pre-trained LLMs has been discov- ered by numerous studies (Zhang et al., 2024b; Xiao et al., 2023b; Liu et al., 2023b; 2024c). In our study, we conduct zero-shot inference on the PG-19 test set (Rae et al., 2019) with Llama2-7B-i288 model. By visualizing the sparsity across different attention heads, demon- strated in Figure 3a, we observe that with a context length of 120K, it is possible to recover over 96% of the attention score with merely 4K tokens across almost all layers.
The presence of sparsity within the attention blocks suggests that a fraction of KV cache could serve as draft cache attain a high acceptance rate during self-speculative decoding. Since KV cache is the bottleneck under this setting, we can load whole model weights with partial KV cache as draft model. Figure 3b demonstrates that utilizing only 1K tokens could theoretically achieve a 97.6% acceptance rate with Top-K selection method. While this scenario represents an optimal theoretical upper bound, practical implementations like H₂O and StreamingLLM exhibit promising results, achieving over 90.5% acceptance rates with 1K KV cache budget. should be noted that we maintain a full cache for the initial two layers for illustration purposes, while no layers skipped in our practical system implementation for efficiency.
Analysis
3.2 Exploiting Contextual Locality for Drafting Efficiency
Observation Our exploration reveals that the information from long context tokens needed by adjacent tokens tends to be similar. In our experiments, with the context length estab- lished at 120K, we instruct the model to generate 256 tokens. By choosing the top-4K indices according to the attention score of the last prefilled token, we use these indices to gather the attention scores for the subsequently generated tokens and assess the score's 's recovery rate for the initially prefilled 120K tokens. As shown in Figure 3c, it leads to high recovery across almost all layers and a slowly decreasing trend as the number of tokens increases.
Insights This observation allows for a single construction of the cache to suffice for multiple decoding steps, thereby amortizing the latency of constructing draft cache and boosting efficiency. As new KV cache are introduced, guided by the understanding that recent words are more strongly correlated with the tokens currently being decoded, these entries will replace the less significant ones. Cache re-building operations can be scheduled at regular intervals or adaptively in response to a drop in the acceptance rate, which ensures that the cache remains dynamically aligned with the evolving context. Notably, both StreamingLLM and H₂O incorporate this principle implicitly. H₂O consistently retains tokens with high scores, and StreamingLLM reuses extensive local information and sink tokens, which both reduce the necessity for complete cache reconstruction.
5

---

Published as a conference paper at COLM 2024
4 TRIFORCE
This section aims to introduce the TRIFORCE, which leverages a retrieval-based KV cache selection policy and hierarchical speculation system. We first argue that our retrieval- based drafting approach is intuitive and lossless compared to existing strategies such as StreamingLLM and H2O. Subsequently, effectively address the dual bottlenecks in speculative decoding, facilitating a substantial improvement in overall speed-up. Finally, TRIFORCE is elaborated in Section 4.3.
introduce the hierarchical system designed to
4.1 Retrieval-based Drafting
Value
Key
In scenarios requiring long-term contextual dependencies, methods like Streaming and H2O underperform due to their cache up- dating strategies, which are ineffective at ac- curately retrieving detailed contextual infor-
0.3
Q.3
0.5
0.1
mation because they inevitably and irrecov- erably discard KV pairs. In our experiment, we challenge Streaming..LM and H2O with needle retrieval task (Liu et al., 2024b; Peng et al., 2023; Liu et al., 2024a). As detailed in Table 1, there is a notable drop in their acceptance rates compared to their performance on the PG-19 dataset, highlighting their limitations. Essentially, StreamingLLM, and HzO operate on a lossy principle, as evicted tokens are permanently discarded, making them poor fit for settings requiring the preservation of full KV cache for the target model.
Chunk Size
Query
The
new
system
for
long
contest
spec
slative
Figure 4: Retrieval-based drafting
The necessity of keeping the entire KV cache in our settings allows us to select KV cache more freely (Singhal et al., 2001). This insight leads us to develop more effective selection policy for lossless approximations. In our approach, demonstrated in Figure 4, KV cache is segmented into small chunks. During the retrieval phase, we calculate the attention between a given query and the average key cache within each chunk. This method effectively highlights the most relevant chunks, enabling us to gather KV cache with fixed budget based on the scores. As illustrated in Table 1, retrieval-basel method excels by actively identifying the most crucial information for the task rather than relying on passive and time- based cache management methods. By focusing on relevance over recency, retrieval-basee policy demonstrates its potential to handle contextually dense datasets.
Table 1: Acceptance rates are shown across various tasks, utilizing 120K context and a 4K budget, while bypassing the initial two layers. There is a notable drop in StreamingLLM and H2O's results on needle retrieval. For reference, Top-K is the theoretical upper bound.
Method
Top-K (Ref.)
Streaming
H2O
Retrieval
PG-19
0.9921
0.9156 0.0519
0.9179
0.9649
Needle Retrieval
0.9989
0.0739
0.9878
4.2 Hierarchical Speculation
While addressing the KV cache bottleneck enhances efficiency, the requirement load whole model weights for drafting reintroduces latency, shifting the bottleneck to model weights again. To tackle this challenge, we implement a hierarchical system, as illustrated in Figure 1. This system employs a secondary, lightweight model with StreamingLLM cache to perform initial speculations for the target model with retrieval-based cache (which serves as a draft model for the target model with full KV cache). By establishing this sequential speculation hierarchy, we effectively reduce drafting latency and accelerate overall inference.
Correctness: The original output distribution is preserved during the final speculation phase, which is identical to the standard speculative decoding algorithm (Leviathan et al., 2023; Chen et al., 2023a), and the proof is trivial.

---

Published as a conference paper at COLM 2024
4.3 Algorithm
TRIFORCE is devised to exploit the bottlenecks associated with both model weights and KV cache to enhance the inference speed of LLMs for long sequence generation. We present the pseudocode for the TRIFORCE in Algorithm 1. starts by prefilling the target model Mp with full cache Cp and draft model M, with StreamingL cache C using given input prefix, and then constructs the retrieval cache Cy. The initialization and update mechanism for the retrieval cache C, guided by the insights of contextual locality discussed in Section 3.2. We first construct C, using the last token of the prefix, arranging tokens by the descending order of importance. In subsequent inferences, we overwrite tokens with the least importance, maintaining the relevance and utility of the cache. A reconstruction of Cy is triggered either when the rolling average acceptance rate drops below a threshold or at designed stride.
The inference progresses iteratively until it reaches the target sequence length T. After each iteration, cache C, and Cg are updated to prepare for the subsequent speculation phase. Each iteration encompasses two speculations: initially, M3 utilizes Cg to predict Mp with C, for Y1 steps until n 2 7/2- Subsequently, these ıI tokens are self-verified (Zhang et al., 2023) by Mp with Cp. This process constructs hierarchy: the first layer of hierarchy employs smaller, faster model My with local context Cy to speculate the large model Mp with partial but high-quality global context C, addressing the model weights bottleneck. The second layer utilizes model M, with retrieval cache for self-speculation, overcoming the bottleneck caused by KV cache. This hierarchical speculation algorithm boosts efficiency by effectively addressing both bottlenecks. System implementation is detailed in Appendix A.
Algorithm 1 A TRIFORCE
1: Input: Prefix [x,,-.. xi],, target model Mp with full cache Cp, draft model Mg with Y2, drafting StreamingLLM cache Cy, target sequence length T, speculation length 72, phase DRAFT, verification phase VERIFY, and correction phase CORRECT;
2: Initialize: Prefill Mp, Mg. construct retrieval cache C, using X,, N <
3: while N < T do
4:
+
5: while n < 7/ do
Set 91, 9n +- DRAFT(M Cg. xsN) Sample 8; ~ qú = L, 71
Run Mg with eviction cache C.
7:
8:
Set PI ... pn, < M,(G,,ENN&sM &5m)
Run Mp with retrieval cache C,
9:
for i = 1 to Y1 do
10:
VERIFY(E;, qi, Pi) then inti + & and t- n+1
11:
12:
else
13:
Luti +- CORRECT(Qi Pi) and Break
14
end if end for
16:
all drafted tokens are accepted, sample next token bat1 ~ p and n +- n1 + end while
18:
19:
Collect Pr.... Pn for 8,, 8n XEN, R5n)
20:
Set Pı. .. , Pn=1 +
D Run Mp with full cache Cp
21:
for = 1 to 17 do
22:
if VERIPY(&;; Pir p:) then XN & & and N t N+1
23:
24:
else
25:
XN+i 4- CORRECT(Pi, pi) and N + N+1 Break
26:
27:
end if
end for
29: If all drafted tokens are accepted, sample next token XN+1 ~ Pa+1 and N + N+1
Update C,, Cg based on the accepted tokens Update KV cache for the next iteration end while
31:
7

---

Published as a conference paper at COLM 2024
Table 2: On-chip results (A100): We indicate the average acceptance rate in parentheses alongside the speedup factor. T means sampling temperature. In the A100 on-chip experi- ments, with prompt length of 122K, and generation length of 256, we evaluate TRIFORCE against the JF68M model with StreamingLLM, cache (Naive Policy). The results clearly demonstrate that TRIFORCE significantly surpasses its performance.
Method
T
Speedup
Naive Policy
TRIFORCE
0.0
2.31x (0.9234)
1.56x (0.4649)
TRIFORCE
0.2
2.25x (0.9203)
1.54x (0.4452)
TRIFORCE
0.4
2.20x (0,9142)
1.47x (0.4256)
TRIFORCE
0.6
2.19x (0.9137)
1.42x (0.4036)
TRIFORCE
0.8
2.08x (0.8986)
1.34x (0.3131)
TRIFORCE
1.0
2.08x (0.9004)
1.29x (0.2872)
TRIFORCE
1.2
2.02x (0.8902)
1.27x (0.2664)
Retrieval w Hierarchy
0.6
1.80 (0.9126)
StreamingLLM w/ Hierarchy
0.6
1.90 (0.8745)
Table 3: Offloading results (RTX 4090): We present latency comparison between TRIFORCE and Auto-regressive (AR) baseline for various models on different GPU setups. The sam- pling temperature is set to 0.6. The results indicate that TRIFORCE achieves significant speedups across a range of models and hardware configurations. The entries marked with an asterisk represent the baseline using DeepSpeed-ZeRO-Infeee
(Aminabadi al., 2022).
GPUs
Target Model
TRIFORCE (ms)
AR (ms)
Speedup
2x RTX 4090s
Llama2-7B-122
108
840
7.78x
2x
RTX 4090s
WWWM-Txx--Cht--11
114
840
7.37 x
2x RTX 4090s
Llama2-13B-128K
226
1794
7.94x
1x
RTX 4090
Llama2-7B-128K
312
1516*
4.86x
1x RTX 4090
LWM-Text-Chat-128K
314
1516*
4.83x
5 Empirical Evaluation
In this section, our goal is to showcase the capabilities of TRIFORCE, scalable and robust speculative decoding algorithm designed to expedite the inference of LLMs for long se- quence generation, which significantly reduces the wall-clock time. We first present our end-to-end system, highlighting the overall speedup achieved, including both on-chip and offloading settings, followed by a comparison with other methods and ablation experiments.
5.1 End-to-end Results
We demonstrate that TRIFORCE accelerates long sequence generation, up to 2.31 x on an A100 in the on-chip setting and x.78x two RTX 4090s with offloading for Llama2-7B-128K.
Setup. Our experiments are based Llama2 and LWM models with 128K context window size (Touvron et al., 2023; Liu et al., 2024a; Peng et al., 2023), which serve as our target models. In this setup, we utilize a 4K retrieval cache as an intermediate draft cache in our hierarchical system, while leveraging the JackFram/ /Llama6MM (JF68M) (Miao et al., 2023) model as the initial draft model. For experiments involving offloading, we aim to maximize memory utilization by filling it up much possible and offloading the remaining KV cache to the CPU (AMD EPYC 9754 2.25 GHz), while keeping the model weights on the GPU, Our evaluation is carried out on the PG-19 (Rae et al., 2019) and NarrativeQA (Kočiský et al., 2018) dataset, each testing on 100 examples, configured to a prompt length of 122K for on-chip settings and 127K for offloading settings, and aiming for generation of 256 tokens. The performance of TRIFORCE is analyzed across various hardware configurations, including on-chip experiments an A A100, and offloading experiments on RTX 4090 GPUs.
8

---

Published as a conference paper at COLM 2024
Naive Policy. Since it is hard to train draft model with long contexts, consider JF68M with LLM cache as a naive policy approach, and its budget is set to 1K. Additionally, we experiment with various temperatures to test its robustness.
Main Results. We evaluate TRIFORCE using different temperatures, as depicted in Table 2. We observe that TRI- FORCE reaches up to 2.31 x speedup for the on-chip setting with minimal 4K KV cache budget for Llama2-7B-1228 For offloading settings, we provide end-to-end results on consumer GPUs for more models, including Llama2-7B- 128K, Llama2-13B-122 128K, and WWM-Text-Chat-122k Re- markably, in Table 3 we demonstrate that TRIFORCE can efficiently serve Llama2-13B with 128K contexts two
24
22
Upper
2,0
Speesuu
1.8
1.8
1.4
TiiForee
Natvo Policy
RTX 4090s, reaching an average time between tokens as low as 0.226 seconds, which is 7.94x faster than a highly optimized offloading system. Moreover, with TRIFORCE, Llama2-7B-128K can be served with 0.108s/ ttoknnoollnn/ half as slow as the auto-regressive baseline on an A100. We also illustrate how TRIFORCE boosts the efficiency of batched inference, a more frequently employed setting in real-world model serving. TRIFORCE achieves 1.9x for a batch size of six, with each sample in the batch having 19K contexts, which is demonstrated in Table 4.
32K
64K
Langth
Figure 5: TRIFORCE's excellent scalability with longer contexts
Table 4: Batching results (A100): TRIFORCE showcases its exceptional capability in effi- ciently handling large batch sizes, consistently exceeding the performance of the JF68M model with Streamingl. cache across all configurations for Llama2-7B-1285
Batch
Budget
Speedup
Naive Policy
(2,56K)
(2,1024)
0.0
1.89x
1.46x
(2,56K)
(2,1024)
0.6
1.75x
1.35x
(6,19K)
(6,768)
0.0
1.90x
1.39x
(6,19K)
(6,768)
0.6
1.76x
1.28x
(10,12K)
(10,768)
0.0
1.72x
1.34x
(10,12K)
(10,768)
0.6
1.61x
1.21x
Analysis. (1) Effectiveness: TRIFORCE's integration of the hierarchical system signifi- cantly enhances speedup, with TRIFORCE showing marked improvements over both the StreamingLLM method with hierarchical speculation and retrieval method without the hier- archical system. (2) Scalability: As depicted in Figure 5, TRIFORCE demonstrates excellent scalability with longer context lengths. This scalability is attributed to its high acceptance rate and the growing gap between the draft and the target model's 's latencies. Theoretically, TRIFORCE could achieve a speedup of up to 13.1 times higher than the naive policy, underscoring its significant scaling potential. (3) Robustness: Unlike vanilla speculative decoding methods, TRIFORCE maintains relatively consistent performance across various temperature settings. It exhibits less temperature sensitivity, maintaining an acceptance rate above 0.9 even when the temperature is set to 1.0, highlighting its stability and reliability.
5.2 Comparison with Other Methods
Method
Speedup
We provide comparison with REST (He et al., 2023) and Skipping Layers (Zhang et al., 2023). Table compares TRIFORCE, REST, and Skipping Layers with Llama2-7B- 128K on an A100 using PG-19 dataset, showing TRIFORCE achieves the best speedup for long sequence generation.
TRIFORCE REST
2.31 x
1.47x
Skipping Layers
1.36x
Table 5: Speedup comparison with REST and Skipping Layers external predefined datastore. In Table 5, Skipping Layers utilizes 68% of the KV cache, whereas TRIFORCE efficiently uses only 3%, addressing the bottleneck in long-context scenarios better.
TRIFORCE retrieves information from the KV cache, al- lowing dynamic adaptation to contexts, while REST uses

---

Published as a conference paper at COLM 2024
Takura
Acceptrne Rute
Triforss Sequula
1.5
TeEl oree Single Sing Serae
ThiForee Indeperdent Siquanoes
2016 1072
1934
25E Budget
Churc Bizo
(a)
(b)
(c)
Figure 6: (a) Analyzing speedup and acceptance rates across different KV cache budgets reveals that 4K budget is optimal, balancing acceptance rates and the drafting overhead. (b) For 4K KV cache budget, excessively small chunk sizes may overfit individual tokens, while overly large chunk sizes could limit selection diversity. (c) TRIFORCE is compatible with tree-based speculations, enhancing the theoretical average number of tokens generated per decoding step of the target model by employing larger speculation budgets.
5.3 Ablation Results
We present extensive ablation studies of TRIFORCE, focusing on three key points: (1) the influence of different KV cache budgets, (2) the impact of chunk size selection, and (3) TRIFORCE's compatibility with tree-based speculative decoding.
5.3.1 KV Cache Budget
As illustrated in Figure 6a, for Llama2-7B-128K, the acceptance rate rises with the cache budget up to 4K, then plateaus towards 1.0. This suggests that increasing the cache size beyond 4K offers diminishing benefits due to drafting latency. Thus, a 4K KV cache budget is optimal for TRIFORCE, balancing high acceptance rates and minimal drafting overhead.
5.3.2 KV Cache Chunk Size
Since utilize contextual locality to reuse the retrieval cache, need to examine the impact KV cache chunk size on performance. Figure 6b shows that smaller chunks may overfit to single tokens, limiting generalization, while larger chunks may dilute high-score tokens with low-score ones, resulting in reduced differentiation among chunks. Large chunks also reduce selection flexibility, constraining diversity within fixed cache budget.
5.3.3 Compatibility with Tree-based Speculative Decoding
We explore the possibility of integrating TRIFORCE with tree-based speculative decoding. Specifically, for Llama2-7B-128K on an A100, we estimate the theoretical number of gen- erated tokens when TRIFORCE is combined with tree structures, including Sequoia (Chen et al., 2024) and Independent Sequences. As depicted in Figure 6c, this integration can potentially improve the end-to-end speedup by utilizing additional speculation budgets.
Conclusion
In this work, we introduced TRIFORCE, a hierarchical speculative decoding system aimed at significantly enhancing the efficiency of serving LLMs with long contexts. Leveraging insights from attention sparsity and contextual locality, TRIFORCE mitigates the dual bottle- necks associated with KV cache and model weights. Our empirical experiments demonstrate TRIFORCE's remarkable performance, including a notable speedup of up to 2.31x on an A100 and 7.78x on two RTX 4090s with offloading, achieving 0.108s/ Ti/8s/token //t//toknnoooo oonnnooolllllee half as slow as the auto-regressive baseline on an A100. These achievements illustrate TRIFORCE's potential to revolutionize the serving of long-context models for long sequence generation.
10

---

Published as a conference paper at COLM 2024
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. Xiv preprint arXiv:2303.08774 2023.
Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, et al. Deepspeed- enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, Pp. 1-15. IEEE, 2022.
inference:
Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Regln-Kelley, and William Brandon. Hydra: Sequentially-dependde draft heads for medusa decoding. areio preprint arXiv:2402.05109, 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. Xrv preprint or Xic:2004..05150, 2020.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple Iim inference acceleration framework with multiple decoding heads. arXiv preprint ar Xitx:401.107744 2024.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arv preprint arXiv:2302.01338. 2023a.
Zhuoming Chen, Avner May, Ruslan Svirschevski, Yuhsun Huang, Max Ryabinin, Zhihao Jia, and Beidi Chen. Sequoia: Scalable, robust, and hardware-aware speculative decoding. ar Xiv preprint ar axXiv:2402.12374, 2024.
Ziyi Chen, Xiaocong Yang, Jiacheng Lin, Chenkai Sun, Jie Huang, and Kevin Chen-Chuan Chang. Cascade speculative drafting for even faster Ilm inference. arXiv preprint arXiv:2312.11462, 2023b.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. Xiv preprint ar axXiv:2307.08991 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359 2022.
Yichao Fu, Peter Bailis, lon Stoica, and Hao Zhang. Break the sequential dependency of llm inference using lookahead decoding. arXiv preprint ar Xix::2402.02057, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for 1lms. arXiv preprint arXivv2310.01801, 2023.
Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He. Rest: Retrieval-based speculative decoding. arXiv preprint arXiv:2311.08222, 2023.
Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, and Yu Wang. Flashdecoding++: Faster large language model inference on gpus. arXiv preprint ar arXiv:2311.01282 2023.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Hasan Genc, Kurt Keutzer, Amir Gholami, and Sophia Shao. Speed: Speculative pipelined execution for efficient decoding. ar Xiv preprint ar Xixv:22310. 12072, 2023.
11

---

Published as a conference paper at COLM 2024
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arv preprint ar Xio:24401.18877 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulmier, et al. Mistral 7b. ariio preprint ar arXi::2310.06825, 2023.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning. Xiv preprint or Xio:2401.00322. 2024.
Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. Advances in Neural Information Processing Systems, 36, 2024.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The NarrativeQA reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018. doi: 10.1162/
a00023,, URL https: / /ttp://////aaathhoooooo org/Q18-1022
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pp. 611-626, 2023.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty. arviv preprint ar arXtt:2401.15077,, 2024.
Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. arXiv preprint A4k::2202.08268,, 2024a.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024b.
Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. LIm-qat: Data-free quantization aware training for large language models. ar Xiv preprint ar Xrii:22305.1788, 2023a.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient ilms at inference time. In International Conference on Machine Learning, pp. 22137-22176. PMLR, 2023b.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anas- tasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for 1lm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024c.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative IIm serving with speculative inference and token tree verification. arviv preprint ar Xi 2023.
Péter Vingelmann NVIDIA and Frank HP Fitzek. Cuda, release: 10.2. 89.
URL
htlps://devvllopr: nvidia. con/cuda-toolle toolkit. Cited, pp. 148, 2020.
12

---

Published as a conference paper at COLM 2024
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An im- perative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint ar Sttt:23990000777 2023.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling trans- former inference. Proceedings of Machine Learning and Systems, 5, 2023.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arv preprint, 2019. URL https://arxiv.org wttp://wwwwg/////......ooo
Varshini Reddy, Rik Koncel-Kedziorski Viet Dac Lai, and Chris Tanner. Docfinqa: A
long-context financial reasoning dataset. arXiv preprint arXiv:2401.06915 2024.
Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorca, Michele Mancusi, Riccardo Marin, and Emanuele Rodolà. Accelerating transformer inference for translation via parallel decoding. Xiv preprint arXiv:2305.10f 2023.
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ton Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with single gpu. In International Conference on Machine Learning, pp. 31094-31116. PMLR, 2023.
Amit Singhal et al. Modern information retrieval: A brief overview. IEEE Data Eng. Bull., 24 (4):35-43, 2001.
Benjamin Spector and Chris Re. Accelerating 11m inference with staged speculative decoding. Xrv preprint arXio:2308.0462, 2023.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31, 2018.
Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. Advances in Neural Information Processing Systems, 36, 2024.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Lean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. ar Xiv preprint ar arXi::2312.11800 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. ar preprint arXiv:2307.0922 2023.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. aix preprint arXiv:1910.033711, 2019.
Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A compre- hensive survey speculative decoding. arXiv preprini ar axX1:22401.07881, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38098 PMLR, 2023a.
13

---

Published as a conference paper at COLM 2024
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. The Twelfth International Conference on Learning Representations, 2023b.
Minghao Yan, Saurabh Agarwal, and Shivaram Venkataraman. Decoding speculative decoding. ar Xiv preprint ar arXiv:2402.01528, 2024.
June Yong Yang, Byeongwook Kim, Jeongin Bae, Beomseok Kwon, Gunho Park, Eunho Yang, Se Jung Kwon, and Dongsoo Lee. No token left behind: Reliable kv cache compression via importance-aware mixed precision quantization. arxiv preprint 2xXx:2402.18099) 2024.
Yuxuan Yue, Zhihang Yuan, Haojie Duanmu, Sifan Zhou, Jianlong Wu, and Liqiang Nie. Wkvquant: Quantizing weight and key/ /value cache for large language models gains more. ar Xiviv preprint ar arXiv:2402.12066 2024.
Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. Draft & verify: Lossless large language model acceleration via self-speculative decoding. arviv preprint ar Xiv::3300 2023.
Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. Soar- ing from 4k to 400k: Extending Ilm's context with activation beacon. arv preprint arXio:2401.03466 2024a.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024b.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang. Zican Dong, et al. A survey of large language models. arxiv preprint ar Xiv::3303.18223, 2023.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: Plug-and-play 2bit kv cache quantization with streaming asymmetric quantization. 2023. doi: 11.314///R.2.22222 372222.. URL https://rgod: net/10. 13140/16. 2. 2.2.28167..337 37282.
Acknowledgments
We thank Yang Zhou, Ranajoy Sadhukhan, Harry Dong, and Jian Chen for their helpful discussions and feedback on early drafts of the paper.
14

---

Published as a conference paper at COLM 2024
A System Implementation
We implement the draft and target models using Transformers (Wolf et al., 2019). Leveraging a predetermined cache budget enables the effective use of PyTorch CUDA graphs (Paszke et al., 2019; NVIDIA & Fitzek, 2020), significantly minimizing the kernel launching overhead during speculative decoding. FlashAttention is used to accelerate attention operations (Dao et al., 2022; Dao, 2023; Kwon et al., 2023; Hong et al., 2023). Notably, we maintain full layer sparsity without omitting the initial two layers for system efficiency, ensuring our approach stays within the fixed KV cache budget. Although this strategy might lead to a lower acceptance rate, the benefit of maintaining a constant drafting cost makes a worthwhile compromise. Additionally, to facilitate faster speculation phase, we also implement two extra speculation cache structures, including our retrieval cache and StreamingLLM, cache.
B Additional Experiments
B.1 TRIFORCE's Scalability for Longer Inputs
For longer inputs, solution is offloading the KV cache to the CPU memory. Below are experiments on single L40 GPU with offloading for longer inputs on LWM-Text models (Liu et al., 2024a) using the PG-19 dataset, showing impressive speedup.
Model
Speedup
LWW--Text2256 11.81
LWM-Text-512K
12.10x
Table 6: TRIFORCE's Scalabiiity for longer inputs
B.2 TRIFORCE'S Scalability for Longer Outputs
In the previous sections of the paper, we focused on a long input and short output to illustrate the efficiency gains in commonly encountered setting. Here we provide additional experiments with longer output sequences up to 2K with Llama2-7B-128K on an A100 GPU using the PG-19 dataset:
Output Sequence Length
Speedup
256
2.31
512
2.28x
768
2.34x
1024
2.32x
1536
2.28x
2048
2.29x
Table 7: TRIFORCE's Scalability for longer outputs
As seen in Table 7, the performance of TRIFORCE remains relatively stable. This stability is due to the retrieval cache being reconstructed fixed intervals, ensuring stays aligned with the evolving context. Furthermore, our retrieved cache can be effectively reused across multiple decoding steps, amortizing the overhead of retrieval. These results demonstrate the scalability of TRIFORCE and its robustness in handling varying output lengths, thereby addressing the concern about its generalizability.
B.3 Ablation of 72
Here we provide additional ablation results of 7ı, 72 in Table with Llama2-7B-128K on an A A100 GPU using the PG-19 dataset. Due the gap between JF68M with StreamingLIMM
15

---

Published as a conference paper at COLM 2024
and Llama-7B-128K with retrieval cache, the acceptance rate of the first speculation phase in our speculation hierarchy is low. As shown in the table, the optimal value for "11 is and for 12 is 6.
Y
Speedup
72
Speedup
1
2.30
1.62x
2
2.31x
2
1.90x
3
2.24x
2.13x
4
2.17x
4
2.22x
5
2.09
5
2.26X
6
2.01x
2.31X
1.92x
7
2.24x
8
1.83X
8
2.20x
(a) Speedup different "1 values with 72 =
(b) Speedup at different 72 values with 7/2
Table 8: Ablation of y1, 72
16


Reviews:
[{"review_id": "PdigecwC5t", "reviewer": "Reviewer_ByTv", "summary": "The paper introduces TRIFORCE, a novel speculative decoding system aimed at accelerating long-sequence generation in large language models (LLMs). The proposed method integrates Retrieval-based Drafting and Hierarchical Speculation to tackle the dual bottlenecks of model weights and key-value (KV) cache. TRIFORCE leverages the original model weights and partial KV cache to create a draft model, which is further speculated by a lightweight model using StreamingLLM cache. The approach is shown to provide significant speedups, particularly for the Llama2-7B-128K model, without compromising the quality of the generated sequences.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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 detailed review and valuable feedback. We appreciate the reviewer’s recognition of the novelty and effectiveness of our method. Below, we address concerns regarding the lack of comparisons or discussions of related work. We hope the reviewer can consider raising your score in light of our response. \n\n**Q1: Discussion of Staged Speculative Decoding**\n\nThank you for the helpful suggestion. We will add a detailed discussion to the revised paper as below.\n\n“Staged speculation techniques [1] have been suggested to further accelerate inference by using a cascade of draft models, and hierarchical speculation shares similarities with these approaches. However, our method focuses on long sequence generation tasks, which presents unique challenges. Our approach uses different drafting methods for each speculation phase to address two distinct bottlenecks in long-context scenarios: model weights and KV cache. This differs from [1], which focuses on solving the same bottleneck (model weights) at every stage for further acceleration.”\n\n**Q2&Q3: Comparison and Discussion of Retrieval-based Drafting and Skipping Layers**\n\nWe appreciate the suggestion to compare TriForce with these two methods. The table below compares TriForce, REST [2], and Skipping Layers [3] with Llama2-7B-128K on an A100 GPU using the PG-19 dataset, showing that TriForce achieves the best speedup for long sequence generation settings.\n|Method|Speedup|\n|-|-|\n|TriForce|2.31x|\n|REST [2]|1.47x|\n|Skipping Layers [3]|1.36x|\n\nIn addition to the efficiency comparison, we also discuss each method as follows:\n\n**REST [2]**: Our retrieval-based drafting differs from REST, as our method retrieves information from the KV cache, while REST retrieves tokens from a predefined datastore. This distinction allows our approach to dynamically adapt to the context during generation, rather than relying on an external datastore.\n\n**Skipping Layers [3]**: In long-context settings where the primary bottleneck is the KV cache, this method shows some limitations. In the table above, it utilizes 68% of the KV cache, which was found to be the optimal combination of skipped layers. In contrast, TriForce can efficiently use just 3% of the KV cache, better addressing the bottleneck in long-context scenarios.\n\n[1] Accelerating LLM Inference with Staged Speculative Decoding\n\n[2] REST: Retrieval-Based Speculative Decoding\n\n[3] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The proposed method is straightforward and easy to follow.\n2. Using partial KV cache as a draft model is novel in the field of speculative decoding.\n3. The method indeed shows notable speedup, particularly in long context scenarios.", "reasons_to_reject": "The main weakness is the lack of comparisons or discussions of related work:\n1. The Hierarchical Speculation in Section 4.2 is not novel. [1] has proposed similar ideas but is not compared or discussed in related work (Section 2.1).\n2. [2] shows similar ideas as Retrieval-based Drafting in section 4.1, but is not compared or discussed in related work.\n3. Skipping layers[3] as a draft model is a similar method to your using partial KV cache, but is not compared or discussed in related work.\n\nMy intention is not that TriForce should beat these methods, since using partial KV cache as a draft model is new and should be useful in some scenarios (e.g. long context). However, I wonder why these related works are missed. I would raise my score if the authors justify the reasons.\n\n[1] Accelerating LLM Inference with Staged Speculative Decoding    \n\n[2] REST: Retrieval-Based Speculative Decoding     \n\n[3] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DdxOgCbbkq", "reviewer": "Reviewer_myEE", "summary": "The paper focuses on improving the decoding efficiency of language models on long context generation, where KV cache IO can be a major bottleneck. The goal is to improve the generation latency without sacrificing accuracy of the language models. The paper proposes TriForce, a speculative decoding method. TriForce is a hierarchical speculative decoding method, where there are two draft models, one is based on a small language model with StreamingLLM KV cache, and the other is based on the original target model but with a small subset of KV cache (obtained by retrieval). It can achieve more than 2x speedup when context is long (e.g., >120K tokens).", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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 encouraging feedback and insightful remarks. We have carefully addressed your questions and hope you will consider raising your score based on our response.\n\n**Q1: Clarification of Figure 3b**\n\n“Top-K” means taking the subset of KV cache with top-K largest attention weights. We will define “acceptance rate” clearly in the revised paper.\n\n**Q2: Clarification of “maintain a full cache for the initial two layers”**\n\nAs shown in Figure 3a, the first two layers are less sparse. Hence, we did not sparsify them in the observation section. For system efficiency, we sparsify all layers in the implemented system. This may lower the acceptance rate, but the constant drafting cost makes it a worthwhile trade-off.\n\n**Q3: Clarification of Figure 3c**\n\nThe x-axis represents the index of generated tokens. “Recovery rate” refers to the ratio of retrieved attention scores to the ground truth sum of attention scores for the prompt tokens ($s_1, \\ldots, s_n$). Assume the generated token index is $m$:\n\n$$\\text{Recovery rate} = \\frac{\\sum_{s_{mi} \\in R} s_{mi}}{\\sum_{i=1}^n s_{mi}}$$\n\nwhere $R$ represents the retrieved attention scores among prompt tokens.\n\n**Q4: Clarification of “JF68M Eviction”**\n\n“JF68M Eviction” refers to using only the JF68M model with StreamingLLM as the draft model.\n\n**Q5: Small draft model concern**\n\nTriForce is a general method for long sequence generation. Many existing self-speculative models can replace the small model for hierarchical speculation, such as Medusa [1] or EAGLE [2].\n\n**Q6: Ablation studies**\n\nWe clarify that our paper includes some ablation studies in Table 2. For example, “StreamingLLM w/ Hierarchy” represents replacing retrieval-based drafting with StreamingLLM. Additionally, we provide a table here.\n|Method|Speedup|\n|-|-|\n|Retrieval w/o Hierarchy|1.80x|\n|StreamingLLM w/o Hierarchy|1.75x|\n|StreamingLLM w/ Hierarchy|1.90x|\n|JF68M StreamingLLM|1.42x|\n\n**Q7: Scalability**\n\nFor longer inputs, a solution is offloading the KV cache to CPU memory. Below are experiments on L40 with offloading for longer inputs on LWM [3] models using PG-19, showing impressive speedup.\n|Model|Speedup|\n|-|-|\n|LWM-Text-256K|11.81x|\n|LWM-Text-512K|12.10x|\n\n[1] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads\n\n[2] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty\n\n[3] World Model on Million-Length Video And Language With Blockwise RingAttention", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Lossless efficient decoding for long-context generation, I believe, has become more and more important. It is good to study this problem, as it is very different scenario compared to what people used to study widely (e.g., speculative decoding with less than one thousand tokens).\n- The general idea that using the model with a subset of KV cache as a draft model makes sense to me, especially in the case of long-context generation. \n- The method is well-motivated by a couple of great observations (section 3 and figure 3). Those observations can be insightful themself, and facilitate other research.\n- The idea of selecting KV cache via retrieval is valid. It is also good that the paper shows that existing methods (StreamingLLM and H2O) perform poorly on the needle retrieval task.\n- The experimental results are fairly strong.", "reasons_to_reject": "- The presentation of the paper can be improved. There are quite a few places, which can be hard to follow without being very familiar with the related work/concepts. (I list some of them in the “Questions” section)\n- This approach requires a small draft model for the hierarchical speculation. This may not easily exist and may also take other resources during inference.\n- I think some ablation studies will be good to add. For example, what if only using StreamingLLM as the draft model, or only using the JF68M model as the draft model.\n- The approach, while being very efficient for 128K-token sequences. However, it may not further extend to much longer input (e.g., 1M), as the memory will be the bottleneck in that case.", "questions_to_authors": "- In Figure 3(b): what does “Top-K” mean? Is this to take the subset of KV cache with top-k largest attention weights? Also, for “acceptance rate”: I would suggest defining what it is before talking about the results. \n- What do you mean by “maintain a full cache for the initial two layers”?\n- In Figure 3(c): what is the x-axis? What does “recovery rate” mean?\n- Table 2/3: What is “JF68M Eviction”?", "extra_scores": {}}, {"review_id": "1djrKs4X6D", "reviewer": "Reviewer_FT2L", "summary": "This research introduces TRIFORCE, a novel hierarchical speculative decoding system that significantly accelerates inference efficiency for large language models on long context settings while fully preserving output quality. By leveraging insights from attention sparsity and contextual locality, the authors propose innovative techniques such as retrieval-based drafting (large model with retrieved cache) and hierarchical speculation (small model with StreamingLLM cache) to effectively mitigate the dual memory bottlenecks posed by KV cache and model weights. Comprehensive experiments on the PG-19 dataset demonstrate the effectiveness and practicality of the proposed approach, achieving up to 2.31x speedup on A100 GPUs and 11.65x speedup in offloading scenarios, as well as strong performance on large-batch inference tasks.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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 the supportive comments and recognizing the novelty of our method and the thorough evaluations. We hope our detailed clarifications and additional experimental results will address the concerns regarding our work.\n\n**Q1: Performance with longer outputs and scalability with increasing length**\n\n\nWe appreciate the suggestion to evaluate TriForce on scenarios with longer output sequences. In our current experiments, we focused on a long input and a short output to illustrate the efficiency gains in a commonly encountered setting. Here we provide additional experiments with longer output sequences up to 2K with Llama2-7B-128K on an A100 GPU using the PG-19 dataset:\n\n| Output Sequence Length | Speedup |\n| - | - |\n|256|2.31x|\n|512|2.28x|\n|768|2.34x|\n|1024|2.32x|\n|1536|2.28x|\n|2048|2.29x|\n\nAs seen in the table, performance remains relatively stable. This stability is due to the retrieval cache being reconstructed at fixed intervals, ensuring it stays aligned with the evolving context. Furthermore, our retrieved cache can be effectively reused across multiple decoding steps, amortizing the overhead of retrieval. These results demonstrate the scalability of TriForce and its robustness in handling varying output lengths, thereby addressing the concern about its generalizability.\n\n**Q2: Ablation of $\\gamma_1, \\gamma_2$**\n\n\nThank you for the suggestion. Here are the values used in our experiments: $\\gamma_1=2, \\gamma_2=6$. We also provide additional ablation results below with Llama2-7B-128K on an A100 GPU using the PG-19 dataset. We will also release our code to the public to facilitate the reproducibility of our work.\n\n\nAblation of $\\gamma_1$ ($\\gamma_2=6$): Due to the gap between JF68M with StreamingLLM and Llama-7B-128K with retrieval cache, the acceptance rate of the first speculation phase in our speculation hierarchy is low. As shown in the table, the optimal value for $\\gamma_1$ is 2.\n|$\\gamma_1$|Speedup|\n| - | - |\n|1|2.30x|\n|**2**|**2.31x**|\n|3|2.24x|\n|4|2.17x|\n|5|2.09x|\n|6|2.01x|\n|7|1.92x|\n|8|1.83x|\n\nAblation of $\\gamma_2$ ($\\gamma_1=2$): The table below shows that $\\gamma_2=6$ is the optimal value observed.\n|$\\gamma_2$|Speedup|\n| - | - |\n|1|1.62x|\n|2|1.90x|\n|3|2.13x|\n|4|2.22x|\n|5|2.26x|\n|**6**|**2.31x**|\n|7|2.24x|\n|8|2.20x|\n|9|2.17x|\n|10|2.13x|\n\nWe hope these additions address your concerns about the generalizability and reproducibility of our work. Let me know if there are any specific details or further adjustments you’d like to make!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Insightful observations in Sections 3.1 and 3.2 effectively motivate the design of TRIFORCE's key techniques.\n\n2. TRIFORCE consistently outperforms existing methods in long-context settings, demonstrating its effectiveness in efficiently handling long sequences.", "reasons_to_reject": "1. The evaluation focuses on a single setting with a long input (120K tokens) and a short output (256 tokens). To better understand TRIFORCE's generalizability, it would be valuable to test its performance on scenarios with longer output sequences and analyze how the method scales with increasing output length.\n\n2. The paper does not specify the hyperparameters γ1 and γ2 values, which control the speculation lengths in the hierarchical system. To facilitate reproducibility and deeper understanding, it is important to list all the hyperparameters used in the experiments and discuss their impact on the system's performance, as the effects of these hyperparameters remain unclear.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: ### Summary
The paper "TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding" introduces a novel approach aimed at improving the efficiency of long-sequence generation in large language models (LLMs). The method, named TriForce, combines Retrieval-based Drafting and Hierarchical Speculation to address the bottlenecks associated with model weights and key-value (KV) cache. The authors demonstrate significant speedups, particularly for the Llama2-7B-128K model, without compromising the quality of the generated sequences.

### Strengths
1. **Innovative Use of KV Cache**: The use of partial KV cache as a draft model is a novel idea in the field of speculative decoding.
2. **Significant Speedup**: The method shows notable speedup, especially in long-context scenarios, achieving up to 2.31x speedup on A100 GPUs and 11.65x in offloading settings on an L40 GPU.
3. **Strong Experimental Results**: The experimental results are robust and demonstrate the effectiveness of the method across various temperatures and settings.

### Weaknesses
1. **Lack of Comparisons to Related Work**: Initial submission lacked comparisons or discussions of related work. This was addressed in the rebuttal, but it would have been beneficial to include these comparisons in the original submission.
2. **Presentation and Clarity**: Some parts of the paper were hard to follow without familiarity with related work and concepts. The authors have addressed this in their rebuttal, clarifying key terms and concepts.

### Recommendations for Authors
1. **Include Comprehensive Related Work**: Ensure that the final manuscript includes detailed comparisons and discussions of related work as provided in the rebuttal.
2. **Clarify Key Concepts**: Improve the clarity of the presentation by defining key terms and concepts, as suggested by the reviewers.

### Final Decision
The paper presents a significant contribution to the field of long-sequence generation in LLMs with its use of hierarchical speculative decoding and retrieval-based drafting. The method's effectiveness is well-supported by strong experimental results. Despite some initial weaknesses in the presentation and lack of comparisons, the authors have addressed these concerns adequately in their rebuttal. Therefore, I recommend accepting this paper for COLM 2024.
Decision: Accept