MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

1. "...generalizable... built on the observation of the attention patterns..."

Thank you for your question. We address it from two angles: the generalization of MInference and the relative stability of dynamic sparse attention patterns across different examples.

For the generalization of MInference,

• To demonstrate the generalizability of MInference, we tested it on most open-source long-context LLMs, including LLaMA-3-8B/70B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K (see general response PDF). MInference consistently achieves good performance and acceleration across models with different pre-training data, training pipelines, RoPE structures, extended long-context methods, and sizes.
• Additionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are not only present in GPT-like LLMs but also in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that "induction heads" exhibit a pattern similar to the "vertical and slash" pattern (see [4]).
• By design, MInference accounts for different dynamic sparse attention patterns, which lends it a certain degree of generalization capability.

[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. [3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024. [4] A Mathematical Framework for Transformer Circuits, 2021.

For the relative stability of dynamic sparse attention patterns across different examples,

• Visualization: Fig.3(a) shows the visualization of the sparse attention pattern distribution of the same head across different tasks and examples, demonstrating good consistency.
• Attention Recall: Fig.3(c) shows the average recall rate of different tasks and examples within the searched pattern as the compression rate changes. The searched pattern achieves a higher recall rate across different tasks.
• Downstream tasks performance: Using a configuration searched from a single example, we achieve performance close to full attention across different scenarios and benchmarks, further indicating the stability of dynamic sparse attention patterns across various tasks and examples.

2. "What if there are even more dynamic attention maps in other LLMs with specific inputs"

Given the extreme sparsity of attention in long-context scenarios, where the significant attention weights are concentrated in a few blocks, even if the graph differs from the vertical and slash pattern and the A-shape, it can still be covered by the block-sparse pattern. Additionally, our experiments in Question 1 demonstrate the ubiquity of this dynamic sparse pattern in current long-context LLMs and the effective of MInference.

3. "How is the search space determined, and how general is it?"

The search space is determined by a specific sparsity rate and adjusted according to the actual FLOPs required for different patterns in the kernel. Based on testing, the search space obtained is quite consistent across different tasks and examples.

4. "Lack of evaluation on larger models"

Thank you for your suggestion. We have included results for LLaMA-3-70B-1M in the supplementary PDF. MInference achieves performance close to or even better than full attention, especially compared to baselines.

Table 1. Performance of different methods with different base models on InfiniteBench.

5. "Decoding"

Thank you for your suggestion. However, extending to decoding would require significant system work for CPU offloading, which we have earmarked for future work. Nonetheless, our preliminary experiments indicate that the extrapolation generality of these patterns is also very good, as shown in:

Table 3. Performance of different methods on InfiniteBench using GLM-4-9B-1M.

6. "Dataset used in Sec 2.2"

In Section 2.2, we used data from InfiniteBench to create Fig.3, including summarization, QA, Math, and retrieval tasks. Both (b) and (c) are results averaged after randomly selecting 10 examples from each subset.

7. "Why does the block-sparse pattern perform so poorly?"

This is mainly because different patterns use different block sizes in the kernel. For example, the Vertical and Slash pattern can use a 1x64 block to process vertical lines instead of 64x64 in block-sparse. If this head pattern is processed using block-sparse, it would result in significant computational waste, leading to a lower attention recall at the same sparsity rate.

8. "I suggest explicitly mentioning which Top-K method is used"

We apologize for any confusion caused by our writing. Yes, your understanding is correct; the Top-K mentioned in Fig.3(c) refers to token-level Top-K. We will add all corresponding granularities of TopK and rewrite Section 2.2 for better comprehension in the next version.

9. "Real FLOPs"

Real FLOPs refer to the FLOP values in the kernel after excluding non-zero calculations. Thank you for your suggestion; we will clarify this meaning in the next version.

10. Error link and typos

Thank you for your detailed review. We will correct these issues in the next version of our paper.

1. "...which may not generalize..."

• To demonstrate the generalizability of MInference, we tested it on most open-source long-context LLMs, including LLaMA-3-8B/70B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K (see general response PDF). MInference consistently achieves good performance and acceleration across models with different pre-training data, training pipelines, RoPE structures, extended long-context methods, and sizes.
• Additionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are not only present in GPT-like LLMs but also in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that "induction heads" exhibit a pattern similar to the "vertical and slash" pattern (see [4]).
• By design, MInference accounts for different dynamic sparse attention patterns, which lends it a certain degree of generalization capability.

[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. [3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024. [4] A Mathematical Framework for Transformer Circuits, 2021.

2. "... assumes that each head will play the same role regardless of the target task..."

Indeed, based on our experiments and observations, these dynamic sparse patterns are relatively fixed across different tasks and examples. This is corroborated by our experimental results, where a configuration searched from a single example can achieve performance close to full attention across various scenarios and benchmarks.

3. "The writing in some sections is poor (e.g., end of Section 2)"

We apologize for any confusion caused by our writing and will rewrite the relevant content, particularly Section 2.2, to improve readability.

4. "...analysis that shows that the attention head patterns are the same regardless of the input task..."

Thank you for your suggestion. We explain the relative stability of head-level dynamic sparse attention patterns from several perspectives:

• Visualization: Fig.3(a) shows the visualization of the sparse attention pattern distribution of the same head across different tasks and examples, demonstrating good consistency.
• Attention Recall: Fig.3(c) shows the average recall rate of different tasks and examples within the searched pattern as the compression rate changes. The searched pattern achieves a higher recall rate across different tasks.
• Downstream tasks performance: Using a configuration searched from a single example, we achieve performance close to full attention across different scenarios and benchmarks, further indicating the stability of dynamic sparse attention patterns across various tasks and examples. We will rewrite the paper to highlight this point and thank you again for your helpful and insightful comments.

Table 2(a). Performance of different methods on RULER.

Thank you for your thorough review, and we apologize for any confusion caused.

1. "Claim that 'we found the three patterns'"

Thank you for your critique. We will revise the wording accordingly. Indeed, we discussed the importance of sparse attention in related works and how it inspired our research. We want to emphasize our understanding of the dynamic nature of these sparse patterns, which led us to propose the training-free MInference method.

2. "Lambda is part of vs?"

Yes, the A-shape pattern is a specific type of "vertical and slash" pattern. However, the A-shape head exhibits a more fixed sparsity pattern across different tasks and examples. Due to its static nature, it does not require online sparse index building and allows for more extreme sparse kernel optimization, resulting in better inference acceleration. We differentiate these three patterns based on kernel speedup and spatial locality.

3. "Compare with H20"

In fact, MInference and H2O optimize two different problems in long-context LLMs inference: prefilling stage latency and KV cache storage, respectively. These methods are orthogonal. Furthermore, we applied MInference with SnapKV in our paper (Table 5 in paper), which is a SoTA KV cache compression method superior to H20. MInference can be used in conjunction with KV cache compression methods to achieve even better performance.

4. "More challenging tasks"

Thank you for your suggestion. However, I would like to clarify that we have tested on very challenging long-context benchmarks, including RULER, InfiniteBench, Needle In A Haystack, and Language Model, with lengths ranging from 128k to 1M tokens.

• RULER includes subtasks such as Multi-Needle and Multi-hop Tracing and is one of the most challenging long-context benchmarks, sharing a similar tasks with BABILong and a similar benchmark ranking.
• We have additionally conducted evaluations on the KV cache compression benchmark[5], including LongBench (refer to Table 2(b) below) and Needle In A Haystack (see Figure 8 in the paper and Figure 1 in the general response PDF), showing that MInference does not negatively impact performance.
• InfiniteBench includes many challenging long-context tasks. For example, KV retrieval task requires LLMs to recall a 36-bit value from several random 36-bit Key-Value pairs by retrieving the value corresponding to a random 36-bit key, which is not as simple as the 5-digit passkey task mentioned in the comments.

Overall, our tests have included similar or even more challenging long-context tasks, effectively reflecting the minimal impact of MInference on the capabilities of long-context LLMs.

Table 2(b). Performance of different methods on LongBench.

[5] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches.

5. "Weaker models"

Besides testing on the powerful LLaMA-3-8B-1M, we also conducted tests on the Yi-9B-200K model, which has an effective context size of only 8k in RULER, showing that MInference can achieve performance close to Full Attention without changing the effective context size. Moreover, we have supplemented our results with GLM-4-9B-1M, Phi-mini-128K, and Qwen2-7B-128K (see general response PDF). Notably, the latter two models have weaker long-context capabilities, with effective context sizes of 4k and 8k in RULER, respectively. They show a significant performance drop in the Needle In A Haystack task when exceeding 100K tokens, but using MInference, they can achieve close or even better performance. We will include these results in future versions. Additionally, the weak LLMs mentioned by the reviewer have context windows smaller than 32K and cannot be tested on benchmarks like InfiniteBench. We will adapt more long-context LLMs in the future.

I don’t want to seem too picky, but I believe the most important thing in research is to be sincere and I am really annoyed by nowadays trends of over-claiming. Your method's contributions are already sufficient, even if it works well only in long texts. No method is perfect.

I have no idea about why you keep trying to emphasize that your method works well on short tasks; such sparsity can be used overall. If you do want to show that MInference has minimal impact on all aspects of the model, then more rigorous testing on standard tasks should be conducted.

Back to your newly updated results. Thank you for the experiments. However, I must admit, I find your gsm8k setting very odd. I have never seen 9shot; most gsm8k experiments are either 5shot, like [1, 2], or 8shot, like [3,4]. And the performance of llama3 seems like it should be better. Using more shots means increase of redundancy.

Given the effectiveness on long texts, at this time, I will not change my score. But I am very much looking forward to more responses from the authors to address my concerns.

[1]QWEN2 TECHNICAL REPORT, https://arxiv.org/pdf/2407.10671 [2]Gemma 2: Improving Open Language Models at a Practical Size, https://arxiv.org/pdf/2408.00118 [3]https://llama.meta.com [4]DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, https://arxiv.org/pdf/2405.04434

We are very grateful for your critique and corrections. We fully agree with and acknowledge your exposition of MInference's limitations in short-context scenarios. In fact, we have never claimed that MInference is suitable for use in short-context scenarios, and we discussed these limitations in Appendix A. We will rewrite this section in the next version of paper to highlight the discussion on limitations in short-context scenarios.

Here we provide some details to support this content:

1. MInference cannot achieve any acceleration in short-context scenarios (prompts < 10k tokens), and due to the dynamic index building, it can be 30%-50% slower than FlashAttention-2. Moreover, since the latency proportion of Attention is not high in short-context, even if the single Attention module has a speedup, the end-to-end speedup is minimal.
2. The sparsity rate of Attention in short-context scenarios is significantly lower than in long-context scenarios. For example, controlling for the same sparsity rate (96.8%) as in Fig.2(b) of the paper, the Attention recall drops from 96.4% at 128K to 89.8% at 4K.
3. If the sparsity ratio from long-context scenarios is used in short-context, MInference causes a noticeable performance degradation. For example, as shown in Table 5, the accuracy of qasper (one of single-document QA tasks in LongBench) drops from 29.64% to 8.04%.

Table 6. Performance of different sparsity rates on qasper in LongBench using LLaMA-3-8B-262K.

"...gsm8k setting being very odd..."

We appreciate your critique and correction. We did not attempt any cherry-picking; we followed the experimental setup for GSM8K as per [1,2]. We will review the relevant experimental settings, align the 8-shot results, and update them in next versions of the paper.

[1] Complexity-Based Prompting for Multi-Step Reasoning, ICLR 2023.

[2] Chain-of-Thought Hub: Measuring LLMs' Reasoning Performance, Github.

Thank you once again for your efforts in reviewing our work. Your input is very helpful for the further improvement of our research, and we will update the relevant content in the next version of our paper.

1. "For larger models and MoE"

Thank you for your suggestion. We have added experimental results for LLaMA-3-70B-1M, as shown in Table 1, where MInference maintains excellent performance in larger models, significantly surpassing StreamingLLM and InfLLM[1], and nearly matching full attention. Due to resource constraints, we could not provide MoE results during the rebuttal period, but based on our experience, modifications to FFN do not affect the sparse attention patterns. We will include MoE results in future versions.

Table 1. Performance of different methods with different base models on InfiniteBench.

[1] InfLLM: Unveiling the Intrinsic Capacity of LLMs for Understanding Extremely Long Sequences with Training-Free Memory, 2024.

2. "Generality ability across models"

We tested MInference on a variety of long-context LLMs, including LLaMA-3-8B-1M, Yi-9B-200K, GLM-4-9B-1M, Phi-3-mini-128K, and Qwen2-7B-128K, all of which maintained performance close to full attention (see general response PDF). Additionally, we observed that three sparse attention patterns, especially the vertical and slash patterns, are present in BERT[2], T5 (both Encoder and Decoder), and MLLM[3]. We also found that "induction heads" exhibit a pattern similar to the "vertical and slash" pattern (see [4]). In conclusion, we believe MInference is highly generalizable across different structures, sizes, pre-training data, training pipelines, RoPE structures, and extended long-context methods.

[2] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021. [3] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024. [4] A Mathematical Framework for Transformer Circuits, 2021.

3. "Compatible with other inference engines, e.g., vLLM"

Yes, we have implemented MInference within vLLM. Our method only requires replacing the corresponding FlashAttention kernel and does not affect the scheduling of vLLM PageAttention or similar features. We will open-source the implementation after the paper review.

4. "CUDA Kernel"

We have utilized Dynamic Sparse Compiler PIT, Triton, and FlashAttention to implement CUDA Kernels for Vertical and Slash Head and Block-sparse Head, ensuring transferability and support for long-context LLMs inference across different GPUs and CUDA versions. Details can be found in Appendix C.4. We have submitted the relevant code with our submission and will open-source the implementation after the review.