From 36c0af6911416e443bb1b5d278a51bd7f54a96e5 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Mon, 24 Aug 2026 19:15:07 +0000 Subject: [PATCH 1/6] feat(rocke): add chunkwise KDA prefill kernels for gfx950 Adds a spec-driven family for chunkwise gated delta-rule linear attention (KDA) prefill: the per-chunk tile builder, a standalone state scan over materialized tiles, and a fused variant that keeps the tiles in LDS. The per-chunk body factors into six state-independent tiles plus a serial state recurrence, and the two halves have opposite parallelism profiles. Both compositions therefore ship: the split path (tiles to HBM, one workgroup per chunk, then one per batch-head for the scan) is the default, and the fused path is retained for the bitwise cross-check and for future arch budgets. The fused path is the design that looks strictly better on paper -- it never writes the tiles out -- and it loses. Holding the tile builder's staging tiles and the scan's operands live at once puts a workgroup over half the LDS budget, so only one is resident per CU, and the scan is a latency-bound chain of small matmuls with nothing left to cover it. That trade is encoded as a rejection rule: is_valid_scan_spec checks LDS against LDS_LIMIT // min_occupancy rather than the full budget, so a spec that fits the hardware but not two workgroups per CU is refused. Correctness rests on two independent float64 oracles: a tile oracle built from pairwise exponent differences with no midpoint factoring, and a token-serial walk of the recurrence that is not chunked at all. The split and fused paths share one emitted scan body and are checked for bitwise equality, not merely agreement within tolerance. Methodology, levers and the attribution technique behind them are in library/builders/gfx950/kda/CASE_STUDY.md. Measurements are not recorded here; they live on the protected internal page. Co-authored-by: Cursor --- .../library/builders/gfx950/kda/CASE_STUDY.md | 227 +++ .../library/builders/gfx950/kda/__init__.py | 0 .../builders/gfx950/kda/kda_chunk_fused.py | 283 +++ .../builders/gfx950/kda/kda_chunk_prep.py | 316 +++ .../builders/gfx950/kda/kda_chunk_split.py | 280 +++ .../library/kernels/gfx950/kda_chunkwise.py | 1775 +++++++++++++++++ .../test_kda_chunkwise_gfx950_numeric.py | 148 ++ .../library/tests/test_kda_chunkwise_spec.py | 173 ++ 8 files changed, 3202 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/CASE_STUDY.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/__init__.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_split.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/kda_chunkwise.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_spec.py 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..17fe3a3c3daa --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/CASE_STUDY.md @@ -0,0 +1,227 @@ +# 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. 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..0747f6ce1b61 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py @@ -0,0 +1,283 @@ +"""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 os +import sys + +_HERE = os.path.dirname(__file__) +_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) +sys.path.insert(0, _RK + "/platform/python") +sys.path.insert(0, _RK + "/library") + +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 = {} + + +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). + """ + 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=3e-2, 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, solve_block=args.sb), + ) + 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): + if with_h0 and not spec.has_initial_state: + continue + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]:") + worst = max(worst, check(spec, B, H, T, gate_low=gate_low)) + + 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 __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..7e33b80c7987 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py @@ -0,0 +1,316 @@ +#!/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 + +_HERE = os.path.dirname(__file__) +_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) +sys.path.insert(0, _RK + "/platform/python") +sys.path.insert(0, _RK + "/library") + +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): + """Launch the prep kernel over ``q/k/g/beta`` already packed by chunk.""" + num_tiles = q.shape[0] + 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, + ) + launcher( + { + "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), + }, + 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..ab9b1ccab7d5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_split.py @@ -0,0 +1,280 @@ +#!/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 os +import sys + +_HERE = os.path.dirname(__file__) +_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) +sys.path.insert(0, _RK + "/platform/python") +sys.path.insert(0, _RK + "/library") + +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 + +import kda_chunk_prep as prep_mod # noqa: E402 +from kda_chunk_fused import make_inputs, ref_token_serial # noqa: E402 + +_LAUNCHER_CACHE: dict = {} + + +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) -> 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. + """ + return KdaChunkPrepSpec( + head_k=spec.head_k, + head_v=spec.head_v, + dtype=spec.dtype, + tile=spec.tile, + ) + + +def run_scan(spec, ws, v, o, ht, bh, nc, h0=None, stream=None): + """Launch the scan over tiles already materialized by the prep kernel.""" + 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, + ) + launcher( + { + "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), + }, + config=cfg, + ) + + +def run_split(spec, q, k, g, beta, v, o, ws, ht, scale, bh, nc, h0=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. + """ + prep_mod.run_prep(prep_spec_of(spec), q, k, g, beta, ws, scale) + run_scan(spec, ws, v, o, ht, bh, nc, 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=3e-2, 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): + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]:") + worst = max(worst, check(spec, B, H, T, gate_low=gate_low)) + + 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 <= 3e-2 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) 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..7829ca54fa0f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/kda_chunkwise.py @@ -0,0 +1,1775 @@ +"""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. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Tuple + +from rocke.core.ir import BF16, F32, I32, IRBuilder, KernelDef, PtrType +from rocke.helpers.atoms import MfmaAtom +from rocke.helpers.spec import SignatureBuilder, kernel_name_join + +# 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} +# gfx950 LDS budget per workgroup. +LDS_LIMIT = 160 * 1024 + + +@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 + # 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 + 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.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" + + @property + def atom(self) -> MfmaAtom: + """The bf16 hero atom. Both C x C products contract over the full head + dim, so the 32x32x16 shape covers one output tile per wave with K + stepping, and ``chunk`` must match its M/N extent.""" + return MfmaAtom.bf16_32x32x16() + + @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: + return kernel_name_join( + self.name, + f"dk{self.head_k}", + f"dv{self.head_v}", + self.dtype, + *self.tile.name_parts(), + ) + + +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)" + + t = spec.tile + atom = spec.atom + if t.chunk != atom.m or t.chunk != atom.n: + return False, ( + f"chunk ({t.chunk}) must equal the MFMA M/N extent " + f"({atom.m}x{atom.n}); the C x C products are one atom tile" + ) + 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})" + ) + # The cumulative sum splits the chunk in half and gives one thread a whole + # (half, channel) column, which is what pins the block to two channels' + # worth of threads. + if t.block_size != 2 * spec.head_k: + return False, ( + f"block_size ({t.block_size}) must be 2*head_k ({2 * spec.head_k}) " + "for the two-pass cumulative sum" + ) + if t.chunk % 2: + return False, f"chunk ({t.chunk}) must be even for the split cumsum" + 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; that requires each vector to + # sit inside one row and to tile the payload exactly. + 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.chunk * spec.head_k) % (8 * t.block_size): + return False, ( + f"C*DK ({t.chunk * spec.head_k}) must be divisible by " + f"8*block_size ({8 * t.block_size}) for the staging sweep" + ) + 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" + + +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): + t = spec.tile + self.b = b + self.spec = spec + self.q_ptr, self.k_ptr, self.g_ptr, self.beta_ptr, self.scale = inputs + 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 + self.ELEM = _DTYPE_IR[spec.dtype] + self.atom = spec.atom + self.N_CD = C * DK + ELEM = self.ELEM + + # ---- LDS ---- + 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.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 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.""" + b = self.b + return b.vec_extract(b.smem_load_vN(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 = b.smem_load_vN( + 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 = b.smem_load_vN(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 + b.smem_store_vN( + smem, + [row, col], + b.vec_pack([b.cast_f32_to(v, self.ELEM) for v in vals_f32], self.ELEM), + 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)): + row = b.add(ew_row, b.const_i32(p * ew_rstep)) + off = b.add(b.mul(tid, b.const_i32(8)), b.const_i32(p * 8 * BLOCK)) + 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_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> 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. + """ + 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 + 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) + q_ptr, k_ptr, g_ptr, beta_ptr = ctx.q_ptr, ctx.k_ptr, ctx.g_ptr, ctx.beta_ptr + tile_cd, tile_c = ch.cd, ch.c + + # ================================================================= + # 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. + for i in range(N_CD // (4 * BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + off = b.mul(vidx, b.const_i32(4)) + row = b.div(off, b.const_i32(DK)) + col = b.mod(off, b.const_i32(DK)) + vec = b.global_load_vN(g_ptr, b.add(tile_cd, off), F32, 4) + b.smem_store_vN(g_lds, [row, col], vec, 4) + + for i in range(N_CD // (8 * BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + off = b.mul(vidx, b.const_i32(8)) + row = b.div(off, b.const_i32(DK)) + col = b.mod(off, b.const_i32(DK)) + gidx = b.add(tile_cd, off) + b.smem_store_vN(k_lds, [row, col], b.global_load_vN(k_ptr, gidx, ELEM, 8), 8) + b.smem_store_vN(q_lds, [row, col], b.global_load_vN(q_ptr, gidx, ELEM, 8), 8) + + with b.scf_if(b.cmp_gt(b.const_i32(C), tid)): + bv = b.global_load_f32(beta_ptr, b.add(tile_c, tid)) + b.smem_store_vN(beta_lds, [tid], bv, 1) + b.sync() + + # ================================================================= + # 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. + d_ch = b.mod(tid, b.const_i32(DK)) + half = b.div(tid, b.const_i32(DK)) + row0 = b.mul(half, b.const_i32(HALF)) + acc = b.const_f32(0.0) + for i in range(HALF): + rr = b.add(row0, b.const_i32(i)) + acc = b.fadd(acc, lds_get(g_lds, [rr, d_ch])) + b.smem_store_vN(g_lds, [rr, d_ch], b.fmul(acc, c_log2e), 1) + b.sync() + + # Second half += first half's total. Two threads per channel, so each takes + # every other row of the upper half. + base_lo = lds_get(g_lds, [b.const_i32(HALF - 1), d_ch]) + for i in range(HALF // 2): + rr = b.add(b.add(b.const_i32(HALF), half), b.const_i32(i * 2)) + cur = lds_get(g_lds, [rr, d_ch]) + b.smem_store_vN(g_lds, [rr, d_ch], b.fadd(cur, base_lo), 1) + b.sync() + + # ================================================================= + # 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) + ] + b.smem_store_vN(gl_lds, [col4], b.vec_pack(gl, F32), 4) + sink.dec(ctx, ch, col4, [ex2(v) for v in gl]) + + # ================================================================= + # 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)): + row = b.add(ew_row, b.const_i32(p * ew_rstep)) + off = b.add(b.mul(tid, b.const_i32(8)), b.const_i32(p * 8 * BLOCK)) + 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() + + def cxc_mfma(a_smem, b_smem): + """One C x C = (C x DK)(C x DK)^T product, K-stepped over the head dim. + + Every wave computes the whole tile. There is only one 32x32 output tile + here, so splitting K across waves would need a cross-wave reduction + through LDS; the product is 32x32x128, small enough that the redundant + issue costs less than the reduction's barriers. + """ + acc = atom.zero_acc(b) + for ks in range(spec.k_steps): + kb = b.add(b.const_i32(ks * atom.k), frag_k_off) + av = b.smem_load_vN(a_smem, lane_m, kb, dtype=ELEM, n=8) + bv = b.smem_load_vN(b_smem, 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() + _emit_gk_gq_pass(ctx, ch, sink) + + # ================================================================= + # 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. + n_vec_zs = (C * C) // 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)) + b.smem_store_vN( + zs_lds, + [b.div(zoff, b.const_i32(C)), b.mod(zoff, b.const_i32(C))], + zero8, + 8, + ) + + with b.scf_if(b.cmp_gt(b.const_i32(64), tid)): + for i in range(atom.c_per_lane): + row, col = atom.lane_to_output(b, lane, 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. + b.smem_store_vN(m_lds, [row, col], tp, 1) + b.smem_store_vN(tb_lds, [row, col], b.cast_f32_to(tp, ELEM), 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). + b.smem_store_vN( + a_lds, + [row, col], + b.select(b.cmp_eq(row, col), bet, b.const_f32(0.0)), + 1, + ) + b.smem_store_vN( + qk_lds, + [row, col], + 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), + ), + 1, + ) + b.sync() + + # ================================================================= + # 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. + BS = spec.tile.solve_block + NB = C // BS + ks_solve = C // 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 = atom.zero_acc(b) + for ks in range(ks_solve): + kb = b.add(b.const_i32(ks * atom.k), frag_k_off) + accu = atom.emit( + b, + b.smem_load_vN(tb_lds, lane_m, kb, dtype=ELEM, n=8), + b.smem_load_vN(zs_lds, 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. + for i in range(bi * 4 * (BS // 8), (bi + 1) * 4 * (BS // 8)): + row, col = atom.lane_to_output(b, lane, i) + b.smem_store_vN( + a_lds, + [row, col], + b.fsub(lds_get(a_lds, [row, col]), b.vec_extract(accu, i)), + 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) + b.smem_store_vN(a_lds, [cr, lane], val, 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], + ) + b.sync() + + # ================================================================= + # 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) + + # ================================================================= + # 8. Kt = (K * gamma_C / Gamma)^T + # ================================================================= + # Transposed output: a thread owns eight consecutive chunk rows at one + # channel, so the eight LDS reads are strided but the global store is a + # single transaction -- the orientation the scan wants as an MFMA operand. + for i in range(N_CD // (8 * BLOCK)): + vidx = b.add(tid, b.const_i32(i * BLOCK)) + 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 = lds_get(gl_lds, [dch]) + vals = [] + for j in range(8): + rj = b.add(r8, b.const_i32(j)) + kv = b.cast_to_f32(lds_get(k_lds, [rj, dch], dtype=ELEM)) + gc = lds_get(g_lds, [rj, dch]) + vals.append(b.cast_f32_to(b.fmul(kv, ex2(b.fsub(glc, gc))), ELEM)) + sink.kt(ctx, ch, off, dch, r8, vals) + + +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] + g: ptr, # [NT, C * DK] per-channel log decay + beta: ptr, # [NT, C] + a_out, gk_out, gq_out, aqk_out, kt_out: ptr, + dec_out: ptr, # [NT, DK] + scale: f32) + + 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] + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.tile.block_size + if spec.tile.waves_per_eu: + b.kernel.attrs["waves_per_eu"] = ( + spec.tile.waves_per_eu, + spec.tile.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) + 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) + + ctx = _ChunkCtx(b, spec, (q_ptr, k_ptr, g_ptr, beta_ptr, scale)) + 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" + tile: KdaTileSpec = KdaTileSpec() + has_initial_state: bool = False + store_final_state: bool = True + 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 state_tiles(self) -> int: + """Atom tiles per wave across the state's ``DK`` extent. + + Each wave owns one ``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.atom.n + + def lds_bytes(self) -> int: + """Prep's tiles plus the three the scan adds. + + GK/GQ land in the X/XQ staging tiles and A/Aqk in the two bf16 C x C + solve tiles, all of which are dead by the time the scan phase runs, so + the fusion only pays for the transposed K tile, the bf16 mirror of the + state, and the ``V~`` tile. + """ + t = self.tile + C, DK, EV = t.chunk, self.head_k, self.head_v + 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",) + 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 + + atom = spec.atom + t = spec.tile + # Every product in the scan body is partitioned by giving each wave one + # atom-row band of the state, 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 spec.head_v != atom.m * waves: + return False, ( + f"head_v ({spec.head_v}) must equal atom.m * waves " + f"({atom.m} * {waves} = {atom.m * waves}); each wave owns one " + "row band of the state" + ) + if spec.head_k % atom.n: + return False, ( + f"head_k ({spec.head_k}) must be a multiple of the atom N extent " + f"({atom.n})" + ) + if t.chunk % atom.k: + return False, ( + f"chunk ({t.chunk}) must be a multiple of the MFMA K step " + f"({atom.k}); it is the contraction extent of three of the products" + ) + lds = spec.lds_bytes() + if lds > LDS_LIMIT: + return False, f"LDS request {lds} B exceeds the {LDS_LIMIT} B budget" + 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): + b = ctx.b + b.smem_store_vN(self.kt_lds, [dch, r8], b.vec_pack(vals, ctx.ELEM), 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, + lane_m, + frag_k_off, + dec_is_log: bool, + ex2=None, + ): + self.b = b + self.ex2 = ex2 + self.atom = atom + self.C, self.DK, self.EV = chunk, head_k, head_v + 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 + self.tid, self.lane, self.lane_m = tid, lane, lane_m + self.frag_k_off = frag_k_off + self.dec_is_log = dec_is_log + self.NS = head_k // atom.n + self.KS_DK, self.KS_C = head_k // atom.k, chunk // atom.k + self.CPL = atom.c_per_lane + self.N_CV = chunk * head_v + # 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 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 + for ks in range(ksteps): + kb = b.add(b.const_i32(ks * self.atom.k), self.frag_k_off) + av = b.smem_load_vN(a_smem, a_row, kb, dtype=self.ELEM, n=8) + bv = b.smem_load_vN(b_smem, b_row, kb, dtype=self.ELEM, n=8) + acc = self.atom.emit(b, av, bv, acc) + return acc + + 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) + b.smem_store_vN( + self.stb_lds, + [ + b.add(self.wrow, row), + b.add(b.const_i32(ti * self.atom.n), col), + ], + b.cast_f32_to(b.vec_extract(state[ti], i), self.ELEM), + 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): + """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. + """ + b, atom = sc.b, sc.atom + lane_m, wrow, CPL = sc.lane_m, sc.wrow, sc.CPL + ELEM = sc.ELEM + cv_base = b.mul(tile, b.const_i32(sc.N_CV)) + + # ---- Z^T = S^T GK^T, then R^T = V^T - Z^T into the V~ tile ---------- + acc_z = sc.gemm( + sc.stb_lds, b.add(wrow, lane_m), sc.gk_lds, lane_m, sc.KS_DK, atom.zero_acc(b) + ) + # A lane's slots are four runs of four consecutive v channels at one fixed + # chunk row, so V arrives in four 4-wide loads and the residual never needs + # an fp32 staging tile. + for grp in range(CPL // 4): + row0, col = sc.slot(4 * grp) + vvec = b.global_load_vN( + sc.v_ptr, + b.add( + cv_base, + b.add(b.mul(col, b.const_i32(sc.EV)), 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, 4 * grp + j), + ) + b.smem_store_vN( + sc.vn_lds, [b.add(wrow, row), col], b.cast_f32_to(res, ELEM), 1 + ) + b.sync() + + # ---- V~^T = R^T A^T ------------------------------------------------- + acc_v = sc.gemm( + sc.vn_lds, b.add(wrow, lane_m), sc.ab_lds, lane_m, sc.KS_C, atom.zero_acc(b) + ) + b.sync() + for i in range(CPL): + row, col = sc.slot(i) + b.smem_store_vN( + sc.vn_lds, + [b.add(wrow, row), col], + b.cast_f32_to(b.vec_extract(acc_v, i), ELEM), + 1, + ) + b.sync() + + # ---- 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. + acc_o = sc.gemm( + sc.gq_lds, lane_m, sc.stb_lds, b.add(wrow, lane_m), sc.KS_DK, atom.zero_acc(b) + ) + acc_o = sc.gemm(sc.aqb_lds, lane_m, sc.vn_lds, b.add(wrow, lane_m), sc.KS_C, acc_o) + # Straight to HBM, no staging tile: a slot's column index is ``lane % 32``, + # so one store instruction covers 32 consecutive v channels per half-wave + # and is already coalesced. + for i in range(CPL): + row, col = sc.slot(i) + b.global_store_vN( + sc.o_ptr, + b.add( + cv_base, + b.add(b.mul(row, b.const_i32(sc.EV)), b.add(wrow, col)), + ), + 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( + b.smem_load_vN( + 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() + sc.publish_state(new_state) + b.sync() + 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 + PDK, PCB = DK + t.pad_dk, C + t.pad_cb + ELEM = _DTYPE_IR[spec.dtype] + atom = spec.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)) + kt_lds = b.smem_alloc(ELEM, [DK, PCB], "kt_s") + stb_lds = b.smem_alloc(ELEM, [EV, PDK], "stb_s") + vn_lds = b.smem_alloc(ELEM, [EV, PCB], "vn_s") + sink = _LdsTileSink(ctx.x_lds, ctx.xq_lds, ctx.tb_lds, ctx.zs_lds, kt_lds) + + 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, + lane_m=ctx.lane_m, + frag_k_off=ctx.frag_k_off, + dec_is_log=True, + ex2=ctx.ex2, + ) + sink = _LdsTileSink(*sc.tiles_for_sink()) + + bh = b.block_id_x() + state_base = b.mul(bh, b.const_i32(EV * DK)) + s_init = ( + sc.load_state(h0_ptr, state_base) if spec.has_initial_state else sc.zero_state() + ) + sc.publish_state(s_init) + b.sync() + + 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) + _emit_chunk_tiles(ctx, tile, sink) + b.sync() + 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_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 + name: str = "rocke_kda_chunk_scan" + + @property + def atom(self) -> MfmaAtom: + return KdaChunkPrepSpec( + head_k=self.head_k, + head_v=self.head_v, + dtype=self.dtype, + tile=self.tile, + ).atom + + @property + def state_tiles(self) -> int: + """Atom tiles per wave across the state's ``DK`` extent.""" + return self.head_k // self.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 + 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: + t = self.tile + return kernel_name_join( + self.name, + f"dk{self.head_k}", + f"dv{self.head_v}", + self.dtype, + f"c{t.chunk}", + f"b{t.block_size}", + ) + + +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}" + + atom = spec.atom + t = spec.tile + # Same single partitioning rule as the fused scan: one atom-row band of the + # state per wave, covering the v extent exactly. Anything else needs a + # second rule for the same accumulator, which is how cross-wave reductions + # creep in. + waves = t.num_waves + if spec.head_v != atom.m * waves: + return False, ( + f"head_v ({spec.head_v}) must equal atom.m * waves " + f"({atom.m} * {waves} = {atom.m * waves}); each wave owns one " + "row band of the state" + ) + if spec.head_k % atom.n: + return False, ( + f"head_k ({spec.head_k}) must be a multiple of the atom N extent " + f"({atom.n})" + ) + if t.chunk % atom.k: + return False, ( + f"chunk ({t.chunk}) must be a multiple of the MFMA K step " + f"({atom.k}); it is the contraction extent of three of the products" + ) + # 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 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" + ) + + 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 + BLOCK = t.block_size + PDK, PCB = DK + t.pad_dk, C + t.pad_cb + ELEM = _DTYPE_IR[spec.dtype] + atom = spec.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) + + 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, PDK], "stb_s") + vn_lds = b.smem_alloc(ELEM, [EV, PCB], "vn_s") + + tid = b.thread_id_x() + lane = b.mod(tid, b.const_i32(64)) + lane_m = b.mod(lane, b.const_i32(32)) + frag_k_off = b.mul(b.div(lane, b.const_i32(32)), b.const_i32(8)) + + sc = _ScanCtx( + b, + atom=atom, + chunk=C, + head_k=DK, + head_v=EV, + 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, + lane_m=lane_m, + frag_k_off=frag_k_off, + dec_is_log=False, + ) + + bh = b.block_id_x() + state_base = b.mul(bh, b.const_i32(EV * DK)) + s_init = ( + sc.load_state(h0_ptr, state_base) if spec.has_initial_state else sc.zero_state() + ) + sc.publish_state(s_init) + b.sync() + + 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() + 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).""" + return (int(bh), 1, 1) + + +def kda_chunk_scan_signature(spec: KdaChunkScanSpec): + return ( + 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") + .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): + return ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("g_ptr", "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") + .build() + ) + + +__all__ = [ + "EXP2_CLAMP", + "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/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..2f70d1cd8665 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py @@ -0,0 +1,148 @@ +# 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 os +import sys + +import pytest + +# The host builders live under library/builders and import each other by module +# name, the way they do when run as scripts. +sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "builders/gfx950/kda", + ), +) + + +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.""" + 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.""" + 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. + """ + 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}" + + +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 + + import kda_chunk_fused as fused + import kda_chunk_split as split + + from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaChunkScanSpec + + 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(), 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 + + import kda_chunk_split as split + + with pytest.raises(ValueError, match="unsupported spec"): + split.make_launcher(KdaChunkScanSpec(head_v=64)) 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..01931ecff7d9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_spec.py @@ -0,0 +1,173 @@ +# 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_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_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_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 + + 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, which is exactly the state + the fused kernel is stuck in. 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 + + 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 From f236799f34b368424be69ec1fd26ecff5c1ba680 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Tue, 25 Aug 2026 05:28:50 +0000 Subject: [PATCH 2/6] optimize kda fused --- .../builders/gfx950/kda/kda_chunk_fused.py | 10 +- .../library/kernels/gfx950/kda_chunkwise.py | 1462 +++++++++++++---- .../test_kda_chunkwise_gfx950_numeric.py | 191 ++- .../library/tests/test_kda_chunkwise_spec.py | 69 +- 4 files changed, 1384 insertions(+), 348 deletions(-) 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 index 0747f6ce1b61..5fdf1ab1cc3d 100644 --- 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 @@ -255,7 +255,15 @@ def main(): spec = KdaChunkFusedSpec( head_k=args.dk, head_v=args.dv, - tile=KdaTileSpec(chunk=args.chunk, solve_block=args.sb), + 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}") 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 index 7829ca54fa0f..49e25ee59845 100644 --- 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 @@ -55,6 +55,25 @@ 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 @@ -63,7 +82,7 @@ from dataclasses import dataclass from typing import Tuple -from rocke.core.ir import BF16, F32, I32, IRBuilder, KernelDef, PtrType +from rocke.core.ir import BF16, F32, I8, I32, IRBuilder, KernelDef, PtrType from rocke.helpers.atoms import MfmaAtom from rocke.helpers.spec import SignatureBuilder, kernel_name_join @@ -76,6 +95,137 @@ _DTYPE_IR = {"bf16": BF16} # 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) @@ -103,6 +253,10 @@ class KdaTileSpec: # 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 @@ -114,6 +268,15 @@ class KdaTileSpec: # 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 @@ -128,6 +291,12 @@ 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 @@ -150,10 +319,14 @@ class KdaChunkPrepSpec: @property def atom(self) -> MfmaAtom: - """The bf16 hero atom. Both C x C products contract over the full head - dim, so the 32x32x16 shape covers one output tile per wave with K - stepping, and ``chunk`` must match its M/N extent.""" - return MfmaAtom.bf16_32x32x16() + """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: @@ -188,6 +361,19 @@ def kernel_name(self) -> str: ) +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": @@ -196,27 +382,54 @@ def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, s return False, f"unsupported dtype {spec.dtype!r} (bf16 only)" 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 != atom.m or t.chunk != atom.n: + if t.chunk not in (16, 32): + return False, "chunk must be 16 or 32 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"chunk ({t.chunk}) must equal the MFMA M/N extent " - f"({atom.m}x{atom.n}); the C x C products are one atom tile" + 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})" ) - # The cumulative sum splits the chunk in half and gives one thread a whole - # (half, channel) column, which is what pins the block to two channels' - # worth of threads. - if t.block_size != 2 * spec.head_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"block_size ({t.block_size}) must be 2*head_k ({2 * spec.head_k}) " - "for the two-pass cumulative sum" + 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.chunk % 2: - return False, f"chunk ({t.chunk}) must be even for the split cumsum" 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 " @@ -225,15 +438,11 @@ def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, s 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; that requires each vector to - # sit inside one row and to tile the payload exactly. + # 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.chunk * spec.head_k) % (8 * t.block_size): - return False, ( - f"C*DK ({t.chunk * spec.head_k}) must be divisible by " - f"8*block_size ({8 * t.block_size}) for the staging sweep" - ) if t.pad_dk % 8: return False, ( f"pad_dk ({t.pad_dk}) must be a multiple of 8 so the padded row " @@ -263,7 +472,9 @@ class _ChunkCtx: work stays a single reusable emission. """ - def __init__(self, b: IRBuilder, spec: KdaChunkPrepSpec, inputs): + def __init__( + self, b: IRBuilder, spec: KdaChunkPrepSpec, inputs, *, overlay: bool = False + ): t = spec.tile self.b = b self.spec = spec @@ -276,25 +487,34 @@ def __init__(self, b: IRBuilder, spec: KdaChunkPrepSpec, inputs): 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 ---- - 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") + 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)) @@ -306,6 +526,34 @@ def __init__(self, b: IRBuilder, spec: KdaChunkPrepSpec, inputs): 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. @@ -318,34 +566,35 @@ def ex2(self, x): def lds_get(self, smem, idx, dtype=F32): """One scalar LDS read.""" - b = self.b - return b.vec_extract(b.smem_load_vN(smem, *idx, dtype=dtype, n=1), 0) + 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 = b.smem_load_vN( - smem, row, b.add(col, b.const_i32(4 * h)), dtype=F32, n=4 - ) + 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 = b.smem_load_vN(smem, row, col, dtype=self.ELEM, n=8) + 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 - b.smem_store_vN( + _st( + b, smem, - [row, col], - b.vec_pack([b.cast_f32_to(v, self.ELEM) for v in vals_f32], self.ELEM), - 8, + row, + col, + value=b.vec_pack( + [b.cast_f32_to(v, self.ELEM) for v in vals_f32], self.ELEM + ), + n=8, ) @@ -434,32 +683,277 @@ def _emit_gk_gq_pass(ctx: _ChunkCtx, ch: "_ChunkOffsets", sink) -> None: 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)): - row = b.add(ew_row, b.const_i32(p * ew_rstep)) - off = b.add(b.mul(tid, b.const_i32(8)), b.const_i32(p * 8 * BLOCK)) - 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_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: + 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 + staged = [] + + 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)) + staged.append( + ( + ctx.g_lds, + b.div(off, b.const_i32(DK)), + b.mod(off, b.const_i32(DK)), + b.global_load_vN(ctx.g_ptr, b.add(tile_cd, off), F32, 4), + 4, + valid, + ) + ) + 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)) + gidx = b.add(tile_cd, off) + staged.append( + ( + ctx.k_lds, + row, + col, + b.global_load_vN(ctx.k_ptr, gidx, ELEM, 8), + 8, + valid, + ) + ) + staged.append( + ( + ctx.q_lds, + row, + col, + b.global_load_vN(ctx.q_ptr, gidx, ELEM, 8), + 8, + valid, + ) + ) + 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 + + +def _emit_stage_commit(ctx: _ChunkCtx, issued) -> None: + """Write what :func:`_emit_stage_issue` loaded into the staging tiles.""" + b = ctx.b + staged, beta = issued + for lds, row, col, value, n, valid in staged: + with b.scf_if(valid) if valid is not None else nullcontext(): + _st(b, lds, row, col, value=value, n=n) + 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 @@ -471,8 +965,7 @@ def _emit_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: lds_put8 = ctx.lds_put8 ch = _ChunkOffsets(ctx, tile) - q_ptr, k_ptr, g_ptr, beta_ptr = ctx.q_ptr, ctx.k_ptr, ctx.g_ptr, ctx.beta_ptr - tile_cd, tile_c = ch.cd, ch.c + v_issued = None # ================================================================= # 1. stage g / k / q into LDS with 128-bit transactions @@ -480,27 +973,9 @@ def _emit_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: # 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. - for i in range(N_CD // (4 * BLOCK)): - vidx = b.add(tid, b.const_i32(i * BLOCK)) - off = b.mul(vidx, b.const_i32(4)) - row = b.div(off, b.const_i32(DK)) - col = b.mod(off, b.const_i32(DK)) - vec = b.global_load_vN(g_ptr, b.add(tile_cd, off), F32, 4) - b.smem_store_vN(g_lds, [row, col], vec, 4) - - for i in range(N_CD // (8 * BLOCK)): - vidx = b.add(tid, b.const_i32(i * BLOCK)) - off = b.mul(vidx, b.const_i32(8)) - row = b.div(off, b.const_i32(DK)) - col = b.mod(off, b.const_i32(DK)) - gidx = b.add(tile_cd, off) - b.smem_store_vN(k_lds, [row, col], b.global_load_vN(k_ptr, gidx, ELEM, 8), 8) - b.smem_store_vN(q_lds, [row, col], b.global_load_vN(q_ptr, gidx, ELEM, 8), 8) - - with b.scf_if(b.cmp_gt(b.const_i32(C), tid)): - bv = b.global_load_f32(beta_ptr, b.add(tile_c, tid)) - b.smem_store_vN(beta_lds, [tid], bv, 1) - b.sync() + if stage_inputs: + _emit_stage_inputs(ctx, ch) + b.sync_lds_only() # ================================================================= # 2. in-place cumulative log decay, scaled to log2 @@ -508,24 +983,38 @@ def _emit_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: # 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)) - half = b.div(tid, b.const_i32(DK)) - row0 = b.mul(half, b.const_i32(HALF)) + 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(HALF): + for i in range(GROUP): rr = b.add(row0, b.const_i32(i)) acc = b.fadd(acc, lds_get(g_lds, [rr, d_ch])) - b.smem_store_vN(g_lds, [rr, d_ch], b.fmul(acc, c_log2e), 1) - b.sync() - - # Second half += first half's total. Two threads per channel, so each takes - # every other row of the upper half. - base_lo = lds_get(g_lds, [b.const_i32(HALF - 1), d_ch]) - for i in range(HALF // 2): - rr = b.add(b.add(b.const_i32(HALF), half), b.const_i32(i * 2)) - cur = lds_get(g_lds, [rr, d_ch]) - b.smem_store_vN(g_lds, [rr, d_ch], b.fadd(cur, base_lo), 1) - b.sync() + _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 @@ -536,9 +1025,17 @@ def _emit_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: lds_get(g_lds, [b.const_i32(C - 1), b.add(col4, b.const_i32(j))]) for j in range(4) ] - b.smem_store_vN(gl_lds, [col4], b.vec_pack(gl, F32), 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 # ================================================================= @@ -561,48 +1058,76 @@ def _emit_chunk_tiles(ctx: _ChunkCtx, tile, sink) -> None: 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)): - row = b.add(ew_row, b.const_i32(p * ew_rstep)) - off = b.add(b.mul(tid, b.const_i32(8)), b.const_i32(p * 8 * BLOCK)) - 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)) + 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: - 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() + 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 cxc_mfma(a_smem, b_smem): - """One C x C = (C x DK)(C x DK)^T product, K-stepped over the head dim. + def tile_output(i): + row, col = atom.lane_to_output(b, lane, i) + return b.add(panel_row, row), b.add(panel_col, col) - Every wave computes the whole tile. There is only one 32x32 output tile - here, so splitting K across waves would need a cross-wave reduction - through LDS; the product is 32x32x128, small enough that the redundant - issue costs less than the reduction's barriers. - """ + 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), frag_k_off) - av = b.smem_load_vN(a_smem, lane_m, kb, dtype=ELEM, n=8) - bv = b.smem_load_vN(b_smem, lane_m, kb, dtype=ELEM, n=8) + 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 @@ -619,8 +1144,10 @@ def cxc_mfma(a_smem, b_smem): # 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() + 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(.) @@ -629,20 +1156,61 @@ def cxc_mfma(a_smem, b_smem): # 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. - n_vec_zs = (C * C) // 8 + 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)) - b.smem_store_vN( + _st( + b, zs_lds, - [b.div(zoff, b.const_i32(C)), b.mod(zoff, b.const_i32(C))], - zero8, - 8, + 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, + ) - with b.scf_if(b.cmp_gt(b.const_i32(64), tid)): + # 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 = atom.lane_to_output(b, lane, i) + row, col = tile_output(i) bet = lds_get(beta_lds, [row]) tp = b.select( b.cmp_gt(row, col), @@ -652,28 +1220,32 @@ def cxc_mfma(a_smem, b_smem): # 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. - b.smem_store_vN(m_lds, [row, col], tp, 1) - b.smem_store_vN(tb_lds, [row, col], b.cast_f32_to(tp, ELEM), 1) + _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). - b.smem_store_vN( + _st( + b, a_lds, - [row, col], - b.select(b.cmp_eq(row, col), bet, b.const_f32(0.0)), - 1, + row, + col, + value=b.select(b.cmp_eq(row, col), bet, b.const_f32(0.0)), + n=1, ) - b.smem_store_vN( + _st( + b, qk_lds, - [row, col], - b.select( + 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), ), - 1, + n=1, ) - b.sync() + b.sync_lds_only() # ================================================================= # 6. A = (I + T')^-1 Diag(beta) by blocked forward substitution @@ -683,10 +1255,20 @@ def cxc_mfma(a_smem, b_smem): # 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. + # 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 // atom.k + 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: @@ -702,26 +1284,49 @@ def cxc_mfma(a_smem, b_smem): # 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 = atom.zero_acc(b) + accu = solve_atom.zero_acc(b) for ks in range(ks_solve): - kb = b.add(b.const_i32(ks * atom.k), frag_k_off) - accu = atom.emit( + kb = b.add(b.const_i32(ks * solve_atom.k), solve_frag_k_off) + accu = solve_atom.emit( b, - b.smem_load_vN(tb_lds, lane_m, kb, dtype=ELEM, n=8), - b.smem_load_vN(zs_lds, lane_m, kb, dtype=ELEM, n=8), + _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. - for i in range(bi * 4 * (BS // 8), (bi + 1) * 4 * (BS // 8)): - row, col = atom.lane_to_output(b, lane, i) - b.smem_store_vN( - a_lds, - [row, col], - b.fsub(lds_get(a_lds, [row, col]), b.vec_extract(accu, i)), - 1, - ) + 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) @@ -740,7 +1345,7 @@ def cxc_mfma(a_smem, b_smem): ), ) zblk.append(val) - b.smem_store_vN(a_lds, [cr, lane], val, 1) + _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. @@ -752,7 +1357,10 @@ def cxc_mfma(a_smem, b_smem): b.const_i32(bi * BS + h * 8), zblk[h * 8 : h * 8 + 8], ) - b.sync() + 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 @@ -774,26 +1382,8 @@ def store_cxc(src, hook): store_cxc(a_lds, sink.a) store_cxc(qk_lds, sink.aqk) - - # ================================================================= - # 8. Kt = (K * gamma_C / Gamma)^T - # ================================================================= - # Transposed output: a thread owns eight consecutive chunk rows at one - # channel, so the eight LDS reads are strided but the global store is a - # single transaction -- the orientation the scan wants as an MFMA operand. - for i in range(N_CD // (8 * BLOCK)): - vidx = b.add(tid, b.const_i32(i * BLOCK)) - 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 = lds_get(gl_lds, [dch]) - vals = [] - for j in range(8): - rj = b.add(r8, b.const_i32(j)) - kv = b.cast_to_f32(lds_get(k_lds, [rj, dch], dtype=ELEM)) - gc = lds_get(g_lds, [rj, dch]) - vals.append(b.cast_f32_to(b.fmul(kv, ex2(b.fsub(glc, gc))), ELEM)) - sink.kt(ctx, ch, off, dch, r8, vals) + if not overlap_solve: + _emit_kt_all_waves(ctx, ch, sink) def build_kda_chunk_prep(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> KernelDef: @@ -862,9 +1452,29 @@ class KdaChunkFusedSpec: head_k: int = 128 head_v: int = 128 dtype: str = "bf16" - tile: KdaTileSpec = KdaTileSpec() + # 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 @@ -881,26 +1491,32 @@ def prep(self) -> KdaChunkPrepSpec: 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 ``atom.m``-row band of ``S^T`` and the full head + 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.atom.n + return self.head_k // self.scan_atom.n def lds_bytes(self) -> int: - """Prep's tiles plus the three the scan adds. + """Exact fused allocation for the selected LDS schedule. - GK/GQ land in the X/XQ staging tiles and A/Aqk in the two bf16 C x C - solve tiles, all of which are dead by the time the scan phase runs, so - the fusion only pays for the transposed K tile, the bf16 mirror of the - state, and the ``V~`` tile. + 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 @@ -915,6 +1531,10 @@ def kernel_name(self) -> str: 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) @@ -926,32 +1546,56 @@ def is_valid_fused_spec( if not ok: return False, why - atom = spec.atom t = spec.tile - # Every product in the scan body is partitioned by giving each wave one - # atom-row band of the state, 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. + 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 spec.head_v != atom.m * waves: + if head_v != atom.m * waves: return False, ( - f"head_v ({spec.head_v}) must equal atom.m * waves " + 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 spec.head_k % atom.n: + if head_k % atom.n: return False, ( - f"head_k ({spec.head_k}) must be a multiple of the atom N extent " - f"({atom.n})" + f"head_k ({head_k}) must be a multiple of the atom N extent " f"({atom.n})" ) - if t.chunk % atom.k: + if t.chunk < atom.k and t.chunk + t.pad_cb < atom.k: return False, ( - f"chunk ({t.chunk}) must be a multiple of the MFMA K step " - f"({atom.k}); it is the contraction extent of three of the products" + 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" ) - lds = spec.lds_bytes() - if lds > LDS_LIMIT: - return False, f"LDS request {lds} B exceeds the {LDS_LIMIT} B budget" return True, "ok" @@ -991,8 +1635,14 @@ 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): - b = ctx.b - b.smem_store_vN(self.kt_lds, [dch, r8], b.vec_pack(vals, ctx.ELEM), 8) + _st( + ctx.b, + self.kt_lds, + dch, + r8, + value=ctx.b.vec_pack(vals, ctx.ELEM), + n=8, + ) class _ScanCtx: @@ -1027,10 +1677,9 @@ def __init__( o_ptr, tid, lane, - lane_m, - frag_k_off, dec_is_log: bool, ex2=None, + v_lds=None, ): self.b = b self.ex2 = ex2 @@ -1048,13 +1697,30 @@ def __init__( ) = tiles self.stb_lds, self.vn_lds = stb_lds, vn_lds self.v_ptr, self.o_ptr = v_ptr, o_ptr - self.tid, self.lane, self.lane_m = tid, lane, lane_m - self.frag_k_off = frag_k_off + # 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, self.KS_C = head_k // atom.k, chunk // atom.k + 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)) @@ -1071,13 +1737,23 @@ def slot(self, 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 = b.smem_load_vN(a_smem, a_row, kb, dtype=self.ELEM, n=8) - bv = b.smem_load_vN(b_smem, b_row, kb, dtype=self.ELEM, n=8) + 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 @@ -1099,14 +1775,13 @@ def publish_state(self, state): for ti in range(self.NS): for i in range(self.CPL): row, col = self.slot(i) - b.smem_store_vN( + _st( + b, self.stb_lds, - [ - b.add(self.wrow, row), - b.add(b.const_i32(ti * self.atom.n), col), - ], - b.cast_f32_to(b.vec_extract(state[ti], i), self.ELEM), - 1, + 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): @@ -1141,7 +1816,7 @@ def store_state(self, ptr, base, state): ) -def _emit_scan_body(sc: _ScanCtx, state, tile): +def _emit_scan_body(sc: _ScanCtx, state, tile, *, prefetch: "_InputPrefetch" = None): """One chunk of the state recurrence; returns the updated state. .. code-block:: text @@ -1154,78 +1829,160 @@ def _emit_scan_body(sc: _ScanCtx, state, tile): 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 cv_base = b.mul(tile, b.const_i32(sc.N_CV)) + # 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, b.add(wrow, lane_m), sc.gk_lds, lane_m, sc.KS_DK, atom.zero_acc(b) - ) - # A lane's slots are four runs of four consecutive v channels at one fixed - # chunk row, so V arrives in four 4-wide loads and the residual never needs - # an fp32 staging tile. - for grp in range(CPL // 4): - row0, col = sc.slot(4 * grp) - vvec = b.global_load_vN( - sc.v_ptr, - b.add( - cv_base, - b.add(b.mul(col, b.const_i32(sc.EV)), 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, 4 * grp + j), - ) - b.smem_store_vN( - sc.vn_lds, [b.add(wrow, row), col], b.cast_f32_to(res, ELEM), 1 - ) - b.sync() + 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, + b.add( + cv_base, + b.add(b.mul(cg, b.const_i32(sc.EV)), 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 ------------------------------------------------- - acc_v = sc.gemm( - sc.vn_lds, b.add(wrow, lane_m), sc.ab_lds, lane_m, sc.KS_C, atom.zero_acc(b) - ) - b.sync() - for i in range(CPL): - row, col = sc.slot(i) - b.smem_store_vN( - sc.vn_lds, - [b.add(wrow, row), col], - b.cast_f32_to(b.vec_extract(acc_v, i), ELEM), - 1, - ) - b.sync() + # 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. - acc_o = sc.gemm( - sc.gq_lds, lane_m, sc.stb_lds, b.add(wrow, lane_m), sc.KS_DK, atom.zero_acc(b) - ) - acc_o = sc.gemm(sc.aqb_lds, lane_m, sc.vn_lds, b.add(wrow, lane_m), sc.KS_C, acc_o) - # Straight to HBM, no staging tile: a slot's column index is ``lane % 32``, - # so one store instruction covers 32 consecutive v channels per half-wave - # and is already coalesced. - for i in range(CPL): - row, col = sc.slot(i) - b.global_store_vN( - sc.o_ptr, - b.add( - cv_base, - b.add(b.mul(row, b.const_i32(sc.EV)), b.add(wrow, col)), - ), - b.cast_f32_to(b.vec_extract(acc_o, i), ELEM), - 1, + # 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, + b.add( + cv_base, + b.add( + b.mul(cidx(jt, row), b.const_i32(sc.EV)), + b.add(wrow, col), + ), + ), + 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 @@ -1236,7 +1993,8 @@ def _emit_scan_body(sc: _ScanCtx, state, tile): for i in range(CPL): _, col = sc.slot(i) d = b.vec_extract( - b.smem_load_vN( + _ld( + b, sc.dec_lds, b.add(b.const_i32(ti * atom.n), col), dtype=F32, @@ -1256,9 +2014,7 @@ def _emit_scan_body(sc: _ScanCtx, state, tile): b.vec_pack(scaled, F32), ) new_state.append(acc) - b.sync() - sc.publish_state(new_state) - b.sync() + b.sync_lds_only() return new_state @@ -1304,9 +2060,8 @@ def build_kda_chunk_fused(spec: KdaChunkFusedSpec, arch: str = "gfx950") -> Kern t = spec.tile C, DK, EV = t.chunk, spec.head_k, spec.head_v BLOCK = t.block_size - PDK, PCB = DK + t.pad_dk, C + t.pad_cb ELEM = _DTYPE_IR[spec.dtype] - atom = spec.atom + atom = spec.scan_atom b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = BLOCK @@ -1324,12 +2079,18 @@ def build_kda_chunk_fused(spec: KdaChunkFusedSpec, arch: str = "gfx950") -> Kern scale = b.param("scale", F32) nc = b.param("nc", I32) - ctx = _ChunkCtx(b, prep, (q_ptr, k_ptr, g_ptr, beta_ptr, scale)) - kt_lds = b.smem_alloc(ELEM, [DK, PCB], "kt_s") - stb_lds = b.smem_alloc(ELEM, [EV, PDK], "stb_s") - vn_lds = b.smem_alloc(ELEM, [EV, PCB], "vn_s") - sink = _LdsTileSink(ctx.x_lds, ctx.xq_lds, ctx.tb_lds, ctx.zs_lds, kt_lds) - + 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, @@ -1352,20 +2113,29 @@ def build_kda_chunk_fused(spec: KdaChunkFusedSpec, arch: str = "gfx950") -> Kern o_ptr=o_ptr, tid=ctx.tid, lane=ctx.lane, - lane_m=ctx.lane_m, - frag_k_off=ctx.frag_k_off, 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() ) - sc.publish_state(s_init) - b.sync() + + # 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), @@ -1377,9 +2147,35 @@ def build_kda_chunk_fused(spec: KdaChunkFusedSpec, arch: str = "gfx950") -> Kern ) with loop as (n, carried): tile = b.add(b.mul(bh, nc), n) - _emit_chunk_tiles(ctx, tile, sink) - b.sync() - b.scf_yield(*_emit_scan_body(sc, list(carried), tile)) + 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) @@ -1448,10 +2244,15 @@ def atom(self) -> MfmaAtom: tile=self.tile, ).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.atom.n + return self.head_k // self.scan_atom.n def lds_bytes(self) -> int: """The six staged tiles plus the state mirror and ``V~``. @@ -1485,6 +2286,7 @@ def kernel_name(self) -> str: self.dtype, f"c{t.chunk}", f"b{t.block_size}", + *((f"sa{t.scan_atom_m}",) if t.scan_atom_m else ()), ) @@ -1497,29 +2299,10 @@ def is_valid_scan_spec( if spec.dtype not in _DTYPE_IR: return False, f"unsupported dtype {spec.dtype}" - atom = spec.atom t = spec.tile - # Same single partitioning rule as the fused scan: one atom-row band of the - # state per wave, covering the v extent exactly. Anything else needs a - # second rule for the same accumulator, which is how cross-wave reductions - # creep in. - waves = t.num_waves - if spec.head_v != atom.m * waves: - return False, ( - f"head_v ({spec.head_v}) must equal atom.m * waves " - f"({atom.m} * {waves} = {atom.m * waves}); each wave owns one " - "row band of the state" - ) - if spec.head_k % atom.n: - return False, ( - f"head_k ({spec.head_k}) must be a multiple of the atom N extent " - f"({atom.n})" - ) - if t.chunk % atom.k: - return False, ( - f"chunk ({t.chunk}) must be a multiple of the MFMA K step " - f"({atom.k}); it is the contraction extent of three of the products" - ) + ok, why = _check_scan_partition(spec.scan_atom, t, spec.head_k, spec.head_v) + 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. @@ -1582,7 +2365,7 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne BLOCK = t.block_size PDK, PCB = DK + t.pad_dk, C + t.pad_cb ELEM = _DTYPE_IR[spec.dtype] - atom = spec.atom + atom = spec.scan_atom b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = BLOCK @@ -1612,8 +2395,6 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne tid = b.thread_id_x() lane = b.mod(tid, b.const_i32(64)) - lane_m = b.mod(lane, b.const_i32(32)) - frag_k_off = b.mul(b.div(lane, b.const_i32(32)), b.const_i32(8)) sc = _ScanCtx( b, @@ -1630,18 +2411,15 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne o_ptr=o_ptr, tid=tid, lane=lane, - lane_m=lane_m, - frag_k_off=frag_k_off, dec_is_log=False, ) 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() ) - sc.publish_state(s_init) - b.sync() def stage(src, dst, rows, cols, base): """One flat ``rows x cols`` HBM tile into its padded LDS tile. @@ -1697,7 +2475,7 @@ def stage(src, dst, rows, cols, base): ), 4, ) - b.sync() + b.sync_lds_only() b.scf_yield(*_emit_scan_body(sc, list(carried), tile)) if spec.store_final_state: 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 index 2f70d1cd8665..11f83e130cc7 100644 --- 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 @@ -110,6 +110,187 @@ def test_fused_path_matches_token_serial(gate_low): assert worst <= TOL, f"gate {gate_low}: 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. + """ + 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 + + 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 + + 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 + + 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 + + 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.""" + 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. @@ -122,13 +303,19 @@ def test_split_and_fused_agree_bitwise(): import kda_chunk_fused as fused import kda_chunk_split as split - from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaChunkScanSpec + 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(), 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" 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 index 01931ecff7d9..3157a0075ea1 100644 --- 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 @@ -103,6 +103,45 @@ 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. @@ -114,6 +153,28 @@ def test_head_v_must_match_wave_count(self): 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 @@ -128,9 +189,8 @@ 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, which is exactly the state - the fused kernel is stuck in. So the occupancy target is a spec-level - rejection rule here, not a tuning note. + 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 @@ -171,3 +231,6 @@ def test_off_default_knobs_reach_the_name(self): 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 From 5f1b6d916fbf9af2e749ceb5fd193302d9cf2bd0 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Wed, 26 Aug 2026 05:52:59 +0000 Subject: [PATCH 3/6] black formatted --- .../library/builders/gfx950/kda/CASE_STUDY.md | 37 ++ .../builders/gfx950/kda/kda_chunk_prep.py | 67 ++- .../builders/gfx950/kda/kda_chunk_split.py | 272 +++++++++- .../library/kernels/gfx950/kda_chunkwise.py | 470 +++++++++++++++--- .../test_kda_chunkwise_gfx950_numeric.py | 64 +++ .../library/tests/test_kda_chunkwise_spec.py | 43 +- 6 files changed, 830 insertions(+), 123 deletions(-) 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 index 17fe3a3c3daa..5850c044fbbb 100644 --- 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 @@ -225,3 +225,40 @@ 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 `/workspace/kda-ref/bench_apples.py` and +`/workspace/kda-ref/sweep_aligned_split.py` (results recorded off-repo 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/kda_chunk_prep.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_prep.py index 7e33b80c7987..92781d16f4f5 100644 --- 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 @@ -86,9 +86,25 @@ def alloc_tiles(num_tiles: int, spec: KdaChunkPrepSpec, device="cuda"): return ws -def run_prep(spec, q, k, g, beta, ws, scale, stream=None): - """Launch the prep kernel over ``q/k/g/beta`` already packed by chunk.""" - num_tiles = q.shape[0] +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 @@ -97,22 +113,35 @@ def run_prep(spec, q, k, g, beta, ws, scale, stream=None): block=(spec.tile.block_size, 1, 1), stream=stream, ) - launcher( - { - "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), - }, - config=cfg, - ) + 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) # --------------------------------------------------------------------- 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 index ab9b1ccab7d5..c691f922312c 100644 --- 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 @@ -73,22 +73,78 @@ def make_launcher(spec: KdaChunkScanSpec) -> KernelLauncher: return _LAUNCHER_CACHE[key] -def prep_spec_of(spec: KdaChunkScanSpec) -> KdaChunkPrepSpec: +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 = KdaTileSpec( + chunk=spec.tile.chunk, + block_size=256, + pad_dk=spec.tile.pad_dk, + pad_c=spec.tile.pad_c, + pad_cb=spec.tile.pad_cb, + solve_block=spec.tile.solve_block, + tile_atom_m=spec.tile.tile_atom_m, + ) return KdaChunkPrepSpec( head_k=spec.head_k, head_v=spec.head_v, dtype=spec.dtype, - tile=spec.tile, + tile=tile, + **raw_kw, ) -def run_scan(spec, ws, v, o, ht, bh, nc, h0=None, stream=None): +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.""" launcher = make_launcher(spec) if stream is None: @@ -98,33 +154,199 @@ def run_scan(spec, ws, v, o, ht, bh, nc, h0=None, stream=None): block=(spec.tile.block_size, 1, 1), stream=stream, ) - launcher( - { - "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), - }, - config=cfg, - ) - - -def run_split(spec, q, k, g, beta, v, o, ws, ht, scale, bh, nc, h0=None): + 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. """ - prep_mod.run_prep(prep_spec_of(spec), q, k, g, beta, ws, scale) - run_scan(spec, ws, v, o, ht, bh, nc, h0=h0) + 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) # --------------------------------------------------------------------- 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 index 49e25ee59845..db2e79e93b19 100644 --- 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 @@ -83,9 +83,12 @@ 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 @@ -316,6 +319,14 @@ class KdaChunkPrepSpec: 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: @@ -352,13 +363,26 @@ def lds_bytes(self) -> int: ) def kernel_name(self) -> str: - return kernel_name_join( + 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: @@ -381,6 +405,22 @@ def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, s 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, ( @@ -462,6 +502,110 @@ def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, s 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, @@ -478,7 +622,8 @@ def __init__( t = spec.tile self.b = b self.spec = spec - self.q_ptr, self.k_ptr, self.g_ptr, self.beta_ptr, self.scale = inputs + 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 @@ -745,7 +890,9 @@ def _emit_stage_issue(ctx: _ChunkCtx, ch): 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)) @@ -755,16 +902,24 @@ def _emit_stage_issue(ctx: _ChunkCtx, ch): 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)) - staged.append( - ( - ctx.g_lds, - b.div(off, b.const_i32(DK)), - b.mod(off, b.const_i32(DK)), - b.global_load_vN(ctx.g_ptr, b.add(tile_cd, off), F32, 4), - 4, - valid, + 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 @@ -775,39 +930,73 @@ def _emit_stage_issue(ctx: _ChunkCtx, ch): off = b.mul(safe, b.const_i32(8)) row = b.div(off, b.const_i32(DK)) col = b.mod(off, b.const_i32(DK)) - gidx = b.add(tile_cd, off) - staged.append( - ( - ctx.k_lds, - row, - col, - b.global_load_vN(ctx.k_ptr, gidx, ELEM, 8), - 8, - valid, - ) - ) - staged.append( - ( - ctx.q_lds, - row, - col, - b.global_load_vN(ctx.q_ptr, gidx, ELEM, 8), - 8, - valid, - ) - ) - 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 + 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 = issued - for lds, row, col, value, n, valid in staged: + 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(): - _st(b, lds, row, col, value=value, n=n) + 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) @@ -1391,13 +1580,15 @@ def build_kda_chunk_prep(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Kernel Kernel signature:: - (q, k: ptr, # [NT, C * DK] - g: ptr, # [NT, C * DK] per-channel log decay + (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) @@ -1405,17 +1596,19 @@ def build_kda_chunk_prep(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Kernel 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"] = spec.tile.block_size - if spec.tile.waves_per_eu: + b.kernel.attrs["max_workgroup_size"] = t.block_size + if t.waves_per_eu: b.kernel.attrs["waves_per_eu"] = ( - spec.tile.waves_per_eu, - spec.tile.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) + 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) @@ -1425,7 +1618,28 @@ def build_kda_chunk_prep(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Kernel dec_ptr = b.param("dec_ptr", PtrType(F32, "global"), writeonly=True, align=16) scale = b.param("scale", F32) - ctx = _ChunkCtx(b, spec, (q_ptr, k_ptr, g_ptr, beta_ptr, scale)) + 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) @@ -1680,11 +1894,18 @@ def __init__( 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, self.EV = chunk, head_k, head_v + 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 if isinstance(v_row_base, int) else v_row_base + self.io = io self.BLOCK = block_size self.ELEM = elem ( @@ -1724,6 +1945,35 @@ def __init__( # 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. @@ -1839,7 +2089,6 @@ def _emit_scan_body(sc: _ScanCtx, state, tile, *, prefetch: "_InputPrefetch" = N ELEM = sc.ELEM NC_T = sc.NC_T srow = b.add(wrow, lane_m) # this wave's S^T operand row - cv_base = b.mul(tile, b.const_i32(sc.N_CV)) # Issued before anything else in the scan so the whole body covers the HBM # latency; the staging tiles are already dead by here. @@ -1905,10 +2154,7 @@ def cidx(jt, col): else: vvec = b.global_load_vN( sc.v_ptr, - b.add( - cv_base, - b.add(b.mul(cg, b.const_i32(sc.EV)), b.add(wrow, row0)), - ), + sc.v_global(tile, cg, b.add(wrow, row0)), ELEM, 4, ) @@ -1973,13 +2219,7 @@ def cidx(jt, col): row, col = sc.slot(i) b.global_store_vN( sc.o_ptr, - b.add( - cv_base, - b.add( - b.mul(cidx(jt, row), b.const_i32(sc.EV)), - b.add(wrow, col), - ), - ), + sc.o_global(tile, cidx(jt, row), col, wrow), b.cast_f32_to(b.vec_extract(acc_o, i), ELEM), 1, ) @@ -2233,6 +2473,11 @@ class KdaChunkScanSpec: 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 @@ -2265,6 +2510,7 @@ def lds_bytes(self) -> int: """ 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) @@ -2272,22 +2518,28 @@ def lds_bytes(self) -> int: + 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) + + 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: t = self.tile - return kernel_name_join( + parts = [ self.name, f"dk{self.head_k}", f"dv{self.head_v}", self.dtype, f"c{t.chunk}", f"b{t.block_size}", - *((f"sa{t.scan_atom_m}",) if t.scan_atom_m else ()), - ) + ] + if t.scan_atom_m: + parts.append(f"sa{t.scan_atom_m}") + if self.value_splits != 1: + parts.append(f"vs{self.value_splits}") + if self.token_major_io: + parts.append("tm") + return kernel_name_join(*parts) def is_valid_scan_spec( @@ -2299,8 +2551,30 @@ def is_valid_scan_spec( 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, spec.head_v) + 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 @@ -2362,6 +2636,7 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne 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] @@ -2384,24 +2659,47 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne 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, PDK], "stb_s") - vn_lds = b.smem_alloc(ELEM, [EV, PCB], "vn_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, + head_v=ev_slice, block_size=BLOCK, elem=ELEM, tiles=(gk_lds, gq_lds, ab_lds, aqb_lds, kt_lds, dec_lds), @@ -2412,10 +2710,15 @@ def build_kda_chunk_scan(spec: KdaChunkScanSpec, arch: str = "gfx950") -> "Kerne tid=tid, lane=lane, dec_is_log=False, + v_row_base=v_row_base, + ev_stride=EV, + io=io, ) - bh = b.block_id_x() - state_base = b.mul(bh, b.const_i32(EV * DK)) + 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() @@ -2486,12 +2789,12 @@ def stage(src, dst, rows, cols, base): def kda_chunk_scan_grid(spec: KdaChunkScanSpec, bh: int) -> Tuple[int, int, int]: - """One workgroup per (batch, head).""" - return (int(bh), 1, 1) + """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): - return ( + sb = ( SignatureBuilder() .ptr("a_ptr", spec.dtype) .ptr("gk_ptr", spec.dtype) @@ -2504,8 +2807,10 @@ def kda_chunk_scan_signature(spec: KdaChunkScanSpec): .ptr("h0_ptr", "f32") .ptr("ht_ptr", "f32") .scalar("nc", "i32") - .build() ) + 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]: @@ -2514,11 +2819,11 @@ def kda_chunk_prep_grid(spec: KdaChunkPrepSpec, num_tiles: int) -> Tuple[int, in def kda_chunk_prep_signature(spec: KdaChunkPrepSpec): - return ( + sb = ( SignatureBuilder() .ptr("q_ptr", spec.dtype) .ptr("k_ptr", spec.dtype) - .ptr("g_ptr", "f32") + .ptr("g_ptr", "bf16" if spec.raw_inputs else "f32") .ptr("beta_ptr", "f32") .ptr("a_ptr", spec.dtype) .ptr("gk_ptr", spec.dtype) @@ -2527,8 +2832,17 @@ def kda_chunk_prep_signature(spec: KdaChunkPrepSpec): .ptr("kt_ptr", spec.dtype) .ptr("dec_ptr", "f32") .scalar("scale", "f32") - .build() ) + 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__ = [ 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 index 11f83e130cc7..195f781a14fd 100644 --- 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 @@ -333,3 +333,67 @@ def test_scan_rejects_a_spec_it_cannot_emit(): 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 + + 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 + + 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) 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 index 3157a0075ea1..87a6dee58e80 100644 --- 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 @@ -58,6 +58,29 @@ 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 @@ -185,6 +208,24 @@ 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_lds_leaves_room_for_two_workgroups(self): """The split path only earns back its tile traffic at 2 WG/CU. @@ -201,7 +242,7 @@ def test_lds_leaves_room_for_two_workgroups(self): 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 + assert "head_v" in why or "v slice" in why def test_staging_alignment_rejections(self): """Staging is ds_write_b128 throughout, so both pitches stay 8-aligned.""" From 6c6186d4e769bd8941f85e25cd7f0029a0874970 Mon Sep 17 00:00:00 2001 From: Raparti Date: Tue, 1 Sep 2026 19:56:19 -0700 Subject: [PATCH 4/6] address review comments, wire in dispatcher, add golden ir parity tests --- .../gfx950/kda/benchmark_chunkwise.py | 95 +++++ .../library/builders/gfx950/kda/ALGORITHM.md | 70 ++++ .../library/builders/gfx950/kda/CASE_STUDY.md | 13 +- .../library/builders/gfx950/kda/README.md | 62 ++++ .../builders/gfx950/kda/kda_chunk_fused.py | 42 ++- .../builders/gfx950/kda/kda_chunk_prep.py | 10 +- .../builders/gfx950/kda/kda_chunk_split.py | 59 +++- .../rocke/library/dispatch/kda/__init__.py | 8 +- .../rocke/library/dispatch/kda/common.py | 76 ++-- .../rocke/library/dispatch/kda/gfx942.py | 2 +- .../rocke/library/dispatch/kda/gfx950.py | 331 ++++++++++++++++++ .../library/kernels/gfx950/kda_chunkwise.py | 70 +++- .../library/tests/dispatch/kda/__init__.py | 1 + .../tests/dispatch/kda/test_gfx942_wiring.py | 7 +- .../tests/dispatch/kda/test_gfx950_wiring.py | 285 +++++++++++++++ .../test_kda_chunkwise_gfx950_numeric.py | 94 +++-- .../library/tests/test_kda_chunkwise_spec.py | 76 ++++ .../platform/dsl_docs/instances/index.md | 1 + .../rocke/platform/dsl_docs/instances/kda.md | 77 +++- .../python/rocke/instances/SUPPORT_MATRIX.md | 1 + .../rocke_representative_ir_sha256.json | 186 +++++++++- .../instances/rocke_ir_parity_harness.py | 75 +++- .../platform/tests/test_rocke_ci_static.py | 1 + 23 files changed, 1507 insertions(+), 135 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/kda/benchmark_chunkwise.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/ALGORITHM.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/README.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/dispatch/kda/gfx950.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/kda/test_gfx950_wiring.py 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 index 5850c044fbbb..d8278e0ac84d 100644 --- 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 @@ -248,9 +248,16 @@ 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 `/workspace/kda-ref/bench_apples.py` and -`/workspace/kda-ref/sweep_aligned_split.py` (results recorded off-repo per -compliance). +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, 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/kda_chunk_fused.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/kda/kda_chunk_fused.py index 5fdf1ab1cc3d..d20225ddd36d 100644 --- 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 @@ -19,13 +19,16 @@ from __future__ import annotations import argparse +import dataclasses import os import sys -_HERE = os.path.dirname(__file__) -_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) -sys.path.insert(0, _RK + "/platform/python") -sys.path.insert(0, _RK + "/library") +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 @@ -41,6 +44,7 @@ from rocke.runtime import KernelLauncher, LaunchConfig, time_launches # noqa: E402 _LAUNCHER_CACHE: dict = {} +TOL = 3e-2 def make_launcher(spec: KdaChunkFusedSpec) -> KernelLauncher: @@ -70,6 +74,8 @@ def run_fused(spec, q, k, g, beta, v, o, ht, scale, bh, nc, h0=None, stream=None ``[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 @@ -179,7 +185,7 @@ def launch_packed(spec, q, k, v, g, beta, h0=None): 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=3e-2, with_h0=False, verbose=True): +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 @@ -275,16 +281,32 @@ def main(): 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): - if with_h0 and not spec.has_initial_state: - continue - print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]:") - worst = max(worst, check(spec, B, H, T, gate_low=gate_low)) + 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 + return 0 if args.no_check or worst <= TOL else 1 if __name__ == "__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 index 92781d16f4f5..e4f218434478 100644 --- 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 @@ -20,10 +20,12 @@ import os import sys -_HERE = os.path.dirname(__file__) -_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) -sys.path.insert(0, _RK + "/platform/python") -sys.path.insert(0, _RK + "/library") +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 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 index c691f922312c..0473c84e4cb7 100644 --- 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 @@ -24,13 +24,16 @@ from __future__ import annotations import argparse +import dataclasses import os import sys -_HERE = os.path.dirname(__file__) -_RK = os.path.abspath(os.path.join(_HERE, "../../../..")) -sys.path.insert(0, _RK + "/platform/python") -sys.path.insert(0, _RK + "/library") +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 @@ -47,10 +50,15 @@ from rocke.helpers.compile import compile_kernel # noqa: E402 from rocke.runtime import KernelLauncher, LaunchConfig, time_launches # noqa: E402 -import kda_chunk_prep as prep_mod # noqa: E402 -from kda_chunk_fused import make_inputs, ref_token_serial # 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: @@ -94,15 +102,7 @@ def prep_spec_of(spec: KdaChunkScanSpec, *, raw: bool = False) -> KdaChunkPrepSp has_dt_bias=True, lower_bound=-5.0, ) - tile = KdaTileSpec( - chunk=spec.tile.chunk, - block_size=256, - pad_dk=spec.tile.pad_dk, - pad_c=spec.tile.pad_c, - pad_cb=spec.tile.pad_cb, - solve_block=spec.tile.solve_block, - tile_atom_m=spec.tile.tile_atom_m, - ) + tile = dataclasses.replace(spec.tile, block_size=256) return KdaChunkPrepSpec( head_k=spec.head_k, head_v=spec.head_v, @@ -146,6 +146,8 @@ def run_scan( 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 @@ -383,7 +385,7 @@ def launch_packed(spec, q, k, v, g, beta, h0=None): 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=3e-2, with_h0=False, verbose=True): +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 @@ -488,14 +490,33 @@ def main(): 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): - print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]:") - worst = max(worst, check(spec, B, H, T, gate_low=gate_low)) + 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 <= 3e-2 else 1 + return 0 if worst <= TOL else 1 if __name__ == "__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 index db2e79e93b19..03bc9f7ba729 100644 --- 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 @@ -96,6 +96,12 @@ 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 @@ -427,8 +433,11 @@ def is_valid_spec(spec: KdaChunkPrepSpec, arch: str = "gfx950") -> Tuple[bool, s f"unsupported tile_atom_m {t.tile_atom_m}; " f"have {sorted(_SCAN_ATOMS)}" ) atom = spec.atom - if t.chunk not in (16, 32): - return False, "chunk must be 16 or 32 for the emitted tile-atom schedules" + 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) @@ -1904,7 +1913,7 @@ def __init__( 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 if isinstance(v_row_base, int) else v_row_base + self.v_row_base = v_row_base self.io = io self.BLOCK = block_size self.ELEM = elem @@ -2481,13 +2490,23 @@ class KdaChunkScanSpec: name: str = "rocke_kda_chunk_scan" @property - def atom(self) -> MfmaAtom: + 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, - ).atom + ) + + @property + def atom(self) -> MfmaAtom: + return self.prep.atom @property def scan_atom(self) -> MfmaAtom: @@ -2524,22 +2543,17 @@ def lds_bytes(self) -> int: ) def kernel_name(self) -> str: - t = self.tile - parts = [ - self.name, - f"dk{self.head_k}", - f"dv{self.head_v}", - self.dtype, - f"c{t.chunk}", - f"b{t.block_size}", - ] - if t.scan_atom_m: - parts.append(f"sa{t.scan_atom_m}") + 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.append(f"vs{self.value_splits}") + parts += (f"vs{self.value_splits}",) if self.token_major_io: - parts.append("tm") - return kernel_name_join(*parts) + parts += ("tm",) + return kernel_name_join(self.name, *parts) def is_valid_scan_spec( @@ -2590,6 +2604,17 @@ def is_valid_scan_spec( 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), @@ -2605,6 +2630,11 @@ def is_valid_scan_spec( 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 @@ -2847,6 +2877,8 @@ def kda_chunk_prep_signature(spec: KdaChunkPrepSpec): __all__ = [ "EXP2_CLAMP", + "KDA_CHUNK_SIZES", + "KDA_DTYPES", "KdaChunkFusedSpec", "KdaChunkPrepSpec", "KdaChunkScanSpec", 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/test_kda_chunkwise_gfx950_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_chunkwise_gfx950_numeric.py index 195f781a14fd..bdd884f968ff 100644 --- 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 @@ -24,21 +24,8 @@ from __future__ import annotations -import os -import sys - import pytest -# The host builders live under library/builders and import each other by module -# name, the way they do when run as scripts. -sys.path.insert( - 0, - os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "builders/gfx950/kda", - ), -) - def _gpu_ready(): """True only on a gfx950 box with ROCm torch. @@ -73,7 +60,7 @@ def _gpu_ready(): @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.""" - import kda_chunk_prep as prep + from builders.gfx950.kda import kda_chunk_prep as prep from kernels.gfx950.kda_chunkwise import KdaChunkPrepSpec @@ -84,7 +71,7 @@ def test_prep_tiles_match_float64_oracle(gate_low): @pytest.mark.parametrize("gate_low", GATES) def test_split_path_matches_token_serial(gate_low): """rocke prep + rocke scan, against the token-serial recurrence.""" - import kda_chunk_split as split + from builders.gfx950.kda import kda_chunk_split as split from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec @@ -100,7 +87,7 @@ def test_fused_path_matches_token_serial(gate_low): this is the check that routing the tiles through HBM instead of LDS did not change the arithmetic. """ - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec @@ -110,6 +97,27 @@ def test_fused_path_matches_token_serial(gate_low): 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. @@ -120,7 +128,7 @@ def test_subtiled_fused_path_matches_token_serial(gate_low): single-atom output onto a tiled one, so this covers index arithmetic the chunk-wide atom never exercises. """ - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -141,7 +149,7 @@ def test_subtiled_and_chunk_wide_fused_agree_to_one_ulp(): """ import torch - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -186,7 +194,7 @@ def test_input_prefetch_is_bitwise_identical(tile_kw): """ import torch - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -219,7 +227,7 @@ def test_fused_lds_overlay_is_bitwise_identical(): """ import torch - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -246,7 +254,7 @@ def test_c32_tile_phase_16x16_panels_agree_to_one_ulp(): """ import torch - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -274,7 +282,7 @@ def test_c32_tile_phase_16x16_panels_agree_to_one_ulp(): def test_c16_fused_matches_token_serial_oracle(): """Literal C=16 is a valid A/B even though it is not the throughput winner.""" - import kda_chunk_fused as fused + from builders.gfx950.kda import kda_chunk_fused as fused from kernels.gfx950.kda_chunkwise import KdaChunkFusedSpec, KdaTileSpec @@ -300,8 +308,8 @@ def test_split_and_fused_agree_bitwise(): """ import torch - import kda_chunk_fused as fused - import kda_chunk_split as split + 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, @@ -329,7 +337,7 @@ def test_scan_rejects_a_spec_it_cannot_emit(): """ from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec - import kda_chunk_split as split + from builders.gfx950.kda import kda_chunk_split as split with pytest.raises(ValueError, match="unsupported spec"): split.make_launcher(KdaChunkScanSpec(head_v=64)) @@ -357,7 +365,7 @@ def test_raw_split_matches_aligned_oracle(gate_low): """Fused raw prep + token-major split scan vs the aligned float64 oracle.""" import torch - import kda_chunk_split as split + from builders.gfx950.kda import kda_chunk_split as split B, H, T = 1, 4, 256 scan, _ = split.aligned_split_specs(1) @@ -383,7 +391,7 @@ 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 - import kda_chunk_split as split + from builders.gfx950.kda import kda_chunk_split as split from kernels.gfx950.kda_chunkwise import KdaChunkScanSpec, KdaTileSpec @@ -397,3 +405,35 @@ def test_value_splits_agree_with_vs1(value_splits, block, scan_atom_m): 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 index 87a6dee58e80..5eea8330ebb1 100644 --- 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 @@ -226,6 +226,35 @@ 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. @@ -244,6 +273,18 @@ def test_head_v_must_match_wave_count(self): 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)): @@ -275,3 +316,38 @@ def test_off_default_knobs_reach_the_name(self): 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/platform/dsl_docs/instances/index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md index 1bca9f176edd..214b529e6118 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 @@ -79,6 +79,7 @@ Direct grouped 4c (same shape, cpg=kpg=4): | `fmha_fwd_fp8.py` | `FmhaFwdFp8Spec` (CK Tile 01 fp8; per-tensor scales) | `instances/attention.md` | | `sage_attention.py` | `SageAttentionSpec`, `SageQuantMode` (CK Tile 49; 4 variants) | `instances/attention.md` | | `sparse_attention.py` | `JengaSparseSpec`, `VsaSparseSpec` (CK Tile 50; jenga + VSA) | `instances/attention.md` | +| `library/kernels/gfx950/kda_chunkwise.py` | `KdaTileSpec`, `KdaChunkPrepSpec`, `KdaChunkScanSpec`, `KdaChunkFusedSpec` | `instances/kda.md` | Runtime entry point: `run_unified_attention_torch(...)`. 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..9df7abe54e30 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 @@ -88,6 +88,7 @@ These emit generic AMDGPU IR; arch only sets the comgr target triple. | `unified_attention_2d` | ✅ | ✅ | ✅ | scalar (no matrix core) | | `unified_attention_3d` | ✅ | ✅ | ✅ | scalar | | `unified_attention_reduce` | ✅ | ✅ | ✅ | scalar | +| `kda_chunkwise` | ❌ | ✅ | ❌ | gfx950-only; dispatcher-wired via `dispatch.kda` | --- 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): From 785463e65651225f1ce2b261f5903e8ade7f5853 Mon Sep 17 00:00:00 2001 From: Raparti Date: Thu, 3 Sep 2026 07:24:42 -0700 Subject: [PATCH 5/6] style(rocke): format gfx950 KDA builders Co-authored-by: Cursor --- .../rocke/library/builders/gfx950/kda/kda_chunk_fused.py | 4 +--- .../rocke/library/builders/gfx950/kda/kda_chunk_split.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) 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 index d20225ddd36d..135e56d01d38 100644 --- 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 @@ -287,9 +287,7 @@ def main(): 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}:" - ) + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]{state_label}:") worst = max( worst, check( 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 index 0473c84e4cb7..889f82a8d78e 100644 --- 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 @@ -497,9 +497,7 @@ def main(): 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}:" - ) + print(f"parity B={B} H={H} T={T} gate in [{gate_low}, 0]{state_label}:") worst = max( worst, check( From 1dd1a44a4e512d74b21d3a1bbe057ef7b44bdc1b Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Thu, 3 Sep 2026 14:23:08 -0700 Subject: [PATCH 6/6] test(rocke): deepen gfx950 KDA golden coverage Co-authored-by: Cursor --- .../tests/golden/kda_gfx950_ir_sha256.json | 233 +++++++++++++++++ .../library/tests/test_kda_gfx950_golden.py | 244 ++++++++++++++++++ .../platform/dsl_docs/instances/index.md | 4 +- .../python/rocke/instances/SUPPORT_MATRIX.md | 24 +- 4 files changed, 490 insertions(+), 15 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/golden/kda_gfx950_ir_sha256.json create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_kda_gfx950_golden.py 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_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 214b529e6118..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 @@ -79,7 +79,6 @@ Direct grouped 4c (same shape, cpg=kpg=4): | `fmha_fwd_fp8.py` | `FmhaFwdFp8Spec` (CK Tile 01 fp8; per-tensor scales) | `instances/attention.md` | | `sage_attention.py` | `SageAttentionSpec`, `SageQuantMode` (CK Tile 49; 4 variants) | `instances/attention.md` | | `sparse_attention.py` | `JengaSparseSpec`, `VsaSparseSpec` (CK Tile 50; jenga + VSA) | `instances/attention.md` | -| `library/kernels/gfx950/kda_chunkwise.py` | `KdaTileSpec`, `KdaChunkPrepSpec`, `KdaChunkScanSpec`, `KdaChunkFusedSpec` | `instances/kda.md` | Runtime entry point: `run_unified_attention_torch(...)`. @@ -97,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/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index 9df7abe54e30..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 @@ -88,7 +88,6 @@ These emit generic AMDGPU IR; arch only sets the comgr target triple. | `unified_attention_2d` | ✅ | ✅ | ✅ | scalar (no matrix core) | | `unified_attention_3d` | ✅ | ✅ | ✅ | scalar | | `unified_attention_reduce` | ✅ | ✅ | ✅ | scalar | -| `kda_chunkwise` | ❌ | ✅ | ❌ | gfx950-only; dispatcher-wired via `dispatch.kda` | --- @@ -99,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 | --- @@ -151,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 ❌