[fsx950223]

Summary

gfx950 tile-kernel and planner changes to the FlyDSL pa_decode, four of which also cover gfx942, plus one host-side change for scalar FP8 scales. MiniMax-M3 TP4 decode and MTP4 verify (Hq16/Hkv1/D128, page 128, ATOM call style) get faster on MI355X: geomean 1.145x over 20 workloads, BF16 MTP4 verify up to 1.55x. Larger page-128 batches gain 1.174x. A 51-shape pa_decode review set, mostly page 16, D256 and MXFP4, gains 1.359x:

• page 16 now prefetches V and keeps it in flight (single-CTA long contexts up to 1.88x);
• one-partition plans write output without the reducer launch;
• the planner sizes tasks by a min-max tile bound, so mixed-length batches no longer end on one sequence's long tasks;
• one-partition and plan-free MXFP4 sequences on sparse grids run two tile groups per CTA (single-CTA 30K context 2.29x, MXFP4 64K 1.41x, MXFP4 2K in a 512K table 1.14x);
• calls with scalar FP8 scales (None or a Python float) no longer launch two scale fills per call (short calls 9.6 -> 6.3 us);
• the tile kernel's leading arguments are the ones its prologue reads first, preloaded into SGPRs for every KV dtype, so a task no longer starts with kernel-argument round trips;
• multi-query (MTP) plans default to four task slots per sequence, clamped to one to two per CU, so verify of up to 64 sequences runs half as many, longer tasks. BF16 and per-token FP8 verify (ATOM's FP8 cache layout) gain from this; FP8 verify with per-tensor scales is ~2% slower than main at 8 sequences and 32K tokens.

Over 24 random-length batches of 64-256 sequences in the MiniMax-M3 configuration the branch is 1.167x; one of them, BF16 decode of 64 sequences with 1K-8K tokens, is 4% slower than main (see Performance).

On gfx942 (measured on MI308X, 80 CU), four of these paths also win and now cover it: one-partition output, the min-max planner, two tile groups (head_dim 128 only) and the packed Q absmax (head_dim up to 256). Over all 786 gfx942 cases of the op test (its 111 cases plus every geometry at batch 1/16/64 and 1K/8K/32K tokens) the branch is 1.024x, one-partition long contexts up to 1.49x and random-length MTP4 batches up to 1.27x; the worst case is 0.969x (see Performance). The other changes stay gfx950-only. gfx1250 kernels are unchanged.

Motivation

On gfx950 four paths left time on the table:

• BF16 MTP4 verify ran the generic multi-query tile. It used the K16 BF16 MFMA, one shared P buffer and 9 barriers per tile. Per KV byte it was 1.4-2x slower than decode.
• BF16 decode issued V only after the softmax barriers and kept K/V lines cacheable. Long contexts stopped at ~5.6 TB/s, VMEM-issue bound.
• Single-tile decode tasks are a dependent-load chain: kernarg, page ids, Q, K, softmax, V. Scalar FP8 page-16, head-256 and BF16 page-16 tiles loaded Q only after the page ids. Per-token FP8 scale staging waited on vmcnt behind K, which drained K and V before the softmax.
• FP8 MTP4 verify is a VALU/CVT chain: 33-41 non-MFMA VALU and 2-3 CVT per MFMA.
• Page-16 tiles staged the next tile's page ids from a block-table load issued after this tile's V. In-order vmcnt made that staging wait drain V, so V overlapped only QK, and the next K went out after the P barrier with only PV to hide it: ~1.7 us per 256-token FP8 tile in one CTA.
• The planner gives every sequence one task and spreads the remaining slots in proportion to tiles. With mixed lengths and about two slots per sequence, long sequences keep the longest tasks. On a page-16 H8 B32 batch of random 1K-8K contexts, the longest task has 16 tiles against an average of 8.5, and the kernel ends on that tail.
• A one-partition or plan-free sequence is a single CTA walking all of its tiles: four waves, one per SIMD, so on sparse grids each tile's latency chain is exposed (FP8 page 16: 1.42 us per 256-token tile; MXFP4: 1.23 us).
• Q quantization reduced |q| over each lane's BF16/FP16 unit by widening every element to f32 (~30 VALU per unit), then across the row's 16 lanes with xor DPP steps of maxnumf (~22 VALU per M-tile), all on the prologue's critical path.
• Head-256 decode tiles compared and selected every score chunk of every tile, although only a task's last tile can hold tokens past the context.
• Multi-query (MTP) tasks quantize and stage every query row and publish their maxima before any KV work, a cost paid once per task. With the default two slots per CU, MTP4 verify ran twice as many, shorter tasks, and the reducer merged twice as many partials. Large batches of uneven lengths still need several slots per sequence to split their longest sequences.
• One-partition plans (max_partitions == 1) still wrote packed partials and launched the reducer to copy them out, a second dispatch on every call.
• The task prologue waited on three dependent scalar loads: the work_info pointer (byte 0x60, past the BF16-only preloaded prefix), the task record, then the remaining kernel arguments, which LLVM sinks into the task branch. FP8 and MXFP4 tiles preloaded nothing.
• Scalar FP8 scales (k_scale/v_scale passed as None or a Python float) were rebuilt on every call with torch.ones / torch.tensor. Each is a fill or host copy, so such a call ran 4 dispatches instead of 2, including inside captured CUDA graphs. That is ~3.4 us on a ~6.5 us call.
• gfx942 runs the same generic paths for the one-partition reducer launch, the proportional planner, the lone four-wave CTA and the f32 Q absmax, so the same costs apply there.

Changes

Every kernel change sits behind a PaDecodeTraits flag that requires is_gfx950, except the four paths that also cover gfx942 (listed under "gfx942" below):

• BF16_MTP_FUSED (BF16, D128, page 128, 2-4 M-tiles, no sliding window, positive softmax scale):
  • Q stays in padded LDS rows, and each M-tile owns a private P/Lsum slot, so one publication barrier suffices.
  • V is carried across M-tiles, and K16 packs pair into gfx950's K32 BF16 MFMA.
  • The resident set fits 256 VGPRs, so 2 CTAs run per CU.
  • The task's last tile is peeled and applies the V tail mask (mask_unowned_kv semantics unchanged).
• BF16_DECODE_NT_KV, DECODE_Q_EARLY, DECODE_LAST_V_AFTER_QK: BF16 single-M decode streams K and V nontemporally, issues Q before the page ids, and loads the last tile's V right after QK. Streaming gives up MALL reuse, so a KV cache that stays resident across calls (one cache replayed hot) can be slower; cold KV, as in serving, gains (see Performance).
• FP8_SCALES_AFTER_Q, FP8_LOOP_SCALES_LATE: per-token FP8 page-128 tiles stage their scales so the scale waits no longer drain K/V.
• FP8_MTP4_LEAN, FP8_MTP4_Q_LDS, FP8_MTP4_PAGES_AHEAD:
  • FP8 MTP4 verify folds the K scale into exp2. Each chunk's P is multiplied by one v_scale / v_max factor, shared by the M-tiles, before the FP8 pack, and PV's output is scaled back by v_max; the PV MFMA itself takes unit scales.
  • It rereads Q from LDS and carries the next tile's page ids, with waves_per_eu=2.
• Page 16, V prefetch: gfx950 scalar-scale FP8 and BF16 page-16 single-M tiles take the existing V-prefetch pipe (PAGE16_VPIPE), which per-token FP8 page 16 already used. Before, they issued V after the softmax barriers, at ~2.7 us (FP8) and ~3.7 us (BF16) per 256-token tile in a single CTA.
• Page 16, PAGE16_PAGES_FIRST / PAGE16_EARLY_K: scalar-scale page-16 tiles load the next tile's page ids ahead of the current V, so staging them no longer drains V. FP8 also issues the next K from those ids right after QK. BF16 keeps K after the P barrier, because its next K needs twice the registers and was 4% slower per tile.
• FP8 page-128 decode, FP8_DECODE_NT_V: V is streamed with nontemporal loads. K stays cached, because streaming it as well slowed short single-tile tasks, and on page 16 streaming fights the V pipe.
• The gfx950 tile launch preloads its leading kernel arguments (see Kernel arguments). It takes AITER_FLYDSL_KERNARG_PRELOAD and min(16, AITER_FLYDSL_KERNARG_PRELOAD_COUNT) from tensor_shim, so AITER_FLYDSL_KERNARG_PRELOAD=0 turns it off as for other FlyDSL kernels; the defaults give 16.

One-partition output, gfx942 and gfx950, FP8/BF16 (pa_decode.py, pa_decode_kernel.py, op_epilog.py):

• When a plan has max_partitions == 1 (capacity equals the batch), sinks are absent and the output has the query's strides, the tile kernel stores normalized rows directly and the reducer is not launched. gfx1250 and MXFP4 already did this.
• Slot s runs packed task s and zero-fills sequence s when the plan gives it no task (reduce_info count 0), as the reducer did. The direct store uses the query's row/head strides for any query length.

Planner (pa_decode_plan.py), min-max tasks on gfx942 and gfx950:

• Batches that fit one row per planner thread give a sequence of t tiles min(ceil(t / T), max_partitions) tasks. T is the smallest bound from ceil(total / capacity) up whose tasks fit the capacity.
• The search runs over [ceil(total / capacity), ceil(total / (capacity - rows))], whose top always fits. It does one block reduction per halving in int32 and is skipped when every slot is already taken (capacity equals the number of nonempty rows).
• Equal-length batches keep their task size and may use fewer slots, which shortens the reducer. A planner refresh costs +0.06 to +0.51 us (3.1-4.6 us before).
• gfx1250, and batches above one row per thread, keep the proportional split.
• gfx950 only: plans with query_length > 1 and no explicit workgroup_budget default to four slots per sequence and KV head, clamped to one to two per CU: one slot per CU up to 64 sequences, two (as before) from 128. Explicit budgets, single-query plans and other architectures are unchanged, and so is prepare_pa_decode_plan, which always passes a budget (2x CU on a tuning-cache miss). Scratch still follows plan.capacity, as documented, which is smaller for batches of up to 64 sequences.
• The budget cannot see the KV dtype or scale mode. The fused BF16 and per-token FP8 MTP4 paths gain from it, and per-token is ATOM's FP8 cache layout (one fp32 per KV head and token). FP8 MTP verify with per-tensor scales is ~2% slower than main at 8 sequences and 32K tokens (0.982x, reviewer measurement); such callers can pass workgroup_budget=2 * num_CUs.
• test_pa_decode_plan asserts the min-max counts when the planner compiles that variant, and the proportional split otherwise.

The direct store keeps main's address expression for single-query kernels, so the MXFP4 plan-free kernels compile as before.

Two tile groups per CTA, gfx950 and gfx942 at head_dim 128 (DUAL_GROUP; traits, context, LDS, pipeline, kernel):

• When a one-partition or plan-free plan has at most one CTA per CU, single-query tiles without a window (scalar FP8, BF16, and MXFP4 when two LDS copies fit 160 KiB) launch eight waves as two four-wave groups.
• Group g walks tiles start + g, + 2, ... with its own copy of the LDS layout. Both run the same iteration count: group 1's last tile may lie past the task and is fully masked. Their CTA barriers therefore pair up.
• After the loop, group 1 publishes (O, max, sum) per lane and group 0 rescales both to the shared maximum and stores, directly or as the task's partial. The head-256 pipe is excluded: it slowed down.
• Dual MXFP4 groups stage only their own tiles' page ids, so the page region of each LDS copy halves, and long tables fit two copies too.
• The sparse-grid test joins the schedule key ahead of its last two fields, which tests read.

On gfx950, and on gfx942 up to head_dim 256, the Q absmax reduces each unit as a packed 16-bit integer max over sign-cleared words, because a 16-bit float's magnitude orders like those bits. The row's 16 lanes then combine the f32 bits with integer max over DPP patterns that fold into the max: rotate by 8 and 4, then the two quad swaps. The result is the same value.

Scalar FP8 page-16, scalar head-256 and BF16 page-16 decode tiles issue Q ahead of the page ids and stage it while K is in flight (DECODE_Q_EARLY), as FP8 head-128 page-128 decode already did. A partial tile's padded rows load the last real row, and the existing row guard never stages them. Per-token head-256 is excluded: its prologue orders the scales ahead of K, and early Q cost it 10%.

Single-query head-256 decode tiles without a window skip the score mask on interior tiles; the peeled last tile keeps it. On the page-16 V pipe the masks cost nothing and the skip measured 1.5% slower, so it is left out there.

Kernel arguments, gfx950 (pa_decode_kernel.py):

• The tile kernel's arguments lead with work_info, the block table, Q, max_blocks_per_seq, num_seqs, the Q strides, pmax (reduce_info on one-partition plans) and K: the 14 dwords that preload into user SGPRs. The preload now covers every KV dtype, not only BF16.
• The task record, page-id and Q loads issue without a kernel-argument wait; K/V and scale pointers are first used after the page ids arrive.
• The kernel body is shared (_pa_decode_tile_body); gfx942 keeps main's argument order and preload setting.

FP8 packing, gfx950 (op_gemm.py): the first cvt_pk_fp8_f32 of each packed word takes an undefined old word instead of 0, because the second conversion writes the other half. That drops a v_mov_b32 0 per word (22 per FP8 MTP4 verify tile); results are bit-identical.

Host side, all architectures (pa_decode.py):

• normalize_scale returns one persistent one-element scale per (device, value) for None and Python-float scales.
• A scale first built while a CUDA graph is capturing is not cached, because its fill has not run yet and its memory belongs to that graph's pool. The next eager call builds and caches it.
• The cache is capped at 64 entries. Tensor scales are passed through as before.

Merge with main: the branch merges main after #7388, which keys the tile kernel on FlyDSL dtypes. The gfx950 paths read main's renamed traits (kv_dtype, query_dtype, OUT_DTYPE) and the KV_COMPUTE_BLOCK constant. The kernel name and grid follow main, and two-group kernels append _dual to the name and launch CTA_THREADS (512) threads. Main's direct-output check, now in the traits schedule, also admits one-partition gfx942 and gfx950 plans.

Tile cache key (traits.py): the specialization key is a PaDecodeCacheKey NamedTuple. PaDecodeSchedule.select() builds it by keyword and PaDecodeTraits.create() reads it by name, so a field that one change adds and another misses fails loudly instead of shifting the later fields. It is still a tuple: hashing, equality and the test's cache_key[-2:] are unchanged, and schedules and traits are field-for-field identical before and after over 5184 gfx942 and 5184 gfx950 configurations.

gfx942 (MI308X, 80 CU). Each gfx950 path was tried on gfx942 against main's gfx942 kernels. Four win there and now cover gfx942 through an is_gfx942 schedule field (outside the cache key):

• one-partition FP8/BF16 output without the reducer launch;
• the min-max planner;
• two tile groups per CTA, for head_dim 128 only: wider gfx942 heads already take up to 512 VGPRs per wave, so an eight-wave CTA of them cannot launch (the pytest's head-1024 cases fail with a launch error otherwise);
• the packed Q absmax, up to head_dim 256: head-1024 tiles already spill, and the changed Q prologue shifted their allocation (FP8 per-tensor transposed-V, 64 sequences at 32K: 0.87x).

The rest stays gfx950-only, because on gfx942 it is hardware-specific (K32 BF16 and scaled FP8 MFMA, MXFP4) or measured no better: BF16 nontemporal K/V (0.973x, up to 12% slower at 16-64 sequences and 8K-32K tokens), FP8 nontemporal V (0.987x, random lengths 0.951x), the MTP default budget (+25-32% at one sequence, -22% at 16 sequences and 32K tokens: 80 CUs need another rule), and the FP8 pack (0.998x). gfx950 and gfx1250 take exactly the previous conditions.

What does not change:

• the API and defaults, and plans on other architectures;
• the reducer, which a separate PR handles;
• gfx942 kernels outside the four paths above;
• gfx1250 kernels, which are gated out by is_gfx950 / is_gfx942;
• gfx950 kernels, compared with the previous head of this PR: the gfx942 extension and review commits leave all 49 PA kernels of the 87-shape set identical on an MI355X.

Performance (if applicable)

Method:

• Paired, in-process A/B. Main's and this branch's pa_decode are loaded side by side in one process. They run on identical inputs and plans, each captured as a CUDA graph of 10 calls. Replays alternate ABBA: 20 pairs, 10 replays per sample. The columns are per-version medians; the speedup is the median paired ratio, and every block of 10 pairs agrees on its sign.
• Inputs: MiniMax-M3 inputs and ATOM call style from zhiding512/pa-benchmark (bench_pa_minimax_m3.py, not part of this PR): tables padded to 256K tokens, no max_context_length, default plan.
• Environment: MI355X (gfx950, 256 CU), ROCm 7.2.4, FlyDSL 0.5.0, main at d02f9d058a. The gfx950 numbers were taken at 59d523033c; the later commits leave all 49 gfx950 PA kernels of the 87 shapes (tile, plan and reducer) identical on the MI355X.
• Noise: a same-tree control reads within +-0.3%.

Hardware	Workload/configuration	Before	After	Speedup
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b1-q1-ctx8192, bf16 KV	7.75 us	7.40 us	+4.7%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b1-q1-ctx8192, fp8 KV	7.50 us	7.30 us	+2.7%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b4-q1-ctx8192, bf16 KV	8.01 us	7.88 us	+1.7%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b4-q1-ctx8192, fp8 KV	7.75 us	7.49 us	+3.4%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b8-q1-ctx8192, bf16 KV	10.36 us	8.95 us	+15.7%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b8-q1-ctx8192, fp8 KV	8.16 us	7.91 us	+3.1%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b32-q1-ctx8192, bf16 KV	23.59 us	23.22 us	+1.6%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b32-q1-ctx8192, fp8 KV	16.55 us	15.08 us	+9.9%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b4-q1-ctx32768, bf16 KV	16.18 us	14.57 us	+11.2%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b4-q1-ctx32768, fp8 KV	10.79 us	9.72 us	+11.1%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b3-q1-ctx225280, bf16 KV	64.52 us	57.68 us	+11.9%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-decode-b3-q1-ctx225280, fp8 KV	33.89 us	32.37 us	+4.8%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx8192, bf16 KV	12.58 us	10.81 us	+16.4%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx8192, fp8 KV	10.98 us	9.95 us	+10.3%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx32768, bf16 KV	25.10 us	16.51 us	+51.9%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx32768, fp8 KV	16.35 us	13.54 us	+20.8%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b8-q4-ctx32768, bf16 KV	38.53 us	28.41 us	+35.7%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b8-q4-ctx32768, fp8 KV	23.20 us	19.51 us	+18.9%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx225280, bf16 KV	89.38 us	57.53 us	+55.3%
MI355X (gfx950, 256 CU)	MiniMax-M3 dense-verify-b3-q4-ctx225280, fp8 KV	44.29 us	38.05 us	+16.4%

Geomean over these 20: 1.145x (min 1.016x). Their FP8 scales are per-token tensors (ATOM's FP8 cache layout), so the host-side change does not apply to them.

Two cases measured by the reviewer (MI355X, main and this PR in separate interleaved processes):

• FP8 MTP4 verify with per-tensor scales, 8 sequences at 32K: 24.30 -> 24.74 us (0.982x) with the new default budget, 1.007x with workgroup_budget=512. With per-token scales the same shape is 1.074x of main at two slots per CU and 1.189x at the new default (table above).
• BF16 decode with a hot cache, 4 sequences at 8K replayed on one KV cache: 7.83 -> 8.14 us (0.962x). Rotated over 24 caches, so the working set exceeds the 256 MB MALL, it is 10.98 -> 8.71 us (1.26x). The nontemporal K/V loads give up MALL reuse across calls; per-layer KV is normally cold in serving. This harness replays one cache and reads that shape at +1.7%.

Shapes this PR does not target, same method:

• 16 larger page-128 MiniMax-M3 batches (b64-b256 decode, b16-b32 verify, b16 at 64K): 1.174x, min 0.985x (BF16 b64 8K, which read 1.007-1.025x in three earlier runs of the same plan; +-2% is the single-run noise here).
• 51 shapes from the pa_decode review benchmark (page 16, D256 page 64, MXFP4, padded 128K tables, budget / max_partitions / hinted plans; FP8 scales passed as None or floats): 1.359x.
  • Best: R7_fp8_d128_p16_B1H1_mp1_tbl32k_ctx30k 2.29x (310.4 -> 135.9 us, page-16 V pipe and two tile groups), W6_fp8_d128_p16_B64H1_mp1_ctx1k 1.99x (16.00 -> 8.03 us, one-partition output), R1_fp8_d128_p16_B1H1_tbl128k_ctx1024 1.59x (9.57 -> 6.03 us, scale fills and early Q), P3_fp8_d256_p64_H1_B1_ctxU1-8k_tbl32k 1.51x.
  • Two tile groups: W3_mxfp4_free_B1H1_ctx=tbl64k 1.40x (314.2 -> 224.4 us), W4_mxfp4_free_B1H1_tbl512k_ctx2k 1.14x. Planner: P2_fp8_d128_p16_H8_B32_ctxU1-8k_tbl32k 1.42x, P1_fp8_d128_p16_H1_B128_ctxU1-8k_tbl32k 1.24x, H2_fp8_d256_p64_H8_B32_ctxU1-8k_tbl32k_hint 1.18x. Head-256 interior masks: X2_fp8_d256_p64_B8H8_mp1_strided_pertok_ctx=tbl12k 1.13x.
  • Lowest: D256 X1 B32 1.016x and BF16 page-16 R2 1.022-1.038x.
• 40 further gfx950 shapes outside both benchmarks (head dim 64 and 256, D128 page 64, GQA 4 and 16, MTP 2-4 off the tuned MiniMax path, sliding windows, per-token FP8 page 16, batch 512, max_partitions and budget plans, MXFP4 B8/B64): 1.182x, min 0.991x (BF16 page-128 GQA-8 H4 B16 with random 1K-16K lengths, 1.039x in an earlier run), 33 of 40 bit-exact.
• 24 random-length large batches in the MiniMax-M3 configuration (Hkv1, GQA 16, D128, page 128; BF16 and per-token FP8; 64, 128 and 256 sequences with lengths uniform in 1K-8K, 1K-32K or 128-2K; decode and MTP4 verify): 1.167x (decode 1.103x, verify 1.382x). The MiniMax-M3 benchmark itself uses equal lengths. FP8 verify of 256 sequences: 93.6 -> 71.2 us (1.32x); BF16 verify: 1.49-1.61x.
  • Slower than main: BF16 decode B64 1K-8K 0.961x (27.2 -> 28.3 us) and B128 128-2K 0.979x; B128 1K-8K is 0.997x. Turning off the BF16 nontemporal K/V streaming speeds up B64 1K-8K by 11% and B128 1K-8K by 4%, but slows B128 1K-32K, B256 1K-8K and B256 1K-32K by 8-12% and B64/B128 128-2K by 3-6%, so the streaming stays. The B128 128-2K gap is not explained by it and was not root-caused.

Over all 87 shapes: 1.272x. Outputs are bit-exact to main on 55 of 87 shapes. The rest are MTP4 verify and FP8 paths whose arithmetic changes, plans whose task boundaries move with the planner (partials are merged in a different grouping), two-group tiles, where FP8 P is quantized against each group's running maximum, and head-256 tiles, where the unmasked score multiply contracts into the exp's FMA. All of them are within the correctness limits below.

The last three commits, each measured on top of the rest of this branch (same method, 87 shapes): the argument preload 1.0073x (scored 1.009x) and the FP8 pack 1.0031x (FP8 MTP4 verify 1.011x), both with every output bit-exact; the MTP default budget 1.0084x, with 12 of the 14 MTP4 verify workloads faster (up to 1.14x) and FP8 b3 225K 0.972x and b16 32K 0.993x. Those two are still faster than main, at 1.16x and 1.14x. Its first form, one slot per CU at any batch size, left 128-256 random-length sequences one or two tasks each: FP8 MTP4 verify of 256 sequences with 1K-8K tokens ran 0.96x of main. Four slots per sequence fix that (1.32x below) and plan all 87 shapes as before.

gfx942 (MI308X)

Method: the same paired in-process ABBA harness, CUDA graph of 10 calls, 10 replays per sample, 10 pairs per case for the full matrix and 12-20 for the probes and re-runs. Cases are the gfx942 cases of op_tests/test_flydsl_pa_decode.py (FP8-FNUZ with per-tensor and per-token scales, plain and transposed V, and BF16): its 111 cases with their own lengths and plans, plus each of its 75 geometry and KV-type combinations (head_dim 64-1024, page 8-128, GQA 4-16, query length 1-8, windows, sinks) at 1, 16 and 64 sequences of 1K, 8K and 32K tokens with the default plan (675). MI308X (gfx942, 80 CU), ROCm 7.2.0, FlyDSL 0.5.0, main at d02f9d058a, this branch at d6fa8700ad with the head-1024 rows re-measured at 8240f4cdd4 (the two differ only in the Q absmax of FP8 heads above 256). A foreign-process monitor on the timing GPU stayed clean; one disturbed probe row was re-run.

Hardware	Workload/configuration	Before	After	Speedup
MI308X (gfx942, 80 CU)	op test decode: FP8 per-tensor, D128, page 16, GQA 8, one partition	19.88 us	15.83 us	+25.6%
MI308X (gfx942, 80 CU)	op test decode: FP8 per-token, D128, page 16, GQA 8, one partition	22.41 us	19.46 us	+15.2%
MI308X (gfx942, 80 CU)	one partition, FP8 per-tensor, page 16, GQA 8, 1 x 30K tokens	424.07 us	284.30 us	+49.1%
MI308X (gfx942, 80 CU)	one partition, BF16, page 16, GQA 8, 1 x 30K tokens	612.07 us	432.92 us	+41.4%
MI308X (gfx942, 80 CU)	one partition, FP8 per-tensor, page 16, GQA 8, 64 x 1K tokens	21.12 us	14.52 us	+45.6%
MI308X (gfx942, 80 CU)	random 1K-8K, BF16 MTP4, page 128, GQA 16, 128 sequences	532.69 us	420.84 us	+26.6%
MI308X (gfx942, 80 CU)	random 1K-8K, FP8 per-token MTP4, page 128, GQA 16, 128 sequences	493.89 us	407.19 us	+21.3%
MI308X (gfx942, 80 CU)	random 1K-8K, FP8 per-token decode, page 128, GQA 16, 128 sequences	133.98 us	116.78 us	+14.7%
MI308X (gfx942, 80 CU)	worst: FP8 per-tensor transposed V, D1024, page 128, GQA 4, sinks, 64 x 1K	162.69 us	167.88 us	-3.1%
MI308X (gfx942, 80 CU)	BF16 MTP4, D512, page 64, GQA 16, 64 x 1K	537.43 us	550.83 us	-2.5%

• All 786 cases: 1.024x (op-test cases 1.058x, grid 1.019x); 441 bit-exact. By KV type: BF16 1.009x, FP8 per-tensor 1.031x (plain V) / 1.034x (transposed), per-token 1.029x / 1.029x. By head_dim: 128 1.032x (464 cases), 256 1.023x, 1024 1.009x, 64 1.006x, 192 1.003x, 512 0.998x.
• Below 0.98x: 5 of 786, all traced to the min-max planner: FP8 head-1024 page-128 with sinks at 0.969-0.979x (64 x 1K and 1 x 32K), and BF16 MTP4 head-512 64 x 1K at 0.975x. This branch with main's planner restored, against this branch, reads 0.971-0.981x on those FP8 rows and 0.974x on the head-512 row, so the planner accounts for nearly all of it. The same comparison gains BF16 head-1024 MTP3/MTP4 at 16-64 sequences up to 1.18x (64 x 1K), so the planner stays on for gfx942.
• 14 random-length batches (lengths uniform in 1K-8K; 16-128 sequences; page 16 and 128; FP8 and BF16; decode and MTP4): 1.104x, min 0.999x.
• 7 one-partition probes (max_partitions=1; 1-64 sequences, 1K-32K tokens): 1.279x, min 1.020x.

Testing

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

op_tests/test_flydsl_pa_decode.py::test_pa_decode_plan now asserts the min-max task counts wherever the planner compiles that variant (gfx942 and gfx950); the rest of the file already covers BF16/FP8 MTP4 page 128, unowned-KV poisoning, one-partition plans, two tile groups and the reducer, and passes unchanged. At the head: gfx950 (MI355X) 355 passed / 37 skipped, gfx942 (MI308X) 285 passed / 107 skipped. MI250X is not targeted. gfx942 was tested on an MI308X (80 CU), not an MI300X; see the self-review.

Documentation

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

These are internal gfx950/gfx942 scheduling and layout changes behind traits flags; the public API and its docstring are unchanged.

Dependencies

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

Merge order: #7709 (gfx942 schedule tuning) conflicts with this PR in traits.py, pipeline.py, op_lds.py and the launch block of pa_decode_kernel.py. This PR should merge first; #7709 is then rebased onto it, keeps block=(traits.CTA_THREADS, 1, 1) (taking BLOCK_THREADS would launch two-group kernels without their second group) and adds its cache-key fields by name, and its gfx942 numbers are re-measured on top of the four paths this PR extends to gfx942.

Breaking Changes

None

PR Self-Review

Reviewed revision

bdabd08

Behavior and intent

On gfx950, the PA decode tile kernel takes new specialized paths for BF16 MTP4 verify, BF16/FP8 single-M decode and FP8 MTP4 verify. They remove barriers, use the K32 BF16 MFMA, stream KV, reorder the dependent loads of single-tile tasks and fold FP8 scales into exp2/MFMA. In addition, scalar-FP8 and BF16 page-16 tiles now prefetch V and load the next page ids ahead of it (FP8 also issues the next K after QK), FP8 page-128 decode streams V, and one-partition FP8/BF16 plans write output from the tile kernel without the reducer. The gfx950 and gfx942 planner gives each sequence min(ceil(t / T), max_partitions) tasks with T the smallest feasible tile bound, instead of one task per sequence plus a proportional share, because mixed-length batches ended on the longest sequences' tasks. One-partition and plan-free sequences on sparse grids run two four-wave tile groups that merge once, because a lone four-wave CTA exposed each tile's latency. The Q absmax uses packed and DPP integer maxima, interior head-256 decode tiles skip the score mask, and more single-M decode families issue Q ahead of the page ids. The tile kernel's arguments are reordered so everything the task prologue reads first is preloaded into SGPRs for every KV dtype, and the FP8 pack no longer zeroes an old word it overwrites. Multi-query plans default to fewer task slots (four per sequence, one to two per CU), because each MTP task pays for staging all of its query rows while large uneven batches still need several tasks per sequence. The old gfx950 paths left verify 1.4-2x slower per byte than decode, long BF16 decode issue-bound, and page-16 tiles at ~2.7-3.7 us per 256-token tile. On gfx942, measured on an 80-CU MI308X, four of these paths also win and now apply there too: one-partition output, the min-max planner, two tile groups at head_dim 128 and the packed Q absmax up to head_dim 256; the rest is gfx950-only because it uses gfx950 instructions or measured no better on gfx942. On every architecture, None or float FP8 scales now reuse one persistent device tensor instead of two fills per call, which had doubled the dispatch count of short calls. The tile cache key becomes a named tuple read by field name, so a mis-ordered key fails loudly. gfx1250 kernels, the API, gfx942 kernels outside the four paths, and the reducer keep their previous behavior.

Validation and unverified scope

Validation:

• gfx950 (MI355X): pytest op_tests/test_flydsl_pa_decode.py at the reviewed revision: 355 passed, 37 skipped. main's version of the file fails 3 test_pa_decode_plan cases on this branch, for example [65-None-None-0-1-100-128], because it asserts the proportional split.
• MiniMax-M3 benchmark correctness, 62 workloads (BF16 and FP8), eager and graph replay with changed inputs and a NaN-filled output: passed at the reviewed revision, with stricter limits than the benchmark's own:
  • BF16 per-row relative L2 within 2x of the base's worst, for standard inputs and for Q scaled by 32 (peaked softmax);
  • FP8 output within 0.035 relative L2 of the base kernel's stored outputs. The base is bitwise deterministic: 124 outputs, two runs, zero diff.
• CUDA graph capture of the cached scales: a standalone script makes the first-ever pa_decode call (scales None, then 0.75) inside torch.cuda.graph capture. It then checks that replay, the next eager call and a second eager call give equal, finite outputs. Passed on gfx950.
• Two tile groups: the pytest's default max_partitions=1 cases (lengths 0, 1, 255, 256, 257, 1027, so even and odd tile counts and a past-the-end group-1 tile) and its MXFP4 plan-free, extreme-scale and exact-equality variant cases run through the two-group CTA and pass.
• One-partition direct output: the pytest's default cases use max_partitions=1 plans for FP8 and BF16, with a context-0 row and a NaN-filled output. A missed zero-fill or a mis-strided direct store would fail them.
• gfx942 (MI308X): the same pytest gave 285 passed, 107 skipped, on main d02f9d058a and on this branch (8240f4cdd4; the head adds only a test comment). Its default max_partitions=1 cases run the gfx942 direct output and two-group CTA, and test_pa_decode_plan checks the gfx942 min-max counts. The first two-group version also ran head_dim 512/1024 in eight-wave CTAs; the pytest's head-1024 cases failed with a launch error, which is why gfx942 two-group tiles are limited to head_dim 128. The 786-case timing run also compares every output against main's: 441 are bit-exact, and the largest difference, relative to the output's maximum magnitude, is 0.034 for FP8 and 0.0076 for BF16. Correctness against a float reference is the pytest's job.
• gfx950 binary identity: compared with 59d523033c (where the gfx950 numbers were measured), the gfx942 extension and review commits leave all 49 gfx950 PA kernels of the 87 shapes (tile, plan, reducer) with identical gpu.binary code on the MI355X. The named cache key gives field-for-field identical schedules and traits over 5184 gfx942 and 5184 gfx950 configurations. gfx1250 tile and reducer kernels compile to identical code (compile-only).
• Formatting: black 26.5.1 (psf/black@stable) and ruff 0.16.0 (the version CI pins) are clean.

Not verified:

• gfx1250 was not run (no hardware). Every kernel change is gated by is_gfx950, is_gfx942 or the architecture string; gfx1250 is checked compile-only.
• MI300X (304 CU) was not run; gfx942 was measured only on the 80-CU MI308X. The two-group test (at most one CTA per CU) and the planner's slot counts depend on the CU count, so MI300X may pick these paths on different shapes.
• On gfx942, the min-max planner makes FP8 head-1024 page-128 tiles with sinks 2-3% slower at 64 x 1K and 1 x 32K, and BF16 MTP4 head-512 2.5% slower at 64 x 1K (5 of 786 cases below 0.98x). This was not fixed, because the same planner gains elsewhere (random-length batches 1.104x).
• FP8 MTP verify with per-tensor scales under the new gfx950 default budget is ~2% slower than main at 8 sequences and 32K tokens (reviewer measurement). It is documented with its workaround, workgroup_budget=2 * num_CUs, but not fixed, because the planner does not see the scale mode.
• No end-to-end ATOM/vLLM serving run.
• BF16 MTP4 verify with page 16/64 keeps the generic path (1.015x and 1.016x in the 40-shape sweep); D256 MTP verify was not timed.
• Planner refresh cost was timed in a CUDA graph (+0.06 to +0.51 us); its end-to-end effect in a serving step is not measured.
• The capture-first scale check is a standalone script, not an op test: pytest's shared process makes "first call in this process" order-dependent.
• Random-length batches were timed only in the MiniMax-M3 configuration (24 shapes above) and in the review set's random-length shapes. The BF16 decode slowdowns at B64 1K-8K and B128 128-2K were not root-caused beyond the streaming ablation.
• Downstream integrations beyond ATOM were not checked for how they size scratch. ATOM's atom/model_ops/base_attention.py keys its scratch on plan.capacity and takes plan_pa_decode's default budget except at head_dim 256, so its D128 MTP plans pick up the new default safely. A caller that computed the capacity itself would now fail pa_decode's workspace shape check for gfx950 MTP plans with the default budget.
• Multi-GPU processes were not run. The scale cache is keyed by device, and a scale is cached only after a sync on its own device's current stream.

PR-specific question

The changed test is test_pa_decode_plan, the canonical planner test. It checks every task record: rows, contiguous tile ranges, counts, padding, refresh inside a CUDA graph, batch buckets and the LDS limit, across windows, partition caps and budgets.

• Before this PR its count assertion encoded the proportional split, and it still does where the planner compiles that variant.
• On gfx942 and gfx950 single-block plans it now recomputes the min-max bound T with an independent host binary search over the same bounds, and requires every row's count to equal min(ceil(t / T), max_partitions).
• Against main the updated test passes, because main never compiles the balanced variant. Against this branch on gfx950, main's assertion fails 3 cases, which is why it had to change. It passes on both architectures at the reviewed revision.
• A device search that stopped one halving early, mis-clamped the int32 bound or took the full-plan shortcut wrongly would fail the new assertion on the pattern lengths (including 0, -1 and 2**31-1).
• The 62-workload correctness gate and the rest of the pytest file cover the decode outputs of the re-partitioned plans.

Other boundaries: the most important kernel boundary is unowned KV in the new peeled last tile of the BF16 MTP4 fused path. V rows past the context, and padded block-table lanes pointing at page 0, may hold NaN or Inf. They must contribute exactly zero, because 0 * NaN is NaN in PV. The fused path masks V on the task's last tile with the same mask_v_chunk the generic path uses, honoring mask_unowned_kv. test_pa_decode_unowned_kv_rows (NaN and Inf fills, pages 16/64/128, query length 4) passes on gfx950. As a check, a copy of this PR with only the fused path's V mask removed fails test_pa_decode_unowned_kv_rows[bf16-128-4-1-16-128-nan] and [...-inf]; this PR passes all 14 of its BF16 page-128 cases. The two-group CTA relies on both groups executing the same CTA barriers. They run the same iteration count, the past-the-end tile is masked, and the merge has a barrier on each side. For MXFP4 a third barrier keeps the four-wave store's LDS writes behind every lane's read of the overlapping merge slots: without it, the MXFP4 plan-free and exact-variant pytest cases produced non-finite outputs. The second boundary is the gfx942/gfx950 one-partition direct output. It relies on plan_pa_decode's invariant that max_partitions == 1 means capacity equals the batch, so each sequence owns at most one task and no reduction is needed. Slot s runs packed task s, and sequence s is zero-filled when its reduce_info count is 0. The NaN-prefilled max_partitions=1 pytest cases with a context-0 row cover both. Plans with sinks or a strided output keep the reducer. The argument reorder is a positional contract between _tile_args and the gfx950 kernel signature: a mismatch passes wrong pointers or strides, which every pytest case and the 62-workload gate would catch, and gfx942 keeps main's argument order. The third boundary is the host-side scale cache under CUDA graph capture. A scale built while capturing is used for that capture only and never cached, because its fill runs only on replay and its memory is in that graph's pool. The capture-first script above exercises exactly that order. On gfx942 the boundary is register capacity: an eight-wave two-group CTA of head_dim 512 or 1024 needs up to 512 VGPRs per wave and cannot launch, so DUAL_GROUP requires head_dim 128 there. An earlier version without that limit failed the pytest's head-1024 cases with a launch error; at the reviewed revision they pass.

Changed op tests

op_tests/test_flydsl_pa_decode.py, only test_pa_decode_plan: its per-row task-count assertion now covers the gfx942/gfx950 min-max planner (bound recomputed on the host) as well as the proportional split. No other test planned and checked task records, so no new test was added; the rest of the file is unchanged.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

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

Extended tests (opt-in via labels):

Label	Tests
ci:mi300x	Run Aiter standard tests on MI300X in PRs; main branch always runs them
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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the Radeon R9700 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
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: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
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 7700 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Reviews:
A reviewer is usually assigned within one working day after the PR is ready for review (see Code Review Process). If nobody has responded for about a week, mention the area contact listed in CONTACTS.md. Want to help review? See REVIEWERS.md.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[aiter-bot]

READY FOR HUMAN REVIEW

PR Self-Review

Revision bdabd08dad5c · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

On gfx950, the PA decode tile kernel takes new specialized paths for BF16 MTP4 verify, BF16/FP8 single-M decode and FP8 MTP4 verify. They remove barriers, use the K32 BF16 MFMA, stream KV, reorder the dependent loads of single-tile tasks and fold FP8 scales into exp2/MFMA. In addition, scalar-FP8 and BF16 page-16 tiles now prefetch V and load the next page ids ahead of it (FP8 also issues the next K after QK), FP8 page-128 decode streams V, and one-partition FP8/BF16 plans write output from the tile kernel without the reducer. The gfx950 and gfx942 planner gives each sequence min(ceil(t / T), max_partitions) tasks with T the smallest feasible tile bound, instead of one task per sequence plus a proportional share, because mixed-length batches ended on the longest sequences' tasks. One-partition and plan-free sequences on sparse grids run two four-wave tile groups that merge once, because a lone four-wave CTA exposed each tile's latency. The Q absmax uses packed and DPP integer maxima, interior head-256 decode tiles skip the score mask, and more single-M decode families issue Q ahead of the page ids. The tile kernel's arguments are reordered so everything the task prologue reads first is preloaded into SGPRs for every KV dtype, and the FP8 pack no longer zeroes an old word it overwrites. Multi-query plans default to fewer task slots (four per sequence, one to two per CU), because each MTP task pays for staging all of its query rows while large uneven batches still need several tasks per sequence. The old gfx950 paths left verify 1.4-2x slower per byte than decode, long BF16 decode issue-bound, and page-16 tiles at ~2.7-3.7 us per 256-token tile. On gfx942, measured on an 80-CU MI308X, four of these paths also win and now apply there too: one-partition output, the min-max planner, two tile groups at head_dim 128 and the packed Q absmax up to head_dim 256; the rest is gfx950-only because it uses gfx950 instructions or measured no better on gfx942. On every architecture, `None` or float FP8 scales now reuse one persistent device tensor instead of two fills per call, which had doubled the dispatch count of short calls. The tile cache key becomes a named tuple read by field name, so a mis-ordered key fails loudly. gfx1250 kernels, the API, gfx942 kernels outside the four paths, and the reducer keep their previous behavior.

Validation and unverified scope

Validation:
- **gfx950 (MI355X):** `pytest op_tests/test_flydsl_pa_decode.py` at the reviewed revision: 355 passed, 37 skipped. main's version of the file fails 3 `test_pa_decode_plan` cases on this branch, for example `[65-None-None-0-1-100-128]`, because it asserts the proportional split.
- **MiniMax-M3 benchmark correctness, 62 workloads (BF16 and FP8), eager and graph replay with changed inputs and a NaN-filled output:** passed at the reviewed revision, with stricter limits than the benchmark's own:
  - BF16 per-row relative L2 within 2x of the base's worst, for standard inputs and for Q scaled by 32 (peaked softmax);
  - FP8 output within 0.035 relative L2 of the base kernel's stored outputs. The base is bitwise deterministic: 124 outputs, two runs, zero diff.
- **CUDA graph capture of the cached scales:** a standalone script makes the first-ever `pa_decode` call (scales `None`, then 0.75) inside `torch.cuda.graph` capture. It then checks that replay, the next eager call and a second eager call give equal, finite outputs. Passed on gfx950.
- **Two tile groups:** the pytest's default `max_partitions=1` cases (lengths 0, 1, 255, 256, 257, 1027, so even and odd tile counts and a past-the-end group-1 tile) and its MXFP4 plan-free, extreme-scale and exact-equality variant cases run through the two-group CTA and pass.
- **One-partition direct output:** the pytest's default cases use `max_partitions=1` plans for FP8 and BF16, with a context-0 row and a NaN-filled output. A missed zero-fill or a mis-strided direct store would fail them.
- **gfx942 (MI308X):** the same pytest gave 285 passed, 107 skipped, on main `d02f9d058a` and on this branch (`8240f4cdd4`; the head adds only a test comment). Its default `max_partitions=1` cases run the gfx942 direct output and two-group CTA, and `test_pa_decode_plan` checks the gfx942 min-max counts. The first two-group version also ran head_dim 512/1024 in eight-wave CTAs; the pytest's head-1024 cases failed with a launch error, which is why gfx942 two-group tiles are limited to head_dim 128. The 786-case timing run also compares every output against main's: 441 are bit-exact, and the largest difference, relative to the output's maximum magnitude, is 0.034 for FP8 and 0.0076 for BF16. Correctness against a float reference is the pytest's job.
- **gfx950 binary identity:** compared with `59d523033c` (where the gfx950 numbers were measured), the gfx942 extension and review commits leave all 49 gfx950 PA kernels of the 87 shapes (tile, plan, reducer) with identical `gpu.binary` code on the MI355X. The named cache key gives field-for-field identical schedules and traits over 5184 gfx942 and 5184 gfx950 configurations. gfx1250 tile and reducer kernels compile to identical code (compile-only).
- **Formatting:** black 26.5.1 (`psf/black@stable`) and ruff 0.16.0 (the version CI pins) are clean.

Not verified:
- gfx1250 was not run (no hardware). Every kernel change is gated by `is_gfx950`, `is_gfx942` or the architecture string; gfx1250 is checked compile-only.
- MI300X (304 CU) was not run; gfx942 was measured only on the 80-CU MI308X. The two-group test (at most one CTA per CU) and the planner's slot counts depend on the CU count, so MI300X may pick these paths on different shapes.
- On gfx942, the min-max planner makes FP8 head-1024 page-128 tiles with sinks 2-3% slower at 64 x 1K and 1 x 32K, and BF16 MTP4 head-512 2.5% slower at 64 x 1K (5 of 786 cases below 0.98x). This was not fixed, because the same planner gains elsewhere (random-length batches 1.104x).
- FP8 MTP verify with per-tensor scales under the new gfx950 default budget is ~2% slower than main at 8 sequences and 32K tokens (reviewer measurement). It is documented with its workaround, `workgroup_budget=2 * num_CUs`, but not fixed, because the planner does not see the scale mode.
- No end-to-end ATOM/vLLM serving run.
- BF16 MTP4 verify with page 16/64 keeps the generic path (1.015x and 1.016x in the 40-shape sweep); D256 MTP verify was

[truncated in readiness report; see the PR description for the full answer]

PR-specific question

The changed test is `test_pa_decode_plan`, the canonical planner test. It checks every task record: rows, contiguous tile ranges, counts, padding, refresh inside a CUDA graph, batch buckets and the LDS limit, across windows, partition caps and budgets.
- Before this PR its count assertion encoded the proportional split, and it still does where the planner compiles that variant.
- On gfx942 and gfx950 single-block plans it now recomputes the min-max bound T with an independent host binary search over the same bounds, and requires every row's count to equal min(ceil(t / T), max_partitions).
- Against main the updated test passes, because main never compiles the balanced variant. Against this branch on gfx950, main's assertion fails 3 cases, which is why it had to change. It passes on both architectures at the reviewed revision.
- A device search that stopped one halving early, mis-clamped the int32 bound or took the full-plan shortcut wrongly would fail the new assertion on the pattern lengths (including 0, -1 and 2**31-1).
- The 62-workload correctness gate and the rest of the pytest file cover the decode outputs of the re-partitioned plans.

Other boundaries: the most important kernel boundary is unowned KV in the new peeled last tile of the BF16 MTP4 fused path. V rows past the context, and padded block-table lanes pointing at page 0, may hold NaN or Inf. They must contribute exactly zero, because 0 * NaN is NaN in PV. The fused path masks V on the task's last tile with the same `mask_v_chunk` the generic path uses, honoring `mask_unowned_kv`. `test_pa_decode_unowned_kv_rows` (NaN and Inf fills, pages 16/64/128, query length 4) passes on gfx950. As a check, a copy of this PR with only the fused path's V mask removed fails `test_pa_decode_unowned_kv_rows[bf16-128-4-1-16-128-nan]` and `[...-inf]`; this PR passes all 14 of its BF16 page-128 cases. The two-group CTA relies on both groups executing the same CTA barriers. They run the same iteration count, the past-the-end tile is masked, and the merge has a barrier on each side. For MXFP4 a third barrier keeps the four-wave store's LDS writes behind every lane's read of the overlapping merge slots: without it, the MXFP4 plan-free and exact-variant pytest cases produced non-finite outputs. The second boundary is the gfx942/gfx950 one-partition direct output. It relies on `plan_pa_decode`'s invariant that `max_partitions == 1` means capacity equals the batch, so each sequence owns at most one task and no reduction is needed. Slot s runs packed task s, and sequence s is zero-filled when its `reduce_info` count is 0. The NaN-prefilled `max_partitions=1` pytest cases with a context-0 row cover both. Plans with sinks or a strided output keep the reducer. The argument reorder is a positional contract between `_tile_args` and the gfx950 kernel signature: a mismatch passes wrong pointers or strides, which every pytest case and the 62-workload gate would catch, and gfx942 keeps main's argument order. The third boundary is the host-side scale cache under CUDA graph capture. A scale built while capturing is used for that capture only and never cached, because its fill runs only on replay and its memory is in that graph's pool. The capture-first script above exercises exactly that order. On gfx942 the boundary is register capacity: an eight-wave two-group CTA of head_dim 512 or 1024 needs up to 512 VGPRs per wave and cannot launch, so `DUAL_GROUP` requires head_dim 128 there. An earlier version without that limit failed the pytest's head-1024 cases with a launch error; at the reviewed revision they pass.

Changed op tests

`op_tests/test_flydsl_pa_decode.py`, only `test_pa_decode_plan`: its per-row task-count assertion now covers the gfx942/gfx950 min-max planner (bound recomputed on the host) as well as the proportional split. No other test planned and checked task records, so no new test was added; the rest of the file is unchanged.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot could not review the current head (for example, quota or service failure)

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[zhiding512]

Advisory review: no correctness defect found; a few things to fix or clarify. Besides the runs below, the DUAL_GROUP barrier pairing, LDS isolation and merge numerics, and the BF16/FP8 MTP fused paths (K32 pairing, tail masks, scale gating) were read through separately; nothing wrong turned up there.

What I ran (MI355X / gfx950, 256 CUs, FlyDSL 0.5.0; this PR at 58bad83 applied on main cb4ac34): MiniMax-M3 TP4 shapes (Hq16 / Hkv1 / D128, page 128), ATOM call style (block table padded to 256K tokens, no max_context_length, default plan, scratch sized from plan.capacity, FP8 with per-tensor scale tensors), CUDA graph of 10 decodes, main and the PR in separate interleaved processes (4 each, 30 samples per process). Geomean 1.082x over 12 shapes; BF16 MTP4 verify b8/32K 1.372x, BF16 decode b3/225K 1.191x. Two shapes are slower than main; see the inline comments on the default budget and on BF16_DECODE_NT_KV. 59d5230 only changes the default budget above 64 sequences, so every shape here plans the same at the current head.

Per-shape latency (us), main vs this PR

shape	main	this PR	speedup
decode b1 ctx8K bf16	7.69	7.37	1.043x
decode b1 ctx8K fp8	7.32	7.00	1.046x
decode b4 ctx8K bf16	7.86	8.16	0.963x
decode b8 ctx8K fp8	7.86	7.52	1.045x
decode b32 ctx8K bf16	23.62	22.98	1.028x
decode b32 ctx8K fp8	15.65	14.48	1.081x
decode b4 ctx32K fp8	10.40	9.64	1.079x
decode b3 ctx225K bf16	65.72	55.17	1.191x
verify b3 q4 ctx8K bf16	12.50	10.63	1.176x
verify b3 q4 ctx8K fp8	11.46	11.07	1.035x
verify b8 q4 ctx32K bf16	38.63	28.15	1.372x
verify b8 q4 ctx32K fp8	24.30	24.74	0.982x

⚠️ Merge order with #7709. #7709 does not apply on top of this PR: a 3-way apply conflicts in op_lds.py (3 hunks), pipeline.py (8), traits.py (7) and pa_decode_kernel.py (1). Both PRs insert cache-key fields ahead of D256_PIPE, with the matching positional unpack in create(), and the one conflict in pa_decode_kernel.py is the launch block, where this PR's block=(traits.CTA_THREADS, 1, 1) sits against #7709's block=(traits.BLOCK_THREADS, 1, 1); taking #7709's side there would launch DUAL_GROUP kernels without their second group. Could you agree an order with #7709, and assert the cache-key layout somewhere so a mis-ordered unpack fails loudly?

Advisory review, AI-assisted; the runtime numbers come from the runs described above.

[zhiding512]

⚠️ For batches up to 64 sequences (one KV head, 256 CUs) this still lands on one slot per CU, for FP8 as well as BF16, and FP8 MTP4 verify ends up slower than main. Measured (setup in the review summary): FP8 verify b8 q4 ctx32K, main 24.30 us vs 24.74 us here (0.982x, plan capacity 256 instead of 512); with workgroup_budget=512 passed explicitly it is 1.007x. The same shape in BF16 gains from the smaller budget (1.277x at 512 slots, 1.372x at 256). plan_pa_decode doesn't know the KV dtype, and ATOM passes no budget for head_dim 128 (flydsl_plan_shape), so FP8 MTP deployments get the slower setting. Should FP8 keep two slots per CU (for example ATOM passing the budget by KV dtype), or is this trade-off intended? If intended, could the description say so?

[fsx950223]

Thanks for measuring this. The difference is the FP8 scale mode: your run used per-tensor scales, and ATOM's FP8 KV cache is per-token (MhaKvPool.scale_views in atom/model_ops/attentions/mha_kv_pool.py returns "one fp32 per (kv head, token)"). Per-token FP8 MTP4 verify takes the lean fused path, and there the one-slot default helps. Same harness as the description (MI355X, paired ABBA in one process, ATOM call style), this branch before vs after the default, vs main:

MiniMax-M3 MTP4 verify, per-token FP8	two slots/CU	one slot/CU (this PR)
b8 q4 ctx32K	1.074x (21.65 us)	1.189x (19.51 us)
b16 q4 ctx8K	1.124x	1.287x
b32 q4 ctx8K	1.095x	1.182x
b3 q4 ctx225K	1.187x	1.164x

So the trade-off is intended for BF16 and per-token FP8, and per-tensor FP8 MTP pays about 2% at b8 32K as you measured. plan_pa_decode cannot see the KV dtype or scale mode, so I kept the default and documented the per-tensor case in the description (Changes, Performance and the unverified list), with the workaround: callers with per-tensor FP8 scales can pass workgroup_budget=2 * num_CUs. The default now also keeps four slots per sequence (59d5230), so batches above 64 sequences are back to two slots per CU; that fixed a random-length B128-B256 regression the equal-length benchmark had hidden.

[zhiding512]

⚠️ Same as in #7701: with True/16 hard-coded, AITER_FLYDSL_KERNARG_PRELOAD=0 (documented in tensor_shim.py L24-33 as turning the preload hint off for all FlyDSL kernels) no longer reaches the gfx950 tile kernel, while the gfx1250 builder in this file still honors it. Both constants are already imported here:

Suggested change

	"amdgpu-kernarg-preload": True,
	"amdgpu-kernarg-preload-count": 16,
	"amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD,
	"amdgpu-kernarg-preload-count": min(16, AITER_FLYDSL_KERNARG_PRELOAD_COUNT),

[fsx950223]

Applied in d6fa8700ad: the gfx950 tile launch now takes AITER_FLYDSL_KERNARG_PRELOAD and min(16, AITER_FLYDSL_KERNARG_PRELOAD_COUNT) from tensor_shim, so AITER_FLYDSL_KERNARG_PRELOAD=0 reaches it like the gfx1250 builder. With the default environment the hint is unchanged (on, 16), and the gfx950 binaries are identical to the previous head.

[zhiding512]

📝 The description says no shape is slower beyond noise, but under hot-cache graph replay decode b4 ctx8K BF16 (MiniMax-M3 TP4) is 7.83 us on main vs 8.14 us here (0.962x; 3 processes x 40 samples per side, spread <= 0.8%). Rotating the graph over 24 KV caches, so the working set exceeds the 256 MB MALL, it becomes 10.98 vs 8.71 us (1.26x). That fits these nontemporal K/V loads giving up MALL reuse (inferred, not profiled). Per-layer KV is normally cold in serving, so this is probably a net win; could you qualify the claim and mention the cache-resident case?

[fsx950223]

Agreed, the claim was too broad. I see the same pattern on random-length batches: BF16 decode of 64 sequences with 1K-8K tokens is 0.961x. Turning the streaming off recovers B64 1K-8K (+11%) and B128 1K-8K (+4%) but costs B128/B256 long contexts 8-12%. Your rotation over 24 caches (1.26x) shows the other side. In d6fa8700ad the BF16_DECODE_NT_KV comment now records the cache-resident loss, and the description replaces "no shape is slower beyond noise" with the cases that are slower (your b4 8K hot-replay 0.962x and the 1.26x rotated result included), noting that per-layer KV is normally cold in serving.

[zhiding512]

📝 This comment (and the description) says a power-of-two V normalization rides the PV MFMA's P scale, but the PV MFMA always receives UNIT_E8M0_SCALES (op_gemm.py L15, L61-63), and the normalization is rcp_f32(v_max + eps) (op_softmax.py L206), multiplied into P before the FP8 pack in pack_fused_lean. Could you update the comment and the description so nobody relies on an exactness property the code doesn't have?

[fsx950223]

You're right, the comment described an earlier design. The PV MFMA always gets UNIT_E8M0_SCALES; P is multiplied by one v_scale / v_max factor per chunk (shared by the M-tiles) before the FP8 pack, and the PV result is scaled back by v_max. Fixed in d6fa8700ad (the FP8_MTP4_LEAN comment) and in the description's Changes section. No exactness property depends on it.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[fsx950223]

Thanks for the review and the independent numbers.

Merge order with #7709. #7709 is my gfx942 PR. Proposed order: #7700 first, then I rebase #7709 onto it:

• The launch block must keep block=(traits.CTA_THREADS, 1, 1), as you noted.
• This PR now also extends four of its paths to gfx942 (80d7bba, 8240f4c): one-partition direct output, the min-max planner, two tile groups at head_dim 128, and the packed Q absmax up to head_dim 256. Head-1024 tiles keep main's absmax, because they already spill and the packed version made one FP8 case 0.87x. On an MI308X the branch is 1.024x over the op test's 786 gfx942 cases (min 0.969x); details are in the description. [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's gfx942 numbers will be re-measured on top of it.
• The description's Dependencies section records this order.

Cache-key layout. In d6fa8700ad the tile specialization key is a NamedTuple: select() builds it by keyword and PaDecodeTraits.create() reads it by name. A field that one change adds and another misses now fails with a TypeError/AttributeError instead of shifting the later fields. It is still a tuple, so hashing, equality and the test's cache_key[-2:] are unchanged. Schedules and traits are field-for-field identical before and after (5184 gfx942 and 5184 gfx950 configurations).

The inline points are answered in their threads.