diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/sdpa_fwd_free_stride_verify.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/sdpa_fwd_free_stride_verify.py new file mode 100644 index 000000000000..cbdff0b58b2d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/sdpa_fwd_free_stride_verify.py @@ -0,0 +1,257 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""On-GPU numeric verify for the layout-general SDPA forward (ticket 2). + +Builds ``kernels.common.fmha_mfma`` with ``free_stride=True`` for a chosen +contiguous layout (BSHD or BHSD), lays Q/K/V/O out in that layout, launches +through the HIP runtime, and compares the output against the torch-free numpy +dense-attention reference. The +point of the test is layout invariance: the same logical problem must give the +same result whether the operands are stored BSHD or BHSD, because the kernel +reads them through runtime ``(batch, head, token)`` strides plus a per-batch +base-pointer rebase. + +It compiles through the comgr / LLVM path (``compile_kernel`` with +``backend="python"``), not the hipcc C++ path: the per-batch rebased O store is +a derived write pointer, which the hipcc backend emits as ``const`` (it rejects +the store), while the LLVM path has no const-qualified pointers. It must run on +a device matching ``--arch``. + + PYTHONPATH=platform/python python3 -m builders.common.sdpa_fwd_free_stride_verify \\ + --arch gfx950 --layout bshd --seqlen-q 128 --seqlen-k 128 --head-size 64 \\ + --heads 8 --kv-heads 2 + +Run both layouts and confirm they agree with the reference (and each other). +The accumulation order of the f32 matmul chain differs from numpy, so parity is +judged within a tolerance, not bit-for-bit. +""" + +from __future__ import annotations + +import argparse +import ctypes +import math +import struct + +import numpy as np + +from rocke.core.arch import ArchTarget +from rocke.helpers.compile import compile_kernel +from rocke.numeric.references import dense_attention_reference +from rocke.runtime.hip_module import Runtime + +from kernels.common._fmha_common import FmhaCommonSpec, FmhaShape +from kernels.common._sdpa_fwd_strides import contiguous_strides +from kernels.common.fmha_mfma import ( + FmhaMfmaSpec, + build_fmha_fwd_mfma, + fmha_fwd_mfma_grid, + is_valid_spec, +) + +_NP_DTYPE = {"f16": np.float16, "bf16": np.float16} # bf16 modeled as f16 bits here + + +def _to_physical(logical: np.ndarray, layout: str) -> np.ndarray: + """C-contiguous physical buffer for ``logical`` (B, H, S, D) in ``layout``.""" + if layout == "bhsd": + return np.ascontiguousarray(logical) + if layout == "bshd": # [B, S, H, D] + return np.ascontiguousarray(logical.transpose(0, 2, 1, 3)) + raise ValueError(f"unsupported layout {layout!r}") + + +def _from_physical(phys_flat: np.ndarray, layout: str, shape) -> np.ndarray: + """Logical (B, H, S, D) view of a flat physical buffer in ``layout``.""" + B, H, S, D = shape + if layout == "bhsd": + return phys_flat.reshape(B, H, S, D) + return phys_flat.reshape(B, S, H, D).transpose(0, 2, 1, 3) # bshd + + +def _reference(ql, kl, vl, *, causal): + """Reference attention over logical (B, H, S, D) operands; the reference + consumes one batch of shape (S, H, D) with KV heads expanded to Q heads.""" + B, Hq = ql.shape[0], ql.shape[1] + Hk = kl.shape[1] + out = [] + for b in range(B): + if Hk != Hq: + kb = np.repeat(kl[b], Hq // Hk, axis=0) + vb = np.repeat(vl[b], Hq // Hk, axis=0) + else: + kb, vb = kl[b], vl[b] + out.append( + dense_attention_reference( + ql[b].transpose(1, 0, 2), + kb.transpose(1, 0, 2), + vb.transpose(1, 0, 2), + causal=causal, + out_dtype=np.float16, + ) + ) + return np.stack(out) # (B, S, Hq, D) + + +def run_verify( + *, + arch, + layout, + heads, + head_size, + seqlen_q, + seqlen_k, + batch, + kv_heads=0, + causal=False, + dtype="f16", + tol=2e-2, + verbose=True, +): + """Build fmha_mfma(free_stride=True) for ``layout``, launch it, and compare + against the numpy reference. Returns ``(passed, max_abs, bad, total)``; + raises ``ValueError`` on an invalid spec. Shared by the CLI and the suite's + GPU-gated numeric test.""" + Hq, D, Hk = heads, head_size, (kv_heads or heads) + Sq, Sk, B = seqlen_q, seqlen_k, batch + + common = FmhaCommonSpec( + FmhaShape(head_size=D, num_query_heads=Hq, num_kv_heads=Hk), + dtype=dtype, + mask_mode="causal" if causal else "none", + ) + spec = FmhaMfmaSpec(common=common, seqlen_q=Sq, seqlen_k=Sk, free_stride=True) + ok, why = is_valid_spec(spec, arch) + if not ok: + raise ValueError(f"invalid fmha_mfma free-stride spec: {why}") + + # comgr / LLVM path (backend="python"): the free-stride forward is + # Python-only for now, and the LLVM path has no const-qualified pointers, + # so the per-batch rebased O store lowers cleanly (unlike the hipcc path). + art = compile_kernel( + build_fmha_fwd_mfma(spec, arch=arch), arch=arch, backend="python" + ) + if verbose: + print(f"[{arch}] built {art.kernel_name} ({art.hsaco_bytes} B)") + + np_dt = _NP_DTYPE[dtype] + rng = np.random.default_rng(0xA11E) + ql = (rng.standard_normal((B, Hq, Sq, D)) * 0.3).astype(np_dt) + kl = (rng.standard_normal((B, Hk, Sk, D)) * 0.3).astype(np_dt) + vl = (rng.standard_normal((B, Hk, Sk, D)) * 0.3).astype(np_dt) + + q_phys = _to_physical(ql, layout) + k_phys = _to_physical(kl, layout) + v_phys = _to_physical(vl, layout) + o_phys = np.zeros_like(q_phys) + + q_b, q_h, q_t = contiguous_strides(layout, heads=Hq, seqlen=Sq, head_size=D) + k_b, k_h, k_t = contiguous_strides(layout, heads=Hk, seqlen=Sk, head_size=D) + v_b, v_h, v_t = k_b, k_h, k_t # V has K's shape/layout + o_b, o_h, o_t = q_b, q_h, q_t # O has Q's shape/layout + scale_log2 = float(1.0 / math.sqrt(D) * math.log2(math.e)) + + rt = Runtime() + module = rt.load_module(art.hsaco) + fn = module.get_function(art.kernel_name) + + def u8(a): + return (ctypes.c_uint8 * int(a.nbytes)).from_buffer(np.ascontiguousarray(a)) + + qd, kd, vd, od = (rt.alloc(x.nbytes) for x in (q_phys, k_phys, v_phys, o_phys)) + rt.memcpy_h2d(qd, u8(q_phys), q_phys.nbytes) + rt.memcpy_h2d(kd, u8(k_phys), k_phys.nbytes) + rt.memcpy_h2d(vd, u8(v_phys), v_phys.nbytes) + rt.memset(od, 0, o_phys.nbytes) + + # fmha_mfma kernarg order: Q,K,V,O, scale_log2, seqlen_q, seqlen_k, the + # eight token/head strides, then the four free-stride batch strides. + packed = struct.pack( + " tol)) + if verbose: + tag = "PASS" if max_abs <= tol else "FAIL" + print( + f"[{arch}] layout={layout} Sq={Sq} Sk={Sk} D={D} Hq={Hq} Hk={Hk} " + f"causal={causal}: max_abs_diff={max_abs:.3e} bad={bad}/{got.size} " + f"tol={tol:.0e} -> {tag}" + ) + return max_abs <= tol, max_abs, bad, got.size + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--arch", default="gfx950") + p.add_argument("--layout", choices=("bhsd", "bshd"), default="bshd") + p.add_argument("--seqlen-q", type=int, default=128) + p.add_argument("--seqlen-k", type=int, default=128) + p.add_argument("--head-size", type=int, default=64) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--kv-heads", type=int, default=0, help="0 -> MHA (== heads)") + p.add_argument("--batch", type=int, default=2) + p.add_argument("--causal", action="store_true") + p.add_argument("--dtype", choices=("f16", "bf16"), default="f16") + p.add_argument("--tol", type=float, default=2e-2) + args = p.parse_args() + + try: + passed, _max_abs, _bad, _total = run_verify( + arch=args.arch, + layout=args.layout, + heads=args.heads, + kv_heads=args.kv_heads, + head_size=args.head_size, + seqlen_q=args.seqlen_q, + seqlen_k=args.seqlen_k, + batch=args.batch, + causal=args.causal, + dtype=args.dtype, + tol=args.tol, + verbose=True, + ) + except ValueError as e: + print(f"[{args.arch}] {e}") + return 2 + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_sdpa_fwd_strides.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_sdpa_fwd_strides.py new file mode 100644 index 000000000000..dcb80fbdf703 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_sdpa_fwd_strides.py @@ -0,0 +1,41 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host-side element strides for the free-stride SDPA forward ABI. + +The generalized forward kernel takes per-tensor ``(batch, head, token)`` +element strides and assumes a unit head-dim stride (the ``StridedTensorAddr`` +contract). This module is the single source of truth for the stride values the +launcher feeds in for the two contiguous 4D orders; it emits no kernel IR. + +General (non-contiguous) and packed-QKV strides are supplied by the caller +(hipDNN hands them straight from the framework tensors), so only the contiguous +``bhsd`` / ``bshd`` orders are derived here. The head dim is always innermost +with unit stride, so it is not part of the returned triple. +""" + +from __future__ import annotations + +_LAYOUTS = ("bhsd", "bshd") + + +def contiguous_strides(layout: str, *, heads: int, seqlen: int, head_size: int): + """``(stride_batch, stride_head, stride_token)`` element strides of one + contiguous 4D tensor. + + ``layout`` is ``"bhsd"`` ([B, H, S, D], heads-major) or ``"bshd"`` + ([B, S, H, D], tokens-major). ``heads`` / ``seqlen`` are this tensor's own + counts (query vs KV tensors differ), ``head_size`` is D. The head dim is + innermost with unit stride and is not returned. + """ + if layout not in _LAYOUTS: + raise ValueError(f"layout must be one of {_LAYOUTS}, got {layout!r}") + for name, val in (("heads", heads), ("seqlen", seqlen), ("head_size", head_size)): + if val <= 0: + raise ValueError(f"{name} must be > 0, got {val}") + h, s, d = heads, seqlen, head_size + if layout == "bhsd": + # [B, H, S, D]: batch = H*S*D, head = S*D, token(seq) = D + return (h * s * d, s * d, d) + # bshd -- [B, S, H, D]: batch = S*H*D, head = D, token(seq) = H*D + return (s * h * d, d, h * d) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/fmha_mfma.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/fmha_mfma.py index 44f0cd3709d4..cae16f2424cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/fmha_mfma.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/fmha_mfma.py @@ -36,7 +36,7 @@ from dataclasses import dataclass from typing import Tuple -from rocke.core.ir import KernelDef +from rocke.core.ir import I64, KernelDef from rocke.helpers.mfma_attention import ( MFMA_ATTN_BLOCK_K, MFMA_ATTN_BLOCK_M, @@ -45,6 +45,9 @@ from rocke.helpers.spec import kernel_name_join from ._fmha_common import FmhaCommonSpec, FmhaKernelBuilder, validate_common_spec +# Element size (bytes) for the free-stride per-batch base-pointer rebase. +_ELEM_BYTES = {"f16": 2, "fp16": 2, "bf16": 2} + __all__ = [ "FmhaMfmaSpec", @@ -68,10 +71,15 @@ class FmhaMfmaSpec: seqlen_q: int seqlen_k: int name: str = "rocke_fmha_fwd_mfma" + # Layout-general addressing: per-tensor batch strides + a base-pointer rebase + # instead of the (BSHD-contiguous) seqlen fold, so BSHD/BHSD/general/packed + # all work. Default False keeps the shipped kernel byte-identical; the + # free-stride path is Python-only (outside the byte-identity gate) for now. + free_stride: bool = False def kernel_name(self) -> str: s = self.common.shape - return kernel_name_join( + parts = [ self.name, f"H{s.head_size}", f"HQ{s.num_query_heads}", @@ -80,7 +88,10 @@ def kernel_name(self) -> str: f"Q{self.seqlen_q}", f"K{self.seqlen_k}", self.common.mask_mode, - ) + ] + if self.free_stride: + parts.append("fs") + return kernel_name_join(*parts) def _mma_family(arch: str) -> str: @@ -175,8 +186,12 @@ def is_valid_spec(spec: FmhaMfmaSpec, arch: str = "gfx950") -> Tuple[bool, str]: return True, "ok" -def _declare_params(kb: FmhaKernelBuilder) -> None: - """The MFMA FMHA kernel ABI (shared between build + sig).""" +def _declare_params(kb: FmhaKernelBuilder, *, free_stride: bool = False) -> None: + """The MFMA FMHA kernel ABI (shared between build + sig). + + ``free_stride`` appends the per-tensor batch strides the layout-general + addressing path reads; the default (seqlen-fold) ABI is unchanged. + """ kb.add_tensor("Q", readonly=True) kb.add_tensor("K", readonly=True) kb.add_tensor("V", readonly=True) @@ -185,6 +200,9 @@ def _declare_params(kb: FmhaKernelBuilder) -> None: kb.add_scalar("seqlen_q", "i32") kb.add_scalar("seqlen_k", "i32") kb.add_strides("q", "k", "v", "o") + if free_stride: + for t in ("q", "k", "v", "o"): + kb.add_scalar(f"stride_{t}_batch", "i32") def build_fmha_fwd_mfma(spec: FmhaMfmaSpec, arch: str = "gfx950") -> KernelDef: @@ -235,7 +253,7 @@ def build_fmha_fwd_mfma(spec: FmhaMfmaSpec, arch: str = "gfx950") -> KernelDef: kb = FmhaKernelBuilder(spec.kernel_name(), s) kb.block_size(wave_size) # one wave per CTA (64 on CDNA, 32 on RDNA) - _declare_params(kb) + _declare_params(kb, free_stride=spec.free_stride) kb.decode_grid(has_batch_axis=True) b = kb.builder @@ -250,6 +268,56 @@ def build_fmha_fwd_mfma(spec: FmhaMfmaSpec, arch: str = "gfx950") -> KernelDef: q_tile_idx = kb.q_token # reuses block_id_x; semantically a tile, not a token q_tile_local = b.mul(q_tile_idx, b.const_i32(MFMA_ATTN_BLOCK_M)) + if spec.free_stride: + # Layout-general path (Python-only; outside the byte-identity gate). + # Rebase each tensor's base pointer to the CTA's batch slice; the body + # then addresses head/token through the runtime strides, so BSHD, BHSD, + # general B/H/S strides and packed QKV all work with batch carried by + # the pointer. The default seqlen-fold path below is byte-identical and + # is NOT entered for free-stride specs. Safe because the body loads + # through raw pointers (no buffer span is sized from the base). + elem = _ELEM_BYTES[s.dtype] + + def _batch_base(name): + off = b.mul( + b.zext(batch_idx, I64), + b.zext(kb.scalar(f"stride_{name}_batch"), I64), + ) + return b.global_ptr_add( + kb.tensor(name.upper()), b.mul(off, b.const_i64(elem)) + ) + + masked = s.mask_mode in ("causal", "sliding_window") + mfma_attention_fwd_inner_body( + b, + Q=_batch_base("q"), + K=_batch_base("k"), + V=_batch_base("v"), + O=_batch_base("o"), + head_size=s.head_size, + seqlen_k=seqlen_k, + q_tile_base=q_tile_local, + q_pos_base=q_tile_local if masked else None, + head_idx=head_idx, + kv_head_idx=kv_head_idx, + stride_q_token=kb.stride_token("q"), + stride_q_head=kb.stride_head("q"), + stride_k_token=kb.stride_token("k"), + stride_k_head=kb.stride_head("k"), + stride_v_token=kb.stride_token("v"), + stride_v_head=kb.stride_head("v"), + stride_o_token=kb.stride_token("o"), + stride_o_head=kb.stride_head("o"), + scale_log2=kb.scalar("scale_log2"), + dtype=s.dtype, + mask_mode=s.mask_mode, + sliding_window=s.sliding_window, + causal_ctx_offset=b.const_i32(0), + arch=arch, + ) + b.ret() + return kb.kernel + # Per-batch shift in rows (for Q and O). The helper multiplies # the row index by stride_q_token / stride_o_token internally, so # we only need to pass the row offset -- the previous duplicate @@ -321,5 +389,5 @@ def fmha_fwd_mfma_grid(spec: FmhaMfmaSpec, *, batch: int) -> Tuple[int, int, int def fmha_fwd_mfma_signature(spec: FmhaMfmaSpec): kb = FmhaKernelBuilder("rocke_fmha_fwd_mfma_sig_probe", spec.common) - _declare_params(kb) + _declare_params(kb, free_stride=spec.free_stride) return kb.signature() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride.py new file mode 100644 index 000000000000..a52d932e9480 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride.py @@ -0,0 +1,75 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU build tests for the fmha_mfma free-stride (layout-general) path. + +fmha_mfma gains an opt-in ``free_stride`` mode: per-tensor batch strides plus a +base-pointer rebase so one kernel reads any layout (BSHD / BHSD / general / +packed), replacing the BSHD-contiguous seqlen fold. The default path must stay +byte-identical, so these tests lock the ABI contract: ``free_stride=False`` is +unchanged and ``free_stride=True`` appends exactly the four batch strides and +tags the kernel name. + +On-GPU numeric parity is the verify harness +(``builders/common/sdpa_fwd_free_stride_verify.py``); byte-identity of the +default path is the repo's own byte-identity gate. +""" + +from __future__ import annotations + +import pytest + +from kernels.common._fmha_common import FmhaCommonSpec, FmhaShape +from kernels.common.fmha_mfma import ( + FmhaMfmaSpec, + build_fmha_fwd_mfma, + fmha_fwd_mfma_signature, +) + +_BATCH = ("stride_q_batch", "stride_k_batch", "stride_v_batch", "stride_o_batch") + + +def _spec(free_stride=False, **over): + shape = FmhaShape( + head_size=over.pop("head_size", 64), + num_query_heads=over.pop("num_query_heads", 8), + num_kv_heads=over.pop("num_kv_heads", 8), + ) + common = FmhaCommonSpec(shape=shape, mask_mode=over.pop("mask_mode", "none")) + return FmhaMfmaSpec( + common=common, seqlen_q=128, seqlen_k=128, free_stride=free_stride + ) + + +def test_default_abi_has_no_batch_strides(): + names = [p["name"] for p in fmha_fwd_mfma_signature(_spec())] + assert len(names) == 15 + assert not any(n in _BATCH for n in names) + + +def test_default_name_untagged(): + assert not _spec().kernel_name().endswith("_fs") + + +def test_free_stride_abi_appends_batch_strides(): + names = [p["name"] for p in fmha_fwd_mfma_signature(_spec(free_stride=True))] + assert len(names) == 19 + assert tuple(names[-4:]) == _BATCH + + +def test_free_stride_name_tagged(): + assert _spec(free_stride=True).kernel_name().endswith("_fs") + + +@pytest.mark.parametrize("arch", ["gfx950", "gfx942"]) +@pytest.mark.parametrize("mask", ["none", "causal"]) +def test_free_stride_builds(arch, mask): + kd = build_fmha_fwd_mfma(_spec(free_stride=True, mask_mode=mask), arch) + assert kd.name.endswith("_fs") + + +def test_free_stride_gqa_builds(): + kd = build_fmha_fwd_mfma( + _spec(free_stride=True, num_query_heads=8, num_kv_heads=2), "gfx950" + ) + assert kd is not None diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride_numeric.py new file mode 100644 index 000000000000..d2bdf853a6db --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_fmha_mfma_free_stride_numeric.py @@ -0,0 +1,69 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""On-GPU numeric parity for the fmha_mfma free-stride (layout-general) path. + +Runs only on a CDNA device (gfx942 / gfx950); skips elsewhere. For each +contiguous layout it builds ``fmha_mfma(free_stride=True)``, launches it, and +checks the output against the numpy dense-attention reference. The point is +layout invariance: BSHD and BHSD must both match the reference (and therefore +each other), proving the per-batch base-pointer rebase reads and writes every +layout correctly. The build/ABI contract is covered CPU-only by +``test_fmha_mfma_free_stride.py``. +""" + +from __future__ import annotations + +import pytest + +from rocke.runtime.hip_module import get_device_arch + +_ARCH = get_device_arch() +_CDNA = ("gfx942", "gfx950") + +pytestmark = pytest.mark.skipif( + _ARCH not in _CDNA, reason=f"needs a CDNA GPU (gfx942/gfx950); got {_ARCH}" +) + + +@pytest.mark.parametrize("layout", ["bhsd", "bshd"]) +@pytest.mark.parametrize("causal", [False, True]) +def test_free_stride_matches_reference(layout, causal): + from builders.common.sdpa_fwd_free_stride_verify import run_verify + + passed, max_abs, bad, total = run_verify( + arch=_ARCH, + layout=layout, + heads=8, + kv_heads=2, # GQA + head_size=64, + seqlen_q=128, + seqlen_k=128, + batch=2, + causal=causal, + verbose=False, + ) + assert ( + passed + ), f"{layout} causal={causal}: max_abs_diff={max_abs:.3e} bad={bad}/{total}" + + +def test_bshd_and_bhsd_agree(): + # Same logical problem, two layouts: the kernel outputs must match bit-for- + # bit-close, which is the core free-stride invariant. + from builders.common.sdpa_fwd_free_stride_verify import run_verify + + common = dict( + arch=_ARCH, + heads=8, + kv_heads=2, + head_size=64, + seqlen_q=128, + seqlen_k=128, + batch=2, + causal=False, + verbose=False, + ) + ok_a, a, _, _ = run_verify(layout="bhsd", **common) + ok_b, b, _, _ = run_verify(layout="bshd", **common) + assert ok_a and ok_b, f"bhsd={a:.3e} bshd={b:.3e}" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_sdpa_free_stride_layouts.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_sdpa_free_stride_layouts.py new file mode 100644 index 000000000000..cc24d027da49 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_sdpa_free_stride_layouts.py @@ -0,0 +1,159 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host contract for free-stride attention forward (ticket 2 / layout axis). + +The generalized forward kernel must read Q/K/V from any memory layout -- BSHD, +BHSD, general B/H/S strides and packed QKV -- addressed by per-tensor +``(batch, head, token)`` element strides with a unit head-dim stride (the +contract of the ``StridedTensorAddr`` layer). + +These host tests pin that contract with no GPU and no kernel: they reconstruct +the logical ``(B, H, S, D)`` tensor from a backing buffer using the same scalar +element-address formula the kernel emits, + + addr = base + b*stride_batch + h*stride_head + s*stride_token + d*stride_dim + +and confirm every layout recovers the same values and the same attention +output. This is exactly what the kernel ABI (the stride params we pass) has to +satisfy, so it is the definition of done for the addressing work and the host +reference the emitter is matched against later. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from kernels.common._sdpa_fwd_strides import contiguous_strides +from rocke.numeric.references import dense_attention_reference + + +def _addr_gather(flat, base, strides, shape): + """Reconstruct a tensor of ``shape`` from a flat element buffer using the + kernel's scalar address formula ``base + sum(idx_i * stride_i)``.""" + s_b, s_h, s_t, s_d = strides + g = np.indices(shape) + addr = base + g[0] * s_b + g[1] * s_h + g[2] * s_t + g[3] * s_d + return flat[addr] + + +def _view_base_strides(view, backing): + """Element ``(flat, base, strides)`` of ``view`` into a C-contiguous + ``backing``: byte offsets and strides divided by the item size.""" + it = backing.itemsize + base = (view.ctypes.data - backing.ctypes.data) // it + strides = tuple(s // it for s in view.strides) + return backing.reshape(-1), base, strides + + +def _logical(B=2, H=4, S=16, D=16, seed=0xF00D): + rng = np.random.default_rng(seed) + return (rng.standard_normal((B, H, S, D)) * 0.3).astype(np.float16) + + +# Each builder returns ``(view, backing)`` where ``view`` presents the logical +# ``(B, H, S, D)`` order and ``backing`` is the C-contiguous buffer behind it. +def _bhsd(x): + b = np.ascontiguousarray(x) # [B, H, S, D]: token stride = D + return b, b + + +def _bshd(x): + # Store tokens-major [B, S, H, D]; the logical view transposes heads back. + b = np.ascontiguousarray(x.transpose(0, 2, 1, 3)) + return b.transpose(0, 2, 1, 3), b + + +def _padded(x): + # Pad the head dim: a non-contiguous slice, so batch/head/token strides all + # carry the padding while the head dim keeps unit stride. + B, H, S, D = x.shape + b = np.zeros((B, H, S, D + 8), x.dtype) + b[..., :D] = x + return b[..., :D], b + + +LAYOUTS = {"bhsd": _bhsd, "bshd": _bshd, "padded": _padded} + + +@pytest.mark.parametrize("name", list(LAYOUTS)) +def test_layout_strides_recover_logical(name): + x = _logical() + view, backing = LAYOUTS[name](x) + flat, base, strides = _view_base_strides(view, backing) + # The addressing layer assumes a unit head-dim stride. + assert strides[3] == 1, f"{name}: head-dim stride must be 1, got {strides[3]}" + got = _addr_gather(flat, base, strides, x.shape) + np.testing.assert_array_equal(got, x) + + +def test_packed_qkv_views_recover_logical(): + # Q, K, V interleaved in one buffer [B, H, S, 3, D]; each is a strided view + # with its own base offset (0, D, 2D) into the shared buffer. + B, H, S, D = 2, 4, 16, 16 + rng = np.random.default_rng(0xBEE5) + qkv = [ + (rng.standard_normal((B, H, S, D)) * 0.3).astype(np.float16) for _ in range(3) + ] + packed = np.zeros((B, H, S, 3, D), np.float16) + for i, t in enumerate(qkv): + packed[..., i, :] = t + for i, want in enumerate(qkv): + view = packed[..., i, :] + flat, base, strides = _view_base_strides(view, packed) + assert strides[3] == 1 + assert base == i * D + np.testing.assert_array_equal( + _addr_gather(flat, base, strides, (B, H, S, D)), want + ) + + +@pytest.mark.parametrize("name", ["bhsd", "bshd"]) +def test_contiguous_strides_match_numpy_views(name): + # The launcher's analytical (batch, head, token) strides must equal what + # numpy computes for the same contiguous layout (the Step 1 ground truth). + x = _logical() + _B, H, S, D = x.shape + view, backing = LAYOUTS[name](x) + _flat, _base, np_strides = _view_base_strides(view, backing) + got = contiguous_strides(name, heads=H, seqlen=S, head_size=D) + assert got == np_strides[:3] # (batch, head, token); dim is unit, excluded + assert np_strides[3] == 1 + + +def test_contiguous_strides_rejects_bad_input(): + with pytest.raises(ValueError, match="layout"): + contiguous_strides("thd", heads=4, seqlen=16, head_size=16) + with pytest.raises(ValueError, match="heads"): + contiguous_strides("bhsd", heads=0, seqlen=16, head_size=16) + + +def _attention_bhsd(ql, kl, vl, *, causal): + """Attention over logical ``(B, H, S, D)`` operands; the reference consumes + one batch of shape ``(S, H, D)``.""" + return np.stack( + [ + dense_attention_reference( + ql[b].transpose(1, 0, 2), + kl[b].transpose(1, 0, 2), + vl[b].transpose(1, 0, 2), + causal=causal, + ) + for b in range(ql.shape[0]) + ] + ) + + +@pytest.mark.parametrize("name", list(LAYOUTS)) +def test_attention_output_is_layout_invariant(name): + ql, kl, vl = _logical(seed=1), _logical(seed=2), _logical(seed=3) + ref = _attention_bhsd(ql, kl, vl, causal=False) + + def through(x): + view, backing = LAYOUTS[name](x) + flat, base, strides = _view_base_strides(view, backing) + return _addr_gather(flat, base, strides, x.shape) + + got = _attention_bhsd(through(ql), through(kl), through(vl), causal=False) + np.testing.assert_allclose(got, ref, atol=1e-5, rtol=0)