XAttention: Block Sparse Attention with Antidiagonal Scoring

🔼 The figure illustrates how XAttention’s novel antidiagonal scoring method effectively captures important information within attention blocks. By summing values along antidiagonals (with a specified stride), XAttention identifies blocks containing both vertical and diagonal patterns – crucial indicators of significant relationships between tokens. This method is superior to simple pooling because it directly addresses and avoids missing these key attention patterns, leading to more precise block selection and higher efficiency in sparse attention computation.

read the caption

Figure 2: XAttention’s antidiagonal pattern intersects both vertical and slash patterns within a block, enabling efficient detection of these patterns and guiding effective sparse attention computation.

🔼 This figure presents a qualitative comparison of video generation results obtained from four different methods using the first prompt from the VBench dataset. The four methods are: (1) Full Attention (used as a baseline for comparison), (2) XAttention without any warmup period (with τ = 0.95), (3) XAttention with a 5-step warmup period (τ = 0.9), and (4) XAttention with a 5-step warmup period (τ = 0.95). Each row displays selected frames from a video generated by one of the four methods, allowing for a visual comparison of the quality and fidelity of the generated videos. The key takeaway is that XAttention, especially when using a warmup period, generates videos with high visual fidelity, closely matching the quality of those produced using the full attention baseline.

read the caption

Figure 3: Qualitative comparison of video generation results on the VBench benchmark using the first prompt in the VBench dataset. Rows show frames from videos generated using: (1) Full Attention (baseline), (2) XAttention with no warmup and (τ𝜏\tauitalic_τ = 0.95), (3) XAttention with 5 warmup steps and (τ𝜏\tauitalic_τ = 0.9), and (4) XAttention with 5 warmup steps and (τ𝜏\tauitalic_τ = 0.95). XAttention with warmup achieves high visual fidelity to the full attention baseline.

🔼 This figure compares the speedup achieved by various attention mechanisms against FlashAttention (as implemented in FlashInfer) across different context lengths. The x-axis represents the sequence length (in tokens), and the y-axis displays the speedup factor relative to FlashAttention. The results demonstrate that XAttention consistently outperforms other sparse attention methods (MInference, SeerAttention, FlexPrefill), achieving the highest speedup, reaching up to 13.5x at a context length of 256K tokens. This highlights XAttention’s efficiency in handling very long sequences.

read the caption

Figure 4: Speedup comparison of attention methods across context lengths, relative to FlashInfer’s implementation of FlashAttention. XAttention consistently outperforms other sparse attention methods, achieving up to 13.5x speedup at 256K tokens.

🔼 Figure 5 is a bar chart comparing the time spent on pattern search and attention computation during the prefill stage of different sparse attention methods. XAttention significantly reduces the time required for pattern search while maintaining a similar attention density compared to other methods, resulting in substantial speedup in overall attention computation.

read the caption

Figure 5: Breakdown of prefill attention time. Xattention significantly reduces pattern selection time while maintaining density, achieving substantial acceleration compared to existing methods.

🔼 This table presents a comparison of the performance of several attention mechanisms on various real-world tasks within the LongBench benchmark. The benchmark utilizes the Llama-3.1-8B-Instruct language model. The attention methods compared include XAttention (the proposed method), FlashAttention (a baseline for dense attention), MInference, FlexPrefill, and SeerAttention (other existing block-sparse methods). The results showcase XAttention’s superior performance in terms of accuracy across different LongBench tasks, particularly when employing stride 8 and a precisely predicted minimum threshold. The table details the performance (scores) achieved by each method on individual LongBench tasks (single-document QA, multi-document QA, 2WikiMQA, MuSiQue, GovReport, QMSum, VCSum, MultiNews, TREC, few-shot learning, TriviaQA, SAMSum, LSHT, LCC, RB-P, and Code), thus offering a comprehensive comparison across diverse and complex long-context natural language understanding scenarios.

read the caption

Table 2: Comparison of different attention methods on real-world LongBench tasks using the Llama-3.1-8B-Instruct model. XAttention, configured with stride 8 and Precisely Predicted Minimum Threshold, achieves the best average scores against all baselines.

🔼 Table 3 presents a comparison of different attention mechanisms on the QwenVL-2-7B model for the video understanding task within the Video-MME dataset. Specifically, it compares the performance of Full Attention (the baseline), XAttention (using a stride of 16 and a threshold of 0.9), MInference, and FlexPrefill across three video lengths: short, medium, and long. The results show XAttention’s performance relative to the baseline and other sparse attention methods. The table highlights XAttention’s superior performance on long videos and its overall best average performance among the sparse attention methods.

read the caption

Table 3: Comparison of different methods on QwenVL-2-7B in the Video-MME video understanding task. XAttention is configured with Stride S=16𝑆16S=16italic_S = 16 and Threshold τ=0.9𝜏0.9\tau=0.9italic_τ = 0.9. XAttention outperforms Full Attention on long video tasks and achieves the best average performance among all sparse attention methods.

🔼 This table presents a quantitative evaluation of XAttention’s performance on the HunyuanVideo model for video generation. The experiment uses the VBench benchmark and incorporates a 5-step full-attention warmup phase. The results show the Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), Learned Perceptual Image Patch Similarity (LPIPS), and density for two different threshold settings (τ). Higher τ values lead to better video quality (higher PSNR and SSIM, lower LPIPS) but at the cost of slightly lower sparsity (higher density). Both τ settings show comparable results to the full-attention baseline, indicating that XAttention maintains good performance with significant computational savings.

read the caption

Table 4: Quantitative results of applying XAttention to the HunyuanVideo model on the VBench benchmark, using a 5-step full-attention warmup. Higher (τ𝜏\tauitalic_τ) yields better fidelity (higher PSNR, higher SSIM, lower LPIPS) at the cost of slightly reduced sparsity (higher density). Both (τ𝜏\tauitalic_τ) settings demonstrate high similarity to the full attention baseline.

🔼 This table presents the density of the attention mechanism across various context lengths (4k, 8k, 16k, 32k, 64k, 128k tokens) using a stride of 8. Density refers to the proportion of attention computation performed, with lower density implying higher sparsity. The results show that using a stride of 8 leads to lower sparsity (more computation) compared to strides of 4 or 16. Furthermore, as the context length increases, the density of the attention decreases, demonstrating that the method achieves greater sparsity with longer sequences.

read the caption

Table 5: Density on Different Context Lengths. Stride S=8𝑆8S=8italic_S = 8 achieves lower sparsity, and as context length increases, sparsity generally increases (lower density).

🔼 Table 6 presents an ablation study comparing three different patterns used for predicting attention block importance in the XAttention model: random, diagonal, and antidiagonal. The study measures the average accuracy and density (sparsity) achieved by each pattern while maintaining the same computational cost. The results demonstrate that the antidiagonal pattern outperforms random and diagonal patterns, achieving both the highest accuracy and the lowest density, which translates to superior efficiency without compromising performance.

read the caption

Table 6: Comparison of different patterns. For the same computation, the antidiagonal achieves the lowest density and the highest score.

🔼 This table presents an ablation study on the impact of different stride sizes (S) used in the XAttention algorithm on the accuracy of identifying important attention blocks. The results show that using excessively large strides negatively affects the ability to differentiate slash patterns of varying lengths, ultimately leading to decreased overall accuracy. In essence, it explores the trade-off between computational efficiency (larger strides mean less computation) and the accuracy of identifying important attention blocks.

read the caption

Table 7: Comparison of different Strides. Excessively long strides fail to distinguish slash patterns with different lengths, leading to decreased accuracy.

🔼 This table compares three different block selection algorithms used in XAttention for sparse attention computation: Top-K, Top-Ratio, and the proposed Threshold-based selection (Dynamic Sparsity). It shows the average density and performance (accuracy) of each method across various stride sizes (4, 8, and 16), which determine the sparsity level. The goal is to find the optimal balance between computational efficiency and accuracy in identifying important attention blocks.

read the caption

Table 8: Comparison of different selection algorithms.

🔼 Table 9 presents a comparison of results obtained using a fixed threshold (T=0.9) versus a dynamically predicted minimum threshold for attention mechanism in the XAttention model. It demonstrates how the dynamic threshold method enhances both the accuracy of the model and its sparsity (resulting in lower density and faster inference) across different stride sizes (S=4, 8, 16). The table highlights the improved efficiency and accuracy achieved by using the dynamic programming approach for threshold prediction.

read the caption

Table 9: Minimum Threshold Prediction yields improvements in both accuracy and sparsity, translating to faster inference.