Skip to content

feat(rocke): GDN decode kernel and dispatch for gfx950 - #23

Closed
AviralGoelAMD wants to merge 1 commit into
users/avirgoel/rocke/gdn-stack-integrationfrom
users/avirgoel/rocke/gdn-decode-pr3
Closed

AviralGoelAMD wants to merge 1 commit into
users/avirgoel/rocke/gdn-stack-integrationfrom
users/avirgoel/rocke/gdn-decode-pr3

Conversation

@AviralGoelAMD

@AviralGoelAMD AviralGoelAMD commented Sep 16, 2026

Copy link
Copy Markdown
Owner

ISSUE ID : AICK-2228

Adds the GDN (Gated Delta Network) decode kernel for gfx950 — the single-token step run once per generated token — plus dispatch wiring, tests and docs. 21 files, +3282.

Why: softmax decode re-reads the whole KV cache, so cost grows with sequence length. GDN keeps a fixed-size state per head: decay it, add a rank-1 update from the new key/value, read the output. Memory per sequence is constant.

Work split one workgroup per (sequence, value head, v-sub-block)
State tile register-resident — no LDS, no barriers
Cross-lane xor butterfly only: quad_perm at offsets 1–2, ds_swizzle above
State storage paged pool via read_indices / write_indices; negative index = idle lane
Pool addressing 64-bit byte offset, so a deep pool does not wrap signed 32-bit
Tile knobs batch-banded table from a correctness-gated sweep; tune.py regenerates it
Support check dispatch calls the kernel's own is_valid_spec — one rule, not two copies that drift

Stack: PR 3 of 4 — quad_perm (ROCm#12070) → dispatch.core hoist (ROCm#12090) → this → GDN prefill. Needs both, so it sits on an integration of develop + those two and builds standalone.

Verification — head 4abab8e95d1

check result
CPU suites 60 passed, 0 skipped
Golden IR, llvm20/22/23 21/21 SHAs byte-identical
AICK-2228 criteria on gfx950 6/6 PASS
Host-validation lanes 23 passed, 0 skipped

Output and state checked against an independent fp32 oracle inside TOL = 1e-2; idle lanes bit-identical; deep-pool lane allocated past 2³². Performance numbers live in the protected results page.

Caveat, scoping all of the above: rocke_engine cannot be built on that node, so the run exercised Python-lowered code objects. The C++ path is covered by the byte-identity gate, not by execution.

For reviewers

Every gate was made to fail, not just pass — the numeric gate rejects a wrong output and a wrong state; the tuner gate rejects a corrupted reference; the new untouched-page check rejects a spilled write while that tile's written-pages term stays inside tolerance, which attributes the rejection to the new check.

The test fixture changed. The pool used to be exactly batch deep with read_indices == write_indices, so every page was written every launch and a correct value written to the wrong slot was invisible. Same injected fault, two geometries:

pool a write that aliases its read slot
old, fitted 1.949e-03 — identical to a clean run, missed
new, disjoint slots 9.780e-01 — rejected

Pool now defaults to 2 * batch + 1 with disjoint slots; check() and the tuner compare untouched pages against their pre-launch contents. make_inputs still exposes pool_depth / disjoint_writes, but the detecting shape is the default.

@github-actions github-actions Bot added documentation Improvements or additions to documentation project: hip-kernel-provider rocke: platform labels Sep 16, 2026
@AviralGoelAMD
AviralGoelAMD changed the base branch from users/avirgoel/rocke/gdn-dispatch-core to develop September 16, 2026 02:42
@AviralGoelAMD
AviralGoelAMD changed the base branch from develop to users/avirgoel/rocke/gdn-dispatch-core September 16, 2026 02:42
@AviralGoelAMD
AviralGoelAMD force-pushed the users/avirgoel/rocke/gdn-decode-pr3 branch from e15cac7 to e10dfc2 Compare September 16, 2026 02:49
@AviralGoelAMD
AviralGoelAMD changed the base branch from users/avirgoel/rocke/gdn-dispatch-core to users/avirgoel/rocke/gdn-stack-integration September 16, 2026 02:49
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.
@AviralGoelAMD
AviralGoelAMD force-pushed the users/avirgoel/rocke/gdn-decode-pr3 branch from e10dfc2 to 4abab8e Compare September 16, 2026 04:42
@AviralGoelAMD

Copy link
Copy Markdown
Owner Author

Superseded by the upstream PR: ROCm#12172 (same content, rebased onto develop so the diff is the 21 decode files only). Closing this working copy.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant