[CI Perf] Prune tests in tests/kernels/attention/

[mgoin]

Purpose

Greatly reduce the number of test cases generated by:

• Only test for bfloat16
• Only test one shape for MHA or GQA
• Reduce to one or two head sizes
• Reduce to one block size
• Reduce soft capping to two cases
• Reduce sliding windows to two cases

This is still not aggressive enough IMO because we are still testing all permutations. Best practice would be to just come up with fixed test points and applying those directly, rather than permutations.

Counts for each test:

pytest --collect-only -qq tests/kernels/attention

# main
tests/kernels/attention/test_aiter_flash_attn.py: 192
tests/kernels/attention/test_attention.py: 1850
tests/kernels/attention/test_attention_selector.py: 31
tests/kernels/attention/test_cache.py: 2201
tests/kernels/attention/test_cascade_flash_attn.py: 198
tests/kernels/attention/test_encoder_decoder_attn.py: 32
tests/kernels/attention/test_flash_attn.py: 9216
tests/kernels/attention/test_flashinfer.py: 768
tests/kernels/attention/test_flashinfer_trtllm_attention.py: 192
tests/kernels/attention/test_flashmla.py: 64
tests/kernels/attention/test_lightning_attn.py: 10
tests/kernels/attention/test_merge_attn_states.py: 648
tests/kernels/attention/test_mha_attn.py: 27
tests/kernels/attention/test_mla_decode_cpu.py: 6
tests/kernels/attention/test_prefix_prefill.py: 1728
tests/kernels/attention/test_rocm_attention_selector.py: 1
tests/kernels/attention/test_triton_decode_attention.py: 96
tests/kernels/attention/test_triton_unified_attention.py: 1152

# pr
tests/kernels/attention/test_aiter_flash_attn.py: 32
tests/kernels/attention/test_attention.py: 278
tests/kernels/attention/test_attention_selector.py: 31
tests/kernels/attention/test_cache.py: 1108
tests/kernels/attention/test_cascade_flash_attn.py: 198
tests/kernels/attention/test_encoder_decoder_attn.py: 32
tests/kernels/attention/test_flash_attn.py: 1024
tests/kernels/attention/test_flashinfer.py: 144
tests/kernels/attention/test_flashinfer_trtllm_attention.py: 24
tests/kernels/attention/test_flashmla.py: 64
tests/kernels/attention/test_lightning_attn.py: 10
tests/kernels/attention/test_merge_attn_states.py: 648
tests/kernels/attention/test_mha_attn.py: 27
tests/kernels/attention/test_mla_decode_cpu.py: 6
tests/kernels/attention/test_prefix_prefill.py: 384
tests/kernels/attention/test_rocm_attention_selector.py: 1
tests/kernels/attention/test_triton_decode_attention.py: 96
tests/kernels/attention/test_triton_unified_attention.py: 128

Test Plan

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request aims to optimize CI by pruning the test matrix in various attention kernel tests. The changes primarily involve reducing the number of parameterized values for tests, such as number of heads, head sizes, block sizes, and data types.

While most of the pruning seems reasonable for CI optimization, I've identified a recurring high-risk change across multiple test files: the removal of torch.float16 from the tested data types. Given that float16 is a very common data type in production, removing it from the test suite could lead to undetected regressions. I have added comments to suggest re-introducing torch.float16 to the test matrix in the affected files to maintain test coverage for this critical data type.

[yewentao256]

LGTM, thanks for the work!