[ROCm][MLA] Support block-size > 1 for AITER MLA backend

[ganyi1996ppo]

Purpose

The AITERMLABackend now only support block-size=1 scenario for inference. This constrain may lead to some serious host overhead when we are about to allocate or free cache blocks for long context requests cause there might exist large amount of blocks to operate. Thanks to the insights of @gyu-amd .

In this PR, we remapping the block_table to 1 block size case every step in AITERMLAMetadataBuilder to alleviate the host overhead during allocate and deallocate blocks This change also helps to support a wider range of block size for AITERMLABackend, makes the AITERMLABackend on ROCm platform aligns with the vllm's official usgae and more flexible .

Test Plan

Verified on gsm8k for accuracy, performance improvement will also be attached later
test script:

export VLLM_USE_V1=1
export SAFETENSORS_FAST_GPU=1
export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_USE_TRITON_FLASH_ATTN=0
export NCCL_DEBUG=WARN
export VLLM_RPC_TIMEOUT=1800000
export VLLM_ROCM_USE_AITER_ASMMOE=1
export VLLM_ROCM_USE_AITER_MHA=0
export VLLM_ROCM_USE_TRITON_ROPE=1

model_path="deepseek-r1-FP8-Dynamic"
vllm serve $model_path \
  --tensor-parallel-size 8 \
  --max-num-batched-tokens 32768 \
  --trust-remote-code \
  --no-enable-prefix-caching \
  --disable-log-requests \
  --gpu_memory_utilization 0.9 \
  --block-size 128 \
  --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}'

Test Result

# gsm8k test
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9507|±  |0.0060|
|     |       |strict-match    |     5|exact_match|↑  |0.9484|±  |0.0061|

---

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 enables support for block sizes greater than 1 for the AITER MLA backend on ROCm, which was previously a limitation. The approach of remapping the block table to token-level indices to match the expectation of the underlying AITER kernel is sound. The implementation is largely correct, but I've identified a potential issue with a hardcoded device string that could lead to runtime errors in multi-GPU environments. Addressing this will improve the robustness of the change.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Recompute last page lengths after token-level remapping

After expanding each block table entry into per-token indices, the code still derives paged_kv_last_page_len from the original page_size (seq_lens % page_size, falling back to page_size). Once the remapping is done, each entry represents a single token, so the last-page length for any non-empty request should always be 1. Keeping the old computation causes the decode kernel to believe that the final page contains page_size tokens (e.g. 128) and it will read that many elements starting from the last token’s index, potentially stepping past the valid token range when block_size > 1. This defeats the goal of supporting larger block sizes and can lead to out-of-bounds accesses or garbage attention results for any request longer than one token.

Useful? React with 👍 / 👎.

[ganyi1996ppo]

@HAIAI Please kindly help to review this PR.

[tjtanaa]

Just a sharing of the performance metric of this amazing optimization PR. There is improvement even in the original support block-size=1. 🚀

Here's a comparison table of the benchmark results on DeepSeek-R1 PTPC FP8:

General Metrics

Metric	Before PR (Block-size 1)	After PR (Block-size 1)	After PR (Block-size 16)	Best Performance
General Performance
Successful requests	320	320	320	All equal
Benchmark duration (s)	359.51	354.92	360.31	Block-size 1
Total generated tokens	298,762	294,597	300,636	Block-size 16
Request throughput (req/s)	0.89	0.90	0.89	Block-size 1
Output token throughput (tok/s)	831.03	830.05	834.37	Block-size 16
Peak output token throughput (tok/s)	1,056.00	1,088.00	1,088.00	Block-size 1 & 16
Total token throughput (tok/s)	4,020.30	4,060.55	4,016.48	Block-size 1

Latency Metrics

Latency Metrics	Before PR	After PR (Block-size 1)	After PR (Block-size 16)	Best Performance
Mean TTFT (ms)	1,923.96	1,522.36	1,742.48	After PR (Block-size 1)
Median TTFT (ms)	1,686.80	1,411.06	1,655.13	After PR (Block-size 1)
P99 TTFT (ms)	5,553.48	5,530.39	5,531.85	After PR (Block-size 1)
Mean TPOT (ms)	57.56	53.10	58.82	After PR (Block-size 1)
Median TPOT (ms)	35.44	36.28	35.15	After PR (Block-size 16)
P99 TPOT (ms)	721.81	805.45	647.23	After PR (Block-size 16)
Mean ITL (ms)	35.03	35.79	34.86	Before PR
Median ITL (ms)	31.29	31.43	31.02	After PR (Block-size 16)
P99 ITL (ms)	209.77	211.47	210.17	Before PR

Workload

#!/bin/bash
PORT=8000
SEED=0
CONCURRENCY=32
NREQUESTS=$(($CONCURRENCY * 10))
ISL=3584
OSL=1024
vllm bench serve --backend vllm \
--model EmbeddedLLM/deepseek-r1-FP8-Dynamic \
--dataset-name random \
--num-prompts ${NREQUESTS} \
--random-input ${ISL} \
--random-output ${OSL} \
--seed ${SEED} \
--max-concurrency ${CONCURRENCY} --port ${PORT}

[ganyi1996ppo]

Just a sharing of the performance metric of this amazing optimization PR. There is improvement even in the original support block-size=1. 🚀

@tjtanaa Thanks for the benchmark metric you shared! I'm actually quite surprise that the block-size=1 get performance boost, we tested this PR with block-size=128 and see slight performance boost over previous result with block-size=1

[tjtanaa]

@ganyi1996ppo can you share the performance value of your experiment for block-size=128 vs block-size=1?

[ganyi1996ppo]

hi @tjtanaa, We tested the performance with 8k/1k and 32/1k configure for different scenario, the block-size=128 seems always get better perf compared with block-size=1 case when prefix-cache is off. Although the perf improvement is almost negligible on tps

block-size=128 with 8k/128:

block-size=1 with 8k/128:

block-size=128 with 32/1k

block-size=1 with 32/1k

[HAIAI]

LGTM

[gshtras]

Does this change require any particular AITER version or branch?

[ganyi1996ppo]

Does this change require any particular AITER version or branch?

@gshtras No specific aiter version is required, it just maps block-size > 1 block table to block-size=1 block table during construct metadata. So if your aiter works fine with previous MLA, it works fine with this PR

[mergify]

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

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