{"id":1053873,"date":"2024-07-04T07:55:27","date_gmt":"2024-07-04T14:55:27","guid":{"rendered":"https:\/\/www.microsoft.com\/en-us\/research\/?post_type=msr-project&#038;p=1053873"},"modified":"2024-08-21T23:25:22","modified_gmt":"2024-08-22T06:25:22","slug":"minference-million-tokens-prompt-inference-for-long-context-llms","status":"publish","type":"msr-project","link":"https:\/\/www.microsoft.com\/en-us\/research\/project\/minference-million-tokens-prompt-inference-for-long-context-llms\/","title":{"rendered":"MInference: Million-Tokens Prompt Inference for Long-context LLMs"},"content":{"rendered":"<section class=\"mb-3 moray-highlight\">\n\t<div class=\"card-img-overlay mx-lg-0\">\n\t\t<div class=\"card-background  has-background-catalina-blue card-background--inset-right\">\n\t\t\t<img loading=\"lazy\" decoding=\"async\" width=\"1792\" height=\"1024\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light.webp\" class=\"attachment-full size-full\" alt=\"MInference Background\" style=\"\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light.webp 1792w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light-300x171.webp 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light-1024x585.webp 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light-768x439.webp 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light-1536x878.webp 1536w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/DALL\u00b7E-2024-07-04-22.47.18-A-modern-high-tech-office-environment-with-numerous-employees-focused-on-their-computers.-The-room-is-large-and-spacious-illuminated-by-natural-light-240x137.webp 240w\" sizes=\"(max-width: 1792px) 100vw, 1792px\" \/>\t\t<\/div>\n\t\t<!-- Foreground -->\n\t\t<div class=\"card-foreground d-flex mt-md-n5 my-lg-5 px-g px-lg-0\">\n\t\t\t<!-- Container -->\n\t\t\t<div class=\"container d-flex mt-md-n5 my-lg-5 \">\n\t\t\t\t<!-- Card wrapper -->\n\t\t\t\t<div class=\"w-100 w-lg-col-5\">\n\t\t\t\t\t<!-- Card -->\n\t\t\t\t\t<div class=\"card material-md-card py-5 px-md-5\">\n\t\t\t\t\t\t<div class=\"card-body \">\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\t\n\n<h1 class=\"wp-block-heading\" id=\"minference\">MInference<\/h1>\n\n\n\n<p><strong><em>Million-Tokens<\/em><\/strong> Prompt Inference for Long-context LLMs<\/p>\n\n\t\t\t\t\t\t<\/div>\n\t\t\t\t\t<\/div>\n\t\t\t\t<\/div>\n\t\t\t<\/div>\n\t\t<\/div>\n\t<\/div>\n<\/section>\n\n\n\n\n\n<div style=\"height:30px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<div class=\"wp-block-columns is-layout-flex wp-container-core-columns-is-layout-1 wp-block-columns-is-layout-flex\">\n<div class=\"wp-block-column is-layout-flow wp-block-column-is-layout-flow\">\n<figure class=\"wp-block-video\"><video controls src=\"https:\/\/github.com\/microsoft\/MInference\/assets\/30883354\/52613efc-738f-4081-8367-7123c81d6b19\"><\/video><\/figure>\n<\/div>\n\n\n\n<div class=\"wp-block-column is-layout-flow wp-block-column-is-layout-flow\">\n<h2 class=\"wp-block-heading\" id=\"what-is-minference\">What is MInference?<\/h2>\n\n\n\n<p><strong>MInference 1.0<\/strong>&nbsp;leverages the dynamic sparse nature of LLMs&#8217; attention, which exhibits some static patterns, to speed up the pre-filling for long-context LLMs. It first determines offline which sparse pattern each head belongs to, then approximates the sparse index online and dynamically computes attention with the optimal custom kernels. This approach achieves up to a&nbsp;<strong>10x speedup<\/strong>&nbsp;for pre-filling on an A100 while maintaining accuracy.<\/p>\n\n\n\n<p><strong>MInference 1.0<\/strong> was presented at the Microsoft booth and ES-FoMo at ICML&#8217;24.<\/p>\n\n\n\n<div class=\"wp-block-buttons is-layout-flex wp-block-buttons-is-layout-flex\">\n<div class=\"wp-block-button\"><a data-bi-type=\"button\" class=\"wp-block-button__link wp-element-button\" href=\"https:\/\/www.microsoft.com\/en-us\/research\/publication\/minference-1-0-accelerating-pre-filling-for-long-context-llms-via-dynamic-sparse-attention\/\">Read the paper<\/a><\/div>\n\n\n\n<div class=\"wp-block-button is-style-fill\"><a data-bi-type=\"button\" class=\"wp-block-button__link wp-element-button\" href=\"https:\/\/huggingface.co\/spaces\/microsoft\/MInference\" target=\"_blank\" rel=\"noreferrer noopener\">View the demo<\/a><\/div>\n\n\n\n<div class=\"wp-block-button is-style-fill-github\"><a data-bi-type=\"button\" class=\"wp-block-button__link wp-element-button\" href=\"https:\/\/github.com\/microsoft\/minference\" target=\"_blank\" rel=\"noreferrer noopener\">Get the code<\/a><\/div>\n<\/div>\n<\/div>\n<\/div>\n\n\n\n<hr class=\"wp-block-separator has-alpha-channel-opacity\" \/>\n\n\n\n<figure class=\"wp-block-image aligncenter size-large\"><img loading=\"lazy\" decoding=\"async\" width=\"1024\" height=\"584\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px-1024x584.png\" alt=\"MInference slide onepager: MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention\" class=\"wp-image-1068771\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px-1024x584.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px-300x171.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px-768x438.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px-240x137.png 240w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_Pmx510-_703vitXUtP-K4_1400px.png 1400w\" sizes=\"(max-width: 1024px) 100vw, 1024px\" \/><\/figure>\n\n\n\n<hr class=\"wp-block-separator has-alpha-channel-opacity\" \/>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"insights\">Insights<\/h3>\n\n\n\n<ol class=\"wp-block-list\">\n<li>Attention, especially in long-context scenarios, is&nbsp;<strong>sparse<\/strong>&nbsp;and&nbsp;<strong>dynamic<\/strong>, i.e., the sparse patterns are largely different across inputs.<\/li>\n\n\n\n<li>This&nbsp;<strong>dynamic sparsity presents three unique spatial aggregation patterns<\/strong>&nbsp;that&nbsp;<strong>persist<\/strong>&nbsp;for all inputs:&nbsp;<strong><em>A-shape<\/em><\/strong>,&nbsp;<strong><em>Vertical-Slash<\/em><\/strong>, and&nbsp;<strong><em>Block-Sparse<\/em><\/strong>.<\/li>\n\n\n\n<li>These dynamic sparse indices can be approximated with minimal overhead online and speed up attention inference using a custom optimized GPU kernel.<\/li>\n<\/ol>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"why-minference\">Why&nbsp;<em>MInference<\/em>?<\/h3>\n\n\n\n<p>Long-context LLM inference faces&nbsp;<strong>two major challenges<\/strong>:&nbsp;<strong>1)<\/strong>&nbsp;long pre-filling stage attention latency, and&nbsp;<strong>2)<\/strong>&nbsp;high storage and transfer costs for KV cache. Previous efficient methods for long-context LLMs have focused on KV-cache compression, static sparse attention (e.g., model compression, SSM, linear attention), or distributed serving. However, these methods struggle to achieve acceptable latency for million-token level prompts with low cost and a single A100 GPU.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-full\"><img loading=\"lazy\" decoding=\"async\" width=\"1084\" height=\"380\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px.png\" alt=\"MInference: Figure 3.1. (a) Visualization of attention weights from different attention heads. For different prompts and tasks, the pattern of the same head is relatively consistent, but the sparse indices are dynamically changing. (b) Distance of the top-10 nearest non-zero element in the attention matrix. (c) Attention recall distribution using our identified patterns, where FLOPs in the kernel refer to the real FLOPs required for sparse attention computing using on GPUs. Here, a 1x64 block size is used for the Vertical-Slash pattern, and a 64x64 block size is used for others on GPUs. All visualizations are based on LLaMA-3-8B-Instruct-262K.\" class=\"wp-image-1068762\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px.png 1084w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px-300x105.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px-1024x359.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px-768x269.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-1_1084px-240x84.png 240w\" sizes=\"(max-width: 1084px) 100vw, 1084px\" \/><figcaption class=\"wp-element-caption\">Figure 3.1. (a) Visualization of attention weights from different attention heads. For different prompts and tasks, the pattern of the same head is relatively consistent, but the sparse indices are dynamically changing.  (b) Distance of the top-10 nearest non-zero element in the attention matrix.  (c) Attention recall distribution using our identified patterns, where FLOPs in the kernel refer to the real FLOPs required for sparse attention computing using on GPUs. Here, a 1&#215;64 block size is used for the Vertical-Slash pattern, and a 64&#215;64 block size is used for others on GPUs. All visualizations are based on LLaMA-3-8B-Instruct-262K.<\/figcaption><\/figure>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<p>To address these issues, we propose MInference, where the name reflects our ambition to enable million-token inference on a single A100 machine. MInference is a&nbsp;<strong>training-free<\/strong>&nbsp;efficient method for the pre-filling stage of long-context LLMs based on dynamic sparse attention. Specifically, we leverage the static spatial aggregation patterns of dynamic sparse attention, as shown above in Figure 3.1, and classify the dynamic sparse patterns into three types:&nbsp;<strong><em>A-shape<\/em><\/strong>,&nbsp;<strong><em>Vertical-Slash<\/em><\/strong>, and&nbsp;<strong><em>Block-Sparse<\/em><\/strong>. MInference first determines the optimal dynamic sparse pattern for each head offline using the Kernel-Aware Sparse Pattern Search algorithm, as illustrated below in Algorithm 1. During inference, it dynamically approximates the dynamic sparse indices based on the head&#8217;s pattern, as shown below in Algorithms 2 and 3. Finally, we perform efficient dynamic sparse attention computation using our optimized GPU kernel, significantly reducing the pre-filling stage latency for long-context LLMs.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-full\"><img loading=\"lazy\" decoding=\"async\" width=\"1400\" height=\"576\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px.png\" alt=\"MInference: diagram showing three columns for Algorithm 1, Algorithm 2, and Alogrithm 3\" class=\"wp-image-1068774\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px.png 1400w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px-300x123.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px-1024x421.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px-768x316.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_alg-1400px-240x99.png 240w\" sizes=\"(max-width: 1400px) 100vw, 1400px\" \/><\/figure>\n\n\n\n<div style=\"height:10px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<p>For example, with the&nbsp;<strong><em>Vertical-Slash<\/em><\/strong>&nbsp;pattern, we first use the attention calculation between the last Q and K to estimate the optimal indices of vertical lines and slash lines. Then, we utilize the dynamic sparse compiler PIT and Triton to construct the Vertical-Slash FlashAttention kernel, accelerating the attention computation. For the&nbsp;<strong><em>A-shape<\/em><\/strong>,&nbsp;<strong><em>Vertical-Slash<\/em><\/strong>, and&nbsp;<strong><em>Block-Sparse<\/em><\/strong>&nbsp;patterns, we first use the mean pooling of Q and K in attention calculations. By leveraging the commutative property of mean pooling and MatMul, we estimate the block-sparse indices. Then, we use Triton to construct the Block-Sparse FlashAttention kernel, accelerating the attention computation. For detailed kernel implementation, please refer to Appendix C.4 and the code.<\/p>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"our-main-contributions-are-four-fold\">Our main contributions are&nbsp;four-fold:<\/h3>\n\n\n\n<ol class=\"wp-block-list\">\n<li>We propose a dynamic sparse attention method,&nbsp;<strong>MInference<\/strong>, to accelerate the pre-filling stage of long-context LLMs by up to&nbsp;<strong>10x for 1M<\/strong>&nbsp;token prompts while&nbsp;<strong>maintaining the capabilities of LLMs, especially their retrieval abilities<\/strong>, as demonstrated in tasks like Needle in a Haystack.<\/li>\n\n\n\n<li>We classify dynamic sparse attention in LLMs into&nbsp;<strong>three patterns<\/strong>&nbsp;and design the Kernel-Aware Sparse Pattern Search algorithm to find the optimal head pattern offline.<\/li>\n\n\n\n<li>We introduce an online approximate method and&nbsp;<strong>optimized GPU kernels<\/strong>&nbsp;to accelerate LLM inference with&nbsp;<strong>minimal overhead<\/strong>. We also propose&nbsp;<strong>an optimal inference codebase that enables 1M token pre-filling inference on a single A100 using LLaMA-style models<\/strong>.<\/li>\n\n\n\n<li>We evaluate MInference across&nbsp;<strong>four benchmarks<\/strong>: InfiniteBench, RULER, PG-19, and Needle in a Haystack, with token lengths ranging from 128k to 1M, to assess the actual context processing capabilities of LLMs. Experimental results reveal that MInference can maintain or slightly improve actual context processing capabilities, while also outperforming in terms of cost efficiency and system latency.<\/li>\n<\/ol>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"experiments-results-in-long-context-benchmarks\">Experiments results in long-context benchmarks<\/h3>\n\n\n\n<p>We tested MInference across a range of scenarios, including QA, coding, retrieval-based tasks, multi-hop QA, summarization, and math tasks. The&nbsp;<strong>RULER benchmark<\/strong>&nbsp;includes several complex multi-hop or multi-needle tasks, effectively reflecting the actual context window size of LLMs. As shown in Table 1, our method effectively preserves the actual context window processing capability of LLMs and even slightly&nbsp;<strong>extends the actual context window size to 32K<\/strong>.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-large\"><img loading=\"lazy\" decoding=\"async\" width=\"1024\" height=\"311\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px-1024x311.png\" alt=\"MInference: Table 1. Performance (%) of different models and different methods on RULER evaluated at lengths from 4k to 128k.\" class=\"wp-image-1068777\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px-1024x311.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px-300x91.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px-768x233.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px-240x73.png 240w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table1_1400px.png 1400w\" sizes=\"(max-width: 1024px) 100vw, 1024px\" \/><figcaption class=\"wp-element-caption\">Table 1. Performance (%) of different models and different methods on RULER evaluated at lengths from 4k to 128k.<\/figcaption><\/figure>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<p>We also tested MInference on a broader range of tasks using the&nbsp;<strong>InfiniteBench<\/strong>, which has an&nbsp;<strong>average token length of 214K<\/strong>, as shown in Table 2. Compared to the SoTA baselines, MInference consistently maintains performance across all tasks. Notably,&nbsp;<strong>in the more challenging retrieval tasks<\/strong>&nbsp;like KV retrieval task, all baselines fail to make accurate predictions, with accuracy rates below 1.2%. However, MInference successfully retains the ability to handle dynamic KV pair retrieval.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-large\"><img loading=\"lazy\" decoding=\"async\" width=\"1024\" height=\"407\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px-1024x407.png\" alt=\"MInference: Table 2. Performance of different methods with different base models on InfiniteBench.\" class=\"wp-image-1068780\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px-1024x407.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px-300x119.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px-768x305.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px-240x95.png 240w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_table2_1400px.png 1400w\" sizes=\"(max-width: 1024px) 100vw, 1024px\" \/><figcaption class=\"wp-element-caption\">Table 2. Performance of different methods with different base models on InfiniteBench.<\/figcaption><\/figure>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<p>To further evaluate performance across different context lengths and positions of key information within prompts, we tested various models and methods using the&nbsp;<strong>Needle in a Haystack task<\/strong>. As shown in Fig.(1), MInference performs well across different models, context windows, and positions within the prompt, maintaining or even slightly improving performance compared to the original models. In the case of LLaMA-3-8B and GLM-4-9B-1M, MInference achieves&nbsp;<strong>full green performance<\/strong>&nbsp;for context windows&nbsp;<strong>up to 1M<\/strong>. In comparison, StreamingLLM and InfLLM experience a performance drop to below 20% in the middle segments of prompts even in the 70K context windows.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-full\"><img loading=\"lazy\" decoding=\"async\" width=\"1400\" height=\"2411\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px.png\" alt=\"MInference: Figure 1. Needle In A Haystack results using LLaMA-3-8B-Instruct-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3Mini-128K, and Qwen2-7-128K.\" class=\"wp-image-1068768\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px.png 1400w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-174x300.png 174w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-595x1024.png 595w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-768x1323.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-892x1536.png 892w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-1189x2048.png 1189w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig1_1400px-105x180.png 105w\" sizes=\"(max-width: 1400px) 100vw, 1400px\" \/><figcaption class=\"wp-element-caption\">Figure 1. Needle In A Haystack results using LLaMA-3-8B-Instruct-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3Mini-128K, and Qwen2-7-128K.<\/figcaption><\/figure>\n\n\n\n<div style=\"height:30px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<p>We also tested MInference on the&nbsp;<strong>language model tasks using PG-19<\/strong>, which includes tokens up to 100k. As shown in Fig.(2), MInference effectively maintains the perplexity of LLaMA-3-8B and Yi-9B-200K, while all baselines experience varying degrees of perplexity drop. Additionally, it can be observed that StreamingLLM with dilated and strided configurations better maintain perplexity performance compared to the standard StreamingLLM.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-large\"><img loading=\"lazy\" decoding=\"async\" width=\"1024\" height=\"296\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px-1024x296.png\" alt=\"MInference: Figure 2. Perplexity results on PG-19 using different models and methods.\" class=\"wp-image-1068765\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px-1024x296.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px-300x87.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px-768x222.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px-240x69.png 240w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig2_1400px.png 1400w\" sizes=\"(max-width: 1024px) 100vw, 1024px\" \/><figcaption class=\"wp-element-caption\">Figure 2. Perplexity results on PG-19 using different models and methods.<\/figcaption><\/figure>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"latency-breakdown-and-sparsity-pattern-in-the-kernel\">Latency breakdown and sparsity pattern in the kernel<\/h3>\n\n\n\n<p>Figure 3.2 shows the&nbsp;<strong>micro-benchmark<\/strong>&nbsp;results of the three attention patterns proposed in this paper, as well as FlashAttention. It can be seen that Vertical-Slash is the slowest among the three patterns, but it still achieves a&nbsp;<strong>13x speedup<\/strong>&nbsp;compared to FlashAttention under 1M context windows.<\/p>\n\n\n\n<p>Figure 4 shows the&nbsp;<strong>sparse indices in the kernel<\/strong>&nbsp;of the Vertical-Slash head. The vertical lines are computed using 1&#215;64 blocks through PIT FlashAttention, while the slash lines are computed using 64&#215;64 blocks through Block-level FlashAttention.<\/p>\n\n\n\n<div style=\"height:30px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<div class=\"wp-block-columns is-layout-flex wp-container-core-columns-is-layout-2 wp-block-columns-is-layout-flex\">\n<div class=\"wp-block-column is-layout-flow wp-block-column-is-layout-flow\" style=\"flex-basis:45%\">\n<figure class=\"wp-block-image aligncenter size-full\"><img loading=\"lazy\" decoding=\"async\" width=\"700\" height=\"390\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-2_700px.png\" alt=\"MInference: Figure 3.2. The latency breakdown of a single attention kernel for three patterns and FlashAttention across different context windows in a single A100, including the index time for dynamic sparse approximation and building dynamic sparsity. At 10k tokens, the latency of the four kernels is very close and all are less than 1ms. At 1M tokens, the latency for A-shape is 164ms.\" class=\"wp-image-1068759\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-2_700px.png 700w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-2_700px-300x167.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig3-2_700px-240x134.png 240w\" sizes=\"(max-width: 700px) 100vw, 700px\" \/><figcaption class=\"wp-element-caption\">Figure 3.2. The latency breakdown of a single attention kernel for three patterns and FlashAttention across different context windows in a single A100, including the index time for dynamic sparse approximation and building dynamic sparsity. At 10k tokens, the latency of the four kernels is very close and all are less than 1ms. At 1M tokens, the latency for A-shape is 164ms.<\/figcaption><\/figure>\n<\/div>\n\n\n\n<div class=\"wp-block-column is-layout-flow wp-block-column-is-layout-flow\" style=\"flex-basis:10%\"><\/div>\n\n\n\n<div class=\"wp-block-column is-layout-flow wp-block-column-is-layout-flow\" style=\"flex-basis:45%\">\n<figure class=\"wp-block-image aligncenter size-full\"><img loading=\"lazy\" decoding=\"async\" width=\"391\" height=\"390\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px.png\" alt=\"MInference: Figure 4. The dynamic sparse mask in the kernel of the vertical-slash pattern schematic using LLaMA-3-8B in the summarization task, where the yellow areas indicate the parts actually involved in computation. The slash lines are covered using 64x64 block sizes, while the vertical lines are covered using 1x64 block sizes.\" class=\"wp-image-1073997\" style=\"object-fit:cover\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px.png 391w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px-300x300.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px-150x150.png 150w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px-180x180.png 180w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig4_390px-360x360.png 360w\" sizes=\"(max-width: 391px) 100vw, 391px\" \/><figcaption class=\"wp-element-caption\">Figure 4. The dynamic sparse mask in the kernel of the vertical-slash pattern schematic using LLaMA-3-8B in the summarization task, where the yellow areas indicate the parts actually involved in computation. The slash lines are covered using 64&#215;64 block sizes, while the vertical lines are covered using 1&#215;64 block sizes.<\/figcaption><\/figure>\n<\/div>\n<\/div>\n\n\n\n<div style=\"height:20px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n\n<h3 class=\"wp-block-heading\" id=\"bibtex\">BibTeX<\/h3>\n\n\n\n<p>If you find this project helpful, please cite the following papers:<\/p>\n\n\n\n<pre class=\"wp-block-code\"><code>@article{jiang2024minference,\n    title={MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention},\n    author={Jiang, Huiqiang and Li, Yucheng and Zhang, Chengruidong and Wu, Qianhui and Luo, Xufang and Ahn, Surin and Han, Zhenhua and Abdi, Amir H and Li, Dongsheng and Lin, Chin-Yew and Yang, Yuqing and Qiu, Lili},\n    journal={arXiv preprint arXiv:2407.02490},\n    year={2024}\n}<\/code><\/pre>\n\n\n\n\n\n<p><strong>Q: How to effectively evaluate the impact of dynamic sparse attention on the capabilities of long-context LLMs?<\/strong><\/p>\n\n\n\n<p>To evaluate long-context LLM capabilities using models like LLaMA-3-8B-Instruct-1M and GLM-4-9B-1M, we tested: 1) context window with RULER, 2) general tasks with InfiniteBench, 3) retrieval tasks with Needle in a Haystack, and 4) language model prediction with PG-19.<\/p>\n\n\n\n<p>We found traditional methods perform poorly in retrieval tasks, with difficulty levels as follows:&nbsp;<strong>KV retrieval > Needle in a Haystack > Retrieval.Number > Retrieval PassKey<\/strong>. The main challenge is the semantic difference between needles and the haystack. Traditional methods excel when this difference is larger, as in passkey tasks. KV retrieval requires higher retrieval capabilities since any key can be a target, and multi-needle tasks are even more complex.<\/p>\n\n\n\n<p>We will continue to update our results with more models and datasets in future versions.<\/p>\n\n\n\n<p><strong>Q: Does this dynamic sparse attention pattern only exist in long-context LLMs that are not fully trained?<\/strong><\/p>\n\n\n\n<p>Firstly, attention is dynamically sparse, a characteristic inherent to the mechanism. We selected state-of-the-art long-context LLMs, GLM-4-9B-1M and LLaMA-3-8B-Instruct-1M, with effective context windows of 64K and 16K. With MInference, these can be extended to 64K and 32K, respectively. We will continue to adapt our method to other advanced long-context LLMs and update our results, as well as explore the theoretical basis for this dynamic sparse attention pattern.<\/p>\n\n\n\n<p><strong>Q: Does this dynamic sparse attention pattern only exist in Auto-regressive LMs or RoPE based LLMs?<\/strong><\/p>\n\n\n\n<p>Similar vertical and slash line sparse patterns have been discovered in BERT[1] and multi-modal LLMs[2]. Our analysis of T5&#8217;s attention patterns, shown in the figure, reveals these patterns persist across different heads, even in bidirectional attention.<\/p>\n\n\n\n<p>[1] SparseBERT: Rethinking the Importance Analysis in Self-Attention, ICML 2021.<br>[2] LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference, 2024.<\/p>\n\n\n\n<figure class=\"wp-block-image aligncenter size-full is-resized\"><img loading=\"lazy\" decoding=\"async\" width=\"1400\" height=\"704\" src=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px.png\" alt=\"MInference: Figure 5. The sparse pattern in T5 Encoder.\" class=\"wp-image-1068750\" style=\"width:804px;height:auto\" srcset=\"https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px.png 1400w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px-300x151.png 300w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px-1024x515.png 1024w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px-768x386.png 768w, https:\/\/www.microsoft.com\/en-us\/research\/uploads\/prod\/2024\/07\/MInference_fig5_1400px-240x121.png 240w\" sizes=\"(max-width: 1400px) 100vw, 1400px\" \/><figcaption class=\"wp-element-caption\">Figure 5. The sparse pattern in T5 Encoder.<\/figcaption><\/figure>\n\n\n\n<p><strong>Q: What is the relationship between MInference, SSM, Linear Attention, and Sparse Attention?<\/strong><br>All four approaches (MInference, SSM, Linear Attention, and Sparse Attention) efficiently optimize attention complexity in Transformers, each introducing inductive bias differently. The latter three require training from scratch. Recent works like Mamba-2 and Unified Implicit Attention Representation unify SSM and Linear Attention as static sparse attention, with Mamba-2 itself being a block-wise sparse method. While these approaches show potential due to sparse redundancy in attention, static sparse attention may struggle with dynamic semantic associations in complex tasks. In contrast, dynamic sparse attention is better suited for managing these relationships.<\/p>\n\n\n\n<div style=\"height:30px\" aria-hidden=\"true\" class=\"wp-block-spacer\"><\/div>\n\n\n","protected":false},"excerpt":{"rendered":"<p>Million-Tokens Prompt Inference for Long-context LLMs MInference 1.0&nbsp;leverages the dynamic sparse nature of LLMs&#8217; attention, which exhibits some static patterns, to speed up the pre-filling for long-context LLMs. It first determines offline which sparse pattern each head belongs to, then approximates the sparse index online and dynamically computes attention with the optimal custom kernels. This [&hellip;]<\/p>\n","protected":false},"featured_media":1053933,"template":"","meta":{"msr-url-field":"","msr-podcast-episode":"","msrModifiedDate":"","msrModifiedDateEnabled":false,"ep_exclude_from_search":false,"footnotes":""},"research-area":[13556,13545,13547],"msr-locale":[268875],"msr-impact-theme":[],"msr-pillar":[],"class_list":["post-1053873","msr-project","type-msr-project","status-publish","has-post-thumbnail","hentry","msr-research-area-artificial-intelligence","msr-research-area-human-language-technologies","msr-research-area-systems-and-networking","msr-locale-en_us","msr-archive-status-active"],"msr_project_start":"2024-02-29","related-publications":[1053816],"related-downloads":[1057272],"related-videos":[],"related-groups":[],"related-events":[],"related-opportunities":[],"related-posts":[],"related-articles":[],"tab-content":[],"slides":[],"related-researchers":[{"type":"user_nicename","display_name":"Huiqiang Jiang","user_id":40807,"people_section":"Related people","alias":"hjiang"},{"type":"guest","display_name":"Yucheng Li","user_id":1053924,"people_section":"Related people","alias":""},{"type":"user_nicename","display_name":"Chengruidong Zhang","user_id":42018,"people_section":"Related people","alias":"chengzhang"},{"type":"user_nicename","display_name":"Qianhui Wu","user_id":40741,"people_section":"Related people","alias":"qianhuiwu"},{"type":"user_nicename","display_name":"Xufang Luo","user_id":40324,"people_section":"Related people","alias":"xufluo"},{"type":"guest","display_name":"Surin Ahn","user_id":1053927,"people_section":"Related people","alias":""},{"type":"guest","display_name":"Amir H. Abdi","user_id":1053930,"people_section":"Related people","alias":""},{"type":"user_nicename","display_name":"Dongsheng Li","user_id":39402,"people_section":"Related people","alias":"dongsli"},{"type":"user_nicename","display_name":"Chin-Yew Lin","user_id":31493,"people_section":"Related people","alias":"cyl"},{"type":"user_nicename","display_name":"Yuqing Yang","user_id":40654,"people_section":"Related people","alias":"yuqyang"},{"type":"user_nicename","display_name":"Lili Qiu","user_id":41320,"people_section":"Related people","alias":"liliqiu"}],"msr_research_lab":[199560],"msr_impact_theme":[],"_links":{"self":[{"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-project\/1053873"}],"collection":[{"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-project"}],"about":[{"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/types\/msr-project"}],"version-history":[{"count":25,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-project\/1053873\/revisions"}],"predecessor-version":[{"id":1078965,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-project\/1053873\/revisions\/1078965"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/media\/1053933"}],"wp:attachment":[{"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/media?parent=1053873"}],"wp:term":[{"taxonomy":"msr-research-area","embeddable":true,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/research-area?post=1053873"},{"taxonomy":"msr-locale","embeddable":true,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-locale?post=1053873"},{"taxonomy":"msr-impact-theme","embeddable":true,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-impact-theme?post=1053873"},{"taxonomy":"msr-pillar","embeddable":true,"href":"https:\/\/www.microsoft.com\/en-us\/research\/wp-json\/wp\/v2\/msr-pillar?post=1053873"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}