[samremes]

Summary

Adds the NHD page-64 KV cache layout to the gfx950 FP8 head-256 paged-attention decode pipe, and a per-shape default work budget for that pipe. The packed pipe (#6655, with #6617's page-16, transposed-V and sliding-window extensions) and the partition reducer (#6619) are already on main. This PR adds only what main does not have: NHD, the default budget, a small public capability API for downstream callers, and softmax_scale and K/V scale checks.

Motivation

• vLLM and the ASM prefill kernels write K/V as [num_blocks, page, num_kv_heads, head_dim] (NHD), typically with K and V as the two halves of one token row. On main, pa_decode raises NotImplementedError for a 4D key_cache, so a caller has to repack the cache into the packed 5D layout. With NHD, decode reads that cache in place.
• [FlyDSL] Add packed D256 PA decode pipeline for gfx950 #6655 set the prepare_pa_decode_plan cache-miss budget to 2×CU for every shape. For the dense page-64 head-256 pipe shapes, 1×CU is faster (median 0.92–0.99 of the 2×CU time, table below).
• Downstream callers probe private symbols (kernels.pa_decode.traits.d256_page64_pipe, inspect.signature(pa_decode)) to decide whether NHD is available. They need a public API instead.

Changes

• NHD layout. A 4D key_cache selects NHD: K and V are [num_blocks, 64, num_kv_heads, 256] views with a contiguous head axis and arbitrary other strides. K is read with NHD addressing. V is copied straight into LDS with nontemporal buffer_load ... lds into a row-major, XOR-swizzled layout and read back with transposed LDS reads. Two V buffers overlap the next tile's copy with PV. A grid wider than the CU count uses one buffer only when two single-buffered CTAs, including the staged page ids, still fit in 160 KiB of LDS. Hardware out-of-bounds zero-fill covers partial-page tails, which can hold FP8 NaN, so NHD skips the V byte masking that main's mask_unowned_kv applies to the packed layouts (the K/V scale masking still applies).
• NHD folds the softmax scale into exp2 instead of scaling scores before masking, as the packed pipe does. On NHD this is up to 8% faster on the attention kernel at long context, while packed is faster with its current order.
• NHD supports per-tensor and per-token K/V scales and positive sliding windows, including MTP windows (the window support and the K/V transpose-flag selection come from @fsx950223's commits in [FlyDSL] Extend packed D256 PA decode pipe to page16, transposed V and sliding windows #6617).
• Packed caches keep main's kernel code paths: every NHD branch is behind a compile-time nhd_layout flag. With explicit budgets, packed outputs are bit-identical to main (table below).
• Default budget. On a tuning-cache miss, prepare_pa_decode_plan uses 1×CU for the gfx950 FP8 head-256 pipe with page 64 and no sliding window, packed or NHD. Page 16, windowed pipe shapes and every non-pipe path (including all of gfx942 and gfx1250) keep 2×CU. The tuning search also gains the CU count as a candidate.
• Tuning. prepare_pa_decode_plan reads KV heads and page size from NHD caches and adds layout (packed or nhd) to the tuning key. The tuning CLI has --layout nhd.
• NHD contract. gfx950, FP8 KV, head 256, page 64, and a uniform query length with at most 32 query rows (query_length * query_group_size, so query length 1 to 4 at GQA 8). Anything else raises NotImplementedError, including page 16, which the packed pipe supports but the one-page-per-warp V copy does not. K and V need the same shape, 16-byte aligned base pointers, block/token/head strides that are multiples of 16 bytes, contiguous channels, and non-overlapping token rows; violations raise ValueError. Strides and the in-page span must fit int32. Wide addressing is chosen from the highest byte either view can reach and also widens the KV-head term, so caches with KV heads outermost work past 2 GiB in total, up to 2 GiB per head plane.
• Integration API, exported from aiter.ops.flydsl.pa_decode and aiter.ops.flydsl:
  • pa_decode_nhd_max_query_length(device, *, kv_dtype, head_dim, block_size, query_group_size) -> int: the longest uniform query length pa_decode takes for an NHD cache, or 0 when NHD cannot serve the configuration. pa_decode raises for any longer query.
  • pa_decode_default_workgroup_budget(device, *, kv_dtype, head_dim, block_size, query_length, query_group_size, sliding_window=0) -> int: the budget prepare_pa_decode_plan uses on a tuning-cache miss. Pass it to plan_pa_decode; that function's own default stays 2×CU because it does not see the cache shape.
  • Both read the architecture from get_gfx_runtime(), as pa_decode does, and compare FP8 against dtypes.fp8; the device supplies the CU count.
  • There is no layout keyword: pa_decode derives trans_k from key_cache.dim(). max_context_length=None specializes on the block-table width; a value must cover every context that will run, including later graph replays.
• gfx1250. The NHD strides are keyword-only PaDecodeContext arguments and are not passed to the gfx1250 launcher, so the gfx1250 context call and launcher keep main's argument lists.
• Lazy exports. from aiter.ops.flydsl import plan_pa_decode, pa_decode (or any other PA decode export first, including prepare_pa_decode_plan on main) bound pa_decode to the submodule, because importing it binds the module on the package under the export's name. The package __getattr__ now drops such shadowing bindings after each lazy load.
• softmax_scale. pa_decode raises ValueError unless it is finite and positive. A negative scale gave masked tokens weight on the paths that scale after masking (NHD and the generic per-tensor path).
• K/V scales. For the same reason, FP8 key_scale/value_scale given as host scalars must be finite and positive, or pa_decode raises ValueError. The check sits in the shared scale normalization, so it covers packed and NHD. Device tensors are not read back (no host sync, graph-capture safe); the docstring states they must be positive.

Performance (if applicable)

MI355X / gfx950, 256 CUs, FlyDSL 0.5.0, PyTorch 2.14 ROCm 10.1. Re-measured at 837c255a90 (merges main 0a17569473, which includes #7388's dtype rework). Attention plus reduce under CUDA graph replay (HIP events, so the gap between the two kernels is included), 2 GiB read-only flush before every replay, median of three runs. 8 query heads, 1 KV head, head 256, FP8 KV.

NHD vs packed plain V. Before is the packed layout a caller has to repack into on main; after is NHD read in place. Both hold the same quantized data and plan. Budgets are this PR's defaults: 1×CU for the dense rows, 2×CU for the windowed row. Per-token outputs are bit-identical; per-tensor outputs agree within 8e-6 because of the different scale order.

Hardware	Workload/configuration	Before: packed plain V (µs)	After: NHD (µs)	Change
MI355X	per-tensor, batch 1, 128K, Q1	25.2	25.7	+2.1%
MI355X	per-tensor, batch 4, 128K, Q1	56.5	54.3	−3.8%
MI355X	per-tensor, batch 1, 1M, Q1	98.9	98.1	−0.8%
MI355X	per-tensor, batch 4, 128K, Q4	60.2	60.3	+0.1%
MI355X	per-token, batch 1, 128K, Q1	26.0	26.2	+0.8%
MI355X	per-token, batch 4, 128K, Q4	62.6	60.5	−3.4%
MI355X	per-token, batch 32, 8K, Q1	34.8	33.9	−2.6%
MI355X	per-token, batch 4, 128K, Q4, window 8193	17.9	19.8	+10.7%

NHD is within 2.1% of packed plain V on the dense shapes and faster at batch 4 and above, with no repacking of the cache. The windowed shape costs about 11%. NHD times match the d7a7d90923 run within 0.5 µs; that is inside its round-to-round spread.

Default budget. Not re-run at this head: the default-selection logic is unchanged, and the packed and NHD kernels it chooses between time as before (tables above and below). From d7a7d90923, FlyDSL 0.4.0: 1×CU / 2×CU time over every sweep point where the two budgets give different plans (batch 2–256, contexts 8K–1M; batch 1 plans are identical), both budgets alternated in one process for three rounds. Measured at d7a7d90923; every point uses the pipe, and 1×CU and 2×CU outputs agree within 8e-4. Four points on 1×CU families exceed 1.01: plain V at batch 32 and 8K (1.012), and GQA 4 at batch 4 and 32 at 8K (1.015, 1.027); every family's median is below 1.0 for the 1×CU defaults. The two unchanged 2×CU families reach 1.024 and 1.038 and keep 2×CU.

Hardware	Workload/configuration	Points	Before: 2×CU	After: 1×CU / 2×CU, median (min–max)	Default
MI355X	packed per-tensor, plain V, page 64, Q1	18	1.000	0.988 (0.971–1.012)	1×CU
MI355X	packed per-token, transposed V, page 64, Q4	18	1.000	0.960 (0.885–0.996)	1×CU
MI355X	packed per-token, transposed V, page 64, Q1, 8 KV heads, GQA 4	10	1.000	0.992 (0.969–1.027)	1×CU
MI355X	NHD per-tensor, Q1	18	1.000	0.927 (0.896–1.009)	1×CU
MI355X	NHD per-token, Q4	18	1.000	0.919 (0.844–1.006)	1×CU
MI355X	packed per-token, transposed V, page 16, Q2	18	1.000	1.002 (0.985–1.038)	2×CU (unchanged)
MI355X	packed per-tensor, transposed V, page 64, Q1, window 4096	18	1.000	1.001 (0.983–1.024)	2×CU (unchanged)

Packed layouts vs main, explicit budgets. Batch 4, 128K, 13 configurations, main 0a17569473 and this PR alternated in separate processes for three rounds. Every output is bit-identical to main; PR / main time is 0.990–1.006 (geomean 1.000).

Hardware	Workload/configuration	Before: main (µs)	After: PR (µs)	Change
MI355X	page 64, plain V, per-tensor, Q1, budget 256	56.52	56.40	−0.2%
MI355X	page 64, plain V, per-tensor, Q1, budget 512	56.60	56.96	+0.6%
MI355X	page 64, plain V, per-token, Q1, budget 256	56.92	57.08	+0.3%
MI355X	page 64, transposed V, per-token, Q4, budget 256	60.00	60.12	+0.2%
MI355X	page 64, transposed V, per-token, Q4, budget 512	65.60	65.56	−0.1%
MI355X	page 64, transposed V, per-tensor, Q1, window 4096, budget 256	15.32	15.16	−1.0%
MI355X	page 64, transposed V, per-tensor, Q1, window 4096, budget 512	15.32	15.28	−0.3%
MI355X	page 64, transposed V, per-token, Q4, window 8193, budget 512	18.20	18.16	−0.2%
MI355X	page 64, plain V, per-token, Q4, window 8193, budget 256	18.00	17.88	−0.7%
MI355X	page 16, transposed V, per-token, Q2, budget 256	55.76	55.92	+0.3%
MI355X	page 16, transposed V, per-token, Q2, budget 512	57.76	58.08	+0.6%
MI355X	page 16, plain V, per-tensor, Q1, budget 512	57.04	57.08	+0.1%
MI355X	page 16, transposed V, per-token, Q1, window 4096, budget 512	15.92	15.96	+0.3%

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X
• Tested on MI355X (gfx950)

Documentation

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

Docstrings: pa_decode (NHD selection, softmax_scale, K/V scale positivity), the two capability functions, and prepare_pa_decode_plan (NHD geometry, layout key, cache-miss default). User guide: aiter/ops/flydsl/configs/pa/README.md documents the new cache-miss default and --layout nhd. There is no separate performance guide for this op.

Dependencies

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

Breaking Changes

• pa_decode now raises ValueError for a host-scalar FP8 key_scale/value_scale that is zero, negative, NaN or infinite. Migration: pass the positive dequantization scale; device-tensor scales are not checked and must already be positive.
• pa_decode now raises ValueError for a softmax_scale that is zero, negative, NaN or infinite. Migration: pass the positive attention scale (for example head_dim ** -0.5); a non-positive scale gave wrong results or NaN on the paths that scale after masking.
• The PA decode tuning key gains layout. Budgets tuned and saved before this PR no longer match, so those shapes take the default budget until re-tuned with python -m aiter.ops.flydsl.pa_decode_tuning. No tuned JSON ships in the repository.
• On a tuning-cache miss, dense page-64 head-256 FP8 pipe shapes on gfx950 plan with 1×CU instead of 2×CU, so plan.capacity can be smaller. Callers that size scratch from plan.capacity, as pa_decode requires, are unaffected; callers passing an explicit workgroup_budget see no change.
• A 4D key_cache used to raise NotImplementedError; on gfx950 it now runs the NHD path. gfx942 and gfx1250 still raise.

Made with Cursor

PR Self-Review

Reviewed revision

a240ddc

Behavior and intent

pa_decode on gfx950 accepts a 4D NHD [num_blocks, 64, num_kv_heads, 256] FP8 K/V cache for head 256 and page 64, so vLLM's cache layout no longer has to be repacked into the packed 5D layout before decode; on main a 4D cache raises NotImplementedError. Packed caches keep main's kernel code, behind a compile-time nhd_layout flag, and produce bit-identical output to main. On a tuning-cache miss, dense page-64 head-256 FP8 pipe shapes now plan with one CTA per CU instead of two, because #6655's uniform 2×CU default is slower for those shapes (median 0.92–0.99 of the 2×CU time). pa_decode_nhd_max_query_length and pa_decode_default_workgroup_budget replace downstream probes of private symbols. pa_decode rejects a non-finite or non-positive softmax_scale, and non-finite or non-positive host-scalar FP8 K/V scales, which gave masked tokens weight on paths that scale after masking. prepare_pa_decode_plan raises ValueError for a cache with zero KV heads instead of ZeroDivisionError. NHD skips main's mask_unowned_kv V byte masking because its V DMA already zero-fills past the context; K/V scale masking still applies. The NHD strides do not change the gfx1250 context call or launcher. Importing a sibling export before pa_decode from aiter.ops.flydsl no longer returns the submodule.

Validation and unverified scope

On MI355X (gfx950, 256 CUs), FlyDSL 0.5.0 (main's pin), PyTorch 2.14 ROCm 10.1, at 837c255a90 (which merges main 0a17569473):

• python -m pytest op_tests/test_flydsl_pa_decode.py -q: 408 passed, 38 skipped at 837c255a90, fresh FlyDSL cache. Later commits only fold the K/V scale cases into the existing softmax_scale test and switch it to packed caches, so it reaches scale validation on every arch (438 tests collected: 400 run, 38 skipped; the test passes with pa_decode's arch patched to gfx942, gfx950 and gfx1250). Every skip is an architecture or support gate: main's MXFP4 and gfx1250-only cases, and the gfx942-only NHD rejection variant.
• The K/V scale branch of test_pa_decode_softmax_scale_rejected fails on all 4 cases without the check and passes with it.
• black --check (26.3.0) and ruff check (0.16.0) on the changed files: passed.
• Performance: the NHD-vs-packed and packed-vs-main tables were re-measured at this head (3 rounds each; main and PR alternated in separate processes). All 13 packed outputs are bit-identical to main, and PR / main time is 0.990–1.006. The default-budget sweep is from d7a7d90923 and was not re-run.
• Carried forward from an earlier revision because the code it covers has not changed: the zero-KV-head ValueError check and the tuning CLI smoke run (batch 4, 32K, page 64), where the search also picks the CU count.

Not run or unverified:

• gfx942 / MI300X was not run at this revision. NHD and the 1×CU default are gfx950-only there, so gfx942 keeps 2×CU and the generic schedule. vedenev-amd ran an earlier revision (cdab845a) on MI300X: 124 passed, 5 skipped.
• gfx1250 was not run; no device was available. The fix that keeps its argument lists unchanged was checked against the code only.
• MI250X (gfx90a) is not a pa_decode target; it supports gfx942, gfx950 and gfx1250 only.
• The default-budget sweep was not re-run at this head.
• vLLM dispatch and serving replay were not run; this PR only adds the AITER kernel and API.

PR-specific question

No tuned configuration file changes; the only file under configs/ is aiter/ops/flydsl/configs/pa/README.md, which is documentation. Two configuration behaviors change. First, the _pa_decode_autotuner key in pa_decode_tuning.py (num_seqs, num_kv_heads, query_group_size, head_dim, block_size, query_length, query_dtype, kv_dtype, per_token_kv, trans_v, sink_dtype, sliding_window, max_partitions, num_cu, implementation) gains layout, which is packed for 5D and nhd for 4D K caches. It is consumed only by prepare_pa_decode_plan lookups in configs/pa/<arch>/_pa_decode_autotuner.json, and no such file ships in the repository. Second, the cache-miss default changes only for keys where d256_page64_pipe holds and there is no window: arch gfx950, kv_dtype dtypes.fp8, head_dim 256, block_size 64, query_length * query_group_size <= 32, sliding_window <= 0. Every other key keeps 2×CU. Evidence that unrelated shapes keep their selection: test_pa_decode_default_workgroup_budget asserts 2×CU for window 4096, page 16, head 128, 64 query rows (query length 4, GQA 16) and BF16 KV, and 1×CU only on gfx950 for the pipe shapes. test_prepare_pa_decode_plan_default_budget checks the resulting plans for page 64 (packed and NHD), page 16, a window, head 128 and 64 query rows; forcing 2×CU everywhere fails only the two page-64 cases. With explicit budgets, all 13 packed configurations in the table above give bit-identical output to main.

Changed op tests

op_tests/test_flydsl_pa_decode.py is the canonical op test for FlyDSL PA decode, so the new coverage lives there rather than in a new file. None of it is covered by an existing case, because main has no NHD path, no per-shape default and no capability API.

• nhd-d256-page64* cases in the main decode test: NHD output against the FP32 reference for per-tensor and per-token scales; query lengths 1, 3 and 4 with a graph capture at 4; window 257 and window 8193 at budget 1 (multi-tile windowed tasks); NaN-poisoned K and V partial-page tails; two KV heads in one token row; and a head-major pool past 2 GiB with three KV heads, which exercises wide addressing and the one-V-buffer path reused across tiles. Each also pins the selected specialization: wide addressing, and the V buffer count where the case sets one.
• test_pa_decode_nhd_rejected: alignment, stride and overlap rejections, BF16, 40 query rows, page 16 and page 128, plus the gfx942-only rejection variant. Removing the page-16 check fails the page-16 case.
• test_pa_decode_nhd_max_query_length and test_pa_decode_default_workgroup_budget: the public capability API on and off gfx950.
• test_prepare_pa_decode_plan_default_budget: cache-miss plans for page 64 (packed and NHD), page 16, a window, head 128 and 64 query rows.
• test_pa_decode_softmax_scale_rejected: zero, negative, NaN and infinite values, passed as softmax_scale and as a host-scalar key_scale, raise ValueError before any kernel compiles.
• test_pa_decode_export_import_order: imports each other PA decode export before pa_decode in a fresh interpreter and checks pa_decode is callable, and skips when FlyDSL is not installed; all five cases fail without the lazy-loader fix.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

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

Extended tests (opt-in via labels):

Label	Tests
ci: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; main branch always runs both MI35X and MI300X
multigpu	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: 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 5949 --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.

[Copilot]

Copilot review overview

🔵 Needs a closer look

The asynchronous GPU pipeline needs human verification, and gfx942 validation remains pending for the shared reducer change.

Review effort: Balanced
Findings: 1

Open (1)

• Include staged page-ID LDS footprint in buffer selection · New

What changed in this PR

Optimizes FlyDSL paged-attention decode for gfx950 head-256 workloads and adds direct NHD cache support without repacking.

Changes:

• Adds streamed K loads, staged page IDs, and prefetched V.
• Supports NHD page-64 caches with alignment and stride checks.
• Extends split scheduling, head-256 reduction, and correctness tests.

File	Description
op_tests/​test_flydsl_pa_decode.py	Adds NHD, scheduling, reducer, and contract tests.
aiter/​ops/​flydsl/​pa_decode.py	Validates NHD inputs and configures scheduling.
aiter/​ops/​flydsl/​kernels/​pa_decode_reduce.py	Extends shared-memory reduction to static head-256 workloads.
aiter/​ops/​flydsl/​kernels/​pa_decode_kernel.py	Implements optimized packed and NHD data pipelines.

---

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

[vedenev-amd]

Claude: Ran op_tests/test_flydsl_pa_decode.py from cdab845 on gfx942 (MI300X): 124 passed, 5 skipped (NHD, gfx950-only), 0 failed. All 106 cases shared with main pass on both. The new gfx942-relevant cases pass, including d256-m1-page64-np130 (reducer above 64 partitions) and nhd_rejected[arch].

[vedenev-amd]

Claude generated:
Ran op_tests/test_flydsl_pa_decode.py at cdab845 on gfx950: 128 passed, 1 skipped. I also ran extra NHD cases (unstaged page ids, single/double buffer, 4 KV heads, 75 repeated runs) and NHD under a CUDA graph with plan refresh: all match. Perf vs main (cold cache): faster on all measured shapes, head 256 and head 128, batch 1–256. A few comments inline.

[samremes]

Thanks for the measurements. Pushed 18957ba3c4.

The NHD V buffer count now comes from the LDS footprint. One buffer is used only when the grid is wider than the CU count and two single-buffered CTAs fit in 160 KiB. At 4096 staged ids a single-buffered CTA is 86992 bytes, so that shape keeps two buffers. Where two CTAs do fit, one buffer stays, and it is still the faster choice. I measured the same direction you did at batch 128, 512K, budget 2×CU: 5187 µs with one buffer and 5072 µs with two.

The CU-count default stays. At that shape it is 5064 µs packed and 5103 µs NHD, against 5073 µs packed and 5072 µs NHD at 2×CU after the buffer change. Staging is under 1% there because the kernel is near the memory-bandwidth limit. The warning now fires only when the block-table width set the bound, which is where passing max_context_length helps (about 5–11% at batch 32–128 with 8K–32K tokens). That point is in the description.

max_context_length stays the caller's contract, including for graph replay. A value below a real context reads past the staged page ids: NHD faults, and packed can return a wrong result without faulting. I did not add a device-side check or a clamp. The docstring now states that every context that will run, including later replays, must be at most that value.

barrier_keep_dma is now s_waitcnt lgkmcnt(0) plus s_barrier. Those match the inline asm, including the memory ordering. fx.rocdl.barrier() waits for the in-flight V copy and is slower, so I did not use it.

The gfx942 run you posted (124 passed, 5 skipped at cdab845a) is noted in the test plan. The skips are the gfx950-only NHD cases.

[Copilot]

🟡 Changes recommended

Padded reducer iterations can turn valid infinite outputs into NaNs through zero-times-infinity arithmetic.

2 open findings

• Prevent padded zero-times-infinity values from poisoning accumulation · New
• Include staged page-ID LDS footprint in buffer selection

🧠 Review effort: Balanced

[samremes]

Addressed the latest Copilot review.

Zero-weight padded slots turn an inf partial into NaN: fixed in 901a0ec. I reproduced it. With ±inf in the last active partition's partial output, 18957ba returned NaN on both D128 and D256. Main's guarded branch returns ±inf there. The LDS reducer now reads partial logits through the per-sequence bounded buffer descriptor that the NP≤64 planned reducer already uses. Rows at or past the sequence's partition count are out of bounds and load as 0, so padded slots add 0 * 0 and active rows keep IEEE behavior. The address clamp is gone. The path needs packed contiguous partial rows, which pa_decode already requires of temporary_output. A direct reducer launch with other strides now raises ValueError instead of running. The ISA for the compact D128 and striped reducers is unchanged.

Reducer time on gfx950 (interleaved A/B, medians, before → after):

• D128 reducer-only, b1 NP128/NP200/NP256: 4.73/5.04/5.29 → 4.56/4.81/5.18 µs; b2/b4 NP256: 4.57/4.55 → 4.50/4.59 µs.
• D256 1M packed, b1/b2/b4: 5.99/4.97/4.66 → 5.97/4.72/4.49 µs.

I also tried selecting the loaded value (+0.3–0.7 µs) and v_mul_legacy_f32 (up to +0.3 µs, because it loses the packed v_pk_mul_f32 pairing). Both were slower, so I did not use them.

Staged page-ID LDS in the NHD buffer count: already addressed by 18957ba, no new change. traits.py picks one V buffer only when the grid is wider than the CU count and 2 * _nhd_single_cta_bytes(stage_pages) <= 160 KiB. That footprint includes the staged page IDs, and PaDecodeTraits.create asserts that the laid-out single-buffer footprint matches it.

op_tests/test_flydsl_pa_decode.py: 113 passed on gfx950.

[Copilot]

🔵 Needs a closer look

NHD validation permits overlapping head or block rows despite the documented cache-layout contract.

1 open finding

• Include staged page-ID LDS footprint in buffer selection

1 resolved since last review

• Prevent padded zero-times-infinity values from poisoning accumulation

Previously missed (1)

In code that hasn't changed since last review

Validate non-overlap across all token, head, and block strides

aiter/​ops/​flydsl/​pa_decode.py:402

The overlap check only covers the token stride. A multi-head view with stride(2) == 16 (or a block stride that aliases another block) passes every check here even though its 256-byte token/head rows overlap, contradicting the NHD contract and allowing different logical cache rows to address the same bytes. Please validate non-overlap across every stepped block/token/head axis (while retaining support for both token-major and heads-outermost layouts), and add malformed head/block-stride cases to the contract test.

🧠 Review effort: Balanced

[samremes]

The two items still listed as open are settled.

Staged page-ID LDS in the NHD buffer count: already addressed in 18957ba. One V buffer is used only when the grid is wider than the CU count and two single-buffered CTAs still fit in 160 KiB. That footprint includes the staged page ids, and the layout asserts that the calculated size matches.

NHD overlap on head and block strides: won't fix. The token-row check is the one the DMA and partial-page zero-fill depend on. Block and head axes are read-only, so aliasing them is a valid view: addresses stay inside the allocation, and the kernel matches a reference on that same view. Rejecting those strides would also reject a stride-0 block broadcast without fixing a bug.

[Copilot]

🟡 Changes recommended

Runtime architecture selection bypasses the shared helper, and the single-buffer NHD path lacks automated correctness coverage.

3 open findings

• Add gfx950 coverage for one-buffer wide-grid partial-page tails · New
• Use shared runtime architecture contract for plan selection · New
• Include staged page-ID LDS footprint in buffer selection

🧠 Review effort: Balanced

[Copilot]

@copilot resolve the merge conflicts in this pull request

Resolved in merge commit 9ba54f5, bringing in latest main while preserving NHD support and main’s MXFP4/gfx1250 changes. Syntax, whitespace, and secret checks passed. GPU tests could not run here: pytest, PyTorch, FlyDSL, and a ROCm compute device are unavailable.

[Copilot]

🔵 Needs a closer look

The compiler-sensitive NHD LDS-DMA synchronization and wait-counter changes need target-GPU validation with pinned FlyDSL 0.4.1, which the supplied 0.4.0 results and static-only review do not provide.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Group-size division occurs before positive KV-head validation, causing zero-head caches to raise ZeroDivisionError instead of ValueError and bypass callers’ geometry-error handling.

1 open issue

• Validate KV head count before computing group size · New

_{Review effort: Balanced}

[fsx950223]

No blocking defect found at 63658e8. Targeted gfx950 correctness tests passed with the pinned FlyDSL 0.4.2, including NHD tails, wide addressing, both V-buffer modes, MTP graph replay, default budgets, and packed/generic/MXFP4 regressions. Approval remains pending gfx942 regression validation of the shared launcher and plan-preparation changes; host authorization prevented testing there.

Validation:

• PASS on local: git -C checkout rev-parse HEAD; git -C source merge-base HEAD f705d6b7a0637b411a3c177b2fc523f9ce31d0ca
  Verified exact head 63658e8 and merge base f705d6b. Complete diff preserved in exact.diff.
• PASS on local: python3 -m compileall -q checkout/aiter/ops/flydsl checkout/op_tests/test_flydsl_pa_decode.py; git -C checkout diff --check f705d6b7a0637b411a3c177b2fc523f9ce31d0ca HEAD
  Both checks exited 0.
• PASS on gfx950 / fsx950223_pr_review_gfx950: flock -w 60 /workspace/gpu.lock timeout 600 python3 -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'nhd or default_budget or default_workgroup_budget or softmax_scale_rejected or export_import_order or d256-page64 or d256-page16'
  FlyDSL 0.4.0: exit 0; 50 passed, 1 skipped, 384 deselected in 45.07s. Exact checkout and gfx950 verified inside container. Log: gfx950-tests-retry.log.
• PASS on gfx950 / fsx950223_pr_review_gfx950: flock -w 60 /workspace/gpu.lock timeout 600 /workspace/codex-pr-review-5949-20261106/venv/bin/python -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'nhd or default_budget or default_workgroup_budget or softmax_scale_rejected or export_import_order or d256-page64 or d256-page16'
  Pinned FlyDSL 0.4.2: exit 0; 50 passed, 1 skipped, 384 deselected in 44.97s. The skipped case is the rejection variant for non-gfx950 hardware. Log: gfx950-tests-pinned-retry.log.
• PASS on gfx950 / fsx950223_pr_review_gfx950: flock -w 60 /workspace/gpu.lock timeout 600 /workspace/codex-pr-review-5949-20261106/venv/bin/python -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'unowned_kv_rows or mxfp4-g16-bf16-q or head192-np80 or d256-unstaged'
  Exit 0; 77 passed, 16 unsupported MXFP4 cases skipped, 342 deselected in 119.67s. Log: gfx950-tests-shared.log.
• SKIP on gfx942 / hjbog-srdc-39.amd.com: ssh -o BatchMode=yes -o ConnectTimeout=10 -o RequestTTY=no -o ControlMaster=no -o ControlPath=none hjbog-srdc-39.amd.com 'docker exec sixifang_dev python3 -c "import torch; print(torch.cuda.get_device_properties(0)); print(torch.cuda.get_device_properties(0).gcnArchName)"'
  Conductor authorization check ended with connection closed. Architecture verification and GPU tests could not run. Log: gfx942-env.log.

Remaining validation:

• gfx942 regression validation remains unavailable and should be retried when host authorization is restored.
• gfx1250 hardware was unavailable; its context and launcher compatibility was inspected statically.
• Performance tables and tuning-search performance were not independently reproduced.
• The reviewed SHA had a successful Copilot review check but no GPU CI results.
• Detailed commands, isolated cache settings, setup failures, and test results are preserved in the task directory.

Automated Codex review of commit 63658e814b9b.

[Copilot]

🔵 Needs a closer look

Scheduler-sensitive LDS DMA and explicit wait ordering lack GPU validation with pinned FlyDSL 0.4.2, while reported runs used 0.4.0 and did not independently exercise the preserved gfx1250 launcher ABI.

No open issues 🎉

1 resolved since last review

• Validate KV head count before computing group size

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The NHD DMA pipeline’s alias metadata and explicit wait ordering need target-GPU validation with pinned FlyDSL 0.4.2, because this review is static-only and supplied GPU results used 0.4.0.

No open issues 🎉

_{Review effort: Balanced}

[samremes]

@fsx950223 Thanks. I don't have access to a working gfx942 machine either, so gfx942 is still pending.

Could CI cover it? On pull requests, Aiter Test runs only on MI35X; the MI300X shards run only on main, so it would check gfx942 after merge. On gfx942 this PR keeps the generic schedule and the 2×CU default. The changes there are host-side: plan preparation and the launcher arguments.

Since your review, 23d21e5a6b fixes the Copilot finding (zero KV heads now raise ValueError). The full op_tests/test_flydsl_pa_decode.py passes on gfx950 at that commit: 397 passed, 38 skipped.

[fsx950223]

No blocking defects found at 2241feb. Reviewed the complete diff, callers, tests, prior feedback, and CI evidence. Targeted correctness tests passed on gfx950 and gfx942 with the pinned FlyDSL 0.4.2, including NHD addressing, partial-page tails, buffer selection, windows, graph replay, and host API changes. The failed CI shard reports a name error in an unchanged MoE kernel outside this diff.

Validation:

• PASS on local: git -C source rev-parse HEAD; git -C source diff --check f705d6b7 HEAD
  HEAD matches the supplied SHA exactly; whitespace checks passed. All 10 changed Python files also parsed successfully.
• PASS on dzm-mi355-gpu34-dev / allocation 9275 / fsx950223_pr_review_gfx950: flock -w 60 /workspace/gpu.lock timeout 600 ../venv042/bin/python -m pytest op_tests/test_flydsl_pa_decode.py -q -rA -k 'nhd or default_workgroup_budget or default_budget or softmax_scale or export_import_order or d256_page64 or d256-page64'
  Exit 0: 50 passed, 1 skipped, 384 deselected. Verified gfx950, FlyDSL 0.4.2, and imports from the isolated exact-SHA checkout. The skip is the gfx942-only NHD rejection variant. Log: gfx950-pinned-tests.log.
• PASS on hjbog-srdc-39.amd.com / sixifang_dev: flock -w 60 ../gpu.lock timeout 600 ../venv042/bin/python -m pytest op_tests/test_flydsl_pa_decode.py -q -rA -k 'nhd or default_workgroup_budget or default_budget or softmax_scale or export_import_order or d256_page64 or d256-page64'
  Exit 0: 35 passed, 16 skipped, 384 deselected. Verified gfx942, FlyDSL 0.4.2, and checkout imports. Skips cover NHD execution, two partition counts above the device's 80-CU limit, and inapplicable rejection variants. Log: gfx942-pinned-tests.log.
• PASS on dzm-mi355-gpu34-dev / fsx950223_pr_review_gfx950: flock -w 120 /workspace/gpu.lock timeout 180 ../venv042/bin/python -m aiter.ops.flydsl.pa_decode_tuning --batch-size 1 --num-query-heads 8 --num-kv-heads 1 --head-dim 256 --query-length 1 --context-length 257 --block-size 64 --layout nhd --max-partitions 4 --output-dir ../tuning-smoke
  Exit 0: candidate correctness validation and tuning completed; configuration saved with plan capacity 4. Log: gfx950-tuning-smoke.log.

Remaining validation:

• gfx1250 execution was not tested; its preserved launcher and context argument lists were checked statically.
• Performance improvements and vLLM serving integration were not independently measured.
• Validation was targeted rather than the full op suite. Detailed commands, environment evidence, results, and CI logs are preserved in the task directory.

Automated Codex review of commit 2241feb7037a.

[Copilot]

🟡 Changes recommended

The NHD per-tensor path still accepts negative K scales that can turn masked scores into the row maximum and underflow a live token’s probability to zero, requiring a fix and regression coverage.

1 open issue

• Negative per-tensor K scale corrupts NHD masked softmax · New

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The NHD DMA path and manual wait-counter ordering require exact-head GPU validation with pinned FlyDSL 0.4.2, while reported results cover earlier revisions and toolchains.

1 open issue

• Negative per-tensor K scale corrupts NHD masked softmax

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The new NHD DMA and wait-counter pipeline needs current-head correctness and latency validation on gfx950 with pinned FlyDSL 0.5.0, as supplied GPU results cover older revisions on 0.4.0.

1 open issue

• Negative per-tensor K scale corrupts NHD masked softmax

_{Review effort: Balanced}

[samremes]

Update (837c255a90): merged main again (including #7388's dtype rework and the FlyDSL 0.5.0 pin) and addressed Copilot's scale comment: pa_decode now rejects non-finite or non-positive host-scalar FP8 K/V scales, with a regression test. On MI355X: 408 passed, 38 skipped; packed outputs are bit-identical to main at 0.990–1.006× main's time. Perf tables in the description are re-measured.

@coderfeli @fsx950223 could you take another look?

[Copilot]

🟡 Changes recommended

The NHD scale cases fail on supported gfx942 and gfx1250 architectures before reaching the intended key_scale validation, so those variants need to be skipped outside gfx950 while retaining packed coverage.

2 open issues

• Skip NHD scale cases on architectures without NHD pipe support · New
• Negative per-tensor K scale corrupts NHD masked softmax

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The scale-validation test selects unsupported NHD caches on gfx942 and gfx1250, causing all four parameterizations to raise NotImplementedError before reaching the expected ValueError.

2 open issues

• Use valid packed K/V views for scale-validation tests · New
• Negative per-tensor K scale corrupts NHD masked softmax

1 resolved since last review

• Skip NHD scale cases on architectures without NHD pipe support

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The NHD path’s explicit DMA/LDS ordering, alias metadata, and wait-count pipelining need kernel-maintainer confirmation because target-GPU correctness and performance could not be independently validated.

No open issues 🎉

2 resolved since last review

• Use valid packed K/V views for scale-validation tests
• Negative per-tensor K scale corrupts NHD masked softmax

_{Review effort: Balanced}

[fsx950223]

Reviewed exact head 92d4293. No blocking defect found. Targeted tests passed on gfx950 and gfx942 with pinned FlyDSL 0.5.0, covering NHD addressing, partial-page tails, both V-buffer schedules, sliding windows, graph replay, capability APIs, default budgets, and scale validation. The failed CI shard concerns an unchanged MoE test.

Validation:

• PASS on local: git -C exact-source rev-parse HEAD; AST-parse all changed Python files
  HEAD matches the supplied SHA; all 10 changed Python files parsed successfully.
• PASS on dzm-mi355-gpu34-dev / gfx950: bash final950.sh: timeout 900 python3 -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'nhd or default_workgroup_budget or default_budget or softmax_scale_rejected or export_import_order or d256-page64' -rA
  Exit 0: 50 passed, 1 architecture-gated skip, 387 deselected. Verified gfx950, exact SHA, checkout imports, and FlyDSL 0.5.0. Used isolated caches and /workspace/gpu.lock. Log: gfx950-final-tests.log.
• PASS on hjbog-srdc-39.amd.com / gfx942: bash final942.sh: timeout 600 python3 -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'nhd or default_workgroup_budget or default_budget or softmax_scale_rejected or export_import_order or d256-page64' -rA
  Exit 0: 35 passed, 16 architecture-gated skips, 387 deselected. Verified gfx942, exact SHA, checkout imports, and FlyDSL 0.5.0. Log: gfx942-final-tests.log.
• PASS on dzm-mi355-gpu34-dev / gfx950: bash base950-final.sh: timeout 600 python3 -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'd256-page64' -rA
  Exit 0: 5 passed against supplied base 0a17569, with separate checkout and caches. Matching packed cases also passed at head. Log: gfx950-base-final-tests.log.
• PASS on dzm-mi355-gpu34-dev / gfx950: bash tune950.sh: timeout 300 python3 -m aiter.ops.flydsl.pa_decode_tuning --layout nhd --batch-size 4 --context-length 8192 --head-dim 256 --num-query-heads 8 --num-kv-heads 1 --block-size 64 --query-length 1 --output-dir /workspace/codex-pr-review-5949-20261112/tune-results
  Exit 0: validated and timed four candidates, including budget 256; saved the selected budget 512 under the task directory. Log: gfx950-final-tune.log.
• PASS on dzm-mi355-gpu34-dev / gfx950: bash repro950.sh: timeout 120 python3 /workspace/codex-pr-review-5949-20261112/repro-small-scale.py
  Exit 0: one-token attention returned the expected value 1.0 for NHD and packed layouts with positive key scales 1.0, 1e-20, and 1e-30. Log: gfx950-small-scale.log.

Remaining validation:

• No gfx1250 device was available; its preserved constructor and launcher interfaces were checked statically.
• The full performance sweep and downstream vLLM serving integration were not independently rerun.
• Initial harness setup failures were corrected before successful validation; setup and final test logs remain in the task directory.

Automated Codex review of commit 92d429387d01.

[Copilot]

🔵 Needs a closer look

The NHD pipeline’s explicit DMA wait-count scheduling and LDS aliasing require specialist runtime validation, with gfx950 correctness and performance and gfx1250 compatibility still unverified.

No open issues 🎉

_{Review effort: Balanced}

[fsx950223]

Reviewed exact head a240ddc. No concrete blocking defect found. Targeted gfx942 tests passed, covering shared launcher changes, packed decode, capability APIs, default-plan preparation, lazy exports, and scalar-scale validation. Approval remains pending targeted gfx950 validation of NHD addressing, partial-page tails, and both DMA buffer schedules: all eight configured gfx950 devices were occupied. The failed CI shard concerns an unchanged FlyDSL HGEMM graph-replay test.

Validation:

• PASS on local: git -C checkout rev-parse HEAD; AST-parse all changed Python files
  HEAD matches the supplied SHA; the isolated checkout is clean and all 10 changed Python files parsed successfully. Log: local-validation.log.
• PASS on hjbog-srdc-39.amd.com / sixifang_dev / gfx942: bash remote942.sh: flock -w 60 /tmp/codex-pr-review-5949-20261113/gpu.lock timeout 600 python3 -m pytest op_tests/test_flydsl_pa_decode.py -q -k 'nhd or default_workgroup_budget or default_budget or softmax_scale_rejected or export_import_order or d256-page64' -rA
  Exit 0: 35 passed, 16 skipped, 387 deselected in 35.07 seconds. Verified exact head SHA, checkout imports, FlyDSL 0.5.0, and gfx942 inside the container. Used idle GPU 0 and isolated source and compilation caches. Skips comprise 10 gfx950-only NHD cases, two cases exceeding the device's 80-CU partition limit, and four architecture-gated rejection variants. Log: gfx942-tests.log.
• SKIP on dzm-mi355-gpu34-dev / fsx950223_pr_review_gfx950 / gfx950: Discover allocation with squeue; srun --overlap --jobid=9581 --ntasks=1 --cpus-per-task=1 podman exec -i fsx950223_pr_review_gfx950 rocm-smi --showuse --showmemuse
  Environment probes exited 0 and verified gfx950 inside the configured container. No idle device was available: all eight GPUs had 89–90% VRAM allocated; the final activity check showed 100% utilization on every GPU. No GPU test was launched. Logs: gfx950-allocation.log, gfx950-environment.log, gfx950-activity-recheck.log.

Remaining validation:

• Exact-head gfx950 correctness validation remains pending and should be retried when an idle device is available.
• No gfx1250 device was available; its preserved constructor and launcher interfaces were checked statically.
• Performance sweeps and downstream vLLM serving integration were not independently rerun.

Automated Codex review of commit a240ddcf95c2.