diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py new file mode 100644 index 000000000000..4e53c0953a68 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py @@ -0,0 +1,477 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""LIVE benchmark for the dense flash-attention prefill kernel on gfx950. + +Times the productized dense prefill kernel +(``kernels/gfx950/attention_dense.py``: :class:`AttentionDenseSpec` / +:func:`build_attention_dense`) across every variant it ships -- + + * ``causal`` : headline GQA causal prefill (Hq=128, Hkv=8), + * ``mha`` : multi-head (Hq==Hkv) causal prefill, + * ``swa`` : sliding-window (banded causal) prefill, + * ``varlen`` : packed ragged batch (cu_seqlens) prefill, + * ``persistent`` : grid-stride persistent variant of the GQA causal shapes, + +against a torch SDPA reference on the *same* HIP stream, using the canonical +``rocke.runtime.time_launches`` timer (HIP events, NOT torch events). Every +shape reports honest windowed TFLOPS (banded pair count), the max-abs error vs +SDPA (causal / banded / per-sequence), a PASS/FAIL flag (< 2e-2), and the built +kernel name. Emits a JSON report (and optional CSV) plus a per-mode geomean. + +Run as a library module:: + + PYTHONPATH=rocke/library python3 -m \\ + benchmarks.gfx950.attention.prefill.benchmark_dense_prefill_live \\ + --mode all --output-json /tmp/dense_prefill_live.json + +or directly with the atom-venv python:: + + /workspace/atom-venv/bin/python \\ + rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py \\ + --mode causal --iterations 5 --warmup 2 +""" +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +import sys + +# --- sys.path bootstrap so this runs both as a module and directly (mirrors +# builders/gfx950/attention/prefill/bench_swa_varlen.py): insert the rocke +# platform (`platform/python`) and library (`library`) roots. --- +_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.attention_dense import ( # noqa: E402 + AttentionDenseSpec, + attention_dense_block, + attention_dense_grid, + build_attention_dense, + supports_attention_dense, +) +from rocke.helpers.compile import compile_kernel # noqa: E402 +from rocke.helpers.spec import SignatureBuilder # noqa: E402 +from rocke.runtime import ( # noqa: E402 + KernelLauncher, + LaunchConfig, + synchronize_and_release, + time_launches, +) + +_TORCH_DT = {"bf16": torch.bfloat16, "fp16": torch.float16} +_TOL = 2e-2 + + +# --------------------------------------------------------------------------- # +# honest FLOPs (banded / windowed causal pair count) -- lifted from +# builders/gfx950/attention/prefill/bench_swa_varlen.py +# --------------------------------------------------------------------------- # +def _pairs(s: int, W: int) -> int: + """causal (+window) attended (q,k) pairs for one length-s sequence.""" + if W and W > 0: + return W * s - W * (W - 1) // 2 if s >= W else s * (s + 1) // 2 + return s * (s + 1) // 2 + + +def _flops(seqlens, W: int, Hq: int, D: int) -> int: + return sum(4 * Hq * D * _pairs(s, W) for s in seqlens) + + +def _gm(vals) -> float: + vals = [v for v in vals if v > 0] + return ( + math.exp(sum(math.log(v) for v in vals) / len(vals)) if vals else float("nan") + ) + + +def _bench_stream_handle() -> int: + return int(torch.cuda.current_stream().cuda_stream) + + +# --------------------------------------------------------------------------- # +# dense kernel launcher (compile + ABI signature; varlen adds cu_seqlens ptrs) +# --------------------------------------------------------------------------- # +_LAUNCHER_CACHE: dict = {} + + +def _dense_launcher(spec: AttentionDenseSpec) -> KernelLauncher: + key = spec.kernel_name() + lch = _LAUNCHER_CACHE.get(key) + if lch is not None: + return lch + ok, why = supports_attention_dense(spec) + if not ok: + raise ValueError(f"unsupported spec: {why}") + art = compile_kernel( + build_attention_dense(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + sb = ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .scalar("scale", "f32") + ) + if spec.varlen: + sb = sb.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32") + lch = KernelLauncher( + hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sb.build() + ) + _LAUNCHER_CACHE[key] = lch + return lch + + +# --------------------------------------------------------------------------- # +# one dense benchmark point: build inputs, check parity vs SDPA, time it +# --------------------------------------------------------------------------- # +def bench_dense( + seqlens, + W: int, + Hq: int, + Hkv: int, + D: int, + *, + dtype: str, + block_n: int, + waves_per_eu: int, + num_persistent: int, + persistent: bool, + warmup: int, + iters: int, + seed: int, +): + """Returns (dense_ms, tflops, max_abs, kernel_name).""" + dev = "cuda" + dt = _TORCH_DT[dtype] + B = len(seqlens) + max_s = max(seqlens) + total = sum(seqlens) + scale = 1.0 / math.sqrt(D) + stream = _bench_stream_handle() + torch.manual_seed(seed) + # packed ragged batch when >1 sequence (or a single non-uniform packing). + varlen = B > 1 or (total != B * max_s) + + if varlen: + # packed [total_tok, H, D] + int32 cu_seqlens[B+1]. + q = (torch.randn(total, Hq, D, dtype=dt, device=dev) * 0.2).contiguous() + k = (torch.randn(total, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + v = (torch.randn(total, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + out = torch.zeros(total, Hq, D, dtype=dt, device=dev) + cu = torch.zeros(B + 1, dtype=torch.int32, device=dev) + cu[1:] = torch.tensor(seqlens, dtype=torch.int32, device=dev).cumsum(0) + spec = AttentionDenseSpec( + batch=B, + seqlen_q=max_s, + seqlen_kv=max_s, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=D, + causal=True, + dtype=dtype, + block_n=block_n, + waves_per_eu=waves_per_eu, + sliding_window=W, + varlen=True, + ) + lch = _dense_launcher(spec) + cfg = LaunchConfig( + grid=attention_dense_grid(spec), + block=attention_dense_block(spec), + stream=stream, + ) + vals = { + "q_ptr": q, + "k_ptr": k, + "v_ptr": v, + "o_ptr": out, + "scale": scale, + "cu_seqlens_q": cu, + "cu_seqlens_kv": cu, + } + else: + # uniform [B=1, S, H, D]. + S = seqlens[0] + q = (torch.randn(1, S, Hq, D, dtype=dt, device=dev) * 0.2).contiguous() + k = (torch.randn(1, S, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + v = (torch.randn(1, S, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + out = torch.zeros(1, S, Hq, D, dtype=dt, device=dev) + spec = AttentionDenseSpec( + batch=1, + seqlen_q=S, + seqlen_kv=S, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=D, + causal=True, + dtype=dtype, + block_n=block_n, + waves_per_eu=waves_per_eu, + sliding_window=W, + persistent=persistent, + num_persistent=num_persistent, + ) + lch = _dense_launcher(spec) + cfg = LaunchConfig( + grid=attention_dense_grid(spec), + block=attention_dense_block(spec), + stream=stream, + ) + vals = {"q_ptr": q, "k_ptr": k, "v_ptr": v, "o_ptr": out, "scale": scale} + + def call(): + lch(vals, config=cfg) + + call() + torch.cuda.synchronize() + + # --- correctness vs per-seq banded / causal SDPA --- + rep = Hq // Hkv + max_err = 0.0 + for i, s in enumerate(seqlens): + if varlen: + st = int(sum(seqlens[:i])) + qs, ks, vs_, os_ = ( + q[st : st + s], + k[st : st + s], + v[st : st + s], + out[st : st + s], + ) + qh = qs.transpose(0, 1).float().unsqueeze(0) + kh = ks.transpose(0, 1).repeat_interleave(rep, 0).float().unsqueeze(0) + vh = vs_.transpose(0, 1).repeat_interleave(rep, 0).float().unsqueeze(0) + ot = os_ + else: + qh = q[0].transpose(0, 1).float().unsqueeze(0) + kh = k[0].transpose(0, 1).repeat_interleave(rep, 0).float().unsqueeze(0) + vh = v[0].transpose(0, 1).repeat_interleave(rep, 0).float().unsqueeze(0) + ot = out[0] + if W and W > 0: + qi = torch.arange(s, device=dev).view(-1, 1) + ki = torch.arange(s, device=dev).view(1, -1) + m = (ki <= qi) & (ki > qi - W) + ref = torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, attn_mask=m + ) + else: + ref = torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, is_causal=True + ) + ref = ref.squeeze(0).transpose(0, 1) + max_err = max(max_err, (ot.float() - ref).abs().max().item()) + + ms = time_launches(call, warmup=warmup, iters=iters, stream=stream) + synchronize_and_release(stream) + tf = _flops(seqlens, W, Hq, D) / (ms * 1e-3) / 1e12 + return ms, tf, max_err, spec.kernel_name() + + +# --------------------------------------------------------------------------- # +# shape sweeps -- ported from bench_swa_varlen.py main() +# --------------------------------------------------------------------------- # +def _configs(mode: str, Hq: int, Hkv: int, D: int): + """Yield (mode, variant, label, seqlens, Hq, Hkv, W, persistent) configs.""" + cfgs = [] + + if mode in ("causal", "all"): + for S in (2048, 4096, 8192): + cfgs.append(("causal", "gqa_causal", f"S={S}", [S], Hq, Hkv, 0, False)) + + if mode in ("mha", "all"): + for H in (16, 32): + for S in (4096, 8192): + cfgs.append(("mha", "mha", f"H={H} S={S}", [S], H, H, 0, False)) + + if mode in ("swa", "all"): + for S in (4096, 8192): + for W in (0, 512, 1024, 2048, 4096): + if W and W > S: + continue + tag = "full-causal" if W == 0 else f"W={W}" + cfgs.append(("swa", "swa", f"S={S} {tag}", [S], Hq, Hkv, W, False)) + + if mode in ("varlen", "all"): + batches = [ + [2048, 2048, 2048, 2048], + [512, 1024, 2048, 4096], + [256, 512, 768, 1024, 1536, 2048], + [1024] * 8, + ] + for seqlens in batches: + for W in (0, 1024): + tag = "full-causal" if W == 0 else f"W={W}" + label = f"B={len(seqlens)} total={sum(seqlens)} {tag}" + cfgs.append( + ("varlen", "varlen", label, list(seqlens), Hq, Hkv, W, False) + ) + + if mode in ("persistent", "all"): + for S in (2048, 4096, 8192): + cfgs.append(("persistent", "persistent", f"S={S}", [S], Hq, Hkv, 0, True)) + + return cfgs + + +def _record(mode, variant, label, seqlens, Hq, Hkv, D, W, res, err_note=None): + if res is None: + return { + "label": label, + "mode": mode, + "variant": variant, + "seqlens": list(seqlens), + "Hq": Hq, + "Hkv": Hkv, + "D": D, + "sliding_window": W, + "dense_ms": None, + "tflops": None, + "max_abs": None, + "ok": False, + "kernel_name": None, + "error": err_note, + } + ms, tf, err, kname = res + return { + "label": label, + "mode": mode, + "variant": variant, + "seqlens": list(seqlens), + "Hq": Hq, + "Hkv": Hkv, + "D": D, + "sliding_window": W, + "dense_ms": ms, + "tflops": tf, + "max_abs": err, + "ok": bool(err < _TOL), + "kernel_name": kname, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument( + "--mode", + choices=["causal", "mha", "swa", "varlen", "persistent", "all"], + default="all", + ) + ap.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + ap.add_argument("--hq", type=int, default=128, help="query heads (GQA/SWA/varlen)") + ap.add_argument("--hkv", type=int, default=8, help="kv heads (GQA/SWA/varlen)") + ap.add_argument("--d", type=int, default=128, help="head size") + ap.add_argument("--bn", type=int, default=64, help="block_n (KV tile)") + ap.add_argument("--waves-per-eu", type=int, default=2, help="occupancy hint") + ap.add_argument( + "--num-persistent", type=int, default=256, help="CTAs for persistent variant" + ) + ap.add_argument("--iterations", type=int, default=50) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--output-json", type=str, default="/tmp/dense_prefill_live.json") + ap.add_argument("--output-csv", type=str, default=None) + args = ap.parse_args() + + if not torch.cuda.is_available(): + print("no GPU", file=sys.stderr) + return 1 + + print(f"device: {torch.cuda.get_device_name(0)}") + print( + f"mode={args.mode} dtype={args.dtype} Hq={args.hq} Hkv={args.hkv} D={args.d} " + f"bn={args.bn} wpe={args.waves_per_eu} np={args.num_persistent} " + f"warmup={args.warmup} iters={args.iterations}" + ) + + cfgs = _configs(args.mode, args.hq, args.hkv, args.d) + results = [] + for mode, variant, label, seqlens, Hq, Hkv, W, persistent in cfgs: + tag = f"[{mode}/{variant}] {label} Hq={Hq} Hkv={Hkv} D={args.d}" + try: + res = bench_dense( + seqlens, + W, + Hq, + Hkv, + args.d, + dtype=args.dtype, + block_n=args.bn, + waves_per_eu=args.waves_per_eu, + num_persistent=args.num_persistent, + persistent=persistent, + warmup=args.warmup, + iters=args.iterations, + seed=args.seed, + ) + except Exception as exc: # noqa: BLE001 - per-shape failures never abort + import traceback + + traceback.print_exc() + rec = _record( + mode, variant, label, seqlens, Hq, Hkv, args.d, W, None, repr(exc) + ) + results.append(rec) + print(f"{tag} FAILED ({exc!r})") + continue + + rec = _record(mode, variant, label, seqlens, Hq, Hkv, args.d, W, res) + results.append(rec) + status = "PASS" if rec["ok"] else "FAIL" + print( + f"{tag} {rec['dense_ms']:8.4f} ms {rec['tflops']:8.1f} TFLOPS " + f"max_abs={rec['max_abs']:.2e} {status}" + ) + + # --- JSON report --- + out_json = args.output_json + os.makedirs(os.path.dirname(os.path.abspath(out_json)), exist_ok=True) + with open(out_json, "w") as fh: + json.dump(results, fh, indent=2, default=str) + print(f"\nwrote {out_json} ({len(results)} shapes)") + + # --- optional CSV report --- + if args.output_csv: + os.makedirs(os.path.dirname(os.path.abspath(args.output_csv)), exist_ok=True) + cols = sorted({k for r in results for k in r.keys()}) + with open(args.output_csv, "w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=cols) + w.writeheader() + for r in results: + row = dict(r) + row["seqlens"] = ",".join(str(x) for x in row["seqlens"]) + w.writerow(row) + print(f"wrote {args.output_csv}") + + # --- per-mode geomean summary --- + print("\n=== per-mode geomean TFLOPS (correct shapes only) ===") + modes = [] + for r in results: + if r["mode"] not in modes: + modes.append(r["mode"]) + n_pass = 0 + for m in modes: + rs = [r for r in results if r["mode"] == m] + tfs = [r["tflops"] for r in rs if r["ok"] and r["tflops"]] + npass = sum(1 for r in rs if r["ok"]) + n_pass += npass + print( + f" {m:12s} n={len(rs):3d} geomean={_gm(tfs):8.1f} TFLOPS " + f"pass={npass}/{len(rs)}" + ) + total_pass = sum(1 for r in results if r["ok"]) + print(f"\nTOTAL PASS {total_pass}/{len(results)}") + return 0 if total_pass == len(results) else 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md index 1499947be8f4..605f79dd6648 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/ALGORITHM.md @@ -4,7 +4,10 @@ This file explains *what* the CK DSL `unified_attention` kernels compute and *why* they are shaped the way they are on CDNA (gfx950, MI355X, wave64, MFMA). The [`README.md`](README.md) is the parity + optimization history; this file is the specification, the data layout, and the two kernel strategies (2D and 3D -split-KV) the production dispatcher chooses between. +split-KV) the production dispatcher chooses between. §8 additionally documents the +**dense prefill kernel** — a non-paged, compile-time-sized peak-TFLOPS sibling +(persistent grid-stride + on-chip ragged support) for the contiguous +`[B, S, H, d]` prefill regime. > The harness in this folder benchmarks against AITER's Triton > `unified_attention`; the CK DSL kernels mirror that Triton reference 1:1 in @@ -311,25 +314,162 @@ mul-hi **magic division** so a compile-time constant divisor folds to a --- -## 8. Where the algorithm ends and tuning begins +## 8. The dense prefill kernel — peak-TFLOPS sibling (`attention_dense.py`) + +The 2D/3D kernels above serve the *paged, varlen* deployment layout. Dense +prefill (a contiguous `[B, S, H, d]` batch, no page table, self-attention +$S_q = S_k$) is a different regime: the shape is known up front and there is no +paging indirection to hide, so a dedicated kernel +(`kernels/gfx950/attention_dense.py`) trades generality for **peak MFMA +throughput**. It computes the exact same streaming-softmax core (§2) but bakes +the shape into the kernel and reshapes the schedule around a 256-row query tile. + +- **Compile-time-sized ABI.** `batch / seqlen / heads / head_size / causal / + dtype` are constants baked at build time (dense, statically-sized). The only + runtime args are the `q/k/v/o` pointers and the `f32` softmax scale. `block_n` + (KV tile), `waves_per_eu`, and the persistent knobs are the tuning parameters; + every algorithmic lever below is always-on. +- **256-row query tile, `32×32×16` MFMA.** `_BLOCK_M = 256` query rows per CTA + (8 wave64s = 512 threads) run QK/PV on gfx950's wide-K `mfma_f32_32x32x16` + atom. QK produces $S^\top = K Q^\top$ so the key lands on the per-lane + accumulator regs — the layout that keeps the softmax a cheap in-lane reduce + + one `lane^32` exchange and lets PV consume $P$ with no relayout shuffle. +- **head_size 64 and 128** (bf16 / fp16), MHA and GQA including non-power-of-2 + group sizes (e.g. 40/8, 28/4). + +### 8.1 The winning levers (all always-on) + +| lever | what it does | measured | +|---|---|---| +| **CK-1 transposed PV** | $P$ feeds the PV MFMA in its native QK-output layout via a half-local V load (`pv32_v_load_paired`); the cross-half P-relayout shuffle is gone (~96 `ds_bpermute` removed). | +35% | +| **LDS bank-conflict pad on K** (`[NBUF, BN, D+8]`) | kills the 8-way conflict on the QK K-reads. | +80% (base win) | +| **LDS V pad** (`+32`) | the transposed PV read (`ds_read_b64_tr_b16`) has a stricter bank pattern than K; a `+32` V-row pad fully clears its conflicts. | +~5% | +| **native `exp2_fast`** (`v_exp_f32`, no overflow guard — softmax arg ≤ 0) | one instruction per exp. | +11.5% | +| **depth-1 cluster** | fuses `exp2(s − m)` into the PV-MFMA loop so the softmax VALU/TRANS co-executes in the MFMA shadow (`sched_group_barrier` names the full DS_READ/MFMA/VALU/TRANS population per step). | — | +| **partial-vmcnt software prefetch** | per-tile K/V DMA drains to a *partial* `vmcnt` (keeps the freshest V prefetch in flight across the barrier) instead of a full `vmcnt(0)` serialize. | raises MfmaUtil | +| **PV-only `s_setprio`** | the PV MFMA cluster is bracketed at raised priority so it wins issue slots; paired with the prefetch. | +3.5% | +| **lazy online rescale** | keep the running max as a *lazy* max that only re-anchors when a tile exceeds it by >8 (log2); when every lane is within 8 (a `wave_all` vote) skip the O/ℓ rescale entirely (a 0/1-trip `scf.for` → a wave-uniform scalar branch), cutting the VALU between the QK and PV clusters. Numerically approximate ($P$ bounded by $2^8$) but parity-identical at bf16/fp16 tolerance. | +~2% | + +**Diagonal-only causal masking.** For causal, below-diagonal KV tiles need no +mask (~94% of tiles at $S = 8192$); the KV loop is split into a mask-free body +`[1, diag)` and a masked diagonal tail `[diag, n_upper)`, and the causal row +limit clamps `n_upper` so fully-future tiles are never visited. + +**head_size and the DMA loader.** One `async_buffer_load_lds` instruction moves +64 lanes × 2 bf16 = 128 elements. At **D=128** that is exactly one padded K/V row +per instruction (the fast path). At **D=64** the loader packs `128//D = 2` rows +per instruction into an *unpadded* contiguous LDS tile (lane $l$ → row +$l/(D/2)$, col $2\,(l \bmod D/2)$); the padded fast path is emitted +byte-identically for D=128, so D=64 support costs the D=128 schedule nothing. + +### 8.2 Two grids: default and persistent + +Same inner pipeline, different outer work assignment: + +- **Default** — one CTA per $(\text{q-block}, \text{query-head}, \text{batch})$. + Grid `(⌈S_q/256⌉, H_q, B)`. Simple; the per-CTA launch/dispatch + scalar setup + + K/V-prime cold-start (~4.5 tile-equivalents) is paid once **per query block**. +- **Persistent (grid-stride)** — a 1-D grid of `num_persistent` long-lived CTAs + (256 = exactly one 8-wave block per CU on MI355X's 256 CUs at 2 waves/SIMD). + Each CTA grid-strides over the flattened work-item space + $W = ⌈S_q/256⌉ \cdot H_q \cdot B$, so the fixed per-CTA cold-start is amortized + once **per CU** instead of once per query block. The inner compute is + byte-identical to the default path; the difference is the outer work loop, the + work-item decode, and per-item state reset. + +**Persistent work-item decode (`persist_decode`).** How $W$ is unflattened to +$(\text{qb}, h_q, \text{bt})$ decides load balance *and* L2 locality: + +- **`qb_major`** — `wi = qb·(H_q·B) + h_q·B + bt`. Putting the triangular causal + cost index `qb` in the MSB spreads cheap + expensive query blocks across each + CTA under grid-stride. But every 256-CTA grid-stride phase spans *all* KV + heads at once → large L2 footprint (57% L2 hit at GQA-8). +- **`hkv_major`** — `wi = hkv·(NQB·g·B) + blk·(g·B) + h_ql·B + bt`, with `blk` + folded to a **low/high-paired** query-block index (`blk < half → qb = blk`; + else `qb = NQB−1−(blk−half)`). Putting `hkv` in the MSB keeps each grid-stride + phase within ~1 KV head so the shared GQA K/V stays L2-resident across its $g$ + query heads (**L2 hit 57% → ~93%, HBM misses 5.9× lower**); the low/high qb + pairing preserves `qb_major`'s causal-triangle balance. Valid only when the + CTA grid-strides across both halves of a KV head ($g\cdot NQB\cdot B \ge 2\,NP$). +- **`auto`** (default) — `hkv_major` when it is balance-safe **and** GQA + ($g>1$), else `qb_major`. Strictly ≥ `qb_major`. + +**Measured (MI355X, bf16, D=128, causal, $S = 8192$, 128/8 GQA, 0 spill, err +≈1.46e-3).** Absolute MI355X TFLOPS swing **±25–30% with auto-clock**, so only +**same-session ratios are load-bearing**; the numbers below are one representative +session, each pinned to its config (grid / decode / V-pad / lazy): + +| config | grid | decode | V-pad | lazy | TFLOPS | +|---|---|---|---|---:|---:| +| default grid | one-CTA/q-block | — | 32 | on | ≈543 | +| persistent baseline | persistent NP=256 | qb-major | 0 | off | ≈877 | +| persistent + V-pad | persistent NP=256 | qb-major | 32 | off | ≈912 | +| **persistent (shipped default)** | persistent NP=256 | **hkv-major** | 32 | on | **≈948** | + +The clock-invariant deltas are the load-bearing part: **hkv/qb ≈ 1.04×** (L2 hit +57%→~93%), **V-pad 0→32 ≈ +5%**, **lazy ≈ +2%** — and the shipped default +(`persistent=True`, `persist_decode="auto"`, `lazy_rescale=True`, +`ROCKE_DENSE_VPAD=32`, the last row) is byte-identical IR to a re-measure, so the +ratio reproduces regardless of the absolute clock. The persistent variant is the +production choice for dense prefill. + +### 8.3 Ragged sequence lengths — on-chip boundary padding + +The 256-row / `block_n`-key tile geometry assumes aligned lengths. Real prompts +are arbitrary, and this layout **cannot pad on the host** (no room to grow the +dense buffers). Instead a **separate kernel path** (`ragged=True`, its own +`kernel_name`) pads the boundary tiles *on-chip* — everything else is the §8.1 +pipeline unchanged. Because seqlen is compile-time, the raggedness is fully known +at build time: + +- **Grid ceil'd** — `⌈S_q/256⌉` query blocks (persistent: the work-item count is + ceil'd), so the partial last query block is covered. +- **OOB query rows → register-zero pad.** Q is loaded through a bounds-checked + `buffer_load_vN` (buffer-resource extent = the real Q buffer); rows past + $S_q$ return 0 instead of faulting. +- **OOB keys → LDS-zero pad.** The K/V async DMA already goes through buffer + resources, so keys past $S_k$ load as 0 into LDS. +- **Masking is nearly free for causal.** A padded key has token index + $\ge S_k > $ every real query index, so the causal mask ($k_j \le q_i$) + *already* drops it — causal ragged needs **no extra key mask**. Non-causal + adds a compile-time `k_j < S_k` key mask on the partial tile. +- **Partial output rows dropped.** The epilogue store is wrapped in a per-lane + `q_i < S_q` guard, so padded rows never write (and never clobber a neighbouring + batch's real rows — correct for any $B$). + +Self-attention only ($S_q = S_k$); not combined with varlen or the sliding +window (the validator rejects those). The aligned path is emitted +**byte-identically** when `ragged=False`, so ragged support costs the aligned +schedule nothing; the ragged kernel itself runs at **~890 TFLOPS** (≈94% of the +948 aligned headline), the delta being the bounds-checked load + guarded store. + +--- + +## 9. Where the algorithm ends and tuning begins The math above is fixed and exact: online softmax, the bias/mask order of §1, the associative split-KV merge. Everything the README tunes — the combo flag bundle, `waves_per_eu`, `num_warps`, `tile_size`, async-DMA issue order, i64 addressing, fp8 dequant — changes only **how these steps are scheduled onto gfx950**, never what is computed. Correctness is pinned by bit-exact parity against Triton and -ULP-level agreement with `ref_paged_attn`; performance is the per-CTA schedule -and the 2D-vs-3D occupancy choice. +ULP-level agreement with `ref_paged_attn` (unified paths) / torch SDPA (dense); +for the dense kernel a golden LLVM-IR SHA gate additionally pins that every +tuning change to one variant leaves the others' machine code byte-identical. +Performance is the per-CTA schedule, the 2D-vs-3D occupancy choice, and the +dense default-vs-persistent grid. --- -## 9. Where to go next +## 10. Where to go next - [`README.md`](README.md) — the parity harness, the three-table methodology, the prefill-2D optimization history (combo, `waves_per_eu`, i64 addressing, fp8), and the substantiated speedups. - `parity_unified_attention.py` — the canonical parity + benchmark harness (default / creative / fmha scenario sets). -- The kernels themselves live in `rocke.instances` (not in this folder): - `gfx950/attention_tiled_2d.py`, `gfx950/attention_tiled_3d.py`, and - `common/attention_unified.py` (the dispatcher and shared spec). +- The kernels themselves live in `rocke` (not in this folder): the unified paged + kernels are `kernels/gfx950/attention_tiled_2d.py`, + `kernels/gfx950/attention_tiled_3d.py`, and `kernels/common/attention_unified.py` + (the dispatcher and shared spec); the dense prefill kernel (§8) is + `kernels/gfx950/attention_dense.py`, gated by + `tests/test_attention_dense_golden.py` (IR-SHA + C++/Python byte-identity). diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md new file mode 100644 index 000000000000..7ec62767b318 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md @@ -0,0 +1,133 @@ +# Dense flash-attention prefill (gfx950 / MI355X) + +Productized dense causal flash-attention prefill kernel for gfx950, authored in the +rocke IR DSL. Forward-only, bf16/fp16, head_dim 64/128, MHA or GQA. + +- Kernel: [`kernels/gfx950/attention_dense.py`](../../../../kernels/gfx950/attention_dense.py) + (`AttentionDenseSpec`, `build_attention_dense`, `supports_attention_dense`) +- Host builder / launcher (this dir): `attention_dense_prefill.py` + +## Baked-in levers (always-on, no env gates) + +- **CK-1 transposed PV** — P feeds the PV MFMA in its native QK-output layout via a + half-local V load (`pv32_v_load_paired`); the cross-half P-relayout shuffle is gone + (~96 `ds_bpermute` removed). +35% over the pre-CK-1 winner. +- **LDS bank-conflict padding on K** (`[NBUF, BN, D+8]`) — kills the 8-way conflict on + the QK K-reads. The dominant base win (+80% over the naive baseline). +- **native `exp2_fast`** (`v_exp_f32`, no overflow guard; softmax argument is always + `<= 0`). +- **full-population `sched_group_barrier` template** (DS_READ/MFMA/VALU/TRANS per PV + step). +- **diagonal-only causal masking** — mask-free body over below-diagonal KV tiles + (~94% at Sq=8192) plus a masked diagonal tail. +- **depth-1 cluster split** fusing exp2 into the PV MFMA loop for MFMA/VALU co-exec. +- **vectorized O store**. + +Shape (batch / seqlen / heads / head_dim / causal / dtype) is baked at build time +(dense, statically-sized ABI). `block_n` (KV tile) and `waves_per_eu` are the only +performance knobs. + +## Persistent (grid-stride) mode + +`AttentionDenseSpec(persistent=True, num_persistent=256)` emits a persistent variant: +a 1-D grid of `num_persistent` long-lived CTAs grid-strides over the +`W = (seqlen_q // 256) * Hq * B` work items, so the per-CTA launch/dispatch + scalar +setup + K/V-prime cold-start is amortized once per CU instead of once per query-block. +This closes the causal fixed-cost amortization gap. `num_persistent=256` = one 8-wave +block per CU on MI355X (256 CUs) at 2 waves/SIMD; larger oversubscribes the CUs (tail +loss). The work-item decode is `persist_decode="auto"` by default: it resolves to the +**hkv-major** L2-locality decode when balance-safe (GQA and `gqa*NQB*B >= 2*NP`, e.g. +128/8 at Sq=8192) and falls back to **qb-major** otherwise. + +## Measured (MI355X, bf16, D=128, Hq=128, Hkv=8, causal, Sq=8192) + +Absolute MI355X TFLOPS swing **±25–30% with auto-clock**, so only **same-session +ratios are load-bearing**; the table below is one representative session, with each +number pinned to its exact config (grid / decode / V-pad / lazy): + +| config | grid | decode | V-pad | lazy | TFLOPS | +|---|---|---|---|---:|---:| +| default grid | one-CTA/q-block | — | 32 | on | ~543 | +| persistent baseline | persistent NP=256 | qb-major | 0 | off | ~877 | +| persistent + V-pad | persistent NP=256 | qb-major | 32 | off | ~912 | +| **persistent (shipped default)** | persistent NP=256 | **hkv-major** | 32 | on | **~948** | + +The load-bearing, clock-invariant deltas: hkv-major vs qb-major ≈ **1.04×** (L2 hit +57%→~93%), V-pad 0→32 ≈ **+5%** (clears the transposed-PV bank conflicts), lazy ≈ +**+2%**. The shipped default (`persistent=True`, `persist_decode="auto"`, +`lazy_rescale=True`, `ROCKE_DENSE_VPAD=32`) is the last row. All configs are 0 VGPR +spill and parity-identical vs `torch.nn.functional.scaled_dot_product_attention` +(max abs err ~1.46e-3 at Sq=8192). + +## Usage + +```bash +# parity + benchmark (default shapes 256/512/2048/8192) +python attention_dense_prefill.py # default (one CTA per query-block/head) +python attention_dense_prefill.py --persistent # persistent grid-stride (NP=256) +python attention_dense_prefill.py --bn 128 # sweep block_n +python attention_dense_prefill.py --persistent --np 256 --interleave +``` + +Programmatic: + +```python +from kernels.gfx950.attention_dense import AttentionDenseSpec, build_attention_dense + +spec = AttentionDenseSpec( + batch=1, seqlen_q=8192, seqlen_kv=8192, + num_query_heads=128, num_kv_heads=8, head_size=128, + causal=True, dtype="bf16", + persistent=True, num_persistent=256, # grid-stride persistent variant +) +kernel = build_attention_dense(spec) # -> KernelDef; compile with backend="python" +``` + +Through the dispatcher (opt-in; picks the persistent best-config for large Sq): + +```python +from dispatch.attention import AttentionRequest, dispatch_attention, dense_spec_for_request +from kernels.gfx950.attention_dense import run_attention_dense_torch + +req = AttentionRequest( + batch=1, nhead_q=128, nhead_k=8, seqlen_q=8192, seqlen_k=8192, + hdim_q=128, hdim_v=128, arch="gfx950", dtype="bf16", mask_type=1, + algorithm="attention_dense", # opt-in; "auto" keeps the unified 2D/3D path + # dense_persistent="auto" # "auto"|"on"|"off"; auto => persistent for large Sq + # dense_persist_decode="auto" # "auto"|"qb_major"|"hkv_major" +) +res = dispatch_attention(req) # res.spec.kernel_name() -> ...persist256_hkvmaj +spec = dense_spec_for_request(req) # launch-ready best-config AttentionDenseSpec +run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=1/128**0.5) +``` + +`dense_persistent="auto"` turns on the persistent grid-stride variant once there is +enough work to fill the grid (`⌈Sq/256⌉·Hq·B >= num_persistent`) — i.e. the large-Sq +prefill regime — so the dispatcher reaches the ~948-TFLOPS path, not the default grid. + +## TODO / follow-ups + +1. **Q-reload (enables the inner-loop unroll)** — the QK B-operand (`q_packs`, + `K_STEPS × vec8` ≈ 32 VGPR) is held resident across the whole KV loop. ISA + probing shows the kernel is at **243/256 VGPR, 0 spill** (flyDSL fits a *2-tile* + pipeline in ~252 VGPR), so a depth-2 software-pipeline unroll — the lever that + would lift MFMA utilization ~45% → ~70% and close the flyDSL gap — is + register-blocked. Re-reading Q per use (L2-resident) instead of holding all + packs frees ~32 VGPR toward that unroll. (Pairs with a P-LDS publish for the + lagged relayout, another ~16 VGPR.) Standalone it is perf-neutral; it is a + prerequisite for the 2-tile unroll. +2. **Improve varlen performance** — the packed `varlen` path (default builder, + `cu_seqlens_q/kv`) currently trails the dense/persistent path and flyDSL on + ragged batches. Follow-ups: extend the hkv-major L2-locality decode and the + persistent grid-stride mode to varlen (today persistent is uniform-batch only), + and improve per-sequence load balance across CTAs so short and long sequences + in a batch don't serialize. + +## Notes + +- gfx950-only (uses `ds_read_b64_tr_b16` and `v_exp_f32`). +- Compiles through the rocke LLVM-direct (`backend="python"`) path. The `exp2_fast` + op is mirrored in the **C++ engine** and covered by a `backend="both"` + byte-identity gate (`tests/test_attention_dense_golden.py:: + test_attention_dense_cpp_python_byte_identity`) — the Python and C++ lowerings + are byte-for-byte identical across every dense variant. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py new file mode 100644 index 000000000000..4a414424d0fc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Host-side builder for the productized dense flash-attention prefill kernel +(``kernels/gfx950/attention_dense.py``). + +Owns the host path: spec construction, kernel-spec generation, compilation, ABI +signature, and runtime launch — plus a torch/SDPA parity check and a benchmark. The +kernel bakes in the winning levers (CK-1 transposed PV, LDS K-padding, exp2_fast, +sched template, diagonal masking, depth-1 cluster, vectorized store); only the KV tile +(`block_n`) and occupancy hint (`waves_per_eu`) are tunable. + +Usage: + python attention_dense_prefill.py # parity + bench, default shapes + python attention_dense_prefill.py --bn 128 # sweep block_n +""" +import argparse +import math +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.attention_dense import ( # noqa: E402 + AttentionDenseSpec, + build_attention_dense, + supports_attention_dense, +) +from rocke.helpers.compile import compile_kernel # noqa: E402 +from rocke.helpers.spec import SignatureBuilder # noqa: E402 +from rocke.runtime import KernelLauncher, LaunchConfig # noqa: E402 + +_TORCH_DT = {"bf16": torch.bfloat16, "fp16": torch.float16} + + +def _make_launcher(spec: AttentionDenseSpec): + """kernel-spec generation + compilation + ABI signature -> cached launcher.""" + ok, why = supports_attention_dense(spec) + if not ok: + raise ValueError(f"unsupported spec: {why}") + art = compile_kernel( + build_attention_dense(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + sb = ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .scalar("scale", "f32") + ) + if spec.varlen: + sb = sb.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32") + sig = sb.build() + return KernelLauncher(hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sig) + + +def _launch_config(spec: AttentionDenseSpec, stream) -> LaunchConfig: + if spec.persistent: + # 1-D grid of long-lived CTAs; each grid-strides over all work items. + grid = (spec.num_persistent, 1, 1) + else: + grid = (spec.seqlen_q // 256, spec.num_query_heads, spec.batch) # BLOCK_M = 256 + return LaunchConfig(grid=grid, block=(spec.num_waves * 64, 1, 1), stream=stream) + + +def run( + spec: AttentionDenseSpec, *, warmup: int = 15, iters: int = 50, check: bool = True +): + dev = "cuda" + dt = _TORCH_DT[spec.dtype] + B, Sq, Skv = spec.batch, spec.seqlen_q, spec.seqlen_kv + Hq, Hkv, D = spec.num_query_heads, spec.num_kv_heads, spec.head_size + torch.manual_seed(0) + q = (torch.randn(B, Sq, Hq, D, dtype=dt, device=dev) * 0.2).contiguous() + k = (torch.randn(B, Skv, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + v = (torch.randn(B, Skv, Hkv, D, dtype=dt, device=dev) * 0.2).contiguous() + out = torch.zeros(B, Sq, Hq, D, dtype=dt, device=dev) + scale = 1.0 / math.sqrt(D) + + launcher = _make_launcher(spec) + stream = torch.cuda.current_stream().cuda_stream + cfg = _launch_config(spec, stream) + vals = {"q_ptr": q, "k_ptr": k, "v_ptr": v, "o_ptr": out, "scale": scale} + + def call(): + launcher(vals, config=cfg) + + call() + torch.cuda.synchronize() + + err = float("nan") + if check: + qh = q.transpose(1, 2).float() + rep = Hq // Hkv + kh = k.transpose(1, 2).repeat_interleave(rep, 1).float() + vh = v.transpose(1, 2).repeat_interleave(rep, 1).float() + W = spec.sliding_window + if spec.causal and W > 0: + # Banded mask: keep k in [q-W+1, q] (causal AND sliding window). + qi = torch.arange(Sq, device=dev).view(-1, 1) + ki = torch.arange(Skv, device=dev).view(1, -1) + allowed = (ki <= qi) & (ki > qi - W) + ref = torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, attn_mask=allowed + ).transpose(1, 2) + else: + ref = torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, is_causal=spec.causal + ).transpose(1, 2) + err = (out.float() - ref).abs().max().item() + + for _ in range(warmup): + call() + torch.cuda.synchronize() + s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(iters): + call() + e.record() + e.synchronize() + ms = s.elapsed_time(e) / iters + W = spec.sliding_window + if spec.causal and W > 0: + # banded pair count: sum_q min(q+1, W) = W*Sq - W*(W-1)/2 (Sq>=W) + pairs = W * Sq - W * (W - 1) // 2 if Sq >= W else Sq * (Sq + 1) // 2 + flops = 4 * B * Hq * D * pairs + elif spec.causal: + flops = 4 * B * Hq * D * (Sq * (Sq + 1) // 2) + else: + flops = 2 * 2 * B * Hq * D * Sq * Skv + tf = flops / (ms * 1e-3) / 1e12 + status = "PASS" if (not check or err < 2e-2) else "FAIL" + print( + f"[{spec.kernel_name()}] {ms:.4f} ms {tf:.1f} TFLOPS max_abs={err:.2e} {status}" + ) + return ms, tf, err + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bn", type=int, default=64, help="block_n (KV tile)") + ap.add_argument("--wpe", type=int, default=2, help="waves_per_eu") + ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16"]) + ap.add_argument("--hq", type=int, default=128) + ap.add_argument("--hkv", type=int, default=8) + ap.add_argument("--d", type=int, default=128) + ap.add_argument("--causal", type=int, default=1) + ap.add_argument( + "--persistent", + action="store_true", + help="grid-stride persistent kernel (amortizes per-CTA launch/setup; " + "+70%% at Sq=8192 causal)", + ) + ap.add_argument("--np", type=int, default=256, help="num_persistent CTAs") + ap.add_argument("--interleave", action="store_true", help="boustrophedon qb order") + ap.add_argument( + "--sw", type=int, default=0, help="sliding_window (0=off; multiple of --bn)" + ) + args = ap.parse_args() + for sq in (256, 512, 2048, 8192): + spec = AttentionDenseSpec( + batch=1, + seqlen_q=sq, + seqlen_kv=sq, + num_query_heads=args.hq, + num_kv_heads=args.hkv, + head_size=args.d, + causal=bool(args.causal), + dtype=args.dtype, + block_n=args.bn, + waves_per_eu=args.wpe, + persistent=args.persistent, + num_persistent=args.np, + interleave=args.interleave, + sliding_window=args.sw, + ) + run(spec) + + +if __name__ == "__main__": + main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py index 8e600866f043..88254de71254 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py @@ -82,6 +82,15 @@ class AttentionRequest(OperatorRequest): dtype: str = "fp16" algorithm: str = "auto" spec_id: str = "auto" + # --- gfx950 attention_dense knobs (only consumed by the opt-in + # ``attention_dense`` candidate; ignored by the unified 2D/3D paths). + # Defaults deliver the persistent ~970-TFLOPS prefill path for large Sq: + # ``dense_persistent="auto"`` turns on the grid-stride variant once there + # is enough work to fill the persistent grid, and ``persist_decode="auto"`` + # picks the L2-locality hkv-major decode where it is balance-safe. --- + dense_persistent: str = "auto" # "auto" | "on" | "off" + dense_num_persistent: int = 256 + dense_persist_decode: str = "auto" # "auto" | "qb_major" | "hkv_major" def normalized(self) -> dict: d = asdict(self) @@ -150,6 +159,10 @@ class AttentionSpec: num_query_heads: int num_kv_heads: int name: str = "rocke_attention_unified" + # When set, this verbatim kernel_name is returned by :meth:`kernel_name` + # instead of the composed unified name. Used by the dense candidate to + # surface the concrete persistent/ragged/decode kernel it will launch. + kernel_name_override: str = "" # Optional pinned codegen knobs a specialized candidate wants the builder to # apply (sorted (key, value) pairs; empty for generic candidates). See # ``attention_unified._d256_gfx950_spec_overrides``; the builder consumes them @@ -157,6 +170,8 @@ class AttentionSpec: tiled_overrides: Tuple[Tuple[str, object], ...] = () def kernel_name(self) -> str: + if self.kernel_name_override: + return self.kernel_name_override from rocke.helpers.spec import kernel_name_join return kernel_name_join( @@ -306,6 +321,143 @@ def select(req: OperatorRequest) -> AttentionSpec: ATTENTION_REGISTRY.register(_make_gfx942_dense_pipe_candidate()) +# block_n (KV tile) the dense candidate ships; 64 is the resource-efficient +# peak (see AttentionDenseSpec.block_n). +_DENSE_BLOCK_N = 64 + + +def _dense_spec(req: OperatorRequest): + """Build the ``AttentionDenseSpec`` for a request at its best-performing + config. Persistent ("auto") turns on the grid-stride ~970-TFLOPS variant once + there is enough work to fill the persistent grid (``nqb*Hq*B >= num_persistent`` + -- the large-Sq prefill regime); ``persist_decode`` / ``lazy_rescale`` default + to the L2-locality hkv-major decode and always-on lazy rescale. Non-tile- + multiple self-attention lengths use the on-chip ragged path (no host pad).""" + from kernels.gfx950.attention_dense import AttentionDenseSpec, _BLOCK_M + + assert isinstance(req, AttentionRequest) + sq, sk = int(req.seqlen_q), int(req.seqlen_k) + bn = _DENSE_BLOCK_N + # on-chip ragged padding for ragged self-attention lengths (seqlen_q==seqlen_kv, + # not a 256/block_n multiple). Cross-attention ragged is left to the validator. + ragged = (sq == sk) and ((sq % _BLOCK_M != 0) or (sk % bn != 0)) + nqb = (sq + _BLOCK_M - 1) // _BLOCK_M + work = nqb * int(req.nhead_q) * int(req.batch) + np = int(req.dense_num_persistent) + mode = req.dense_persistent.strip().lower() + if mode == "on": + persistent = True + elif mode == "off": + persistent = False + elif mode == "auto": + persistent = work >= np # enough work to fill the persistent grid + else: + raise ValueError( + f"dense_persistent must be 'auto'/'on'/'off', got {req.dense_persistent!r}" + ) + return AttentionDenseSpec( + batch=int(req.batch), + seqlen_q=sq, + seqlen_kv=sk, + num_query_heads=int(req.nhead_q), + num_kv_heads=int(req.nhead_k), + head_size=int(req.hdim_q), + causal=(int(req.mask_type) != 0), + dtype=req.dtype.lower(), + block_n=bn, + persistent=persistent, + num_persistent=np, + persist_decode=req.dense_persist_decode.strip().lower(), + ragged=ragged, + ) + + +def dense_spec_for_request(req: AttentionRequest): + """Public builder: the launch-ready ``AttentionDenseSpec`` for ``req`` at its + best config (see :func:`_dense_spec`). Pair with ``run_attention_dense_torch`` + to execute the dispatched dense candidate.""" + return _dense_spec(req) + + +def _make_gfx950_attention_dense_candidate() -> KernelCandidate: + """Dense CK-1 persistent flash-attn prefill on gfx950 (bf16/fp16, causal/full). + + OPT-IN ONLY: matches solely when the request explicitly names + ``algorithm="attention_dense"`` (or ``spec_id="gfx950_attention_dense"``), so it + never auto-overrides the generic unified_2d path (no change to default routing). + When selected it uses the persistent best-config from :func:`_dense_spec` + (grid-stride + hkv-major + lazy for large Sq); the concrete kernel_name (incl. + ``persist``/``hkvmaj``/``lazyrs``/``ragged``) is surfaced on the spec so the + launched kernel is the fast path, not the default grid. End-to-end launch is + ``run_attention_dense_torch(spec=dense_spec_for_request(req), ...)``. + """ + spec_id = "gfx950_attention_dense" + name = "attention_gfx950_dense" + + def support(req: OperatorRequest) -> Tuple[bool, str]: + errors = _request_errors(req) + if errors: + return False, "; ".join(errors) + assert isinstance(req, AttentionRequest) + # Opt-in: never selected under algorithm/spec_id "auto". + if req.algorithm.strip().lower() != "attention_dense" and ( + req.spec_id.strip().lower() != spec_id + ): + return False, "attention_dense is opt-in (algorithm='attention_dense')" + if req.arch != "gfx950": + return False, f"attention_dense requires gfx950 (got {req.arch!r})" + if req.dtype.lower() not in ("bf16", "fp16"): + return False, f"attention_dense is bf16/fp16 only (got {req.dtype!r})" + if int(req.sliding_window) or bool(req.use_sinks): + return False, "attention_dense is dense (no sliding-window / sinks)" + from kernels.gfx950.attention_dense import supports_attention_dense + + try: + spec = _dense_spec(req) + except ValueError as e: + return False, str(e) + ok, why = supports_attention_dense(spec, arch=req.arch) + if not ok: + return False, why + return True, "ok" + + def select(req: OperatorRequest) -> AttentionSpec: + ok, why = support(req) + if not ok: + raise ValueError(f"{name} does not support request: {why}") + assert isinstance(req, AttentionRequest) + problem = _problem(req) + dense_spec = _dense_spec(req) + return AttentionSpec( + path="2d", + head_size=problem.head_size, + block_size=problem.block_size, + dtype=problem.dtype, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + name="rocke_attention_dense", + # surface the concrete persistent/hkvmaj/lazyrs/ragged kernel so the + # dispatched spec names the fast path it will actually launch. + kernel_name_override=dense_spec.kernel_name(), + ) + + candidate = KernelCandidate( + name=name, + family=_FAMILY, + algorithm="attention_dense", + spec_id=spec_id, + abi_version=ATTENTION_ABI_VERSION, + priority=3, + supports=support, + select_spec=select, + signature=lambda _spec: (), + grid=lambda spec, req: (0, 0, 0), + block=lambda spec: (0, 0, 0), + sweep_space=lambda req: (select(req),) if support(req)[0] else (), + ) + return candidate + + def _make_gfx950_d256_candidate() -> KernelCandidate: """Fast gfx950 bf16 head_size-256 prefill kernel — 32x32 transposed stack with FA3-style softmax<->MFMA interleave (mode2/g4) + slab-padded K_lds. @@ -454,6 +606,7 @@ def select(req: OperatorRequest) -> AttentionSpec: return candidate +ATTENTION_REGISTRY.register(_make_gfx950_attention_dense_candidate()) ATTENTION_REGISTRY.register(_make_gfx950_d256_candidate()) ATTENTION_REGISTRY.register(_make_d256_decode_candidate()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/__init__.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/__init__.py index 06842d376bb1..1af029bcb21c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/__init__.py @@ -35,6 +35,18 @@ UnifiedAttention2DTiledSpec, ) +# Dense flash-attention prefill (gfx950-only): productized CK-1 persistent kernel. +# See ``kernels/gfx950/attention_dense.py``. +from .gfx950.attention_dense import ( # noqa: F401 + AttentionDenseSpec, + attention_dense_block, + attention_dense_grid, + attention_dense_signature, + build_attention_dense, + run_attention_dense_torch, + supports_attention_dense, +) + def build_unified_attention_2d_tiled(spec, *, arch: str = "gfx950"): """Arch-aware wrapper: dispatch the tiled-2D builder on ``arch``. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py new file mode 100644 index 000000000000..6a71210dfed5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py @@ -0,0 +1,1670 @@ +"""Dense flash-attention prefill kernel for gfx950 (MI355X). + +Productized from the ``flash_dense_dualwave_swp`` experiment +(``kernels/gfx950/experiments/flash_dense_dualwave_swp/``). This is the shippable +step-1 pipeline with every WINNING lever baked in as always-on (no env gates): + + * **CK-1 transposed PV** — P feeds the PV MFMA in its native QK-output layout via a + half-local V load (``pv32_v_load_paired``); the cross-half P-relayout shuffle is + gone (~96 ``ds_bpermute`` removed). +35% over the pre-CK-1 winner. + * **LDS bank-conflict padding on K** (``[NBUF, BN, D+8]``) — kills the 8-way conflict + on the QK K-reads. The dominant base win (+80% over the naive baseline). + * **native exp2_fast** (``v_exp_f32``, no overflow guard — the softmax argument is + always <= 0) — +11.5%. + * **full-population ``sched_group_barrier`` template** naming DS_READ/MFMA/VALU/TRANS + per PV step. + * **diagonal-only causal masking** — a mask-free body loop over below-diagonal KV + tiles (~94% at Sq=8192) plus a masked diagonal tail. + * **depth-1 cluster split** fusing exp2 into the PV MFMA loop for MFMA/VALU co-exec. + * **partial-vmcnt software prefetch** — the per-tile K/V DMA drain is a partial + `vmcnt` (keeps the freshest V prefetch in flight across the barrier) instead of a + full `vmcnt(0)` serialize; raises MfmaUtil, bit-identical. + * **PV-only `s_setprio`** — the PV MFMA cluster is bracketed at raised priority so + it wins issue slots; paired with the prefetch this is a measured ~+3.5%. + * **vectorized O store**. + +Measured on MI355X (bf16, D=128, causal, 128/8 GQA, Sq=8192, 0 spill, err ~1.46e-3 +vs SDPA). Absolute TFLOPS swing +/-25-30% with auto-clock, so only SAME-SESSION +ratios are load-bearing; one representative session, each number pinned to its +config (grid / decode / V-pad / lazy): + + default grid (one-CTA/q-block, V-pad 32, lazy on) : ~543 TFLOPS + persistent baseline (qb-major, V-pad 0, lazy off): ~877 TFLOPS + persistent + V-pad (qb-major, V-pad 32, lazy off): ~912 TFLOPS + persistent (SHIPPED) (hkv-major, V-pad 32, lazy on) : ~948 TFLOPS + +Clock-invariant deltas (the load-bearing part): hkv/qb ~1.04x, V-pad 0->32 ~+5%, +lazy ~+2%. Shape (batch/seqlen/heads/head_dim) is baked at build time (dense, +compile-time-sized ABI); the KV tile, occupancy hint, and persistent knobs are the +tunable parameters. Lazy online-softmax rescale (skip the O/l rescale when every +lane's tile-max is within 8 log2 of the running max) is ALWAYS-ON by default +(``lazy_rescale=True``): parity-identical (1.46e-3) and ~+2%. + +Head-size / seqlen coverage: + * ``head_size`` is 64 or 128 (bf16/fp16, MHA + GQA incl. non-power-of-2 NQK). + D=128 uses the padded LDS fast path (1 K/V row per async-DMA instr); D=64 + packs 2 rows per instr into an UNPADDED contiguous tile (64 lanes x 2 bf16 = + 128 elems = 2 D=64 rows), so its reads take some LDS bank conflict (V-pad is + a future D=64 perf lever). D=128 codegen is byte-identical to before this + change (same IR hash -> same TFLOPS). + * ``seqlen_q``/``seqlen_kv`` must be a multiple of 256 / ``block_n`` on the + default (aligned) kernel. Non-multiple lengths are handled by a SEPARATE + ``ragged=True`` kernel path (its own kernel_name) that pads the boundary + tiles ON-CHIP: OOB query rows load as 0 via a bounds-checked buffer load + (register pad), OOB keys load as 0 into LDS (LDS pad), the grid/work-item + count is ceil'd to cover the partial last query block, and the partial O rows + are dropped by a guarded store. Causal needs no key mask (padded ktok >= + seqlen_kv > every real query, so causal drops them); non-causal adds a + ktok0 each query token q attends to + # keys k in [q-W+1, q] (causal AND within-window). The valid KV region per + # 256-row block is a slope-1 parallelogram band, so the KV loop is pruned + # to ~(W/block_n + block_m/block_n) tiles instead of the full causal + # triangle. Requires causal=True and W % block_n == 0. + sliding_window: int = 0 + # ragged: separate kernel path for sequence lengths that are NOT a multiple + # of the tile geometry (seqlen_q % 256 != 0 and/or seqlen_kv % block_n != 0). + # Instead of host-side zero-padding (which we cannot do), the boundary tiles + # are padded ON-CHIP: OOB query rows load as 0 via a bounds-checked buffer + # load (register pad), OOB key rows load as 0 into LDS (LDS pad), the ceil'd + # grid covers the partial last query block, and the partial O rows are + # dropped by a guarded store. Causal masking already excludes the padded + # keys (their token index >= seqlen_kv > every real query), so causal needs + # NO extra key mask; non-causal adds a ktok -20%). + waves_per_eu: int = 2 + # persistent: emit the grid-stride PERSISTENT variant instead of one CTA per + # (query-block, head, batch). A 1-D grid of ``num_persistent`` long-lived CTAs + # grid-strides over the W = (seqlen_q//256)*Hq*B work items, so the per-CTA + # launch/dispatch + scalar setup + K/V-prime cold-start (~4.5 tile-equivalents, + # plan.md "CAUSAL GAP = FIXED-COST AMORTIZATION") is paid once per CU instead of + # once per query-block. Inner compute is byte-identical to the default path. + # Measured MI355X Sq=8192 causal: 512 -> 853 TFLOPS (+70%), 0 spill, err 1.46e-3. + persistent: bool = False + # num_persistent: number of long-lived CTAs when ``persistent``. 256 = exactly one + # 8-wave block per CU on MI355X (256 CUs) at 2 waves/SIMD; larger oversubscribes + # the CUs -> a serialized 2nd block -> tail loss (304 measured -20%). + num_persistent: int = 256 + # interleave: boustrophedon query-block ordering that reverses qb on alternating + # (hq,bt) planes to spread the triangular causal load across CTAs. A large-Sq + # lever (helps Sq>=16384) that slightly hurts small Sq; only used when persistent. + interleave: bool = False + # persist_decode: work-item -> (qb, hq, bt) decode for the persistent grid. + # "qb_major" (default): wi = qb*(Hq*B) + hq*B + bt. Balances the triangular + # causal load across CTAs, but every 256-CTA grid-stride phase spans ALL + # kv-heads at once -> large L2 footprint (measured 57% L2 hit @ GQA-8). + # "hkv_major": wi = hkv*(NQB*gqa*B) + blk*(gqa*B) + hql*B + bt, with blk + # folded to a low/high-paired qb. Concentrates each grid-stride phase on + # ~1 kv-head so the shared GQA K/V stays L2-resident across its gqa query + # heads (measured L2 hit 57% -> 93%, HBM misses 5.9x lower, matching the + # non-persistent grid). Only balances the causal triangle when each CTA + # grid-strides across BOTH halves of a kv-head, i.e. gqa*NQB*B >= 2*NP; + # otherwise each CTA gets a fixed qb (severe imbalance) -> slower. + # "auto" (default): hkv_major when it is balance-safe AND GQA + # (gqa>1 and gqa*NQB*B >= 2*num_persistent), else qb_major. Strictly >= + # qb_major (falls back where hkv_major would lose). + persist_decode: str = "auto" + # lazy_rescale: adaptive online-softmax rescale. Keep the running max as a + # LAZY max that only re-anchors when a tile's max exceeds it by > 8 (log2); + # when every lane is within 8 (wave_all vote) skip the O/l rescale entirely + # (a 0/1-trip scf.for compiles the skip to a wave-uniform scalar branch), + # cutting the per-tile VALU between the QK and PV MFMA clusters (raises + # MFMA utilization). P is then bounded by exp2(8)=256 (safe for fp32 accum + # / bf16 P) rather than <=1, so this is a numerically APPROXIMATE lever + # (still within bf16/fp16 tolerance). ALWAYS-ON by default (parity-identical + # at 1.46e-3, ~+2% TFLOPS); set False only to disable for A/B. + lazy_rescale: bool = True + + def __post_init__(self) -> None: + if self.dtype not in _DTYPE_IR: + raise ValueError( + f"dtype must be one of {sorted(_DTYPE_IR)}, got {self.dtype}" + ) + # head_size must be 64 or 128: the QK/PV MFMA tiling needs a multiple of + # 32, and the async K/V DMA (64 lanes x 2 bf16 = 128 elems/instr) needs + # 128 % head_size == 0 so it packs a whole number of rows per instr. + if self.head_size not in (64, 128): + raise ValueError(f"head_size must be 64 or 128, got {self.head_size}") + if self.block_n % 32 != 0: + raise ValueError(f"block_n must be a multiple of 32, got {self.block_n}") + if self.ragged: + if self.seqlen_q <= 0 or self.seqlen_kv <= 0: + raise ValueError("ragged requires positive seqlen_q/seqlen_kv") + if self.seqlen_q != self.seqlen_kv: + raise ValueError( + "ragged is self-attention only (seqlen_q == seqlen_kv), got " + f"{self.seqlen_q} != {self.seqlen_kv}" + ) + if self.varlen: + raise ValueError("ragged is not supported with varlen") + if self.sliding_window > 0: + raise ValueError("ragged is not supported with sliding_window") + else: + if self.seqlen_q % _BLOCK_M != 0: + raise ValueError( + f"seqlen_q must be a multiple of {_BLOCK_M}, got {self.seqlen_q}" + ) + if self.seqlen_kv % self.block_n != 0: + raise ValueError( + f"seqlen_kv must be a multiple of block_n={self.block_n}, got {self.seqlen_kv}" + ) + if self.num_kv_heads == 0 or self.num_query_heads % self.num_kv_heads != 0: + raise ValueError( + f"num_query_heads ({self.num_query_heads}) must be a positive multiple " + f"of num_kv_heads ({self.num_kv_heads})" + ) + if self.block_n % 32 != 0 or self.block_n <= 0: + raise ValueError( + f"block_n must be a positive multiple of 32, got {self.block_n}" + ) + if self.persistent and self.num_persistent <= 0: + raise ValueError( + f"num_persistent must be positive, got {self.num_persistent}" + ) + if self.persist_decode not in ("qb_major", "hkv_major", "auto"): + raise ValueError( + f"persist_decode must be 'qb_major', 'hkv_major', or 'auto', got " + f"{self.persist_decode}" + ) + if self.sliding_window < 0: + raise ValueError(f"sliding_window must be >= 0, got {self.sliding_window}") + if self.sliding_window > 0: + if not self.causal: + raise ValueError("sliding_window>0 requires causal=True") + if self.sliding_window % self.block_n != 0: + raise ValueError( + f"sliding_window ({self.sliding_window}) must be a multiple of " + f"block_n={self.block_n}" + ) + if self.varlen: + if self.persistent: + raise ValueError("varlen is not supported with persistent=True") + if not self.causal: + raise ValueError("varlen requires causal=True") + + @property + def num_waves(self) -> int: + return _BLOCK_M // 32 + + @property + def dtype_ir(self): + return _DTYPE_IR[self.dtype] + + @property + def num_queries_per_kv(self) -> int: + return self.num_query_heads // self.num_kv_heads + + @property + def resolved_persist_decode(self) -> str: + """Resolve persist_decode='auto' to 'hkv_major' (GQA L2-locality win, + when balance-safe) or 'qb_major'. hkv_major balances the causal triangle + only when each CTA grid-strides across both halves of a kv-head + (gqa*NQB*B >= 2*num_persistent) and there is GQA sharing (gqa>1).""" + if self.persist_decode != "auto": + return self.persist_decode + gqa = self.num_query_heads // self.num_kv_heads + nqb = (self.seqlen_q + _BLOCK_M - 1) // _BLOCK_M # ceil (ragged partial) + per_hkv = gqa * nqb * self.batch + if gqa > 1 and per_hkv >= 2 * self.num_persistent: + return "hkv_major" + return "qb_major" + + def kernel_name(self) -> str: + parts = [ + "rocke_attention_dense", + f"d{self.head_size}", + f"hq{self.num_query_heads}", + f"kv{self.num_kv_heads}", + f"bn{self.block_n}", + self.dtype, + f"sq{self.seqlen_q}", + f"sk{self.seqlen_kv}", + "causal" if self.causal else "full", + ] + if self.ragged: + parts.append("ragged") + if self.sliding_window > 0: + parts.append(f"swa{self.sliding_window}") + if self.varlen: + parts.append("varlen") + if self.lazy_rescale: + parts.append("lazyrs") + if self.persistent: + parts.append(f"persist{self.num_persistent}") + if self.resolved_persist_decode == "hkv_major": + parts.append("hkvmaj") + if self.interleave: + parts.append("intl") + return kernel_name_join(*parts) + + +def supports_attention_dense( + spec: AttentionDenseSpec, *, arch: str = "gfx950" +) -> Tuple[bool, str]: + """Return (ok, reason). The kernel is gfx950-only and dense (no paging/bias).""" + if arch != "gfx950": + return False, f"attention_dense is gfx950-only (got {arch})" + try: + AttentionDenseSpec(**{f.name: getattr(spec, f.name) for f in spec.__dataclass_fields__.values()}) # type: ignore[attr-defined] + except ValueError as e: + return False, str(e) + return True, "" + + +def build_attention_dense( + spec: AttentionDenseSpec, *, arch: str = "gfx950" +) -> KernelDef: + """Emit the dense flash-attention prefill kernel described by ``spec``.""" + if arch != "gfx950": + raise NotImplementedError(f"attention_dense is gfx950-only (got {arch})") + + if spec.persistent: + return _build_attention_dense_persistent(spec) + + B = spec.batch + Sq = spec.seqlen_q + Skv = spec.seqlen_kv + Hq = spec.num_query_heads + Hkv = spec.num_kv_heads + D = spec.head_size + causal = spec.causal + dtype = spec.dtype_ir + + BLOCK_M = _BLOCK_M + WAVES = spec.num_waves + BN = spec.block_n + NBUF = _NBUF + PAD = _LDS_PAD + W = spec.sliding_window + Wt = W // BN # window length in KV tiles (0 when disabled) + varlen = spec.varlen + RAGGED = spec.ragged + LAZY_RESCALE = spec.lazy_rescale + + K_STEPS = D // 16 + D_TILES = D // 32 + N_SUB = BN // 32 + KK_STEPS = BN // 16 + gqa = Hq // Hkv + stride_q_tok = Hq * D + stride_k_tok = Hkv * D + # DMA row packing: one async_buffer_load_lds instr moves 64 lanes x 2 bf16 = + # 128 elems. D==128 => 1 row/instr (the padded fast path, byte-identical). + # D<128 => pack 128//D rows/instr into an UNPADDED contiguous LDS tile (the + # transpose/K reads still get correct pitch from the [BN, LDROW] tensor). + ROWS_PER_INSTR = 128 // D + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = WAVES * 64 + b.kernel.attrs["waves_per_eu"] = int(spec.waves_per_eu) + + q = b.param( + "q_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + k = b.param( + "k_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + v = b.param( + "v_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + o = b.param( + "o_ptr", PtrType(dtype, "global"), noalias=True, writeonly=True, align=16 + ) + scale = b.param("scale", F32) + if varlen: + cu_q = b.param( + "cu_seqlens_q", PtrType(I32, "global"), noalias=True, readonly=True, align=4 + ) + cu_kv = b.param( + "cu_seqlens_kv", + PtrType(I32, "global"), + noalias=True, + readonly=True, + align=4, + ) + qk_scale = b.fmul(scale, b.const_f32(LOG2E)) + + _exp2 = b.exp2_fast # native v_exp_f32 (softmax arg always <= 0) + + tid = b.thread_id_x() + wave = b.div(tid, b.const_i32(64)) + lane = b.mod(tid, b.const_i32(64)) + lane_m = b.mod(lane, b.const_i32(32)) + lane_h = b.div(lane, b.const_i32(32)) + d_base = b.mul(lane_h, b.const_i32(8)) + neg_inf = b.const_f32(-1e30) + + qb = b.block_id_x() + hq = b.block_id_y() + bt = b.block_id_z() + hkv = b.div(hq, b.const_i32(gqa)) + q_tok0 = b.add(b.mul(qb, b.const_i32(BLOCK_M)), b.mul(wave, b.const_i32(32))) + + if varlen: + # Per-sequence token bases from cu_seqlens; early-exit q-blocks that fall + # past this sequence's length (packed [total_tok, H, D] layout). + q_seq0 = b.global_load_i32(cu_q, bt) + q_seq1 = b.global_load_i32(cu_q, b.add(bt, b.const_i32(1))) + seqlen_q_b = b.sub(q_seq1, q_seq0) + kv_seq0 = b.global_load_i32(cu_kv, bt) + kv_seq1 = b.global_load_i32(cu_kv, b.add(bt, b.const_i32(1))) + seqlen_kv_b = b.sub(kv_seq1, kv_seq0) + with b.scf_if(b.cmp_ge(b.mul(qb, b.const_i32(BLOCK_M)), seqlen_q_b)): + b.ret() + q_base = b.add( + b.mul(q_seq0, b.const_i32(stride_q_tok)), b.mul(hq, b.const_i32(D)) + ) + k_base = b.add( + b.mul(kv_seq0, b.const_i32(stride_k_tok)), b.mul(hkv, b.const_i32(D)) + ) + else: + q_base = b.add( + b.mul(b.mul(bt, b.const_i32(Sq)), b.const_i32(stride_q_tok)), + b.mul(hq, b.const_i32(D)), + ) + k_base = b.add( + b.mul(b.mul(bt, b.const_i32(Skv)), b.const_i32(stride_k_tok)), + b.mul(hkv, b.const_i32(D)), + ) + + # --- LDS allocation: PAD on K (bank-conflict fix), +PAD_V on V (transposed + # PV read bank-conflict pad). Row padding requires 1 row/instr (a padded + # row is not contiguous with the next); the packed D<128 loader must use + # an unpadded pitch (LDROW==D) so its multi-row DMA lands row-aligned. --- + LDROW = D + PAD if ROWS_PER_INSTR == 1 else D + VROW = D + _LDS_PAD_V if ROWS_PER_INSTR == 1 else D + K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") + V_lds = b.smem_alloc(dtype, [NBUF, BN, VROW], name_hint="Vlds") + + # Q packs (B operand), scaled once by qk_scale = softmax_scale * log2(e). + # ragged: a bounds-checked buffer load returns 0 for OOB query rows (the + # partial last block), so padded rows are register-zero (their output is + # dropped by the guarded store). Aligned: direct global load (unchanged IR). + q_rsrc = b.buffer_rsrc(q, b.const_i32(B * Sq * Hq * D * 2)) if RAGGED else None + q_tok = b.add(q_tok0, lane_m) + q_packs = [] + for ks in range(K_STEPS): + col = b.add(b.const_i32(ks * 16), d_base) + addr = b.add(b.add(q_base, b.mul(q_tok, b.const_i32(stride_q_tok))), col) + if RAGGED: + raw = b.buffer_load_vN( + q_rsrc, b.mul(addr, b.const_i32(2)), b.const_i32(0), dtype, 8 + ) + else: + raw = b.global_load_vN(q, addr, dtype, 8, align=16) + elems = [ + b.cast_f32_to(b.fmul(b.cast_to_f32(b.vec_extract(raw, j)), qk_scale), dtype) + for j in range(8) + ] + q_packs.append(b.vec_pack(elems, dtype)) + + # ragged: ceil so the partial last KV tile is visited (its OOB keys load 0 + # into LDS and are masked out); aligned: exact. + n_ktiles = ((Skv + BN - 1) // BN) if RAGGED else (Skv // BN) + n_per = BLOCK_M // BN + + K_BYTES_PER_BUF = BN * LDROW * 2 + K_LDROW_BYTES = LDROW * 2 + V_BYTES_PER_BUF = BN * VROW * 2 + V_LDROW_BYTES = VROW * 2 + ROWS_PER_WAVE = BN // WAVES + WAVE_BYTES = 64 * 16 + V_DMA_PASSES = (BN * D) // (WAVES * 64 * 8) + zero_soff = b.const_i32(0) + K_lds_addr = b.smem_addr_of(K_lds) + V_lds_addr = b.smem_addr_of(V_lds) + k_rsrc = b.buffer_rsrc(k, b.const_i32(B * Skv * Hkv * D * 2)) + v_rsrc = b.buffer_rsrc(v, b.const_i32(B * Skv * Hkv * D * 2)) + v_wave_off_i64 = b.zext(b.to_sgpr_u32(b.mul(wave, b.const_i32(WAVE_BYTES))), I64) + + def _async_load(rsrc, lds_base, buf_val, tile_key0, bytes_per_buf, ldrow_bytes): + """Async DMA one K/V tile into the [BN, LDROW] LDS layout. + + ROWS_PER_INSTR==1 (D==128): one instr per padded row -- 64 lanes x 2 bf16 + fill exactly one D-wide row (the original byte-identical fast path). + ROWS_PER_INSTR>1 (D<128): pack ROWS_PER_INSTR rows per instr into the + unpadded contiguous tile (lane l -> row l//(D/2), col 2*(l%(D/2))).""" + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(bytes_per_buf)) + if ROWS_PER_INSTR == 1: + for r in range(ROWS_PER_WAVE): + row = b.add(b.mul(wave, b.const_i32(ROWS_PER_WAVE)), b.const_i32(r)) + row_lds_off = b.add( + buf_off, b.zext(b.mul(row, b.const_i32(ldrow_bytes)), I64) + ) + row_base = b.smem_ptr_add(lds_base, row_lds_off) + gkey = b.add(tile_key0, row) + gcol = b.mul(lane, b.const_i32(2)) + voff = b.add( + b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), gcol + ) + b.async_buffer_load_lds_addr( + rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 + ) + else: + lanes_per_row = D // 2 + sub_row = b.div(lane, b.const_i32(lanes_per_row)) + col = b.mul(b.mod(lane, b.const_i32(lanes_per_row)), b.const_i32(2)) + for it in range(ROWS_PER_WAVE // ROWS_PER_INSTR): + row0 = b.add( + b.mul(wave, b.const_i32(ROWS_PER_WAVE)), + b.const_i32(it * ROWS_PER_INSTR), + ) + row_lds_off = b.add( + buf_off, b.zext(b.mul(row0, b.const_i32(ldrow_bytes)), I64) + ) + row_base = b.smem_ptr_add(lds_base, row_lds_off) + gkey = b.add(b.add(tile_key0, row0), sub_row) + voff = b.add(b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), col) + b.async_buffer_load_lds_addr( + rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 + ) + + def async_load_k(lds_base, buf_val, tile_key0): + _async_load( + k_rsrc, lds_base, buf_val, tile_key0, K_BYTES_PER_BUF, K_LDROW_BYTES + ) + + def async_load_v(lds_base, buf_val, tile_key0): + _async_load( + v_rsrc, lds_base, buf_val, tile_key0, V_BYTES_PER_BUF, V_LDROW_BYTES + ) + + def load_tile(buf_val, tile_idx): + tk0 = b.mul(tile_idx, b.const_i32(BN)) + async_load_k(K_lds_addr, buf_val, tk0) + async_load_v(V_lds_addr, buf_val, tk0) + + # ---- per-tile compute closures ---- + + def do_qk(kbuf): + """QK MFMA: S^T = K@Q^T. mfma(a=K, bv=Q) => key on the 16 per-lane accumulator + regs (+lane^32), query on lane%32 -- the layout that keeps softmax a cheap + in-lane reduce + one lane^32 exchange, and lets CK-1's transposed PV consume P + with no relayout shuffle.""" + s_reg = [] + for nsub in range(N_SUB): + acc = b.zero_vec_f32(16) + krow = b.add(b.const_i32(nsub * 32), lane_m) + for ks in range(K_STEPS): + col = b.add(b.const_i32(ks * 16), d_base) + k_pack = b.smem_load_vN(K_lds, kbuf, krow, col, dtype=dtype, n=8) + acc = mfma_32x32x16_for_dtype(b, dtype, k_pack, q_packs[ks], acc) + s_reg.append([b.vec_extract(acc, i) for i in range(16)]) + return s_reg + + def do_mask(s_reg, tile_idx, lower=False, upper=True): + """Apply causal (upper: ktok<=q) and/or sliding-window (lower: + ktok>q-W) masks in-place on the QK-output layout. W is compile-time so + the lower threshold folds to an immediate. No relayout (reuses the same + lane->ktok/query_tok maps as causal).""" + if not causal: + return + tile_key0 = b.mul(tile_idx, b.const_i32(BN)) + query_tok = b.add(q_tok0, _mfma_32x32_c_col(b, lane, 0)) + # lower bound key: q - W + 1 (keep iff ktok > q - W) + win_lo = b.sub(query_tok, b.const_i32(W)) if lower else None + for nsub in range(N_SUB): + sub_base = b.add(tile_key0, b.const_i32(nsub * 32)) + for i in range(16): + ktok = b.add(sub_base, _mfma_32x32_c_row(b, lane, i)) + if upper: + s_reg[nsub][i] = b.select( + b.cmp_le(ktok, query_tok), s_reg[nsub][i], neg_inf + ) + if lower: + s_reg[nsub][i] = b.select( + b.cmp_gt(ktok, win_lo), s_reg[nsub][i], neg_inf + ) + + def do_kbound_mask(s_reg, tile_idx): + """ragged non-causal: force scores of padded keys (ktok >= seqlen_kv, the + OOB rows of the partial last KV tile) to -inf. Causal doesn't need this + (padded ktok >= seqlen_kv > every real query, so causal already drops + them). seqlen_kv is compile-time -> the bound folds to an immediate.""" + tile_key0 = b.mul(tile_idx, b.const_i32(BN)) + for nsub in range(N_SUB): + sub_base = b.add(tile_key0, b.const_i32(nsub * 32)) + for i in range(16): + ktok = b.add(sub_base, _mfma_32x32_c_row(b, lane, i)) + s_reg[nsub][i] = b.select( + b.cmp_lt(ktok, b.const_i32(Skv)), s_reg[nsub][i], neg_inf + ) + + def softmax_max(s_reg, m_i): + local_max = neg_inf + for nsub in range(N_SUB): + for i in range(16): + local_max = b.fmax(local_max, s_reg[nsub][i]) + tile_max = b.fmax(local_max, b.warp_shuffle_xor(local_max, 32)) + if LAZY_RESCALE: + m_diff = b.fsub(tile_max, m_i) + below_i32 = b.select( + b.fcmp("ole", m_diff, b.const_f32(_LAZY_RESCALE_THRESHOLD)), + b.const_i32(1), + b.const_i32(0), + ) + skip = b.cmp_ne(b.wave_all(below_i32), b.const_i32(0)) + m_new = b.select(skip, m_i, b.fmax(m_i, tile_max)) + else: + skip = None + m_new = b.fmax(m_i, tile_max) + alpha = _exp2(b.fsub(m_i, m_new)) + return m_new, alpha, skip + + def relayout_p(p): + """CK-1 half-local P feed: assemble the PV B-operand from lane-local P regs + only (a bf16 cast + pack, NO cross-half warp_shuffle_xor/select). Pairs with + the half-local V load in ``read_v`` so the K-axis stays aligned.""" + packs = [] + for kk_step in range(KK_STEPS): + elems = [] + for kk in range(8): + local_in_group = kk % 4 + band = kk // 4 + key_idx = kk_step * 16 + band * 8 + local_in_group + p_tile = key_idx // 32 + row_static = key_idx % 32 + preg = (row_static // 8) * 4 + (row_static % 4) + elems.append(b.cast_f32_to(p[p_tile][preg], dtype)) + packs.append(b.vec_pack(elems, dtype)) + return packs + + def read_v(dt, kk_step, vbuf): + """CK-1 half-local transposed V A-operand (matches ``relayout_p``).""" + return pv32_v_load_paired( + b, + V_lds=V_lds, + v_buf=vbuf, + n=dt, + k=kk_step, + lane_half32=lane_h, + lane_col32=lane_m, + dtype=dtype, + ) + + def do_pv(o_acc_in, p_packs, vbuf): + out = [] + for dt in range(D_TILES): + acc_o = o_acc_in[dt] + for kk_step in range(KK_STEPS): + acc_o = mfma_32x32x16_for_dtype( + b, dtype, read_v(dt, kk_step, vbuf), p_packs[kk_step], acc_o + ) + out.append(acc_o) + return out + + def rescale_o(o_acc, alpha): + return [ + b.vec_pack( + [b.fmul(b.vec_extract(o_acc[dt], i), alpha) for i in range(16)], F32 + ) + for dt in range(D_TILES) + ] + + def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): + """Depth-1 cluster: interleave exp2(s - m_new) into the PV MFMA loop so the + softmax VALU/TRANS co-executes in the MFMA shadow. The full per-step + instruction population (DS_READ/MFMA/VALU/TRANS) is named to sched_group_barrier + so the IGLP grouping matches the real stream.""" + exp_per = -(-(N_SUB * 16) // (D_TILES * KK_STEPS)) + slots = [(nsub, i) for nsub in range(N_SUB) for i in range(16)] + p_vals = [[None] * 16 for _ in range(N_SUB)] + it = iter(slots) + out = [] + for dt in range(D_TILES): + acc_o = o_acc_in[dt] + for kk_step in range(KK_STEPS): + acc_o = mfma_32x32x16_for_dtype( + b, dtype, read_v(dt, kk_step, vbuf), p_packs[kk_step], acc_o + ) + n_emit = 0 + for _ in range(exp_per): + slot = next(it, None) + if slot is None: + break + nsub, i = slot + p_vals[nsub][i] = _exp2(b.fsub(s_reg[nsub][i], m_new)) + n_emit += 1 + b.sched_group_barrier(DS_READ, 2, 0) + b.sched_group_barrier(MFMA, 1, 0) + b.sched_group_barrier(VALU, max(1, n_emit), 0) + b.sched_group_barrier(TRANS, max(1, n_emit), 0) + out.append(acc_o) + for slot in it: + nsub, i = slot + p_vals[nsub][i] = _exp2(b.fsub(s_reg[nsub][i], m_new)) + l_local = b.const_f32(0.0) + for nsub in range(N_SUB): + for i in range(16): + l_local = b.fadd(l_local, p_vals[nsub][i]) + l_tile = b.fadd(l_local, b.warp_shuffle_xor(l_local, 32)) + return out, p_vals, l_tile + + n_ktiles_val = ( + b.div(seqlen_kv_b, b.const_i32(BN)) if varlen else b.const_i32(n_ktiles) + ) + if causal: + n_upper = b.add(b.mul(qb, b.const_i32(n_per)), b.const_i32(n_per)) + n_upper = b.select(b.cmp_lt(n_upper, n_ktiles_val), n_upper, n_ktiles_val) + else: + n_upper = n_ktiles_val + + # Sliding-window: first KV tile any row in this block attends to. Valid band + # is [start_tile, n_upper); tiles < start_tile are fully outside the window + # (all -inf) so they are never visited (the KV-loop prune). W==0 keeps + # start_tile=0 -> full causal, byte-identical to the always-on path. + if causal and W > 0: + _diag0 = b.mul(qb, b.const_i32(n_per)) + _lo_raw = b.sub(_diag0, b.const_i32(Wt)) + start_tile = b.select( + b.cmp_gt(_lo_raw, b.const_i32(0)), _lo_raw, b.const_i32(0) + ) + else: + start_tile = b.const_i32(0) + start_buf = b.mod(start_tile, b.const_i32(NBUF)) + start_buf1 = b.mod(b.add(start_tile, b.const_i32(1)), b.const_i32(NBUF)) + + # Prologue: prime the K/V double buffer and compute the first (start) tile. + load_tile(start_buf, start_tile) + load_tile(start_buf1, b.add(start_tile, b.const_i32(1))) + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + # ragged non-causal needs the key-pad mask (ktok 0: + do_mask(s0, start_tile, lower=True, upper=True) + else: + do_mask(s0, start_tile) + if RAG_KBOUND: + do_kbound_mask(s0, start_tile) + m0, _alpha0, _skip0 = softmax_max(s0, neg_inf) + # tile-0 softmax exp + relayout only; PV lags by one tile (fused into the loop). + p0_vals = [ + [_exp2(b.fsub(s0[nsub][i], m0)) for i in range(16)] for nsub in range(N_SUB) + ] + l0_local = b.const_f32(0.0) + for nsub in range(N_SUB): + for i in range(16): + l0_local = b.fadd(l0_local, p0_vals[nsub][i]) + l0 = b.fadd(l0_local, b.warp_shuffle_xor(l0_local, 32)) + o0 = [b.zero_vec_f32(16) for _ in range(D_TILES)] + pk0 = relayout_p(p0_vals) + + iter_args = ( + [("m", m0), ("l", l0)] + + [(f"o{dt}", o0[dt]) for dt in range(D_TILES)] + + [(f"pk{kk}", pk0[kk]) for kk in range(KK_STEPS)] + ) + + _rs_ctr = [0] + + def emit_loop_body(j, carry, mask_lower=False, mask_upper=False, mask_kbound=False): + m_i = carry[0] + l_i = carry[1] + o_acc = list(carry[2 : 2 + D_TILES]) + p_prev = list(carry[2 + D_TILES : 2 + D_TILES + KK_STEPS]) + kbuf = b.mod(j, b.const_i32(NBUF)) + vbuf_prev = b.mod(b.add(j, b.const_i32(NBUF - 1)), b.const_i32(NBUF)) + pbuf = b.mod(b.add(j, b.const_i32(1)), b.const_i32(NBUF)) + + # PF (partial-vmcnt prefetch): keep the freshest V(j) DMA in flight so it + # overlaps compute instead of a full vmcnt(0) serialize (bit-identical). + b.s_waitcnt(vmcnt=V_DMA_PASSES) + b.s_barrier_bare() + s = do_qk(kbuf) + if mask_lower or mask_upper: + do_mask(s, j, lower=mask_lower, upper=mask_upper) + if mask_kbound: + do_kbound_mask(s, j) + m_new, alpha, skip = softmax_max(s, m_i) + b.sched_barrier(0) # depth-1 fence: m_new region-live-in + b.s_setprio(1) # PV-only s_setprio (paired with PF ~+3.5%) + o_acc, p_vals, l_tile = pv_fused_exp(o_acc, p_prev, vbuf_prev, s, m_new) + b.s_setprio(0) + if LAZY_RESCALE: + _rs_ctr[0] += 1 + tg = _rs_ctr[0] + trips = b.select(skip, b.const_i32(0), b.const_i32(1)) + rs_args = [(f"ro{dt}_{tg}", o_acc[dt]) for dt in range(D_TILES)] + rs_args.append((f"rl_{tg}", l_i)) + rs = b.scf_for_iter( + b.const_i32(0), trips, b.const_i32(1), rs_args, iv_name=f"rs{tg}" + ) + with rs as (_iv, rc): + o_sc = rescale_o(list(rc[:D_TILES]), alpha) + b.scf_yield(*o_sc, b.fmul(rc[D_TILES], alpha)) + o_acc = list(rs.results[:D_TILES]) + l_new = b.fadd(rs.results[D_TILES], l_tile) + else: + l_new = b.fadd(b.fmul(l_i, alpha), l_tile) + o_acc = rescale_o(o_acc, alpha) + p_packs = relayout_p(p_vals) + b.s_barrier_bare() + load_tile(pbuf, b.add(j, b.const_i32(1))) + b.scf_yield(m_new, l_new, *o_acc, *p_packs) + + if causal and W > 0: + # Sliding-window three-phase band loop (prologue already did start_tile): + # L: [start+1, mid_lo) window-edge tiles (masked) + # M: [mid_lo, mid_hi) interior (mask-free: both bounds hold for all rows) + # R: [mid_hi, n_upper) causal-edge tiles (masked) + # Boundary phases apply BOTH bounds (robust for W= seqlen_q) via a per-lane guard so + # they never write (and never clobber a neighbouring batch's real rows). A + # buffer store's OOB-drop only protects the last batch's overflow, so use an + # explicit predicate that is correct for any batch. + o_store_ctx = ( + b.scf_if(b.cmp_lt(qtok, b.const_i32(Sq))) if RAGGED else _nullcontext() + ) + with o_store_ctx: + for dt in range(D_TILES): + for g in range(4): + d0 = b.add(b.const_i32(dt * 32 + g * 8), d_half) + addr = b.add(q_row_byte, d0) + vals = [ + b.cast_f32_to( + b.fmul(b.vec_extract(o_acc[dt], g * 4 + kk), rcp_l), dtype + ) + for kk in range(4) + ] + b.global_store_vN(o, addr, b.vec_pack(vals, dtype), 4, align=8) + b.ret() + return b.kernel + + +def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: + """Persistent (grid-stride) variant of the dense flash-attention kernel. + + Launches a 1-D grid of ``spec.num_persistent`` long-lived CTAs; each CTA + grid-strides over the flattened work-item space ``W = (Sq//BLOCK_M)*Hq*B`` and + runs the byte-identical inner step-1 CK-1 pipeline per work item, so the per-CTA + launch/dispatch + scalar setup + K/V-prime cold-start is amortized once per CU + instead of once per query-block (see the ``persistent`` spec field). Every + algorithmic lever is the same always-on set as the default build; the only + differences are the outer work loop, the qb-major work decode (load-balances the + causal triangle), the per-work-item state reset, and ``exp_per=1`` (keeps the + extra loop-carried index math within 256 VGPR at 0 spill; numerically identical + to the default's ``exp_per=2`` — pure emission ordering).""" + B = spec.batch + Sq = spec.seqlen_q + Skv = spec.seqlen_kv + Hq = spec.num_query_heads + Hkv = spec.num_kv_heads + D = spec.head_size + causal = spec.causal + dtype = spec.dtype_ir + + BLOCK_M = _BLOCK_M + WAVES = spec.num_waves + BN = spec.block_n + NBUF = _NBUF + PAD = _LDS_PAD + NP = spec.num_persistent + INTERLEAVE = spec.interleave + + K_STEPS = D // 16 + D_TILES = D // 32 + N_SUB = BN // 32 + KK_STEPS = BN // 16 + gqa = Hq // Hkv + stride_q_tok = Hq * D + stride_k_tok = Hkv * D + # DMA row packing (see default builder): 1 row/instr for D==128 (padded fast + # path), else pack 128//D rows/instr into an unpadded contiguous LDS tile. + ROWS_PER_INSTR = 128 // D + RAGGED = spec.ragged + # ragged: ceil both the KV tiles and the query-block count so the partial + # last block/tile is covered (padded rows/keys are handled on-chip). + n_ktiles = ((Skv + BN - 1) // BN) if RAGGED else (Skv // BN) + n_per = BLOCK_M // BN + NQB = ((Sq + BLOCK_M - 1) // BLOCK_M) if RAGGED else (Sq // BLOCK_M) + W = NQB * Hq * B # total work items + SW = spec.sliding_window # sliding-window length (0 = disabled) + SWt = SW // BN # window length in KV tiles + LAZY_RESCALE = spec.lazy_rescale + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = WAVES * 64 + b.kernel.attrs["waves_per_eu"] = int(spec.waves_per_eu) + + q = b.param( + "q_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + k = b.param( + "k_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + v = b.param( + "v_ptr", PtrType(dtype, "global"), noalias=True, readonly=True, align=16 + ) + o = b.param( + "o_ptr", PtrType(dtype, "global"), noalias=True, writeonly=True, align=16 + ) + scale = b.param("scale", F32) + qk_scale = b.fmul(scale, b.const_f32(LOG2E)) + _exp2 = b.exp2_fast + + # ----- CTA-invariant scalar setup (paid ONCE per persistent CTA) ----- + tid = b.thread_id_x() + wave = b.div(tid, b.const_i32(64)) + lane = b.mod(tid, b.const_i32(64)) + lane_m = b.mod(lane, b.const_i32(32)) + lane_h = b.div(lane, b.const_i32(32)) + d_base = b.mul(lane_h, b.const_i32(8)) + neg_inf = b.const_f32(-1e30) + + # 1 row/instr => padded pitch (bank-conflict fix); packed D<128 => unpadded. + LDROW = D + PAD if ROWS_PER_INSTR == 1 else D + VROW = (D + _LDS_PAD_V) if ROWS_PER_INSTR == 1 else D + K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") + V_lds = b.smem_alloc(dtype, [NBUF, BN, VROW], name_hint="Vlds") + + K_BYTES_PER_BUF = BN * LDROW * 2 + K_LDROW_BYTES = LDROW * 2 + V_BYTES_PER_BUF = BN * VROW * 2 + V_LDROW_BYTES = VROW * 2 + ROWS_PER_WAVE = BN // WAVES + WAVE_BYTES = 64 * 16 + V_DMA_PASSES = (BN * D) // (WAVES * 64 * 8) + zero_soff = b.const_i32(0) + K_lds_addr = b.smem_addr_of(K_lds) + V_lds_addr = b.smem_addr_of(V_lds) + k_rsrc = b.buffer_rsrc(k, b.const_i32(B * Skv * Hkv * D * 2)) + v_rsrc = b.buffer_rsrc(v, b.const_i32(B * Skv * Hkv * D * 2)) + # ragged: bounds-checked Q load (OOB partial-block rows -> 0 register pad). + q_rsrc = b.buffer_rsrc(q, b.const_i32(B * Sq * Hq * D * 2)) if RAGGED else None + v_wave_off_i64 = b.zext(b.to_sgpr_u32(b.mul(wave, b.const_i32(WAVE_BYTES))), I64) + + # ----- persistent grid-stride loop over the flattened work-item space ----- + cta_id = b.block_id_x() + outer = b.scf_for(cta_id, b.const_i32(W), b.const_i32(NP), iv_name="wi") + with outer as wi: + # Cross-work-item LDS reuse safety: drain the previous item's trailing DMA + # and barrier so all waves finished the previous epilogue reads before we + # reissue into the shared K/V buffers. + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + + if spec.resolved_persist_decode == "hkv_major": + # hkv-MAJOR + causal-balanced decode: + # wi = hkv*(NQB*gqa*B) + blk*(gqa*B) + hql*B + bt + # * hkv in the MSB -> each grid-stride phase (NP consecutive wi) stays + # within ~1 kv-head, so the shared GQA K/V is L2-resident across its + # gqa query heads (recovers the non-persistent grid's locality: + # measured L2 hit 57% -> ~90%+ vs qb_major). + # * `blk` (0..NQB-1) is folded to a query-block index that PAIRS a low + # and a high qb per CTA: blk qb=blk (cheap), blk>=half -> + # qb=NQB-1-(blk-half) (expensive), so a CTA that grid-strides over + # both halves of a kv-head does qb=X and qb=NQB-1-X (constant causal + # cost) -> keeps qb_major's load balance while gaining L2 locality. + half = NQB // 2 + bt = b.mod(wi, b.const_i32(B)) + rem = b.div(wi, b.const_i32(B)) # hkv*(NQB*gqa) + blk*gqa + hql + hql = b.mod(rem, b.const_i32(gqa)) + r2 = b.div(rem, b.const_i32(gqa)) # hkv*NQB + blk + blk = b.mod(r2, b.const_i32(NQB)) + hkv = b.div(r2, b.const_i32(NQB)) + hq = b.add(b.mul(hkv, b.const_i32(gqa)), hql) + # qb = blk 1: + odd = b.cmp_eq(b.mod(rem, b.const_i32(2)), b.const_i32(1)) + qb = b.select(odd, b.sub(b.const_i32(NQB - 1), qb0), qb0) + else: + qb = qb0 + hkv = b.div(hq, b.const_i32(gqa)) + + q_tok0 = b.add(b.mul(qb, b.const_i32(BLOCK_M)), b.mul(wave, b.const_i32(32))) + q_base = b.add( + b.mul(b.mul(bt, b.const_i32(Sq)), b.const_i32(stride_q_tok)), + b.mul(hq, b.const_i32(D)), + ) + k_base = b.add( + b.mul(b.mul(bt, b.const_i32(Skv)), b.const_i32(stride_k_tok)), + b.mul(hkv, b.const_i32(D)), + ) + + q_tok = b.add(q_tok0, lane_m) + q_packs = [] + for ks in range(K_STEPS): + col = b.add(b.const_i32(ks * 16), d_base) + addr = b.add(b.add(q_base, b.mul(q_tok, b.const_i32(stride_q_tok))), col) + if RAGGED: + raw = b.buffer_load_vN( + q_rsrc, b.mul(addr, b.const_i32(2)), b.const_i32(0), dtype, 8 + ) + else: + raw = b.global_load_vN(q, addr, dtype, 8, align=16) + elems = [ + b.cast_f32_to( + b.fmul(b.cast_to_f32(b.vec_extract(raw, j)), qk_scale), dtype + ) + for j in range(8) + ] + q_packs.append(b.vec_pack(elems, dtype)) + + def _async_load(rsrc, lds_base, buf_val, tile_key0, bytes_per_buf, ldrow_bytes): + """Async DMA one K/V tile (see default builder ``_async_load``). + ROWS_PER_INSTR==1 (D==128): incremental one-instr-per-padded-row fast + path (byte-identical). ROWS_PER_INSTR>1 (D<128): pack rows per instr + into the unpadded contiguous tile.""" + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(bytes_per_buf)) + if ROWS_PER_INSTR == 1: + row0 = b.mul(wave, b.const_i32(ROWS_PER_WAVE)) + row_lds_off = b.add( + buf_off, b.zext(b.mul(row0, b.const_i32(ldrow_bytes)), I64) + ) + gcol = b.mul(lane, b.const_i32(2)) + voff = b.add( + b.add( + k_base, + b.mul(b.add(tile_key0, row0), b.const_i32(stride_k_tok)), + ), + gcol, + ) + for r in range(ROWS_PER_WAVE): + row_base = b.smem_ptr_add(lds_base, row_lds_off) + b.async_buffer_load_lds_addr( + rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 + ) + if r + 1 < ROWS_PER_WAVE: + row_lds_off = b.add(row_lds_off, b.const_i64(ldrow_bytes)) + voff = b.add(voff, b.const_i32(stride_k_tok)) + else: + lanes_per_row = D // 2 + sub_row = b.div(lane, b.const_i32(lanes_per_row)) + col = b.mul(b.mod(lane, b.const_i32(lanes_per_row)), b.const_i32(2)) + for it in range(ROWS_PER_WAVE // ROWS_PER_INSTR): + row0 = b.add( + b.mul(wave, b.const_i32(ROWS_PER_WAVE)), + b.const_i32(it * ROWS_PER_INSTR), + ) + row_lds_off = b.add( + buf_off, b.zext(b.mul(row0, b.const_i32(ldrow_bytes)), I64) + ) + row_base = b.smem_ptr_add(lds_base, row_lds_off) + gkey = b.add(b.add(tile_key0, row0), sub_row) + voff = b.add( + b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), col + ) + b.async_buffer_load_lds_addr( + rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 + ) + + def async_load_k(lds_base, buf_val, tile_key0): + _async_load( + k_rsrc, lds_base, buf_val, tile_key0, K_BYTES_PER_BUF, K_LDROW_BYTES + ) + + def async_load_v(lds_base, buf_val, tile_key0): + _async_load( + v_rsrc, lds_base, buf_val, tile_key0, V_BYTES_PER_BUF, V_LDROW_BYTES + ) + + def load_tile(buf_val, tile_idx): + tk0 = b.mul(tile_idx, b.const_i32(BN)) + async_load_k(K_lds_addr, buf_val, tk0) + async_load_v(V_lds_addr, buf_val, tk0) + + def do_qk(kbuf): + s_reg = [] + for nsub in range(N_SUB): + acc = b.zero_vec_f32(16) + krow = b.add(b.const_i32(nsub * 32), lane_m) + for ks in range(K_STEPS): + col = b.add(b.const_i32(ks * 16), d_base) + k_pack = b.smem_load_vN(K_lds, kbuf, krow, col, dtype=dtype, n=8) + acc = mfma_32x32x16_for_dtype(b, dtype, k_pack, q_packs[ks], acc) + s_reg.append([b.vec_extract(acc, i) for i in range(16)]) + return s_reg + + def do_mask(s_reg, tile_idx, lower=False, upper=True): + if not causal: + return + tile_key0 = b.mul(tile_idx, b.const_i32(BN)) + query_tok = b.add(q_tok0, _mfma_32x32_c_col(b, lane, 0)) + win_lo = b.sub(query_tok, b.const_i32(SW)) if lower else None + for nsub in range(N_SUB): + sub_base = b.add(tile_key0, b.const_i32(nsub * 32)) + for i in range(16): + ktok = b.add(sub_base, _mfma_32x32_c_row(b, lane, i)) + if upper: + s_reg[nsub][i] = b.select( + b.cmp_le(ktok, query_tok), s_reg[nsub][i], neg_inf + ) + if lower: + s_reg[nsub][i] = b.select( + b.cmp_gt(ktok, win_lo), s_reg[nsub][i], neg_inf + ) + + def do_kbound_mask(s_reg, tile_idx): + """ragged non-causal: -inf the padded keys (ktok >= seqlen_kv) of the + partial last KV tile. Causal drops them for free.""" + tile_key0 = b.mul(tile_idx, b.const_i32(BN)) + for nsub in range(N_SUB): + sub_base = b.add(tile_key0, b.const_i32(nsub * 32)) + for i in range(16): + ktok = b.add(sub_base, _mfma_32x32_c_row(b, lane, i)) + s_reg[nsub][i] = b.select( + b.cmp_lt(ktok, b.const_i32(Skv)), s_reg[nsub][i], neg_inf + ) + + RAG_KBOUND = RAGGED and (not causal) and (Skv % BN != 0) + + def softmax_max(s_reg, m_i): + local_max = neg_inf + for nsub in range(N_SUB): + for i in range(16): + local_max = b.fmax(local_max, s_reg[nsub][i]) + tile_max = b.fmax(local_max, b.warp_shuffle_xor(local_max, 32)) + if LAZY_RESCALE: + # Lazy max: only re-anchor when some lane's tile_max exceeds the + # running max by > threshold; else keep m_i (skip the rescale). + m_diff = b.fsub(tile_max, m_i) + below_i32 = b.select( + b.fcmp("ole", m_diff, b.const_f32(_LAZY_RESCALE_THRESHOLD)), + b.const_i32(1), + b.const_i32(0), + ) + skip = b.cmp_ne(b.wave_all(below_i32), b.const_i32(0)) + m_new = b.select(skip, m_i, b.fmax(m_i, tile_max)) + else: + skip = None + m_new = b.fmax(m_i, tile_max) + alpha = _exp2(b.fsub(m_i, m_new)) + return m_new, alpha, skip + + def softmax_stats(s_reg, m_i): + m_new, alpha, _skip = softmax_max(s_reg, m_i) + p = [ + [_exp2(b.fsub(s_reg[nsub][i], m_new)) for i in range(16)] + for nsub in range(N_SUB) + ] + l_local = b.const_f32(0.0) + for nsub in range(N_SUB): + for i in range(16): + l_local = b.fadd(l_local, p[nsub][i]) + l_tile = b.fadd(l_local, b.warp_shuffle_xor(l_local, 32)) + return m_new, alpha, p, l_tile + + def relayout_p(p): + packs = [] + for kk_step in range(KK_STEPS): + elems = [] + for kk in range(8): + local_in_group = kk % 4 + band = kk // 4 + key_idx = kk_step * 16 + band * 8 + local_in_group + p_tile = key_idx // 32 + row_static = key_idx % 32 + preg = (row_static // 8) * 4 + (row_static % 4) + elems.append(b.cast_f32_to(p[p_tile][preg], dtype)) + packs.append(b.vec_pack(elems, dtype)) + return packs + + def read_v(dt, kk_step, vbuf): + return pv32_v_load_paired( + b, + V_lds=V_lds, + v_buf=vbuf, + n=dt, + k=kk_step, + lane_half32=lane_h, + lane_col32=lane_m, + dtype=dtype, + ) + + def do_pv(o_acc_in, p_packs, vbuf): + out = [] + for dt in range(D_TILES): + acc_o = o_acc_in[dt] + for kk_step in range(KK_STEPS): + acc_o = mfma_32x32x16_for_dtype( + b, dtype, read_v(dt, kk_step, vbuf), p_packs[kk_step], acc_o + ) + out.append(acc_o) + return out + + def rescale_o(o_acc, alpha): + return [ + b.vec_pack( + [b.fmul(b.vec_extract(o_acc[dt], i), alpha) for i in range(16)], + F32, + ) + for dt in range(D_TILES) + ] + + def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): + exp_per = ( + 1 # one exp2 per PV-MFMA step -> 256 VGPR / 0 spill (see docstring) + ) + slots = [(nsub, i) for nsub in range(N_SUB) for i in range(16)] + p_vals = [[None] * 16 for _ in range(N_SUB)] + it = iter(slots) + out = [] + for dt in range(D_TILES): + acc_o = o_acc_in[dt] + for kk_step in range(KK_STEPS): + acc_o = mfma_32x32x16_for_dtype( + b, dtype, read_v(dt, kk_step, vbuf), p_packs[kk_step], acc_o + ) + n_emit = 0 + for _ in range(exp_per): + slot = next(it, None) + if slot is None: + break + nsub, i = slot + p_vals[nsub][i] = _exp2(b.fsub(s_reg[nsub][i], m_new)) + n_emit += 1 + b.sched_group_barrier(DS_READ, 2, 0) + b.sched_group_barrier(MFMA, 1, 0) + b.sched_group_barrier(VALU, max(1, n_emit), 0) + b.sched_group_barrier(TRANS, max(1, n_emit), 0) + out.append(acc_o) + for slot in it: + nsub, i = slot + p_vals[nsub][i] = _exp2(b.fsub(s_reg[nsub][i], m_new)) + l_local = b.const_f32(0.0) + for nsub in range(N_SUB): + for i in range(16): + l_local = b.fadd(l_local, p_vals[nsub][i]) + l_tile = b.fadd(l_local, b.warp_shuffle_xor(l_local, 32)) + return out, p_vals, l_tile + + _rs_ctr = [0] + + def emit_loop_body( + j, carry, mask_lower=False, mask_upper=False, mask_kbound=False + ): + m_i = carry[0] + l_i = carry[1] + o_acc = list(carry[2 : 2 + D_TILES]) + p_prev = list(carry[2 + D_TILES : 2 + D_TILES + KK_STEPS]) + pbuf = b.mod(b.add(j, b.const_i32(1)), b.const_i32(NBUF)) + kbuf = b.mod(j, b.const_i32(NBUF)) + vbuf_prev = b.mod(b.add(j, b.const_i32(NBUF - 1)), b.const_i32(NBUF)) + + # PF (partial-vmcnt prefetch): keep the freshest V(j) DMA in flight + # (drain only V(j-1)+K(j), both older) so DMA overlaps compute instead + # of a full vmcnt(0) serialize. Bit-identical, raises MfmaUtil. + b.s_waitcnt(vmcnt=V_DMA_PASSES) + b.s_barrier_bare() + s = do_qk(kbuf) + if mask_lower or mask_upper: + do_mask(s, j, lower=mask_lower, upper=mask_upper) + if mask_kbound: + do_kbound_mask(s, j) + m_new, alpha, skip = softmax_max(s, m_i) + b.sched_barrier(0) + # PV-only s_setprio: the PV MFMA cluster wins issue slots; paired with + # PF this converts to ~+3.5% (Sq=8192 causal, ~852 -> ~877 TFLOPS). + b.s_setprio(1) + o_acc, p_vals, l_tile = pv_fused_exp(o_acc, p_prev, vbuf_prev, s, m_new) + b.s_setprio(0) + if LAZY_RESCALE: + # Skip the O/l rescale via a wave-uniform 0/1-trip loop when the + # max didn't move (>threshold): 0 trips -> pass o_acc/l_i through + # unscaled; 1 trip -> scale by alpha. Compiles to a scalar branch. + # Unique names per emit (called once per KV-loop phase). + _rs_ctr[0] += 1 + tg = _rs_ctr[0] + trips = b.select(skip, b.const_i32(0), b.const_i32(1)) + rs_args = [(f"ro{dt}_{tg}", o_acc[dt]) for dt in range(D_TILES)] + rs_args.append((f"rl_{tg}", l_i)) + rs = b.scf_for_iter( + b.const_i32(0), trips, b.const_i32(1), rs_args, iv_name=f"rs{tg}" + ) + with rs as (_iv, rc): + o_sc = rescale_o(list(rc[:D_TILES]), alpha) + b.scf_yield(*o_sc, b.fmul(rc[D_TILES], alpha)) + o_acc = list(rs.results[:D_TILES]) + l_new = b.fadd(rs.results[D_TILES], l_tile) + else: + l_new = b.fadd(b.fmul(l_i, alpha), l_tile) + o_acc = rescale_o(o_acc, alpha) + p_packs = relayout_p(p_vals) + b.s_barrier_bare() + load_tile(pbuf, b.add(j, b.const_i32(1))) + b.scf_yield(m_new, l_new, *o_acc, *p_packs) + + if causal: + n_upper = b.add(b.mul(qb, b.const_i32(n_per)), b.const_i32(n_per)) + n_upper = b.select( + b.cmp_lt(n_upper, b.const_i32(n_ktiles)), + n_upper, + b.const_i32(n_ktiles), + ) + else: + n_upper = b.const_i32(n_ktiles) + + # Sliding-window start tile (see default builder). SW==0 -> start_tile=0. + if causal and SW > 0: + _diag0 = b.mul(qb, b.const_i32(n_per)) + _lo_raw = b.sub(_diag0, b.const_i32(SWt)) + start_tile = b.select( + b.cmp_gt(_lo_raw, b.const_i32(0)), _lo_raw, b.const_i32(0) + ) + else: + start_tile = b.const_i32(0) + start_buf = b.mod(start_tile, b.const_i32(NBUF)) + start_buf1 = b.mod(b.add(start_tile, b.const_i32(1)), b.const_i32(NBUF)) + + load_tile(start_buf, start_tile) + load_tile(start_buf1, b.add(start_tile, b.const_i32(1))) + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + s0 = do_qk(start_buf) + if causal and SW > 0: + do_mask(s0, start_tile, lower=True, upper=True) + else: + do_mask(s0, start_tile) + if RAG_KBOUND: + do_kbound_mask(s0, start_tile) + m0, _alpha0, p0, l0 = softmax_stats(s0, neg_inf) + o0 = [b.zero_vec_f32(16) for _ in range(D_TILES)] + pk0 = relayout_p(p0) + + iter_args = ( + [("m", m0), ("l", l0)] + + [(f"o{dt}", o0[dt]) for dt in range(D_TILES)] + + [(f"pk{kk}", pk0[kk]) for kk in range(KK_STEPS)] + ) + + if causal and SW > 0: + # Sliding-window three-phase band loop (prologue did start_tile). + diag_start = b.mul(qb, b.const_i32(n_per)) + a = b.add(start_tile, b.const_i32(1)) + left_end = b.add(diag_start, b.const_i32(n_per - SWt)) + + def _clamp(x, lo, hi): + x = b.select(b.cmp_lt(x, lo), lo, x) + x = b.select(b.cmp_lt(x, hi), x, hi) + return x + + mid_lo = _clamp(left_end, a, n_upper) + mid_hi = _clamp(diag_start, mid_lo, n_upper) + + phL = b.scf_for_iter(a, mid_lo, b.const_i32(1), iter_args, iv_name="swl") + with phL as (j, carry): + emit_loop_body(j, carry, mask_lower=True, mask_upper=True) + mid_args = [ + (name + "_m", val) for (name, _), val in zip(iter_args, phL.results) + ] + phM = b.scf_for_iter( + mid_lo, mid_hi, b.const_i32(1), mid_args, iv_name="swm" + ) + with phM as (j, carry): + emit_loop_body(j, carry) + rgt_args = [ + (name + "_r", val) for (name, _), val in zip(iter_args, phM.results) + ] + loop = b.scf_for_iter( + mid_hi, n_upper, b.const_i32(1), rgt_args, iv_name="swr" + ) + with loop as (j, carry): + emit_loop_body(j, carry, mask_lower=True, mask_upper=True) + elif causal: + diag_start = b.mul(qb, b.const_i32(n_per)) + body_upper = b.select(b.cmp_lt(diag_start, n_upper), diag_start, n_upper) + body = b.scf_for_iter( + b.const_i32(1), body_upper, b.const_i32(1), iter_args, iv_name="nb" + ) + with body as (j, carry): + emit_loop_body(j, carry) + tail_args = [ + (name + "_t", val) for (name, _), val in zip(iter_args, body.results) + ] + tail_lo = b.select( + b.cmp_lt(diag_start, b.const_i32(1)), b.const_i32(1), diag_start + ) + loop = b.scf_for_iter( + tail_lo, n_upper, b.const_i32(1), tail_args, iv_name="nt" + ) + with loop as (j, carry): + emit_loop_body(j, carry, mask_upper=True) + else: + loop = b.scf_for_iter( + b.const_i32(1), n_upper, b.const_i32(1), iter_args, iv_name="nkt" + ) + with loop as (j, carry): + emit_loop_body(j, carry, mask_kbound=RAG_KBOUND) + + res = loop.results + l_i = res[1] + o_acc = list(res[2 : 2 + D_TILES]) + p_prev = list(res[2 + D_TILES : 2 + D_TILES + KK_STEPS]) + + # PF: drain the last iter's in-flight V prefetch before the epilogue do_pv. + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + last_vbuf = b.mod(b.add(n_upper, b.const_i32(NBUF - 1)), b.const_i32(NBUF)) + o_acc = do_pv(o_acc, p_prev, last_vbuf) + + # Epilogue: recompute (bt, hq) from the live loop IV so they need not cross + # the KV loop (keeps the loop-carried live set minimal -> 0 spill). Must + # mirror the work-item decode used at the top of the loop. + rcp_l = b.rcp(l_i) + bt_e = b.mod(wi, b.const_i32(B)) + if spec.resolved_persist_decode == "hkv_major": + rem_e = b.div(wi, b.const_i32(B)) + hql_e = b.mod(rem_e, b.const_i32(gqa)) + hkv_e = b.div(b.div(rem_e, b.const_i32(gqa)), b.const_i32(NQB)) + hq_e = b.add(b.mul(hkv_e, b.const_i32(gqa)), hql_e) + else: + hq_e = b.mod(b.div(wi, b.const_i32(B)), b.const_i32(Hq)) + o_base = b.add( + b.mul(b.mul(bt_e, b.const_i32(Sq)), b.const_i32(stride_q_tok)), + b.mul(hq_e, b.const_i32(D)), + ) + qtok = b.add(q_tok0, _mfma_32x32_c_col(b, lane, 0)) + q_row_byte = b.add(o_base, b.mul(qtok, b.const_i32(stride_q_tok))) + d_half = b.mul(lane_h, b.const_i32(4)) + # ragged: guard padded query rows (qtok >= seqlen_q) so they never write. + o_store_ctx = ( + b.scf_if(b.cmp_lt(qtok, b.const_i32(Sq))) if RAGGED else _nullcontext() + ) + with o_store_ctx: + for dt in range(D_TILES): + for g in range(4): + d0 = b.add(b.const_i32(dt * 32 + g * 8), d_half) + addr = b.add(q_row_byte, d0) + vals = [ + b.cast_f32_to( + b.fmul(b.vec_extract(o_acc[dt], g * 4 + kk), rcp_l), dtype + ) + for kk in range(4) + ] + b.global_store_vN(o, addr, b.vec_pack(vals, dtype), 4, align=8) + + b.ret() + return b.kernel + + +# --------------------------------------------------------------------------- # +# Public launch geometry + ABI (promoted from the prefill builder so the kernel +# is dispatchable / framework-callable without the host script). +# --------------------------------------------------------------------------- # + + +def attention_dense_grid(spec: AttentionDenseSpec) -> Tuple[int, int, int]: + """Launch grid for ``spec``. Persistent = 1-D grid of ``num_persistent`` CTAs; + default = one CTA per (query-block, query-head, batch).""" + if spec.persistent: + return (spec.num_persistent, 1, 1) + nqb = (spec.seqlen_q + _BLOCK_M - 1) // _BLOCK_M # ceil: ragged partial block + return (nqb, spec.num_query_heads, spec.batch) + + +def attention_dense_block(spec: AttentionDenseSpec) -> Tuple[int, int, int]: + """CTA block dims: ``num_waves`` wave64s (= 512 threads).""" + return (spec.num_waves * 64, 1, 1) + + +def attention_dense_signature(spec: AttentionDenseSpec): + """ABI signature for :class:`KernelLauncher`. q/k/v/o pointers + f32 scale, + plus the two ``cu_seqlens`` i32 pointers when ``spec.varlen`` (the kernel + emits a 7-arg ABI in that case -- see :func:`build_attention_dense`).""" + from rocke.helpers.spec import SignatureBuilder + + sig = ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .scalar("scale", "f32") + ) + if spec.varlen: + sig = sig.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32") + return sig.build() + + +_DENSE_LAUNCHER_CACHE: dict = {} + + +def align_up(n: int, mult: int) -> int: + """Round ``n`` up to the next multiple of ``mult`` (kernel tile alignment).""" + return ((int(n) + mult - 1) // mult) * mult + + +def run_attention_dense_torch( + *, + spec: AttentionDenseSpec, + q, + k, + v, + out, + scale: float, + stream: int = 0, + arch: str = "gfx950", + cu_seqlens_q=None, + cu_seqlens_kv=None, +): + """High-level framework entry: compile (cached) + launch the dense prefill + kernel on torch tensors. ``q``/``k``/``v``/``out`` are dense contiguous + tensors ([B, S, H, D] for q/out, [B, Skv, Hkv, D] for k/v); ``scale`` is the + softmax scale (1/sqrt(D)). Returns ``out``. torch is imported lazily by the + launcher — this module stays torch-free at import time. + + Arbitrary (non-256-multiple) sequence lengths are served WITHOUT host + padding by the in-kernel ragged path: build ``spec`` with ``ragged=True`` + and the TRUE (un-rounded) ``seqlen_q``/``seqlen_kv`` and pass the true-length + q/k/v/out tensors. The kernel pads the boundary tiles on-chip (register-zero + OOB query rows, LDS-zero OOB keys) and drops the partial O rows; the grid is + ceil-sized automatically. See the ``ragged`` spec field. + + Varlen (``spec.varlen``): the kernel emits a 7-arg ABI (packed + ``[total_tok, H, D]`` q/k/v/o + two int32 ``cu_seqlens`` [batch+1]); pass both + ``cu_seqlens_q`` and ``cu_seqlens_kv`` or a ``ValueError`` is raised (they are + required — never silently launch the 5-arg ABI against a 7-arg kernel).""" + ok, why = supports_attention_dense(spec, arch=arch) + if not ok: + raise NotImplementedError(f"attention_dense unsupported for spec: {why}") + if spec.varlen and (cu_seqlens_q is None or cu_seqlens_kv is None): + raise ValueError( + "varlen=True requires cu_seqlens_q and cu_seqlens_kv (int32 [batch+1]); " + "the varlen kernel has a 7-arg ABI and cannot be launched with q/k/v/o/scale" + ) + if not spec.varlen and (cu_seqlens_q is not None or cu_seqlens_kv is not None): + raise ValueError("cu_seqlens_* provided but spec.varlen is False") + from rocke.helpers.compile import compile_kernel + from rocke.runtime import KernelLauncher, LaunchConfig + + key = spec.kernel_name() + launcher = _DENSE_LAUNCHER_CACHE.get(key) + if launcher is None: + art = compile_kernel( + build_attention_dense(spec, arch=arch), + arch=arch, + backend="python", + capture_ir_text=False, + ) + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=attention_dense_signature(spec), + ) + _DENSE_LAUNCHER_CACHE[key] = launcher + vals = {"q_ptr": q, "k_ptr": k, "v_ptr": v, "o_ptr": out, "scale": float(scale)} + if spec.varlen: + vals["cu_seqlens_q"] = cu_seqlens_q + vals["cu_seqlens_kv"] = cu_seqlens_kv + launcher( + vals, + config=LaunchConfig( + grid=attention_dense_grid(spec), + block=attention_dense_block(spec), + stream=int(stream), + ), + ) + return out diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_pipe_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_pipe_wiring.py index d559c25ad71a..a7c5cc4fb5d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_pipe_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_pipe_wiring.py @@ -291,11 +291,13 @@ def test_bf16_ring_enabled_for_d64_prefill(self): ) self.assertTrue(au._enable_gfx942_flash_k_sliced_ring(p)) - def test_bf16_ring_disabled_for_d128(self): - # D128 bf16 ring is excluded (kernel generator bug). + def test_bf16_ring_enabled_for_d128(self): + # D128 bf16 now shares the sliced-K ring path: the earlier generator-bug + # exclusion was fixed on develop (ring uses the fp16-flash nw4+cfvst + # geometry; verified numerically correct, byte-identity gate green). with _Gfx942Arch(): p = self._bf16_problem(seqlen_q=2048, seqlen_k=2048) - self.assertFalse(au._enable_gfx942_flash_k_sliced_ring(p)) + self.assertTrue(au._enable_gfx942_flash_k_sliced_ring(p)) def test_bf16_mask_limit_enabled_for_prefill(self): with _Gfx942Arch(): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json new file mode 100644 index 000000000000..1283453442bd --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json @@ -0,0 +1,157 @@ +{ + "flavors": { + "llvm20": { + "cases": { + "attention_dense/bn128_sq512": { + "bytes": 610799, + "sha256": "67feb714296db5af1d332df6427ea306727f0ec72215982d6048da5c398cc6b9" + }, + "attention_dense/default_causal_sq512": { + "bytes": 341939, + "sha256": "1dc714eeaa2e69c3198eeadc725956f2c5257db2f6263f524dc3930ed707651d" + }, + "attention_dense/fp16_h64_sq512": { + "bytes": 210000, + "sha256": "4c7f461c2e10edd51ce0386ba4cf2a8ac4a798656c7ac93b730dc2effb6eb85e" + }, + "attention_dense/lazy_off_sq512": { + "bytes": 337008, + "sha256": "281d6f975a9056969fb059357faf048a7a11d56557ebba075652aa81c8a6b5cb" + }, + "attention_dense/noncausal_sq512": { + "bytes": 217781, + "sha256": "db33f2ff07c424d2b3e759aed66e4b0d29056e7e4780ce11fdcc8d199da99334" + }, + "attention_dense/persist_h64_sq512": { + "bytes": 214702, + "sha256": "b4daea7ca97abd002b6a50d217a43b0a443263d8cb11c5192193b329caf7e51c" + }, + "attention_dense/persist_hkvmaj_sq512": { + "bytes": 327546, + "sha256": "f49a0e2f5415bce83f55d251c2a5a7be00fcbd999db3b038023f50ab53113f74" + }, + "attention_dense/persist_intl_sq512": { + "bytes": 327020, + "sha256": "b1bde5a27e66647c342d5ee4b2e010e52106d33d22d901cedf46f59e3ce9fca8" + }, + "attention_dense/persist_lazy_off_sq512": { + "bytes": 321142, + "sha256": "7f5251827f8e9e4853c16b0908950cbf1b878686007643438c9beb241254d6a3" + }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 326109, + "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" + }, + "attention_dense/persist_ragged_sq500": { + "bytes": 328078, + "sha256": "0291e1ac1a08bb6bb05274285754b15137b9ad855cf02e43c856a040e111a8aa" + }, + "attention_dense/persist_swa_w128_sq512": { + "bytes": 444213, + "sha256": "3a2039ebb82f66fbf375f8ff9b636ce5cda48702a914541b854d5997ce2de222" + }, + "attention_dense/persistent_causal_sq512": { + "bytes": 326109, + "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" + }, + "attention_dense/ragged_causal_sq500": { + "bytes": 343927, + "sha256": "616a5fa62b3dbe3d66a21d84029874d93773a517ad69575230afdf46ca5fa8a4" + }, + "attention_dense/ragged_full_sq500": { + "bytes": 248062, + "sha256": "411b18f74bf637dfce4b19e8972163cd9fb202da9943112e9ecd8ddc0b2306e3" + }, + "attention_dense/ragged_h64_sq500": { + "bytes": 214541, + "sha256": "97c05c4bdd44b4b5c46c4e0e16272fdb35a1f68c81f4537fcf61a3faf4290593" + }, + "attention_dense/swa_w128_sq512": { + "bytes": 464090, + "sha256": "cbde8ed9be5ac5b89b68305c8b4092630bf5ee56ccd58cd2c8e70228879b1f1d" + }, + "attention_dense/varlen_sq512": { + "bytes": 344011, + "sha256": "afacdff3e337450ae8db2ec59bf87d5aea6949a23e83256bc4af8b97f63da80a" + } + } + }, + "llvm22": { + "cases": { + "attention_dense/bn128_sq512": { + "bytes": 610903, + "sha256": "b75fea02477fec773dc639131b94c33efda4ea2217f7822622c3b69dba42c165" + }, + "attention_dense/default_causal_sq512": { + "bytes": 342043, + "sha256": "0698cf388935622a76524c20f300b383a946ecbdb512e93c662db29ad574a1f0" + }, + "attention_dense/fp16_h64_sq512": { + "bytes": 210100, + "sha256": "b920ada5e3c48c8083a8c8d8c904fb685ac49eac66a7d4fbb67329a0cbb4f014" + }, + "attention_dense/lazy_off_sq512": { + "bytes": 337112, + "sha256": "11a2458a4ec86709fcec37da59d48c1de6a8b55ca4f23e5fbb27f13e5388222d" + }, + "attention_dense/noncausal_sq512": { + "bytes": 217881, + "sha256": "83333cfcfd4d2e3c26ba2cf6e61faffc8dce46ad43cc1216d7efd68ad0653e54" + }, + "attention_dense/persist_h64_sq512": { + "bytes": 214802, + "sha256": "7d4c6fce84a638779797c53d6b57f60e2619076aaeb0ec9e557180be8bc902a0" + }, + "attention_dense/persist_hkvmaj_sq512": { + "bytes": 327650, + "sha256": "8a0e7f52f54566f579cf9f0f20d074b0a57f5079c242f6a6b98d35ab9cd000a4" + }, + "attention_dense/persist_intl_sq512": { + "bytes": 327124, + "sha256": "6d9f80a763c0d4d992159487222d0ca30e8ff588e1a9c0347969e94e6b4a4373" + }, + "attention_dense/persist_lazy_off_sq512": { + "bytes": 321246, + "sha256": "6c805ce2f02d540afb7f4e36f6cb55edb19d3e0526d2b2880e8a0c2beb1bea4e" + }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 326213, + "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" + }, + "attention_dense/persist_ragged_sq500": { + "bytes": 328227, + "sha256": "5ce0e47b2f540e77344c8439293ac8b2041e242001ef5017d71cdb423bf22230" + }, + "attention_dense/persist_swa_w128_sq512": { + "bytes": 444317, + "sha256": "6fe644368ff3c588edfcf80eb74e6ae9264458a9032b40ca78fdfb93e6f6c403" + }, + "attention_dense/persistent_causal_sq512": { + "bytes": 326213, + "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" + }, + "attention_dense/ragged_causal_sq500": { + "bytes": 344076, + "sha256": "24245154a81457eef7ea50f91a5f75e362cd3ef8fee69ca827449efbd8e6ff90" + }, + "attention_dense/ragged_full_sq500": { + "bytes": 248205, + "sha256": "653696a9ad4e05b199b22c126434ac7391cfa09f4e320d84563d6083b523218f" + }, + "attention_dense/ragged_h64_sq500": { + "bytes": 214684, + "sha256": "8b37f349d2fbdf2e01ed3005b12b96690ba8fc3883e1aa0da01c1831ec863e52" + }, + "attention_dense/swa_w128_sq512": { + "bytes": 464194, + "sha256": "35c4abc272427bc1f92ca1cfe06256e7e1a02e8a2690d21cec105c32269fc1c0" + }, + "attention_dense/varlen_sq512": { + "bytes": 344115, + "sha256": "a9bf509c16f6fde382fe7f778b3508f2b385b67011ed0f778f81035a12372998" + } + } + } + }, + "schema": "attention_dense.ir_golden_sha256/v1" +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_golden.py new file mode 100644 index 000000000000..c30caa3da832 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_golden.py @@ -0,0 +1,212 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Golden LLVM-IR byte-stability test for the gfx950 dense flash-attn kernel. + +Hashes the Python-lowered LLVM IR (SHA256) of representative ``attention_dense`` +specs and compares against a checked-in per-flavor golden fixture, catching any +unintended codegen drift. Pure text lowering — no GPU / no comgr required. + +NOT WIRED INTO CI (by request): this file lives under ``library/tests/`` (which +``platform/tests/run_all.py``'s CI gate does NOT collect — it only pytests +``platform/tests/``), registers NO byte-identity ``*_emit`` parity pair, adds NO +case to ``rocke_ir_parity_harness.cases()`` / the platform golden, and adds NO +``add_test(...)`` CMake entry. Run it manually: + + cd rocke/library + PYTHONPATH=../platform/python:. python -m pytest tests/test_attention_dense_golden.py + +Re-bless after an intended codegen change: + + cd rocke/library + PYTHONPATH=../platform/python:. python tests/test_attention_dense_golden.py --write +""" +import hashlib +import json +import sys +from pathlib import Path + +_GOLDEN = Path(__file__).resolve().parent / "golden" / "attention_dense_ir_sha256.json" +_FLAVORS = ("llvm20", "llvm22") + + +def _cases(): + """cid -> zero-arg builder returning a KernelDef. Small Sq keeps the IR compact + while still exercising the full pipeline (both grid variants).""" + from kernels.gfx950.attention_dense import ( + AttentionDenseSpec, + build_attention_dense, + ) + + base = dict( + batch=1, + seqlen_q=512, + seqlen_kv=512, + num_query_heads=128, + num_kv_heads=8, + head_size=128, + causal=True, + dtype="bf16", + ) + + def mk(**over): + d = dict(base) + d.update(over) + return lambda: build_attention_dense(AttentionDenseSpec(**d)) + + return { + # --- default (one CTA per q-block/head) grid --- + "attention_dense/default_causal_sq512": mk(), + "attention_dense/swa_w128_sq512": mk(sliding_window=128), + "attention_dense/varlen_sq512": mk(varlen=True), + "attention_dense/lazy_off_sq512": mk(lazy_rescale=False), + "attention_dense/fp16_h64_sq512": mk(dtype="fp16", head_size=64), + "attention_dense/bn128_sq512": mk(block_n=128), + "attention_dense/noncausal_sq512": mk(causal=False), + # --- persistent (grid-stride) grid + decode variants --- + "attention_dense/persistent_causal_sq512": mk( + persistent=True, num_persistent=256 + ), + "attention_dense/persist_qbmaj_sq512": mk( + persistent=True, num_persistent=256, persist_decode="qb_major" + ), + "attention_dense/persist_hkvmaj_sq512": mk( + persistent=True, num_persistent=256, persist_decode="hkv_major" + ), + "attention_dense/persist_intl_sq512": mk( + persistent=True, num_persistent=256, interleave=True + ), + "attention_dense/persist_lazy_off_sq512": mk( + persistent=True, num_persistent=256, lazy_rescale=False + ), + "attention_dense/persist_swa_w128_sq512": mk( + persistent=True, num_persistent=256, sliding_window=128 + ), + # D=64 packed-row DMA loader (2 rows/instr, unpadded LDS) on the + # persistent builder -- locks the head_size=64 fix (fp16_h64 above only + # exercises the default builder). + "attention_dense/persist_h64_sq512": mk( + persistent=True, num_persistent=256, head_size=64 + ), + # ragged (non-256 seqlen) in-kernel path: on-chip boundary padding. + # causal (no key mask), non-causal (ktok {got}") + assert not drift, "attention_dense IR drift vs golden:\n " + "\n ".join(drift) + + +def test_attention_dense_cpp_python_byte_identity(): + """Stream-4 parity gate: the C++ engine (rocke_engine) lowers every + dense-prefill variant to byte-identical LLVM IR vs the Python lowerer. + + Both sides go through ``_lower_llvm_via_backend`` so they resolve the same + llvm flavor; ``ROCKE_CPP_STRICT=1`` disables the silent python fallback so a + missing/stale C++ engine surfaces as a skip (not a false pass). Requires the + C++ ``exp2_fast`` op (added for the dense softmax hot path). + """ + import os + + import pytest + + try: + from rocke.helpers.compile import _lower_llvm_via_backend + except Exception as e: # pragma: no cover + pytest.skip(f"backend lowering unavailable: {e}") + + prev = os.environ.get("ROCKE_CPP_STRICT") + os.environ["ROCKE_CPP_STRICT"] = "1" + mism = [] + try: + for cid, build in _cases().items(): + k = build() + py = _lower_llvm_via_backend(k, arch="gfx950", backend="python", spec=None) + try: + cpp = _lower_llvm_via_backend( + k, arch="gfx950", backend="cpp", spec=None + ) + except Exception as e: # C++ engine not built / opcode gap + pytest.skip(f"C++ engine unavailable ({cid}): {str(e)[:140]}") + if py != cpp: + mism.append(cid) + finally: + if prev is None: + os.environ.pop("ROCKE_CPP_STRICT", None) + else: + os.environ["ROCKE_CPP_STRICT"] = prev + assert not mism, "attention_dense cpp/python IR byte-mismatch:\n " + "\n ".join( + mism + ) + + +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_attention_dense_ir_matches_golden() + print("PASS") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index 9928abf74400..a4f45e363c06 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -470,6 +470,7 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = { /* math.* */ "math.exp2", + "math.exp2_fast", "math.log2", "math.rcp", "math.rcp_fast", @@ -675,6 +676,7 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = { /* math.* */ /* math.exp2 */ true, + /* math.exp2_fast*/ true, /* math.log2 */ true, /* math.rcp */ true, /* math.rcp_fast */ true, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp index 9837b63720cf..910fe60eb2fc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp @@ -422,6 +422,11 @@ rocke_value_t* rocke_b_exp2(rocke_ir_builder_t* b, rocke_value_t* a) return rocke_i_unop(b, ROCKE_OP_MATH_EXP2, a, "exp2"); } +rocke_value_t* rocke_b_exp2_fast(rocke_ir_builder_t* b, rocke_value_t* a) +{ + return rocke_i_unop(b, ROCKE_OP_MATH_EXP2_FAST, a, "exp2f"); +} + rocke_value_t* rocke_b_log2(rocke_ir_builder_t* b, rocke_value_t* a) { return rocke_i_unop(b, ROCKE_OP_MATH_LOG2, a, "log2"); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp index 8086a7e22b44..6655d76d8dd0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp @@ -609,6 +609,11 @@ static void _op_math_exp2(rocke_lower_t* L, const rocke_op_t* op) { ll_math_f32_unary(L, op, "exp2", "exp2.f32", "exp2.f32"); } +/* Python _op_math_exp2_fast -> llvm.amdgcn.exp2.f32 (native v_exp_f32, no guard). */ +static void _op_math_exp2_fast(rocke_lower_t* L, const rocke_op_t* op) +{ + ll_math_f32_unary(L, op, "exp2_fast", "amdgcn.exp2.f32", "amdgcn.exp2.f32"); +} /* Python _op_math_log2. */ static void _op_math_log2(rocke_lower_t* L, const rocke_op_t* op) { @@ -738,6 +743,7 @@ void rocke_ll_register_arith(void) rocke_ll_set_handler(ROCKE_OP_ARITH_UMUL_HI_I32, _op_arith_umul_hi_i32); rocke_ll_set_handler(ROCKE_OP_MATH_EXP2, _op_math_exp2); + rocke_ll_set_handler(ROCKE_OP_MATH_EXP2_FAST, _op_math_exp2_fast); rocke_ll_set_handler(ROCKE_OP_MATH_LOG2, _op_math_log2); rocke_ll_set_handler(ROCKE_OP_MATH_RCP, _op_math_rcp); rocke_ll_set_handler(ROCKE_OP_MATH_RCP_FAST, _op_math_rcp_fast); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp index d98b507c08d6..eb0033cc356e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp @@ -77,6 +77,7 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"workgroup.z", "declare i32 @llvm.amdgcn.workgroup.id.z()"}, {"s.barrier", "declare void @llvm.amdgcn.s.barrier()"}, {"exp2.f32", "declare float @llvm.exp2.f32(float)"}, + {"amdgcn.exp2.f32", "declare float @llvm.amdgcn.exp2.f32(float)"}, {"log2.f32", "declare float @llvm.log2.f32(float)"}, {"sqrt.f32", "declare float @llvm.sqrt.f32(float)"}, {"rsqrt.f32", "declare float @llvm.amdgcn.rsq.f32(float)"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index 837dc317ebae..d6520a8b7a2a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -264,6 +264,7 @@ static int is_unary_same_type(rocke_opcode_t o) case ROCKE_OP_ARITH_FABS: case ROCKE_OP_ARITH_NOT: case ROCKE_OP_MATH_EXP2: + case ROCKE_OP_MATH_EXP2_FAST: case ROCKE_OP_MATH_LOG2: case ROCKE_OP_MATH_RCP: case ROCKE_OP_MATH_RCP_FAST: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h index 33309461c0dd..d78a4bb78ff1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h @@ -251,6 +251,7 @@ typedef enum rocke_opcode /* math.* */ ROCKE_OP_MATH_EXP2, + ROCKE_OP_MATH_EXP2_FAST, ROCKE_OP_MATH_LOG2, ROCKE_OP_MATH_RCP, ROCKE_OP_MATH_RCP_FAST, @@ -682,6 +683,7 @@ rocke_value_t* /* ----- math ----- */ rocke_value_t* rocke_b_exp2(rocke_ir_builder_t* b, rocke_value_t* a); +rocke_value_t* rocke_b_exp2_fast(rocke_ir_builder_t* b, rocke_value_t* a); rocke_value_t* rocke_b_log2(rocke_ir_builder_t* b, rocke_value_t* a); rocke_value_t* rocke_b_rcp(rocke_ir_builder_t* b, rocke_value_t* a); rocke_value_t* rocke_b_rcp_fast(rocke_ir_builder_t* b, rocke_value_t* a); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index bbf4064ab14e..e7722452a834 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -512,6 +512,20 @@ def fmin(self, a: Value, b: Value) -> Value: def exp2(self, a: Value) -> Value: return self._op("math.exp2", [a], [a.type], result_name_hint="exp2").result + def exp2_fast(self, a: Value) -> Value: + """Native single-instruction base-2 exp: ``llvm.amdgcn.exp2.f32`` -> one + ``v_exp_f32``. Unlike ``exp2`` (``llvm.exp2.f32``), it emits NO + overflow/underflow guard (the ~5-VALU range-reduction clamp), so the + caller must guarantee the argument stays in the hardware's valid range. + + Safe for softmax where the argument is always <= 0 (``s - m_new`` and + ``m_i - m_new``): no overflow, and ``v_exp_f32`` already flushes large + negatives to 0. Matches FlyDSL's ``rocdl.exp2`` emission. f32 only. + """ + return self._op( + "math.exp2_fast", [a], [a.type], result_name_hint="exp2f" + ).result + def log2(self, a: Value) -> Value: return self._op("math.log2", [a], [a.type], result_name_hint="log2").result @@ -3525,6 +3539,7 @@ def scf_if(self, cond: Value): "arith.sitofp_f32", "arith.cvt_fp8_to_f32", "math.exp2", + "math.exp2_fast", "math.log2", "math.rcp", "math.rcp_fast", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index 6c2585510610..ec7c3ec1f669 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -273,6 +273,7 @@ def _resolve_llvm_flavor() -> str: "ptr addrspace(1) nocapture, ptr addrspace(3) nocapture, i32 immarg, i32 immarg)" ), "exp2.f32": "declare float @llvm.exp2.f32(float)", + "amdgcn.exp2.f32": "declare float @llvm.amdgcn.exp2.f32(float)", "log2.f32": "declare float @llvm.log2.f32(float)", "sqrt.f32": "declare float @llvm.sqrt.f32(float)", "rsqrt.f32": "declare float @llvm.amdgcn.rsq.f32(float)", @@ -1646,6 +1647,17 @@ def _op_math_exp2(self, op: Op) -> None: f" {op.result.name} = call float @llvm.exp2.f32(float {self._operand(v)})" ) + def _op_math_exp2_fast(self, op: Op) -> None: + (v,) = op.operands + if v.type.name != "f32": + raise NotImplementedError("math.exp2_fast currently supports f32") + # Native single-instruction exp2 (v_exp_f32), no overflow guard. + self._need("amdgcn.exp2.f32") + self._current().emit( + f" {op.result.name} = call float " + f"@llvm.amdgcn.exp2.f32(float {self._operand(v)})" + ) + def _op_math_log2(self, op: Op) -> None: (v,) = op.operands if v.type.name != "f32":