[fsx950223]

This extends the gfx950 D256 FP8 PA decode pipe to page sizes 16/64, packed transposed V, per-token KV scales, and positive sliding windows, with at most 32 query rows. Non-transposed 4D K caches are explicitly unsupported: the host, compiler and plan-preparation entry points raise NotImplementedError; trans_k=True remains the default compilation argument.

The PR targets samremes/pa-decode-longctx and merges that branch through 8cf11692f9. Merging this PR incorporates these changes into #5949.

Changes

• Share K/V operands across up to two query M-tiles. Reuse per-token score scaling, V-scale normalization and both packed V loaders, with scale staging before K prefetch.
• Enable the D256 pipe for page16 as well as page64. Read vector-width page IDs from the staged LDS list, with scalar block-table fallback when outside the staged bound. Use a separate scalar block-table view so page16 staging indexes page IDs correctly.
• Size staged page IDs from the visible sliding-window tile union, including unaligned MTP boundaries and ragged plans. Preserve empty-row behavior by scaling before masking.
• Use the CU-count default plan budget for eligible D256 pipe shapes. Saved autotune results take precedence. gfx942 retains the generic path.
• Remove the NHD implementation, its DMA/gather/LDS scheduling, buffer-selection policy, six runtime stride arguments, tuning CLI layout option and NHD test cases. Keep packed K with either plain or transposed V. Remove from_table_width and its warning while retaining the 4096-page-ID staging limit and block-table fallback.
• Retain the upstream reducer and the simplified six-argument d256_pipe. Supplementary validation scripts remain outside this PR.

Validation

Final production code at 06fa78657:

• Full op_tests/test_flydsl_pa_decode.py suite: 117 passed on gfx950/MI355X; 113 passed, 4 skipped on gfx942/MI308X (partition caps exceed the device CU count).
• 320 packed D256/page16 configurations passed against independent FP32 attention at atol=rtol=0.005, including 32 graph-replay cases. Covers BF16/FP16 queries, per-tensor/per-token scales, plain/transposed V, QL1–4 at GQA8, windows 0/1/17/257/8193, small/large budgets, single/multi-tile plans, empty/short/tail contexts and sinks. Maximum reference error: 0.0016555.
• 20 additional page16 boundary configurations passed, including unstaged page IDs, a scheduling bound below the real context, multiple KV heads, GQA4/16/32, QL8 at GQA4 and graph replay. Their random V amplitudes are scaled by 1/16 to isolate addressing and scheduling correctness within the existing FP32 tolerance.
• 9 unsupported-layout checks passed on each GPU, covering both V transpose modes at host, schedule, compiler, plan preparation and tuning-geometry entry points.
• A supplementary unscaled random GQA16/QL2/window17 input exceeds the existing FP32 tolerance: maximum error 0.00748747. The previous commit's generic page16 path and the new pipe produce bit-identical outputs on that input; this is a pre-existing FP8 precision limitation.
• Black, Ruff and git diff --check passed. All changed Python files in the final commit match the validated source snapshot.

Tests use AITER_AOT_IMPORT=1 and equivalent Torch reference quantizer formulas in the validation harness to avoid unrelated CK imports. PA kernels and FP32 references execute from source.

Regression comparison

Compare final code with PR head 7072a730e using the same inputs, scales, page tables, scratch and work-plan metadata. BF16 queries/output, D256/page64, both packed V layouts and both scale modes; five alternating sweeps with 50 calls per configuration. Contexts include 8K/128K, QL1/4, windows 0/257/8193 and multiple KV heads.

Device	Configurations	Cold time change	Warm time change
gfx950 / MI355X	48	−1.102% to +0.445%	−1.725% to +0.173%
gfx942 / MI308X	24	−0.185% to +0.499%	−0.304% to +0.156%

All 72 configurations have bit-identical outputs and unchanged packed scheduling/storage traits after excluding removed NHD fields. No material page64 regression was observed.

gfx950 uses HIP event elapsed time around attention+reduce graph replay, including the gap between the kernels; repeated Torch profiler sessions dropped GPU events. gfx942 uses the sum of profiler attention+reduce kernel times. Cold samples use a 2 GiB read-only flush before the measured region. Both exclude compilation, allocation and plan refresh. These comparisons validate page64 stability; page16 speedup over the prior generic path is not measured, and serving throughput is not measured.

[samremes]

#5949 now has per-token scales on the head-256 pipe for packed and NHD V, and 8cf11692f9 put the reducer back to main's. Please rebase onto samremes/pa-decode-longctx for the remaining pieces: transposed V and sliding windows.

[samremes]

The kernel changes that survive a rebase onto 8cf11692f9 are correct and faster: transposed V is 16–20% faster and windows are 9–18% faster, with outputs matching the generic path to within 1.2e-4.

• Rebase drops per-token NHD. One line in PaDecodeSchedule.select (per_token_kv or sliding_window > 0) merges without a conflict and fails our three per-token NHD tests. Keep only the window rejection; per-token NHD is already supported on our branch.

• The new paths have no tests. The third commit deletes every transposed-V, window, and varied-scale pipe case, and the "108 comparisons / 72 FP32 checks" in the description were run at 0c2db7343, before that deletion. A minimal set of six pipe cases plus the host-side staging-bound check passes on the rebased tree and grows the suite from 125 to 134 tests. The six cases cover both scale modes and both V layouts on the new pipe paths:

  • per-tensor, transposed V
  • per-token, transposed V, narrow grid (PER_TOKEN_M1 true), with a graph replay that changes the K and V scales
  • per-token, transposed V, wide grid (PER_TOKEN_M1 false), the only path where the new REUSE_KV_PAGES term changes the kernel
  • per-token, transposed V, window 257 (single-tile plan)
  • per-token, transposed V, window 8193 at budget 1 (multi-tile plan, window-bounded page staging)
  • per-token, plain V, window 1, with V rows outside the window poisoned so a leak into the shared V scale fails

  The host check covers the staging bound for windows 1, 257, and 8193. The rest of the deleted cases repeat paths already on our branch.

• Drop the XCD-helper commit. It edits the reducer, which our branch no longer has, and after the rebase the import is unused.

• Two transposed-V per-token kernels. A grid-width rule tuned for head 128 still compiles two variants, and the new REUSE_KV_PAGES term only exists to restore page-id reuse on one of them. One per-token policy for the pipe would remove that split.

• The scale-before-K reorder is unmeasured. The first tile uses a different gate from the loop, and the prologue barrier already waits for every load, so the claimed wait saving is not visible. Gate both sites on REUSE_KV_PAGES.

• Per-tensor softmax also changed. or D256_PIPE moves the per-tensor pipe onto scale-before-mask. The output differs by at most 1.5e-5 and the time is unchanged, so either keep it and say why, or drop that term.

• The window tile bound is copied three times. The formula is correct. One helper in traits.py should feed the staging bound, the planner, and the reducer.

[fsx950223]

@samremes, thanks for the detailed review and the measurements. Following up on your comment, the current head is de63b3169, which merges samremes/pa-decode-longctx through 8cf11692f9.

• Per-token NHD is preserved. Both NHD checks now reject positive windows without rejecting per-token scales. The upstream LDS sizing and DMA guard are retained. All 125 upstream tests pass on gfx950, including the three per-token NHD cases; gfx942 has 114 passed and 11 skipped.
• The XCD-helper edits no longer survive in the PR diff. The reducer and test file match your branch exactly, and the unused import is gone. The earlier helper commit remains in the merge history, but contributes no remaining reducer changes.
• The coverage gap remains. I removed the additional cases in the earlier cleanup. The current 125-test suite covers your existing D256 paths, but does not add the transposed-V/window coverage you listed. Sixteen supplementary packed layout/scale/QL/window cases passed locally, with graph replay changing queries; they do not replace the proposed in-tree cases or the replay that changes K/V scales. I agree your six-case set plus the staging-bound checks covers the missing paths without restoring the whole matrix. The description now separates current-head validation from the historical 108 output comparisons / 72 FP32 checks and performance measurements.
• The policy split and prefetch gates are still outstanding. PER_TOKEN_M1 still depends on grid width for transposed V, and the first-tile gate still differs from the loop's REUSE_KV_PAGES gate. I have not measured the scale-before-K reorder independently, so the timings should not be attributed to that reorder or to a prologue wait saving.
• Scale-before-mask is retained for the per-tensor pipe. It uses the existing -inf masking and safe exponentiation handling for fully masked MTP/window rows, including zero query scales. I am keeping it for that behavior and consistency across the pipe's scale modes; there is no measured performance benefit to claim.
• The repeated window bound is still outstanding. The merge preserves the correct formula but does not consolidate it into a shared helper.

One supplementary QL3/window=1 varying-scale input exceeds the FP32 tolerance (maximum absolute error 0.01171875). I checked it against both 98789fb87 and your 8cf11692f9: all three outputs are bit-identical. Windows 255/256/8193 also match exactly and pass the tolerance. This does not indicate a merge regression, and the limitation is recorded in the description.