[ROCm][Perf] New design on ROCm AITER MHA backend Implementation

[ganyi1996ppo]

Purpose

The current AiterFlashAttentionImpl fetches K/V every run, which creates unnecessary memory pressure and non-trivial latency—especially with long prompts. This PR:

• Removes redundant KV fetches from the attention backend
• Introduces a phase-aware execution (decode, pure prefill, chunk prefill) and reorders inputs to [decode:chunk_prefill:pure_prefill] for token-contiguous memory access.
• Rewrites the “fetch KV” Triton kernel for better occupancy in chunked prefill and similar scenarios.

Design and implementation

Phase-aware path:

• decode
• chunk prefill (cp)
• pure prefill (pp)

Input reordering to [decode:cp:pp] ensures tokens are contiguous in memory, improving kernel locality and occupancy. The reorder occurs in both Scheduler's scheduling phase and ModelRunner's state updating phase. We add this split_prefill_from_chunk to the SchedulerConfig to control this behavior, which will be turned on if both VLLM_ROCM_USE_AITER and VLLM_ROCM_USE_AITER_MHA are set.

Compared with the old one, this solution is more memory efficient and fast, especially on the long prompt scenario. Here is the Performance Measured on Qwen3, Mi308:

Long prompt, short output (2k prompt, 16 output): ~4.x throughput improvement.
Short prompt, long output (128 prompt, 1k output): ~2.x throughput improvement.
Extremely long prompt (192k prompt, 2k output): ~5.x throughput improvement.

Test Plan

acc : lm_eval test for accuracy verification
perf : vllm bench test

Test Result

2k prompt 16 output case:

old impl

new impl

128 prompt 1k output case:

old impl

new impl

acc verification

We test this PR on Qwen3-30B-A3B-FP8 on gsm8k with lm_eval, and here is the result:

# previous implementation

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7437|±  |0.0120|
|     |       |strict-match    |     5|exact_match|↑  |0.8817|±  |0.0089|

# new implementation

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7703|±  |0.0116|
|     |       |strict-match    |     5|exact_match|↑  |0.8878|±  |0.0087|

---

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.

[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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a new, more performant MHA backend implementation for ROCm. The changes include removing redundant KV fetches, introducing phase-aware execution (decode, pure prefill, chunk prefill), reordering inputs for better memory access, and rewriting the Triton kernel for fetching KV cache. The performance improvements demonstrated are significant. However, I have identified a few critical bugs in the implementation that need to be addressed. These include incorrect scaling in a Triton kernel, an invalid tensor view operation that will lead to a runtime error, and a logical error in the batch reordering logic. Addressing these issues is crucial for the correctness and stability of the new backend.

[wuhuikx]

Could you please help clarify, which Qwen3 model and datatype are you using? Could you please also append the accuracy results?

[wuhuikx]

cc @wuhuikx @sunway513

[ganyi1996ppo]

Could you please help clarify, which Qwen3 model and datatype are you using? Could you please also append the accuracy results?

Thanks for the suggestion, just update the PR description with model and accuracy verification.

[ganyi1996ppo]

hi @LucasWilkinson , please take a look. Does this PR requires further change?

[LucasWilkinson]

hi @LucasWilkinson , please take a look. Does this PR requires further change?

Apologies for the delay! (FYI OOO currently till next week); and apologies for the review thrash (I think its important to discuss changes that touch core components so I appreciate the patience). Per discussions (#25763 (comment)) can you please rebase this PR off #27367 and remove ReorderSpec. This should unblock and greatly simplify this PR (we can always add ReorderSpec in a follow-up PR post GPU model runner refactor if the spec-decode benefits are needed). Apologies for the inconvenience!

[ganyi1996ppo]

Hi @LucasWilkinson, we found the reorder algorithm in #27367 may generate incorrect result when one position been reordered twice time or more, I put a new testcase to the test_batch_reordering.py for such scenario. I also updated the algorithm in the PR to fix this issue, please take a look and see if its good enough.

[LucasWilkinson]

Hi @LucasWilkinson, we found the reorder algorithm in #27367 may generate incorrect result when one position been reordered twice time or more, I put a new testcase to the test_batch_reordering.py for such scenario. I also updated the algorithm in the PR to fix this issue, please take a look and see if its good enough.

Apologies silly bugs 🤦; does this work for you? #27739

[ganyi1996ppo]

Apologies silly bugs 🤦; does this work for you? #27739

@LucasWilkinson , this fix works fine on my environment, thanks!

Apologies silly bugs 🤦; does this work for you? #27739

@LucasWilkinson , this fix works fine on my environment, thanks!

Thanks for testing that! Appreciate all the back and forth, and patience!

Once that lands this PR should be purely in the rocm/AITER backend so should be quick to review and get in. I am OOO this week (despite being online more than I intended haha); I'll try to do the final review. Otherwise @gshtras can also review it as the AMD code owner (now that it's purely an AMD change)

[ganyi1996ppo]

Once that lands this PR should be purely in the rocm/AITER backend so should be quick to review and get in. I am OOO this week (despite being online more than I intended haha); I'll try to do the final review. Otherwise @gshtras can also review it as the AMD code owner (now that it's purely an AMD change)

Thanks @LucasWilkinson for your time and thoughtful review ! That's exactly what we wanted.

[ganyi1996ppo]

@gshtras Rebase is done, please also take a look.

[LucasWilkinson]

overall LGTM; left a couple final nits, not super familiar with AMD backend so didn't review it that closely but seems ok (would still be nice if @gshtras takes a look but not a requirement)

[LucasWilkinson]

nit: lets stick to the extend/prefill naming instead of prefill/pure_prefill

[ganyi1996ppo]

Done

[LucasWilkinson]

nit: could refactor this to

extend_slice = slice(num_decode_tokens, num_decode_tokens + num_extend_tokens)
extend_querys = query[extend_slice]
extend_keys = key[extend_slice]
extend_values = value[extend_slice]
extend_outputs = output[extend_slice]

to use less lines

[ganyi1996ppo]

Done