TheRock/rocm-libraries/projects/composablekernel/dispatcher/tests/fmha_smoke_matrix.py

#!/usr/bin/env python3

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""
FMHA smoke test matrix generator.

Generates the same test cases as smoke_test_fwd.sh and smoke_test_bwd.sh
from the CK Tile 01_fmha example, for automated parity testing.
"""

from dataclasses import dataclass
from typing import List, Set, Tuple


@dataclass
class TestCase:
    name: str = ""
    direction: str = "fwd"
    prec: str = "fp16"
    mode: int = 0
    batch: int = 2
    nhead_q: int = 2
    nhead_k: int = -1
    hdim_q: int = 128
    hdim_v: int = -1
    seqlen_q: int = 128
    seqlen_k: int = 128
    bias: str = "n"
    mask: str = "0"
    lse: int = 0
    p_drop: float = 0.0
    perm: int = 1
    num_splits: int = 1
    page_block_size: int = 0
    cache_batch_idx: int = 0
    s_kpad: str = ""
    q_eff_lens: str = ""
    kv_eff_lens: str = ""
    s_qpad: str = ""
    rotary_dim: int = 0
    rotary_interleaved: int = 0
    deterministic: int = 0
    dbias: int = 0

    def effective_nhead_k(self):
        return self.nhead_k if self.nhead_k > 0 else self.nhead_q

    def effective_hdim_v(self):
        return self.hdim_v if self.hdim_v > 0 else self.hdim_q


def generate_fwd_fp16_bf16_matrix() -> List[TestCase]:
    """Generate the run_fp16_bf16_tests matrix from smoke_test_fwd.sh."""
    cases = []
    idx = 0
    for prec in ["fp16", "bf16"]:
        for mode in [1, 0]:
            for perm in [0, 1]:
                for hdim in [32, 64, 128, 256]:
                    for lse in [0, 1]:
                        for bias in ["n", "e", "a"]:
                            for p_drop in [0.0, 0.2]:
                                base = dict(
                                    prec=prec,
                                    mode=mode,
                                    perm=perm,
                                    lse=lse,
                                    bias=bias,
                                    p_drop=p_drop,
                                )
                                subcases = [
                                    dict(
                                        batch=2,
                                        nhead_q=2,
                                        nhead_k=1,
                                        hdim_q=16,
                                        hdim_v=hdim,
                                        seqlen_q=55,
                                        seqlen_k=256,
                                        mask="0",
                                    ),
                                    dict(
                                        batch=1,
                                        nhead_q=3,
                                        hdim_q=hdim,
                                        seqlen_q=100,
                                        seqlen_k=51,
                                        mask="0",
                                    ),
                                    dict(
                                        batch=2,
                                        nhead_q=1,
                                        hdim_q=16,
                                        hdim_v=hdim,
                                        seqlen_q=99,
                                        seqlen_k=256,
                                        mask="1",
                                    ),
                                    dict(
                                        batch=1,
                                        nhead_q=2,
                                        nhead_k=1,
                                        hdim_q=hdim,
                                        seqlen_q=1024,
                                        seqlen_k=256,
                                        mask="2",
                                    ),
                                    dict(
                                        batch=2,
                                        nhead_q=1,
                                        hdim_q=hdim,
                                        hdim_v=24,
                                        seqlen_q=3,
                                        seqlen_k=99,
                                        mask="2",
                                    ),
                                    dict(
                                        batch=3,
                                        nhead_q=2,
                                        nhead_k=1,
                                        hdim_q=hdim,
                                        seqlen_q=200,
                                        seqlen_k=520,
                                        mask="t:128,30",
                                    ),
                                    dict(
                                        batch=2,
                                        nhead_q=1,
                                        hdim_q=hdim,
                                        seqlen_q=99,
                                        seqlen_k=32,
                                        mask="b:4,35",
                                    ),
                                    dict(
                                        batch=1,
                                        nhead_q=2,
                                        nhead_k=1,
                                        hdim_q=hdim,
                                        seqlen_q=33,
                                        seqlen_k=0,
                                        mask="2",
                                    ),
                                    dict(
                                        batch=1,
                                        nhead_q=2,
                                        nhead_k=1,
                                        hdim_q=hdim,
                                        seqlen_q=1,
                                        seqlen_k=10,
                                        s_kpad="32",
                                        mask="2",
                                    ),
                                ]
                                for sc in subcases:
                                    idx += 1
                                    c = TestCase(
                                        name=f"fwd_{idx:04d}_{prec}_m{mode}_h{hdim}",
                                        direction="fwd",
                                        **base,
                                        **sc,
                                    )
                                    cases.append(c)
    return cases


def generate_bwd_matrix() -> List[TestCase]:
    """Generate the bwd smoke test matrix from smoke_test_bwd.sh."""
    cases = []
    idx = 0
    base_shapes = [
        dict(batch=1, nhead_q=4, nhead_k=2, seqlen_q=259, seqlen_k=259, mask="0"),
        dict(batch=2, nhead_q=2, seqlen_q=516, seqlen_k=253, mask="0"),
        dict(batch=1, nhead_q=4, nhead_k=1, seqlen_q=500, seqlen_k=251, mask="1"),
        dict(batch=1, nhead_q=2, seqlen_q=900, seqlen_k=258, mask="2"),
        dict(batch=2, nhead_q=1, seqlen_q=987, seqlen_k=219, mask="t:128,30"),
        dict(batch=2, nhead_q=3, nhead_k=1, seqlen_q=244, seqlen_k=499, mask="b:4,35"),
    ]
    for prec in ["fp16", "bf16"]:
        for perm in [0, 1]:
            for hdim in [32, 64, 128, 256]:
                for mode in [0, 1]:
                    for bias in ["n", "a"]:
                        for p_drop in [0.0, 0.2]:
                            for shape in base_shapes:
                                idx += 1
                                cases.append(
                                    TestCase(
                                        name=f"bwd_{idx:04d}_{prec}_h{hdim}",
                                        direction="bwd",
                                        prec=prec,
                                        mode=mode,
                                        perm=perm,
                                        hdim_q=hdim,
                                        hdim_v=hdim,
                                        bias=bias,
                                        p_drop=p_drop,
                                        lse=1,
                                        **shape,
                                    )
                                )
    return cases


def generate_splitkv_matrix() -> List[TestCase]:
    """Generate the splitkv smoke test matrix (same subcases as fwd, with num_splits > 1)."""
    cases = []
    idx = 0
    for prec in ["fp16", "bf16"]:
        for mode in [0]:  # splitkv only supports batch mode in smoke test
            for perm in [0, 1]:
                for hdim in [64, 128, 256]:
                    for num_splits in [2, 3]:
                        for bias in ["n"]:
                            subcases = [
                                dict(
                                    batch=2,
                                    nhead_q=2,
                                    nhead_k=1,
                                    seqlen_q=55,
                                    seqlen_k=256,
                                    mask="0",
                                ),
                                dict(
                                    batch=1,
                                    nhead_q=3,
                                    seqlen_q=100,
                                    seqlen_k=51,
                                    mask="0",
                                ),
                                dict(
                                    batch=1,
                                    nhead_q=2,
                                    nhead_k=1,
                                    seqlen_q=1024,
                                    seqlen_k=256,
                                    mask="2",
                                ),
                                dict(
                                    batch=3,
                                    nhead_q=2,
                                    nhead_k=1,
                                    seqlen_q=200,
                                    seqlen_k=520,
                                    mask="t:128,30",
                                ),
                            ]
                            for sc in subcases:
                                idx += 1
                                cases.append(
                                    TestCase(
                                        name=f"splitkv_{idx:04d}_{prec}_h{hdim}_s{num_splits}",
                                        direction="fwd_splitkv",
                                        prec=prec,
                                        mode=mode,
                                        perm=perm,
                                        hdim_q=hdim,
                                        hdim_v=hdim,
                                        lse=1,
                                        bias=bias,
                                        p_drop=0.0,
                                        num_splits=num_splits,
                                        page_block_size=128,
                                        cache_batch_idx=1,
                                        **sc,
                                    )
                                )
    return cases


def generate_padding_matrix() -> List[TestCase]:
    """Generate padding edge-case test cases."""
    cases = []
    idx = 0
    for prec in ["fp16"]:
        for hdim in [32, 64, 128]:
            edge_shapes = [
                dict(batch=1, nhead_q=1, seqlen_q=1, seqlen_k=1, mask="0"),
                dict(batch=1, nhead_q=1, seqlen_q=1, seqlen_k=256, mask="0"),
                dict(batch=1, nhead_q=1, seqlen_q=255, seqlen_k=1, mask="0"),
                dict(batch=1, nhead_q=2, seqlen_q=3, seqlen_k=5, mask="1"),
                dict(batch=2, nhead_q=1, seqlen_q=17, seqlen_k=33, mask="2"),
            ]
            for shape in edge_shapes:
                idx += 1
                cases.append(
                    TestCase(
                        name=f"pad_{idx:04d}_{prec}_h{hdim}",
                        direction="fwd",
                        prec=prec,
                        mode=0,
                        perm=1,
                        hdim_q=hdim,
                        hdim_v=hdim,
                        bias="n",
                        lse=0,
                        p_drop=0.0,
                        **shape,
                    )
                )
    return cases


def generate_fp8_matrix() -> List[TestCase]:
    """Generate fp8 smoke test cases (fp8bf16 and fp8fp32)."""
    cases = []
    idx = 0
    for prec in ["fp8bf16"]:
        for mode in [0]:
            for perm in [1]:
                for hdim in [64, 128]:
                    for mask in ["0", "2"]:
                        idx += 1
                        cases.append(
                            TestCase(
                                name=f"fp8_{idx:04d}_{prec}_h{hdim}",
                                direction="fwd",
                                prec=prec,
                                mode=mode,
                                perm=perm,
                                hdim_q=hdim,
                                hdim_v=hdim,
                                batch=2,
                                nhead_q=4,
                                nhead_k=4,
                                seqlen_q=128,
                                seqlen_k=128,
                                bias="n",
                                mask=mask,
                                lse=0,
                                p_drop=0.0,
                            )
                        )
    return cases


def unique_kernel_configs(cases: List[TestCase]) -> Set[Tuple]:
    """Extract unique kernel configs needed to run the test cases."""
    configs = set()
    for c in cases:
        dv = c.effective_hdim_v()
        mask_cat = (
            "no" if c.mask == "0" else ("causal" if c.mask in ["1", "2"] else "window")
        )
        bias_cat = c.bias
        configs.add(
            (
                c.direction,
                c.prec,
                c.hdim_q,
                dv,
                mask_cat,
                bias_cat,
                bool(c.lse),
                c.p_drop > 0,
            )
        )
    return configs


def to_ck_cli_args(case: TestCase) -> list:
    """Convert a TestCase to CK Tile CLI arguments."""
    nk = case.effective_nhead_k()
    dv = case.effective_hdim_v()
    args = [
        f"-prec={case.prec}",
        f"-mode={case.mode}",
        f"-b={case.batch}",
        f"-h={case.nhead_q}",
    ]
    if nk != case.nhead_q:
        args.append(f"-h_k={nk}")
    args += [f"-d={case.hdim_q}"]
    if dv != case.hdim_q:
        args.append(f"-d_v={dv}")
    args += [
        f"-s={case.seqlen_q}",
        f"-s_k={case.seqlen_k}",
        f"-bias={case.bias}",
        f"-mask={case.mask}",
        f"-lse={case.lse}",
        f"-p_drop={case.p_drop}",
        f"-iperm={case.perm}",
        f"-operm={case.perm}",
        "-v=1",
        "-warmup=0",
        "-repeat=1",
    ]
    if case.s_kpad:
        args.append(f"-s_kpad={case.s_kpad}")
    if case.num_splits > 1:
        args.append(f"-num_splits={case.num_splits}")
    if case.page_block_size > 0:
        args.append(f"-page_block_size={case.page_block_size}")
    if case.cache_batch_idx:
        args.append(f"-cache_batch_idx={case.cache_batch_idx}")
    return args


if __name__ == "__main__":
    fwd = generate_fwd_fp16_bf16_matrix()
    bwd = generate_bwd_matrix()
    skv = generate_splitkv_matrix()
    pad = generate_padding_matrix()
    fp8 = generate_fp8_matrix()

    all_cases = fwd + bwd + skv + pad + fp8
    all_configs = unique_kernel_configs(all_cases)

    print(f"Forward:   {len(fwd):5d} cases")
    print(f"Backward:  {len(bwd):5d} cases")
    print(f"SplitKV:   {len(skv):5d} cases")
    print(f"Padding:   {len(pad):5d} cases")
    print(f"FP8:       {len(fp8):5d} cases")
    print(f"Total:     {len(all_cases):5d} cases, {len(all_configs)} unique configs")