[Gluon][MLA] Split-major grid and blocked stage-2 reduce for the small-nhead decode regime for kimi-k3

[amd-ethany]

On top of #4507, this PR is worth 6.22% of inter-token latency at batch 1 and decays to nothing by
batch 48, because what it shortens is a serial chain whose depth is the split count — and the
split count is num_CUs / batch. A small, split-count-dependent win on the stage that becomes the
next bottleneck, not a second large one.

Motivation

Once split parallelism is restored by #4507, the stage-2 reduce becomes the next constraint in the small-batch
corner. _mla_softmax_reducev_kernel walks the splits serially, each step's rescale depending on the
previous e_max, on a (batch, nhead, qlen) grid that is only 12 workgroups at batch=1, nhead=12.
Its cost therefore grows with the split count, reaching a third of the batch-1 decode budget at 256
splits — which is what made raising NUM_KV_SPLITS look like a bad trade before this.

Stage-1 has two smaller issues in the same regime. Its grid is (batch, split, head_block*q_pos) and
Triton linearizes axis 0 fastest, so one split's workgroups scatter across all 8 XCDs, each pulling
its own copy of KV pages that a prefix-cached decode batch shares ~94% of. Separately, the loop body
opened with a full wait_group(0) drain, serializing the page-number prefetch behind the KV stream.

Technical Details

1. Split-major stage-1 grid. Swap grid axes 0 and 1 for the bh16 regimes, so the linear
workgroup id is split + NUM_KV_SPLITS*batch + .... All requests' workgroups for a given split then
land on the same XCD and share its L2, and under prefix caching they are reading the same rows. The
swap is taken only when NUM_KV_SPLITS % NUM_XCD == 0, which is exactly the condition under which
the id stays congruent mod NUM_XCD for every batch; otherwise the ordering is left alone. It is a
relabeling of which workgroup computes which (batch, split) item — same KV range per program, same
reduction order, bit-identical output.

2. Blocked, head-dim-tiled stage-2 reduce. Two re-associations of the same online-softmax merge:
merge BLOCK_S splits per iteration as one [BLOCK_S, BLOCK_D] tile (shortens the serial chain by
BLOCK_S, gives the memory system BLOCK_S independent loads), and tile the 512-wide head dim
across D_SPLIT programs packed into grid axis 2 alongside q_pos. Each head-dim slice re-derives
the same scalar lse statistics from the same Mid_lse vector — a few hundred duplicated bytes — and
owns a disjoint slice of o, so slices never communicate. BLOCK_S=1, D_SPLIT=1 is exactly the
kernel as it stands today; the launcher picks BLOCK_S = min(256, next_pow2(NUM_KV_SPLITS)) and
D_SPLIT = 8, backing D_SPLIT off when the head dim does not divide evenly or the slice would be
under 16 elements.

Two invariants carry over. The empty-split guard keeps its semantics: the loop starts on the block
boundary at or below LOOP_START and masks the uninitialised slots off inside the block, so
stage-1's untouched slots are never read. The MTP NaN guard is applied per row of the tile, so a
fully causal-masked split (NaN logits with lse=-inf) contributes exactly zero.

3. Two partial waits in the stage-1 loop. At the top of the loop body the outstanding groups are
exactly the four the previous trip committed — page, kv0, kpe, kv1 — and only the page copy is
needed there. Retire just the oldest (wait_group(3 if HAS_PE else 2)) and move the rest of the wait
down to the k_c local load, which retires exactly the three groups it depends on. Two KV blocks are
then resident at once at no extra LDS. This is the same granularity the two epilogues already use.
The prologue leaves the same set outstanding as a steady-state trip, so trip 0 accounts identically,
and the set outstanding at loop exit is unchanged, so epilogue 1 is unaffected.

Test Plan

End-to-end, concurrency sweep (primary). 8x MI355X (gfx950), Kimi-K3 TP8, bf16 KV, 12 heads/rank.
aiperf, mean measured ISL 68,088.7 / OSL 350.0, concurrency 1/8/16/24/48 at 5/40/80/120/240 requests,
4 profile runs per point with a 15 s cooldown, seed 42, prefix caching on,
cudagraph_mode=FULL_AND_PIECEWISE. Three legs — stock, #4507, #4507 + this PR — each a fresh server
on the same node in the same container, each confirmed by server-log grep to have actually loaded (or
not loaded) the patched module.

End-to-end, long context (secondary). Same hardware, 327,600-token context, batch 8. vLLM serving
benchmark: ISL 18,264 plus a 308,736-token shared random prefix, OSL 1,200, concurrency 8, 16 prompts,
seed 42. Reference leg is #4507 alone; candidate is #4507 plus this change. Both legs serialized on a
GPU lock, both to completion, cudagraph captured on both (est. graph memory 1.22 GiB each), engagement
8/8 workers on both.

Kernel-level. Frozen, implementation-independent fp32 online-softmax reference at tol 2e-2 under
capture-once / replay-many, including ragged batches
(seq_lens = [1, 63, 64, 65, 4096, 4097, 65536, 327600]) and prefix-boundary batches
([308736, 308737, 312831, 312832, 312833, 320000, 327599, 327600]).

aiter's own MLA tests. Both configs pass on the rebased branch, with identical verdicts to
unpatched main run 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

All three changes are exercised, 1/1 verdict on both legs. The
second config lands in the bh64 regime, the regression surface for the wait change since the
stage-1 loop is shared by every regime: 10/10 verdicts on both legs.

Isolation. The stage-2 knob flip described above, to separate the two levers.

Accuracy. GSM8K, 200 problems, on both legs of the same paired session.

Style. ruff check (0.16.0, the pinned version) and black --check pass on the changed file.

Test Result

1. Concurrency sweep at 68k context — this PR's increment over #4507

Reference leg is #4507 alone, candidate is #4507 plus this PR, 4 trials each, ratios formed within a
trial index so prefix-cache history is matched on both sides. splits is what #4507 selects
(num_CUs / batch) and is also the depth of the chain this PR shortens.

conc	splits	split-major grid	#4507 ITL P50	+ this PR	delta	per-trial deltas
1	256	off — batch 1	26.26 ms	24.73 ms	+6.22%	+6.4, +6.4, +5.3, +6.7
8	32	on	39.73 ms	38.75 ms	+2.54%	+3.2, +4.9, +1.9, +0.4
16	16	on	54.84 ms	54.08 ms	+1.40%	+2.0, +1.7, +2.1, −0.2
24	10	off — 10 % 8 ≠ 0	65.53 ms	64.91 ms	+0.96%	−0.4, +4.0, +1.1, −0.8
48	5	off — 5 % 8 ≠ 0	101.11 ms	101.17 ms	−0.06%	−1.1, +2.0, −0.9, −0.2

The shape is what the mechanism predicts: the win tracks the split count, and by 5 splits there is no
chain left to shorten. Only the top two rows have every trial on the same side of zero, and at 32
splits one trial is just +0.4%, so I would defend +6.22% at 256 splits and treat everything from 16
splits down as within noise.

Work parity: completed requests, mean ISL and mean OSL identical across all three legs at every
point, zero errors. The candidate range is disjoint from stock at every concurrency.

The grid order is gated off at batch 1, so it was checked on its own rather than through the sweep:
flipping only that axis at batch 8, 32 splits and 327,600 tokens moves the decode kernel 0.4568 →
0.3655 ms, with L2 hit 0.71% → 82.18% and HBM traffic 3.033 → 0.544 GB (rocprofv3, 5 dispatches).
L2→CU request volume is identical either way; only the order changes. The analytic ceiling for that
page table is 82.46% hit, leaving 0.28 pp on the axis.

2. Long context, single point — 327,600 tokens

The original measurement, at 4.8x the context and 2 repeats per leg:

	ref (#4507 only)	candidate (+ this PR)	delta
End-to-end throughput, median	155.47 tok/s	172.35 tok/s	+10.86%
End-to-end, run 1 (cold prefix)	114.04 tok/s	125.95 tok/s	+10.4%
End-to-end, run 2 (warm prefix)	196.90 tok/s	218.75 tok/s	+11.1%
Median TPOT (cold / warm)	49.07 / 39.32 ms	40.64 / 34.46 ms	—
Stage-2 reduce, per call	87.8 µs	4.0 µs	22x
Isolated geomean vs unpatched	—	39.15x	—
GSM8K (200 problems)	0.905	0.910	within tolerance
KV footprint	71.32 GiB / 2,731,470 tok	identical	parity

The larger number here is consistent with the sweep rather than in tension with it: this point runs
at batch 8 over a 327,600-token context, so stage-2 is a far larger share of a much longer decode
than it is at 68k. The binary it was measured on also carried a stage-1 tail-mask peel that is held
back from this PR, since MTP on main makes that mask live for QLEN > 1, so some of the +10.86%
belongs to the peel. It also rests on 2 repeats against the sweep's 4, which is why the sweep leads.

The reduction order changes relative to the reference leg, so bitwise equality is not expected;
GSM8K moved 0.905 → 0.910, inside the 0.05 tolerance.

AI assistance

This came out of GEAK, AMD's agentic GPU-kernel optimization
framework, running an end-to-end optimization pass against a live Kimi-K3 TP8 server on 8x MI355X.

Everything in this PR was verified manually before submission.

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: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 4509 --add-label <label>

[vgokhale]

Can you benchmark higher concurrencies? Can you benchmark the other models to confirm they don't regress?

[Dewei-Wang-sh]

hi can you check #4555, the stage2 reduce is optimized and deprecated kv_seq_len for split.

[amd-ethany]

This PR is out-dated, plz check the follow-up PR in #4758