TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

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.

Q1: Performance with longer outputs and scalability with increasing length

We 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:

As 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.

Q2: Ablation of $\gamma_1, \gamma_2$

Thank 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.

Ablation 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.

Ablation of $\gamma_2$ ($\gamma_1=2$): The table below shows that $\gamma_2=6$ is the optimal value observed.

We 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!

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.

Q1: Clarification of Figure 3b

“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.

Q2: Clarification of “maintain a full cache for the initial two layers”

As 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.

Q3: Clarification of Figure 3c

The 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$:

$\text{Recovery rate} = \frac{\sum_{s_{mi} \in R} s_{mi}}{\sum_{i=1}^n s_{mi}}$

where $R$ represents the retrieved attention scores among prompt tokens.

Q4: Clarification of “JF68M Eviction”

“JF68M Eviction” refers to using only the JF68M model with StreamingLLM as the draft model.

Q5: Small draft model concern

TriForce 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].

Q6: Ablation studies

We 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.

Q7: Scalability

For 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.

[1] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

[2] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

[3] World Model on Million-Length Video And Language With Blockwise RingAttention

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.

Q1: Discussion of Staged Speculative Decoding

Thank you for the helpful suggestion. We will add a detailed discussion to the revised paper as below.

“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.”

Q2&Q3: Comparison and Discussion of Retrieval-based Drafting and Skipping Layers

We 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.

In addition to the efficiency comparison, we also discuss each method as follows:

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.

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.

[1] Accelerating LLM Inference with Staged Speculative Decoding

[2] REST: Retrieval-Based Speculative Decoding

[3] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding