1. Bibliographic Information

1.1. Title

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

1.2. Authors

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. Affiliations include Microsoft Corporation and the University of Surrey.

1.3. Journal/Conference

The paper was published as a preprint on arXiv, with the original source link indicating its submission to arXiv. arXiv is a widely respected open-access repository for preprints of scientific papers in various fields, including computer science and artificial intelligence. While not a peer-reviewed journal or conference in its preprint form, papers on arXiv often undergo subsequent peer review and publication in prestigious venues.

1.4. Publication Year

2024

1.5. Abstract

The paper addresses the significant computational challenges of Large Language Model (LLM) inference, particularly the pre-filling stage for long-context LLMs, which suffers from the quadratic complexity of attention. This quadratic complexity results in unacceptably long latencies (e.g., 30 minutes for a 1M token prompt on an 8B LLM using a single A100 GPU). Existing sparse attention methods often fail to maintain accuracy or efficiency for long contexts or introduce substantial overhead. To overcome this, the authors introduce MInference (Milliontokens Inference), a dynamic sparse calculation method designed to accelerate long-sequence pre-filling. MInference identifies three unique patterns in long-context attention matrices—the A-shape, Vertical-Slash, and Block-Sparse—that enable efficient sparse computation on GPUs. It determines the optimal pattern for each attention head offline and dynamically builds sparse indices during inference based on the assigned pattern. The method utilizes optimized GPU kernels for these sparse attention calculations. MInference is training-free, applicable to existing LLMs, and evaluated on a wide range of tasks (InfiniteBench, RULER, PG-19, Needle In A Haystack) and models (LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, Qwen2-128K). The results demonstrate up to a 10x speedup for pre-filling on an A100 while maintaining accuracy.

1.6. Original Source Link

https://arxiv.org/abs/2407.02490

2. Executive Summary

2.1. Background & Motivation

The rapid advancement of Large Language Models (LLMs) has led to their deployment in a multitude of complex real-world applications, especially with the advent of long-context processing capabilities (e.g., context windows up to 1M or even 10M tokens). These extended contexts enable LLMs to handle tasks like repository-level code understanding, long-document question-answering, and long-horizon agent tasks.

However, a major bottleneck in deploying these long-context LLMs is the computational cost of inference, particularly during the pre-filling stage. This stage involves processing the entire input prompt before generating the first output token. The root cause of this high cost is the attention mechanism, which has a quadratic complexity with respect to the sequence length. As illustrated in the paper, processing a 1M token prompt for an 8B LLM can take up to 30 minutes on a single A100 GPU, leading to an unacceptable Time To First Token (TTFT) experience for users. The self-attention computation alone accounts for over 90% of the total pre-filling latency.

Prior research has noted that attention matrices are often highly sparse, meaning only a small fraction of token-pair interactions are truly significant. This observation led to fixed sparse attention methods (e.g., Longformer, BigBird). However, the exact sparse patterns are highly dynamic and vary significantly across different inputs and tasks, preventing these fixed methods from being directly applied to long-context LLMs without expensive re-training or fine-tuning. While some dynamic sparse attention methods exist, they often introduce significant computational overhead for estimating attention patterns, making them inefficient for long-context scenarios.

The core problem is therefore to accelerate the pre-filling stage of long-context LLMs by effectively leveraging attention sparsity in a dynamic and computationally efficient manner, without requiring re-training or extensive fine-tuning, thereby reducing latency and enabling wider deployment.

2.2. Main Contributions / Findings

The paper introduces MInference (Milliontokens Inference), a novel approach to significantly accelerate the pre-filling stage of long-context LLM inference. Its primary contributions and findings are:

• Identification of Novel Sparse Attention Patterns: MInference identifies three recurrent and distinct sparse attention patterns in long-context LLMs: A-shape, Vertical-Slash, and Block-Sparse. These patterns aggregate spatial information efficiently and can be leveraged for highly optimized sparse computation.
• Kernel-Aware Optimal Sparse Pattern Search: The paper proposes an offline kernel-aware search method to assign the most suitable sparse pattern and its optimal configuration (e.g., number of lines, blocks) to each attention head. This search considers the actual computational cost on GPUs, not just theoretical FLOPs, ensuring practical efficiency.
• Efficient Dynamic Sparse Mask Approximation: During inference, MInference performs an efficient online approximation to dynamically build sparse masks for each attention head based on its assigned pattern and the specific input. This minimizes the overhead associated with identifying dynamic sparsity, a common challenge in prior dynamic sparse attention methods.
• Optimized GPU Kernels: The authors developed three highly optimized GPU kernels, built upon PIT (Permutation Invariant Transformation), Triton, and FlashAttention, specifically tailored for the A-shape, Vertical-Slash, and Block-Sparse patterns. These kernels enable extremely efficient sparse attention calculations.
• Training-Free and General Applicability: The proposed technique can be directly applied to existing LLMs without any modifications to their pre-training setup or requiring additional fine-tuning, making it a practical plug-and-play solution.
• Significant Performance Improvement: Extensive experiments demonstrate that MInference effectively reduces pre-filling latency by up to 10x for 1M token contexts on a single A100 GPU (e.g., from 30 minutes to 3 minutes), while rigorously maintaining or even slightly improving the LLM's accuracy across various long-context benchmarks (InfiniteBench, RULER, PG-19, Needle In A Haystack) and diverse models (LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, Qwen2-128K).
• Compatibility: MInference is shown to be compatible with other inference optimization techniques, such as KV cache compression methods like SnapKV.

3. Prerequisite Knowledge & Related Work

3.1. Foundational Concepts

3.1.1. Large Language Models (LLMs)

Large Language Models (LLMs) are advanced artificial intelligence models, typically based on the Transformer architecture, that are trained on massive datasets of text and code. Their primary function is to understand and generate human-like language. They can perform a wide range of natural language processing tasks, including question answering, summarization, translation, and code generation. LLMs are characterized by their vast number of parameters (often billions) and their ability to leverage extensive context to generate coherent and relevant responses.

3.1.2. Transformer Architecture

The Transformer is a neural network architecture introduced in 2017, which revolutionized natural language processing. It primarily relies on the self-attention mechanism to process sequences of data, unlike previous models that used recurrent neural networks (RNNs) or convolutional neural networks (CNNs). A Transformer typically consists of an encoder and a decoder stack, though many modern LLMs primarily use a decoder-only architecture for generative tasks. Each layer within the Transformer stack usually contains a multi-head self-attention sub-layer and a feed-forward network (FFN) sub-layer.

3.1.3. Self-Attention Mechanism

Self-attention is the core mechanism in Transformers that allows the model to weigh the importance of different words in the input sequence when encoding or decoding a particular word. It computes a weighted sum of value vectors, where the weights are determined by the similarity between query and key vectors.

For an input sequence of length $S$, each token (word) is first converted into an embedding vector. These embeddings are then transformed into three different vectors for each token:

• Query (Q): Represents what information the current token is "looking for."

• Key (K): Represents what information the current token "has" to offer.

• Value (V): Represents the actual information content of the token.

  These $Q$, $K$, and $V$ vectors are typically represented as matrices. For a sequence of length $S$ and hidden dimension $d_h$, $Q$, $K$, and $V$ would be matrices of shape $S \times d_h$. The attention mechanism is then calculated as follows:

$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_h}}\right)V $

Where:

• $Q$ is the query matrix.
• $K$ is the key matrix.
• $K^T$ is the transpose of the key matrix.
• $QK^T$ computes the similarity scores between all query-key pairs. This results in an attention score matrix of size $S \times S$.
• $d_h$ is the dimension of the key/query vectors. Dividing by $\sqrt{d_h}$ is a scaling factor to prevent the dot products from becoming too large, which could push the softmax function into regions with very small gradients.
• $\mathrm{softmax}(\cdot)$ normalizes the attention scores so that they sum to 1, effectively turning them into probability distributions.
• $V$ is the value matrix. The softmax output (attention weights) is multiplied by $V$ to get the final output, which is a weighted sum of the value vectors. This output is also an $S \times d_h$ matrix.

3.1.4. Quadratic Complexity of Attention

The quadratic complexity refers to how the computational cost of the self-attention mechanism scales with the length of the input sequence. Specifically, computing the $QK^T$ matrix involves multiplying two $S \times d_h$ and $d_h \times S$ matrices, resulting in an $S \times S$ matrix. This operation has a computational cost proportional to $S^2 \cdot d_h$. As the sequence length ($S$) increases, the computational cost grows quadratically, making it a major bottleneck for long-context LLMs. For instance, if the sequence length doubles, the attention computation cost quadruples.

3.1.5. Pre-filling Stage

In LLM inference, the process is typically divided into two main stages:

1. Pre-filling (or Prompt Processing) Stage: This is when the entire input prompt (e.g., a long document or conversation history) is fed into the model. The model computes the Key (K) and Value (V) states for all tokens in the prompt and stores them in a KV cache. This stage involves a single, large self-attention computation over the entire input sequence, which is where the quadratic complexity becomes most problematic. The output of this stage is primarily the KV cache and the logit for the first generated token.

2. Decoding (or Generation) Stage: After the pre-filling stage, the model generates tokens one by one. For each new token, the self-attention computation only involves the newly generated token attending to all previous tokens (including the prompt tokens stored in the KV cache). This stage is typically less computationally intensive per token but can involve many steps for long generations.

   The pre-filling stage is the focus of this paper because its quadratic complexity makes it the dominant factor in Time To First Token (TTFT) for long-context LLMs.

3.1.6. Sparse Attention

Sparse attention is a modification of the standard self-attention mechanism where, instead of computing attention scores between all possible pairs of tokens, only a subset of these interactions is calculated. This is based on the empirical observation that in many cases, not all tokens are equally important for understanding the context of a given token; many attention scores are close to zero. By strategically selecting which Q-K pairs to compute, sparse attention aims to reduce the computational cost from $O(S^2)$ to a more efficient scaling (e.g., $O(S \cdot \text{log}S)$ or $O(S)$) while retaining most of the model's performance.

3.1.7. GPU Kernels

A GPU kernel is a function that can be executed on a Graphics Processing Unit (GPU). GPUs are designed for highly parallel computations, making them ideal for the matrix multiplications and other operations involved in neural networks. Optimized GPU kernels are specialized implementations of these functions that are carefully crafted to maximize performance on specific GPU architectures, leveraging techniques like shared memory, register usage, and instruction-level parallelism. FlashAttention and Triton (which MInference builds upon) are examples of efforts to create highly efficient GPU kernels for attention.

3.1.8. FLOPs (Floating-Point Operations per Second)

FLOPs refer to the number of floating-point operations (e.g., additions, multiplications) performed per second. In the context of neural networks, FLOPs are a common metric to measure the computational intensity of a model or an operation. Reducing the FLOPs for a given task, especially computationally expensive ones like attention, directly contributes to faster inference and lower energy consumption.

3.2. Previous Works

The paper discusses several categories of previous work relevant to sparse attention and scaling context windows.

3.2.1. Static Sparse Patterns

These methods impose a predefined, fixed pattern on the attention matrix, meaning certain query-key interactions are always ignored.

• Sliding Windows: Methods like Mistral 7B [JSM+23] and Phi-3-Mini-128K [AJA+24] use a fixed window size, where each token only attends to a limited number of previous tokens within that window.
• Dilated Attention: $[CGRS19, SGR+21, DMD+23]$ introduce gaps within the attention window, allowing tokens to attend to distant tokens at fixed intervals.
• Mixed Sparse Patterns: Longformer [BPC20] and BigBird [ZGD+20] combine global attention (a few tokens attend to all others) with local windowed attention.
• Limitation: These methods typically require pre-training the model from scratch with the new sparse pattern, making them difficult to apply to existing, ready-to-use LLMs without significant computational cost. MInference aims to be training-free.

3.2.2. Dynamic Sparse Attention

These methods attempt to identify and compute only the most important attention scores dynamically for each input, rather than using a fixed pattern.

• Prediction with Low-Rank Hidden Dimensions: Some approaches [LOC+22, RCHG+24, WZH21] predict sparse patterns by using low-rank hidden states or post-statistical methods.
• Limitation: The paper argues that these existing dynamic sparse attention methods often introduce substantial computational overhead in the estimation step, making them less suitable for long-context LLMs where efficiency is paramount. MInference differentiates itself by designing an efficient online approximation for pattern estimation.

3.2.3. Scaling Context Windows of LLMs

This category focuses on methods to enable LLMs to handle longer input sequences.

• Staged Pre-training [NXH+23, FPN+24]: Training models initially on shorter contexts and then continuing training on longer contexts.
• Modifying/Interpolating Position Embeddings [PSL22, CWCT23, PQFS24, DZ Z+24]: Adjusting how positional information is encoded to extrapolate to longer sequences (e.g., RoPE extension methods).
• External Memory Modules [BANG23, TSP+23, XZH+24]: Using separate memory units to store and retrieve context, effectively extending the context window without increasing the attention matrix size. InfLLM [XZH+24] is an example.
• Distributed Computations [LZA24]: Spreading the computational load across multiple devices (e.g., Ring Attention).
• Limitation: These methods primarily focus on enabling long-context capabilities but do not inherently alleviate the high inference costs in the pre-filling stage due to the quadratic nature of attention.

3.2.4. Long-Context LLM Inference Optimizations

This category directly addresses the high computational cost of attention and KV cache storage.

• Pre-filling Optimizations:
  • State Space Models (SSMs) [GGR22, GD23] and linear attention methods [SDH+23, PAA+23]: Offer sub-quadratic scaling, but often require training from scratch or significant architectural changes.
  • Memory-based methods [MFG24, HBK+24] and Hybrid methods [LLB+24, RLL+24].
  • Prompt compression methods [LDGL23, JWL+23, JW+24, PWJ+24]: Reduce the input tokens before feeding to the LLM.
  • kNN or cluster-based sparse attention [MEL24, XZH+24, LCE+24]: Accelerate inference but may suffer from reduced accuracy or limited speedup, or be restricted to CPUs.
• Decoding Stage Optimizations: Focus on KV cache management (compression, dropping, offloading, quantization) or speculative decoding. These are relevant for the decoding stage but do not address the pre-filling bottleneck.
• Baselines used in the paper:
  • StreamingLLM [XTC+24]: A prominent training-free approach that uses attention sinks (a few global tokens that always attend to all other tokens) combined with sliding window local attention. This is analogous to the A-shape pattern identified by MInference.
  • InfLLM [XZH+24]: Uses a memory unit to process streaming long sequences, combining global and local attention.
  • H2O [ZS Z+24]: Another KV cache compression method which focuses on identifying and retaining "heavy-hitter" tokens.

3.3. Technological Evolution

The evolution of attention mechanisms has been driven by the increasing need to process longer sequences efficiently.

1. Dense Attention (Original Transformer): The initial approach, where every token attends to every other token, offered powerful contextual understanding but suffered from quadratic complexity.

2. Static Sparse Attention: Recognizing the sparsity in attention matrices, researchers introduced fixed patterns (e.g., Longformer, BigBird) to reduce computation. While effective, these required retraining and lacked adaptability to diverse inputs.

3. Dynamic Sparse Attention (Early Forms): Attempts were made to dynamically identify important tokens $[LOC+22]$, but these often came with significant overhead for the sparsity estimation itself, negating some of the computational benefits.

4. Hardware-Aware Optimizations: Concurrently, efforts like FlashAttention [Dao24] focused on optimizing the dense attention calculation on GPUs to reduce memory access overhead and increase parallelism, even without changing the algorithmic complexity.

5. Long-Context LLM Specific Optimizations: As LLMs grew, specialized methods for context window extension and inference acceleration emerged (e.g., StreamingLLM, InfLLM). These often combine fixed sparse patterns with some form of memory management.

   MInference fits into this timeline by bridging the gap between dynamic sparse attention and long-context LLM inference efficiency. It builds upon the understanding of attention sparsity and hardware-aware optimization (FlashAttention, Triton) to create a training-free dynamic sparse attention method that is specifically optimized for the unique patterns observed in long-context LLMs, with minimal overhead for dynamic mask estimation.

3.4. Differentiation Analysis

MInference differentiates itself from existing methods primarily in its targeted approach to dynamic sparsity for long-context LLMs without requiring re-training or fine-tuning.

• Compared to Static Sparse Attention (e.g., Longformer, BigBird, StreamingLLM's A-shape component):

  • Core Difference: Static methods use predetermined, fixed attention patterns. MInference dynamically identifies and adapts its sparse pattern to the specific input sequence.
  • Innovation: MInference allows for dynamic mask building for Vertical-Slash and Block-Sparse heads, which better captures the context-dependent nature of attention. While MInference includes an A-shape (static) pattern, it intelligently combines it with dynamic ones based on a kernel-aware search. This is crucial, as the ablation study shows static indices significantly degrade performance in dynamic tasks like KV retrieval.

• Compared to Prior Dynamic Sparse Attention Methods (e.g., [LOC+22, RCHG+24]):

  • Core Difference: Previous dynamic methods often rely on complex calculations (e.g., low-rank approximations, statistical methods) to predict sparsity, incurring large computational overhead themselves.
  • Innovation: MInference is "designed specifically for long-context scenarios with minimal overhead in estimation." It uses efficient online approximation techniques (e.g., $matmul on last_q queries$, mean pooling for blocks) to build the dynamic sparse mask, ensuring the overhead doesn't negate the speedup.

• Compared to General Long-Context LLM Scaling Methods (e.g., RoPE extension, external memory):

  • Core Difference: Many of these methods enable long-context but don't directly address the quadratic complexity of attention computation in the pre-filling stage. Some even require re-training.
  • Innovation: MInference is a training-free "plugin" that directly optimizes the computational bottleneck of pre-filling for existing long-context LLMs. It's complementary to methods that expand context windows.

• Hardware Optimization: MInference also emphasizes kernel-aware optimization and develops custom GPU kernels based on Triton and FlashAttention for its identified sparse patterns. This ensures that the theoretical FLOPs reduction translates into real-world latency improvements.

4. Methodology

4.1. Principles

The core principle behind MInference is to exploit the inherent sparsity of attention weights in Large Language Models (LLMs), particularly for long-context inputs, to significantly accelerate the pre-filling stage. While attention matrices are known to be sparse, their exact sparse patterns are dynamic and context-dependent. MInference aims to efficiently identify these dynamic patterns and perform computations only on the most important parts of the attention matrix using optimized GPU kernels, thereby drastically reducing FLOPs and latency without sacrificing accuracy. This is achieved by:

1. Categorizing common dynamic sparse attention patterns observed in long-context LLMs.
2. Assigning the optimal sparse pattern for each attention head offline, based on a kernel-aware search that considers real computational costs.
3. Dynamically approximating the sparse indices (mask) for each head during inference with minimal overhead.
4. Executing sparse attention calculations using specially optimized GPU kernels tailored for these patterns.

4.2. Core Methodology In-depth (Layer by Layer)

4.2.1. Problem Formulation

The paper formalizes the problem of accelerating the pre-filling stage of long-context LLMs using sparse attention. The attention matrix $A$ with a dynamic sparse mask $M$ is formulated as:

$ A ( M ) = \mathrm { S o f t m a x } \left( \frac { 1 } { \sqrt { d } } Q K ^ { \top } - c ( 1 - M ) \right) $

Where:

• A(M) represents the attention matrix computed using the dynamic sparse mask $M$.

• $M_{i,j} \in \{0, 1\}$ is a binary dynamic sparse mask element. If $M_{i,j} = 1$, the attention weight at position (i, j) is computed; if $M_{i,j} = 0$, it is masked out.

• $Q$ is the query matrix.

• $K^\top$ is the transpose of the key matrix.

• $d$ is the dimension of the query and key vectors (often denoted as $d_h$).

• $c$ is a large constant (e.g., 1e5). When $M_{i,j} = 0$, the term $c(1-M_{i,j})$ becomes $c$, resulting in a very small (approaching zero) value after the softmax function, effectively pruning unimportant attention weights.

  The goal of the dynamic sparse attention system is to achieve significant speedup with minimal overhead while preserving as much of the original attention weights as possible. This objective is formally expressed as:

$ \begin{array} { r l } & { \operatorname* { m i n } \quad | \boldsymbol { A } ( \boldsymbol { M } ) - \boldsymbol { A } _ { \mathrm { d e n s e } } | , } \ & { \operatorname* { m i n } t _ { \mathrm { s p a r s e } } ( \boldsymbol { M } ) + t _ { \mathrm { o v e r h e a d } } ( \boldsymbol { M } ) , } \end{array} $

Where:

• $|\boldsymbol{A}(\boldsymbol{M}) - \boldsymbol{A}_{\mathrm{dense}}|$ represents the difference (e.g., L1 or L2 norm) between the sparse attention matrix A(M) and the dense attention matrix $A_{\mathrm{dense}}$. The objective is to minimize this difference, ensuring accuracy is maintained.
• $t_{\mathrm{sparse}}(\boldsymbol{M})$ represents the time spent on the dynamic sparse attention computation itself.
• $t_{\mathrm{overhead}}(\boldsymbol{M})$ represents the time spent on estimating the approximate dynamic sparse pattern (i.e., building the mask $M$). The objective is to minimize the sum of these two times to maximize overall efficiency.

4.2.2. Attention Sparsity Patterns

Through extensive analysis of long-context prompts, the authors identified three general sparse attention patterns that occur in LLMs. These patterns allow for efficient sparse computation on GPUs.

The following are the characteristics and differences between these patterns (Table 1 from the original paper):

1. A-shape pattern:

   • Description: Attention weights for these heads are primarily concentrated on initial tokens (global tokens) and within local windows. This pattern is relatively stable and structured.
   • Characteristics (from Table 1): Static structured spatial distribution, Low latency on GPU, Zero time to build the index (as it's static).
   • Relevance: This pattern is similar to approaches like StreamingLLM, which maintain a few global tokens and local sliding windows.

2. Vertical-Slash (VS) pattern:

   • Description: Attention weights are concentrated on specific tokens (forming "vertical lines" in the attention matrix) and tokens at fixed intervals (forming "slash lines"). The positions of these vertical and slash lines dynamically change with the context content. This dynamic nature makes them difficult to capture with only local windows or A-shape patterns.
   • Characteristics (from Table 1): Dynamic structured spatial distribution, Medium latency on GPU, Small time to build the index.
   • Visualization: Figure 3(a) and Figure 4 (image/4.jpg) illustrate this pattern. An example dynamic sparse mask for a Vertical-Slash pattern is shown in Figure 7 (image/7.jpg).

3. Block-Sparse pattern:

   • Description: This is the most dynamic sparsity pattern, exhibiting a more dispersed distribution of attention weights. Despite its dynamism, the weights show spatial clustering, meaning non-zero attention values tend to be close to each other. The authors found that distances between nearest non-zero neighbors are generally concentrated around 5 (Figure 3(b)). This suggests that attention can be effectively approximated by selecting important blocks.

   • Characteristics (from Table 1): Dynamic structured spatial distribution, Low latency on GPU, Small time to build the index.

   • Visualization: Figure 3(a) and Figure 4 (image/4.jpg) illustrate this pattern.

     Figure 3 (image/3.jpg) illustrates these patterns, their spatial clustering, and their recall efficiency:

• Figure 3(a) visualizes attention weights from different attention heads showcasing the A-shape, Vertical-Slash, and Block-Sparse patterns. The yellow areas indicate computed parts.
• Figure 3(b) shows that for a 128k prompt, the distances between non-zero attention weights and their top-K nearest non-zero neighbors are concentrated around 5, indicating strong spatial clustering.
• Figure 3(c) demonstrates that MInference's identified patterns are significantly more efficient than Top-K methods in retrieving attention scores for a given FLOPs budget. For example, Top-K methods struggle with the Block-Sparse pattern, while MInference retrieves scores more efficiently and accurately.

4.2.3. Kernel-Aware Optimal Sparse Pattern Search

To achieve the best accuracy under a limited FLOPs budget, MInference employs an offline Kernel-Aware Optimal Sparse Pattern Search method. This step determines which of the three sparse patterns (A-shape, Vertical-Slash, Block-Sparse) will be used for each attention head in the LLM, along with its optimal configuration (e.g., number of vertical/slash lines, number of top-k blocks).

The search process is detailed in Algorithm 1:

Algorithm 1 Kernel-Aware Sparse Pattern Search

Input: $Q, K, V \in \mathbb{R}^{S \times d_h}$, patterns search space $\rho$, target FLOPs $t$, initialized search space $\sigma$

# Build kernel-aware search space for i ← 1 to |σ| do t_i ← FLOPs_in_kernel(σ_i) while |t_i - t| > ϵ do σ_i ← ChangeSpace(σ_i, p_i) t_i ← FLOPs_in_kernel(σ_i) end while ρ ← ρ ∪ σ_i end for # Search for optimal head pattern p_best ← ∅ y ← Softmax(QKᵀ / √d) for i ← 1 to |ρ| do y_i ← SparseAttention(QKᵀ / √d, ρ_i) p_best ← argmin(y_i − y, p_best) end for return p_best

Explanation of Algorithm 1:

1. Input:

   • $Q, K, V \in \mathbb{R}^{S \times d_h}$: The query, key, and value matrices for a reference example. $S$ is the sequence length, $d_h$ is the head dimension.
   • $patterns search space ρ$: A collection of candidate sparse patterns and their configurations.
   • target FLOPs t: A predefined computational budget in terms of FLOPs for sparse attention.
   • $initialized search space σ$: An initial set of pattern configurations to explore.

2. Build kernel-aware search space:

   • The algorithm iterates through an $initialized search space σ$ (e.g., different initial settings for Vertical-Slash or Block-Sparse patterns).
   • For each candidate configuration $\sigma_i$, it calculates its actual FLOPs using $FLOPs_in_kernel(σ_i)$. This is kernel-aware because it measures the real FLOPs within the GPU kernels, which is crucial for accurate latency prediction, unlike conceptual estimations.
   • It then iteratively adjusts the configuration $\sigma_i$ using ChangeSpace until its FLOPs ($t_i$) are within a small epsilon ($\epsilon$) range of the target FLOPs t. This ensures all candidates in the final $search space ρ$ have comparable computational costs.
   • The adjusted configurations are added to the $patterns search space ρ$.

3. Search for optimal head pattern:

   • p_best is initialized as empty, which will store the optimal pattern and its settings.

   • $y ← Softmax(QKᵀ / √d)$: The full, dense attention output $y$ is computed for the reference example. This serves as the ground truth.

   • The algorithm then iterates through all candidate sparse patterns and settings ($\rho_i$) in the built $search space ρ$.

   • For each $\rho_i$, it computes the sparse attention output $y_i$ using $SparseAttention(QKᵀ / √d, ρ_i)$.

   • $p_best ← argmin(y_i − y, p_best)$: The objective criterion for selecting the best pattern is the recall of the attention output. The pattern that minimizes the difference (e.g., L2 norm or a similar metric indicating reconstruction error) between its sparse attention output $y_i$ and the dense attention output $y$ is chosen as p_best. This approach implicitly leverages information from the $V$ matrix and FlashAttention (to reduce GPU memory) for an end-to-end selection.

     This offline search is performed once per model to assign a fixed pattern type and its parameters to each attention head. The distribution of optimal patterns across LLaMA-3-8B layers is shown in Figure 11 (image/11.jpg), indicating that Vertical-Slash is the most common pattern.

4.2.4. Sparsity Indices Approximation and Dynamic Sparse Attention Calculation

During the actual inference stage, MInference performs an online estimation of the attention matrix to dynamically determine the spatial distribution of sparse indices for each head, based on its pre-assigned pattern and the specific input. After obtaining these dynamic sparse masks, optimized GPU kernels are used for efficient computation. It's important to note that for A-shape heads, the sparse mask is static, so there is no overhead in building dynamic masks; only sparse calculation is performed.

The overall architecture of the three sparse methods in MInference is illustrated in Figure 4 (image/4.jpg), showing how sparse computation (left) is combined with dynamic approximation (right) for Vertical-Slash and Block-Sparse heads, while A-shape relies solely on sparse computation with a fixed pattern.

4.2.4.1. Vertical-Slash Head

For Vertical-Slash heads, the process involves approximating the vertical and slash lines and then performing sparse attention computation.

Algorithm 2 Vertical-Slash Head

Input: $Q, K, V \in \mathbb{R}^{S \times d_h}, k_v, k_s \in \mathbb{N}$

# Approximate vertical and slash pattern (last_q = 64)
 ← softmax(Q_[-last_q:] Kᵀ / √d + m_casual)

# Indices of top k_v vertical line, sum in vertical
i_v ← argtopk(sum_v(Â), k_v)

# Indices of top k_s slash line, sum in slash
i_s ← argtopk(sum_s(Â), k_s)

# Build sparse attention index
i_vs ← sparseformat(i_v, i_s)

# Final dynamic sparse attention scores (only index block)
A ← softmax(sparse(QKᵀ, i_vs)/√d)

# Sparse mixed scores and values
y ← sparse(AV, i_vs)

return y

Explanation of Algorithm 2:

1. Input:

   • $Q, K, V \in \mathbb{R}^{S \times d_h}$: Query, Key, and Value matrices for the current input.
   • $k_v$: Number of top vertical lines to retain.
   • $k_s$: Number of top slash lines to retain.

2. Approximate vertical and slash pattern:

   • $Â ← softmax(Q_[-last_q:] Kᵀ / √d + m_casual)$: Due to the continuous nature of vertical and slash lines (meaning they often span the full sequence length), an efficient approximation is used. Instead of computing the full attention matrix, a partial attention matrix $\widehat{A}$ is estimated by multiplying only the $last_q$ query vectors ($Q_{[-last\_q:]}$, typically $last_q = 64$) with all key vectors $K^\top$. This provides a snapshot of where strong attention patterns might lie. m_casual is the causal mask.

3. Indices of top $k_v$ vertical line, sum in vertical:

   • $i_v \leftarrow \mathrm{argtopk}(\mathrm{sum}_v(\widehat{A}), k_v)$: From the approximated matrix $\widehat{A}$, the algorithm sums attention scores vertically ($sum_v$) to find the key token positions that receive the most attention. The $top-k_v$ such key indices are selected as $i_v$. These form the vertical lines.

4. Indices of top $k_s$ slash line, sum in slash:

   • $i_s \leftarrow \mathrm{argtopk}(\mathrm{sum}_s(\widehat{A}), k_s)$: Similarly, attention scores are summed along slash lines ($sum_s$) in $\widehat{A}$ to identify the $top-k_s$ slash line indices $i_s$. These slash lines capture attention to tokens at fixed intervals.

5. Build sparse attention index:

   • $i_{vs} \leftarrow \mathrm{sparseformat}(i_v, i_s)$: The identified vertical and slash indices are combined and converted into a sparse format $i_{vs}$ suitable for sparse computation on GPUs. Figure 7 (image/7.jpg) illustrates this combined mask.

6. Final dynamic sparse attention scores:

   • $A \leftarrow \mathrm{softmax}(\mathrm{sparse}(QK^\top, i_{vs})/\sqrt{d})$: The actual attention weights $A$ are computed, but only for the entries specified by the sparse index $i_{vs}$.

   • $y \leftarrow \mathrm{sparse}(AV, i_{vs})$: Finally, the sparse attention output $y$ is computed by multiplying the sparse attention weights $A$ with the value matrix $V$, again only for the active sparse indices.

     The implementation of Vertical-Slash involves two custom kernels: Vertical-Slash sparse index kernel (Algorithm 4 in Appendix C.4.2) and Vertical-Slash sparse FlashAttention kernel (Algorithm 5 in Appendix C.4.2).

Algorithm 4 Vertical-Slash Sparse Index Kernel (simplified summary from Appendix C.4.2) This algorithm builds the index for each row of blocks. It performs a point-range two-way merge algorithm where vertical indices are treated as points and slash indices are converted to ranges based on the row index. The output consists of merged ranges (represented by block indexes) and separate column indexes. The time complexity to build an index for a row is $O(k_v + k_s)$.

Algorithm 5 Vertical-Slash Sparse FlashAttention Kernel (simplified summary from Appendix C.4.2) This is a hybrid kernel combining block-sparse attention and PIT (Permutation Invariant Transformation) sparse attention. A thread block first loops through the block indexes (the "block part") and then loops through the column indexes grouped by block size (the "PIT part"). The latency of this hybrid kernel is linearly related to the total area of the computed blocks and columns.

4.2.4.2. Block-Sparse Head

For Block-Sparse heads, mean pooling is used to efficiently approximate block-level attention weights.

Algorithm 3 Block-Sparse Head

Input: $Q, K, V \in \mathbb{R}^{S \times d_h}, k_b \in \mathbb{N}$

# Approximate block-sparse pattern (block_size = 64)
Q̂ ← MeanPooling(Q, block_size)
K̂ ← MeanPooling(K, block_size)
 ← softmax(Q̂K̂ᵀ / √d + m_casual)

# Indices of top k_b blocks
î_b ← argtopk(Â, k_b)

# Build sparse attention index
î_b ← sparseformat(î_b)

# Final dynamic sparse attention scores (only index block)
A ← softmax(sparse(QKᵀ, î_b)/d)

# Sparse mixed scores and values
y ← sparse(AV, î_b)

return y

Explanation of Algorithm 3:

1. Input:

   • $Q, K, V \in \mathbb{R}^{S \times d_h}$: Query, Key, and Value matrices.
   • $k_b$: Number of top blocks to retain.

2. Approximate block-sparse pattern:

   • $\widehat{Q} \leftarrow \mathrm{MeanPooling}(Q, \mathrm{block\_size})$: Mean pooling is applied to the query matrix $Q$ in blocks of block_size (e.g., block_size = 64) to obtain a downsampled $\widehat{Q}$.
   • $\widehat{K} \leftarrow \mathrm{MeanPooling}(K, \mathrm{block\_size})$: Similarly, mean pooling is applied to the key matrix $K$ to obtain $\widehat{K}$.
   • $\widehat{A} \leftarrow \mathrm{softmax}(\widehat{Q}\widehat{K}^\top / \sqrt{d} + m_{casual})$: The downsampled $\widehat{Q}$ and $\widehat{K}$ are multiplied to get an estimated block-level attention weights matrix $\widehat{A}$. Since mean pooling and matrix multiplication are approximately commutative, this efficiently approximates the block-sparse pattern of the actual attention weights with minimal overhead.

3. Indices of top $k_b$ blocks:

   • $\hat{\iota}_b \leftarrow \mathrm{argtopk}(\widehat{A}, k_b)$: From the estimated block-level attention weights $\widehat{A}$, the $top-k_b$ blocks with the highest attention scores are selected as indices $\hat{\iota}_b$.

4. Build sparse attention index:

   • $\hat{\iota}_b \leftarrow \mathrm{sparseformat}(\hat{\iota}_b)$: These block indices are converted into a sparse format $\hat{\iota}_b$ suitable for GPU kernel computation.

5. Final dynamic sparse attention scores:

   • $A \leftarrow \mathrm{softmax}(\mathrm{sparse}(QK^\top, \hat{\iota}_b)/\sqrt{d})$: The actual attention weights $A$ are computed, but only for the entries within the selected $top-k_b$ blocks.

   • $y \leftarrow \mathrm{sparse}(AV, \hat{\iota}_b)$: The sparse attention output $y$ is computed by multiplying the sparse attention weights $A$ with the value matrix $V$, again only for the active sparse blocks.

     The Block-Sparse kernel implementation (Appendix C.4.1) is based on the Triton version of FlashAttention. It loops through the top-K blocks in a row, with latency linearly related to the number of blocks, providing a speedup ratio $s_p = \frac{S}{2B \times k_b}$, where $S$ is sequence length, $B$ is block size, and $k_b$ is the number of blocks.

Figure 4: The three sparse methods in MInference.

该图像是一个示意图，展示了MInference中的三种稀疏计算方法，包括Λ形头、垂直斜杠头和块稀疏头。这些方法用于提高长序列处理的效率，其中左侧为稀疏计算示意，右侧为动态近似方法。

5. Experimental Setup

5.1. Datasets

The experiments used a diverse set of long-context benchmarks to thoroughly evaluate MInference's effectiveness across various tasks and context lengths.

5.1.1. InfiniteBench

• Source: $[ZCH+24]$
• Scale & Characteristics: Consists of 10 distinct tasks, covering a wide range of long-context scenarios. The average context length for tasks in InfiniteBench is approximately 214K tokens, with a total of 3,992 examples.
• Domain & Tasks:
  • Retrieval tasks: PassKey retrieval, Number retrieval, KV retrieval. These test the model's ability to extract specific information from long texts.
  • Realistic tasks: Question-answering (based on novels, drama scripts, Chinese texts), coding (debugging large code repositories), dialogue, and summarization (entire novel summarization).
  • Math reasoning: Identifying the largest/smallest number in arrays.
• Why chosen: This benchmark offers a comprehensive assessment of long-context LLMs across different modalities of understanding and generation.

5.1.2. RULER

• Source: $[HSK+24]$
• Scale & Characteristics: A challenging synthetic benchmark with 13 complex tasks grouped into 4 categories. It includes subsets with various prompt lengths, ranging from 4K to 128K tokens, with 2,600 examples per length.
• Domain & Tasks:
  • Retrieval: Single Needle-in-a-Haystack (S-NIAH), Multi-keys Needle-in-a-Haystack (MK-NIAH), Multi-values Needle-in-a-Haystack (MV-NIAH), Multi-queries Needle-in-a-Haystack (MQ-NIAH). These tasks test the retrieval of specific information amidst distractors.
  • Multi-hop Tracing: Variable Tracking (VT), requiring models to trace variable bindings.
  • Aggregation: Common Words Extraction (CWE) and Frequent Words Extraction (FWE), testing summarization and frequency analysis.
  • Question Answering (QA): Extends short-context QA by adding distracting paragraphs, challenging the model to find relevant information.
• Why chosen: Designed to reveal the actual context window size of models and test complex multi-hop reasoning and aggregation capabilities in long contexts.

5.1.3. Needle In A Haystack

• Source: [Kam23]
• Scale & Characteristics: This is a long-context retrieval benchmark where a "needle" (specific, targeted information) is embedded within a large "haystack" (complex body of text). The test varies both context length (from 1K to 1M tokens) and needle depth (position of the crucial information). 750 examples were used.
• Why chosen: Directly measures an LLM's ability to identify and utilize specific information across extremely long contexts and different positions, highlighting long-range dependency understanding.

5.1.4. PG-19

• Source: $[RPJ+20]$
• Scale & Characteristics: A dataset suitable for long-context language modeling tasks, containing texts up to 500K tokens in length. For experiments, 1,000 random samples longer than 100K tokens were used.
• Why chosen: Evaluates language modeling performance (using perplexity) in long-context scenarios, which is a fundamental capability of LLMs.

5.2. Evaluation Metrics

For each evaluation metric, a complete explanation is provided:

5.2.1. Accuracy

• Conceptual Definition: Accuracy is a common metric used to evaluate the correctness of a model's predictions. In the context of LLMs and tasks like question answering or retrieval, it measures the percentage of predictions that exactly match the ground truth or are semantically equivalent. For Needle In A Haystack, it typically refers to whether the model correctly retrieves the embedded "needle."
• Mathematical Formula: $ \mathrm{Accuracy} = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}} $
• Symbol Explanation:
  • Number of Correct Predictions: The count of instances where the model's output is identical or considered equivalent to the true answer.
  • Total Number of Predictions: The total number of evaluation instances or examples.

5.2.2. Perplexity (PPL)

• Conceptual Definition: Perplexity is a measure of how well a probability model predicts a sample. In language modeling, it quantifies how well a language model predicts a sequence of words. A lower perplexity score indicates that the model is better at predicting the next word in a sequence, suggesting a higher quality language model. It can be interpreted as the geometric mean of the inverse probabilities of the words in the test set.
• Mathematical Formula: Given a test set of $N$ words, $W = (w_1, w_2, \dots, w_N)$, perplexity is defined as: $ \mathrm{PPL}(W) = P(w_1, w_2, \dots, w_N)^{-\frac{1}{N}} = \sqrt[N]{\frac{1}{P(w_1, w_2, \dots, w_N)}} $ Using the chain rule of probability, P(w_1, \dots, w_N) = \prod_{i=1}^N P(w_i | w_1, \dots, w_{i-1}), so: $ \mathrm{PPL}(W) = \sqrt[N]{\prod_{i=1}^N \frac{1}{P(w_i | w_1, \dots, w_{i-1})}} = \exp\left(-\frac{1}{N}\sum_{i=1}^N \log P(w_i | w_1, \dots, w_{i-1})\right) $ Note that -\frac{1}{N}\sum_{i=1}^N \log P(w_i | w_1, \dots, w_{i-1}) is the cross-entropy loss (or negative log-likelihood) per word.
• Symbol Explanation:
  • $W$: The test sequence of words.
  • $N$: The total number of words in the test sequence.
  • $w_i$: The $i$-th word in the sequence.
  • $P(w_1, w_2, \dots, w_N)$: The joint probability of the entire sequence $W$ according to the language model.
  • $P(w_i | w_1, \dots, w_{i-1})$: The probability of the $i$-th word given all preceding words in the sequence, as predicted by the language model.
  • $\exp(\cdot)$: The exponential function.
  • $\log(\cdot)$: The natural logarithm.

5.2.3. Latency

• Conceptual Definition: Latency refers to the time delay between the initiation of a request and the completion of the response. In this context, it specifically measures the end-to-end time taken for the pre-filling stage of LLM inference to process an input prompt and become ready to generate the first output token. Lower latency implies faster response times and a better user experience.
• Measurement: Typically measured in milliseconds (ms) or minutes.

5.2.4. Speedup

• Conceptual Definition: Speedup is a metric used to quantify the performance improvement achieved by an optimization. It is calculated as the ratio of the execution time of the unoptimized (baseline) method to the execution time of the optimized method. A speedup of 10x means the optimized method is 10 times faster.
• Mathematical Formula: $ \mathrm{Speedup} = \frac{\text{Latency}{\text{Baseline}}}{\text{Latency}{\text{Optimized}}} $
• Symbol Explanation:
  • $\text{Latency}_{\text{Baseline}}$: The time taken by the original or unoptimized method (e.g., FlashAttention for dense computation).
  • $\text{Latency}_{\text{Optimized}}$: The time taken by the proposed MInference method.

5.3. Baselines

The paper compares MInference against five training-free sparse attention approaches as baselines, all focusing on optimizing the pre-filling stage. During the decoding stage, all baselines retain dense computation.

1. StreamingLLM [XTC+24]:

   • Description: This method combines attention sinks (a fixed number of initial tokens that always attend to all other tokens) with sliding window local attention.
   • Configuration: For experiments, it used 1k global tokens (attention sinks) and 4k local windows. This baseline effectively represents the A-shape pattern identified by MInference.

2. StreamingLLM w/ dilated [BPC20]:

   • Description: An extension of StreamingLLM that incorporates dilated local windows, meaning attention within the local window is applied with intervals.
   • Configuration: Used 1k global tokens and 8k dilated attention windows with an interval of 1.

3. StreamingLLM w/ strided [CGRS19]:

   • Description: Another variant of StreamingLLM that retains local windows while also adding dilated attention (presumably in a different configuration or alongside the local window).
   • Configuration: Used 1k global tokens, 2k local windows, and 4k dilated attention windows with an interval of 1.

4. InfLLM [XZH+24]:

   • Description: This method utilizes a memory unit to process streaming long sequences, aiming to unveil the intrinsic capacity of LLMs for understanding extremely long sequences without training.
   • Configuration: Following the original paper, 128 global tokens and 8k local windows were used.

5. Ours w/ static:

   • Description: This is an ablation variant of MInference itself, where the Vertical-Slash and Block-Sparse heads (which are normally dynamic) use static sparse indices instead of dynamically built ones.
   • Purpose: To demonstrate the necessity and effectiveness of MInference's dynamic strategy.

5.3.1. Models Used

The experiments were conducted on several state-of-the-art long-context LLMs:

• LLaMA-3-8B-Instruct-262k: A LLaMA-3 variant optimized for long contexts, leveraging NTK-aware interpolation and Ring Attention.
• LLaMA-3-8B-Instruct-1048k: Similar to the 262K version but supporting up to 1M tokens.
• GLM-4-9B-1M: An LLM improved for a 1M context window.
• Yi-9B-200K: An LLM balancing long-context performance with general capabilities.
• Phi-3-Mini-128K: A smaller but powerful model with up to 128K context, powered by LongRoPE [DZZ+24]. (Tested on Needle In A Haystack)
• Qwen2-7B-128K: A recently updated Qwen series model with up to 128K context. (Tested on Needle In A Haystack)
• LLaMA-3-70B-1M: A larger LLM (tested for scaling-up analysis).

5.3.2. Implementation and Hardware Details

• Decoding Strategy: Greedy decoding was used for all experiments to ensure stable results.
• Implementation: A custom implementation in PyTorch, built upon FlashAttention [Dao24], Triton [TKC19], and the dynamic sparse compiler PIT [ZJZ+23].
• Hardware: All latency experiments were conducted on a single Nvidia A100 GPU using bfloat16 precision.
• Configuration:
  • Target FLOPs t: Set to be equivalent to 1k global tokens and 4k local window tokens in the A-shape pattern.
  • $last_q$: Set to 64 in the Vertical-Slash pattern for approximation.
  • block_size: Set to 64 in the Block-Sparse pattern for approximation.
  • The Kernel-Aware Optimal Sparse Pattern Search took approximately 15 minutes on a single A100, using a single sample from KV retrieval synthetic data with 30k token inputs as a validation set. The same optimal sparse pattern configuration was used for both LLaMA-3-8B-Instruct-262K and LLaMA-3-8B-Instruct-1M.
• Custom Optimizations for Single A100 (for >50k tokens):
  1. Tensor Splitting: Attention was split by head and MLP by sequence dimension to maintain 100% GPU utilization for long-context scenarios.
  2. Reduction of Intermediate Variables: Minimized memory allocation by removing the attention mask and implementing causal mask logic directly within the kernel.
  3. Elimination of Unnecessary Computations: Only logits for the last token in the pre-filling stage were retained for the LM Head Linear layer, as these are the only meaningful outputs.

Figure 10: The latency breakdown of a single attention kernel for three patterns and FlashAttention [Dao24] across different context windows in a single A100, including the index time for dynamic sparse approximation an buildiyamparsiyAt 0k tokens, the atey the four kernes is very cose an al arle h 1ms. At 1M tokens, the latency for A-shape is $1 6 4 \\mathrm { m s }$ .

该图像是一个图表，展示了在单一 A100 GPU 上，不同上下文窗口中三种模式（A-shape、Vertical-Slash、Block-Sparse）以及 FlashAttention 的延迟情况（单位：毫秒）。在 1M tokens 的情况下，A-shape 模式的延迟为 164 ext{ ms}。

Figure 11: u isrutrear e pats tmosWeu hetal sr a configuration for both LLaMA-3-8B-Instruct-262K and LLaMA-3-8B-Instruct-1M.

该图像是一个示意图，展示了在不同层次（Layer）中，LLama-3-8B-Instruct-262K/1M和Yi-9B-200K模型各注意力头（Head Number）的分布情况，标记了三种稀疏模式：A-shape、Vertical-Slash和Block-Sparse。

Figure 12: Fur The istriutiosparsy the kere across differnt context windows rer the proot th keehacybloovecpar earsy henusasAt with a causal mask.

该图像是图表，展示了不同上下文窗口大小下的稀疏性。X轴为上下文窗口大小，Y轴为内核中的稀疏性。数据中包含三种模式：A-shape、Vertical-Slash和Block-Sparse，结果表明稀疏性在窗口大小增加时趋于稳定。

6. Results & Analysis

6.1. Core Results Analysis

6.1.1. InfiniteBench Performance

The results on InfiniteBench demonstrate MInference's superior overall performance compared to baseline methods, particularly in maintaining accuracy across diverse long-context tasks.

The following are the results from Table 2 of the original paper:

• Overall Performance: MInference achieves the best average performance across all models (e.g., 38.8 for LLaMA-3-8B-262K, 37.7 for Yi-9B-200K, 47.0 for GLM-4-9B-1M). It matches or slightly surpasses the original full attention baseline in some tasks, despite the significant acceleration.
• Retrieval Tasks: MInference performs exceptionally well on retrieval-related tasks (Retr.PassKey, Retr.Num, Retr.KV), often achieving 100% accuracy where baselines like StreamingLLM and InfLLM struggle (e.g., StreamingLLM scores 0.8% on Retr.KV for LLaMA-3-8B-262K, while MInference achieves 12.8%, which is closer to the 14.4% of the full model). This highlights the effectiveness of MInference's dynamic sparse patterns in capturing crucial long-range dependencies for retrieval.
• Other Tasks: MInference maintains strong performance in natural language tasks (summarization, QA, code debugging) and Math.Find, indicating its versatility.
• Comparison to Baselines: StreamingLLM variants (with dilated/strided attention) generally perform poorly, especially on retrieval tasks, suggesting that extending local windows with fixed intervals is insufficient for complex long-context understanding. InfLLM shows better performance than StreamingLLM but still falls short of MInference on average and on certain retrieval tasks.
• Importance of Dynamism: The Ours w/ static variant consistently underperforms Ours (e.g., 31.9 vs 38.8 for LLaMA-3-8B-262K), especially in Retr.KV (dropping from 12.8% to 0.2%), confirming the critical need for dynamic sparse indices.

6.1.2. RULER Performance

RULER results further validate MInference's capability in preserving long-context performance even in challenging multi-hop or aggregation tasks.

The following are the results from Table 3 of the original paper:

• Effective Context Window: MInference significantly extends the effective context window (context length with performance over 85%) for base models. For LLaMA-3-8B-262K, it increases from 16K (full attention) to 32K. For GLM-4-9B-1M, it extends from 64K (full attention) to 64K (matching the full model's effective context). This indicates MInference's ability to retain long-range dependencies.
• Performance at Scale: MInference consistently outperforms all baselines across increasing context lengths. For LLaMA-3-8B-262K, MInference achieves 87.0% average accuracy, surpassing the full model's 84.4% and InfLLM's 62.9%.
• Robustness: Even for lengths beyond 32K, where full attention models start to decline, MInference maintains strong performance, demonstrating its robustness in handling complex long-context challenges that StreamingLLM and its variants fail to address (often dropping to very low accuracy scores at longer contexts).

6.1.3. Language Modeling Performance (PG-19)

MInference yields the best perplexity results among sparse attention approaches on the PG-19 dataset, showing minimal degradation compared to the full attention baseline.

Figure 5 (image/5.jpg) shows the perplexity results:

• For LLaMA-3-8B-262K, MInference's perplexity is only 0.2 higher than full attention at 100K tokens, but 0.75 lower than StreamingLLM.
• For Yi-9B-200K, MInference's perplexity is 0.2 higher than full attention, but 0.25 lower than StreamingLLM. This indicates that MInference effectively preserves language modeling capabilities in long contexts while achieving significant speedups.

6.1.4. Needle In A Haystack Performance

MInference effectively preserves the model's ability to retrieve information across various context windows (1K to 1M tokens) and needle depths.

Figure 1(a) (image/1.jpg) shows MInference's performance on LLaMA-3-8B-1M, which is strong across the board. In contrast, StreamingLLM (Figure 6, image/6.jpg) and InfLLM (Figure 8, image/8.jpg) experience a sharp decline in performance once the critical information (needle) falls outside the range of their global tokens or local windows. Figure 9 (image/9.jpg) further shows consistent performance for MInference across GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, and Qwen2-7B-128K models. There is even a slight performance improvement around the 100k context length for Yi-9B-200K and Phi-3-Mini-128K, which might be due to a more focused attention on important tokens.

6.1.5. Latency and Breakdown

MInference achieves substantial latency reductions in the pre-filling stage.

Figure 1(b) (image/1.jpg) and Figure 10 (image/10.jpg) illustrate the latency and its breakdown.

• Speedup:
  • 100K tokens: 1.8x speedup
  • 300K tokens: 4.1x speedup
  • 500K tokens: 6.8x speedup
  • 1M tokens: 10x speedup
• Real-world Impact: For a 1M token prompt, MInference reduces pre-filling latency from 30 minutes to 3 minutes on a single A100 GPU. With distributed setups (8x A100 GPUs using tensor parallel and context parallel), this can be further reduced to 22 seconds.
• Overhead Analysis: The overhead for dynamic sparse index building accounts for approximately 5%-20% of the total time, with the remaining time dedicated to dynamic sparse calculation. This overhead is higher for Block-Sparse due to MeanPooling and block-level matmul. Memory overhead for sparse indexing is small ($<160MB$ for LLaMA-3-8B in 1M context).
• Kernel Performance (Figure 10): Block-Sparse is the fastest pattern (30x speedup over FlashAttention for 1M contexts), followed by Vertical-Slash (13x speedup) and A-shape (50% slower than Vertical-Slash at 1M).

6.1.6. Integration with KV Cache Compression Methods

MInference is compatible with KV cache compression techniques.

The following are the results from Table 5 of the original paper:

• When combined with SnapKV [LHY+24], a KV cache compression method, MInference shows minimal change in performance for most tasks on InfiniteBench, with the average score even slightly increasing (37.3 vs 36.0 for LLaMA-3 w/ SnapKV). This demonstrates its practical value as an optimization for serving long-context LLMs.

6.1.7. Scaling-up on Larger LLMs

MInference maintains strong performance on larger models.

The following are the results from Table 6 of the original paper:

• On LLaMA-3-70B-1M, MInference achieves 47.0% average accuracy, surpassing LLaMA-3-70B-262K's 46.5% and InfLLM's 39.2%.
• It matches or slightly improves performance compared to full attention in dynamic tasks like KV retrieval (39.0% vs 34.0%). This confirms its scalability and effectiveness beyond 8B models.

6.2. Ablation Studies / Parameter Analysis

6.2.1. Contributions of Different Components

An ablation study was conducted to evaluate the contributions of MInference's different components.

The following are the results from Table 4 of the original paper:

• Ours w/ static: (from Table 2) This variant uses static sparse masks for Vertical-Slash and Block-Sparse patterns. It showed significant performance degradation (31.9% average for LLaMA-3-8B-262K compared to 38.8% for full Ours), especially in dynamic tasks like KV retrieval (accuracy dropped from 12.8% to 0.2%). This strongly emphasizes the necessity of the dynamic strategy and dynamically built sparse indices.

• Ours w/ only A-shape: This is equivalent to StreamingLLM. As seen in Table 2, its average performance (23.8%) is significantly lower than full MInference (38.8%), indicating that A-shape alone can only capture information within local windows and struggles with long-range dependencies outside this scope.

• Ours w/ only block-sparse: Using only the Block-Sparse pattern (average 18.7%) results in a significant performance decline, similar to StreamingLLM. This suggests that while Block-Sparse is effective, it cannot solely capture all necessary attention patterns.

• Ours w/ only vertical-slash: This variant (average 37.1%) manages to preserve most of the performance, indicating a good balance between dynamicity and StreamingLLM's pattern. However, it still falls behind the full version of MInference (38.8%), especially in KV retrieval (5.0% vs 12.8%), highlighting the complementary nature of all three patterns.

  The following are the results from Table 8 of the original paper:

  Methods	En.Sum	En.QA	En.MC	En.Dia	Zh.QA	Code.Debug	Math.Find	Retr.PassKey	Retr.Num	Retr.KV	Avg.
  Ours	20.5	12.9	65.9	7.5	12.5	22.3	33.1	100.0	100.0	12.8	38.8
  Ours w/ only vertical	13.7	6.2	30.1	2.0	6.5	7.9	1.7	65.4	52.7	0.0	18.6
  Ours w/ only slash	18.4	11.5	60.1	3.0	11.4	22.1	28.4	100.0	100.0	4.2	35.9

• Ours w/ only vertical: Using only vertical lines (and a top-1 slash line) results in a significant performance drop (average 18.6%), especially in retrieval tasks (Retr.KV 0.0%). This is similar to the performance of only block-sparse, suggesting vertical lines alone are insufficient.

• Ours w/ only slash: Using only slash lines (and a top-1 vertical line) preserves most of the performance (average 35.9%), but still suffers in highly dynamic tasks like KV retrieval (4.2%). This further indicates that MInference's strength lies in the intelligent combination of all identified patterns.

6.2.2. Pattern Distribution

Figure 11 (image/11.jpg) shows the distribution of the optimal head configuration obtained through the kernel-aware search.

• Dominant Pattern: The Vertical-Slash pattern is the most prevalent, accounting for over 90% of the attention heads in LLaMA-3-8B.
• Distribution of other patterns: The Block-Sparse pattern is primarily found in intermediate to later layers, while the A-shape pattern is concentrated in the middle layers.
• Generalizability: The optimal patterns generally align across different models. The fact that the same configuration works very well for both LLaMA-3-8B-Instruct-262K and LLaMA-3-8B-Instruct-1M demonstrates the generalizability of the optimal sparse pattern.

6.2.3. Sparsity in Kernel Distribution

Figure 12 (image/12.jpg) illustrates the sparsity distribution of the three patterns during actual kernel computation across different context windows.

• High Sparsity: When context windows exceed 200K tokens, the actual sparsity for all three patterns surpasses 90%.
• Speedup Potential: For context windows exceeding 500K tokens, the sparsity relative to FlashAttention exceeds 95%, with a theoretical speedup of over 15x. Even with a 20% index-building overhead, this high sparsity ensures that the kernel achieves a speedup of over 8x. This demonstrates that the FLOPs reduction achieved by MInference is significant and effectively translates into speedup.

6.3. Case Studies

6.3.1. EN.SUM Task (Summarization)

Table 9 presents a comparison of generated summaries for the EN.SUM task from InfiniteBench (200K input length) using LLaMA-3-8B-262K.

The following are the results from Table 9 of the original paper:

• Original Summary: Provides a detailed, coherent, and factually accurate narrative.
• StreamingLLM: Generates a seemingly coherent but factually incorrect summary, introducing details not present in the original story (e.g., boat trip, fishermen, drowned men).
• StreamingLLM w/ dilated & strided: Produce largely incoherent and nonsensical outputs, indicating complete failure.
• Ours w/ static: Generates repetitive and meaningless text (def solverome2 def solverome2...), again highlighting the failure of static patterns in dynamic tasks.
• Ours (MInference): Produces a detailed and coherent narrative that closely aligns with the original story, accurately depicting main events and themes, demonstrating its superiority in generating high-quality summaries.

6.3.2. KV Retrieval Task

Table 10 presents a comparison of retrieval performance for the Retrieve.KV task (200K input length) using LLaMA-3-8B-262K.

The following are the results from Table 10 of the original paper:

• Original: Demonstrates perfect retrieval for both examples, correctly predicting the exact UUID strings.
• StreamingLLM: Generates plausible-looking but factually incorrect UUIDs, failing the retrieval task.
• StreamingLLM w/ dilated & strided, and Ours w/ static: Produce repetitive or nonsensical character sequences, indicating complete failure to retrieve the required key-value pairs.
• Ours (MInference): Performs on par with the original, accurately retrieving and predicting the exact key-value pairs for both examples. This confirms MInference's superior capability in handling KV retrieval tasks with precision and reliability.

7. Conclusion & Reflections

7.1. Conclusion Summary

This paper effectively addresses the critical challenge of high computational cost and unacceptable latency during the pre-filling stage of long-context Large Language Models (LLMs). The authors propose MInference, a novel method that leverages dynamic sparse attention with identified spatial aggregation patterns to significantly accelerate this stage. Key to MInference is the categorization of attention heads into three distinct patterns: A-shape (static, local/global windows), Vertical-Slash (dynamic, specific token/interval attention), and Block-Sparse (dynamic, spatially clustered blocks). An offline kernel-aware optimal sparse pattern search is employed to assign the most efficient pattern to each head. During inference, a fast approximation approach is used to build dynamic sparse masks for Vertical-Slash and Block-Sparse heads, enabling efficient sparse attention calculations via optimized GPU kernels.

Experimental results across a wide range of long-context benchmarks (InfiniteBench, RULER, PG-19, Needle In A Haystack) and diverse LLM architectures (LLaMA-3, GLM-4, Yi, Phi-3, Qwen2) demonstrate MInference's effectiveness. It achieves up to a 10x speedup for pre-filling on a single A100 GPU for 1M token prompts, reducing latency from 30 minutes to 3 minutes, all while maintaining or even slightly improving the model's accuracy and long-context capabilities. MInference is a training-free solution, making it highly practical for immediate deployment on existing LLMs.

7.2. Limitations & Future Work

The authors acknowledge several limitations and suggest future research directions:

• Overhead for Shorter Contexts: As the context length decreases, the time required to build the dynamic index becomes a more significant proportion of the total attention computation time. For example, with a 10k context, the index building overhead can increase from 5% to 30%, causing the overall end-to-end latency to approach that of FlashAttention (which performs dense computation efficiently). This overhead decreases as prompts lengthen, making MInference most beneficial for truly long-context scenarios.
• Performance vs. Sparsity Rate: When using a higher sparsity rate (i.e., retaining fewer attention weights), the model's performance may noticeably decline. This suggests a trade-off that needs careful tuning based on specific application requirements.
• Generalizability to Other LLMs: The paper notes that similar dynamic sparse attention patterns have been observed in multi-modal LLMs [WWL+24] and encoder-decoder LLMs [RSR+20]. This suggests that MInference could be applied to accelerate the pre-filling stage in these other LLM variants, which is a promising direction for future work. The analysis of T5-style Encoder Attention in Figure 13 (image/13.jpg) further supports this claim.

7.3. Personal Insights & Critique

7.3.1. Personal Insights

MInference presents a highly practical and impactful solution for long-context LLM inference. The core insight that attention sparsity is not only present but also exhibits identifiable, dynamic patterns is powerful. This moves beyond the brute-force static windowing or global token approaches, offering a more nuanced and adaptive strategy. The kernel-aware search and the development of specialized GPU kernels are crucial. Many research papers propose algorithmic improvements that look good on paper (e.g., in FLOPs), but fail to translate to real-world latency gains due to hardware limitations (e.g., memory access patterns). By being kernel-aware, MInference ensures that the algorithmic savings are realized efficiently on actual GPUs.

The training-free nature is its strongest practical advantage. LLM training is incredibly expensive, so any method that can plug into existing, pre-trained models without demanding re-training offers immediate value to developers and deployers. The ability to reduce pre-filling latency from minutes to seconds fundamentally changes the user experience for long-context LLMs, making them viable for interactive applications. The compatibility with KV cache compression methods further enhances its utility, suggesting it can be part of a broader suite of inference optimizations.

7.3.2. Critique

While MInference is impressive, a few points warrant critical reflection:

• Robustness of Offline Pattern Search: The kernel-aware optimal sparse pattern search is performed offline using a single reference example (30k token KV retrieval synthetic data). While the paper states this exhibits "strong generalization and stability across different lengths and domains," the robustness of this offline search across an extremely diverse range of real-world inputs and tasks for which LLMs are used might be a subtle point of failure. Could a pattern optimized for one type of long-context data (e.g., code) be suboptimal for another (e.g., legal documents)? An online or adaptive search that periodically refines patterns or weights different patterns based on input characteristics could be an interesting, albeit more complex, future direction.
• Fixed Hyperparameters for Approximation: The $last_q = 64$ for Vertical-Slash and block_size = 64 for Block-Sparse are fixed hyperparameters. While chosen for efficiency, the optimal $last_q$ or block_size might vary depending on the LLM architecture, context length, or even the specific attention head. More dynamic or adaptive mechanisms for these approximation parameters could yield further gains.
• Generalization of Identified Patterns: The paper states it "identifies three unique patterns." While compelling evidence is presented for their existence in tested LLMs, the universality across all LLMs (especially new architectures or those with different positional encoding schemes) is an implicit assumption. Further architectural analysis could solidify the theoretical basis for why these three patterns are truly fundamental.
• Trade-off at Shorter Contexts: The acknowledged limitation regarding index building overhead for shorter contexts means MInference isn't a universal attention optimization but rather specifically targeted at extremely long contexts. For moderate long contexts (e.g., 10K-50K tokens), the benefits might be less pronounced compared to pure FlashAttention or other methods with lower setup costs. Clearly defining the sweet spot for MInference's applicability is important for users.

7.3.3. Transferability & Application

The methods and conclusions of MInference are highly transferable:

• Multi-modal LLMs: The authors explicitly mention the discovery of similar vertical and slash patterns in multi-modal LLMs. This opens the door for MInference to accelerate computationally intensive cross-attention or self-attention layers in models like LLaVA or InternVL, especially for long image sequences or dense captioning tasks.
• Encoder-Decoder LLMs: The analysis of T5-style Encoder Attention showing similar patterns suggests MInference could benefit encoder-decoder architectures, which are prevalent in tasks like machine translation and summarization where both input and output sequences can be long.
• Edge Devices/Resource-Constrained Environments: By significantly reducing FLOPs and latency, MInference could make long-context LLMs more feasible on edge devices or in other resource-constrained environments, democratizing access to powerful AI capabilities.
• Cost Reduction: The 10x speedup translates directly into 10x cost reduction in GPU inference time, which is a major factor in the operational expenses of LLM deployment, making long-context LLMs more economically viable.