[Bug]: Flash attention with sliding window

[heheda12345]

Your current environment

N/A

🐛 Describe the bug

pytest -vs test_correctness_sliding_window.py::test_sliding_window_retrival\[1-5-bigcode/starcoder2-3b\] starts to fail on H100 since #16998. This test can pass if I change it to num_heads_kv=self.num_heads_q like here https://github.com/heheda12345/vllm/tree/fa3_hack .

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[heheda12345]

CC @LucasWilkinson