[ROCm] Fix cu_seqlens_q off-by-one in AITER FA speculative decode path

[Bortlesboat]

In the speculative decode path of AiterFlashAttentionImpl (when decode_max_query_len > 1), cu_seqlens_q is sliced as query_start_loc[:num_decodes], giving num_decodes elements. But cu_seqlens_q for unified_attention is a cumulative length array that needs num_seqs + 1 entries (including the leading 0).

The correct pattern is already used in the fallback path 33 lines later (line 1228-1229), which correctly does [:num_decodes + 1].

Also fixed descale_shape which computed shape[0] - 1 on the wrong-length slice, producing num_decodes - 1 instead of num_decodes. Simplified to use num_decodes directly, matching line 1232.

No duplicate PRs found.

Testing:

• uvx pre-commit run ruff-format --files vllm/v1/attention/backends/rocm_aiter_fa.py (passes locally)
• Previous CI failure only requested the matching ruff format whitespace change on this line; this follow-up commit applies that exact diff.

AI assistance was used for code search and CI log review; all changes reviewed by the human submitter.

Signed-off-by: Bortlesboat bortstheboat@gmail.com

Reference: The correct cu_seqlens_q slicing is confirmed by the upstream AITER unified_attention implementation:
https://github.com/ROCm/aiter/blob/bc5ea32c19c604cda2f4781a97cb04d5fc494543/aiter/ops/triton/_triton_kernels/attention/unified_attention.py#L426

[gemini-code-assist]

Code Review

This pull request modifies the ROCm AITER Flash Attention backend to update the calculation of descale_shape and the slicing of cu_seqlens_q passed to the unified_attention function. These changes ensure the correct number of decodes and sequence lengths are used. I have no feedback to provide.

[AndreasKaratzas]

While this sounds theoretically correct, I think that there is a difference in the behavior of flash attention on ROCm VS upstream. I will enable the spec decode tests just to be sure, but they are currently passing.

[tjtanaa]

LGTM. Thanks for catching this. Add this link as a proof for review.

https://github.com/ROCm/aiter/blob/bc5ea32c19c604cda2f4781a97cb04d5fc494543/aiter/ops/triton/_triton_kernels/attention/unified_attention.py#L426

[mergify]

Hi @Bortlesboat, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[Bortlesboat]

Added the AITER reference link to the description — that unified_attention.py line makes it clear cu_seqlens_q needs num_seqs + 1 entries. Thanks for the pointer.

[Bortlesboat]

Remaining red on this PR is buildkite/ci/pr/multi-modal-models-standard-4-other-plus-whisper on commit 01f4f2b.

I re-checked the pipeline config and test layout: that lane is the broad multimodal + Whisper suite gated on any vllm/ change, while this PR only touches vllm/v1/attention/backends/rocm_aiter_fa.py. The targeted spec-decode lane (buildkite/ci/pr/v1-spec-decode) is green on the same commit.

I don't have Buildkite access from here to manually rebuild that step, but this looks unrelated to the speculative-decode fix. Could an authorized maintainer retry the failing Buildkite step?

[Bortlesboat]

Quick recheck on this one: the branch is still the same approved change, buildkite/amd-ci is green, and buildkite/ci/pr/v1-spec-decode is green. The only remaining red is still buildkite/ci/pr/multi-modal-models-standard-4-other-plus-whisper from build 60608. If that lane still looks unrelated on your side, could someone with Buildkite access rerun it or mark this ready?

[AndreasKaratzas]

@Bortlesboat probably rebase the branch. if the issue persists, we will post the PR on slack

[Bortlesboat]

Rebased onto current main per the note above. Patch content is still the same narrow rocm_aiter_fa.py fix, and CI is rerunning on 7f9f7aa.

[Bortlesboat]

Quick update after the rebase rerun on 7f9f7aa: buildkite/ci/pr/v1-spec-decode is green, the earlier multimodal lane is green now too, and the only red in the main PR build is buildkite/ci/pr/cpu-language-generation-and-pooling-model-tests. buildkite/amd-ci is still running.

This patch is still just the small vllm/v1/attention/backends/rocm_aiter_fa.py fix, so that CPU language-generation/pooling failure looks unrelated from my side. If it looks the same with Buildkite access, could someone retry that lane when convenient?

[Bortlesboat]

Merged current main to refresh the branch on the latest base. The functional diff is still the same 4-line vllm/v1/attention/backends/rocm_aiter_fa.py off-by-one fix, and CI is rerunning on 4d55a91.