From 831201a286b604ff78873ca005cc8059eb9cdcbb Mon Sep 17 00:00:00 2001 From: Parthiv Aneesh Date: Thu, 29 Oct 2026 18:58:48 -0700 Subject: [PATCH 1/6] perf(rocke): add opt-in gfx942 BF16 CFVST and KDB/VPF Add three opt-in knobs to the gfx942 dense attention kernel for BF16 D=128: use_cfvst (CFVST V transpose), lds_num_buffers=2 (K double buffer) and v_reg_prefetch (V register prefetch). Defaults are unchanged; use_cfvst=None resolves to off. KDB/VPF require persistent=True and block_n=64. - dispatch: bf16_cfvst_supported() is the single cohort predicate; lds_num_buffers and v_reg_prefetch become tunable axes - tests: LDS budget, tuning space, numerics and re-blessed goldens - bench: 7-way ablation harness with replay identity, clean-tree provenance and static resource capture - docs: case study, README opt-in table, dispatch AGENTS.md migration note, runbook entry --- .../prefill/bench_rocke_7way_specs.py | 1828 +++++++++++++++++ .../gfx942/attention/prefill/README.md | 15 +- ...42_dense_cfvst_kdb_vprefetch_case_study.md | 553 +++++ .../rocke/library/dispatch/AGENTS.md | 53 +- .../rocke/library/dispatch/attention/axes.py | 9 +- .../library/dispatch/attention/dense_rules.py | 21 +- .../library/kernels/gfx942/attention_dense.py | 567 +++-- .../dispatch/attention/test_tuning_space.py | 279 ++- .../attention_dense_gfx942_ir_sha256.json | 180 ++ .../tests/test_attention_dense_gfx942.py | 598 +++++- .../test_attention_dense_gfx942_golden.py | 239 ++- .../tests/test_attention_dense_gfx942_lds.py | 69 +- .../test_attention_dense_gfx942_numeric.py | 475 ++++- .../optimization/optimization_runbook.md | 8 + 14 files changed, 4721 insertions(+), 173 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py new file mode 100644 index 000000000000..07439024ca50 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py @@ -0,0 +1,1828 @@ +#!/usr/bin/env python3 +"""Eight-arm persistent gfx942 BF16 D128 attention: Step-0 winner A vs seven-stage B ablation. + +A uses the seven validated shape-specific Step-0 winners from the +2026-10-28 staged sweep (no CFVST/KDB/VPF). B[1..7] always start from the CURRENT +per-shape production-dispatch baseline, independently of A. Thus this is a +fair new same-source A/B measurement, NOT a claim about cross-commit timings. + + A sweep_winner per-shape best preexisting levers (new pipeline off) + 1 baseline production dispatch (CFVST/KDB/VPF off) + 2 cfvst baseline + CFVST + 3 cfvst_kdb CFVST + K LDS double buffering + 4 cfvst_vpf CFVST + V register prefetch (KDB off) + 5 pipeline CFVST + KDB + V register prefetch + 6 pipeline_pad full pipeline + lds_row_pad=4 + 7 pipeline_pad_dsplit previous + causal_diag_split=True + +Stages 3 and 4 are separate branches from CFVST alone. Stage 5 combines +both branches. A is the ONLY arm receiving lever-sweep geometry, scheduling, +and grid knobs. The seven B configurations do not inherit A's knobs. +The exact single-knob B deltas are checked against production dispatch. +The rebuilt A must compile to the recorded Step-0 winner kernel name, or the +run aborts; resolved specs, HSACO hashes and artifacts are kept. + +Methodology: torch-free NumPy/rocKE, one process and explicit HIP stream, +shared Q/K/V and chunked NumPy FP32 causal GQA reference over BF16 inputs, +BF16 correctness before timing and after every timed batch, +round-robin warmup, five randomized rounds, 50 launches per arm per round, +median latency and same-round paired min/max speedups. +Every arm is selected through tuning_spec_with_knobs and replayed from its +recorded (spec_id, tuning_id, knobs) against its dispatch-time spec_hash before +timing. A dirty or untracked checkout is refused unless --allow-dirty is +passed. provenance.json records the source SHA-256 of HASHED_SOURCES; +--verify-provenance OUT_DIR re-checks them and the replay at a later checkout, +and refuses any non-Markdown change since the measured commit. +static_resources.csv records registers, spills, scratch, LDS and static ISA +counts from the exact HSACOs that are timed. +Absolute latency results are kept in local CSVs. +Default cohort: twelve persistent BF16 D128 causal shapes. Seven have recorded +Step-0 winners and run all eight arms. B1 S1024/2048 (Hkv 8 and 16) and B16 +S8192 Hkv8 have no Step-0 winner, so they run the seven production-derived B +arms without A: their B ablations against production dispatch are reported, +and no A/B ratio is fabricated for them. +Custom --shape values must be among these twelve shapes and select persistent +dispatch. +""" + +from __future__ import annotations + +import argparse +import ctypes +import csv +from contextlib import contextmanager +from dataclasses import asdict, replace +from datetime import datetime, timezone +from importlib import metadata as importlib_metadata +import hashlib +import json +import math +import numpy as np +from pathlib import Path +import random +import re +import shlex +import shutil +import statistics +import subprocess +import sys +from dataclasses import dataclass + +B_VARIANTS = ( + "baseline", + "cfvst", + "cfvst_kdb", + "cfvst_vpf", + "pipeline", + "pipeline_pad", + "pipeline_pad_dsplit", +) +VARIANTS = ("A", *B_VARIANTS) +STAGE_NUMBERS = {"A": 0, **{name: i for i, name in enumerate(B_VARIANTS, start=1)}} + +STAGE_DESCRIPTIONS = { + "A": "Per-shape best existing-lever sweep winner (CFVST/KDB/VPF off)", + "baseline": "Production-dispatch baseline, without CFVST/KDB/VPF", + "cfvst": "Baseline + conflict-free V store transpose", + "cfvst_kdb": "CFVST + K LDS double buffering (no V prefetch)", + "cfvst_vpf": "CFVST + V register prefetch (no K double buffering)", + "pipeline": "CFVST + K LDS double buffering + V register prefetch", + "pipeline_pad": "Full pipeline + lds_row_pad=4", + "pipeline_pad_dsplit": "Full pipeline + lds_row_pad=4 + causal_diag_split=True", +} + +# Dependency graph (not a linear chain): both KDB and VPF extend CFVST. +# The full pipeline is cross-checked against both upstream branches. +PARENT_VARIANT = { + "A": None, + "baseline": None, + "cfvst": "baseline", + "cfvst_kdb": "cfvst", + "cfvst_vpf": "cfvst", + "pipeline": "cfvst_kdb", + "pipeline_pad": "pipeline", + "pipeline_pad_dsplit": "pipeline_pad", +} + +STAGE_OVERRIDES = { + "A": {}, # Separate sweep-winner specification; not a B stage. + "baseline": {}, + "cfvst": {"use_cfvst": True}, + "cfvst_kdb": {"use_cfvst": True, "lds_num_buffers": 2}, + "cfvst_vpf": {"use_cfvst": True, "v_reg_prefetch": True}, + "pipeline": {"use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True}, + "pipeline_pad": { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + "lds_row_pad": 4, + }, + "pipeline_pad_dsplit": { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + "lds_row_pad": 4, + "causal_diag_split": True, + }, +} + + +# Exact final A/B-tested Step-0 winners from the 2026-10-28 staged sweep, +# NOT the older 12-row sweep. The sweep's source HEAD is recorded. +SWEEP_SOURCE_COMMIT = "ceeae550e7aa95b291aff5cb75ec4d626a072bf3" + +# Every arm is selected through dispatch under this candidate. +SPEC_ID = "gfx942_dense" + +# Files whose bytes determine what is benchmarked, relative to the rocke root. +# provenance.json records their SHA-256 so --verify-provenance can prove that a +# later checkout (e.g. a docs-only commit on top) still measures the same code. +HASHED_SOURCES = ( + "library/kernels/gfx942/attention_dense.py", + "library/kernels/common/attention_dense_spec.py", + "library/dispatch/attention/axes.py", + "library/dispatch/attention/dense_rules.py", + "library/dispatch/attention/gfx942_dense.py", + "library/builders/gfx942/attention/prefill/attention_dense_prefill.py", + "library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py", +) + +# Shapes without a Step-0 winner in the newer sweep. B=16, S=8192, Hkv=8 has +# no completed final A/B validation; B=1, S=1024/2048 are not in that sweep's +# cohort. Run the seven production-derived B arms only; do not invent a sweep +# winner or an A arm for these shapes. +MISSING_SWEEP_WINNERS = { + (1, 1024, 8): "not in newer sweep cohort", + (1, 1024, 16): "not in newer sweep cohort", + (1, 2048, 8): "not in newer sweep cohort", + (1, 2048, 16): "not in newer sweep cohort", + (16, 8192, 8): "missing_files (no completed final A/B validation in newer sweep)", +} + +# Shared fields match the newer report's recorded full winner specifications. +# None and auto values are intentional: they preserve the original policy +# inputs rather than inventing resolved values. The new pipeline is held OFF. +A_COMMON = dict( + block_m=256, + block_n=64, + block_size=0, + causal=True, + causal_bottom_right=False, + chiplet_num_xcds=8, + dtype="bf16", + head_size=128, + persistent=True, + num_persistent=304, + persist_decode="auto", + nonpersist_decode="auto", + num_kv_blocks=0, + num_query_heads=32, + paged=False, + ragged=False, + sliding_window=0, + use_sinks=False, + varlen=False, + lds_row_pad=4, + lds_k_group_pad=8, + lds_num_buffers=1, + iglp=True, + iglp_mode=1, + pv_loop_order="k_major", + causal_diag_split=True, + o_store_width=4, + lazy_rescale=True, + interleave=False, + use_cfvst=None, + use_exp2_fast=None, + use_v_swizzle=None, + v_row_pad=None, + pv_sched_fence_mask=None, + emit_lse=False, + v_reg_prefetch=False, # new field: absent from old report; must be OFF. +) + +# Only shape-specific winning knobs differ across the seven completed rows. +# All shapes: Hq=32, D=128, BF16, causal, 304 persistent workgroups. +WINNING_LEVERS = { + (1, 4096, 8): dict(waves_per_eu=3, pv_priority=1), + (1, 4096, 16): dict(waves_per_eu=2, pv_priority=1), + (1, 8192, 8): dict(waves_per_eu=3, pv_priority=1), + (1, 8192, 16): dict(waves_per_eu=3, pv_priority=2), + (1, 16384, 8): dict(waves_per_eu=3, pv_priority=3), + (16, 4096, 8): dict(waves_per_eu=2, pv_priority=3), + (16, 4096, 16): dict(waves_per_eu=2, pv_priority=3), +} + +# Expected ORIGINAL kernel symbols are used for auditing only. The current +# emitter may produce a different symbol if policy or implementation changed. +_A_PREFIX = "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_" +_A_SUFFIX = "_krowpad4_iglp1_iglpm1_pvkmaj_prio" +EXPECTED_A_KERNEL = { + (1, 4096, 8): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "1_dsplit", + (1, 4096, 16): _A_PREFIX + "bthkvmin_gfx942_wpe2" + _A_SUFFIX + "1_dsplit", + (1, 8192, 8): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "1_dsplit", + (1, 8192, 16): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "2_dsplit", + (1, 16384, 8): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "3_dsplit", + (16, 4096, 8): _A_PREFIX + "hkvmaj_gfx942_wpe2" + _A_SUFFIX + "3_dsplit", + (16, 4096, 16): _A_PREFIX + "gfx942_wpe2" + _A_SUFFIX + "3_dsplit", +} + +assert set(EXPECTED_A_KERNEL) == set(WINNING_LEVERS) +assert not set(MISSING_SWEEP_WINNERS) & set(WINNING_LEVERS) + + +# speedup = median(numerator)/median(denominator) +# Paired min/max = extrema of same-round numerator/denominator ratios. +# Each effect is measured relative to its correct parent, not list position. +COMPARISON_PAIRS = ( + ("baseline", "cfvst", "cfvst_only"), + ("cfvst", "cfvst_kdb", "kdb_given_cfvst"), + ("cfvst", "cfvst_vpf", "vpf_given_cfvst"), + ("cfvst_kdb", "pipeline", "vpf_given_cfvst_kdb"), + ("cfvst_vpf", "pipeline", "kdb_given_cfvst_vpf"), + ("baseline", "cfvst_kdb", "cfvst_plus_kdb_total"), + ("baseline", "cfvst_vpf", "cfvst_plus_vpf_total"), + ("baseline", "pipeline", "pipeline_total"), + ("pipeline", "pipeline_pad", "lds_pad"), + ("pipeline_pad", "pipeline_pad_dsplit", "causal_split"), + ("baseline", "pipeline_pad_dsplit", "total"), + # A is the best earlier Step-0 lever winner; each B stage uses a separate + # production-derived spec. A/B results are remeasured on the SAME checkout. + *(("A", name, f"A_vs_{name}") for name in B_VARIANTS), +) +EFFECT_NAMES = tuple(effect for _, _, effect in COMPARISON_PAIRS) + + +def changed_fields(before, after) -> set[str]: + """Confirm cumulative stages differ only in their advertised knobs.""" + b, a = asdict(before), asdict(after) + return {key for key in b if b[key] != a[key]} + + +def make_request(shape): + """Exact production request: mirror the gfx942 dense builder's dense_request + (BF16 causal, square self-attention), rather than hand-building a spec.""" + from dispatch.attention import AttentionRequest + + return AttentionRequest( + batch=shape.batch, + nhead_q=shape.hq, + nhead_k=shape.hkv, + seqlen_q=shape.seqlen, + seqlen_k=shape.seqlen, + hdim_q=shape.d, + hdim_v=shape.d, + arch="gfx942", + mask_type=1, # causal; same as dense_request(..., causal=True) + dtype=shape.dtype, + algorithm="attention_dense", + spec_id=SPEC_ID, + ) + + +def tuned_arm(req, production, overrides: dict): + """Select one arm through dispatch, as a pinned request would. + + The knobs are the arm's field delta against the production spec, passed to + ``tuning_spec_with_knobs``; the result must be exactly the spec the overrides + describe, so the ablation measures what dispatch hands out. + """ + from dispatch.attention import tuning_spec_with_knobs + + expected = replace(production, **overrides) + want = asdict(expected) + have = asdict(production) + knobs = {key: want[key] for key in want if want[key] != have[key]} + tuned = tuning_spec_with_knobs(req, SPEC_ID, knobs) + if tuned.kernel_spec != expected: + got = asdict(tuned.kernel_spec) + raise RuntimeError( + f"dispatch built a different spec than {sorted(knobs)} describe: " + f"{sorted(key for key in got if got[key] != want[key])}" + ) + return tuned + + +def replay_record(req, tuned) -> dict: + """What a tuning store keeps to rebuild an arm: the candidate, ``tuning_id`` + and canonical knobs, checked on replay against the dispatch-time + ``spec_hash`` (dispatch ARCHITECTURE.md section 11.1).""" + from rocke.dispatch.core import spec_identity + + return { + "candidate": tuned.candidate_name, + "algorithm": req.algorithm, + "spec_id": SPEC_ID, + "tuning_id": tuned.tuning_id, + "config_key": tuned.config_key, + "knobs": dict(tuned.knobs), + "spec_hash": spec_identity(tuned), + } + + +def replay_spec_hash(req, record: dict) -> str: + """Dispatch ``req`` pinned to a stored record, round-tripped through JSON as + a tuning store would keep it, and return the result's ``spec_hash``.""" + from dispatch.attention import dispatch_attention + + stored = json.loads(json.dumps(record)) + pinned = replace( + req, + algorithm=stored["algorithm"], + spec_id=stored["spec_id"], + tuning_id=stored["tuning_id"], + tuning_knobs=stored["knobs"], + ) + return dispatch_attention(pinned).kernel_id.spec_hash + + +@dataclass(frozen=True) +class Shape: + batch: int + seqlen: int + hq: int + hkv: int + d: int + dtype: str + + @property + def label(self) -> str: + return ( + f"b{self.batch}_s{self.seqlen}_hq{self.hq}_hkv{self.hkv}" + f"_d{self.d}_{self.dtype}" + ) + + +DEFAULT_SHAPES = ( + Shape(1, 1024, 32, 8, 128, "bf16"), + Shape(1, 1024, 32, 16, 128, "bf16"), + Shape(1, 2048, 32, 8, 128, "bf16"), + Shape(1, 2048, 32, 16, 128, "bf16"), + Shape(1, 4096, 32, 8, 128, "bf16"), + Shape(1, 4096, 32, 16, 128, "bf16"), + Shape(1, 8192, 32, 8, 128, "bf16"), + Shape(1, 8192, 32, 16, 128, "bf16"), + Shape(1, 16384, 32, 8, 128, "bf16"), + Shape(16, 4096, 32, 8, 128, "bf16"), + Shape(16, 4096, 32, 16, 128, "bf16"), + Shape(16, 8192, 32, 8, 128, "bf16"), +) + + +def parse_shape(value: str) -> Shape: + try: + parts = [s.strip() for s in value.split(",")] + if len(parts) != 6: + raise ValueError("expected exactly six comma-separated entries") + b, s, hq, hkv, d = (int(x) for x in parts[:5]) + dtype = parts[5].lower() + except ValueError as exc: + raise argparse.ArgumentTypeError( + "--shape must be B,S,Hq,Hkv,D,dtype, e.g. 1,4096,32,8,128,bf16" + ) from exc + # The pipeline supports BF16 D128. + if dtype != "bf16" or d != 128: + raise argparse.ArgumentTypeError("seven-way ablation requires bf16 D128") + if min(b, s, hq, hkv) <= 0 or hq % hkv: + raise argparse.ArgumentTypeError( + "positive extents and Hq divisible by Hkv required" + ) + if s % 256 or s % 64: + raise argparse.ArgumentTypeError( + "S must be divisible by block_m=256 and block_n=64" + ) + return Shape(b, s, hq, hkv, d, dtype) + + +def locate_rocke(explicit: str | None) -> Path: + starts = [] + if explicit: + starts.append(Path(explicit).expanduser()) + starts.extend((Path.cwd(), Path(__file__).resolve().parent)) + for start in starts: + for path in (start, *start.parents): + for candidate in (path, path / "dnn-providers/hip-kernel-provider/rocke"): + if (candidate / "library/kernels/gfx942/attention_dense.py").is_file(): + return candidate.resolve() + raise SystemExit("ROCKE root not found: pass --rocke-root") + + +def add_module_paths(rocke: Path) -> None: + for p in (rocke / "library", rocke / "platform/python", rocke): + if str(p) not in sys.path: + sys.path.insert(0, str(p)) + + +def fp32_reference(q, k, v, *, scale: float, chunk: int): + """Chunked causal GQA SDPA in NumPy FP32 over the exact BF16 inputs. + + Batches and KV heads are processed separately. Peak score matrix is + (Hq/Hkv, chunk, S) so largest shapes do not allocate an S-by-S tensor. + """ + batch, sq, hq, d = q.shape + _, sk, hkv, _ = k.shape + gqa = hq // hkv + result = np.empty((batch, sq, hq, d), dtype=np.float32) + key_indices = np.arange(sk)[None, :] + for b in range(batch): + for h in range(hkv): + first, last = h * gqa, (h + 1) * gqa + kt = k[b, :, h, :].T + val = v[b, :, h, :] + for qs in range(0, sq, chunk): + qe = min(qs + chunk, sq) + query = q[b, qs:qe, first:last, :].transpose(1, 0, 2) + scores = np.matmul(query, kt) + scores *= np.float32(scale) + # Bottom-left aligned causal mask: key_index <= query_index. + masked = key_indices > np.arange(qs, qe)[:, None] + np.copyto(scores, -np.inf, where=masked[None, :, :]) + scores -= np.max(scores, axis=-1, keepdims=True) + np.exp(scores, out=scores) + scores /= np.sum(scores, axis=-1, keepdims=True) + result[b, qs:qe, first:last, :] = np.matmul(scores, val).transpose( + 1, 0, 2 + ) + return result + + +def validate(actual_bits, reference, *, atol: float, rtol: float, chunk: int = 256): + """Compare every BF16 output element to FP32 reference, chunked on host.""" + from rocke.numeric.bf16 import bf16_to_f32 + + if actual_bits.shape != reference.shape: + raise AssertionError( + f"output shape {actual_bits.shape} != reference {reference.shape}" + ) + if not np.isfinite(reference).all(): + raise AssertionError("FP32 reference contains nonfinite values") + max_abs = 0.0 + max_rel = 0.0 + for b in range(reference.shape[0]): + for qs in range(0, reference.shape[1], chunk): + sl = slice(qs, qs + chunk) + actual32 = bf16_to_f32(actual_bits[b, sl]) + expected = reference[b, sl] + np.testing.assert_allclose( + actual32, expected, atol=atol, rtol=rtol, equal_nan=False + ) + diff = np.abs(actual32 - expected) + max_abs = max(max_abs, float(diff.max())) + max_rel = max( + max_rel, float((diff / np.maximum(np.abs(expected), 1e-6)).max()) + ) + return max_abs, max_rel + + +def _device(array): + """Torch-free ROCm transfer: NumPy -> rocKE DeviceMem.""" + from rocke.runtime.hip_module import Runtime + from rocke.runtime.host_buffers import as_u8_buffer + from rocke.runtime.launcher import DeviceMem + + array = np.ascontiguousarray(array) + mem = DeviceMem(array.nbytes) + Runtime().memcpy_h2d(mem.ptr(), as_u8_buffer(array), array.nbytes) + return mem + + +def _download_bf16(device_buffer, shape): + """Read raw BF16 uint16 bits using rocKE's HIP host-copy utility.""" + from rocke.runtime.hip_module import Runtime + from rocke.runtime.host_buffers import as_u8_buffer + + bits = np.empty(shape, dtype=np.uint16) + Runtime().memcpy_d2h(as_u8_buffer(bits), device_buffer.ptr(), bits.nbytes) + return bits + + +@contextmanager +def _hip_stream(): + """A dedicated HIP stream, avoiding framework-dependent stream lookup.""" + from rocke.runtime import hip_module + from rocke.runtime.hip_module import Runtime + + hip = hip_module._resolve_hip() # Same HIP DSO used by rocKE. + create = hip.hipStreamCreate + create.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + create.restype = ctypes.c_int + destroy = hip.hipStreamDestroy + destroy.argtypes = [ctypes.c_void_p] + destroy.restype = ctypes.c_int + handle = ctypes.c_void_p() + hip_module._check(create(ctypes.byref(handle)), "hipStreamCreate") + if not handle.value: + raise RuntimeError("hipStreamCreate returned null HIP stream") + stream = int(handle.value) + try: + yield stream + finally: + try: + Runtime().wait_stream(stream) + finally: + hip_module._check(destroy(handle), "hipStreamDestroy") + + +def timed_batch(launcher, values, config, stream: int, repeat: int) -> float: + """HIP events around exactly N launches; preserves old launch fencing. + + Equivalent to the original torch.cuda.Event harness: no changed + LaunchConfig fence setting, no compilation or correctness checks timed. + """ + from rocke.runtime.hip_module import Runtime + + first = Runtime().event() + last = Runtime().event() + try: + first.record(stream=stream) + for _ in range(repeat): + launcher(values, config=config) + last.record(stream=stream) + last.synchronize() + return float(first.elapsed_to(last)) / repeat + finally: + first.destroy() + last.destroy() + + +def write_csv(path: Path, columns: tuple[str, ...], rows: list[dict]) -> None: + with path.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=columns) + writer.writeheader() + writer.writerows(rows) + + +def _command_output(cmd: list[str], *, cwd: Path | None = None) -> str | None: + try: + return subprocess.check_output( + cmd, cwd=cwd, text=True, stderr=subprocess.STDOUT, timeout=12 + ).strip() + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired): + return None + + +def _package_version(package: str) -> str | None: + try: + return importlib_metadata.version(package) + except importlib_metadata.PackageNotFoundError: + return None + + +def _git_provenance(root: Path) -> dict: + repo = _command_output(["git", "rev-parse", "--show-toplevel"], cwd=root) + head = _command_output(["git", "rev-parse", "HEAD"], cwd=root) + branch = _command_output(["git", "branch", "--show-current"], cwd=root) + status = _command_output( + ["git", "status", "--porcelain=v1", "--untracked-files=all"], cwd=root + ) + return { + "repository_root": repo, + "head_sha": head, + "branch": branch, + "working_tree_dirty": None if status is None else bool(status), + } + + +def _capture_git_checkout(root: Path, *, allow_dirty: bool) -> dict: + """Record Git provenance. A dirty, untracked or unknown checkout is refused + unless ``allow_dirty``: results from it cannot be tied to a commit.""" + state = _git_provenance(root) + clean = state["head_sha"] is not None and state["working_tree_dirty"] is False + print( + f"Git HEAD={state['head_sha'] or 'unknown'}; " + f"dirty={state['working_tree_dirty']}" + + (" (allowed by --allow-dirty)" if not clean and allow_dirty else ""), + flush=True, + ) + if not clean and not allow_dirty: + raise SystemExit( + "Refusing to benchmark: the checkout is dirty, has untracked files, or " + "is not a Git work tree. Commit first, or pass --allow-dirty for a " + "development run whose results cannot be cited." + ) + return state + + +def _llvm_tool(name: str) -> str: + """An LLVM binutil from PATH, else from ``llvm/bin`` beside the loaded COMGR. + + The fallback disassembles with the toolchain that compiled the kernels. + """ + from rocke.runtime.comgr import loaded_compiler_info + + found = shutil.which(name) + if found: + return found + compiler = loaded_compiler_info() + if compiler is not None and compiler.comgr_path: + tool = Path(compiler.comgr_path).resolve().parent / "llvm" / "bin" / name + if tool.is_file(): + return str(tool) + raise SystemExit(f"{name} not found on PATH or beside the loaded COMGR library") + + +def static_resources(hsaco: bytes, hsaco_path: Path) -> dict: + """Compiler-reported registers, spills, scratch and LDS from the code-object + notes, plus static ISA instruction counts from the disassembly. Taken from + the exact HSACO that is timed, so the tables share the timings' provenance.""" + from rocke.analysis.isa import analyze_hsaco + from rocke.benchmark.perf.occupancy import resources + + notes = resources(hsaco, "gfx942") + if not notes: + raise RuntimeError(f"no code-object notes in {hsaco_path}") + analysis = analyze_hsaco( + hsaco_path, + objdump=_llvm_tool("llvm-objdump"), + readelf=_llvm_tool("llvm-readelf"), + ) + return { + "vgpr": notes["vgpr"], + "agpr": notes["agpr"], + "sgpr": notes["sgpr"], + "vgpr_spill": notes["vgpr_spill"], + "sgpr_spill": notes["sgpr_spill"], + "scratch_bytes": analysis.resources.scratch_bytes, + "lds_bytes": notes["lds_bytes"], + "occupancy_waves_per_simd": notes["occupancy"], + **analysis.isa.as_dict(), + } + + +def _source_hashes(rocke: Path) -> dict[str, str]: + """SHA-256 of every file in HASHED_SOURCES, keyed by rocke-relative path.""" + return {rel: _file_sha256(rocke / rel) for rel in HASHED_SOURCES} + + +def _file_sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def _loaded_hip_runtime_info() -> dict: + """Identify the HIP runtime actually used by rocKE, without loading another. + + rocKE caches the ctypes handle in hip_module._resolve_hip(). Resolve the + version symbol on THAT handle. Linux /proc/self/maps identifies the file + backing the symbol's executable mapping, not a second libamdhip64 found + via LD_LIBRARY_PATH or an arbitrary package version. + """ + from rocke.runtime import hip_module + + hip = hip_module._resolve_hip() + version_fn = hip.hipRuntimeGetVersion + version_fn.argtypes = [ctypes.POINTER(ctypes.c_int)] + version_fn.restype = ctypes.c_int + encoded_version = ctypes.c_int(-1) + status = version_fn(ctypes.byref(encoded_version)) + if status != 0 or encoded_version.value <= 0: + raise RuntimeError( + "hipRuntimeGetVersion on rocKE's loaded HIP library failed: " + f"status={status}, version={encoded_version.value}" + ) + + symbol_ptr = ctypes.cast(version_fn, ctypes.c_void_p).value + if not symbol_ptr: + raise RuntimeError("Cannot resolve hipRuntimeGetVersion function address") + + # Resolve the actual mapped shared object from the symbol's instruction + # address. /proc/self/maps is available on the Linux ROCm benchmark host. + mapped_path = None + try: + with Path("/proc/self/maps").open(encoding="utf-8") as maps: + for line in maps: + fields = line.rstrip("\n").split(maxsplit=5) + if len(fields) < 6: + continue + lo, hi = (int(x, 16) for x in fields[0].split("-", 1)) + if lo <= symbol_ptr < hi: + if fields[5].startswith("/"): + mapped_path = fields[5] + break + except OSError as exc: + raise RuntimeError("Cannot inspect /proc/self/maps for HIP library") from exc + if mapped_path is None: + raise RuntimeError( + "Cannot identify the mapped HIP shared object containing " + "hipRuntimeGetVersion; refusing ambiguous runtime provenance" + ) + + # HIP_VERSION encoding, not CUDA's encoding: + # major*10_000_000 + minor*100_000 + patch. + raw = encoded_version.value + major = raw // 10_000_000 + minor = (raw // 100_000) % 100 + patch = raw % 100_000 + return { + "version": f"{major}.{minor}.{patch}", + "version_encoded": raw, + "loaded_library_path": mapped_path, + "rocke_library_opened_as": str(hip._name), + "query": "hipRuntimeGetVersion from rocke.runtime.hip_module's cached HIP handle", + } + + +# These are ROCm components that can be loaded by rocKE or the +# compiler. The package version is NOT necessarily the version of every library +# loaded into the process, so capture the actual mapped shared objects too. +_ROCM_SHARED_OBJECT = re.compile( + r"^lib(?:amdhip64|hsa-runtime64|amd_comgr|hiprtc|hipblas(?:lt)?|" + r"rocblas|rocsolver|rocsparse|rocfft|hipfft|hiprand|rocrand|" + r"miopen|rccl|roctracer64|rocprofiler(?:-sdk)?(?:64)?|" + r"rocprim|hipcub)(?:[.-]|$)", + flags=re.IGNORECASE, +) + + +def _rocm_sdk_package_versions() -> dict[str, str]: + """Installed ROCm wheel metadata, not a claim about loaded DSOs.""" + distributions = ( + "rocm", + "rocm-sdk", + "rocm-sdk-core", + "rocm-sdk-runtime", + "rocm-sdk-libraries", + "rocm-hip-runtime", + "rocm-comgr", + "_rocm_sdk_core", + ) + return { + name: version + for name in distributions + if (version := _package_version(name)) is not None + } + + +def _loaded_rocm_shared_objects() -> list[dict]: + """Actual ROCm shared objects mapped into THIS running Linux process. + + Read /proc/self/maps to get mapped file identity instead of searching PATH, + LD_LIBRARY_PATH, or an installation directory. The same basename can be + loaded from different directories, so preserve every distinct mapping. + This does not infer an unexposed per-library ROCm release version. + """ + found: dict[tuple[str, str, str], dict] = {} + try: + with Path("/proc/self/maps").open(encoding="utf-8") as maps: + for line in maps: + fields = line.rstrip("\n").split(maxsplit=5) + if len(fields) < 6: + continue + mapped_path = fields[5] + if not mapped_path.startswith("/"): + continue + # A mapped library might have been unlinked after loading. + name = Path(mapped_path.removesuffix(" (deleted)")).name + if not _ROCM_SHARED_OBJECT.match(name): + continue + key = (mapped_path, fields[3], fields[4]) + found[key] = { + "name": name, + "mapped_path": mapped_path, + "device": fields[3], + "inode": fields[4], + } + except OSError as exc: + raise RuntimeError("Cannot inspect /proc/self/maps for ROCm libraries") from exc + return sorted(found.values(), key=lambda r: (r["name"], r["mapped_path"])) + + +def _environment( + device_id: int, rocke: Path, script_file: Path, git_state: dict +) -> dict: + # Hardware identity comes from rocKE's HIP runtime, not PyTorch properties. + from rocke.core.lower_llvm import _resolve_llvm_flavor + from rocke.runtime.comgr import loaded_compiler_info + from rocke.runtime.hip_module import ( + get_device_name, + get_device_num_cus, + get_device_target_id, + ) + + gpu = { + "name": get_device_name(device_id), + "architecture": get_device_target_id(device_id), + "cu_count_reported_by_hip": get_device_num_cus(device_id), + } + # Record the loaded COMGR compiler rather than an environment override or + # an unrelated clang executable found on PATH. + compiler = loaded_compiler_info() + llvm_flavor = _resolve_llvm_flavor() + return { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "gpu": gpu, + "software": { + "python_version": sys.version.split()[0], + "python_executable": sys.executable, + "numpy": np.__version__, + "hip_runtime": _loaded_hip_runtime_info(), + "triton": _package_version("triton"), + "rocm_sdk_package_versions": _rocm_sdk_package_versions(), + }, + "compile_backend": "python", + "llvm_flavor": llvm_flavor, + "loaded_compiler": { + "llvm_version": ( + ".".join(map(str, compiler.llvm_version)) + if compiler is not None and compiler.llvm_version is not None + else None + ), + "detection_source": compiler.source if compiler is not None else None, + "comgr_library": compiler.comgr_path if compiler is not None else None, + "query_library": ( + compiler.query_library_path if compiler is not None else None + ), + }, + "git": git_state, + "source_files": { + "attention_dense.py": str( + rocke / "library/kernels/gfx942/attention_dense.py" + ), + "attention_dense_sha256": _file_sha256( + rocke / "library/kernels/gfx942/attention_dense.py" + ), + "benchmark_script": str(script_file), + "benchmark_script_sha256": _file_sha256(script_file), + "builder_sha256": _file_sha256( + rocke + / "library/builders/gfx942/attention/prefill/attention_dense_prefill.py" + ), + }, + "source_sha256": _source_hashes(rocke), + } + + +def verify_provenance(out_dir: Path, rocke: Path) -> int: + """Check that a previous run still describes the current checkout. + + Passes only when the run was measured on a clean commit, every file in + HASHED_SOURCES has the bytes it had then, and every arm's stored replay + record (spec_id, tuning_id, knobs) still dispatches to its recorded + spec_hash. No GPU is needed. Returns the process exit status. + """ + provenance = json.loads((out_dir / "provenance.json").read_text()) + resolved = json.loads((out_dir / "resolved_specs.json").read_text()) + problems: list[str] = [] + + measured = provenance["git"] + current = _git_provenance(rocke) + print( + f"Measured at {measured['head_sha']} (dirty={measured['working_tree_dirty']}); " + f"verifying at {current['head_sha']} (dirty={current['working_tree_dirty']})" + ) + if measured["working_tree_dirty"] is not False or not measured["head_sha"]: + problems.append("the run was measured on a dirty or unknown checkout") + if provenance.get("run", {}).get("dirty_checkout_allowed"): + problems.append("the run was started with --allow-dirty") + + # Everything the kernels depend on (platform compiler, dispatch, emitters) + # must be byte-identical; only Markdown may change after the measured commit. + changed = None + if measured["head_sha"]: + changed = _command_output( + ["git", "diff", "--name-only", measured["head_sha"]], cwd=rocke + ) + if changed is None: + problems.append("cannot diff the checkout against the measured commit") + else: + for rel in changed.splitlines(): + if not rel.endswith(".md"): + problems.append(f"{rel}: changed since the measured commit") + print( + f" {len(changed.splitlines())} file(s) changed since the measured commit" + ) + + recorded = provenance.get("source_sha256", {}) + for rel, digest in _source_hashes(rocke).items(): + if rel not in recorded: + problems.append(f"{rel}: not recorded by the run") + elif recorded[rel] != digest: + problems.append(f"{rel}: SHA-256 changed since the run") + else: + print(f" same bytes {rel}") + + for row in resolved: + req = make_request(Shape(**row["shape_fields"])) + for name, record in row["replay"].items(): + replayed = replay_spec_hash(req, record) + if replayed != record["spec_hash"]: + problems.append( + f"{row['shape']}/{name}: {record['tuning_id']} now replays to " + f"spec_hash {replayed}, recorded {record['spec_hash']}" + ) + print(f" replayed {row['shape']}: {len(row['replay'])} arm(s)") + + for problem in problems: + print(f"MISMATCH: {problem}") + print("Provenance verified." if not problems else "Provenance NOT verified.") + return 1 if problems else 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--rocke-root", help="Directory containing library/kernels/gfx942" + ) + parser.add_argument( + "--shape", + type=parse_shape, + action="append", + help="B,S,Hq,Hkv,128,bf16 (repeatable; overrides default twelve persistent shapes)", + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--repeat", type=int, default=50) + parser.add_argument("--seed", type=int, default=942) + parser.add_argument("--reference-chunk", type=int, default=128) + parser.add_argument("--rtol", type=float, default=1e-2) + parser.add_argument("--atol", type=float, default=1e-2) + parser.add_argument("--out-dir", default="~/rocke_sevenway_ablation_results") + parser.add_argument( + "--allow-dirty", + action="store_true", + help="Run on a dirty/untracked checkout (development only; recorded in " + "provenance.json, and the results cannot be cited)", + ) + parser.add_argument( + "--verify-provenance", + metavar="OUT_DIR", + help="Check a previous run's provenance.json and resolved_specs.json " + "against the current checkout (source hashes and dispatch replay), then exit", + ) + args = parser.parse_args() + if args.verify_provenance: + rocke = locate_rocke(args.rocke_root) + add_module_paths(rocke) + return verify_provenance(Path(args.verify_provenance).expanduser(), rocke) + if args.rounds < 3 or args.repeat <= 0 or args.warmup < 0: + parser.error("require rounds >= 3, repeat > 0 and warmup >= 0") + if args.reference_chunk <= 0: + parser.error("reference-chunk must be positive") + if args.atol < 0 or args.rtol < 0: + parser.error("atol and rtol must be nonnegative") + + rocke = locate_rocke(args.rocke_root) + measured_git = _capture_git_checkout(rocke, allow_dirty=args.allow_dirty) + add_module_paths(rocke) + + import kernels.gfx942.attention_dense as dense + from rocke.helpers.compile import compile_kernel + from rocke.runtime import KernelLauncher, LaunchConfig + from rocke.runtime.hip_module import ( + current_device, + get_device_arch, + get_device_count, + ) + from dispatch.attention import AttentionRequest, attention_tuning_spec + + if get_device_count() == 0: + raise SystemExit("ROCm GPU is not visible through HIP") + for required in ("Gfx942AttentionDenseSpec", "attention_dense_runtime_args"): + if not hasattr(dense, required): + raise SystemExit(f"Expected updated attention_dense.py; missing {required}") + source_file = Path(dense.__file__).resolve() + expected_file = (rocke / "library/kernels/gfx942/attention_dense.py").resolve() + if source_file != expected_file: + raise SystemExit( + f"Loaded wrong emitter! expected {expected_file}, got {source_file}" + ) + needed_knobs = { + knob for overrides in STAGE_OVERRIDES.values() for knob in overrides + } + needed_knobs.update(A_COMMON) + needed_knobs.update( + {"persistent", "num_persistent", "persist_decode", "nonpersist_decode"} + ) + for overrides in WINNING_LEVERS.values(): + needed_knobs.update(overrides) + for knob in sorted(needed_knobs): + if knob not in dense.Gfx942AttentionDenseSpec.__dataclass_fields__: + raise SystemExit(f"Current Gfx942AttentionDenseSpec is missing knob {knob}") + rng = random.Random(args.seed) + shapes = args.shape if args.shape else list(DEFAULT_SHAPES) + # All twelve are runnable: a shape without a Step-0 winner gets no A arm. + requested_shapes = list(shapes) + no_winner_shapes = [] + shapes = [] + for shape in requested_shapes: + key = (shape.batch, shape.seqlen, shape.hkv) + if shape.hq != 32: + parser.error(f"No recorded sweep cohort for {shape.label}: Hq must be 32") + if key in MISSING_SWEEP_WINNERS: + reason = MISSING_SWEEP_WINNERS[key] + print(f"NO STEP-0 WINNER {shape.label}: {reason}; B arms only", flush=True) + no_winner_shapes.append( + {"shape": shape.label, "reason": reason, "status": "b_arms_only"} + ) + shapes.append(shape) + elif key in WINNING_LEVERS: + shapes.append(shape) + else: + parser.error( + f"Shape is absent from the twelve-workload cohort: {shape.label}" + ) + if not shapes: + parser.error("No runnable shapes selected") + out_dir = Path(args.out_dir).expanduser().resolve() + out_dir.mkdir(parents=True, exist_ok=True) + device_id = current_device() + if get_device_arch(device_id) != "gfx942": + raise SystemExit("This benchmark requires a gfx942 GPU") + environment = _environment(device_id, rocke, Path(__file__).resolve(), measured_git) + hip_runtime = environment["software"]["hip_runtime"] + print( + f"ROCKE_ROOT={rocke} | GPU={environment['gpu']['name']} " + f"({environment['gpu']['cu_count_reported_by_hip']} CUs) | " + f"Git={environment['git']['head_sha']} " + f"(dirty={environment['git']['working_tree_dirty']}) | " + f"LLVM={environment['llvm_flavor']}" + ) + print( + f"HIP runtime={hip_runtime['version']} | DSO={hip_runtime['loaded_library_path']}" + ) + # Stream is released even if compilation/numerical validation fails. + with _hip_stream() as stream_ptr: + return _run_shapes( + args, + rocke, + environment, + shapes, + out_dir, + rng, + stream_ptr, + no_winner_shapes=no_winner_shapes, + requested_shapes=requested_shapes, + ) + + +def _run_shapes( + args, + rocke, + environment, + shapes, + out_dir, + rng, + stream_ptr, + *, + no_winner_shapes, + requested_shapes, +) -> int: + import kernels.gfx942.attention_dense as dense + from dispatch.attention import attention_tuning_spec + from rocke.helpers.compile import compile_kernel + from rocke.numeric.bf16 import bf16_to_f32, f32_to_bf16 + from rocke.runtime import KernelLauncher, LaunchConfig + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import DeviceMem + + rt = Runtime() + print( + f"HIP stream={stream_ptr} | full ablation stages={list(VARIANTS)}; " + "a shape without a Step-0 winner runs the B arms only" + ) + + raw_rows: list[dict] = [] + validation_rows: list[dict] = [] + summary_rows: list[dict] = [] + ratio_rows: list[dict] = [] + comparison_rows: list[dict] = [] + resolved_rows: list[dict] = [] + static_rows: list[dict] = [] + print(f"Running {len(shapes)} shape(s): {[s.label for s in shapes]}") + for shape_index, shape in enumerate(shapes, start=1): + print(f"\n=== Shape {shape_index}/{len(shapes)}: {shape.label} ===", flush=True) + req = make_request(shape) + production = attention_tuning_spec(req, SPEC_ID).kernel_spec + if not isinstance(production, dense.Gfx942AttentionDenseSpec): + raise RuntimeError( + f"{shape.label}: dispatch returned {type(production).__name__}, " + "not Gfx942AttentionDenseSpec" + ) + if not production.persistent: + raise RuntimeError( + f"{shape.label}: production dispatch selected persistent=False; " + "KDB/VPF are gated to persistent=True. This eight-arm benchmark " + "requires a persistent shape." + ) + + # A is benchmarked only when a validated sweep winner exists; the B + # stages derive from production dispatch and run on every shape. + key = (shape.batch, shape.seqlen, shape.hkv) + no_winner = key in MISSING_SWEEP_WINNERS + variant_names = B_VARIANTS if no_winner else VARIANTS + # Every arm, A included, is selected through dispatch from its knob + # delta and replayed from its stored record before anything is timed. + tuned = {} + a_overrides = None + original_a_name = None + actual_a_name = None + if not no_winner: + a_overrides = A_COMMON | WINNING_LEVERS[key] + tuned["A"] = tuned_arm(req, production, a_overrides) + for name in B_VARIANTS: + tuned[name] = tuned_arm(req, production, STAGE_OVERRIDES[name]) + specs = {name: t.kernel_spec for name, t in tuned.items()} + replay = {name: replay_record(req, t) for name, t in tuned.items()} + for name, record in replay.items(): + replayed = replay_spec_hash(req, record) + if replayed != record["spec_hash"]: + raise RuntimeError( + f"{shape.label}/{name}: replaying {record['tuning_id']} gives " + f"spec_hash {replayed}, dispatch-time {record['spec_hash']}" + ) + if not no_winner: + original_a_name = EXPECTED_A_KERNEL[key] + actual_a_name = specs["A"].kernel_name() + if actual_a_name != original_a_name: + raise RuntimeError( + f"{shape.label}: rebuilt A is not the recorded Step-0 winner; " + f"recorded={original_a_name}, rebuilt={actual_a_name}" + ) + if ( + specs["A"].resolved_use_cfvst() + or specs["A"].lds_num_buffers != 1 + or specs["A"].v_reg_prefetch + ): + raise RuntimeError( + f"{shape.label}: A unexpectedly enables CFVST/KDB/VPF" + ) + + if ( + production.resolved_use_cfvst() + or production.lds_num_buffers != 1 + or production.v_reg_prefetch + ): + raise RuntimeError( + f"{shape.label}: production dispatch already enables some " + "CFVST/KDB/VPF pipeline knobs. Cannot label it a clean baseline." + ) + expected_steps = ( + ("baseline", "cfvst", {"use_cfvst"}), + ("cfvst", "cfvst_kdb", {"lds_num_buffers"}), + ("cfvst", "cfvst_vpf", {"v_reg_prefetch"}), + ("cfvst_kdb", "pipeline", {"v_reg_prefetch"}), + ("cfvst_vpf", "pipeline", {"lds_num_buffers"}), + ("pipeline", "pipeline_pad", {"lds_row_pad"}), + ("pipeline_pad", "pipeline_pad_dsplit", {"causal_diag_split"}), + ) + for before, after, expected in expected_steps: + actual = changed_fields(specs[before], specs[after]) + if actual != expected: + raise RuntimeError( + f"{shape.label}: {before} -> {after} changed {sorted(actual)}, " + f"expected {sorted(expected)}. Current dispatch may already " + "enable a lever; refusing a misleading ablation." + ) + # Check resolved policy too: default/None fields might silently resolve + # differently than their literal dataclass values. + for name in B_VARIANTS: + spec = specs[name] + expected = STAGE_OVERRIDES[name] + observed_pipeline = ( + bool(spec.resolved_use_cfvst()), + int(spec.lds_num_buffers), + bool(spec.v_reg_prefetch), + ) + expected_pipeline = ( + bool(expected.get("use_cfvst", False)), + int(expected.get("lds_num_buffers", 1)), + bool(expected.get("v_reg_prefetch", False)), + ) + if observed_pipeline != expected_pipeline: + raise RuntimeError( + f"{shape.label}/{name}: resolved pipeline knobs " + f"{observed_pipeline} differ from expected {expected_pipeline}" + ) + for name, spec in specs.items(): + ok, reason = dense.supports_attention_dense(spec, arch="gfx942") + if not ok: + raise RuntimeError(f"{shape.label}/{name} rejected: {reason}") + + # Trace full specs, including policy resolutions that dataclass fields + # otherwise encode as None. + def describe_spec(spec): + return { + "fields": asdict(spec), + "resolved_use_cfvst": spec.resolved_use_cfvst(), + "resolved_use_exp2_fast": spec.resolved_use_exp2_fast(), + "resolved_persist_decode": spec.resolved_persist_decode, + "resolved_v_row_pad": spec.resolved_v_row_pad(), + "resolved_use_v_swizzle": spec.resolved_use_v_swizzle(), + "kernel_name": spec.kernel_name(), + "grid": dense.attention_dense_grid(spec), + "block": dense.attention_dense_block(spec), + } + + resolved_rows.append( + { + "shape": shape.label, + "shape_fields": asdict(shape), + "request": asdict(req), + "source": "dispatch.attention.tuning_spec_with_knobs(req, 'gfx942_dense', knobs) per arm", + # Per arm: candidate, spec_id, tuning_id, canonical knobs and the + # dispatch-time spec_hash; each was replayed before timing. + "replay": replay, + "production_baseline": True, + "production_dispatch_kernel": production.kernel_name(), + "execution_mode": ( + "b_arms_only" if no_winner else "eight_arm_ablation" + ), + "missing_step0_winner_reason": MISSING_SWEEP_WINNERS.get(key), + "variant_names": list(variant_names), + "A_role": ( + None + if no_winner + else "per-shape Step-0 lever-sweep winner, independently rebuilt" + ), + "A_original_expected_kernel": original_a_name, + "A_current_kernel": actual_a_name, + "A_kernel_name_matches_original": ( + None if no_winner else actual_a_name == original_a_name + ), + "A_overrides": a_overrides, + "B_stage_baseline": "current production dispatch", + "stage_descriptions": { + name: STAGE_DESCRIPTIONS[name] for name in variant_names + }, + "stage_overrides": { + name: STAGE_OVERRIDES[name] for name in variant_names + }, + "parent_variant": { + name: PARENT_VARIANT[name] for name in variant_names + }, + "comparison_pairs": [ + pair + for pair in COMPARISON_PAIRS + if pair[0] in variant_names and pair[1] in variant_names + ], + "arms": {name: describe_spec(spec) for name, spec in specs.items()}, + } + ) + for name, spec in specs.items(): + print(f" {STAGE_NUMBERS[name]} {name:20s}: {spec.kernel_name()}") + print( + " No Step-0 winner: B stages only, from current production dispatch" + if no_winner + else " A source: Step-0 winner; B stages: current production dispatch " + "with pipeline-only overrides" + ) + print( + f" dispatch: persistent={production.persistent}, " + f"NP={production.num_persistent}, " + f"decode={production.resolved_persist_decode}, " + f"exp2_fast={production.resolved_use_exp2_fast()}, " + f"bm/bn={production.block_m}/{production.block_n}, " + f"wpe={production.waves_per_eu}" + ) + names = {name: spec.kernel_name() for name, spec in specs.items()} + if len(set(names.values())) != len(variant_names): + raise RuntimeError(f"Compile-time kernel-name collision: {names}") + + # Match the original two BF16-rounding points: random BF16, multiply + # by 0.2, round back to BF16. The NumPy RNG differs from torch's RNG. + input_rng = np.random.default_rng(0) + q_shape = (shape.batch, shape.seqlen, shape.hq, 128) + kv_shape = (shape.batch, shape.seqlen, shape.hkv, 128) + + def bf16_random_scaled(dims): + original = input_rng.standard_normal(dims, dtype=np.float32) + initial_bf16 = bf16_to_f32(f32_to_bf16(original)) + return f32_to_bf16(initial_bf16 * np.float32(0.2)) + + q_bits = bf16_random_scaled(q_shape) + k_bits = bf16_random_scaled(kv_shape) + v_bits = bf16_random_scaled(kv_shape) + q = _device(q_bits) + k = _device(k_bits) + v_tensor = _device(v_bits) + outputs = {name: DeviceMem(q_bits.nbytes) for name in variant_names} + scale = 1.0 / math.sqrt(128) + launchers = {} + configurations = {} + values = {} + compiled_artifacts = {} + print(f" Compiling {len(variant_names)} variant(s) and saving HSACOs...") + for name in variant_names: + spec = specs[name] + artifact = compile_kernel( + dense.build_attention_dense(spec, arch="gfx942"), + arch="gfx942", + backend="python", + capture_ir_text=False, + ) + hsaco_dir = out_dir / "hsaco" / shape.label + hsaco_dir.mkdir(parents=True, exist_ok=True) + + hsaco_path = hsaco_dir / f"{name}.hsaco" + hsaco_path.write_bytes(artifact.hsaco) + if artifact.kernel_name != names[name]: + raise RuntimeError(f"Artifact/kernel name mismatch: {name}") + compiled_artifacts[name] = { + "kernel_name": artifact.kernel_name, + "hsaco_sha256": hashlib.sha256(artifact.hsaco).hexdigest(), + "lds_bytes": dense._lds_bytes(spec), + } + static_rows.append( + { + "shape": shape.label, + "variant": name, + "kernel_name": artifact.kernel_name, + "hsaco_sha256": compiled_artifacts[name]["hsaco_sha256"], + **static_resources(artifact.hsaco, hsaco_path), + } + ) + launchers[name] = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=dense.attention_dense_signature(spec), + ) + configurations[name] = LaunchConfig( + grid=dense.attention_dense_grid(spec), + block=dense.attention_dense_block(spec), + stream=stream_ptr, + ) + values[name] = { + "q_ptr": q, + "k_ptr": k, + "v_ptr": v_tensor, + "o_ptr": outputs[name], + "scale": scale, + # REQUIRED for updated runtime-shape kernel ABI and fast division. + **dense.attention_dense_runtime_args(spec), + } + + resolved_rows[-1]["compiled_artifacts"] = compiled_artifacts + print("Building one chunked FP32 NumPy SDPA reference...") + reference = fp32_reference( + bf16_to_f32(q_bits), + bf16_to_f32(k_bits), + bf16_to_f32(v_bits), + scale=scale, + chunk=args.reference_chunk, + ) + del q_bits, k_bits, v_bits + + def check_and_record(name: str, stage: str, round_id: int | str = ""): + """Validate against the same FP32 reference and record the exact gate.""" + try: + bits = _download_bf16(outputs[name], reference.shape) + max_abs, max_rel = validate( + bits, reference, atol=args.atol, rtol=args.rtol + ) + except (AssertionError, RuntimeError) as exc: + raise RuntimeError( + f"ABORT: correctness failed: {shape.label}/{name}, " + f"stage={stage}, round={round_id}" + ) from exc + validation_rows.append( + dict( + shape=shape.label, + variant=name, + stage=stage, + round=round_id, + passed=1, + max_abs=f"{max_abs:.9g}", + max_rel=f"{max_rel:.9g}", + atol=args.atol, + rtol=args.rtol, + ) + ) + return max_abs, max_rel + + print(f" Validating all {len(variant_names)} compiled kernels...") + for name in variant_names: + # Validate through exactly the launcher/config used for timed calls. + # NaN poison catches elements not written by a kernel. + rt.memset(outputs[name].ptr(), 0xFF, outputs[name].nbytes()) + rt.sync() # memset runs on default HIP stream + launchers[name](values[name], config=configurations[name]) + rt.wait_stream(stream_ptr) + max_abs, max_rel = check_and_record(name, "pre_timing") + print(f" PASS {name:20s} max_abs={max_abs:.6g} max_rel={max_rel:.6g}") + # Keep the FP32 reference alive for post-batch validation. + rt.wait_stream(stream_ptr) + + print(f"Round-robin warmup: {args.warmup} launches per variant") + for _ in range(args.warmup): + for name in variant_names: + launchers[name](values[name], config=configurations[name]) + rt.wait_stream(stream_ptr) + + samples = {name: [] for name in variant_names} + for round_id in range(1, args.rounds + 1): + shuffled = list(variant_names) + rng.shuffle(shuffled) + print(f"Round {round_id}: {shuffled}") + for position, name in enumerate(shuffled, 1): + ms = timed_batch( + launchers[name], + values[name], + configurations[name], + stream_ptr, + args.repeat, + ) + # Timed events have already synchronized. Only the last output + # of this batch survives; validate it outside the timed window. + max_abs, max_rel = check_and_record(name, "post_batch", round_id) + samples[name].append(ms) + raw_rows.append( + dict( + shape=shape.label, + round=round_id, + order=position, + variant=name, + tuning_id=replay[name]["tuning_id"], + spec_hash=replay[name]["spec_hash"], + use_cfvst=int(specs[name].resolved_use_cfvst()), + lds_num_buffers=specs[name].lds_num_buffers, + v_reg_prefetch=int(specs[name].v_reg_prefetch), + lds_row_pad=specs[name].lds_row_pad, + causal_diag_split=int(specs[name].causal_diag_split), + repeat=args.repeat, + ms=f"{ms:.9f}", + ) + ) + print(f" {name:20s} {ms:.6f} ms | PASS post-batch (abs={max_abs:.4g})") + + del reference + rt.wait_stream(stream_ptr) + median = {name: statistics.median(times) for name, times in samples.items()} + print( + " Median ms: " + ", ".join(f"{name}={median[name]:.6f}" for name in median) + ) + for name in variant_names: + x = median[name] + summary_rows.append( + dict( + shape=shape.label, + variant=name, + kernel_name=names[name], + tuning_id=replay[name]["tuning_id"], + knobs=json.dumps(replay[name]["knobs"], sort_keys=True), + spec_hash=replay[name]["spec_hash"], + median_ms=f"{x:.9f}", + min_ms=f"{min(samples[name]):.9f}", + max_ms=f"{max(samples[name]):.9f}", + speedup_vs_baseline=f"{median['baseline']/x:.9f}", + parent_variant=PARENT_VARIANT[name] or "", + speedup_vs_parent=( + f"{median[PARENT_VARIANT[name]]/x:.9f}" + if PARENT_VARIANT[name] is not None + else "1.000000000" + ), + stage=STAGE_NUMBERS[name], + description=STAGE_DESCRIPTIONS[name], + rounds=args.rounds, + repeat=args.repeat, + ) + ) + + if any(x <= 0 for x in median.values()): + raise RuntimeError("GPU timing returned nonpositive latency") + pair_metrics = {} + for numerator, denominator, effect in COMPARISON_PAIRS: + if numerator not in median or denominator not in median: + continue # baseline-only shape has no comparable arms + before, after = median[numerator], median[denominator] + paired = [x / y for x, y in zip(samples[numerator], samples[denominator])] + speedup = before / after + reduction = 100.0 * (before - after) / before + metrics = dict( + shape=shape.label, + effect=effect, + numerator=numerator, + denominator=denominator, + speedup=f"{speedup:.9f}", + latency_reduction_pct=f"{reduction:.6f}", + paired_min=f"{min(paired):.9f}", + paired_max=f"{max(paired):.9f}", + ) + ratio_rows.append(metrics) + pair_metrics[effect] = metrics + comparison_rows.append( + { + "batch": shape.batch, + "seqlen": shape.seqlen, + "hkv": shape.hkv, + **{f"{name}_kernel": names[name] for name in variant_names}, + **{ + f"{name}_median_ms": f"{median[name]:.9f}" for name in variant_names + }, + **{ + f"{effect}_speedup": pair_metrics[effect]["speedup"] + for effect in pair_metrics + }, + **{ + f"{effect}_reduction_pct": pair_metrics[effect][ + "latency_reduction_pct" + ] + for effect in pair_metrics + }, + } + ) + print( + " Incremental, branch-specific and cumulative results" + + (" (no A: missing Step-0 winner):" if no_winner else ":"), + flush=True, + ) + for effect in pair_metrics: + r = pair_metrics[effect] + print( + f" {effect:28s} {float(r['speedup']):.5f}x " + f"({float(r['latency_reduction_pct']):+.2f}% latency reduction; " + f"paired {float(r['paired_min']):.5f}-{float(r['paired_max']):.5f}x)", + flush=True, + ) + rt.wait_stream(stream_ptr) + del outputs, values, q, k, v_tensor, launchers, configurations + + (out_dir / "missing_sweep_winners.json").write_text( + json.dumps(no_winner_shapes, indent=2) + "\n" + ) + (out_dir / "resolved_specs.json").write_text( + json.dumps(resolved_rows, indent=2, default=str) + "\n" + ) + # Self-contained eight-arm schemas; no external analyzer assumptions. + exports = ( + ( + "ablation_comparison.csv", + "batch seqlen hkv " + + " ".join(f"{name}_kernel" for name in VARIANTS) + + " " + + " ".join(f"{name}_median_ms" for name in VARIANTS) + + " " + + " ".join(f"{effect}_speedup" for effect in EFFECT_NAMES) + + " " + + " ".join(f"{effect}_reduction_pct" for effect in EFFECT_NAMES), + comparison_rows, + ), + ( + "ablation_runs.csv", + "shape round order variant tuning_id spec_hash use_cfvst lds_num_buffers v_reg_prefetch " + "lds_row_pad causal_diag_split repeat ms", + raw_rows, + ), + ( + "ablation_summary.csv", + "shape variant kernel_name tuning_id knobs spec_hash median_ms min_ms max_ms " + "speedup_vs_baseline parent_variant speedup_vs_parent stage description rounds repeat", + summary_rows, + ), + ( + "validation.csv", + "shape variant stage round passed max_abs max_rel atol rtol", + validation_rows, + ), + ( + "ablation_ratios.csv", + "shape effect numerator denominator speedup latency_reduction_pct paired_min paired_max", + ratio_rows, + ), + ) + for filename, schema, rows in exports: + write_csv(out_dir / filename, tuple(schema.split()), rows) + if static_rows: + write_csv(out_dir / "static_resources.csv", tuple(static_rows[0]), static_rows) + # Speedup-only Markdown; raw absolute timings remain in local CSV files. + ratios_by_shape = {} + for r in ratio_rows: + ratios_by_shape.setdefault(r["shape"], {})[r["effect"]] = r + + def fmt(r, effect): + if effect not in r: + return "—" # no A arm for a shape without a Step-0 winner + x = r[effect] + return ( + f"{float(x['speedup']):.5f}x " + f"[{float(x['paired_min']):.5f}, " + f"{float(x['paired_max']):.5f}]" + ) + + lines = [ + "# Step-0 versus B1–B7 ablation", + "", + "## Variant definitions", + "A is the recorded shape-specific existing-lever sweep winner (when available), rebuilt on the current emitter. ", + "B1-B7 are derived separately from the CURRENT production dispatch. ", + "A's optimization knobs are **not** added to any B stage. ", + "Original A symbols and current compiled A symbols are recorded in resolved_specs.json.", + "Shapes without a Step-0 winner run B1-B7 only, and their A columns are —.", + "", + "| # | Variant | Knobs enabled | Parent comparison |", + "|---:|---|---|---|", + ] + for name in VARIANTS: + overrides = STAGE_OVERRIDES[name] + knobs = ( + "Shape-specific existing-lever winner (see resolved_specs.json)" + if name == "A" + else ", ".join(f"`{k}={v}`" for k, v in overrides.items()) + or "Production dispatch" + ) + parent = PARENT_VARIANT[name] or "—" + lines.append(f"| {STAGE_NUMBERS[name]} | `{name}` | {knobs} | `{parent}` |") + lines += [ + "", + "## A versus the production-derived B stages (new same-source A/B)", + "| B | S | Hkv | A / B1 baseline | A / B2 CFVST | A / B3 CFVST+KDB | A / B4 CFVST+VPF | A / B5 full | A / B6 full+pad | A / B7 full+pad+split |", + "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for sh in shapes: + r = ratios_by_shape.get(sh.label, {}) + effects = tuple(f"A_vs_{name}" for name in B_VARIANTS) + lines.append( + f"| {sh.batch} | {sh.seqlen} | {sh.hkv} | " + + " | ".join(fmt(r, effect) for effect in effects) + + " |" + ) + # A shape without a Step-0 winner appears as a row of — in the A/B table. + lines += [ + "", + "## Direct B ablations (production baseline / CFVST anchor)", + "| B | S | Hkv | CFVST / baseline | CFVST+KDB / CFVST | CFVST+VPF / CFVST | Full / CFVST+KDB | Full / CFVST+VPF |", + "|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for sh in shapes: + r = ratios_by_shape.get(sh.label, {}) + effects = ( + "cfvst_only", + "kdb_given_cfvst", + "vpf_given_cfvst", + "vpf_given_cfvst_kdb", + "kdb_given_cfvst_vpf", + ) + lines.append( + f"| {sh.batch} | {sh.seqlen} | {sh.hkv} | " + + " | ".join(fmt(r, effect) for effect in effects) + + " |" + ) + lines += [ + "", + "## Cumulative, padding and causal-split ablation", + "| B | S | Hkv | Full / baseline | + LDS pad (6/5) | + Causal split (7/6) | Total (7/1) |", + "|---:|---:|---:|---:|---:|---:|---:|", + ] + for sh in shapes: + r = ratios_by_shape.get(sh.label, {}) + effects = ("pipeline_total", "lds_pad", "causal_split", "total") + lines.append( + f"| {sh.batch} | {sh.seqlen} | {sh.hkv} | " + + " | ".join(fmt(r, effect) for effect in effects) + + " |" + ) + lines += [ + "", + "Speedup = numerator-arm median latency divided by denominator-arm median " + "latency. Bracketed range = min/max of same-round paired latency ratios. " + "A value >1x means the denominator arm is faster.", + "", + "## Compiled kernels", + "| B | S | Hkv | " + + " | ".join(f"{STAGE_NUMBERS[name]}: {name}" for name in VARIANTS) + + " |", + "|---:|---:|---:|" + "---|" * len(VARIANTS), + ] + for r in comparison_rows: + lines.append( + f"| {r['batch']} | {r['seqlen']} | {r['hkv']} | " + + " | ".join( + f"`{r[f'{name}_kernel']}`" if f"{name}_kernel" in r else "—" + for name in VARIANTS + ) + + " |" + ) + if no_winner_shapes: + lines += [ + "", + "## Shapes without a Step-0 winner (B1-B7 only, no A)", + "| Shape | Reason |", + "|---|---|", + ] + lines += [f"| {row['shape']} | {row['reason']} |" for row in no_winner_shapes] + (out_dir / "ablation_comparison.md").write_text("\n".join(lines) + "\n") + # A-focused compact report retains the expected symbol, current symbol, + # paired A/B min/max ratios and medians without mixing old sweep timings. + a_vs_b_rows = [] + for sh in shapes: + key = (sh.batch, sh.seqlen, sh.hkv) + if key not in EXPECTED_A_KERNEL: + continue # no A arm, so no A/B7 comparison + r = next( + row + for row in comparison_rows + if (row["batch"], row["seqlen"], row["hkv"]) == key + ) + p = ratios_by_shape[sh.label] + a_vs_b_rows.append( + { + "batch": sh.batch, + "seqlen": sh.seqlen, + "hkv": sh.hkv, + "A_original_kernel": EXPECTED_A_KERNEL[key], + "A_compiled_kernel": r["A_kernel"], + "A_kernel_name_matches_original": r["A_kernel"] + == EXPECTED_A_KERNEL[key], + "B7_kernel": r["pipeline_pad_dsplit_kernel"], + "A_median_ms": r["A_median_ms"], + "B7_median_ms": r["pipeline_pad_dsplit_median_ms"], + "speedup_A_over_B7": p["A_vs_pipeline_pad_dsplit"]["speedup"], + "latency_reduction_pct": p["A_vs_pipeline_pad_dsplit"][ + "latency_reduction_pct" + ], + "paired_min": p["A_vs_pipeline_pad_dsplit"]["paired_min"], + "paired_max": p["A_vs_pipeline_pad_dsplit"]["paired_max"], + } + ) + a_vs_b_fields = ( + "batch", + "seqlen", + "hkv", + "A_original_kernel", + "A_compiled_kernel", + "A_kernel_name_matches_original", + "B7_kernel", + "A_median_ms", + "B7_median_ms", + "speedup_A_over_B7", + "latency_reduction_pct", + "paired_min", + "paired_max", + ) + write_csv(out_dir / "A_vs_B7.csv", a_vs_b_fields, a_vs_b_rows) + environment["run"] = { + "shapes": [asdict(sh) for sh in shapes], + "requested_shapes": [asdict(sh) for sh in requested_shapes], + "no_step0_winner_shapes": no_winner_shapes, + "variant_counts_by_shape": { + sh.label: ( + len(B_VARIANTS) + if (sh.batch, sh.seqlen, sh.hkv) in MISSING_SWEEP_WINNERS + else len(VARIANTS) + ) + for sh in shapes + }, + "sweep_winner_source_commit": SWEEP_SOURCE_COMMIT, + "full_ablation_variant_count": len(VARIANTS), + "correctness_atol": args.atol, + "correctness_rtol": args.rtol, + "reference": "chunked NumPy FP32 causal GQA SDPA over BF16-rounded inputs", + "input_dtype": "bf16", + "warmup_per_arm": args.warmup, + "rounds": args.rounds, + "launches_per_round_per_arm": args.repeat, + "randomization_seed": args.seed, + "input_seed": 0, + "input_rng": "numpy.random.default_rng(0)", + "input_encoding": "NumPy FP32, BF16 RNE, multiply 0.2, BF16 RNE", + "hip_stream_handle": stream_ptr, + "timing": ( + f"rocKE Runtime HIP events around {args.repeat} direct KernelLauncher " + "calls, unchanged launcher fencing" + ), + "validation_scope": "pre-timing and final output after every timed batch", + "validation_checks_by_shape": { + sh.label: ( + len(B_VARIANTS) + if (sh.batch, sh.seqlen, sh.hkv) in MISSING_SWEEP_WINNERS + else len(VARIANTS) + ) + * (1 + args.rounds) + for sh in shapes + }, + "all_validation_passed": all(v["passed"] == 1 for v in validation_rows), + "dispatch_baseline_for_all_shapes": True, + "A_is_per_shape_step0_winner_where_available": True, + "A_config_matches_newer_seven_winner_report": True, + "no_A_comparison_for_shapes_without_step0_winner": True, + "A_expected_kernels": {str(k): v for k, v in EXPECTED_A_KERNEL.items()}, + "B_stages_derive_only_from_production": True, + "stage_descriptions": STAGE_DESCRIPTIONS, + "stage_overrides": STAGE_OVERRIDES, + "parent_variant": PARENT_VARIANT, + "comparison_pairs": COMPARISON_PAIRS, + "dirty_checkout_allowed": args.allow_dirty, + "exact_python_command": shlex.join([sys.executable, *sys.argv]), + "working_directory": str(Path.cwd()), + } + # Snapshot after compilation/launch so lazily loaded ROCm DSOs appear. + environment["software"][ + "rocm_libraries_loaded_at_end" + ] = _loaded_rocm_shared_objects() + (out_dir / "provenance.json").write_text( + json.dumps(environment, indent=2, default=str) + "\n" + ) + print(f"\nOutputs in {out_dir}") + print( + "Reproducibility: provenance.json, resolved_specs.json, validation.csv, " + "ablation_ratios.csv, ablation_comparison.csv, ablation_comparison.md, A_vs_B7.csv, " + "static_resources.csv, missing_sweep_winners.json; " + "keep latency CSVs internal" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md index 2797c1afbf3b..287af145748e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md @@ -185,12 +185,23 @@ magnitudes live outside the repo per `AGENTS.md`. | Persistent grid-stride | all | `num_persistent` CTAs grid-stride over decoded work items; `persist_decode` orders them (auto: `hq_minor_swz` for large causal MHA, `bt_hkv_minor` for causal below `chiplet_num_xcds` batches) | **shipped** — default at D128 | | Default-grid block order | causal | `nonpersist_decode`: batch, then kv head, fastest with longest-first query blocks (`bt_hkv_minor`); Swizzled Head-first (`hq_minor_swz`) for large causal MHA | **shipped** — auto; see the [case study](../../../../../platform/dsl_docs/architecture/attention_thread_block_mapping_case_study.md) | +### Opt-in (off by default) + +Measurements, static register and ISA counts, and methodology are in the +[bf16 D128 CFVST / K double buffering / V prefetch case study](gfx942_dense_cfvst_kdb_vprefetch_case_study.md). + +| Lever | Config | Mechanism | Verdict | +|---|---|---|---| +| Conflict-free V at **bf16 D128** (`use_cfvst=True`) | bf16 D128 causal, either grid, with or without `causal_diag_split` | same `perm_b32` store transpose as fp16 D128; fills the `waves_per_eu=2` VGPR budget with little headroom on the `.1k` MFMA schedule | **opt-in** — not a default | +| K-only LDS double buffering (`lds_num_buffers=2`) | bf16 D128 CFVST causal, `persistent=True`, `block_n=64` only | allocates two K LDS slots while keeping one V LDS slot | **opt-in** — independent of V register prefetch. Not validated at `block_n=32` | +| V register prefetch (`v_reg_prefetch=True`) | bf16 D128 CFVST causal, `persistent=True`, `block_n=64` only | loads V for the next tile into registers before the current PV; V LDS stays single-buffered | **opt-in** — independent of K LDS buffering; can also be combined with `lds_num_buffers=2` | + ### Evaluated and rejected | Lever | Why it lost | |---|---| | Conflict-free V at **D64** | D64 is VGPR-bound; the register round-trip costs more than the LDS-instruction saving | -| Conflict-free V at **bf16 D128** | spills over the `waves_per_eu=2` cap on the `.1k` MFMA schedule | +| Conflict-free V at **bf16 D128** as a default | fills the `waves_per_eu=2` VGPR budget on the `.1k` MFMA schedule; kept off by default and available as an opt-in (see the Opt-in table) | | `exp2_fast` at **bf16 D128** (former holdout) | initially rejected on a measured spill over the `waves_per_eu=2` cap; **revisited and adopted** (see Adopted table) — re-measurement reads 0 scratch and lower VGPR on both grids, so the original reading no longer reproduces | | `waves_per_eu=3` at **fp16 D64** | reaches 2 WG/CU but loses more ILP than the second workgroup buys | | Drop K/V LDS pads to reach 2 WG/CU at D128 | D128 stays 1 WG/CU even unpadded (register floor co-limits), so it only reintroduces bank conflicts — **catastrophic** | @@ -200,7 +211,7 @@ magnitudes live outside the repo per `AGENTS.md`. | cfvst load/store chunking | does not bound register pressure: the later chunks' loads carry no dependency on the earlier ones, so LLVM hoists them above the intervening full `s_waitcnt(vmcnt=0)`, which then covers them anyway | | PV-only `s_setprio` | **proven-negative** on gfx942 `attention_tiled_2d`. Note this is one lever in one placement — it is *not* a verdict on the scheduling-intrinsic family, and hand-written `sched_group_barrier` remains open | | Diagonal two-phase causal peel | **proven-negative** on gfx942 `attention_tiled_2d`; gated on the bound shifting to compute, which it has not | -| partial-vmcnt software prefetch | N/A — it is a double-buffering lever and this kernel is single-buffered (NBUF=1), so there is no prefetch to partially overlap. The older "NBUF=2 does not fit 64 KB LDS at D128" reason is **retired**: it only holds at the shipped `block_n=64` (a K-only double buffer, or `block_n=32`, does fit), so what actually closes the door is the measured `block_n=32` verdict in the row above, not the LDS arithmetic | +| partial-vmcnt software prefetch | N/A on the default build — it is a double-buffering lever and the default build is single-buffered (NBUF=1), so there is no prefetch to partially overlap. A full K+V NBUF=2 fits 64 KB LDS at D128 only at `block_n=32`, which the `block_n=32` row above closes by measurement. The opt-in `lds_num_buffers=2` path doubles only K at `block_n=64`; `v_reg_prefetch=True` independently brings next-tile V into registers. Both paths wait with full `vmcnt(0)`, not a partial one | ## Bench diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md new file mode 100644 index 000000000000..4e85c6f0c869 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md @@ -0,0 +1,553 @@ +# gfx942 bf16 D128 causal prefill: CFVST, K LDS double buffering, and V register prefetch + +This change makes CFVST available as an opt-in and implements opt-in K LDS Double Buffering and V register pre-fetching on D128 bf16 persistent=true causal prefill attention gfx942 kernels. All three are off by default. + +## Problem + +Benchmarking, profiling and tracing highlighted bottlenecks in: + +1. V read from LDS: V was stored in traditional layout in LDS. Narrow reads. +2. Allocation of LDS per WG (34 KiB, 34,816 B, in the production baseline). Max limit per CU is 64 KiB. This limits WG occupancy in CU to 1 while leaving 30 KiB of LDS unused. +3. Stalls at sync point for loading K and V tiles before MFMA. + +The concerned optimization points were V LDS layout, register usage, LDS usage, software pipelining. + +## Step-0 Lever Sweep + +Lever sweep was performed to check whether performance could be improved by tuning levers to affect optimization points: **LDS usage/layout, software pipelining/scheduling, and register pressure**. Available levers are listed below along with reasons for selection decision and values swept. + +#### Configuration + +| Parameter | Value | +|---|---| +| GPU | (`gfx942`, 304 CUs) | +| Repository / branch | [ROCm/rocm-libraries](https://github.com/ROCm/rocm-libraries), `develop` | +| Source commit | `ceeae550e7aa95b291aff5cb75ec4d626a072bf3`| +| Data type | **BF16** | +| Head dimension (`D`) | **128** | +| Sequence length (`Sq = Skv = S`) | **4096, 8192, 16384** | +| Batch size (`B`) | **1, 16** | +| Query heads (`Hq`) | **32** | +| KV heads (`Hkv`) | **8, 16** | +| Attention | Dense, causal prefill | + +Eight `(B, S, Hkv)` workloads were swept independently. The completed sweep report records this source commit with LLVM flavor `llvm23`. + +#### Lever inventory and exclusions + +| Lever | Swept? | Group | Reason if excluded | +|---|---|---|---| +| `use_cfvst` | No | — | Not supported when enabled for BF16 D128 at this commit; keep the non-CFVST baseline. | +| `lds_row_pad` | **Yes** | A | — | +| `lds_k_group_pad` | No | — | Only changes the D64 K-LDS row-group layout; D128 uses `lds_row_pad`. | +| `persistent` | No | — | Selects persistent vs nonpersistent grid scheduling, not the inner K/V LDS/register pipeline; preserve production dispatch. | +| `block_m` | **Yes** | A | — | +| `block_n` | **Yes** | A | — | +| `num_persistent` | No | — | Changes persistent workgroup count and device-level utilization, not the inner software pipeline; preserve production dispatch. | +| `persist_decode` | No | — | Changes grid-level mapping of work to query blocks/heads/batches; preserve production dispatch. | +| `interleave` | No | — | Changes query-block work ordering across CTAs, not instruction interleaving within the K/V pipeline; preserve production dispatch. | +| `use_v_swizzle` | No | — | Only affects the CFVST transposed-V LDS path, which is disabled for BF16 D128. | +| `v_row_pad` | No | — | Only affects the CFVST transposed-V LDS layout; has no effective tuning impact on the selected baseline path. | +| `use_exp2_fast` | **Yes** | C | — | +| `iglp` | **Yes** | B | — | +| `iglp_mode` | **Yes** | B | — | +| `pv_sched_fence_mask` | **Yes** | B | — | +| `pv_priority` | **Yes** | B | — | +| `pv_loop_order` | **Yes** | A, B | — | +| `o_store_width` | **Yes** | D | — | +| `causal_diag_split` | **Yes** | A, C | — | +| `waves_per_eu` | **Yes** | A, B, C | — | +| `lds_num_buffers` | No | — | Only `1` is implemented/valid for the kernel at this revision; no useful sweep axis. | +| `lazy_rescale` | No | — | Inactive in the gfx942 kernel body; only changes the kernel-name token. | +| `emit_lse` | No | — | Requests an additional LSE output, rather than tuning the same output contract. | + +**Count: 12 swept.** The grid knobs were **not forced off**; each workload retained its dispatch-resolved settings. CFVST, K LDS double buffering, and V register prefetch were not enabled in the baseline sweep. + +#### Grouping the 12 swept knobs + +| Group | Purpose | Knobs | +|---|---|---| +| **A — Geometry / LDS** | Tile size, LDS layout, occupancy-related interactions | `block_m`, `block_n`, `lds_row_pad`, `waves_per_eu`, `pv_loop_order`, `causal_diag_split` | +| **B — Instruction scheduling** | IGLP, barriers, PV ordering and priority | `iglp`, `iglp_mode`, `pv_sched_fence_mask`, `pv_priority`, `pv_loop_order`, `waves_per_eu` | +| **C — Softmax / causal computation** | Exponential policy, diagonal handling and wave interactions | `use_exp2_fast`, `causal_diag_split`, `waves_per_eu` | +| **D — Output stores** | Store vector width | `o_store_width` | + +There are **12 unique knobs**, not 16: `waves_per_eu`, `pv_loop_order`, and `causal_diag_split` appear in multiple groups so their interactions can be revisited. + +#### Staged sweep methodology + +The sweep followed **A1 → A2 → C → B1 → B2 → D → R1 → R2**, independently for each workload. + +| Stage | Procedure | +|---|---| +| **A1 — Geometry screening** | Sweep `block_m={64,128,256}`, `block_n={32,64}`, `lds_row_pad={0,4,8,16}`, `waves_per_eu={1,2,4}` (72 combinations before legality/deduplication). Shortlist promising complete configurations. | +| **A2 — Geometry interactions** | On A1 finalists, retune `waves_per_eu={1,2,3,4}`, `pv_loop_order={d_major,k_major}`, `causal_diag_split={False,True}`. | +| **C — Computation** | On A2 finalists, tune `use_exp2_fast={Auto,False,True}`, `causal_diag_split`, and `waves_per_eu`. | +| **B1 — Scheduling** | On C finalists, test legal `iglp`/`iglp_mode`/`pv_sched_fence_mask` combinations together with `pv_loop_order` and `pv_priority={0,1,2,3}`. IGLP and explicit scheduling fences are mutually exclusive. | +| **B2 — Scheduling interactions** | On B1 finalists, retune `waves_per_eu` and `pv_loop_order`. | +| **D — Output** | On B2 finalists, test `o_store_width={1,2,4}`. | +| **R1, R2 — Cross-group refinement** | Two bounded rounds of single-knob neighborhood changes around the strongest **full configurations**, allowing adjustments to earlier group choices. | + +At every stage, the sweep rejects unsupported/duplicate effective specs, numerically checks legal candidates against FP32 causal SDPA (`atol=rtol=1e-2`), screens GPU latency, then remeasures promising candidates against the production baseline. **Several top full configurations** are carried forward, rather than choosing one winner per independent knob. This is a **bounded staged search**, not an exhaustive Cartesian-product sweep. + +#### Production dispatch default vs winning kernels + +These are the best **final A/B-tested** configurations, compared with each workload's production-dispatch baseline, as measured by the sweep at its source commit. Results table 3 remeasures the same winners at the benchmarked HEAD. Each finalist passed a BF16-output-versus-FP32 causal SDPA numerical check (`atol=rtol=1e-2`) before timing. The timing statistic is the median of five alternating paired baseline/candidate rounds, with 50 timed launches per round. + +| B | S | Hkv | Speedup | Latency reduction | Paired min–max | Baseline kernel | Winning kernel | +|---:|---:|---:|---:|---:|---|---|---| +| 1 | 4096 | 8 | **1.10141×** | **9.21%** | 1.09785–1.10610× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe3_krowpad4_iglp1_iglpm1_pvkmaj_prio1_dsplit` | +| 1 | 4096 | 16 | **1.09689×** | **8.83%** | 1.09584–1.09764× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2_krowpad4_iglp1_iglpm1_pvkmaj_prio1_dsplit` | +| 1 | 8192 | 8 | **1.10407×** | **9.43%** | 1.10167–1.10657× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe3_krowpad4_iglp1_iglpm1_pvkmaj_prio1_dsplit` | +| 1 | 8192 | 16 | **1.10338×** | **9.37%** | 1.10162–1.10493× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe3_krowpad4_iglp1_iglpm1_pvkmaj_prio2_dsplit` | +| 1 | 16384 | 8 | **1.09043×** | **8.29%** | 1.08865–1.09340× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe3_krowpad4_iglp1_iglpm1_pvkmaj_prio3_dsplit` | +| 16 | 4096 | 8 | **1.10144×** | **9.21%** | 1.10002–1.10216× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2_krowpad4_iglp1_iglpm1_pvkmaj_prio3_dsplit` | +| 16 | 4096 | 16 | **1.08896×** | **8.17%** | 1.08857–1.08966× | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_gfx942_wpe2` | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_gfx942_wpe2_krowpad4_iglp1_iglpm1_pvkmaj_prio3_dsplit` | +| 16 | 8192 | 8 | — | — | — | — | **Missing files** | + +This shape has no completed Step-0 result, so it is compared against production dispatch only. + +## Optimization 1: Enable CFVST + +V was stored in traditional layout in LDS. This caused narrow loads. The hypothesis was that transposed layout would enable wide loads. This feature was already present (CFVST) but disabled for bf16 due to register spillage concerns. Performance improved substantially with it enabled, so it is offered as an opt-in. + +CFVST transposes V into a layout that enables wider LDS reads during the PV MFMA stage. + +**Comparison:** Production default vs. CFVST enabled +**Workload:** B=1, S=4096, Hq=32, Hkv=8, D=128, BF16, causal, gfx942. + +| Metric | Production default | CFVST enabled | +|---|---|---| +| Static `ds_read_u16` instructions | Baseline | **−100%** | +| Static `ds_read2st64_b64` instructions | Baseline | Introduced (0 → 16; percentage undefined) | +| Executed LDS instructions | Baseline | **−71.23%** | + +Static counts are from the benchmark run's HSACOs; executed LDS instructions are `SQ_INSTS_LDS` per dispatch from `rocprofv3 --pmc`, both at the benchmarked HEAD. + +**Conclusion:** CFVST improves performance through wider V LDS reads and reduced LDS instruction traffic. It is **not** justified as a bank-conflict reduction optimization. + +### Register Usage and Static Occupancy (§R.4) + +Static counts from the exact HSACOs timed in the benchmark run (`static_resources.csv`; LLVM23, BM256/BN64, persistent-304, 512 threads/WG), for B=1, S=4096, Hkv=8. Instruction columns are static counts in the code object, not executed counts. + +| Arm | VGPR | SGPR | VGPR / SGPR spills | Scratch | LDS bytes | `v_mfma` | `ds_read` | `ds_write` | `buffer_load` | `s_barrier` | `s_waitcnt` | VALU | SALU | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| A (Step-0 winner) | 249 | 106 | 0 / 1 | 0 B | 33792 | 128 | 288 | 0 | 0 | 5 | 86 | 1645 | 493 | +| B1 baseline | 194 | 44 | 0 / 0 | 0 B | 34816 | 64 | 144 | 0 | 0 | 3 | 86 | 1194 | 469 | +| B2 CFVST | 225 | 45 | 0 / 0 | 0 B | 33792 | 64 | 32 | 8 | 8 | 3 | 31 | 1225 | 428 | +| B3 CFVST + KDB | 230 | 78 | 0 / 0 | 0 B | 51200 | 128 | 64 | 16 | 16 | 5 | 52 | 1600 | 600 | +| B4 CFVST + VPF | 230 | 76 | 0 / 0 | 0 B | 33792 | 128 | 64 | 16 | 16 | 5 | 57 | 1606 | 573 | +| B5 CFVST + KDB + VPF | 230 | 78 | 0 / 0 | 0 B | 51200 | 128 | 64 | 16 | 16 | 5 | 59 | 1604 | 585 | +| B6 Full + pad4 | 230 | 78 | 0 / 0 | 0 B | 50176 | 128 | 64 | 16 | 16 | 5 | 59 | 1604 | 586 | +| B7 Full + pad4 + diagonal split | 232 | 81 | 0 / 0 | 0 B | 50176 | 192 | 96 | 24 | 24 | 7 | 88 | 1964 | 661 | + +Every B arm is spill-free with 0 B scratch on all twelve shapes. `A` has one SGPR spill on six of its seven shapes. Every arm reports 2 waves/SIMD. Across shapes, the B arms change only in their SGPR, VALU, and SALU counts, following each shape's production block order. + +The SGPR V addressing (the wave index scalarized before the V buffer offsets) is used only by the KDB and VPF pipeline paths; CFVST alone and the default kernel are unchanged. It has no separate knob, so its delta is included in the CFVST → CFVST + KDB and CFVST → CFVST + VPF rows. The gfx950 dense kernel computes its wave index in VGPRs and does not share this pattern. + +## Optimization 2 and 3: K LDS Double Buffering + V register pre-fetching + +Analysis of stalls using threadtrace found a large number of waiting at sync points. One such point was waiting for K and V tiles to be loaded before MFMA. This wait restricted MFMA and memory loading parallelism. The hypothesis was that K LDS Double Buffering + V register pre-fetching would introduce parallelism and latency hiding. Optimizing 1 would leave other the bottleneck. The optimization improved performance substantially, so it is offered as an opt-in. + +### Synchronization-Stall Analysis + +Advanced Thread Trace (ATT) was used to evaluate synchronization overhead for gfx942 BF16 D128 causal prefill (B=1, S=4096, Hq=32, Hkv=8): `rocprofv3 --att` on CU 1, one traced dispatch per kernel at the benchmarked HEAD. Stalls are the `s_barrier` and `s_waitcnt` stall cycles in the trace's `code.json`. + +**CFVST + KDB + VPF vs. CFVST** + +| Metric | Reduction | +|---|---:| +| BARRIER hits | 0.00% | +| BARRIER issue stalls | **10.70%** | +| WAITCNT issue stalls | **27.23%** | +| **Combined issue stalls** | **21.31%** | + +The combined pipeline reorganizes K/V loading through a prologue, pipelined main loop, and epilogue. ATT records the same aggregate number of barrier hits across both variants, but fewer BARRIER- and WAITCNT-attributed issue stalls. + +These results are consistent with reduced synchronization overhead and improved latency hiding through K double buffering and V register prefetch. + +## Results + + +All workloads use **BF16, Hq=32, D=128, persistent=True, causal attention** on (gfx942). + +### 1. Ablation: Each Configuration vs. Step-0 Winner + +| B | S | Hkv | CFVST | CFVST + KDB | CFVST + VPF | CFVST + KDB + VPF | Full + LDS pad4 + causal split | CFVST min–max | CFVST + KDB min–max | CFVST + VPF min–max | CFVST + KDB + VPF min–max | Full + LDS pad4 + causal split min–max | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1024 | 8 | — | — | — | — | — | — | — | — | — | — | +| 1 | 1024 | 16 | — | — | — | — | — | — | — | — | — | — | +| 1 | 2048 | 8 | — | — | — | — | — | — | — | — | — | — | +| 1 | 2048 | 16 | — | — | — | — | — | — | — | — | — | — | +| 1 | 4096 | 8 | 1.37% | 3.64% | 5.98% | 7.81% | 11.44% | 1.00% to 57.80% | 3.29% to 59.00% | 5.53% to 59.98% | 7.50% to 60.61% | 11.17% to 61.87% | +| 1 | 4096 | 16 | 1.16% | 2.93% | 6.17% | 7.34% | 10.21% | 0.10% to 1.78% | 1.72% to 3.02% | 5.73% to 6.69% | 6.77% to 7.57% | 9.41% to 15.02% | +| 1 | 8192 | 8 | 1.38% | 3.78% | 6.76% | 8.43% | 12.72% | 1.11% to 1.62% | 3.39% to 3.90% | 6.71% to 7.11% | 8.27% to 8.58% | 12.53% to 13.78% | +| 1 | 8192 | 16 | 1.31% | 3.67% | 6.13% | 8.27% | 13.00% | 1.08% to 1.49% | 3.40% to 4.15% | 5.31% to 6.77% | 7.97% to 8.61% | 12.75% to 13.00% | +| 1 | 16384 | 8 | 1.40% | 2.90% | 6.80% | 7.96% | 13.43% | 0.89% to 1.83% | 2.46% to 3.57% | 6.54% to 7.03% | 7.57% to 8.32% | 13.15% to 13.63% | +| 16 | 4096 | 8 | 0.55% | 1.72% | 5.49% | 6.85% | 10.65% | 0.42% to 0.72% | 1.55% to 1.98% | 5.32% to 5.51% | 6.66% to 6.90% | 10.40% to 10.74% | +| 16 | 4096 | 16 | 0.82% | 3.43% | 6.19% | 7.63% | 10.52% | 0.62% to 0.99% | 3.30% to 3.59% | 5.78% to 6.35% | 7.39% to 7.71% | 10.26% to 10.69% | +| 16 | 8192 | 8 | — | — | — | — | — | — | — | — | — | — | + +### 2. Ablation: Each Configuration vs. Production Dispatch + +| B | S | Hkv | CFVST | CFVST + KDB | CFVST + VPF | CFVST + KDB + VPF | Full + LDS pad4 + causal split | CFVST min–max | CFVST + KDB min–max | CFVST + VPF min–max | CFVST + KDB + VPF min–max | Full + LDS pad4 + causal split min–max | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1024 | 8 | 8.74% | 11.92% | 15.57% | 15.96% | 18.98% | 6.51% to 11.12% | 8.96% to 12.98% | 14.32% to 16.58% | 14.02% to 18.54% | -405.47% to 20.04% | +| 1 | 1024 | 16 | 9.27% | 11.30% | 14.62% | 15.64% | 16.78% | 7.98% to 11.34% | 10.86% to 13.14% | 13.91% to 15.96% | 15.41% to 17.55% | 16.07% to 18.67% | +| 1 | 2048 | 8 | 10.94% | 13.24% | 15.70% | 17.90% | 20.99% | 10.42% to 11.35% | 12.24% to 14.11% | 15.42% to 16.29% | 17.32% to 18.37% | 19.95% to 21.62% | +| 1 | 2048 | 16 | 11.01% | 13.15% | 15.13% | 17.12% | 20.21% | 9.93% to 11.89% | 11.96% to 14.04% | 10.72% to 16.43% | 15.44% to 17.95% | 19.67% to 20.53% | +| 1 | 4096 | 8 | 9.52% | 11.61% | 13.75% | 15.43% | 18.76% | 8.74% to 10.01% | 11.20% to 12.31% | 12.92% to 14.96% | 15.02% to 18.43% | 18.42% to 19.10% | +| 1 | 4096 | 16 | 9.58% | 11.21% | 14.17% | 15.24% | 17.86% | 9.16% to 9.73% | 9.84% to 12.31% | 13.34% to 14.82% | 14.50% to 16.03% | 17.31% to 21.74% | +| 1 | 8192 | 8 | 10.12% | 12.30% | 15.02% | 16.54% | 20.46% | 9.76% to 10.33% | 11.84% to 12.48% | 14.94% to 15.27% | 16.35% to 16.64% | 20.24% to 21.35% | +| 1 | 8192 | 16 | 9.89% | 12.04% | 14.28% | 16.24% | 20.56% | 9.65% to 9.98% | 11.83% to 12.46% | 13.52% to 14.93% | 15.87% to 16.66% | 20.32% to 20.66% | +| 1 | 16384 | 8 | 10.41% | 11.77% | 15.31% | 16.37% | 21.34% | 9.89% to 10.69% | 11.32% to 12.25% | 15.03% to 15.49% | 15.97% to 16.58% | 21.11% to 21.41% | +| 16 | 4096 | 8 | 9.36% | 10.42% | 13.85% | 15.10% | 18.56% | 9.36% to 9.50% | 10.21% to 10.67% | 13.65% to 14.06% | 14.96% to 15.26% | 18.37% to 18.79% | +| 16 | 4096 | 16 | 8.87% | 11.27% | 13.80% | 15.12% | 17.78% | 8.77% to 9.07% | 11.17% to 11.45% | 13.64% to 14.01% | 15.03% to 15.37% | 17.69% to 17.84% | +| 16 | 8192 | 8 | 10.67% | 12.46% | 15.70% | 17.19% | 20.53% | 10.48% to 10.69% | 12.34% to 12.69% | 15.55% to 15.73% | 17.01% to 17.35% | 19.97% to 20.64% | + +### 3. Final Pipeline: Comparison of Both Baselines + +| B | S | Hkv | Step-0 vs. Production | Full Pipeline vs. Step-0 | Full Pipeline vs. Production | Step-0 vs. Production min–max | Full Pipeline vs. Step-0 min–max | Full Pipeline vs. Production min–max | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1024 | 8 | — | — | 18.98% | — | — | -405.47% to 20.04% | +| 1 | 1024 | 16 | — | — | 16.78% | — | — | 16.07% to 18.67% | +| 1 | 2048 | 8 | — | — | 20.99% | — | — | 19.95% to 21.62% | +| 1 | 2048 | 16 | — | — | 20.21% | — | — | 19.67% to 20.53% | +| 1 | 4096 | 8 | 8.26% | 11.44% | 18.76% | -113.87% to 8.93% | 11.17% to 61.87% | 18.42% to 19.10% | +| 1 | 4096 | 16 | 8.53% | 10.21% | 17.86% | 7.90% to 9.65% | 9.41% to 15.02% | 17.31% to 21.74% | +| 1 | 8192 | 8 | 8.86% | 12.72% | 20.46% | 8.75% to 8.93% | 12.53% to 13.78% | 20.24% to 21.35% | +| 1 | 8192 | 16 | 8.69% | 13.00% | 20.56% | 8.58% to 8.81% | 12.75% to 13.00% | 20.32% to 20.66% | +| 1 | 16384 | 8 | 9.13% | 13.43% | 21.34% | 9.00% to 9.35% | 13.15% to 13.63% | 21.11% to 21.41% | +| 16 | 4096 | 8 | 8.85% | 10.65% | 18.56% | 8.80% to 9.12% | 10.40% to 10.74% | 18.37% to 18.79% | +| 16 | 4096 | 16 | 8.12% | 10.52% | 17.78% | 7.94% to 8.35% | 10.26% to 10.69% | 17.69% to 17.84% | +| 16 | 8192 | 8 | — | — | 20.53% | — | — | 19.97% to 20.64% | + +**Outliers retained:** in one round each, an arm ran 2.3–6.3× slower than its own median (B=1, S=1024, Hkv=8, Full + LDS pad4 + causal split; B=1, S=4096, Hkv=8, the Step-0 winner `A`), while the other arms in that round did not. The min–max ranges include those rounds as measured; the medians are unaffected. + +## Results and Ablation Methodology + +The benchmark compares each shape's **Step-0 existing-lever sweep winner (`A`)** against seven independently compiled configurations derived from the **current production dispatch**. It uses twelve persistent-304 BF16 D128 causal attention shapes. Seven have a Step-0 winner; B=1, S=1024/2048 (Hkv 8 and 16) and B=16, S=8192, Hkv=8 do not, so they are compared with production dispatch only. + +For each of the twelve shapes (the five shapes without a Step-0 winner skip `A` and every comparison with it): + +1. Start one **torch-free** rocKE/NumPy Python process with one ROCm stack, the Python compilation backend, and LLVM23. Refuse a dirty or untracked checkout. Record the GPU, the loaded HIP runtime/COMGR libraries, Git HEAD, and the SHA-256 of the kernel, spec, dispatch, builder, and harness sources. +2. Create one dedicated HIP stream using `rocke.runtime.hip_module` and use that stream for every kernel launch and timing event. +3. Build the shape's `AttentionRequest` and resolve the production specification using `attention_tuning_spec(req, "gfx942_dense").kernel_spec`. Require persistent dispatch (`persistent=True`, `num_persistent=304`); reject unsupported shapes. +4. Build the separate **`A` comparator** through `tuning_spec_with_knobs(req, "gfx942_dense", knobs)` from the recorded shape-specific Step-0 winning *existing* levers, and abort unless it compiles to the recorded winner kernel name. Keep `use_cfvst=False`, `lds_num_buffers=1`, and `v_reg_prefetch=False`. The previous winner is rebuilt using the **current checkout**, not timed using historical measurements. +5. Build seven **B-stage** specifications through `tuning_spec_with_knobs` from their knob deltas against the *current production-dispatch specification*, **not** from `A`. Preserve all non-ablation production knobs except the explicitly stated LDS-padding and causal-split overrides: + + | Arm | Variant | `use_cfvst` | `lds_num_buffers` (K) | `v_reg_prefetch` | `lds_row_pad` | `causal_diag_split` | + |---:|---|---:|---:|---:|---|---| + | A | `A` (Step-0 winner) | `False` | 1 | `False` | Shape-specific winner | Shape-specific winner | + | B1 | `baseline` | `False` | 1 | `False` | Production dispatch | Production dispatch | + | B2 | `cfvst` | `True` | 1 | `False` | Production dispatch | Production dispatch | + | B3 | `cfvst_kdb` | `True` | 2 | `False` | Production dispatch | Production dispatch | + | B4 | `cfvst_vpf` | `True` | 1 | `True` | Production dispatch | Production dispatch | + | B5 | `pipeline` | `True` | 2 | `True` | Production dispatch | Production dispatch | + | B6 | `pipeline_pad` | `True` | 2 | `True` | 4 | Production dispatch | + | B7 | `pipeline_pad_dsplit` | `True` | 2 | `True` | 4 | `True` | + +6. Assert that the A configuration disables the new pipeline and that production dispatch starts with the three new knobs disabled. Verify every successive B-stage change modifies exactly its advertised field(s), including resolved knob values. Record each arm's `spec_id`, `tuning_id`, canonical knobs, and dispatch-time `spec_hash`, replay each record through `dispatch_attention`, and abort on any `spec_hash` mismatch. +7. Check `supports_attention_dense` for every arm. Require distinct compile-time kernel names; retain the recorded and rebuilt Step-0 winner names for auditing. +8. Run the Python kernel emitter and compile **one independent HSACO per arm**; record their kernel names, code-object hashes, specifications, and static register, spill, LDS, and ISA instruction counts. No runtime GPU-kernel variant-selection branch is used. +9. Generate one shared BF16-rounded Q/K/V input set. Allocate a separate output buffer and construct a direct `KernelLauncher` plus `LaunchConfig` for each compiled kernel. +10. Compute one independent **chunked NumPy FP32 causal GQA SDPA** reference from the shared BF16-rounded inputs. No PyTorch import or `torch.testing.assert_close` is used. +11. **Before timing**, fill each output with `0xFF` bytes (NaN poison), invoke its kernel, synchronize, and compare **every output element** against the FP32 reference using `numpy.testing.assert_allclose(atol=1e-2, rtol=1e-2, equal_nan=False)`. Record maximum absolute and relative errors; abort on failure. +12. Perform **10 round-robin warmup launches per variant**, then synchronize the HIP stream. +13. Execute **five measurement rounds**, randomly shuffling the order of all arms within each round. +14. For each variant in each round, record a HIP start event, make **50 direct `KernelLauncher` calls**, record a HIP end event, synchronize that event, and divide elapsed GPU-event time by 50 to obtain per-launch latency. +15. After each timed batch, validate its final output against the same FP32 reference with the same tolerances, **outside the GPU-event timing interval**. Record errors and abort on failure. Validation checks the last output of each 50-launch batch, not all 50 intermediate outputs. +16. Save the shape, variant, measurement round, randomized execution order, configuration knobs, latency, and correctness result. Preserve the resolved specs, kernel names, HSACO hashes, and environment provenance. + +**Correctness coverage:** Each of the seven shapes with a Step-0 winner has eight initial checks and `8 × 5 = 40` post-batch checks, or **48 checks per shape**. The five shapes without one run B1–B7 only, so each has `7 + 7 × 5 = 42`. Across the twelve shapes there are **546 checks** (91 initial and 455 post-batch). + +17. Collect five per-round latency samples per variant and compute each variant's independent median: + + `T_variant = median(t_variant,1, ..., t_variant,5)` + +18. Use the shape-specific **Step-0 winner (`A`) as the reference denominator** for the main tables. For each candidate B2, B3, B4, B5, or B7, compute its speedup as: + + `Speedup_A_to_candidate = T_A / T_candidate` + +19. Compute the main percentage latency reduction **against `A`**, not against production dispatch: + + `Reduction_A_to_candidate (%) = 100 × (1 - T_candidate / T_A)` + +20. For the five-column ablation table, report B2 (`CFVST`), B3 (`CFVST+KDB`), B4 (`CFVST+VPF`), B5 (`CFVST+KDB+VPF`), and B7 (`CFVST+KDB+VPF+LDS pad4+causal split`), **each compared directly with `A`**. +21. For the final full-optimization table, compare only **B7 versus `A`**. Calculate the percentage reduction for each matched round: + + `R_r = 100 × (1 - t_B7,r / t_A,r)` + +22. Report the observed same-round range: + + `Min = min(R_1, R_2, R_3, R_4, R_5)` + `Max = max(R_1, R_2, R_3, R_4, R_5)` + +23. Report the central B7-vs-A reduction from the **ratio of the two medians**, not the median of the five round-level reductions. +24. Independently calculate *within-B-stage* comparisons to isolate pipeline effects: `baseline → cfvst`, `cfvst → cfvst_kdb`, `cfvst → cfvst_vpf`, `cfvst_kdb → pipeline`, `cfvst_vpf → pipeline`, `pipeline → pipeline_pad`, and `pipeline_pad → pipeline_pad_dsplit`. +25. Export `ablation_runs.csv`, `ablation_summary.csv`, `ablation_ratios.csv`, `ablation_comparison.csv`, `ablation_comparison.md`, `A_vs_B7.csv`, `validation.csv`, `static_resources.csv`, `missing_sweep_winners.json`, `resolved_specs.json`, and `provenance.json` for reproduction and review. +26. Run `--verify-provenance` on the output at the PR head: it re-checks the source hashes, refuses any non-Markdown change since the measured commit, and replays every arm's record. + +**Interpretation:** `A` and B1–B7 are measured within the same session and compiled from the same current source, so direct **Step-0-winner-versus-optimized-production** comparisons are valid. However, the B stages **do not inherit `A`'s geometry or scheduling knobs**: the results do **not** quantify the incremental benefit of adding CFVST/KDB/VPF on top of the Step-0 winner itself. **Losses:** B=1, S=1024/2048 have no Step-0 winner in this sweep, so they are compared with production dispatch only. In the earlier twelve-shape Step-0 sweep at `c8980b69`, their winners were non-persistent configurations (`block_m=128`, `block_n=32`) that outperformed the CFVST/KDB/VPF pipeline, which runs on the persistent grid only. This is a loss at those shapes. + +#### Benchmark Environment and Provenance + +| Component | Configuration | +| :--- | :--- | +| Hardware target | gfx942, 304 CUs | +| Benchmarked source HEAD | `cc84e2621e8fd9ce83f9bf43dbc1564c690c290d` (clean working tree; `--verify-provenance` passed) | +| Kernel source SHA-256 | `8bf80337c529fdab6e3adf65d942bf423eac9623b324c6cd650d0987adb71d2e` | +| Benchmark harness SHA-256 | `888b484dab6a0dc76e5544de1660047e03956b10f9f898cf8cb3cb228ec61130` | +| ROCm stack | 10.0.0 | +| HIP | `7.15.26333` | +| Python | `3.12.3` | +| Triton | `3.8.0+git4cff872c.rocm10.0.0` | +| ROCKE compiler backend | Python | +| LLVM flavor | `llvm23` | + +#### Shape-by-shape default kernel identity built by `attention_tuning_spec(req, "gfx942_dense").kernel_spec` + +| dtype | B | S | Hq | Hkv | D | Exact baseline kernel name | persistent | exp2_fast (resolved) | decode (resolved) | grid | block | +|---|---:|---:|---:|---:|---:|---|---|---|---|---|---| +| bf16 | 1 | 1024 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 1024 | 32 | 16 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 2048 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 2048 | 32 | 16 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 4096 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 4096 | 32 | 16 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 8192 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 8192 | 32 | 16 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 1 | 16384 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2` | `true` | `true` | `bt_hkv_minor` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 16 | 4096 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2` | `true` | `true` | `hkv_major` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 16 | 4096 | 32 | 16 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_gfx942_wpe2` | `true` | `true` | `qb_major` | `[304, 1, 1]` | `[512, 1, 1]` | +| bf16 | 16 | 8192 | 32 | 8 | 128 | `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_hkvmaj_gfx942_wpe2` | `true` | `true` | `hkv_major` | `[304, 1, 1]` | `[512, 1, 1]` | + +**Benchmark harness:** `library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py` + +**Kernel source:** `library/kernels/gfx942/attention_dense.py` + +#### Reproducibility + +Clone repo to gpu-bench/. + +**`setup.sh`** — install the ROCm Python environment and rocKE: + +```bash +#!/usr/bin/env bash +set -Eeuo pipefail + +REPO="$(git rev-parse --show-toplevel)" +ROOT="${GPU_BENCH_ROOT:-/root/gpu-bench}" +ROCKE="$REPO/dnn-providers/hip-kernel-provider/rocke" +VENV="$ROOT/venvs" +ENV="$ROOT/env" + +AMD_INDEX=https://stable.repo.amd.com/rocm/whl-next/ + +mkdir -p "$VENV" "$ENV" + +# ---------------------------------------------------------------------- +# Host dependencies +# ---------------------------------------------------------------------- + +apt-get update +apt-get install -y \ + git python3.12 python3.12-venv python3.12-dev python3-pip \ + cmake ninja-build build-essential pkg-config libnuma-dev + +# ---------------------------------------------------------------------- +# ROCKE virtual environment +# ---------------------------------------------------------------------- + +ROCKE_VENV="$VENV/rocke" + +if [ ! -x "$ROCKE_VENV/bin/python" ] || \ + ! "$ROCKE_VENV/bin/python" -m pip --version >/dev/null 2>&1; then + rm -rf "$ROCKE_VENV" + python3.12 -m venv "$ROCKE_VENV" +fi + +"$ROCKE_VENV/bin/python" -m pip install -U \ + pip setuptools wheel packaging ninja psutil + +"$ROCKE_VENV/bin/python" -m pip install \ + --index-url "$AMD_INDEX" \ + 'rocm[devel,device-gfx942,libraries]==10.0.0' \ + 'torch[device-gfx942]==2.13.0+rocm10.0.0' + +"$ROCKE_VENV/bin/rocm-sdk" init >/dev/null + +# ---------------------------------------------------------------------- +# Install ROCKE +# ---------------------------------------------------------------------- + +"$ROCKE_VENV/bin/python" -m pip install -e "$ROCKE/platform" +"$ROCKE_VENV/bin/python" -m pip install -e "$ROCKE/library" --no-deps + +# ---------------------------------------------------------------------- +# ROCKE environment activation script +# ---------------------------------------------------------------------- + +cat > "$ENV/rocke_env.sh" </dev/null + +export ROCM_HOME="\$(rocm-sdk path --root)" +export ROCM_PATH="\$ROCM_HOME" +export HIP_PATH="\$ROCM_HOME" +export PATH="\$(rocm-sdk path --bin):\$PATH" + +export HIP_PLATFORM=amd +export ROCKE_BACKEND=python +export ROCKE_LLVM_FLAVOR=llvm23 + +LIB="\$(python -c 'import importlib.util, pathlib; s=importlib.util.find_spec("_rocm_sdk_core"); print(pathlib.Path(list(s.submodule_search_locations)[0]) / "lib")')" + +export ROCKE_COMGR_LIB="\$LIB/libamd_comgr.so.3" +export ROCKE_HIP_LIB="\$LIB/libamdhip64.so.7" +export LD_LIBRARY_PATH="\$LIB:\${LD_LIBRARY_PATH:-}" + +unset LIB +EOF + +bash -n "$ENV/rocke_env.sh" + +echo +echo "ROCKE setup complete." +echo "ROCKE env: source $ENV/rocke_env.sh" +``` + +Run the setup script from its directory: + +```bash +chmod +x setup.sh +./setup.sh +``` + +Run the 7-way benchmark on a committed, clean checkout (a dirty tree is refused), then verify the run's provenance: + +```bash +ROOT="${GPU_BENCH_ROOT:-/root/gpu-bench}" +source "$ROOT/env/rocke_env.sh" +cd "$(git rev-parse --show-toplevel)/dnn-providers/hip-kernel-provider/rocke" + +BENCH=library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py +python "$BENCH" --warmup 10 --rounds 5 --repeat 50 --atol 1e-2 --rtol 1e-2 \ + --out-dir "$ROOT/rocke_7way_specs_results" +python "$BENCH" --verify-provenance "$ROOT/rocke_7way_specs_results" +``` + +Results are written outside the repository, to the `--out-dir` directory. + +### Global-Memory Coalescing and ISA Verification (§R.5) + +Verified on **gfx942**, BF16 D128 causal attention, BM256/BN64, persistent-304, LLVM23, B=1, S=4096, Hkv=8, on the HSACOs timed in the benchmark run. + +| Global access | Modeled positions | Coalesced | Cache lines | Footprint | ISA | +|---|---:|---|---:|---:|---| +| K async → LDS | 64 | Yes | 2 | 1.0× | `buffer_load_dword ... lds` (32-bit) | +| V CFVST x0 | 32 | Yes | 2 | 1.0× | `buffer_load_dword` (32-bit) | +| V CFVST x1 | 32 | Yes | 2 | 1.0× | `buffer_load_dword` (32-bit) | + +`analyze_coalescing` (with `vector_transactions` and `addr_fn_from_strides`, `line_bytes_for("gfx942") = 128`) reports for every access: `load [token stride 1024, head_dim stride 1]: contiguous axis = head_dim; 1 instr, VW=2 elems (b128-ideal), lines/instr 2..2 (128B lines), footprint 1x, FULLY COALESCED`. **`assert_asm_backed` passed** for K in all seven B-arm HSACOs and for V x0/x1 in the six CFVST arms; no wider or narrower K/V global load is emitted. + +x0 and x1 are separate loads. In the KDB/VPF pipeline they share one VGPR lane offset and differ only in their SGPR `soffset`, one runtime token stride apart. + +**Result:** No modeled global-memory coalescing regression in the verified configurations. Evidence covers the tested K/V access paths, not every tuning configuration. + +## Runbook and Architecture Compliance + +The CFVST, K LDS double-buffering (KDB), and V register-prefetching (VPF) optimizations follow the CK DSL Optimization Runbook (`platform/dsl_docs/optimization/optimization_runbook.md`) and the gfx942-specific architecture constraints (`platform/dsl_docs/optimization/arch/gfx942.md`, §21). + +### Per-Lever Mapping + +| Lever | Implementation | Runbook sections | §13.3 Attention checklist | gfx942 architecture reference | +|---|---|---|---|---| +| `use_cfvst=True` | `perm_b32` store-path transpose into transposed V LDS, with optional XOR swizzling | §6.3 LDS, §6.4 bank conflicts, §6.5 registers, §7.2 operand layouts | Tile Q / K / V; shared memory capacity; register pressure from accumulators | §21.2 LDS geometry and banks; §21.3 `v_perm_b32` cross-lane primitive; §21.4 VGPR constraints; §21.8 transpose-read limitations | +| `lds_num_buffers=2` | Double-buffer K in LDS to overlap next-tile K loading with current-tile computation; V LDS remains single-buffered | §6.3 LDS, §8.1 software pipelining, §8.2 async copy, §8.3 waits and barriers | Tile Q / K / V; shared memory capacity; register pressure from accumulators | §21.2 64 KiB LDS capacity; §21.4 occupancy constraints; §21.6 async LDS load width; §21.8 ISA constraints | +| `v_reg_prefetch=True` | Prefetch next-tile V into registers to overlap memory loading with current-tile computation | §6.5 registers, §8.1 software pipelining, §8.2 async copy, §8.3 waits and barriers | Tile Q / K / V; register pressure from accumulators | §21.4 VGPR pressure and occupancy; §21.6 global buffer-load widths; §21.8 ISA constraints | + +### Architecture and Configuration Gating + +The optimizations are implemented in `library/kernels/gfx942/attention_dense.py`, using `Gfx942AttentionDenseSpec`. + +Architecture eligibility is enforced by the existing `arch != "gfx942"` checks in `supports_attention_dense(spec, arch=...)` and `build_attention_dense(spec, arch=...)`. The support predicate returns `(False, reason)` for other architectures, while the builder raises `NotImplementedError`. Dispatch also restricts the gfx942 dense candidate to gfx942 through its architecture-specific registration and spec factory. + +The optimization-specific eligibility conditions are enforced by `supports_attention_dense()`: + +- **CFVST (`use_cfvst=True`):** For BF16, requires D128 causal attention without sliding window or attention sinks, on either grid. Causal diagonal splitting is supported. Existing FP16 CFVST behavior is preserved. +- **KDB (`lds_num_buffers=2`):** Requires BF16 D128 causal attention, CFVST enabled, `persistent=True`, `block_n=64`, no sliding window, and no attention sinks. +- **VPF (`v_reg_prefetch=True`):** Has the same eligibility conditions as KDB but is independently controlled. Enabling VPF does not require enabling KDB. + +The combined K/V LDS allocation is checked against the gfx942 capacity using `attention_lds_capacity_bytes(arch)` and `_lds_bytes(spec)`. Unsupported configurations are rejected before kernel emission. + +For BF16, the default configuration retains `use_cfvst=None` (policy resolves to off), `lds_num_buffers=1`, and `v_reg_prefetch=False`. All three optimizations remain opt-in and do not change the default kernel behavior. + +## gfx942 Barrier / LGKM ISA Evidence + +**Kernel:** `rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_bthkvmin_gfx942_wpe2_kdb1_vpf1_krowpad4_vrowpad0_cfvst1_dsplit` +**Workload:** BF16, B=1, Sq=Skv=4096, Hq=32, Hkv=8, D=128; gfx942, 304 CUs. +**Artifact:** `rocke_7way_specs_results/hsaco/b1_s4096_hq32_hkv8_d128_bf16/pipeline_pad_dsplit.hsaco` +**Evidence:** complete `.text` disassembly generated with `llvm-objdump -d` (not just emitter/IR), from the HSACO timed in the benchmark run at HEAD `cc84e2621e8fd9ce83f9bf43dbc1564c690c290d`. + +### Barrier inventory + +All **seven** `s_barrier` sites in this code object are accounted for: + +| Region | Last relevant LDS op | `s_waitcnt` | `s_barrier` | Finding | +| :--- | :--- | :--- | :--- | :--- | +| Pre-output / reduction | `ds_read2st64_b64` @ `0x2CF0` | `lgkmcnt(0)` @ `0x2D18` | `0x2D1C` | LDS drained | +| Pre-loop / initialization entry | No LDS ops since earlier drain on either entry path | `vmcnt(0)` @ `0x3960` | `0x3964` | No adjacent `lgkmcnt(0)`; see exception below | +| Initial LDS staging | `ds_write_b32` @ `0x3BA0` | `lgkmcnt(0)` @ `0x3BA8` | `0x3BAC` | LDS writes drained | +| Main-loop barrier 1 | `ds_read2st64_b64` @ `0x49F4` | `lgkmcnt(0)` @ `0x49FC` | `0x4A00` | LDS reads drained | +| Main-loop barrier 2 | `ds_write_b32` @ `0x4A78` | `lgkmcnt(0)` @ `0x4A8C` | `0x4A90` | LDS writes drained | +| Alternate-loop barrier 1 | `ds_read2st64_b64` @ `0x59DC` | `lgkmcnt(0)` @ `0x59E4` | `0x59E8` | LDS reads drained | +| Alternate-loop barrier 2 | `ds_write_b32` @ `0x5A60` | `lgkmcnt(0)` @ `0x5A70` | `0x5A74` | LDS writes drained | + +### Key `llvm-objdump` excerpts + +```asm +# Main loop: read-side barrier +0x49F4: ds_read2st64_b64 v[184:187], v132 offset0:82 offset1:90 +0x49FC: s_waitcnt lgkmcnt(0) +0x4A00: s_barrier +0x4A04: s_waitcnt vmcnt(0) +0x4A08: ds_write_b32 v117, v2 offset:33792 + +# Main loop: write-side barrier +0x4A70: ds_write_b32 v123, v2 offset:33792 +0x4A78: ds_write_b32 v124, v3 offset:33920 +0x4A8C: s_waitcnt lgkmcnt(0) +0x4A90: s_barrier + +# Alternate loop: same synchronization sequence +0x59DC: ds_read2st64_b64 v[192:195], v132 offset0:82 offset1:90 +0x59E4: s_waitcnt lgkmcnt(0) +0x59E8: s_barrier +0x59EC: s_waitcnt vmcnt(0) +0x59F0: ds_write_b32 v117, v2 offset:33792 +0x5A60: ds_write_b32 v124, v4 offset:33920 +0x5A70: s_waitcnt lgkmcnt(0) +0x5A74: s_barrier + +# Pre-output/reduction synchronization +0x2CF0: ds_read2st64_b64 v[170:173], v132 offset0:82 offset1:90 +0x2D18: s_waitcnt lgkmcnt(0) +0x2D1C: s_barrier +0x2D48: global_store_dwordx2 v[6:7], v[4:5], off + +# Initialization LDS staging +0x3BA0: ds_write_b32 v124, v2 offset:33920 +0x3BA8: s_waitcnt lgkmcnt(0) +0x3BAC: s_barrier +``` + +**Exception — `0x3964`:** + +```asm +0x3960: s_waitcnt vmcnt(0) +0x3964: s_barrier +``` + +There is **no immediately adjacent `lgkmcnt(0)`** at `0x3964`. The two paths visible in this disassembly reach it after a prior `lgkmcnt(0)` at `0x1BE4` (prologue path branching at `0x1EFC` to `0x38D0`) or `0x2D18` (post-output path); neither intervening instruction region contains another `ds_*` operation. This supports an already-drained LGKM state on those paths. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md index f1cba07155ea..3025588ecafe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md @@ -184,14 +184,45 @@ untunable (`lds_num_buffers`), or swept. gfx942 dense uses the same walk over `_GFX942_DENSE_AXES` (`("gfx942", "dense")`). It registers one candidate, so `persistent` is a knob there alongside `block_m` and `block_n`, and the persistent-only knobs -are pruned per spec. Its knobs are the K row and -V row pads, the D64 K group pad, the conflict-free-V store and its swizzle, -the persistent and non-persistent block orders, exp2, `iglp` / `iglp_mode`, -the PV fence mask, `pv_priority`, `pv_loop_order`, the bf16 O store width, -and `causal_diag_split`. The -LDS-saving knobs lead as enablers so an over-budget tile can still reach the -setting that fits. Untunable here: `lds_num_buffers` (only 1 is implemented) -and `lazy_rescale` (the body never reads it). +are pruned per spec. Its knobs include the K and V row pads, the D64 K group +pad, conflict-free V store (`use_cfvst`) and its swizzle, K LDS buffering +(`lds_num_buffers`), V register prefetch (`v_reg_prefetch`), the persistent +and non-persistent block orders, exp2, `iglp` / `iglp_mode`, the PV fence +mask, `pv_priority`, `pv_loop_order`, the bf16 O store width, and +`causal_diag_split`. LDS-saving and prerequisite axes lead as enablers so +an otherwise invalid configuration can become valid before later knobs +are considered. Only `lazy_rescale` is untunable on gfx942, because the +kernel body does not read it. + +The three performance levers are **independent opt-ins** for supported +BF16 D128 causal attention (`sliding_window=0`, no sinks), with or without +`causal_diag_split`. `use_cfvst=True` enables conflict-free V storage on +either grid; `lds_num_buffers=2` enables **K LDS double buffering only**; and +`v_reg_prefetch=True` enables **V register prefetch only**. KDB and VPF also +require CFVST, `persistent=True` and `block_n=64`, as enforced by the +kernel's validator; turning on either one does not implicitly enable the +other. + +The production tuning space remains **one knob at a time**, using CFVST +as a prerequisite when necessary. It can offer CFVST alone, CFVST + KDB, +and CFVST + VPF without requiring the triple combination. The triple +CFVST + KDB + VPF combination remains reachable through an explicit knob +pin or the full tuning space. The BF16 D128 default remains +`use_cfvst=None` (policy resolves to off), `lds_num_buffers=1`, and +`v_reg_prefetch=False`; these tuning choices do not alter the shipped +default or make the opt-in dense candidate an automatic dispatch winner. +Bare `tuning_id` replay searches production points only; replay a +full-space-only combination with its recorded knob dictionary. + +**Tuning-id migration for `gfx942_dense`.** Adding the defaulted +`v_reg_prefetch` field changes the defaults fingerprint in `config_key`, so +every `gfx942_dense` `tuning_id` recorded before it is rekeyed: replaying one +raises `PinRefused` instead of rebuilding a different kernel. No such id is +stored in this tree. A downstream tuning store must drop its old +`gfx942_dense` ids and re-sweep. `TUNING_ID_VERSION` is not bumped: the +defaults fingerprint already separates the old and new ids, and +`lds_num_buffers` was untunable on gfx942 before, so no stored id names it. + `registered_attention_combos(req)` is the multi-engine bench entry: it probes `ATTENTION_EXECUTION_REGISTRY` for `req.arch` and flattens each candidate's `sweep_space` (dense, WMMA, and unified tuning). Routing-only @@ -201,7 +232,7 @@ unified path labels are omitted. tier, that opt-in check is the only thing keeping them off the default path — a tier-3 candidate whose `support()` forgets it would silently claim all traffic for its arch. `Capability` cannot express the check (it constrains the request's -*selector*, not its shape), so it stays in the predicate and every tier-3 +**selector**, not its shape), so it stays in the predicate and every tier-3 candidate has to carry it. ## Layout @@ -238,7 +269,7 @@ the two unified paths and `d256_decode` -- which is not the same as portable: each still lists its arches explicitly, because there is no family wildcard. The two generic candidates declare every known arch on purpose: they select a -*path*, and `attention_unified` picks the concrete backend downstream from the +**path**, and `attention_unified` picks the concrete backend downstream from the running device (wave64 MFMA on gfx942/gfx950, wave32 WMMA on gfx1250, the arch-neutral scalar kernel elsewhere). They are an explicit, tested exception to the wave-size consistency invariant. @@ -284,7 +315,7 @@ Follow the `_make_d256_decode_candidate()` pattern in `generic.py`: [Definition of Done](../../KERNEL_AUTHORING.md#a-new-kernel-family) item. Point `build` at the real builder if the candidate has one. The unified - paths do not: they return an `AttentionSpec` naming a *path*, and no builder + paths do not: they return an `AttentionSpec` naming a **path**, and no builder consumes that. A standalone kernel like `attention_gfx1250_wmma` returns its builder's own spec instead, which is why it can declare `build`, a real grid, and a real signature where the unified candidates declare none. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py index ca8ac6783bf2..afea7886dd65 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/axes.py @@ -182,7 +182,7 @@ # ``lazy_rescale``) the body never reads the field and only the name changes. DENSE_UNTUNABLE_KNOBS: Mapping[str, frozenset] = { "gfx950": frozenset({"lds_num_buffers"}), - "gfx942": frozenset({"lds_num_buffers", "lazy_rescale"}), + "gfx942": frozenset({"lazy_rescale"}), } # Axes whose values are relative to the base spec (request or policy derived), # so a choice may equal a dataclass default; the base-equal choice is pruned. @@ -281,7 +281,10 @@ def _decode_axis(name: str, registry: Mapping, supported) -> KnobAxis: # as the tile (block_m, block_n). The LDS-saving knobs # lead as enablers: they can bring an otherwise over-budget tile under the LDS # limit (e.g. D64 block_n=256 fits only with lds_k_group_pad=0). Then -# prerequisites-first: exp2 policy reads persistent, v_row_pad policy reads +# prerequisites-first: BF16 CFVST is the required enabling choice for +# independent K double buffering and V register prefetch. The production +# tuning walk pairs each knob with that prerequisite, without requiring the +# combined KDB+VPF variant. Exp2 policy reads persistent, v_row_pad policy reads # block_n, a non-derived v_row_pad needs the swizzle off, interleave is read # only on the resolved qb_major decode, nonpersist_decode only on the # non-persistent grid, and iglp_mode / the PV fence follow iglp. @@ -289,6 +292,8 @@ def _decode_axis(name: str, registry: Mapping, supported) -> KnobAxis: _values("use_cfvst", None, (True, False), enabler=True), _values("lds_row_pad", 8, _GFX942_LDS_ROW_PADS, enabler=True), _values("lds_k_group_pad", 8, _DENSE_LDS_PADS, enabler=True), + _values("lds_num_buffers", 1, (2,)), + _flag("v_reg_prefetch"), _choices_axis("persistent", (True, False)), _choices_axis("block_m", _GFX942_BLOCK_M), _choices_axis("block_n", _GFX942_BLOCK_N), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py index 6be75f06dce0..6e0dc1412c5a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/dense_rules.py @@ -93,6 +93,16 @@ class _DenseTuningArch: production_crossed: Tuple[str, ...] = () +def _gfx942_cfvst_tuning_scope(s) -> bool: + """CFVST can be tuned where the default enables it or where the kernel + implements opt-in BF16 CFVST (``bf16_cfvst_supported``: D128 causal, with or + without causal_diag_split, on either grid). + """ + from kernels.gfx942.attention_dense import bf16_cfvst_supported + + return s.resolved_use_cfvst() or bf16_cfvst_supported(s) + + _DENSE_ARCH: Mapping[str, _DenseTuningArch] = { "gfx950": _DenseTuningArch( scope={ @@ -130,11 +140,12 @@ class _DenseTuningArch: "interleave": lambda s: bool(s.causal), "lds_row_pad": lambda s: int(s.head_size) == 128, "lds_k_group_pad": lambda s: int(s.head_size) < 128, - # Policy turns cfvst (and its swizzle) on only where it is legal; - # the knobs can only turn it off. - "use_cfvst": lambda s: s.resolved_use_cfvst(), - "use_v_swizzle": lambda s: s.resolved_use_cfvst(), - "v_row_pad": lambda s: s.resolved_use_cfvst(), + # FP16 D128 CFVST is policy-enabled. BF16 D128 causal CFVST + # (diagonal split included) is opt-in and must also be reachable + # through tuning. + "use_cfvst": _gfx942_cfvst_tuning_scope, + "use_v_swizzle": _gfx942_cfvst_tuning_scope, + "v_row_pad": _gfx942_cfvst_tuning_scope, "causal_diag_split": lambda s: bool(s.causal) and int(s.sliding_window) == 0, }, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index 9aef59557672..d239a59ac80a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py @@ -111,26 +111,28 @@ occupancy/latency-bound at 1 WG/CU and nowhere near compute-bound -- and the peel was proven negative on gfx942 tiled_2d via register pressure, which D128 (at its register floor) is the most exposed to. Keep the single-phase mask. - * partial-vmcnt software prefetch: N/A here, but NOT for the reason previously - recorded. It is a double-buffered pipelining lever and this kernel is - single-buffered (NBUF=1), so there is no prefetch to partially overlap. The old - "NBUF=2 does not fit 64 KB LDS at D128" justification is only true at the shipped - ``block_n=64`` (fp16 33 KB, 2x overflows): at ``block_n=32`` the K tile halves (V - stays width-64), D128 LDS drops to ~24.5 KB and NBUF=2 DOES fit. - That door is closed by measurement instead -- ``block_n=32`` was GPU-timed in the - Step-0 funnel and is proven-negative for fp16 D128 and only part-dependently - positive for bf16 D128 (halving the KV tile doubles the tile/grid count and the - extra loop/barrier overhead outweighs the LDS relief while occupancy stays - pinned), so it is not wired and the prefetch it would unlock is unreachable on - the shipped config. Independently, bounding the live set with partial ``vmcnt=N`` - waits entangles with the K async DMA, which shares the same counter -- see the - reverted cfvst-chunking experiment recorded in the ``V_ITEMS`` NOTE below. + * partial-vmcnt software prefetch: N/A on the default build. It is a + double-buffered pipelining lever and the default build is single-buffered + (NBUF=1), so there is no prefetch to partially overlap. A full K+V NBUF=2 does not + fit 64 KB LDS at D128 ``block_n=64``; at ``block_n=32`` it fits, but ``block_n=32`` + was GPU-timed in the Step-0 funnel and is proven-negative for fp16 D128 and only + part-dependently positive for bf16 D128 (halving the KV tile doubles the tile/grid + count and the extra loop/barrier overhead outweighs the LDS relief while occupancy + stays pinned). The opt-in BF16 D128 CFVST path on the persistent grid can + double only K at ``block_n=64`` and/or prefetch next V into registers. Both are + separate opt-in tuning knobs, off by default; the combined path fits and waits + with a full ``vmcnt(0)``, not a partial one. Bounding the + live set with partial ``vmcnt=N`` waits entangles with the K async DMA, which + shares the same counter -- see the reverted cfvst-chunking experiment recorded in + the ``V_ITEMS`` NOTE below. P1 conflict-free V (:func:`_use_cfvst`): V is stored TRANSPOSED via a perm_b32 store-path transpose so the PV A-operand read is a contiguous ds_read_b64 instead of P0's 4 element-wise ds_read_u16 (128 -> ~32 LDS instrs/tile at D128). Gated to -**D128 fp16**: D64 is VGPR-bound so the register round-trip regresses it, and D128 -bf16 spills over the waves-per-eu=2 cap (a P2/P3 register-headroom item). The lever +**D128 fp16** by default: D64 is VGPR-bound so the register round-trip regresses it, +and D128 bf16 CFVST leaves little VGPR headroom under the waves-per-eu=2 cap, so it is +opt-in through ``use_cfvst=True`` (register usage is recorded in the gfx942 CFVST / KDB / +VPF case study). The lever is numerically identity-preserving (same MFMA operands; only the LDS layout and read width change) -- the cohort's ``max_abs`` is unchanged from the naive path. The conflict-free vehicle was proven correct cross-part before being lifted (see the @@ -140,6 +142,13 @@ single LDS buffer) and reads V element-wise; the remaining perf levers (P2-P4) layer on top. It is validated against an fp32 SDPA reference across the in-scope cohort. +Opt-in bf16 D128: ``use_cfvst=True`` enables the V transpose for +causal attention on either grid, including ``causal_diag_split=True``. On the +persistent grid, ``lds_num_buffers=2`` independently uses two K LDS slots and +``v_reg_prefetch=True`` loads next-tile V registers before current PV. V LDS +stays single-buffered for all combinations. All +shipped defaults use ``lds_num_buffers=1`` and ``v_reg_prefetch=False``. + :func:`supports_attention_dense` is the SINGLE gate: it rejects every spec :func:`build_attention_dense` cannot emit -- including the modes deferred to later phases -- so ``supports_attention_dense(spec)[0] is True`` implies the build @@ -219,7 +228,7 @@ LOG2E = 1.4426950408889634 _DTYPE_IR = {"bf16": BF16, "fp16": F16} -# Pipeline constants (mirror gfx950; this body is NBUF=1, a single LDS buffer). +# Pipeline constants (gfx942 BF16 D128 K-buffer count and V prefetch are opt-in). # The shipped geometry starts at 256 query rows per CTA = 8 wave64s. ``block_m`` # is inherited from AttentionDenseSpec, so the grid and body read the same # immutable value and cannot disagree silently (rows written twice/never). @@ -281,7 +290,8 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): --------------------------------------------------------- ``use_cfvst`` / ``use_exp2_fast`` / ``v_row_pad`` / ``use_v_swizzle`` default to ``None``, which resolves through the shipping policy (:func:`_use_cfvst`, - :func:`_use_exp2_fast`, :func:`_v_row_pad`, :func:`_use_v_swizzle`). That + :func:`_use_exp2_fast`, :func:`_v_row_pad`, and the spec's + :meth:`resolved_use_v_swizzle`). That tri-state is load-bearing: a harness that omits a VALUED field freezes the config at whatever the default happened to be the day it was written and then silently reports a stale verdict (that is exactly how a real regression got mis-reported @@ -354,7 +364,7 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): # v_row_pad: V^T row (token axis) width pad, in elements, for the P1 conflict-free-V # store. V_lds is transposed to [D, block_n+pad] so the PV A-operand read is a - # contiguous ds_read_b64. On the cfvst path (fp16-D128) the pad is DERIVED from + # contiguous ds_read_b64. On the cfvst path (D128) the pad is DERIVED from # block_n by :func:`_v_row_pad` so the row is the smallest pow2 the XOR bank- # conflict swizzle needs (:func:`_v_swizzle_width`): pad 64->0, 32->32, 128->0; 8 # (the original unswizzled layout) elsewhere. The pad VALUE is no longer OPEN -- @@ -367,19 +377,18 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): # :func:`_lds_bytes` budget so the two cannot drift. v_row_pad: int | None = None - # use_cfvst: force the P1 conflict-free perm_b32 store-path V transpose on/off. - # None (default) -> :func:`_use_cfvst`, which owns the measured verdict. - # Forcing it ON where the policy says OFF is REJECTED by - # supports_attention_dense (D64 regresses under it -- VGPR-bound; bf16 D128 - # spills past the waves-per-eu cap), so this field can only turn cfvst off, or - # pin it on where it already is. + # use_cfvst: force the P1 conflict-free perm_b32 V transpose on/off. + # None -> the :func:`_use_cfvst` policy (fp16 D128). True is also legal at + # bf16 D128 causal attention (no sliding window or sinks), with or without + # causal_diag_split; False is always legal; D64 rejects True. KDB and VPF + # opt-ins require bf16 D128 CFVST. use_cfvst: bool | None = None # use_v_swizzle: force the V^T XOR bank-conflict swizzle on/off, INDEPENDENTLY of # the pad value -- so the pad-value sweep can vary ``v_row_pad`` without silently # toggling the swizzle (an explicit non-derived pad used to turn it off by making - # the row non-pow2). None (default) -> :func:`_use_v_swizzle` (the cfvst path). - # Only the transposed-V (fp16-D128) store has a V_lds to swizzle; forcing it ON + # the row non-pow2). None (default) -> enabled on the resolved CFVST path. + # Only the transposed-V (D128) store has a V_lds to swizzle; forcing it ON # off that path, or with a pad that leaves V_LDROW non-pow2 / < 64, is REJECTED # by supports_attention_dense rather than silently mis-addressing LDS. use_v_swizzle: bool | None = None @@ -435,10 +444,16 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): # ``lazy_rescale`` field has no effect here: this body always rescales, and # only the shared kernel_name() reads it (the ``lazyrs`` token). # - # lds_num_buffers: K/V LDS buffers. Only 1 is implemented: the tile-start - # vmcnt(0) and the trailing sync_lds_only() WAR guard assume one buffer, and - # NBUF=2 only fits at block_n=32, which was measured negative. + # lds_num_buffers: K LDS buffer count, independent of V register prefetch. + # 1 keeps a single K slot; 2 allows K[j+1] DMA to overlap PV[j] in a + # different slot. V LDS is always single-buffered. The opt-in pipeline + # requires persistent=True and BF16 D128 causal CFVST at block_n=64, + # with or without causal_diag_split. lds_num_buffers: int = field(default=1, kw_only=True) + # v_reg_prefetch: issue the next V tile's register loads before PV[j] + # rather than after PV[j]. Independent of the K LDS buffer count; same + # scope as lds_num_buffers=2. + v_reg_prefetch: bool = field(default=False, kw_only=True) # iglp_mode: the iglp_opt argument emitted when ``iglp`` is on (0 or 1). iglp_mode: int = field(default=_DEFAULT_IGLP_MODE, kw_only=True) # pv_loop_order: PV MFMA traversal. "d_major" keeps the output tile outer, @@ -457,8 +472,8 @@ class Gfx942AttentionDenseSpec(AttentionDenseSpec): pv_sched_fence_mask: int | None = field(default=None, kw_only=True) # causal_diag_split: causal without a sliding window only -- an unmasked KV loop # over below-diagonal tiles plus a masked diagonal tail, instead of masking - # every tile. Off by default: the peel is recorded negative (register - # pressure on the gfx942 tiled_2d kernel); kept to re-measure on dense. + # every tile. The CFVST KDB/VPF pipeline also preserves K/V publication across + # the two loop phases. Off by default; kept to re-measure on dense. causal_diag_split: bool = field(default=False, kw_only=True) # emit_lse: also write the natural-log softmax normalizer (LSE) as FP32 # [B, Hq, Sq, 1] through a trailing ``lse_ptr`` kernel argument. Off by default; @@ -485,13 +500,29 @@ def __post_init__(self) -> None: raise ValueError( f"emit_lse must be a bool, got {type(self.emit_lse).__name__}" ) + # Reject, never coerce: bool("false") is True, so a truthiness read would + # silently build the CFVST kernel the caller asked to turn off. + if self.use_cfvst is not None and type(self.use_cfvst) is not bool: + raise ValueError( + f"use_cfvst must be None or a bool, got {type(self.use_cfvst).__name__}" + ) + if type(self.lds_num_buffers) is not int: + raise ValueError( + "lds_num_buffers must be an int, got " + f"{type(self.lds_num_buffers).__name__}" + ) + if type(self.v_reg_prefetch) is not bool: + raise ValueError( + "v_reg_prefetch must be a bool, got " + f"{type(self.v_reg_prefetch).__name__}" + ) if self.causal_bottom_right: raise ValueError( "gfx942 attention_dense: causal_bottom_right not yet supported" ) def resolved_use_cfvst(self) -> bool: - """Resolved conflict-free-V decision (``None`` -> :func:`_use_cfvst`).""" + """None -> the :func:`_use_cfvst` policy; True also enables bf16 D128.""" if self.use_cfvst is None: return _use_cfvst(self.head_size, self.dtype) return bool(self.use_cfvst) @@ -502,6 +533,11 @@ def resolved_v_row_pad(self) -> int: Read by both the builder's ``V_lds`` alloc and the :func:`_lds_bytes` budget, so the two cannot drift.""" if self.v_row_pad is None: + # The default resolver for every configuration, including explicit + # FP16 use_cfvst=False. Only bf16 D128 with use_cfvst=True needs the + # transposed V-row pitch. + if self.resolved_use_cfvst() and not _use_cfvst(self.head_size, self.dtype): + return _v_swizzle_width(self.block_n) - self.block_n return _v_row_pad(self.head_size, self.dtype, self.block_n) return int(self.v_row_pad) @@ -509,11 +545,12 @@ def resolved_use_v_swizzle(self) -> bool: """Resolved V^T bank-conflict swizzle. Gated on the cfvst path (only the transposed-V store has a V_lds to swizzle), and within it independent of ``v_row_pad`` so the pad sweep never silently toggles it. ``None`` -> policy - (:func:`_use_v_swizzle`).""" + (enabled when resolved CFVST is active).""" if not self.resolved_use_cfvst(): return False if self.use_v_swizzle is None: - return _use_v_swizzle(self.head_size, self.dtype) + # Explicit BF16 CFVST uses the same conflict-free swizzle as FP16. + return True return bool(self.use_v_swizzle) def resolved_use_exp2_fast(self) -> bool: @@ -687,6 +724,7 @@ def _as_gfx942_spec(spec: AttentionDenseSpec) -> Gfx942AttentionDenseSpec: "AttentionDenseSpec", "Gfx942AttentionDenseSpec", "supports_attention_dense", + "bf16_cfvst_supported", "build_attention_dense", "attention_dense_grid", "attention_dense_block", @@ -718,16 +756,22 @@ def _tuning_name_tags(spec: "Gfx942AttentionDenseSpec") -> str: :meth:`Gfx942AttentionDenseSpec.kernel_name` itself. """ parts: list[str] = [] + # Both IR-changing controls have independent cache/name tags. Neither + # changes the default symbol, which has one K buffer and no V prefetch. + if _use_kdb(spec): + parts.append("kdb1") + if _use_vpf(spec): + parts.append("vpf1") if spec.lds_row_pad != _DEFAULT_LDS_ROW_PAD: parts.append(f"krowpad{spec.lds_row_pad}") vp = spec.resolved_v_row_pad() - if vp != _v_row_pad(spec.head_size, spec.dtype, spec.block_n): + if vp != _default_v_row_pad(spec): parts.append(f"vrowpad{vp}") cfvst = spec.resolved_use_cfvst() - if cfvst != _use_cfvst(spec.head_size, spec.dtype): + if cfvst != _default_use_cfvst(spec): parts.append("cfvst1" if cfvst else "cfvst0") swz = spec.resolved_use_v_swizzle() - if swz != (spec.resolved_use_cfvst() and _use_cfvst(spec.head_size, spec.dtype)): + if swz != spec.resolved_use_cfvst(): parts.append("vswz1" if swz else "vswz0") e2f = spec.resolved_use_exp2_fast() if e2f != _use_exp2_fast(spec.head_size, spec.dtype, spec.persistent): @@ -839,15 +883,56 @@ def _use_cfvst(head_size: int, dtype: str) -> bool: (plan §6.1), so the register round-trip + perm temps raise pressure without relieving the actual bottleneck. * D128 bf16: the bf16 ``.1k`` MFMA schedule keeps more registers live, so cfvst - pushes VGPR over the waves-per-eu=2 cap (256) and SPILLS. fp16 does not - (it drops slightly, 176 -> 173). The flash x8 ladder is fp16-first anyway - (tiled_2d keeps bf16 D128 on the narrow path); a bf16 cfvst register-pressure - fix is a P3 occupancy item. + leaves little headroom under the waves-per-eu=2 VGPR cap. It is opt-in through + ``use_cfvst=True``. ``lds_num_buffers=2`` / ``v_reg_prefetch=True`` are + admitted only on the persistent grid; the non-persistent grid spilled with + them and is rejected. Register usage per variant and toolchain is recorded + in the gfx942 CFVST / KDB / VPF case study. The flash x8 ladder is + fp16-first anyway. Gate on rows-per-DMA == 1 (D128) AND fp16 so the LDS budget and the body agree. """ return _rows_per_instr(head_size) == 1 and dtype == "fp16" +def _default_use_cfvst(spec: "Gfx942AttentionDenseSpec") -> bool: + """The :func:`_use_cfvst` default policy, used when comparing kernel-name tags.""" + return _use_cfvst(spec.head_size, spec.dtype) + + +def _default_v_row_pad(spec: "Gfx942AttentionDenseSpec") -> int: + """The default V-row pad, independent of manual override.""" + return _v_row_pad(spec.head_size, spec.dtype, spec.block_n) + + +def _use_kdb(spec: "Gfx942AttentionDenseSpec") -> bool: + """Whether K uses two LDS slots rather than one (opt-in, off by default).""" + return spec.lds_num_buffers == 2 + + +def _use_vpf(spec: "Gfx942AttentionDenseSpec") -> bool: + """Whether to issue next-tile V VMEM loads before the current PV (opt-in).""" + return spec.v_reg_prefetch + + +def bf16_cfvst_supported(spec: Gfx942AttentionDenseSpec) -> bool: + """Whether opt-in CFVST is implemented for this problem at BF16: D128 causal + attention without a sliding window or sinks, on either grid, with or without + ``causal_diag_split``. + + The single definition of the BF16 CFVST cohort: :func:`supports_attention_dense` + gates on it, and the gfx942 dense tuning scope offers the CFVST axes from it, so + the two cannot drift. The KDB/VPF pipeline additionally needs the persistent grid + and ``block_n=64``. + """ + return ( + spec.dtype == "bf16" + and spec.head_size == 128 + and spec.causal + and spec.sliding_window == 0 + and not spec.use_sinks + ) + + def _v_swizzle_width(block_n: int) -> int: """V^T row width (elements) the XOR bank-conflict swizzle needs: the smallest pow2 >= max(64, block_n). SINGLE source of truth -- the derived pad (:func:`_v_row_pad`), @@ -867,12 +952,6 @@ def _v_swizzle_width(block_n: int) -> int: return max(64, 1 << (block_n - 1).bit_length()) -def _use_v_swizzle(head_size: int, dtype: str) -> bool: - """Policy for the V^T bank-conflict swizzle: rides the :func:`_use_cfvst` gate, - since only the transposed-V (fp16-D128) store has a V_lds to swizzle.""" - return _use_cfvst(head_size, dtype) - - def _v_row_pad(head_size: int, dtype: str, block_n: int) -> int: """V^T row (token-axis) pad, in elements, on the cfvst path (fp16-D128); 8 (no swizzle, original layout) everywhere else. On the cfvst path the row is widened to @@ -994,7 +1073,7 @@ def _k_group_stride(spec: AttentionDenseSpec) -> int: def _lds_bytes(spec: "Gfx942AttentionDenseSpec") -> int: - """Total LDS footprint: ``K_lds[1, block_n, row_stride] + V_lds[1, D, block_n+pad]``. + """LDS footprint: one or two K buffers and one V buffer. K keeps the natural ``[token, dim]`` layout with :func:`_lds_row_stride` (async DMA target). V is TRANSPOSED to ``[dim, token]`` for the P1 conflict-free store, so @@ -1011,13 +1090,14 @@ def _lds_bytes(spec: "Gfx942AttentionDenseSpec") -> int: ``resolved_v_row_pad`` is read HERE and in the builder's ``V_lds`` allocation; both sites go through the same resolver, so the budget cannot silently under-count the - allocation. + allocation. ``v_reg_prefetch`` changes VGPR scheduling only, not LDS size. """ if _k_group_pad_active(spec): rpi = _rows_per_instr(spec.head_size) k_bytes = (spec.block_n // rpi) * _k_group_stride(spec) * 2 else: k_bytes = spec.block_n * _lds_row_stride(spec) * 2 + k_bytes *= spec.lds_num_buffers # Only K is double-buffered. if spec.resolved_use_cfvst(): # V transposed to [dim, token+pad] for the conflict-free store (D128). v_bytes = spec.head_size * (spec.block_n + spec.resolved_v_row_pad()) * 2 @@ -1158,18 +1238,21 @@ def supports_attention_dense( _wpe = spec.resolved_waves_per_eu() if _wpe <= 0: return False, f"waves_per_eu must be positive, got {_wpe}" - # cfvst forced ON where the policy says OFF. Both cases the policy excludes are - # known-bad, not merely untuned: D64 (rows-per-DMA > 1) is VGPR-bound, so the - # register round-trip regresses it AND the transposed store no longer matches the - # packed async-DMA layout the naive path relies on; bf16 D128 spills past the - # waves-per-eu cap on the .1k schedule. Turning cfvst OFF is always legal (it is - # the naive path), so only the ON direction is gated. See _use_cfvst. - if spec.resolved_use_cfvst() and not _use_cfvst(spec.head_size, spec.dtype): + # BF16 CFVST supports causal diagonal splitting as well as the unsplit path. + # Preserve all existing FP16 CFVST configurations. + if spec.dtype == "bf16" and spec.resolved_use_cfvst(): + if not bf16_cfvst_supported(spec): + return False, ( + "NOT_YET_IMPLEMENTED: BF16 CFVST requires D128 causal attention " + "without sliding_window or sinks" + ) + # CFVST is supported for D128 (bf16 and fp16). D64 uses packed DMA rows + # and remains on the naive-V path; explicit CFVST there is rejected. + if spec.resolved_use_cfvst() and _rows_per_instr(spec.head_size) != 1: return False, ( f"use_cfvst=True is rejected at D{spec.head_size}/{spec.dtype}: " - f"the conflict-free-V store is measured-negative there (D64 is VGPR-bound " - f"and regresses; bf16 D128 spills past the waves-per-eu cap), and it is " - f"only tile-exact on the rows-per-DMA==1 fp16 path -- see _use_cfvst" + f"the conflict-free-V store is supported only for rows-per-DMA==1 " + f"(D128); D64 uses the naive-V path -- see _use_cfvst" ) # use_v_swizzle forced ON where there is no transposed-V buffer to swizzle. Only the # cfvst path builds V_lds[dim, token]; on the naive path the flag is a no-op, so @@ -1178,13 +1261,23 @@ def supports_attention_dense( return False, ( f"use_v_swizzle=True is rejected at D{spec.head_size}/{spec.dtype}: " f"the XOR bank-conflict swizzle only exists on the conflict-free-V " - f"(fp16-D128) path -- there is no transposed V_lds to swizzle otherwise" + f"(D128) path -- there is no transposed V_lds to swizzle otherwise" ) # Performance-only codegen knobs: ranges, plus the one scheduling conflict. - if spec.lds_num_buffers != 1: + if spec.lds_num_buffers not in (1, 2): + return False, "lds_num_buffers must be 1 or 2" + # These are independent opt-in scheduling levers, but both currently share + # the same validated BF16 D128 CFVST scope. In particular, block_n=64 only: + # the other tile sizes have not been validated for these pipeline modes. + if (_use_kdb(spec) or _use_vpf(spec)) and not ( + spec.persistent + and bf16_cfvst_supported(spec) + and spec.resolved_use_cfvst() + and spec.block_n == 64 + ): return False, ( - f"lds_num_buffers={spec.lds_num_buffers}: only 1 is implemented (the " - f"tile-start wait and trailing WAR barrier assume one K/V buffer)" + "NOT_YET_IMPLEMENTED: lds_num_buffers=2 or v_reg_prefetch=True requires " + "persistent=True, bf16 D128 causal CFVST at block_n=64, without sliding_window or sinks" ) if spec.iglp_mode not in (0, 1): return False, f"iglp_mode must be 0 or 1, got {spec.iglp_mode}" @@ -1304,11 +1397,11 @@ def build_attention_dense( ) -> KernelDef: """Emit the gfx942 dense flash-attention prefill kernel for ``spec``. - Delegates to :func:`_build_attention_dense_single_buffer`, the single-LDS-buffer - (NBUF=1) body (32x32x8 atom + K-loop doubling). NBUF=1 is the only axis that name - describes: the body carries the full shipped lever set (cfvst, exp2_fast, fused - rescale, per-config waves-per-eu, persistent grid-stride), so it is not a - correctness-first stand-in for a later tuned body. Every scope restriction -- + Delegates to :func:`_build_attention_dense_single_buffer`, which builds the + default single-buffer body and the independently selected BF16 D128 K-double- + buffer / V-register-prefetch variants. It retains CFVST, exp2_fast, fused + rescale, per-config waves-per-eu and persistent grid-stride. Every scope + restriction -- including the modes deferred to later follow-ups -- lives in :func:`supports_attention_dense`, which is consulted here, so this function is a thin gate-plus-delegate and cannot reject a spec that ``supports`` accepted. @@ -1367,11 +1460,14 @@ def _build_attention_dense_single_buffer( ) -> KernelDef: """gfx942 dense prefill body: 32x32x8 atom + K-loop doubling. - Non-pipelined (NBUF=1): per KV tile load→wait→QK→mask→softmax→PV, single LDS - buffer, simple online-softmax rescale. The per-work-item compute is factored into + The opt-in BF16 D128 persistent-grid pipeline exposes K LDS double + buffering (lds_num_buffers=2) and V register prefetch (v_reg_prefetch=True) + as independent tuning knobs, off by default. At (1, False) the original single-buffer + body is unchanged. Each tile computes QK then softmax/PV. + The per-work-item compute is factored into ``_run_work_item(qb, hq, bt)``, shared by the default grid (one CTA per work item) and the P4 persistent grid-stride path (one CTA strides over many). Perf levers - are all on: conflict-free V/cfvst (P1, D128 fp16), exp2_fast + fused rescale (P2), + are all on: conflict-free V/cfvst (P1, D128), exp2_fast + fused rescale (P2), per-config waves-per-eu (P3), persistent grid-stride + qb/hkv-major decode (P4). The transposed-QK architecture (S^T=K@Q^T so P feeds PV lane-locally) is kept from the gfx950 sibling — only the MFMA atom and K/V fragment widths change. @@ -1478,7 +1574,13 @@ def _stride_k(): qk_scale = b.fmul(scale, b.const_f32(LOG2E)) tid = b.thread_id_x() - wave = b.div(tid, b.const_i32(64)) + # In the CFVST KDB/VPF pipeline, the wave index is uniform across a wave. + # Scalarize at full EXEC before /64 so subsequent V buffer offsets can + # stay on the SGPR/SALU path; leave all nonpipeline IR unchanged. + if _use_kdb(spec) or _use_vpf(spec): + wave = b.div(b.to_sgpr_u32(tid), b.const_i32(64)) + else: + 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)) @@ -1493,7 +1595,11 @@ def _stride_k(): # LDS-budget check in supports_attention_dense so the two cannot drift. The one # exception is when _k_group_pad_active(spec): D64 K_lds then takes the padded # [1, block_n // rows_per_instr, _k_group_stride] 2-row-group layout below. - USE_CFVST = spec.resolved_use_cfvst() # P1 conflict-free V: D128 fp16 only + USE_CFVST = spec.resolved_use_cfvst() # P1 V transpose at D128. + USE_KDB = _use_kdb(spec) + USE_VPF = _use_vpf(spec) + USE_PIPELINE = USE_KDB or USE_VPF + K_BUFS = spec.lds_num_buffers LDROW = _lds_row_stride(spec) # Hypothesis #3 (D64 K-LDS bank conflicts): K_lds with a row-group boundary pad, # sized by the shared spec.lds_k_group_pad (0 disables it -- that is the A/B probe). @@ -1505,7 +1611,7 @@ def _stride_k(): dtype, [1, BN // K_GROUP, K_GROUP_STRIDE], name_hint="Klds" ) else: - K_lds = b.smem_alloc(dtype, [1, BN, LDROW], name_hint="Klds") + K_lds = b.smem_alloc(dtype, [K_BUFS, BN, LDROW], name_hint="Klds") if USE_CFVST: # V stored TRANSPOSED (P1 conflict-free V): [dim, token] with token inner and # padded by resolved_v_row_pad, so the PV A-operand read is a contiguous @@ -1670,8 +1776,7 @@ def read_v(dt, kk): conflicts live on the STORE (8-way -> 4-way: within a wave t0 is constant and d0 = 2*lane, so (d0 & mask) takes only 8 values), so further conflict work belongs there, not on this read. - naive (D64 / bf16-D128): element-wise from V_lds[key, dim] (bank-heavy, but - those configs are VGPR-bound / spill under cfvst, not LDS-read-bound).""" + naive (D64): element-wise from V_lds[key, dim].""" if USE_CFVST: dim_row = b.add(b.const_i32(dt * 32), lane_m) key0 = b.add(b.const_i32(kk * 8), d_base) @@ -1746,11 +1851,20 @@ def _run_work_item(qb, hq, bt, hkv=None): addr = b.add(b.add(q_base, b.mul(q_tok, _stride_q())), col) q_packs.append(b.global_load_vN(q, addr, dtype, 4, align=8)) - def _async_load(rsrc, lds_base, tile_key0, group_pad=False): + def _async_load(rsrc, lds_base, tile_key0, group_pad=False, kbuf=None): + # Only the pipelined K path passes kbuf; the default emission is + # unchanged (no extra IR operations and the same LDS addresses). + buf_off = ( + b.zext(b.mul(kbuf, b.const_i32(BN * LDROW * 2)), I64) + if kbuf is not None + else None + ) 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.zext(b.mul(row, b.const_i32(K_LDROW_BYTES)), I64) + if buf_off is not None: + row_lds_off = b.add(buf_off, row_lds_off) 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)) @@ -1804,6 +1918,11 @@ def load_tile(tile_idx): if not USE_CFVST: _async_load(v_rsrc, V_lds_addr, tk0) + def load_pipeline_k(kbuf, tile_idx): + """Issue K(tile_idx) DMA into the selected K LDS slot.""" + tk0 = b.mul(tile_idx, b.const_i32(BN)) + _async_load(k_rsrc, K_lds_addr, tk0, kbuf=kbuf) + def _cfvst_load_v(tile_idx): """Issue the cfvst VMEM loads and keep each thread's V tile in VGPRs. @@ -1835,7 +1954,49 @@ def _cfvst_load_v(tile_idx): payload.append((d0, t0, b.bitcast(x0, I32), b.bitcast(x1, I32))) return payload - def do_qk(): + def _cfvst_load_v_pipeline(tile_idx): + """Load packed V with one SGPR tile/wave base and a VGPR lane offset. + + For BF16 D128 / BN64, V_DIM_PAIRS=64 and + blk = it * (WAVES * 64) + tid, + d0 = 2 * lane, + t0 = 2 * wave + 2 * it * WAVES. + The original unrolled loader rebuilt per-thread byte addresses and + div/mod coordinates for each item. Use buffer soffset for the uniform + tile/wave part and buffer voffset for the reusable lane offset. + Reconstruct V^T publication coordinates after PV rather than keeping + them live across the prefetch window. + """ + if D != 128 or V_DIM_PAIRS != 64: + raise ValueError("pipeline V SGPR addressing requires D=128") + wave_key0 = b.add( + b.mul(tile_idx, b.const_i32(BN)), + b.mul(wave, b.const_i32(2)), + ) + # Do the uniform-to-scalar conversion just once per tile. + soff0 = b.to_sgpr_u32( + b.mul(b.add(k_base, b.mul(wave_key0, _stride_k())), b.const_i32(2)) + ) + stride_bytes = b.mul(_stride_k(), b.const_i32(2)) + lane_byte = b.mul(lane, b.const_i32(4)) + values = [] + for it in range(V_ITEMS): + soff_it = ( + soff0 + if it == 0 + else b.add( + soff0, + b.mul(stride_bytes, b.const_i32(it * 2 * WAVES)), + ) + ) + x0 = b.buffer_load_vN(v_rsrc, lane_byte, soff_it, dtype, 2) + x1 = b.buffer_load_vN( + v_rsrc, lane_byte, b.add(soff_it, stride_bytes), dtype, 2 + ) + values.append((b.bitcast(x0, I32), b.bitcast(x1, I32))) + return values + + def do_qk(kbuf=None): """S^T = K@Q^T via the 32x32x8 atom (K-doubled).""" s_reg = [] for nsub in range(N_SUB): @@ -1853,7 +2014,7 @@ def do_qk(): if KPAD_D64: k_pack = b.smem_load_vN( K_lds, - b.const_i32(0), + b.const_i32(0) if kbuf is None else kbuf, k_group, b.add(k_within_off, col), dtype=dtype, @@ -1861,7 +2022,12 @@ def do_qk(): ) else: k_pack = b.smem_load_vN( - K_lds, b.const_i32(0), krow, col, dtype=dtype, n=4 + K_lds, + b.const_i32(0) if kbuf is None else kbuf, + krow, + col, + dtype=dtype, + n=4, ) acc = mfma_32x32x8_for_dtype(b, dtype, k_pack, q_packs[ks], acc) # Into the log2 domain in fp32, so masking, the row max and @@ -1956,45 +2122,11 @@ def do_mask(s_reg, tile_idx, lower=False, upper=True): (f"o{dt}", o0[dt]) for dt in range(D_TILES) ] - def tile_body(j, carry, masked): - m_i = carry[0] - l_i = carry[1] - o_acc = list(carry[2 : 2 + D_TILES]) - - if spec.iglp: - # Runbook lever 7: one canned-scheduler hint at the loop-body top. - # Only meaningful on the cfvst path (in-loop ds_write to interleave). - b.iglp_opt(spec.iglp_mode) - - if USE_CFVST: - load_tile(j) # K async DMA -> K_lds - v_payload = _cfvst_load_v( - j - ) # V register buffer_loads (natural [token,dim]) - # Drain BOTH the K async DMA and the V register loads (all vmcnt) - # before the V perm+store reads those registers and before QK reads - # K_lds. - b.s_waitcnt(vmcnt=0) - _cfvst_store_v(v_payload) # perm_b32 -> ds_write V_lds[dim,token] - # Publish K_lds (DMA) AND V_lds (ds_write) to all waves. This MUST - # drain lgkmcnt for the V ds_write and cannot be a bare s_barrier: - # gfx90a+ set FeatureBackOffBarrier, so LLVM would sink the lgkm wait - # past the barrier to the ds_read consumer (the P0 read-after-barrier - # race, now on the tile-START barrier because V publication is an - # in-loop ds_write). sync_lds_only() emits the lgkm drain BEFORE the - # barrier (CK Tile block_sync_lds). vmcnt is already 0, no full sync. - b.sync_lds_only() - else: - # Naive (D64): K and V both async-DMA'd into their natural layout. - # Drain the DMA (vmcnt) then a plain barrier -- V is not an in-loop - # ds_write here, so the tile-start barrier needs no lgkm drain (the - # DMA landed the data; the trailing sync_lds_only guards the - # read-before-overwrite). - load_tile(j) - b.s_waitcnt(vmcnt=0) - b.s_barrier_bare() - - s = do_qk() + def softmax_pv(s, j, m_i, l_i, o_acc, *, masked, lower): + """One KV tile's mask -> online softmax -> PV update, shared by the + single-buffer ``tile_body`` and the K-double-buffer pipeline so both + honour the same exp2, fused-rescale and PV scheduling knobs. Returns + the updated ``(m, l, o_acc)``; the caller owns loads and barriers.""" # Every visited tile gets the causal upper bound; W > 0 additionally # applies the window lower bound (redundant no-op compares on tiles # fully inside the window). W == 0 -> lower=False, byte-identical. @@ -2003,7 +2135,7 @@ def tile_body(j, carry, masked): # signature parity with the gfx950 do_mask, which does use upper=False. # masked=False only for below-diagonal tiles of the diagonal split. if masked: - do_mask(s, j, lower=(SW > 0), upper=True) + do_mask(s, j, lower=lower, upper=True) # tile max over keys (both lane-halves) for this query. local_max = neg_inf @@ -2064,6 +2196,50 @@ def tile_body(j, carry, masked): o_acc = do_pv(o_acc, p_packs) if PV_PRIORITY: b.s_setprio(0) + return m_new, l_new, o_acc + + def tile_body(j, carry, masked): + m_i = carry[0] + l_i = carry[1] + o_acc = list(carry[2 : 2 + D_TILES]) + + if spec.iglp: + # Runbook lever 7: one canned-scheduler hint at the loop-body top. + # Only meaningful on the cfvst path (in-loop ds_write to interleave). + b.iglp_opt(spec.iglp_mode) + + if USE_CFVST: + load_tile(j) # K async DMA -> K_lds + v_payload = _cfvst_load_v( + j + ) # V register buffer_loads (natural [token,dim]) + # Drain BOTH the K async DMA and the V register loads (all vmcnt) + # before the V perm+store reads those registers and before QK reads + # K_lds. + b.s_waitcnt(vmcnt=0) + _cfvst_store_v(v_payload) # perm_b32 -> ds_write V_lds[dim,token] + # Publish K_lds (DMA) AND V_lds (ds_write) to all waves. This MUST + # drain lgkmcnt for the V ds_write and cannot be a bare s_barrier: + # gfx90a+ set FeatureBackOffBarrier, so LLVM would sink the lgkm wait + # past the barrier to the ds_read consumer (the P0 read-after-barrier + # race, now on the tile-START barrier because V publication is an + # in-loop ds_write). sync_lds_only() emits the lgkm drain BEFORE the + # barrier (CK Tile block_sync_lds). vmcnt is already 0, no full sync. + b.sync_lds_only() + else: + # Naive (D64): K and V both async-DMA'd into their natural layout. + # Drain the DMA (vmcnt) then a plain barrier -- V is not an in-loop + # ds_write here, so the tile-start barrier needs no lgkm drain (the + # DMA landed the data; the trailing sync_lds_only guards the + # read-before-overwrite). + load_tile(j) + b.s_waitcnt(vmcnt=0) + b.s_barrier_bare() + + s = do_qk() + m_new, l_new, o_acc = softmax_pv( + s, j, m_i, l_i, o_acc, masked=masked, lower=(SW > 0) + ) # Tile done; the next iteration refills the SINGLE K/V buffer (K via async # DMA, V via the perm_b32 ds_write), so every wave's LDS reads must have # LANDED (not just issued) before any wave starts writing. This MUST drain @@ -2092,7 +2268,142 @@ def tile_body(j, carry, masked): # hardcoded to "tile.sync" and unroll defaults False. Pin it so a future # unroll=True (P3) or a sync_lds_only->sync swap cannot silently delete it. # Verified byte-identical codegen with and without the flag today. - if causal and not SW and DIAG_SPLIT: + if USE_PIPELINE: + # tile_body serves the shipped single-buffer configurations unchanged. + # This BF16 CFVST pipeline independently schedules K[j+1] DMA and + # V[j+1] register loads across softmax/PV[j]. KDB overlaps its next + # DMA only with a distinct K slot; VPF overlaps V loads in registers. + def publish_prefetched_v(values): + # Prefetch carries only packed V data across softmax/PV, not the + # transposition coordinates. Recompute coordinates at publication. + payload = [] + for it, (x0, x1) in enumerate(values): + # blk=it*THREADS+tid, with THREADS=64*WAVES and + # V_DIM_PAIRS=64. Unroll-time offsets eliminate div/mod. + t0 = b.add( + b.mul(wave, b.const_i32(2)), + b.const_i32(it * 2 * WAVES), + ) + d0 = b.mul(lane, b.const_i32(2)) + payload.append((d0, t0, x0, x1)) + _cfvst_store_v(payload) + + # Prologue: prepare K[0] and publish V[0]. All waves rendezvous + # after draining their VMEM and after V's LDS writes complete. + zero = b.const_i32(0) + load_pipeline_k(zero, zero) + v0 = _cfvst_load_v_pipeline(zero) + b.s_waitcnt(vmcnt=0) + publish_prefetched_v(v0) + b.sync_lds_only() + + # Main loop deliberately excludes the final KV tile. n_up=1 makes + # this a zero-trip loop; its prologue and epilogue still execute. + last_j = b.sub(n_up, b.const_i32(1)) + + def pipeline_step(j, carry, *, masked): + """Compute tile j and publish tile j+1, retaining the LDS + rendezvous and pipeline scheduling across split-loop boundaries.""" + m_i = carry[0] + l_i = carry[1] + o_acc = list(carry[2 : 2 + D_TILES]) + if spec.iglp: + b.iglp_opt(spec.iglp_mode) + + # The current K and V tiles are fully published by the prologue + # or the preceding iteration. With one K slot, do not overwrite + # it until *all* waves have completed current-tile LDS reads. + kbuf = b.mod(j, b.const_i32(K_BUFS)) + s = do_qk(kbuf) + next_j = b.add(j, b.const_i32(1)) + next_kbuf = b.mod(next_j, b.const_i32(K_BUFS)) + if USE_KDB: + # Distinct LDS slot: K[j+1] DMA overlaps current softmax/PV. + load_pipeline_k(next_kbuf, next_j) + if USE_VPF: + # Only V data (not addresses) stays live across softmax/PV. + v_next = _cfvst_load_v_pipeline(next_j) + m_new, l_new, o_acc = softmax_pv( + s, j, m_i, l_i, o_acc, masked=masked, lower=False + ) + + # WAR: finish all waves' K/V LDS reads before a single K slot + # or the (always single-buffered) V slot can be overwritten. + # LDS-only sync leaves early K DMA / V VMEM loads in flight. + b.sync_lds_only() + if not USE_KDB: + # Single K slot: only overwrite it after the cross-wave WAR. + load_pipeline_k(next_kbuf, next_j) + if not USE_VPF: + # No V prefetch: issue V[j+1] loads only after current PV. + v_next = _cfvst_load_v_pipeline(next_j) + # Drain next K async-DMA and next V VMEM loads before V LDS + # publication; then rendezvous before any next-tile QK/PV. + b.s_waitcnt(vmcnt=0) + publish_prefetched_v(v_next) + b.sync_lds_only() + b.scf_yield(m_new, l_new, *o_acc) + + if causal and DIAG_SPLIT: + # At kt < qb * (BLOCK_M/BN), every key precedes every query + # row, so the causal upper mask is unnecessary. Split only the + # processing loop (which excludes last_j). It continuously + # publishes K/V tile j+1 across the loop transition, and the + # final tile remains safely masked even when it is below the + # diagonal (e.g. seqlen_kv < seqlen_q). + diag_start = b.mul(qb, b.const_i32(n_per)) + body_upper = b.select(b.cmp_lt(diag_start, last_j), diag_start, last_j) + body = b.scf_for_iter( + zero, + body_upper, + b.const_i32(1), + iter_args, + iv_name="kbp", + elide_trailing_barrier=False, + ) + with body as (j, carry): + pipeline_step(j, carry, masked=False) + tail_args = [ + (name + "_t", val) + for (name, _), val in zip(iter_args, body.results) + ] + loop = b.scf_for_iter( + body_upper, + last_j, + b.const_i32(1), + tail_args, + iv_name="ktp", + elide_trailing_barrier=False, + ) + with loop as (j, carry): + pipeline_step(j, carry, masked=True) + else: + # Preserve the non-split pipeline's original emitted operations. + loop = b.scf_for_iter( + zero, + last_j, + b.const_i32(1), + iter_args, + iv_name="ktp", + elide_trailing_barrier=False, + ) + with loop as (j, carry): + pipeline_step(j, carry, masked=True) + + # Epilogue: do not issue K/V loads beyond n_up-1. The last LDS + # barrier also protects CTA-persistent storage reuse by the next item. + loop_res = loop.results + m_i = loop_res[0] + l_i = loop_res[1] + o_acc = list(loop_res[2 : 2 + D_TILES]) + last_kbuf = b.mod(last_j, b.const_i32(K_BUFS)) + s = do_qk(last_kbuf) + m_i, l_i, o_acc = softmax_pv( + s, last_j, m_i, l_i, o_acc, masked=True, lower=False + ) + b.sync_lds_only() + + elif causal and not SW and DIAG_SPLIT: # Tiles below this query block's diagonal (kt < qb*n_per) hold only keys # <= every query row of the block, so the causal mask is a no-op there; # only [qb*n_per, n_up) is masked. The single K/V buffer spans both @@ -2134,9 +2445,13 @@ def tile_body(j, carry, masked): with loop as (j, carry): tile_body(j, carry, masked=True) - res = loop.results - l_i = res[1] - o_acc = list(res[2 : 2 + D_TILES]) + # Both branches leave the FINAL running max in m_i: the pipeline folds its + # last tile in after the loop, so loop.results[0] is stale on that path. + if not USE_PIPELINE: + res = loop.results + m_i = res[0] + l_i = res[1] + o_acc = list(res[2 : 2 + D_TILES]) # Epilogue: O[query,dim] = (P@V)/l, from the transposed C[dim,query] accum. if spec.emit_lse: @@ -2185,7 +2500,7 @@ def tile_body(j, carry, masked): ) # Both half-waves hold the same query row; one of them writes its LSE. with b.scf_if(b.cmp_eq(lane_h, b.const_i32(0))): - lse_store(b, lse, lse_index, res[0], l_i, row_has_mass=row_has_mass) + lse_store(b, lse, lse_index, m_i, l_i, row_has_mass=row_has_mass) _emit_o_store( b, o, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py index f8bae6411308..df6aac08577d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tuning_space.py @@ -7,6 +7,7 @@ import dataclasses import unittest from dataclasses import replace +from unittest import mock from itertools import islice from dispatch.attention import ( @@ -35,7 +36,10 @@ tuning_axes, ) from dispatch.attention.common import _problem -from dispatch.attention.dense_rules import resolve_dense_num_persistent +from dispatch.attention.dense_rules import ( + _dense_axes_for, + resolve_dense_num_persistent, +) from dispatch.attention.gfx950_unified import GFX950_TUNING_VARIANTS from dispatch.attention.tuning_specs import _SEMANTIC_FIELDS from dispatch.attention.unified_rules import ( @@ -635,7 +639,7 @@ def test_lds_model_is_the_lowered_pool(self): ) pool = sum( int(n) - for n in re.findall(r"addrspace\(3\) global \[(\d+) x i8\]", ir) + for n in re.findall(r"addrspace\(3\)\s+global\s+\[(\d+) x i8\]", ir) ) self.assertEqual(_gfx950_2d_lds_bytes(spec.kernel_spec), pool) @@ -1010,10 +1014,9 @@ def test_untunable_knobs_have_nothing_to_sweep(self): from rocke.core.lower_llvm import lower_kernel_to_llvm from kernels.gfx942.attention_dense import build_attention_dense + # These pipeline knobs are tunable now, not classified as untunable. + self.assertEqual(DENSE_UNTUNABLE_KNOBS["gfx942"], frozenset({"lazy_rescale"})) base = gfx942_dense_spec(_gfx942_request()) - ok, why = supports_gfx942(replace(base, lds_num_buffers=2), arch="gfx942") - self.assertFalse(ok) - self.assertIn("lds_num_buffers", why) def ir(spec): kernel = build_attention_dense(spec, arch="gfx942") @@ -1032,11 +1035,277 @@ def test_production_varies_every_declared_knob(self): _gfx942_request( dtype="bf16", hdim_q=64, hdim_v=64, nhead_q=8, nhead_k=1 ), + _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128), ) ) missing = _dense_swept_fields("gfx942") - changed self.assertFalse(missing, sorted(missing)) + def test_bf16_cfvst_kdb_vpf_sweep_and_explicit_selection(self): + req = _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128) + default = attention_tuning_spec(req, "gfx942_dense") + base = default.kernel_spec + + # New tuning choices must not alter the default kernel spec. + self.assertIsNone(base.use_cfvst) + self.assertFalse(base.resolved_use_cfvst()) + self.assertEqual(base.lds_num_buffers, 1) + self.assertFalse(base.v_reg_prefetch) + + # Production-level tuning stays one knob at a time, pairing CFVST + # with a dependent knob when the dependent knob alone is illegal. + specs = _gfx942_specs(req, level="production") + modes = { + ( + spec.kernel_spec.resolved_use_cfvst(), + spec.kernel_spec.lds_num_buffers, + spec.kernel_spec.v_reg_prefetch, + ) + for spec in specs + } + for expected in ( + (False, 1, False), # default + (True, 1, False), # CFVST only + (True, 2, False), # CFVST + KDB + (True, 1, True), # CFVST + VPF + ): + self.assertIn(expected, modes) + self.assertNotIn((True, 2, True), modes) + + # Explicit tuning pins can opt into every legal independent mode, + # including the combined mode outside the curated production sweep. + choices = ( + {"use_cfvst": True}, + {"use_cfvst": True, "lds_num_buffers": 2}, + {"use_cfvst": True, "v_reg_prefetch": True}, + { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + }, + ) + for knobs in choices: + with self.subTest(knobs=knobs): + selected = tuning_spec_with_knobs(req, "gfx942_dense", knobs) + spec = selected.kernel_spec + self.assertTrue(spec.resolved_use_cfvst()) + self.assertEqual(spec.lds_num_buffers, knobs.get("lds_num_buffers", 1)) + self.assertEqual( + spec.v_reg_prefetch, knobs.get("v_reg_prefetch", False) + ) + # Bare IDs are searchable only in the production sweep. + # Full-only IDs need their explicit knobs for replay. + if not (spec.lds_num_buffers == 2 and spec.v_reg_prefetch): + replayed = attention_tuning_spec( + req, "gfx942_dense", selected.tuning_id + ) + self.assertEqual(replayed, selected) + + fp16_req = _gfx942_request(dtype="fp16") + for knob in ({"lds_num_buffers": 2}, {"v_reg_prefetch": True}): + with self.subTest(rejected=knob): + with self.assertRaises(ValueError): + tuning_spec_with_knobs(fp16_req, "gfx942_dense", knob) + + def test_ids_recorded_before_v_reg_prefetch_are_refused(self): + """The tuning-id migration. Adding the defaulted v_reg_prefetch field + (default False; lds_num_buffers keeps its default 1) adds one entry to the + gfx942 dense defaults fingerprint, so every gfx942_dense config_key changes. + An id recorded before the field existed is refused, never silently rebuilt, + and its knobs re-canonicalize to the current id.""" + from dispatch.attention import dense_rules, gfx942_dense + + req = _gfx942_request(dtype="bf16", seqlen_q=4096, seqlen_k=4096) + current = gfx942_dense._make_gfx942_attention_dense_candidate() + shipped = current.select_spec( + replace(req, algorithm=current.algorithm, spec_id="gfx942_dense") + ) + self.assertEqual(shipped.kernel_spec.lds_num_buffers, 1) + self.assertFalse(shipped.kernel_spec.v_reg_prefetch) + + real_defaults = dense_rules._dense_defaults + + def defaults_without_v_reg_prefetch(base, recorded): + return tuple( + (name, value) + for name, value in real_defaults(base, recorded) + if name != "v_reg_prefetch" + ) + + for knobs in ({}, {"pv_priority": 2}): + with self.subTest(knobs=knobs): + # Record the id as the fingerprint stood before the field existed. + with mock.patch.object( + dense_rules, "_dense_defaults", defaults_without_v_reg_prefetch + ): + old = gfx942_dense._make_gfx942_attention_dense_candidate() + pin = replace( + req, + algorithm=old.algorithm, + spec_id="gfx942_dense", + tuning_knobs=knobs, + ) + stored = old.select_spec(pin) + fresh = tuning_spec_with_knobs(req, "gfx942_dense", knobs) + self.assertNotEqual(stored.config_key, fresh.config_key) + self.assertEqual(stored.kernel_spec, fresh.kernel_spec) + pin = replace( + pin, tuning_id=stored.tuning_id, tuning_knobs=stored.knobs + ) + ok, why = current.admits(pin) + self.assertFalse(ok) + self.assertIn("defaults they are relative to", why) + self.assertIn(stored.tuning_id, why) + self.assertIn(fresh.tuning_id, why) + + def test_diag_split_keeps_bf16_cfvst_axes_reachable(self): + """A diagonal-split base keeps the CFVST axes the kernel accepts.""" + req = _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128) + base = attention_tuning_spec(req, "gfx942_dense").kernel_spec + for persistent in (False, True): + with self.subTest(persistent=persistent): + split = replace(base, persistent=persistent, causal_diag_split=True) + axes = {axis.name for axis in _dense_axes_for(split, "gfx942")} + self.assertTrue({"use_cfvst", "use_v_swizzle", "v_row_pad"} <= axes) + selected = tuning_spec_with_knobs( + req, + "gfx942_dense", + { + "persistent": persistent, + "causal_diag_split": True, + "use_cfvst": True, + }, + ) + ok, why = supports_gfx942(selected.kernel_spec, arch="gfx942") + self.assertTrue(ok, why) + self.assertTrue(selected.kernel_spec.causal_diag_split) + self.assertTrue(selected.kernel_spec.resolved_use_cfvst()) + + def test_nonpersistent_bf16_cfvst_remains_a_legal_explicit_pin(self): + """The new persistent gate applies to KDB/VPF, not CFVST on its own.""" + req = _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128) + choice = tuning_spec_with_knobs( + req, "gfx942_dense", {"persistent": False, "use_cfvst": True} + ) + kernel = choice.kernel_spec + self.assertFalse(kernel.persistent) + self.assertTrue(kernel.resolved_use_cfvst()) + self.assertEqual(kernel.lds_num_buffers, 1) + self.assertFalse(kernel.v_reg_prefetch) + ok, why = supports_gfx942(kernel, arch="gfx942") + self.assertTrue(ok, why) + + def test_pipeline_tuning_pins_reject_out_of_scope_modes(self): + """Tuning must not advertise configurations that the kernel gate rejects. + + Exercise the real public pinning path, including combinations with the + prerequisite CFVST explicitly enabled. Nonpersistent CFVST alone is + tested positively in a separate test. + """ + bf16 = _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128) + fp16 = _gfx942_request(dtype="fp16", hdim_q=128, hdim_v=128) + modes = ( + {"lds_num_buffers": 2}, + {"v_reg_prefetch": True}, + {"lds_num_buffers": 2, "v_reg_prefetch": True}, + ) + invalid = ( + [ + # BF16 D128: the prerequisite CFVST and persistent grid are both required. + (bf16, {"persistent": False, "use_cfvst": True, **mode}) + for mode in modes + ] + + [(bf16, mode) for mode in modes] # Missing BF16 CFVST. + + [(fp16, {"use_cfvst": True, **mode}) for mode in modes] + + [(bf16, {"use_cfvst": True, "block_n": 32, **mode}) for mode in modes] + + [ + ( + _gfx942_request( + dtype="bf16", hdim_q=128, hdim_v=128, sliding_window=128 + ), + {"use_cfvst": True, **mode}, + ) + for mode in modes + ] + + [ + ( + _gfx942_request( + dtype="bf16", hdim_q=128, hdim_v=128, use_sinks=True + ), + {"use_cfvst": True, **mode}, + ) + for mode in modes + ] + + [ + ( + _gfx942_request(dtype="bf16", hdim_q=128, hdim_v=128, mask_type=0), + {"use_cfvst": True, **mode}, + ) + for mode in modes + ] + ) + for req, knobs in invalid: + with self.subTest(dtype=req.dtype, mask=req.mask_type, knobs=knobs): + with self.assertRaises(ValueError): + tuning_spec_with_knobs(req, "gfx942_dense", knobs) + + def test_full_pipeline_pad4_dsplit_explicit_pin_replays(self): + """The measured full optimization is selectable without hand-building a spec. + + The curated production sweep intentionally excludes the combined mode; + full combination pins still need a stable selection/replay identity. + """ + req = _gfx942_request( + dtype="bf16", + hdim_q=128, + hdim_v=128, + seqlen_q=4096, + seqlen_k=4096, + batch=1, + nhead_q=32, + nhead_k=8, + ) + knobs = { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + "lds_row_pad": 4, + "causal_diag_split": True, + } + selected = tuning_spec_with_knobs(req, "gfx942_dense", knobs) + spec = selected.kernel_spec + self.assertTrue(spec.persistent) + self.assertEqual(spec.block_n, 64) + self.assertTrue(spec.resolved_use_cfvst()) + self.assertEqual(spec.lds_num_buffers, 2) + self.assertTrue(spec.v_reg_prefetch) + self.assertEqual(spec.lds_row_pad, 4) + self.assertTrue(spec.causal_diag_split) + self.assertEqual(spec.resolved_persist_decode, "bt_hkv_minor") + ok, why = supports_gfx942(spec, arch="gfx942") + self.assertTrue(ok, why) + name = selected.kernel_name() + for marker in ("cfvst1", "kdb1", "vpf1", "krowpad4", "dsplit"): + self.assertIn(marker, name) + replay = tuning_spec_with_knobs(req, "gfx942_dense", dict(selected.knobs)) + self.assertEqual(replay.kernel_spec, spec) + self.assertEqual(replay.tuning_id, selected.tuning_id) + self.assertEqual(replay.kernel_name(), name) + + def test_pipeline_axis_values_remain_declared(self): + """Both independent pipeline levers and their prerequisite remain tunable.""" + axes = {axis.name: axis for axis in tuning_axes("gfx942", "dense")} + for key, expected in ( + ("use_cfvst", True), + ("lds_num_buffers", 2), + ("v_reg_prefetch", True), + ): + self.assertIn(key, axes) + self.assertTrue( + any((key, expected) in choice for choice in axes[key].choices), + (key, expected), + ) + def test_production_specs_are_legal_distinct_and_start_at_the_shipped_spec(self): for req in (_gfx942_request(), _gfx942_request(seqlen_q=8192, seqlen_k=8192)): specs = _gfx942_specs(req) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json index 3bd882d10b2c..adfff8e670f7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json @@ -3,6 +3,14 @@ "flavors": { "llvm20": { "cases": { + "attention_dense_gfx942/cfvst_d128_bf16_causal": { + "bytes": 124240, + "sha256": "11d726f6232c66d61554d0072238d3e5cb3b1b837bc633793e2bfcc5c0794cb9" + }, + "attention_dense_gfx942/cfvst_dsplit_d128_bf16_causal": { + "bytes": 208216, + "sha256": "f2ffea15bb3441572b464bf7c40c838753c59a2f49b5206028b60a122696d80c" + }, "attention_dense_gfx942/default_d128_bf16_causal": { "bytes": 171590, "sha256": "3a21804add31a71453e5f86874912285e3f66bb0cb8c3b481f44e6b55ad00812" @@ -59,6 +67,18 @@ "bytes": 178429, "sha256": "15390d0c3bb2f04b83a10166f160c44a5553d8298057f68e64f820155e7add4a" }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdb_d128_bf16": { + "bytes": 231585, + "sha256": "45502ab2cdec4176a4b27161d3bb14d07156bc03d61f0315a9db8784eb767f92" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdbvpf_d128_bf16": { + "bytes": 231883, + "sha256": "c149d0d3ee04c0ccbd5e7e0c3f999615d198a645cbb4f9d50b41223881af7e97" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_vpf_d128_bf16": { + "bytes": 231593, + "sha256": "6710a2392b8cec43a99fde47bdb2bc8ebc27e8099ac329cf69fe0f6f79ee9377" + }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { "bytes": 185931, "sha256": "0695fa68401f74c016c4b677a89bcb09582a8b2692d50dcc71c92fe44851f9c2" @@ -67,6 +87,34 @@ "bytes": 186480, "sha256": "b6a7bbe360dff6bdc8f98c81f730736002d60c5d48129479df3e55900bf80f85" }, + "attention_dense_gfx942/persist_bthkvmin_cfvst_d128_bf16_causal": { + "bytes": 127239, + "sha256": "195f266ab2ae68fe8a4f90ecd4f3e7e06ac0c31d001344dc77f760e2373296af" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_d128_bf16_causal": { + "bytes": 224606, + "sha256": "509a964ea2f0d6d96d8c5a5440663e44047f13697dd518b77417453b23fc0e0b" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_dsplit_d128_bf16": { + "bytes": 309367, + "sha256": "f03f97d2030a50b4f00b6964b5df15fd09d5ebae9457c8af7d4de20bd7db1bf1" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_d128_bf16_causal": { + "bytes": 224904, + "sha256": "79e1829fd58e2c00a2959d582136024f33ff0f97f0ad1c1664746f8077aec468" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_pad4_dsplit": { + "bytes": 310814, + "sha256": "e405da0694fa6c72796e2988043bf1dcd17db640cf34cb29d8a60189e7b0e136" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_d128_bf16_causal": { + "bytes": 224614, + "sha256": "df133d48346cf0a6b55abf9c1a4bf5dfa39465eea7229172c2a7d1d1538c1e12" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_dsplit_d128_bf16": { + "bytes": 309285, + "sha256": "b276e84937a40403dfeff10c8dbb5041c44232202b53f596540e1693ba514037" + }, "attention_dense_gfx942/persist_d128_bf16_intl": { "bytes": 174216, "sha256": "279d5d2fbfbeaf1e36529843e5958901109c561011497045684cfccd9dc5124d" @@ -87,6 +135,14 @@ "bytes": 101084, "sha256": "fb41aa30f46608819cb1a10327aaaf99613b4dd45c6b87bb399ab62ff85c9dae" }, + "attention_dense_gfx942/persist_kdb_d128_bf16_qbmaj": { + "bytes": 223771, + "sha256": "9ec5db6979b0c8631a671074075870de79083f72e1a170cdad8140213a10a11c" + }, + "attention_dense_gfx942/persist_kdbvpf_d128_bf16_qbmaj": { + "bytes": 224071, + "sha256": "11ed81620c6a0be3c0a4e0df2c3b1347a888e92616a2295f59cf3f541127e853" + }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { "bytes": 174592, "sha256": "4cacee0fc5604453e35fb33740b6e030a0768f851e825d578a030e46b01bef39" @@ -103,6 +159,10 @@ "bytes": 178969, "sha256": "ffb28779bc551be8722dcc2ab8b141f02599b5cc7668b3a1e2b7b82a1d560c17" }, + "attention_dense_gfx942/persist_vpf_d128_bf16_qbmaj": { + "bytes": 223777, + "sha256": "3a070b951657ef7efa2c79abcf77c4d12f196d659d77d8c5ff84fa5d0055cf46" + }, "attention_dense_gfx942/sink_d128_bf16_causal": { "bytes": 172103, "sha256": "eff00bf729567a82bcd771814ca4e702b9420cb93585b55a61418bc5660822ec" @@ -135,6 +195,14 @@ }, "llvm22": { "cases": { + "attention_dense_gfx942/cfvst_d128_bf16_causal": { + "bytes": 124336, + "sha256": "146717e6efea705cc91a1b69fff21515c3ad5c6560ee558970c6e7c38917779a" + }, + "attention_dense_gfx942/cfvst_dsplit_d128_bf16_causal": { + "bytes": 208312, + "sha256": "6b47bef0a2b1739c58c19ff20e2aca08275241ab4e99a932f519dca8fc8c2922" + }, "attention_dense_gfx942/default_d128_bf16_causal": { "bytes": 171686, "sha256": "2658a44f0016de96754d17f3dff885a6c6d018ef70629ee728410050218b3c24" @@ -191,6 +259,18 @@ "bytes": 178525, "sha256": "360f96a913ca34db8880643a4101efa9a1559272d347ee87f425cd500bb8b390" }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdb_d128_bf16": { + "bytes": 231681, + "sha256": "c1491cb845b5e6cf29135cc355a12be2e86aae18c8eb9d5b94d1d5b80acfe2e4" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdbvpf_d128_bf16": { + "bytes": 231979, + "sha256": "686e4618b627e8a31f7744e0f848d1c15711069000226247bb40f50bf9bd7bbf" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_vpf_d128_bf16": { + "bytes": 231689, + "sha256": "300efd8b52d20ea3ba28dc8c6d1296e05c5cb81890d3d1e7cdf0b01a944d1476" + }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { "bytes": 186027, "sha256": "dd2d507d110839f77636a9705f2b2d4b064351f2f203a8ebc1177bed01d66ad5" @@ -199,6 +279,34 @@ "bytes": 186576, "sha256": "85ba825b459ad10fd73ac96c80a7b5f3ef51d0d0a88f5c2d99e7b646b57c8068" }, + "attention_dense_gfx942/persist_bthkvmin_cfvst_d128_bf16_causal": { + "bytes": 127335, + "sha256": "8501b01c3690e83e186fb01ddb045851bcfc663e3f0b862744f8cdf00e97716e" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_d128_bf16_causal": { + "bytes": 224702, + "sha256": "eef106838ab84c78bdb4d30f8886183feff075f440d4e063d73865bed13737f0" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_dsplit_d128_bf16": { + "bytes": 309467, + "sha256": "b61590394320c0f0f0c30e232593e9a443b402d98d5625688f5bcb494f231b32" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_d128_bf16_causal": { + "bytes": 225000, + "sha256": "193d8d0d90bbe366beae52544849bc259d14d0d4ffb568ebc54c089961d9cc57" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_pad4_dsplit": { + "bytes": 310914, + "sha256": "042e235745f9ffe3cb641927b8a84476627d155647aeb02f6da6b749ad0f1767" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_d128_bf16_causal": { + "bytes": 224710, + "sha256": "f8956dd216eac888658c0e3c36bef86668ce4626141c9a0b67059a4d4ac40ca5" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_dsplit_d128_bf16": { + "bytes": 309385, + "sha256": "20c88db88f97337df48af8236bcde19743fa293284c95b2cecc43009d607d494" + }, "attention_dense_gfx942/persist_d128_bf16_intl": { "bytes": 174312, "sha256": "8b187788a121447658b3a60dc296b659f2ac992140c8274e1107aee5e33acac5" @@ -219,6 +327,14 @@ "bytes": 101180, "sha256": "9795b13b2179739e6504054478a30732b4fb5346cb9f5cab5926a49c5ba8c185" }, + "attention_dense_gfx942/persist_kdb_d128_bf16_qbmaj": { + "bytes": 223867, + "sha256": "6e82ca7999e9468c23bc4e54af299fd34300d9ce5d5f39f27863854f7dfd304b" + }, + "attention_dense_gfx942/persist_kdbvpf_d128_bf16_qbmaj": { + "bytes": 224167, + "sha256": "cb2df14d8292bfe6588e3f22987698db1aa2eac0203ad088ff5fffeca9ad24bd" + }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { "bytes": 174688, "sha256": "8c4f15cff00adee735dcec47d96c6b1917a8609182a624fd18c53b197b00beaf" @@ -235,6 +351,10 @@ "bytes": 179065, "sha256": "9218c9179f0b72d2c6eab8f32495e91c70520d737c21201a1fda0b4ce9c33cb6" }, + "attention_dense_gfx942/persist_vpf_d128_bf16_qbmaj": { + "bytes": 223873, + "sha256": "d18fc4e99f92662a1f6da224fc438cbe81423c4a4910b13ab98c97e98a15e969" + }, "attention_dense_gfx942/sink_d128_bf16_causal": { "bytes": 172199, "sha256": "48fdf25db27ad6317f33e4d40c5a286f74377471fbfec608160680f809fd15d3" @@ -267,6 +387,14 @@ }, "llvm23": { "cases": { + "attention_dense_gfx942/cfvst_d128_bf16_causal": { + "bytes": 124400, + "sha256": "352287ab92656fbd74e009f22b9ec1757fd5ca6f562811afd0f21c27bc39d9da" + }, + "attention_dense_gfx942/cfvst_dsplit_d128_bf16_causal": { + "bytes": 208376, + "sha256": "e7f78d1691e3021ec7808a458292368f340db46f30bd0c4bfc6600fa82369f0d" + }, "attention_dense_gfx942/default_d128_bf16_causal": { "bytes": 171750, "sha256": "bc12b9a1aba7d9bfae67d64c24339d006cec293c40e31a9f41e33aefbe7981eb" @@ -323,6 +451,18 @@ "bytes": 178589, "sha256": "31d3c09d639960a9803fdc847515ec3dab0a852082091c90f9f13ef396de7c0b" }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdb_d128_bf16": { + "bytes": 231745, + "sha256": "26383cdc12fb8506d3b77bda406488da38d93758a86e4fb59ea94f7a67480209" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_kdbvpf_d128_bf16": { + "bytes": 232043, + "sha256": "ff06bd3ce3aef1268cf17d766836c9aebadba604a1e9f52fd4f5fd9b7485a211" + }, + "attention_dense_gfx942/lse_persist_bthkvmin_vpf_d128_bf16": { + "bytes": 231753, + "sha256": "2787856f3e07c430f01d3380e3e4284a0f83112e7e46da33bd5031b0a1a1a5ab" + }, "attention_dense_gfx942/lse_persist_empty_window_d128_bf16": { "bytes": 186091, "sha256": "a724043ad5acf38e53968e888c0f8e02019c5c858dd1bc2047941fddf34c5a87" @@ -331,6 +471,34 @@ "bytes": 186640, "sha256": "3631253c777bc46831d8b85fa421f81cffe0df9460a83a9a1742c3c0bfdad2df" }, + "attention_dense_gfx942/persist_bthkvmin_cfvst_d128_bf16_causal": { + "bytes": 127399, + "sha256": "a166503270793efbf64be8b727922602e4d93bd345c570f8b8a24727d7591992" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_d128_bf16_causal": { + "bytes": 224766, + "sha256": "afd5165061b5402e232212205bb0afe7f7b50c5e782d612a04c8781254902459" + }, + "attention_dense_gfx942/persist_bthkvmin_kdb_dsplit_d128_bf16": { + "bytes": 309531, + "sha256": "f09682701409db1cffcd71b2e13abe2ab6580b6ec329577af7fa588492925b34" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_d128_bf16_causal": { + "bytes": 225064, + "sha256": "54d46da1faf7c979fade1bf721333edff76265ae494e0f726615d711e5ad7eb7" + }, + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_pad4_dsplit": { + "bytes": 310978, + "sha256": "cd399d182f2997930266b0c854604a735d14c7f83ab540d1fc92a8fb60a29c35" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_d128_bf16_causal": { + "bytes": 224774, + "sha256": "78ad4db33411ac97346bc9a2d0a178299b7f7ac8fa5909cc70533a72b4063394" + }, + "attention_dense_gfx942/persist_bthkvmin_vpf_dsplit_d128_bf16": { + "bytes": 309449, + "sha256": "3f7fd5c5f9c37f824916946187a3747fba4333301c033ef0b3a48a72589ea7c3" + }, "attention_dense_gfx942/persist_d128_bf16_intl": { "bytes": 174376, "sha256": "6e231a1399c80e2918e18ca1e92080523df45163c09b5fd2aab621a8326fe11c" @@ -351,6 +519,14 @@ "bytes": 101244, "sha256": "f480aeab86891c6110ff294f7997309ad4fec1bd62c3c75e1f12404a8a3c4ea7" }, + "attention_dense_gfx942/persist_kdb_d128_bf16_qbmaj": { + "bytes": 223931, + "sha256": "2316c9c6d5ac17d8d6f424bf71f644a137fa76e84762a5533bc5c7c2444cf948" + }, + "attention_dense_gfx942/persist_kdbvpf_d128_bf16_qbmaj": { + "bytes": 224231, + "sha256": "6fae801542dcc1de2b27ef59e46fa421182600357613fa26134074f1c841010d" + }, "attention_dense_gfx942/persist_sink_d128_bf16_qbmaj": { "bytes": 174752, "sha256": "0848ea40055a8d4b810794512cec6c04917a1306636300d02ea4e28245d71283" @@ -367,6 +543,10 @@ "bytes": 179129, "sha256": "7499a1560311c9ecc61e06e8df7fda97ea285e5ad75da3625a58e8c49ae086db" }, + "attention_dense_gfx942/persist_vpf_d128_bf16_qbmaj": { + "bytes": 223937, + "sha256": "d8e510fcb45e5dfa11f891cdf8c1c72cbb3d19855560b5c4e749214e94f81e44" + }, "attention_dense_gfx942/sink_d128_bf16_causal": { "bytes": 172263, "sha256": "31d29f2bfc8d35c82af0911a0ef50ffdeaca07cb4198ffc135bcaaa355da1916" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index 9818d76d292c..02e3ccedecfb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -44,6 +44,7 @@ attention_dense_grid, attention_dense_runtime_args, attention_dense_signature, + bf16_cfvst_supported, build_attention_dense, gfx942_kernel_name, run_attention_dense_torch, @@ -53,6 +54,9 @@ _FASTDIV_DIVIDEND_LIMIT, _k_group_pad_active, _k_group_stride, + _lds_bytes, + _use_kdb, + _use_vpf, _use_exp2_fast, _v_row_pad, _v_swizzle_width, @@ -267,7 +271,6 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): "block_size", "num_kv_blocks", "causal_bottom_right", - "lds_num_buffers", } ) @@ -280,7 +283,7 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): # Second LEGAL values per field. Every candidate is filtered through # supports_attention_dense before use, so a candidate that is illegal for a given # base is skipped rather than asserted on; listing several per field keeps coverage -# when one base rejects one of them (e.g. use_cfvst=True is legal only at fp16 D128). +# when one base rejects one of them (e.g. the KDB/VPF depth needs BF16 D128 CFVST). _SPEC_PERTURBATIONS = { # The shape fields (batch, both seqlens, both head counts) are _RUNTIME_PARAM_FIELDS # on BOTH grids: every perturbation must move only the kernarg values, never the @@ -327,7 +330,8 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): "use_v_swizzle": (False, True), "use_exp2_fast": (False, True), "iglp": (True, False), - "lds_num_buffers": (2,), # unbuildable: only NBUF=1 is implemented + "lds_num_buffers": (2,), # BF16 D128 CFVST enables KDB + "v_reg_prefetch": (True,), # independent V register prefetch "iglp_mode": (1,), # legal only with iglp=True -> the *_iglp base "pv_loop_order": ("k_major",), "o_store_width": (1, 2), @@ -340,16 +344,57 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): _PERTURBATIONS = {**_SPEC_PERTURBATIONS, **_PRIVATE_PERTURBATIONS} # Base configurations the perturbations are applied to, one field at a time. Chosen -# so every structurally distinct arm of the builder is a base: cfvst on (fp16 D128), -# cfvst off at D128 (bf16, exp2_fast), both D64 dtypes (packed 2-rows-per-DMA + the +# so every structurally distinct arm of the builder is a base: CFVST on +# (FP16 D128 and opt-in BF16 D128), naive V on BF16 D128, and both +# D64 dtypes (packed 2-rows-per-DMA + the # K row-group pad + the wpe=4 tune), and the P4 persistent grid -- without which # num_persistent / interleave / persist_decode are inert and their coverage vacuous, # and the runtime-shape claim would be checked on the default grid only. # The iglp base is what makes iglp_mode (read only under iglp=True) reachable. +# The *_kdbvpf bases enable BOTH independent pipeline levers explicitly; changing +# lds_num_buffers no longer implicitly enables V prefetch. Every other knob is +# also perturbed inside the combined pipeline body, +# including iglp_mode (iglp base) and the persistent-grid fields (persistent base). _INJECTIVITY_BASES = { "d128_fp16_cfvst": dict(head_size=128, dtype="fp16"), "d128_fp16_cfvst_iglp": dict(head_size=128, dtype="fp16", iglp=True), "d128_bf16_naive": dict(head_size=128, dtype="bf16"), + "d128_bf16_cfvst": dict(head_size=128, dtype="bf16", use_cfvst=True), + "persistent_d128_bf16_cfvst": dict( + head_size=128, + dtype="bf16", + use_cfvst=True, + persistent=True, + num_persistent=228, + seqlen_q=4096, + seqlen_kv=4096, + num_query_heads=32, + ), + "d128_bf16_kdbvpf": dict( + head_size=128, + dtype="bf16", + use_cfvst=True, + persistent=True, + num_persistent=228, + seqlen_q=4096, + seqlen_kv=4096, + num_query_heads=32, + lds_num_buffers=2, + v_reg_prefetch=True, + ), + "d128_bf16_kdbvpf_iglp": dict( + head_size=128, + dtype="bf16", + use_cfvst=True, + persistent=True, + num_persistent=228, + seqlen_q=4096, + seqlen_kv=4096, + num_query_heads=32, + lds_num_buffers=2, + v_reg_prefetch=True, + iglp=True, + ), "d64_fp16": dict(head_size=64, dtype="fp16"), "d64_bf16": dict(head_size=64, dtype="bf16"), "persistent_d128_fp16": dict( @@ -361,6 +406,18 @@ def test_kernel_name_pins_the_two_shipped_cache_collisions(): seqlen_kv=4096, num_query_heads=32, ), + "persistent_d128_bf16_kdbvpf": dict( + head_size=128, + dtype="bf16", + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + persistent=True, + num_persistent=228, + seqlen_q=4096, + seqlen_kv=4096, + num_query_heads=32, + ), } @@ -946,7 +1003,8 @@ def test_no_sinks_emits_no_sink_param_or_load(): def test_supports_rejects_over_budget_lds(): - """block_n=128 at D128/bf16 needs 2*128*(128+8)*2 = 69632 B > the 64 KB gfx942 LDS. + """Naive-V D128/BF16 with block_n=128 needs 2*128*(128+8)*2 = 69632 B, + exceeding the 64 KiB gfx942 LDS capacity. Without this gate it reaches comgr and dies with an opaque CODEGEN abort. The pad is set EXPLICITLY on the spec rather than inherited from the field default: @@ -955,7 +1013,14 @@ def test_supports_rejects_over_budget_lds(): flips to ACCEPT and this test's premise would silently invert into a tautology. The pad=0 boundary is pinned deliberately in the test below. """ - spec = _spec(block_n=128, head_size=128, dtype="bf16", lds_row_pad=8) + spec = _spec( + block_n=128, + head_size=128, + dtype="bf16", + lds_row_pad=8, + use_cfvst=False, + lds_num_buffers=1, + ) ok, why = supports_attention_dense(spec, arch="gfx942") assert not ok and "LDS" in why # the arithmetic the rejection rests on, spelled out so a capacity/pad change is loud @@ -977,7 +1042,14 @@ def test_lds_budget_boundary_is_accepted_at_exactly_capacity(): from kernels.gfx942.attention_dense import _lds_bytes - spec = _spec(block_n=128, head_size=128, dtype="bf16", lds_row_pad=0) + spec = _spec( + block_n=128, + head_size=128, + dtype="bf16", + lds_row_pad=0, + use_cfvst=False, + lds_num_buffers=1, + ) capacity = attention_lds_capacity_bytes("gfx942") assert _lds_bytes(spec) == capacity, "premise: exactly at capacity" ok, why = supports_attention_dense(spec, arch="gfx942") @@ -986,6 +1058,480 @@ def test_lds_budget_boundary_is_accepted_at_exactly_capacity(): assert build_attention_dense(spec, arch="gfx942") is not None +# --------------------------------------------------------------------------- # +# Opt-in BF16 D128 CFVST and K-only double buffering / V register prefetch. +# lds_num_buffers=1 stays byte-for-byte upstream for existing configurations. +# --------------------------------------------------------------------------- # +def test_bf16_cfvst_kdb_vpf_are_separate_opt_in_variants(): + persistent = dict(persistent=True, num_persistent=304) + baseline = _spec(**persistent) + cfvst_only = _spec(use_cfvst=True, **persistent) + kdb_only = _spec(use_cfvst=True, lds_num_buffers=2, **persistent) + vpf_only = _spec(use_cfvst=True, v_reg_prefetch=True, **persistent) + combined = _spec( + use_cfvst=True, lds_num_buffers=2, v_reg_prefetch=True, **persistent + ) + fp16 = _spec(dtype="fp16", **persistent) + + assert not baseline.resolved_use_cfvst() + assert not _use_kdb(baseline) and not _use_vpf(baseline) + assert cfvst_only.resolved_use_cfvst() + assert not _use_kdb(cfvst_only) and not _use_vpf(cfvst_only) + assert _use_kdb(kdb_only) and not _use_vpf(kdb_only) + assert not _use_kdb(vpf_only) and _use_vpf(vpf_only) + assert _use_kdb(combined) and _use_vpf(combined) + assert fp16.resolved_use_cfvst() + assert not _use_kdb(fp16) and not _use_vpf(fp16) + + assert baseline.resolved_v_row_pad() == 8 + for spec in (cfvst_only, kdb_only, vpf_only, combined): + assert spec.resolved_v_row_pad() == 0 + assert spec.resolved_use_v_swizzle() + + cases = (baseline, cfvst_only, kdb_only, vpf_only, combined, fp16) + names = {gfx942_kernel_name(spec) for spec in cases} + assert len(names) == len(cases), "distinct variants must not share a cache name" + assert "_cfvst1" not in gfx942_kernel_name(baseline) + assert "_cfvst1" in gfx942_kernel_name(cfvst_only) + assert "_kdb1" not in gfx942_kernel_name(cfvst_only) + assert "_vpf1" not in gfx942_kernel_name(cfvst_only) + assert "_kdb1" in gfx942_kernel_name(kdb_only) + assert "_vpf1" not in gfx942_kernel_name(kdb_only) + assert "_kdb1" not in gfx942_kernel_name(vpf_only) + assert "_vpf1" in gfx942_kernel_name(vpf_only) + assert "_kdb1" in gfx942_kernel_name(combined) + assert "_vpf1" in gfx942_kernel_name(combined) + assert "_kdb1" not in gfx942_kernel_name(fp16) + assert "_vpf1" not in gfx942_kernel_name(fp16) + + for spec in cases: + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert ok, reason + built = build_attention_dense(spec, arch="gfx942") + assert built.name == gfx942_kernel_name(spec) + + +@pytest.mark.parametrize( + "overrides", + [ + {"use_cfvst": True, "lds_num_buffers": 2}, + {"use_cfvst": True, "v_reg_prefetch": True}, + {"use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True}, + ], +) +def test_cfvst_kdb_vpf_pipeline_requires_persistent(overrides): + spec = _spec(persistent=False, **overrides) + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert not ok and "persistent=True" in reason, reason + with pytest.raises(ValueError, match="unsupported gfx942 attention_dense spec"): + build_attention_dense(spec, arch="gfx942") + + +@pytest.mark.parametrize( + "overrides, expected_bytes", + [ + ({}, 34816), # naive V, K[1, 64, 136] + V[1, 64, 136] + ({"use_cfvst": True}, 33792), # K[1, 64, 136] + V^T[128, 64] + ( + { + "use_cfvst": True, + "lds_num_buffers": 2, + "persistent": True, + "num_persistent": 304, + }, + 51200, + ), + ( + { + "use_cfvst": True, + "v_reg_prefetch": True, + "persistent": True, + "num_persistent": 304, + }, + 33792, + ), + ( + { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + "persistent": True, + "num_persistent": 304, + }, + 51200, + ), + ({"dtype": "fp16"}, 33792), # unchanged upstream FP16 CFVST + ], + ids=[ + "bf16-naive", + "bf16-cfvst", + "bf16-kdb", + "bf16-vpf", + "bf16-kdb-vpf", + "fp16-upstream", + ], +) +def test_cfvst_kdb_vpf_lds_budgets(overrides, expected_bytes): + spec = _spec(**overrides) + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert ok, reason + assert _lds_bytes(spec) == expected_bytes + + +@pytest.mark.parametrize( + "overrides, expected_message", + [ + ( + {"use_cfvst": "false"}, + "use_cfvst must be None or a bool, got str", + ), + ( + {"use_cfvst": 1}, + "use_cfvst must be None or a bool, got int", + ), + ( + {"lds_num_buffers": 2.0}, + "lds_num_buffers must be an int, got float", + ), + ( + {"lds_num_buffers": True}, + "lds_num_buffers must be an int, got bool", + ), + ( + {"v_reg_prefetch": "false"}, + "v_reg_prefetch must be a bool, got str", + ), + ( + {"v_reg_prefetch": 1}, + "v_reg_prefetch must be a bool, got int", + ), + ( + {"v_reg_prefetch": None}, + "v_reg_prefetch must be a bool, got NoneType", + ), + ], + ids=[ + "cfvst-str", + "cfvst-int", + "nbuf-float", + "nbuf-bool", + "vpf-str", + "vpf-int", + "vpf-none", + ], +) +def test_cfvst_kdb_vpf_knobs_reject_non_exact_types(overrides, expected_message): + """Invalid knob types must raise parameter-specific ValueErrors.""" + with pytest.raises(ValueError) as exc_info: + _spec(**overrides) + + assert type(exc_info.value) is ValueError + assert str(exc_info.value) == expected_message + + +def test_cfvst_kdb_vpf_rejects_over_budget_lds(): + """Two K slots and one transposed V slot exceed 64 KiB with a wide K row pad. + + lds_num_buffers=2 is gated to block_n=64, so the budget is reached through the + K pad: lds_row_pad=64 lands exactly on 65536 B (accepted), 68 is over.""" + spec = _spec( + block_n=64, + head_size=128, + dtype="bf16", + lds_row_pad=68, + use_cfvst=True, + lds_num_buffers=2, + persistent=True, + num_persistent=304, + ) + # 2 * 64 * (128 + 68) * 2 + 128 * 64 * 2 = 66560 B. + assert _lds_bytes(spec) == 66560 + ok, _ = supports_attention_dense( + dataclasses.replace(spec, lds_row_pad=64), arch="gfx942" + ) + assert ok, "65536 B is exactly the gfx942 LDS capacity and must fit" + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert not ok and "LDS" in reason, reason + with pytest.raises(ValueError, match="unsupported gfx942 attention_dense spec"): + build_attention_dense(spec, arch="gfx942") + + +@pytest.mark.parametrize( + "overrides", + [ + {"lds_num_buffers": 2}, # no CFVST + {"v_reg_prefetch": True}, # no CFVST + {"dtype": "fp16", "lds_num_buffers": 2}, + {"dtype": "fp16", "v_reg_prefetch": True}, + {"use_cfvst": True, "lds_num_buffers": 2, "causal": False}, + {"use_cfvst": True, "lds_num_buffers": 2, "sliding_window": 128}, + {"use_cfvst": True, "use_sinks": True}, # bf16 CFVST alone: not validated + {"use_cfvst": True, "lds_num_buffers": 2, "use_sinks": True}, + { + "use_cfvst": True, + "lds_num_buffers": 2, + "use_sinks": True, + "emit_lse": True, + "persistent": True, + "num_persistent": 304, + }, + {"use_cfvst": True, "lds_num_buffers": 2, "block_n": 32}, # gated to 64 + {"use_cfvst": True, "v_reg_prefetch": True, "block_n": 32}, + {"use_cfvst": True, "v_reg_prefetch": True, "sliding_window": 128}, + {"use_cfvst": True, "lds_num_buffers": 3}, + ], +) +def test_cfvst_kdb_vpf_rejects_unsupported_modes(overrides): + spec = _spec(**{"persistent": True, "num_persistent": 304, **overrides}) + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert not ok and reason, (overrides, reason) + with pytest.raises(ValueError, match="unsupported gfx942 attention_dense spec"): + build_attention_dense(spec, arch="gfx942") + + +@pytest.mark.parametrize("persistent", [False, True]) +@pytest.mark.parametrize("diag_split", [False, True]) +def test_bf16_cfvst_supported_accepts_the_cohort(persistent, diag_split): + """bf16_cfvst_supported is the single BF16 CFVST cohort definition: BF16 D128 + causal on either grid, with or without causal_diag_split. The CFVST gate in + supports_attention_dense admits exactly these specs.""" + spec = _spec(persistent=persistent, causal_diag_split=diag_split) + assert spec.dtype == "bf16" and spec.head_size == 128 and spec.causal + assert bf16_cfvst_supported(spec) + ok, why = supports_attention_dense( + dataclasses.replace(spec, use_cfvst=True), arch="gfx942" + ) + assert ok, why + + +@pytest.mark.parametrize( + "change", + [ + {"head_size": 64}, + {"causal": False}, + {"sliding_window": 256}, + {"use_sinks": True}, + ], + ids=["d64", "non_causal", "sliding_window", "sinks"], +) +def test_bf16_cfvst_supported_rejects_each_broken_condition(change): + """Breaking any one cohort condition leaves the predicate false, and the gate + then refuses BF16 use_cfvst=True rather than building it.""" + cohort = _spec() + assert bf16_cfvst_supported(cohort) + spec = dataclasses.replace(cohort, **change) + assert not bf16_cfvst_supported(spec) + ok, why = supports_attention_dense( + dataclasses.replace(spec, use_cfvst=True), arch="gfx942" + ) + assert not ok and why + + +def test_bf16_cfvst_supported_is_bf16_only(): + """FP16 D128 is outside the BF16 cohort; its CFVST stays policy-driven.""" + spec = _spec(dtype="fp16") + assert not bf16_cfvst_supported(spec) + assert spec.resolved_use_cfvst() + + +@pytest.mark.parametrize( + "overrides", + [ + {"use_cfvst": True, "lds_num_buffers": 2}, + {"use_cfvst": True, "v_reg_prefetch": True}, + { + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + }, + ], + ids=["kdb", "vpf", "kdb_vpf"], +) +@pytest.mark.parametrize( + "decode", + ["qb_major", "bt_hkv_minor"], +) +def test_cfvst_kdb_vpf_accepts_causal_diag_split(overrides, decode): + spec = _spec( + persistent=True, + num_persistent=304, + persist_decode=decode, + causal_diag_split=True, + **overrides, + ) + + ok, reason = supports_attention_dense(spec, arch="gfx942") + assert ok, reason + + kernel = build_attention_dense(spec, arch="gfx942") + assert kernel.name == gfx942_kernel_name(spec) + assert "_dsplit" in kernel.name + + unsplit = dataclasses.replace(spec, causal_diag_split=False) + assert gfx942_kernel_name(unsplit) != kernel.name + + +@pytest.mark.parametrize("decode", ["qb_major", "bt_hkv_minor"]) +def test_cfvst_full_pad4_dsplit_winner_is_supported_and_distinct(decode): + """The exact opt-in KDB+VPF+pad4+split combination stays buildable.""" + spec = _spec( + persistent=True, + num_persistent=304, + persist_decode=decode, + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + lds_row_pad=4, + causal_diag_split=True, + ) + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, why + name = gfx942_kernel_name(spec) + assert all(tag in name for tag in ("cfvst1", "kdb1", "vpf1", "krowpad4", "dsplit")) + assert _lds_bytes(spec) == 50176 + kernel = build_attention_dense(spec, arch="gfx942") + assert kernel.name == name + assert ( + gfx942_kernel_name(dataclasses.replace(spec, causal_diag_split=False)) != name + ) + + +@pytest.mark.parametrize( + "lds_num_buffers,v_reg_prefetch", [(2, False), (1, True), (2, True)] +) +@pytest.mark.parametrize("diag_split", [False, True]) +def test_cfvst_pipeline_loop_keeps_lds_handoff_barriers( + lds_num_buffers, v_reg_prefetch, diag_split +): + """Both the split and unsplit pipeline keep cross-wave overwrite/publication guards. + + This examines emitted builder operations, not source-text substrings. Numeric + concurrency tests remain necessary because a correct-looking barrier sequence + alone cannot prove all scheduling races absent on the GPU. + """ + spec = _spec( + persistent=True, + num_persistent=3, + num_query_heads=16, + num_kv_heads=4, + seqlen_q=512, + seqlen_kv=512, + use_cfvst=True, + lds_num_buffers=lds_num_buffers, + v_reg_prefetch=v_reg_prefetch, + causal_diag_split=diag_split, + ) + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, why + kernel = build_attention_dense(spec, arch="gfx942") + + def walk(op): + yield op + for region in getattr(op, "regions", ()): + for child in region.ops: + yield from walk(child) + + loops = [ + op + for root in kernel.body.ops + for op in walk(root) + if op.name == "scf.for" and op.attrs.get("elide_trailing_barrier") is False + ] + assert len(loops) >= (2 if diag_split else 1) + for loop in loops: + names = [op.name for op in loop.regions[0].ops] + assert names[-2:] == ["tile.sync_lds_only", "scf.yield"], names + assert names.count("tile.sync_lds_only") >= 2, names + # Current V LDS is single-buffered. Next-tile publication is after + # current-tile read completion, and waits for the next V/K VMEM loads. + first_guard = names.index("tile.sync_lds_only") + drain = names.index("tile.s_waitcnt", first_guard) + final_publish = len(names) - 2 + assert first_guard < drain < final_publish, names + + +@pytest.mark.parametrize( + "request_changes,knobs", + [ + ({}, {"use_cfvst": True, "persistent": False, "lds_num_buffers": 2}), + ({}, {"use_cfvst": True, "persistent": False, "v_reg_prefetch": True}), + ( + {}, + { + "use_cfvst": True, + "persistent": True, + "lds_num_buffers": 2, + "block_n": 32, + }, + ), + ({"dtype": "fp16"}, {"persistent": True, "lds_num_buffers": 2}), + ({}, {"persistent": True, "lds_num_buffers": 2}), + ], + ids=[ + "kdb_nonpersistent", + "vpf_nonpersistent", + "kdb_bn32", + "fp16_kdb", + "kdb_without_cfvst", + ], +) +def test_cfvst_pipeline_invalid_explicit_dispatch_is_rejected(request_changes, knobs): + """Illegal modes must be rejected through tuning selection, not only support().""" + from dispatch.attention import AttentionRequest, tuning_spec_with_knobs + + req_kwargs = dict( + batch=1, + nhead_q=32, + nhead_k=8, + seqlen_q=2048, + seqlen_k=2048, + hdim_q=128, + hdim_v=128, + arch="gfx942", + dtype="bf16", + mask_type=1, + ) + req_kwargs.update(request_changes) + with pytest.raises(ValueError): + tuning_spec_with_knobs(AttentionRequest(**req_kwargs), "gfx942_dense", knobs) + + +def test_cfvst_full_pad4_dsplit_reachable_by_explicit_tuning(): + """Real dispatch tuning selection reaches the measured combined configuration.""" + from dispatch.attention import AttentionRequest, tuning_spec_with_knobs + + req = AttentionRequest( + batch=1, + nhead_q=32, + nhead_k=8, + seqlen_q=4096, + seqlen_k=4096, + hdim_q=128, + hdim_v=128, + arch="gfx942", + dtype="bf16", + mask_type=1, + ) + selected = tuning_spec_with_knobs( + req, + "gfx942_dense", + { + "persistent": True, + "use_cfvst": True, + "lds_num_buffers": 2, + "v_reg_prefetch": True, + "lds_row_pad": 4, + "causal_diag_split": True, + }, + ) + spec = selected.kernel_spec + assert spec.persistent and spec.resolved_use_cfvst() + assert (spec.lds_num_buffers, spec.v_reg_prefetch, spec.lds_row_pad) == (2, True, 4) + assert spec.causal_diag_split and spec.block_n == 64 + assert supports_attention_dense(spec, arch="gfx942")[0] + assert build_attention_dense(spec, arch="gfx942").name == gfx942_kernel_name(spec) + + def test_supports_accepts_lds_budget_that_fits(): # Same block_n at D64 halves the footprint (unpadded rows) and must still pass. ok, why = supports_attention_dense(_spec(block_n=128, head_size=64), arch="gfx942") @@ -1138,10 +1684,10 @@ def test_gfx942_auto_decode_cannot_leak_to_gqa_pair(): dict(lds_row_pad=2), # REJECTED: not a multiple of 4 elements dict(v_row_pad=16), # accepted dict(v_row_pad=-4), # REJECTED: negative - # --- private: use_cfvst (the one tri-state with a rejected direction) --- + # --- private: use_cfvst (legal opt-in for BF16 D128; rejected on D64) --- dict(dtype="fp16", use_cfvst=False), # accepted: OFF is always legal dict(head_size=64, use_cfvst=True), # REJECTED: policy says off at D64 - dict(dtype="bf16", head_size=128, use_cfvst=True), # REJECTED: spills + dict(dtype="bf16", head_size=128, use_cfvst=True), # ACCEPTED: CFVST-only # --- private: use_v_swizzle (independent of the pad; cfvst-gated) --- dict(dtype="fp16", use_v_swizzle=False), # accepted: OFF is always legal dict(head_size=64, use_v_swizzle=True), # REJECTED: no V^T path at D64 @@ -1153,8 +1699,23 @@ def test_gfx942_auto_decode_cannot_leak_to_gqa_pair(): dict(dtype="fp16", iglp=True), dict(iglp=True, pv_sched_fence_mask=0), # REJECTED: iglp owns the schedule # --- private: performance-only codegen knobs --- - dict(lds_num_buffers=1), # accepted: the only implemented depth - dict(lds_num_buffers=2), # REJECTED: NBUF=2 is not implemented + dict(lds_num_buffers=1), # accepted: upstream single K buffer + dict(lds_num_buffers=2), # REJECTED: BF16 without CFVST + dict( + use_cfvst=True, lds_num_buffers=2, persistent=True, num_persistent=304 + ), # ACCEPTED: KDB alone + dict( + use_cfvst=True, v_reg_prefetch=True, persistent=True, num_persistent=304 + ), # ACCEPTED: VPF alone + dict( + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + ), # ACCEPTED: both + dict(v_reg_prefetch=True), # REJECTED: BF16 without CFVST + dict(dtype="fp16", lds_num_buffers=2), # REJECTED: pipeline is BF16-only dict(iglp=True, iglp_mode=1), # accepted dict(iglp_mode=1), # REJECTED: needs iglp=True dict(pv_loop_order="k_major"), # accepted @@ -1918,6 +2479,7 @@ def test_iglp_true_builds_lowers_and_emits_the_intrinsic(dtype, d): # --------------------------------------------------------------------------- # _CODEGEN_KNOB_DEFAULTS = { "lds_num_buffers": 1, + "v_reg_prefetch": False, "iglp_mode": 0, "pv_loop_order": "d_major", "o_store_width": 4, @@ -1925,7 +2487,16 @@ def test_iglp_true_builds_lowers_and_emits_the_intrinsic(dtype, d): "pv_sched_fence_mask": None, "causal_diag_split": False, } -_CODEGEN_KNOB_TOKENS = ("_iglpm", "_pvkmaj", "_osw", "_prio", "_fence", "_dsplit") +_CODEGEN_KNOB_TOKENS = ( + "_kdb1", + "_vpf1", + "_iglpm", + "_pvkmaj", + "_osw", + "_prio", + "_fence", + "_dsplit", +) def _mfma_calls(ir: str) -> int: @@ -2039,6 +2610,7 @@ def test_lazy_rescale_is_inert_on_gfx942(): "overrides, reason", [ (dict(lds_num_buffers=2), "lds_num_buffers"), + (dict(v_reg_prefetch=True), "v_reg_prefetch"), (dict(lds_num_buffers=0), "lds_num_buffers"), (dict(iglp_mode=1), "iglp_mode only applies with iglp=True"), (dict(iglp=True, iglp_mode=2), "iglp_mode must be 0 or 1"), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_golden.py index 0ec21295e4fc..09673e1b507d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_golden.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_golden.py @@ -15,7 +15,11 @@ Covers the acceptance matrix: D64/D128 x bf16/fp16 x default/persistent x GQA, causal and full, plus sliding-window (W128/W256, both grids, D64/D128) and sinks (alone and combined). Every case is in the gfx942 supported set; varlen and ragged are rejected on -gfx942, so unlike the gfx950 sibling they are absent here. +gfx942, so unlike the gfx950 sibling they are absent here. The opt-in bf16 D128 CFVST +and independent K-double-buffer (``lds_num_buffers=2``) and V-register-prefetch +(``v_reg_prefetch=True``) knobs are pinned individually and together on supported +persistent grids (including LSE). CFVST alone also covers the nonpersistent grid; +no shipped default reaches the new pipeline paths. Both D64 K-LDS layouts are pinned so drift on either is caught: * ``default_d64_*`` -- specs built DIRECTLY with ``lds_k_group_pad=0``: the UNPADDED @@ -65,7 +69,7 @@ def _cases(): """cid -> zero-arg builder returning a KernelDef. Small Sq keeps the IR compact - while still exercising the full pipeline (both grid variants, both decodes). + while covering CFVST on the default grid and KDB/VPF on persistent decodes. Two families of builders: * ``mk`` builds an ``AttentionDenseSpec`` DIRECTLY. The D64 cases pass @@ -100,9 +104,10 @@ def _cases(): def mk(**over): d = dict(base) d.update(over) - spec_type = ( - Gfx942AttentionDenseSpec if d.get("emit_lse", False) else AttentionDenseSpec + gfx942_only = d.get("emit_lse", False) or any( + k in d for k in ("use_cfvst", "lds_num_buffers", "v_reg_prefetch") ) + spec_type = Gfx942AttentionDenseSpec if gfx942_only else AttentionDenseSpec return lambda: build_attention_dense(spec_type(**d), arch=_ARCH) def mk_dispatch(**over): @@ -273,6 +278,121 @@ def mk_dispatch(**over): seqlen_kv=64, sliding_window=64, ), + # --- Opt-in BF16 D128 CFVST with independently enabled K double buffering + # and V register prefetch. CFVST alone can use the default grid; KDB/VPF + # require the persistent grid. Pin both persistent decodes and LSE; + # V LDS stays single-buffered. --- + "attention_dense_gfx942/cfvst_d128_bf16_causal": mk(use_cfvst=True), + "attention_dense_gfx942/persist_bthkvmin_kdb_d128_bf16_causal": mk( + use_cfvst=True, + lds_num_buffers=2, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/persist_bthkvmin_vpf_d128_bf16_causal": mk( + use_cfvst=True, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_d128_bf16_causal": mk( + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/persist_kdb_d128_bf16_qbmaj": mk( + use_cfvst=True, + lds_num_buffers=2, + persistent=True, + num_persistent=304, + persist_decode="qb_major", + ), + "attention_dense_gfx942/persist_vpf_d128_bf16_qbmaj": mk( + use_cfvst=True, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="qb_major", + ), + "attention_dense_gfx942/persist_kdbvpf_d128_bf16_qbmaj": mk( + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="qb_major", + ), + # The last KV tile is folded in after the loop; the LSE epilogue must + # use the final row max independently for both pipeline levers. + "attention_dense_gfx942/lse_persist_bthkvmin_kdb_d128_bf16": mk( + emit_lse=True, + use_cfvst=True, + lds_num_buffers=2, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/lse_persist_bthkvmin_vpf_d128_bf16": mk( + emit_lse=True, + use_cfvst=True, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/lse_persist_bthkvmin_kdbvpf_d128_bf16": mk( + emit_lse=True, + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + # The measured final optimization: KDB + VPF + K pad4 + diagonal split. + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_pad4_dsplit": mk( + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + lds_row_pad=4, + causal_diag_split=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + # Diagonal split with each pipeline lever alone, and BF16 CFVST alone with + # the split and on the persistent grid: each flag-on path has its own pin. + "attention_dense_gfx942/persist_bthkvmin_kdb_dsplit_d128_bf16": mk( + use_cfvst=True, + lds_num_buffers=2, + causal_diag_split=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/persist_bthkvmin_vpf_dsplit_d128_bf16": mk( + use_cfvst=True, + v_reg_prefetch=True, + causal_diag_split=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), + "attention_dense_gfx942/cfvst_dsplit_d128_bf16_causal": mk( + use_cfvst=True, + causal_diag_split=True, + ), + "attention_dense_gfx942/persist_bthkvmin_cfvst_d128_bf16_causal": mk( + use_cfvst=True, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + ), } @@ -317,6 +437,117 @@ def test_attention_dense_gfx942_ir_matches_golden(): ) +def test_pipeline_golden_entries_cover_each_independent_mode_and_llvm_flavor(): + """Pin non-vacuous fixture coverage without generating unverified new hashes. + + This checks the *existing* golden matrix (including the pad4 + diagonal-split + winner) independently of lowering. The main golden test separately verifies + that the current lowered LLVM IR matches each recorded digest. + """ + from rocke.core.ir_golden import GOLDEN_ALIASES, GOLDEN_PROFILE_FLAVORS + + doc = json.loads(_GOLDEN.read_text()) + expected = set(_cases()) + required = { + "attention_dense_gfx942/cfvst_d128_bf16_causal", + "attention_dense_gfx942/persist_bthkvmin_kdb_d128_bf16_causal", + "attention_dense_gfx942/persist_bthkvmin_vpf_d128_bf16_causal", + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_d128_bf16_causal", + "attention_dense_gfx942/persist_kdb_d128_bf16_qbmaj", + "attention_dense_gfx942/persist_vpf_d128_bf16_qbmaj", + "attention_dense_gfx942/persist_kdbvpf_d128_bf16_qbmaj", + "attention_dense_gfx942/lse_persist_bthkvmin_kdb_d128_bf16", + "attention_dense_gfx942/lse_persist_bthkvmin_vpf_d128_bf16", + "attention_dense_gfx942/lse_persist_bthkvmin_kdbvpf_d128_bf16", + "attention_dense_gfx942/persist_bthkvmin_kdbvpf_pad4_dsplit", + "attention_dense_gfx942/persist_bthkvmin_kdb_dsplit_d128_bf16", + "attention_dense_gfx942/persist_bthkvmin_vpf_dsplit_d128_bf16", + "attention_dense_gfx942/cfvst_dsplit_d128_bf16_causal", + "attention_dense_gfx942/persist_bthkvmin_cfvst_d128_bf16_causal", + } + assert required <= expected, sorted(required - expected) + assert set(doc["flavors"]) == set(GOLDEN_PROFILE_FLAVORS) + assert doc["aliases"] == GOLDEN_ALIASES + for flavor, value in doc["flavors"].items(): + cases = value["cases"] + assert set(cases) == expected, ( + f"{flavor}: missing={sorted(expected - set(cases))}, " + f"extra={sorted(set(cases) - expected)}" + ) + for case_id in required: + entry = cases[case_id] + assert len(entry["sha256"]) == 64 and entry["bytes"] > 0, ( + flavor, + case_id, + entry, + ) + + +def test_independent_pipeline_diag_split_changes_lowered_ir(): + """CPU codegen proof for each split-mode pipeline branch. + + Split-only KDB/VPF lowerings are compared with their unsplit twins after + normalizing the symbol, so the split must change the pipeline body itself; + each split mode is also pinned by its own golden SHA-256 entry. + """ + from kernels.gfx942.attention_dense import ( + Gfx942AttentionDenseSpec, + build_attention_dense, + supports_attention_dense, + ) + from rocke.core.lower_llvm import ( + _lower_kernel_to_llvm_python, + _resolve_llvm_flavor, + ) + + base = dict( + batch=1, + seqlen_q=512, + seqlen_kv=512, + num_query_heads=16, + num_kv_heads=4, + head_size=128, + dtype="bf16", + causal=True, + block_n=64, + persistent=True, + num_persistent=304, + persist_decode="bt_hkv_minor", + use_cfvst=True, + ) + fingerprints = {} + for nbuf, vpf, label in ( + (2, False, "kdb"), + (1, True, "vpf"), + (2, True, "kdb_vpf"), + ): + fingerprints[label] = {} + for split in (False, True): + spec = Gfx942AttentionDenseSpec( + **base, + lds_num_buffers=nbuf, + v_reg_prefetch=vpf, + causal_diag_split=split, + ) + ok, why = supports_attention_dense(spec, arch=_ARCH) + assert ok, f"{label} split={split}: {why}" + kd = build_attention_dense(spec, arch=_ARCH) + llvm = _lower_kernel_to_llvm_python( + kd, arch=_ARCH, llvm_flavor=_resolve_llvm_flavor() + ) + # Exclude symbol renaming: require a genuine change to the IR body. + normalized = llvm.replace(kd.name, "KERNEL") + fingerprints[label][split] = hashlib.sha256( + normalized.encode("utf-8") + ).hexdigest() + assert ( + fingerprints[label][True] != fingerprints[label][False] + ), f"{label} split and unsplit unexpectedly emit identical IR" + assert ( + len({fingerprints[mode][True] for mode in fingerprints}) == 3 + ), "independent split-mode pipeline IRs unexpectedly aliased" + + 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/library/tests/test_attention_dense_gfx942_lds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_lds.py index e25911a19c93..f80ae052a3fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_lds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_lds.py @@ -29,7 +29,7 @@ with the library lane: cd rocke/library - PYTHONPATH=../platform/python:. python -m pytest \ + PYTHONPATH=../platform/python:. python -m pytest \\ tests/test_attention_dense_gfx942_lds.py """ @@ -117,6 +117,11 @@ def _lds_pool_bytes(spec) -> int: dict(lds_row_pad=16, v_row_pad=16), dict(use_cfvst=False), dict(use_cfvst=True), + # KDB/VPF are legal only on persistent BF16 D128 CFVST at BN64. + # Keep these within the filtered cohort and ensure anti-vacuity catches drift. + dict(use_cfvst=True, persistent=True, lds_num_buffers=2), + dict(use_cfvst=True, persistent=True, v_reg_prefetch=True), + dict(use_cfvst=True, persistent=True, lds_num_buffers=2, v_reg_prefetch=True), dict(use_cfvst=False, v_row_pad=32), dict(block_m=128), ) @@ -200,6 +205,36 @@ def test_cohort_is_not_vacuous(): assert {s.dtype for s in specs} == {"bf16", "fp16"} assert {s.block_n for s in specs} >= {32, 64, 128} assert {s.resolved_use_cfvst() for s in specs} == {False, True} + assert {s.persistent for s in specs} == {False, True} + assert any( + s.dtype == "bf16" + and s.head_size == 128 + and s.resolved_use_cfvst() + and s.lds_num_buffers == 2 + and s.persistent + and s.block_n == 64 + for s in specs + ), "cohort must exercise the BF16 D128 CFVST KDB allocation" + assert any( + sp.dtype == "bf16" + and sp.head_size == 128 + and sp.resolved_use_cfvst() + and sp.lds_num_buffers == 1 + and sp.v_reg_prefetch + and sp.persistent + and sp.block_n == 64 + for sp in specs + ), "cohort must exercise V register prefetch without K double buffering" + assert any( + sp.dtype == "bf16" + and sp.head_size == 128 + and sp.resolved_use_cfvst() + and sp.lds_num_buffers == 2 + and sp.v_reg_prefetch + and sp.persistent + and sp.block_n == 64 + for sp in specs + ), "cohort must exercise the combined KDB + VPF pipeline" assert {s.lds_row_pad for s in specs} > {_DEFAULT_LDS_ROW_PAD} # v_row_pad is tri-state: None ("ask the policy") is its shipped default. assert {s.v_row_pad for s in specs} > {None} @@ -299,3 +334,35 @@ def test_capacity_boundary_allocation_is_exactly_capacity(): assert ok, why assert _lds_bytes(spec) == _CAPACITY assert _lds_pool_bytes(spec) == _CAPACITY + + +def test_bf16_d128_kdb_and_vpf_have_independent_lds_footprints(): + """KDB allocates one additional K slot; VPF adds no LDS slot. + V LDS remains single-buffered for every mode. + """ + cfvst = _spec( + dtype="bf16", + head_size=128, + block_n=64, + lds_row_pad=8, + use_cfvst=True, + persistent=True, + num_persistent=304, + lds_num_buffers=1, + v_reg_prefetch=False, + ) + kdb = dataclasses.replace(cfvst, lds_num_buffers=2) + vpf = dataclasses.replace(cfvst, v_reg_prefetch=True) + both = dataclasses.replace(kdb, v_reg_prefetch=True) + for spec in (cfvst, kdb, vpf, both): + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, why + assert _lds_bytes(cfvst) == _lds_pool_bytes(cfvst) == 33792 + assert _lds_bytes(kdb) == _lds_pool_bytes(kdb) == 51200 + assert _lds_bytes(vpf) == _lds_pool_bytes(vpf) == 33792 + assert _lds_bytes(both) == _lds_pool_bytes(both) == 51200 + extra_k_slot = 64 * (128 + 8) * 2 + assert _lds_pool_bytes(kdb) - _lds_pool_bytes(cfvst) == extra_k_slot + assert _lds_pool_bytes(both) - _lds_pool_bytes(vpf) == extra_k_slot + assert _lds_pool_bytes(vpf) == _lds_pool_bytes(cfvst) + assert _lds_pool_bytes(both) == _lds_pool_bytes(kdb) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py index 813a033e587a..cbdb5c4e69fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942_numeric.py @@ -32,6 +32,7 @@ from kernels.gfx942.attention_dense import ( _as_gfx942_spec, run_attention_dense_torch, + supports_attention_dense, ) @@ -222,7 +223,10 @@ def _sdpa_reference(q, k, v, scale, *, causal=True, attn_mask=None): def test_dense_numeric_vs_fp32_sdpa(dtype, d, hq, hkv, persistent, causal): import torch - tol = 2e-2 if dtype == "fp16" else 4e-2 + # Element-wise atol + rtol*|ref| rather than a raw max-abs bound: rounding the + # fp32 reference to the output dtype alone consumes most of a fixed absolute + # budget on randn inputs, which would leave the check seed-sensitive. + tol = 2e-3 if dtype == "fp16" else 1e-2 tdt = getattr(torch, _TORCH_DT[dtype]) B, S = 1, 512 scale = 1.0 / math.sqrt(d) @@ -241,9 +245,15 @@ def test_dense_numeric_vs_fp32_sdpa(dtype, d, hq, hkv, persistent, causal): ref = _sdpa_reference(q, k, v, scale, causal=causal) max_abs = (ref - out.float()).abs().max().item() - assert max_abs < tol, ( - f"{dtype} D{d} GQA{hq}/{hkv} {'causal' if causal else 'full'} " - f"{'persist' if persistent else 'default'}: max_abs={max_abs:.3e} >= {tol}" + torch.testing.assert_close( + out.float(), + ref, + atol=tol, + rtol=tol, + msg=lambda m: ( + f"{dtype} D{d} GQA{hq}/{hkv} {'causal' if causal else 'full'} " + f"{'persist' if persistent else 'default'}: max_abs={max_abs:.3e}\n{m}" + ), ) @@ -1243,6 +1253,223 @@ def test_codegen_knob_is_bit_identical_to_default(dtype, d, persistent, override ) +# Persistent grid size for independent CFVST, KDB and VPF tests. At the shipped +# NP (304) every +# shape below has fewer work items than workgroups, so no workgroup ever runs a +# second item and the LDS hand-off between items (the next item's prologue writing +# K slot 0 and V_lds after the previous item's epilogue) is never exercised. +_MULTI_ITEM_NP = 3 + + +def _multi_item_persistent(spec): + """``spec`` on a small persistent grid: more than two work items per workgroup + and a ragged last round, the same preconditions as + ``test_one_persistent_binary_serves_every_shape_and_head_config``.""" + spec = dataclasses.replace(spec, num_persistent=_MULTI_ITEM_NP) + work = (spec.seqlen_q // spec.block_m) * spec.num_query_heads * spec.batch + assert work > 2 * spec.num_persistent, (work, spec.num_persistent) + assert work % spec.num_persistent, (work, spec.num_persistent) + return spec + + +@requires_gfx942_gpu +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", ["bf16", "fp16"]) +@pytest.mark.parametrize("persistent", [False, True]) +@pytest.mark.parametrize( + "sq,skv", + [ + (256, 64), # n_up=1: zero-trip pipeline main loop + (256, 192), # n_up=3: the loop runs, then the last tile is in K slot 0 + (512, 512), # Multiple KV tiles (even n_up: last tile in K slot 1) + ], +) +def test_cfvst_pipeline_strict_numeric(dtype, persistent, sq, skv): + import torch + + B, Hq, Hkv, D = 1, 16, 4, 128 + tol = 1e-2 if dtype == "bf16" else 2e-3 + dt = getattr(torch, _TORCH_DT[dtype]) + scale = 1.0 / math.sqrt(D) + + torch.manual_seed(42) + + q = torch.randn(B, sq, Hq, D, device="cuda", dtype=dt) + k = torch.randn(B, skv, Hkv, D, device="cuda", dtype=dt) + v = torch.randn_like(k) + + # Explicit global causal mask, including non-square Sq/Skv. + qi = torch.arange(sq, device="cuda")[:, None] + ki = torch.arange(skv, device="cuda")[None, :] + mask = ki <= qi + + ref = _sdpa_reference( + q, + k, + v, + scale, + causal=False, + attn_mask=mask, + ) + + base = _as_gfx942_spec(_spec(dtype, D, Hq, Hkv, persistent, batch=B, sq=sq)) + + # BF16 CFVST alone can run on either grid; KDB/VPF require persistent. + if dtype == "bf16": + variants = [ + {"use_cfvst": False, "lds_num_buffers": 1, "v_reg_prefetch": False}, + {"use_cfvst": True, "lds_num_buffers": 1, "v_reg_prefetch": False}, + ] + if persistent: + variants.extend( + [ + {"use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": False}, + {"use_cfvst": True, "lds_num_buffers": 1, "v_reg_prefetch": True}, + {"use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True}, + ] + ) + else: + variants = [{}] # Preserve the existing FP16 path. + + for knobs in variants: + spec = dataclasses.replace(base, seqlen_kv=skv, **knobs) + if persistent: + spec = _multi_item_persistent(spec) + + out = torch.full_like(q, float("nan")) + + run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=scale) + torch.cuda.synchronize() + + assert torch.isfinite(out).all(), ( + f"Non-finite output: {dtype}, {knobs}, " + f"persistent={persistent}, Sq={sq}, Skv={skv}" + ) + + torch.testing.assert_close( + out.float(), + ref, + atol=tol, + rtol=tol, + ) + + +# Representative BF16 D128 CFVST numeric tests covering tile, scheduling, +# output-store, and persistent-grid configurations with independently enabled +# K double buffering and V register prefetch. Not an exhaustive sweep. +# Each row is (id, (Hq, Hkv), Skv, persistent, knobs). +_CFVST_KNOB_ROWS = [ + ("bn32", (16, 4), 512, False, {"block_n": 32}), + ("bm128", (16, 4), 512, False, {"block_m": 128}), + # block_m=64 -> n_up = 1..8 across query blocks, so odd n_up > 1 is covered. + ("bm64", (16, 4), 512, False, {"block_m": 64}), + ("iglp0", (16, 4), 512, False, {"iglp": True}), + ("iglp1", (16, 4), 512, False, {"iglp": True, "iglp_mode": 1}), + ("kmajor", (16, 4), 512, False, {"pv_loop_order": "k_major"}), + ("ostore1", (16, 4), 512, False, {"o_store_width": 1}), + ("ostore2", (16, 4), 512, False, {"o_store_width": 2}), + ("pvprio", (16, 4), 512, False, {"pv_priority": 1}), + ("pvfence", (16, 4), 512, False, {"pv_sched_fence_mask": 0}), + ("persist_intl", (16, 4), 512, True, {"interleave": True}), + ("persist_qbmaj", (16, 4), 512, True, {"persist_decode": "qb_major"}), + ("persist_hkvmaj", (16, 4), 512, True, {"persist_decode": "hkv_major"}), + ("skv_gt_sq", (16, 4), 1024, False, {}), + ("mha", (16, 16), 512, False, {}), + ("gqa8", (32, 4), 512, False, {}), +] + + +# KDB/VPF require BOTH persistent=True and block_n=64. Keep the CFVST-only +# coverage on either grid; never generate an illegal positive pipeline test. +_CFVST_PIPELINE_MODES = ( + (1, False, "cfvst_only"), + (2, False, "kdb_only"), + (1, True, "vpf_only"), + (2, True, "kdb_vpf"), +) +_CFVST_KNOB_CASES = [ + pytest.param(*row[1:], nbuf, vpf, id=f"{row[0]}-{mode}") + for row in _CFVST_KNOB_ROWS + for nbuf, vpf, mode in _CFVST_PIPELINE_MODES + if (nbuf == 1 and not vpf) or (row[3] and row[4].get("block_n", 64) == 64) +] + + +@requires_gfx942_gpu +@pytest.mark.gpu +@pytest.mark.parametrize( + "heads,skv,persistent,knobs,lds_num_buffers,v_reg_prefetch", + _CFVST_KNOB_CASES, +) +def test_cfvst_knob_matrix_numeric( + heads, skv, persistent, knobs, lds_num_buffers, v_reg_prefetch +): + import torch + + hq, hkv = heads + sq, d = 512, 128 + scale = 1.0 / math.sqrt(d) + torch.manual_seed(5) + q = torch.randn(1, sq, hq, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, skv, hkv, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + qi = torch.arange(sq, device="cuda")[:, None] + ki = torch.arange(skv, device="cuda")[None, :] + ref = _sdpa_reference(q, k, v, scale, causal=False, attn_mask=ki <= qi) + + base = _as_gfx942_spec(_spec("bf16", d, hq, hkv, persistent, batch=1, sq=sq)) + spec = dataclasses.replace( + base, + seqlen_kv=skv, + use_cfvst=True, + lds_num_buffers=lds_num_buffers, + v_reg_prefetch=v_reg_prefetch, + **knobs, + ) + if persistent: + spec = _multi_item_persistent(spec) + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, f"gate no longer accepts a row it used to open: {why}" + out = torch.full_like(q, float("nan")) + run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=scale) + torch.cuda.synchronize() + torch.testing.assert_close( + out.float(), + ref, + atol=1e-2, + rtol=1e-2, + ) + # KDB and VPF are performance-only changes. Compare their outputs + # against the same CFVST kernel with both optimizations disabled. + if lds_num_buffers == 2 or v_reg_prefetch: + baseline_spec = dataclasses.replace( + spec, + lds_num_buffers=1, + v_reg_prefetch=False, + ) + + baseline_out = torch.full_like(q, float("nan")) + + run_attention_dense_torch( + spec=baseline_spec, + q=q, + k=k, + v=v, + out=baseline_out, + scale=scale, + ) + torch.cuda.synchronize() + + assert torch.equal( + out.view(torch.int16), + baseline_out.view(torch.int16), + ), ( + f"KDB/VPF changed BF16 output bits: " + f"{knobs}, nbuf={lds_num_buffers}, " + f"vpf={v_reg_prefetch}" + ) + + @requires_gfx942_gpu @pytest.mark.gpu class TestDenseLseNumericGfx942: @@ -1274,6 +1501,56 @@ def test_output_and_lse(self, dtype, d, heads, persistent, causal): assert (lse - lse_ref).abs().max().item() < _tolerance(dtype) assert torch.equal(out, off_out) + @pytest.mark.parametrize( + "sq,skv", [(256, 64), (512, 512)], ids=["zero_trip", "multi_tile"] + ) + @pytest.mark.parametrize( + "persistent,lds_num_buffers,v_reg_prefetch", + [ + (False, 1, False), + (True, 1, False), + (True, 2, False), + (True, 1, True), + (True, 2, True), + ], + ids=[ + "ordinary_cfvst", + "persistent_cfvst", + "persistent_kdb", + "persistent_vpf", + "persistent_kdb_vpf", + ], + ) + def test_cfvst_pipeline_lse_modes( + self, persistent, sq, skv, lds_num_buffers, v_reg_prefetch + ): + """LSE uses the final tile's row max in all independent pipeline modes. + + A stale max shifts LSE by (m_new - m_old) * ln2; at n_up=1 + (zero_trip) the pipeline loop does not run at all. Both grids are + exercised, including the persistent multi-work-item handoff. + """ + spec = dataclasses.replace( + _lse_spec("bf16", 128, 8, 2, persistent, batch=1, sq=512), + seqlen_q=sq, + seqlen_kv=skv, + use_cfvst=True, + lds_num_buffers=lds_num_buffers, + v_reg_prefetch=v_reg_prefetch, + ) + if persistent: + spec = _multi_item_persistent(spec) + scale = 1.0 / math.sqrt(128) + torch.manual_seed(7) + q = torch.randn(1, sq, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, skv, 2, 128, device=q.device, dtype=torch.bfloat16) + v = torch.randn_like(k) + out, lse, off_out = _run_lse_and_off(spec, q, k, v, scale) + ref, lse_ref = _lse_reference(q, k, v, scale, causal=True) + assert (out.float() - ref).abs().max().item() < _tolerance("bf16") + assert (lse - lse_ref).abs().max().item() < 1e-3 + assert torch.equal(out, off_out) + @pytest.mark.parametrize("dtype", ["fp16", "bf16"]) @pytest.mark.parametrize("d", [64, 128]) @pytest.mark.parametrize( @@ -1594,5 +1871,195 @@ def test_lse_runtime_shape_reuse(self, persistent, decode): _DENSE_LAUNCHER_CACHE[key] = entry +# All following cases are within the BF16 D128 causal BN64 persistent gate. +# Unlike the unsplit pipeline, diagonal split has TWO processing loops; K/V +# publication and the running softmax state cross that loop boundary. +@pytest.mark.parametrize( + "lds_num_buffers,v_reg_prefetch,lds_row_pad", + [ + (1, False, 8), # CFVST alone: no pipeline, split loops only. + (2, False, 8), + (1, True, 8), + (2, True, 8), + (2, True, 4), # Measured full optimization (pad4 + dsplit). + ], + ids=["cfvst_only", "kdb", "vpf", "kdb_vpf", "winner_pad4"], +) +@pytest.mark.parametrize( + "sq,skv", + [(256, 64), (512, 192), (512, 512)], + ids=["single_tile", "odd_k_tiles", "split_loop_boundary"], +) +@requires_gfx942_gpu +@pytest.mark.gpu +def test_cfvst_pipeline_diag_split_numeric( + sq, skv, lds_num_buffers, v_reg_prefetch, lds_row_pad +): + """Validate CFVST alone, each pipeline and the pad4/dsplit winner on GPU. + + NP=3 forces work-item LDS reuse and a ragged grid; the same Q/K/V inputs + exercise split and unsplit code paths against a true FP32 causal oracle. + """ + hq, hkv, d = 16, 4, 128 + scale = 1.0 / math.sqrt(d) + torch.manual_seed(47) + q = torch.randn(1, sq, hq, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, skv, hkv, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + qi = torch.arange(sq, device="cuda")[:, None] + ki = torch.arange(skv, device="cuda")[None, :] + ref = _sdpa_reference(q, k, v, scale, causal=False, attn_mask=ki <= qi) + + base = _as_gfx942_spec(_spec("bf16", d, hq, hkv, True, sq=sq)) + split = _multi_item_persistent( + dataclasses.replace( + base, + seqlen_kv=skv, + use_cfvst=True, + lds_num_buffers=lds_num_buffers, + v_reg_prefetch=v_reg_prefetch, + lds_row_pad=lds_row_pad, + causal_diag_split=True, + ) + ) + ok, why = supports_attention_dense(split, arch="gfx942") + assert ok, why + unsplit = dataclasses.replace(split, causal_diag_split=False) + out = torch.full_like(q, float("nan")) + control = torch.full_like(q, float("nan")) + for spec, output in ((split, out), (unsplit, control)): + run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=output, scale=scale) + torch.cuda.synchronize() + assert torch.isfinite(out).all() + torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(control.float(), ref, atol=1e-2, rtol=1e-2) + assert torch.equal(out.view(torch.int16), control.view(torch.int16)), ( + "causal_diag_split changed BF16 output bits in the pipeline: " + f"Sq={sq} Skv={skv} KDB={lds_num_buffers == 2} " + f"VPF={v_reg_prefetch} pad={lds_row_pad}" + ) + + +@pytest.mark.parametrize("sq,skv", [(256, 64), (512, 192), (512, 512)]) +@requires_gfx942_gpu +@pytest.mark.gpu +def test_cfvst_pipeline_pad4_dsplit_lse_numeric(sq, skv): + """The full split pipeline must compute LSE from the FINAL tile's row max.""" + hq, hkv, d = 16, 4, 128 + scale = 1.0 / math.sqrt(d) + torch.manual_seed(53) + q = torch.randn(1, sq, hq, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, skv, hkv, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + spec = _multi_item_persistent( + dataclasses.replace( + _lse_spec("bf16", d, hq, hkv, True, sq=sq), + seqlen_kv=skv, + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + lds_row_pad=4, + causal_diag_split=True, + ) + ) + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, why + out, lse, off_out = _run_lse_and_off(spec, q, k, v, scale) + ref, lse_ref = _lse_reference(q, k, v, scale, causal=True) + torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(lse, lse_ref, atol=1e-3, rtol=1e-3) + assert torch.equal(out, off_out) + + +@pytest.mark.parametrize("block_m", [64, 512]) +@requires_gfx942_gpu +@pytest.mark.gpu +def test_cfvst_pipeline_extreme_legal_block_m_numeric(block_m): + """Keep the BN64 pipeline valid at both ends of the implemented BM tuning axis.""" + hq, hkv, d, sq = 16, 4, 128, 512 + scale = 1.0 / math.sqrt(d) + torch.manual_seed(59) + q = torch.randn(1, sq, hq, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, sq, hkv, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + spec = _multi_item_persistent( + dataclasses.replace( + _as_gfx942_spec(_spec("bf16", d, hq, hkv, True, sq=sq)), + block_m=block_m, + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + ) + ) + ok, why = supports_attention_dense(spec, arch="gfx942") + assert ok, why + out = torch.full_like(q, float("nan")) + run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=scale) + torch.cuda.synchronize() + torch.testing.assert_close( + out.float(), _sdpa_reference(q, k, v, scale), atol=1e-2, rtol=1e-2 + ) + + +@requires_gfx942_gpu +@pytest.mark.gpu +def test_cfvst_kdb_vpf_pipeline_reuses_runtime_shape_binary(): + """Same combined pipeline launcher, different B/S/Hq/Hkv and distinct inputs.""" + from kernels.common.attention_dense_spec import attention_dense_cache_key + from kernels.gfx942.attention_dense import _DENSE_LAUNCHER_CACHE + + shape_rows = [(1, 512, 16, 4), (2, 1024, 32, 8)] + specs = [] + for batch, sq, hq, hkv in shape_rows: + base = _as_gfx942_spec( + _spec( + "bf16", + 128, + hq, + hkv, + True, + batch=batch, + sq=sq, + persist_decode="qb_major", + ) + ) + spec = dataclasses.replace( + base, + num_persistent=3, + use_cfvst=True, + lds_num_buffers=2, + v_reg_prefetch=True, + ) + assert supports_attention_dense(spec, arch="gfx942")[0] + assert spec.runtime_shape + specs.append(spec) + keys = [attention_dense_cache_key(s, arch="gfx942") for s in specs] + assert keys[0] == keys[1], "same runtime-shape kernel must have one cache key" + old_entry = _DENSE_LAUNCHER_CACHE.pop(keys[0], None) + try: + launchers = [] + for i, spec in enumerate(specs): + torch.manual_seed(61 + i) + shape = (spec.batch, spec.seqlen_q, spec.num_query_heads, 128) + kv_shape = (spec.batch, spec.seqlen_kv, spec.num_kv_heads, 128) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + k = torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + out = torch.full_like(q, float("nan")) + scale = 1.0 / math.sqrt(128) + run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=scale) + torch.cuda.synchronize() + torch.testing.assert_close( + out.float(), _sdpa_reference(q, k, v, scale), atol=1e-2, rtol=1e-2 + ) + launchers.append(_launcher_for(spec)) + assert launchers[0] is not None + assert launchers[0] is launchers[1], "pipeline recompiled across runtime shapes" + finally: + _DENSE_LAUNCHER_CACHE.pop(keys[0], None) + if old_entry is not None: + _DENSE_LAUNCHER_CACHE[keys[0]] = old_entry + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index f77b3809d998..9ba48728ef81 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -2520,6 +2520,14 @@ An optimization is done when: good enough. - If compiler flags produce wrong answers, remove them even if another codebase uses them safely. +- If attention PV reads V from LDS with narrow loads, a store-side + transpose that allows wide LDS reads may pay for its extra + registers (`use_cfvst` on gfx942 bf16 D128). +- If the loop waits for the next K/V tile at its barriers and LDS has + room for a second K slot but not a second V slot, double-buffer K + only and prefetch next-tile V into registers (`lds_num_buffers=2`, + `v_reg_prefetch=True`). Each is independent; see the + [gfx942 bf16 D128 case study](../../../library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md). - If the HIP-debug backend disagrees with LLVM-direct on a kernel, one of the two is missing an op lowering — file the issue, don't ignore it. From 3715cddcbce73e77145b9224cc4b62f9f7eb16ab Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sun, 1 Nov 2026 04:06:24 -0500 Subject: [PATCH 2/6] docs(rocke): state BF16 D128 CFVST is the existing default _use_cfvst returns True for every D128 spec regardless of dtype (since #14367), so use_cfvst=None resolves ON for BF16 D128. Correct the case study, dispatch AGENTS.md, prefill README, module docstring and a test comment that said it resolves off / is opt-in. Only K double buffering and V register prefetch are new opt-in levers in this PR. Tests are unchanged; they already assert the real default. --- .../rocke/library/builders/gfx942/attention/prefill/README.md | 4 +--- .../prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md | 4 ++-- .../hip-kernel-provider/rocke/library/dispatch/AGENTS.md | 2 +- .../rocke/library/kernels/gfx942/attention_dense.py | 3 ++- .../rocke/library/tests/test_attention_dense_gfx942.py | 2 +- 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md index 287af145748e..7d8e35ad580c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md @@ -176,7 +176,7 @@ magnitudes live outside the repo per `AGENTS.md`. | Lever | Config | Mechanism | Verdict | |---|---|---|---| | 32×32×8 atom + K-loop doubling | all | CDNA3 has no 32×32×16 fp16/bf16 atom; C-layout is identical so softmax/epilogue port unchanged | **shipped** — enablement | -| Conflict-free V (`perm_b32` store transpose) | **D128 fp16** | V stored transposed → PV A-operand read is one contiguous `ds_read_b64` instead of 4 element-wise `ds_read_u16` | **shipped** — large; identity-preserving | +| Conflict-free V (`perm_b32` store transpose) | **D128** fp16 and bf16 (`use_cfvst=None` resolves on; bf16 since #14367) | V stored transposed → PV A-operand read is one contiguous `ds_read_b64` instead of 4 element-wise `ds_read_u16` | **shipped** — large; identity-preserving | | `exp2_fast` | all | softmax args are provably ≤ 0, so `llvm.exp2`'s guarded range reduction is dead work; bf16 D128 (the last holdout) re-enabled after re-measurement — 0 scratch and lower VGPR on both grids, numerically identical | **shipped** — dominant on the VALU-bound path | | Fused softmax rescale | all | exp2 → accumulate → cast → pack in one pass instead of materializing a full f32 `p_vals` matrix | **shipped** — pure live-range relief, bit-identical | | Per-config `waves_per_eu` | **bf16 D64** → 4 | forces the allocator low enough that a second workgroup co-resides (1 → 2 WG/CU) | **shipped** — large at long sequences | @@ -192,7 +192,6 @@ Measurements, static register and ISA counts, and methodology are in the | Lever | Config | Mechanism | Verdict | |---|---|---|---| -| Conflict-free V at **bf16 D128** (`use_cfvst=True`) | bf16 D128 causal, either grid, with or without `causal_diag_split` | same `perm_b32` store transpose as fp16 D128; fills the `waves_per_eu=2` VGPR budget with little headroom on the `.1k` MFMA schedule | **opt-in** — not a default | | K-only LDS double buffering (`lds_num_buffers=2`) | bf16 D128 CFVST causal, `persistent=True`, `block_n=64` only | allocates two K LDS slots while keeping one V LDS slot | **opt-in** — independent of V register prefetch. Not validated at `block_n=32` | | V register prefetch (`v_reg_prefetch=True`) | bf16 D128 CFVST causal, `persistent=True`, `block_n=64` only | loads V for the next tile into registers before the current PV; V LDS stays single-buffered | **opt-in** — independent of K LDS buffering; can also be combined with `lds_num_buffers=2` | @@ -201,7 +200,6 @@ Measurements, static register and ISA counts, and methodology are in the | Lever | Why it lost | |---|---| | Conflict-free V at **D64** | D64 is VGPR-bound; the register round-trip costs more than the LDS-instruction saving | -| Conflict-free V at **bf16 D128** as a default | fills the `waves_per_eu=2` VGPR budget on the `.1k` MFMA schedule; kept off by default and available as an opt-in (see the Opt-in table) | | `exp2_fast` at **bf16 D128** (former holdout) | initially rejected on a measured spill over the `waves_per_eu=2` cap; **revisited and adopted** (see Adopted table) — re-measurement reads 0 scratch and lower VGPR on both grids, so the original reading no longer reproduces | | `waves_per_eu=3` at **fp16 D64** | reaches 2 WG/CU but loses more ILP than the second workgroup buys | | Drop K/V LDS pads to reach 2 WG/CU at D128 | D128 stays 1 WG/CU even unpadded (register floor co-limits), so it only reintroduces bank conflicts — **catastrophic** | diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md index 4e85c6f0c869..07941199f0f1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md @@ -1,6 +1,6 @@ # gfx942 bf16 D128 causal prefill: CFVST, K LDS double buffering, and V register prefetch -This change makes CFVST available as an opt-in and implements opt-in K LDS Double Buffering and V register pre-fetching on D128 bf16 persistent=true causal prefill attention gfx942 kernels. All three are off by default. +This change implements opt-in K LDS Double Buffering and V register pre-fetching on the CFVST body of D128 bf16 persistent=true causal prefill attention gfx942 kernels. CFVST itself is already the D128 default for both dtypes (`use_cfvst=None` resolves on via `_use_cfvst`, since #14367); this change does not alter it. Both new levers are off by default. ## Problem @@ -482,7 +482,7 @@ The optimization-specific eligibility conditions are enforced by `supports_atten The combined K/V LDS allocation is checked against the gfx942 capacity using `attention_lds_capacity_bytes(arch)` and `_lds_bytes(spec)`. Unsupported configurations are rejected before kernel emission. -For BF16, the default configuration retains `use_cfvst=None` (policy resolves to off), `lds_num_buffers=1`, and `v_reg_prefetch=False`. All three optimizations remain opt-in and do not change the default kernel behavior. +For BF16, the default configuration retains `use_cfvst=None` (policy resolves to on at D128), `lds_num_buffers=1`, and `v_reg_prefetch=False`. K double buffering and V register prefetch remain opt-in and do not change the default kernel behavior. ## gfx942 Barrier / LGKM ISA Evidence diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md index 3025588ecafe..afab4bb103cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md @@ -208,7 +208,7 @@ as a prerequisite when necessary. It can offer CFVST alone, CFVST + KDB, and CFVST + VPF without requiring the triple combination. The triple CFVST + KDB + VPF combination remains reachable through an explicit knob pin or the full tuning space. The BF16 D128 default remains -`use_cfvst=None` (policy resolves to off), `lds_num_buffers=1`, and +`use_cfvst=None` (policy resolves to on at D128), `lds_num_buffers=1`, and `v_reg_prefetch=False`; these tuning choices do not alter the shipped default or make the opt-in dense candidate an automatic dispatch winner. Bare `tuning_id` replay searches production points only; replay a diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index b8147a3eff35..72642e3d4190 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py @@ -140,7 +140,8 @@ single LDS buffer) and reads V element-wise; the remaining perf levers (P2-P4) layer on top. It is validated against an fp32 SDPA reference across the in-scope cohort. -Opt-in bf16 D128 causal CFVST on the persistent grid: ``lds_num_buffers=2`` +Opt-in pipeline levers on the bf16 D128 causal CFVST body (CFVST is the D128 default) +on the persistent grid: ``lds_num_buffers=2`` independently uses two K LDS slots and ``v_reg_prefetch=True`` loads next-tile V registers before current PV. V LDS stays single-buffered for all combinations. All shipped defaults use ``lds_num_buffers=1`` and ``v_reg_prefetch=False``. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index 4c1fdabf94df..741c60713ba1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -1718,7 +1718,7 @@ def test_gfx942_auto_decode_cannot_leak_to_gqa_pair(): dict(lds_row_pad=2), # REJECTED: not a multiple of 4 elements dict(v_row_pad=16, use_v_swizzle=False), # accepted (non-pow2 row: swizzle off) dict(v_row_pad=-4), # REJECTED: negative - # --- private: use_cfvst (legal opt-in for BF16 D128; rejected on D64) --- + # --- private: use_cfvst (D128 default; pinning ON is legal there, rejected on D64) --- dict(dtype="fp16", use_cfvst=False), # accepted: OFF is always legal dict(head_size=64, use_cfvst=True), # REJECTED: policy says off at D64 dict(dtype="bf16", head_size=128, use_cfvst=True), # accepted: policy on at D128 From 63edb97f9748af7907298aca1a0e3eaa2a71bed4 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sun, 1 Nov 2026 04:08:06 -0500 Subject: [PATCH 3/6] docs(rocke): only KDB/VPF are opt-ins in dispatch AGENTS.md CFVST is the existing D128 default (use_cfvst=None resolves on); describe it as the prerequisite body, and list the production modes as CFVST off, CFVST + KDB, CFVST + VPF. --- .../rocke/library/dispatch/AGENTS.md | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md index afab4bb103cd..4006c7fe32e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md @@ -194,18 +194,19 @@ an otherwise invalid configuration can become valid before later knobs are considered. Only `lazy_rescale` is untunable on gfx942, because the kernel body does not read it. -The three performance levers are **independent opt-ins** for supported -BF16 D128 causal attention (`sliding_window=0`, no sinks), with or without -`causal_diag_split`. `use_cfvst=True` enables conflict-free V storage on -either grid; `lds_num_buffers=2` enables **K LDS double buffering only**; and -`v_reg_prefetch=True` enables **V register prefetch only**. KDB and VPF also -require CFVST, `persistent=True` and `block_n=64`, as enforced by the -kernel's validator; turning on either one does not implicitly enable the -other. - -The production tuning space remains **one knob at a time**, using CFVST -as a prerequisite when necessary. It can offer CFVST alone, CFVST + KDB, -and CFVST + VPF without requiring the triple combination. The triple +Two performance levers are **independent opt-ins** for supported BF16 D128 +causal attention (`sliding_window=0`, no sinks), with or without +`causal_diag_split`: `lds_num_buffers=2` enables **K LDS double buffering +only**, and `v_reg_prefetch=True` enables **V register prefetch only**. Both +run on the CFVST body, which is already the D128 default +(`use_cfvst=None` resolves on via `_use_cfvst`; `use_cfvst=False` opts out +on either grid). KDB and VPF also require `persistent=True` and +`block_n=64`, as enforced by the kernel's validator; turning on either one +does not implicitly enable the other. + +The production tuning space remains **one knob at a time** on the default +CFVST body. It offers CFVST off, CFVST + KDB, and CFVST + VPF without +requiring the triple combination. The triple CFVST + KDB + VPF combination remains reachable through an explicit knob pin or the full tuning space. The BF16 D128 default remains `use_cfvst=None` (policy resolves to on at D128), `lds_num_buffers=1`, and From 74e0809b9e81bb2fd7274602a799cd9371d79c46 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sun, 1 Nov 2026 04:15:39 -0500 Subject: [PATCH 4/6] fix(rocke): rebase the 7-way bench on CFVST-on production; gate KDB/VPF C++ parity Since #14367 production resolves BF16 D128 CFVST on, so the bench aborted on all 12 shapes before timing. B1 now pins use_cfvst=False (the CFVST-off reference), B2 is production dispatch, downstream stages inherit the policy, and A pins use_cfvst=False to rebuild the pre-#14367 Step-0 winners (symbol gains the cfvst0 token). All 12 shapes pass the pre-timing guards. Date the case-study results to cc84e262 (pre-#14367) instead of implying they were measured at the PR head. test_attention_ir_cpp_parity now also lowers the 12 gfx942 KDB/VPF golden cases through both engines. --- .../prefill/bench_rocke_7way_specs.py | 57 +++++++++++-------- ...42_dense_cfvst_kdb_vprefetch_case_study.md | 7 ++- .../tests/test_attention_ir_cpp_parity.py | 16 ++++++ 3 files changed, 53 insertions(+), 27 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py index 07439024ca50..40158dd54c54 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py @@ -2,13 +2,16 @@ """Eight-arm persistent gfx942 BF16 D128 attention: Step-0 winner A vs seven-stage B ablation. A uses the seven validated shape-specific Step-0 winners from the -2026-10-28 staged sweep (no CFVST/KDB/VPF). B[1..7] always start from the CURRENT -per-shape production-dispatch baseline, independently of A. Thus this is a -fair new same-source A/B measurement, NOT a claim about cross-commit timings. - - A sweep_winner per-shape best preexisting levers (new pipeline off) - 1 baseline production dispatch (CFVST/KDB/VPF off) - 2 cfvst baseline + CFVST +2026-10-28 staged sweep (no CFVST/KDB/VPF). That sweep predates #14367, so A +pins ``use_cfvst=False`` to rebuild the recorded winner. B[1..7] always start +from the CURRENT per-shape production-dispatch spec, independently of A. Since +#14367, production resolves CFVST ON for BF16 D128, so B1 forces it off to give +the CFVST-off reference and B2 IS production dispatch. Thus this is a fair new +same-source A/B measurement, NOT a claim about cross-commit timings. + + A sweep_winner per-shape best preexisting levers (CFVST/KDB/VPF off) + 1 baseline production dispatch with CFVST forced off (use_cfvst=False) + 2 cfvst production dispatch (CFVST on by policy; KDB/VPF off) 3 cfvst_kdb CFVST + K LDS double buffering 4 cfvst_vpf CFVST + V register prefetch (KDB off) 5 pipeline CFVST + KDB + V register prefetch @@ -82,8 +85,8 @@ STAGE_DESCRIPTIONS = { "A": "Per-shape best existing-lever sweep winner (CFVST/KDB/VPF off)", - "baseline": "Production-dispatch baseline, without CFVST/KDB/VPF", - "cfvst": "Baseline + conflict-free V store transpose", + "baseline": "Production dispatch with CFVST forced off (use_cfvst=False)", + "cfvst": "Production dispatch: CFVST on by policy, KDB/VPF off", "cfvst_kdb": "CFVST + K LDS double buffering (no V prefetch)", "cfvst_vpf": "CFVST + V register prefetch (no K double buffering)", "pipeline": "CFVST + K LDS double buffering + V register prefetch", @@ -104,21 +107,22 @@ "pipeline_pad_dsplit": "pipeline_pad", } +# Field deltas against the production spec. Production leaves use_cfvst=None, +# which _use_cfvst resolves ON at BF16 D128 (#14367); dispatch canonicalizes a +# redundant use_cfvst=True pin back to None, so only the CFVST-off stage pins it. STAGE_OVERRIDES = { "A": {}, # Separate sweep-winner specification; not a B stage. - "baseline": {}, - "cfvst": {"use_cfvst": True}, - "cfvst_kdb": {"use_cfvst": True, "lds_num_buffers": 2}, - "cfvst_vpf": {"use_cfvst": True, "v_reg_prefetch": True}, - "pipeline": {"use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True}, + "baseline": {"use_cfvst": False}, + "cfvst": {}, + "cfvst_kdb": {"lds_num_buffers": 2}, + "cfvst_vpf": {"v_reg_prefetch": True}, + "pipeline": {"lds_num_buffers": 2, "v_reg_prefetch": True}, "pipeline_pad": { - "use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True, "lds_row_pad": 4, }, "pipeline_pad_dsplit": { - "use_cfvst": True, "lds_num_buffers": 2, "v_reg_prefetch": True, "lds_row_pad": 4, @@ -192,7 +196,7 @@ o_store_width=4, lazy_rescale=True, interleave=False, - use_cfvst=None, + use_cfvst=False, # the sweep predates #14367: its winners ran CFVST off. use_exp2_fast=None, use_v_swizzle=None, v_row_pad=None, @@ -213,10 +217,13 @@ (16, 4096, 16): dict(waves_per_eu=2, pv_priority=3), } -# Expected ORIGINAL kernel symbols are used for auditing only. The current -# emitter may produce a different symbol if policy or implementation changed. +# Expected kernel symbols, used for auditing only: the recorded Step-0 winner +# symbols plus the ``cfvst0`` token, which kernel_name() adds because A now pins +# use_cfvst=False explicitly (the recorded winners resolved CFVST off from None, +# before #14367). The current emitter may produce a different symbol if policy +# or implementation changed. _A_PREFIX = "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_" -_A_SUFFIX = "_krowpad4_iglp1_iglpm1_pvkmaj_prio" +_A_SUFFIX = "_krowpad4_cfvst0_iglp1_iglpm1_pvkmaj_prio" EXPECTED_A_KERNEL = { (1, 4096, 8): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "1_dsplit", (1, 4096, 16): _A_PREFIX + "bthkvmin_gfx942_wpe2" + _A_SUFFIX + "1_dsplit", @@ -1151,13 +1158,14 @@ def _run_shapes( ) if ( - production.resolved_use_cfvst() + not production.resolved_use_cfvst() or production.lds_num_buffers != 1 or production.v_reg_prefetch ): raise RuntimeError( - f"{shape.label}: production dispatch already enables some " - "CFVST/KDB/VPF pipeline knobs. Cannot label it a clean baseline." + f"{shape.label}: production dispatch is not CFVST-on with KDB/VPF " + "off. B2 must equal production dispatch; refusing a misleading " + "ablation." ) expected_steps = ( ("baseline", "cfvst", {"use_cfvst"}), @@ -1187,7 +1195,8 @@ def _run_shapes( bool(spec.v_reg_prefetch), ) expected_pipeline = ( - bool(expected.get("use_cfvst", False)), + # Unpinned use_cfvst resolves ON at BF16 D128 (_use_cfvst). + bool(expected.get("use_cfvst", True)), int(expected.get("lds_num_buffers", 1)), bool(expected.get("v_reg_prefetch", False)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md index 07941199f0f1..4ddb6038583c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/gfx942_dense_cfvst_kdb_vprefetch_case_study.md @@ -168,6 +168,7 @@ These results are consistent with reduced synchronization overhead and improved ## Results +> **Dated results.** Every timing, static count and trace below was measured at `cc84e2621e8fd9ce83f9bf43dbc1564c690c290d`, before this branch merged #14367. At that commit production dispatch resolved BF16 D128 CFVST **off**, so B1 was production dispatch and B2 added CFVST. Since #14367 production resolves CFVST **on**: the harness now pins `use_cfvst=False` for B1 and A, and B2 *is* production dispatch. The arm definitions (resolved CFVST/KDB/VPF per arm) are unchanged, but kernel sources have changed since `cc84e262`, so these numbers are not re-measured at the PR head and `--verify-provenance` against it reports the source changes. All workloads use **BF16, Hq=32, D=128, persistent=True, causal attention** on (gfx942). @@ -239,8 +240,8 @@ For each of the twelve shapes (the five shapes without a Step-0 winner skip `A` | Arm | Variant | `use_cfvst` | `lds_num_buffers` (K) | `v_reg_prefetch` | `lds_row_pad` | `causal_diag_split` | |---:|---|---:|---:|---:|---|---| | A | `A` (Step-0 winner) | `False` | 1 | `False` | Shape-specific winner | Shape-specific winner | - | B1 | `baseline` | `False` | 1 | `False` | Production dispatch | Production dispatch | - | B2 | `cfvst` | `True` | 1 | `False` | Production dispatch | Production dispatch | + | B1 | `baseline` | `False` (pinned) | 1 | `False` | Production dispatch | Production dispatch | + | B2 | `cfvst` | `True` (production policy since #14367) | 1 | `False` | Production dispatch | Production dispatch | | B3 | `cfvst_kdb` | `True` | 2 | `False` | Production dispatch | Production dispatch | | B4 | `cfvst_vpf` | `True` | 1 | `True` | Production dispatch | Production dispatch | | B5 | `pipeline` | `True` | 2 | `True` | Production dispatch | Production dispatch | @@ -295,7 +296,7 @@ For each of the twelve shapes (the five shapes without a Step-0 winner skip `A` | Component | Configuration | | :--- | :--- | | Hardware target | gfx942, 304 CUs | -| Benchmarked source HEAD | `cc84e2621e8fd9ce83f9bf43dbc1564c690c290d` (clean working tree; `--verify-provenance` passed) | +| Benchmarked source HEAD | `cc84e2621e8fd9ce83f9bf43dbc1564c690c290d` (clean working tree; `--verify-provenance` passed at that commit; predates the #14367 merge, see *Dated results*) | | Kernel source SHA-256 | `8bf80337c529fdab6e3adf65d942bf423eac9623b324c6cd650d0987adb71d2e` | | Benchmark harness SHA-256 | `888b484dab6a0dc76e5544de1660047e03956b10f9f898cf8cb3cb228ec61130` | | ROCm stack | 10.0.0 | diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py index b685bbbb0e9e..791703b13334 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py @@ -174,6 +174,22 @@ def test_attention_ir_cpp_python_byte_identity(): } ) + # Opt-in K LDS double buffering / V register prefetch (bf16 D128 persistent + # CFVST): the gfx942 golden's pipeline cases, which the harness and the LSE + # cases above build with both levers off. + import test_attention_dense_gfx942_golden as gfx942_golden + + pipeline = { + cid: build + for cid, build in gfx942_golden._cases().items() + if "_kdb" in cid or "_vpf" in cid + } + assert pipeline, "gfx942 golden declares no KDB/VPF cases" + for cid, build in pipeline.items(): + cases.append( + {"family": "attention_dense", "case_id": cid, "arch": "gfx942", "build": build} + ) + prev = os.environ.get("ROCKE_CPP_STRICT") os.environ["ROCKE_CPP_STRICT"] = "1" mism = [] From 41cad87666a25e98e957829d303e3e6cb25cfdd0 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sun, 1 Nov 2026 04:16:41 -0500 Subject: [PATCH 5/6] style(rocke): black test_attention_ir_cpp_parity.py --- .../rocke/library/tests/test_attention_ir_cpp_parity.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py index 791703b13334..e482aca798e6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_ir_cpp_parity.py @@ -187,7 +187,12 @@ def test_attention_ir_cpp_python_byte_identity(): assert pipeline, "gfx942 golden declares no KDB/VPF cases" for cid, build in pipeline.items(): cases.append( - {"family": "attention_dense", "case_id": cid, "arch": "gfx942", "build": build} + { + "family": "attention_dense", + "case_id": cid, + "arch": "gfx942", + "build": build, + } ) prev = os.environ.get("ROCKE_CPP_STRICT") From 6c567235e1ad7d19f3f454ba3fadc814a231ad28 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sun, 1 Nov 2026 04:23:04 -0500 Subject: [PATCH 6/6] fix(rocke): keep historical Step-0 A symbols separate from the rebuilt-A guard EXPECTED_A_KERNEL again holds the seven recorded symbols verbatim, so A_original_expected_kernel / A_kernel_name_matches_original stay truthful (the rebuilt symbol now differs by the cfvst0 token). The guard compares against CURRENT_EXPECTED_A_KERNEL (historical + cfvst0), which is also exported alongside a separate match flag. --- .../prefill/bench_rocke_7way_specs.py | 43 +++++++++++++++---- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py index 40158dd54c54..57bd628a2d82 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/bench_rocke_7way_specs.py @@ -217,13 +217,10 @@ (16, 4096, 16): dict(waves_per_eu=2, pv_priority=3), } -# Expected kernel symbols, used for auditing only: the recorded Step-0 winner -# symbols plus the ``cfvst0`` token, which kernel_name() adds because A now pins -# use_cfvst=False explicitly (the recorded winners resolved CFVST off from None, -# before #14367). The current emitter may produce a different symbol if policy -# or implementation changed. +# The seven HISTORICAL Step-0 winner symbols, exactly as recorded by the sweep +# (pre-#14367, use_cfvst=None resolving off). Exported verbatim for provenance. _A_PREFIX = "rocke_attention_dense_d128_bn64_bf16_causal_lazyrs_persist304_" -_A_SUFFIX = "_krowpad4_cfvst0_iglp1_iglpm1_pvkmaj_prio" +_A_SUFFIX = "_krowpad4_iglp1_iglpm1_pvkmaj_prio" EXPECTED_A_KERNEL = { (1, 4096, 8): _A_PREFIX + "bthkvmin_gfx942_wpe3" + _A_SUFFIX + "1_dsplit", (1, 4096, 16): _A_PREFIX + "bthkvmin_gfx942_wpe2" + _A_SUFFIX + "1_dsplit", @@ -234,6 +231,19 @@ (16, 4096, 16): _A_PREFIX + "gfx942_wpe2" + _A_SUFFIX + "3_dsplit", } + +def _current_a_symbol(historical: str) -> str: + """The symbol the rebuilt A must compile to at this checkout: the historical + symbol plus the ``cfvst0`` token kernel_name() emits for the explicit + use_cfvst=False pin A now needs (CFVST resolves on by default since #14367). + The guard compares against this; provenance keeps the historical name too.""" + assert historical.count("_krowpad4_") == 1, historical + return historical.replace("_krowpad4_", "_krowpad4_cfvst0_") + + +CURRENT_EXPECTED_A_KERNEL = { + key: _current_a_symbol(name) for key, name in EXPECTED_A_KERNEL.items() +} assert set(EXPECTED_A_KERNEL) == set(WINNING_LEVERS) assert not set(MISSING_SWEEP_WINNERS) & set(WINNING_LEVERS) @@ -1143,10 +1153,11 @@ def _run_shapes( if not no_winner: original_a_name = EXPECTED_A_KERNEL[key] actual_a_name = specs["A"].kernel_name() - if actual_a_name != original_a_name: + if actual_a_name != CURRENT_EXPECTED_A_KERNEL[key]: raise RuntimeError( f"{shape.label}: rebuilt A is not the recorded Step-0 winner; " - f"recorded={original_a_name}, rebuilt={actual_a_name}" + f"expected={CURRENT_EXPECTED_A_KERNEL[key]} " + f"(historical={original_a_name}), rebuilt={actual_a_name}" ) if ( specs["A"].resolved_use_cfvst() @@ -1247,10 +1258,18 @@ def describe_spec(spec): else "per-shape Step-0 lever-sweep winner, independently rebuilt" ), "A_original_expected_kernel": original_a_name, + "A_current_expected_kernel": ( + None if no_winner else CURRENT_EXPECTED_A_KERNEL[key] + ), "A_current_kernel": actual_a_name, "A_kernel_name_matches_original": ( None if no_winner else actual_a_name == original_a_name ), + "A_kernel_name_matches_current_expected": ( + None + if no_winner + else actual_a_name == CURRENT_EXPECTED_A_KERNEL[key] + ), "A_overrides": a_overrides, "B_stage_baseline": "current production dispatch", "stage_descriptions": { @@ -1731,9 +1750,12 @@ def fmt(r, effect): "seqlen": sh.seqlen, "hkv": sh.hkv, "A_original_kernel": EXPECTED_A_KERNEL[key], + "A_current_expected_kernel": CURRENT_EXPECTED_A_KERNEL[key], "A_compiled_kernel": r["A_kernel"], "A_kernel_name_matches_original": r["A_kernel"] == EXPECTED_A_KERNEL[key], + "A_kernel_name_matches_current_expected": r["A_kernel"] + == CURRENT_EXPECTED_A_KERNEL[key], "B7_kernel": r["pipeline_pad_dsplit_kernel"], "A_median_ms": r["A_median_ms"], "B7_median_ms": r["pipeline_pad_dsplit_median_ms"], @@ -1750,8 +1772,10 @@ def fmt(r, effect): "seqlen", "hkv", "A_original_kernel", + "A_current_expected_kernel", "A_compiled_kernel", "A_kernel_name_matches_original", + "A_kernel_name_matches_current_expected", "B7_kernel", "A_median_ms", "B7_median_ms", @@ -1807,6 +1831,9 @@ def fmt(r, effect): "A_config_matches_newer_seven_winner_report": True, "no_A_comparison_for_shapes_without_step0_winner": True, "A_expected_kernels": {str(k): v for k, v in EXPECTED_A_KERNEL.items()}, + "A_current_expected_kernels": { + str(k): v for k, v in CURRENT_EXPECTED_A_KERNEL.items() + }, "B_stages_derive_only_from_production": True, "stage_descriptions": STAGE_DESCRIPTIONS, "stage_overrides": STAGE_OVERRIDES,