[ROCm] [Bugfix] Preserve CPU query offsets during capture

[akii96]

Motivation

On ROCm, full CUDA graph capture can fail during server startup for mixed-attention models when a later metadata builder reads a query length of zero. The root cause is that RocmAttentionMetadataBuilder clears both device and CPU query offsets to prevent an invalid memory access in the prefix prefill kernel. The device offsets must be zeroed for safe graph capture, but the CPU offsets are not used by that kernel and remain shared with subsequent metadata builders. Clearing the CPU copy therefore removes valid batch information those builders need.

Proposed Fix

Preserve the CPU query offsets while continuing to zero the device offsets.

This keeps the existing prefix prefill capture behavior unchanged while allowing subsequent metadata builders to calculate valid query lengths. The change is scoped to full graph capture through the ROCm metadata builder shared by ROCM_ATTN and ROCM_AITER_UNIFIED_ATTN.

Test results

Tested on 4 GPUs, MI325, model -> EmbeddedLLM/MiniMax-M3-FP8-dynamic

Before the change, CUDA graph profiling failed:

File "vllm/models/minimax_m3/common/indexer.py", line 317, in build
    assert decode_query_len > 0
AssertionError
EngineCore failed to start.

After the change:

Graph capturing finished in 8 secs, took 0.19 GiB
Application startup complete.
POST /v1/chat/completions HTTP/1.1 200 OK

Reproduction command

export HIP_VISIBLE_DEVICES=0,1,2,3
export VLLM_ROCM_USE_AITER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=0
export VLLM_ROCM_USE_AITER_MOE=0
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1

vllm serve EmbeddedLLM/MiniMax-M3-FP8-dynamic \
  --served-model-name minimax-m3 \
  --tensor-parallel-size 4 \
  --block-size 128 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.92 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 32 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --kv-cache-dtype fp8 \
  --language-model-only \
  --port 8000 \
  --attention-backend ROCM_AITER_UNIFIED_ATTN

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[akii96]

I noticed that the InfX recipe for amd/MiniMax-M3-MXFP4 used trition attention as the attention backend and was curious if they had the same issue as the one my PR is trying to fix and sure enough it does!

Environment:

• Image: vllm/vllm-openai-rocm:nightly-ac7509e2b1db40fec2f03dde1ed4e9dfdc2338c9
• ROCm 7.2.3
• vLLM 0.27.2rc1.dev77+gac7509e2b.rocm723
• AITER 0.1.19

Serving command

export VLLM_USE_BREAKABLE_CUDAGRAPH=1
export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export SAFETENSORS_FAST_GPU=1

vllm serve amd/MiniMax-M3-MXFP4 \
  --tensor-parallel-size 4 \
  --distributed-executor-backend mp \
  --trust-remote-code \
  --block-size 128 \
  --no-enable-prefix-caching \
  --language-model-only \
  --max-model-len 133120 \
  --gpu-memory-utilization 0.95 \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 128 \
  --enable-chunked-prefill \
  --async-scheduling \
  --attention-backend ROCM_AITER_UNIFIED_ATTN \
  --moe-backend aiter \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --no-enable-log-requests

With my fix the AITER UA backend now completes CUDA graph capture and serves normally. Before the fix, startup failed at assert decode_query_len > 0.

Performance versus TRITON_ATTN, averaged over three runs:

Workload	Metric	Triton	AITER UA	Delta
60k/600, C16	Output throughput	274.75 tok/s	314.79 tok/s	+14.57%
60k/600, C16	TTFT	11,807 ms	9,823 ms	-16.80%
60k/600, C16	TPOT	38.56 ms	34.45 ms	-10.65%
60k/600, C16	E2E latency	34,903 ms	30,460 ms	-12.73%
8k/1k, C64	Output throughput	2,374.93 tok/s	2,391.24 tok/s	+0.69%
8k/1k, C64	TTFT	4,127 ms	4,029 ms	-2.36%
8k/1k, C64	TPOT	22.79 ms	22.70 ms	-0.42%
8k/1k, C64	E2E latency	26,898 ms	26,705 ms	-0.72%

Full GSM8K: 1,319 examples, 5-shot, chat template, few-shot as multiturn, and max_gen_toks=16384.

Filter	Triton	AITER UA	Delta
Flexible extract	94.01% (1240/1319)	93.93% (1239/1319)	-0.08 pp
Strict match	94.09% (1241/1319)	94.01% (1240/1319)	-0.08 pp

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84156 for commit a3ad4c0aa1d6.

[github-actions]

✅ @akii96, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[mergify]

Hi @akii96, 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.

[tjtanaa]

@akii96 please fix precommit

[akii96]

@tjtanaa it seems more related to what mergify bot did, the original 1 commit PR was passing precommit already 🤔 I am going to hit to rebase branch in 15 mins unless you have an alternate solution to the issue

[akii96]

@tjtanaa do you have direction on how to proceed here 😅? the mergify bot yesterday failed on precommit when it tried to rebase the branch to main (my initial commit was passing precommit). So I rebased it manually but took a while to run CI checks and now it is 1 day behind. Did I do something wrong here?

cc @AndreasKaratzas potentially for inputs on CI issues

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84643 for commit 8f90293f4ae0.

[akii96]

@tjtanaa CI looks green! Maybe the merge goes through this time! 😄🤞