XAttention: Block Sparse Attention with Antidiagonal Scoring

1. Antidiagonal Selection: Performance and Insights

Antidiagonal selection offers significant advantages over other patterns:

• Retains all token information while simultaneously detecting both vertical and slash patterns critical in LLM prefill
• Diagonal patterns miss slash patterns with probability $\frac{\text{stride} - 1}{\text{stride}}$
• Other patterns (horizontal lines, pooling methods) lose token-level information
• Empirically validated superiority: outperforms diagonal patterns (Table 6) and pooling-based methods used by FlexPrefill

2. Precomputation Process

Precomputation occurs before computing the full attention for $Q\times K^T$. Based on block selection results, we determine which blocks require full attention computation. The process is detailed in Algorithm 1.

3. Minimum Threshold Prediction

Minimum Threshold Prediction is an optional algorithm which finds the optimal threshold for each head, further optimizing efficiency and accuracy. This process:

• Is conducted offline (not during prefill)
• Does not increase runtime computational complexity
• Profiles configurations in advance, determining minimum thresholds per head

4. Baseline Selection Justification

Our selected baselines represent SOTA methods for sparse prefill attention. Previous approaches:

• Have high precomputation complexity and overhead (Figure 5)
• Limited applicability (FlexPrefill/MInference don't support video generation or chunk prefill)
• Require additional fine-tuning (SeerAttention)

5. Stride Selection Strategy

Stride is indeed critical for balancing efficiency and accuracy:

Different stride values (tested on RULER with Llama3.1 8B):

These results show stride = 8 maintains accuracy comparable to stride = 4 with less overhead, and performance remains effective for stride ≤ 16.

According to our complexity analysis:

• Precomputation: $\frac{2n^2d}{\text{stride}} + \frac{3n^2}{\text{stride}^2}$
• Full attention: $4n^2d + 4n^2 + 6nd^2$

With stride = 4, precomputation cost is < 1/8 of full compute With stride = 8, precomputation cost is < 1/16 of full compute

(As shown in Figure 5, while the actual latency closely aligns with theoretical predictions, it may be slightly higher due to factors such as vector reordering, memory overhead, and the implementation details of the Triton kernel.)

6. Predefined Sparsity Implementation

Fixed sparsity can be achieved by setting a block limit k and selecting the top-k highest-scoring blocks. This approach is equivalent to the Top-k strategy (Table 8).

However, we believe fixed sparsity across all inputs is suboptimal since information density varies between requests. XAttention's dynamic sparsity determination enables better generalization and accuracy across diverse scenarios.

1. Model Generalizability

We tested XAttention across diverse architectures with consistent results:

LLMs Accuracy (RULER)

LLMs Speed-up

Video Generation (Wan 2.1 14B)

---

2. Hyperparameter Sensitivity

Threshold = 0.9 demonstrates strong generalizability across models:

Same threshold (0.9) on different models:

Trade-off curve with LLama 3.1 8B (stride=4):

---

3. Video Generation Quality

These quantitative metrics are well-acknowledged methods used in efficient visual content generation works to compare the difference between original generated content and the efficient generated content, such as used in Distrifusion (Li et al., https://arxiv.org/pdf/2402.19481) and Sparse VideoGen (Xi et al., https://arxiv.org/pdf/2502.01776). They are extremely strict compared to content level scores since they require pixel-level exactness, and the scores XAttention achieved is a level which very similar to the original generation. Beyond quantitative metrics, we've included qualitative samples in the Figure 3 comparing between XAttention-generated videos and full attention videos, confirming they are virtually indistinguishable.

1. Algorithm Clarification

XAttention precomputes attention within the antidiagonal pattern and uses these scores to guide block-sparse attention selection. Our method:

• Requires no fine-tuning
• Achieves lowest overhead among prefill acceleration approaches (Figure 5)
• Delivers up to 13.5× speedup over FlashAttention (Figure 4)

---

2. Complexity Analysis

Precomputation complexity: $\frac{2n^2d}{\text{stride}} \quad (\text{approximate attention}) \ + \ \frac{3n^2}{\text{stride}^2} \quad (\text{approximate softmax})$

Full attention complexity: $4n^2d + 3n^2 + 4nd^2$

With stride = 4, precomputation cost is < 1/8 of full compute With stride = 16, precomputation cost is < 1/32 of full compute.

Figure 5 shows a breakdown of precompute time that closely aligns with this theoretical result. (While the actual latency closely aligns with theoretical predictions, it may be slightly higher due to factors such as vector reordering, memory overhead, and the implementation details of the Triton kernel.)

---

3. Threshold DP Design Choice

As described in Section 2.3, this is an optional component for optimizing sparsity by assigning different thresholds per head. XAttention also works effectively with global thresholds on LLMs, VLMs, and video generation.

---

4. Antidiagonal Pattern Effectiveness

The antidiagonal pattern:

• Retains all token information
• Simultaneously detects vertical and slash patterns with $1/\text{stride}$ probability
• Prevents information loss compared to horizontal lines or pooling methods
• Empirically outperforms both diagonal patterns (Table 6) and pooling methods (e.g., FlexPrefill)

---

5. Two-Level Loop Structure

In Algorithm 1:

• Outer loop: blocks (blocknum)
• Inner loop: stride slices (stride)

Despite appearing as nested loops, the actual complexity is O(n), not O(n²), where n represents sequence length.

---

6. Softmax Approximation

Line 120 explains why we normalize attention scores - to create a probability distribution for threshold selection. The Softmax complexity here is $\left(\frac{\text{stride}}{\text{blocksize}}\right)^2$ of full attention (e.g., $\frac{1}{1024}$ with stride = 4).

Top-k could replace Softmax, but Table 8 demonstrates this performs worse than our threshold method.

---

7. Implementation Details

As noted in line 20 of the abstract, our method is:

• Plug-and-play
• Requires no finetuning
• Specifically designed for prefill stage

We used both Block Sparse Attention (https://github.com/mit-han-lab/Block-Sparse-Attention) CUDA kernel and Triton kernel for implementation.

---

8. Figure Presentation

Thank you for these suggestions. We will improve figure clarity in the revision.

1. Antidiagonal Selection: Performance and Insights

Antidiagonal selection offers significant advantages over other patterns:

• Retains all token information while simultaneously detecting both vertical and slash patterns critical in LLM prefill
• Diagonal patterns miss slash patterns with probability $\frac{\text{stride} - 1}{\text{stride}}$
• Other patterns (horizontal lines, pooling methods) lose token-level information
• Empirically validated superiority: outperforms diagonal patterns (Table 6) and pooling-based methods used by FlexPrefill

---

2. Novel Contribution

We respectfully disagree with the assessment of limited contributions. While block-sparse attention exists, effectively implementing it without accuracy/efficiency degradation remains challenging:

• FlexPrefill and MInference: Suffer from high precomputation costs and rely on last block for pattern detection, preventing chunk prefill and adaptation to non-text tasks like video
• SeerAttention: Requires additional parameter training, limiting generalizability. Also it has poor empirical performance.
• Our work: First to identify antidiagonal pattern as an effective block importance proxy, creating a simple yet powerful method that significantly improves the efficiency-accuracy trade-off

---

3. FlashAttention Comparison

As shown in the Figure 4 of our submitted paper, XAttention achieves up to 13.5× speedup over FlashInfer FlashAttention (one of the fastest implementations):

4. References and Formatting

Thank you for noting these issues. We will address them in the revision.

---

5. Appendices

The main paper effectively covers our contributions. We will include an appendix showing additional visual samples, theoretical analysis, and experimental results in the next revision.