[ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 - #62590

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

MI300X (gfx942) numbers for #62590 @ f87fca3

Setup: MI300X, 304 CUs, one GPU (on_gfx942() is True). Images: vllm/vllm-openai-rocm:nightly sha256:5800b3b4cd79dc…c84f (0.33.0rc1.dev573, torch 2.13+rocm10.0, triton 3.8.0; the 7-line hunk applied to ops/qsa.py with patch, offset 14) and :v0.33.0-rocm72 sha256:4fe7bc121cbe2f…a182 (torch 2.13, HIP 7.2.53211). Host ROCm 7.0.2. The PR's bench_gqa.py, stock vs stock+hunk.

Dispatch: TP1 (24q/2kv) and TP2 (12q/1kv) have BLOCK_M=16 and use BLOCK_N=64 up to 512 programs, 32 above. TP4 (6q/1kv, BLOCK_M=8) stays at 64 at 512, 513 and 1024 rows.

Speedup on nightly (rocm72 within ~1%):

Max abs diff vs stock: 4.9e-4 for TP1/TP2, 0 for TP4. Against an fp32 torch reference I ran 9 shapes x (bf16, fp8-quantised values stored as bf16) x (default dispatch, forced BLOCK_N=64). The BLOCK_N=32 path is off by at most 3.8e-4; the overall max is 5.07e-4 (TP4, which the PR doesn't touch). All within the 2e-2 tolerance.

The flaky CUDA-graph test I hit in the QSA tests on MI300X is unrelated to this PR; see the #62589 comment.

Caveats: the kernel has no native fp8 path (it asserts bf16), synthetic activations, single seed, host load 7-22 while timing (GPU use 0-2%), the PR's 5x TP4 slowdown at BLOCK_N=32 not re-measured, the PR branch itself not built on nightly.

@eppaneamd @tc3oliver: #65341 and #62590 change the same base_programs > 512 branch of qsa_sparse_paged_attention. I measured both PRs, a few combinations, and the full launch-profile grid on MI300X (gfx942) and MI355X (gfx950), at kernel level only. Neither PR alone is the best profile on gfx942. A combined rule is faster than both PRs (or equal to the better one) in every cell I measured on both archs, and it passes the reference test. I suggest landing it as one PR.

Proposed combined hunk

@tc3oliver add @eppaneamd as a co-author in this PR and add the following change

    else:
        block_n, target_splits, partial_warps = 64, 1, 2
        if current_platform.is_rocm():
            from vllm.platforms.rocm import get_cdna_version

            # One warp per program is faster on CDNA3 and CDNA4, except BF16
            # K/V with BLOCK_M >= 16 on CDNA3, where BLOCK_N=32 is faster
            # (BLOCK_N=32 with one warp faults there).
            cdna_version = get_cdna_version()
            if cdna_version == 3 and block_m >= 16 and not is_fp8:
                block_n = 32
            elif cdna_version in (3, 4):
                partial_warps = 1

get_cdna_version() returns 3 exactly on gfx942 and 4 exactly on gfx950. On CDNA4 (gfx950) this is exactly #65341. On CDNA3 (gfx942) it is #62590 for BF16 K/V at TP1/TP2, and one warp everywhere else.

Results

Speedup over main at 8,192 rows, with the min-max over 520 to 16,384 rows in parentheses. "Bitwise" means the combined output equals main's.

What this means for each PR:

• @tc3oliver ([ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590, this PR): your BF16 numbers reproduce on MI300X (1.4-2.1x at TP1/TP2). With fp8 K/V, BLOCK_N=32 gives only 1.2x. At TP2 with 520 rows it is a 0.78x regression, while one warp at BLOCK_N=64 gives 1.7-1.8x, bitwise equal to main. The kernel does take fp8 K/V: it asserts only that q is bf16, and is_fp8 selects the e4m3 path. So "no native fp8 path" in your MI300X comment doesn't hold, and the fp8 case needs its own profile. If your Qwen3.8-Flash-Next-FP8 end-to-end runs used --kv-cache-dtype fp8, they were on that path. One warp also gives a bitwise-equal 1.2-1.3x at TP4/TP8 on gfx942, which [ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590 leaves unchanged.
• @eppaneamd ([ROCm][Perf][Qwen4Exp] Use one warp for QSA sparse paged attention prefill on gfx950 #65341): your gfx950 claim reproduces on MI355X: one warp is fastest and bitwise equal at all four shapes and both dtypes. I also measured TP1 (2 KV heads), which the PR lists as unmeasured: 1.43-1.75x at bf16 and 1.18-1.59x at fp8. The same rule is a 1.2-1.8x win on gfx942 for everything except BF16 at BLOCK_M >= 16, where [ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590's BLOCK_N=32 is better (1.9x vs 1.55x). You can't simply apply both changes on both archs: BLOCK_N=32 with one warp is HSA_STATUS_ERROR_MEMORY_FAULT on gfx942 with BF16 and BLOCK_M=16. Hence the elif.
• Tests: tp2_split1 ([ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590) and tp2_split1_prefill ([ROCm][Perf][Qwen4Exp] Use one warp for QSA sparse paged attention prefill on gfx950 #65341) are the same case, (513, 12, 1, 1792), so keep one. With the combined qsa.py copied over the installed file, pytest tests/models/qwen4_exp/test_qsa_amd.py -k sparse_paged_attention_matches_reference gives 18 passed on both MI300X and MI355X. The get_cdna_version() form above is bitwise equal to the measured on_gfx942()/on_gfx950() form on both archs (5 shapes × 2 row counts × bf16/fp8). That covers 9 cases × bf16/fp8: the existing cases plus tp2_split1, tp1_split1_prefill, and two extra cases I added, (4096, 12, 1, 1792) and (2048, 3, 1, 1024). main passes the same 18.

Method

• Image: vllm/vllm-openai-rocm:nightly-40941f2c4ba96fba0a7e282a5355e89cd74d9ac5 (image id 65812e58df91, same on both nodes). The baseline is its own qsa.py; each variant is that file with only the launcher hunk changed. The [ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590 and [ROCm][Perf][Qwen4Exp] Use one warp for QSA sparse paged attention prefill on gfx950 #65341 hunks match the PR diffs exactly ([ROCm][Perf][Qwen4Exp] Tune sparse GQA prefill profile on gfx942 and gfx950 #62590 applies with an offset of 14).
• Hardware: one MI300X and one MI355X.
• Accuracy: against an fp64 reference on 64 rows per cell (first 16, last 16 and 32 random rows), rel_l2 for every variant equals main's: 2.26e-3 for bf16, 1.67e-3 to 1.69e-3 for fp8. Outputs with one warp are bitwise equal to main; BLOCK_N=32 differs by at most 2.0e-3 absolute, from accumulation order.

Not covered and possible follow-up

• No end-to-end numbers; this is kernel level only. MI325X and MI350X were not measured. The 24q/1kv shape (BLOCK_M=32, not a Qwen3.8-Flash-Next shape) is in the sweep only: there, one warp is 0.5x on gfx942 at bf16, which the block_m >= 16 arm also covers.
• Headroom: BLOCK_N=16 with one warp at BLOCK_M=16 is a further 1.2-1.3x over the combined rule at ≥ 2,048 rows on both archs. On MI355X at bf16 TP2 it is 2.0-2.1x vs 1.7x, and on MI300X at bf16 TP2 2.3-2.6x vs 1.9-2.1x. It is slower at 520-1,024 rows (MI300X bf16 TP2: 1.12x vs 1.39x at 520), so it would need a row-count threshold. I'd leave it for a separate PR.

@mjkvaak-amd @eppaneamd Thanks for the benchmarks, the analysis and the combined proposal.

I've folded both PRs into d902859 and added @eppaneamd as co-author. In the >512-program profile, gfx942 keeps BLOCK_N=32 for BF16 K/V with BLOCK_M >= 16 and uses one warp otherwise, and gfx950 uses one warp for every shape. Other platforms are unchanged. The reference test now covers BF16 and FP8 K/V for TP1/TP2/TP4/TP8, on both sides of the 512-program boundary, and has one copy of the duplicate TP2 case.

My earlier note that the kernel had no FP8 K/V path was wrong. The description is corrected.

So far I've only run ruff and a CPU check of the dispatch logic, on a Mac. The new commit hasn't run on an AMD GPU yet. I'll run the correctness test on gfx942 later and post the results here. Could you review it and, when you have time, check gfx950?

Ran d902859 on MI325X (gfx942), nightly-76c6f3f2 image, main vs this commit, same GPU and inputs, 3 rounds.

pytest -k sparse_paged_attention_matches_reference: 24 passed. Whole test_qsa_amd.py: 37 passed. No HSA fault or NaN.

Launch config matches the commit: BF16 TP2 gets BLOCK_N=32 with 2 warps, everything else above 512 programs gets BLOCK_N=64 with 1 warp (main: 64, 2 warps).

qsa_sparse_paged_attention latency in ms, median of 3 rounds (rounds within ~1%), FP8 with real fp8 K/V:

No cell is slower than main. Output matches main except BF16 TP2, which differs by up to 2e-3 (test tolerance is 2e-2). Kernel microbenchmark only, not end-to-end, gfx942 only.

Timed as 30 back-to-back launches per sample, so host launch overhead is included.