[ROCm][FEAT] Support Decode Context Parallelism for ROCM_AITER_FA Backend

[vllmellm]

Purpose

This PR enables Decode Context Parallelism (DCP) support for the ROCM_AITER_FA backend in vLLM, enabling improved performance for decoding operations on AMD GPUs with the AITER Flash Attention implementation.

Test Plan

Evaluation on GSM8K benchmark dataset.

lm_eval --model local-completions \
--tasks gsm8k \
--model_args model=Qwen/Qwen3-235B-A22B-Instruct-2507,base_url=http://localhost:9099/v1/completions \
--trust_remote_code \
--num_fewshot 5 \
--batch_size 128

Test Result

With DCP Enabled

Configuration: Qwen/Qwen3-235B-A22B-Instruct-2507 with DCP enabled (TP=8, dcp=2)

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9136	±	0.0077
		strict-match	5	exact_match	↑	0.9045	±	0.0081

Without DCP

Configuration: Qwen/Qwen3-235B-A22B-Instruct-2507 without DCP (TP=8)

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9100	±	0.0166
		strict-match	5	exact_match	↑	0.9033	±	0.0171

Benchmark Result

Performance benchmarking was conducted using vllm bench serve with the following command:

vllm bench serve \
    --model Qwen/Qwen3-235B-A22B-Instruct-2507 \
    --backend vllm \
    --endpoint /v1/completions \
    --host localhost \
    --port 9099 \
    --num-prompts 160 \
    --max-concurrency 16 \
    --dataset-name random \
    --random-input-len <INPUT_LEN> \
    --random-output-len <OUTPUT_LEN>

Benchmark Setting 1: 2K Input / 1K Output

Configuration: 160 prompts, max concurrency 16, input length 2000, output length 1000

Metric	DCP1 - FULL_AND_PIECEWISE	DCP1 - PIECEWISE	DCP2 (DCP enabled)
Successful requests	160	160	160
Failed requests	0	0	0
Maximum request concurrency	16	16	16
Benchmark duration (s)	210.46	209.47	841.21
Total input tokens	320000	320000	320000
Total generated tokens	160000	160000	160000
Request throughput (req/s)	0.76	0.76	0.19
Output token throughput (tok/s)	760.23	763.82	190.20
Peak output token throughput (tok/s)	848.00	848.00	208.00
Peak concurrent requests	32.00	32.00	29.00
Total Token throughput (tok/s)	2280.69	2291.46	570.61
Mean TTFT (ms)	834.35	847.72	837.53
Median TTFT (ms)	776.04	967.38	911.24
P99 TTFT (ms)	1534.03	1519.85	1949.84
Mean TPOT (ms)	20.23	20.11	83.35
Median TPOT (ms)	20.21	20.11	83.33
P99 TPOT (ms)	21.21	21.10	84.15
Mean ITL (ms)	20.23	20.11	83.35
Median ITL (ms)	19.41	19.19	82.46
P99 ITL (ms)	20.49	20.52	87.20

Benchmark Setting 2: 1K Input / 4K Output

Configuration: 160 prompts, max concurrency 16, input length 1000, output length 4000

Metric	DCP1 - FULL_AND_PIECEWISE	DCP1 - PIECEWISE	DCP2 (DCP enabled)
Successful requests	160	160	160
Failed requests	0	0	0
Maximum request concurrency	16	16	16
Benchmark duration (s)	790.25	782.16	3301.72
Total input tokens	160000	160000	160000
Total generated tokens	640000	640000	640000
Request throughput (req/s)	0.20	0.20	0.05
Output token throughput (tok/s)	809.87	818.25	193.84
Peak output token throughput (tok/s)	848.00	864.00	208.00
Peak concurrent requests	32.00	32.00	32.00
Total Token throughput (tok/s)	1012.33	1022.81	242.30
Mean TTFT (ms)	531.77	572.44	702.43
Median TTFT (ms)	492.98	503.75	634.58
P99 TTFT (ms)	878.07	827.54	979.62
Mean TPOT (ms)	19.63	19.42	82.39
Median TPOT (ms)	19.55	19.39	82.48
P99 TPOT (ms)	20.38	20.16	83.03
Mean ITL (ms)	19.63	19.42	82.39
Median ITL (ms)	19.44	19.20	81.90
P99 ITL (ms)	20.64	20.47	86.83

Performance Analysis

Decode Stage Speculation

The benchmark results show that DCP2 significantly underperforms compared to both DCP1 configurations, particularly in the decode stage (as evidenced by the 4x higher TPOT metrics). Based on the implementation in rocm_aiter_fa.py, we identify the following performance bottleneck:

Root Cause: Incompatibility with Paged KV Cache

The DCP decode path uses aiter.flash_attn_varlen_func, which requires contiguous KV cache. Since vLLM uses paged KV cache, the implementation must gather cache blocks into contiguous buffers via cp_mha_gather_cache on every decode step. This involves block table lookups and memory copies from scattered locations.

In contrast, the non-DCP path uses torch.ops.aiter.paged_attention_v1, which works directly with paged KV cache without gathering overhead.

The cache gathering overhead outweighs the benefits of decode context parallelism, causing the observed 4x throughput degradation. Future optimization should use a paged attention kernel that returns lse information.

---

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.

[mergify]

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

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

[hmellor]

Closing as stale: no author activity since the last push on 2025-11-26, and the approach is superseded by changes to the AITER FA backend and the DCP opt-in gate added in #55780.