[ROCm][Perf] Enable shuffle kv cache layout and assembly paged attention kernel for AiterFlashAttentionBackend

[ganyi1996ppo]

Purpose

This PR add assembly paged attention kernel in aiter to AiterFlashAttentionBackend. We verify this implementation on Qwen3-30B-A3B-FP8 Mi308 and observed at least about 20% thoughput gain and obvious latency reduction on tpot.

This PR add this flag USING_SHUFFLE_LAYOUT to control whether to enable assembly paged attention, and set it False by default to prevent any unexpected circumstance. This flag will be removed after assembly paged attention fully verified by the users.

Test Plan

gsm8k for accuracy
vllm bench for performance

Test Result

# normal aiter paged attention ============ Serving Benchmark Result ============ Successful requests: 200 Failed requests: 0 Maximum request concurrency: 200 Benchmark duration (s): 95.28 Total input tokens: 768000 Total generated tokens: 204800 Request throughput (req/s): 2.10 Output token throughput (tok/s): 2149.35 Peak output token throughput (tok/s): 4000.00 Peak concurrent requests: 200.00 Total Token throughput (tok/s): 10209.43 ---------------Time to First Token---------------- Mean TTFT (ms): 21644.01 Median TTFT (ms): 21321.70 P99 TTFT (ms): 41420.74 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 71.65 Median TPOT (ms): 72.00 P99 TPOT (ms): 90.40 ---------------Inter-token Latency---------------- Mean ITL (ms): 71.65 Median ITL (ms): 53.10 P99 ITL (ms): 1744.72 ================================================== # assembly aiter paged attention ============ Serving Benchmark Result ============ Successful requests: 200 Failed requests: 0 Maximum request concurrency: 200 Benchmark duration (s): 79.77 Total input tokens: 768000 Total generated tokens: 204800 Request throughput (req/s): 2.51 Output token throughput (tok/s): 2567.23 Peak output token throughput (tok/s): 5400.00 Peak concurrent requests: 200.00 Total Token throughput (tok/s): 12194.33 ---------------Time to First Token---------------- Mean TTFT (ms): 21566.08 Median TTFT (ms): 21239.65 P99 TTFT (ms): 41279.14 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 56.62 Median TPOT (ms): 56.95 P99 TPOT (ms): 75.42 ---------------Inter-token Latency---------------- Mean ITL (ms): 56.62 Median ITL (ms): 37.89 P99 ITL (ms): 1742.49 ================================================== 

accuracy result

# normal |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.8264|± |0.0104| | | |strict-match | 5|exact_match|↑ |0.8901|± |0.0086| # asm |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.8294|± |0.0104| | | |strict-match | 5|exact_match|↑ |0.8848|± |0.0088| 

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a new shuffle KV cache layout and an assembly paged attention kernel for the AiterFlashAttentionBackend on ROCm, which demonstrates a significant performance improvement. The implementation adds new Triton kernels and conditional logic to support this new layout. My review identified a critical typo in a new Triton kernel that would cause it to fail. I have provided a code suggestion to correct this issue.