X-Attention Appears to Be Underperforming in Needle-in-a-Haystack Benchmark

[GITHUBear]

Hello, X-Attention teams!
After reading your paper, I became very interested in your work and proceeded to integrate X-Attention into vLLM. You can find my integration branch here: https://github.com/GITHUBear/vllm/tree/experiment , along with the relevant commit GITHUBear/vllm@a3abc73

In this implementation:

1. I replaced the prefill phase of vLLM's Flash-Attention backend with X-Attention's prefill interface
2. Addressed several compatibility issues between X-Attention and vLLM while retaining its core implementation:

• Limited support for batch_size > 1
• Incompatibility with GQA scenarios where query head count exceeds key head count
• Lack of support for KVCache block tables

I conducted tests using the qwen-2.5-7b-1M model as the baseline. Below are the testing environment details and results. Given the significant discrepancy between the baseline results and X-Attention's performance, I would greatly appreciate your team's expertise in reviewing this matter and providing actionable recommendations.

Env Info

• Python: Python 3.11.11
• vLLM: 0.9.1
• cuda: cuda_12.8
• GPU: NVIDIA L20
• torch: 2.7.0+cu126

Benchmark

You can run the Needle-in-a-Haystack benchmark with run_needle.sh

1. Baseline: FullAttention with FA2

python ./needle_test.py --model_name Qwen/Qwen2.5-7B-Instruct-1M --max_length 500000 --min_length 10000 --trust_remote_code --enable_chunked_prefill --tensor_parallel_size 4

2. X-Attention mode with default hyperparameters：

• stride: int = 16
• threshold: float = 0.9
• block_size: int = 128
• chunk_size: int = 2048

python ./needle_test.py --model_name Qwen/Qwen2.5-7B-Instruct-1M --max_length 500000 --min_length 10000 --trust_remote_code --enable_chunked_prefill --tensor_parallel_size 4 --enforce_eager --sparse_prefill_type 1 --run_name x_attn

Benchmark Result

• Baseline:

• X-Attention:

[Guangxuan-Xiao]

Thanks for the detailed integration and benchmark results — really impressive work!

To help narrow down the performance discrepancy, could you try the following adjustments and let us know if they help?

• Disable the chunked prefill (--enable_chunked_prefill → off)
• Start with a smaller context length, e.g., reduce max_length to something like 128k or 256k for initial testing
• Increase the sparsity threshold to --threshold 0.95
• Try a smaller stride size, such as --stride 8

These changes may better align with the assumptions under which XAttention performs optimally. Curious to see if that brings the results closer to baseline.

[GITHUBear]

@Guangxuan-Xiao Thank you very much for your suggestions. By following the recommended parameter settings, we obtained the expected results as shown in the figure below.

However, we observed that the vLLM chunked prefill feature significantly impacts the accuracy of X-Attention. We speculate that this is because when X-Attention is enabled, repeated chunked prefill operations cause cumulative deviations in KVCache except for layer 0. But interestingly, in our initial tests we set the chunk size to 130k, which only results in approximately 3-4 chunked prefill operations. Theoretically, this shouldn't lead to significant accuracy deviations?

Given that chunked prefill is an effective method for reducing GPU memory usage, we're committed to improving X-Attention's performance under chunked prefill conditions.

However, before proceeding, we are very curious about what factors are causing X-Attention to underperform in chunked prefill scenarios?