[CI/Build] Refactor Attention backend for test_prefix_prefill from xformers to SDPA

[zhewenl]

Purpose

1. This PR will help [Core] Remove xformers dependency #28287 (comment) by changing the ground truth for attention backend from xformers to pytorch SDPA
2. Fixes some incompatibilities on AMD:
   a. The ROCm paged attention kernel expects 32-bit int tensors, but the test passes 64-bit torch.long tensors.
   b. ROCm paged attention kernel only supports auto, fp8, and fp8_e4m3 KV cache dtypes.

Test Plan

H100(https://gist.github.com/zhewenl/5ada1e5f360c4d230bc8b6eff47effcc):

pytest -s -v 'tests/kernels/attention/test_prefix_prefill.py'
...
================================================================== warnings summary ==================================================================
<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================== 192 passed, 192 skipped, 2 warnings in 208.65s (0:03:28) ==============================================

MI300 (failing with some numeric issues, will track in a separate issue/PR): https://gist.github.com/zhewenl/3224057e57aad300341c8a0d66bd9878

pytest -s -v 'tests/kernels/attention/test_prefix_prefill.py'
...
================================================================= short test summary info ==================================================================
FAILED tests/kernels/attention/test_prefix_prefill.py::test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:0-auto-dtype0-128-1-64] - AssertionError: Tensor-likes are not close!
FAILED tests/kernels/attention/test_prefix_prefill.py::test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:0-fp8-dtype0-128-1-64] - AssertionError: Tensor-likes are not close!
FAILED tests/kernels/attention/test_prefix_prefill.py::test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:1-auto-dtype0-128-1-64] - AssertionError: Tensor-likes are not close!
FAILED tests/kernels/attention/test_prefix_prefill.py::test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:1-fp8-dtype0-128-1-64] - AssertionError: Tensor-likes are not close!

Results (299.93s (0:04:59)):
     156 passed
       4 failed
         - tests/kernels/attention/test_prefix_prefill.py:101 test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:0-auto-dtype0-128-1-64]
         - tests/kernels/attention/test_prefix_prefill.py:101 test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:0-fp8-dtype0-128-1-64]
         - tests/kernels/attention/test_prefix_prefill.py:101 test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:1-auto-dtype0-128-1-64]
         - tests/kernels/attention/test_prefix_prefill.py:101 test_contexted_kv_attention[chunked_prefill_paged_decode-0-cuda:1-fp8-dtype0-128-1-64]
     224 skipped

[DarkLight1337]

cc @ywang96

[ywang96]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the attention backend for test_prefix_prefill from xformers to PyTorch's Scaled Dot Product Attention (SDPA), which is a great move for standardization and potentially performance. The changes also include important compatibility fixes for ROCm, such as using int32 for certain tensors and skipping unsupported fp8_e5m2 configurations. The implementation of the SDPA reference is well-structured. I've included one suggestion to improve the performance of the newly added test utility function for creating attention masks, which should help reduce the overall test execution time.

[ywang96]

b6f852f saves CI about 10 seconds (better than nothing I guess)

[ywang96]

Thanks for working on this!