MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding

We thank the reviewer for the supportive feedback. We are glad that you appreciated our work. And we have run some additional experiments, combined with theoretical analysis, to address your concern of real-world application of MagicDec.

Q1. Potential limitations in real-world application

We thank the reviewer for highlighting the practical constraints of real-world configurations in cloud environments. We agree that very long prompts in large batches perhaps remain uncommon in typical LLM usage patterns. However, MagicDec can also get gain for moderate length input prompt and medium batch size without adding any cost or hurting accuracy. As shown in Figure 3 (Page 5) of our paper, for 8xA100 GPU with a batch size of 256, speculative decoding becomes beneficial once the context length exceeds 3297 tokens, a threshold we term the critical sequence length. For modern GPUs in cloud centers like the H100, this critical length is even lower due to the higher FLOPs-to-memory bandwidth ratio. This critical prompt length is not a very large value and is more common in real-world applications.

To demonstrate practical applicability, we tested MagicDec using Llama-3.2-1B with streamingLLM KV for speculation of Llama-3.1-8B. The draft budget is 256. We used a prompt length of 3072 and a batch size of 128. Results are shown below:

For this moderate sequence length and batch size, MagicDec demonstrates speedup over standard decoding. As batch size or sequence length increases, the speedup becomes even more pronounced. The growth of sequence length is the current trend.

Recently, with the emergence of long-context models and applications like retrieval-augmented generation and in-context learning, the input context lengths have increased significantly in real-world serving systems. Production systems like Anthropic have reported this surge in context length of the input prompts [1]. Moreover, they often augment the user prompts with extra context to generate better responses. For instance, OpenAI o1 has built-in reasoning capabilities and it usually augments user prompts internally with chain-of-thought prompting. Hence, we believe that MagicDec will become more and more relevant in coming years.

[1] Anthropic: Prompt Caching with Claude. https://www.anthropic.com/news/prompt-caching

Dear Reviewer 3wSP,

Thank you for your thoughtful and constructive feedback on our work. We have carefully addressed your comments and revised the work and manuscript accordingly. As the discussion period nears its end, we would greatly appreciate any additional questions or points of clarification. If our responses have satisfactorily addressed your concerns, we kindly ask you to consider reflecting this in your score.

Thanks again for your time and expertise.

Q4. Lack of Discussion and Limitation: Performance Analysis on Lower-End GPUs

Thanks for the suggestion. In our revised draft, we have added the limitations of our work $\text{\textcolor{blue}{(Section 6, Page 10)}}$, that can be summarized as follows:

• Higher-end GPUs benefit more from MagicDec because of their better peak FLOPs to memory bandwidth ratio.
• Our current work only focuses on improving decoding performance.
• Our work does not exhaustively study all the SOTA KV compression algorithms, restricting ourselves to fewer algorithms that are representative of a broader class of KV compression methods.

Performance on Commodity Machines

MagicDec works better with GPUs with high peak FLOPs to memory bandwidth ratio. Because LLM inference is more memory-bound on such devices, restricting verification to target decoding cost ratio to a value close to 1. We have compared MagicDec’s speedups on customer-level machine 4090 with that on H100. For these experiments, we used Llama-3.2-1B with StreamingLLM KV as the draft model to speculate Llama-3.1-8B.

If the inference system utilizes an even resource-constrained device with a much smaller peak FLOPs to memory bandwidth ratio, MagicDec would suggest to use autoregressive decoding instead if required.

A Possible Utilization of MagicDec on Commodity Machines We present a way to adapt MagicDec to serve LLMs using commodity machines. To support batch inference of long sequences on these lower-end devices, we can adopt a distributed attention strategy proposed by recent works like [5, 6, 7] in order to optimize resource utilization. These approaches involve offloading part of the KV cache to other devices, including CPU. As described before in Q3, MagicDec can provide new opportunities to reduce the communication overhead in these methods by locally accommodating the compressed KV cache for drafting and only utilizing distributed attention for verification purposes. The nature of speculative decoding can also decrease the run times of the large target model with full KV cache, thus amortizing the communication overhead in offloading or distributed settings.

Optimizing Decoding Performance Only

Our current work only focuses on decoding performance, without any concern for prefilling performance. However, considering a disaggregated serving system (Zhong et al. [1], Qin et al. [2]) that decouples prefill and decoding stages, optimizing decoding performance alone is also quite beneficial.

Non-Exhaustive Study of KV Compression Methods

Our method is not an exhaustive collection of all the state-of-the-art KV compression techniques. Instead, we pick some representative solutions (static and dynamic KV selection strategies) based on their empirical acceptance rates to demonstrate the trade-off analysis among different kinds of KV compression methods in different batch size and sequence length regimes.

Future Plan

We plan to conduct an extensive study across all the state-of-the-art SD variants, including the layer-skip methods like (Elhoushi et al. [3], Zhang et al. [4], Xia et al. [5]), KV quantization (Liu et al. [6], Hooper et al. [7]), low rank decomposition (Singhania et al. [8]), etc., that are likely to mitigate the KV bottleneck problem in a large batch size, long-context regime.

References

[1] DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving (Zhong et al., 2024) [2] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving (Qin et al., 2024) [3] LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding (Elhoushi et al., 2024) [4] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding (Zhang et al., 2023) [5] SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration (Xia et al., 2024) [6] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (Liu et al., 2024) [7] KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (Hooper et al., 2024) [8] Loki: Low-rank Keys for Efficient Sparse Attention(Singhania et al., 2024)

Q3. Lack of Case Study or Worse-Case Analysis

We appreciate the reviewer’s concern about the variability in token acceptance rate as it is an important consideration in our speedup analysis. Our analysis shows that for a fixed setting, i.e. model, hardware, draft cost etc., different sequence length requires different minimum acceptance rates to achieve any speedup with self-speculation, with longer context lengths having more relaxed requirements.

The following table illustrates for a given draft KV cache budget, the minimum acceptance rate required by different sequence lengths to see any speedup (we consider batch size 4 here) for self-speculation. The admissibility of the draft budgets for each setting is based on the empirical acceptance rates obtained for PG-19 documents. The min_acceptance_rate for the small draft model with compressed KV cache will be much lower than these values because of lower draft_cost.

This indicates that a single KV cache budget can not be applicable for all sequence lengths.

Challenges for Heterogeneous Batches and Ways to Mitigate Them

As you have pointed out, the phenomenon discussed above is an important consideration in a real-world setting where heterogeneous batches can appear with sequences of varying lengths. Our analysis above can guide us to overcome this challenge in the following ways:

• Allocating Different Draft KV Cache Budgets for Sequences of Different Lengths:
  Because MagicDec supports PagedAttention, different draft KV cache budgets can be easily allocated to different requests in the batch.
• Request-Scheduling Algorithms:
  We can develop intelligent request scheduling algorithms based on their sequence lengths and task type (often dictates the acceptance rate). Recent works like Fu et al. [1], Srivatsa et al. [2] have developed new algorithms to schedule requests across multiple computing instances for better load balancing and higher throughput. Similarly, we can support different drafting algorithms with different draft cost budgets on different computing nodes and route the incoming requests to suitable nodes based on their requirements.
• Adding More Flexibility to Real-World Distributed Serving Systems:
  Recent distributed serving systems (Lin et al. [3]) have looked into distributing attention computation by offloading KV caches to a shared GPU pool in the same computing cluster. This approach has been effective for better load balancing and improving throughput of real-world heterogeneous batches, although at the cost of some communication overhead of routing query vectors to perform attention. In this scenario, MagicDec can provide new opportunities by prioritizing retention of draft KV cache in local memory and only utilizing the offloaded KV cache for verification purposes.

References

[1] Efficient LLM Scheduling by Learning to Rank (Fu et al., 2024) [2] Preble: Efficient Distributed Prompt Scheduling for LLM Serving (Srivatsa et al., 2024) [3] Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin et al., 2024)

Q2. Evaluation on Limited Spectrum of Models

Thanks for suggesting us to evaluate MagicDec on a broader spectrum of models. We have added evaluations for Mistral and Qwen series models to show the trends seen for Llama models also translate to the former $\text{\textcolor{blue}{(Appendix A.5, Page 15)}}$.

MagicDec achieves impressive speedups for Mistral-7B-v0.3, Qwen-2.5-7B and Qwen2.5-32B even at large batch sizes

We utilize self-speculation with SnapKV based KV selection for Mistral-7B-v0.3 and Qwen-2.5-7B models. For Qwen2.5-32B, we utilize Qwen-2.5-7B with streamingLLM KV cache as the draft model. We report the speedups obtained with the optimal speculation lengths.

Note: we use 8xH100s for Mistral-7B-v0.3 and Qwen-2.5-32B models and 4xH100s for Qwen-2.5-7B because it has 4 KV heads.

Similar to Llama models, we can see increasing speedups with increasing batch size for Mistral and Qwen models

Mistral-7B-v0.3 Speedups with SnapKV-Based Self-Speculation (Using 8xH100s)

Qwen-2.5-7B Speedups with SnapKV-Based Self-Speculation (Using 4xH100s)

Qwen-2.5-7B Speedups with Qwen-2.5-7B Draft Using StreamingLLM Cache (Using 4xH100s)

We appreciate your supportive comments and constructive suggestions. We have updated our manuscript to clarify the confusion about “missing baselines” [Q1] and have added comprehensive evaluations across different models [Q2] to illustrate the generalizability of our approach. In addition, we have added more case studies [Q3] regarding “worst case analysis” and stated the limitations of our current setup [Q4]. We hope our detailed clarification with further experiment results will clear the doubts about the significance of our work.

Q1. Missing discussion and comparisons of reasonable baselines: Triforce and other self-speculation methods, inference pipelines like VLLM

We thank the reviewer for suggesting these self-speculation works. We first want to clarify that our work illustrates how speculative decoding can be made useful even in a large batch setting and provides a generalized framework to evaluate different speculative decoding algorithms in a long-context large batch size regime. MagicDec requires a drafting strategy whose KV cache loading cost increases slower than the target model with increasing batch size and sequence length. Hence, any KV compression method for draft KV cache is an ideal candidate for MagicDec.

From that perspective, the KV retrieval algorithm of Triforce does indeed fall under the suite of algorithms that MagicDec can work with. Hence, it actually complements our efforts rather than serving as a baseline.

We also appreciate the other useful self-speculation methods suggested, such as Xia et al. [1] and Zhang et al. [2], which utilize layer-skipping strategies. These methods can also reduce draft KV cache cost and hence can be added on top of the retrieval algorithms that MagicDec studies. In summary, we view these suggested methods as complementary efforts that can expand MagicDec’s search space of suitable KV compression algorithms rather than as baselines.

Finally, MagicDec can provide additional performance improvements when deployed on efficient inference frameworks like VLLM by optimizing the inference-time memory loading cost. In fact, continuous batching-based serving systems allow MagicDec to perform at its full potential. Because continuous batching techniques have high memory utilization, they represent the exact setting where MagicDec can be most effective. Although our current implementation does not support continuous batching, we are planning to integrate MagicDec into VLLM in our future work.

[1] SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration (Xia et al. 2024)

[2] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding (Zhang et al.)

Dear Reviewer XfUb,

Thank you for your thoughtful and constructive feedback on our work. We have carefully addressed your comments and revised the work and manuscript accordingly. As the discussion period nears its end, we would greatly appreciate any additional questions or points of clarification. If our responses have satisfactorily addressed your concerns, we kindly ask you to consider reflecting this in your score.

Thanks again for your time and expertise.

Dear Authors,

Thanks for clarifying the questions to address my concerns, and I decided to keep my score.

Dear Reviewer XfUb,

Thank you for sharing your concerns. We would like to further clarify them.

(1) Lack of Novelty

As we clarified in our response to Reviewer PRVw, the main contribution of MagicDec is its identification that speculative decoding can accelerate large-batch inference in long-context serving. This challenges the conventional wisdom that speculative decoding does not provide speedup for large-batch inference. Reviewer UdZh, 3wSP and PRVw all acknowledged our novelty.

• Through the analysis of speculative decoding speedup and LLM inference performance, MagicDec first identifies for long-context serving, speculative decoding can accelerate large batch inference. More interestingly, the speedup even increases with batch size.
• MagicDec proposes the key to achieve high speedup is keeping the draft cost growing independently with sequence length. KV cache is the performance bottleneck that scales with both batch size and context length, so compressing the KV cache of the draft model can be a good way to limit the draft cost. Thus identifying KV compression as a necessary tool for efficient drafting is our main novelty rather than proposing a new compression method. There are potentially several ways to achieve that including small draft models with compressed KV cache, original model speculating itself with compressed KV cache or skipping its own layers etc.
• Finally, the primary goal of KV compression has traditionally been to preserve model accuracy. However, it remains unclear whether higher model accuracy directly correlates with higher token acceptance rates. For instance, while Llama-3.1-70B is more accurate than Llama-3.1-8B, it exhibits a lower token acceptance rate when speculating the latter. Interestingly, MagicDec suggests some KV compression algorithms can indeed achieve token acceptance rates when used in drafting stages.

Hence, all the existing KV compression techniques are indeed multiple ways to help us achieve the goal – keep draft cost constant with sequence length. MagicDec is a general framework that guides how to choose the optimal drafting strategy or KV compression methods based on draft cost, acceptance rate and hardware.

(2) Experimental Validation

1. We provided the experimental results on various tasks in our paper. We compare the acceptance rate of SnapKV-based self-speculation $\text{\textcolor{blue}{ (Figure 5, Section 4.1, Page 7)}}$ and the end-to-end speedup $\text{\textcolor{blue}{ (Table 2, Section 5.2, Page 9)}}$ comparison between PG-19, Ruler tasks(niah-multikey-3, cwe, qa-1). These results demonstrate the generalizability and practical applicability of MagicDec. In addition, we have not added results on InfiniteBench in particular, because most of the tasks in this benchmark have very small generated token lengths, which is not suitable to evaluate acceptance rates.

Update: We evaluated the acceptance rates of SnapKV based acceptance rates of Llama-3.1-8B on math-calc and longbook-sum-eng (the two subtasks of infinitebench with sufficiently long average output tokens). We got acceptance rates of 93.2% and 92.6% respectively for prefill length of 32k and draft budget 2k, which are both higher than the acceptance rate on PG-19. Based on our analysis, under the same prefill length, batch size and draft budget, the speedup of MagicDec on math-calc and longbook-sum-eng tasks of Infinitebench is expected to be higher than the speedup on PG-19 when using the same KV-compression method due to the higher acceptance rate. For instance, we achieved a 2.34x speedup for math-calc task and 2.29x speedup for longbook-sum-eng task with Llama-3.1-8B model (prefill length=32k, batch size=128, draft budget=2k, gamma=4) on 8 A100s. Both these speedups are higher than what we achieved for PG-19 task on 8 H100s.

2. We also implemented SnapKV-based self-speculation on the state-of-the-art LLM inference framework MLC-LLM, and compared the end-to-end speedup results in $\text{\textcolor{blue}{ (Table 4 and Table5, Section A.3, Page 14)}}$, which also highlights the generalizability of our method.

Dear Reviewer XfUb,

If our latest responses have addressed your additional concerns, we kindly ask you to consider adjusting your rating back to positive. Thanks for your time and valuable insights.

Comments

Comment 1: Comment on Lines 52-74 and Figure 1a

Thanks for your suggestion. The KV store time is a very small portion of the total inference time. We have updated Figure 1a $\text{\textcolor{blue}{(Section 1, Page 2)}}$, isolating KV load time and store time to further clarify our claim. For a sequence with prefix length equal to 16000, during decoding the KV load time is approximately 16000 times larger than the store time, as each time we need to load all the previous tokens’ KV cache, while only need to store the key and value states of the new generated token. Thus, the bottleneck of inference is exactly KV load time.

Comment 2: Comment on Lines 203-211 ("Expected Generation Length Per Step")

Thanks for mentioning this interesting work. You are right this work also uses expected generation length and draft cost to assess whether a draft model is good or not. We think the distributed speculative inference proposed in this paper is perfectly complementary with our work. The distributed speculative inference overlaps verification cost. With the high acceptance rate and low draft cost offered by compressed KV-based drafting, the speedup could be higher when applied to long-context serving. We have added the discussion of this work in our related work section $\text{\textcolor{blue}{(Section 2, Page 3)}}$.

Thank you for providing such thoughtful and supportive feedback. We are glad that you found our analysis novel and insightful, and our contribution highly valuable. Based on your excellent suggestion, we have streamlined the analysis section discussing speculation decoding speedup factors and KV caching. We hope that we have been able to improve the readability of our paper, especially for the broader research community.

Additionally, we are thankful for your insightful questions and further suggestions to improve our paper. We hope that responses answer some of your questions, and would look forward to any further comments.

Q1: Can you clarify the potential trade-offs in performance if MagicDec were applied to significantly smaller LLM models?

Interestingly, we find that smaller LLMs have a sharper growth in speedup with batch size compared to larger LLMs. In addition, the critical sequence length beyond which they attain higher speedups is also lower compared to the latter. For instance, while the critical sequence length for Llama-3.1-8B is ~4000, that of Llama-3.2-1B is just ~2000. This is because of two reasons:

• The verification-to-decoding cost ratio is lower for smaller LLMs. For a given batch size and hardware, smaller LLMs become memory-bound for a shorter sequence length because of a smaller hidden state dimension. The more memory-bound the target model is, the smaller is the verification-to-decoding cost ratio.
• The draft-to-target cost ratio is lower as well. Because the parameter and compute cost shared by the draft and the target model are much smaller than the cost the draft model optimizes for, the KV loading cost.

Here is a comparison between the theoretical speedups achieved by Llama-3.2-1B model and Llama-3.1-8B model using self-speculation.

Speedups achieved by Llama-3.2-1B model

Speedups achieved by Llama-3.1-8B model

However, there could be one potential disadvantage for smaller LLMs. Even though self-speculation is usually sufficient for them, in general, they do not have suitable small draft model options. Hence, they might lose a little bit of flexibility in terms of choosing the best drafting strategy.

Q2: Is the performance gain from MagicDec sustained across varying types of long-context tasks, particularly those that require variable batch sizes or non-standard hardware configurations?

Since MagicDec presents an algorithm that can be deployed in different task and hardware settings, we believe that performance gains can be sustained across a broad application space. Moreover, MagicDec is flexible enough to adapt to varying types of task loads. For instance: If some long-context task requires variable batch sizes, then MagicDec would dynamically choose the appropriate draft KV budget and KV compression algorithm to achieve the best performance for the corresponding batch size. In case of non-standard hardware, MagicDec’s effectiveness depends on the peak FLOPs-to-memory bandwidth ratio of the device. For instance, for CPU-based inference, our method would suggest using regular autoregressive decoding instead.

Please let us know if this answers your question. We would be happy to discuss this in more detail.

Dear Reviewer UdZh,

Thank you for your thoughtful and constructive feedback on our work. We have carefully addressed your comments and revised the work and manuscript accordingly. As the discussion period nears its end, we would greatly appreciate any additional questions or points of clarification. If our responses have satisfactorily addressed your concerns, we kindly ask you to consider reflecting this in your score.

Thanks again for your time and expertise.

Dear Reviewer UdZh,

Thank you once again for your thoughtful feedback and the time you’ve dedicated to reviewing our work. As the extended discussion period draws to a close, we want to ensure that all your concerns have been fully addressed. If there are any remaining points requiring further clarification, please don’t hesitate to let us know.

We deeply appreciate your time and valuable input, which have been instrumental in improving our work.

Q5. The applications for the very long contexts prompts in large batch size still need further justification.

We thank the reviewer for raising the concern of real-world application for very long-context and large batch size prompts. As we have mentioned in our response to Reviwer 3wSP, the emergence of long-context models, retrieval-augmented generation (RAG), and in-context learning has dramatically increased prompt lengths in recent years. (Anthropic [1]). Even if the user prompts are not too long, the production systems internally increase the context length with retrieved documents or chain-of-thought prompting. For instance, recently OpenAI o1 has started using COT prompting to improve their responses.

But, we agree that prompt length like 100k with large batch size may not be as common in current real-world applications. However, MagicDec can also get gain for moderate length input prompt and medium batch size without adding any cost or hurting accuracy. As shown in Figure 3 (Page 5) of our paper, for 8xA100 GPU with a batch size of 256, speculative decoding becomes beneficial once the context length exceeds 3297 tokens, a threshold we term the critical sequence length. For modern GPUs in cloud centers like the H100, this critical length is even lower due to the higher FLOPs-to-memory bandwidth ratio. This critical prompt length is not a very large value, and is more common in real-world applications.

To demonstrate practical applicability, we tested MagicDec using Llama-3.2-1B with streamingLLM KV cache for speculation of Llama-3.1-8B. The draft budget is 256. We used a prompt length of 3072 and a batch size of 128. Results are shown below:

For this moderate sequence length and batch size, MagicDec demonstrates speedup over standard decoding without hurting generation quality. And as batch size or sequence length increases, the speedup becomes even more pronounced.

[1] Anthropic: Prompt Caching with Claude. https://www.anthropic.com/news/prompt-caching

Q4. Some important design details can be more explicitly clarified.

We thank the reviewer for posing these questions. We have added clarifications for the details mentioned by the reviewer with blue color in the revised paper. Specifically:

Where does the time breakdown such as Figure 1(a) come from?

The time breakdown is based on the analysis results from [1], which is a tool for visualizing LLMs and analyzing the performance on different hardware platforms. We have used both roofline modeling (as suggested by you) and additive modeling (does not consider latency-hiding optimizations) in our analysis. However, in our paper, we report the time breakdown with additive modeling only, as neither of the two theoretical modeling approaches is fully accurate. As you have mentioned, the modeling type does not contradict the main takeaways of the paper.
[1] Yuan Z, Shang Y, Zhou Y, et al. LLM Inference Unveiled: Survey and Roofline Model Insights. arXiv preprint arXiv:2402.16363, 2024.

Is KV compression used for both the draft and target model?

The KV compression is only used for the draft model, while the draft model can be a smaller model or the target model itself (self-speculation). For self-speculation, we actually have two independent KV caches, but use the same model weights. During the prefill phase of the target model, we use the KV compression method to generate a compressed KV cache with a certain KV budget. During the drafting stage, we use the LLM and the compressed KV cache to generate draft tokens. During the verification stage, we use the LLM with the full KV cache to verify the draft tokens.

What is the baseline for results in Figure 4?

In Figure 4, we don’t compare the results with any baseline. We can provide more explanations for Figure 4. Figure 4(a) and Figure 4(b) show that for different prompt lengths, the memory footprints of a small draft model will be close to or even surpass the target model when batch size increases. The increase of the draft KV cache leads to this phenomenon, demonstrating the inefficiency of a small draft model in large batch size inference regimes and the necessity of KV compression for the draft model. Figure 4(c) compares the acceptance rates of different KV compression methods when applied as the drafter in speculative decoding. Top-K in this context represents the theoretical upper bound of sparse attention approximation methods.

As for this lossy KV compression selection method, perhaps consider adding the evaluation that shows the selection in KV compression strategy has advantage (e.g., speedups) over the original speculative decoding.

We have added the comparison between the KV compression-based speculative decoding and the original speculative decoding to Figure 7(b) $\text{\textcolor{blue}{(Page 10, Evaluation Section)}}$ in the revised paper. The results are shown below. The input prompt length is 8192. Hardware: 8xH100. We use Llama-3.2-1B as the drafter to do speculation for Llama-3.1-8B.

Llama-3.2-1B, full KV cache

Llama-3.2-1B, compressed KV with Constant Budget 512

We can see from these results that Llama-3.2-1B with compressed KV cache outperforms original speculative decoding for each batch size. The main reason for this is that compressed KV cache limits the growth of draft cost while still keeping a high acceptance rate.

How the framework guides the selection of KV compression method remains unclear to me.

Thanks for pointing out the unclarity. We have updated the Method section $\text{\textcolor{blue}{(Section 4, Pages 7-8)}}$ of the revised paper to better illustrate how to choose the optimal drafting strategy from: Different draft model sizes, Draft KV budgets, KV compression methods.

Q2: The paper jumps to some conclusions too fast without enough explanations

Thanks for pointing out the writing issue. We have revised our paper to make our explanations more streamlined and easy to follow. The modifications are shown in blue color in the Introduction section. Specifically:

What does Figure 1(b) try to show?

Figure 1(b) illustrates the comparison of throughput between MagicDec and standard autoregressive decoding at a given token-wise latency budget. It is well-known that simultaneously improving throughput and latency is challenging, especially when model quality cannot be sacrificed. This plot exhibits MagicDec’s ability to achieve better throughput and latency across the spectrum for long-context requests.

Maybe consider adding more details about the effects of increasing batch size with the original speculative decoding.

Thanks for figuring this out. We have added additional experiment results of original speculative decoding in the paper to show the effect of increasing batch size, which is shown in Figure 7(a) $\text{\textcolor{blue}{(Fig. 7, Page 10)}}$. The detailed results are shown below. T_D stands for draft cost.T_V stands for verification cost. T_SD is the average speculative decoding latency (ms), while T_Auto is the baseline autoregressive decoding latency. We use Llama-3.2-1B as the draft model to speculate Llama-3.1-8B. The results show that speedup decreases as batch size increases.

Q3: Some terms are vague in the texts.

We thank the reviewer for figuring out these issues. We have updated more explanation for the terms mentioned by the reviewer in the paper (highlighted blue). Specifically:

What does “self-speculation” in the texts mean?

Self-speculation in our work refers to leveraging the same LLM to perform speculative decoding by utilizing a compressed KV cache as a draft mechanism. During the drafting stage, we use the LLM with the compressed KV cache to generate several tokens. During the verification stage, we use the LLM with the full KV cache to verify these drafted tokens.

What does “KV budget” mean?

KV budget in our paper means the size of the KV cache for each sequence after compression.

We thank the reviewer's suggestions and questions. We have updated the manuscript to clarify the novelty of MagicDec [Q1], add more explanation for Figure 1(b) and details of increasing batch size with original speculative decoding [Q2]. We also added more explanations for some terms in the paper [Q3] and clarified the details mentioned by the reviewer [Q4]. We add some experiments on smaller batch sizes and shorter context length settings to show the effectiveness of MagicDec for not very long prompts [Q5]. We hope our detailed clarification with further experiment results will clear the doubts about the significance of our work.

Q1: The idea of compression KV is already in the literature, so the main new ideas are not all that large.

Thanks for pointing out the lack of clarity about the primary novelty of our work. We want to first clarify the main contributions of MagicDec. The conventional wisdom says that speculative decoding can not provide speedup for large batch inference, primarily because the token verification process becomes too expensive in the compute-bound regime. We are the first to present the limitations of this existing wisdom.

• Through the analysis of speculative decoding speedup and LLM inference performance, MagicDec first identifies for long-context serving, speculative decoding can accelerate large batch inference. More interestingly, the speedup even increases with batch size.
• MagicDec proposes the key to achieve high speedup is keeping the draft cost growing independently with sequence length. KV cache is the performance bottleneck that scales with both batch size and context length, so compressing the KV cache of the draft model can be a good way to limit the draft cost. Thus identifying KV compression as a necessary tool for efficient drafting is our main novelty rather than proposing a new compression method. There are potentially several ways to achieve that including small draft models with compressed KV cache, original model speculating itself with compressed KV cache or skipping its own layers (as suggested by R1) etc.
• Finally, the primary goal of KV compression has traditionally been to preserve model accuracy. However, it remains unclear whether higher model accuracy directly correlates with higher token acceptance rates. For instance, while Llama-3.1-70B is more accurate than Llama-3.1-8B, it exhibits a lower token acceptance rate when speculating the latter. Interestingly, MagicDec suggests some KV compression algorithms can indeed achieve token acceptance rates when used in drafting stages.

Hence, all the existing KV compression techniques are indeed multiple ways to help us achieve the goal – keep draft cost constant with sequence length. MagicDec is a general framework that guides how to choose the optimal drafting strategy or KV compression methods based on draft cost, acceptance rate and hardware.

Dear Reviewer PRVw,

Thank you for your thoughtful and constructive feedback on our work. We have carefully addressed your comments and revised the work and manuscript accordingly. As the discussion period nears its end, we would greatly appreciate any additional questions or points of clarification. If our responses have satisfactorily addressed your concerns, we kindly ask you to consider reflecting this in your score.

Thanks again for your time and expertise.

Dear Reviewer PRVw,

Thank you once again for your thoughtful feedback and the time you’ve dedicated to reviewing our work. As the extended discussion period draws to a close, we want to ensure that all your concerns have been fully addressed. If there are any remaining points requiring further clarification, please don’t hesitate to let us know.

We deeply appreciate your time and valuable input, which have been instrumental in improving our work.