Sequoia: Scalable and Robust Speculative Decoding

Thank you very much for your thoughtful review. We are glad that the reviewer found our algorithm novel and the evaluation thorough and sound. We have tried to carefully address your questions. We hope the reviewer may consider raising their score in light of our response.

Q1: Batch serving setting

Thank you for this question. As you have correctly pointed out, Sequoia gives the largest gains for small batch sizes. This is because at small batch sizes, decoding is more memory-bound, and thus one can afford to verify a large tree without increasing verification time; by discovering the optimal structure for this large tree, Sequoia gives large speedups over baselines. Conversely, for larger batch sizes, the speculation budget is smaller, and the relative gains of Sequoia relative to other speculation tree structures (chains, k-ary trees, etc) is much smaller—we show this in Figure 1(a). In a real inference system with continuous batching, one could use Sequoia to determine the optimal tree for any batch size (as in Liu et al, 2024 [1]), thus attaining strong performance across the board. Furthermore, as you correctly pointed out, Sequoia can be combined with more powerful draft models (like Eagle) to attain even larger speedups (discussed briefly in Section 3.1.1).

[1] Xiaoxuan Liu, Cade Daniel, Langxiang Hu, Woosuk Kwon, Zhuohan Li, Xiangxi Mo, Alvin Cheung, Zhijie Deng, Ion Stoica, Hao Zhang. Optimizing Speculative Decoding for Serving Large Language Models Using Goodput. CoRR abs/2406.14066, 2024.

Q2: Performance decreasing with temperature

Thank you for your insightful question!

In this work, we use sampling without replacement to make the performance more robust across temperature. However, you are correct that decreasing is still observed when temperature goes up. Generally, high temperature means a higher degree of randomness, which is more difficult for draft models to guess. Improving the performance at higher temperatures is an interesting direction for future work—perhaps methods that better align draft and target models (e.g., via distillation [2]), or more advanced sampling algorithms [3, 4], could yield improvements here.

[2] Zhou et al. DistillSpec: Improving Speculative Decoding via Knowledge Distillation. ICLR 2024.

[3] Sun et al. Block Verification Accelerates Speculative Decoding. Efficient Systems for Foundation Models workshop, ICML 2024.

[4] Qin et al. Multi-Token Joint Speculative Decoding for Accelerating Large Language Model Inference. Arxiv 2024.

Thank you very much for your thoughtful review. We are glad that the reviewer appreciates our dynamic programming based tree search algorithm as well as the robustness of sampling and verification algorithms. We have tried to carefully address your questions. We hope the reviewer can consider raising your score in light of our response.

Q1: Time cost for dynamic programming and building optimal tree

Thank you for raising this concern. The time for building a Sequoia tree has two parts:

• Offline Dynamic Programming: Fixing the draft/target model and the workload, our dynamic programming can be conducted before inference (or offline). It is a one time preprocessing effort and does not introduce extra overhead during inference. Furthermore, our dynamic programming can run very fast, taking less than 20s to generate a tree as large as 1536, and less than 5s for trees of size 64 or 128.

• Tree building during inference: Time for building the tree is decided by the inference cost of the draft model and the depth of the tree (as we analyzed in Section 4.1: Hardware Optimizer). In this sense, building a Sequoia-optimal tree of depth D should cost the same time as building a k-array tree of depth D.

Q2: Sequoia’s performance at various batch sizes

Thank you for this question. As you have correctly pointed out, Sequoia gives the largest gains for small batch sizes. This is because at small batch sizes, decoding is more memory-bound, and thus one can afford to verify a large tree without increasing verification time; by discovering the optimal structure for this large tree, Sequoia gives large speedups over baselines. Conversely, for larger batch sizes, the speculation budget is smaller, and the relative gains of Sequoia relative to other speculation tree structures (chains, k-ary trees, etc) is much smaller—we show this in Figure 1(a). In a real inference system with continuous batching, one could use Sequoia to determine the optimal tree for any batch size (as in Liu et al, 2024 [1]), thus attaining strong performance across the board.

[1] Xiaoxuan Liu, Cade Daniel, Langxiang Hu, Woosuk Kwon, Zhuohan Li, Xiangxi Mo, Alvin Cheung, Zhijie Deng, Ion Stoica, Hao Zhang. Optimizing Speculative Decoding for Serving Large Language Models Using Goodput. CoRR abs/2406.14066, 2024.

Q3: Different trends for SpecInfer, Top-K and Sequoia

Thank you for raising this thoughtful question!

In short, SpecInfer’s performance gets worse at low temperature, whereas Sequoia and Top-K sampling improve at lower temperatures, because these methods use different sampling and verification algorithms. SpecInfer uses sampling with replacement, Top-K uses top-k sampling and Sequoia uses sampling without replacement. As a result, at low temperatures SpecInfer is the only method that would repeatedly sample an incorrect token, leading to poor performance.

We will now give a more detailed answer to your question, by first discussing how top-k sampling and SpecInfer sampling compare at both low and high temperatures, and then explaining how Sequoia is able to get “the best of both worlds” and perform well across both low and high temperatures.

In a low temp regime, the token with largest probability in the target model’s output will take up over 95%. This token is very likely to appear in the top-k tokens of draft models (even if it is not the top-1 token). In this case, Top-K sampling can easily get accepted. However, for SpecInfer’s sampling with replacement, since the temperature is low, SpecInfer will keep drafting the same token. When the drafted token is not exactly the top-1 token of the target model, SpecInfer will get rejected. To summarize, in the low temp regime, Top-K requires the top-1 token of target model is among the top-k tokens of draft model, while SpecInfer requires the top-1 of the target model to be exactly the same as the top-1 of the draft model, leading to worse performance. Top-K method benefits from its cover property (Section 3.2.2).

In a high temp regime, let’s consider an extreme case (temp $\rightarrow \infty$), i.e. the output of target model and draft model are totally random. In this case, it’s nearly impossible for Top-K to get accepted, as its total acceptance chance is just K / V, where V is the vocabulary size and K is the number of proposed tokens. However, for SpecInfer, all the draft tokens will be accepted, since the target and draft token probability is the same (1/V). SpecInfer method benefits from its optimal transport property (Section 3.2.2).

Sequoia with sampling without replacement, behaves more like Top-K in low temp and more like SpecInfer in high temp as it has both cover property and optimal transport property. That’s the reason why Sequoia is able to perform well across a wide range of temperatures. The trend of acceptance rates of three methods is shown in Figure 3, which verifies our above claims.

Thank you very much for your thoughtful review. We are glad that the reviewer found our work easy-to-follow and having comprehensive empirical results as well as theoretical analysis. We have tried to carefully address your questions. We hope the reviewer can consider raising your score in light of our response.

Q1: Clarifying the offloading setting

Thank you for pointing out our unclear description. In our offloading setting, we performed a layer-wise offloading for the target model (default setting of deepspeed-zero) and the draft model is on-device. So your understanding is correct, in this setting, the draft model is very fast (24.2ms) compared to the offloaded target model (5.7s). With Sequoia, we can accelerate offloading based decoding from 5.7s/token to 0.6s/token, which is more tolerable for running a large model on a consumer GPU. Furthermore, the setting of offloading (a big gap between FLOPs and memory bandwidth) also simulates the trend of future hardware, as shown in Figure 1(a).

Q2: Comparison with SpecInfer

Thank you for your question about our SpecInfer baseline. Below we show, through 3 sets of experiments, that Sequoia outperforms SpecInfer when we perform thorough sweeps for the SpecInfer tree configuration:

(1) Sweep of SpecInfer tree structure

Here, we use JF68M as the draft model, and Llama2-7b as the target, and sweep a very wide range of SpecInfer tree structures (for both greedy and stochastic):

Greedy Decoding (T=0.0):

Stochastic Decoding (T=0.6):

Sequoia attains speedups of 4.04x (greedy) and 3.18x (stochastic), outperforming all tree configurations of SpecInfer, including 5x8 tree (3.45x greedy, 2.47x stochastic).

(2) More SpecInfer results for A100

For on-device settings, we add the results for the 8x8 tree and 16x8 tree (SpecInfer) as follows,

Greedy Decoding (C4, T=0)

Stochastic Decoding (C4, T=0.6)

SpecInfer’s performance already degrades by enlarging the tree from 5x8 to 8x8. For SpecInfer, although the accepted tokens will marginally increase, the cost of verification/draft proposing will increase more.

(3) More SpecInfer results for L40 offloading

In Table 2 of the submission, we compare Sequoia trees of size 768 to SpecInfer trees of size 768, composed of 16 independent sequences of length 48 ("16x48"). Here we additionally compare to SpecInfer trees of shape 32x24 and 24x32:

We will include a larger sweep in the revised version (these experiments are time consuming and require specific CPU/PCIEs that are often not available on cloud servers). Thank you for understanding!

Q3: Standalone throughput of the draft and target models

Thank you for pointing out this missing part. We will include the numbers in the revised version.

JF68M: 0.5ms

Llama-2-7b: 24.2ms

Llama-2-70b: 5.7s

Our system is implemented in Huggingface with CUDAGraph and Static KV Cache (not as optimized as frameworks such as vLLM and deepspeed).

Q4: Accuracy of acceptance vector measurement

Empirically, this measurement is accurate. Below, we show that when we measure the average acceptance rate vectors across 10 different groups of 200 prompts each, the variance across these groups is quite small:

Setting: JF68M for Llama-2-7b, CNN, T=0.6, width=8

• Variance Vector: [0.0067, 0.0030, 0.0024, 0.0015, 0.0010, 0.0011, 0.0008, 0.0007]
• Acceptance rate Vector: [0.5608, 0.1077, 0.0539, 0.0343, 0.0236, 0.0186, 0.0146, 0.0122]
• Relative Error [1.2%, 2.8%, 4.4%, 4.5%, 4.2%, 5.9%, 5.3%, 5.4%]

Q5: dynamic programming algorithm should consider draft/target throughput

Yes. We discussed this in Section 4.1 (Hardware Aware Optimizer). In the case you mentioned, the optimizer will choose the shallowest tree (depth = 0) even if the acceptance rate is high, which means we do not use speculative decoding at all.

Thank you very much for your thoughtful review. We are glad that the reviewer found our work to have strong theoretical guarantees for scalability and also have good empirical results. Also thanks for noticing that our dynamic programming algorithm can be pre-computed offline. We have tried to carefully address your questions. We hope the reviewer can consider raising your score in light of our response.

Q1: Differences between Sequoia and existing tree-based methods

Thank you for raising this question.

Sequoia is designed to be a scalable and robust speculative decoding method that can yield large speedups in settings where the speculation budget is large (e.g., small batch serving, and offloading). In these settings, we show how we can optimally construct very large trees (e.g., 128 tokens for on-chip, and 768 for offloading), and that the number of accepted tokens keeps growing as the tree size grows—in Figures 1,4,6 (and theorem 3.6) we show that the number of accepted tokens can grow roughly logarithmically with the number of speculated tokens in the tree. We believe that Sequoia could be even more powerful on future hardware—the gap between computation and memory bandwidth is getting higher (see Figure 1), thus allowing a larger speculation budget.

Previous works like SpecInfer, Eagle, and Medusa apply shallow and small trees without studying the problem of scalability. By leveraging Sequoia’s tree search algorithm, our empirical results show that Sequoia scales well to arbitrary tree sizes. On the offloading setting on L40, which has a large budget, Sequoia can outperform SpecInfer by 51% on average.

In our work, in addition to improving the structure of the tree, we show that an improved sampling and verification algorithm can outperform others (SpecInfer, TopK) in terms of scalability and robustness. As described in Figure 4 and Table 5 in appendix, we can see that Sequoia achieves the largest speedups across all temperatures, attaining up to 1.65× and 1.27× speedup relative to SpecInfer and top-k sampling, respectively.

Q2: Speed up over SOTA

Thank you for raising this issue. In the last columns of Table 1/2, we have presented the speedups of SpecInfer, which is a very strong baseline.

In Table 1, we show that compared to SpecInfer, Sequoia attains speedups of 5% to 30% in the A100 on-chip setting (avg 22%). In Table 2, we show that Sequoia attains speedups over SpecInfer of 36% to 62% in the L40 offloading setting (avg 51%). More A100 on-device results can be found in Table 4 in appendix. We will add relative speedups compared to SpecInfer in the revised version.

In addition, since our method does not improve/train draft models, we only compare with draft model agnostic baseline. And the draft model improvement methods, like Eagle, Glide, Medusa are orthogonal to us and can be combined to achieve a better performance.

Q3: Reasons for not comparing with multiple-draft version of SpecInfer

Thank you for raising this problem! We have the following three reasons.

1. Draft model availability: Often only a single draft model is available (e.g., Llama3-8B as a draft model for Llama3-70B), and it would thus require significant time/energy for practitioners to train additional draft models.
2. Performance: As compared in SpecInfer (Appendix A, Table 4), most of the time one draft model outperforms multiple draft models.
3. System: SpecInfer serves each draft model on one GPU. However, every experiment in Sequoia is conducted on a single GPU, which means if we want to use the multiple-draft version, we need to run these draft models sequentially. This will further reduce the performance of the multi-draft version of SpecInfer.

Q4: Positional Acceptance Assumption

Validity of the positional acceptance assumption

The positional acceptance assumption states that the probability of a verification algorithm accepting a token $t$ which is the $k^{th}$ child of an already accepted token depends only on the value of $k$. This is a simplifying assumption, given that the probability of a specific token getting accepted additionally depends on the token and context. However, we can consider the average acceptance rates, across many tokens/contexts, for sampled tokens, as a function of their position $k$. In practice, we observe that this average is quite stable. For example (JF68M for Llama-2-7b, CNN, T=0.6, width=8), we show that when we measure the average acceptance rate vectors across 10 different groups of 200 prompts each, the variance across these groups is quite small..

• Variance Vector: [0.0067, 0.0030, 0.0024, 0.0015, 0.0010, 0.0011, 0.0008, 0.0007]
• Acceptance rate Vector: [0.5608, 0.1077, 0.0539, 0.0343, 0.0236, 0.0186, 0.0146, 0.0122]
• Relative Error [1.2%, 2.8%, 4.4%, 4.5%, 4.2%, 5.9%, 5.3%, 5.4%]

Impact of positional acceptance assumption

• Theoretically: we pointed out that, in tree based speculation methods, the probability of getting accepted is a function of the position of the speculated token. As a result, the optimal tree cannot be a balanced one (e.g. k-independent chains). This is also the intuition why we need and why we can search for an optimal tree.

• Empirically: With expected acceptance probability for each position (i.e. our acceptance vector, vector P in algorithm 1), we can pre-compute expected accepted tokens to search for the optimal tree structure (algorithm 1).

  For each experiment, we sample a subset of 200 sentences to calculate acceptance vectors and feed into algorithm 1 for tree searching.