feat(rocke): GDN decode kernel and dispatch for gfx950 - #12172
AviralGoelAMD wants to merge 8 commits into
Conversation
github-actions
Bot
added
documentation
project: hip-kernel-provider
rocke: platform
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
AviralGoelAMD
added
the
rocke-lib: attention
| - gfx942/gfx950 cells use a portable f16 16x16x16 config; an instance marked ❌ | ||
| for a CDNA arch lacks the specific atom that config selects (e.g. `mfma_gemm` | ||
| and `direct_conv_16c` need the CDNA4 16x16x32 atom absent on gfx942). | ||
| - **`gdn_decode` is gfx950-only by registration, not by capability.** The kernel |
There was a problem hiding this comment.
This footnote states the opposite of what the validator does.
is_valid_spec rejects on spec.wave_size != target.wave_size at L171-182 of library/kernels/gfx950/gdn_decode.py — strictly before the max_threads_per_block check at L197. GdnDecodeSpec.wave_size defaults to 64 and is never derived from arch, so every wave32 target (gfx1151, gfx1201) is rejected regardless of thread-block size.
This PR's own test_a_wave32_target_is_rejected_for_a_wave64_spec (test_gdn_decode_spec.py L87-95) exists precisely because gfx1151 used to be wrongly accepted, and its comment says so. The instance page added by this same PR also now states the validator "checks the target wave size" — so the two docs in this changeset contradict each other.
"arch-neutral SSA" is also misleading given the emitter comment right above the check, which notes the xor-butterfly arithmetic would be silently wrong on a wave32 target. Please reword to: gated by registration and a wave64 target match.
| ) | ||
| if not validate_indices: | ||
| return | ||
| for name in ("read_indices", "write_indices"): |
There was a problem hiding this comment.
Nothing requires or checks that two active sequences write distinct pool slots.
_validate_decode_inputs (L261-271) only bounds the range via min()/max(), and no docstring, README, or gdn.md line states uniqueness as a caller contract. If a scheduler ever emits write_indices = [3, 3, ...] with both lanes active, the workgroups for b_i=0 and b_i=1 both store to state[3, hv, v_row, :] with no ordering — one sequence's recurrent state is silently clobbered by the other. That's the "right for one token and wrong forever after" failure this PR's own commit message warns about.
Duplicate read indices are fine, and read_pool == write_pool for the same sequence is the normal in-place case. Only write-vs-write across lanes races. The validator already pays for a device→host sync here, so this is free:
w = inp["write_indices"]
live = w[w >= 0]
if live.numel() != int(live.unique().numel()):
raise ValueError("write_indices must be unique across active sequences")At minimum, state it as a precondition in the kernel docstring next to the -1 sentinel contract.
There was a problem hiding this comment.
resolved, I added this validation.
| batch = 8 | ||
| pool_depth = make_inputs(spec, batch, device=DEVICE)["state"].shape[0] | ||
|
|
||
| for name, bad in ( |
There was a problem hiding this comment.
No test sets read_indices[i] = -1 with a valid write_indices[i] (or the converse).
The kernel's active predicate ANDs both non-negative checks (L287-290 simple, L465-468 warp-tiled), so a mismatched pair should deactivate the whole lane. But every -1 test (test_gdn_decode_prepare.py L47-54, test_gdn_decode_gfx950_numeric.py L154-177) sets both tensors to -1 together, and _validate_decode_inputs (builders/gfx950/gdn/gdn_decode.py L261) validates the two index tensors in independent loops — so a mismatched pair reaches the kernel with nothing flagging it.
If the AND ever became an OR, or the skip were split across the read and write phases, a lane with read=-1, write=5 would write and corrupt live slot 5, silently. Add a mismatched case asserting the target slot stays bit-identical to its pre-launch contents. While here, the below-sentinel case ("write_indices", -2) is also missing from this table — only read_indices is covered.
There was a problem hiding this comment.
Yep. added as you suggested.
| pytest.skip(f"no gdn_decode golden recorded for llvm flavor {flavor!r}") | ||
| drift = [] | ||
| for cid, build in _cases().items(): | ||
| want = recorded["cases"].get(cid, {}).get("sha256") |
There was a problem hiding this comment.
_build_doc() writes {"error": str(exc)[:160]} instead of a hash when emission raises. The companion test test_every_shipped_configuration_is_recorded only checks that the case key exists, so an {"error": ...} entry counts as "recorded" and satisfies it.
Failure scenario: someone re-blesses with --write on a node where one flavor or one tile fails to lower (exactly the kind of partially-broken node the PR description describes). The fixture records an error string for that case; from then on both tests pass while that configuration is never hashed again. Drift in it becomes undetectable.
if want is None:
drift.append(f"{cid}: no sha256 recorded ({recorded['cases'].get(cid)})")
continueThere was a problem hiding this comment.
I fixed the sha generation and its validation test
| v_lane = b.div(w_tid, b.const_i32(WTK)) | ||
| warp_k_start = b.mul(k_lane, b.const_i32(VPT)) | ||
| gv_start = b.add(b.mul(wid, b.const_i32(WTV)), v_lane) | ||
| k_lane0 = b.mul(v_lane, b.const_i32(WTK)) # first lane of this WTK group |
There was a problem hiding this comment.
k_lane0 is assigned and never read. Since IRBuilder is side-effecting, this isn't a dead local — it emits a real instruction into every warp-tiled kernel.
There was a problem hiding this comment.
deleted this unused code and updated the sha
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
ISSUE ID : AICK-2228 Single-token GDN decode over a paged recurrent state. One workgroup owns one (sequence, value head, v-sub-block); the state tile is register-resident, so there is no LDS and no barrier in the step, and the only cross-lane traffic is the xor butterfly that forms the key-value dot products -- quad_perm at offsets 1-2, ds_swizzle above. Serving surface: the state lives in a paged pool addressed by read_indices / write_indices, so continuous batching can hand the kernel a different physical page each step. A negative index marks an idle lane and leaves its state slot untouched. The pool base advances by a 64-bit byte offset so a deep pool does not wrap signed 32-bit arithmetic. Tile selection (num_warps, warp_threads_k, blocks_per_v_dim) is a table banded on batch, produced by an exhaustive correctness-gated sweep over the legal tile space; tune.py regenerates it, so the table is a recorded search result rather than a hand-guess. Dispatch's support check calls the kernel's own is_valid_spec, so the spec the kernel can emit and the spec dispatch may select are one rule rather than two that drift. Tests: spec validation and emission, host prepare, golden IR hashes for every dispatched tile, dispatch wiring through the real registry, and an on-silicon numeric test that gates the written state as well as the output -- a decode step can be right for one token and wrong forever after if the state is wrong. Depends on the quad_perm intrinsic and the hoisted make_kernel_id; see the PR description for the stack.
c541c2d to
9842b47
Compare
3 participants
ISSUE ID : AICK-2228
What this adds
Adds the gfx950 GDN (Gated Delta Network) single-token decode kernel, dispatcher registration, and validation coverage. Decode advances a fixed-size recurrent state instead of rereading an ever-growing KV cache; this PR is decode-only.
What changed
library/kernels/gfx950/gdn_decode.pylibrary/dispatch/gdn/library/builders/gfx950/gdn/library/tests/develop; its former stack dependencies are now included.Why it works
Each active lane reads one recurrent-state page, applies decay plus a rank-1 update, writes the next state page, and emits one output token. The host guard protects raw page addressing before launch, while the device predicate skips any lane whose read or write index is negative. The dispatcher uses the kernel validator as its final support authority, preventing dispatch rules from drifting from emitter legality.
How we validated
test_gdn_decode_golden.pyandtest_gdn_decode_spec.py: 28 passed, 6 subtests passed on rebased head9842b47.test_gdn_decode_prepare.py: 10 passed.test_mismatched_skip_index_leaves_write_page_untouched: 1 passed on MARKHAM gfx950.The full fp32-oracle numeric suite remains pending on gfx950. The on-silicon result above proves the mismatched skip lane leaves its valid write page bit-identical; it does not replace full output/state oracle coverage.
Notes