A single card A100 can achieve millions of token inferences, which is 10 times faster. This is Microsoft’s official large model inference acceleration.

This Microsoft research allows developers to process over 1M input text 10x faster on a single card machine.

Large Language Model (LLM) has entered the era of long context processing, and the context window it supports has soared from the previous 128K to the 10M token level.

However, due to the quadratic complexity of the attention mechanism, it may take several minutes for the model to process the input prompt (i.e., the pre-population stage) and start generating the first token. As a result, it takes too long to generate the first token, which seriously affects the user experience, which also greatly limits the widespread application of long-context LLM.

For example (as shown in Figure 2a), when serving LLaMA-3-8B on a single machine equipped with A100, if it prompts that there are 300,000 tokens, the model will take 6 minutes to complete the preset. During the pre-filling phase, if the prompt increases to 1 million tokens, this number will increase to 30 minutes.

The overhead of self-attention computation accounts for more than 90% of the total prefill delay, making it the main bottleneck when LLM handles long contexts. Existing accelerated pre-population methods often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs.

In order to solve the above problems, researchers from Microsoft and the University of Surrey proposed a sparse calculation method designed to accelerate pre-filling in long sequence processing: MInference (Milliontokens Inference).

• Paper address: https://arxiv.org/pdf/2407.02490
• Paper homepage: https://hqjiang.com/minference.html
• Paper title: MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

MInference can be directly applied to existing LLMs without modifications to pre-training settings or additional fine-tuning.

By testing various downstream tasks (including InfiniteBench, RULER, PG-19 and Needle In A Haystack) and models (including LLaMA-3-1M, Yi-200K, GLM-4-1M, Phi-3- 128K and Qwen2-128K), experiments demonstrate that MInference can effectively reduce pre-populated inference latency on the A100 by up to 10x while maintaining accuracy.

                                                                                                                                                                                   ^{Using MInference 1.0, long context LLM (such as LLaMA-3-8B-1M, GLM-4-1M) achieves a 10x improvement in inference speed on a single A100, and is more accurate.}

Method introduction

The authors proposed MInference, a name that reflects their ambition to achieve million token inference on an A100 machine.

MInference is a training-free and efficient method for the pre-population stage of long-context LLM based on dynamic sparse attention.

Researchers believe that attention, especially in long contexts, is sparse and dynamic, that is, in different inputs, the sparse pattern is very different. This dynamic sparsity exhibits three unique spatial aggregation patterns that apply to all inputs: A-shape, Vertical-Slash, and Block-Sparse.

MInference first determines the best dynamic sparse pattern offline for each head using a kernel-aware sparse pattern search algorithm, as shown in Algorithm 1. During inference, it dynamically approximates the dynamic sparsity index based on the pattern of the head, as shown in Algorithms 2, 3. Finally, the authors use optimized GPU kernels to perform efficient dynamic sparse attention computation, significantly reducing the pre-population stage latency of long-context LLM.

For example, for the "vertical - diagonal line" mode, the author first uses the attention calculation between the last Q and K to estimate the best index of the vertical line and diagonal line. They then leveraged the dynamic sparse compilers PIT and Triton to build a vertical-slash FlashAttention kernel that accelerates attention computation. For the A-shaped, vertical-slope, and block-sparse patterns, the authors first use mean pooling of Q and K in the attention computation. Using mean pooling and the exchange property of MatMul, the blocky-sparse index can be estimated. They then used the Triton building block sparse FlashAttention kernel to accelerate attention computation. See Appendix C.4 and code for detailed implementation of the kernel.

Evaluation results on long context benchmarks

The authors tested MInference in a range of scenarios, including QA, encoding, retrieval-based tasks, multi-hop QA, summarization and math tasks . The RULER benchmark includes several complex multi-hop or multi-needle tasks, effectively reflecting the actual context window size of LLM. As shown in Table 1, MInference effectively preserves the actual context window handling capabilities of LLM and even slightly extends the actual context window size to 32K.

The authors also tested MInference on a wider range of tasks using InfiniteBench with an average token length of 214K, as shown in Table 2. Compared to the SoTA baseline, MInference consistently maintains good performance across all tasks. Notably, in more challenging retrieval tasks such as the KV retrieval task, all baselines failed to make accurate predictions, with an accuracy below 1.2%. However, MInference successfully retains the ability to handle dynamic KV pair retrieval.

To further evaluate the performance of different context lengths and different positions of key information in the prompt, the authors tested various models and methods using the "needle in the haystack" task. As shown in Figure 1, MInference performs well under different models, context windows, and cue information locations, and its performance remains the same or even slightly improved compared to the original model. In the case of LLaMA-3-8B and GLM-4-9B-1M, MInference achieves fully green performance in context windows up to 1M. In contrast, StreamingLLM and InfLLM performance drops below 20% in the middle of the prompt, even in a 70K context window.

The authors also tested MInference on a language model task using PG-19, which included up to 100k tokens.如图 2 所示，MInference 有效地保持了 LLaMA-3-8B 和 Yi-9B-200K 的困惑度，而所有基线都出现了不同程度的困惑度下降。此外，与标准的 StreamingLLM 相比，使用膨胀和步长配置的 StreamingLLM 更好地保持了困惑度性能。

延迟和内核中的稀疏模式 

图 3 展示了本文提出的三种注意力模式以及 FlashAttention 的微基准测试结果。可以看出，Vertical-Slash 是三种模式中最慢的，但在 1M 上下文窗口下，相比 FlashAttention 仍然实现了 13 倍的加速。

图 4 展示了 Vertical-Slash 头部内核中的稀疏索引。垂直线通过 PIT FlashAttention 使用 1x64 块计算，而斜线通过块级 FlashAttention 使用 64x64 块计算。

^{参考链接：https://hqjiang.com/minference.html}

The above is the detailed content of A single card A100 can achieve millions of token inferences, which is 10 times faster. This is Microsoft’s official large model inference acceleration.. For more information, please follow other related articles on the PHP Chinese website!