Skip to content

feat(hip-kernel-provider): KDA decode from the GDN decode emitter - #20

Open
AviralGoelAMD wants to merge 1 commit into
users/avirgoel/rocke/gdn-prefillfrom
users/avirgoel/rocke/kda-decode-gfx950
Open

AviralGoelAMD wants to merge 1 commit into
users/avirgoel/rocke/gdn-prefillfrom
users/avirgoel/rocke/kda-decode-gfx950

Conversation

@AviralGoelAMD

@AviralGoelAMD AviralGoelAMD commented Sep 14, 2026

Copy link
Copy Markdown
Owner

ISSUE ID : AICK-2240

What this adds

Adds KDA decode on gfx950 through the existing GDN decode emitter. rocKE already had KDA prefill, but no KDA path for the single-token decode step used during generation.

The scope is D128 decode. gfx942 support and surrounding conv1d/RMSNorm fusion remain separate follow-ups.

What changed

Area / files Crux of change
library/kernels/gfx950/gdn_decode.py Adds per-channel KDA decay through gate_kind, while retaining scalar GDN decay.
library/builders/gfx950/gdn/{gdn_decode,tune}.py Adds KDA inputs, an independent fp32 reference, full-pool correctness checks, and work-axis tuning.
library/dispatch/gdn/{common,gfx950}.py Carries gate_kind, keeps separate GDN/KDA tile tables, and prevents cross-kind cache or routing collisions.
library/tests/ + golden fixture Covers MHA/GQA, fused and precomputed gate modes, routing, input validation, resources, and LLVM IR stability.
Benchmark and GDN/KDA docs Adds the KDA decode case and documents the shared contract and supported scope.
  • GDN keeps its original batch-keyed table; KDA uses work = batch × num_v_heads because its table was measured across head-count geometries.
  • The precomputed-decay mode is benchmark-only. Dispatch ships the fused gate path.

Why it works

GDN and KDA share the same recurrent decode update; only the forget gate differs. GDN applies one scalar decay per head, while KDA applies one decay per key channel. The emitter loads each KDA decay value once per channel chunk and reuses it across state rows, so the recurrence, paged state handling, and output path stay shared instead of being copied into a second kernel.

How we validated

Validation ran on dependency-complete snapshot bd12bb8d0d5 on real gfx950 with zero skips. Its 20 PR file blobs are byte-identical to live PR head 096ed04140c:

Gate Result
KDA on-device numerics 19 passed; includes direct-spec coverage and three request → dispatch → selected-kernel → launch → oracle cases
GDN regressions 15 decode + 15 prefill passed
Host and dispatch coverage 121 passed, plus 119 subtests, including CPU-only dispatch selection and wiring
Accuracy Worst KDA output 1.53e-4; full-state-pool 2.69e-3; tolerance 1e-2
Golden IR GDN 7/7 unchanged; KDA 6/6 matched in LLVM 20, 22, and 23

The first rebase run caught a host validation conflict that rejected valid KDA gate buffers after their mode-specific checks. The corrected path is pinned by a regression test that fails when the duplicate scalar-GDN check is restored.

Notes

@AviralGoelAMD AviralGoelAMD self-assigned this Sep 14, 2026
@github-actions github-actions Bot added documentation Improvements or additions to documentation rocke: platform labels Sep 14, 2026
ISSUE ID : AICK-2240

Rebases the KDA decode work onto the completed GDN stack. The shared gfx950
decode emitter now supports scalar GDN and per-channel KDA gates; GDN keeps
its batch-keyed routing and KDA uses its measured work-keyed tile table.

Preserves the newer GDN pool, input-validation, untouched-page and tuner gates
while adding KDA fused and benchmark-only precomputed-decay modes.
@AviralGoelAMD
AviralGoelAMD changed the base branch from users/avirgoel/rocke/gdn-decode-tier1 to users/avirgoel/rocke/gdn-prefill September 17, 2026 16:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment