feat(rocke): add chunkwise KDA prefill kernels for gfx950

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/kda/benchmark_chunkwise.py

@@ -0,0 +1,95 @@
+"""Benchmark the gfx950 chunkwise KDA builder compositions.
+
+Run from ``rocke/library``:
+
+    python -m benchmarks.gfx950.kda.benchmark_chunkwise --path both
+    python -m benchmarks.gfx950.kda.benchmark_chunkwise \
+        --path split --shapes 8x8x1024 8x16x2048
+
+The builder modules own input construction, launch configuration, and timing.
+This driver provides a stable benchmark scenario without wiring the family into
+dispatch.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+
+def _shapes(raw: list[str]) -> list[tuple[int, int, int]]:
+    result = []
+    for value in raw:
+        try:
+            shape = tuple(int(part) for part in value.split("x"))
+        except ValueError as exc:
+            raise argparse.ArgumentTypeError(
+                f"shape must be BxHxT, got {value!r}"
+            ) from exc
+        if len(shape) != 3 or any(dim <= 0 for dim in shape):
+            raise argparse.ArgumentTypeError(
+                f"shape must be positive BxHxT, got {value!r}"
+            )
+        if shape[2] % 32:
+            raise argparse.ArgumentTypeError(
+                f"T must be divisible by the default chunk size 32, got {value!r}"
+            )
+        result.append(shape)
+    return result
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument(
+        "--path",
+        choices=("split", "fused", "both"),
+        default="both",
+        help="builder composition to benchmark",
+    )
+    parser.add_argument(
+        "--shapes",
+        nargs="+",
+        default=["8x8x1024", "8x16x2048", "32x16x2048"],
+        metavar="BxHxT",
+    )
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=30)
+    args = parser.parse_args()
+    shapes = _shapes(args.shapes)
+
+    if args.path in ("split", "both"):
+        from builders.gfx950.kda import kda_chunk_split as split
+        from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec
+
+        print("== split ==")
+        spec = KdaChunkScanSpec()
+        for batch, heads, tseq in shapes:
+            split.bench(
+                spec,
+                batch,
+                heads,
+                tseq,
+                warmup=args.warmup,
+                iters=args.iters,
+            )
+
+    if args.path in ("fused", "both"):
+        from builders.gfx950.kda import kda_chunk_fused as fused
+        from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec
+
+        print("== fused ==")
+        spec = KdaChunkFusedSpec()
+        for batch, heads, tseq in shapes:
+            fused.bench(
+                spec,
+                batch,
+                heads,
+                tseq,
+                warmup=args.warmup,
+                iters=args.iters,
+            )
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/ALGORITHM.md

@@ -0,0 +1,70 @@
+# Chunkwise KDA algorithm
+
+KDA is a gated delta-rule linear-attention recurrence. For one head, let
+`S` have shape `DK x DV`, let `q` and `k` be `DK`-vectors, and let `v` be a
+`DV`-vector. For each token:
+
+```text
+S <- Diag(exp(g)) S
+u <- beta (v - k^T S)
+S <- S + k u^T
+o <- scale q^T S
+```
+
+The token recurrence is serial, but tokens can be grouped into chunks of `C`.
+Within a chunk, define the cumulative per-channel decay
+`Gamma_i = exp(sum(j <= i, g_j))` and the whole-chunk decay
+`gamma_C = Gamma_(C-1)`. The chunk body is factored into six tiles:
+
+```text
+A    = (I + StrictTril(Diag(beta) Akk))^-1 Diag(beta)
+       Akk_ij = k_i . (k_j * Gamma_i / Gamma_j)
+GK   = K * Gamma
+GQ   = Q * Gamma * scale
+Aqk  = Tril(GQ (K / Gamma)^T)
+Kt   = (K * gamma_C / Gamma)^T
+dec  = gamma_C
+```
+
+The state-dependent part then becomes:
+
+```text
+Vt = A (V - GK S)
+O  = GQ S + Aqk Vt
+S  = Diag(dec) S + Kt^T Vt
+```
+
+`A`, `GK`, `GQ`, `Aqk`, `Kt`, and `dec` depend only on the chunk inputs.
+Their construction is parallel over chunks. The state update remains serial
+over chunks for each `(batch, head)` pair, while matrix products are partitioned
+over workgroup waves.
+
+## Stable decay factorization
+
+The ratio `Gamma_i / Gamma_j` can overflow when formed directly. The emitter
+uses a midpoint row `CREF = C // 2`:
+
+```text
+K * exp(Gc - Gref)   and   K * exp(Gref - Gc)
+```
+
+Their product reconstructs the required ratio while keeping each exponential
+within a bounded half-chunk range. The cumulative gate is maintained in the
+base-2 exponent domain, and exponent inputs are clamped to the finite hardware
+range.
+
+## Two compositions
+
+The split composition uses one workgroup per chunk for tile construction. It
+writes the six tiles to global memory, then uses one workgroup per `(batch, head)`
+for the ordered state scan. This separates the tile and scan occupancy
+requirements.
+
+The fused composition uses one workgroup per `(batch, head)` and keeps the
+tiles in LDS while walking the chunks. It has less global-memory traffic but a
+larger live LDS footprint. The two paths share the scan body and are checked for
+bitwise equality where their layouts and accumulation order match.
+
+Both compositions support a supplied initial state `h0` and an optional final
+state `ht`. The numeric tests compare this state-carrying path with the
+un-chunked token-serial reference above.