[Gluon][MLA] Deeper async-copy pipeline in the bh16 stage-1 decode loop

[amd-ethany]

This is a follow-up PR to PR #4555, #4509, #4507.

On this baseline the PR is worth about +10% of the MLA decode kernel at concurrency 1, 24 and 48 —
commit 3 carries the first, the wait change carries the other two at +10.5% each — and it costs 2.2%
at concurrency 8. Every one of those figures reproduces in every round of every run. The regression is
the more important half of that sentence and is quantified in section 3.

Every number here is kernel-level — one decode op, device time, cold.

Motivation

PR #4555 restored split parallelism and shortened the stage-2 merge, which moved the bottleneck rather
than removing it. Its policy gives each request roughly 256/batch splits, trading split count against
how far each stage-1 program has to walk, and that leaves a gap at each end of the trade. One PR covers
both because each gap only appears at its own end of the split range.

High concurrency: few splits, so a long loop, and stage-1 waiting on memory. At 48 requests each gets
5 splits, so one program walks 213 KV blocks — and that loop drained every outstanding copy on every
iteration when only one of them was needed yet. Only one KV block is ever in flight, so a 213-iteration
loop runs latency-bound instead of bandwidth-bound. Commit 1 closes this, the larger of the two gaps.

Batch 1: splits are plentiful, but stage-2 has almost nothing to run them on. Its grid is one
workgroup per request, head and query position — 12 of them on a 256-CU part at nhead 12 — so shortening
the merge chain cannot help when the idle machine is what costs the time. Commit 3 widens that grid
instead, but only while the machine is in fact idle, which is why it is gated: tiling unconditionally
made nhead-64/128 shapes, whose grid already runs into the thousands of programs, 4.2% slower.

Technical Details

1. Two partial waits in the stage-1 loop. The top of the loop needs only one of the four copies the
previous iteration issued, the page number, but it drained all four. Retire just that one there and defer
the rest to the point where their data is actually read. Two KV blocks are then in flight at no extra LDS,
which is the granularity both epilogues already use. The prologue leaves the same copies outstanding as a
steady-state iteration, and the set at loop exit is unchanged, so neither the first iteration nor the
epilogues change behaviour. This is the one change in code that every regime runs, bh64 included.

2. Split-major stage-1 grid. Swap two grid axes for bh16 so that all requests' workgroups for a
given split get consecutive ids. They then land on one XCD and share its L2, and under prefix caching
they are reading the same KV rows. Only taken when batch > 1 and the split count is a multiple of the XCD
count, which is what makes that hold for every request rather than some. Pure relabelling — same KV range
per program, same reduction order, bit-identical output.

3. Head-dim-tiled stage-2 reduce. Split the 512-wide head dimension across several programs, each
owning a disjoint slice of the output and re-deriving the same scalar statistics independently, so slices
never communicate and the untiled case is today's kernel. Because each slice repeats that derivation,
tiling buys occupancy by duplicating work, so the launcher tiles only while the grid would otherwise
leave CUs idle — up to 8 ways, and less when the head dimension will not divide evenly. Stage-2's warp
count follows the tiling, with the untiled path keeping main's 8. That coupling matters: at batch 1 the
warp count is the difference between +33.6% and +18.5% of stage-2, and at concurrency 48 between −4.5%
and −19.1%.

Test Plan

One process, one kernel file. All three commits are switches on a single generated kernel, so every
leg runs against the same allocations — no JIT-cache or allocator drift of the kind that swapping whole
files between processes invites. Both endpoints are checked against the real kernel files every run:
everything off reproduces main to within 1 ulp, everything on reproduces #4509 rebased bitwise.

Timing. Device time from CUDA events over a captured graph, with a 512 MB cache flush before each cold
sample. 5 rounds, legs rotated between them, each ratio formed inside its own round; the reported value is
the median, and every number quoted had all 5 rounds on the same side of zero.

Workload. 68,000 tokens, a 63,240-token prefix shared across 8 requests, nhead 12 (96 at TP8), page
size 1, and a full-size KV pool so the 64-bit addressing path is taken. Concurrency 1–48 at whatever split
count the policy picks, plus a sweep that sets the split count independently of batch.

aiter's own MLA tests, both configs, against unpatched main back to back in the same container:

python op_tests/test_mla.py -c 327600 -b 8 -n 12,1 -d bf16 -kvd bf16
python op_tests/test_mla.py -c 16384 -b 64 128 -n 64,1 128,1 -d bf16 -kvd bf16

Both pass, 6 of 6 Gluon verdicts on each leg of the second config.

Which regime gets what. Commit 2 applies to bh16 only. Commits 1 and 3 are in code every regime
runs, so the second command above is the one that covers bh64 (nhead 64/128) — and there commit 3's gate
leaves stage-2 exactly as main has it, with timings level against unpatched main. That gate is the
whole point: tiling unconditionally, as this PR was first written, ran bh64 4.2% slower. Commit 1 is not
implicated, since reverting its waits leaves bh64 where it was. Attributed by reverting one commit at a
time, with a second main process as the control for what this test's timing can resolve.

Test Result

MI355X (gfx950, 256 CU), rocm/vllm-dev@sha256:5aa7e626.

1. Per-lever, whole decode op, 5-round median

trips is loop iterations per stage-1 program, SEQ_LEN/splits/BLOCK_N — the variable that governs
commit 1. The commit 3 column was measured before its occupancy guard, which keeps the full 8-way tiling
at concurrency 1–4 and reduces it above that, where the column is already ≤ +0.12%.

conc	splits	trips	commit 3	commit 1	commit 2	all three
1	256	5	+8.60%	+1.12%	off — batch 1	+10.08%
2	128	9	+2.99%	−0.05%	+0.10%	+3.05%
4	64	17	+0.63%	−0.35%	+0.41%	+1.15%
8	32	34	+0.07%	−2.16%	+0.51%	−0.36%
16	16	67	+0.12%	+5.47%	−0.68%	+5.23%
24	10	107	+0.08%	+10.54%	off — 10 % 8 ≠ 0	+10.45%
48	5	213	+0.02%	+10.54%	off — 5 % 8 ≠ 0	+10.30%

Two levers, active at opposite ends for unrelated reasons. Commit 3 is worth +8.6% at one request and
is spent by 4, since the split count halves with every doubling of concurrency and there is no merge
cost to attack below ~128 splits. Commit 1 is flat-to-negative low and +10.5% high, and nothing helps
at concurrency 4–8. The two ~10% ends differ 8× in absolute terms: 58.1 → 52.8 µs at concurrency 1
versus 436.8 → 396.0 µs at 48. Commit 1 at concurrency 24 / 48 across three independent runs:
+10.54 / +10.54%, +10.07 / +10.39%, +11.33 / +10.94%; concurrency 8 regressed in all three at
−2.16%, −2.56%, −2.07%.

2. Mechanism: the gain tracks loop depth, not concurrency

Split count forced at fixed batch, commit 1 only — the commit message's claim tested rather than
restated:

splits	tok/split	trips	batch 8	batch 48
2	34,000	532	+12.03%	+6.11%
4	17,000	266	+13.97%	+9.13%
8	8,500	133	+15.95%	+9.54%
16	4,250	67	+6.99%	+11.24%
32	2,125	34	−1.68%	+7.46%
64	1,062	17	−1.38%	+4.89%
128	531	9	+0.64%	+5.05%
256	265	5	+1.79%	+4.54%

The gain appears where the theory says it should and collapses where it says it should, and it is not a
function of batch: at 34 trips batch 8 loses 1.7% while batch 48 gains 7.5%. Batch 8 at 32 splits is
exactly one wave — 8 × 32 = 256 workgroups on 256 CUs — so there is neither trip depth nor a second wave
to hide the pipeline ramp behind, and the production policy happens to place concurrency 8 in that
pocket; the forced sweep reproduces −1.68% where the policy run measures −2.16%. I would rather a
reviewer judge this table than the headline, since a single point estimate could be a scheduling
accident while a dose-response curve with a sign flip where the mechanism predicts one is harder to
explain away.

3. Why commit 2 is in this PR

Commit 2 is here to make commit 1's concurrency-8 regression tolerable, and paired with it that is what
happens: commit 1 alone measures −2.56% there (per-round −2.1 −2.8 −2.7 −1.7 −1.6), with commit 2 −0.33%
(+0.1 −1.0 −0.4 −0.1 +0.2), so commit 2 recovers 2.2 of the 2.6 points. All three together read −0.46%.
The trade is 0.8 points at concurrency 16 (+5.7% → +4.9%), which I took because a regression at a
production point is worth more than an equal gain elsewhere — but it is a judgement call.

AI assistance

This came out of GEAK, AMD's agentic GPU-kernel optimization
framework, running an end-to-end pass against a live Kimi-K3 TP8 server on 8× MI355X. The
measurements above were produced by and reviewed manually.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4758 --add-label <label>