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

Summary

This paper presents a sparse calculation technique for the attention mechanism in long-context large language models during the pre-filling stage. Specifically, the technique builds on the authors' observation of three patterns in the attention map: the A-shape pattern, the Vertical-Slash pattern, and the Block-Sparse pattern. In the proposed method, the authors first determine the optimal pattern for each attention head offline and then dynamically decide the hyperparameters for each pattern (e.g., sparse indices) on-the-fly. It is worth noting that the A-shape pattern is more regular in terms of the distribution of non-zero indices compared to the Vertical-Slash and Block-Sparse patterns. To provide a low-cost estimation of the non-zero indices in the Vertical-Slash and Block-Sparse patterns, the authors predict these indices using the matrix multiplication results between the last query vector and the key matrix and the matrix multiplication results between a mean-pooled query matrix and the key matrix, respectively. Results on three different pre-trained LLMs and four datasets indicate that the proposed method effectively maintains information from long prompts while reducing pre-filling latency compared to current state-of-the-art methods.

Strengths

1. Impressive results: Based on the reported experimental results, the proposed method can reduce the pre-filling latency for a 1M context by ten times without compromising accuracy. 2. Well-motivated: I appreciate the analysis and high-quality figures in Section 2, which explicitly illustrate the key takeaways for readers and effectively motivate the proposed method. 3. Comprehensive summary in the related works section.

Weaknesses

1. How to ensure that the observations on the three LLMs in this paper are generalizable to different LLMs? Since the entire method is built on the observation of the attention patterns in these three LLMs, the generalizability of this observation is crucial for the robustness of this work. For example, the block-sparse pattern is the most dynamic among the three. What if there are even more dynamic attention maps in other LLMs with specific inputs? How is the search space determined, and how general is it? The offline Kernel-Aware Optimal Sparse Pattern Search determines the pattern for each head offline, and the pattern does not change during inference regardless of the input. Is it always true that the attention pattern for a specific head remains nearly the same? 2. Lack of evaluation on larger models: For longer contexts, larger models seem to be a better option to ensure quality. However, the models evaluated in this paper are relatively small (<10B parameters). Thus, it is unclear whether the improvements still hold for larger models. 3. Lack of discussion on usage for the generation/decoding stage: Pre-filling is important because it determines the first token latency, but generation/decoding is also important because it determines the throughput. It would be beneficial to add some discussion on whether the observations and the proposed method are still applicable for the generation/decoding stage. For example, while the A-shape pattern might still exist, the sparse indices approximation in the vertical-slash pattern might not work for generation/decoding.

Questions

1. What is the dataset of input data used for the analysis in Section 2.2? How was it chosen? 2. In Figure 3(c), why does the block-sparse pattern perform so poorly in attention compared to the A-shape and Vertical-Slash patterns? Based on my understanding, the block-sparse pattern is a more general and dynamic case for the A-shape and Vertical-Slash patterns, as shown in Figure 4. Could the block size be an important factor causing this issue? 3. In Figure 3(c) and the corresponding description, I suggest explicitly mentioning which Top-K method is used. Based on my understanding, the Top-K method involves selecting the top-K tokens for all heads, while the other methods select the top-K tokens for each head. Presenting them together may confuse readers into thinking that Top-K means the top-K for each head. 4. What does "real FLOPs" mean in Line 137? FLOPs itself is a "conceptual estimation" compared to real measured latency. I believe the authors may want to refer to the FLOPs after discarding the zero values in the attention pattern. If so, please be more specific. 5. The LLaMA-3-8B-Instruct-262k is linked to LLaMA-3-70B-Instruct-262k in Line 176. Please clarify which model is used here. 6. There are minor consistency issues, such as StreamLLLM and StreamingLLM.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discussed the limitation and social impact of their work.

Summary

A key challenge for LLM inference with processing long context lengths is time-to-first token for long prompts. This paper introduces a sparse attention method designed to accelerate prefill with long context lengths. They utilize a strategy that incorporates three different types of sparsity (A-shape, vertical-slash, and block-sparse). They calibrate for which sparsity pattern to apply to different heads, thereby adapting to the varying sparsity patterns across LLM heads. Their method is training-free and obtains significant latency gains in terms of time to first token.

Strengths

- Time to first token is a significant problem that is under-addressed in existing KV cache compression works, as many applications have long input prompts and short generation lengths - Their approach is able to adapt dynamically to different inputs for certain attention patterns (which have different sparsity patterns), as well as to handle differing behaviors across attention heads - They provide efficient online kernel implementations for each sparsity pattern (both for constructing the sparsity pattern for the two dynamic patterns as well as for sparse computation) - They observe significant prefill speedups attained with minimal accuracy degradation - They provide detailed evaluation across a range of long context tasks, as well as ablation for each type of sparsity pattern (and additional analysis showing the distribution of sparsity patterns across attention heads, etc.)

Weaknesses

- Their approach is heuristic-based (as it is based on fixed types of attention patterns), which may not generalize to new models that are released (some of these patterns could be specific to models that employ particular positional encodings or other model architecture features, which may limit generalizability) - The offline kernel search assumes that each head will play the same role regardless of the target task (or alternatively that calibration data will always be available for the target task in order to pre-determine which general attention pattern to use and what budget to allocate to that head) - The writing in some sections is poor (eg. end of Section 2)

Questions

- Do you have any analysis that shows that the attention head patterns are the same regardless of the input task (or that their search converges to the same solution for the same head across different tasks)?

Rating

7

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes

Summary

Targeting at the time-consuming profiling stage of long contexts, this paper proposed an efficient prefilling stage spare attention mechanism. It's based on the observation to common attention patterns. The proposed method can be integrated into most existing LLMs, such as LLama3 and Yi-9B. It achieves good performance compared to the original model while significantly reduce the computation times.

Strengths

1: The motivation is clear and the paper is easy to understand. 2: The proposed method seems simple yet effective. 3: The experiments are comprehensive.

Weaknesses

1: It’s tricky to claim that “we found the three patterns”. Those patterns have been found many years ago, since the years of Bert / transformer models were proposed. Even the idea of leveraging such sparsity of attention has been invented many times for different transformer based models. Some other major concerns can be found in the "Question" section.

Questions

1. Is the slash-vertical pattern distinguishable from lambda-shaped patterns? By my understanding, lambda-shaped patterns are just a special case of the slash-vertical pattern. Perhaps if you replace all lambda-shaped patterns with the slash-vertical pattern, it would still work well. 2. Could you please do a comparison between the proposed MInference and H2O[1]? From my understanding, without considering the real efficiency, it’s hard for a method relying on attention scores/similarity to obtain significantly better performance than H2O. 3. From my own experience, when the tasks become more challenging, the performance of efficient attention methods might show quite different effectiveness. For example, many long context methods cannot pass the passkey retrieval test with a 100-digit passkey, while with a 5-digit passkey, most methods work perfectly. Could you please construct experiments on more challenging tasks? Especially on those most recently proposed ones, for example: [2][3]. 4. Also, Llama-3 is much more robust to perturbation and error accumulation. Could you please construct experiments with some weaker LLMs, such as Mistralv0.2, Llama2, and gemma-1.1? References: [1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models [2] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches [3] BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack

Rating

7

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes

Summary

This paper proposes MInference, a method to accelerate the pre-filling stage for long-context LLM generation. The key method leveraged by MInference is dynamic sparsification, which consists of three sparse patterns observed in attention matrices: the A-shape pattern, the Vertical-Slash pattern, and the Block-Sparse pattern. A sparse pattern search method is developed to minimize the sparsification error for each attention. MInference has been tested across various LLMs and benchmarks, demonstrating that it significantly accelerates LLM generation in long-context settings with no performance drop.

Strengths

- The paper is well-written and well-motivated. - Studying how to accelerate LLM generation for long-context settings is of great practical importance. - The experimental results of the proposed MInference method are promising.

Weaknesses

- The experimented model scales are relatively small (e.g., LLaMA-3-8B and Yi-9B) and are only constrained to dense LLMs (i.e., not MoE). - It is not clear how general the observed sparse patterns are across LLMs of various scales. - The implementation details on how the proposed MInference interacts with existing CUDA kernels (e.g., flashattention) are not very clear.

Questions

- How does MInference scale up to LLMs with larger sizes, e.g., 70B? And how does MInference perform for MoE-based LLMs, e.g., Mixtral 8x7B? - How general are the observed sparse patterns across LLMs? - Is MInference compatible with other inference engines, e.g., vLLM? - Are the Vertical-Slash Head and Block-Sparse Head part of CUDA kernels in MInference, or does MInference leverage other implementations of those sparse head kernels?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations of the proposed method have been thoroughly discussed. And there is no potential negative societal impact of this work.

Thanks for the rebuttal!

Thank you to the authors for the detailed rebuttal! I truly appreciate your efforts in putting everything together within such a short period. I have raised my score to 7. One minor issue: For this statement: "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." --> I understand that the Slash pattern’s block size is smaller than that of block-sparse, i.e., 1x64 vs. 64x64. However, the logical connection between "significant computational waste" and "a lower attention recall at the same sparsity rate" is not entirely clear to me. Specifically, I believe that "significant computational waste" would lead to lower efficiency rather than lower recall.

Rebuttal Response

I appreciate the author's detailed responses addressing my comments. Their response sufficiently justifies the claims that their approach generalizes to a range of LLMs, as well as that particular heads play a similar role regardless of the task. I will therefore keep my initial score.

Response by Authors

Thank you for your patience as additional experiments took some time to complete. And thank you for your insightful suggestions. Your input has been instrumental in enhancing the completeness of our paper, and we will include the relevant discussions in the updated version. 1. _"...show some to supporting results on standard short context tasks..."_ We add the results in GSM8K with 9-shot[1], which has an average length of 2.8K tokens, as shown in Table 4. The table uses a sparsity equivalent to 128 + 128 A-shape FLOPs, with an average sparsity of 84% in the kernel. Information redundancy is lower in short contexts than in long contexts, making them more sensitive to compression. As can be seen, our method experiences a slight drop in performance but remains close to the results of full attention. Compared to other baselines, there is a clear and uniform improvement. This demonstrates that, in addition to complex retrieval or knowledge-intensive tasks, our method also effectively preserves the reasoning capabilities of LLMs. |Methods|GSM8K| |-|-| |LLaMA-3-8B-262K|63.8| |w/o few-shot|37.2| |StreamingLLM|53.2| |InfLLM|56.4| |**MInference**|61.7| Table 4. Performance of different methods on GSM8K with 9-shot ICL using LLaMA-3-8B-262K. [1] Complexity-Based Prompting for Multi-Step Reasoning, ICLR 2023. However, we would like to emphasize that our method is primarily designed for long-context scenarios. In short contexts, the latency proportion of attention is not high, and the overhead from dynamic sparse index building becomes more significant. Furthermore, we already have tested our method on extremely challenging long-context LLMs benchmarks, especially the RULER benchmark, which includes Multi-hop Tracing and Multi-Needle tasks that require reasoning abilities to complete. 2. _"...provide this ratio for some real examples?"_ Apologies for the oversight in my previous response. As illustrated in Fig.12 of our paper, the average block-wise sparsity ratio of "78%-90% around 100k" is 88%, which allows MInference to achieve a 3x-4x speedup in single attention operations and a 1.8x speedup on an end-to-end basis. The variability in sparsity across different inputs mainly stems from the density of the slash lines, which causes differences sparsity in the blocks within the kernel. Thank you for the correction. The average number of tokens in LongBench is 12k; we mistakenly cited data from InfLLM. We have calculated the average sparsity rate in the kernel for LongBench, which ranges from 55% to 92%, depending on the number of tokens in the task. For example, tasks with longer contexts like narrativeqa have an average sparsity rate above 90%, while shorter context tasks like lcc are closer to 60%. Moreover, in terms of performance, our method outperforms other methods in the baselines and is closer to the full attention performance.