PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, Lili Qiu

arXiv:2407.02490v2cs.CLcs.LG

TL;DR

Long-context LLM pre-filling is bottlenecked by quadratic attention costs, motivating MInference's dynamic sparse-attention approach. By exploiting recurring attention patterns and input-specific sparse masks, it achieves substantial pre-filling speedups while maintaining accuracy, with limitations at shorter contexts and higher sparsity.

  • Problem

    Quadratic attention makes long-context pre-filling prohibitively slow, while existing acceleration methods may not preserve acceptable accuracy or efficiency.

  • Method

    MInference assigns A-shape, Vertical-Slash, or Block-Sparse patterns to attention heads, builds dynamic sparse masks for each input, and applies sparse attention kernels.

  • Results

    Across long-context models and benchmarks, MInference maintains or improves accuracy and delivers up to 10× faster pre-filling on 1M-token contexts.

  • Takeaways & Limitations

    MInference reduces the computation and latency of long-context pre-filling while remaining applicable to existing LLMs without pre-training changes or additional fine-tuning.

  • Takeaways & Limitations

    At 10K contexts, dynamic-index construction grows from 5% to 30% of latency, and higher sparsity rates may noticeably reduce performance.

Abstract

from arXiv · show

The computational challenges of Large Language Model (LLM) inference remain a significant barrier to their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on a single A100 GPU. Existing methods for speeding up prefilling often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs. To address this gap, we introduce MInference (Milliontokens Inference), a sparse calculation method designed to accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices-the A-shape, Vertical-Slash, and Block-Sparsethat can be leveraged for efficient sparse computation on GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse indices based on the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels to significantly reduce the latency in the pre-filling stage of long-context LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications to the pre-training setup or additional fine-tuning. By evaluating on a wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, and Qwen2-128K, we demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy. Our code is available at https://aka.ms/MInference.

1 Introduction

Long-context LLMs face severe pre-filling latency because attention scales quadratically, while MInference uses dynamic sparse attention to reduce computation and accelerate inference without sacrificing accuracy.

  • Long-context windows support applications including repository-level code understanding, long-document question answering, self-play reasoning, and long-horizon agent tasks.
  • 30 minutes are required to pre-fill a 1M-token prompt for LLaMA-3-8B on one A100, with self-attention exceeding 90% of total latency.At 300K tokens, pre-filling takes 6 minutes on the same setup.
  • MInference reduces 95% of attention FLOPs using dynamic sparse attention tailored to long-context inference.It identifies three sparse patterns, assigns one per attention head, and builds input-dependent sparse masks online.
  • 10× speedup reduces 1M-token pre-filling from 30 minutes to 3 minutes on a single A100 while maintaining or improving accuracy across long-context evaluations.Experiments cover multiple long-context models and benchmarks, including InfiniteBench, RULER, Needle In A Haystack, and PG-19.

2 Attention Heads: Dynamic, Sparse, and Characteristic

Long-context attention is sparse but input-dependent, and its nonzero weights form recurring spatial patterns that can guide efficient GPU computation.

  • Dynamic sparsity: 96.8% of total attention is recalled by retaining only the top 4k columns in a 128k × 128k attention matrix.This indicates that each token attends to a limited subset of the long sequence.
  • Dynamic sparsity: Reusing top-4k columns from another prompt lowers attention recall, showing that sparse indices vary substantially with context.
  • Pattern structure: The same attention head tends to preserve its pattern across prompts and tasks while its sparse indices dynamically change.
  • Pattern structure: Attention sparsity is categorized into A-shape, Vertical-Slash, and Block-Sparse patterns with distinct spatial distributions.These patterns are identified across long-context prompts of varying lengths and tasks.
  • Pattern structure: A-shape heads concentrate attention on initial tokens and local windows, whereas Vertical-Slash heads target specific tokens and fixed-interval positions that change with context.The Vertical-Slash distribution is difficult to capture with local windows or A-shape masks.
  • Pattern structure: Block-Sparse attention is more dispersed but retains spatial clustering, with nearest nonzero weights generally about 5 positions apart.
  • Efficient computation: At equal FLOPs, the identified patterns achieve higher attention-score recall than competing sparse methods, including for Block-Sparse heads.Prior Top-K methods struggle with Block-Sparse attention because they focus on specific tokens globally.

3 MInference 1.0

MInference accelerates long-context pre-filling by assigning each attention head an offline-selected sparse pattern, dynamically building input-dependent indices, and executing sparse attention with optimized GPU kernels.

  • MInference pipeline: MInference uses three steps: offline pattern identification per head, dynamic sparse-index construction, and sparse attention calculation with optimized GPU kernels.A-shape heads use static masks, avoiding dynamic-mask construction overhead.
  • Problem formulation: The sparse-mask formulation sets masked attention logits near zero after softmax by subtracting a large constant when M_i,j = 0.The mask entries M_i,j are binary, and c can be 1e5.
  • Kernel-aware pattern search: The method searches candidate patterns and settings under a target FLOPs budget, selecting the configuration with the highest attention-output recall.The search uses FlashAttention to reduce GPU memory overhead and includes information from the V matrix.
  • Dynamic index construction: MInference dynamically estimates attention patterns from the current input and builds sparse indices according to each head’s assigned pattern.Vertical-Slash heads estimate vertical and slash indices from recent query vectors, while Block-Sparse heads estimate block-level weights after mean pooling Q and K.

4 Experiments

Experiments evaluate MInference across long-context models and benchmarks, showing substantial pre-filling speedups while preserving performance across retrieval, reasoning, language modeling, and other tasks.

  • Evaluation setup: MInference is evaluated on InfiniteBench, RULER, Needle In A Haystack, and PG-19 using multiple long-context LLMs and context lengths up to 1M tokens.The evaluation covers retrieval, QA, multi-hop reasoning, aggregation, summarization, coding, and language modeling.
  • InfiniteBench: MInference achieves the best overall performance on InfiniteBench and matches or slightly surpasses full attention on some tasks.Its performance remains close to the original model on retrieval-related tasks, unlike some baseline methods.
  • RULER: MInference maintains long-context performance on RULER and achieves effective context windows of 32K and 64K in LLaMA-3-8B-262K and GLM-4-9B-1M.It outperforms original full attention for testing lengths beyond 32K.
  • Language modeling: For 100K-token PG-19 prompts, MInference perplexity is 0.2 higher than full attention and lower than StreamingLLM by 0.25 on Yi-9B-200K and 0.75 on LLaMA-3-262K.The method yields the best results among the compared sparse approaches while remaining close to full attention.
  • Needle In A Haystack: MInference retains retrieval performance across context windows from 1K to 1M tokens in Needle In A Haystack, while some baselines decline when critical information exceeds global and local windows.The comparison concerns information placed at different positions across varying context lengths.
  • Ablations: Static sparse indices sharply reduce performance in dynamic tasks, while removing any of the three patterns also causes performance degradation.The only-Vertical-Slash variant preserves most performance but remains below the full method.
  • Latency: At 100K, 300K, 500K, and 1M tokens, MInference achieves 1.8×, 4.1×, 6.8×, and 10× speedups, reducing 1M-token pre-filling from 30 minutes to 3 minutes on one A100.Dynamic sparse-index building accounts for about 5%-20% of the latency overhead.
  • Compatibility and scaling: MInference remains compatible with SnapKV compression and maintains nearly unchanged performance on most tasks, while also performing strongly on larger LLaMA-3-70B-1M models.On dynamic KV-retrieval tasks, it can match or slightly improve upon full attention.

5 Related Works

Related work addresses sparse attention, context-window scaling, and long-context inference through distinct strategies spanning model training, attention computation, memory, and decoding.

  • Sparse Attention: Sparse-attention methods include static patterns, cluster-based approaches, and dynamic sparse attention.Static patterns include sliding-window, dilated, and mixed sparse attention; cluster-based methods use hashing or k-nearest neighbors.
  • Sparse Attention: Many static and cluster-based sparse methods require pre-training from scratch, limiting their direct use as plugins for ready-to-use LLMs.This distinguishes them from approaches designed to operate on existing models without retraining.
  • Scaling Context Windows of LLMs: Context-window scaling methods include staged pre-training, position-embedding modification or interpolation, external memory, and distributed computation.These approaches focus on enabling pretrained LLMs to handle longer contexts.
  • Long-Context LLM Inference: Long-context inference research separates pre-filling optimizations from decoding optimizations addressing KV-cache storage and computation.Prefilling methods include state-space, linear-attention, memory-based, hybrid, and prompt-compression approaches.
  • Long-Context LLM Inference: Decoding optimizations include KV reuse, static or dynamic cache dropping, cache offloading, restoration methods, and hierarchical speculative decoding.These methods primarily target decoding-time memory or computation rather than pre-filling attention.

6 Conclusion

MInference accelerates long-context pre-filling by assigning attention heads sparse spatial patterns, dynamically building masks, and applying optimized sparse kernels. Experiments show up to 10× speedup while maintaining long-context capabilities.

  • 6 Conclusion: MInference categorizes attention heads into A-shape, Vertical-Slash, and Block-Sparse patterns.It searches for each head’s optimal pattern, builds input-dependent sparse masks, and applies sparse attention kernels.
  • 6 Conclusion: MInference achieves up to 10× speedup for 1-million-token prompts on a single A100 GPU.Latency decreases from 30 minutes to 3 minutes per prompt.
  • 6 Conclusion: MInference maintains long-context capabilities across InfiniteBench, RULER, language modeling, and Needle In A Haystack benchmarks.The method was evaluated on multiple long-context tasks and models.
  • 6 Conclusion: Similar dynamic sparse attention patterns also appear in multimodal and encoder-decoder LLMs.The conclusion reports these patterns beyond the evaluated decoder-only setting.

A Limitations

MInference’s indexing overhead is more consequential at shorter context lengths, while increasing sparsity can reduce model performance. The method is therefore most advantageous for sufficiently long prompts and moderate sparsity.

  • A Limitations: At 10k context, dynamic-index construction rises from 5% to 30% of latency, bringing end-to-end time close to FlashAttention.The overhead proportion decreases as prompt length increases.
  • A Limitations: Higher sparsity rates may noticeably reduce model performance.This limits how aggressively computation can be reduced without accuracy loss.
  • A Limitations: Index-building overhead becomes more significant as context length decreases and attention computation time falls.The limitation is less pronounced for longer prompts.

C.1 Dataset Details

The evaluation spans diverse long-context benchmarks, models, kernels, sparse patterns, and latency analyses. These details cover dataset composition, model selection, single-GPU execution, and implementation behavior.

  • C.1 Dataset Details: InfiniteBench contains 10 tasks spanning summarization, question answering, code debugging, numerical identification, and retrieval.Its tasks average 214k tokens and include 3,992 examples.
  • C.1 Dataset Details: RULER evaluates 13 complex tasks across retrieval, multi-hop tracing, aggregation, and question answering categories.Its retrieval tasks include single-needle, multi-key, multi-value, and multi-query variants.
  • C.1 Dataset Details: Needle In A Haystack tests retrieval of targeted information embedded in large text across context lengths and document depths.The study scales it to 1M-token contexts with 750 examples.
  • C.1 Dataset Details: PG-19 measures long-text language modeling with perplexity on 1,000 random samples longer than 100K tokens.The test set contains texts up to 500K tokens.
  • C.2 Additional Implementation Details: Experiments use LLaMA-3, Yi, Phi-3, Qwen2, and GLM-4 long-context models with greedy decoding.The models support context windows from 128K to 1M tokens.
  • C.3 Single A100 Implementation Details: Single-A100 execution uses tensor splitting, reduced intermediate allocation, and removal of unnecessary LM-head computations.These optimizations enable 1M-prompt inference on one A100 GPU.
  • C.4.2 Vertical-Slash Attention: The Vertical-Slash index kernel merges vertical columns and slash ranges with per-row complexity O(kv + ks).Its sparse FlashAttention kernel combines block-sparse processing with PIT column processing.
  • D.2 Latency Breakdown: Vertical-Slash reaches 13× speedup and Block-Sparse reaches 30× speedup over FlashAttention at 1M context.Block-Sparse is fastest among the three proposed patterns, while Vertical-Slash has the highest relative latency.

D.3 Additional Ablation Study

The ablation study separates vertical and slash components of Vertical-Slash attention. Vertical-only masking substantially harms retrieval, whereas slash-only masking preserves most performance but weakens highly dynamic retrieval.

  • D.3 Additional Ablation Study: The ablations compare Vertical-Slash variants using only vertical lines or only slash lines.Each variant retains one line from the other component because of kernel limitations.
  • D.3 Additional Ablation Study: Vertical-only masking causes a significant performance drop, especially on retrieval tasks.Its performance is similar to using only block-sparse attention.
  • D.3 Additional Ablation Study: Slash-only masking retains most performance but drops 2.9% on average versus the full method in highly dynamic KV retrieval.This indicates that both components contribute to robust performance across task types.

E Pattern Distribution

The optimal sparse-head configurations are dominated by Vertical-Slash patterns, while Block-Sparse and A-shape patterns concentrate in later and middle layers, respectively. The same configuration transfers effectively across two LLaMA model versions.

  • Pattern distribution: More than 90% of optimal attention-head patterns are Vertical-Slash, but using only this pattern harms highly dynamic KV-retrieval tasks.The ablation indicates that pattern diversity remains important despite Vertical-Slash dominance.
  • Pattern distribution: Block-Sparse patterns primarily occur in intermediate-to-later layers, whereas A-shape patterns are concentrated in middle layers.
  • Cross-model configuration: Using one configuration for two LLaMA versions yields nearly perfect Needle In A Haystack results for the 1M model, supporting pattern generalizability.
  • Cross-model configuration: The same optimal sparse-pattern configuration works for LLaMA-3-8B-Instruct-262K and LLaMA-3-8B-Instruct-1M.

F Sparsity in Kernel Distribution

Kernel-level sparsity rises substantially with context length across the three sparse patterns. Beyond 200K tokens, actual sparsity exceeds 90%, and beyond 500K it exceeds 95% relative to FlashAttention.

  • Kernel sparsity: Beyond 200K tokens, actual kernel sparsity exceeds 90% for all three sparse patterns.This measure reflects computation remaining after block coverage.
  • Kernel sparsity: Beyond 500K tokens, sparsity relative to FlashAttention exceeds 95%, corresponding to a theoretical speedup over 15×.
  • Kernel sparsity: Accounting for 20% index-building overhead, contexts above 200K still support kernel speedups over 8×.
  • Broader attention settings: The analysis also examines sparse patterns in T5-style bidirectional encoder attention on Flan-UL2 summarization.

H Case Study

Case studies compare generated outputs across summarization and KV retrieval, while the accompanying algorithms construct Vertical-Slash indices and execute sparse Flash Attention. The examples contrast faithful retrieval or summarization with failures from several baselines and static patterns.

  • Summarization: StreamingLLM produces a coherent-looking but factually incorrect summary by introducing elements absent from the original story.
  • KV retrieval: The original method achieves perfect retrieval on both examples, whereas StreamingLLM variants and a static sparse pattern fail significantly.
  • Summarization: Table 9 compares generated summaries from different methods on InfiniteBench EN.SUM with 200K-token inputs.
  • Vertical-Slash implementation: The Vertical-Slash index algorithm sorts vertical and slash indices, merges their ranges, and returns block and column index structures.
  • Vertical-Slash implementation: The index construction extends ranges, records covered blocks, and returns cblk, iblk, ccol, and icol for kernel execution.
  • Vertical-Slash implementation: Algorithm 5 takes Q, K, V and block or column indices as inputs, then combines block-sparse and PIT sparse Flash Attention operations.
  • KV retrieval: Table 10 compares LLaMA-3-8B-Instruct-262K methods on the 200K-token Retrieve.KV task.
Loading 2407.02490v2…