[fsx950223]

Summary

Brings three of #7709's gfx942 pa_decode schedule policies to gfx950, each re-tuned on MI355X:

• K/V streaming for long head rows;
• MTP query splitting over idle CUs;
• split-major launches on sparse grids.

gfx942 and gfx1250 kernels are bit-for-bit unchanged.

On MI355X, over all 817 cases of the unit-test-geometry matrix described below: geomean 1.19x against the base branch (unit-test cases 1.22x, grid cases 1.19x). The 333 cases whose kernels change run 1.54x faster, and none is slower than 0.98x.

This PR is stacked.

• Its base, fsx950223/pa-stack-7700-7701-7709, is the merge of [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700, [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 and [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709 (see "Base branch" below). Please review only the three commits here.
• aiter's CI, pre-checks and self-review bot run only on PRs into main, so this PR gets them once it is retargeted to main after its bases land.
• Draft: the base is behind. Everything here was built and measured on [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 59d523033c, [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 02dc5d11b8 and [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709 65eed8399b. Since then [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 has moved to bdabd08dad: it extends four gfx950 paths to gfx942, adopts a named cache key, and keeps main's Q absmax above D256 on gfx942. [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 has moved to a8b9f032b9. The stack will be rebuilt and re-measured once those heads settle.
• [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700's new gfx942 paths (one-partition direct output, the min-max planner, dual-group CTAs, packed Q absmax) overlap [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's gfx942 schedule policies. Whichever lands second has to resolve that interplay, and it needs gfx942 re-measurement.

Motivation

#7709 found four schedule problems on gfx942. To see which also hold on gfx950, I opened each of #7709's gfx942-only gates on gfx950 and measured it on its own, against main and against #7700+#7701. Three hold:

1. Spills at long head rows. The tile holds a whole next-tile K prefetch plus the tile's V gather in VGPRs. Once a lane's K row reaches 1 KiB (BF16 D512+, FP8 D1024), gfx950 spills as gfx942 does. Single and paired MTP M-tiles with 384/256-byte rows drop to one wave per SIMD on busy grids. Streaming runs these 1.8-7.7x faster.
2. Idle CUs under MTP. With few sequences, one CTA's serial latency over all query positions is the decode time while most CUs idle.
3. Plan-order dispatch on small grids. The tile grid dispatches x first, so later query splits and KV heads queue behind every empty plan slot.

Three gfx942 policies did not carry over and stay gfx942-only:

• waves_per_eu=2 on busy grids: [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700's gfx950 MTP plan holds capacity at one or two slots per CU, so it changes no case.
• The busy BF16 D128 query split: 0.6-0.9x on gfx950.
• V prefetch, early Q and early page ids: 1.01-1.03x overall but with cases down to 0.93x. [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 already has its own early-Q and page-lookahead variants for its D128 shapes.

Changes

All changes are in aiter/ops/flydsl/kernels/pa_decode/{traits,op_epilog}.py and pa_decode_kernel.py, one commit each:

1. 372b2bc843 K/V streaming on gfx950 (STREAM_KV).
   • It selects on gfx950 under [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's row-size rules. gfx950's fused BF16 MTP and head-256 pipes keep their own schedules, and streamed tiles skip the dual-group CTA.
   • The busy-grid flags it reads are now arch-neutral and renamed: gfx942_busy -> busy_grid, gfx942_two_per_cu -> two_per_cu. All their other consumers stay gfx942-only.
2. 50e1d38283 MTP query split over idle gfx950 CUs, using [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's rules with two gfx950 differences:
   • Per-token FP8 keeps its fused MTP4 tiles. Splitting them ran 0.75x.
   • The FP8 split also requires at most eight KV tiles per CU. This uses tile_demand, the sequence x head x tile count before the plan's capacity cap. gfx950 MTP plans fix capacity regardless of load, so a capacity bound no longer means idle CUs (B64/32K per-tensor ran 0.87x without the cap).
   • The busy BF16 D128 split stays gfx942-only.
3. cbffdd201b Split-major launches on sparse gfx950 grids.
   • gfx950 uses split-major only while at most a quarter of the CUs hold work. On fuller grids slot-major ran up to 1.25x faster.
   • It applies to FP8, and to BF16 up to D128 (LDS-forced BF16 D192/D256 splits ran 0.91-0.93x). MXFP4 keeps slot-major.
   • Two gfx950-only paths assumed the slot-major axes and now follow the grid:
     • The BF16 D128 planned output's slot count read grid_dim.x. With a split-major grid that is heads x splits, so partials were stored at wrong offsets.
     • The one-partition plan's empty-row zero fill read block_id("y") as the KV head.

Not changed:

• the API, plan_pa_decode and its budgets;
• the reducer;
• softmax math;
• MXFP4 paths;
• every gfx942 and gfx1250 binary.

Base branch

fsx950223/pa-stack-7700-7701-7709 is built from:

• [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 at 59d523033c;
• merged with [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 at 02dc5d11b8;
• merged with [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709 at 8274dfbc42;
• merged with [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's race fix at 65eed8399b.

The #7709 merge resolves 19 conflict hunks in traits.py, pipeline.py, op_lds.py and pa_decode_kernel.py as a union of both sides: #7709's policies stay gfx942-only and #7700's paths gfx950-only. A compile-only replay of every matrix case gives:

• gfx950 and gfx1250: identical to [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700+[FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701.
• gfx942: identical to [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709 except 18 padded-row specializations. There the merged Q-row clamp takes [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700's form, which selects the same rows. On MI308X those geometries measure 1.0003x (191 cases).

While porting streaming, I found a first-tile LDS race in #7709's streamed path. It is fixed in #7709 itself (65eed8399b), not here.

Performance (if applicable)

How these numbers were measured:

• Hardware: MI355X (gfx950, 256 CUs), ROCm 7.2.4, torch 2.10, FlyDSL 0.5.0. Base = 03894b767e (the stack branch), after = cbffdd201b.
• Matrix (817 cases): every geometry of op_tests/test_flydsl_pa_decode.py that runs on gfx950 at its unit-test lengths (115 cases), plus each geometry at batch {1, 16, 64} x context {1K, 8K, 32K} with the default plan (702 cases).
• Method: paired in-process A/B, both trees loaded in one process. CUDA-graph replays of 10 calls, 10 replays x 20 ABBA pairs per case; median ratio.
• Noise: the 484 cases whose kernels are binary-identical read geomean 1.0000, 1st-99th percentile 0.984-1.018.

Hardware	Workload/configuration	Before	After	Change
MI355X	Full matrix, 817 cases (geomean)	-	-	1.19x
MI355X	The 333 changed cases (geomean)	-	-	1.54x (min 0.984x)
MI355X	BF16 D1024 MTP4 GQA16, B16 ctx1K	202.2 us	26.2 us	7.72x
MI355X	BF16 D1024 MTP4 GQA16, B16 ctx8K	1662.0 us	389.5 us	4.27x
MI355X	BF16 D512 MTP4 GQA16, B16 ctx8K	326.5 us	91.5 us	3.57x
MI355X	FP8 D1024 GQA4, B16 ctx8K	116.7 us	49.6 us	2.35x
MI355X	FP8 MTP4 GQA16 page128, UT mtp4-g16-page128-direct	20.2 us	9.3 us	2.17x
MI355X	BF16 D64 MTP4 GQA16 page128, UT ...-direct-d64	20.3 us	12.4 us	1.64x
MI355X	FP8 per-token MTP4 window 257, B1 ctx8K	9.2 us	6.4 us	1.45x
MI355X	Worst changed case: FP8 D256 MTP4 2 heads page128, B64 ctx32K	362.6 us	368.5 us	0.98x

Per commit, each against its parent, over the cases it changes:

Hardware	Commit / workload	Change
MI355X	1: K/V streaming, 96 cases	geomean 2.67x (BF16 3.88x, FP8 1.83-2.07x), min 0.984x
MI355X	2: MTP query split, 164 cases	geomean 1.19x (UT 1.51x), min 1.004x
MI355X	3: sparse split-major, 166 cases	geomean 1.13x (grid 1.18x, UT 1.00x), min 0.982x

By family over all 817 cases:

• BF16: grid 1.35x, UT 1.41x;
• FP8 per-tensor: grid 1.11-1.16x, UT 1.21-1.33x;
• FP8 per-token: grid 1.09-1.11x, UT 1.03-1.06x.

Known regressions: no changed case is below 0.98x. Three unchanged cases read 0.974-0.976x; their binaries are identical, so those readings are noise.

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X

Results:

• gfx950 (MI355X): op_tests/test_flydsl_pa_decode.py passes in full at cbffdd201b: 355 passed, 37 skipped.
• Negative control: with commit 3's two axis fixes reverted, 9 cases fail (the MTP4 -direct, -direct-d64 and BF16 -cluster cases), so the existing tests cover them.
• gfx942 (MI308X, 80 CUs): 285 passed, 107 skipped. These are the same counts as [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709, and the gfx942 binaries are identical to the base.

No unit test is added: the existing cases reach every new gfx950 specialization. Not tested on MI250X (gfx90a is unsupported by this kernel), MI300X (304 CUs) or MI325X.

Documentation

• Docstrings updated
• User guide updated
• Performance guide updated
• No documentation change is needed, with the reason stated below

The API, the plan and the set of supported shapes are unchanged; the change is internal to the gfx950 tile schedule.

Dependencies

• No new third-party dependencies added
• New dependencies are listed and justified below

Breaking Changes

None

PR Self-Review

Reviewed revision

cbffdd2

Behavior and intent

On gfx950 the pa_decode tile schedule now uses three policies #7709 introduced for gfx942:

• K/V streaming for long head rows removes the next-tile K prefetch spill.
• The MTP query split spreads one sequence's query positions over idle CUs.
• Split-major launches stop small grids' later splits and heads from queuing behind empty plan slots.

Each is re-tuned on MI355X where the gfx942 rule lost: per-token FP8 keeps its fused MTP4 tiles, the FP8 split is bounded by the uncapped tile count, and split-major is limited to sparse grids and to FP8 or BF16 up to D128. Without these policies, gfx950 spilled at D512/D1024 and left CUs idle under MTP (1.19x geomean over the matrix). gfx942 and gfx1250 binaries are unchanged.

Validation and unverified scope

Evidence on MI355X:

• pytest op_tests/test_flydsl_pa_decode.py: 355 passed, 37 skipped. The negative control with commit 3's axis fixes reverted fails 9 cases.
• Paired in-process A/B: per commit and end-to-end over all 817 matrix cases. Geomean 1.19x, no changed case below 0.98x, 0 errors, 0 non-finite.
• Determinism: every one of the 333 changed cases repeats bit-exactly over 20 launches (det_check.py).
• Accuracy: 19 changed cases are not bit-exact against the base (streamed accumulation order). On each, the max error against the fp32 reference equals the base's.

Evidence on gfx942:

• pytest on MI308X: 285 passed, 107 skipped.
• Compile-only replay (COMPILE_ONLY=1) of every matrix case: gfx942 and gfx1250 binaries identical to the base branch.

Unverified:

• Not run on MI350X or other gfx950 SKUs with fewer CUs; the CU/4 and 8-tiles-per-CU bounds were measured at 256 CUs.
• Shapes outside the matrix are not covered.
• Callers that pass no max_context_length get no tile bound, so their FP8 MTP rows do not split on gfx950.
• CI and the self-review bot have not run, because the base is not main.
• Not yet rebuilt on [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 bdabd08dad / [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 a8b9f032b9; every number above is for the base SHAs listed in the summary.

PR-specific question

The most important boundary is the split-major grid on gfx950: it swaps which block-id axis names the plan slot. Any gfx950-only path that still assumes slot-major axes silently reads the wrong slot or head.

Two such paths existed:

• the BF16 D128 planned output's slot count (grid_dim.x);
• the one-partition empty-row zero fill (block_id("y")).

Opening split-major on gfx950 without fixing them produced relative errors of 1-19 on BF16 D128 two-head shapes. Both now select the axis from SPLIT_MAJOR, and their IR is unchanged when it is off. The remaining direct block_id/grid_dim reads are gfx1250-only pipelines, or the MXFP4 direct path, which keeps slot-major. Evidence: with the two fixes reverted, the existing UT fails 9 cases; with them, all pass; every changed case is deterministic and matches the base's fp32-reference error.

The second boundary is the cache key. The renamed busy flags keep their positions ahead of D256_PIPE, so create()'s positional unpacking and the test's cache_key[-2:] are unchanged; the pipe-policy tests pass.

Changed op tests

None. The existing test_pa_decode cases reach every new gfx950 specialization: the head1024/head512/head192 streamed rows, the MTP4/MTP8 -direct, -merge, -probability_tail and window cases for the query split, and the -direct/-cluster cases for split-major. The negative control shows they fail when the split-major axis fixes are missing.

🤖 Generated with Claude Code