TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx1151/attention/survey.py

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""25-shape survey of the best single-wave WMMA FMHA kernels on gfx1151.

For each shape it runs BOTH best single-wave kernels -- ``fmha_pipelined`` (software
pipelining, the D64 record) and ``fmha_singlewave`` with ``fuse_k`` (the D128 record)
-- keeps the faster, and also times PyTorch SDPA (the only available reference on
this APU is the MATH fallback; flash/efficient are runtime-disabled on gfx1151).
Reports TFLOP/s and % of the ~59 TF f16 WMMA peak for each.
"""

from __future__ import annotations


from .bench_v_staging import _find_objdump
from .fmha_singlewave import SingleWaveCfg
from .fmha_pipelined import PipelinedCfg
from .tune import Shape
from . import tune as opt_tune
from . import sp_tune

PEAK_TF = 59.0  # Radeon 8060S f16 WMMA peak: 40 CU * 512 FLOP/clk * 2.9 GHz


# 25 shapes: vary head_size, seqlen, batch, heads, GQA, causal.
SHAPES = [
    # D64 sweep (seqlen / batch / causal)
    Shape(batch=4, heads=8, seqlen_q=512, seqlen_k=512, head_size=64),
    Shape(batch=4, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=64),
    Shape(batch=4, heads=8, seqlen_q=2048, seqlen_k=2048, head_size=64),
    Shape(batch=8, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=64),
    Shape(batch=2, heads=16, seqlen_q=1024, seqlen_k=1024, head_size=64),
    Shape(batch=4, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=64, causal=True),
    Shape(batch=4, heads=8, seqlen_q=2048, seqlen_k=2048, head_size=64, causal=True),
    Shape(batch=2, heads=8, kv_heads=2, seqlen_q=1024, seqlen_k=1024, head_size=64),
    # D128 sweep
    Shape(batch=4, heads=8, seqlen_q=512, seqlen_k=512, head_size=128),
    Shape(batch=4, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=128),
    Shape(batch=4, heads=8, seqlen_q=2048, seqlen_k=2048, head_size=128),
    Shape(batch=8, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=128),
    Shape(batch=2, heads=16, seqlen_q=1024, seqlen_k=1024, head_size=128),
    Shape(batch=1, heads=32, seqlen_q=1024, seqlen_k=1024, head_size=128),
    Shape(batch=4, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=128, causal=True),
    Shape(batch=4, heads=8, seqlen_q=2048, seqlen_k=2048, head_size=128, causal=True),
    Shape(batch=2, heads=8, kv_heads=2, seqlen_q=1024, seqlen_k=1024, head_size=128),
    Shape(batch=2, heads=32, kv_heads=8, seqlen_q=2048, seqlen_k=2048, head_size=128),
    # larger sequence / occupancy
    Shape(batch=1, heads=16, seqlen_q=4096, seqlen_k=4096, head_size=64),
    Shape(batch=1, heads=16, seqlen_q=4096, seqlen_k=4096, head_size=128),
    Shape(batch=1, heads=16, seqlen_q=4096, seqlen_k=4096, head_size=128, causal=True),
    Shape(batch=8, heads=16, seqlen_q=512, seqlen_k=512, head_size=64),
    Shape(batch=8, heads=16, seqlen_q=512, seqlen_k=512, head_size=128),
    Shape(batch=4, heads=8, seqlen_q=4096, seqlen_k=4096, head_size=128),
    Shape(batch=2, heads=16, seqlen_q=2048, seqlen_k=2048, head_size=64, causal=True),
]


def _sdpa_tf(shape: Shape):
    import torch
    from torch.nn.attention import SDPBackend, sdpa_kernel
    import time

    dev = "cuda"
    B, H, Hk = shape.batch, shape.heads, shape.kvh
    Sq, Sk, D = shape.seqlen_q, shape.seqlen_k, shape.head_size
    q = torch.randn(B, H, Sq, D, device=dev, dtype=torch.float16) * 0.3
    k = torch.randn(B, Hk, Sk, D, device=dev, dtype=torch.float16) * 0.3
    v = torch.randn(B, Hk, Sk, D, device=dev, dtype=torch.float16) * 0.3
    if Hk != H:
        k = k.repeat_interleave(H // Hk, dim=1)
        v = v.repeat_interleave(H // Hk, dim=1)

    def f():
        with sdpa_kernel(SDPBackend.MATH):
            return torch.nn.functional.scaled_dot_product_attention(
                q, k, v, is_causal=shape.causal
            )

    for _ in range(5):
        f()
    torch.cuda.synchronize()
    t = time.perf_counter()
    N = 30
    for _ in range(N):
        f()
    torch.cuda.synchronize()
    ms = (time.perf_counter() - t) / N * 1e3
    flops = 4.0 * B * H * Sq * Sk * D * (0.5 if shape.causal else 1.0)
    return flops / (ms * 1e-3) / 1e12


def _best_mine(shape: Shape, objdump):
    """Run fmha_pipelined and fmha_singlewave(fuse_k auto); return the faster verified one."""
    results = []
    mask = "causal" if shape.causal else "none"
    sp = PipelinedCfg(
        head_size=shape.head_size,
        num_query_heads=shape.heads,
        num_kv_heads=shape.kv_heads,
        mask_mode=mask,
    )
    opt = SingleWaveCfg(
        head_size=shape.head_size,
        num_query_heads=shape.heads,
        num_kv_heads=shape.kv_heads,
        mask_mode=mask,
        fuse_k=None,
    )
    try:
        r = sp_tune.verify_and_time(sp, shape, objdump=objdump)
        if r["ok"]:
            results.append(("sp", r))
    except Exception as e:  # noqa: BLE001
        print(f"    sp FAIL: {str(e)[:60]}")
    try:
        r = opt_tune.verify_and_time(opt, shape, objdump=objdump)
        if r["ok"]:
            results.append(("opt", r))
    except Exception as e:  # noqa: BLE001
        print(f"    opt FAIL: {str(e)[:60]}")
    if not results:
        return None
    return max(results, key=lambda kr: kr[1]["tflops"])


def main():
    objdump = _find_objdump()
    print(f"{'shape':46s} {'mine':>16s} {'%pk':>5s}  {'sdpa':>9s}  {'speedup':>7s}")
    rows = []
    for sh in SHAPES:
        tag = (
            f"B{sh.batch} H{sh.heads}"
            + (f"/{sh.kvh}" if sh.kvh != sh.heads else "")
            + f" S{sh.seqlen_q} D{sh.head_size}"
            + ("c" if sh.causal else "")
        )
        best = _best_mine(sh, objdump)
        try:
            sdpa = _sdpa_tf(sh)
        except Exception as e:  # noqa: BLE001
            sdpa = float("nan")
            print(f"    sdpa FAIL: {str(e)[:60]}")
        if best is None:
            print(f"{tag:46s} {'BUILD/VERIFY FAIL':>16s}")
            continue
        kind, r = best
        tf = r["tflops"]
        pk = tf / PEAK_TF * 100.0
        spd = tf / sdpa if sdpa == sdpa and sdpa > 0 else float("nan")
        rows.append((tag, kind, tf, pk, sdpa, spd))
        print(f"{tag:46s} {tf:7.2f}TF({kind:3s}) {pk:4.1f}  {sdpa:7.2f}TF  {spd:6.1f}x")

    if rows:
        best = max(rows, key=lambda x: x[2])
        avg_pk = sum(x[3] for x in rows) / len(rows)
        avg_spd = sum(x[5] for x in rows if x[5] == x[5]) / len(rows)
        print(
            f"\nbest: {best[0]} -> {best[2]:.2f} TF ({best[3]:.1f}% of {PEAK_TF} peak)"
        )
        print(
            f"avg: {avg_pk:.1f}% of peak, {avg_spd:.1f}x vs SDPA over {len(rows)} shapes"
        )
        n25 = sum(1 for x in rows if x[3] >= 25.0)
        print(f"shapes >=25% peak: {n25}/{len(rows)}")


if __name__ == "__main__":
    raise SystemExit(main())