diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/kda/benchmark_chunkwise.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/kda/benchmark_chunkwise.py new file mode 100644 index 000000000000..9079c71fcf06 --- /dev/null +++ b/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()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/ALGORITHM.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/ALGORITHM.md new file mode 100644 index 000000000000..400b913b056c --- /dev/null +++ b/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. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/CASE_STUDY.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/CASE_STUDY.md new file mode 100644 index 000000000000..d8278e0ac84d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/CASE_STUDY.md @@ -0,0 +1,271 @@ +# Chunkwise KDA on gfx950 — optimization case study + +How the chunkwise gated delta-rule linear attention (KDA) prefill kernels in +`kernels/gfx950/kda_chunkwise.py` were built and tuned against an existing +reference implementation on a 32-CU gfx950 partition. + +Per [`AGENTS.md`](../../../../AGENTS.md) §Compliance this document records +**methodology and levers only**. No measured numbers, ratios, or achieved +bandwidths appear here; those live on the protected Confluence page. Everything +below is replayable — each claim names the harness that produced it, so the +numbers can be regenerated on demand rather than quoted. + +## 1. What the workload is + +KDA is a gated delta-rule linear attention. Tokens are grouped into chunks of +`C`, and per chunk the token-serial recurrence factors into six +**state-independent** tiles plus a short **state-dependent** recurrence. The +full contract, including the overflow-avoiding factorization, is in the module +docstring of `kernels/gfx950/kda_chunkwise.py`. + +The split matters more than any individual optimization: + +- The six tiles (`A`, `GK`, `GQ`, `Aqk`, `Kt`, `dec`) depend only on that + chunk's `q/k/g/beta`. Fully parallel over the sequence. +- The recurrence carrying the state `S` across chunks is serial per + (batch, head), and parallel only across heads. + +Two halves with opposite parallelism profiles, which is what makes the fusion +question (§5) interesting rather than obvious. + +## 2. Harnesses + +Nothing was tuned without a correctness gate attached to the same run. Three +independent references, each checking a different claim: + +| Harness | Reference | What it establishes | +| --- | --- | --- | +| `kda_chunk_prep.py` | float64 tile oracle | The six tiles, individually | +| `kda_chunk_split.py` | token-serial float64 walk | The chunkwise factorization itself | +| `kda_chunk_fused.py` | token-serial float64 walk | Same, for the fused path | + +The tile oracle is deliberately built by a different route than the kernel: it +forms the pairwise exponent *difference* directly, with no midpoint factoring +and no `C x C x DK` reuse. So agreement checks the factorization rather than +restating it. The token-serial oracle is not chunked at all. + +`tests/test_kda_chunkwise_gfx950_numeric.py` runs all three across the gate +range, plus a bitwise-equality check between the split and fused paths (they +share one emitted scan body, so tolerance-based agreement would be too weak a +statement). + +The gate range is swept to `-5`, the reference `gate_lower_bound`. This is not +padding: a 32-token chunk accumulates up to 160 nats there, which overflows +fp32 if the ratio `Gamma_i / Gamma_j` is formed directly. Only that end of the +sweep exercises the factoring and the clamps. + +## 3. Method: attribute before optimizing + +Every lever below came from an attribution step, not a guess. Two techniques +did the work. + +### 3.1 Phase ablation + +A build-time switch dropped individual phases (cumsum, elementwise, the `C x C` +MFMAs, the triangular solve, each store) and the kernel was re-timed. Each +ablation breaks correctness, so it only ever answers "what does this phase +cost" — and because the compiler dead-codes anything whose only consumer was +ablated, each delta is an *upper* bound on the named phase. + +This is what ranked the phases and pointed at the triangular solve and the +elementwise sweep as the two worth restructuring. The scaffolding was removed +once it had served its purpose; it is recorded here rather than kept in the +kernel, since a phase switch that breaks correctness has no business shipping. + +### 3.2 LDS ballast: separating occupancy from instruction count + +A wall-clock number cannot distinguish "starved of concurrency" from +"serialized on a dependency chain". To separate them, an **unused** LDS +allocation was added to hold the instruction stream fixed while varying only +the workgroups resident per CU. + +**This probe was initially broken and gave a confidently wrong answer.** An LDS +allocation with no accesses is eliminated outright, leaving +`group_segment_fixed_size` unchanged — so the probe measured nothing while +appearing to show that occupancy was irrelevant. That conclusion survived +several rounds of tuning before it was caught. + +The fix is one keep-alive store from one thread. With the allocation actually +live, dropping from two workgroups per CU to one turned out to cost a large +multiple on the tile work — which reframed the whole fusion question in §5. + +**Lesson, and the reason this is written down:** a probe that reports "no +effect" is indistinguishable from a probe that does not run. Verify the +mechanism moved before trusting a null result. For an LDS probe, that means +reading `group_segment_fixed_size` out of the code object, not reasoning about +the source. + +## 4. Levers that paid off, in the order they were found + +1. **128-bit LDS access throughout.** Every staging tile's padded row pitch is + constrained to a multiple of 8 elements so operands move as `ds_read_b128` / + `ds_write_b128` rather than scalar transactions. The pad is also bank-conflict + padding: the MFMA operand reads are at a row stride of `DK`, and an unpadded + 128-element bf16 row puts every lane in the same 16 banks. + + The `pad_cb` knob exists because of a bug caught here: a `C x C` bf16 tile at + a pitch that is a multiple of 4 but not 8 puts odd rows on an 8-byte + boundary, silently violating the `ds_read_b128` alignment contract. That is + now an `is_valid_spec` rejection rather than a comment. + +2. **One fused elementwise sweep.** The separate passes producing `GK`, `GQ` and + the three MFMA operands all read the same `q`, `k` and cumulative-gate values. + Merging them into a single sweep, where a thread owns eight consecutive + channels of one row, cut both the LDS traffic and the exponential count — + each `exp2` result feeds several consumers. + +3. **Blocked triangular solve.** The solve's arithmetic splits into per-block + forward substitution (irreducibly scalar VALU) and a rank update against + already-solved blocks (a matmul, so MFMA). The substitution shrinks as the + *square* of the block size, so smaller blocks move more of the `O(C^3)` work + onto the MFMA pipe at the cost of one more block step. `solve_block` exposes + the trade; the degenerate `solve_block == chunk` is the original unblocked + scalar solve. + + This step introduced a race: `ds_write_b128` followed by a `ds_read_b128` + from a *different* lane needs an explicit `s_waitcnt lgkmcnt(0)` between + them. It presented as a correctness failure only at the smaller block sizes, + where the cross-lane distance is shortest. + +4. **Keeping every product in `A B^T` form.** The state recurrence is written + transposed throughout, which puts the contraction on the fastest axis of both + operands for all five of its matrix products. No operand ever needs an LDS + transpose. This is a layout decision, not an optimization applied afterward — + it is why §5's scan is as cheap as it is. + +5. **No fp32 staging tile for the residual.** In `R^T = V^T - Z^T`, the + accumulator's lane mapping puts each lane's slots at one chunk row and four + runs of four consecutive `v` channels. So `V` is subtracted straight into the + accumulator with four short vector loads, and only the bf16 result reaches + LDS. The same observation lets `O` go straight to HBM from the accumulators: + a slot's column index is `lane % 32`, so one store instruction is already + coalesced across 32 consecutive channels. + +## 5. The fusion that lost, and why that is the interesting result + +The obvious design fuses both halves: one workgroup per (batch, head) walks its +chunks in order, tiles are produced and consumed in LDS, and the only HBM +traffic is the inputs and the output. It is strictly less traffic than +materializing the tiles — the split path writes every tile once and reads it +back — and `build_kda_chunk_fused` implements exactly that. + +It is measurably slower, and the ablation harness plus the (repaired) ballast +probe explain why without ambiguity: + +- Ablating the scan showed the fused kernel's **tile** phase alone costs far + more than the standalone tile kernel doing identical work. +- A looped tile-only kernel writing to HBM destinations ruled out the chunk loop + structure and the LDS-versus-HBM destination switch. +- Forcing the standalone tile kernel down to one workgroup per CU with ballast + reproduced the fused kernel's tile-phase cost almost exactly. + +So the cost is entirely occupancy. Holding the tile builder's staging tiles and +the scan's operands live simultaneously puts the workgroup over half the LDS +budget, so one workgroup fits per CU. The scan is a latency-bound chain of small +matmuls with a serial dependency between chunks; at one workgroup per CU there +is no second workgroup to cover that latency, and it bleeds through into the +tile phase too. + +**The conclusion generalizes past this kernel:** minimizing memory traffic and +maximizing occupancy are competing objectives, and when the consumer is +latency-bound rather than bandwidth-bound, occupancy wins. Paying for a tile +round trip to keep two workgroups resident is the cheaper side of that trade. + +This is encoded in the spec rather than left as a comment: +`is_valid_scan_spec` checks `lds_bytes()` against `LDS_LIMIT // +min_occupancy`, not against the full budget. A scan spec that would fit the +hardware but not two-per-CU is **rejected**, because it would not be worth +running. + +## 6. Where each path wins + +Both paths ship, because their grids differ: + +| Path | Tile grid | Scan grid | +| --- | --- | --- | +| Split (`build_kda_chunk_prep` + `build_kda_chunk_scan`) | one WG per chunk | one WG per (batch, head) | +| Fused (`build_kda_chunk_fused`) | one WG per (batch, head) | same WG | + +The split path is the default and is the one that beats the reference at batch +scale. Its scan grid is `BH`, so it needs `BH` comfortably above the CU count. +Below that the scan is parallelism-starved and a reference implementation that +splits the `v` dimension to manufacture more workgroups will win instead — a +real limitation, and the honest place for it is here rather than hidden in a +shape sweep. Splitting `v` in this scan would require a second partitioning rule +for the same accumulator, which is how cross-wave reductions creep in; it has +not been implemented. + +The fused path is retained because it is the right shape of answer for a +future arch with either a larger LDS budget or cheaper cross-chunk latency, and +because the bitwise agreement test makes it a strong check on the split path. + +## 7. Reproducing + +```bash +cd rocke/library/builders/gfx950/kda + +python kda_chunk_prep.py # six tiles vs the float64 oracle +python kda_chunk_split.py # end to end, default path +python kda_chunk_fused.py # end to end, fused path +python kda_chunk_split.py --no-check --shapes 32x16x2048 +``` + +The pytest lanes: + +```bash +cd rocke/library +python -m pytest tests/test_kda_chunkwise_spec.py # specs + compile, no GPU +python -m pytest tests/test_kda_chunkwise_gfx950_numeric.py -m gpu +``` + +Timing note: the benchmarks wrap launches in `time_launches`, which runs under +`no_fence()`. Timing fenced launches instead charges every iteration a host +stream sync (`LaunchConfig.fence` defaults to `True`), which is not how the +kernels run in a pipeline — the scan consumes the tiles on the same stream, so +FIFO order already covers the dependency. At short sequences that sync is a +large fraction of the measurement, and an early version of these benchmarks +understated the kernels because of it. + +## 8. Aligned raw split path (stack-gap closure) + +The production gap on equal-length C=32 shapes was traced to **host +preprocessing and layout packing**, not to the core recurrence math. The aligned +closure keeps the ten-shape contract fixed (`bf16` token-major I/O, `DK=DV=128`, +`C=32`, `Tseq % 32 == 0`, V-first fp32 final state) and adds a **default-off** +raw producer on the split path: + +- `KdaChunkPrepSpec.raw_inputs` plus fused `q/k` L2 norm, `lower_bound * + sigmoid(exp(A_log) * (g + dt_bias))`, and `sigmoid(beta)` at the load/commit + seam. Prepared chunk-packed signatures stay byte-identical when raw mode is off. +- `KdaChunkScanSpec.value_splits in {1,2,4,8}` launches `(BH * value_splits)` + workgroups with disjoint V/state slices, mirroring FlashKDA K2's independent + `ceil(V/BW)` grid dimension. +- Token-major scan I/O removes the head-major pack from the hot path. + +### Step 0 sweep space (methodology only) + +Batch-compile the Cartesian product of legal `value_splits`, prep/scan tile +`block_size`, `scan_atom_m`, padding knobs, and `solve_block`, then correctness- +prune against the aligned float64 oracle (`kda_chunk_split.ref_aligned_raw`) and +the token-serial walk. Survivors are timed with identical event windows and at +least five median blocks via the in-repo benchmark scenario: + +```bash +cd rocke/library +python -m benchmarks.gfx950.kda.benchmark_chunkwise --path split \ + --shapes 8x8x1024 8x16x2048 32x16x2048 +``` + +The builder-local `bench()` helpers expose the same timing hooks for focused +experiments. Keep measured results outside the repository per compliance. + +Before profiling, inspect HSACO resources and ISA for each survivor: VGPR/SGPR/ +LDS occupancy, vector load/store widths, barrier/wait counts, shuffle reductions, +and transcendental growth on the raw producer. Profile one **underfilled** +`BH=12` shape and one filled shape to separate grid starvation, transform cost, +LDS stalls, and lost prefetch overlap. Revert levers that do not move the +bottleneck named by ATT. + +Do not hard-code a `value_splits` cutover until the sweep shows a stable +`BH/CU` boundary. Raw monolithic fusion stays gated until the best raw split +path misses the end-to-end target on filled shapes. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/README.md new file mode 100644 index 000000000000..8284b7cb863f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/README.md @@ -0,0 +1,62 @@ +# gfx950 chunkwise KDA + +This directory contains the spec-driven builders and correctness harnesses for +chunkwise gated delta-rule linear-attention prefill: + +- `kda_chunk_prep.py` builds the six state-independent per-chunk tiles. +- `kda_chunk_split.py` materializes those tiles, then runs the serial state scan. +- `kda_chunk_fused.py` builds and consumes the tiles in one workgroup. + +The split path is the default composition. The fused path is retained as an +alternative schedule and as an independent cross-check. Both builders accept +an optional non-zero initial state and can store the final state. + +## Run the builders + +From `rocke/library` with the editable `rocke` and library packages available: + +```bash +python builders/gfx950/kda/kda_chunk_prep.py +python builders/gfx950/kda/kda_chunk_split.py +python builders/gfx950/kda/kda_chunk_fused.py +``` + +Use `--no-check` for benchmarking only and `--shapes BxHxT,...` to select +benchmark shapes. The checks compare against independent float64 references. + +The reusable benchmark scenario runs both compositions: + +```bash +python -m benchmarks.gfx950.kda.benchmark_chunkwise --path both +``` + +This family is currently builder-only. It is intentionally not registered in +the dispatcher. + +## Tests + +The CPU lane validates admission rules and compiles all three builders through +comgr when available: + +```bash +python -m pytest tests/test_kda_chunkwise_spec.py +``` + +The GPU lane requires a gfx950 device and checks the tile oracle, both +compositions, non-zero initial state, value splits, and split/fused agreement: + +```bash +python -m pytest tests/test_kda_chunkwise_gfx950_numeric.py -m gpu +``` + +KDA is a family in the shared CPU-only IR parity harness. From the `rocke/` +root, re-bless or verify all representative families with: + +```bash +PYTHONPATH=platform/python:library \ + python platform/tests/instances/rocke_ir_parity_harness.py \ + --write platform/tests/golden/rocke_representative_ir_sha256.json +PYTHONPATH=platform/python:library \ + python -m pytest platform/tests/test_rocke_ci_static.py \ + -k ir_cases_match_golden_sha256 +``` diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/__init__.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py new file mode 100644 index 000000000000..135e56d01d38 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py @@ -0,0 +1,311 @@ +"""Host builder for the fused chunkwise KDA forward on gfx950. + +Builds, validates and benchmarks :func:`build_kda_chunk_fused`: one workgroup +per (batch, head) walks that head's chunks in order, so the six per-chunk tiles +are produced and consumed in LDS and the only HBM traffic is q/k/g/beta in and +o (plus the final state) out. + +Correctness is checked against a token-serial float64 oracle, which shares no +algebra with the kernel: the kernel evaluates the chunkwise factorization, +while the oracle just walks the gated delta rule one token at a time. Agreement +therefore tests the factorization itself, not only its implementation. + +Run directly for a parity sweep plus a benchmark:: + + python kda_chunk_fused.py # parity + bench + python kda_chunk_fused.py --no-check # bench only +""" + +from __future__ import annotations + +import argparse +import dataclasses +import os +import sys + +try: + import rocke # noqa: F401 +except ImportError: # running as a bare script outside the editable install + _HERE = os.path.dirname(__file__) + _RK = os.path.abspath(os.path.join(_HERE, "../../../..")) + sys.path[:0] = [_RK + "/library", _RK + "/platform/python"] + +import torch # noqa: E402 + +from kernels.gfx950.kda_chunkwise import ( # noqa: E402 + KdaChunkFusedSpec, + KdaTileSpec, + build_kda_chunk_fused, + is_valid_fused_spec, + kda_chunk_fused_grid, + kda_chunk_fused_signature, +) +from rocke.helpers.compile import compile_kernel # noqa: E402 +from rocke.runtime import KernelLauncher, LaunchConfig, time_launches # noqa: E402 + +_LAUNCHER_CACHE: dict = {} +TOL = 3e-2 + + +def make_launcher(spec: KdaChunkFusedSpec) -> KernelLauncher: + key = spec.kernel_name() + if key not in _LAUNCHER_CACHE: + ok, why = is_valid_fused_spec(spec) + if not ok: + raise ValueError(f"unsupported spec: {why}") + art = compile_kernel( + build_kda_chunk_fused(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + _LAUNCHER_CACHE[key] = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=kda_chunk_fused_signature(spec), + ) + return _LAUNCHER_CACHE[key] + + +def run_fused(spec, q, k, g, beta, v, o, ht, scale, bh, nc, h0=None, stream=None): + """Launch the fused kernel over inputs already packed by chunk. + + ``q/k/g`` are ``[BH*NC, C*DK]``, ``beta`` is ``[BH*NC, C]``, ``v``/``o`` are + ``[BH*NC, C*DV]``, and ``h0``/``ht`` are ``[BH, DV, DK]`` fp32 (the state + kept transposed, which is the orientation both of its consumers want). + """ + if h0 is None: + assert not spec.has_initial_state, "h0=None but the kernel reads h0_ptr" + launcher = make_launcher(spec) + if stream is None: + stream = torch.cuda.current_stream().cuda_stream + cfg = LaunchConfig( + grid=kda_chunk_fused_grid(spec, bh), + block=(spec.tile.block_size, 1, 1), + stream=stream, + ) + launcher( + { + "q_ptr": q, + "k_ptr": k, + "g_ptr": g, + "beta_ptr": beta, + "v_ptr": v, + "o_ptr": o, + # unread when the kernel was built with has_initial_state=False + "h0_ptr": ht if h0 is None else h0, + "ht_ptr": ht, + "scale": float(scale), + "nc": int(nc), + }, + config=cfg, + ) + + +# --------------------------------------------------------------------- +# reference + parity +# --------------------------------------------------------------------- + + +def ref_token_serial(q, k, v, g, beta, scale, h0=None): + """Token-serial float64 oracle for the gated delta rule. + + Per token, with ``S`` the ``DK x DV`` state: + + .. code-block:: text + + S <- Diag(exp(g_t)) S + u <- beta_t (v_t - k_t^T S) + S <- S + k_t u^T + o_t = scale * q_t^T S + + Deliberately not chunked: this is the recurrence the chunkwise + factorization claims to compute, so it is an independent check of that + claim rather than a re-derivation of it. + """ + B, H, T, DK = q.shape + DV = v.shape[-1] + qd, kd, vd = (x.double() for x in (q, k, v)) + gd, bd = g.double(), beta.double() + o = torch.zeros(B, H, T, DV, dtype=torch.float64, device=q.device) + S = ( + torch.zeros(B, H, DK, DV, dtype=torch.float64, device=q.device) + if h0 is None + else h0.double().clone() + ) + for t in range(T): + S = S * gd[:, :, t, :].exp().unsqueeze(-1) + kt = kd[:, :, t, :] + kv = torch.einsum("bhd,bhde->bhe", kt, S) + u = bd[:, :, t].unsqueeze(-1) * (vd[:, :, t, :] - kv) + S = S + kt.unsqueeze(-1) * u.unsqueeze(-2) + o[:, :, t, :] = scale * torch.einsum("bhd,bhde->bhe", qd[:, :, t, :], S) + return o, S + + +def make_inputs(B, H, T, DK, DV, gate_low=-0.5, seed=0, device="cuda"): + gen = torch.Generator(device=device).manual_seed(seed) + kw = dict(device=device, generator=gen) + q = torch.nn.functional.normalize( + torch.randn(B, H, T, DK, dtype=torch.float32, **kw), dim=-1 + ).bfloat16() + k = torch.nn.functional.normalize( + torch.randn(B, H, T, DK, dtype=torch.float32, **kw), dim=-1 + ).bfloat16() + v = (torch.randn(B, H, T, DV, dtype=torch.float32, **kw) * 0.2).bfloat16() + g = gate_low * torch.rand(B, H, T, DK, dtype=torch.float32, **kw) + beta = torch.rand(B, H, T, dtype=torch.float32, **kw) + return [x.contiguous() for x in (q, k, v, g, beta)] + + +def launch_packed(spec, q, k, v, g, beta, h0=None): + """Reshape dense [B,H,T,D] inputs into the kernel's chunk-major views.""" + B, H, T, DK = q.shape + DV = v.shape[-1] + C = spec.tile.chunk + BH, NC = B * H, T // C + nt = BH * NC + o = torch.empty(nt, C * DV, dtype=torch.bfloat16, device=q.device) + ht = torch.zeros(BH, DV, DK, dtype=torch.float32, device=q.device) + h0t = None if h0 is None else h0.transpose(-1, -2).contiguous().view(BH, DV, DK) + run_fused( + spec, + q.view(nt, C * DK), + k.view(nt, C * DK), + g.view(nt, C * DK), + beta.view(nt, C), + v.view(nt, C * DV), + o, + ht, + DK**-0.5, + BH, + NC, + h0=h0t, + ) + return o.view(B, H, T, DV), ht.view(B, H, DV, DK).transpose(-1, -2) + + +def check(spec, B, H, T, gate_low=-0.5, tol=TOL, with_h0=False, verbose=True): + DK, DV = spec.head_k, spec.head_v + q, k, v, g, beta = make_inputs(B, H, T, DK, DV, gate_low) + h0 = None + if with_h0: + gen = torch.Generator(device="cuda").manual_seed(7) + h0 = 0.1 * torch.randn( + B, H, DK, DV, dtype=torch.float32, device="cuda", generator=gen + ) + o_got, ht_got = launch_packed(spec, q, k, v, g, beta, h0=h0) + torch.cuda.synchronize() + o_ref, s_ref = ref_token_serial(q, k, v, g, beta, DK**-0.5, h0=h0) + + worst = 0.0 + for name, got, ref in (("out", o_got, o_ref), ("final_state", ht_got, s_ref)): + d = (got.double() - ref).abs().max().item() + den = max(ref.abs().max().item(), 1e-30) + rel = d / den + worst = max(worst, rel) + if verbose: + print( + f" {name:12s} max_abs={d:.3e} rel={rel:.3e} " + f"ref_absmax={den:.3e} finite={bool(torch.isfinite(got).all())}" + ) + verdict = "PASS" if worst <= tol else "FAIL" + if verbose: + print(f" -> worst rel {worst:.3e} tol {tol:.1e} {verdict}") + return worst + + +def bench(spec, B, H, T, warmup=10, iters=30): + DK, DV = spec.head_k, spec.head_v + q, k, v, g, beta = make_inputs(B, H, T, DK, DV) + C = spec.tile.chunk + BH, NC = B * H, T // C + nt = BH * NC + o = torch.empty(nt, C * DV, dtype=torch.bfloat16, device="cuda") + ht = torch.zeros(BH, DV, DK, dtype=torch.float32, device="cuda") + args = ( + spec, + q.view(nt, C * DK), + k.view(nt, C * DK), + g.view(nt, C * DK), + beta.view(nt, C), + v.view(nt, C * DV), + o, + ht, + DK**-0.5, + BH, + NC, + ) + ms = time_launches(lambda: run_fused(*args), warmup=warmup, iters=iters) + # Essential traffic only: q/k/v in, g/beta in, o out. The fused kernel adds + # nothing else, which is the whole point of the design. + rd = nt * (2 * C * DK * 2 + C * DV * 2 + C * DK * 4 + C * 4) + wr = nt * C * DV * 2 + gbps = (rd + wr) / (ms * 1e-3) / 2**30 + print(f" {ms:.4f} ms {gbps:.1f} GiB/s (B={B} H={H} T={T}, {nt} chunks)") + return ms + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--shapes", default="8x8x1024,8x16x2048,32x16x2048") + ap.add_argument("--check-shape", default="2x4x256") + ap.add_argument("--chunk", type=int, default=32) + ap.add_argument("--dk", type=int, default=128) + ap.add_argument("--dv", type=int, default=128) + ap.add_argument("--sb", type=int, default=8) + ap.add_argument("--no-check", action="store_true") + ap.add_argument("--iters", type=int, default=30) + args = ap.parse_args() + + spec = KdaChunkFusedSpec( + head_k=args.dk, + head_v=args.dv, + tile=KdaTileSpec( + chunk=args.chunk, + block_size=512, + pad_dk=16, + pad_cb=16, + tile_atom_m=16, + solve_block=args.sb, + scan_atom_m=16, + ), + ) + ok, why = is_valid_fused_spec(spec) + print(f"{spec.kernel_name()} lds={spec.lds_bytes()} B valid={ok} {why}") + if not ok: + return 1 + + worst = 0.0 + if not args.no_check: + B, H, T = (int(x) for x in args.check_shape.split("x")) + for gate_low in (-0.1, -0.5, -2.0, -5.0): + for with_h0 in (False, True): + check_spec = ( + spec + if not with_h0 + else dataclasses.replace(spec, has_initial_state=True) + ) + state_label = " with h0" if with_h0 else "" + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]{state_label}:") + worst = max( + worst, + check( + check_spec, + B, + H, + T, + gate_low=gate_low, + with_h0=with_h0, + ), + ) + + for s in args.shapes.split(","): + B, H, T = (int(x) for x in s.split("x")) + bench(spec, B, H, T, iters=args.iters) + + return 0 if args.no_check or worst <= TOL else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py new file mode 100644 index 000000000000..e4f218434478 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py @@ -0,0 +1,347 @@ +#!/usr/bin/env python3 +"""Host-side builder for the chunkwise KDA per-chunk tile kernel +(``kernels/gfx950/kda_chunkwise.py``). + +Owns the host path: spec construction, compilation, ABI signature, workspace +allocation and launch -- plus a torch parity check of all six output tiles and a +throughput report. + +The six tiles this kernel emits are exactly the state-independent part of the +chunkwise KDA body, so it can be validated (and benchmarked) on its own, ahead of +the serial state scan that consumes them. + +Usage: + python kda_chunk_prep.py # parity + bench, default shapes + python kda_chunk_prep.py --shapes 8x16x2048 + python kda_chunk_prep.py --chunk 32 --no-check +""" + +import argparse +import os +import sys + +try: + import rocke # noqa: F401 +except ImportError: # running as a bare script outside the editable install + _HERE = os.path.dirname(__file__) + _RK = os.path.abspath(os.path.join(_HERE, "../../../..")) + sys.path[:0] = [_RK + "/library", _RK + "/platform/python"] + +import torch # noqa: E402 + +from kernels.gfx950.kda_chunkwise import ( # noqa: E402 + KdaChunkPrepSpec, + KdaTileSpec, + build_kda_chunk_prep, + is_valid_spec, + kda_chunk_prep_grid, + kda_chunk_prep_signature, +) +from rocke.helpers.compile import compile_kernel # noqa: E402 +from rocke.runtime import KernelLauncher, LaunchConfig, time_launches # noqa: E402 + +_LAUNCHER_CACHE = {} + + +def make_launcher(spec: KdaChunkPrepSpec) -> KernelLauncher: + key = spec.kernel_name() + if key not in _LAUNCHER_CACHE: + ok, why = is_valid_spec(spec) + if not ok: + raise ValueError(f"unsupported spec: {why}") + art = compile_kernel( + build_kda_chunk_prep(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + _LAUNCHER_CACHE[key] = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=kda_chunk_prep_signature(spec), + ) + return _LAUNCHER_CACHE[key] + + +def alloc_tiles(num_tiles: int, spec: KdaChunkPrepSpec, device="cuda"): + """The six per-chunk tiles, carved from one pool. + + Six separate allocations show up on the host critical path at short + sequences, and the tiles are always produced and consumed together. + """ + C, DK = spec.tile.chunk, spec.head_k + shapes = [ + ("a", (num_tiles, C, C)), + ("gk", (num_tiles, C, DK)), + ("gq", (num_tiles, C, DK)), + ("aqk", (num_tiles, C, C)), + ("kt", (num_tiles, DK, C)), + ] + counts = [s[0] * s[1] * s[2] for _, s in shapes] + pool = torch.empty(sum(counts), dtype=torch.bfloat16, device=device) + ws, off = {}, 0 + for (name, shape), cnt in zip(shapes, counts): + ws[name] = pool[off : off + cnt].view(*shape) + off += cnt + ws["dec"] = torch.empty(num_tiles, DK, dtype=torch.float32, device=device) + ws["_pool"] = pool + return ws + + +def run_prep( + spec, + q, + k, + g, + beta, + ws, + scale, + stream=None, + *, + batch=None, + heads=None, + tseq=None, + nc=None, + a_log=None, + dt_bias=None, +): + """Launch the prep kernel over ``q/k/g/beta`` already packed by chunk or raw.""" + num_tiles = q.shape[0] if not spec.raw_inputs else batch * heads * nc + launcher = make_launcher(spec) + if stream is None: + stream = torch.cuda.current_stream().cuda_stream + cfg = LaunchConfig( + grid=kda_chunk_prep_grid(spec, num_tiles), + block=(spec.tile.block_size, 1, 1), + stream=stream, + ) + args = { + "q_ptr": q, + "k_ptr": k, + "g_ptr": g, + "beta_ptr": beta, + "a_ptr": ws["a"], + "gk_ptr": ws["gk"], + "gq_ptr": ws["gq"], + "aqk_ptr": ws["aqk"], + "kt_ptr": ws["kt"], + "dec_ptr": ws["dec"], + "scale": float(scale), + } + if spec.raw_inputs: + if dt_bias is None: + dt_bias = torch.zeros( + heads * spec.head_k, dtype=torch.float32, device=q.device + ) + args.update( + { + "a_log_ptr": a_log, + "dt_bias_ptr": dt_bias, + "batch": int(batch), + "heads": int(heads), + "tseq": int(tseq), + "nc": int(nc), + } + ) + launcher(args, config=cfg) + + +# --------------------------------------------------------------------- +# reference + parity +# --------------------------------------------------------------------- + + +def ref_prep_tiles(q, k, g, beta, scale, nt_block=32): + """float64 torch oracle for the six tiles. q/k/g: [NT, C, DK], beta: [NT, C]. + + ``Akk`` and ``Aqk`` are formed from the pairwise exponent *difference* + ``Gcum_i - Gcum_j`` rather than the ratio ``Gamma_i / Gamma_j``, so the + oracle never materializes the overflowing intermediate the kernel is + designed to avoid -- and it does so by a different route (explicit + ``C x C x DK`` outer product, no midpoint factoring), which is what makes it + an independent check of the factoring rather than a restatement of it. + + Only the retained triangle is meaningful: for ``i >= j`` the exponent is + already ``<= 0``, so clamping it there is exact and keeps the discarded + upper triangle from going to inf. Blocked over NT to bound the + ``C x C x DK`` temporary. + """ + NT, C, _ = q.shape + ar = torch.arange(C, device=q.device) + lower = (ar[:, None] > ar[None, :]).unsqueeze(0) + lower_eq = (ar[:, None] >= ar[None, :]).unsqueeze(0) + zero = torch.zeros((), device=q.device, dtype=torch.float64) + + out = {n: [] for n in ("a", "gk", "gq", "aqk", "kt", "dec")} + for s in range(0, NT, nt_block): + e = min(s + nt_block, NT) + qb, kb = q[s:e].double(), k[s:e].double() + gcum = g[s:e].double().cumsum(dim=1) + gl = gcum[:, -1:, :] + gam = gcum.exp() + bb = beta[s:e].double() + + # [n, i, j, d] pairwise exponent, clamped where it is discarded anyway. + dif = (gcum[:, :, None, :] - gcum[:, None, :, :]).clamp(max=0.0).exp() + akk = (kb[:, :, None, :] * kb[:, None, :, :] * dif).sum(-1) + # ``dif`` already carries the whole Gamma_i / Gamma_j ratio, so the + # query side enters as bare q * scale -- folding Gamma_i in via ``gq`` + # here would count the decay twice. + aqk = (qb[:, :, None, :] * scale * kb[:, None, :, :] * dif).sum(-1) + gq = qb * gam * scale + + tprime = torch.where(lower, bb[:, :, None] * akk, zero) + eye = torch.eye(C, device=q.device, dtype=torch.float64).expand_as(tprime) + out["a"].append( + torch.linalg.solve_triangular( + eye + tprime, torch.diag_embed(bb), upper=False + ) + ) + out["aqk"].append(torch.where(lower_eq, aqk, zero)) + out["gk"].append(kb * gam) + out["gq"].append(gq) + out["kt"].append((kb * (gl - gcum).exp()).transpose(1, 2)) + out["dec"].append(gl.squeeze(1).exp()) + return {n: torch.cat(v, 0) for n, v in out.items()} + + +def make_inputs(num_tiles, C, DK, gate_low=-0.5, seed=0, device="cuda"): + gen = torch.Generator(device=device).manual_seed(seed) + kw = dict(device=device, generator=gen) + q = torch.nn.functional.normalize( + torch.randn(num_tiles, C, DK, dtype=torch.float32, **kw), dim=-1 + ).to(torch.bfloat16) + k = torch.nn.functional.normalize( + torch.randn(num_tiles, C, DK, dtype=torch.float32, **kw), dim=-1 + ).to(torch.bfloat16) + # Per-channel log decay in the reference range (gate_lower_bound = -5), which is + # what stresses the factored C x C construction. + g = gate_low * torch.rand(num_tiles, C, DK, dtype=torch.float32, **kw) + beta = torch.rand(num_tiles, C, dtype=torch.float32, **kw) + return q, k, g, beta + + +def check(spec, num_tiles, gate_low=-0.5, tol=2e-2, verbose=True): + C, DK = spec.tile.chunk, spec.head_k + q, k, g, beta = make_inputs(num_tiles, C, DK, gate_low) + scale = DK**-0.5 + ws = alloc_tiles(num_tiles, spec) + run_prep( + spec, + q.reshape(num_tiles, -1), + k.reshape(num_tiles, -1), + g.reshape(num_tiles, -1), + beta, + ws, + scale, + ) + torch.cuda.synchronize() + + ref = ref_prep_tiles(q.float(), k.float(), g, beta, scale) + pairs = [(n, ws[n].float(), ref[n].float()) for n in ref] + # A's diagonal is beta and dominates its magnitude, so scoring the whole + # tile against its own absmax would hide the off-diagonal terms -- which are + # the part that actually came out of the solve. Score them separately. + pairs.append(("a_low", ws["a"].float().tril(-1), ref["a"].float().tril(-1))) + + worst, worst_name = 0.0, "" + for name, got, exp in pairs: + d = (got - exp).abs() + denom = exp.abs().max().clamp_min(1e-30) + rel = (d.max() / denom).item() + if verbose: + print( + f" {name:<5} max_abs={d.max().item():.3e} rel={rel:.3e} " + f"ref_absmax={exp.abs().max().item():.3e} " + f"finite={torch.isfinite(got).all().item()}" + ) + if rel > worst: + worst, worst_name = rel, name + ok = worst <= tol + print( + f" -> worst rel {worst:.3e} on {worst_name!r} tol {tol:.1e} " + f"{'PASS' if ok else 'FAIL'}" + ) + return ok + + +def bench(spec, num_tiles, warmup=15, iters=50): + C, DK = spec.tile.chunk, spec.head_k + q, k, g, beta = make_inputs(num_tiles, C, DK) + ws = alloc_tiles(num_tiles, spec) + args = ( + spec, + q.reshape(num_tiles, -1), + k.reshape(num_tiles, -1), + g.reshape(num_tiles, -1), + beta, + ws, + DK**-0.5, + ) + # time_launches runs the loop under no_fence and owns the drain. Timing + # fenced launches instead would charge every iteration a host stream sync + # (LaunchConfig.fence defaults True), which is not how the kernel runs in a + # pipeline: the scan consumes the tiles on the same stream, so FIFO order + # already covers the dependency. At short sequences the sync is a large + # fraction of the measurement. + ms = time_launches(lambda: run_prep(*args), warmup=warmup, iters=iters) + + tile_el = 2 * C * C + 2 * C * DK + DK * C + rd = num_tiles * (2 * C * DK * 2 + C * DK * 4 + C * 4) + wr = num_tiles * (tile_el * 2 + DK * 4) + gbps = (rd + wr) / (ms * 1e-3) / 2**30 + print(f" {ms:.4f} ms {gbps:.1f} GiB/s ({num_tiles} chunks)") + return ms + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--shapes", + default="1x8x2048,8x8x1024,8x16x2048", + help="BxHxT triples; chunk count is B*H*T/chunk", + ) + ap.add_argument("--chunk", type=int, default=32) + ap.add_argument("--dk", type=int, default=128) + ap.add_argument("--dv", type=int, default=128) + ap.add_argument("--pad-dk", type=int, default=8) + ap.add_argument("--pad-c", type=int, default=4) + ap.add_argument("--wpe", type=int, default=0) + ap.add_argument("--no-check", action="store_true") + ap.add_argument("--iters", type=int, default=50) + args = ap.parse_args() + + spec = KdaChunkPrepSpec( + head_k=args.dk, + head_v=args.dv, + tile=KdaTileSpec( + chunk=args.chunk, + block_size=2 * args.dk, + pad_dk=args.pad_dk, + pad_c=args.pad_c, + waves_per_eu=args.wpe, + ), + ) + ok, why = is_valid_spec(spec) + print(f"{spec.kernel_name()} lds={spec.lds_bytes()} B valid={ok} {why}") + if not ok: + return 1 + + failed = False + if not args.no_check: + # Sweep the gate range: near-zero is the typical regime, -5 is the + # reference gate lower bound and saturates the factored exponents. + for gl in (-0.1, -0.5, -2.0, -5.0): + print(f"parity (128 chunks, gate in [{gl}, 0]):") + failed |= not check(spec, 128, gate_low=gl) + + for s in args.shapes.split(","): + B, H, T = (int(x) for x in s.split("x")) + nt = B * H * (T // args.chunk) + print(f"bench B={B} H={H} T={T}:") + bench(spec, nt, iters=args.iters) + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_split.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_split.py new file mode 100644 index 000000000000..889f82a8d78e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_split.py @@ -0,0 +1,521 @@ +#!/usr/bin/env python3 +"""Host builder for the split-path chunkwise KDA forward on gfx950. + +Two kernels: :func:`build_kda_chunk_prep` builds the six per-chunk tiles, one +workgroup per chunk and fully parallel over the sequence, then +:func:`build_kda_chunk_scan` walks each (batch, head)'s chunks in order, +staging that chunk's tiles from HBM into LDS and running the state recurrence. + +The alternative is the fused kernel in ``kda_chunk_fused.py``, which keeps the +tiles in LDS and never writes them out. That is strictly less traffic and +measurably slower: the combined footprint puts one workgroup over half the LDS +budget, so the scan's latency-bound chain of small matmuls runs with no second +workgroup to cover it. This path pays the tile round trip to keep two resident. + +Correctness is checked end to end against the same token-serial float64 oracle +the fused builder uses, so the two paths are held to one contract. + +Run directly for a parity sweep plus a benchmark:: + + python kda_chunk_split.py # parity + bench + python kda_chunk_split.py --no-check # bench only +""" + +from __future__ import annotations + +import argparse +import dataclasses +import os +import sys + +try: + import rocke # noqa: F401 +except ImportError: # running as a bare script outside the editable install + _HERE = os.path.dirname(__file__) + _RK = os.path.abspath(os.path.join(_HERE, "../../../..")) + sys.path[:0] = [_RK + "/library", _RK + "/platform/python"] + +import torch # noqa: E402 + +from kernels.gfx950.kda_chunkwise import ( # noqa: E402 + KdaChunkPrepSpec, + KdaChunkScanSpec, + KdaTileSpec, + build_kda_chunk_scan, + is_valid_scan_spec, + is_valid_spec, + kda_chunk_scan_grid, + kda_chunk_scan_signature, +) +from rocke.helpers.compile import compile_kernel # noqa: E402 +from rocke.runtime import KernelLauncher, LaunchConfig, time_launches # noqa: E402 + +if __package__: + from . import kda_chunk_prep as prep_mod # noqa: E402 + from .kda_chunk_fused import make_inputs, ref_token_serial # noqa: E402 +else: + import kda_chunk_prep as prep_mod # noqa: E402 + from kda_chunk_fused import make_inputs, ref_token_serial # noqa: E402 + +_LAUNCHER_CACHE: dict = {} +TOL = 3e-2 + + +def make_launcher(spec: KdaChunkScanSpec) -> KernelLauncher: + key = spec.kernel_name() + if key not in _LAUNCHER_CACHE: + ok, why = is_valid_scan_spec(spec) + if not ok: + raise ValueError(f"unsupported spec: {why}") + art = compile_kernel( + build_kda_chunk_scan(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + _LAUNCHER_CACHE[key] = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=kda_chunk_scan_signature(spec), + ) + return _LAUNCHER_CACHE[key] + + +def prep_spec_of(spec: KdaChunkScanSpec, *, raw: bool = False) -> KdaChunkPrepSpec: + """The tile builder that produces this scan's inputs. + + Both halves are derived from one scan spec so the tile layouts cannot drift + apart: the scan's staging copies assume exactly the flat per-chunk layouts + the prep kernel's global sink writes. + + Raw prep keeps the default 256-thread tile builder even when the scan uses a + narrower block for ``value_splits > 1``. + """ + raw_kw = {} + tile = spec.tile + if raw: + raw_kw = dict( + raw_inputs=True, + fuse_qk_l2norm=True, + fuse_gate=True, + fuse_beta_sigmoid=True, + has_dt_bias=True, + lower_bound=-5.0, + ) + tile = dataclasses.replace(spec.tile, block_size=256) + return KdaChunkPrepSpec( + head_k=spec.head_k, + head_v=spec.head_v, + dtype=spec.dtype, + tile=tile, + **raw_kw, + ) + + +def aligned_split_specs(value_splits: int = 1): + """Default aligned-C32 split family knobs for the ten benchmark shapes.""" + if value_splits == 8: + tile = KdaTileSpec(chunk=32, block_size=64, scan_atom_m=16) + elif value_splits == 4: + tile = KdaTileSpec(chunk=32, block_size=64) + elif value_splits == 2: + tile = KdaTileSpec(chunk=32, block_size=128) + else: + tile = KdaTileSpec(chunk=32, block_size=256) + scan = KdaChunkScanSpec( + tile=tile, + value_splits=value_splits, + token_major_io=True, + ) + return scan, prep_spec_of(scan, raw=True) + + +def run_scan( + spec, + ws, + v, + o, + ht, + bh, + nc, + h0=None, + stream=None, + *, + batch=None, + heads=None, + tseq=None, +): + """Launch the scan over tiles already materialized by the prep kernel.""" + if h0 is None: + assert not spec.has_initial_state, "h0=None but the kernel reads h0_ptr" + launcher = make_launcher(spec) + if stream is None: + stream = torch.cuda.current_stream().cuda_stream + cfg = LaunchConfig( + grid=kda_chunk_scan_grid(spec, bh), + block=(spec.tile.block_size, 1, 1), + stream=stream, + ) + args = { + "a_ptr": ws["a"], + "gk_ptr": ws["gk"], + "gq_ptr": ws["gq"], + "aqk_ptr": ws["aqk"], + "kt_ptr": ws["kt"], + "dec_ptr": ws["dec"], + "v_ptr": v, + "o_ptr": o, + # unread when the kernel was built with has_initial_state=False + "h0_ptr": ht if h0 is None else h0, + "ht_ptr": ht, + "nc": int(nc), + } + if spec.token_major_io: + args.update({"batch": int(batch), "heads": int(heads), "tseq": int(tseq)}) + launcher(args, config=cfg) + + +def run_split( + spec, + q, + k, + g, + beta, + v, + o, + ws, + ht, + scale, + bh, + nc, + h0=None, + *, + raw=False, + batch=None, + heads=None, + tseq=None, + a_log=None, + dt_bias=None, +): + """Both kernels, back to back on one stream. + + No fence between them: they run in FIFO order on the same stream, so the + scan already sees the tiles the prep kernel wrote. + """ + pspec = prep_spec_of(spec, raw=raw) + prep_mod.run_prep( + pspec, + q, + k, + g, + beta, + ws, + scale, + batch=batch, + heads=heads, + tseq=tseq, + nc=nc, + a_log=a_log, + dt_bias=dt_bias, + ) + run_scan( + spec, + ws, + v, + o, + ht, + bh, + nc, + h0=h0, + batch=batch, + heads=heads, + tseq=tseq, + ) + + +def run_raw_split( + spec, + q, + k, + v, + g, + beta, + a_log, + dt_bias, + o, + ht, + scale, + batch, + heads, + tseq, + h0=None, +): + """Aligned raw token-major split path: fused prep + value-split scan.""" + C = spec.tile.chunk + BH, NC = batch * heads, tseq // C + nt = BH * NC + ws = prep_mod.alloc_tiles(nt, prep_spec_of(spec, raw=True)) + run_split( + spec, + q, + k, + g, + beta, + v, + o, + ws, + ht, + scale, + BH, + NC, + h0=h0, + raw=True, + batch=batch, + heads=heads, + tseq=tseq, + a_log=a_log, + dt_bias=dt_bias, + ) + return ws + + +def launch_raw( + spec, + q, + k, + v, + g, + beta, + a_log, + dt_bias, + h0=None, +): + """Token-major [B,T,H,D] in/out with V-first final state.""" + B, T, H, DK = q.shape + DV = v.shape[-1] + o = torch.empty_like(v) + ht = torch.zeros(B * H, DV, DK, dtype=torch.float32, device=q.device) + h0t = None + if h0 is not None: + h0t = h0.transpose(-1, -2).contiguous().view(B * H, DV, DK) + run_raw_split( + spec, + q, + k, + v, + g, + beta, + a_log, + dt_bias, + o, + ht, + DK**-0.5, + B, + H, + T, + h0=h0t, + ) + return o, ht.view(B, H, DV, DK).transpose(-1, -2) + + +def ref_aligned_raw( + q, + k, + v, + g, + beta, + a_log, + dt_bias, + scale, + lower_bound=-5.0, + h0=None, +): + """Float64 oracle for the aligned raw contract (Aiter preprocessing + KDA).""" + import torch.nn.functional as F + + qn = F.normalize(q.float(), dim=-1) + kn = F.normalize(k.float(), dim=-1) + B, T, H, DK = q.shape + db = dt_bias.view(H, DK) + al = a_log.view(H) + gf = g.float() + gate = lower_bound * torch.sigmoid( + torch.exp(al)[None, None, :, None] * (gf + db[None, None, :, :]) + ) + bb = torch.sigmoid(beta.float()) + qbh = qn.permute(0, 2, 1, 3).to(torch.bfloat16) + kbh = kn.permute(0, 2, 1, 3).to(torch.bfloat16) + vbh = v.permute(0, 2, 1, 3) + gbh = gate.permute(0, 2, 1, 3) + bbh = bb.permute(0, 2, 1) + return ref_token_serial(qbh, kbh, vbh, gbh, bbh, scale, h0=h0) + + +# --------------------------------------------------------------------- +# parity +# --------------------------------------------------------------------- + + +def launch_packed(spec, q, k, v, g, beta, h0=None): + """Reshape dense [B,H,T,D] inputs into the kernels' chunk-major views.""" + B, H, T, DK = q.shape + DV = v.shape[-1] + C = spec.tile.chunk + BH, NC = B * H, T // C + nt = BH * NC + ws = prep_mod.alloc_tiles(nt, prep_spec_of(spec)) + o = torch.empty(nt, C * DV, dtype=torch.bfloat16, device=q.device) + ht = torch.zeros(BH, DV, DK, dtype=torch.float32, device=q.device) + h0t = None if h0 is None else h0.transpose(-1, -2).contiguous().view(BH, DV, DK) + run_split( + spec, + q.view(nt, C * DK), + k.view(nt, C * DK), + g.view(nt, C * DK), + beta.view(nt, C), + v.view(nt, C * DV), + o, + ws, + ht, + DK**-0.5, + BH, + NC, + h0=h0t, + ) + return o.view(B, H, T, DV), ht.view(B, H, DV, DK).transpose(-1, -2) + + +def check(spec, B, H, T, gate_low=-0.5, tol=TOL, with_h0=False, verbose=True): + DK, DV = spec.head_k, spec.head_v + q, k, v, g, beta = make_inputs(B, H, T, DK, DV, gate_low) + h0 = None + if with_h0: + gen = torch.Generator(device="cuda").manual_seed(7) + h0 = 0.1 * torch.randn( + B, H, DK, DV, dtype=torch.float32, device="cuda", generator=gen + ) + o_got, ht_got = launch_packed(spec, q, k, v, g, beta, h0=h0) + torch.cuda.synchronize() + o_ref, s_ref = ref_token_serial(q, k, v, g, beta, DK**-0.5, h0=h0) + + worst = 0.0 + for name, got, ref in (("out", o_got, o_ref), ("final_state", ht_got, s_ref)): + d = (got.double() - ref).abs().max().item() + den = max(ref.abs().max().item(), 1e-30) + rel = d / den + worst = max(worst, rel) + if verbose: + print( + f" {name:12s} max_abs={d:.3e} rel={rel:.3e} " + f"ref_absmax={den:.3e} finite={bool(torch.isfinite(got).all())}" + ) + verdict = "PASS" if worst <= tol else "FAIL" + if verbose: + print(f" -> worst rel {worst:.3e} tol {tol:.1e} {verdict}") + return worst + + +def bench(spec, B, H, T, warmup=10, iters=30, split=True): + """Time the pair, and each half alone, on one shape. + + The halves are reported too because they answer different questions: the + prep time is bandwidth against a known roofline, while the scan time is a + serial chain whose only lever is occupancy. + """ + DK, DV = spec.head_k, spec.head_v + q, k, v, g, beta = make_inputs(B, H, T, DK, DV) + C = spec.tile.chunk + BH, NC = B * H, T // C + nt = BH * NC + ws = prep_mod.alloc_tiles(nt, prep_spec_of(spec)) + o = torch.empty(nt, C * DV, dtype=torch.bfloat16, device="cuda") + ht = torch.zeros(BH, DV, DK, dtype=torch.float32, device="cuda") + qf, kf, gf = (x.view(nt, C * DK) for x in (q, k, g)) + bf, vf = beta.view(nt, C), v.view(nt, C * DV) + scale = DK**-0.5 + + def do_prep(): + prep_mod.run_prep(prep_spec_of(spec), qf, kf, gf, bf, ws, scale) + + def do_scan(): + run_scan(spec, ws, vf, o, ht, BH, NC) + + def do_both(): + do_prep() + do_scan() + + ms = time_launches(do_both, warmup=warmup, iters=iters) + # Essential traffic only: q/k/v/g/beta in, o out. The tile round trip this + # path adds on purpose is excluded, so the number is comparable to the + # fused kernel's and to FlyDSL's. + rd = nt * (2 * C * DK * 2 + C * DV * 2 + C * DK * 4 + C * 4) + wr = nt * C * DV * 2 + gbps = (rd + wr) / (ms * 1e-3) / 2**30 + print(f" {ms:.4f} ms {gbps:.1f} GiB/s (B={B} H={H} T={T}, {nt} chunks)") + if split: + p = time_launches(do_prep, warmup=warmup, iters=iters) + s = time_launches(do_scan, warmup=warmup, iters=iters) + print(f" prep {p:.4f} ms scan {s:.4f} ms sum {p + s:.4f} ms") + return ms + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--shapes", default="8x8x1024,8x16x2048,32x16x2048") + ap.add_argument("--check-shape", default="2x4x256") + ap.add_argument("--chunk", type=int, default=32) + ap.add_argument("--dk", type=int, default=128) + ap.add_argument("--dv", type=int, default=128) + ap.add_argument("--sb", type=int, default=8) + ap.add_argument("--no-check", action="store_true") + ap.add_argument("--iters", type=int, default=30) + args = ap.parse_args() + + spec = KdaChunkScanSpec( + head_k=args.dk, + head_v=args.dv, + tile=KdaTileSpec(chunk=args.chunk, solve_block=args.sb), + ) + ok, why = is_valid_scan_spec(spec) + print(f"{spec.kernel_name()} lds={spec.lds_bytes()} B valid={ok} {why}") + if not ok: + return 1 + pspec = prep_spec_of(spec) + ok, why = is_valid_spec(pspec) + print(f"{pspec.kernel_name()} lds={pspec.lds_bytes()} B valid={ok} {why}") + if not ok: + return 1 + + worst = 0.0 + if not args.no_check: + B, H, T = (int(x) for x in args.check_shape.split("x")) + for gate_low in (-0.1, -0.5, -2.0, -5.0): + for with_h0 in (False, True): + check_spec = ( + spec + if not with_h0 + else dataclasses.replace(spec, has_initial_state=True) + ) + state_label = " with h0" if with_h0 else "" + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]{state_label}:") + worst = max( + worst, + check( + check_spec, + B, + H, + T, + gate_low=gate_low, + with_h0=with_h0, + ), + ) + + for s in args.shapes.split(","): + B, H, T = (int(x) for x in s.split("x")) + bench(spec, B, H, T, iters=args.iters) + + return 0 if worst <= TOL else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/__init__.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/__init__.py index 5fd9be29ce75..c835dd88043e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/__init__.py @@ -3,7 +3,7 @@ """Chunkwise Kimi Delta Attention dispatcher family (path-level selection). -Backed by :mod:`kernels.gfx942.kda_chunkwise`. This module owns only the +Backed by the gfx942 and gfx950 KDA kernels. This module owns only the assembly -- the registry, the entry point, and the re-exports that make ``dispatch.kda`` one import for callers. What each candidate *is* lives in the arch module that owns it. See :mod:`.common` for the scope of the dispatch @@ -29,7 +29,7 @@ stable_json_hash, ) -from . import gfx942 +from . import gfx942, gfx950 from .common import ( FAMILY, KDA_ABI_VERSION, @@ -40,7 +40,6 @@ KDA_PARTITION_HEAD_V, KdaRequest, _request_errors, - _selector_matches, ) _FAMILY = FAMILY @@ -51,7 +50,7 @@ KDA_REGISTRY = CandidateRegistry( _FAMILY, dim_vocabulary=KDA_DIM_VOCABULARY, require_build=True ) -for _module in (gfx942,): +for _module in (gfx942, gfx950): _module.register(KDA_REGISTRY) @@ -124,6 +123,7 @@ def dispatch_kda(req: KdaRequest, *, ranker: Ranker | None = None) -> DispatchRe f"algorithm={candidate.algorithm}", f"spec_id={candidate.spec_id}", f"v_partitions={req.v_partitions} workgroups={req.workgroups}", + f"chunks={req.num_chunks}", f"spec_hash={kid.spec_hash}", f"request_hash={kid.request_hash}", ), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/common.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/common.py index 7110308937f6..94db76fb8720 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/common.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/common.py @@ -4,39 +4,33 @@ """Shared pieces of the KDA family: the request and the gates every candidate re-uses. -Arch-neutral by construction even though the only arch today is gfx942: the -arch module imports this one and nothing here imports it, which keeps the +Arch-neutral by construction: each arch module imports this one and nothing +here imports an arch module, which keeps the registry assembly in ``__init__`` free of import-order dependence. SCOPE -- what this dispatcher decides ------------------------------------- Which of the two chunkwise prefill paths runs: the **fused** kernel (one -workgroup per (batch, head, V partition) walks that head's chunks, keeping the -six per-chunk tiles in LDS), or the **split** path (a tile-builder kernel that -is one workgroup per chunk, then a state-scan kernel over the materialized -tiles). Both compute the same recurrence; they trade HBM traffic against -parallelism. See :mod:`kernels.gfx942.kda_chunkwise`. - -The request carries the *logical* value width ``head_v`` (128 for a Kimi Delta -Attention head). The partition into ``KDA_PARTITION_HEAD_V``-channel workgroups -is an implementation fact of gfx942's LDS budget, so the candidates apply it -when they build the spec and when they report their grid -- exactly as -``builders/gfx942/kda`` does on the host. +workgroup per recurrence stream walks that stream's chunks, keeping the six +per-chunk tiles in LDS), or the **split** path (a tile-builder kernel that is +one workgroup per chunk, then a state-scan kernel over the materialized tiles). +Both compute the same recurrence and trade HBM traffic against parallelism. + +The request carries the *logical* value width ``head_v``. gfx942 partitions a +logical head into ``KDA_PARTITION_HEAD_V``-channel workgroups to fit its LDS +budget; gfx950 keeps the whole value width in one recurrence stream and may +apply the scan spec's own ``value_splits`` in its grid helper. DEFERRED -- the fused/split crossover ------------------------------------- -The kernel docstrings state the trade in the right direction (split wins once -there is enough batch/head parallelism to fill the device, fused wins for -prefill at batch scale) but there is no measured crossover, so there is no -honest threshold to encode. The split candidates are therefore **opt-in** and -fused is the default, rather than a heuristic that would look measured and is -not. Naming ``algorithm="chunk_prep"`` / ``"chunk_scan"`` selects the split -halves. - -Also deferred: ``bind``. The launch-side pack for these kernels lives in -``builders/gfx942/kda/hostpack.py`` and is already reachable through the -manifest runner, which is a different seam; wiring it as a ``ProblemBinding`` -is a separate change, so the registry does not set ``require_binding``. +The available measurements do not establish a stable cross-architecture +threshold, so the split candidates are **opt-in** and fused is the default. +Naming ``algorithm="chunk_prep"`` / ``"chunk_scan"`` selects the split halves. + +Also deferred: ``bind``. gfx942 launches through its manifest host pack while +gfx950 launches through ``KernelLauncher`` with torch tensors. The split path +also needs two launches and a tile workspace, which is not represented by one +:class:`~rocke.dispatch.core.ProblemBinding`. """ from __future__ import annotations @@ -65,8 +59,10 @@ class KdaRequest(OperatorRequest): """Normalized chunkwise Kimi Delta Attention prefill request. ``head_v`` is the logical value width of one attention head, not the - per-workgroup partition. ``seqlen`` is the padded per-sequence length; this - family has no varlen path, so a ragged batch must be padded by the caller. + per-workgroup partition used on gfx942. ``seqlen`` is the padded + per-sequence length; this family has no varlen path, so a ragged batch must + be padded by the caller. An omitted ``chunk_size`` selects the tuned + architecture default: 16 on gfx942 and 32 on gfx950. """ batch: int @@ -75,7 +71,7 @@ class KdaRequest(OperatorRequest): arch: str head_k: int = 128 head_v: int = 128 - chunk_size: int = 16 + chunk_size: int | None = None op: str = "kda" dtype: str = "bf16" algorithm: str = "auto" @@ -86,6 +82,7 @@ class KdaRequest(OperatorRequest): def normalized(self) -> dict: d = asdict(self) d["dtype"] = self.dtype.lower() + d["chunk_size"] = self.effective_chunk_size return d def dims(self) -> dict[str, int]: @@ -95,7 +92,7 @@ def dims(self) -> dict[str, int]: "seqlen": int(self.seqlen), "head_k": int(self.head_k), "head_v": int(self.head_v), - "chunk_size": int(self.chunk_size), + "chunk_size": self.effective_chunk_size, "num_chunks": self.num_chunks, } @@ -107,22 +104,31 @@ def features(self) -> frozenset[str]: active.add("final_state") return frozenset(active) + @property + def effective_chunk_size(self) -> int: + """Requested chunk, or the architecture's tuned default.""" + if self.chunk_size is not None: + return int(self.chunk_size) + return 16 if self.arch == "gfx942" else 32 + @property def num_chunks(self) -> int: """Chunks per sequence. Zero when ``seqlen`` does not tile exactly.""" - chunk = int(self.chunk_size) + chunk = self.effective_chunk_size if chunk <= 0 or int(self.seqlen) % chunk: return 0 return int(self.seqlen) // chunk @property def v_partitions(self) -> int: - """Workgroups per logical head, from gfx942's value-channel partition.""" - return int(self.head_v) // KDA_PARTITION_HEAD_V + """Workgroups per logical head before any spec-local value splits.""" + if self.arch == "gfx942": + return int(self.head_v) // KDA_PARTITION_HEAD_V + return 1 @property def workgroups(self) -> int: - """Independent (batch, head, V partition) scan streams.""" + """Independent recurrence streams for the selected architecture.""" return int(self.batch) * int(self.num_heads) * self.v_partitions @@ -145,9 +151,11 @@ def _request_errors(req: OperatorRequest) -> list[str]: errors: list[str] = [] if req.op != "kda": errors.append(f"unsupported op {req.op!r}") - for field in ("batch", "num_heads", "seqlen", "head_k", "head_v", "chunk_size"): + for field in ("batch", "num_heads", "seqlen", "head_k", "head_v"): if int(getattr(req, field)) <= 0: errors.append(f"{field} must be positive") + if req.effective_chunk_size <= 0: + errors.append("chunk_size must be positive") try: ArchTarget.from_gfx(req.arch) except KeyError as e: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx942.py index e333a7d2e333..bb22b2fba825 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx942.py @@ -90,7 +90,7 @@ def _tile(req: KdaRequest) -> KdaTileSpec: - return KdaTileSpec(chunk=int(req.chunk_size)) + return KdaTileSpec(chunk=req.effective_chunk_size) def _fused_spec(req: OperatorRequest) -> KdaChunkFusedSpec: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx950.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx950.py new file mode 100644 index 000000000000..9efc7273e9d3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx950.py @@ -0,0 +1,331 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""gfx950 chunkwise KDA candidates (CDNA4, bf16, K-packed bf16 MFMA). + +``kernels/gfx950/kda_chunkwise.py`` has carried specs, validators, builders, +grid helpers and signatures since it landed, and nothing referenced them from +dispatch -- the family was compilable but unreachable. This module is the +registration that makes it selectable. + +Three candidates, one per emitted kernel, because each is separately +compilable with its own grid and ABI: + +* ``kda_gfx950_chunk_fused`` -- the whole prefill in one kernel. Default. +* ``kda_gfx950_chunk_prep`` -- split path, phase 1 (per-chunk tiles to HBM). +* ``kda_gfx950_chunk_scan`` -- split path, phase 2 (state scan over tiles). + +The split halves are opt-in and must be launched in that order; see +:mod:`.common` for why there is no automatic fused/split routing yet. + +The two paths do not share a tile schedule, and neither is transcribed here. +The fused kernel carries its own tuned one (a 512-thread panel mapping over +wider row pitches, which is what keeps eight waves resident while the state +lives in registers); the split scan is deliberately capped so two workgroups +fit per CU. Each candidate therefore starts from its own spec's default tile +and overrides only the requested chunk. +""" + +from __future__ import annotations + +import dataclasses +from typing import Tuple + +from kernels.gfx950.kda_chunkwise import ( + KDA_CHUNK_SIZES, + KDA_DTYPES, + KdaChunkFusedSpec, + KdaChunkPrepSpec, + KdaChunkScanSpec, + KdaTileSpec, + build_kda_chunk_fused, + build_kda_chunk_prep, + build_kda_chunk_scan, + is_valid_fused_spec, + is_valid_scan_spec, + is_valid_spec, + kda_chunk_fused_grid, + kda_chunk_fused_signature, + kda_chunk_prep_grid, + kda_chunk_prep_signature, + kda_chunk_scan_grid, + kda_chunk_scan_signature, +) +from rocke.dispatch.core import ( + CandidateRegistry, + Capability, + KernelCandidate, + OperatorRequest, + ShapeRange, +) + +from .common import ( + FAMILY, + KDA_ABI_VERSION, + KdaRequest, + _request_errors, + _selector_matches, +) + +# Declared coverage, restating the validators' *data* gates so a request can be +# filtered without constructing a spec, and taking the numbers from the kernel +# rather than transcribing them. Everything the validators compute -- the LDS +# budget, the cumsum row-group divisibility, the wave/panel cover, the scan +# partition -- stays in the residual predicate, because it is a function of the +# spec rather than a property of the request. +# +# ``num_chunks >= 1`` is how a seqlen that does not tile the chunk is rejected: +# KdaRequest.num_chunks reports 0 for a ragged length, and this family has no +# varlen path. There is no head_v relation to declare: what makes a value width +# legal is that the workgroup's waves cover it exactly at the scan atom's M +# extent, and both of those come from the spec's tile. +_SHARED_SHAPES = ( + ShapeRange("chunk_size", allowed=KDA_CHUNK_SIZES), + ShapeRange("head_k", multiple_of=8), + ShapeRange("head_v", multiple_of=8), + ShapeRange("num_chunks", min=1), + ShapeRange("batch", min=1), + ShapeRange("num_heads", min=1), +) +# Both state flags describe the *problem*, so every candidate that can take +# part in serving it declares them -- including the tile builder, which is +# state-independent and simply leaves them to the scan half. Withholding them +# there would make the split path unable to serve any problem that wants a +# final state, which is the opposite of true: the split path is how you get one +# without the fused kernel. +_STATE_FEATURES = frozenset({"initial_state", "final_state"}) + +# The kernel's own schedules, read rather than copied, so a retune of either +# path moves the dispatcher with it. +_FUSED_TILE = KdaChunkFusedSpec().tile +_SPLIT_TILE = KdaTileSpec() + + +def _fused_tile(req: KdaRequest) -> KdaTileSpec: + return dataclasses.replace(_FUSED_TILE, chunk=req.effective_chunk_size) + + +def _split_tile(req: KdaRequest) -> KdaTileSpec: + return dataclasses.replace(_SPLIT_TILE, chunk=req.effective_chunk_size) + + +def _fused_spec(req: OperatorRequest) -> KdaChunkFusedSpec: + assert isinstance(req, KdaRequest) + return KdaChunkFusedSpec( + head_k=int(req.head_k), + head_v=int(req.head_v), + dtype=req.dtype.lower(), + tile=_fused_tile(req), + has_initial_state=bool(req.has_initial_state), + store_final_state=bool(req.store_final_state), + ) + + +def _prep_spec(req: OperatorRequest) -> KdaChunkPrepSpec: + assert isinstance(req, KdaRequest) + return KdaChunkPrepSpec( + head_k=int(req.head_k), + head_v=int(req.head_v), + dtype=req.dtype.lower(), + tile=_split_tile(req), + ) + + +def _scan_spec(req: OperatorRequest) -> KdaChunkScanSpec: + assert isinstance(req, KdaRequest) + return KdaChunkScanSpec( + head_k=int(req.head_k), + head_v=int(req.head_v), + dtype=req.dtype.lower(), + tile=_split_tile(req), + has_initial_state=bool(req.has_initial_state), + store_final_state=bool(req.store_final_state), + ) + + +def _scan_validator(spec: KdaChunkScanSpec, arch: str) -> Tuple[bool, str]: + """A scan is selectable only if the tile builder that feeds it is too. + + The scan validates the staging copies and the state partition; the tile + builder additionally validates the grouped cumulative sum, whose row-group + rules the scan has no equivalent of. So a spec can pass here and leave the + pair unusable -- ``head_k=32`` at the shipped split schedule is one -- and + admitting that would move the failure to whoever tried to launch the two + halves in order. + """ + ok, why = is_valid_scan_spec(spec, arch=arch) + if not ok: + return False, why + ok, why = is_valid_spec(spec.prep, arch=arch) + if not ok: + return False, f"tile builder for this scan is unbuildable: {why}" + return True, "ok" + + +def _make_candidate( + *, + name: str, + algorithm: str, + spec_id: str, + priority: int, + capability: Capability, + spec_for, + validator, + builder, + grid_for, + signature_for, + opt_in_reason: str = "", +) -> KernelCandidate: + def support(req: OperatorRequest) -> Tuple[bool, str]: + errors = _request_errors(req) + if errors: + return False, "; ".join(errors) + assert isinstance(req, KdaRequest) + if ( + opt_in_reason + and req.algorithm.strip().lower() != algorithm + and (req.spec_id.strip().lower() != spec_id) + ): + return False, opt_in_reason + ok, why = _selector_matches(req, candidate) + if not ok: + return False, why + # Capability already cleared arch, dtype and chunk, so the spec + # construction below cannot raise. Only the residual, spec-computed + # gates run here. + return validator(spec_for(req), arch=req.arch) + + def select(req: OperatorRequest): + ok, why = candidate.admits(req) + if not ok: + raise ValueError(f"{name} does not support request: {why}") + return spec_for(req) + + candidate = KernelCandidate( + name=name, + family=FAMILY, + algorithm=algorithm, + spec_id=spec_id, + abi_version=KDA_ABI_VERSION, + priority=priority, + capability=capability, + _supports=support, + select_spec=select, + build=builder, + grid=grid_for, + block=lambda spec: (spec.tile.block_size, 1, 1), + signature=signature_for, + sweep_space=lambda req: (select(req),) if candidate.admits(req)[0] else (), + ) + return candidate + + +def _fused_grid(spec: KdaChunkFusedSpec, req: OperatorRequest): + assert isinstance(req, KdaRequest) + return kda_chunk_fused_grid(spec, req.workgroups) + + +def _scan_grid(spec: KdaChunkScanSpec, req: OperatorRequest): + assert isinstance(req, KdaRequest) + # The helper applies the spec's own ``value_splits`` v-band multiplier, so + # this stays the count of (batch, head) recurrences. + return kda_chunk_scan_grid(spec, req.workgroups) + + +def _prep_grid(spec: KdaChunkPrepSpec, req: OperatorRequest): + assert isinstance(req, KdaRequest) + # One workgroup per chunk -- the same ``BH * NC`` tile count the host + # builder allocates and packs. + return kda_chunk_prep_grid(spec, req.workgroups * req.num_chunks) + + +_SPLIT_OPT_IN = ( + "the gfx950 KDA split path is opt-in (algorithm='chunk_prep' / " + "'chunk_scan'); the fused/split crossover has been measured but not at " + "enough shapes to encode a threshold, so default routing stays on the " + "fused kernel" +) + + +def _capability() -> Capability: + return Capability( + arches=("gfx950",), + dtypes=KDA_DTYPES, + shapes=_SHARED_SHAPES, + supports_features=_STATE_FEATURES, + ) + + +def _fused_candidate() -> KernelCandidate: + """The default gfx950 prefill path. + + Default rather than opt-in because this is the family's first routing + decision, not a re-route of a benchmarked path: nothing dispatched KDA + before. It keeps the six per-chunk tiles in LDS, so the only HBM traffic is + q/k/g/beta in and o plus the final state out, and it is one launch rather + than two -- which is also why it is the half that can carry a ``bind`` + when the launch seam moves. + """ + return _make_candidate( + name="kda_gfx950_chunk_fused", + algorithm="chunk_fused", + spec_id="gfx950_chunk_fused", + priority=10, + capability=_capability(), + spec_for=_fused_spec, + validator=is_valid_fused_spec, + builder=build_kda_chunk_fused, + grid_for=_fused_grid, + signature_for=kda_chunk_fused_signature, + ) + + +def _prep_candidate() -> KernelCandidate: + """Split path phase 1: the state-independent per-chunk tile builder. + + Fully parallel over the sequence, one workgroup per chunk, which is where + the split path's parallelism advantage at small ``B*H`` comes from. + ``KdaChunkPrepSpec`` has no state fields at all, so the request's state + flags do not reach this kernel -- the scan half applies them. It still + declares the state features, for the reason given at ``_STATE_FEATURES``. + """ + return _make_candidate( + name="kda_gfx950_chunk_prep", + algorithm="chunk_prep", + spec_id="gfx950_chunk_prep", + priority=20, + capability=_capability(), + spec_for=_prep_spec, + validator=is_valid_spec, + builder=build_kda_chunk_prep, + grid_for=_prep_grid, + signature_for=kda_chunk_prep_signature, + opt_in_reason=_SPLIT_OPT_IN, + ) + + +def _scan_candidate() -> KernelCandidate: + """Split path phase 2: the state scan over materialized tiles. + + Consumes what ``kda_gfx950_chunk_prep`` wrote, so it is only meaningful + after that kernel has run on the same problem. + """ + return _make_candidate( + name="kda_gfx950_chunk_scan", + algorithm="chunk_scan", + spec_id="gfx950_chunk_scan", + priority=20, + capability=_capability(), + spec_for=_scan_spec, + validator=_scan_validator, + builder=build_kda_chunk_scan, + grid_for=_scan_grid, + signature_for=kda_chunk_scan_signature, + opt_in_reason=_SPLIT_OPT_IN, + ) + + +def register(registry: CandidateRegistry) -> None: + registry.register(_fused_candidate()) + registry.register(_prep_candidate()) + registry.register(_scan_candidate()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/kda_chunkwise.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/kda_chunkwise.py new file mode 100644 index 000000000000..03bc9f7ba729 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/kda_chunkwise.py @@ -0,0 +1,2899 @@ +"""Chunkwise gated delta-rule linear attention (KDA) prefill kernels for gfx950. + +KDA is a gated delta-rule linear attention: a per-channel decay gate plus a +delta-rule state write, evaluated chunkwise so the token-serial recurrence +collapses into a handful of dense matmuls per chunk. + +Contract +-------- +Per head, tokens are grouped into chunks of ``C`` rows. Writing the per-channel +cumulative log decay within a chunk as ``Gamma_i = exp(sum_{j<=i} g_j)`` and the +whole-chunk decay as ``gamma_C = Gamma_{C-1}``, the chunk body factorizes into +six state-independent tiles + +.. code-block:: text + + A = (I + StrictTril(Diag(beta) Akk))^-1 Diag(beta) C x C + Akk_ij = k_i . (k_j * Gamma_i / Gamma_j) + GK = K * Gamma C x DK + GQ = Q * Gamma * scale C x DK + Aqk = Tril(GQ (K / Gamma)^T) (i >= j) C x C + Kt = (K * gamma_C / Gamma)^T DK x C + dec = gamma_C DK + +and a serial walk over chunks carrying the state ``S`` (DK x DV) in fp32 + +.. code-block:: text + + Vt = A (V - GK S) + O = GQ S + Aqk Vt + S = Diag(dec) S + Kt^T Vt + +Only ``S`` is serial, so the tile construction above is one workgroup per chunk +and fully parallel over the sequence. + +Numerics +-------- +``Akk`` and ``Aqk`` both need the ratio ``Gamma_i / Gamma_j``, which spans the +whole chunk's decay range and overflows fp32 if formed directly: at the +reference ``gate_lower_bound = -5`` a 32-token chunk accumulates up to 160 nats. +Both are +therefore built factored against the chunk's midpoint row ``CREF = C // 2``, + +.. code-block:: text + + Akk = (K * e^(Gc - Gref)) (K * e^(Gref - Gc))^T + +so each factor's exponent is bounded by half the chunk range, and the product +reconstructs the ratio exactly. Every exponential is additionally clamped to the +fp32 exp2 range. The cumulative sum is kept scaled by ``log2(e)`` throughout so +the hardware ``v_exp_f32`` (base 2) is used directly with no extra multiply. + +Layout +------ +Inputs arrive already packed by chunk: ``tile = bh * NC + n`` indexes a chunk of +one (batch, head). q/k are bf16, the gate is fp32 per channel (not pre-summed), +beta is fp32 per token, and q/k are expected pre-normalized -- the L2 norm, gate +activation and beta sigmoid belong to the host-side pack pass, not here. + +Synchronization +--------------- +Every phase boundary here is an LDS-visibility rendezvous, so all of them use +``sync_lds_only`` (``lgkmcnt(0)`` + ``s_barrier``) rather than the full +``sync``. There is no direct-to-LDS path in this kernel: a global load always +lands in a register and reaches LDS through a ``ds_write``, so waiting on +``lgkmcnt`` transitively waits on the load that feeds it. The global *stores* +are kernel outputs, consumed across a dispatch boundary rather than across one +of these barriers. + +This is the minimum legal wait, not a faster one. Dropping ``vmcnt(0)`` off all +14 barriers is exactly neutral end to end, and ATT says why: the stall is +conserved, moving from the vmcnt bucket (22.3% -> 18.7% of total) into the lgkm +bucket (23.7% -> 27.5%) with the barrier bucket and the total both unchanged. +The drain was never redundant work -- it was the same load dependency, charged +to a different counter. Cheapening these waits is therefore not a lever on this +kernel; the barrier waits on wave arrival, and wave arrival waits on the loads +either way. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Tuple + +from rocke.core.ir import BF16, F32, I8, I32, IRBuilder, KernelDef, PtrType +from rocke.helpers.activations import _sigmoid_via_exp2 +from rocke.helpers.atoms import MfmaAtom +from rocke.helpers.spec import SignatureBuilder, kernel_name_join + +_RAW_VALUE_SPLITS = (1, 2, 4, 8) + +# Cumulative log decay is kept in the log2 domain so exp2 is a bare v_exp_f32. +LOG2E = 1.4426950408889634 +# v_exp_f32 saturates past this; the clamp keeps a saturated gate finite instead +# of turning a whole chunk into NaN. +EXP2_CLAMP = 126.0 + +_DTYPE_IR = {"bf16": BF16} +# Declared coverage, exported so a dispatch candidate can state what it serves +# by importing these rather than transcribing them. The validators below read +# the same names, so the two cannot drift apart. +KDA_DTYPES: Tuple[str, ...] = tuple(sorted(_DTYPE_IR)) +# The only chunk lengths the emitted tile-atom schedules cover. +KDA_CHUNK_SIZES: Tuple[int, ...] = (16, 32) +# gfx950 LDS budget per workgroup. +LDS_LIMIT = 160 * 1024 +# The K-packed bf16 atoms, by M/N extent. Both are gfx950 shapes; the K extent +# is whatever that shape packs, since the scan's contractions are all K-stepped. +_SCAN_ATOMS = { + 32: MfmaAtom.bf16_32x32x16, + 16: MfmaAtom.bf16_16x16x32, +} + + +class _LdsTile: + """A 2-D (or 1-D) view onto a 1-D ``i8`` LDS pool. + + The fused kernel's tile phase and scan phase are mutually exclusive in + time, but they live in one ``scf.for``, so the LDS packer's loop liveness + will not overlay them. A single pool plus byte-offset views lets the + scan-only ``S^T`` mirror and ``V~`` reuse tile-only storage explicitly. + """ + + __slots__ = ("pool", "off", "elem_bytes", "cols", "rank") + + def __init__(self, pool, off: int, elem, shape: Tuple[int, ...]): + self.pool = pool + self.off = off + self.elem_bytes = 4 if elem.name == "f32" else 2 + self.rank = len(shape) + self.cols = int(shape[-1]) if shape else 1 + + def index(self, b: IRBuilder, *idx): + if self.rank == 1: + (col,) = idx + return b.add( + b.const_i32(self.off), b.mul(col, b.const_i32(self.elem_bytes)) + ) + row, col = idx + return b.add( + b.const_i32(self.off), + b.add( + b.mul(row, b.const_i32(self.cols * self.elem_bytes)), + b.mul(col, b.const_i32(self.elem_bytes)), + ), + ) + + +def _ld(b: IRBuilder, smem, *idx, dtype, n: int): + """Vector LDS load that accepts either a typed alloc or an ``_LdsTile``.""" + if isinstance(smem, _LdsTile): + return b.smem_load_vN(smem.pool, smem.index(b, *idx), dtype=dtype, n=n) + return b.smem_load_vN(smem, *idx, dtype=dtype, n=n) + + +def _st(b: IRBuilder, smem, *idx, value, n: int) -> None: + """Vector LDS store that accepts either a typed alloc or an ``_LdsTile``.""" + if isinstance(smem, _LdsTile): + b.smem_store_vN(smem.pool, [smem.index(b, *idx)], value, n) + else: + b.smem_store_vN(smem, list(idx), value, n) + + +def _fused_lds_layout(head_k: int, head_v: int, tile: "KdaTileSpec") -> dict: + """Byte offsets of every fused LDS tile inside one pool. + + Three groups have different lifetimes: + + * V/Y, Kt, and GK/GQ/A/Aqk/dec survive from the tile phase into the scan. + Kt must be dedicated even though the scan consumes it later: helper waves + produce it during the solve while ``m/a/qk`` are still live. + * g/k/q/m/a/qk/beta are tile-only once the post-solve copies finish. + * the ``S^T`` mirror and ``V~`` are scan-only. Together they fit inside the + tile-only span and reuse it exactly after the tile/scan rendezvous. + + Default-shape accounting:: + + V/Y + Kt + GK/GQ/A/Aqk/dec 41,984 B + tile-only region 46,720 B + S^T mirror + V~ (overlaid) 45,056 B + pool 88,704 B + + This is down from 133,760 B. It is intentionally not the tempting 78,464-B + layout: that version overlapped both V and Kt before their last use and + produced NaNs. + + The barrier-free next-input prefetch is also incompatible with this layout: + it writes g/k/q/beta during the scan, exactly while the overlaid scan extras + occupy those addresses. ``is_valid_fused_spec`` rejects that combination. + """ + C, DK, EV = tile.chunk, head_k, head_v + pdk, pc, pcb = DK + tile.pad_dk, C + tile.pad_c, C + tile.pad_cb + off = 0 + + def take(n: int) -> int: + nonlocal off + o = off + off += n + return o + + # Live from tile production into the scan. + lay = { + "y": take(2 * C * pdk), + "kt": take(2 * DK * pcb), + "x": take(2 * C * pdk), + "xq": take(2 * C * pdk), + "tb": take(2 * C * pcb), + "zs": take(2 * C * pcb), + "gl": take(4 * DK), + } + + # Tile-only, one contiguous reuse span. + tile_only = off + lay.update( + { + "g": take(4 * C * DK), + "k": take(2 * C * DK), + "q": take(2 * C * DK), + "m": take(4 * C * pc), + "a": take(4 * C * pc), + "qk": take(4 * C * pc), + "beta": take(4 * C), + } + ) + pool_end = off + + # Scan-only views reuse the tile-only span. + stb_n, vn_n = 2 * EV * pdk, 2 * EV * pcb + lay["stb"], lay["vn"] = tile_only, tile_only + stb_n + if lay["vn"] + vn_n > pool_end: + raise ValueError( + f"scan-only tiles ({stb_n + vn_n} B) exceed reusable tile-only " + f"region ({pool_end - tile_only} B)" + ) + off = pool_end + lay["pool"] = off + return lay + + +@dataclass(frozen=True) +class KdaTileSpec: + """Tiling and LDS-layout knobs for the chunkwise KDA kernels. + + ``chunk`` is the algorithmic chunk length: it sets the ``C x C`` triangular + solve cost (O(C^2) serial) against the number of chunks (O(T/C) serial in + the scan), so it is the primary throughput knob rather than a free choice. + + The pads are bank-conflict padding, in elements, on the LDS row pitch. The + MFMA operand reads are ``ds_read_b128`` at a row stride of ``DK``; an + unpadded 128-element bf16 row puts every lane's read in the same 16 banks, + so the pitch is padded to spread them. + """ + + chunk: int = 32 + block_size: int = 256 + # LDS row padding, in elements, for the (C x DK) MFMA operand staging tiles. + pad_dk: int = 8 + # LDS row padding, in elements, for the scalar-accessed fp32 (C x C) tiles. + pad_c: int = 4 + # LDS row padding, in elements, for the bf16 (C x C) tiles that feed the + # rank-update MFMA. These are read with ds_read_b128, so the padded row + # pitch must stay a multiple of 8 elements (16 B) or odd rows land on an + # 8-byte boundary and silently break the read's alignment contract. + pad_cb: int = 8 + # M/N extent of the tile phase's C x C products. 32 keeps the original + # one-atom tile; 16 maps the eight waves onto the four 16x16 inner panels + # (two waves per panel) while preserving the algorithmic C=32 chunk. + tile_atom_m: int = 32 + # solve_block: row-block size of the triangular solve. The solve's arithmetic + # splits into per-block substitution (serial, scalar VALU) and the rank + # update against already-solved blocks (a matmul, so MFMA). Only the + # substitution part is irreducibly scalar, and it shrinks as the square of + # the block size, so smaller blocks move more of the O(C^3) work onto the + # MFMA pipe -- at the cost of one more block step. ``solve_block == chunk`` + # is the degenerate single-block case: one unblocked scalar substitution and + # no MFMA. Must be a multiple of 8 (the accumulator holds a contiguous run + # of 8 output rows per group of 4 slots, which is what lets a block step + # write back only its own rows) and must divide ``chunk``. + solve_block: int = 8 + # M/N extent of the atom the *state scan* uses, which need not be the tile + # phase's. The C x C tile products want an atom as wide as the chunk, but the + # scan's partitioning rule is independent: one wave owns one ``atom.m``-row + # band of S^T, so a narrower atom subtiles the state into shorter bands. + # Twice as many waves then cover the v extent and each lane carries half as + # many state accumulators -- which is the only way to fit the fused kernel's + # loop-carried state under the 256-VGPR ceiling that two waves per SIMD + # needs. 0 = reuse the chunk-wide atom. + scan_atom_m: int = 0 + waves_per_eu: int = 0 # 0 = leave the occupancy hint off + + @property + def wave_size(self) -> int: + return 64 + + @property + def num_waves(self) -> int: + return self.block_size // self.wave_size + + def name_parts(self) -> Tuple[str, ...]: + parts = (f"c{self.chunk}", f"b{self.block_size}", f"sb{self.solve_block}") + if (self.pad_dk, self.pad_c) != (8, 4): + parts += (f"p{self.pad_dk}x{self.pad_c}",) + if self.pad_cb != 8: + parts += (f"pcb{self.pad_cb}",) + if self.scan_atom_m: + parts += (f"sa{self.scan_atom_m}",) + if self.tile_atom_m != 32: + parts += (f"ta{self.tile_atom_m}",) + if self.waves_per_eu: + parts += (f"wpe{self.waves_per_eu}",) + return parts + + +@dataclass(frozen=True) +class KdaChunkPrepSpec: + """Compile-time spec for the state-independent per-chunk tile builder. + + One workgroup per chunk, grid ``(BH * NC, 1, 1)``. Every shape field is + baked into the kernel as a constant, so the ABI carries pointers and the + softmax scale only. + """ + + head_k: int = 128 + head_v: int = 128 + dtype: str = "bf16" + tile: KdaTileSpec = KdaTileSpec() + name: str = "rocke_kda_chunk_prep" + # Default-off raw token-major path: inputs are [B,T,H,D] / [B,T,H] and the + # producer may fuse the Aiter-equivalent q/k L2 norm, gate, and beta sigmoid. + raw_inputs: bool = False + fuse_qk_l2norm: bool = False + fuse_gate: bool = False + fuse_beta_sigmoid: bool = False + has_dt_bias: bool = False + lower_bound: float = -5.0 + + @property + def atom(self) -> MfmaAtom: + """The bf16 atom used by the tile phase's C x C products.""" + make = _SCAN_ATOMS.get(self.tile.tile_atom_m) + if make is None: + raise ValueError( + f"no bf16 tile atom with M/N extent {self.tile.tile_atom_m}; " + f"have {sorted(_SCAN_ATOMS)}" + ) + return make() + + @property + def k_steps(self) -> int: + return self.head_k // self.atom.k + + def lds_bytes(self) -> int: + c, t = self.tile.chunk, self.tile + pdk, pc = self.head_k + t.pad_dk, c + t.pad_c + return ( + 4 * c * self.head_k # g_cum fp32 (C x DK) + + 2 * c * self.head_k # k_s bf16 + + 2 * c * self.head_k # q_s bf16 + + 2 * c * pdk # x_s bf16 Akk A operand + + 2 * c * pdk # y_s bf16 shared B operand + + 2 * c * pdk # xq_s bf16 Aqk A operand + + 4 * c * pc # m_mat fp32 T', in-block substitution + + 4 * c * pc # a_mat fp32 solve result -> A + + 4 * c * pc # aqk_mat fp32 Aqk + + 2 * c * (c + t.pad_cb) # tb_s bf16 rank-update A operand + + 2 * c * (c + t.pad_cb) # zs_s bf16 solved-so-far, transposed + + 4 * c # beta_s fp32 + + 4 * self.head_k # gl_s fp32 whole-chunk log decay + ) + + def kernel_name(self) -> str: + parts = [ + self.name, + f"dk{self.head_k}", + f"dv{self.head_v}", + self.dtype, + *self.tile.name_parts(), + ] + if self.raw_inputs: + parts.append("raw") + if self.fuse_qk_l2norm: + parts.append("nl2") + if self.fuse_gate: + parts.append("fg") + if self.fuse_beta_sigmoid: + parts.append("fb") + if self.has_dt_bias: + parts.append("db") + if self.lower_bound != -5.0: + parts.append(f"lb{self.lower_bound:g}") + return kernel_name_join(*parts) + + +def _scan_atom(tile: KdaTileSpec, default: MfmaAtom) -> MfmaAtom: + """The atom the state scan runs on, per ``KdaTileSpec.scan_atom_m``.""" + if not tile.scan_atom_m: + return default + make = _SCAN_ATOMS.get(tile.scan_atom_m) + if make is None: + raise ValueError( + f"no bf16 atom with M/N extent {tile.scan_atom_m}; " + f"have {sorted(_SCAN_ATOMS)}" + ) + return make() + + +def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, str]: + """Return ``(ok, reason)`` for a prep spec on ``arch``.""" + if arch != "gfx950": + return False, f"kda_chunk_prep is gfx950-only (got {arch})" + if spec.dtype not in _DTYPE_IR: + return False, f"unsupported dtype {spec.dtype!r} (bf16 only)" + + if spec.has_dt_bias and not spec.fuse_gate: + return False, "has_dt_bias requires fuse_gate" + for flag, name in ( + (spec.fuse_qk_l2norm, "fuse_qk_l2norm"), + (spec.fuse_gate, "fuse_gate"), + (spec.fuse_beta_sigmoid, "fuse_beta_sigmoid"), + ): + if flag and not spec.raw_inputs: + return False, f"{name} requires raw_inputs=True" + if spec.raw_inputs: + if not (spec.fuse_qk_l2norm or spec.fuse_gate or spec.fuse_beta_sigmoid): + return False, ( + "raw_inputs requires at least one fused preprocessing flag " + "(fuse_qk_l2norm, fuse_gate, fuse_beta_sigmoid)" + ) + + t = spec.tile + if t.tile_atom_m not in _SCAN_ATOMS: + return False, ( + f"unsupported tile_atom_m {t.tile_atom_m}; " f"have {sorted(_SCAN_ATOMS)}" + ) + atom = spec.atom + if t.chunk not in KDA_CHUNK_SIZES: + return False, ( + f"chunk must be one of {list(KDA_CHUNK_SIZES)} for the emitted " + "tile-atom schedules" + ) + if t.chunk % atom.m or t.chunk % atom.n: + return False, (f"tile atom ({atom.m}x{atom.n}) must divide chunk ({t.chunk})") + panels = (t.chunk // atom.m) * (t.chunk // atom.n) + if t.num_waves % panels: + return False, ( + f"waves ({t.num_waves}) must be a multiple of C x C panels " + f"({panels}) so every panel has equal wave coverage" + ) + if spec.head_k % atom.k: + return False, ( + f"head_k ({spec.head_k}) must be a multiple of the MFMA K step " + f"({atom.k})" + ) + solve_atom = MfmaAtom.bf16_16x16x32() if t.chunk == 16 else MfmaAtom.bf16_32x32x16() + if t.chunk + t.pad_cb < solve_atom.k: + return False, ( + f"chunk + pad_cb ({t.chunk + t.pad_cb}) must cover solve MFMA K " + f"({solve_atom.k}) so the C=16 tail can be zero padded" + ) + # The cumulative sum gives one thread a whole (row-group, channel) column, + # so the block has to be a whole number of copies of the channel extent, and + # the chunk has to split evenly into that many groups. + if t.block_size % spec.head_k or t.block_size < 2 * spec.head_k: + return False, ( + f"block_size ({t.block_size}) must be a multiple of head_k " + f"({spec.head_k}), at least 2x, for the grouped cumulative sum" + ) + n_groups = t.block_size // spec.head_k + if t.chunk % n_groups: + return False, ( + f"chunk ({t.chunk}) must split evenly into the {n_groups} cumsum " + f"row groups implied by block_size/head_k" + ) + # The fold deals rows out round-robin across a channel's threads and folds a + # compile-time set of group totals into each, which needs every step's rows + # to land inside one group. + if (t.chunk // n_groups) % n_groups: + return False, ( + f"cumsum group height ({t.chunk // n_groups}) must be a multiple of " + f"the group count ({n_groups}) so the fold's offsets stay uniform" + ) + if t.solve_block % 8 or t.chunk % t.solve_block: + return False, ( + f"solve_block ({t.solve_block}) must be a multiple of 8 and divide " + f"chunk ({t.chunk})" + ) + if t.block_size % t.wave_size: + return False, f"block_size ({t.block_size}) must be a wave multiple" + + # Every global access is a 128-bit transaction; vector starts must stay + # inside rows. C=16 leaves half the 512-thread block inactive in the bf16 + # staging sweep, which the emitter predicates explicitly. + if spec.head_k % 8 or t.chunk % 8: + return False, "head_k and chunk must be multiples of 8 for 128-bit access" + if t.pad_dk % 8: + return False, ( + f"pad_dk ({t.pad_dk}) must be a multiple of 8 so the padded row " + "pitch keeps ds_read_b128 alignment" + ) + if t.pad_c % 4: + return False, f"pad_c ({t.pad_c}) must be a multiple of 4" + if (t.chunk + t.pad_cb) % 8: + return False, ( + f"chunk + pad_cb ({t.chunk + t.pad_cb}) must be a multiple of 8 so " + "the bf16 C x C row pitch keeps ds_read_b128 alignment" + ) + + lds = spec.lds_bytes() + if lds > LDS_LIMIT: + return False, f"LDS request {lds} B exceeds the {LDS_LIMIT} B budget" + return True, "ok" + + +def _reduce16_fadd(b: IRBuilder, v): + """Sum ``v`` across the 16-lane row subgroup (XOR butterfly).""" + for mask in (1, 2, 4, 8): + v = b.fadd(v, b.warp_shuffle_xor(v, mask)) + return v + + +def _l2norm_scale8(b: IRBuilder, vals8): + """One in-register L2 norm over eight bf16-as-f32 values in a 16-lane group.""" + ss = b.const_f32(0.0) + for x in vals8: + ss = b.fadd(ss, b.fmul(x, x)) + ss = _reduce16_fadd(b, ss) + return b.rsqrt(b.fadd(ss, b.const_f32(1e-6))) + + +def _apply_gate_fused( + b: IRBuilder, + raw_g, + a_log, + dt_bias, + head, + col, + lower_bound, + *, + has_dt_bias: bool, + dk: int, +): + """``lower_bound * sigmoid(exp(A_log[h]) * (g + dt_bias[h,d]))``.""" + g = b.cast_to_f32(raw_g) + if has_dt_bias: + g = b.fadd( + g, + b.global_load_f32(dt_bias, b.add(b.mul(head, b.const_i32(dk)), col)), + ) + a = b.global_load_f32(a_log, head) + x = b.fmul(b.exp2(b.fmul(b.const_f32(LOG2E), a)), g) + return b.fmul(b.const_f32(lower_bound), _sigmoid_via_exp2(b, x)) + + +class _RawTokenAddr: + """Token-major [B,T,H,D] / [B,T,H] addressing for one chunk tile.""" + + def __init__( + self, b: IRBuilder, *, heads, tseq, nc, chunk, dk, a_log=None, dt_bias=None + ): + self.b = b + self.H = heads + self.T = tseq + self.nc = nc + self.C = chunk + self.DK = dk + self.a_log = a_log + self.dt_bias = dt_bias + self.stride_token_qk = b.mul(heads, b.const_i32(dk)) + self.stride_batch_qk = b.mul(tseq, self.stride_token_qk) + self.stride_token_beta = heads + self.stride_batch_beta = b.mul(tseq, heads) + + def _parts(self, tile, row): + b = self.b + bh = b.div(tile, self.nc) + chunk_n = b.mod(tile, self.nc) + batch = b.div(bh, self.H) + head = b.mod(bh, self.H) + token = b.add(b.mul(chunk_n, b.const_i32(self.C)), row) + return batch, head, token + + def qk_off(self, tile, row, col): + b = self.b + batch, head, token = self._parts(tile, row) + return b.add( + b.add( + b.mul(batch, self.stride_batch_qk), + b.mul(token, self.stride_token_qk), + ), + b.add(b.mul(head, b.const_i32(self.DK)), col), + ) + + def beta_off(self, tile, row): + b = self.b + batch, head, token = self._parts(tile, row) + return b.add( + b.add( + b.mul(batch, self.stride_batch_beta), + b.mul(token, self.stride_token_beta), + ), + head, + ) + + def v_off(self, tile, chunk_row, ev, ev_dim): + b = self.b + batch, head, token = self._parts(tile, chunk_row) + stride_token = b.mul(self.H, b.const_i32(ev_dim)) + stride_batch = b.mul(self.T, stride_token) + return b.add( + b.add( + b.mul(batch, stride_batch), + b.mul(token, stride_token), + ), + b.add(b.mul(head, b.const_i32(ev_dim)), ev), + ) + + +class _ChunkCtx: + """Everything the per-chunk tile emitter needs that does not depend on + *which* chunk is being built: the LDS tiles, the thread/lane decomposition, + and the shared access helpers. + + Split out from the emitter so a kernel that walks many chunks in one + workgroup pays for this setup once, outside its loop, while the per-chunk + work stays a single reusable emission. + """ + + def __init__( + self, b: IRBuilder, spec: KdaChunkPrepSpec, inputs, *, overlay: bool = False + ): + t = spec.tile + self.b = b + self.spec = spec + self.q_ptr, self.k_ptr, self.g_ptr, self.beta_ptr, self.scale = inputs[:5] + self.raw = inputs[5] if len(inputs) > 5 else None + self.C = C = t.chunk + self.DK = DK = spec.head_k + self.BLOCK = t.block_size + self.PDK = DK + t.pad_dk + self.PC = C + t.pad_c + self.PCB = C + t.pad_cb + self.CREF = C // 2 + self.HALF = C // 2 + # Row groups for the cumulative sum: one thread owns a whole + # (row-group, channel) column, so the number of groups is however many + # copies of the channel extent the block holds. + self.NG = self.BLOCK // DK + self.GROUP = C // self.NG + self.ELEM = _DTYPE_IR[spec.dtype] + self.atom = spec.atom + self.N_CD = C * DK + ELEM = self.ELEM + + # ---- LDS ---- + if overlay: + self._init_overlay_lds() + else: + self.g_lds = b.smem_alloc(F32, [C, DK], "g_cum") + self.k_lds = b.smem_alloc(ELEM, [C, DK], "k_s") + self.q_lds = b.smem_alloc(ELEM, [C, DK], "q_s") + self.x_lds = b.smem_alloc(ELEM, [C, self.PDK], "x_s") + self.y_lds = b.smem_alloc(ELEM, [C, self.PDK], "y_s") + self.xq_lds = b.smem_alloc(ELEM, [C, self.PDK], "xq_s") + self.m_lds = b.smem_alloc(F32, [C, self.PC], "m_mat") + self.a_lds = b.smem_alloc(F32, [C, self.PC], "a_mat") + self.qk_lds = b.smem_alloc(F32, [C, self.PC], "aqk_mat") + self.tb_lds = b.smem_alloc(ELEM, [C, self.PCB], "tb_s") + self.zs_lds = b.smem_alloc(ELEM, [C, self.PCB], "zs_s") + self.beta_lds = b.smem_alloc(F32, [C], "beta_s") + self.gl_lds = b.smem_alloc(F32, [DK], "gl_s") + self.kt_lds = self.stb_lds = self.vn_lds = None + + self.tid = b.thread_id_x() + self.lane = lane = b.mod(self.tid, b.const_i32(64)) + self.lane_m = b.mod(lane, b.const_i32(32)) + self.lane_h = lane_h = b.div(lane, b.const_i32(32)) + self.frag_k_off = b.mul(lane_h, b.const_i32(8)) + + self.c_clamp = b.const_f32(-EXP2_CLAMP) + self.c_log2e = b.const_f32(LOG2E) + self.c_clamp_hi = b.const_f32(EXP2_CLAMP) + + def _init_overlay_lds(self) -> None: + """One i8 pool; tile-only buffers and scan extras share the dead region.""" + b, spec, t = self.b, self.spec, self.spec.tile + C, DK, EV = self.C, self.DK, spec.head_v + ELEM = self.ELEM + lay = _fused_lds_layout(DK, EV, t) + pool = b.smem_alloc(I8, [lay["pool"]], "lds_pool") + + def view(key, elem, shape): + return _LdsTile(pool, lay[key], elem, shape) + + self.g_lds = view("g", F32, (C, DK)) + self.k_lds = view("k", ELEM, (C, DK)) + self.q_lds = view("q", ELEM, (C, DK)) + self.y_lds = view("y", ELEM, (C, self.PDK)) + self.m_lds = view("m", F32, (C, self.PC)) + self.a_lds = view("a", F32, (C, self.PC)) + self.qk_lds = view("qk", F32, (C, self.PC)) + self.beta_lds = view("beta", F32, (C,)) + self.x_lds = view("x", ELEM, (C, self.PDK)) + self.xq_lds = view("xq", ELEM, (C, self.PDK)) + self.tb_lds = view("tb", ELEM, (C, self.PCB)) + self.zs_lds = view("zs", ELEM, (C, self.PCB)) + self.gl_lds = view("gl", F32, (DK,)) + self.stb_lds = view("stb", ELEM, (EV, self.PDK)) + self.kt_lds = view("kt", ELEM, (DK, self.PCB)) + self.vn_lds = view("vn", ELEM, (EV, self.PCB)) + + def ex2(self, x): + """exp2 with the argument clamped into the fp32 exponent range. + + A saturated gate can drive the factored exponents past the fp32 range; + clamping keeps the tile finite instead of letting one channel turn the + whole chunk into NaN. + """ + b = self.b + return b.exp2(b.fmin(b.fmax(x, self.c_clamp), self.c_clamp_hi)) + + def lds_get(self, smem, idx, dtype=F32): + """One scalar LDS read.""" + return self.b.vec_extract(_ld(self.b, smem, *idx, dtype=dtype, n=1), 0) + + def lds_get8_f32(self, smem, row, col): + """Eight consecutive fp32 from LDS as two ``ds_read_b128``.""" + b = self.b + out = [] + for h in range(2): + v = _ld(b, smem, row, b.add(col, b.const_i32(4 * h)), dtype=F32, n=4) + out += [b.vec_extract(v, j) for j in range(4)] + return out + + def lds_get8_elem(self, smem, row, col): + """Eight consecutive bf16 from LDS as one ``ds_read_b128``, as f32.""" + b = self.b + v = _ld(b, smem, row, col, dtype=self.ELEM, n=8) + return [b.cast_to_f32(b.vec_extract(v, j)) for j in range(8)] + + def lds_put8(self, smem, row, col, vals_f32): + """Eight f32 truncated to bf16 and written as one ``ds_write_b128``.""" + b = self.b + _st( + b, + smem, + row, + col, + value=b.vec_pack( + [b.cast_f32_to(v, self.ELEM) for v in vals_f32], self.ELEM + ), + n=8, + ) + + +class _ChunkOffsets: + """The one chunk's base offsets into each flat per-chunk array.""" + + def __init__(self, ctx: _ChunkCtx, tile): + b, C, DK = ctx.b, ctx.C, ctx.DK + self.tile = tile + self.cd = b.mul(tile, b.const_i32(ctx.N_CD)) + self.cc = b.mul(tile, b.const_i32(C * C)) + self.c = b.mul(tile, b.const_i32(C)) + self.dk = b.mul(tile, b.const_i32(DK)) + + +class _GlobalTileSink: + """Tile destination: HBM, for a separate scan kernel to read back. + + The hooks fire at the point each tile's values become available, so the + store keeps the 128-bit shape the producing pass already had its values in. + """ + + # GK/GQ go straight out of the fused elementwise pass; nothing is competing + # for the destination, so there is no reason to defer them. + deferred_gk_gq = False + + def __init__(self, a_ptr, gk_ptr, gq_ptr, aqk_ptr, kt_ptr, dec_ptr): + self.a_ptr = a_ptr + self.gk_ptr = gk_ptr + self.gq_ptr = gq_ptr + self.aqk_ptr = aqk_ptr + self.kt_ptr = kt_ptr + self.dec_ptr = dec_ptr + + def dec(self, ctx, ch, col4, dec4): + b = ctx.b + b.global_store_vN(self.dec_ptr, b.add(ch.dk, col4), b.vec_pack(dec4, F32), 4) + + def gk_gq(self, ctx, ch, row, col, off, gk8, gq8): + b, ELEM = ctx.b, ctx.ELEM + gidx = b.add(ch.cd, off) + b.global_store_vN( + self.gk_ptr, + gidx, + b.vec_pack([b.cast_f32_to(v, ELEM) for v in gk8], ELEM), + 8, + ) + b.global_store_vN( + self.gq_ptr, + gidx, + b.vec_pack([b.cast_f32_to(v, ELEM) for v in gq8], ELEM), + 8, + ) + + def _cxc(self, ctx, ch, out_ptr, off, row, col, vals): + b, ELEM = ctx.b, ctx.ELEM + b.global_store_vN( + out_ptr, + b.add(ch.cc, off), + b.vec_pack([b.cast_f32_to(v, ELEM) for v in vals], ELEM), + 8, + ) + + def a(self, ctx, ch, off, row, col, vals): + self._cxc(ctx, ch, self.a_ptr, off, row, col, vals) + + def aqk(self, ctx, ch, off, row, col, vals): + self._cxc(ctx, ch, self.aqk_ptr, off, row, col, vals) + + def kt(self, ctx, ch, off, dch, r8, vals): + b = ctx.b + b.global_store_vN(self.kt_ptr, b.add(ch.cd, off), b.vec_pack(vals, ctx.ELEM), 8) + + +def _emit_gk_gq_pass(ctx: _ChunkCtx, ch: "_ChunkOffsets", sink) -> None: + """GK = K * Gamma and GQ = Q * Gamma * scale, as a standalone sweep. + + Same access shape as the main elementwise pass -- one thread owns eight + consecutive channels of a row, so every LDS transaction is 128-bit -- but + emitted late, for a sink whose GK/GQ destination is only free once the + chunk's MFMA operands have been consumed. + """ + b = ctx.b + tid, scale = ctx.tid, ctx.scale + DK, BLOCK, N_CD = ctx.DK, ctx.BLOCK, ctx.N_CD + ew_col = b.mod(b.mul(tid, b.const_i32(8)), b.const_i32(DK)) + ew_row = b.div(b.mul(tid, b.const_i32(8)), b.const_i32(DK)) + ew_rstep = (BLOCK * 8) // DK + for p in range((N_CD + 8 * BLOCK - 1) // (8 * BLOCK)): + vidx = b.add(tid, b.const_i32(p * BLOCK)) + guard = ( + b.scf_if(b.cmp_gt(b.const_i32(N_CD // 8), vidx)) + if N_CD % (8 * BLOCK) + else nullcontext() + ) + with guard: + row = b.add(ew_row, b.const_i32(p * ew_rstep)) + off = b.mul(vidx, b.const_i32(8)) + g8 = ctx.lds_get8_f32(ctx.g_lds, row, ew_col) + k8 = ctx.lds_get8_elem(ctx.k_lds, row, ew_col) + q8 = ctx.lds_get8_elem(ctx.q_lds, row, ew_col) + gk8, gq8 = [], [] + for j in range(8): + eg = ctx.ex2(g8[j]) + gk8.append(b.fmul(k8[j], eg)) + gq8.append(b.fmul(b.fmul(q8[j], scale), eg)) + sink.gk_gq(ctx, ch, row, ew_col, off, gk8, gq8) + + +def _emit_kt_slot(ctx: _ChunkCtx, ch, sink, vidx) -> None: + """One thread's eight-row Kt column at flat slot ``vidx``.""" + b, C, ELEM = ctx.b, ctx.C, ctx.ELEM + off = b.mul(vidx, b.const_i32(8)) + dch = b.div(off, b.const_i32(C)) + r8 = b.mod(off, b.const_i32(C)) + glc = ctx.lds_get(ctx.gl_lds, [dch]) + vals = [] + for j in range(8): + rj = b.add(r8, b.const_i32(j)) + kv = b.cast_to_f32(ctx.lds_get(ctx.k_lds, [rj, dch], dtype=ELEM)) + gc = ctx.lds_get(ctx.g_lds, [rj, dch]) + vals.append(b.cast_f32_to(b.fmul(kv, ctx.ex2(b.fsub(glc, gc))), ELEM)) + sink.kt(ctx, ch, off, dch, r8, vals) + + +def _emit_kt_all_waves(ctx: _ChunkCtx, ch, sink) -> None: + """Kt with the original uniform ``tid + i*BLOCK`` mapping, every wave.""" + b, tid, BLOCK, N_CD = ctx.b, ctx.tid, ctx.BLOCK, ctx.N_CD + for i in range(N_CD // (8 * BLOCK)): + _emit_kt_slot(ctx, ch, sink, b.add(tid, b.const_i32(i * BLOCK))) + + +def _emit_stage_issue(ctx: _ChunkCtx, ch): + """Issue the HBM loads for one chunk's g / k / q / beta. + + Returns each loaded value next to the staging slot it belongs in, so a + caller can place the ``ds_write``s further along and let whatever runs in + between cover the load latency. Slots are ``tid + i*BLOCK``, which is exact + for every legal block size. + + ``beta`` is loaded on every thread against a clamped index rather than under + the ``tid < C`` predicate its store needs: the value has to outlive the issue + point, and a predicated load would leave it trapped inside the ``scf.if``. + The redundant lanes re-read 128 B that is already in cache. + """ + b = ctx.b + C, DK, BLOCK, N_CD = ctx.C, ctx.DK, ctx.BLOCK, ctx.N_CD + ELEM = ctx.ELEM + tid = ctx.tid + tile_cd, tile_c = ch.cd, ch.c + tile = ch.tile + staged = [] + raw = ctx.raw + + for i in range((N_CD + 4 * BLOCK - 1) // (4 * BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + valid = None + safe = vidx + if N_CD % (4 * BLOCK): + valid = b.cmp_gt(b.const_i32(N_CD // 4), vidx) + safe = b.select(valid, vidx, b.const_i32(N_CD // 4 - 1)) + off = b.mul(safe, b.const_i32(4)) + row = b.div(off, b.const_i32(DK)) + col4 = b.mod(off, b.const_i32(DK)) + if raw is not None: + gidx = raw.qk_off(tile, row, col4) + gval = b.global_load_vN(ctx.g_ptr, gidx, ELEM, 4) + staged.append((ctx.g_lds, row, col4, gval, 4, valid, "g")) + else: + staged.append( + ( + ctx.g_lds, + row, + col4, + b.global_load_vN(ctx.g_ptr, b.add(tile_cd, off), F32, 4), + 4, + valid, + None, + ) + ) + for i in range((N_CD + 8 * BLOCK - 1) // (8 * BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + valid = None + safe = vidx + if N_CD % (8 * BLOCK): + valid = b.cmp_gt(b.const_i32(N_CD // 8), vidx) + safe = b.select(valid, vidx, b.const_i32(N_CD // 8 - 1)) + off = b.mul(safe, b.const_i32(8)) + row = b.div(off, b.const_i32(DK)) + col = b.mod(off, b.const_i32(DK)) + if raw is not None: + gidx = raw.qk_off(tile, row, col) + kval = b.global_load_vN(ctx.k_ptr, gidx, ELEM, 8) + qval = b.global_load_vN(ctx.q_ptr, gidx, ELEM, 8) + else: + gidx = b.add(tile_cd, off) + kval = b.global_load_vN(ctx.k_ptr, gidx, ELEM, 8) + qval = b.global_load_vN(ctx.q_ptr, gidx, ELEM, 8) + staged.append((ctx.k_lds, row, col, kval, 8, valid, "k")) + staged.append((ctx.q_lds, row, col, qval, 8, valid, "q")) + if raw is not None: + bcol = b.select(b.cmp_gt(b.const_i32(C), tid), tid, b.const_i32(C - 1)) + beta = b.global_load_f32(ctx.beta_ptr, raw.beta_off(tile, bcol)) + else: + bcol = b.select(b.cmp_gt(b.const_i32(C), tid), tid, b.const_i32(C - 1)) + beta = b.global_load_f32(ctx.beta_ptr, b.add(tile_c, bcol)) + return staged, beta, tile + + +def _emit_stage_commit(ctx: _ChunkCtx, issued) -> None: + """Write what :func:`_emit_stage_issue` loaded into the staging tiles.""" + b = ctx.b + staged, beta, tile = issued + spec = ctx.spec + raw = ctx.raw + bh = None + head = None + if raw is not None: + bh = b.div(tile, raw.nc) + head = b.mod(bh, raw.H) + + for lds, row, col, value, n, valid, kind in staged: + with b.scf_if(valid) if valid is not None else nullcontext(): + if kind == "g" and raw is not None and spec.fuse_gate: + for j in range(4): + raw_g = b.vec_extract(value, j) + col_j = b.add(col, b.const_i32(j)) + gate = _apply_gate_fused( + b, + raw_g, + raw.a_log, + raw.dt_bias, + head, + col_j, + spec.lower_bound, + has_dt_bias=spec.has_dt_bias, + dk=ctx.DK, + ) + _st(b, lds, row, col_j, value=gate, n=1) + elif kind in ("k", "q") and raw is not None and spec.fuse_qk_l2norm: + vals = [b.cast_to_f32(b.vec_extract(value, j)) for j in range(8)] + inv = _l2norm_scale8(b, vals) + vals = [b.fmul(v, inv) for v in vals] + _st( + b, + lds, + row, + col, + value=b.vec_pack( + [b.cast_f32_to(v, ctx.ELEM) for v in vals], ctx.ELEM + ), + n=8, + ) + else: + _st(b, lds, row, col, value=value, n=n) + if spec.fuse_beta_sigmoid and raw is not None: + beta = _sigmoid_via_exp2(b, beta) + with b.scf_if(b.cmp_gt(b.const_i32(ctx.C), ctx.tid)): + _st(b, ctx.beta_lds, ctx.tid, value=beta, n=1) + + +def _emit_stage_inputs(ctx: _ChunkCtx, ch) -> None: + """g / k / q / beta for ``ch`` from HBM into their staging tiles.""" + _emit_stage_commit(ctx, _emit_stage_issue(ctx, ch)) + + +class _InputPrefetch: + """Stage the *next* chunk's g / k / q / beta from inside this chunk's scan. + + All four staging tiles are dead for the whole scan phase -- g and k are last + read by Kt in the solve window, q by the elementwise sweep, beta by the T' + construction -- so the next chunk's inputs are written straight back into + them and no second buffer is needed. + + The hand-off also needs no barrier of its own. The tile-phase rendezvous + ahead of the scan separates the write from this chunk's last read (WAR), and + the scan's own rendezvous separate it from the next chunk's first read + (RAW). Both already exist for other reasons. The chunk loop therefore drops + its post-staging barrier outright, and the load retires behind the scan's + matmuls instead of at the head of the next chunk. + """ + + def __init__(self, ctx: _ChunkCtx, ch_next): + self.ctx = ctx + self.ch = ch_next + + def issue(self): + return _emit_stage_issue(self.ctx, self.ch) + + def commit(self, issued) -> None: + _emit_stage_commit(self.ctx, issued) + + +def _emit_v_issue(ctx: _ChunkCtx, ch, v_ptr): + """Issue this chunk's V loads well before their LDS commit. + + The fused 512-thread shape has exactly one 128-bit V vector per thread. + Keep the generic loop for other legal compile-time shapes; tail threads + read the final vector again and predicate only the eventual LDS write so + the loaded SSA value can live outside an ``scf.if``. + """ + b = ctx.b + C, EV, BLOCK = ctx.C, ctx.spec.head_v, ctx.BLOCK + n_v = (C * EV) // 8 + cv = b.mul(ch.tile, b.const_i32(C * EV)) + issued = [] + for i in range((n_v + BLOCK - 1) // BLOCK): + vidx = b.add(ctx.tid, b.const_i32(i * BLOCK)) + safe = b.select(b.cmp_gt(b.const_i32(n_v), vidx), vidx, b.const_i32(n_v - 1)) + off = b.mul(safe, b.const_i32(8)) + issued.append( + ( + vidx, + off, + b.global_load_vN(v_ptr, b.add(cv, off), ctx.ELEM, 8), + ) + ) + return issued + + +def _emit_v_commit(ctx: _ChunkCtx, v_lds, issued) -> None: + """Park hoisted V vectors after both C x C MFMAs release the V tile.""" + b = ctx.b + EV = ctx.spec.head_v + n_v = (ctx.C * EV) // 8 + for vidx, off, value in issued: + with b.scf_if(b.cmp_gt(b.const_i32(n_v), vidx)): + _st( + b, + v_lds, + b.div(off, b.const_i32(EV)), + b.mod(off, b.const_i32(EV)), + value=value, + n=8, + ) + + +def _emit_idle_during_solve(ctx: _ChunkCtx, ch, sink) -> None: + """Work the non-wave-0 threads do while wave 0 runs the triangular solve. + + Those waves would otherwise sit on the trailing ``s_barrier``. Kt does not + read the solve tiles (it is ``(K * gamma_C / Gamma)^T`` off ``k``, ``g``, + ``gl``), and V is not a tile-phase operand at all, so both are safe to + issue in this window. Wave 0 does not participate: the original Kt loop + dealt slots out as ``tid + i*BLOCK``, so the coverage is remapped onto + ``BLOCK - 64`` workers. + + V used to be loaded here as well, but that put a ``vmcnt`` drain directly + after each load. It is now issued before the elementwise/MFMA section and + committed once the two C x C MFMAs release Y, leaving this window + exclusively for Kt. + """ + b = ctx.b + tid, BLOCK, N_CD = ctx.tid, ctx.BLOCK, ctx.N_CD + workers = BLOCK - 64 + wid = b.sub(tid, b.const_i32(64)) + + n_kt = N_CD // 8 + for i in range((n_kt + workers - 1) // workers): + vidx = b.add(wid, b.const_i32(i * workers)) + with b.scf_if(b.cmp_gt(b.const_i32(n_kt), vidx)): + _emit_kt_slot(ctx, ch, sink, vidx) + + +def _emit_chunk_tiles( + ctx: _ChunkCtx, + tile, + sink, + v_ptr=None, + v_lds=None, + *, + overlap_solve: bool = True, + stage_inputs: bool = True, +) -> None: + """Emit the six state-independent tiles for the chunk indexed by ``tile``. + + Reads ``q/k/g/beta`` for the chunk out of HBM and hands each finished tile + to ``sink``. Nothing here depends on the state recurrence, so this is the + whole parallel part of a chunkwise KDA forward. + + ``overlap_solve`` (prep, and fused at ``block_size >= 512``) runs Kt on the + idle waves during the wave-0 solve. A 256-thread fused group has only three + helper waves; the extra control flow in that already-fat kernel lost more + than the overlap saved, so that path keeps the uniform post-solve Kt. + ``v_ptr`` is fused-only and only set with ``overlap_solve``. + + ``stage_inputs=False`` says the inputs are already in their staging tiles -- + put there by the previous chunk's scan, see :class:`_InputPrefetch` -- so + both the staging loads and the barrier that published them are skipped. + """ + b, spec = ctx.b, ctx.spec + C, DK, BLOCK = ctx.C, ctx.DK, ctx.BLOCK + CREF, HALF, N_CD = ctx.CREF, ctx.HALF, ctx.N_CD + ELEM, atom, scale = ctx.ELEM, ctx.atom, ctx.scale + tid, lane, lane_m, frag_k_off = ctx.tid, ctx.lane, ctx.lane_m, ctx.frag_k_off + # The triangular rank update covers the whole algorithmic chunk on wave 0. + # C=16 uses the K=32 atom and zero-pads the contraction's upper half. + solve_atom = MfmaAtom.bf16_16x16x32() if C == 16 else MfmaAtom.bf16_32x32x16() + solve_lane_m = b.mod(lane, b.const_i32(solve_atom.m)) + solve_frag_k_off = b.mul( + b.div(lane, b.const_i32(solve_atom.m)), + b.const_i32(solve_atom.a_per_lane), + ) + g_lds, k_lds, q_lds = ctx.g_lds, ctx.k_lds, ctx.q_lds + x_lds, y_lds, xq_lds = ctx.x_lds, ctx.y_lds, ctx.xq_lds + m_lds, a_lds, qk_lds = ctx.m_lds, ctx.a_lds, ctx.qk_lds + tb_lds, zs_lds = ctx.tb_lds, ctx.zs_lds + beta_lds, gl_lds = ctx.beta_lds, ctx.gl_lds + c_log2e = ctx.c_log2e + ex2, lds_get = ctx.ex2, ctx.lds_get + lds_get8_f32, lds_get8_elem = ctx.lds_get8_f32, ctx.lds_get8_elem + lds_put8 = ctx.lds_put8 + + ch = _ChunkOffsets(ctx, tile) + v_issued = None + + # ================================================================= + # 1. stage g / k / q into LDS with 128-bit transactions + # ================================================================= + # A 128-bit vector never straddles a row: DK is a multiple of both 4 (fp32) + # and 8 (bf16), so (row, col) decomposition of the flat vector index is + # exact and the LDS destination stays lane-contiguous. + if stage_inputs: + _emit_stage_inputs(ctx, ch) + b.sync_lds_only() + + # ================================================================= + # 2. in-place cumulative log decay, scaled to log2 + # ================================================================= + # One thread owns a whole (half-chunk, channel) column, so the running sum + # stays in a register and the only cross-thread step is folding the first + # half's total into the second. + NG, GROUP = ctx.NG, ctx.GROUP + d_ch = b.mod(tid, b.const_i32(DK)) + grp = b.div(tid, b.const_i32(DK)) + row0 = b.mul(grp, b.const_i32(GROUP)) + acc = b.const_f32(0.0) + for i in range(GROUP): + rr = b.add(row0, b.const_i32(i)) + acc = b.fadd(acc, lds_get(g_lds, [rr, d_ch])) + _st(b, g_lds, rr, d_ch, value=b.fmul(acc, c_log2e), n=1) + b.sync_lds_only() + + # Fold each group's running total into every later group. A group's local + # total sits at its last row, so the offsets are read up front -- the writes + # below land on some of those same rows once NG > 2, and the correct offset + # is the pre-fold local total. + totals = [ + lds_get(g_lds, [b.const_i32((j + 1) * GROUP - 1), d_ch]) for j in range(NG - 1) + ] + if NG > 2: + b.sync_lds_only() + # Only rows above the first group need folding, and they are dealt out + # round-robin across the NG threads of a channel, so every thread writes the + # same count. ``NG`` divides ``GROUP``, so all NG rows of one step fall in + # the same group and the offset for a step is a compile-time sum. + for i in range((C - GROUP) // NG): + base = GROUP + i * NG + rr = b.add(b.const_i32(base), grp) + off = totals[0] + for j in range(1, base // GROUP): + off = b.fadd(off, totals[j]) + _st(b, g_lds, rr, d_ch, value=b.fadd(lds_get(g_lds, [rr, d_ch]), off), n=1) + b.sync_lds_only() + + # ================================================================= + # 3. dec = gamma_C, and cache the whole-chunk log decay for Kt + # ================================================================= + with b.scf_if(b.cmp_gt(b.const_i32(DK // 4), tid)): + col4 = b.mul(tid, b.const_i32(4)) + gl = [ + lds_get(g_lds, [b.const_i32(C - 1), b.add(col4, b.const_i32(j))]) + for j in range(4) + ] + _st(b, gl_lds, col4, value=b.vec_pack(gl, F32), n=4) + sink.dec(ctx, ch, col4, [ex2(v) for v in gl]) + + # V is independent of the remaining tile work. Issue it here -- after the + # cumsum's temporaries die, but before the long elementwise + dual-MFMA + # section -- and consume it only after those MFMAs release Y. This leaves + # hundreds of instructions to hide the old ~508-cycle HBM round trip + # without extending four VGPRs across the entire tile phase. + if v_ptr is not None and v_lds is not None: + v_issued = _emit_v_issue(ctx, ch, v_ptr) + + # ================================================================= + # 4. fused elementwise pass: GK / GQ out, and all three MFMA operands + # ================================================================= + # Every one of these five tiles is a pointwise function of the same + # (k, q, Gcum) element, so they are built in a single sweep: g/k/q are read + # once, and one thread owns eight consecutive channels of a row so every + # LDS and global access in the pass is a 128-bit transaction. + # + # GK = K * Gamma -> global + # GQ = Q * Gamma * scale -> global + # X = K * e^(Gc - Gref) -> LDS, Akk's A operand + # Y = K * e^(Gref - Gc) -> LDS, shared B operand + # XQ = Q * scale * e^(Gc - Gref) -> LDS, Aqk's A operand + # + # X/Y/XQ are factored against the chunk's midpoint row so that X Y^T + # reconstructs k_i . (k_j * Gamma_i / Gamma_j) with each factor's exponent + # bounded by half the chunk's decay range instead of all of it. + ew_col = b.mod(b.mul(tid, b.const_i32(8)), b.const_i32(DK)) + ew_row = b.div(b.mul(tid, b.const_i32(8)), b.const_i32(DK)) + ew_rstep = (BLOCK * 8) // DK + gref8 = lds_get8_f32(g_lds, b.const_i32(CREF), ew_col) + + for p in range((N_CD + 8 * BLOCK - 1) // (8 * BLOCK)): + vidx = b.add(tid, b.const_i32(p * BLOCK)) + guard = ( + b.scf_if(b.cmp_gt(b.const_i32(N_CD // 8), vidx)) + if N_CD % (8 * BLOCK) + else nullcontext() + ) + with guard: + row = b.add(ew_row, b.const_i32(p * ew_rstep)) + off = b.mul(vidx, b.const_i32(8)) + g8 = lds_get8_f32(g_lds, row, ew_col) + k8 = lds_get8_elem(k_lds, row, ew_col) + q8 = lds_get8_elem(q_lds, row, ew_col) + + gk8, gq8, x8, y8, xq8 = [], [], [], [], [] + for j in range(8): + gc, kv = g8[j], k8[j] + qv = b.fmul(q8[j], scale) + dref = b.fsub(gc, gref8[j]) + emr = ex2(dref) + erm = ex2(b.fsub(gref8[j], gc)) + if not sink.deferred_gk_gq: + eg = ex2(gc) + gk8.append(b.fmul(kv, eg)) + gq8.append(b.fmul(qv, eg)) + x8.append(b.fmul(kv, emr)) + y8.append(b.fmul(kv, erm)) + xq8.append(b.fmul(qv, emr)) + + if not sink.deferred_gk_gq: + sink.gk_gq(ctx, ch, row, ew_col, off, gk8, gq8) + lds_put8(x_lds, row, ew_col, x8) + lds_put8(y_lds, row, ew_col, y8) + lds_put8(xq_lds, row, ew_col, xq8) + b.sync_lds_only() + + panelized = atom.m < C + if panelized: + # Four 16x16 panels, assigned round-robin to eight waves. Duplicating + # each panel once keeps every wave useful and avoids a reduction: the + # duplicate waves write identical values to the same LDS addresses. + wave = b.div(tid, b.const_i32(64)) + panel = b.mod(wave, b.const_i32((C // atom.m) ** 2)) + panel_m = b.div(panel, b.const_i32(C // atom.n)) + panel_n = b.mod(panel, b.const_i32(C // atom.n)) + panel_row = b.mul(panel_m, b.const_i32(atom.m)) + panel_col = b.mul(panel_n, b.const_i32(atom.n)) + tile_lane_m = b.mod(lane, b.const_i32(atom.m)) + tile_frag_k_off = b.mul( + b.div(lane, b.const_i32(atom.m)), b.const_i32(atom.a_per_lane) + ) + else: + panel_row = panel_col = b.const_i32(0) + tile_lane_m = b.mod(lane, b.const_i32(atom.m)) + tile_frag_k_off = b.mul( + b.div(lane, b.const_i32(atom.m)), + b.const_i32(atom.a_per_lane), + ) + + def tile_output(i): + row, col = atom.lane_to_output(b, lane, i) + return b.add(panel_row, row), b.add(panel_col, col) + + def cxc_mfma(a_smem, b_smem): + """This wave's panel of ``C x C = (C x DK)(C x DK)^T``.""" + acc = atom.zero_acc(b) + for ks in range(spec.k_steps): + kb = b.add(b.const_i32(ks * atom.k), tile_frag_k_off) + av = _ld(b, a_smem, b.add(panel_row, tile_lane_m), kb, dtype=ELEM, n=8) + bv = _ld(b, b_smem, b.add(panel_col, tile_lane_m), kb, dtype=ELEM, n=8) + acc = atom.emit(b, av, bv, acc) + return acc + + # Both products share the Y operand and neither depends on the other, so + # they are issued back to back off the one staging barrier -- the second + # MFMA chain fills the first's latency, and the operand rebuild plus the two + # extra barriers the sequential version needed are gone. + acc_kk = cxc_mfma(x_lds, y_lds) + acc_qk = cxc_mfma(xq_lds, y_lds) + + if sink.deferred_gk_gq: + # A fused consumer has no spare LDS for GK/GQ, so they are built into the + # X/XQ staging tiles now that both MFMAs have consumed them. The barrier + # is what makes that safe: every wave issues the products above, so a + # wave running ahead would otherwise overwrite an operand another wave is + # still reading. The recomputed exp2 is cheaper than the tiles it saves. + b.sync_lds_only() + _emit_gk_gq_pass(ctx, ch, sink) + if v_issued is not None: + _emit_v_commit(ctx, v_lds, v_issued) + + # ================================================================= + # 5. T' = StrictTril(Diag(beta) Akk), RHS = Diag(beta), Aqk = Tril(.) + # ================================================================= + # Both accumulators are already in the atom's output layout, so the masks + # and the beta scale are applied in-register on the way to LDS. + # ``zs`` accumulates the solved columns for the rank update and must read as + # zero for rows not yet solved, so the unsolved part contributes nothing. + zs_zero_cols = ctx.PCB if C < solve_atom.k else C + n_vec_zs = (C * zs_zero_cols) // 8 + zero8 = b.vec_pack([b.cast_f32_to(b.const_f32(0.0), ELEM)] * 8, ELEM) + with b.scf_if(b.cmp_gt(b.const_i32(n_vec_zs), tid)): + zoff = b.mul(tid, b.const_i32(8)) + _st( + b, + zs_lds, + b.div(zoff, b.const_i32(zs_zero_cols)), + b.mod(zoff, b.const_i32(zs_zero_cols)), + value=zero8, + n=8, + ) + # The C=16 solve/scan contract over a K=32 atom. Keep only the valid lower + # half of T'; the padded columns must read as exact zero. + if ctx.PCB > C: + n_pad = (C * (ctx.PCB - C)) // 8 + with b.scf_if(b.cmp_gt(b.const_i32(n_pad), tid)): + poff = b.mul(tid, b.const_i32(8)) + _st( + b, + tb_lds, + b.div(poff, b.const_i32(ctx.PCB - C)), + b.add( + b.const_i32(C), + b.mod(poff, b.const_i32(ctx.PCB - C)), + ), + value=zero8, + n=8, + ) + if hasattr(sink, "kt_lds"): + n_kt_pad = (DK * (ctx.PCB - C)) // 8 + for i in range((n_kt_pad + BLOCK - 1) // BLOCK): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + with b.scf_if(b.cmp_gt(b.const_i32(n_kt_pad), vidx)): + poff = b.mul(vidx, b.const_i32(8)) + _st( + b, + sink.kt_lds, + b.div(poff, b.const_i32(ctx.PCB - C)), + b.add( + b.const_i32(C), + b.mod(poff, b.const_i32(ctx.PCB - C)), + ), + value=zero8, + n=8, + ) + + # The chunk-wide atom is redundantly computed by every wave, so wave 0 is + # its sole writer. In panel mode every wave owns one panel (duplicated once) + # and all waves must publish their panel. + writer = nullcontext() if panelized else b.scf_if(b.cmp_gt(b.const_i32(64), tid)) + with writer: + for i in range(atom.c_per_lane): + row, col = tile_output(i) + bet = lds_get(beta_lds, [row]) + tp = b.select( + b.cmp_gt(row, col), + b.fmul(bet, b.vec_extract(acc_kk, i)), + b.const_f32(0.0), + ) + # fp32 for the in-block substitution (it multiplies the solved + # values directly, so it is the precision-critical copy) and bf16 + # for the rank-update MFMA operand. + _st(b, m_lds, row, col, value=tp, n=1) + _st(b, tb_lds, row, col, value=b.cast_f32_to(tp, ELEM), n=1) + # The substitution reads its starting value straight out of a_mat, + # so seed it with the right-hand side Diag(beta) here; each later + # block overwrites its own rows with (RHS - rank update). + _st( + b, + a_lds, + row, + col, + value=b.select(b.cmp_eq(row, col), bet, b.const_f32(0.0)), + n=1, + ) + _st( + b, + qk_lds, + row, + col, + value=b.select( + b.cmp_gt(b.add(row, b.const_i32(1)), col), # row >= col + b.vec_extract(acc_qk, i), + b.const_f32(0.0), + ), + n=1, + ) + b.sync_lds_only() + + # ================================================================= + # 6. A = (I + T')^-1 Diag(beta) by blocked forward substitution + # ================================================================= + # Per block: rank-update the block's rows against every already-solved row + # (a matmul, issued on MFMA), then substitute within the block (scalar, and + # the only irreducibly serial part). The whole loop runs on wave 0 alone, so + # the LDS hand-offs between the two halves need no s_barrier -- they are + # ordered by the wave's own lgkmcnt -- and the other waves wait once at the + # end instead of twice per block. Those waiting waves emit Kt (and, on the + # fused path, prefetch V into the dead Y tile) in the same window: both are + # independent of the solve, so they come off the critical path without an + # extra barrier. + # + # Raising wave 0 to ``s_setprio(3)`` across this region measures as nothing + # (9.599 ms vs 9.587 ms at 32x16x2048, inside a ~0.02 ms spread). Priority + # only decides issue arbitration between waves that are both ready, and the + # substitution below is bound by its own serial lgkmcnt round trips, so + # there is no issue slot for wave 0 to win off the one helper sharing its + # SIMD. Do not retry without first making the solve issue-bound. + BS = spec.tile.solve_block + NB = C // BS + ks_solve = (C + solve_atom.k - 1) // solve_atom.k + with b.scf_if(b.cmp_gt(b.const_i32(64), tid)): + for bi in range(NB): + if bi > 0: + # The previous block wrote its solved rows into zs, and this + # rank update reads them at a row another lane wrote. Being on + # one wave orders the two only once the write has retired, and + # the LDS addresses are computed independently enough that the + # backend does not pair them on its own -- so drain lgkmcnt + # explicitly. Cheaper than an s_barrier, and correctness here + # does not survive without it once blocks are adjacent. + b.s_waitcnt(lgkmcnt=0) + # U = T' @ Zs over the full C x C x C shape. Rows outside this + # block and columns of already-solved-but-irrelevant rows cost + # nothing to include: Zs is zero wherever a row is unsolved, so + # the extra lanes of the product are exact zeros. + accu = solve_atom.zero_acc(b) + for ks in range(ks_solve): + kb = b.add(b.const_i32(ks * solve_atom.k), solve_frag_k_off) + accu = solve_atom.emit( + b, + _ld(b, tb_lds, solve_lane_m, kb, dtype=ELEM, n=8), + _ld(b, zs_lds, solve_lane_m, kb, dtype=ELEM, n=8), + accu, + ) + # Slots [4*bi*t, 4*(bi+1)*t) are exactly this block's rows, so + # the update folds straight into a_mat's seeded right-hand side + # and the substitution below needs no separate staging tile. + if solve_atom.m == 32: + update_slots = range(bi * 4 * (BS // 8), (bi + 1) * 4 * (BS // 8)) + for i in update_slots: + row, col = solve_atom.lane_to_output(b, lane, i) + _st( + b, + a_lds, + row, + col, + value=b.fsub( + lds_get(a_lds, [row, col]), + b.vec_extract(accu, i), + ), + n=1, + ) + else: + for i in range(solve_atom.c_per_lane): + row, col = solve_atom.lane_to_output(b, lane, i) + with b.scf_if(b.cmp_ge(row, b.const_i32(bi * BS))): + with b.scf_if(b.cmp_lt(row, b.const_i32((bi + 1) * BS))): + _st( + b, + a_lds, + row, + col, + value=b.fsub( + lds_get(a_lds, [row, col]), + b.vec_extract(accu, i), + ), + n=1, + ) + # Same hand-off in the other direction: the substitution below + # reads rows this loop just wrote from a different lane. + b.s_waitcnt(lgkmcnt=0) + + with b.scf_if(b.cmp_gt(b.const_i32(C), lane)): + zblk = [] + for r in range(bi * BS, (bi + 1) * BS): + cr = b.const_i32(r) + val = lds_get(a_lds, [cr, lane]) + for j in range(bi * BS, r): + val = b.fsub( + val, + b.fmul( + lds_get(m_lds, [cr, b.const_i32(j)]), + zblk[j - bi * BS], + ), + ) + zblk.append(val) + _st(b, a_lds, cr, lane, value=val, n=1) + # Transposed, so a thread's whole solved block is contiguous and + # goes out as one ds_write_b128 -- and lands in the (n, k) order + # the next rank update's B operand wants. + if bi + 1 < NB: + for h in range(BS // 8): + lds_put8( + zs_lds, + lane, + b.const_i32(bi * BS + h * 8), + zblk[h * 8 : h * 8 + 8], + ) + if overlap_solve: + with b.scf_if(b.cmp_ge(tid, b.const_i32(64))): + _emit_idle_during_solve(ctx, ch, sink) + b.sync_lds_only() + + # ================================================================= + # 7. A and Aqk out, 128-bit stores + # ================================================================= + n_vec_cc = (C * C) // 8 + + def store_cxc(src, hook): + """One C x C fp32 LDS tile out as bf16, 128-bit per thread.""" + for i in range(max(1, n_vec_cc // BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + with b.scf_if(b.cmp_gt(b.const_i32(n_vec_cc), vidx)): + off = b.mul(vidx, b.const_i32(8)) + row = b.div(off, b.const_i32(C)) + col = b.mod(off, b.const_i32(C)) + vals = [ + lds_get(src, [row, b.add(col, b.const_i32(j))]) for j in range(8) + ] + hook(ctx, ch, off, row, col, vals) + + store_cxc(a_lds, sink.a) + store_cxc(qk_lds, sink.aqk) + if not overlap_solve: + _emit_kt_all_waves(ctx, ch, sink) + + +def build_kda_chunk_prep(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> KernelDef: + """Build the IR for the per-chunk tile builder. + + Kernel signature:: + + (q, k: ptr, # [NT, C * DK] or raw [B,T,H,D] + g: ptr, # [NT, C * DK] per-channel log decay, or raw bf16 + beta: ptr, # [NT, C] + a_out, gk_out, gq_out, aqk_out, kt_out: ptr, + dec_out: ptr, # [NT, DK] + scale: f32) + + Raw mode adds ``a_log``, optional ``dt_bias``, and ``batch``/``heads``/``tseq``/``nc``. + + Grid ``(NT, 1, 1)`` where ``NT = BH * NC``; block ``(block_size, 1, 1)``. + """ + ok, why = is_valid_spec(spec, arch=arch) + if not ok: + raise ValueError(f"invalid kda_chunk_prep spec for {arch}: {why}") + + ELEM = _DTYPE_IR[spec.dtype] + t = spec.tile + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = t.block_size + if t.waves_per_eu: + b.kernel.attrs["waves_per_eu"] = ( + t.waves_per_eu, + t.waves_per_eu, + ) + + q_ptr = b.param("q_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + k_ptr = b.param("k_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + g_elem = ELEM if spec.raw_inputs else F32 + g_ptr = b.param("g_ptr", PtrType(g_elem, "global"), readonly=True, align=16) + beta_ptr = b.param("beta_ptr", PtrType(F32, "global"), readonly=True, align=4) + a_ptr = b.param("a_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + gk_ptr = b.param("gk_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + gq_ptr = b.param("gq_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + aqk_ptr = b.param("aqk_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + kt_ptr = b.param("kt_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + dec_ptr = b.param("dec_ptr", PtrType(F32, "global"), writeonly=True, align=16) + scale = b.param("scale", F32) + + raw = None + if spec.raw_inputs: + a_log_ptr = b.param("a_log_ptr", PtrType(F32, "global"), readonly=True, align=4) + dt_bias_ptr = b.param( + "dt_bias_ptr", PtrType(F32, "global"), readonly=True, align=4 + ) + batch = b.param("batch", I32) + heads = b.param("heads", I32) + tseq = b.param("tseq", I32) + nc = b.param("nc", I32) + raw = _RawTokenAddr( + b, + heads=heads, + tseq=tseq, + nc=nc, + chunk=t.chunk, + dk=spec.head_k, + a_log=a_log_ptr, + dt_bias=dt_bias_ptr, + ) + + ctx = _ChunkCtx(b, spec, (q_ptr, k_ptr, g_ptr, beta_ptr, scale, raw)) + sink = _GlobalTileSink(a_ptr, gk_ptr, gq_ptr, aqk_ptr, kt_ptr, dec_ptr) + _emit_chunk_tiles(ctx, b.block_id_x(), sink) + + b.ret() + return b.kernel + + +@dataclass(frozen=True) +class KdaChunkFusedSpec: + """Compile-time spec for the fused per-chunk-tiles + state-scan kernel. + + One workgroup owns a whole (batch, head) and walks its chunks in order, so + the six per-chunk tiles are built and consumed in LDS and never reach HBM. + That is the entire point of the fusion: on the split path the tiles are + written once by the tile builder and read back once per v-split by the scan, + and that round-trip is several times the traffic the problem actually needs. + + The trade is parallelism. The split tile builder has one workgroup per + chunk; this kernel has one per (batch, head), so it needs ``BH`` to be + comfortably above the CU count to fill the device. It is the right choice + for prefill at batch scale, not for a single short sequence. + """ + + head_k: int = 128 + head_v: int = 128 + dtype: str = "bf16" + # Fused-only tuned schedule. The larger row pitches improve LDS bank + # distribution while the 512-thread panel mapping keeps eight waves + # resident. Do not copy these pads to the split scan: its LDS budget is + # deliberately capped for two workgroups per CU. + tile: KdaTileSpec = KdaTileSpec( + block_size=512, + pad_dk=16, + pad_cb=16, + tile_atom_m=16, + scan_atom_m=16, + ) + has_initial_state: bool = False + store_final_state: bool = True + # Stage each chunk's inputs from the previous chunk's scan phase rather than + # at the head of its own tile phase. Costs no LDS (the staging tiles are + # dead across the scan) and removes the per-chunk staging barrier; see + # :class:`_InputPrefetch`. Worth 7-9% at 32x16x2048 and bitwise identical, + # so it is on by default; the flag exists to keep the A/B measurable. + prefetch_inputs: bool = True + # Reuse tile-only LDS for the scan extras. This cannot be combined with the + # barrier-free input prefetch because that prefetch deliberately writes the + # same staging addresses while the scan is live. + overlay_lds: bool = False + name: str = "rocke_kda_chunk_fused" + + @property + def prep(self) -> KdaChunkPrepSpec: + """The per-chunk tile builder whose emission this kernel reuses.""" + return KdaChunkPrepSpec( + head_k=self.head_k, + head_v=self.head_v, + dtype=self.dtype, + tile=self.tile, + ) + + @property + def atom(self) -> MfmaAtom: + return self.prep.atom + + @property + def scan_atom(self) -> MfmaAtom: + """The atom the state scan runs on; see ``KdaTileSpec.scan_atom_m``.""" + return _scan_atom(self.tile, self.atom) + + @property + def state_tiles(self) -> int: + """Atom tiles per wave across the state's ``DK`` extent. + + Each wave owns one ``scan_atom.m``-row band of ``S^T`` and the full head + dimension, which is what makes every one of the five per-chunk products + partition the same way and keeps the state in that wave's registers. + """ + return self.head_k // self.scan_atom.n + + def lds_bytes(self) -> int: + """Exact fused allocation for the selected LDS schedule. + + The normal schedule uses typed allocations. The overlay schedule uses + one byte pool with explicit views so loop liveness cannot keep + tile-only and scan-only buffers simultaneously live. + """ + t = self.tile + C, DK, EV = t.chunk, self.head_k, self.head_v + if self.overlay_lds: + return _fused_lds_layout(DK, EV, t)["pool"] + return ( + self.prep.lds_bytes() + + 2 * DK * (C + t.pad_cb) # kt_s bf16 (DK x C), B operand + + 2 * EV * (DK + t.pad_dk) # stb_s bf16 mirror of S^T + + 2 * EV * (C + t.pad_cb) # vn_s bf16 residual, then V~^T + ) + + def kernel_name(self) -> str: + parts = (f"dk{self.head_k}", f"dv{self.head_v}", self.dtype) + parts += self.tile.name_parts() + if self.has_initial_state: + parts += ("h0",) + if not self.store_final_state: + parts += ("noht",) + if not self.prefetch_inputs: + parts += ("nopf",) + if self.overlay_lds: + parts += ("ovl",) + return kernel_name_join(self.name, *parts) + + +def is_valid_fused_spec( + spec: KdaChunkFusedSpec, arch: str = "gfx950" +) -> Tuple[bool, str]: + """Return ``(ok, reason)`` for a fused spec on ``arch``.""" + ok, why = is_valid_spec(spec.prep, arch=arch) + if not ok: + return False, why + + t = spec.tile + if spec.overlay_lds and spec.prefetch_inputs: + return False, ( + "overlay_lds aliases g/k/q/beta with live scan extras; " + "set prefetch_inputs=False" + ) + if spec.overlay_lds and t.chunk != 32: + return False, "overlay_lds is only modeled for the C32 fused layout" + ok, why = _check_scan_partition(spec.scan_atom, t, spec.head_k, spec.head_v) + if not ok: + return False, why + lds = spec.lds_bytes() + if lds > LDS_LIMIT: + return False, f"LDS request {lds} B exceeds the {LDS_LIMIT} B budget" + return True, "ok" + + +def _check_scan_partition( + atom: MfmaAtom, t: KdaTileSpec, head_k: int, head_v: int +) -> Tuple[bool, str]: + """Validate the one rule that partitions every product in the scan body. + + Each wave owns one ``atom.m``-row band of ``S^T`` and the full head + dimension, so the v extent has to be exactly covered by the waves. Anything + else needs a second partitioning rule for the same accumulator, which is how + cross-wave reductions creep in. + """ + waves = t.num_waves + if head_v != atom.m * waves: + return False, ( + f"head_v ({head_v}) must equal scan atom.m * waves " + f"({atom.m} * {waves} = {atom.m * waves}); each wave owns one " + "row band of the state" + ) + if head_k % atom.n: + return False, ( + f"head_k ({head_k}) must be a multiple of the atom N extent " f"({atom.n})" + ) + if t.chunk < atom.k and t.chunk + t.pad_cb < atom.k: + return False, ( + f"chunk + pad_cb ({t.chunk + t.pad_cb}) must cover scan MFMA K " + f"({atom.k}) for zero padding" + ) + # The chunk extent lands on the N side of two products and the M side of a + # third, so tiling it needs a square atom that divides it. + if atom.m != atom.n or t.chunk % atom.m: + return False, ( + f"scan atom ({atom.m}x{atom.n}) must be square and divide chunk " + f"({t.chunk}); the chunk extent is tiled on both operand sides" + ) + return True, "ok" + + +class _LdsTileSink: + """Tile destination: LDS, for a scan fused into the same workgroup. + + Each tile lands in bf16 in the exact layout its consumer wants as an MFMA + operand. Nothing is allocated for GK/GQ, A or Aqk: they overwrite staging + tiles that are dead by the time they are produced (the X/XQ MFMA operands + and the two bf16 C x C tiles the triangular solve used for its rank + updates), which is what keeps the fused kernel inside the LDS budget. + ``dec`` needs no destination at all -- the emitter already caches the + whole-chunk log decay in ``gl_s``, and the state decay exponentiates it + where it is applied. + """ + + deferred_gk_gq = True + + def __init__(self, gk_lds, gq_lds, ab_lds, aqb_lds, kt_lds): + self.gk_lds = gk_lds + self.gq_lds = gq_lds + self.ab_lds = ab_lds + self.aqb_lds = aqb_lds + self.kt_lds = kt_lds + + def dec(self, ctx, ch, col4, dec4): + pass + + def gk_gq(self, ctx, ch, row, col, off, gk8, gq8): + ctx.lds_put8(self.gk_lds, row, col, gk8) + ctx.lds_put8(self.gq_lds, row, col, gq8) + + def a(self, ctx, ch, off, row, col, vals): + ctx.lds_put8(self.ab_lds, row, col, vals) + + def aqk(self, ctx, ch, off, row, col, vals): + ctx.lds_put8(self.aqb_lds, row, col, vals) + + def kt(self, ctx, ch, off, dch, r8, vals): + _st( + ctx.b, + self.kt_lds, + dch, + r8, + value=ctx.b.vec_pack(vals, ctx.ELEM), + n=8, + ) + + +class _ScanCtx: + """LDS tiles, lane decomposition and helpers for the state-scan body. + + The scan body is identical whether the six per-chunk tiles were just built + in this workgroup's LDS or staged in from HBM, so both kernels share it. + The only differences are how the tiles arrive and whether the whole-chunk + decay arrives in the log domain (the tile emitter caches the log; the + materialized ``dec`` tile is already exponentiated), which is what + ``dec_is_log`` selects. + + Every wave owns one ``atom.m``-row band of ``S^T`` and the full head + dimension. That single rule partitions all five products, which is what + keeps the state in that wave's accumulators with no cross-wave reduction. + """ + + def __init__( + self, + b: IRBuilder, + *, + atom: MfmaAtom, + chunk: int, + head_k: int, + head_v: int, + block_size: int, + elem, + tiles, + stb_lds, + vn_lds, + v_ptr, + o_ptr, + tid, + lane, + dec_is_log: bool, + ex2=None, + v_lds=None, + v_row_base=0, + ev_stride=None, + io=None, + ): + self.b = b + self.ex2 = ex2 + self.atom = atom + self.C, self.DK = chunk, head_k + self.EV = head_v + self.EV_stride = head_v if ev_stride is None else ev_stride + self.v_row_base = v_row_base + self.io = io + self.BLOCK = block_size + self.ELEM = elem + ( + self.gk_lds, + self.gq_lds, + self.ab_lds, + self.aqb_lds, + self.kt_lds, + self.dec_lds, + ) = tiles + self.stb_lds, self.vn_lds = stb_lds, vn_lds + self.v_ptr, self.o_ptr = v_ptr, o_ptr + # Fused parks V in the dead Y tile during the solve; the residual then + # reads this instead of ``v_ptr``. None on the split scan. + self.v_lds = v_lds + self.tid, self.lane = tid, lane + self.dec_is_log = dec_is_log + self.NS = head_k // atom.n + self.KS_DK = head_k // atom.k + self.KS_C = (chunk + atom.k - 1) // atom.k + self.CPL = atom.c_per_lane + self.N_CV = chunk * head_v + # Atom tiles spanning the chunk extent. The three products whose output + # carries a C extent (Z^T, V~^T and O) are one atom tile per wave only + # while the atom is as wide as the chunk; a narrower atom -- which is + # how the state subtiles down to fewer registers per lane -- splits + # them into this many tiles instead. + self.NC_T = chunk // atom.m + # Operand lane mapping, derived from *this* atom rather than inherited + # from the tile phase's: the two need not be the same atom, and an A + # operand's row/K split is a function of the atom's M extent and + # per-lane K width. Holds for every supported MFMA shape. + self.lane_m = b.mod(lane, b.const_i32(atom.m)) + self.frag_k_off = b.mul( + b.div(lane, b.const_i32(atom.m)), b.const_i32(atom.a_per_lane) + ) + # This wave's band of S^T rows. + self.wrow = b.mul(b.div(tid, b.const_i32(64)), b.const_i32(atom.m)) + + def cv_base(self, tile): + """Flat chunk base for v/o in the chunk-packed layout.""" + return self.b.mul(tile, self.b.const_i32(self.C * self.EV_stride)) + + def global_ev(self, local_ev): + b = self.b + if isinstance(self.v_row_base, int): + if self.v_row_base == 0: + return local_ev + return b.add(b.const_i32(self.v_row_base), local_ev) + return b.add(self.v_row_base, local_ev) + + def v_global(self, tile, chunk_row, local_ev): + """Absolute index of one v/o element.""" + b = self.b + gev = self.global_ev(local_ev) + if self.io is not None: + return self.io.v_off(tile, chunk_row, gev, self.EV_stride) + return b.add( + self.cv_base(tile), + b.add(b.mul(chunk_row, b.const_i32(self.EV_stride)), gev), + ) + + def o_global(self, tile, chunk_row, ev_col, v_row): + """Absolute index of one output element (chunk row x v channel).""" + b = self.b + local_ev = b.add(v_row, ev_col) + return self.v_global(tile, chunk_row, local_ev) + + def slot(self, i): + """Slot ``i``'s (row, col) inside the atom's output tile. + + Recomputed at each use rather than hoisted: the mapping is a couple of + VALU ops off ``lane``, but holding all ``c_per_lane`` pairs live spans + the whole chunk loop at two registers each, and this kernel's occupancy + is register-sensitive. + """ + return self.atom.lane_to_output(self.b, self.lane, i) + + def gemm(self, a_smem, a_row, b_smem, b_row, ksteps, acc): + """``acc += A B^T`` over ``ksteps`` atom steps, both operands from LDS.""" + b = self.b + apl = self.atom.a_per_lane + for ks in range(ksteps): + kb = b.add(b.const_i32(ks * self.atom.k), self.frag_k_off) + av = _ld(b, a_smem, a_row, kb, dtype=self.ELEM, n=apl) + bv = _ld(b, b_smem, b_row, kb, dtype=self.ELEM, n=apl) + acc = self.atom.emit(b, av, bv, acc) + return acc + + def crow(self, j, off=None): + """Row ``lane_m`` of chunk-extent atom tile ``j``, i.e. ``j*m + lane_m``. + + The operand row for tile ``j`` of a product whose output spans the chunk + extent. ``off`` shifts within the tile for the state's ``DK`` extent. + """ + base = j * self.atom.m if off is None else off + return self.b.add(self.b.const_i32(base), self.lane_m) + + def state_idx(self, base, i, ti): + """Global (ev, dk) offset of accumulator slot ``i`` of state tile ``ti``.""" + b = self.b + row, col = self.slot(i) + ev = b.add(self.wrow, row) + dk = b.add(b.const_i32(ti * self.atom.n), col) + return b.add(base, b.add(b.mul(ev, b.const_i32(self.DK)), dk)) + + def publish_state(self, state): + """S^T -> its bf16 mirror, the operand form both consumers read. + + One scalar write per element: a lane's slots are consecutive *rows* of + ``S^T``, so in the operand layout they sit a row pitch apart and cannot + be packed. The mirror itself is unavoidable -- the accumulator and + A-operand lane mappings are transposes of each other, so this round + trip through LDS is what performs the relayout. + """ + b = self.b + for ti in range(self.NS): + for i in range(self.CPL): + row, col = self.slot(i) + _st( + b, + self.stb_lds, + b.add(self.wrow, row), + b.add(b.const_i32(ti * self.atom.n), col), + value=b.cast_f32_to(b.vec_extract(state[ti], i), self.ELEM), + n=1, + ) + + def tiles_for_sink(self): + """The five tile destinations a fused tile emitter writes into. + + ``dec`` is excluded: the emitter already caches the whole-chunk log + decay itself, so the sink has nothing to do for it. + """ + return (self.gk_lds, self.gq_lds, self.ab_lds, self.aqb_lds, self.kt_lds) + + def zero_state(self): + return [self.atom.zero_acc(self.b) for _ in range(self.NS)] + + def load_state(self, ptr, base): + """Seed the accumulators from an ``[BH, DV, DK]`` fp32 state.""" + b = self.b + out = [] + for ti in range(self.NS): + vals = [ + b.global_load_f32(ptr, self.state_idx(base, i, ti)) + for i in range(self.CPL) + ] + out.append(b.vec_pack(vals, F32)) + return out + + def store_state(self, ptr, base, state): + b = self.b + for ti in range(self.NS): + for i in range(self.CPL): + b.global_store_vN( + ptr, self.state_idx(base, i, ti), b.vec_extract(state[ti], i), 1 + ) + + +def _emit_scan_body(sc: _ScanCtx, state, tile, *, prefetch: "_InputPrefetch" = None): + """One chunk of the state recurrence; returns the updated state. + + .. code-block:: text + + Z^T = S^T GK^T EV x C + R^T = V^T - Z^T EV x C (in-register) + V~^T = R^T A^T EV x C + O = GQ S + Aqk V~ C x EV + S^T <- Diag(dec) S^T + V~^T Kt^T EV x DK + + Working transposed keeps every product in ``A B^T`` form with the + contraction on the fastest axis, so no operand ever needs an LDS transpose. + + ``prefetch`` stages the next chunk's inputs across this body: the loads go + out first and the writes land after the ``V~`` rendezvous, so the scan's own + matmuls sit between them and the closing barrier publishes the result. + """ + b, atom = sc.b, sc.atom + lane_m, wrow, CPL = sc.lane_m, sc.wrow, sc.CPL + ELEM = sc.ELEM + NC_T = sc.NC_T + srow = b.add(wrow, lane_m) # this wave's S^T operand row + + # Issued before anything else in the scan so the whole body covers the HBM + # latency; the staging tiles are already dead by here. + issued = prefetch.issue() if prefetch is not None else None + + # C=16 uses the 16x16x32 atom. V~/residual has only 16 real chunk columns, + # so zero the padded K half before any of the three C-contracted products. + if sc.C < atom.k: + pad = atom.k - sc.C + n_pad = (sc.EV * pad) // 8 + for p in range((n_pad + sc.BLOCK - 1) // sc.BLOCK): + vidx = b.add(sc.tid, b.const_i32(p * sc.BLOCK)) + with b.scf_if(b.cmp_gt(b.const_i32(n_pad), vidx)): + poff = b.mul(vidx, b.const_i32(8)) + _st( + b, + sc.vn_lds, + b.div(poff, b.const_i32(pad)), + b.add( + b.const_i32(sc.C), + b.mod(poff, b.const_i32(pad)), + ), + value=b.vec_pack( + [b.cast_f32_to(b.const_f32(0.0), ELEM)] * 8, + ELEM, + ), + n=8, + ) + + # The mirror is a derived copy of the state the caller already carries in + # registers, so it is rebuilt here rather than left behind by the previous + # chunk. That is what keeps its live range inside the scan phase: published + # at the end instead, it would span the tile phase and the LDS packer could + # not alias it onto the staging tiles that are dead by then. + sc.publish_state(state) + b.sync_lds_only() + + # The chunk extent lands on the N side of two products and the M side of + # the third, so tiling it both ways needs a square atom. + assert atom.m == atom.n, "chunk-extent tiling needs a square atom" + + def cidx(jt, col): + """Global chunk index of ``col`` inside chunk-extent tile ``jt``.""" + return col if jt == 0 else b.add(b.const_i32(jt * atom.m), col) + + # ---- Z^T = S^T GK^T, then R^T = V^T - Z^T into the V~ tile ---------- + acc_z = [ + sc.gemm(sc.stb_lds, srow, sc.gk_lds, sc.crow(jt), sc.KS_DK, atom.zero_acc(b)) + for jt in range(NC_T) + ] + # A lane's slots are runs of four consecutive v channels at one fixed chunk + # row, so V arrives in 4-wide loads and the residual never needs an fp32 + # staging tile. On the fused path those loads hit the Y tile idle waves + # filled during the solve; the split scan still reads V from HBM. Every + # residual tile has to land before the next product reads them, so the + # whole chunk extent is written before the barrier. + for jt in range(NC_T): + for grp in range(CPL // 4): + row0, col = sc.slot(4 * grp) + cg = cidx(jt, col) + if sc.v_lds is not None: + vvec = _ld(b, sc.v_lds, cg, b.add(wrow, row0), dtype=ELEM, n=4) + else: + vvec = b.global_load_vN( + sc.v_ptr, + sc.v_global(tile, cg, b.add(wrow, row0)), + ELEM, + 4, + ) + for j in range(4): + row, _ = sc.slot(4 * grp + j) + res = b.fsub( + b.cast_to_f32(b.vec_extract(vvec, j)), + b.vec_extract(acc_z[jt], 4 * grp + j), + ) + _st( + b, + sc.vn_lds, + b.add(wrow, row), + cg, + value=b.cast_f32_to(res, ELEM), + n=1, + ) + b.sync_lds_only() + + # ---- V~^T = R^T A^T ------------------------------------------------- + # Contracts over the full chunk extent, so every tile reads the whole + # residual: all of them are computed before any overwrites it in place. + acc_v = [ + sc.gemm(sc.vn_lds, srow, sc.ab_lds, sc.crow(jt), sc.KS_C, atom.zero_acc(b)) + for jt in range(NC_T) + ] + b.sync_lds_only() + for jt in range(NC_T): + for i in range(CPL): + row, col = sc.slot(i) + _st( + b, + sc.vn_lds, + b.add(wrow, row), + cidx(jt, col), + value=b.cast_f32_to(b.vec_extract(acc_v[jt], i), ELEM), + n=1, + ) + b.sync_lds_only() + + # The next chunk's inputs land here: late enough that the loads issued at + # the top of this body have had the Z and V~ matmul groups to retire behind, + # and still ahead of the rendezvous that closes the body, which is what + # publishes them to every wave. + if prefetch is not None: + prefetch.commit(issued) + + # ---- O = GQ S + Aqk V~ ---------------------------------------------- + # This wave owns a band of output *columns* here (the same band of v + # channels it owns of the state), so GQ and Aqk are the A operands and the + # state mirror and V~ are the B operands -- which puts the chunk extent on + # the M side, so a chunk-extent tile picks the A operand's row. + for jt in range(NC_T): + acc_o = sc.gemm( + sc.gq_lds, sc.crow(jt), sc.stb_lds, srow, sc.KS_DK, atom.zero_acc(b) + ) + acc_o = sc.gemm(sc.aqb_lds, sc.crow(jt), sc.vn_lds, srow, sc.KS_C, acc_o) + # Straight to HBM, no staging tile: a slot's column index is the lane's + # position in the atom's N extent, so one store instruction covers a + # contiguous run of v channels per lane group and is already coalesced. + for i in range(CPL): + row, col = sc.slot(i) + b.global_store_vN( + sc.o_ptr, + sc.o_global(tile, cidx(jt, row), col, wrow), + b.cast_f32_to(b.vec_extract(acc_o, i), ELEM), + 1, + ) + + # ---- S^T <- Diag(dec) S^T + V~^T Kt^T ------------------------------- + # dec is the whole-chunk decay per k channel, i.e. the state's column + # index, so it is read once per accumulator slot. + new_state = [] + for ti in range(sc.NS): + scaled = [] + for i in range(CPL): + _, col = sc.slot(i) + d = b.vec_extract( + _ld( + b, + sc.dec_lds, + b.add(b.const_i32(ti * atom.n), col), + dtype=F32, + n=1, + ), + 0, + ) + if sc.dec_is_log: + d = sc.ex2(d) + scaled.append(b.fmul(b.vec_extract(state[ti], i), d)) + acc = sc.gemm( + sc.vn_lds, + b.add(wrow, lane_m), + sc.kt_lds, + b.add(b.const_i32(ti * atom.n), lane_m), + sc.KS_C, + b.vec_pack(scaled, F32), + ) + new_state.append(acc) + b.sync_lds_only() + return new_state + + +def build_kda_chunk_fused(spec: KdaChunkFusedSpec, arch: str = "gfx950") -> KernelDef: + """Build the IR for the fused chunkwise KDA forward. + + Kernel signature:: + + (q, k: ptr, # [NT, C * DK] NT = BH * NC + g: ptr, # [NT, C * DK] per-channel log decay + beta: ptr, # [NT, C] + v: ptr, # [NT, C * DV] + o: ptr, # [NT, C * DV] + h0: ptr, # [BH, DV, DK] S^T, read iff has_initial_state + ht: ptr, # [BH, DV, DK] S^T, written iff store_final_state + scale: f32, + nc: i32) # chunks per (batch, head) + + Grid ``(BH, 1, 1)``; block ``(block_size, 1, 1)``. + + Per chunk the body is the tile emission followed by + + .. code-block:: text + + Z^T = S^T GK^T EV x C + R^T = V^T - Z^T EV x C (in-register, see below) + V~^T = R^T A^T EV x C + O = GQ S + Aqk V~ C x EV + S^T <- Diag(dec) S^T + V~^T Kt^T EV x DK + + Working transposed is what keeps this cheap: every product is then + ``A B^T`` with the contraction on the fastest axis, so no operand ever + needs an LDS transpose. ``R^T`` never reaches LDS as fp32 either -- the + accumulator's lane mapping puts each lane's 16 slots at one chunk row and + four runs of four consecutive v channels, so ``V`` is subtracted straight + into the accumulator with four short vector loads per lane. + """ + ok, why = is_valid_fused_spec(spec, arch=arch) + if not ok: + raise ValueError(f"invalid kda_chunk_fused spec for {arch}: {why}") + + prep = spec.prep + t = spec.tile + C, DK, EV = t.chunk, spec.head_k, spec.head_v + BLOCK = t.block_size + ELEM = _DTYPE_IR[spec.dtype] + atom = spec.scan_atom + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = BLOCK + if t.waves_per_eu: + b.kernel.attrs["waves_per_eu"] = (t.waves_per_eu, t.waves_per_eu) + + q_ptr = b.param("q_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + k_ptr = b.param("k_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + g_ptr = b.param("g_ptr", PtrType(F32, "global"), readonly=True, align=16) + beta_ptr = b.param("beta_ptr", PtrType(F32, "global"), readonly=True, align=4) + v_ptr = b.param("v_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + o_ptr = b.param("o_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + h0_ptr = b.param("h0_ptr", PtrType(F32, "global"), readonly=True, align=16) + ht_ptr = b.param("ht_ptr", PtrType(F32, "global"), writeonly=True, align=16) + scale = b.param("scale", F32) + nc = b.param("nc", I32) + + ctx = _ChunkCtx( + b, + prep, + (q_ptr, k_ptr, g_ptr, beta_ptr, scale), + overlay=spec.overlay_lds, + ) + if spec.overlay_lds: + kt_lds, stb_lds, vn_lds = ctx.kt_lds, ctx.stb_lds, ctx.vn_lds + else: + kt_lds = b.smem_alloc(ELEM, [DK, C + t.pad_cb], "kt_s") + stb_lds = b.smem_alloc(ELEM, [EV, DK + t.pad_dk], "stb_s") + vn_lds = b.smem_alloc(ELEM, [EV, C + t.pad_cb], "vn_s") + sc = _ScanCtx( + b, + atom=atom, + chunk=C, + head_k=DK, + head_v=EV, + block_size=BLOCK, + elem=ELEM, + tiles=( + ctx.x_lds, # GK overwrites the X staging tile + ctx.xq_lds, # GQ overwrites the XQ staging tile + ctx.tb_lds, # A overwrites the solve's rank-update operand + ctx.zs_lds, # Aqk overwrites the solved-so-far tile + kt_lds, + ctx.gl_lds, # the whole-chunk decay, still in the log domain + ), + stb_lds=stb_lds, + vn_lds=vn_lds, + v_ptr=v_ptr, + o_ptr=o_ptr, + tid=ctx.tid, + lane=ctx.lane, + dec_is_log=True, + ex2=ctx.ex2, + # V is issued before the elementwise/MFMA section and parked in Y once + # its C x C consumers finish. Restrict this to the 512-thread schedule: + # the 256-thread path did not benefit from the extra live vectors. + v_lds=ctx.y_lds if (ctx.PDK >= EV and BLOCK >= 512) else None, + ) + sink = _LdsTileSink(*sc.tiles_for_sink()) + + bh = b.block_id_x() + state_base = b.mul(bh, b.const_i32(EV * DK)) + # No initial publish: the scan body mirrors whatever state it is handed. + s_init = ( + sc.load_state(h0_ptr, state_base) if spec.has_initial_state else sc.zero_state() + ) + + # With the prefetch on, every chunk finds its inputs already staged by its + # predecessor's scan, so chunk 0 is the one that has to stage its own and is + # peeled out here. This is the only staging barrier left in the kernel. + pf = spec.prefetch_inputs + if pf: + _emit_stage_inputs(ctx, _ChunkOffsets(ctx, b.mul(bh, nc))) + b.sync_lds_only() + + loop = b.scf_for_iter( + b.const_i32(0), + nc, + b.const_i32(1), + [(f"s{ti}", s_init[ti]) for ti in range(sc.NS)], + iv_name="chunk", + elide_trailing_barrier=False, + ) + with loop as (n, carried): + tile = b.add(b.mul(bh, nc), n) + overlap = BLOCK >= 512 + _emit_chunk_tiles( + ctx, + tile, + sink, + v_ptr=v_ptr if sc.v_lds is not None else None, + v_lds=sc.v_lds, + overlap_solve=overlap, + stage_inputs=not pf, + ) + # The scan does not consume tile outputs until after publish_state()'s + # rendezvous, so that later barrier can publish both the tile outputs + # and the state mirror. The explicit overlay is the exception: its + # mirror writes reuse tile-only addresses and must not race outstanding + # reads from the final A/Aqk copy. + if spec.overlay_lds: + b.sync_lds_only() + prefetch = None + if pf: + # The last chunk has no successor, so it re-stages itself rather + # than reading off the end of the inputs. Clamping beats a branch: + # the loads are redundant, not predicated, and this is one chunk in + # a sequence of NC. + nxt = b.add(n, b.const_i32(1)) + nxt = b.select(b.cmp_gt(nc, nxt), nxt, n) + prefetch = _InputPrefetch( + ctx, _ChunkOffsets(ctx, b.add(b.mul(bh, nc), nxt)) + ) + b.scf_yield(*_emit_scan_body(sc, list(carried), tile, prefetch=prefetch)) + + if spec.store_final_state: + sc.store_state(ht_ptr, state_base, loop.results) + + b.ret() + return b.kernel + + +def kda_chunk_fused_grid(spec: KdaChunkFusedSpec, bh: int) -> Tuple[int, int, int]: + """One workgroup per (batch, head).""" + return (int(bh), 1, 1) + + +def kda_chunk_fused_signature(spec: KdaChunkFusedSpec): + return ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("g_ptr", "f32") + .ptr("beta_ptr", "f32") + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .ptr("h0_ptr", "f32") + .ptr("ht_ptr", "f32") + .scalar("scale", "f32") + .scalar("nc", "i32") + .build() + ) + + +@dataclass(frozen=True) +class KdaChunkScanSpec: + """Compile-time spec for the standalone state scan over materialized tiles. + + The second half of the split path: :func:`build_kda_chunk_prep` writes the + six per-chunk tiles to HBM, and this kernel walks one (batch, head)'s chunks + in order, staging each chunk's tiles into LDS and running the same + recurrence the fused kernel runs in registers. + + Fusing the two would remove that HBM round trip entirely, which sounds + strictly better and is not: holding the tile builder's staging tiles and the + scan's operands live at once puts the workgroup over half the LDS budget, so + only one fits per CU. The scan is a latency-bound chain of small matmuls, and + at one workgroup per CU there is no second workgroup to cover it. Paying for + the tile traffic to keep two resident is the cheaper side of that trade, + which is why ``lds_bytes`` here is checked against *half* the budget rather + than all of it -- occupancy is the spec's binding constraint, not capacity. + """ + + head_k: int = 128 + head_v: int = 128 + dtype: str = "bf16" + tile: KdaTileSpec = KdaTileSpec() + has_initial_state: bool = False + store_final_state: bool = True + # Workgroups resident per CU that the LDS request must leave room for. + min_occupancy: int = 2 + # Split the v extent across ``value_splits`` independent workgroups per + # (batch, head), mirroring FlashKDA K2's ``ceil(V/BW)`` grid dimension. + value_splits: int = 1 + # Read/write token-major [B,T,H,D] tensors instead of chunk-packed views. + token_major_io: bool = False + name: str = "rocke_kda_chunk_scan" + + @property + def prep(self) -> KdaChunkPrepSpec: + """The tile builder whose output this scan stages. + + Derived rather than passed alongside, because the scan's staging + copies assume exactly the layouts that builder's global sink writes: + two specs that could drift apart would drift into a silent misread. + """ + return KdaChunkPrepSpec( + head_k=self.head_k, + head_v=self.head_v, + dtype=self.dtype, + tile=self.tile, + ) + + @property + def atom(self) -> MfmaAtom: + return self.prep.atom + + @property + def scan_atom(self) -> MfmaAtom: + """The atom the state scan runs on; see ``KdaTileSpec.scan_atom_m``.""" + return _scan_atom(self.tile, self.atom) + + @property + def state_tiles(self) -> int: + """Atom tiles per wave across the state's ``DK`` extent.""" + return self.head_k // self.scan_atom.n + + def lds_bytes(self) -> int: + """The six staged tiles plus the state mirror and ``V~``. + + Every tile is staged in the exact layout its consumer wants as an MFMA + operand, so staging is a straight copy and the scan body is identical to + the fused one. Unlike the fused kernel there is nothing to overlap them + with, so each is its own allocation -- which is still the smaller + footprint, because none of the tile builder's staging tiles exist here. + """ + t = self.tile + C, DK, EV = t.chunk, self.head_k, self.head_v + ev = EV // self.value_splits + PDK, PCB = DK + t.pad_dk, C + t.pad_cb + return ( + 2 * C * PDK # gk_s bf16 (C x DK) + + 2 * C * PDK # gq_s bf16 (C x DK) + + 2 * C * PCB # a_s bf16 (C x C) + + 2 * C * PCB # aqk_s bf16 (C x C) + + 2 * DK * PCB # kt_s bf16 (DK x C) + + 2 * ev * PDK # stb_s bf16 mirror of S^T + + 2 * ev * PCB # vn_s bf16 (EV x C) + + 4 * DK # dec_s fp32 (DK) + ) + + def kernel_name(self) -> str: + parts = (f"dk{self.head_k}", f"dv{self.head_v}", self.dtype) + parts += self.tile.name_parts() + if self.has_initial_state: + parts += ("h0",) + if not self.store_final_state: + parts += ("noht",) + if self.value_splits != 1: + parts += (f"vs{self.value_splits}",) + if self.token_major_io: + parts += ("tm",) + return kernel_name_join(self.name, *parts) + + +def is_valid_scan_spec( + spec: KdaChunkScanSpec, arch: str = "gfx950" +) -> Tuple[bool, str]: + """Return ``(ok, reason)`` for a standalone scan spec on ``arch``.""" + if arch != "gfx950": + return False, f"unsupported arch {arch}" + if spec.dtype not in _DTYPE_IR: + return False, f"unsupported dtype {spec.dtype}" + + if spec.value_splits not in _RAW_VALUE_SPLITS: + return False, ( + f"value_splits must be one of {_RAW_VALUE_SPLITS}, " + f"got {spec.value_splits}" + ) + if spec.head_v % spec.value_splits: + return False, ( + f"head_v ({spec.head_v}) must divide evenly by value_splits " + f"({spec.value_splits})" + ) + ev_slice = spec.head_v // spec.value_splits + if ev_slice % spec.scan_atom.n: + return False, ( + f"v slice ({ev_slice}) must divide scan atom n ({spec.scan_atom.n})" + ) + num_waves = spec.tile.num_waves + if ev_slice != num_waves * spec.scan_atom.m: + return False, ( + f"v slice ({ev_slice}) must equal waves ({num_waves}) * scan atom m " + f"({spec.scan_atom.m}); adjust block_size, scan_atom_m, or value_splits" + ) + + t = spec.tile + ok, why = _check_scan_partition(spec.scan_atom, t, spec.head_k, ev_slice) + if not ok: + return False, why + # Staging is 128-bit per thread throughout, so both padded row pitches have + # to keep 8-element alignment and each tile has to divide evenly across the + # workgroup's 8-element slots. + if (spec.head_k + t.pad_dk) % 8: + return False, ( + f"padded head_k pitch ({spec.head_k + t.pad_dk}) must be a multiple " + "of 8 elements; the staging copies are ds_write_b128" + ) + if (t.chunk + t.pad_cb) % 8: + return False, ( + f"padded chunk pitch ({t.chunk + t.pad_cb}) must be a multiple of 8 " + "elements; the staging copies are ds_write_b128" + ) + for rows, cols in ( + (t.chunk, spec.head_k), + (t.chunk, t.chunk), + (spec.head_k, t.chunk), + ): + n_slot = rows * cols // 8 + if n_slot > t.block_size and n_slot % t.block_size: + return False, ( + f"tile {rows}x{cols} needs {n_slot} slots, not a multiple of " + f"{t.block_size}" + ) + for what, n in ( + ("C x DK", t.chunk * spec.head_k), + ("C x C", t.chunk * t.chunk), + ("DK x C", spec.head_k * t.chunk), + ): + if n % 8: + return False, ( + f"{what} tile ({n} elements) must be a multiple of 8; staging " + "moves one 128-bit slot per thread" + ) + if spec.head_k % 4: + return False, ( + f"dec tile ({spec.head_k}) must be a multiple of 4; it is staged as " + "one fp32 4-vector per thread" + ) + if spec.head_k > 4 * t.block_size: + return False, ( + f"dec tile ({spec.head_k}) exceeds one guarded pass of " + f"{4 * t.block_size} elements for block_size ({t.block_size})" + ) + + lds = spec.lds_bytes() + budget = LDS_LIMIT // spec.min_occupancy + if lds > budget: + return False, ( + f"LDS request {lds} B exceeds the {budget} B budget for " + f"{spec.min_occupancy} workgroups per CU; the split path only pays " + "for its tile traffic if it keeps that occupancy" + ) + return True, "ok" + + +def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "KernelDef": + """Serial state scan over per-chunk tiles already materialized in HBM. + + One workgroup per (batch, head) walks that head's chunks in order. Per + chunk it stages the six tiles from HBM into LDS and runs + :func:`_emit_scan_body`, so the recurrence is bit-for-bit the fused + kernel's -- the only difference is where the operands came from. + + The state stays in accumulators across the whole loop, so the sequence + length costs nothing in registers, and ``dec`` arrives already + exponentiated (the tile builder stored it that way), which is the one place + this body diverges from the fused one. + """ + ok, why = is_valid_scan_spec(spec, arch=arch) + if not ok: + raise ValueError(f"invalid kda_chunk_scan spec for {arch}: {why}") + + t = spec.tile + C, DK, EV = t.chunk, spec.head_k, spec.head_v + ev_slice = EV // spec.value_splits + BLOCK = t.block_size + PDK, PCB = DK + t.pad_dk, C + t.pad_cb + ELEM = _DTYPE_IR[spec.dtype] + atom = spec.scan_atom + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = BLOCK + if t.waves_per_eu: + b.kernel.attrs["waves_per_eu"] = (t.waves_per_eu, t.waves_per_eu) + + a_ptr = b.param("a_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + gk_ptr = b.param("gk_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + gq_ptr = b.param("gq_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + aqk_ptr = b.param("aqk_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + kt_ptr = b.param("kt_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + dec_ptr = b.param("dec_ptr", PtrType(F32, "global"), readonly=True, align=16) + v_ptr = b.param("v_ptr", PtrType(ELEM, "global"), readonly=True, align=16) + o_ptr = b.param("o_ptr", PtrType(ELEM, "global"), writeonly=True, align=16) + h0_ptr = b.param("h0_ptr", PtrType(F32, "global"), readonly=True, align=16) + ht_ptr = b.param("ht_ptr", PtrType(F32, "global"), writeonly=True, align=16) + nc = b.param("nc", I32) + + io = None + if spec.token_major_io: + batch = b.param("batch", I32) + heads = b.param("heads", I32) + tseq = b.param("tseq", I32) + io = _RawTokenAddr( + b, + heads=heads, + tseq=tseq, + nc=nc, + chunk=C, + dk=DK, + ) + + gk_lds = b.smem_alloc(ELEM, [C, PDK], "gk_s") + gq_lds = b.smem_alloc(ELEM, [C, PDK], "gq_s") + ab_lds = b.smem_alloc(ELEM, [C, PCB], "a_s") + aqb_lds = b.smem_alloc(ELEM, [C, PCB], "aqk_s") + kt_lds = b.smem_alloc(ELEM, [DK, PCB], "kt_s") + dec_lds = b.smem_alloc(F32, [DK], "dec_s") + stb_lds = b.smem_alloc(ELEM, [ev_slice, PDK], "stb_s") + vn_lds = b.smem_alloc(ELEM, [ev_slice, PCB], "vn_s") + + tid = b.thread_id_x() + lane = b.mod(tid, b.const_i32(64)) + + wg = b.block_id_x() + if spec.value_splits == 1: + bh = wg + v_row_base = b.const_i32(0) + else: + bh = b.div(wg, b.const_i32(spec.value_splits)) + v_split = b.mod(wg, b.const_i32(spec.value_splits)) + v_row_base = b.mul(v_split, b.const_i32(ev_slice)) + + sc = _ScanCtx( + b, + atom=atom, + chunk=C, + head_k=DK, + head_v=ev_slice, + block_size=BLOCK, + elem=ELEM, + tiles=(gk_lds, gq_lds, ab_lds, aqb_lds, kt_lds, dec_lds), + stb_lds=stb_lds, + vn_lds=vn_lds, + v_ptr=v_ptr, + o_ptr=o_ptr, + tid=tid, + lane=lane, + dec_is_log=False, + v_row_base=v_row_base, + ev_stride=EV, + io=io, + ) + + state_base = b.add( + b.mul(bh, b.const_i32(EV * DK)), + b.mul(v_row_base, b.const_i32(DK)), + ) + # No initial publish: the scan body mirrors whatever state it is handed. + s_init = ( + sc.load_state(h0_ptr, state_base) if spec.has_initial_state else sc.zero_state() + ) + + def stage(src, dst, rows, cols, base): + """One flat ``rows x cols`` HBM tile into its padded LDS tile. + + Both sides are 128-bit: the source row length is a multiple of 8, so a + thread's eight consecutive elements never straddle a row and the only + difference between the two addresses is the destination's pad. A tile + smaller than one workgroup sweep (the ``C x C`` pair, at half) just + leaves the upper threads idle rather than giving them a second, narrower + access pattern. + """ + n_slot = rows * cols // 8 + for i in range(max(1, n_slot // BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + guard = ( + nullcontext() + if n_slot >= BLOCK + else b.scf_if(b.cmp_gt(b.const_i32(n_slot), vidx)) + ) + with guard: + off = b.mul(vidx, b.const_i32(8)) + b.smem_store_vN( + dst, + [b.div(off, b.const_i32(cols)), b.mod(off, b.const_i32(cols))], + b.global_load_vN(src, b.add(base, off), ELEM, 8), + 8, + ) + + loop = b.scf_for_iter( + b.const_i32(0), + nc, + b.const_i32(1), + [(f"s{ti}", s_init[ti]) for ti in range(sc.NS)], + iv_name="chunk", + elide_trailing_barrier=False, + ) + with loop as (n, carried): + tile = b.add(b.mul(bh, nc), n) + cd = b.mul(tile, b.const_i32(C * DK)) + cc = b.mul(tile, b.const_i32(C * C)) + stage(gk_ptr, gk_lds, C, DK, cd) + stage(gq_ptr, gq_lds, C, DK, cd) + stage(kt_ptr, kt_lds, DK, C, cd) + stage(a_ptr, ab_lds, C, C, cc) + stage(aqk_ptr, aqb_lds, C, C, cc) + with b.scf_if(b.cmp_gt(b.const_i32(DK // 4), tid)): + col4 = b.mul(tid, b.const_i32(4)) + b.smem_store_vN( + dec_lds, + [col4], + b.global_load_vN( + dec_ptr, b.add(b.mul(tile, b.const_i32(DK)), col4), F32, 4 + ), + 4, + ) + b.sync_lds_only() + b.scf_yield(*_emit_scan_body(sc, list(carried), tile)) + + if spec.store_final_state: + sc.store_state(ht_ptr, state_base, loop.results) + + b.ret() + return b.kernel + + +def kda_chunk_scan_grid(spec: KdaChunkScanSpec, bh: int) -> Tuple[int, int, int]: + """One workgroup per (batch, head), times ``value_splits`` v bands.""" + return (int(bh) * int(spec.value_splits), 1, 1) + + +def kda_chunk_scan_signature(spec: KdaChunkScanSpec): + sb = ( + SignatureBuilder() + .ptr("a_ptr", spec.dtype) + .ptr("gk_ptr", spec.dtype) + .ptr("gq_ptr", spec.dtype) + .ptr("aqk_ptr", spec.dtype) + .ptr("kt_ptr", spec.dtype) + .ptr("dec_ptr", "f32") + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .ptr("h0_ptr", "f32") + .ptr("ht_ptr", "f32") + .scalar("nc", "i32") + ) + if spec.token_major_io: + sb = sb.scalar("batch", "i32").scalar("heads", "i32").scalar("tseq", "i32") + return sb.build() + + +def kda_chunk_prep_grid(spec: KdaChunkPrepSpec, num_tiles: int) -> Tuple[int, int, int]: + """One workgroup per chunk. ``num_tiles = BH * NC``.""" + return (int(num_tiles), 1, 1) + + +def kda_chunk_prep_signature(spec: KdaChunkPrepSpec): + sb = ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("g_ptr", "bf16" if spec.raw_inputs else "f32") + .ptr("beta_ptr", "f32") + .ptr("a_ptr", spec.dtype) + .ptr("gk_ptr", spec.dtype) + .ptr("gq_ptr", spec.dtype) + .ptr("aqk_ptr", spec.dtype) + .ptr("kt_ptr", spec.dtype) + .ptr("dec_ptr", "f32") + .scalar("scale", "f32") + ) + if spec.raw_inputs: + sb = ( + sb.ptr("a_log_ptr", "f32") + .ptr("dt_bias_ptr", "f32") + .scalar("batch", "i32") + .scalar("heads", "i32") + .scalar("tseq", "i32") + .scalar("nc", "i32") + ) + return sb.build() + + +__all__ = [ + "EXP2_CLAMP", + "KDA_CHUNK_SIZES", + "KDA_DTYPES", + "KdaChunkFusedSpec", + "KdaChunkPrepSpec", + "KdaChunkScanSpec", + "KdaTileSpec", + "LOG2E", + "build_kda_chunk_fused", + "build_kda_chunk_prep", + "build_kda_chunk_scan", + "is_valid_fused_spec", + "is_valid_scan_spec", + "is_valid_spec", + "kda_chunk_fused_grid", + "kda_chunk_fused_signature", + "kda_chunk_prep_grid", + "kda_chunk_prep_signature", + "kda_chunk_scan_grid", + "kda_chunk_scan_signature", +] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/__init__.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/__init__.py index 1df485718439..fd6725c0f729 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/__init__.py @@ -1,2 +1,3 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT +"""KDA dispatch tests.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx942_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx942_wiring.py index 5335176157db..2fc2755af7a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx942_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx942_wiring.py @@ -50,7 +50,7 @@ def _candidate(name: str): class TestRegistration(unittest.TestCase): def test_every_kernel_is_registered(self): - self.assertEqual({c.name for c in kda_candidates()}, set(_ALL)) + self.assertLessEqual(set(_ALL), {c.name for c in kda_candidates()}) def test_identity(self): for name, algorithm, spec_id in ( @@ -84,6 +84,9 @@ class TestRouting(unittest.TestCase): def test_default_routing_is_the_fused_kernel(self): self.assertEqual(dispatch_kda(_req()).candidate.name, _FUSED) + def test_unspecified_chunk_uses_the_gfx942_default(self): + self.assertEqual(dispatch_kda(_req(chunk_size=None)).spec.tile.chunk, 16) + def test_named_algorithm_selects_a_split_half(self): for name, algorithm in ((_PREP, "chunk_prep"), (_SCAN, "chunk_scan")): with self.subTest(candidate=name): @@ -121,7 +124,7 @@ def test_no_candidate_admits_an_arch_it_did_not_declare(self): def test_registry_serves_them_only_to_gfx942(self): self.assertEqual({c.name for c in KDA_REGISTRY.for_arch("gfx942")}, set(_ALL)) - self.assertEqual(KDA_REGISTRY.for_arch("gfx950"), ()) + self.assertFalse(set(_ALL) & {c.name for c in KDA_REGISTRY.for_arch("gfx950")}) class TestCapabilityGates(unittest.TestCase): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx950_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx950_wiring.py new file mode 100644 index 000000000000..f997f9990382 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx950_wiring.py @@ -0,0 +1,285 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Wiring for the gfx950 chunkwise KDA candidates. + +The kernels had specs, validators, builders, grid helpers and signatures long +before anything referenced them from dispatch. These tests pin what registering +them did -- and, just as much, what it deliberately did not do: the split path +is reachable only when the caller names it, because the fused/split crossover +has not been measured at enough shapes to encode a threshold. + +CPU-only. Numeric coverage lives in ``test_kda_chunkwise_gfx950_numeric.py``. +""" + +from __future__ import annotations + +import unittest + +from dispatch.kda import ( + KDA_REGISTRY, + KdaRequest, + dispatch_kda, + kda_candidates, + kda_sweep_space, +) + +_FUSED = "kda_gfx950_chunk_fused" +_PREP = "kda_gfx950_chunk_prep" +_SCAN = "kda_gfx950_chunk_scan" +_ALL = (_FUSED, _PREP, _SCAN) + +_BATCH, _HEADS, _SEQLEN, _CHUNK = 2, 8, 1024, 32 + + +def _req(**kw) -> KdaRequest: + base = dict( + batch=_BATCH, + num_heads=_HEADS, + seqlen=_SEQLEN, + arch="gfx950", + head_k=128, + head_v=128, + chunk_size=_CHUNK, + ) + base.update(kw) + return KdaRequest(**base) + + +def _candidate(name: str): + return KDA_REGISTRY.get(name) + + +class TestRegistration(unittest.TestCase): + def test_every_kernel_is_registered(self): + self.assertLessEqual(set(_ALL), {c.name for c in kda_candidates()}) + + def test_identity(self): + for name, algorithm, spec_id in ( + (_FUSED, "chunk_fused", "gfx950_chunk_fused"), + (_PREP, "chunk_prep", "gfx950_chunk_prep"), + (_SCAN, "chunk_scan", "gfx950_chunk_scan"), + ): + with self.subTest(candidate=name): + c = _candidate(name) + self.assertEqual(c.algorithm, algorithm) + self.assertEqual(c.spec_id, spec_id) + self.assertEqual(c.capability.arches, ("gfx950",)) + + def test_every_candidate_declares_a_build(self): + # The whole point of the registration: the three builders were + # unreachable from dispatch before these candidates existed. + for name in _ALL: + with self.subTest(candidate=name): + self.assertIsNotNone(_candidate(name).build) + + def test_the_family_refuses_an_unbuildable_candidate(self): + self.assertTrue(KDA_REGISTRY.require_build) + + def test_coverage_is_queryable_without_a_request(self): + coverage = KDA_REGISTRY.coverage() + self.assertEqual(coverage["family"], "kda_chunkwise") + self.assertTrue(all(c["buildable"] for c in coverage["candidates"])) + + +class TestRouting(unittest.TestCase): + def test_default_routing_is_the_fused_kernel(self): + self.assertEqual(dispatch_kda(_req()).candidate.name, _FUSED) + + def test_unspecified_chunk_uses_the_gfx950_default(self): + self.assertEqual(dispatch_kda(_req(chunk_size=None)).spec.tile.chunk, 32) + + def test_named_algorithm_selects_a_split_half(self): + for name, algorithm in ((_PREP, "chunk_prep"), (_SCAN, "chunk_scan")): + with self.subTest(candidate=name): + result = dispatch_kda(_req(algorithm=algorithm)) + self.assertEqual(result.candidate.name, name) + + def test_named_spec_id_selects_a_split_half(self): + result = dispatch_kda(_req(spec_id="gfx950_chunk_scan")) + self.assertEqual(result.candidate.name, _SCAN) + + def test_the_refusal_explains_the_opt_in(self): + for name in (_PREP, _SCAN): + with self.subTest(candidate=name): + ok, why = _candidate(name).admits(_req()) + self.assertFalse(ok) + self.assertIn("opt-in", why) + + def test_sweep_space_offers_only_what_is_reachable(self): + # auto routes to fused alone, so the sweep space is one spec -- not the + # three it would be if the opt-in gate were advisory. + self.assertEqual(len(kda_sweep_space(_req())), 1) + + +class TestArchGate(unittest.TestCase): + def test_no_candidate_admits_an_arch_it_did_not_declare(self): + from rocke.core.arch import known_arches + + for name in _ALL: + for arch in known_arches(): + if arch == "gfx950": + continue + with self.subTest(candidate=name, arch=arch): + req = _req(arch=arch, algorithm=_candidate(name).algorithm) + self.assertFalse(_candidate(name).admits(req)[0]) + + def test_registry_serves_them_only_to_gfx950(self): + self.assertEqual({c.name for c in KDA_REGISTRY.for_arch("gfx950")}, set(_ALL)) + self.assertFalse(set(_ALL) & {c.name for c in KDA_REGISTRY.for_arch("gfx942")}) + + +class TestCapabilityGates(unittest.TestCase): + """Data gates, checked before any spec is constructed.""" + + def test_fp16_is_rejected(self): + ok, why = _candidate(_FUSED).admits(_req(dtype="fp16")) + self.assertFalse(ok) + self.assertIn("fp16", why) + + def test_a_seqlen_that_does_not_tile_the_chunk_is_rejected(self): + # No varlen path here, so a ragged length must be padded by the caller + # rather than silently truncated to whole chunks. + ok, why = _candidate(_FUSED).admits(_req(seqlen=1000)) + self.assertFalse(ok) + self.assertIn("num_chunks", why) + + def test_an_unsupported_chunk_length_is_rejected(self): + self.assertFalse(_candidate(_FUSED).admits(_req(chunk_size=64))[0]) + + def test_the_split_path_can_serve_a_problem_that_wants_a_final_state(self): + # Regression: withholding the state features from the tile builder made + # the whole split path unreachable for the default request, since + # store_final_state defaults on. The flags describe the problem; both + # halves take part in serving it. + for name in (_PREP, _SCAN): + with self.subTest(candidate=name): + req = _req( + algorithm=_candidate(name).algorithm, + has_initial_state=True, + store_final_state=True, + ) + self.assertTrue(_candidate(name).admits(req)[0]) + + def test_the_state_flags_do_not_reach_the_tile_builder(self): + # KdaChunkPrepSpec has no state fields: the tile phase is + # state-independent and the scan half is what applies them. + spec = dispatch_kda(_req(algorithm="chunk_prep", has_initial_state=True)).spec + self.assertFalse(hasattr(spec, "has_initial_state")) + self.assertTrue( + dispatch_kda( + _req(algorithm="chunk_scan", has_initial_state=True) + ).spec.has_initial_state + ) + + +class TestResidualGates(unittest.TestCase): + """Gates the validators compute from the spec, not from the request.""" + + def test_c16_is_admitted_by_capability_then_split_on_the_schedule(self): + # chunk=16 is a declared chunk length, so the prefilter passes it for + # every candidate. What separates them is the tile schedule each one + # starts from: the fused kernel's narrower 16x16 atom tiles a 16-token + # chunk, the split path's chunk-wide 32x32 atom does not. Pinning both + # halves of that split is the point -- a capability that also encoded + # the atom would have to recompute the spec to do it. + for name in _ALL: + with self.subTest(candidate=name): + req = _req(chunk_size=16, algorithm=_candidate(name).algorithm) + self.assertTrue(_candidate(name).capability.check(req)[0]) + + self.assertTrue(_candidate(_FUSED).admits(_req(chunk_size=16))[0]) + for name in (_PREP, _SCAN): + with self.subTest(candidate=name): + req = _req(chunk_size=16, algorithm=_candidate(name).algorithm) + ok, why = _candidate(name).admits(req) + self.assertFalse(ok) + self.assertIn("16", why) + + def test_a_scan_whose_tile_builder_is_unbuildable_is_refused(self): + # head_k=32 passes every rule the scan itself has, and fails the tile + # builder's cumsum row-group fold. Admitting the scan alone would leave + # a kernel nothing can feed. + req = _req(head_k=32, algorithm="chunk_scan") + self.assertTrue(_candidate(_SCAN).capability.check(req)[0]) + ok, why = _candidate(_SCAN).admits(req) + self.assertFalse(ok) + self.assertIn("tile builder", why) + + +class TestGeometry(unittest.TestCase): + def test_fused_grid_is_one_workgroup_per_head(self): + # gfx950 holds a whole DV=128 state mirror, so a logical head is not + # partitioned the way it is on gfx942. + result = dispatch_kda(_req()) + self.assertEqual(result.grid, (_BATCH * _HEADS, 1, 1)) + self.assertEqual(result.block, (512, 1, 1)) + + def test_prep_grid_is_one_workgroup_per_chunk(self): + result = dispatch_kda(_req(algorithm="chunk_prep")) + self.assertEqual(result.grid, (_BATCH * _HEADS * (_SEQLEN // _CHUNK), 1, 1)) + self.assertEqual(result.block, (256, 1, 1)) + + def test_scan_grid_matches_the_fused_grid(self): + fused = dispatch_kda(_req()).grid + scan = dispatch_kda(_req(algorithm="chunk_scan")).grid + self.assertEqual(fused, scan) + + def test_the_spec_carries_the_requested_head_widths(self): + spec = dispatch_kda(_req()).spec + self.assertEqual(spec.head_k, 128) + self.assertEqual(spec.head_v, 128) + + def test_each_path_keeps_its_own_tile_schedule(self): + # The fused kernel's tuned schedule must not leak into the split scan, + # whose LDS budget is capped for two workgroups per CU. + fused = dispatch_kda(_req()).spec.tile + scan = dispatch_kda(_req(algorithm="chunk_scan")).spec.tile + self.assertEqual((fused.block_size, fused.scan_atom_m), (512, 16)) + self.assertEqual((scan.block_size, scan.scan_atom_m), (256, 0)) + self.assertEqual((fused.chunk, scan.chunk), (_CHUNK, _CHUNK)) + + +class TestSignatures(unittest.TestCase): + def test_each_signature_matches_its_kernel_abi(self): + from kernels.gfx950.kda_chunkwise import ( + kda_chunk_fused_signature, + kda_chunk_prep_signature, + kda_chunk_scan_signature, + ) + + for algorithm, signature_fn in ( + ("chunk_fused", kda_chunk_fused_signature), + ("chunk_prep", kda_chunk_prep_signature), + ("chunk_scan", kda_chunk_scan_signature), + ): + with self.subTest(algorithm=algorithm): + result = dispatch_kda(_req(algorithm=algorithm)) + expected = [a["name"] for a in signature_fn(result.spec)] + self.assertEqual([a["name"] for a in result.signature], expected) + + +class TestBuild(unittest.TestCase): + def test_it_builds_what_it_selects(self): + for algorithm, needle in ( + ("chunk_fused", "fused"), + ("chunk_prep", "prep"), + ("chunk_scan", "scan"), + ): + with self.subTest(algorithm=algorithm): + kernel = dispatch_kda(_req(algorithm=algorithm)).build() + self.assertIn(needle, kernel.name) + self.assertIn("dk128", kernel.name) + self.assertIn("dv128", kernel.name) + + def test_the_built_kernel_takes_the_declared_signature(self): + result = dispatch_kda(_req()) + names = [p.name for p in result.build().params] + self.assertEqual(names, [a["name"] for a in result.signature]) + + def test_initial_state_reaches_the_kernel_name(self): + kernel = dispatch_kda(_req(has_initial_state=True)).build() + self.assertIn("h0", kernel.name) + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/kda_gfx950_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/kda_gfx950_ir_sha256.json new file mode 100644 index 000000000000..7b831b552e31 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/kda_gfx950_ir_sha256.json @@ -0,0 +1,233 @@ +{ + "flavors": { + "llvm20": { + "cases": { + "kda_gfx950/fused_c16_sa16": { + "bytes": 203100, + "sha256": "d546f54dd172beebb60b35a388bc9a15ccfc2c045c893de41ede34066c269270" + }, + "kda_gfx950/fused_c32_h0": { + "bytes": 293249, + "sha256": "355f5bae9494a13d171b2d590b7f2d26f7601d57ab6c75c68312dfdc13eeb7a1" + }, + "kda_gfx950/fused_c32_h0_noht": { + "bytes": 276626, + "sha256": "46abdeed2bae56c48472d90fcf5478029d544c42b53fdcdd2315fc8bc1fafdf3" + }, + "kda_gfx950/fused_c32_noht": { + "bytes": 258841, + "sha256": "7110d3877ea0e85c3945844a35b2c0362ef721774cbdc61c1666d5cdf7bdafc0" + }, + "kda_gfx950/fused_c32_overlay": { + "bytes": 527767, + "sha256": "991b6a4071c41e09082d9e6ab3afe05bf3306afb0cf21df0e2c00746d6567b0d" + }, + "kda_gfx950/fused_c32_sa16": { + "bytes": 275408, + "sha256": "92a389f0a982b5874e5428cef873fd4754fd6ff4f7cc16a9112f00ee1c8830a8" + }, + "kda_gfx950/fused_c32_sa32": { + "bytes": 413168, + "sha256": "25d1221f03ef379221180d0bd81e09576d1bcbf305b1c434eb3104453d962284" + }, + "kda_gfx950/split_c16_prep": { + "bytes": 107246, + "sha256": "f516b233bbbf4ba3333f69e6fde158a045085c37028961c54d046b2736ded04e" + }, + "kda_gfx950/split_c16_scan_sa16": { + "bytes": 83793, + "sha256": "3416fee700952edc54c17e6a709df1096f0df2110e3e4f46e47a5ebef2f611d5" + }, + "kda_gfx950/split_c32_prep": { + "bytes": 228596, + "sha256": "9d9a75f52ef0c75b10d9dab706dbcf0e062ae1f913f38feae7180548ede3fd62" + }, + "kda_gfx950/split_c32_prep_raw_fusions": { + "bytes": 268306, + "sha256": "06770cbb78b8ebf6d3c6778f232a642eb59c72a211a7b6c4257e656993dce067" + }, + "kda_gfx950/split_c32_scan_h0": { + "bytes": 211995, + "sha256": "49bb7f8e3043d4fd4d44e4a5ed4024b7187530e610a267b14bcc5c453d6b743b" + }, + "kda_gfx950/split_c32_scan_h0_noht": { + "bytes": 175211, + "sha256": "db30504ff7f8a24f5993be118964f0129bf60a71eb284f50ef9dc706ef1b99be" + }, + "kda_gfx950/split_c32_scan_noht": { + "bytes": 137303, + "sha256": "eb26bc293efa65b1741ab36400db80f6f905d617bd5cbef2ca349a8e8ea5cc0b" + }, + "kda_gfx950/split_c32_scan_sa32": { + "bytes": 173959, + "sha256": "c789c33410b0e19a458c4fff0742340465ef84ae5b8f8c810dbfd8443056d3a9" + }, + "kda_gfx950/split_c32_scan_vs2": { + "bytes": 185398, + "sha256": "6cdd3470a3ca827cfa825ecb90c02439ce2f864216404b1ced86b14bb8206d22" + }, + "kda_gfx950/split_c32_scan_vs4": { + "bytes": 192986, + "sha256": "bc8c2ff78e5ae80bbb0d9b2ee8db2e0f9bbb05f22d4744455d72b1a4d4ce974e" + }, + "kda_gfx950/split_c32_scan_vs8": { + "bytes": 112923, + "sha256": "68e555bb083fe21e7e21145af5b30fb18743412b599716fdb7bbc7756f057f6a" + } + } + }, + "llvm22": { + "cases": { + "kda_gfx950/fused_c16_sa16": { + "bytes": 203107, + "sha256": "adc548a3b17baa6fbad52d8d1bdaab80b26386fac5002794f029ed2d7f0de18b" + }, + "kda_gfx950/fused_c32_h0": { + "bytes": 293256, + "sha256": "abca427c72836be12edd23a9fba7380ab9a10b31c90a3960b175fb1ec73093bc" + }, + "kda_gfx950/fused_c32_h0_noht": { + "bytes": 276633, + "sha256": "c63a5841abdb9df9f2039365106ed495155d8da85e3b220a7167e702cf76ff7b" + }, + "kda_gfx950/fused_c32_noht": { + "bytes": 258848, + "sha256": "efb4b9c7aa7d772cf9581db3401497c604a0358d6cf6d99356439566cf213a0e" + }, + "kda_gfx950/fused_c32_overlay": { + "bytes": 527774, + "sha256": "d4f0b0fc2cc416d895918bc015c1437bd9a71579a8788b42509d6a9d0c3af335" + }, + "kda_gfx950/fused_c32_sa16": { + "bytes": 275415, + "sha256": "05aeab52f2d3c076d09ef025965f04ae8d393fcff6233a67037a09c4fc34041f" + }, + "kda_gfx950/fused_c32_sa32": { + "bytes": 413175, + "sha256": "dc042a809a7f423cd91fd664355f179e3e0814fdc4ab87741f5a50c95abbad59" + }, + "kda_gfx950/split_c16_prep": { + "bytes": 107253, + "sha256": "2e3d21a32aa62be07e62e290e1a9cc7bd40cc10efb2dc433239d4c61cba720cc" + }, + "kda_gfx950/split_c16_scan_sa16": { + "bytes": 83800, + "sha256": "cf2e5524b1cdf8b00679a6d585e19797893ffbfac05becd528e5f963936d4c65" + }, + "kda_gfx950/split_c32_prep": { + "bytes": 228603, + "sha256": "0abf19d288d02ed05b997ac2ebfcd100f3f05c0aecb74b55f7670ab8c006ffd7" + }, + "kda_gfx950/split_c32_prep_raw_fusions": { + "bytes": 268313, + "sha256": "556c14b34332821a5cabfedf0a6922dafef83875cf21737286172eda5a7aa27b" + }, + "kda_gfx950/split_c32_scan_h0": { + "bytes": 212002, + "sha256": "88ce5e5959ef1a90b50eeb6070507697c0651ab79c7899c32980ddd35edbf064" + }, + "kda_gfx950/split_c32_scan_h0_noht": { + "bytes": 175218, + "sha256": "f7601e73a1b4038e02f394588e8add93dd8d4722c63756e0855d9480562a17a4" + }, + "kda_gfx950/split_c32_scan_noht": { + "bytes": 137310, + "sha256": "295699374721b0dd374a1a62738e9d2e8e834bf76f77f6e4f03c856d50ce9c55" + }, + "kda_gfx950/split_c32_scan_sa32": { + "bytes": 173966, + "sha256": "6141220badadc913363758769b31c38240b7138f1beb0585d27472b8e29e7fbb" + }, + "kda_gfx950/split_c32_scan_vs2": { + "bytes": 185405, + "sha256": "4817504e04899dc548298dde55bfba6a52bc766ea65c6bd0406ecd936194d66f" + }, + "kda_gfx950/split_c32_scan_vs4": { + "bytes": 192993, + "sha256": "1fc20a32293f07b3c3bf0f8510bbb963b82ef7cb0ef0502e582851c0f25f9d09" + }, + "kda_gfx950/split_c32_scan_vs8": { + "bytes": 112930, + "sha256": "d83c10ee4f4c74080574210073ae14eb46df0600d73b8978b039d8180345f4de" + } + } + }, + "llvm23": { + "cases": { + "kda_gfx950/fused_c16_sa16": { + "bytes": 203111, + "sha256": "b89e565e36085b37f9cc952cabae6d1e4de5aaf90c2ba6c548ee844a963fb6ea" + }, + "kda_gfx950/fused_c32_h0": { + "bytes": 293260, + "sha256": "cfbad14ed945b973e5a658626a27a7355eae6f6d620c37c837ef42babb0fe317" + }, + "kda_gfx950/fused_c32_h0_noht": { + "bytes": 276637, + "sha256": "0019441b808ce23834e09626eccf253bb2defbce1e309d4f30fa019bb92c7faf" + }, + "kda_gfx950/fused_c32_noht": { + "bytes": 258852, + "sha256": "22d24890e2ff79b8885362a826b49e65275e5dc5abf3c2fdeb72598778e438f7" + }, + "kda_gfx950/fused_c32_overlay": { + "bytes": 527778, + "sha256": "9afda08570596db7ef731bfc2726594ff042224caed75204e630a97f34cb6c9c" + }, + "kda_gfx950/fused_c32_sa16": { + "bytes": 275419, + "sha256": "dadd9be83684192ce94bad5a0b994e3d94578eecd25f293d12e8a812ae1ea9ec" + }, + "kda_gfx950/fused_c32_sa32": { + "bytes": 413179, + "sha256": "c2f6d9d7a704b587f52890d389e13357a6ddf1cf53a0f71a08cff2a6e9990ba1" + }, + "kda_gfx950/split_c16_prep": { + "bytes": 107257, + "sha256": "70786a5846257ffd3f6b9504e77e70a488b5cbe08990ce90d8b903b8b97486df" + }, + "kda_gfx950/split_c16_scan_sa16": { + "bytes": 83804, + "sha256": "b9c13afb6ce3040a401206b617f468c20735edb909a9c0420e17c146946ff100" + }, + "kda_gfx950/split_c32_prep": { + "bytes": 228607, + "sha256": "9770db795eef0dc7d2e00b06690cad2467debd8bf1ebae1decb246a5d05e8bef" + }, + "kda_gfx950/split_c32_prep_raw_fusions": { + "bytes": 268317, + "sha256": "07064cbf14a40f2a0437e0e76107cc5af4858fb9d9cd9daca4621bd5339e8970" + }, + "kda_gfx950/split_c32_scan_h0": { + "bytes": 212006, + "sha256": "c518d16d087eacb40ebaf6ea29a9bde1f35c79f947c8ab4ebc7ec3c5fef46452" + }, + "kda_gfx950/split_c32_scan_h0_noht": { + "bytes": 175222, + "sha256": "98324f82e31e28e91e25954944d4dc25581fcdd8648e36f02411236f0b8c70e8" + }, + "kda_gfx950/split_c32_scan_noht": { + "bytes": 137314, + "sha256": "f0b6de2006132daecc5831c4970363a9c10bc2bc51dc571a38f7b90fc7dca2a6" + }, + "kda_gfx950/split_c32_scan_sa32": { + "bytes": 173970, + "sha256": "3a5994d628f0685de3b5ee259e73355981f3e40add204aca5f996a7eaacf6978" + }, + "kda_gfx950/split_c32_scan_vs2": { + "bytes": 185409, + "sha256": "f042bf63ea2d9255fcf3ef3dfff2a91a3c15e011e222fa04d8b8af380cd7ff12" + }, + "kda_gfx950/split_c32_scan_vs4": { + "bytes": 192997, + "sha256": "1980619583b3d103d659620ac0f8b904a6fc080a6bf7e25fe3792ef6ff3ddaab" + }, + "kda_gfx950/split_c32_scan_vs8": { + "bytes": 112934, + "sha256": "000fb3bcf7d12a98061eca22f0b74efd27913fa94850ef0748ed9f92003c89b2" + } + } + } + }, + "schema": "kda_gfx950.ir_golden_sha256/v1" +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py new file mode 100644 index 000000000000..bdd884f968ff --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py @@ -0,0 +1,439 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""On-GPU numeric lane for the chunkwise KDA family on gfx950. + +Two references, deliberately independent of each other and of the kernels: + +- The six per-chunk tiles are checked against a float64 torch oracle that + builds them from the pairwise exponent *difference*, with no midpoint + factoring -- so it never forms the overflowing ``Gamma_i / Gamma_j`` the + kernel is designed around, and agreement is a check of that factoring rather + than a restatement of it. +- Both full forward paths are checked against a token-serial float64 walk of + the gated delta rule. That oracle is not chunked at all, so agreement tests + the chunkwise factorization itself. + +The gate range is swept to -5, the reference ``gate_lower_bound``: a 32-token chunk +accumulates up to 160 nats there, which is the regime that saturates the +factored exponents and the only one where the clamping actually matters. + +Every test is marked ``gpu`` and skipped off a gfx950. Select with +``run_all.py --gpu`` (or ``pytest -m gpu``); the default CPU lane excludes it. +""" + +from __future__ import annotations + +import pytest + + +def _gpu_ready(): + """True only on a gfx950 box with ROCm torch. + + Gate on ``gcnArchName`` (the ISA target), NOT the marketing name. + """ + try: + import torch + except Exception: # noqa: BLE001 + return False + if not torch.cuda.is_available(): + return False + arch = torch.cuda.get_device_properties(0).gcnArchName.lower() + return "gfx950" in arch + + +requires_gfx950_gpu = pytest.mark.skipif( + not _gpu_ready(), reason="needs a gfx950 GPU with ROCm torch" +) + +pytestmark = [pytest.mark.gpu, requires_gfx950_gpu] + +# -0.1 is the typical regime; -5.0 is the reference gate lower bound and is what +# saturates the factored exponents. +GATES = [-0.1, -0.5, -2.0, -5.0] +# bf16 operands and an fp32 accumulator against a float64 oracle. The chunkwise +# path also carries a triangular solve, so the tolerance is a bf16 tolerance +# with room for its conditioning, not a tight fp32 one. +TOL = 3e-2 + + +@pytest.mark.parametrize("gate_low", GATES) +def test_prep_tiles_match_float64_oracle(gate_low): + """All six per-chunk tiles, over enough chunks to cover every code path.""" + from builders.gfx950.kda import kda_chunk_prep as prep + + from kernels.gfx950.kda_chunkwise import KdaChunkPrepSpec + + spec = KdaChunkPrepSpec() + assert prep.check(spec, 128, gate_low=gate_low, tol=2e-2, verbose=False) + + +@pytest.mark.parametrize("gate_low", GATES) +def test_split_path_matches_token_serial(gate_low): + """rocke prep + rocke scan, against the token-serial recurrence.""" + from builders.gfx950.kda import kda_chunk_split as split + + from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec + + worst = split.check(KdaChunkScanSpec(), 2, 4, 256, gate_low=gate_low, verbose=False) + assert worst <= TOL, f"gate {gate_low}: worst rel {worst:.3e}" + + +@pytest.mark.parametrize("gate_low", GATES) +def test_fused_path_matches_token_serial(gate_low): + """One kernel, tiles never leaving LDS, against the same oracle. + + The two paths run the same emitted scan body over the same tile math, so + this is the check that routing the tiles through HBM instead of LDS did not + change the arithmetic. + """ + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec + + worst = fused.check( + KdaChunkFusedSpec(), 2, 4, 256, gate_low=gate_low, verbose=False + ) + assert worst <= TOL, f"gate {gate_low}: worst rel {worst:.3e}" + + +@pytest.mark.parametrize("path", ["split", "fused"]) +@pytest.mark.parametrize("gate_low", GATES) +def test_nonzero_initial_state_matches_token_serial(path, gate_low): + """Both compositions must preserve a non-zero state carried into prefill.""" + if path == "split": + from builders.gfx950.kda import kda_chunk_split as module + from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec + + spec = KdaChunkScanSpec(has_initial_state=True) + else: + from builders.gfx950.kda import kda_chunk_fused as module + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec + + spec = KdaChunkFusedSpec(has_initial_state=True) + + worst = module.check( + spec, 2, 4, 256, gate_low=gate_low, with_h0=True, verbose=False + ) + assert worst <= TOL, f"{path}, gate {gate_low}: initial-state worst rel {worst:.3e}" + + +@pytest.mark.parametrize("gate_low", GATES) +def test_subtiled_fused_path_matches_token_serial(gate_low): + """The state-subtiled fused kernel, against the same oracle. + + Subtiling narrows the scan atom so each wave owns a shorter band of ``S^T`` + and twice as many waves cover the v extent, which halves each lane's + loop-carried state. That moves every one of the five products off a + single-atom output onto a tiled one, so this covers index arithmetic the + chunk-wide atom never exercises. + """ + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + spec = KdaChunkFusedSpec(tile=KdaTileSpec(block_size=512, scan_atom_m=16)) + worst = fused.check(spec, 2, 4, 256, gate_low=gate_low, verbose=False) + assert worst <= TOL, f"gate {gate_low}: worst rel {worst:.3e}" + + +def test_subtiled_and_chunk_wide_fused_agree_to_one_ulp(): + """Subtiling must repartition the arithmetic without changing the answer. + + Not bit-exact, and it should not be: the narrower atom packs twice the K per + instruction, so a contraction that took eight accumulation steps takes four, + and fp32 summation is not associative. What must hold is that this is the + *only* difference -- a last-bit effect on a small minority of elements. A + real indexing error in the retiled products would show up here as a large + difference on many elements, which no oracle tolerance would localize. + """ + import torch + + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + B, H, T = 2, 4, 256 + wide = KdaChunkFusedSpec(tile=KdaTileSpec()) + sub = KdaChunkFusedSpec(tile=KdaTileSpec(block_size=512, scan_atom_m=16)) + q, k, v, g, beta = fused.make_inputs(B, H, T, wide.head_k, wide.head_v) + o_w, ht_w = fused.launch_packed(wide, q, k, v, g, beta) + o_s, ht_s = fused.launch_packed(sub, q, k, v, g, beta) + torch.cuda.synchronize() + + # One bf16 ULP is 2^-8 relative; allow half that against the tile's peak + # magnitude, which is far tighter than the 3e-2 oracle tolerance. + for name, a, b in (("output", o_w, o_s), ("final state", ht_w, ht_s)): + diff = (a.float() - b.float()).abs().max().item() + scale = a.float().abs().max().item() + frac = (a != b).float().mean().item() + assert diff <= 2e-3 * scale, ( + f"{name} differs by {diff:.3e} ({diff / scale:.3e} of peak), " + "more than rounding: the retiled indexing is wrong" + ) + assert frac <= 0.10, ( + f"{name}: {frac:.1%} of elements differ; rounding alone should " + "touch only a small minority" + ) + + +@pytest.mark.parametrize("tile_kw", [{}, {"block_size": 512, "scan_atom_m": 16}]) +def test_input_prefetch_is_bitwise_identical(tile_kw): + """Staging a chunk's inputs one chunk early must not change the arithmetic. + + The prefetch moves *when* g/k/q/beta reach LDS, never what lands there, so + this is a bitwise check rather than a tolerance one. + + It is also the race detector. The tiles being written are the same ones the + current chunk staged into, and the hand-off deliberately adds no barrier of + its own -- it leans on the rendezvous the scan already has. A missed + dependency means some chunk reads its successor's inputs, and the earlier + attempt at this optimization showed that failure drifts in rather than + appearing on the first launch: the buffers have to already hold a previous + run's data. Hence the repeat. + """ + import torch + + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + B, H, T = 2, 4, 256 + tile = KdaTileSpec(**tile_kw) + # Both spelled out: the prefetch is the default, so leaving either implicit + # would compare a spec against itself. + base = KdaChunkFusedSpec(tile=tile, prefetch_inputs=False) + pf = KdaChunkFusedSpec(tile=tile, prefetch_inputs=True) + assert base.kernel_name() != pf.kernel_name() + q, k, v, g, beta = fused.make_inputs(B, H, T, base.head_k, base.head_v) + + for it in range(3): + o_b, ht_b = fused.launch_packed(base, q, k, v, g, beta) + o_p, ht_p = fused.launch_packed(pf, q, k, v, g, beta) + torch.cuda.synchronize() + assert torch.equal(o_b, o_p), ( + f"launch {it}: prefetched output differs from the in-phase staging; " + "a chunk read inputs the hand-off had not published" + ) + assert torch.equal(ht_b, ht_p), f"launch {it}: final states diverge" + + +def test_fused_lds_overlay_is_bitwise_identical(): + """Explicit byte-pool views may change addresses, never values. + + This catches lifetime mistakes in the overlay itself. Kt is produced while + the solve tiles are live and V survives into residual formation, so either + one being placed in the reusable tile-only span quickly produces NaNs. + """ + import torch + + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + B, H, T = 2, 4, 256 + tile = KdaTileSpec(block_size=512, scan_atom_m=16) + typed = KdaChunkFusedSpec(tile=tile, prefetch_inputs=False, overlay_lds=False) + overlay = KdaChunkFusedSpec(tile=tile, prefetch_inputs=False, overlay_lds=True) + q, k, v, g, beta = fused.make_inputs(B, H, T, typed.head_k, typed.head_v) + o_t, ht_t = fused.launch_packed(typed, q, k, v, g, beta) + o_o, ht_o = fused.launch_packed(overlay, q, k, v, g, beta) + torch.cuda.synchronize() + + assert torch.isfinite(o_o).all() and torch.isfinite(ht_o).all() + assert torch.equal(o_t, o_o), "LDS overlay changed the fused output" + assert torch.equal(ht_t, ht_o), "LDS overlay changed the final state" + + +def test_c32_tile_phase_16x16_panels_agree_to_one_ulp(): + """Inner panels repartition C x C work without changing the algorithm. + + The 16x16x32 atom takes four K steps where 32x32x16 takes eight, so fp32 + accumulation order changes and bitwise identity is not expected. Differences + must remain last-bit effects rather than panel indexing errors. + """ + import torch + + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + B, H, T = 2, 4, 256 + common = dict(block_size=512, scan_atom_m=16) + atom32 = KdaChunkFusedSpec(tile=KdaTileSpec(**common, tile_atom_m=32)) + panels16 = KdaChunkFusedSpec(tile=KdaTileSpec(**common, tile_atom_m=16)) + q, k, v, g, beta = fused.make_inputs(B, H, T, atom32.head_k, atom32.head_v) + o_32, ht_32 = fused.launch_packed(atom32, q, k, v, g, beta) + o_16, ht_16 = fused.launch_packed(panels16, q, k, v, g, beta) + torch.cuda.synchronize() + + for name, a, b in (("output", o_32, o_16), ("final state", ht_32, ht_16)): + diff = (a.float() - b.float()).abs().max().item() + scale = a.float().abs().max().item() + frac = (a != b).float().mean().item() + assert diff <= 2e-3 * scale, ( + f"{name} differs by {diff:.3e} ({diff / scale:.3e} of peak), " + "more than atom accumulation order can explain" + ) + assert ( + frac <= 0.10 + ), f"{name}: {frac:.1%} of elements differ; panel indexing is suspect" + + +def test_c16_fused_matches_token_serial_oracle(): + """Literal C=16 is a valid A/B even though it is not the throughput winner.""" + from builders.gfx950.kda import kda_chunk_fused as fused + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec + + spec = KdaChunkFusedSpec( + tile=KdaTileSpec( + chunk=16, + block_size=512, + pad_cb=16, + tile_atom_m=16, + scan_atom_m=16, + ) + ) + worst = fused.check(spec, 2, 4, 256, verbose=False) + assert worst <= TOL, f"C16 fused worst relative error {worst:.3e}" + + +def test_split_and_fused_agree_bitwise(): + """The two paths share one emitted scan body, so they must agree exactly. + + Not merely within tolerance: same math, same order, same rounding. A + difference here means the tile round trip through HBM lost something the + LDS path kept, which no tolerance-based check would localize. + """ + import torch + + from builders.gfx950.kda import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_split as split + + from kernels.gfx950.kda_chunkwise import ( + KdaChunkFusedSpec, + KdaChunkScanSpec, + KdaTileSpec, + ) + + B, H, T = 2, 4, 256 + spec_s = KdaChunkScanSpec() + q, k, v, g, beta = fused.make_inputs(B, H, T, spec_s.head_k, spec_s.head_v) + o_s, ht_s = split.launch_packed(spec_s, q, k, v, g, beta) + o_f, ht_f = fused.launch_packed( + KdaChunkFusedSpec(tile=KdaTileSpec()), q, k, v, g, beta + ) + torch.cuda.synchronize() + assert torch.equal(o_s, o_f), "outputs diverge between the split and fused paths" + assert torch.equal(ht_s, ht_f), "final states diverge" + + +def test_scan_rejects_a_spec_it_cannot_emit(): + """The admission rule is part of the contract, not a convenience. + + Checked on the GPU lane too because this is the guard that stops a spec + whose lane mapping has no valid emission from reaching the launcher. + """ + from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec + + from builders.gfx950.kda import kda_chunk_split as split + + with pytest.raises(ValueError, match="unsupported spec"): + split.make_launcher(KdaChunkScanSpec(head_v=64)) + + +def _make_raw_inputs(B, H, T, DK, DV, gate_low=-0.5, seed=0): + import torch + + gen = torch.Generator(device="cuda").manual_seed(seed) + kw = dict(device="cuda", generator=gen) + q = torch.randn(B, T, H, DK, dtype=torch.bfloat16, **kw) + k = torch.randn(B, T, H, DK, dtype=torch.bfloat16, **kw) + v = (torch.randn(B, T, H, DV, dtype=torch.float32, **kw) * 0.2).bfloat16() + g = (gate_low * torch.rand(B, T, H, DK, dtype=torch.float32, **kw)).bfloat16() + beta = torch.randn(B, T, H, dtype=torch.float32, **kw) + a_log = torch.randn(H, dtype=torch.float32, device="cuda", generator=gen) * 0.1 + dt_bias = ( + torch.randn(H * DK, dtype=torch.float32, device="cuda", generator=gen) * 0.01 + ) + return q, k, v, g, beta, a_log, dt_bias + + +@pytest.mark.parametrize("gate_low", GATES) +def test_raw_split_matches_aligned_oracle(gate_low): + """Fused raw prep + token-major split scan vs the aligned float64 oracle.""" + import torch + + from builders.gfx950.kda import kda_chunk_split as split + + B, H, T = 1, 4, 256 + scan, _ = split.aligned_split_specs(1) + q, k, v, g, beta, a_log, dt_bias = _make_raw_inputs(B, H, T, 128, 128, gate_low) + o_got, ht_got = split.launch_raw(scan, q, k, v, g, beta, a_log, dt_bias) + torch.cuda.synchronize() + o_ref, s_ref = split.ref_aligned_raw(q, k, v, g, beta, a_log, dt_bias, 128**-0.5) + o_ref = o_ref.permute(0, 2, 1, 3) + worst = max( + (o_got.double() - o_ref).abs().max().item() + / (o_ref.abs().max().item() + 1e-30), + (ht_got.double() - s_ref).abs().max().item() + / (s_ref.abs().max().item() + 1e-30), + ) + assert worst <= TOL, f"gate {gate_low}: worst rel {worst:.3e}" + + +@pytest.mark.parametrize( + "value_splits,block,scan_atom_m", + [(1, 256, 0), (2, 128, 0), (4, 64, 0), (8, 64, 16)], +) +def test_value_splits_agree_with_vs1(value_splits, block, scan_atom_m): + """Every legal value_splits geometry must match the unsplit scan.""" + import torch + + from builders.gfx950.kda import kda_chunk_split as split + + from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec, KdaTileSpec + + B, H, T = 1, 4, 256 + q, k, v, g, beta, a_log, dt_bias = _make_raw_inputs(B, H, T, 128, 128) + base, _ = split.aligned_split_specs(1) + o0, ht0 = split.launch_raw(base, q, k, v, g, beta, a_log, dt_bias) + tile = KdaTileSpec(chunk=32, block_size=block, scan_atom_m=scan_atom_m) + spec = KdaChunkScanSpec(tile=tile, value_splits=value_splits, token_major_io=True) + o1, ht1 = split.launch_raw(spec, q, k, v, g, beta, a_log, dt_bias) + torch.cuda.synchronize() + assert torch.allclose(o0.float(), o1.float(), rtol=0, atol=3e-2) + assert torch.allclose(ht0.float(), ht1.float(), rtol=0, atol=3e-2) + + +def test_dispatched_specs_match_the_token_serial_oracle(): + """What the dispatcher selects must be what the oracle validates. + + The registry derives its specs from a request instead of reusing the + builders' defaults, so this is what stops the two from drifting: a + selection that compiles but computes something else passes every CPU + wiring test there is. + """ + from builders.gfx950.kda import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_split as split + + from dispatch.kda import KdaRequest, dispatch_kda + + B, H, T = 2, 4, 256 + + def request(**kw): + return KdaRequest(batch=B, num_heads=H, seqlen=T, arch="gfx950", **kw) + + fused_spec = dispatch_kda(request()).spec + assert fused.check(fused_spec, B, H, T, verbose=False) <= TOL + + h0_spec = dispatch_kda(request(has_initial_state=True)).spec + assert fused.check(h0_spec, B, H, T, with_h0=True, verbose=False) <= TOL + + scan_spec = dispatch_kda(request(algorithm="chunk_scan")).spec + prep_spec = dispatch_kda(request(algorithm="chunk_prep")).spec + # The halves are dispatched independently, so their tile layouts have to + # agree or the scan stages something the tile builder never wrote. + assert split.prep_spec_of(scan_spec) == prep_spec + assert split.check(scan_spec, B, H, T, verbose=False) <= TOL diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_spec.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_spec.py new file mode 100644 index 000000000000..5eea8330ebb1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_spec.py @@ -0,0 +1,353 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU lane for the chunkwise KDA family (``kernels/gfx950/kda_chunkwise.py``). + +Covers the three specs' admission rules and takes each builder through comgr, +asserting the emitted code object fits the resource budget. No GPU: the +numeric contract lives in ``test_kda_chunkwise_gfx950_numeric.py``. + +The rejection cases are the interesting half. Each one encodes a structural +assumption the emitted IR depends on -- a lane mapping, an ``ds_read_b128`` +alignment contract, or an occupancy target -- so a spec that violates one must +be refused rather than silently miscompiled or quietly slow. +""" + +from __future__ import annotations + +import pytest + +from kernels.gfx950.kda_chunkwise import ( + KdaChunkFusedSpec, + KdaChunkPrepSpec, + KdaChunkScanSpec, + KdaTileSpec, + build_kda_chunk_fused, + build_kda_chunk_prep, + build_kda_chunk_scan, + is_valid_fused_spec, + is_valid_scan_spec, + is_valid_spec, +) + +ARCH = "gfx950" + + +def _compile_or_skip(kernel, *, arch: str = ARCH): + """Compile through comgr, skipping only when the toolchain is missing. + + A failed compile is a real defect and propagates; an absent toolchain is an + environment fact and skips. + """ + try: + from rocke.helpers.compile import compile_kernel + except Exception as e: # noqa: BLE001 # pragma: no cover - env-dependent + pytest.skip(f"comgr toolchain unavailable: {e}") + try: + return compile_kernel(kernel, arch=arch, capture_ir_text=False) + except ImportError as e: # pragma: no cover - env-dependent + pytest.skip(f"comgr toolchain unavailable: {e}") + + +def _tile(**kw) -> KdaTileSpec: + return KdaTileSpec(**kw) + + +class TestPrepSpec: + def test_default_is_admitted(self): + ok, why = is_valid_spec(KdaChunkPrepSpec(), arch=ARCH) + assert ok, why + + def test_raw_default_off_matches_prepared_name(self): + base = KdaChunkPrepSpec().kernel_name() + assert "raw" not in base + + def test_raw_fusion_requires_raw_inputs(self): + bad = KdaChunkPrepSpec(fuse_gate=True) + ok, why = is_valid_spec(bad, arch=ARCH) + assert not ok + assert "raw_inputs" in why + + def test_raw_aligned_contract_is_admitted(self): + spec = KdaChunkPrepSpec( + raw_inputs=True, + fuse_qk_l2norm=True, + fuse_gate=True, + fuse_beta_sigmoid=True, + has_dt_bias=True, + ) + ok, why = is_valid_spec(spec, arch=ARCH) + assert ok, why + assert "raw" in spec.kernel_name() + assert spec.kernel_name() != KdaChunkPrepSpec().kernel_name() + + def test_lds_within_half_budget(self): + """The prep kernel's whole optimization story is 2 workgroups per CU.""" + assert KdaChunkPrepSpec().lds_bytes() <= 160 * 1024 // 2 + + @pytest.mark.parametrize( + "kw,needle", + [ + # C is the M and N extent of every C x C product, so it is pinned to + # the atom rather than free. + (dict(chunk=64), "atom"), + # A thread owns one (half-chunk, channel) column of the cumsum. + (dict(block_size=128), "block_size"), + # Odd rows of a bf16 C x C tile would land off a 16 B boundary and + # silently break the rank update's ds_read_b128 alignment. + (dict(pad_cb=4), "8"), + # A block step writes back its own rows out of a contiguous run of 8 + # accumulator slots. + (dict(solve_block=4), "8"), + (dict(solve_block=12), "divide"), + ], + ) + def test_rejections(self, kw, needle): + ok, why = is_valid_spec(KdaChunkPrepSpec(tile=_tile(**kw)), arch=ARCH) + assert not ok, f"{kw} should be rejected" + assert needle in why, f"{kw}: unhelpful reason {why!r}" + + def test_unsupported_arch_and_dtype(self): + assert not is_valid_spec(KdaChunkPrepSpec(), arch="gfx942")[0] + assert not is_valid_spec(KdaChunkPrepSpec(dtype="fp16"), arch=ARCH)[0] + + def test_builds_and_fits(self): + spec = KdaChunkPrepSpec() + art = _compile_or_skip(build_kda_chunk_prep(spec)) + assert art.hsaco_bytes > 0 + + def test_invalid_spec_raises(self): + with pytest.raises(ValueError, match="invalid kda_chunk_prep spec"): + build_kda_chunk_prep(KdaChunkPrepSpec(tile=_tile(chunk=64))) + + +class TestFusedSpec: + def test_default_is_admitted(self): + ok, why = is_valid_fused_spec(KdaChunkFusedSpec(), arch=ARCH) + assert ok, why + + def test_c32_16x16_inner_panels_are_admitted(self): + tile = _tile(block_size=512, scan_atom_m=16, tile_atom_m=16) + spec = KdaChunkFusedSpec(tile=tile) + ok, why = is_valid_fused_spec(spec, arch=ARCH) + assert ok, why + assert "ta16" in spec.kernel_name() + + def test_c16_padded_atom_schedule_is_admitted(self): + tile = _tile( + chunk=16, + block_size=512, + pad_cb=16, + tile_atom_m=16, + scan_atom_m=16, + ) + spec = KdaChunkFusedSpec(tile=tile) + ok, why = is_valid_fused_spec(spec, arch=ARCH) + assert ok, why + assert spec.lds_bytes() == 87_104 + + # K=32 consumes 16 real chunk columns plus 16 explicit zero columns. + bad = KdaChunkFusedSpec( + tile=_tile( + chunk=16, + block_size=512, + pad_cb=8, + tile_atom_m=16, + scan_atom_m=16, + ) + ) + ok, why = is_valid_fused_spec(bad, arch=ARCH) + assert not ok + assert "zero pad" in why + + def test_unknown_tile_atom_is_rejected(self): + ok, why = is_valid_spec(KdaChunkPrepSpec(tile=_tile(tile_atom_m=24)), arch=ARCH) + assert not ok + assert "tile_atom_m" in why + + def test_head_v_must_match_wave_count(self): + """Each wave owns one atom-row band of the state and nothing else. + + That single rule partitions all five products in the scan body, which is + what keeps the state in registers with no cross-wave reduction, so a + ``head_v`` the waves do not cover exactly has no valid emission. + """ + ok, why = is_valid_fused_spec(KdaChunkFusedSpec(head_v=64), arch=ARCH) + assert not ok + assert "head_v" in why + + def test_overlay_budget_and_prefetch_exclusion(self): + """The explicit pool aliases only buffers with disjoint lifetimes. + + Barrier-free input prefetch deliberately writes the staging tiles while + the scan is live, so combining it with that aliasing would corrupt the + state mirror/residual. Keep the invalid combination out of codegen. + """ + tile = KdaTileSpec() + bad = KdaChunkFusedSpec(tile=tile, overlay_lds=True) + ok, why = is_valid_fused_spec(bad, arch=ARCH) + assert not ok + assert "prefetch_inputs=False" in why + + overlay = KdaChunkFusedSpec(tile=tile, prefetch_inputs=False, overlay_lds=True) + ok, why = is_valid_fused_spec(overlay, arch=ARCH) + assert ok, why + assert overlay.lds_bytes() == 88_704 + assert ( + overlay.lds_bytes() + < KdaChunkFusedSpec(tile=tile, prefetch_inputs=False).lds_bytes() + ) + + def test_builds_and_fits(self): + art = _compile_or_skip(build_kda_chunk_fused(KdaChunkFusedSpec())) + assert art.hsaco_bytes > 0 + + +class TestScanSpec: + def test_default_is_admitted(self): + ok, why = is_valid_scan_spec(KdaChunkScanSpec(), arch=ARCH) + assert ok, why + + @pytest.mark.parametrize( + "value_splits,block,scan_atom_m", + [(1, 256, 0), (2, 128, 0), (4, 64, 0), (8, 64, 16)], + ) + def test_value_split_geometries(self, value_splits, block, scan_atom_m): + tile = KdaTileSpec(chunk=32, block_size=block, scan_atom_m=scan_atom_m) + spec = KdaChunkScanSpec( + tile=tile, value_splits=value_splits, token_major_io=True + ) + ok, why = is_valid_scan_spec(spec, arch=ARCH) + assert ok, why + if value_splits != 1: + assert f"vs{value_splits}" in spec.kernel_name() + + def test_token_major_flag_reaches_name(self): + spec = KdaChunkScanSpec(token_major_io=True) + assert "tm" in spec.kernel_name() + + def test_raw_prep_preserves_scan_tile_knobs(self): + pytest.importorskip("torch") + from builders.gfx950.kda.kda_chunk_split import prep_spec_of + + tile = KdaTileSpec( + block_size=64, + pad_dk=16, + pad_c=8, + pad_cb=16, + solve_block=16, + tile_atom_m=16, + scan_atom_m=16, + waves_per_eu=2, + ) + spec = KdaChunkScanSpec(tile=tile) + prep = prep_spec_of(spec, raw=True) + + assert prep.tile.block_size == 256 + for field in ( + "pad_dk", + "pad_c", + "pad_cb", + "solve_block", + "tile_atom_m", + "scan_atom_m", + "waves_per_eu", + ): + assert getattr(prep.tile, field) == getattr(tile, field) + + def test_lds_leaves_room_for_two_workgroups(self): + """The split path only earns back its tile traffic at 2 WG/CU. + + The scan is a latency-bound chain of small matmuls; at one workgroup per + CU there is no second workgroup to cover it. So the occupancy target is a + spec-level rejection rule here, not a tuning note. + """ + spec = KdaChunkScanSpec() + assert spec.lds_bytes() <= 160 * 1024 // spec.min_occupancy + ok, why = is_valid_scan_spec(KdaChunkScanSpec(min_occupancy=3), arch=ARCH) + assert not ok + assert "workgroups per CU" in why + + def test_head_v_must_match_wave_count(self): + ok, why = is_valid_scan_spec(KdaChunkScanSpec(head_v=64), arch=ARCH) + assert not ok + assert "head_v" in why or "v slice" in why + + def test_staging_tiles_must_fit_whole_workgroup_passes(self): + spec = KdaChunkScanSpec(head_k=96) + ok, why = is_valid_scan_spec(spec, arch=ARCH) + assert not ok + assert "slots" in why + + def test_dec_tile_must_fit_one_guarded_pass(self): + spec = KdaChunkScanSpec(head_k=2048) + ok, why = is_valid_scan_spec(spec, arch=ARCH) + assert not ok + assert "dec tile" in why + + def test_staging_alignment_rejections(self): + """Staging is ds_write_b128 throughout, so both pitches stay 8-aligned.""" + for kw in (dict(pad_dk=4), dict(pad_cb=4)): + ok, why = is_valid_scan_spec(KdaChunkScanSpec(tile=_tile(**kw)), arch=ARCH) + assert not ok, kw + assert "8" in why + + def test_builds_and_fits(self): + art = _compile_or_skip(build_kda_chunk_scan(KdaChunkScanSpec())) + assert art.hsaco_bytes > 0 + + +class TestSpecNaming: + def test_names_are_distinct_and_carry_the_shape(self): + names = { + KdaChunkPrepSpec().kernel_name(), + KdaChunkFusedSpec().kernel_name(), + KdaChunkScanSpec().kernel_name(), + } + assert len(names) == 3 + for n in names: + assert "dk128" in n and "dv128" in n and "c32" in n + + def test_off_default_knobs_reach_the_name(self): + """Two specs that emit different code must not share a cache key.""" + a = KdaChunkPrepSpec().kernel_name() + b = KdaChunkPrepSpec(tile=_tile(solve_block=32)).kernel_name() + assert a != b + c = KdaChunkPrepSpec(tile=_tile(pad_cb=16)).kernel_name() + assert a != c + assert "pcb16" in c + + scan = KdaChunkScanSpec() + scan_h0 = KdaChunkScanSpec(has_initial_state=True, store_final_state=False) + assert scan.kernel_name() != scan_h0.kernel_name() + assert "h0" in scan_h0.kernel_name() + assert "noht" in scan_h0.kernel_name() + + fused = KdaChunkFusedSpec() + fused_h0 = KdaChunkFusedSpec( + has_initial_state=True, + store_final_state=False, + prefetch_inputs=False, + overlay_lds=True, + tile=_tile( + block_size=512, + pad_dk=16, + pad_cb=16, + tile_atom_m=16, + scan_atom_m=16, + solve_block=32, + waves_per_eu=2, + ), + ) + assert fused.kernel_name() != fused_h0.kernel_name() + for needle in ( + "h0", + "noht", + "nopf", + "ovl", + "sb32", + "p16x4", + "ta16", + "wpe2", + ): + assert needle in fused_h0.kernel_name() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_gfx950_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_gfx950_golden.py new file mode 100644 index 000000000000..4d3d8541e899 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_gfx950_golden.py @@ -0,0 +1,244 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Golden LLVM-IR stability test for the gfx950 KDA kernel family. + +This is the per-family deep lane for chunkwise KDA. It pins both chunk sizes, +both scan atom widths, the split and fused paths, state ABI flags, and every +supported value-split geometry. Raw-input fusion and the fused LDS overlay are +also represented. The test only lowers Python to LLVM IR; no GPU or comgr is +required. + +Run or re-bless from ``rocke/library``:: + + python tests/run_all.py --only kda_gfx950 + python tests/test_kda_gfx950_golden.py --write + +Review the IR change before re-blessing the hashes. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path +from typing import Callable + +_TESTS = Path(__file__).resolve().parent +_LIBRARY = _TESTS.parent +_PLATFORM_PYTHON = _LIBRARY.parent / "platform" / "python" +_GOLDEN = _TESTS / "golden" / "kda_gfx950_ir_sha256.json" +_FLAVORS = ("llvm20", "llvm22", "llvm23") +_ARCH = "gfx950" + +# The library path must precede tests/ so tests/dispatch cannot shadow the real +# dispatch package when this file is executed directly. +for _path in (str(_LIBRARY), str(_PLATFORM_PYTHON)): + if _path not in sys.path: + sys.path.insert(0, _path) + + +def _cases() -> dict[str, Callable]: + """Return builders spanning the gfx950 KDA code-generation knob space.""" + from kernels.gfx950.kda_chunkwise import ( + KdaChunkFusedSpec, + KdaChunkPrepSpec, + KdaChunkScanSpec, + KdaTileSpec, + build_kda_chunk_fused, + build_kda_chunk_prep, + build_kda_chunk_scan, + ) + + cases: dict[str, Callable] = {} + kernel_names: dict[str, str] = {} + + def add(cid, spec, builder): + name = spec.kernel_name() + assert name not in kernel_names, ( + f"KDA kernel-name collision: {cid!r} and {kernel_names[name]!r} " + f"both use {name!r}" + ) + kernel_names[name] = cid + cases[cid] = lambda spec=spec, builder=builder: builder(spec, arch=_ARCH) + + # C16 needs an explicit zero-padded atom schedule. Reuse it in both halves + # of the split path and in the fused path so all three emissions are pinned. + c16 = KdaTileSpec( + chunk=16, + block_size=512, + pad_cb=16, + tile_atom_m=16, + scan_atom_m=16, + ) + + add("kda_gfx950/split_c32_prep", KdaChunkPrepSpec(), build_kda_chunk_prep) + add( + "kda_gfx950/split_c16_prep", + KdaChunkPrepSpec(tile=c16), + build_kda_chunk_prep, + ) + add( + "kda_gfx950/split_c32_prep_raw_fusions", + KdaChunkPrepSpec( + raw_inputs=True, + fuse_qk_l2norm=True, + fuse_gate=True, + fuse_beta_sigmoid=True, + has_dt_bias=True, + ), + build_kda_chunk_prep, + ) + + # The default standalone scan is C32/SA32/value_splits=1. + add("kda_gfx950/split_c32_scan_sa32", KdaChunkScanSpec(), build_kda_chunk_scan) + add( + "kda_gfx950/split_c16_scan_sa16", + KdaChunkScanSpec(tile=c16), + build_kda_chunk_scan, + ) + add( + "kda_gfx950/split_c32_scan_h0", + KdaChunkScanSpec(has_initial_state=True), + build_kda_chunk_scan, + ) + add( + "kda_gfx950/split_c32_scan_noht", + KdaChunkScanSpec(store_final_state=False), + build_kda_chunk_scan, + ) + add( + "kda_gfx950/split_c32_scan_h0_noht", + KdaChunkScanSpec(has_initial_state=True, store_final_state=False), + build_kda_chunk_scan, + ) + for value_splits, block_size, scan_atom_m in ( + (2, 128, 0), + (4, 64, 0), + (8, 64, 16), + ): + add( + f"kda_gfx950/split_c32_scan_vs{value_splits}", + KdaChunkScanSpec( + tile=KdaTileSpec( + block_size=block_size, + scan_atom_m=scan_atom_m, + ), + value_splits=value_splits, + token_major_io=True, + ), + build_kda_chunk_scan, + ) + + # The tuned fused default is C32/SA16; an untuned C32 tile pins SA32. + add("kda_gfx950/fused_c32_sa16", KdaChunkFusedSpec(), build_kda_chunk_fused) + add( + "kda_gfx950/fused_c32_sa32", + KdaChunkFusedSpec(tile=KdaTileSpec()), + build_kda_chunk_fused, + ) + add( + "kda_gfx950/fused_c16_sa16", + KdaChunkFusedSpec(tile=c16), + build_kda_chunk_fused, + ) + add( + "kda_gfx950/fused_c32_h0", + KdaChunkFusedSpec(has_initial_state=True), + build_kda_chunk_fused, + ) + add( + "kda_gfx950/fused_c32_noht", + KdaChunkFusedSpec(store_final_state=False), + build_kda_chunk_fused, + ) + add( + "kda_gfx950/fused_c32_h0_noht", + KdaChunkFusedSpec(has_initial_state=True, store_final_state=False), + build_kda_chunk_fused, + ) + add( + "kda_gfx950/fused_c32_overlay", + KdaChunkFusedSpec( + tile=KdaTileSpec(), + prefetch_inputs=False, + overlay_lds=True, + ), + build_kda_chunk_fused, + ) + + return cases + + +def _current_flavor() -> str: + from rocke.core.lower_llvm import _resolve_llvm_flavor + + return _resolve_llvm_flavor() + + +def _sha_for(build: Callable, flavor: str) -> tuple[str, int]: + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + llvm = _lower_kernel_to_llvm_python(build(), arch=_ARCH, llvm_flavor=flavor) + data = llvm.encode("utf-8") + return hashlib.sha256(data).hexdigest(), len(data) + + +def _build_doc() -> dict: + cases = _cases() + return { + "schema": "kda_gfx950.ir_golden_sha256/v1", + "flavors": { + flavor: { + "cases": { + cid: {"sha256": sha, "bytes": nbytes} + for cid, build in cases.items() + for sha, nbytes in [_sha_for(build, flavor)] + } + } + for flavor in _FLAVORS + }, + } + + +def test_kda_gfx950_ir_matches_golden(): + assert _GOLDEN.exists(), ( + "missing gfx950 KDA golden fixture; generate it with " + f"`python {Path(__file__).name} --write`" + ) + golden = json.loads(_GOLDEN.read_text()) + assert golden.get("schema") == "kda_gfx950.ir_golden_sha256/v1" + + flavor = _current_flavor() + assert flavor in golden.get("flavors", {}), ( + f"no gfx950 KDA golden recorded for LLVM flavor {flavor!r}; " + "review and re-bless the fixture" + ) + + cases = _cases() + recorded = golden["flavors"][flavor]["cases"] + assert set(recorded) == set(cases), ( + "gfx950 KDA golden case set drifted: " + f"recorded={sorted(recorded)}, current={sorted(cases)}" + ) + + drift = [] + for cid, build in cases.items(): + want = recorded[cid]["sha256"] + got, nbytes = _sha_for(build, flavor) + if got != want: + drift.append( + f"{cid}: {want} -> {got} " + f"({recorded[cid]['bytes']} -> {nbytes} bytes)" + ) + assert not drift, "gfx950 KDA LLVM IR drift vs golden:\n " + "\n ".join(drift) + + +if __name__ == "__main__": + if "--write" in sys.argv: + _GOLDEN.parent.mkdir(parents=True, exist_ok=True) + _GOLDEN.write_text(json.dumps(_build_doc(), indent=2, sort_keys=True) + "\n") + print(f"wrote {_GOLDEN}") + else: + test_kda_gfx950_ir_matches_golden() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md index 1bca9f176edd..74cff9aa4aab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md @@ -96,9 +96,10 @@ code with the family above. | File | Spec | Doc | |-----------------------------------|-------------------------------------------------------------------|------------------------------| | `gfx942/kda_chunkwise.py` | `KdaChunkFusedSpec`, `KdaChunkPrepSpec`, `KdaChunkScanSpec`, `KdaTileSpec` | `instances/kda.md` | +| `gfx950/kda_chunkwise.py` | `KdaChunkFusedSpec`, `KdaChunkPrepSpec`, `KdaChunkScanSpec`, `KdaTileSpec` | `instances/kda.md` | Three kernels: a fused prefill, and a two-phase split path (per-chunk tile -builder, then state scan). gfx942 / bf16 only; prefill only, no varlen. +builder, then state scan). gfx942 and gfx950 are bf16-only; prefill only, no varlen. Dispatch is `library/dispatch/kda/` (`dispatch_kda`), which defaults to the fused kernel and keeps the split halves opt-in. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/kda.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/kda.md index 158dfa34d058..0ab63dd121d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/kda.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/kda.md @@ -1,4 +1,79 @@ -# Chunkwise KDA Instances +# Chunkwise KDA + +The gfx950 chunkwise gated delta-rule linear-attention prefill implementation +is provided by +[`kda_chunkwise.py`](../../../library/kernels/gfx950/kda_chunkwise.py). +The host builders are in +[`library/builders/gfx950/kda`](../../../library/builders/gfx950/kda/). + +## Specs and compositions + +`KdaTileSpec` owns the chunk, workgroup, MFMA, solve, padding, and occupancy +knobs. The operation has three build entry points: + +- `build_kda_chunk_prep`: constructs the six state-independent tiles for each + chunk. +- `build_kda_chunk_scan`: stages those tiles and applies the serial state + recurrence for each `(batch, head)`. +- `build_kda_chunk_fused`: constructs and consumes the tiles in one workgroup. + +The split composition is the default. The fused composition is useful for +cross-checking the shared scan body and for future LDS/resource schedules. +Both support an optional initial state and final-state output. + +## Numerical contract + +The implementation uses a midpoint decay factorization so per-channel gate +ratios remain finite over a chunk. The tile builder and both complete paths are +checked against independent float64 references. The split and fused paths are +also checked for bitwise equality when their accumulation order is identical. + +## Validation + +Run the CPU admission/build checks from `rocke/library`: + +```bash +python -m pytest tests/test_kda_chunkwise_spec.py +``` + +Run the gfx950 numeric checks on a matching GPU: + +```bash +python -m pytest tests/test_kda_chunkwise_gfx950_numeric.py -m gpu +``` + +Run the CPU wiring checks for the dispatcher family: + +```bash +python -m pytest tests/dispatch/kda/test_gfx950_wiring.py +``` + +## Dispatch + +`dispatch.kda` registers one candidate per emitted kernel: + +| Candidate | `algorithm` | Kernel | +| --- | --- | --- | +| `kda_gfx950_chunk_fused` | `chunk_fused` | `build_kda_chunk_fused` | +| `kda_gfx950_chunk_prep` | `chunk_prep` | `build_kda_chunk_prep` | +| `kda_gfx950_chunk_scan` | `chunk_scan` | `build_kda_chunk_scan` | + +```python +from dispatch.kda import KdaRequest, dispatch_kda + +result = dispatch_kda(KdaRequest(batch=8, num_heads=16, seqlen=2048, arch="gfx950")) +``` + +The fused kernel is the default. The split halves are opt-in by `algorithm` +or `spec_id` and must be launched in that order: the scan consumes the tiles +the tile builder writes. There is no automatic fused/split routing, because +the crossover has been measured at too few shapes to encode a threshold; see +`dispatch/kda/common.py`. `bind` is not wired, so launching still goes through +the host builders. + +The standalone benchmark scenario is `benchmarks.gfx950.kda.benchmark_chunkwise`. + +## gfx942 implementation Kimi Delta Attention prefill: a **gated delta-rule linear recurrence**, not softmax attention. Each step applies a per-channel forget gate to a diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index c48fa65248fb..40d0a118abc3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -98,9 +98,9 @@ as described in the notes. | Instance | gfx942 | gfx950 | gfx1151 | Notes | |---|:--:|:--:|:--:|---| -| `kda_chunk_fused` | ✅ | ❌ | ❌ | bf16 only; fused prefill, one workgroup per (batch, head, V partition) | -| `kda_chunk_prep` | ✅ | ❌ | ❌ | bf16 only; split path phase 1, one workgroup per chunk | -| `kda_chunk_scan` | ✅ | ❌ | ❌ | bf16 only; split path phase 2, consumes what prep wrote | +| `kda_chunk_fused` | ✅ | ✅ | ❌ | bf16 only; fused prefill; gfx942 partitions V, gfx950 owns a full head | +| `kda_chunk_prep` | ✅ | ✅ | ❌ | bf16 only; split path phase 1, one workgroup per chunk | +| `kda_chunk_scan` | ✅ | ✅ | ❌ | bf16 only; split path phase 2, consumes what prep wrote | --- @@ -150,15 +150,14 @@ as described in the notes. supported. fp8/bf8 output needs the CDNA-only `v_cvt_pk_{fp8,bf8}_f32` conversion, so fp8/bf8 specs are rejected by the validator on non-CDNA families. -- **KDA (`kda_chunk_*`)** is gfx942-only by construction, not by omission: all - three validators reject any other `arch` (`test_kda_chunkwise_gfx942_spec.py`, - and the dispatch arch gate in `library/tests/dispatch/kda/`). CDNA3 lacks the - K-packed bf16 atoms and transposing LDS reads a CDNA4 path would use, so the - swizzled store/load pairing here assumes `a_per_lane == 4`. The ❌ cells are - therefore a refusal, not an untested gap. The gfx942 ✅ is GPU-numeric-verified - against a torch reference (`test_kda_chunkwise_gfx942_numeric.py`, needs a - local gfx942 device); the emitted IR is additionally pinned by - `test_kda_gfx942_golden.py`, which needs no device. +- **KDA (`kda_chunk_*`)** has separate gfx942 and gfx950 emitters, each with an + architecture gate and schedule matched to that ISA. gfx942 partitions a + logical value head into 64-channel workgroups to fit its LDS budget; gfx950 + uses CDNA4 K-packed bf16 atoms and supports both full-head and value-split + scans. Both ✅ columns are GPU-numeric-verified by their architecture-specific + tests, and their emitted IR is pinned by `test_kda_gfx942_golden.py` and + `test_kda_gfx950_golden.py`. gfx1151 remains an explicit refusal, not an + untested gap. - All other ✅ cells remain compile-verified only (HSACO produced for the target; not yet GPU-numeric-verified). - gfx942/gfx950 cells use a portable f16 16x16x16 config; an instance marked ❌ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 7dbae5a67b95..1431fe990680 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -866,6 +866,62 @@ "sha256": "e301169b6e63d0ddfccbaf62df00abedbcb880f32b5505a8400b29e89ac0e41d", "status": "ok" }, + "kda_chunkwise/gfx950/fused_default": { + "arch": "gfx950", + "bytes": 275408, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16", + "sha256": "92a389f0a982b5874e5428cef873fd4754fd6ff4f7cc16a9112f00ee1c8830a8", + "status": "ok" + }, + "kda_chunkwise/gfx950/fused_h0_noht": { + "arch": "gfx950", + "bytes": 276626, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16_h0_noht", + "sha256": "46abdeed2bae56c48472d90fcf5478029d544c42b53fdcdd2315fc8bc1fafdf3", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_default": { + "arch": "gfx950", + "bytes": 228596, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "9d9a75f52ef0c75b10d9dab706dbcf0e062ae1f913f38feae7180548ede3fd62", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_raw": { + "arch": "gfx950", + "bytes": 268306, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8_raw_nl2_fg_fb_db", + "sha256": "06770cbb78b8ebf6d3c6778f232a642eb59c72a211a7b6c4257e656993dce067", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_default": { + "arch": "gfx950", + "bytes": 173959, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "c789c33410b0e19a458c4fff0742340465ef84ae5b8f8c810dbfd8443056d3a9", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_h0_noht": { + "arch": "gfx950", + "bytes": 175211, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8_h0_noht", + "sha256": "db30504ff7f8a24f5993be118964f0129bf60a71eb284f50ef9dc706ef1b99be", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_vs4": { + "arch": "gfx950", + "bytes": 184711, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b64_sb8_vs4", + "sha256": "0b397b8339942ebdd73cdb99e4316c2f5abab915c071f5d2621941b77fc6766f", + "status": "ok" + }, "moe_sort/gfx1151/hist_t32_k2_e8": { "arch": "gfx1151", "bytes": 3011, @@ -924,7 +980,7 @@ "gfx1250": 4, "gfx90a": 5, "gfx942": 24, - "gfx950": 63 + "gfx950": 70 }, "expected_failures": 0, "families": { @@ -937,10 +993,10 @@ "deep_fused_conv": 9, "gemm": 13, "grouped_gemm": 1, - "kda_chunkwise": 6, + "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 114 + "ok_cases": 121 } }, "llvm22": { @@ -1809,6 +1865,62 @@ "sha256": "968ebfe2238263f08600d5f447bafc7a31e91c5c2803d72f56af6906c4518420", "status": "ok" }, + "kda_chunkwise/gfx950/fused_default": { + "arch": "gfx950", + "bytes": 275415, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16", + "sha256": "05aeab52f2d3c076d09ef025965f04ae8d393fcff6233a67037a09c4fc34041f", + "status": "ok" + }, + "kda_chunkwise/gfx950/fused_h0_noht": { + "arch": "gfx950", + "bytes": 276633, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16_h0_noht", + "sha256": "c63a5841abdb9df9f2039365106ed495155d8da85e3b220a7167e702cf76ff7b", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_default": { + "arch": "gfx950", + "bytes": 228603, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "0abf19d288d02ed05b997ac2ebfcd100f3f05c0aecb74b55f7670ab8c006ffd7", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_raw": { + "arch": "gfx950", + "bytes": 268313, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8_raw_nl2_fg_fb_db", + "sha256": "556c14b34332821a5cabfedf0a6922dafef83875cf21737286172eda5a7aa27b", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_default": { + "arch": "gfx950", + "bytes": 173966, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "6141220badadc913363758769b31c38240b7138f1beb0585d27472b8e29e7fbb", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_h0_noht": { + "arch": "gfx950", + "bytes": 175218, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8_h0_noht", + "sha256": "f7601e73a1b4038e02f394588e8add93dd8d4722c63756e0855d9480562a17a4", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_vs4": { + "arch": "gfx950", + "bytes": 184718, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b64_sb8_vs4", + "sha256": "788425adf2304fc0a9fdffb293c8afff7732a0df85019314e8a0096c8e7ba5d1", + "status": "ok" + }, "moe_sort/gfx1151/hist_t32_k2_e8": { "arch": "gfx1151", "bytes": 3018, @@ -1867,7 +1979,7 @@ "gfx1250": 4, "gfx90a": 5, "gfx942": 24, - "gfx950": 63 + "gfx950": 70 }, "expected_failures": 0, "families": { @@ -1880,10 +1992,10 @@ "deep_fused_conv": 9, "gemm": 13, "grouped_gemm": 1, - "kda_chunkwise": 6, + "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 114 + "ok_cases": 121 } }, "llvm23": { @@ -2752,6 +2864,62 @@ "sha256": "1b409ef851d96bfa15a5b5dbb700aa3d808f4adaa45bd31499f3a9d2e237ce4c", "status": "ok" }, + "kda_chunkwise/gfx950/fused_default": { + "arch": "gfx950", + "bytes": 275419, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16", + "sha256": "dadd9be83684192ce94bad5a0b994e3d94578eecd25f293d12e8a812ae1ea9ec", + "status": "ok" + }, + "kda_chunkwise/gfx950/fused_h0_noht": { + "arch": "gfx950", + "bytes": 276637, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_fused_dk128_dv128_bf16_c32_b512_sb8_p16x4_pcb16_sa16_ta16_h0_noht", + "sha256": "0019441b808ce23834e09626eccf253bb2defbce1e309d4f30fa019bb92c7faf", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_default": { + "arch": "gfx950", + "bytes": 228607, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "9770db795eef0dc7d2e00b06690cad2467debd8bf1ebae1decb246a5d05e8bef", + "status": "ok" + }, + "kda_chunkwise/gfx950/prep_raw": { + "arch": "gfx950", + "bytes": 268317, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_prep_dk128_dv128_bf16_c32_b256_sb8_raw_nl2_fg_fb_db", + "sha256": "07064cbf14a40f2a0437e0e76107cc5af4858fb9d9cd9daca4621bd5339e8970", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_default": { + "arch": "gfx950", + "bytes": 173970, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8", + "sha256": "3a5994d628f0685de3b5ee259e73355981f3e40add204aca5f996a7eaacf6978", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_h0_noht": { + "arch": "gfx950", + "bytes": 175222, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b256_sb8_h0_noht", + "sha256": "98324f82e31e28e91e25954944d4dc25581fcdd8648e36f02411236f0b8c70e8", + "status": "ok" + }, + "kda_chunkwise/gfx950/scan_vs4": { + "arch": "gfx950", + "bytes": 184722, + "family": "kda_chunkwise", + "kernel_name": "rocke_kda_chunk_scan_dk128_dv128_bf16_c32_b64_sb8_vs4", + "sha256": "16ad41e2198c8aa36421c0927923a22b029627e37917a0338396a0f7cb99372e", + "status": "ok" + }, "moe_sort/gfx1151/hist_t32_k2_e8": { "arch": "gfx1151", "bytes": 3022, @@ -2810,7 +2978,7 @@ "gfx1250": 4, "gfx90a": 5, "gfx942": 24, - "gfx950": 63 + "gfx950": 70 }, "expected_failures": 0, "families": { @@ -2823,10 +2991,10 @@ "deep_fused_conv": 9, "gemm": 13, "grouped_gemm": 1, - "kda_chunkwise": 6, + "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 114 + "ok_cases": 121 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 3e759b85f587..e7b98ec81e28 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -428,6 +428,41 @@ def _build(): return _build +def build_kda_chunkwise_gfx950(kind, arch, **over): + """gfx950 chunkwise KDA kernel from a representative spec. + + ``kind`` selects one of the three emitted kernels. Spec overrides keep the + case table compact while covering state flags, fused preprocessing, and a + non-default scan partition. + """ + + def _build(): + from kernels.gfx950.kda_chunkwise import ( + KdaChunkFusedSpec, + KdaChunkPrepSpec, + KdaChunkScanSpec, + KdaTileSpec, + build_kda_chunk_fused, + build_kda_chunk_prep, + build_kda_chunk_scan, + ) + + if kind == "prep": + return build_kda_chunk_prep(KdaChunkPrepSpec(**over), arch=arch) + if kind == "scan": + scan_over = dict(over) + tile_over = scan_over.pop("tile", {}) + tile = KdaTileSpec(**tile_over) if tile_over else KdaTileSpec() + return build_kda_chunk_scan( + KdaChunkScanSpec(tile=tile, **scan_over), arch=arch + ) + if kind == "fused": + return build_kda_chunk_fused(KdaChunkFusedSpec(**over), arch=arch) + raise ValueError(f"unknown KDA kernel kind {kind!r}") + + return _build + + def build_kda_chunkwise_gfx942(kind, arch, **over): """Build one representative gfx942 chunkwise KDA kernel.""" @@ -2081,8 +2116,44 @@ def add(family, case_id, arch, build): ), ) - # Chunkwise KDA: all three emitted kernels, state-sensitive ABI variants, - # and DK64 fused coverage for the guarded partial Kt pass. + # Chunkwise KDA: each emitted kernel plus the ABI/resource-sensitive + # variants that change preprocessing, state pointers, or scan geometry. + for _case_id, _kind, _over in ( + ("prep_default", "prep", {}), + ( + "prep_raw", + "prep", + { + "raw_inputs": True, + "fuse_qk_l2norm": True, + "fuse_gate": True, + "fuse_beta_sigmoid": True, + "has_dt_bias": True, + }, + ), + ("scan_default", "scan", {}), + ( + "scan_h0_noht", + "scan", + {"has_initial_state": True, "store_final_state": False}, + ), + ("scan_vs4", "scan", {"tile": {"block_size": 64}, "value_splits": 4}), + ("fused_default", "fused", {}), + ( + "fused_h0_noht", + "fused", + {"has_initial_state": True, "store_final_state": False}, + ), + ): + add( + "kda_chunkwise", + f"kda_chunkwise/gfx950/{_case_id}", + "gfx950", + build_kda_chunkwise_gfx950(_kind, "gfx950", **_over), + ) + + # gfx942 KDA: all three emitted kernels, state-sensitive ABI variants, and + # DK64 fused coverage for the guarded partial Kt pass. for _case_id, _kind, _over in ( ("prep_default", "prep", {}), ("scan_default", "scan", {}), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke_ci_static.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke_ci_static.py index 885576959795..b4cf95d17239 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke_ci_static.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke_ci_static.py @@ -108,6 +108,7 @@ def test_ir_cases_cover_all_static_architectures(self): self.assertIn("moe", family_by_arch["gfx950"]) self.assertIn("deep_fused_conv", family_by_arch["gfx950"]) + self.assertIn("kda_chunkwise", family_by_arch["gfx950"]) self.assertIn("kda_chunkwise", family_by_arch["gfx942"]) def test_ir_case_ids_are_unique(self):