FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference

R W1: Unclear Cost Analysis of Pre-calculation on JS Divergence

We note that the pre-calculation of JS divergence involves the following computational costs:

• Representative Attention Score Computation: $O(bnd)$, where $b$ is the block size, $d$ is the head dimension, and $n$ is the sequence length.
• JS Divergence Computation: $O(bn)$, including block-wise attention distribution estimation and JS divergence computation.

We will incorporate these detailed cost analyses into our next revision to provide a clearer and more comprehensive representation.

R Q1: JS Divergence Usage and Efficiency

While the Wasserstein distance is a potential alternative, we selected Jensen-Shannon divergence (JS divergence) for both theoretical and practical reasons. JS divergence is computationally efficient and provides bounded (0 to 1), symmetric measurements, making it well-suited for our application.

We use JS divergence to evaluate whether block-wise attention estimation sufficiently approximates the true distribution. This enables flexible switching between query-aware and vertical-slash patterns. As shown in Figure 6, the additional computational cost introduced by the pattern search component is minimal, contributing less than 5% to the total latency.

R Q2: Model Structure Impact

FlexPrefill has demonstrated robust performance across diverse model architectures, including Llama, GLM, Yi, and Qwen, as shown in our experimental results. This consistent performance across different architectures suggests that our method's adaptive nature allows it to effectively accommodate various model structures.

R W2: insufficient baselines & Q3: Automatic Selection Comparison

While automatic selection methods like MoA are interesting, we focused our comparison with MInference as it represents a strong baseline. Moreover, we conducted additional experiments comparing FlexPrefill with MoA across different sequence lengths:

For Llama-3-262k-instruct:

These results demonstrate that FlexPrefill not only achieves better average performance (85.74% vs 84.58%) but also provides speedup (2.43x vs 1.31x at 128k tokens). FlexPrefill's dynamic pattern selection operates efficiently while delivering better performance-speed trade-offs.

R Q4: JS Distance Formulation

Our implementation follows the original Jensen-Shannon divergence definition using the mean distribution $m=(\bar{a}+\hat{a})/2$.

While symmetrised KL divergence could also measure distribution distances, we chose JS divergence for its symmetry and boundedness properties, which provide more stable pattern determination across different attention scenarios

R Q1: Question about the comprehensiveness of Figure 4 and comparison curves

Figure 4 demonstrates our method's key advantage of flexible performance-latency trade-offs through $\gamma$ adjustment. The single point for MInference reflects a practical limitation - as shown in Table 7 of the MInference (NeurIPS 2024) paper, their method only supports six predefined combinations of patterns and ratios, each requiring extensive offline search to determine optimal configurations. We selected their paper's recommended configuration as a representative comparison point. In contrast, FlexPrefill offers continuous adjustment through $\gamma$, enabling fine-grained control over the performance-latency trade-off without requiring additional offline search processes.

We have chosen not to include the performance of "Fixed Budget Query-Aware" due to its significantly degraded results. Notably, its poorer performance underscores the importance of our pattern selection mechanism, which is a key innovation of our FlexPrefill method.

R Q2: Question about normalization differences between Algorithms 3 and 4

The apparent difference in normalization between Algorithms 3 and 4 represents two mathematically equivalent implementations of the same concept. While Algorithm 4 explicitly shows normalization across all attention scores, Algorithm 3's mean operation over vertical and slash lines achieves the same effect through a different computational path.

We plan to revise the algorithm descriptions to better highlight this equivalence in the final version.

R Q3: The rationale behind the transition from Step 6 to Step 7 The transition from Step 6 to Step 7 follows from the properties of the Lagrangian dual problem. We now provide the key intuition:

In Step 6, the objective maximizes over λ with the expression: $\sum\limits(\exp(x_{j})/\sum\limits \exp(x_j')-\lambda)+\lambda\gamma_S$ for all $j$ where $\exp(x_j)/\sum\limits \exp(x_j')\ge\lambda$

For any fixed $\lambda$, this naturally defines a set $S$ of indices where $\exp(x_j)/\sum\limits\exp(x_j') \ge \lambda$. The optimal $\lambda^*$ will be the smallest value that satisfies the constraint in Step 7, because any larger value would reduce the objective in Step 6. Therefore, minimizing $|S|$ subject to the sum constraint in Step 7 is equivalent to finding the optimal $\lambda^*$ in Step 6.

The relationship $\gamma_a = 1 - \lambda^*\gamma_S$ emerges from complementary slackness in the KKT conditions, where $\lambda$ represents the optimal trade-off between set size and attention score coverage.

We will add more detailed intermediate steps in the final version to make this derivation more transparent.

R Presentation:

We express our sincere gratitude to your insightful comments. Regarding the presentation issues raised, we will carefully incorporate all suggested improvements in our next revision.

R W1: Insufficient Baseline Comparisons:

We here provide a direct comparison with HiP-Attention using results from their original paper for Llama 3.1 8B:

This comparison demonstrates that FlexPrefill achieves both better average performance (89.18% vs 85.4%) and better computational efficiency (2.4x vs 1.7x speedup at 128k tokens). Notably, FlexPrefill maintains significantly better performance on longer sequences, with a 20.89 percentage point advantage at 128k tokens, while also achieving higher speedup.

We note that the implementation of SampleAttention is not publicly available at this time. Once it becomes accessible, we will incorporate it into our comparison for a more comprehensive evaluation.

R W2: Concerns about potential latency increase from unbalanced computation allocation across attention heads.

We conducted comprehensive measurements for Llama 3.1 8B with 128k context and 32 heads to analyze this impact. The following experiment used the same average computational budget (11%).

We note that Unbalanced Budget's computational budget was significantly unbalanced, varied from 1.74% to 43.33% across attention heads.

Our empirical measurements demonstrate that the practical impact on latency is minimal. Our experiments show that the additional latency overhead is quite modest - approximately 2ms per attention call compared to balanced budget approaches.

R W3: Questions the justification for defaulting to vertical-slash patterns when attention distribution estimation shows a discrepancy.

Our choice of using vertical-slash patterns as a fallback is motivated by empirical observations of attention patterns in large language models, as illustrated in Figure 2 of our paper. Through extensive analysis, we observed that attention heads predominantly exhibit either Query-Aware or Vertical-Slash patterns.

While we acknowledge that the distribution discrepancy alone doesn't theoretically necessitate the use of vertical-slash patterns, as stated in line 199 of our paper, in such cases, "we fall back to a more conservative approach", our empirical results support this as a robust fallback strategy that maintains model performance while ensuring computational efficiency.

Thank you for your interest and acknowledgment to our work.

After carefully reviewing the SampleAttention [1] paper, we would like to clarify the distinctions between our works. We have positioned these observations in Section 2 of our paper (Lines 119-141), specifically in paragraphs "Attention Sparse Patterns Exhibit Variability" and "Varying Samples Require Adaptive Sparsity Ratio", as motivational context rather than claiming them as novel discoveries. We acknowledge that other sparse attention works, including MInference and SampleAttention, were similarly motivated by these phenomena, and we will include comprehensive references to provide a complete background.

Regarding methodology, it is important to note that our "Query-Aware Sparse Pattern Determination" differs fundamentally from "Query-Guided Attention Sampling."

• Our term "Query-Aware" refers to a specific pattern defined in lines 174-175, which is designed to capture variable sparse patterns contingent upon query positions.

• "Query-Aware Pattern Determination" in our work involves applying the last query block to analyze differences between block-wise attention distributions (as detailed in lines 184-185), which then informs our selection of the appropriate sparse attention pattern.

This approach stands in clear contrast to SampleAttention's methodology, where "Query-Guided" refers to their KV selection process. Specifically, SampleAttention uses query guidance primarily for selecting key-value indices, computing exact scores for selected queries in Stage-1 to determine KV indices of interest in Stage-2. This fundamental difference in purpose and implementation distinguishes our methodological contributions.

While we share similar goals for sparse attention, our implementation differs significantly.

• We first determine patterns before KV selection, as we employ different query-key pair selection strategies across patterns, heads, and inputs. This contrasts with SampleAttention's fixed "local window pattern" and "column stripe pattern" approach for all heads.

• Furthermore, while SampleAttention samples a subset of queries to minimize CRA, we apply different strategies for different heads.

  • For "query-aware head" we utilize all queries with a pooling operation to estimate the attention map and select different Q-K pairs for different queries.

  • For the remaining heads, we apply the last block query to estimate vertical and slash lines.

For a clear understanding of the differences between the two approaches, we encourage readers to compare Figure 9 in FlexPrefill with Figure 1 in SampleAttention, which visually illustrates the distinct sparse attention mask of each method.

We uploaded our implementation code to the supplementary materials and highlighted that, a key distinction of FlexPrefill lies in its flexibility. FlexPrefill can perform online pattern determination and query-key index selection without offline search. This allows FlexPrefill to be integrated into existing LLM frameworks through a single function call replacing standard attention mechanisms, such as in vLLM and transformers.

Regarding benchmarking, we acknowledge that we haven't made direct comparisons with SampleAttention due to its implementation not being publicly available. We would welcome guidance on implementation details to ensure accurate reproduction and incorporate it into our comparative evaluation.

[1] SampleAttention

R Q2: Integration Feasibility

FlexPrefill is designed for seamless integration into existing LLM frameworks through a single function call that can replace standard attention mechanisms. This design makes it readily compatible with popular frameworks like transformers and vLLM. Unlike MInference, which requires model-specific pattern search, our method provides a unified interface that works consistently across different models and scenarios.

R Q3: Full Attention Latency

In Figure 4, the average latency of the full attention model is reported as follows: 4.5 ms (averaged across all input lengths), 20.58 ms (for 128k input length), and 1.19 ms (for 32k input length).

Our proposed method, as illustrated in Figure 4, demonstrates varying latency depending on the parameter $\gamma$, while maintaining comparable performance. Specifically, our method achieves improved latency: 1.7 ms (averaged across all input lengths), 6 ms (for 128k input length), and 0.9 ms (for 32k input length).

We will include a detailed side-by-side comparison with Flexprefill in our next revision to provide a clearer comparison and further highlight the efficiency of our approach.

R Q6: Threshold Selection

Our threshold parameters offer flexible control over the speed-accuracy trade-off. The cumulative attention threshold $\gamma$ allows direct tuning of the computation-performance balance - higher $\gamma$ values favor accuracy while lower values prioritize speed. The sparse pattern threshold $\tau$ was empirically determined. This flexibility enables users to adjust parameters based on their computational resources and performance requirements.

R W1: Detailed computational complexity analysis to justify overhead of multi-step sparse attention mechanism & Q1: Computational Complexity Analysis.

We here provide a clear justification for the overhead associated with FlexPrefill's multi-step sparse attention mechanism. As outlined in Figure 6 of our paper, the analysis spans context lengths ranging from 8K to 512K tokens and demonstrates that the incurred overhead is minimal compared to the computational savings achieved.

1. Computational Complexity Breakdown

The computational complexity of FlexPrefill is composed of the following components:

• Representative Attention Score Computation: $O(bnd)$, where $b$ is the block size, $d$ is the hidden dimension, and $n$ is the sequence length.
• Pattern Search: $O(bn)$, which involves computing JS divergence and dynamically selecting the most suitable attention pattern.
• Sparse Index Construction: $O(n \log n)$, required for efficiently organizing sparse attention indices.
• Sparse Attention Computation: $O(\alpha n^2 d)$, where $\alpha$ represents the sparsity ratio.

In contrast, the standard attention mechanism has a computational cost of $O(n^2 d)$. By comparison, FlexPrefill introduces a small overhead (approximately $O(\alpha n^2 d)+O(n\log n)$), which is substantially outweighed by the computational savings enabled by the sparsity mechanism.

2. Empirical Justification for Long Sequences

Our empirical evidence further supports the efficiency of FlexPrefill, particularly for long sequences ($n \geq 128k$). As demonstrated in Figure 6, the overhead becomes increasingly negligible as the dominant computation shifts to sparse attention.

R W2: Need for more detailed quantification of computational savings compared to baselines & Q4: Computational Savings.

We conducted comprehensive experiments to quantify computational costs across four LLMs, measuring both the average time of a single attention computation and accuracy:

For Llama:

For GLM:

For Yi:

For Qwen:

Our results reveal computational advantages while maintaining or improving performance.

At 64k tokens, FlexPrefill demonstrates better efficiency compared to MInference, which shows higher latency than full attention. For Llama, FlexPrefill achieves 85.40ms latency compared to MInference's 152.56ms while maintaining higher accuracy (89.18% vs 85.42%).

In long-sequence scenarios (128K tokens), FlexPrefill continues to show performance-efficiency trade-offs. For Qwen, we achieve 173ms latency while maintaining 70.75% accuracy, compared to MInference's 306.47ms.

Across all sequence lengths, FlexPrefill consistently demonstrates practical speed advantages while maintaining or improving model performance.

R W3: Limited comparison scope with recent works in sparse attention & Q5: Comparisons with Recent Works.

We appreciate this suggestion and have conducted additional experiments incorporating recent baselines. Below are comprehensive comparison results on the RULER dataset (We note that LM-Infinite does not support GLM and Yi):

Our choice of MInference as a primary baseline was deliberate. As demonstrated in their work, MInference already outperforms many recent methods in similar evaluation settings. By showing efficiency improvements over MInference, our method implicitly demonstrates advantages over these other approaches.

R W4: Questions about performance benefits and computational trade-offs compared to Full Attention and baselines.

Our experimental results demonstrate that FlexPrefill achieves comparable or better performance than both Full Attention and MInference while providing computational advantages.

As shown in Table 1 in our original paper, for the Llama model, FlexPrefill maintains an average accuracy of 89.18% compared to Full Attention's 88.70% and MInference's 85.42%. This performance is achieved while delivering speedups. Similar patterns are observed across other models.

We acknowledge the importance of providing a detailed comparison of computational costs. Please refer to R W2 for further details.

Thank you for addressing my concerns! I will update my rating accordingly.

R W1: Limited Use of Query-Aware Patterns

We appreciate the insightful observation regarding the role of Query-Aware patterns and would like to elucidate FlexPrefill's advantages by focusing on two key aspects: Pattern Determination Strategy and Cumulative Attention Mechanism.

1. Pattern Determination Strategy

While Query-Aware patterns may appear less frequently, their integration with Vertical-Slash patterns is crucial for enhancing model performance. Unlike MInference, which employs a fixed pattern approach, FlexPrefill dynamically selects the most appropriate pattern for each attention head, leveraging a flexible and adaptive strategy.

As illustrated in Figure 9, our empirical analysis reveals that Query-Aware attention heads can capture more complex attention patterns, complementing the Vertical-Slash patterns without sacrificing overall model capabilities. Furthermore, the importance of Query-Aware patterns extends beyond the specific examples analyzed in Figure 8, which focuses on four particular inputs. In R Q2, we provide a more comprehensive statistical analysis, demonstrating the distribution of Query-Aware patterns across different layers and their significance as a crucial component of the model's design.

2. Cumulative Attention Mechanism

FlexPrefill introduces a cumulative-attention-based index selection mechanism, which goes beyond merely choosing attention patterns. By dynamically adjusting the cumulative attention threshold (denoted as $\gamma$), FlexPrefill achieves better performance while maintaining faster inference speeds compared to MInference.

As shown in Figure 4, the adaptability of the cumulative attention mechanism allows FlexPrefill to transition from initial underperformance to surpassing MInference with careful tuning of $\gamma$. This is reflected in the purple curve, where FlexPrefill demonstrates its ability to optimize the trade-off between performance and computational efficiency.

In conclusion, this dual innovation—dynamic pattern determination and cumulative attention—positions FlexPrefill as a more adaptable and efficient alternative to MInference. By allowing for optimal trade-offs between computational cost and model effectiveness across various scenarios, FlexPrefill achieves a performance-efficiency balance that meets diverse operational requirements.

R Q1: Selection method for representative query subset Q^

As detailed in our implementation section (lines 322-323), we utilize the last block_size query vectors as the representative query set $\hat{Q}$. This choice is motivated by two key factors:

1. Performance Efficacy: Our ablation studies in Table 6 demonstrate that using the last block consistently achieves better performance (89.18% avg accuracy) compared to alternative selections (e.g., middle block: 38.18% avg accuracy)
2. Complete Key Coverage: The last queries have visibility to all previous keys, enabling more comprehensive pattern assessment

R Q2: Methodology behind Figure 8 generation

Figure 8 in our paper presents an illustrative example to demonstrate how attention patterns vary across 4 different input scenarios. To provide a more comprehensive view of pattern distribution, we conducted extensive analysis across the entire RULER dataset using the Qwen2 model to obtain statistical results for Query-Aware head distributions over different layers.

Our analysis reveals that the pattern distribution shows significant layer-wise variation. Early layers (1-7) show a higher proportion of Query-Aware patterns (averaging around 46%), while deeper layers (21-28) tend to rely more on Vertical-Slash patterns (Query-Aware patterns around 16%).

R Q3: Global versus query-wise index selection in Algorithm 4

We conducted comprehensive experiments comparing global index selection (with flatten operation) versus query-wise selection across different models. The results demonstrate comparable performance:

For average performance across all context lengths:

For long-context scenarios (128k tokens):

Given these comparable results, we opted for the global approach with flatten operation primarily for implementation efficiency. This choice simplifies the implementation of the attention mechanism while maintaining performance parity.