feat(rocke): GDN decode kernel and dispatch for gfx950 - #23
Closed
AviralGoelAMD wants to merge 1 commit into
Closed
AviralGoelAMD wants to merge 1 commit into
AviralGoelAMD wants to merge 1 commit into
Conversation
AviralGoelAMD
changed the base branch from
users/avirgoel/rocke/gdn-dispatch-core
to
develop
September 16, 2026 02:42
AviralGoelAMD
changed the base branch from
develop
to
users/avirgoel/rocke/gdn-dispatch-core
September 16, 2026 02:42
AviralGoelAMD
force-pushed
the
users/avirgoel/rocke/gdn-decode-pr3
branch
from
September 16, 2026 02:49
e15cac7 to
e10dfc2
Compare
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
force-pushed
the
users/avirgoel/rocke/gdn-decode-pr3
branch
from
September 16, 2026 04:42
e10dfc2 to
4abab8e
Compare
Owner
Author
|
Superseded by the upstream PR: ROCm#12172 (same content, rebased onto |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
(sequence, value head, v-sub-block)quad_permat offsets 1–2,ds_swizzleaboveread_indices/write_indices; negative index = idle lanetune.pyregenerates itis_valid_spec— one rule, not two copies that driftStack: 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
4abab8e95d1Output 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_enginecannot 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
batchdeep withread_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:1.949e-03— identical to a clean run, missed9.780e-01— rejectedPool now defaults to
2 * batch + 1with disjoint slots;check()and the tuner compare untouched pages against their pre-launch contents.make_inputsstill exposespool_depth/disjoint_writes, but the detecting shape is the default.