Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation

Summary

We sincerely thank the reviewer for the effort in providing constructive comments. We addressed all the feedback from the reviewer, and believe that with these changes, the results in our paper become significantly stronger (added experiment result numbers in reviewer USDK’s rebuttal).

Per-feedback detailed response

I cannot find any comparison between baseline.

In the original submission, we included two RAG baselines where we put the conclusion in the experiment section 4.4 and detailed results in the Appendix B. We now move all results in T1 (reviewer USDK's comment) to the main body, along with two more baselines as suggested by the reviewer:

1. We evaluate LLMLingua 2 using LongBench. Similar to RAG, LLMLingua does not support query-aware compression, and therefore, we separate the context from the query and send it to the compression model. In T1 (reviewer USDK's comment), SpecPrefill outperforms LLMLingua and RAG Llama.

2. We benchmark MInference's quality on LongBench and its inference latency in vLLM for a more comprehensive view. In T1 (reviewer USDK's comment), the quality score of MInference can already be matched very closely with SpecPrefill with 10% token keep rate (0.3% loss) and is outperformed when we increase the rate to 30% and higher. In T2 (reviewer USDK's comment) where we vary the batch size x sequence length products from 128 * 4K to 4 * 128K, we found that SpecPrefill is particularly effective in efficiency compared to MInference with similar or better quality (getting 2.54x - 6.54x relative speedup compared to MInference). The main reason why MInference might fall short is due to the overhead in dynamic index building, which gets significant when the batch size is large (this is also observed in paper https://arxiv.org/abs/2502.05431 and confirmed by the author in issue https://github.com/microsoft/MInference/issues/18).

I think an ablation study on N should be in the paper.

We agree that ablation on N will provide more insights and are running a full ablation on N for the camera ready version. For our current results as in Figure 2 and Figure 5, we showcase the benefits of using LAH with steps = 8 (with LAH as the suffix in model names), which proves its usefulness especially for shorter context tasks.

I strongly suggest authors submit the code as a supplementary to find potential practical issues rather than showing it in only papers.

We agree on this and we will upload code whenever OpenReview allows uploading again. We’re confident that our experiments are reproducible and implementations are highly optimized (see Appendix E).

I cannot fully agree that MLP is a significant bottleneck, so speeding up attention is not important, as declared in lines 108 and 121. I partially agree that MLP is a bottleneck in a somewhat short context (less than 8K tokens).

Sorry for the confusion! We fully agree that attention is a significant bottleneck (in both short-context and especially long context) and SpecPrefill also helps skip the attention part as well, as shown in T2 (reviewer USDK’s comment) when we move on to 128K. What we hoped to convey was that there are many existing techniques to optimize attention in various regimes but optimizing MLP hasn't attracted too much attention, which is where SpecPrefill stands out. We revised the wordings to make this point clear.

I think we should try at least up to 128K tokens... The RULER and InfintieBench benchmark will be a great choice to test the effectiveness.

We understand the reviewer’s concern on longer context performance. We increase the RULER results in the original paper from 32K to 64K and observe consistent trends (T3 reviewer USDK’s comment). We will present the result on 128K for the camera-ready version and believe the same conclusion to hold (it is still currently running due to large models and limited cards).

Essential References Not Discussed.

We thank the reviewer for providing relevant works:

1. MInference, SnapKV, HiP Attention, SampleAttention, DuoAttention, and FlexPrefill are awesome work. However, they focus primarily on the attention module while keeping the MLP part intact. We include those in our camera-ready version for more context.
2. For MInference, we added additional experiments to show that SpecPrefill achieves better speedup with better quality.
3. MInference, SnapKV, and DuoAttention were originally in our draft and we added more context and discussion in our camera-ready version.

We truly agree with the concerns and feedback from the reviewer and hope our responses can help address some of those. With the added results, we are confident that SpecPrefill can be applied and benefit many real world applications and the contributions are novel and fully reproducible. We’d love to hear from the reviewer on further suggestions if any to make the paper better!

Summary

We thank the reviewer for the constructive and insightful feedback! We addressed all feedback and believe that incorporating all comments makes this paper significantly stronger:

Per-feedback detailed response

It would be interesting to see more exploration of how it compares with alternative attention sparsity methods such as Minference…

As the reviewer has mentioned, SpecPrefill not only reduces the computation on the attention but also on the MLP while sparse attention primarily focuses on the attention module alone. In order to provide a more comprehensive comparison between the two routes, we conducted experiments comparing ours against MInference as requested by the reviewer.

We adopt the pre-searched optimal pattern from MInference’s official repo for the 70B model and benchmark’s its quality on LongBench in T1 (see below). As we can see SpecPrefill can achieve 99.7% quality of MInference while keeping only 10% tokens and surpasses it when using 30% or more tokens. Because MInference is an attention based sparsity optimization, we benchmark its latency for a fair comparison in terms of acceleration and quality. In T2 (see below) where we vary the batch size x sequence length products from 128 * 4K to 4 * 128K, we found that SpecPrefill is particularly effective in efficiency compared to MInference with similar or better quality (getting 2.54x - 6.54x relative speedup compared to MInference). The main reason why MInference might fall short is due to the overhead in dynamic index building, which gets significant when the batch size is large. We do want to emphasize that Minference is an excellent work and its advantage starts to shine when the sequence becomes ultra long, at least 128K+ (or the ratio between sequence length and batch size gets large enough). Another trade-offs we want to mention between our methods and sparse attention approaches are the decision of whether we want to compute / store all the KV caches. Besides, MLP generally dominates the prefill part for large batch size + decent context lengths (at least within our tested 128K). This becomes essentially important based on the application scenario (e.g whether the prefill/decode is disaggregated and hence requires KV transmission). Overall, we believe that since both methods require no-training and can be switched on and off on the fly, they are complementary to each other and can hence meet a wider range of requirements when combined together.

With the added results, we are confident that SpecPrefill can be applied and benefit many real world applications and the contributions are novel and fully reproducible. We’d love to hear from the reviewer on further suggestions if any to make the paper better!

Newly added experiment results:

LongBench comparison with baselines (T1):

Latency comparison (T2):

RULER 64K (T3):

Summary

We sincerely thank the reviewer for the time and constructive feedback! We addressed all feedback and believe that it makes this paper significantly stronger (added experiment result numbers in reviewer USDK’s rebuttal). Specifically, we addressed the feedback by:

1. We moved the original main results for our RAG-baselines to the main body and created a table to compare all methods including two newly added baselines: LLMLingua and MInference. SpecPrefill achieves higher quality and better efficiency than all of them.
2. We agree with the application of multi-turn situations and believe that SpecPrefill will be able to solve it with the full knowledge of the speculator. We will add results in the camera-ready version (running the experiments).
3. We added more context and discussions for the essential references in our camera-ready version.

Per-feedback detailed response:

Not a single baseline was included in the experiment… Try comparing spec-prefill against prompt compression methods, and dynamic sparse attention approaches.

In the previous draft, we compared two RAG-based methods in the Appendix. We should have done a better job in making them more visible -- to do this, we will add a summary of baseline comparison in T1 (reviewer USDK's comment) in the main paper body.

In addition, we follow the advice from the reviewer by adding LLMLingua (text-level prompt compression, https://arxiv.org/abs/2310.05736) and MInference (sparse attention, https://arxiv.org/abs/2407.02490) to complement our original RAG baseline.

We evaluate all three baselines on LongBench, and the result is shown in T1 (reviewer USDK's comment). SpecPrefill outperforms RAG and LLMLingua by margins when keeping the same tokens from the prompt (i.e. same tokens ~= similar speedup). Comparing MInference requires a more careful design due to the fact that it is a sparse attention optimization that does not change the computation of the MLP and needs to compute all KVs. Therefore, we benchmark MInference’s prefill latency in T2 (reviewer USDK's comment) where we vary the batch size x sequence length products from 128 * 4K to 4 * 128K. We observe that MInference has substantial overhead when the batch size is large and only gets amortized when the sequence length gets ultra long (at least > 128K with small batch size < 4, this is also observed in paper https://arxiv.org/abs/2502.05431 and confirmed by the author in issue https://github.com/microsoft/MInference/issues/18). Overall, SpecPrefill can achieve 2.54x to 6.54x relative speedup compared to MInference with 99.7% quality and surpasses MInference’s quality when keep rate is at least 30%.

Try test spec-prefill on multi-turn long context benchmarks such as SCBench.

We fully agree that multi-turn tasks are important applications to test SpecPrefill. Unlike many prompt compression methods, SpecPrefill maintains the full KV for the speculator, which makes it possible to revive dropped tokens during the first round of conversation. This allows us to re-estimate token importance when a new turn comes in and let the base model fill-in the missing KVs as needed. We are currently running evaluations and will include results for the camera-ready version (due to new implementation and evaluation time, we can’t finish it during the rebuttal window). In addition, here’s our detailed description of what we’re running:

We outline the high-level algorithm here:

1. Check if it is the first turn. If so we just call our standard SpecPrefill and return the results. If not, move to step 2.
2. Let the speculator first process the new N-token context and estimate important tokens just like in the first prefill.
3. Identify missing KVs based on slot-mapping and position information.
4. For positions missing KVs, we simply keep those ids.
5. For the base model, we will recompute the KVs if some ids are not in the cache. Return the new results.

The total number of KVs that will be computed and stored will not exceed the context length, which means there is no waste of “recomputation”.

Nonetheless, we’d like to mention that our initial starting goal for SpecPrefill is to tackle large-batch offline (no further interaction) prefill acceleration that targets a large traffic of requests such as bulk document summarization, QA, etc.

Essential reference not mentioned.

All of these are awesome add-ons and have been included in appropriate places for reference.

We are confident that SpecPrefill can be applied and benefit many real world applications and the contributions are novel and fully reproducible. We’d love to hear from the reviewer on further suggestions if any to make the paper better!