[Triton/Gluon] [Kernel][Perf] split-K long-context decode for shuffled fp8 SWA path by reger-men · Pull Request #3962 · ROCm/aiter

reger-men · 2026-06-26T17:27:01Z

Summary

Adds aiter/ops/triton/attention/pa_decode_shuffle_perf.py, a flash-decoding
(split-K) variant of the sliding-window shuffled-fp8 paged decode. The one-pass
kernel (companion PR, pa_decode_shuffle.py) runs one program per (sequence, kv
head) and walks the whole windowed key range serially, which underfills the device
at long context. This splits that range into partition_size chunks across the grid
and reduces the per-partition softmax state in a second pass. It is gated to engage
only where it actually helps and never regresses the short-context path.

Motivation

At long context the one-pass per-program key walk gets long while the grid
(num_seqs * num_kv_heads) is too small to fill the GPU, so most CUs sit idle. The
asm decode path solves this with page-split kernels, but the asm path can't take a
window (the gfx942/950 pa/*.co kernels are prebuilt code objects with no .s
source or PA asm generator in the tree, and the registry has only a causal-mask flag,
no window). So the realistic lever for the windowed shuffled-cache path is the triton
kernel, parallelising context the same way the asm page-split kernels do.

Changes

New pa_decode_shuffle_perf.py with
paged_attention_decode_shuffle_swa_perf(...): a partition kernel + a reduce
kernel, plus a one-pass fallback that runs the same algorithm as the companion
one-pass kernel.
The split is bounded to [seq_len - window, seq_len), so a windowed layer only
partitions the keys it actually attends. max_seq_len is taken as a host int
(vLLM already tracks it), so no per-step .item() device sync.
fp32 scratch is reused across decode steps (one buffer per device/shape) so the
split path doesn't allocate every step.
Gating (_pick_n_part + _arch_allows_split): the auto path splits only when the
grid underfills or the serial walk is long (>= a threshold of partitions), AND only
on arches where the split is measured to win end to end (gfx950 by default;
override with AITER_PA_SHUFFLE_SPLIT_ARCH). Short context, a full grid, or a
non-listed arch stays one-pass, so it is identical to the one-pass kernel there.
force_split overrides the gating and is used by the tests to exercise both code
paths on any arch.

Performance

Decode tok/s, clock-locked (perf_determinism), Mistral-family SWA model
(sliding_window=4096), --max-model-len 8192, fp8 shuffled KV, ROCM_AITER_FA.
Before = one-pass kernel (companion PR), After = this split-K kernel.

MI355X (gfx950), ROCm 7.13 (where the split is enabled by default):

Configuration	Before (one-pass)	After (split-K)	Improvement
batch 32, ctx ~590 (< window, stays one-pass)	3573	3559	parity (no split)
batch 8, ctx ~6700 (> window, n_part ~16)	382-385	762-820	~2.0-2.1x

At long context split-K is ~2.0-2.1x over one-pass and recovers decode to ~84-90% of
the bf16 baseline (909-913 tok/s); at short context it stays one-pass (parity). The
numbers above are from the 0.19.1 image with the 0.19.1 dispatch.

Re-confirmed on vLLM built from source at main d980a3cc
(0.23.1rc1.dev480+gd980a3cc6, same node, clock-locked), routing through the corrected
main dispatch, zephyr-7b-beta: short (batch 32, ctx ~590) bf16 3679 / one-pass 3470 /
split-K 3452 (parity, stays one-pass); long (batch 8, ctx ~6700, n_part ~16) bf16
~743 / one-pass ~387 / split-K 731 and 786 on two warm runs (~1.9-2.0x over one-pass,
recovering to ~parity with bf16). The first windowed split-K call pays a one-time
triton JIT for the partition grid (544 tok/s cold vs ~759 warm); one-pass and bf16
are stable to ~1%.

MI300X (gfx942), ROCm 7.2.1 (kept one-pass by the arch gate; numbers are with the
gate forced off, to show why):

Configuration	Before (one-pass)	After (forced split-K)	Improvement
batch 32, ctx <window	~949	~1005	~1.06x (within run-to-run spread)
batch 8, ctx ~6700 (> window)	~244	~228	~0.93x (slightly behind)

Isolated kernel microbenchmark on gfx942 (clock-locked) is a real win: parity at
short window, ~1.3-1.4x at win1024/ctx2048, ~4.6-6.4x at win4096/ctx8192 and full
causal. But on a dense 7B at these shapes the decode-attention kernel is a small slice
of the step next to the QKV/O projections and the MLP, so the isolated win does not
surface end to end and the extra reduce launch makes it parity-to-slightly-behind.
That is why the auto path keeps gfx942 on one-pass. gfx950's bf16 decode is much
faster (909 vs ~276 tok/s at the same long shape), so the underfilled one-pass kernel
is a far larger fraction of the step there and the split has real headroom.

Testing

Numerics harness against the independent fp8 reference, both fp8 dtypes
(e4m3fn, e4m3fnuz), both code paths (force_split on and off), full case matrix
(head 128, GQA 32/8 and MQA, block 16/32, windowed and full, per-tensor and per-token)
plus long-context cases that exercise 4-16 partitions. Split-K reproduces the one-pass
result within the fp8 tolerance, so the reduce is not a correctness regression. This
is the same case matrix and ALL_PASS result as the companion one-pass kernel PR.

MI300X (gfx942): ALL_PASS both dtypes, max abs <= 1.81e-2 vs a 0.18-0.35 quant floor.
MI355X (gfx950): ALL_PASS both dtypes, max abs <= 1.55e-2 vs the same floor.
Tested on MI300X (gfx942)
Tested on MI355X (gfx950)
Performance benchmarks run (tables above)

Documentation

Layout, split scheme, and gating are documented in the module docstrings.

github-actions · 2026-06-26T17:27:13Z

🏷️ 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 3962 --add-label <label>

reger-men · 2026-08-06T17:07:04Z

Stacked on #3959. Only failure is one shard, MI35X shard 6 (moe_2stage fp4 allclose), a different path from the split-K decode this touches, everything else green. @vgokhale mind kicking off a re-run so we can confirm it's just flaky? I'll land #3959 first, then this.

reger-men · 2026-08-13T16:01:36Z

Following up. This is stacked on #3959, and the only red is Standard Tests MI35X shard 6, the moe_2stage fp4 allclose test, a different suite from the split-K decode path this PR changes. @vgokhale could you re-run CI here? It should go green once #3959's wheel gate is fixed.

reger-men · 2026-09-11T13:51:20Z

@vgokhale same as #3959, blocked on the prepare_triton_wheel gate rather than the code. Can you kick the CI or tell me what to fix?

reger-men · 2026-09-18T07:44:43Z

@vgokhale same for this one, it depends on #3959 and is blocked on the MI35X shard-6 job, not the change itself. Can you kick the CI?

reger-men · 2026-10-02T16:18:11Z

@cagrikymk @azaidy same ask here. This is the split-K long-context variant of #3959 and only adds pa_decode_shuffle_perf.py. The one red is Standard Tests MI35X shard 6, the moe_2stage fp4 allclose test, which this file doesnt touch. I can rebase both if you'd rather have fresh CI.

reger-men · 2026-10-12T15:39:42Z

@azaidy @vgokhale same situation as #3959, this adds pa_decode_shuffle_perf.py (the split-K long-context variant) and has no conflicts. The only red is the June MI35X shard 6 run, the moe_2stage fp4 allclose test, which is outside this file. I can rebase both together once someone has time to review.

The one-pass sliding-window shuffled-fp8 decode runs one program per (sequence, kv head) and walks the whole windowed key range serially, which underfills the device at long context while the grid is too small to fill it. The asm page-split path can't take a window, so the lever for the windowed shuffled path is the triton kernel. Add paged_attention_decode_shuffle_swa_perf, a flash-decoding (split-K) variant: split the windowed range into partition_size chunks across the grid, then reduce the per-partition softmax state. The split is bounded to [seq_len - window, seq_len), max_seq_len is a host int (no per-step device sync), and fp32 scratch is reused across steps. Gating keeps it one-pass for short context or a full grid, and only splits on arches where it is measured to win end to end (gfx950 by default, overridable via AITER_PA_SHUFFLE_SPLIT_ARCH); force_split overrides the gate for testing both paths on any arch. On gfx950 long context this is ~2.0-2.1x over one-pass, recovering decode to ~84-90% of bf16; at short context it stays one-pass (parity). On gfx942 the isolated kernel is faster (up to ~5x microbench) but decode attention is a small slice of a dense-7B step, so it is parity-to-slightly-behind end to end; the arch gate keeps gfx942 on one-pass. Numerics ALL_PASS on both arches and both fp8 dtypes, one-pass and forced split, within the fp8 quantization floor. Signed-off-by: adlashab <adlashab@amd.com>

reger-men · 2026-10-15T17:43:24Z

@azaidy @vgokhale same for the split-K variant: rebased onto e8bdaf8, new head 1bff7ae. Main's ruff config flagged two lint points in the file (import order in _arch_allows_split and max_seq_len: int | None), I fixed those and nothing else changed.

Numerics on MI355X today: 16 of 16 cases pass for e4m3fn and e4m3fnuz, one-pass and forced split. In git it doesnt depend on #3959, each PR only adds its own file, so they can go in either order.

Copilot

🟡 Changes recommended

Partition coverage and workspace defects can produce incorrect outputs or out-of-bounds writes.

8 open findings

What changed in this PR

Adds a Triton split-K decode path for sliding-window attention over shuffled fp8 KV caches, targeting long-context GPU utilization.

Changes:

Adds partition and softmax-reduction kernels with a one-pass fallback.
Selects split execution by context length, grid size and architecture.
Adds reusable fp32 workspace and caller-controlled dispatch overrides.

File	Description
`aiter/ops/triton/attention/pa_decode_shuffle_perf.py`	Adds decode kernels, dispatch gates and workspace management.

🧠 Review effort: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+        p = tl.exp(scores - m_new[:, None])
+        alpha = tl.exp(m_i - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None] + tl.dot(p.to(tl.float32), v)


+    )
+    acc_part = tl.load(
+        acc_ws_ptr + acc_off,
+        mask=gmask[:, None, None] & pmask[None, :, None],


+    need_m = (num_seqs, num_kv_heads, grp, n_part)
+    need_a = (num_seqs, num_kv_heads, grp, n_part, head_pow2)
+    if workspace is None:
+        workspace = _WS_CACHE.setdefault(dev.index, {})


+    if workspace.get("shape_m") != need_m:
+        m_ws = torch.empty(need_m, dtype=torch.float32, device=dev)
+        l_ws = torch.empty(need_m, dtype=torch.float32, device=dev)
+        acc_ws = torch.empty(need_a, dtype=torch.float32, device=dev)
+        workspace.update(shape_m=need_m, m=m_ws, l=l_ws, acc=acc_ws)


+        GRP=grp,
+        GRP_POW2=grp_pow2,
+        SCALE_SCALAR=scale_scalar,
+        PART=partition_size,


+    idx = device.index if device.index is not None else torch.cuda.current_device()
+    n = _N_CU.get(idx)
+    if n is None:
+        n = torch.cuda.get_device_properties(idx).multi_processor_count


+    if workspace.get("shape_m") != need_m:
+        m_ws = torch.empty(need_m, dtype=torch.float32, device=dev)
+        l_ws = torch.empty(need_m, dtype=torch.float32, device=dev)
+        acc_ws = torch.empty(need_a, dtype=torch.float32, device=dev)
+        workspace.update(shape_m=need_m, m=m_ws, l=l_ws, acc=acc_ws)


+    return allow
+
+
+def paged_attention_decode_shuffle_swa_perf(


… and NaN paths, add test and bench When the grid underfills the device the auto path caps n_part, but every partition still read partition_size keys, so only n_part * partition_size keys were used (eff 4096, base grid 128, 256 CUs: 4 partitions read 1024 keys). The partition size now grows so n_part partitions cover eff. It is a runtime arg, so a growing context does not compile new kernels. The fp32 scratch is per (device, stream) and only grows. It is checked by element count, so a head 128 call after a head 64 call no longer writes past the acc buffer, and a caller dict that holds tensors of another device gets new ones. The partition dim is rounded up to a power of two and n_part is a runtime arg of the reduce, so changing partition counts reuse the same scratch and the same reduce kernel. A replaced buffer stays referenced because a captured cudagraph can still point at it. NaN paths: V rows of unused slots are zeroed in the edge blocks (0 * NaN in p @ v), the reduce skips partitions with l = 0 because their acc is never written, and rows with seq_len 0 return zeros. V scales use their own strides. CU count comes from get_num_sms() (it honors CU_NUM) and the arch from arch_info.get_arch(). Add op_tests/triton_tests/attention/test_pa_decode_shuffle_perf.py and op_tests/op_benchmarks/triton/bench_pa_decode_shuffle_perf.py. Signed-off-by: adlashab <adlashab@amd.com>

reger-men · 2026-10-20T09:13:02Z

I went through the 8 review bot points on MI355X with a repro against the old head 1bff7ae (KV cache, per-token scales and scratch filled with NaN wherever the kernel shouldnt read). All 8 were real, the fix is in 476b9e6.

Capped partitions (L557): real, and the worst one. With eff 4096, partition 256, base grid 128 on 256 CUs the auto path picked 4 partitions and read 1024 of the 4096 keys, max abs err 2.35 vs an fp32 reference (2 seqs x 20000 keys: 0.62). The partition size now grows so n_part partitions cover eff (block aligned, at least partition_size). Its a runtime arg, so a growing context doesnt compile new kernels. Same cases now 0.004 and 0.001.
Scratch shared across streams (L506): real, two split calls in flight on two streams with the same shape gave max abs err 2.09 on one of them. Scratch is per (device, stream) now, calls on one stream are ordered by the stream and can reuse it.
Undersized acc and other device (L511): real. head 128 after head 64 with the same shape_m reused a 65536-float acc where 131072 are needed, so it wrote past the buffer (err 5.5). Buffers are checked by element count and only grow, and a dict that holds tensors of another device gets new ones. A replaced buffer stays referenced because a captured cudagraph can still point at it.
Shape churn (L511): real, 7 allocations in 7 calls when the partition count alternates 12/16/20. The partition dim is rounded up to a power of two and n_part is a runtime arg of the reduce, now 2 allocations for the same 7 calls and no new reduce kernel per partition count.
NaN in unused V slots (L118): real, seq_len 17 and 1 on 16-slot pages gave NaN rows on both paths. V rows of unused slots are zeroed. Only the window edge blocks have unused slots, so interior blocks skip the select.
Empty partitions (L309): real, NaN scratch came back through 0 * NaN, and seq_len 0 rows were 0 / 0. The reduce masks partitions with l = 0, rows with no keys return zeros.
CU count helper (L364): it uses get_num_sms() now, which honors CU_NUM (CU_NUM=80 gives 80), and arch_info.get_arch() for the arch gate. I didnt extend get_num_sms for an explicit device, the other aiter Triton ops size their grids from it the same way.
Test and bench (L399): added op_tests/triton_tests/attention/test_pa_decode_shuffle_perf.py, 48 cases: both fp8 types, scalar and per-token scales (K and V scales with different strides), 16 and 32 slot pages, one-pass / split / auto, the capped case, seq_len 0 rows, NaN scratch, workspace reuse across head sizes, two streams. And op_tests/op_benchmarks/triton/bench_pa_decode_shuffle_perf.py. 48 passed on MI355X, black and ruff clean.

V scales also use their own strides now.

Perf vs the old head, per call in a cudagraph, MI355X at perf_determinism 1700 MHz, 32 q / 8 kv heads, head 128, 16-slot pages, 2 runs each in ABBA order: where both read the same keys its within -5% to +3% (bs 1 ctx 8192 window 4096: 43.8 -> 45.1 us, bs 32: 327.1 -> 329.9 us, bs 64: 637.9 -> 638.8 us, bs 32 ctx 2048 window 1024: 130.4 -> 124.1 us). Where the old code was capped it read fewer keys, so the new one is slower but correct (bs 8: 59.5 -> 108.0 us, bs 16: 55.3 -> 195.8 us). Thats still 4.5x and 2.5x faster than the one-pass kernel at the same shapes (482.9 and 488.3 us).

github-actions · 2026-11-05T01:33:46Z

No activity for 15 days, so this is now labelled stale. A push or a comment clears it; keep-open exempts it. Nothing is closed automatically.

reger-men requested a review from a team June 26, 2026 17:27

zufayu requested a review from amd-ruitang3 July 13, 2026 06:51

zufayu removed the request for review from amd-ruitang3 July 21, 2026 02:40

zufayu requested a review from vgokhale August 5, 2026 02:49

reger-men mentioned this pull request Aug 6, 2026

[Triton/Gluon] [Kernel] sliding-window decode over shuffled fp8 paged KV #3959

Open

2 of 3 tasks

reger-men force-pushed the swa-shuffle-fp8-paged-decode-splitk branch from a3b19d7 to 1bff7ae Compare October 15, 2026 17:42

github-actions Bot changed the title ~~[Kernel][Perf] split-K long-context decode for shuffled fp8 SWA path~~ [Triton/Gluon] [Kernel][Perf] split-K long-context decode for shuffled fp8 SWA path Oct 15, 2026

github-actions Bot added the Triton/Gluon label Oct 15, 2026

Boss2002n requested a balanced review from Copilot October 15, 2026 21:33

Copilot started reviewing on behalf of Boss2002n October 15, 2026 21:34

View session

Copilot AI reviewed Oct 15, 2026

View reviewed changes

azaidy assigned vgokhale Oct 19, 2026

reger-men force-pushed the swa-shuffle-fp8-paged-decode-splitk branch from b138774 to 476b9e6 Compare October 20, 2026 09:12

github-actions Bot added the stale The PR hasn't been updated in 2 weeks label Nov 5, 2026

Navigation Menu

Repository navigation

[Triton/Gluon] [Kernel][Perf] split-K long-context decode for shuffled fp8 SWA path - #3962

[Triton/Gluon] [Kernel][Perf] split-K long-context decode for shuffled fp8 SWA path#3962
reger-men wants to merge 2 commits into
ROCm:mainfrom
reger-men:swa-shuffle-fp8-paged-decode-splitk

reger-men commented Jun 26, 2026

Uh oh!

github-actions Bot commented Jun 26, 2026

Uh oh!

reger-men commented Aug 6, 2026

Uh oh!

reger-men commented Aug 13, 2026

Uh oh!

reger-men commented Sep 11, 2026

Uh oh!

reger-men commented Sep 18, 2026

Uh oh!

reger-men commented Oct 2, 2026

Uh oh!

reger-men commented Oct 12, 2026

Uh oh!

reger-men commented Oct 15, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

reger-men commented Oct 20, 2026

Uh oh!

github-actions Bot commented Nov 5, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

reger-men commented Jun 26, 2026

Summary

Motivation

Changes

Performance

Testing

Documentation

Uh oh!

github-actions Bot commented Jun 26, 2026

🏷️ CI Guide

Uh oh!

reger-men commented Aug 6, 2026

Uh oh!

reger-men commented Aug 13, 2026

Uh oh!

reger-men commented Sep 11, 2026

Uh oh!

reger-men commented Sep 18, 2026

Uh oh!

reger-men commented Oct 2, 2026

Uh oh!

reger-men commented Oct 12, 2026

Uh oh!

reger-men commented Oct 15, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

reger-men commented Oct 20, 2026

Uh oh!

github-actions Bot commented Nov 5, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants