From 3e7982b8f694bf6ab243869cb50517ab76b8edd7 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Wed, 15 Jul 2026 23:15:24 +0000 Subject: [PATCH 01/11] (rocke): performant dense flash attention implementatoin that gives ~840 tflops --- .../gfx950/attention/prefill/README.md | 79 ++ .../prefill/attention_dense_prefill.py | 168 +++ .../library/kernels/gfx950/attention_dense.py | 997 ++++++++++++++++++ .../rocke/platform/python/rocke/core/ir.py | 15 + .../platform/python/rocke/core/lower_llvm.py | 12 + 5 files changed, 1271 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py 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..153c09ea7ac9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md @@ -0,0 +1,79 @@ +# 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 (qb-major decode for causal +load-balance), 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). + +## Measured (MI355X, bf16, D=128, Hq=128, Hkv=8, causal) + +| Sq | default | persistent (NP=256) | +|----|--------:|--------------------:| +| 2048 | ~410 TFLOPS | ~565 TFLOPS | +| 8192 | ~521 TFLOPS | **~850 TFLOPS** | + +Both paths are 0 VGPR spill and bit-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" +``` + +## 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 currently lowered on the LLVM-direct backend only (not yet mirrored in the + C++ engine); byte-identity coverage for it is a follow-up. 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..54b9a24d9785 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/attention_dense_prefill.py @@ -0,0 +1,168 @@ +#!/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, + ) + 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") + .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() + 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 + flops = ( + 4 * B * Hq * D * (Sq * (Sq + 1) // 2) + if spec.causal + else 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") + 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, + ) + run(spec) + + +if __name__ == "__main__": + main() 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..a0d9b1bd42bb --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py @@ -0,0 +1,997 @@ +"""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. + * **vectorized O store**. + +Measured on MI355X: **521 TFLOPS @ Sq=8192, bf16, D=128, causal, 0 spill**, error +~1.46e-3 vs SDPA. Shape (batch/seqlen/heads/head_dim) is baked at build time (dense, +compile-time-sized ABI); only the KV tile and occupancy hint are tunable knobs. + +Experimental/negative levers from the sweep (step-2 8-cluster, K-staging, per-nsub +staging, score truncation, s_setprio, PV V-prefetch, lazy rescale) are intentionally +NOT carried over — see the experiment's ``plan.md`` for their measured results. +""" + +from dataclasses import dataclass +from typing import Optional, Tuple + +from rocke.core.ir import IRBuilder, KernelDef, PtrType, BF16, F16, F32, I64 +from rocke.helpers.attention import mfma_32x32x16_for_dtype, pv32_v_load_paired +from rocke.helpers.schedule import MFMA, VALU, TRANS, DS_READ +from rocke.helpers.spec import kernel_name_join +from kernels.gfx950.attention_tiled_2d import _mfma_32x32_c_row, _mfma_32x32_c_col + +LOG2E = 1.4426950408889634 +_DTYPE_IR = {"bf16": BF16, "fp16": F16} + +# Baked pipeline constants (NOT tunable knobs — these are load-bearing): +# _BLOCK_M: query rows per CTA. The causal mask + P relayout assume 256; the +# kernel FAULTS at other values until those hardcodes are lifted. +# num_waves = _BLOCK_M // 32 = 8 (block = 512 threads). +# _NBUF=2 double-buffer (NBUF=3 is a measured dead end: 256 VGPR + 58 spills). +# _LDS_PAD=8 bf16 elements of K-row padding (the +80% bank-conflict fix). +_BLOCK_M = 256 +_NBUF = 2 +_LDS_PAD = 8 + + +@dataclass(frozen=True) +class AttentionDenseSpec: + """Compile-time spec for the dense flash-attention prefill kernel. + + Functional fields (batch / seqlen / heads / head_size / causal / dtype) are baked + into the kernel as constants — this is a dense, statically-sized ABI. ``block_n`` + and ``waves_per_eu`` are the only performance knobs; every algorithmic lever is + always-on (see the module docstring). + """ + + # --- functional (compile-time shape) --- + batch: int + seqlen_q: int + seqlen_kv: int + num_query_heads: int + num_kv_heads: int + head_size: int + causal: bool = True + dtype: str = "bf16" + + # --- validated performance knobs --- + # block_n: KV tile length. 64 (66 KB LDS, WPE-tunable) and 128 (135 KB LDS, pins + # the 256-VGPR cap) both match ~peak; 64 is strictly more resource-efficient. + block_n: int = 64 + # waves_per_eu: occupancy hint. 2 is a free win (tighter allocation, still 2 + # waves/SIMD); 3 is a measured trap (VGPR<=170 forces spills -> -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 + + 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}" + ) + if self.head_size % 32 != 0: + raise ValueError( + f"head_size must be a multiple of 32, 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.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}" + ) + + @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 + + 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.persistent: + parts.append(f"persist{self.num_persistent}") + 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 + + 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 + + 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 # 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))) + 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), unpadded V --- + LDROW = D + PAD + K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") + V_lds = b.smem_alloc(dtype, [NBUF, BN, D], name_hint="Vlds") + + # Q packs (B operand), scaled once by qk_scale = softmax_scale * log2(e). + 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) + 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)) + + n_ktiles = Skv // BN + n_per = BLOCK_M // BN + + K_BYTES_PER_BUF = BN * LDROW * 2 + K_LDROW_BYTES = LDROW * 2 + V_BYTES_PER_BUF = BN * D * 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_k(lds_base, buf_val, tile_key0): + """Row-by-row async DMA into padded [BN, LDROW] K_lds layout.""" + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(K_BYTES_PER_BUF)) + 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(K_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( + k_rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 + ) + + def async_load_v(lds_base, buf_val, tile_key0): + """Contiguous async DMA into unpadded [BN, D] V_lds layout.""" + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) + base_off = b.add(v_wave_off_i64, buf_off) + for c in range(V_DMA_PASSES): + wave_base = b.smem_ptr_add( + lds_base, b.add(base_off, b.const_i64(c * WAVES * WAVE_BYTES)) + ) + flat = b.mul( + b.add(b.mul(b.const_i32(c), b.const_i32(WAVES * 64)), tid), + b.const_i32(8), + ) + krow = b.div(flat, b.const_i32(D)) + kcol = b.mod(flat, b.const_i32(D)) + gkey = b.add(tile_key0, krow) + voff = b.add(b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), kcol) + b.async_buffer_load_lds_addr( + v_rsrc, wave_base, b.mul(voff, b.const_i32(2)), zero_soff, 4 + ) + + 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): + 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)) + 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_le(ktok, query_tok), 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)) + m_new = b.fmax(m_i, tile_max) + alpha = _exp2(b.fsub(m_i, m_new)) + return m_new, alpha + + 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 + + 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) + + # Prologue: prime the K/V double buffer and compute tile 0. + load_tile(b.const_i32(0), b.const_i32(0)) + load_tile(b.const_i32(1), b.const_i32(1)) + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + s0 = do_qk(b.const_i32(0)) + do_mask(s0, b.const_i32(0)) + m0, _alpha0 = 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)] + ) + + def emit_loop_body(j, carry, masked): + 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)) + + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + s = do_qk(kbuf) + if masked: + do_mask(s, j) + m_new, alpha = softmax_max(s, m_i) + b.sched_barrier(0) # depth-1 fence: m_new region-live-in + o_acc, p_vals, l_tile = pv_fused_exp(o_acc, p_prev, vbuf_prev, s, m_new) + 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: + # Diagonal-only masking: below-diagonal tiles need no mask (~94% at Sq=8192). + 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, masked=False) + 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, masked=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, masked=False) + + 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]) + + 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: O = (P@V) / l, vectorized bf16 store. + rcp_l = b.rcp(l_i) + o_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)), + ) + 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)) + 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 + n_ktiles = Skv // BN + n_per = BLOCK_M // BN + NQB = Sq // BLOCK_M + W = NQB * Hq * B # total work items + + 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) + + LDROW = D + PAD + K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") + V_lds = b.smem_alloc(dtype, [NBUF, BN, D], name_hint="Vlds") + + K_BYTES_PER_BUF = BN * LDROW * 2 + K_LDROW_BYTES = LDROW * 2 + V_BYTES_PER_BUF = BN * D * 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) + + # ----- 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() + + # qb-MAJOR decode: wi = qb*(Hq*B) + hq*B + bt. Putting qb (the triangular + # causal cost index) in the MSB spreads cheap+expensive query blocks across + # each CTA under grid-stride; a qb-fast decode would alias qb to a constant + # per CTA when NP is a multiple of NQB (32x imbalance). + bt = b.mod(wi, b.const_i32(B)) + rem = b.div(wi, b.const_i32(B)) + hq = b.mod(rem, b.const_i32(Hq)) + qb0 = b.div(rem, b.const_i32(Hq)) + if INTERLEAVE and causal and NQB > 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) + 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_k(lds_base, buf_val, tile_key0): + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(K_BYTES_PER_BUF)) + 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(K_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( + k_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(K_LDROW_BYTES)) + voff = b.add(voff, b.const_i32(stride_k_tok)) + + def async_load_v(lds_base, buf_val, tile_key0): + buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) + base_off = b.add(v_wave_off_i64, buf_off) + for c in range(V_DMA_PASSES): + wave_base = b.smem_ptr_add( + lds_base, b.add(base_off, b.const_i64(c * WAVES * WAVE_BYTES)) + ) + flat = b.mul( + b.add(b.mul(b.const_i32(c), b.const_i32(WAVES * 64)), tid), + b.const_i32(8), + ) + krow = b.div(flat, b.const_i32(D)) + kcol = b.mod(flat, b.const_i32(D)) + gkey = b.add(tile_key0, krow) + voff = b.add( + b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), kcol + ) + b.async_buffer_load_lds_addr( + v_rsrc, wave_base, b.mul(voff, b.const_i32(2)), zero_soff, 4 + ) + + 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): + 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)) + 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_le(ktok, query_tok), 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)) + m_new = b.fmax(m_i, tile_max) + alpha = _exp2(b.fsub(m_i, m_new)) + return m_new, alpha + + def softmax_stats(s_reg, m_i): + m_new, alpha = 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 + + def emit_loop_body(j, carry, masked): + 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)) + + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + s = do_qk(kbuf) + if masked: + do_mask(s, j) + m_new, alpha = softmax_max(s, m_i) + b.sched_barrier(0) + o_acc, p_vals, l_tile = pv_fused_exp(o_acc, p_prev, vbuf_prev, s, m_new) + 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) + + load_tile(b.const_i32(0), b.const_i32(0)) + load_tile(b.const_i32(1), b.const_i32(1)) + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + s0 = do_qk(b.const_i32(0)) + do_mask(s0, b.const_i32(0)) + 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: + 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, masked=False) + 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, masked=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, masked=False) + + 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]) + + 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). + rcp_l = b.rcp(l_i) + bt_e = b.mod(wi, b.const_i32(B)) + 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)) + 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 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 0294db0998ee..4fc946285b4e 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 @@ -536,6 +536,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 @@ -3551,6 +3565,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": From 775ed04a3832bc819f98d82c7532db74e393ecfa Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Thu, 16 Jul 2026 03:10:20 +0000 Subject: [PATCH 02/11] feat(rocke): dense attn +877 TFLOPS (prefetch+setprio), dispatch + golden test - Kernel: add always-on partial-vmcnt software prefetch + PV-only s_setprio to attention_dense (both default and persistent paths). Bit-identical (err 1.46e-3), 0 spill. Sq=8192 causal: default 521->~543, persistent ~852->~877 TFLOPS. - Dispatch: register an opt-in gfx950 attention_dense KernelCandidate (algorithm="attention_dense", priority 3) that never auto-overrides unified_2d; promote attention_dense_grid/_block/_signature + a run_attention_dense_torch framework entry, exported from kernels/__init__. - Test: add a golden LLVM-IR SHA256 byte-stability test (both llvm flavors) under library/tests/, deliberately NOT wired into CI (no parity emit pair, no harness case, no CMake add_test; library/tests is outside the run_all.py gate). Co-authored-by: Cursor --- .../rocke/library/dispatch/attention.py | 90 +++++++++++++ .../rocke/library/kernels/__init__.py | 12 ++ .../library/kernels/gfx950/attention_dense.py | 121 +++++++++++++++++- .../golden/attention_dense_ir_sha256.json | 29 +++++ .../tests/test_attention_dense_golden.py | 117 +++++++++++++++++ 5 files changed, 364 insertions(+), 5 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_golden.py 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 a0bd12d5a44f..0df74f2841d2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py @@ -301,6 +301,96 @@ def select(req: OperatorRequest) -> AttentionSpec: ATTENTION_REGISTRY.register(_make_gfx942_dense_pipe_candidate()) +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). + Launch geometry + ABI are real (``attention_dense_grid`` / ``_block`` / + ``_signature``); end-to-end launch is ``run_attention_dense_torch``. + """ + spec_id = "gfx950_attention_dense" + name = "attention_gfx950_dense" + + def _dense_spec(req: OperatorRequest): + from kernels.gfx950.attention_dense import AttentionDenseSpec + + assert isinstance(req, AttentionRequest) + return AttentionDenseSpec( + batch=int(req.batch), + seqlen_q=int(req.seqlen_q), + seqlen_kv=int(req.seqlen_k), + 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(), + ) + + 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) + 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", + ) + + 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 + + +ATTENTION_REGISTRY.register(_make_gfx950_attention_dense_candidate()) + + def attention_candidates() -> Tuple[KernelCandidate, ...]: return ATTENTION_REGISTRY.candidates() 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 index a0d9b1bd42bb..12aac10bca99 100644 --- 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 @@ -16,11 +16,19 @@ * **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: **521 TFLOPS @ Sq=8192, bf16, D=128, causal, 0 spill**, error -~1.46e-3 vs SDPA. Shape (batch/seqlen/heads/head_dim) is baked at build time (dense, -compile-time-sized ABI); only the KV tile and occupancy hint are tunable knobs. +Measured on MI355X (bf16, D=128, causal, 0 spill, err ~1.46e-3 vs SDPA): +**~543 TFLOPS @ Sq=8192** for the default one-CTA-per-query-block grid, and +**~877 TFLOPS @ Sq=8192** for the persistent grid-stride variant +(``persistent=True``). 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. Experimental/negative levers from the sweep (step-2 8-cluster, K-staging, per-nsub staging, score truncation, s_setprio, PV V-prefetch, lazy rescale) are intentionally @@ -495,14 +503,18 @@ def emit_loop_body(j, carry, masked): 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)) - b.s_waitcnt(vmcnt=0) + # 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 masked: do_mask(s, j) m_new, alpha = 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) l_new = b.fadd(b.fmul(l_i, alpha), l_tile) o_acc = rescale_o(o_acc, alpha) p_packs = relayout_p(p_vals) @@ -540,6 +552,9 @@ def emit_loop_body(j, carry, masked): 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) @@ -894,14 +909,21 @@ def emit_loop_body(j, carry, masked): kbuf = b.mod(j, b.const_i32(NBUF)) vbuf_prev = b.mod(b.add(j, b.const_i32(NBUF - 1)), b.const_i32(NBUF)) - b.s_waitcnt(vmcnt=0) + # 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 masked: do_mask(s, j) m_new, alpha = 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) l_new = b.fadd(b.fmul(l_i, alpha), l_tile) o_acc = rescale_o(o_acc, alpha) p_packs = relayout_p(p_vals) @@ -966,6 +988,9 @@ def emit_loop_body(j, carry, masked): 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) @@ -995,3 +1020,89 @@ def emit_loop_body(j, carry, masked): 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) + return (spec.seqlen_q // _BLOCK_M, 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 (q/k/v/o pointers + f32 scale) for :class:`KernelLauncher`.""" + from rocke.helpers.spec import SignatureBuilder + + return ( + SignatureBuilder() + .ptr("q_ptr", spec.dtype) + .ptr("k_ptr", spec.dtype) + .ptr("v_ptr", spec.dtype) + .ptr("o_ptr", spec.dtype) + .scalar("scale", "f32") + .build() + ) + + +_DENSE_LAUNCHER_CACHE: dict = {} + + +def run_attention_dense_torch( + *, + spec: AttentionDenseSpec, + q, + k, + v, + out, + scale: float, + stream: int = 0, + arch: str = "gfx950", +): + """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.""" + ok, why = supports_attention_dense(spec, arch=arch) + if not ok: + raise NotImplementedError(f"attention_dense unsupported for spec: {why}") + 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)} + 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/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..5b745f34003f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json @@ -0,0 +1,29 @@ +{ + "flavors": { + "llvm20": { + "cases": { + "attention_dense/default_causal_sq512": { + "bytes": 320191, + "sha256": "454a7ab4aec162e1d7dccc3fa8ca5a186ee74508d7763275823bde47723532b3" + }, + "attention_dense/persistent_causal_sq512": { + "bytes": 313398, + "sha256": "9b4e796e22a7475f4339c1cc521337f68b15d837301adb43672ba756d5638873" + } + } + }, + "llvm22": { + "cases": { + "attention_dense/default_causal_sq512": { + "bytes": 320295, + "sha256": "eb781b077fb7944788a71d2636e2d2e139913b3f0f27ef48d2856467dfdb71ac" + }, + "attention_dense/persistent_causal_sq512": { + "bytes": 313502, + "sha256": "1f089fb980de6ef2a7d25bb61c1665adbebeafed36ae73d981092f3a37fcf978" + } + } + } + }, + "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..789de4dadcc2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_golden.py @@ -0,0 +1,117 @@ +# 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", + ) + return { + "attention_dense/default_causal_sq512": lambda: build_attention_dense( + AttentionDenseSpec(**base) + ), + "attention_dense/persistent_causal_sq512": lambda: build_attention_dense( + AttentionDenseSpec(**base, persistent=True, num_persistent=256) + ), + } + + +def _current_flavor(): + from rocke.core.lower_llvm import _resolve_llvm_flavor + + return _resolve_llvm_flavor() + + +def _sha_for(build, flavor): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + llvm = _lower_kernel_to_llvm_python(build(), arch="gfx950", llvm_flavor=flavor) + data = llvm.encode("utf-8") + return hashlib.sha256(data).hexdigest(), len(data) + + +def _build_doc(): + doc = {"schema": "attention_dense.ir_golden_sha256/v1", "flavors": {}} + for flavor in _FLAVORS: + cases = {} + for cid, build in _cases().items(): + try: + sha, nbytes = _sha_for(build, flavor) + cases[cid] = {"sha256": sha, "bytes": nbytes} + except Exception as e: # pragma: no cover - diagnostic + cases[cid] = {"error": str(e)[:160]} + doc["flavors"][flavor] = {"cases": cases} + return doc + + +def test_attention_dense_ir_matches_golden(): + import pytest + + if not _GOLDEN.exists(): + pytest.skip("golden fixture missing; generate with --write") + golden = json.loads(_GOLDEN.read_text()) + flavor = _current_flavor() + gflav = golden.get("flavors", {}).get(flavor) + if not gflav: + pytest.skip(f"no golden recorded for llvm flavor {flavor!r}") + drift = [] + for cid, build in _cases().items(): + want = gflav["cases"].get(cid, {}).get("sha256") + if want is None: + continue + got, _ = _sha_for(build, flavor) + if got != want: + drift.append(f"{cid}: {want} -> {got}") + assert not drift, "attention_dense IR drift vs golden:\n " + "\n ".join(drift) + + +if __name__ == "__main__": + if "--write" in sys.argv: + _GOLDEN.parent.mkdir(parents=True, exist_ok=True) + _GOLDEN.write_text(json.dumps(_build_doc(), indent=2, sort_keys=True) + "\n") + print(f"wrote {_GOLDEN}") + else: + test_attention_dense_ir_matches_golden() + print("PASS") From e2cb92d0981170ad054ebc159d0b276880e9bd43 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Thu, 16 Jul 2026 03:33:26 +0000 Subject: [PATCH 03/11] chore(rocke): re-run PR policy check (JIRA ID : AICK-1489 in description) Co-authored-by: Cursor From 7f8876c49fb2bdc56b3b4d3b4303ed16b845010f Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Thu, 16 Jul 2026 23:03:39 +0000 Subject: [PATCH 04/11] Author: Yaswanth Raparti Date: Thu Jul 16 23:03:39 2026 +0000 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(rocke): dense attn gfx950 — sliding-window, varlen, GQA L2 decode, V-pad Extend the productized dense flash-attention prefill kernel (kernels/gfx950/attention_dense.py) and add profiling/benchmark harnesses: - Sliding-window attention (sliding_window spec field, 0 = byte-identical to the always-on causal path): three-phase band KV loop (window-edge / mask-free interior / causal-edge) with per-lane bounds reusing the QK-output layout (no relayout); prunes the KV loop to ~W/block_n tiles. Both default and persistent builders. Bit-exact vs banded-mask SDPA. - Varlen batching (varlen + cu_seqlens_q/kv): packed [total_tok, H, D] with per-sequence runtime offsets/bounds and scf.if early-exit; composes with SWA. - GQA L2-locality persistent decode (persist_decode=auto, default): hkv-major causal-balanced work decode keeps each grid-stride phase on ~1 KV head so the shared GQA K/V stays L2-resident. Measured L2 hit 57%->93%, HBM misses 5.9x lower, ~+7% TFLOPS @ S=8192 GQA-8 (paired same-session A/B). Balance-safe gating (gqa*NQB*B >= 2*NP) falls back to qb_major elsewhere (no regression). - V-row LDS pad (_LDS_PAD_V=32) for the transposed PV read (ds_read_b64_tr_b16): swept {0:30, 8:29, 16:11, 32:0} conflicts -> +32 fully clears them (K's +8 is insufficient for the transpose read). Row-strided V DMA to carry the pad. ~+5% TFLOPS, 0 conflicts, no occupancy loss. Headline S=8192 GQA-8: ~953 TFLOPS. - Host (attention_dense_prefill.py): --sw flag, banded-mask parity, honest banded/windowed FLOP accounting. - Harnesses: bench_swa_varlen.py (dense vs flyDSL 0.2.4 vs unified_attention across sliding-window/varlen/MHA/GQA), _flydsl_bench.py, _prof_dense.py, _prof_compare.sh (rocprofv3 single-pass counter comparison), plan.md. --- .../gfx950/attention/prefill/_flydsl_bench.py | 140 +++++ .../gfx950/attention/prefill/_prof_compare.sh | 26 + .../gfx950/attention/prefill/_prof_dense.py | 90 +++ .../prefill/attention_dense_prefill.py | 40 +- .../attention/prefill/bench_swa_varlen.py | 524 ++++++++++++++++++ .../builders/gfx950/attention/prefill/plan.md | 248 +++++++++ .../library/kernels/gfx950/attention_dense.py | 471 ++++++++++++---- 7 files changed, 1432 insertions(+), 107 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py new file mode 100644 index 000000000000..cc842b2316a5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Standalone flyDSL flash-attn timer (own process + interpreter). + +Uses flyDSL **0.2.4** (pip-installed into /tmp/flydsl024-venv, torch inherited +from atom-venv via a .pth) + the flydsl-main HEAD kernels (also 0.2.4). This is +the DUALWAVE_SWP path: dense/varlen, GQA/MHA, causal. flyDSL prefill has NO +sliding window. Run this with /tmp/flydsl024-venv/bin/python and +PYTHONPATH=/workspace/flydsl-main. + +Modes: + dense + varlen (causal) +Prints: "OK " or "ERR ". +""" +import sys + +import torch + + +def _time(call, wu, it): + for _ in range(wu): + call() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(it): + call() + e.record() + e.synchronize() + return s.elapsed_time(e) / it + + +def main(): + from kernels.attention.flash_attn_interface import flydsl_flash_attn_func + + dev = "cuda" + DT = torch.bfloat16 + torch.manual_seed(1) + mode = sys.argv[1] + + if mode == "dense": + B, S, Hq, Hkv, D = (int(sys.argv[i]) for i in range(2, 7)) + causal = bool(int(sys.argv[7])) + wu, it = int(sys.argv[8]), int(sys.argv[9]) + q = (torch.randn(B, S, Hq, D, dtype=DT, device=dev) * 0.2).contiguous() + k = (torch.randn(B, S, Hkv, D, dtype=DT, device=dev) * 0.2).contiguous() + v = (torch.randn(B, S, Hkv, D, dtype=DT, device=dev) * 0.2).contiguous() + + def call(): + return flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv) + + o = call() + torch.cuda.synchronize() + rep = Hq // Hkv + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).repeat_interleave(rep, 1).float(), + v.transpose(1, 2).repeat_interleave(rep, 1).float(), + is_causal=causal, + ).transpose(1, 2) + err = (o.float() - ref).abs().max().item() + ms = _time(call, wu, it) + pairs = (S * (S + 1) // 2) if causal else S * S + tf = 4 * B * Hq * D * pairs / (ms * 1e-3) / 1e12 + print(f"OK {ms:.6f} {tf:.3f} {err:.3e}") + return + + if mode == "varlen": + seqlens = [int(x) for x in sys.argv[2].split(",")] + Hq, Hkv, D = int(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5]) + wu, it = int(sys.argv[6]), int(sys.argv[7]) + B = len(seqlens) + total = sum(seqlens) + mx = max(seqlens) + 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() + cu = torch.zeros(B + 1, dtype=torch.int32, device=dev) + cu[1:] = torch.tensor(seqlens, dtype=torch.int32, device=dev).cumsum(0) + + def call(): + return flydsl_flash_attn_func( + q, + k, + v, + causal=True, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=mx, + cross_seqlen=False, + ) + + o = call() + torch.cuda.synchronize() + rep = Hq // Hkv + err = 0.0 + for i, s in enumerate(seqlens): + st = sum(seqlens[:i]) + qh = q[st : st + s].transpose(0, 1).float().unsqueeze(0) + kh = ( + k[st : st + s] + .transpose(0, 1) + .repeat_interleave(rep, 0) + .float() + .unsqueeze(0) + ) + vh = ( + v[st : st + s] + .transpose(0, 1) + .repeat_interleave(rep, 0) + .float() + .unsqueeze(0) + ) + ref = ( + torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, is_causal=True + ) + .squeeze(0) + .transpose(0, 1) + ) + err = max(err, (o[st : st + s].float() - ref).abs().max().item()) + ms = _time(call, wu, it) + pairs = sum(x * (x + 1) // 2 for x in seqlens) + tf = 4 * Hq * D * pairs / (ms * 1e-3) / 1e12 + print(f"OK {ms:.6f} {tf:.3f} {err:.3e}") + return + + print(f"ERR unknown mode {mode!r}") + + +if __name__ == "__main__": + try: + main() + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f"ERR {exc!r}") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh new file mode 100644 index 000000000000..2557c0ac5497 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# rocprofv3 counter comparison: rocke dense vs flyDSL 0.2.4, GQA Hq128/Hkv8 S8192 causal. +# Each metric runs in its own single-pass rocprofv3 invocation (fast, no app replay). +set -u +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT=/tmp/prof_cmp +rm -rf "$OUT"; mkdir -p "$OUT" +ATOM=/workspace/atom-venv/bin/python +FLY=/tmp/flydsl024-venv/bin/python + +METRICS=(MfmaUtil MeanOccupancyPerCU OccupancyPercent FetchSize WriteSize LDSBankConflict MemUnitStalled SQ_WAVES) + +dense_cmd=($ATOM "$HERE/_prof_dense.py" 8192 128 8 128 1 8) +fly_cmd=($FLY "$HERE/_flydsl_bench.py" dense 1 8192 128 8 128 1 1 8) + +for M in "${METRICS[@]}"; do + rocprofv3 --pmc "$M" --kernel-include-regex "rocke_attention_dense" \ + -d "$OUT/dense_$M" --output-format csv -- "${dense_cmd[@]}" \ + >/dev/null 2>&1 + PYTHONPATH=/workspace/flydsl-main rocprofv3 --pmc "$M" \ + --kernel-include-regex "flash_attn_dualwave_swp" \ + -d "$OUT/flydsl_$M" --output-format csv -- "${fly_cmd[@]}" \ + >/dev/null 2>&1 + echo "done $M" +done +echo "ALL_DONE" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py new file mode 100644 index 000000000000..f853c0e3432d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Minimal single-kernel launcher for rocprofv3 profiling of the rocke dense +prefill kernel. Warmup builds/JITs; then launches the kernel `iters` times so +the profiler captures steady-state dispatches. + +Args: S Hq Hkv D causal(0/1) iters +""" +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 ( + AttentionDenseSpec, + build_attention_dense, +) # noqa: E402 +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 + + +def main(): + S = int(sys.argv[1]) if len(sys.argv) > 1 else 8192 + Hq = int(sys.argv[2]) if len(sys.argv) > 2 else 128 + Hkv = int(sys.argv[3]) if len(sys.argv) > 3 else 8 + D = int(sys.argv[4]) if len(sys.argv) > 4 else 128 + causal = bool(int(sys.argv[5])) if len(sys.argv) > 5 else True + iters = int(sys.argv[6]) if len(sys.argv) > 6 else 20 + dev = "cuda" + DT = torch.bfloat16 + torch.manual_seed(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=causal, + dtype="bf16", + persistent=True, + num_persistent=256, + ) + art = compile_kernel( + build_attention_dense(spec), + arch="gfx950", + backend="python", + capture_ir_text=False, + ) + sig = ( + SignatureBuilder() + .ptr("q_ptr", "bf16") + .ptr("k_ptr", "bf16") + .ptr("v_ptr", "bf16") + .ptr("o_ptr", "bf16") + .scalar("scale", "f32") + .build() + ) + lch = KernelLauncher(hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sig) + stream = torch.cuda.current_stream().cuda_stream + cfg = LaunchConfig( + grid=(256, 1, 1), block=(spec.num_waves * 64, 1, 1), stream=stream + ) + vals = { + "q_ptr": q, + "k_ptr": k, + "v_ptr": v, + "o_ptr": out, + "scale": 1.0 / math.sqrt(D), + } + # warmup / JIT + lch(vals, config=cfg) + torch.cuda.synchronize() + sys.stderr.write(f"[prof_dense] kernel={art.kernel_name}\n") + for _ in range(iters): + lch(vals, config=cfg) + torch.cuda.synchronize() + + +if __name__ == "__main__": + main() 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 index 54b9a24d9785..4a414424d0fc 100644 --- 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 @@ -47,15 +47,17 @@ def _make_launcher(spec: AttentionDenseSpec): backend="python", capture_ir_text=False, ) - sig = ( + 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") - .build() ) + 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) @@ -99,9 +101,19 @@ def call(): rep = Hq // Hkv kh = k.transpose(1, 2).repeat_interleave(rep, 1).float() vh = v.transpose(1, 2).repeat_interleave(rep, 1).float() - ref = torch.nn.functional.scaled_dot_product_attention( - qh, kh, vh, is_causal=spec.causal - ).transpose(1, 2) + 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): @@ -114,11 +126,15 @@ def call(): e.record() e.synchronize() ms = s.elapsed_time(e) / iters - flops = ( - 4 * B * Hq * D * (Sq * (Sq + 1) // 2) - if spec.causal - else 2 * 2 * B * Hq * D * Sq * Skv - ) + 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( @@ -144,6 +160,9 @@ def main(): ) 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( @@ -160,6 +179,7 @@ def main(): persistent=args.persistent, num_persistent=args.np, interleave=args.interleave, + sliding_window=args.sw, ) run(spec) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py new file mode 100644 index 000000000000..c99fc9098225 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +"""Benchmark dense (rocke) vs flyDSL vs unified_attention across sliding-window +and varlen-batch prefill shapes on gfx950 / MI355X. + +Backends +-------- +* ``dense`` : the productized rocke dense prefill kernel with the new + ``sliding_window`` (A) and ``varlen`` (B) features + (``kernels/gfx950/attention_dense.py``). +* ``flydsl`` : ``flydsl_flash_attn_func`` (DUALWAVE_SWP). Causal dense / packed + varlen only -- flyDSL prefill has NO sliding-window, so on SWA + rows it runs FULL causal (a "no-window" reference, marked *). +* ``unified``: rocke ``run_unified_attention_torch`` (paged 2D/3D). Supports + both sliding-window and varlen (paged). + +All three are timed with the same torch CUDA-event timer on the current stream. +Reported TFLOPS uses the honest banded pair count (windowed work); flyDSL's SWA +rows compute full-causal work so its TFLOPS bar is on a *different* workload +(hence the *). Correctness for dense is checked vs a banded-mask SDPA reference. + +Usage: + python bench_swa_varlen.py --mode swa + python bench_swa_varlen.py --mode varlen + python bench_swa_varlen.py --mode both --iters 50 +""" +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, +) +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 + +_DT = torch.bfloat16 + + +def _pairs(s, W): + """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, Hq, D): + return sum(4 * Hq * D * _pairs(s, W) for s in seqlens) + + +def _time(fn, warmup, iters): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(iters): + fn() + e.record() + e.synchronize() + return s.elapsed_time(e) / iters + + +# --------------------------------------------------------------------------- # +# rocke dense kernel +# --------------------------------------------------------------------------- # +_LAUNCHER_CACHE = {} + + +def _dense_launcher(spec): + key = spec.kernel_name() + if key in _LAUNCHER_CACHE: + return _LAUNCHER_CACHE[key] + 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 + + +def bench_dense(seqlens, W, Hq, Hkv, D, warmup, iters, persistent): + dev = "cuda" + B = len(seqlens) + max_s = max(seqlens) + total = sum(seqlens) + scale = 1.0 / math.sqrt(D) + stream = torch.cuda.current_stream().cuda_stream + torch.manual_seed(0) + varlen = B > 1 or (total != B * max_s) + + if varlen: + 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="bf16", + sliding_window=W, + varlen=True, + ) + lch = _dense_launcher(spec) + cfg = LaunchConfig( + grid=(max_s // 256, Hq, B), block=(spec.num_waves * 64, 1, 1), 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: + 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="bf16", + sliding_window=W, + persistent=persistent, + num_persistent=256, + ) + lch = _dense_launcher(spec) + if persistent: + grid = (256, 1, 1) + else: + grid = (S // 256, Hq, 1) + cfg = LaunchConfig(grid=grid, block=(spec.num_waves * 64, 1, 1), 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 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(call, warmup, iters) + tf = _flops(seqlens, W, Hq, D) / (ms * 1e-3) / 1e12 + return ms, tf, max_err + + +# --------------------------------------------------------------------------- # +# flyDSL +# --------------------------------------------------------------------------- # +# flyDSL 0.2.4 (DUALWAVE_SWP: dense/varlen, GQA/MHA, causal; NO sliding window). +# pip-installed into /tmp/flydsl024-venv (torch inherited from atom-venv via .pth); +# kernels come from the flydsl-main HEAD checkout. Runs in a subprocess with its +# own interpreter + PYTHONPATH so its top-level `kernels` package doesn't collide +# with rocke's. +_FLYDSL_PY = "/tmp/flydsl024-venv/bin/python" +_FLYDSL_ENV = {**os.environ, "PYTHONPATH": "/workspace/flydsl-main"} + + +def _flydsl_run(args): + import subprocess + + cmd = [_FLYDSL_PY, os.path.join(_HERE, "_flydsl_bench.py"), *args] + try: + out = subprocess.run( + cmd, capture_output=True, text=True, timeout=600, env=_FLYDSL_ENV + ) + except Exception as e: # noqa: BLE001 + return None, None, f"subproc: {e!r}" + line = "" + for ln in out.stdout.splitlines(): + if ln.startswith(("OK ", "ERR ")): + line = ln + if line.startswith("OK "): + p = line.split() + return float(p[1]), float(p[2]), (float(p[3]) if len(p) > 3 else "") + return ( + None, + None, + ( + line[4:] + if line.startswith("ERR ") + else (out.stderr.strip().splitlines() or ["no output"])[-1] + ), + ) + + +def bench_flydsl_dense(B, S, Hq, Hkv, D, warmup, iters, causal=True): + return _flydsl_run( + [ + "dense", + str(B), + str(S), + str(Hq), + str(Hkv), + str(D), + str(int(causal)), + str(warmup), + str(iters), + ] + ) + + +def bench_flydsl_varlen(seqlens, Hq, Hkv, D, warmup, iters): + sl = ",".join(str(x) for x in seqlens) + return _flydsl_run( + ["varlen", sl, str(Hq), str(Hkv), str(D), str(warmup), str(iters)] + ) + + +# --------------------------------------------------------------------------- # +# unified_attention (paged) +# --------------------------------------------------------------------------- # +def bench_unified(seqlens, W, Hq, Hkv, D, warmup, iters): + try: + from rocke.assets import shape_utils_dir + + sys.path.insert(0, str(shape_utils_dir())) + import _ua_shape_utils as u + from kernels import UnifiedAttentionProblem, run_unified_attention_torch + except Exception as e: # noqa: BLE001 + return None, None, f"import: {e!r}" + dev = "cuda" + B = len(seqlens) + total = sum(seqlens) + max_s = max(seqlens) + block_size = 16 + num_blocks = sum((s + block_size - 1) // block_size for s in seqlens) + B + 16 + max_blocks = (max_s + block_size - 1) // block_size + win = (W - 1, 0) if (W and W > 0) else (-1, -1) + try: + shape = u.UAShape( + source_file="synthetic", + line_idx=0, + call_idx=0, + kind="prefill", + all_decode=False, + num_seqs=B, + total_q=total, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=D, + block_size=block_size, + num_blocks=num_blocks, + max_blocks_per_seq=max_blocks, + max_seqlen_q=max_s, + max_seqlen_k=max_s, + softmax_scale=1.0 / math.sqrt(D), + softcap=0.0, + window_size=win, + has_sinks=False, + has_alibi=False, + has_output_scale=False, + q_dtype="torch.bfloat16", + k_dtype="torch.bfloat16", + v_dtype="torch.bfloat16", + out_dtype="torch.bfloat16", + ) + data = u.make_inputs(shape, seed=2, cap_blocks=None) + sw = W if (W and W > 0) else 0 + prob = UnifiedAttentionProblem( + total_q=shape.total_q, + num_seqs=shape.num_seqs, + num_query_heads=Hq, + num_kv_heads=Hkv, + head_size=D, + block_size=block_size, + max_seqlen_q=max_s, + max_seqlen_k=max_s, + dtype="bf16", + sliding_window=sw, + softcap=0.0, + use_sinks=False, + use_alibi=False, + use_qq_bias=False, + use_fp8=False, + num_sms=256, + ) + hip_stream = torch.cuda.current_stream().cuda_stream + + def call(): + run_unified_attention_torch( + problem=prob, + q=data["query"], + k=data["key_cache"], + v=data["value_cache"], + out=data["output"], + cu_seqlens_q=data["cu_seqlens_q"], + seqused_k=data["kv_lens"], + softmax_scale=data["scale"], + block_table=data["block_tables"], + softcap=0.0, + sinks=data.get("sinks"), + alibi_slopes=data.get("alibi_slopes"), + qq_bias=None, + qq_bias_stride_0=0, + backend="auto", + stream=hip_stream, + ) + + call() + torch.cuda.synchronize() + except Exception as e: # noqa: BLE001 + return None, None, f"run: {e!r}" + ms = _time(call, warmup, iters) + tf = _flops(seqlens, W, Hq, D) / (ms * 1e-3) / 1e12 + return ms, tf, "" + + +def _row(label, res): + ms, tf, note = res + if ms is None: + return f" {label:<10} FAILED ({note})" + extra = f" max_abs={note:.2e}" if isinstance(note, float) else "" + return f" {label:<10} {ms:8.4f} ms {tf:8.1f} TFLOPS{extra}" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--mode", choices=["swa", "varlen", "mha", "both", "all"], default="all" + ) + 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("--warmup", type=int, default=10) + ap.add_argument("--iters", type=int, default=50) + ap.add_argument("--persistent", action="store_true", default=True) + args = ap.parse_args() + Hq, Hkv, D = args.hq, args.hkv, args.d + + if args.mode in ("mha", "all"): + print("\n=== DENSE-CAUSAL GQA (Hq=128, Hkv=8) -- headline prefill config ===") + for S in (2048, 4096, 8192): + print(f"\n-- B=1 S={S} Hq={Hq} Hkv={Hkv} D={D} causal --") + print( + _row( + "dense", + bench_dense( + [S], 0, Hq, Hkv, D, args.warmup, args.iters, args.persistent + ), + ) + ) + print( + _row( + "flydsl", + bench_flydsl_dense(1, S, Hq, Hkv, D, args.warmup, args.iters, True), + ) + ) + print( + _row( + "unified", + bench_unified([S], 0, Hq, Hkv, D, args.warmup, args.iters), + ) + ) + + print("\n=== DENSE-CAUSAL MHA ===") + for H in (16, 32): + for S in (4096, 8192): + print(f"\n-- B=1 S={S} H={H} (MHA) D={D} causal --") + print( + _row( + "dense", + bench_dense( + [S], 0, H, H, D, args.warmup, args.iters, args.persistent + ), + ) + ) + print( + _row( + "flydsl", + bench_flydsl_dense( + 1, S, H, H, D, args.warmup, args.iters, True + ), + ) + ) + print( + _row( + "unified", + bench_unified([S], 0, H, H, D, args.warmup, args.iters), + ) + ) + + if args.mode in ("swa", "both", "all"): + print("\n\n=== SLIDING-WINDOW sweep (single-seq prefill, GQA, causal) ===") + print(" (flyDSL v0.2.0 prefill has NO sliding-window -> not shown)") + 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}" + print(f"\n-- S={S} {tag} (Hq={Hq} Hkv={Hkv}) --") + print( + _row( + "dense", + bench_dense( + [S], W, Hq, Hkv, D, args.warmup, args.iters, args.persistent + ), + ) + ) + print( + _row( + "unified", + bench_unified([S], W, Hq, Hkv, D, args.warmup, args.iters), + ) + ) + + if args.mode in ("varlen", "both", "all"): + print("\n\n=== VARLEN sweep (packed ragged batch, GQA, causal) ===") + print(" (flyDSL: W=0 only -- no sliding window in flyDSL prefill)") + 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}" + print( + f"\n-- B={len(seqlens)} total={sum(seqlens)} " + f"seqlens={seqlens} {tag} --" + ) + print( + _row( + "dense", + bench_dense( + seqlens, W, Hq, Hkv, D, args.warmup, args.iters, False + ), + ) + ) + if W == 0: + print( + _row( + "flydsl", + bench_flydsl_varlen( + seqlens, Hq, Hkv, D, args.warmup, args.iters + ), + ) + ) + print( + _row( + "unified", + bench_unified(seqlens, W, Hq, Hkv, D, args.warmup, args.iters), + ) + ) + + +if __name__ == "__main__": + main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md new file mode 100644 index 000000000000..398b2d8d34b9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md @@ -0,0 +1,248 @@ +# Plan: porting sliding-window, batching, and bias to the dense prefill kernel + +Scope: extend the productized gfx950 dense flash-attention prefill kernel +(`kernels/gfx950/attention_dense.py`, host builder `attention_dense_prefill.py`) +with **sliding-window attention (SWA)**, richer **batching** (varlen), and +**bias addition** (ALiBi / additive bias / logit softcap), informed by a +structured read of the flyDSL implementations in `~/flydsl-main`. + +This document records (1) what flyDSL actually does for each feature — tile +layouts, MFMA atom sizes, layout transforms, wave/register/LDS allocation, and +tile/wave/K-loop scheduling — and (2) the concrete change plan for the rocke +dense kernel, with the always-on 877-TFLOPS pipeline preserved as the no-feature +fast path. + +--- + +## 0. Baseline: the rocke dense kernel we are extending + +- `BLOCK_M = 256` query rows/CTA, `block_n = 64` KV tile, `NBUF = 2` double + buffer, `num_waves = 8` (each wave owns 32 query rows), `waves_per_eu = 2`, + 256 VGPR / 0 spill. +- MFMA `32x32x16` for both QK and PV (`mfma_32x32x16_for_dtype`). +- CK-1 transposed PV (P fed to PV MMA in native QK-output layout via half-local + V load `pv32_v_load_paired`; no cross-half P relayout). +- LDS K padded `[NBUF, BN, D+8]` (bank-conflict fix); V unpadded `[NBUF, BN, D]`. +- Diagonal-only causal masking: mask-free body loop + masked diagonal tail. +- Depth-1 cluster fusing `exp2` into the PV MFMA loop; `sched_group_barrier` + template (DS_READ/MFMA/VALU/TRANS). +- Default grid `(Sq//256, Hq, B)`; persistent grid-stride variant loops one 1-D + grid of `num_persistent` CTAs over `W = (Sq//256)*Hq*B`. +- Inputs: dense contiguous `[B, S, H, D]`, uniform seqlen; params + `q,k,v,o,scale`. No varlen/paged, no bias, no sliding window. + +**Key finding:** flyDSL's production gfx950 prefill kernel +(`flash_attn_gfx950.py`, "DUALWAVE_SWP") is the *same design family* as this +kernel — same `BLOCK_M=256`/`BLOCK_N=64`, same `32x32x16` atoms, 8 waves × 32 +rows, P-in-registers, V transpose-read, double-buffered `[K0,V0,K1,V1]`, +`sched_group_barrier` + `s_setprio` + staggered wave barriers, and a lazy online +rescale. So flyDSL is a high-fidelity reference for the pipeline mechanics; the +deltas we need are exactly the three requested features. + +--- + +## 1. flyDSL findings by feature + +### 1a. Core pipeline + batching (`flash_attn_gfx950.py`, `flash_attn_generic.py`) + +- **Tile layout / atoms:** `BLOCK_M=256`, `BLOCK_N=64`, `K_SUB_N=32`, 8 waves × + 32 rows; MFMA `f32_32x32x16_{bf16,f16}` (K=16, v8 packs, v16f32 acc) for QK + and PV; fp8 path uses `mfma_f32_32x32x16_fp8_fp8` with optional wide + `32x32x64` PV. Q register-resident, pre-scaled by `1/sqrt(D)*log2(e)`. +- **Transforms:** K uses either XOR swizzle (`col ^ ((row&mask)<<4)`) or a padded + `u_rk` layout (`+8 bf16`); V is transpose-read with `ds_read_b64_tr_b16` + (padded `+32 bf16`) — matches our `pv32_v_load_paired`. P stays in VGPRs (no + LDS roundtrip); half-wave partners combined via `permlane32_swap`. O store + fuses partners with `permlane32_swap + cvt_pk_bf16_f32` → `buffer_store_dwordx4`. +- **Wave/register/LDS:** WG=512 threads (8 waves), `waves_per_eu=2`, interleaved + double buffer `[K0,V0,K1,V1]` (~68 KiB @ D=128). DUALWAVE staggers two + 4-wave groups with extra barriers to time-multiplex the SW pipeline. +- **K-loop scheduling:** hand-built 8-cluster software pipeline, `j += 2` (2 KV + tiles/iter), PV lags QK by ~1 tile; async `buffer_load_lds` (128B) with split + `s_waitcnt` (lgkmcnt vs vmcnt); `sched_group_barrier` (MFMA/VALU/EXP masks), + `s_setprio`, `s_nop`. **Lazy online rescale** skips O/P rescale when + `max(m_tile - m_row) <= 8.0` (ballot-vs-exec test) — a lever we have not baked. +- **Batching:** 3D grid `(NUM_HEADS_Q, num_q_blocks, grid_z)` with + `grid_z = batch` (or `batch*num_kv_splits` for split-K); GQA via + `kv_head = q_head // group_size`; **varlen** via `cu_seqlens_q/kv` rebased + buffer descriptors (`q_tok_base = cu[b]`, `seqlen_q_b = cu[b+1]-cu[b]`), grid + sized by `max_seqlen_q` with early-exit of empty q-tiles; **paged** via a + block table staged in LDS; **split-K** with an fp32 workspace + combine kernel. + **No persistent/grid-stride** (our persistent mode is a rocke-specific lever). + +### 1b. Sliding window (`pa_decode_swa.py` — decode/paged/fp8 only) + +**flyDSL has SWA only on the paged-decode path; flash-attn prefill is causal-only.** +The technique is directly portable: + +- **Masking math (compile-time `W`):** per query, `causal_bound = C + 1 - L + qi` + (exclusive upper), `seq_start = C - L - W + qi` (inclusive lower); allowed iff + `seq_start <= k < causal_bound`, i.e. `k ∈ [q_abs - W, q_abs]` (**W+1** keys + incl. diagonal; note the `+1` convention). `q_abs = C - L + qi`. +- **Two-level pruning:** (a) tile-level — `seq_start_global = max(0, C-L-W)`, + `tail_start_tile = seq_start_global >> 8` (tiles are 256 tokens), only visit + `[tail_start_tile, num_tiles)`; host bounds partitions with + `ceil((W+L-1)/256)+1`. (b) per-lane range mask on the MFMA `f32x4` output: + `in_range = (tok < causal_bound) & (tok >= seq_start)`, `select(in_range, s, -inf)`. +- **Mask placement / layout:** applied **after QK MFMA, before row-max/exp2**, in + the *native MFMA accumulator layout* — **no relayout/transpose**. Token indices + built inline (`kv_tok_base + td*16 + i`); bounds are SGPR scalars broadcast + across the 4 lanes. +- **No mask-free interior fast path in SWA** (every visited tile is masked). But + flash-attn *causal* has `tile_needs_mask` (skip mask when + `max_kv_col <= q_start`) — the optimization to replicate for the SWA interior. +- **Wave cooperation:** all waves iterate the same tile range; per-wave rows + differ only in the per-lane bound; cooperative K/V loads unchanged. +- **Cost:** +1 compare vs causal-only (two bounds instead of one) per `f32x4`. + +### 1c. Bias / ALiBi / softcap — **not implemented anywhere in flyDSL** + +Full-repo grep confirms: no ALiBi, no additive bias tensor, no QQ-bias, no +softcap in any kernel; the only `alibi_slopes`/`bias` refs are test harnesses +passing `None` to the external aiter reference. flyDSL applies only causal / KV-pad +/ SWA masks + QK softmax scaling (+ fp8 quant scales). So for bias we have **no +flyDSL reference to port** — the rocke `unified_attention` tiled path is the +in-repo reference (it *does* implement ALiBi + QQ-bias + softcap). flyDSL still +tells us the *insertion mechanics*: + +- **Insertion point:** after QK MFMA, **before mask/row-max/exp2**, on the raw + accumulator layout (same place SWA masks) — no relayout. +- **Domain:** gfx950 pre-scales Q by `sm_scale*log2(e)`, so scores are already in + log2 domain; any bias must be added **in that domain** (multiply linear biases + by `log2(e)` = `RCP_LN2`), matching how the rocke tiled path applies ALiBi + (`slope*(key_pos-ctx)*RCP_LN2`) and QQ-bias (`qq_bias[...]*RCP_LN2`). +- **ALiBi:** `S += slope_h * (k_pos - q_pos)` using the same per-lane + `k_pos`/`q_pos` maps SWA/causal already compute (`_mfma_32x32_c_row/_col`). +- **Softcap:** nonlinear `S = cap * tanh(S/cap)` applied **after scale, before + mask** (must precede max). Careful: our exp2 fast path assumes softmax arg + `<= 0` and no overflow guard — softcap changes the score range, revisit. +- **Composition:** per-row constant biases cancel in softmax; **relative** biases + (ALiBi, elementwise, softcap) do not and must be applied before max so online + rescale stays consistent. + +--- + +## 2. Porting plan + +Ordering by value/effort: **(A) sliding window** (cheap, high value, byte-safe +default) → **(B) varlen batching** → **(C) bias/ALiBi/softcap**. Each is gated so +the no-feature path emits today's exact 877-TFLOPS kernel. + +### A. Sliding window + +Design already scoped in prior analysis; flyDSL confirms the mechanics. + +- **Spec:** add `sliding_window: int = 0` to `AttentionDenseSpec` (0 = disabled → + bit-identical to today). Bake into `kernel_name()`; require `W % block_n == 0` + for the clean path (`__post_init__`). Mirror in the persistent builder. +- **Geometry:** both causal upper (`k<=q`) and window lower (`k>=q-W+1`) are + slope-1 diagonals; the valid region per 256-row block is a **parallelogram + band**. With `n_per = BLOCK_M//BN = 4` and `Wt = W//BN`: + `lo_tile = max(0, qb*n_per - Wt)`, `hi_tile = min(n_ktiles, (qb+1)*n_per)`. + Each boundary band is exactly `n_per = 4` tiles wide. +- **Three-phase K-loop:** (L) window-edge tiles `[lo_tile, lo_tile+n_per)` lower- + masked; (M) interior `[lo_tile+n_per, diag_start)` **mask-free** (identical hot + loop → same peak); (R) causal diagonal `[diag_start, hi_tile)` upper-masked + (existing tail). If `W < BLOCK_M` bands overlap → single merged both-bounds + phase. Reuse the existing `select`-guarded bound idiom for empty phases. +- **Masking:** extend `do_mask` with `mask_lower`/`mask_upper` flags; reuse + `_mfma_32x32_c_row/_col` exactly (no new layout). Lower predicate + `cmp_gt(ktok, query_tok - W)` (W compile-time → immediate). +- **Register pinning:** hoist `query_tok` (lane-only, tile-invariant) and + `query_tok-(W-1)` to once-per-block VGPRs; masking stays outside the loop- + carried set (`m,l,o,pk`) → no VGPR growth, preserve 0-spill. +- **Pipeline:** prologue primes `lo_tile`/`lo_tile+1` (not 0/1) and left-masks the + prologue tile; init `m=-inf,l=0,o=0` unchanged (tiles `0 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 + # varlen: packed variable-length batch. Q/K/V/O are packed [total_tok, H, D] + # and per-sequence boundaries come from cu_seqlens_q/cu_seqlens_kv (int32 + # [batch+1]) at runtime. seqlen_q/seqlen_kv are the MAX lengths (grid + # sizing); each sequence's length must be a multiple of block_m (q) and + # block_n (kv). Self-attention only (per-seq seqlen_q == seqlen_kv). Not + # supported with persistent. 0-cost when False (dense uniform path). + varlen: bool = False # --- validated performance knobs --- # block_n: KV tile length. 64 (66 KB LDS, WPE-tunable) and 128 (135 KB LDS, pins @@ -101,6 +125,21 @@ class AttentionDenseSpec: # (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" def __post_init__(self) -> None: if self.dtype not in _DTYPE_IR: @@ -134,6 +173,26 @@ def __post_init__(self) -> None: 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: @@ -147,6 +206,21 @@ def dtype_ir(self): 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 + 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", @@ -159,8 +233,14 @@ def kernel_name(self) -> str: f"sk{self.seqlen_kv}", "causal" if self.causal else "full", ] + if self.sliding_window > 0: + parts.append(f"swa{self.sliding_window}") + if self.varlen: + parts.append("varlen") 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) @@ -203,6 +283,9 @@ def build_attention_dense( 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 K_STEPS = D // 16 D_TILES = D // 32 @@ -229,6 +312,17 @@ def build_attention_dense( "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) @@ -246,19 +340,40 @@ def build_attention_dense( 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))) - 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), unpadded V --- + 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) --- LDROW = D + PAD + VROW = D + _LDS_PAD_V K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") - V_lds = b.smem_alloc(dtype, [NBUF, BN, D], name_hint="Vlds") + 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). q_tok = b.add(q_tok0, lane_m) @@ -278,7 +393,8 @@ def build_attention_dense( K_BYTES_PER_BUF = BN * LDROW * 2 K_LDROW_BYTES = LDROW * 2 - V_BYTES_PER_BUF = BN * D * 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) @@ -306,23 +422,20 @@ def async_load_k(lds_base, buf_val, tile_key0): ) def async_load_v(lds_base, buf_val, tile_key0): - """Contiguous async DMA into unpadded [BN, D] V_lds layout.""" + """Row-by-row async DMA into [BN, VROW] V_lds (row-strided so V can carry + the +PAD_V transposed-PV bank pad); 64 lanes x 2 bf16 = D per row.""" buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) - base_off = b.add(v_wave_off_i64, buf_off) - for c in range(V_DMA_PASSES): - wave_base = b.smem_ptr_add( - lds_base, b.add(base_off, b.const_i64(c * WAVES * WAVE_BYTES)) - ) - flat = b.mul( - b.add(b.mul(b.const_i32(c), b.const_i32(WAVES * 64)), tid), - b.const_i32(8), + 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(V_LDROW_BYTES)), I64) ) - krow = b.div(flat, b.const_i32(D)) - kcol = b.mod(flat, b.const_i32(D)) - gkey = b.add(tile_key0, krow) - voff = b.add(b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), kcol) + 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( - v_rsrc, wave_base, b.mul(voff, b.const_i32(2)), zero_soff, 4 + v_rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 ) def load_tile(buf_val, tile_idx): @@ -348,18 +461,29 @@ def do_qk(kbuf): s_reg.append([b.vec_extract(acc, i) for i in range(16)]) return s_reg - def do_mask(s_reg, tile_idx): + 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)) - s_reg[nsub][i] = b.select( - b.cmp_le(ktok, query_tok), s_reg[nsub][i], neg_inf - ) + 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 softmax_max(s_reg, m_i): local_max = neg_inf @@ -460,21 +584,40 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): 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, b.const_i32(n_ktiles)), n_upper, b.const_i32(n_ktiles) + 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: - n_upper = b.const_i32(n_ktiles) + 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 tile 0. - load_tile(b.const_i32(0), b.const_i32(0)) - load_tile(b.const_i32(1), b.const_i32(1)) + # 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() - s0 = do_qk(b.const_i32(0)) - do_mask(s0, b.const_i32(0)) + s0 = do_qk(start_buf) + if causal and W > 0: + do_mask(s0, start_tile, lower=True, upper=True) + else: + do_mask(s0, start_tile) m0, _alpha0 = softmax_max(s0, neg_inf) # tile-0 softmax exp + relayout only; PV lags by one tile (fused into the loop). p0_vals = [ @@ -494,7 +637,7 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): + [(f"pk{kk}", pk0[kk]) for kk in range(KK_STEPS)] ) - def emit_loop_body(j, carry, masked): + def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): m_i = carry[0] l_i = carry[1] o_acc = list(carry[2 : 2 + D_TILES]) @@ -508,8 +651,8 @@ def emit_loop_body(j, carry, masked): b.s_waitcnt(vmcnt=V_DMA_PASSES) b.s_barrier_bare() s = do_qk(kbuf) - if masked: - do_mask(s, j) + if mask_lower or mask_upper: + do_mask(s, j, lower=mask_lower, upper=mask_upper) m_new, alpha = 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%) @@ -522,7 +665,41 @@ def emit_loop_body(j, carry, masked): load_tile(pbuf, b.add(j, b.const_i32(1))) b.scf_yield(m_new, l_new, *o_acc, *p_packs) - if causal: + 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 KernelDef: n_per = BLOCK_M // BN NQB = 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 b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = WAVES * 64 @@ -654,12 +838,14 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: neg_inf = b.const_f32(-1e30) LDROW = D + PAD + VROW = D + _LDS_PAD_V # V-row pitch (padded for the transposed PV read) K_lds = b.smem_alloc(dtype, [NBUF, BN, LDROW], name_hint="Klds") - V_lds = b.smem_alloc(dtype, [NBUF, BN, D], name_hint="Vlds") + 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 * D * 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) @@ -680,20 +866,45 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: b.s_waitcnt(vmcnt=0) b.s_barrier_bare() - # qb-MAJOR decode: wi = qb*(Hq*B) + hq*B + bt. Putting qb (the triangular - # causal cost index) in the MSB spreads cheap+expensive query blocks across - # each CTA under grid-stride; a qb-fast decode would alias qb to a constant - # per CTA when NP is a multiple of NQB (32x imbalance). - bt = b.mod(wi, b.const_i32(B)) - rem = b.div(wi, b.const_i32(B)) - hq = b.mod(rem, b.const_i32(Hq)) - qb0 = b.div(rem, b.const_i32(Hq)) - if INTERLEAVE and causal and NQB > 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) + 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( @@ -743,25 +954,30 @@ def async_load_k(lds_base, buf_val, tile_key0): voff = b.add(voff, b.const_i32(stride_k_tok)) def async_load_v(lds_base, buf_val, tile_key0): + """Row-by-row async DMA into [BN, VROW] V_lds. Row-strided (like K) + so V can carry the +PAD_V bank-conflict pad for the transposed PV + read; each row's 64 lanes load 2 bf16 each (= D) contiguously.""" buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) - base_off = b.add(v_wave_off_i64, buf_off) - for c in range(V_DMA_PASSES): - wave_base = b.smem_ptr_add( - lds_base, b.add(base_off, b.const_i64(c * WAVES * WAVE_BYTES)) - ) - flat = b.mul( - b.add(b.mul(b.const_i32(c), b.const_i32(WAVES * 64)), tid), - b.const_i32(8), - ) - krow = b.div(flat, b.const_i32(D)) - kcol = b.mod(flat, b.const_i32(D)) - gkey = b.add(tile_key0, krow) - voff = b.add( - b.add(k_base, b.mul(gkey, b.const_i32(stride_k_tok))), kcol - ) + 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(V_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( - v_rsrc, wave_base, b.mul(voff, b.const_i32(2)), zero_soff, 4 + v_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(V_LDROW_BYTES)) + voff = b.add(voff, b.const_i32(stride_k_tok)) def load_tile(buf_val, tile_idx): tk0 = b.mul(tile_idx, b.const_i32(BN)) @@ -780,18 +996,24 @@ def do_qk(kbuf): s_reg.append([b.vec_extract(acc, i) for i in range(16)]) return s_reg - def do_mask(s_reg, tile_idx): + 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)) - s_reg[nsub][i] = b.select( - b.cmp_le(ktok, query_tok), s_reg[nsub][i], neg_inf - ) + 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 softmax_max(s_reg, m_i): local_max = neg_inf @@ -900,7 +1122,7 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): l_tile = b.fadd(l_local, b.warp_shuffle_xor(l_local, 32)) return out, p_vals, l_tile - def emit_loop_body(j, carry, masked): + def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): m_i = carry[0] l_i = carry[1] o_acc = list(carry[2 : 2 + D_TILES]) @@ -915,8 +1137,8 @@ def emit_loop_body(j, carry, masked): b.s_waitcnt(vmcnt=V_DMA_PASSES) b.s_barrier_bare() s = do_qk(kbuf) - if masked: - do_mask(s, j) + if mask_lower or mask_upper: + do_mask(s, j, lower=mask_lower, upper=mask_upper) m_new, alpha = softmax_max(s, m_i) b.sched_barrier(0) # PV-only s_setprio: the PV MFMA cluster wins issue slots; paired with @@ -941,12 +1163,27 @@ def emit_loop_body(j, carry, masked): else: n_upper = b.const_i32(n_ktiles) - load_tile(b.const_i32(0), b.const_i32(0)) - load_tile(b.const_i32(1), b.const_i32(1)) + # 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(b.const_i32(0)) - do_mask(s0, b.const_i32(0)) + 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) m0, _alpha0, p0, l0 = softmax_stats(s0, neg_inf) o0 = [b.zero_vec_f32(16) for _ in range(D_TILES)] pk0 = relayout_p(p0) @@ -957,14 +1194,47 @@ def emit_loop_body(j, carry, masked): + [(f"pk{kk}", pk0[kk]) for kk in range(KK_STEPS)] ) - if causal: + 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, masked=False) + emit_loop_body(j, carry) tail_args = [ (name + "_t", val) for (name, _), val in zip(iter_args, body.results) ] @@ -975,13 +1245,13 @@ def emit_loop_body(j, carry, masked): tail_lo, n_upper, b.const_i32(1), tail_args, iv_name="nt" ) with loop as (j, carry): - emit_loop_body(j, carry, masked=True) + 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, masked=False) + emit_loop_body(j, carry) res = loop.results l_i = res[1] @@ -995,10 +1265,17 @@ def emit_loop_body(j, carry, masked): 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). + # 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)) - hq_e = b.mod(b.div(wi, b.const_i32(B)), b.const_i32(Hq)) + 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)), From 2ad7acd86fad4d41a3d0307a99ccff5f3a631143 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Fri, 17 Jul 2026 00:37:47 +0000 Subject: [PATCH 05/11] (rocke): lds conflicts, lazy o improvements --- .../library/kernels/gfx950/attention_dense.py | 102 +++++++++++++++--- 1 file changed, 90 insertions(+), 12 deletions(-) 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 index 4836f0517430..c519274be31e 100644 --- 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 @@ -66,6 +66,9 @@ import os as _os # noqa: E402 _LDS_PAD_V = int(_os.environ.get("ROCKE_DENSE_VPAD", "32")) +# Lazy-rescale re-anchor threshold in the log2 domain: skip the O/l rescale when +# every lane's (tile_max - running_max) <= this. exp2(8)=256 bounds P safely. +_LAZY_RESCALE_THRESHOLD = 8.0 @dataclass(frozen=True) @@ -140,6 +143,15 @@ class AttentionDenseSpec: # (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); default off. + lazy_rescale: bool = False def __post_init__(self) -> None: if self.dtype not in _DTYPE_IR: @@ -237,6 +249,8 @@ def kernel_name(self) -> str: 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": @@ -286,6 +300,7 @@ def build_attention_dense( W = spec.sliding_window Wt = W // BN # window length in KV tiles (0 when disabled) varlen = spec.varlen + LAZY_RESCALE = spec.lazy_rescale K_STEPS = D // 16 D_TILES = D // 32 @@ -491,9 +506,20 @@ def softmax_max(s_reg, m_i): 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)) - m_new = b.fmax(m_i, tile_max) + 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 + 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 @@ -618,7 +644,7 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): do_mask(s0, start_tile, lower=True, upper=True) else: do_mask(s0, start_tile) - m0, _alpha0 = softmax_max(s0, neg_inf) + 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) @@ -637,6 +663,8 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): + [(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): m_i = carry[0] l_i = carry[1] @@ -653,13 +681,28 @@ def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): s = do_qk(kbuf) if mask_lower or mask_upper: do_mask(s, j, lower=mask_lower, upper=mask_upper) - m_new, alpha = softmax_max(s, m_i) + 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) - l_new = b.fadd(b.fmul(l_i, alpha), l_tile) - o_acc = rescale_o(o_acc, alpha) + 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))) @@ -807,6 +850,7 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: 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 @@ -1021,12 +1065,25 @@ def softmax_max(s_reg, m_i): 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)) - m_new = b.fmax(m_i, tile_max) + 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 + return m_new, alpha, skip def softmax_stats(s_reg, m_i): - m_new, alpha = softmax_max(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) @@ -1122,6 +1179,8 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): 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): m_i = carry[0] l_i = carry[1] @@ -1139,15 +1198,34 @@ def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): s = do_qk(kbuf) if mask_lower or mask_upper: do_mask(s, j, lower=mask_lower, upper=mask_upper) - m_new, alpha = softmax_max(s, m_i) + 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) - l_new = b.fadd(b.fmul(l_i, alpha), l_tile) - o_acc = rescale_o(o_acc, alpha) + 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))) From a0977c226b72deda58846b4ec8078f968ea65fa3 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Fri, 17 Jul 2026 01:12:56 +0000 Subject: [PATCH 06/11] (rocke): update readme --- .../gfx950/attention/prefill/README.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 index 153c09ea7ac9..77ea25240402 100644 --- 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 @@ -71,6 +71,24 @@ spec = AttentionDenseSpec( kernel = build_attention_dense(spec) # -> KernelDef; compile with backend="python" ``` +## 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`). From b778e4cca0c4d22913fa769dd26539895f74298a Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Fri, 17 Jul 2026 01:50:17 +0000 Subject: [PATCH 07/11] (rocke): precommit hook format --- .../prefill/benchmark_dense_prefill_live.py | 477 ++++++++++++++++++ .../golden/attention_dense_ir_sha256.json | 104 +++- .../tests/test_attention_dense_golden.py | 78 ++- .../rocke/platform/cpp/core/ir/core_types.cpp | 2 + .../rocke/platform/cpp/core/ir/ir_arith.cpp | 5 + .../platform/cpp/core/lower_llvm/arith.cpp | 6 + .../platform/cpp/core/lower_llvm/data.cpp | 1 + .../rocke/platform/cpp/core/verify.cpp | 1 + .../rocke/platform/cpp/include/rocke/ir.h | 2 + 9 files changed, 664 insertions(+), 12 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_dense_prefill_live.py 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/tests/golden/attention_dense_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_ir_sha256.json index 5b745f34003f..6e1741429b04 100644 --- 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 @@ -2,25 +2,113 @@ "flavors": { "llvm20": { "cases": { + "attention_dense/bn128_sq512": { + "bytes": 604858, + "sha256": "000a9408578a9728c7e9baa80428d50b4d4cc17dc3fc1d79e3ad36c6e8642ac5" + }, "attention_dense/default_causal_sq512": { - "bytes": 320191, - "sha256": "454a7ab4aec162e1d7dccc3fa8ca5a186ee74508d7763275823bde47723532b3" + "bytes": 337008, + "sha256": "281d6f975a9056969fb059357faf048a7a11d56557ebba075652aa81c8a6b5cb" + }, + "attention_dense/fp16_h64_sq512": { + "bytes": 227134, + "sha256": "70e5e3a9a5923a9e1209e39b800c6309a41265eaf1f45dc3df9d87f35645cb06" + }, + "attention_dense/lazy_rescale_sq512": { + "bytes": 341939, + "sha256": "1dc714eeaa2e69c3198eeadc725956f2c5257db2f6263f524dc3930ed707651d" + }, + "attention_dense/noncausal_sq512": { + "bytes": 214925, + "sha256": "c35254ad8eb1e051a5735c6314a6ccfaaeff0ea003db2b3b9be8804da8c314c2" + }, + "attention_dense/persist_hkvmaj_sq512": { + "bytes": 322579, + "sha256": "4ee6f48d6cbc461f71e041922f5dab82bfb5c42749087bc595973cb963cca526" + }, + "attention_dense/persist_intl_sq512": { + "bytes": 322053, + "sha256": "4b26b91f3677719a588dde62844cf6a42a4346817800b0e6be4700240e9d4946" + }, + "attention_dense/persist_lazy_sq512": { + "bytes": 326109, + "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" + }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 321142, + "sha256": "7f5251827f8e9e4853c16b0908950cbf1b878686007643438c9beb241254d6a3" + }, + "attention_dense/persist_swa_w128_sq512": { + "bytes": 437202, + "sha256": "4ccce534163778d590ad2e70a0c258e30330ec9fd80dd202fb91b4b6be034ec7" }, "attention_dense/persistent_causal_sq512": { - "bytes": 313398, - "sha256": "9b4e796e22a7475f4339c1cc521337f68b15d837301adb43672ba756d5638873" + "bytes": 321142, + "sha256": "7f5251827f8e9e4853c16b0908950cbf1b878686007643438c9beb241254d6a3" + }, + "attention_dense/swa_w128_sq512": { + "bytes": 457079, + "sha256": "cc49b513ebedd8c9691c41311cd2a1340e6c6c1066eab680bda1d852c7d0712d" + }, + "attention_dense/varlen_sq512": { + "bytes": 339044, + "sha256": "f0fdae668ea39aa9a8a0d5cd425a5d22a1607c0a768e79e4fd4c35e31c358ea7" } } }, "llvm22": { "cases": { + "attention_dense/bn128_sq512": { + "bytes": 604962, + "sha256": "fc497f9f50d107c9634b528447a9b8af35d5c50192f940409e936f084f12ff79" + }, "attention_dense/default_causal_sq512": { - "bytes": 320295, - "sha256": "eb781b077fb7944788a71d2636e2d2e139913b3f0f27ef48d2856467dfdb71ac" + "bytes": 337112, + "sha256": "11a2458a4ec86709fcec37da59d48c1de6a8b55ca4f23e5fbb27f13e5388222d" + }, + "attention_dense/fp16_h64_sq512": { + "bytes": 227234, + "sha256": "39cd4557fdda2cce4813b22ef11787d94804f5c420a2f1924cb398a00fbeb5f0" + }, + "attention_dense/lazy_rescale_sq512": { + "bytes": 342043, + "sha256": "0698cf388935622a76524c20f300b383a946ecbdb512e93c662db29ad574a1f0" + }, + "attention_dense/noncausal_sq512": { + "bytes": 215025, + "sha256": "408840a8a2cac5bf8e38759694f344ac4f454e5e397d8d2e31f14a9e124f14d2" + }, + "attention_dense/persist_hkvmaj_sq512": { + "bytes": 322683, + "sha256": "cf7ea81a10abc6ab2e0be03eddacb7eaa316ecf7ef817c352f57ee1054ff49f8" + }, + "attention_dense/persist_intl_sq512": { + "bytes": 322157, + "sha256": "816f23613fa06f85f548cf0f797b389d34b616617db80ee67a5419ea4bc78bed" + }, + "attention_dense/persist_lazy_sq512": { + "bytes": 326213, + "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" + }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 321246, + "sha256": "6c805ce2f02d540afb7f4e36f6cb55edb19d3e0526d2b2880e8a0c2beb1bea4e" + }, + "attention_dense/persist_swa_w128_sq512": { + "bytes": 437306, + "sha256": "83b993a204700a4132aadaedb9e13291f8ea6ca9810f976676f337471b45d967" }, "attention_dense/persistent_causal_sq512": { - "bytes": 313502, - "sha256": "1f089fb980de6ef2a7d25bb61c1665adbebeafed36ae73d981092f3a37fcf978" + "bytes": 321246, + "sha256": "6c805ce2f02d540afb7f4e36f6cb55edb19d3e0526d2b2880e8a0c2beb1bea4e" + }, + "attention_dense/swa_w128_sq512": { + "bytes": 457183, + "sha256": "13f345abeb46865f08f4575956942e85e9e6f27d34166ba524b6b13d92d190b5" + }, + "attention_dense/varlen_sq512": { + "bytes": 339148, + "sha256": "0a0d6b81445a01b63b4ef793bb850afd81cbaafc521cdf74a0278c3677be0cf9" } } } 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 index 789de4dadcc2..ba974e47beb7 100644 --- 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 @@ -48,12 +48,39 @@ def _cases(): causal=True, dtype="bf16", ) + + def mk(**over): + d = dict(base) + d.update(over) + return lambda: build_attention_dense(AttentionDenseSpec(**d)) + return { - "attention_dense/default_causal_sq512": lambda: build_attention_dense( - AttentionDenseSpec(**base) + # --- 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_rescale_sq512": mk(lazy_rescale=True), + "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/persistent_causal_sq512": lambda: build_attention_dense( - AttentionDenseSpec(**base, persistent=True, num_persistent=256) + "attention_dense/persist_intl_sq512": mk( + persistent=True, num_persistent=256, interleave=True + ), + "attention_dense/persist_lazy_sq512": mk( + persistent=True, num_persistent=256, lazy_rescale=True + ), + "attention_dense/persist_swa_w128_sq512": mk( + persistent=True, num_persistent=256, sliding_window=128 ), } @@ -107,6 +134,49 @@ def test_attention_dense_ir_matches_golden(): 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) 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); From fa3e798e75bcb35a7af256370cf5efa541e9bb45 Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Fri, 17 Jul 2026 01:59:31 +0000 Subject: [PATCH 08/11] =?UTF-8?q?(rocke):=20dense=20attn=20=E2=80=94=20ena?= =?UTF-8?q?ble=20lazy=20online=20rescale=20by=20default=20Flip=20Attention?= =?UTF-8?q?DenseSpec.lazy=5Frescale=20default=20False->True=20so=20the=20p?= =?UTF-8?q?arity-identical=20(1.46e-3)=20lazy=20O/l=20rescale=20is=20alway?= =?UTF-8?q?s-on:=20persistent=20Sq=3D8192=20GQA=20reaches=20~980=20TFLOPS?= =?UTF-8?q?=20out-of-the-box=20(was=20~877).=20Re-bless=20the=20golden=20I?= =?UTF-8?q?R=20fixture=20and=20keep=20lazy-off=20branch=20coverage;=20C++/?= =?UTF-8?q?Python=20backend=3Dboth=20remains=2013/13=20byte-identical.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../library/kernels/gfx950/attention_dense.py | 14 ++- .../golden/attention_dense_ir_sha256.json | 112 +++++++++--------- .../tests/test_attention_dense_golden.py | 6 +- 3 files changed, 69 insertions(+), 63 deletions(-) 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 index c519274be31e..288f4e4f3278 100644 --- 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 @@ -30,9 +30,14 @@ (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% (bf16, Sq=8192, GQA -> ~970 TFLOPS with the +hkv-major decode + V-pad). + Experimental/negative levers from the sweep (step-2 8-cluster, K-staging, per-nsub -staging, score truncation, s_setprio, PV V-prefetch, lazy rescale) are intentionally -NOT carried over — see the experiment's ``plan.md`` for their measured results. +staging, score truncation, PV V-prefetch) are intentionally NOT carried over — see +the experiment's ``plan.md`` for their measured results. """ from dataclasses import dataclass @@ -150,8 +155,9 @@ class AttentionDenseSpec: # 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); default off. - lazy_rescale: bool = False + # (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: 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 index 6e1741429b04..7b5088154acb 100644 --- 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 @@ -3,112 +3,112 @@ "llvm20": { "cases": { "attention_dense/bn128_sq512": { - "bytes": 604858, - "sha256": "000a9408578a9728c7e9baa80428d50b4d4cc17dc3fc1d79e3ad36c6e8642ac5" + "bytes": 610799, + "sha256": "67feb714296db5af1d332df6427ea306727f0ec72215982d6048da5c398cc6b9" }, "attention_dense/default_causal_sq512": { - "bytes": 337008, - "sha256": "281d6f975a9056969fb059357faf048a7a11d56557ebba075652aa81c8a6b5cb" + "bytes": 341939, + "sha256": "1dc714eeaa2e69c3198eeadc725956f2c5257db2f6263f524dc3930ed707651d" }, "attention_dense/fp16_h64_sq512": { - "bytes": 227134, - "sha256": "70e5e3a9a5923a9e1209e39b800c6309a41265eaf1f45dc3df9d87f35645cb06" + "bytes": 230977, + "sha256": "e30bb63f718f6437eaac14bc49b347174b398c2fcc225dd1ca9d46959ddf2fa1" }, - "attention_dense/lazy_rescale_sq512": { - "bytes": 341939, - "sha256": "1dc714eeaa2e69c3198eeadc725956f2c5257db2f6263f524dc3930ed707651d" + "attention_dense/lazy_off_sq512": { + "bytes": 337008, + "sha256": "281d6f975a9056969fb059357faf048a7a11d56557ebba075652aa81c8a6b5cb" }, "attention_dense/noncausal_sq512": { - "bytes": 214925, - "sha256": "c35254ad8eb1e051a5735c6314a6ccfaaeff0ea003db2b3b9be8804da8c314c2" + "bytes": 217781, + "sha256": "db33f2ff07c424d2b3e759aed66e4b0d29056e7e4780ce11fdcc8d199da99334" }, "attention_dense/persist_hkvmaj_sq512": { - "bytes": 322579, - "sha256": "4ee6f48d6cbc461f71e041922f5dab82bfb5c42749087bc595973cb963cca526" + "bytes": 327546, + "sha256": "f49a0e2f5415bce83f55d251c2a5a7be00fcbd999db3b038023f50ab53113f74" }, "attention_dense/persist_intl_sq512": { - "bytes": 322053, - "sha256": "4b26b91f3677719a588dde62844cf6a42a4346817800b0e6be4700240e9d4946" - }, - "attention_dense/persist_lazy_sq512": { - "bytes": 326109, - "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" + "bytes": 327020, + "sha256": "b1bde5a27e66647c342d5ee4b2e010e52106d33d22d901cedf46f59e3ce9fca8" }, - "attention_dense/persist_qbmaj_sq512": { + "attention_dense/persist_lazy_off_sq512": { "bytes": 321142, "sha256": "7f5251827f8e9e4853c16b0908950cbf1b878686007643438c9beb241254d6a3" }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 326109, + "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" + }, "attention_dense/persist_swa_w128_sq512": { - "bytes": 437202, - "sha256": "4ccce534163778d590ad2e70a0c258e30330ec9fd80dd202fb91b4b6be034ec7" + "bytes": 444213, + "sha256": "3a2039ebb82f66fbf375f8ff9b636ce5cda48702a914541b854d5997ce2de222" }, "attention_dense/persistent_causal_sq512": { - "bytes": 321142, - "sha256": "7f5251827f8e9e4853c16b0908950cbf1b878686007643438c9beb241254d6a3" + "bytes": 326109, + "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" }, "attention_dense/swa_w128_sq512": { - "bytes": 457079, - "sha256": "cc49b513ebedd8c9691c41311cd2a1340e6c6c1066eab680bda1d852c7d0712d" + "bytes": 464090, + "sha256": "cbde8ed9be5ac5b89b68305c8b4092630bf5ee56ccd58cd2c8e70228879b1f1d" }, "attention_dense/varlen_sq512": { - "bytes": 339044, - "sha256": "f0fdae668ea39aa9a8a0d5cd425a5d22a1607c0a768e79e4fd4c35e31c358ea7" + "bytes": 344011, + "sha256": "afacdff3e337450ae8db2ec59bf87d5aea6949a23e83256bc4af8b97f63da80a" } } }, "llvm22": { "cases": { "attention_dense/bn128_sq512": { - "bytes": 604962, - "sha256": "fc497f9f50d107c9634b528447a9b8af35d5c50192f940409e936f084f12ff79" + "bytes": 610903, + "sha256": "b75fea02477fec773dc639131b94c33efda4ea2217f7822622c3b69dba42c165" }, "attention_dense/default_causal_sq512": { - "bytes": 337112, - "sha256": "11a2458a4ec86709fcec37da59d48c1de6a8b55ca4f23e5fbb27f13e5388222d" + "bytes": 342043, + "sha256": "0698cf388935622a76524c20f300b383a946ecbdb512e93c662db29ad574a1f0" }, "attention_dense/fp16_h64_sq512": { - "bytes": 227234, - "sha256": "39cd4557fdda2cce4813b22ef11787d94804f5c420a2f1924cb398a00fbeb5f0" + "bytes": 231077, + "sha256": "aeced7d0a10933d2eb2143fd127d7fdbc08ae86ac8c7f8e40fe30dca6dbcf2f3" }, - "attention_dense/lazy_rescale_sq512": { - "bytes": 342043, - "sha256": "0698cf388935622a76524c20f300b383a946ecbdb512e93c662db29ad574a1f0" + "attention_dense/lazy_off_sq512": { + "bytes": 337112, + "sha256": "11a2458a4ec86709fcec37da59d48c1de6a8b55ca4f23e5fbb27f13e5388222d" }, "attention_dense/noncausal_sq512": { - "bytes": 215025, - "sha256": "408840a8a2cac5bf8e38759694f344ac4f454e5e397d8d2e31f14a9e124f14d2" + "bytes": 217881, + "sha256": "83333cfcfd4d2e3c26ba2cf6e61faffc8dce46ad43cc1216d7efd68ad0653e54" }, "attention_dense/persist_hkvmaj_sq512": { - "bytes": 322683, - "sha256": "cf7ea81a10abc6ab2e0be03eddacb7eaa316ecf7ef817c352f57ee1054ff49f8" + "bytes": 327650, + "sha256": "8a0e7f52f54566f579cf9f0f20d074b0a57f5079c242f6a6b98d35ab9cd000a4" }, "attention_dense/persist_intl_sq512": { - "bytes": 322157, - "sha256": "816f23613fa06f85f548cf0f797b389d34b616617db80ee67a5419ea4bc78bed" - }, - "attention_dense/persist_lazy_sq512": { - "bytes": 326213, - "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" + "bytes": 327124, + "sha256": "6d9f80a763c0d4d992159487222d0ca30e8ff588e1a9c0347969e94e6b4a4373" }, - "attention_dense/persist_qbmaj_sq512": { + "attention_dense/persist_lazy_off_sq512": { "bytes": 321246, "sha256": "6c805ce2f02d540afb7f4e36f6cb55edb19d3e0526d2b2880e8a0c2beb1bea4e" }, + "attention_dense/persist_qbmaj_sq512": { + "bytes": 326213, + "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" + }, "attention_dense/persist_swa_w128_sq512": { - "bytes": 437306, - "sha256": "83b993a204700a4132aadaedb9e13291f8ea6ca9810f976676f337471b45d967" + "bytes": 444317, + "sha256": "6fe644368ff3c588edfcf80eb74e6ae9264458a9032b40ca78fdfb93e6f6c403" }, "attention_dense/persistent_causal_sq512": { - "bytes": 321246, - "sha256": "6c805ce2f02d540afb7f4e36f6cb55edb19d3e0526d2b2880e8a0c2beb1bea4e" + "bytes": 326213, + "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" }, "attention_dense/swa_w128_sq512": { - "bytes": 457183, - "sha256": "13f345abeb46865f08f4575956942e85e9e6f27d34166ba524b6b13d92d190b5" + "bytes": 464194, + "sha256": "35c4abc272427bc1f92ca1cfe06256e7e1a02e8a2690d21cec105c32269fc1c0" }, "attention_dense/varlen_sq512": { - "bytes": 339148, - "sha256": "0a0d6b81445a01b63b4ef793bb850afd81cbaafc521cdf74a0278c3677be0cf9" + "bytes": 344115, + "sha256": "a9bf509c16f6fde382fe7f778b3508f2b385b67011ed0f778f81035a12372998" } } } 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 index ba974e47beb7..971f2f1bd342 100644 --- 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 @@ -59,7 +59,7 @@ def mk(**over): "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_rescale_sq512": mk(lazy_rescale=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), @@ -76,8 +76,8 @@ def mk(**over): "attention_dense/persist_intl_sq512": mk( persistent=True, num_persistent=256, interleave=True ), - "attention_dense/persist_lazy_sq512": mk( - persistent=True, num_persistent=256, lazy_rescale=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 From 823566f06ee76fdff8b537859df354c5e9b0dd1a Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti Date: Fri, 17 Jul 2026 02:12:58 +0000 Subject: [PATCH 09/11] (rocke): delete test benches unwanted --- .../gfx950/attention/prefill/_flydsl_bench.py | 140 ----- .../gfx950/attention/prefill/_prof_compare.sh | 26 - .../gfx950/attention/prefill/_prof_dense.py | 90 --- .../attention/prefill/bench_swa_varlen.py | 524 ------------------ .../builders/gfx950/attention/prefill/plan.md | 248 --------- 5 files changed, 1028 deletions(-) delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py deleted file mode 100644 index cc842b2316a5..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_flydsl_bench.py +++ /dev/null @@ -1,140 +0,0 @@ -#!/usr/bin/env python3 -"""Standalone flyDSL flash-attn timer (own process + interpreter). - -Uses flyDSL **0.2.4** (pip-installed into /tmp/flydsl024-venv, torch inherited -from atom-venv via a .pth) + the flydsl-main HEAD kernels (also 0.2.4). This is -the DUALWAVE_SWP path: dense/varlen, GQA/MHA, causal. flyDSL prefill has NO -sliding window. Run this with /tmp/flydsl024-venv/bin/python and -PYTHONPATH=/workspace/flydsl-main. - -Modes: - dense - varlen (causal) -Prints: "OK " or "ERR ". -""" -import sys - -import torch - - -def _time(call, wu, it): - for _ in range(wu): - call() - torch.cuda.synchronize() - s = torch.cuda.Event(enable_timing=True) - e = torch.cuda.Event(enable_timing=True) - s.record() - for _ in range(it): - call() - e.record() - e.synchronize() - return s.elapsed_time(e) / it - - -def main(): - from kernels.attention.flash_attn_interface import flydsl_flash_attn_func - - dev = "cuda" - DT = torch.bfloat16 - torch.manual_seed(1) - mode = sys.argv[1] - - if mode == "dense": - B, S, Hq, Hkv, D = (int(sys.argv[i]) for i in range(2, 7)) - causal = bool(int(sys.argv[7])) - wu, it = int(sys.argv[8]), int(sys.argv[9]) - q = (torch.randn(B, S, Hq, D, dtype=DT, device=dev) * 0.2).contiguous() - k = (torch.randn(B, S, Hkv, D, dtype=DT, device=dev) * 0.2).contiguous() - v = (torch.randn(B, S, Hkv, D, dtype=DT, device=dev) * 0.2).contiguous() - - def call(): - return flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv) - - o = call() - torch.cuda.synchronize() - rep = Hq // Hkv - ref = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2).float(), - k.transpose(1, 2).repeat_interleave(rep, 1).float(), - v.transpose(1, 2).repeat_interleave(rep, 1).float(), - is_causal=causal, - ).transpose(1, 2) - err = (o.float() - ref).abs().max().item() - ms = _time(call, wu, it) - pairs = (S * (S + 1) // 2) if causal else S * S - tf = 4 * B * Hq * D * pairs / (ms * 1e-3) / 1e12 - print(f"OK {ms:.6f} {tf:.3f} {err:.3e}") - return - - if mode == "varlen": - seqlens = [int(x) for x in sys.argv[2].split(",")] - Hq, Hkv, D = int(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5]) - wu, it = int(sys.argv[6]), int(sys.argv[7]) - B = len(seqlens) - total = sum(seqlens) - mx = max(seqlens) - 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() - cu = torch.zeros(B + 1, dtype=torch.int32, device=dev) - cu[1:] = torch.tensor(seqlens, dtype=torch.int32, device=dev).cumsum(0) - - def call(): - return flydsl_flash_attn_func( - q, - k, - v, - causal=True, - num_kv_heads=Hkv, - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - max_seqlen_q=mx, - cross_seqlen=False, - ) - - o = call() - torch.cuda.synchronize() - rep = Hq // Hkv - err = 0.0 - for i, s in enumerate(seqlens): - st = sum(seqlens[:i]) - qh = q[st : st + s].transpose(0, 1).float().unsqueeze(0) - kh = ( - k[st : st + s] - .transpose(0, 1) - .repeat_interleave(rep, 0) - .float() - .unsqueeze(0) - ) - vh = ( - v[st : st + s] - .transpose(0, 1) - .repeat_interleave(rep, 0) - .float() - .unsqueeze(0) - ) - ref = ( - torch.nn.functional.scaled_dot_product_attention( - qh, kh, vh, is_causal=True - ) - .squeeze(0) - .transpose(0, 1) - ) - err = max(err, (o[st : st + s].float() - ref).abs().max().item()) - ms = _time(call, wu, it) - pairs = sum(x * (x + 1) // 2 for x in seqlens) - tf = 4 * Hq * D * pairs / (ms * 1e-3) / 1e12 - print(f"OK {ms:.6f} {tf:.3f} {err:.3e}") - return - - print(f"ERR unknown mode {mode!r}") - - -if __name__ == "__main__": - try: - main() - except Exception as exc: # noqa: BLE001 - import traceback - - traceback.print_exc() - print(f"ERR {exc!r}") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh deleted file mode 100644 index 2557c0ac5497..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_compare.sh +++ /dev/null @@ -1,26 +0,0 @@ -#!/bin/bash -# rocprofv3 counter comparison: rocke dense vs flyDSL 0.2.4, GQA Hq128/Hkv8 S8192 causal. -# Each metric runs in its own single-pass rocprofv3 invocation (fast, no app replay). -set -u -HERE="$(cd "$(dirname "$0")" && pwd)" -OUT=/tmp/prof_cmp -rm -rf "$OUT"; mkdir -p "$OUT" -ATOM=/workspace/atom-venv/bin/python -FLY=/tmp/flydsl024-venv/bin/python - -METRICS=(MfmaUtil MeanOccupancyPerCU OccupancyPercent FetchSize WriteSize LDSBankConflict MemUnitStalled SQ_WAVES) - -dense_cmd=($ATOM "$HERE/_prof_dense.py" 8192 128 8 128 1 8) -fly_cmd=($FLY "$HERE/_flydsl_bench.py" dense 1 8192 128 8 128 1 1 8) - -for M in "${METRICS[@]}"; do - rocprofv3 --pmc "$M" --kernel-include-regex "rocke_attention_dense" \ - -d "$OUT/dense_$M" --output-format csv -- "${dense_cmd[@]}" \ - >/dev/null 2>&1 - PYTHONPATH=/workspace/flydsl-main rocprofv3 --pmc "$M" \ - --kernel-include-regex "flash_attn_dualwave_swp" \ - -d "$OUT/flydsl_$M" --output-format csv -- "${fly_cmd[@]}" \ - >/dev/null 2>&1 - echo "done $M" -done -echo "ALL_DONE" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py deleted file mode 100644 index f853c0e3432d..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/_prof_dense.py +++ /dev/null @@ -1,90 +0,0 @@ -#!/usr/bin/env python3 -"""Minimal single-kernel launcher for rocprofv3 profiling of the rocke dense -prefill kernel. Warmup builds/JITs; then launches the kernel `iters` times so -the profiler captures steady-state dispatches. - -Args: S Hq Hkv D causal(0/1) iters -""" -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 ( - AttentionDenseSpec, - build_attention_dense, -) # noqa: E402 -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 - - -def main(): - S = int(sys.argv[1]) if len(sys.argv) > 1 else 8192 - Hq = int(sys.argv[2]) if len(sys.argv) > 2 else 128 - Hkv = int(sys.argv[3]) if len(sys.argv) > 3 else 8 - D = int(sys.argv[4]) if len(sys.argv) > 4 else 128 - causal = bool(int(sys.argv[5])) if len(sys.argv) > 5 else True - iters = int(sys.argv[6]) if len(sys.argv) > 6 else 20 - dev = "cuda" - DT = torch.bfloat16 - torch.manual_seed(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=causal, - dtype="bf16", - persistent=True, - num_persistent=256, - ) - art = compile_kernel( - build_attention_dense(spec), - arch="gfx950", - backend="python", - capture_ir_text=False, - ) - sig = ( - SignatureBuilder() - .ptr("q_ptr", "bf16") - .ptr("k_ptr", "bf16") - .ptr("v_ptr", "bf16") - .ptr("o_ptr", "bf16") - .scalar("scale", "f32") - .build() - ) - lch = KernelLauncher(hsaco=art.hsaco, kernel_name=art.kernel_name, signature=sig) - stream = torch.cuda.current_stream().cuda_stream - cfg = LaunchConfig( - grid=(256, 1, 1), block=(spec.num_waves * 64, 1, 1), stream=stream - ) - vals = { - "q_ptr": q, - "k_ptr": k, - "v_ptr": v, - "o_ptr": out, - "scale": 1.0 / math.sqrt(D), - } - # warmup / JIT - lch(vals, config=cfg) - torch.cuda.synchronize() - sys.stderr.write(f"[prof_dense] kernel={art.kernel_name}\n") - for _ in range(iters): - lch(vals, config=cfg) - torch.cuda.synchronize() - - -if __name__ == "__main__": - main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py deleted file mode 100644 index c99fc9098225..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/bench_swa_varlen.py +++ /dev/null @@ -1,524 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark dense (rocke) vs flyDSL vs unified_attention across sliding-window -and varlen-batch prefill shapes on gfx950 / MI355X. - -Backends --------- -* ``dense`` : the productized rocke dense prefill kernel with the new - ``sliding_window`` (A) and ``varlen`` (B) features - (``kernels/gfx950/attention_dense.py``). -* ``flydsl`` : ``flydsl_flash_attn_func`` (DUALWAVE_SWP). Causal dense / packed - varlen only -- flyDSL prefill has NO sliding-window, so on SWA - rows it runs FULL causal (a "no-window" reference, marked *). -* ``unified``: rocke ``run_unified_attention_torch`` (paged 2D/3D). Supports - both sliding-window and varlen (paged). - -All three are timed with the same torch CUDA-event timer on the current stream. -Reported TFLOPS uses the honest banded pair count (windowed work); flyDSL's SWA -rows compute full-causal work so its TFLOPS bar is on a *different* workload -(hence the *). Correctness for dense is checked vs a banded-mask SDPA reference. - -Usage: - python bench_swa_varlen.py --mode swa - python bench_swa_varlen.py --mode varlen - python bench_swa_varlen.py --mode both --iters 50 -""" -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, -) -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 - -_DT = torch.bfloat16 - - -def _pairs(s, W): - """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, Hq, D): - return sum(4 * Hq * D * _pairs(s, W) for s in seqlens) - - -def _time(fn, warmup, iters): - for _ in range(warmup): - fn() - torch.cuda.synchronize() - s = torch.cuda.Event(enable_timing=True) - e = torch.cuda.Event(enable_timing=True) - s.record() - for _ in range(iters): - fn() - e.record() - e.synchronize() - return s.elapsed_time(e) / iters - - -# --------------------------------------------------------------------------- # -# rocke dense kernel -# --------------------------------------------------------------------------- # -_LAUNCHER_CACHE = {} - - -def _dense_launcher(spec): - key = spec.kernel_name() - if key in _LAUNCHER_CACHE: - return _LAUNCHER_CACHE[key] - 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 - - -def bench_dense(seqlens, W, Hq, Hkv, D, warmup, iters, persistent): - dev = "cuda" - B = len(seqlens) - max_s = max(seqlens) - total = sum(seqlens) - scale = 1.0 / math.sqrt(D) - stream = torch.cuda.current_stream().cuda_stream - torch.manual_seed(0) - varlen = B > 1 or (total != B * max_s) - - if varlen: - 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="bf16", - sliding_window=W, - varlen=True, - ) - lch = _dense_launcher(spec) - cfg = LaunchConfig( - grid=(max_s // 256, Hq, B), block=(spec.num_waves * 64, 1, 1), 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: - 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="bf16", - sliding_window=W, - persistent=persistent, - num_persistent=256, - ) - lch = _dense_launcher(spec) - if persistent: - grid = (256, 1, 1) - else: - grid = (S // 256, Hq, 1) - cfg = LaunchConfig(grid=grid, block=(spec.num_waves * 64, 1, 1), 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 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(call, warmup, iters) - tf = _flops(seqlens, W, Hq, D) / (ms * 1e-3) / 1e12 - return ms, tf, max_err - - -# --------------------------------------------------------------------------- # -# flyDSL -# --------------------------------------------------------------------------- # -# flyDSL 0.2.4 (DUALWAVE_SWP: dense/varlen, GQA/MHA, causal; NO sliding window). -# pip-installed into /tmp/flydsl024-venv (torch inherited from atom-venv via .pth); -# kernels come from the flydsl-main HEAD checkout. Runs in a subprocess with its -# own interpreter + PYTHONPATH so its top-level `kernels` package doesn't collide -# with rocke's. -_FLYDSL_PY = "/tmp/flydsl024-venv/bin/python" -_FLYDSL_ENV = {**os.environ, "PYTHONPATH": "/workspace/flydsl-main"} - - -def _flydsl_run(args): - import subprocess - - cmd = [_FLYDSL_PY, os.path.join(_HERE, "_flydsl_bench.py"), *args] - try: - out = subprocess.run( - cmd, capture_output=True, text=True, timeout=600, env=_FLYDSL_ENV - ) - except Exception as e: # noqa: BLE001 - return None, None, f"subproc: {e!r}" - line = "" - for ln in out.stdout.splitlines(): - if ln.startswith(("OK ", "ERR ")): - line = ln - if line.startswith("OK "): - p = line.split() - return float(p[1]), float(p[2]), (float(p[3]) if len(p) > 3 else "") - return ( - None, - None, - ( - line[4:] - if line.startswith("ERR ") - else (out.stderr.strip().splitlines() or ["no output"])[-1] - ), - ) - - -def bench_flydsl_dense(B, S, Hq, Hkv, D, warmup, iters, causal=True): - return _flydsl_run( - [ - "dense", - str(B), - str(S), - str(Hq), - str(Hkv), - str(D), - str(int(causal)), - str(warmup), - str(iters), - ] - ) - - -def bench_flydsl_varlen(seqlens, Hq, Hkv, D, warmup, iters): - sl = ",".join(str(x) for x in seqlens) - return _flydsl_run( - ["varlen", sl, str(Hq), str(Hkv), str(D), str(warmup), str(iters)] - ) - - -# --------------------------------------------------------------------------- # -# unified_attention (paged) -# --------------------------------------------------------------------------- # -def bench_unified(seqlens, W, Hq, Hkv, D, warmup, iters): - try: - from rocke.assets import shape_utils_dir - - sys.path.insert(0, str(shape_utils_dir())) - import _ua_shape_utils as u - from kernels import UnifiedAttentionProblem, run_unified_attention_torch - except Exception as e: # noqa: BLE001 - return None, None, f"import: {e!r}" - dev = "cuda" - B = len(seqlens) - total = sum(seqlens) - max_s = max(seqlens) - block_size = 16 - num_blocks = sum((s + block_size - 1) // block_size for s in seqlens) + B + 16 - max_blocks = (max_s + block_size - 1) // block_size - win = (W - 1, 0) if (W and W > 0) else (-1, -1) - try: - shape = u.UAShape( - source_file="synthetic", - line_idx=0, - call_idx=0, - kind="prefill", - all_decode=False, - num_seqs=B, - total_q=total, - num_query_heads=Hq, - num_kv_heads=Hkv, - head_size=D, - block_size=block_size, - num_blocks=num_blocks, - max_blocks_per_seq=max_blocks, - max_seqlen_q=max_s, - max_seqlen_k=max_s, - softmax_scale=1.0 / math.sqrt(D), - softcap=0.0, - window_size=win, - has_sinks=False, - has_alibi=False, - has_output_scale=False, - q_dtype="torch.bfloat16", - k_dtype="torch.bfloat16", - v_dtype="torch.bfloat16", - out_dtype="torch.bfloat16", - ) - data = u.make_inputs(shape, seed=2, cap_blocks=None) - sw = W if (W and W > 0) else 0 - prob = UnifiedAttentionProblem( - total_q=shape.total_q, - num_seqs=shape.num_seqs, - num_query_heads=Hq, - num_kv_heads=Hkv, - head_size=D, - block_size=block_size, - max_seqlen_q=max_s, - max_seqlen_k=max_s, - dtype="bf16", - sliding_window=sw, - softcap=0.0, - use_sinks=False, - use_alibi=False, - use_qq_bias=False, - use_fp8=False, - num_sms=256, - ) - hip_stream = torch.cuda.current_stream().cuda_stream - - def call(): - run_unified_attention_torch( - problem=prob, - q=data["query"], - k=data["key_cache"], - v=data["value_cache"], - out=data["output"], - cu_seqlens_q=data["cu_seqlens_q"], - seqused_k=data["kv_lens"], - softmax_scale=data["scale"], - block_table=data["block_tables"], - softcap=0.0, - sinks=data.get("sinks"), - alibi_slopes=data.get("alibi_slopes"), - qq_bias=None, - qq_bias_stride_0=0, - backend="auto", - stream=hip_stream, - ) - - call() - torch.cuda.synchronize() - except Exception as e: # noqa: BLE001 - return None, None, f"run: {e!r}" - ms = _time(call, warmup, iters) - tf = _flops(seqlens, W, Hq, D) / (ms * 1e-3) / 1e12 - return ms, tf, "" - - -def _row(label, res): - ms, tf, note = res - if ms is None: - return f" {label:<10} FAILED ({note})" - extra = f" max_abs={note:.2e}" if isinstance(note, float) else "" - return f" {label:<10} {ms:8.4f} ms {tf:8.1f} TFLOPS{extra}" - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument( - "--mode", choices=["swa", "varlen", "mha", "both", "all"], default="all" - ) - 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("--warmup", type=int, default=10) - ap.add_argument("--iters", type=int, default=50) - ap.add_argument("--persistent", action="store_true", default=True) - args = ap.parse_args() - Hq, Hkv, D = args.hq, args.hkv, args.d - - if args.mode in ("mha", "all"): - print("\n=== DENSE-CAUSAL GQA (Hq=128, Hkv=8) -- headline prefill config ===") - for S in (2048, 4096, 8192): - print(f"\n-- B=1 S={S} Hq={Hq} Hkv={Hkv} D={D} causal --") - print( - _row( - "dense", - bench_dense( - [S], 0, Hq, Hkv, D, args.warmup, args.iters, args.persistent - ), - ) - ) - print( - _row( - "flydsl", - bench_flydsl_dense(1, S, Hq, Hkv, D, args.warmup, args.iters, True), - ) - ) - print( - _row( - "unified", - bench_unified([S], 0, Hq, Hkv, D, args.warmup, args.iters), - ) - ) - - print("\n=== DENSE-CAUSAL MHA ===") - for H in (16, 32): - for S in (4096, 8192): - print(f"\n-- B=1 S={S} H={H} (MHA) D={D} causal --") - print( - _row( - "dense", - bench_dense( - [S], 0, H, H, D, args.warmup, args.iters, args.persistent - ), - ) - ) - print( - _row( - "flydsl", - bench_flydsl_dense( - 1, S, H, H, D, args.warmup, args.iters, True - ), - ) - ) - print( - _row( - "unified", - bench_unified([S], 0, H, H, D, args.warmup, args.iters), - ) - ) - - if args.mode in ("swa", "both", "all"): - print("\n\n=== SLIDING-WINDOW sweep (single-seq prefill, GQA, causal) ===") - print(" (flyDSL v0.2.0 prefill has NO sliding-window -> not shown)") - 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}" - print(f"\n-- S={S} {tag} (Hq={Hq} Hkv={Hkv}) --") - print( - _row( - "dense", - bench_dense( - [S], W, Hq, Hkv, D, args.warmup, args.iters, args.persistent - ), - ) - ) - print( - _row( - "unified", - bench_unified([S], W, Hq, Hkv, D, args.warmup, args.iters), - ) - ) - - if args.mode in ("varlen", "both", "all"): - print("\n\n=== VARLEN sweep (packed ragged batch, GQA, causal) ===") - print(" (flyDSL: W=0 only -- no sliding window in flyDSL prefill)") - 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}" - print( - f"\n-- B={len(seqlens)} total={sum(seqlens)} " - f"seqlens={seqlens} {tag} --" - ) - print( - _row( - "dense", - bench_dense( - seqlens, W, Hq, Hkv, D, args.warmup, args.iters, False - ), - ) - ) - if W == 0: - print( - _row( - "flydsl", - bench_flydsl_varlen( - seqlens, Hq, Hkv, D, args.warmup, args.iters - ), - ) - ) - print( - _row( - "unified", - bench_unified(seqlens, W, Hq, Hkv, D, args.warmup, args.iters), - ) - ) - - -if __name__ == "__main__": - main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md deleted file mode 100644 index 398b2d8d34b9..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/plan.md +++ /dev/null @@ -1,248 +0,0 @@ -# Plan: porting sliding-window, batching, and bias to the dense prefill kernel - -Scope: extend the productized gfx950 dense flash-attention prefill kernel -(`kernels/gfx950/attention_dense.py`, host builder `attention_dense_prefill.py`) -with **sliding-window attention (SWA)**, richer **batching** (varlen), and -**bias addition** (ALiBi / additive bias / logit softcap), informed by a -structured read of the flyDSL implementations in `~/flydsl-main`. - -This document records (1) what flyDSL actually does for each feature — tile -layouts, MFMA atom sizes, layout transforms, wave/register/LDS allocation, and -tile/wave/K-loop scheduling — and (2) the concrete change plan for the rocke -dense kernel, with the always-on 877-TFLOPS pipeline preserved as the no-feature -fast path. - ---- - -## 0. Baseline: the rocke dense kernel we are extending - -- `BLOCK_M = 256` query rows/CTA, `block_n = 64` KV tile, `NBUF = 2` double - buffer, `num_waves = 8` (each wave owns 32 query rows), `waves_per_eu = 2`, - 256 VGPR / 0 spill. -- MFMA `32x32x16` for both QK and PV (`mfma_32x32x16_for_dtype`). -- CK-1 transposed PV (P fed to PV MMA in native QK-output layout via half-local - V load `pv32_v_load_paired`; no cross-half P relayout). -- LDS K padded `[NBUF, BN, D+8]` (bank-conflict fix); V unpadded `[NBUF, BN, D]`. -- Diagonal-only causal masking: mask-free body loop + masked diagonal tail. -- Depth-1 cluster fusing `exp2` into the PV MFMA loop; `sched_group_barrier` - template (DS_READ/MFMA/VALU/TRANS). -- Default grid `(Sq//256, Hq, B)`; persistent grid-stride variant loops one 1-D - grid of `num_persistent` CTAs over `W = (Sq//256)*Hq*B`. -- Inputs: dense contiguous `[B, S, H, D]`, uniform seqlen; params - `q,k,v,o,scale`. No varlen/paged, no bias, no sliding window. - -**Key finding:** flyDSL's production gfx950 prefill kernel -(`flash_attn_gfx950.py`, "DUALWAVE_SWP") is the *same design family* as this -kernel — same `BLOCK_M=256`/`BLOCK_N=64`, same `32x32x16` atoms, 8 waves × 32 -rows, P-in-registers, V transpose-read, double-buffered `[K0,V0,K1,V1]`, -`sched_group_barrier` + `s_setprio` + staggered wave barriers, and a lazy online -rescale. So flyDSL is a high-fidelity reference for the pipeline mechanics; the -deltas we need are exactly the three requested features. - ---- - -## 1. flyDSL findings by feature - -### 1a. Core pipeline + batching (`flash_attn_gfx950.py`, `flash_attn_generic.py`) - -- **Tile layout / atoms:** `BLOCK_M=256`, `BLOCK_N=64`, `K_SUB_N=32`, 8 waves × - 32 rows; MFMA `f32_32x32x16_{bf16,f16}` (K=16, v8 packs, v16f32 acc) for QK - and PV; fp8 path uses `mfma_f32_32x32x16_fp8_fp8` with optional wide - `32x32x64` PV. Q register-resident, pre-scaled by `1/sqrt(D)*log2(e)`. -- **Transforms:** K uses either XOR swizzle (`col ^ ((row&mask)<<4)`) or a padded - `u_rk` layout (`+8 bf16`); V is transpose-read with `ds_read_b64_tr_b16` - (padded `+32 bf16`) — matches our `pv32_v_load_paired`. P stays in VGPRs (no - LDS roundtrip); half-wave partners combined via `permlane32_swap`. O store - fuses partners with `permlane32_swap + cvt_pk_bf16_f32` → `buffer_store_dwordx4`. -- **Wave/register/LDS:** WG=512 threads (8 waves), `waves_per_eu=2`, interleaved - double buffer `[K0,V0,K1,V1]` (~68 KiB @ D=128). DUALWAVE staggers two - 4-wave groups with extra barriers to time-multiplex the SW pipeline. -- **K-loop scheduling:** hand-built 8-cluster software pipeline, `j += 2` (2 KV - tiles/iter), PV lags QK by ~1 tile; async `buffer_load_lds` (128B) with split - `s_waitcnt` (lgkmcnt vs vmcnt); `sched_group_barrier` (MFMA/VALU/EXP masks), - `s_setprio`, `s_nop`. **Lazy online rescale** skips O/P rescale when - `max(m_tile - m_row) <= 8.0` (ballot-vs-exec test) — a lever we have not baked. -- **Batching:** 3D grid `(NUM_HEADS_Q, num_q_blocks, grid_z)` with - `grid_z = batch` (or `batch*num_kv_splits` for split-K); GQA via - `kv_head = q_head // group_size`; **varlen** via `cu_seqlens_q/kv` rebased - buffer descriptors (`q_tok_base = cu[b]`, `seqlen_q_b = cu[b+1]-cu[b]`), grid - sized by `max_seqlen_q` with early-exit of empty q-tiles; **paged** via a - block table staged in LDS; **split-K** with an fp32 workspace + combine kernel. - **No persistent/grid-stride** (our persistent mode is a rocke-specific lever). - -### 1b. Sliding window (`pa_decode_swa.py` — decode/paged/fp8 only) - -**flyDSL has SWA only on the paged-decode path; flash-attn prefill is causal-only.** -The technique is directly portable: - -- **Masking math (compile-time `W`):** per query, `causal_bound = C + 1 - L + qi` - (exclusive upper), `seq_start = C - L - W + qi` (inclusive lower); allowed iff - `seq_start <= k < causal_bound`, i.e. `k ∈ [q_abs - W, q_abs]` (**W+1** keys - incl. diagonal; note the `+1` convention). `q_abs = C - L + qi`. -- **Two-level pruning:** (a) tile-level — `seq_start_global = max(0, C-L-W)`, - `tail_start_tile = seq_start_global >> 8` (tiles are 256 tokens), only visit - `[tail_start_tile, num_tiles)`; host bounds partitions with - `ceil((W+L-1)/256)+1`. (b) per-lane range mask on the MFMA `f32x4` output: - `in_range = (tok < causal_bound) & (tok >= seq_start)`, `select(in_range, s, -inf)`. -- **Mask placement / layout:** applied **after QK MFMA, before row-max/exp2**, in - the *native MFMA accumulator layout* — **no relayout/transpose**. Token indices - built inline (`kv_tok_base + td*16 + i`); bounds are SGPR scalars broadcast - across the 4 lanes. -- **No mask-free interior fast path in SWA** (every visited tile is masked). But - flash-attn *causal* has `tile_needs_mask` (skip mask when - `max_kv_col <= q_start`) — the optimization to replicate for the SWA interior. -- **Wave cooperation:** all waves iterate the same tile range; per-wave rows - differ only in the per-lane bound; cooperative K/V loads unchanged. -- **Cost:** +1 compare vs causal-only (two bounds instead of one) per `f32x4`. - -### 1c. Bias / ALiBi / softcap — **not implemented anywhere in flyDSL** - -Full-repo grep confirms: no ALiBi, no additive bias tensor, no QQ-bias, no -softcap in any kernel; the only `alibi_slopes`/`bias` refs are test harnesses -passing `None` to the external aiter reference. flyDSL applies only causal / KV-pad -/ SWA masks + QK softmax scaling (+ fp8 quant scales). So for bias we have **no -flyDSL reference to port** — the rocke `unified_attention` tiled path is the -in-repo reference (it *does* implement ALiBi + QQ-bias + softcap). flyDSL still -tells us the *insertion mechanics*: - -- **Insertion point:** after QK MFMA, **before mask/row-max/exp2**, on the raw - accumulator layout (same place SWA masks) — no relayout. -- **Domain:** gfx950 pre-scales Q by `sm_scale*log2(e)`, so scores are already in - log2 domain; any bias must be added **in that domain** (multiply linear biases - by `log2(e)` = `RCP_LN2`), matching how the rocke tiled path applies ALiBi - (`slope*(key_pos-ctx)*RCP_LN2`) and QQ-bias (`qq_bias[...]*RCP_LN2`). -- **ALiBi:** `S += slope_h * (k_pos - q_pos)` using the same per-lane - `k_pos`/`q_pos` maps SWA/causal already compute (`_mfma_32x32_c_row/_col`). -- **Softcap:** nonlinear `S = cap * tanh(S/cap)` applied **after scale, before - mask** (must precede max). Careful: our exp2 fast path assumes softmax arg - `<= 0` and no overflow guard — softcap changes the score range, revisit. -- **Composition:** per-row constant biases cancel in softmax; **relative** biases - (ALiBi, elementwise, softcap) do not and must be applied before max so online - rescale stays consistent. - ---- - -## 2. Porting plan - -Ordering by value/effort: **(A) sliding window** (cheap, high value, byte-safe -default) → **(B) varlen batching** → **(C) bias/ALiBi/softcap**. Each is gated so -the no-feature path emits today's exact 877-TFLOPS kernel. - -### A. Sliding window - -Design already scoped in prior analysis; flyDSL confirms the mechanics. - -- **Spec:** add `sliding_window: int = 0` to `AttentionDenseSpec` (0 = disabled → - bit-identical to today). Bake into `kernel_name()`; require `W % block_n == 0` - for the clean path (`__post_init__`). Mirror in the persistent builder. -- **Geometry:** both causal upper (`k<=q`) and window lower (`k>=q-W+1`) are - slope-1 diagonals; the valid region per 256-row block is a **parallelogram - band**. With `n_per = BLOCK_M//BN = 4` and `Wt = W//BN`: - `lo_tile = max(0, qb*n_per - Wt)`, `hi_tile = min(n_ktiles, (qb+1)*n_per)`. - Each boundary band is exactly `n_per = 4` tiles wide. -- **Three-phase K-loop:** (L) window-edge tiles `[lo_tile, lo_tile+n_per)` lower- - masked; (M) interior `[lo_tile+n_per, diag_start)` **mask-free** (identical hot - loop → same peak); (R) causal diagonal `[diag_start, hi_tile)` upper-masked - (existing tail). If `W < BLOCK_M` bands overlap → single merged both-bounds - phase. Reuse the existing `select`-guarded bound idiom for empty phases. -- **Masking:** extend `do_mask` with `mask_lower`/`mask_upper` flags; reuse - `_mfma_32x32_c_row/_col` exactly (no new layout). Lower predicate - `cmp_gt(ktok, query_tok - W)` (W compile-time → immediate). -- **Register pinning:** hoist `query_tok` (lane-only, tile-invariant) and - `query_tok-(W-1)` to once-per-block VGPRs; masking stays outside the loop- - carried set (`m,l,o,pk`) → no VGPR growth, preserve 0-spill. -- **Pipeline:** prologue primes `lo_tile`/`lo_tile+1` (not 0/1) and left-masks the - prologue tile; init `m=-inf,l=0,o=0` unchanged (tiles ` Date: Fri, 17 Jul 2026 03:13:31 +0000 Subject: [PATCH 10/11] (rocke): black formatting --- .../builders/gfx950/attention/ALGORITHM.md | 143 +++++- .../library/kernels/gfx950/attention_dense.py | 422 +++++++++++++----- .../golden/attention_dense_ir_sha256.json | 48 +- .../tests/test_attention_dense_golden.py | 25 ++ 4 files changed, 503 insertions(+), 135 deletions(-) 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..d0baed521bf5 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,149 @@ 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):** default grid ≈543 TFLOPS; **persistent ≈877 TFLOPS**, and +**≈948 TFLOPS with the `auto`→`hkv_major` decode + V-pad + lazy rescale** +(byte-identical IR ⇒ this number is schedule, not luck). 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/kernels/gfx950/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_dense.py index 288f4e4f3278..1b030fbf984d 100644 --- 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 @@ -35,11 +35,30 @@ parity-identical (1.46e-3) and ~+2% (bf16, Sq=8192, GQA -> ~970 TFLOPS with the hkv-major decode + V-pad). +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 + ktok= seqlen_kv > every real query), so causal needs + # NO extra key mask; non-causal adds a ktok None: raise ValueError( f"dtype must be one of {sorted(_DTYPE_IR)}, got {self.dtype}" ) - if self.head_size % 32 != 0: - raise ValueError( - f"head_size must be a multiple of 32, got {self.head_size}" - ) + # 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.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.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 " @@ -233,7 +278,7 @@ def resolved_persist_decode(self) -> str: if self.persist_decode != "auto": return self.persist_decode gqa = self.num_query_heads // self.num_kv_heads - nqb = self.seqlen_q // _BLOCK_M + 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" @@ -251,6 +296,8 @@ def kernel_name(self) -> str: 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: @@ -306,6 +353,7 @@ def build_attention_dense( 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 @@ -315,6 +363,11 @@ def build_attention_dense( 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 @@ -390,26 +443,39 @@ def build_attention_dense( ) # --- LDS allocation: PAD on K (bank-conflict fix), +PAD_V on V (transposed - # PV read bank-conflict pad) --- - LDROW = D + PAD - VROW = D + _LDS_PAD_V + # 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) - raw = b.global_load_vN(q, addr, dtype, 8, align=16) + 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)) - n_ktiles = Skv // BN + # 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 @@ -426,38 +492,57 @@ def build_attention_dense( 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): - """Row-by-row async DMA into padded [BN, LDROW] K_lds layout.""" - buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(K_BYTES_PER_BUF)) - 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(K_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( - k_rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 - ) + _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): - """Row-by-row async DMA into [BN, VROW] V_lds (row-strided so V can carry - the +PAD_V transposed-PV bank pad); 64 lanes x 2 bf16 = D per row.""" - buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) - 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(V_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( - v_rsrc, row_base, b.mul(voff, b.const_i32(2)), zero_soff, 1 - ) + _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)) @@ -506,6 +591,20 @@ def do_mask(s_reg, tile_idx, lower=False, upper=True): 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): @@ -645,11 +744,16 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): 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 = [ @@ -671,7 +775,7 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): _rs_ctr = [0] - def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): + 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]) @@ -687,6 +791,8 @@ def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): 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%) @@ -771,7 +877,7 @@ def _clamp(x, lo, hi): 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) + emit_loop_body(j, carry, mask_kbound=RAG_KBOUND) res = loop.results l_i = res[1] @@ -798,17 +904,25 @@ def _clamp(x, lo, hi): 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)) - 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) + # ragged: drop padded query rows (qtok >= 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 @@ -850,9 +964,15 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: gqa = Hq // Hkv stride_q_tok = Hq * D stride_k_tok = Hkv * D - n_ktiles = Skv // BN + # 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 + 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 @@ -887,8 +1007,9 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: d_base = b.mul(lane_h, b.const_i32(8)) neg_inf = b.const_f32(-1e30) - LDROW = D + PAD - VROW = D + _LDS_PAD_V # V-row pitch (padded for the transposed PV read) + # 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") @@ -904,6 +1025,8 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: 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 ----- @@ -971,7 +1094,12 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: 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) - raw = b.global_load_vN(q, addr, dtype, 8, align=16) + 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 @@ -980,54 +1108,63 @@ def _build_attention_dense_persistent(spec: AttentionDenseSpec) -> KernelDef: ] 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): - buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(K_BYTES_PER_BUF)) - 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(K_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, + _async_load( + k_rsrc, lds_base, buf_val, tile_key0, K_BYTES_PER_BUF, K_LDROW_BYTES ) - for r in range(ROWS_PER_WAVE): - row_base = b.smem_ptr_add(lds_base, row_lds_off) - b.async_buffer_load_lds_addr( - k_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(K_LDROW_BYTES)) - voff = b.add(voff, b.const_i32(stride_k_tok)) def async_load_v(lds_base, buf_val, tile_key0): - """Row-by-row async DMA into [BN, VROW] V_lds. Row-strided (like K) - so V can carry the +PAD_V bank-conflict pad for the transposed PV - read; each row's 64 lanes load 2 bf16 each (= D) contiguously.""" - buf_off = b.mul(b.zext(buf_val, I64), b.const_i64(V_BYTES_PER_BUF)) - 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(V_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, + _async_load( + v_rsrc, lds_base, buf_val, tile_key0, V_BYTES_PER_BUF, V_LDROW_BYTES ) - for r in range(ROWS_PER_WAVE): - row_base = b.smem_ptr_add(lds_base, row_lds_off) - b.async_buffer_load_lds_addr( - v_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(V_LDROW_BYTES)) - voff = b.add(voff, b.const_i32(stride_k_tok)) def load_tile(buf_val, tile_idx): tk0 = b.mul(tile_idx, b.const_i32(BN)) @@ -1065,6 +1202,20 @@ def do_mask(s_reg, tile_idx, lower=False, upper=True): 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): @@ -1187,7 +1338,9 @@ def pv_fused_exp(o_acc_in, p_packs, vbuf, s_reg, m_new): _rs_ctr = [0] - def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): + 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]) @@ -1204,6 +1357,8 @@ def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): 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 @@ -1268,6 +1423,8 @@ def emit_loop_body(j, carry, mask_lower=False, mask_upper=False): 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) @@ -1335,7 +1492,7 @@ def _clamp(x, lo, hi): 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) + emit_loop_body(j, carry, mask_kbound=RAG_KBOUND) res = loop.results l_i = res[1] @@ -1367,17 +1524,22 @@ def _clamp(x, lo, hi): 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)) - 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) + # 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 @@ -1394,7 +1556,8 @@ def attention_dense_grid(spec: AttentionDenseSpec) -> Tuple[int, int, int]: default = one CTA per (query-block, query-head, batch).""" if spec.persistent: return (spec.num_persistent, 1, 1) - return (spec.seqlen_q // _BLOCK_M, spec.num_query_heads, spec.batch) + 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]: @@ -1420,6 +1583,11 @@ def attention_dense_signature(spec: AttentionDenseSpec): _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, @@ -1435,7 +1603,15 @@ def run_attention_dense_torch( 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.""" + 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. + """ ok, why = supports_attention_dense(spec, arch=arch) if not ok: raise NotImplementedError(f"attention_dense unsupported for spec: {why}") 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 index 7b5088154acb..1283453442bd 100644 --- 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 @@ -11,8 +11,8 @@ "sha256": "1dc714eeaa2e69c3198eeadc725956f2c5257db2f6263f524dc3930ed707651d" }, "attention_dense/fp16_h64_sq512": { - "bytes": 230977, - "sha256": "e30bb63f718f6437eaac14bc49b347174b398c2fcc225dd1ca9d46959ddf2fa1" + "bytes": 210000, + "sha256": "4c7f461c2e10edd51ce0386ba4cf2a8ac4a798656c7ac93b730dc2effb6eb85e" }, "attention_dense/lazy_off_sq512": { "bytes": 337008, @@ -22,6 +22,10 @@ "bytes": 217781, "sha256": "db33f2ff07c424d2b3e759aed66e4b0d29056e7e4780ce11fdcc8d199da99334" }, + "attention_dense/persist_h64_sq512": { + "bytes": 214702, + "sha256": "b4daea7ca97abd002b6a50d217a43b0a443263d8cb11c5192193b329caf7e51c" + }, "attention_dense/persist_hkvmaj_sq512": { "bytes": 327546, "sha256": "f49a0e2f5415bce83f55d251c2a5a7be00fcbd999db3b038023f50ab53113f74" @@ -38,6 +42,10 @@ "bytes": 326109, "sha256": "b3dd90bf381adf2dffac99288b78de56ae3137dddf8f6349418338ae4df6baed" }, + "attention_dense/persist_ragged_sq500": { + "bytes": 328078, + "sha256": "0291e1ac1a08bb6bb05274285754b15137b9ad855cf02e43c856a040e111a8aa" + }, "attention_dense/persist_swa_w128_sq512": { "bytes": 444213, "sha256": "3a2039ebb82f66fbf375f8ff9b636ce5cda48702a914541b854d5997ce2de222" @@ -46,6 +54,18 @@ "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" @@ -67,8 +87,8 @@ "sha256": "0698cf388935622a76524c20f300b383a946ecbdb512e93c662db29ad574a1f0" }, "attention_dense/fp16_h64_sq512": { - "bytes": 231077, - "sha256": "aeced7d0a10933d2eb2143fd127d7fdbc08ae86ac8c7f8e40fe30dca6dbcf2f3" + "bytes": 210100, + "sha256": "b920ada5e3c48c8083a8c8d8c904fb685ac49eac66a7d4fbb67329a0cbb4f014" }, "attention_dense/lazy_off_sq512": { "bytes": 337112, @@ -78,6 +98,10 @@ "bytes": 217881, "sha256": "83333cfcfd4d2e3c26ba2cf6e61faffc8dce46ad43cc1216d7efd68ad0653e54" }, + "attention_dense/persist_h64_sq512": { + "bytes": 214802, + "sha256": "7d4c6fce84a638779797c53d6b57f60e2619076aaeb0ec9e557180be8bc902a0" + }, "attention_dense/persist_hkvmaj_sq512": { "bytes": 327650, "sha256": "8a0e7f52f54566f579cf9f0f20d074b0a57f5079c242f6a6b98d35ab9cd000a4" @@ -94,6 +118,10 @@ "bytes": 326213, "sha256": "0574a132367b90ef4a980afa9d482746c973556545032a61f1e9b66c891aa7e1" }, + "attention_dense/persist_ragged_sq500": { + "bytes": 328227, + "sha256": "5ce0e47b2f540e77344c8439293ac8b2041e242001ef5017d71cdb423bf22230" + }, "attention_dense/persist_swa_w128_sq512": { "bytes": 444317, "sha256": "6fe644368ff3c588edfcf80eb74e6ae9264458a9032b40ca78fdfb93e6f6c403" @@ -102,6 +130,18 @@ "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" 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 index 971f2f1bd342..c30caa3da832 100644 --- 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 @@ -82,6 +82,31 @@ def mk(**over): "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 Date: Tue, 21 Jul 2026 01:52:47 +0000 Subject: [PATCH 11/11] address review comments --- .../builders/gfx950/attention/ALGORITHM.md | 21 +++- .../gfx950/attention/prefill/README.md | 70 +++++++++---- .../rocke/library/dispatch/attention.py | 99 +++++++++++++++---- .../library/kernels/gfx950/attention_dense.py | 55 ++++++++--- 4 files changed, 192 insertions(+), 53 deletions(-) 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 d0baed521bf5..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 @@ -395,10 +395,23 @@ $(\text{qb}, h_q, \text{bt})$ decides load balance *and* L2 locality: ($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):** default grid ≈543 TFLOPS; **persistent ≈877 TFLOPS**, and -**≈948 TFLOPS with the `auto`→`hkv_major` decode + V-pad + lazy rescale** -(byte-identical IR ⇒ this number is schedule, not luck). The persistent variant -is the production choice for dense prefill. +≈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 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 index 77ea25240402..7ec62767b318 100644 --- 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 @@ -31,21 +31,33 @@ performance knobs. `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 (qb-major decode for causal -load-balance), 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). - -## Measured (MI355X, bf16, D=128, Hq=128, Hkv=8, causal) - -| Sq | default | persistent (NP=256) | -|----|--------:|--------------------:| -| 2048 | ~410 TFLOPS | ~565 TFLOPS | -| 8192 | ~521 TFLOPS | **~850 TFLOPS** | - -Both paths are 0 VGPR spill and bit-identical vs `torch.nn.functional. -scaled_dot_product_attention` (max abs err ~1.46e-3 at Sq=8192). +`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 @@ -71,6 +83,28 @@ spec = AttentionDenseSpec( 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`, @@ -93,5 +127,7 @@ kernel = build_attention_dense(spec) # -> KernelDef; compile with backend= - 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 currently lowered on the LLVM-direct backend only (not yet mirrored in the - C++ engine); byte-identity coverage for it is a follow-up. + 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/dispatch/attention.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py index 0df74f2841d2..de101419c534 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,8 +159,14 @@ 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 = "" 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( @@ -301,33 +316,79 @@ 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). - Launch geometry + ABI are real (``attention_dense_grid`` / ``_block`` / - ``_signature``); end-to-end launch is ``run_attention_dense_torch``. + 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 _dense_spec(req: OperatorRequest): - from kernels.gfx950.attention_dense import AttentionDenseSpec - - assert isinstance(req, AttentionRequest) - return AttentionDenseSpec( - batch=int(req.batch), - seqlen_q=int(req.seqlen_q), - seqlen_kv=int(req.seqlen_k), - 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(), - ) - def support(req: OperatorRequest) -> Tuple[bool, str]: errors = _request_errors(req) if errors: @@ -361,6 +422,7 @@ def select(req: OperatorRequest) -> AttentionSpec: 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, @@ -369,6 +431,9 @@ def select(req: OperatorRequest) -> AttentionSpec: 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( 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 index 1b030fbf984d..6a71210dfed5 100644 --- 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 @@ -23,17 +23,22 @@ it wins issue slots; paired with the prefetch this is a measured ~+3.5%. * **vectorized O store**. -Measured on MI355X (bf16, D=128, causal, 0 spill, err ~1.46e-3 vs SDPA): -**~543 TFLOPS @ Sq=8192** for the default one-CTA-per-query-block grid, and -**~877 TFLOPS @ Sq=8192** for the persistent grid-stride variant -(``persistent=True``). 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% (bf16, Sq=8192, GQA -> ~970 TFLOPS with the -hkv-major decode + V-pad). +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). @@ -1566,18 +1571,22 @@ def attention_dense_block(spec: AttentionDenseSpec) -> Tuple[int, int, int]: def attention_dense_signature(spec: AttentionDenseSpec): - """ABI signature (q/k/v/o pointers + f32 scale) for :class:`KernelLauncher`.""" + """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 - return ( + 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") - .build() ) + if spec.varlen: + sig = sig.ptr("cu_seqlens_q", "i32").ptr("cu_seqlens_kv", "i32") + return sig.build() _DENSE_LAUNCHER_CACHE: dict = {} @@ -1598,6 +1607,8 @@ def run_attention_dense_torch( 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 @@ -1611,10 +1622,21 @@ def run_attention_dense_torch( 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 @@ -1634,6 +1656,9 @@ def run_attention_dense_torch( ) _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(