Distributed Speculative Inference (DSI): Speculation Parallelism for Provably Faster Lossless Language Model Inference

Response to Reviewer LmXh

We sincerely thank you for your detailed and thoughtful review. We are glad that you found our proposed speculation parallelism novel and interesting and that you appreciated the rigor in our theoretical analysis and simulations.

Acknowledged Strengths:

• Novelty and Inspiration: We appreciate your recognition that our proposed speculation parallelism is novel and could inspire future work on speculative inference and LLM inference acceleration.

• Rigor in Analysis and Validation: Thank you for highlighting the rigor in both our theoretical analysis and the careful design of numerical simulations validating our approach.

• Transparency: We are pleased that you appreciated our provision of source code for the experiments, which demonstrates our commitment to transparency and reproducibility.

Responses to Weaknesses and Questions:

Weakness: My main concern is that the practical value of this work is not yet evident enough. If I understand correctly, the implementation of DSI in the current work is more or less a proof-of-concept, rather than one that is ready for application. More specifically, only the outline of the multi-thread system is implemented, while the LLM part (e.g., forward passes of the drafter or target model) is not actually implemented but rather replaced by WAIT commands with delays that mimic the actual latencies of real LLM forward passes, for the purpose of simulations.

Response: Your understanding of our proof-of-concept is correct. What we have are theoretical guarantees backed by simulations. The simulations show the expected speedups for well-known models (Table 2) and for any possible drafter (Figure 2) in the single-node setup.

Weakness: Combining LLM inference with the implemented multi-thread system can be challenging, as these two might not be completely decoupled. One particular aspect I can think of is the compatibility with KV caches, which have been supported by both non-SI and SI, and enabled by default in most cases. With multiple target LLMs (or verifiers) on different GPUs, each at its own pace, managing and synchronizing their KV caches seems like a non-trivial task and requires quite some engineering work, which has not been implemented in this work.

Response: DSI essentially constructs and verifies token trees on the fly. Efficient KV cache management of token trees has already been developed in SpecInfer, where tree paths can share common prefixes. Practitioners who have access to a node with 8 GPUs can apply SpecInfer’s KV cache management as-is to achieve the expected speedups reported in this paper. While it might require some engineering effort to implement SpecInfer’s KV cache management, it is a solved research problem and has been shown to add negligible latency. Therefore, the orchestration (DSI) is essentially decoupled from the underlying computation of forwards (including KV cache management) not only theoretically but also practically.

Miao, Xupeng, et al. "SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification." arXiv preprint arXiv:2305.09781 (2023).

Question: Do you think DSI with a batch size larger than 1 is possible, at least in theory? Batch inference is already quite a challenge for SI, and things seem to get much more complicated with the proposed speculation parallelism.

Response: Great question, and a very exciting future direction with huge potential! It has recently been shown that SI is effective in the multi-request setup (where the batch size is often larger than one) for sufficiently long sequences ¹. No fundamental limitations prevent practitioners from applying DSI to multi-request setups. This paper covers all the foundational concepts of such an implementation of DSI. However, we believe that analyzing or simulating the expected speedups is not straightforward. Unfortunately, we do not currently have access to a node with 8 GPUs to further research this direction.

Minor Writing Issues:

• The paragraph between Lines 422–487 is a bit long; might be better to divide it into several paragraphs.
• Line 128 mentions "implementing (3) above," while Eq. (3) is actually located in the appendix.
• Line 303: Symbols f1 and f2 should be fj and fm instead.
• Line 931: "upper bound" should be "lower bound," since some components of latency in reality are neglected in the estimate.

Response: Thank you for carefully catching these typos. We will revise the paper to address all these issues and improve the clarity of the paragraph between Lines 422–487.

We appreciate your insightful comments and constructive feedback, which have already helped us improve the paper. Thank you again for your thoughtful review.

I'd like to thank the authors for their responses. My overall perspective on this work hasn't changed much, thus I'm inclined to keep my rating for now.

We sincerely thank you for your thoughtful review and valuable feedback. We are glad that you found the paper well-written and easy to follow and appreciated the clarity and effectiveness of our proposed method.

Acknowledged Strengths:

• Clarity and Readability: We appreciate your positive remarks regarding the clarity and ease of understanding of the paper.

• Effectiveness of the Idea: We are pleased that you found our proposed idea clear and supported by the evaluation results.

Responses to Weaknesses and Questions:

Weakness: The largest target model is capped at 13B. How does the model work on larger models?

Response: The largest target model we evaluate is 15B (StarCoder). The guarantees of DSI hold for any model size, with the expected speedup provided in Figure 2. As the target model size increases while the drafter remains fixed, the "drafter latency" becomes smaller because it represents the ratio between the latencies of the two models.

Weakness: How does the data communication defect the approach? The KV-cache also seems to move between GPUs. How about synchronization issues?

Response: Each server maintains its own KV cache. The servers collaboratively process a token tree with shared prefixes. Therefore, the implementation directly follows the token tree KV cache management method of SpecInfer. Synchronizations occur at every draft rejection.

Miao, Xupeng, et al. "SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification." arXiv preprint arXiv:2305.09781 (2023).

Weakness: The analysis of the energy consumption compared to the original speculative decoding method seems to be missing. More justification may be needed given the energy inefficiency of the method.

Response: How do you define “energy consumption” in this context, and what analysis do you believe is missing?

We hope that our responses address your concerns and provide clarity. Thank you once again for your time and consideration. We look forward to any further suggestions you might have.

We sincerely thank you for your thoughtful review and valuable feedback. We are pleased that you recognize the strengths of our approach, particularly the effectiveness of Speculation Parallelism in transforming the blocking operation of Speculative Inference (SI) into a non-blocking one, and that our method is backed by proof.

Acknowledged Strengths:

• Clarity up to the Method Section: We appreciate your positive comment that the paper is well-written up to the method section.

• Effectiveness of Speculation Parallelism: We're glad that you acknowledge how our proposed Speculation Parallelism effectively solves the blocking operation of SI, making it non-blocking.

• Method Backed by Proof: Your recognition that our method is supported by theoretical proofs is encouraging.

Responses to Weaknesses and Questions:

Weakness: Evaluations are very weak. Can we get the experimental result by varying lookahead values on different machines?

Response: The reported results are based on simulations. They rely on estimated TTFT and TPOT—as measured “offline,” namely, before running the simulations. Changing machines will likely change the TTFT and TPOT, and therefore the effective “drafter latency.” Figure 2 provides the expected speedups for any drafter latency, and therefore covers all possible single-node machine configurations with 8 GPUs.

Weakness: Actual improvement measurement seems missing.

Response: Could you please clarify what specific measurements you believe are missing? The actual improvement measurements are reported in our results section, demonstrating the speedups achieved by DSI over SI and standard methods.

Weakness: Evaluating distributed evaluation on a single machine doesn’t make sense. Authors really need to find a way to run the evaluation on multi-GPU machines. Distribution involves communication latency as well, so in practice there are many things involved, but these practical concerns are not discussed in the paper.

Response: The theoretical guarantees of DSI hold for both single- and multi-node setups. Unfortunately, we do not have the resources to evaluate a complete implementation of DSI over even a single node with 8 GPUs. Consequently, evaluating a multi-node implementation of DSI (orchestrating more than 8 GPUs) is not planned.

Weakness: Evaluation writing is also very confusing.

Response: We will rewrite this part to improve clarity and ensure that the evaluation methodology and results are clearly communicated.

Question: What kind of communications are involved in distributed speculative parallelism and how does that affect the latency?

Response: The speedups reported in Table 2 are based on simulations where DSI is implemented as a multithreaded system, and therefore incur all the real-world latencies of such systems, including context switching.

We hope that our responses address your concerns and provide the necessary clarifications. We will work on improving the evaluation section to enhance clarity and address the issues you've raised. Thank you once again for your time and consideration.

We sincerely thank you for your thoughtful review and valuable feedback. We are pleased that you recognize the novelty of our approach and its effectiveness in addressing the fundamental limitation of Speculative Inference (SI) as a sequential algorithm.

Acknowledged Strengths:

• Novelty of DSI: We appreciate your acknowledgment that Distributed Speculative Inference (DSI) is a novel approach that not only builds upon SI but also addresses its weaknesses, such as the drafter being blocked by the verifier until generated tokens are verified. Your recognition of DSI's ability to generalize SI to work with multiple drafters in parallel is encouraging.

• Empirical Evidence of Performance: We're glad that you find our empirical evidence compelling in demonstrating that DSI outperforms both traditional SI and non-SI methods. Your positive remarks on our experimental results motivate us to further our research in this area.

Responses to Weaknesses and Questions:

Weakness: The approach requires more compute and memory (multiple KV caches) than SI, with the amount of extra compute and memory required increasing with the number and size of the drafters, and with an increase in input context length. Therefore, I believe that the increase in cost/resources also needs to be considered in addition to the reduction in latency to make a fair assessment.

Lines 273-291 compare DSI, tensor parallelism (TP), and pipeline parallelism (PP) over the same budget.

Question 1: I do not understand the rationale behind the formula for acceptance rate on Page 8, line 410. Please clarify.

We estimate the acceptance rate by fitting a geometric distribution as in Mamou, Jonathan, et al. "Accelerating Speculative Decoding using Dynamic Speculation Length." arXiv preprint arXiv:2405.04304 (2024).

Question 2: Why are you using an estimate of the acceptance rate instead of independently simulating acceptance/rejection for each drafter-target pair? For example, you could generate 'lookahead' tokens, verify them, go back to the drafter, and so on until the entire sequence is generated, and then use the estimates of TTFT and TPOT to estimate the total latency.

The reported total latency and the total latency that is computed as suggested are equal in expectation.

Question 3: Please report the value of lookahead used for each row of Table 2, and clarify why that value was chosen.

The lookahead is the minimal value out of 1, 5, and 10 that satisfies Equation 1, as mentioned in Lines 365-369.

Question 4: Why are real-world latencies of multithreading and context switching ignored (Page 10, lines 487-488)? These contribute to the overhead of DSI, and I feel that they should be reported and studied for a fair comparison with SI and non-SI.

Table 2 reports speedups in a simulation where DSI is implemented as a multithreaded system that incurs all the real-world latencies, including context switching. (You referred to Figure 2, a visualization of an additional ablation experiment. This experiment implements an “offline” simulation where we count forwards and multiply them by their estimated latency.)

Question 5: Will it be possible to show any simulation results with multiple drafters? I feel that one of the strengths of DSI is that it can work with multiple drafters, and I would like to see how using multiple drafters compares with using a single drafter.

Thank you. Adding drafters is mathematically guaranteed to increase the acceptance rate, as studied in SpecInfer. Since we use SpecInfer’s verification method as-is, their results immediately apply—suggesting the expected increase in the acceptance rate. Then, Figure 2 provides the expected speedup of DSI for any acceptance rate.

Miao, Xupeng, et al. "SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification." arXiv preprint arXiv:2305.09781 (2023).

We hope that our responses address your concerns and provide the necessary clarifications. We are grateful for your insightful comments, which have helped us improve our manuscript. Thank you once again for your time and consideration.

Thank you for following up and raising this important and thoughtful question.

The total latency in this experiment is defined by the sums of the number of target and drafter forward passes. To calculate these sums, we simulate the process by sampling from the models, executing the verification procedure, and counting the number of required forwards. Since sampling from models is inherently stochastic, the total latency $T$ is a random variable, and we focus on its expectation, $\mathbb{E}[T]$.

There exists an acceptance rate $p$ such that the expected latency satisfies $\mathbb{E}[T] = T$. This acceptance rate can be estimated in two ways. The first method directly computes the likelihood that a token is accepted by the verifier. Alternatively, as explained in Lines 407–410 of the paper, $p$ can be estimated by fitting a geometric distribution to the acceptance process. Specifically, this involves calculating the average number of tokens accepted per iteration and extracting $p$ as the parameter of the geometric distribution.

Both approaches are consistent. As the number of iterations approaches infinity, the estimated acceptance rate converges to the true empirical acceptance rate $p$. Consequently, the estimated total latency based on $p$ converges to the actual total latency. This guarantees that the reported latency and the computed latency are equal in expectation.

We greatly appreciate this question and recognize that additional clarity in this explanation would strengthen the paper. We will revise the manuscript to provide a more detailed and precise discussion of this point. Thank you again for your insightful feedback, which has helped us improve the presentation of our work.