[ROCm][Perf][Attention] Add FlyDSL gfx950 prefill attention backend

[akii96]

Important

Needs ROCm/FlyDSL#1056 merged first, and then picked up by the flydsl install that ships with aiter. Until that lands the backend has no kernel to call.

Motivation

On ROCm the attention kernel that serves long prompts is the same one that serves single token decode steps. Those two jobs want different things: a long prompt has plenty of work to spread across the GPU, while a decode step has almost none and mostly waits on memory. Using one kernel for both gives up throughput on the prompt side.

FlyDSL has a gfx950 kernel built for the prompt side. It is meaningfully faster there and much slower on decode, so it is not a drop in replacement. What is missing is a way to use it only where it wins.

Technical Details

Adds ROCM_FLYDSL_DUALWAVE_ATTN, which sends causal prefill batches to the FlyDSL dualwave kernel and lets everything else, decode included, fall through to the AITER unified attention kernel it subclasses. It reuses the parent's KV cache layout, metadata builder and cache update path, so the two are interchangeable without touching the cache.

The name says dualwave because FlyDSL also ships MLA decode, fp8 and MoE kernels, and this backend is specifically the gfx950 dualwave prefill path.

Which batches go to FlyDSL is decided per batch from the batch shape alone, through four checks:

check	keeps out
q-block occupancy	batches too small to fill the GPU, measured 0.25x to 0.96x below the threshold
minimum context	short contexts, measured 0.68x at ctx 1024
decode rows in the batch	mixed batches the kernel loses on, measured 0.92x at 16 rows. The cap scales with context, since a deeper KV walk amortises more rows
causal q-block ordering	picks the ordering that suits the batch, worth up to 1.23x on all prefill work and negative on batches with a near empty row

Every term is a host side shape fact, so one shape always traces one launch sequence and the decision stays fixed for a CUDA graph capture key. num_decode_rows is derived from query_start_loc_cpu alongside the existing num_short_query_rows, sharing one numpy diff, so there is no device read and no sync.

Enable with VLLM_ROCM_USE_FLYDSL_DUALWAVE_PREFILL_ATTN=1, or select it directly with --attention-backend ROCM_FLYDSL_DUALWAVE_ATTN. The kernels are not in the flydsl wheel, so VLLM_FLYDSL_KERNELS_PATH points at a FlyDSL checkout.

Test Plan

Unit tests in tests/v1/attention/test_rocm_flydsl_attn_dispatch.py replay measured shapes through the dispatch rules and assert the faster kernel is chosen for each.

Then an A/B against a server running only AITER UA, same model and flags on both arms, on amd/MiniMax-M3-MXFP4 at TP=4 with 133k max model len. Low and high concurrency across three ISL/OSL points, followed by gsm8k on each arm.

Test Result

39 passed, 15 warnings in 17.27s

Output token throughput, mean TPOT, mean TTFT:

ISL / OSL	conc	tok/s aiter	tok/s flydsl		TPOT	TTFT
8k / 1k	8	740.5	740.3	flat	flat	flat
8k / 1k	32	1723.1	1723.3	flat	flat	flat
32k / 1k	8	449.0	535.9	+19.3%	-19.0%	-7.0%
32k / 1k	32	771.9	914.4	+18.5%	-20.7%	+6.5%
60k / 600	8	264.0	305.5	+15.7%	-13.6%	-13.6%
60k / 600	32	331.6	397.2	+19.8%	-19.1%	-8.7%

The 8k runs are flat by design. At 8k a scheduler step is one chunk plus more decode rows than the cap allows, so those batches fall back and both arms run the same code. Kernel level, admitting them measures 0.92x.

The one TTFT regression, 32k/1k C32, is a scheduling artefact of ISL == --max-num-batched-tokens. At 32768 both, the scheduler fits exactly one prefill per step with no slack, so prefills serialise and TTFT is set by queue position rather than prefill speed. Varying only that flag, with the same workload and backend:

max-num-batched-tokens	tok/s	mean TTFT	median TTFT
16384	889.2	6039.6	4197.0
32768	914.4	8515.3	7992.6
65536	869.2	10750.9	12148.0

At 16384 this config beats the AITER baseline on every metric: throughput +15.2%, mean TTFT -24.4%, median TTFT -37.2%, median E2EL -14.4%.

gsm8k, full set, 8 shot:

	flexible	strict
aiter	0.9469	0.9477
flydsl	0.9325	0.9333

Reproduction

Both arms use the same flags; only --attention-backend differs. Drop the two FlyDSL env vars and set --attention-backend ROCM_AITER_UNIFIED_ATTN for the baseline arm.

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

export VLLM_ROCM_USE_FLYDSL_DUALWAVE_PREFILL_ATTN=1
export VLLM_FLYDSL_KERNELS_PATH=/path/to/FlyDSL

vllm serve amd/MiniMax-M3-MXFP4 \
  --port 8000 \
  --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_FLYDSL_DUALWAVE_ATTN \
  --moe-backend aiter \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --no-enable-log-requests

Sweep, per ISL/OSL and concurrency:

vllm bench serve \
  --backend openai-chat --endpoint /v1/chat/completions \
  --model amd/MiniMax-M3-MXFP4 --host localhost --port 8000 \
  --dataset-name random \
  --random-input-len 32768 --random-output-len 1024 --random-range-ratio 0 \
  --num-prompts 128 --max-concurrency 32 \
  --ignore-eos --num-warmups 1 --seed 1234 \
  --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,99

gsm8k:

lm_eval --model local-chat-completions \
  --model_args model=amd/MiniMax-M3-MXFP4,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=32,tokenized_requests=False,max_gen_toks=16384 \
  --tasks gsm8k --num_fewshot 8 \
  --apply_chat_template --fewshot_as_multiturn \
  --batch_size 32

TO-DO

The gsm8k difference is 1.4 points at a stderr of about 0.006, measured once per arm. Comparing the two arms token by token on the same prompts, 80 percent of answers diverge somewhere but only 6 percent of those change the final answer, and the divergences sit at a median logprob margin of 0.25 with a quarter of them at exactly zero. That points at near tie tokens being re-rolled rather than reasoning degrading. But need to investigate this over multiple runs so listing this as a to-do item