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

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""Combine every GPU-verifiable single-wave lever in every way, per shape.

The campaign measured each lever in isolation (``tune``/``sp_tune``/``bn_tune``).
This driver runs the FULL cartesian product of every compatible lever across all
three single-wave kernels -- ``fmha_singlewave`` (``bm_tiles`` x ``p_mode`` x
``v_mode`` x ``q_preload`` x ``fuse_k``), ``fmha_pipelined`` (software pipelining x
``sched`` x ``p_xpose`` x ``fuse_k``), and ``fmha_blockn`` (``bn_tiles`` x
``fuse_k``) -- GPU-verifies each against the numpy reference (tol 2e-2), times it,
and reports the single best combination per shape and globally.

It exists to answer "did any *combination* of levers beat what each scores
alone?" The standout positive interaction it surfaced: ``q_preload + fuse_k``
together on D128 (each is neutral / a win *alone*; together they are the best
single-wave config and clear 20% of the 59 TF peak on the compute-heavy shapes).

Run: ``python -m builders.gfx1151.attention.combo`` (sweeps a default set
of representative shapes); ``--quick`` for a 2-shape smoke.
"""

from __future__ import annotations

import argparse
import itertools

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

PEAK_TF = 59.0


# Representative shapes spanning D64/D128, seqlen, GQA, causal. Chosen to include
# the survey's top performers (compute-heavy D128) where 20% is in reach.
DEFAULT_SHAPES = [
    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=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=2, heads=8, kv_heads=2, seqlen_q=1024, seqlen_k=1024, head_size=128),
    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=128, causal=True),
]


def _opt_configs(shape: Shape):
    mask = "causal" if shape.causal else "none"
    # bm_tiles=2 only where seqlen allows the larger BLOCK_M; 1 always.
    bms = [1]
    if shape.seqlen_q % 32 == 0:
        bms.append(2)
    grid = itertools.product(
        bms,
        ["lds"],
        ["gather", "lds_t"],
        [0, 1],
        [0, 1],  # bm, p, v, qpreload, fusek
    )
    for bm, pm, vm, qp, fk in grid:
        yield SingleWaveCfg(
            head_size=shape.head_size,
            num_query_heads=shape.heads,
            num_kv_heads=shape.kv_heads,
            mask_mode=mask,
            bm_tiles=bm,
            p_mode=pm,
            v_mode=vm,
            q_preload=bool(qp),
            fuse_k=bool(fk),
        )


def _sp_configs(shape: Shape):
    mask = "causal" if shape.causal else "none"
    for sched, xp, fk in itertools.product([0, 1], ["lds", "shuffle"], [0, 1]):
        yield PipelinedCfg(
            head_size=shape.head_size,
            num_query_heads=shape.heads,
            num_kv_heads=shape.kv_heads,
            mask_mode=mask,
            sched=bool(sched),
            p_xpose=xp,
            fuse_k=bool(fk),
        )


def _bn_configs(shape: Shape):
    mask = "causal" if shape.causal else "none"
    for bn, fk in itertools.product([2, 4], [0, 1]):
        yield BlockNCfg(
            head_size=shape.head_size,
            num_query_heads=shape.heads,
            num_kv_heads=shape.kv_heads,
            mask_mode=mask,
            bn_tiles=bn,
            fuse_k=bool(fk),
        )


def _tag(shape: Shape):
    return (
        f"B{shape.batch} H{shape.heads}"
        + (f"/{shape.kvh}" if shape.kvh != shape.heads else "")
        + f" S{shape.seqlen_q} D{shape.head_size}"
        + ("c" if shape.causal else "")
    )


def _label(kind, cfg):
    if kind == "opt":
        return (
            f"opt bm{cfg.bm_tiles} p{cfg.p_mode} v{cfg.v_mode} "
            f"qp{int(cfg.q_preload)} fk{cfg.fuse_k}"
        )
    if kind == "sp":
        return f"sp sched{int(cfg.sched)} xp{cfg.p_xpose} fk{cfg.fuse_k}"
    return f"bn bn{cfg.bn_tiles} fk{cfg.fuse_k}"


def sweep_shape(shape: Shape, objdump, *, verbose=True):
    runners = [
        ("opt", _opt_configs(shape), opt_tune.verify_and_time),
        ("sp", _sp_configs(shape), sp_tune.verify_and_time),
        ("bn", _bn_configs(shape), bn_tune.verify_and_time),
    ]
    best = None
    for kind, cfgs, vt in runners:
        for cfg in cfgs:
            try:
                r = vt(cfg, shape, objdump=objdump)
            except Exception as e:  # noqa: BLE001
                if verbose:
                    print(f"    {_label(kind, cfg):46s} BUILD/RUN FAIL: {str(e)[:40]}")
                continue
            ok = r.get("ok")
            tf = r.get("tflops", 0.0)
            if verbose:
                flag = "Y" if ok else "N"
                print(
                    f"    {_label(kind, cfg):46s} {flag} {tf:7.2f} TF "
                    f"spill={r.get('vspill', '-')}"
                )
            if ok and (best is None or tf > best[2]):
                best = (kind, cfg, tf, r)
    return best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--quick", action="store_true", help="only the two compute-heavy D128 shapes"
    )
    ap.add_argument("--verbose", action="store_true", help="print every config")
    args = ap.parse_args()

    shapes = DEFAULT_SHAPES
    if args.quick:
        shapes = [
            Shape(batch=4, heads=8, seqlen_q=2048, seqlen_k=2048, head_size=128),
            Shape(batch=4, heads=8, seqlen_q=1024, seqlen_k=1024, head_size=64),
        ]

    objdump = _find_objdump()
    rows = []
    for sh in shapes:
        print(f"\n### {_tag(sh)}")
        best = sweep_shape(sh, objdump, verbose=args.verbose)
        if best is None:
            print(f"  {_tag(sh):30s} NO VERIFIED CONFIG")
            continue
        kind, cfg, tf, r = best
        pk = tf / PEAK_TF * 100.0
        rows.append((_tag(sh), kind, cfg, tf, pk))
        print(
            f"  BEST {_tag(sh):30s} -> {tf:7.2f} TF ({pk:4.1f}%) "
            f"[{_label(kind, cfg)}]  max_abs={r['max_abs']:.1e}"
        )

    if rows:
        print("\n==== best combination per shape ====")
        for tag, kind, cfg, tf, pk in rows:
            print(f"{tag:30s} {tf:7.2f} TF {pk:5.1f}%  {_label(kind, cfg)}")
        gbest = max(rows, key=lambda x: x[3])
        avg = sum(x[4] for x in rows) / len(rows)
        n20 = sum(1 for x in rows if x[4] >= 20.0)
        print(
            f"\nglobal best: {gbest[0]} -> {gbest[3]:.2f} TF ({gbest[4]:.1f}%) "
            f"[{_label(gbest[1], gbest[2])}]"
        )
        print(f"avg of per-shape best: {avg:.1f}% of {PEAK_TF} peak")
        print(f"shapes >=20% peak: {n20}/{len(rows)}")


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