[yraparti]

Summary

ISSUE ID : AICK-2051

Adds a spec-driven kernel family for chunkwise gated delta-rule linear attention (KDA) prefill on gfx950, under library/kernels/gfx950/kda_chunkwise.py with host builders under library/builders/gfx950/kda/.

The per-chunk body factors into six state-independent tiles (A, GK, GQ, Aqk, Kt, dec) plus a serial state recurrence. The two halves have opposite parallelism profiles — the tiles are fully parallel over the sequence, the recurrence is serial per (batch, head) — so both compositions ship:

Path	Tile grid	Scan grid
Split (build_kda_chunk_prep + build_kda_chunk_scan), default	one WG per chunk	one WG per (batch, head)
Fused (build_kda_chunk_fused)	one WG per (batch, head)	same WG

Three specs, each with its own is_valid_*_spec and build_*, following the platform instance conventions.

The interesting result is a negative one

The fused path is the design that looks strictly better on paper — tiles never leave LDS, so it is strictly less memory traffic — and it loses.

Holding the tile builder's staging tiles and the scan's operands live simultaneously puts a workgroup over half the LDS budget, so only one is resident per CU. The scan is a latency-bound chain of small matmuls with a serial dependency between chunks; at one workgroup per CU there is no second workgroup to cover that latency, and it bleeds into the tile phase as well. Paying for a tile round trip through HBM to keep two workgroups resident is the cheaper side of the trade.

This was isolated, not guessed: ablating the scan showed the fused kernel's tile phase alone costs far more than the standalone tile kernel doing identical work; a looped tile-only kernel writing to HBM ruled out the loop structure and the destination switch; and forcing the standalone kernel down to one workgroup per CU with unused LDS ballast reproduced the fused cost.

The conclusion is encoded in the spec, not left as a comment. is_valid_scan_spec checks lds_bytes() against LDS_LIMIT // min_occupancy rather than the full budget, so a scan spec that would fit the hardware but not two-per-CU is rejected — it would not be worth running.

Levers

1. 128-bit LDS access throughout. Padded row pitches are constrained to multiples of 8 elements so operands move as ds_read_b128/ds_write_b128. The pad doubles as bank-conflict padding: at a row stride of DK, an unpadded 128-element bf16 row puts every lane in the same 16 banks.
2. One fused elementwise sweep. The passes producing GK, GQ and the three MFMA operands all read the same q/k/gate values; merging them cuts both LDS traffic and the exponential count, since each exp2 result feeds several consumers.
3. Blocked triangular solve. The solve splits into per-block forward substitution (irreducibly scalar VALU) and a rank update against solved blocks (MFMA). Substitution shrinks as the square of the block size, so solve_block trades scalar work for one more block step.
4. Every product in A B^T form. The recurrence is written transposed throughout, putting the contraction on the fastest axis of both operands for all five of its matrix products. No operand ever needs an LDS transpose.
5. No fp32 staging for the residual or the output. The accumulator lane mapping puts each lane's slots at one chunk row and four runs of four consecutive v channels, so V is subtracted straight into the accumulator and O goes straight to HBM already coalesced.

Two bugs worth recording, both now spec rejection rules or explicit barriers: a C x C bf16 tile at a pitch that is a multiple of 4 but not 8 puts odd rows off a 16-byte boundary and silently violates the ds_read_b128 alignment contract (hence pad_cb); and the blocked solve's cross-lane ds_write -> ds_read needs an explicit s_waitcnt lgkmcnt(0), which presented as a failure only at the smallest block sizes.

Correctness

Two float64 oracles, deliberately independent of the kernels and of each other:

• Tile oracle — forms pairwise exponent differences directly, with no midpoint factoring, so it never builds the overflowing intermediate the kernel is designed around. Agreement checks the factorization rather than restating it.
• Token-serial oracle — walks the gated delta rule one token at a time, not chunked at all. Agreement tests the chunkwise factorization itself.

The gate range is swept to -5 (the reference gate_lower_bound), where a 32-token chunk accumulates up to 160 nats — the only regime that exercises the factoring and the clamps.

The split and fused paths share one emitted scan body and are checked for bitwise equality, not merely agreement within tolerance.

Test plan

• tests/test_kda_chunkwise_spec.py — spec admission and rejection rules, plus comgr compile of all three builders (no GPU). 20 tests.
• tests/test_kda_chunkwise_gfx950_numeric.py — both oracles across the gate range, plus split-vs-fused bitwise equality. 14 tests, @pytest.mark.gpu, gated on gcnArchName.
• Byte-identity gate GREEN across all families (tools/check_byte_identity.py).
• library/tests/run_all.py CPU lane — no new failures. (Three pre-existing failures in gfx942 attention code are untouched by this change.)

cd rocke/library
python -m pytest tests/test_kda_chunkwise_spec.py
python -m pytest tests/test_kda_chunkwise_gfx950_numeric.py -m gpu

Notes for reviewers

• Methodology, levers and the attribution technique are documented in library/builders/gfx950/kda/CASE_STUDY.md. Per AGENTS.md §Compliance, measured numbers are not in this repo or this PR — they are recorded on the protected internal page.
• The case study also records a methodology failure worth reading: the LDS occupancy probe was initially dead-code eliminated and confidently reported that occupancy was irrelevant. That wrong conclusion survived several rounds of tuning. A probe reporting "no effect" is indistinguishable from a probe that does not run.
• Known limitation: the scan grid is BH, so below roughly the CU count it is parallelism-starved and an implementation that splits the v dimension to manufacture more workgroups will do better. Fixing it needs a second partitioning rule for the same accumulator, which invites cross-wave reductions; it is documented rather than implemented.
• KDA is not yet wired into any dispatch path; this adds the kernels, builders and tests only.

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[AnaghaRaoAMD]

DoD audit against rocke/KERNEL_AUTHORING.md — five boxes unticked for a new kernel family, most importantly the golden IR SHA256 fixture.

Core DoD (every change, KERNEL_AUTHORING.md L95-113)

Item	Status	Evidence
On-GPU numeric parity test, not spec-only geometry	✅	library/tests/test_kda_chunkwise_gfx950_numeric.py — 14 tests against a token-serial float64 oracle that shares no algebra with the kernel. Genuinely strong.
Coverage in both directions: spec → IR and problem → spec	❌	Neither. No IR-level test at all, and no dispatcher/selection test (nothing selects KDA — see below).
Golden: "a new family/case adds an entry"	❌	library/tests/golden/ still contains exactly one file, attention_dense_gfx942_ir_sha256.json. No KDA entry.
run_checks.py green	⚠️	Stages are guard → gate → parity → pytest → numeric. Every KDA test is pytest.mark.gpu (numeric L62; spec tests likewise), so on any machine without gfx950 this PR contributes zero executed coverage.
Docs	⚠️	CASE_STUDY.md is present and good. No README.md harness/results doc, no ALGORITHM.md.

§A — New kernel family (L120-140)

Requirement	Status	Evidence
Spec-driven build_*(), reuse helpers/atoms	✅	KdaTileSpec/KdaChunkFusedSpec/KdaChunkScanSpec + build_kda_chunk_{prep,split,fused}. Clean.
C engine mirror in the same change	➖ exempt	library/kernels/ has no C++ mirror by design.
.py/.c parity pair under tests/parity/	➖ exempt	Follows from the above. library/tests/parity/ holds 20 pairs, all platform-mirrored families; no kda_*_emit.* is owed.
Dispatcher wiring — "A family not in a registry is unreachable."	❌	library/dispatch/ has no KDA reference (grep -ril kda across the tree, excluding the 8 PR files, returns only an incidental substring in kernels/common/fmha_appendkv.py). Arguably KDA shouldn't go in ATTENTION_REGISTRY — it selects from an SDPA UnifiedAttentionProblem and linear attention has no such request shape. But the dense precedent registered opt-in rather than skipping (library/dispatch/AGENTS.md L35-37: attention_gfx942_dense, attention_gfx950_dense, both "(opt-in only)"). Either wire it or state the decision in the PR body — right now the family is reachable only by running a builder by hand.
Golden created for the new family	❌	See below.
End-user visibility — SUPPORT_MATRIX.md	❌	platform/python/rocke/instances/SUPPORT_MATRIX.md "## Attention / FMHA" table (L73, ~14 rows) has no KDA row.
Docs discoverability — instances/index.md row + per-family page	❌	platform/dsl_docs/instances/index.md "## Attention Family" table (L65) has no KDA row; no per-family page.
A benchmark scenario	❌	library/benchmarks/ has 12 drivers, none for KDA. The bench loop inside kda_chunk_fused.py main() is a builder-local sweep, not a registered scenario.

Appendix §A step 8 (L219-257)

Requirement	Status
ALGORITHM.md "if the math is new"	❌ — the math is new here (chunkwise factorization of the gated delta rule); library/builders/gfx950/attention/ carries both this and a README, library/builders/gfx950/kda/ carries neither
README.md harness/results doc	❌
If any of these are not to be addressed in the PR, we need follow up tickets created

[therock-pr-bot]

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

• python -m pip install pre-commit
• pre-commit install
• pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

[assistant-librarian]

TheRock Submodule Bump Activity

Newest first

• 2026-09-08 — Included in TheRock via ROCm/TheRock#7975.