[ROCm][Perf] gfx950: use FlyDSL fp8 MQA logits kernel for the sparse indexer prefill

[amd-sriram]

Purpose

Extends #49544 to gfx950: the sparse indexer's prefill MQA logits use the AITER FlyDSL kernel on gfx950 as well as gfx942 whenever AITER is enabled, instead of the Triton/Gluon kernel.

Needs ROCm/aiter#4538, which is in AITER v0.1.23 (the version vLLM main pins).

Test Plan

• pytest tests/v1/attention/test_rocm_flydsl_mqa_logits_dispatch.py: gfx942 and gfx950 use FlyDSL when AITER is enabled; other archs use Triton/Gluon; AITER disabled uses the torch fallback. 4/4 pass on MI355X.
• End to end A/B on GLM-5.2-MXFP4: throughput sweep, GPQA Diamond and RULER niah_single_2.

Test Result

GLM-5.2-MXFP4, TP4, MI355X. Same build for both runs; the only difference is the prefill MQA-logits kernel (Gluon before, FlyDSL with this change).

Accuracy

Eval	Gluon	FlyDSL
GPQA Diamond	0.899	0.884
RULER niah_single_2, 64k	1.00	1.00
RULER niah_single_2, 128k	1.00	1.00

GPQA is a single run of 198 questions (stderr about 0.02).

ISL 60000 / OSL 600

Conc	Tput/GPU Gluon	Tput/GPU FlyDSL	Δ	TTFT Gluon (ms)	TTFT FlyDSL (ms)	Δ
4	3372.9	3419.9	+1.4%	2866.7	2708.7	−5.5%
8	4250.0	4297.1	+1.1%	3589.9	3559.9	−0.8%
16	4890.4	4964.6	+1.5%	5697.8	5505.0	−3.4%
32	5315.0	5399.4	+1.6%	13285.8	12848.0	−3.3%

ISL 8192 / OSL 1024 and ISL 1024 / OSL 1024 (tput/GPU)

Conc	8k Gluon	8k FlyDSL	Δ	1k Gluon	1k FlyDSL	Δ
4	645.1	646.2	+0.2%	156.6	155.5	−0.7%
8	1058.7	1053.9	−0.4%	270.4	267.4	−1.1%
16	1593.2	1586.5	−0.4%	436.9	434.1	−0.6%
32	2248.1	2232.8	−0.7%	679.3	674.7	−0.7%
64	3118.3	3106.5	−0.4%	1060.7	1051.9	−0.8%
128	4060.3	4042.8	−0.4%	1637.7	1624.8	−0.8%
256	4925.5	4907.5	−0.4%	2424.7	2414.3	−0.4%

Geomean over all 18 points: 1560.4 → 1558.5 tok/s/GPU (−0.1%).

Kernel trace (60k/600 profile, rank 0)

Prefill MQA-logits kernel time: 13,974 ms over 17,738 launches → 9,118 ms over 17,769 launches (−35%, 1.53×). On the same prefill chunk (third 16k chunk of a 60k prompt): 153.4 ms → 98.7 ms.

Gluon (_gluon_fp8_mqa_logits_kernel, split-KV pair per layer):

FlyDSL (fp8_mqa_logits_H32_D128_mfma32x32x64_bkv64_r2_w2_lds3_swizzled_nocl_flydsl):

Both views show the same 6.4 ms window.

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

[github-actions]

✅ @amd-sriram, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /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.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[dllehr-amd]

/ci run

[github-actions]

❌ This PR is 15 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[simondanielsson]

Thanks @amd-sriram! Do we also have the TTFT numbers available for the short ctx prefill (1k and 8k), in addition to the e2e throuhgput?

Also optionally, if we'd like to be even more confident we're not regressing here, we could add a benchmark script in /benchmarks/kernels/ and run it for different shapes on gfx950 and compare it to the gluon baseline

[simondanielsson]

Suggestion: can you please check if we break test_unpaged_mqa_logits_preserves_intervals_and_scale_layout which assumes the triton kernel will be the one running on gfx950? I.e. run this on an mi350 node

pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py

[simondanielsson]

Nit: Does this line test anything given we return out from all of the mocked implementations?

[tjtanaa]

Also optionally, if we'd like to be even more confident we're not regressing here, we could add a benchmark script in /benchmarks/kernels/ and run it for different shapes on gfx950 and compare it to the gluon baseline

Since the script is just to compare AITER kernels perf, it is ideally having the script live in AITER. We should keep the kernel benchmark scripts necessary mainly for kernels that live in repo. Because the benchmark script should be a developer tool and we should be able to use agent to regen them.

However, it is good to disclose the microbenchmark details, and attach the script to the PR description with

Details

tag as a record.

[amd-sriram]

@tjtanaa @simondanielsson Thanks for the comments. I am currently working on migrating this to AITER similar to paged mqa logits - ROCm/aiter#6140. I have got some comments from the AITER team. Once I understand how to write the AITER side top level method, I will port this prefill mqa logits to AITER too. Currently, I will mark this as draft and close it when I have the new vllm side PR.

[amd-sriram]

@tjtanaa @simondanielsson I have create an AITER PR to do backend selection for prefill mqa logits - ROCm/aiter#6185. I have drafted a new vLLM PR to call this AITER function. I will open the vLLM PR when the AITER PR gets merged.