[Perf] [Feat] [ROCm] Add densemha support to ROCm AITER Sparse MLA

[tjtanaa]

Purpose

This PR follows #47327 in optimizing the performance for short sequence prefill of ROCm AITER Sparse MLA backend.

Test Plan

1. Update unit tests

2. Add unit tests to make sure that we cover the cases below for metadata and correctness tests: Permute the following aspects:

   • Prefill length: seqlen < topk_tokens and seqlen >= topk_tokens
   • Workload: Prefill only, Mixed prefill decode, Decode only
   • Dtypes: bfloat16, fp8

3. GSM8K with 30 num-shot

4. Performance gain on DeepSeek V3.2

Test Result

Unit test

pytest tests/v1/attention/test_sparse_mla_backends.py

========== 112 passed, 28 skipped, 238 warnings in 177.00s (0:02:56) ===========

pytest tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py

======================= 26 passed, 18 warnings in 12.35s =======================

Server command:

#!/bin/bash

export VLLM_ROCM_USE_AITER=1

rm -rf ~/.cache/vllm

vllm serve deepseek-ai/DeepSeek-V3.2 \
  --trust-remote-code \
  --kernel-config.enable_flashinfer_autotune=False \
  --tensor-parallel-size 8 \
  --tokenizer-mode deepseek_v32 \
  --tool-call-parser deepseek_v32 \
  --enable-auto-tool-choice \
  --reasoning-parser deepseek_v3 \
| tee dsv32.log

LM-Eval Command

MODEL=deepseek-ai/DeepSeek-V3.2
NUM_FEWSHOT=${NUM_FEWSHOT:-30}
RESULT_TAG=${RESULT_TAG:-dsv32_numshot${NUM_FEWSHOT}_cc256}
lm_eval --model local-completions --model_args model=$MODEL,base_url=http://0.0.0.0:8000/v1/completions,num_concurrent=256,max_retries=10,max_gen_toks=2048,max_length=1048576,timeout=60000,trust_remote_code=True,temperature=0 --batch_size auto --tasks gsm8k --num_fewshot "${NUM_FEWSHOT}" \
  --output_path "./results_${RESULT_TAG}" \
  --log_samples \
| tee "lmeval_${RESULT_TAG}.log"

Few-shot	Metric	All-MQA before	Dense-MHA/AITER FA after	Delta
5	Flexible extract exact match	0.9553 +/- 0.0057	0.9583 +/- 0.0055	+0.0030
5	Strict match exact match	0.9560 +/- 0.0056	0.9591 +/- 0.0055	+0.0030
30	Flexible extract exact match	0.9560 +/- 0.0056	0.9568 +/- 0.0056	+0.0008
30	Strict match exact match	0.9568 +/- 0.0056	0.9575 +/- 0.0056	+0.0008

Perf improvement on DeepSeek-V3.2

Concurrency	Successful requests	Failed requests	Request throughput (req/s)	Output throughput (tok/s)	Total throughput (tok/s)	Mean TTFT (ms)	Mean TPOT (ms)	Mean E2E (ms)
1	5	0	0.077	79.34	158.98	176.88	12.44	12,907.01
2	10	0	0.143	145.92	292.41	122.75	13.60	14,032.09
4	20	0	0.256	261.79	524.60	528.47	14.78	15,644.46
8	40	0	0.456	466.88	935.59	468.60	16.69	17,542.64
16	80	0	0.787	805.88	1,614.90	733.71	19.15	20,324.92
32	160	0	1.225	1,254.85	2,514.60	1,057.46	24.48	26,103.46
64	320	0	1.904	1,949.95	3,907.52	770.40	32.08	33,588.61

Benchmark command:

#!/bin/bash
set -euo pipefail

# After you have launched the server with the command in @launchdeepseekv4graph.sh

BASE_URL=${BASE_URL:-http://127.0.0.1:8001}
RESULT_DIR=${RESULT_DIR:-./dsv4_pr_benchmark}
RESULT_PREFIX=${RESULT_PREFIX:-dsv4_fse_benchmark}
MODEL=${MODEL:-deepseek-ai/DeepSeek-V4-Pro}
CONCURRENCIES=${CONCURRENCIES:-"1 2 4 8 16 32 64"}
INPUT_LEN=${INPUT_LEN:-1024}
OUTPUT_LEN=${OUTPUT_LEN:-1024}

for C in ${CONCURRENCIES}; do
  NUM_PROMPTS=$((C * 5))
  NUM_WARMUPS=$((C * 1))

  vllm bench serve \
    --backend openai-chat \
    --base-url "${BASE_URL}" \
    --endpoint /v1/chat/completions \
    --model "${MODEL}" \
    --dataset-name random \
    --input-len "${INPUT_LEN}" \
    --output-len "${OUTPUT_LEN}" \
    --num-prompts "${NUM_PROMPTS}" \
    --request-rate inf \
    --max-concurrency "${C}" \
    --num-warmups "${NUM_WARMUPS}" \
    --random-range-ratio 0 \
    --ignore-eos \
    --percentile-metrics ttft,tpot,itl,e2el \
    --metric-percentiles 50,90,99 \
    --save-result \
    --result-dir "${RESULT_DIR}" \
    --result-filename "${RESULT_PREFIX}-C${C}.json" \
    --metadata concurrency="${C}" workload=random_${INPUT_LEN}_${OUTPUT_LEN} num_prompts="${NUM_PROMPTS}"
done

---

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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[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.

[tjtanaa]

For Upstream CI:

test_sparse_mla_backends.py from .buildkite/test_areas/attention.yaml with test group named V1 attention (H100-MI300)

test_rocm_aiter_mla_sparse_metadata_sync.py from .buildkite/test_areas/kernels.yaml , with test group named: Kernels Attention Test %N

For AMD CI:

test_sparse_mla_backends.py is covered by V1 attention (H100-MI250), V1 attention (H100-MI300), V1 attention (B200-MI355)

test_rocm_aiter_mla_sparse_metadata_sync.py is covered by Kernels Attention Test %N

[tjtanaa]

When fixing the unit test this issue surfaces.

[tjtanaa]

The issue is surfaced by these two numerical cases:

tests/v1/attention/test_sparse_mla_backends.py::test_sparse_backend_decode_correctness[1.0-1.0-64-2-auto-large_q_pure_prefill-ROCMAiterMLA]

tests/v1/attention/test_sparse_mla_backends.py::test_sparse_backend_decode_correctness[2.0-3.0-64-2-auto-large_q_pure_prefill-ROCMAiterMLA]

The parameters mean:

• BF16 KV cache: auto
• Pure prefill: large_q_pure_prefill
• Prefill length: 256
• topk_tokens: 128
• Simulated TP=2: 64 query heads
• Block size: 64
• Scale pairs: 1.0/1.0 and 2.0/3.0

[AndreasKaratzas]

Is this applied when you run the tests in a bulk? I had a weird case in the past that I dont remember well, where this would not be applied unless I just pytest only this file.

[tjtanaa]

Yes. It does skip the rest.

https://buildkite.com/vllm/amd-ci/builds/11246/list?sid=019f9610-091f-44ae-a7c9-624bafce6453&tab=output

[AndreasKaratzas]

It's LGTM from me, gonna stamp it soon as well (mostly waiting for the answer on my previous comment, although i m sure i m just imagining things .. would like a review from Rohan (@Rohan138) for the src (the sparse mla changes).

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tjtanaa.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork