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

#!/usr/bin/env python3

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

"""
Full FMHA Parity Test -- parallel JIT build, sequential GPU test.

Phase 1: JIT-compile every unique kernel config in parallel (hipcc only, no GPU).
Phase 2: Run each test case sequentially through CK Tile and the dispatcher
          (each dispatcher invocation in its own subprocess for HIP isolation).

Usage:
    python3 full_parity_test.py --max-cases 100
    python3 full_parity_test.py --max-cases 0       # all ~3500 cases
    python3 full_parity_test.py --workers 8          # parallel JIT build
    python3 full_parity_test.py --skip-jit           # reuse previous build
"""

import sys
import os
import time
import argparse
import subprocess
import json
from pathlib import Path
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional, Dict, Tuple
from fmha_smoke_matrix import (
    generate_fwd_fp16_bf16_matrix,
    generate_bwd_matrix,
    generate_splitkv_matrix,
    generate_padding_matrix,
    generate_fp8_matrix,
    to_ck_cli_args,
    TestCase,
)

SCRIPT_DIR = Path(__file__).resolve().parent
DISPATCHER_DIR = SCRIPT_DIR.parent
PYTHON_DIR = DISPATCHER_DIR / "python"

sys.path.insert(0, str(SCRIPT_DIR))


# =========================================================================
# Config dedup + tile lookup
# =========================================================================

HDIM_TILE_TABLE = {
    (32, 32): (128, 64, 16, 32, 32, 32),
    (64, 64): (128, 64, 32, 64, 32, 64),
    (128, 128): (128, 128, 32, 128, 32, 128),
    (192, 128): (128, 128, 32, 128, 32, 192),
    (192, 192): (128, 128, 32, 192, 32, 192),
    (256, 256): (128, 128, 32, 256, 32, 256),
    (80, 96): (128, 128, 16, 96, 32, 96),
    (96, 128): (128, 128, 32, 128, 32, 96),
}


def _round_hdim(d: int) -> int:
    for t in [32, 64, 96, 128, 192, 256]:
        if d <= t:
            return t
    return 256


def _lookup_tile(dq: int, dv: int):
    key = (dq, dv)
    if key in HDIM_TILE_TABLE:
        return HDIM_TILE_TABLE[key]
    sq = max(dq, dv)
    key2 = (sq, sq)
    if key2 in HDIM_TILE_TABLE:
        t = list(HDIM_TILE_TABLE[key2])
        t[3] = dv
        t[5] = sq
        return tuple(t)
    return (128, 64, 16, dv, 32, sq)


def _mask_str(m: str) -> str:
    return "no" if m == "0" else "top_left"


def _bias_str(b: str) -> str:
    return {"n": "no", "e": "bias", "a": "alibi"}.get(b, "no")


def config_key(c: TestCase) -> tuple:
    tdq = _round_hdim(c.hdim_q)
    tdv = _round_hdim(c.effective_hdim_v())
    # GQA (nhead_q != nhead_k) is a runtime property handled via strides,
    # NOT a compile-time kernel variant.  is_group_mode refers to
    # variable-length batching (mode=1), not GQA.
    is_varlen = c.mode == 1
    return (
        c.prec,
        tdq,
        tdv,
        _mask_str(c.mask),
        _bias_str(c.bias),
        bool(c.lse),
        c.p_drop > 0,
        is_varlen,
    )


def config_name(key: tuple) -> str:
    prec, dq, dv, mask, bias, lse, drop, varlen = key
    n = f"{prec}_h{dq}x{dv}_{'grp' if varlen else 'bat'}_{mask}_{bias}"
    if lse:
        n += "_lse"
    if drop:
        n += "_drop"
    return n


# Backward tile tables from CK codegen (gfx9/gfx950, fp16/bf16, tr_load=f)
# Format: tile(9), wave(9), warp(6) -- from fmha_bwd.py KernelComponentFactoryGfx9
BWD_CONFIGS = {
    32: {
        "tile": [32, 128, 32, 32, 32, 32, 64, 32, 32],
        "wave": [1, 4, 1, 4, 1, 1, 2, 2, 1],
        "warp": [16, 16, 32, 16, 16, 16],
    },
    64: {
        "tile": [32, 128, 64, 32, 64, 32, 32, 64, 64],
        "wave": [1, 4, 1, 4, 1, 1, 1, 4, 1],
        "warp": [16, 16, 32, 16, 16, 16],
    },
    96: {
        "tile": [32, 128, 96, 32, 96, 32, 32, 96, 96],
        "wave": [1, 4, 1, 4, 1, 1, 2, 2, 1],
        "warp": [16, 16, 32, 16, 16, 16],
    },
    128: {
        "tile": [16, 128, 128, 16, 128, 16, 32, 128, 128],
        "wave": [1, 4, 1, 4, 1, 1, 1, 4, 1],
        "warp": [16, 16, 32, 16, 16, 16],
    },
    256: {
        "tile": [16, 64, 256, 16, 256, 16, 32, 256, 256],
        "wave": [1, 4, 1, 4, 1, 1, 1, 4, 1],
        "warp": [16, 16, 32, 16, 16, 16],
    },
}


def config_to_codegen_json(key: tuple, arch: str) -> str:
    """Produce the JSON string that generate_fmha_fallback.py expects."""
    prec, dq, dv, mask, bias, lse, drop, is_varlen = key
    tile = _lookup_tile(dq, dv)
    return json.dumps(
        {
            "arch": arch,
            "signature": {
                "family": "fwd",
                "data_type": prec,
                "mode": "group" if is_varlen else "batch",
                "vlayout": "r",
                "hdim_q": dq,
                "hdim_v": dv,
                "mask": mask,
                "bias": bias,
                "lse": lse,
                "dropout": drop,
                "qscale": "no",
                "rope": "none",
                "logits": False,
                "paged_kv": False,
                "fp8_static_quant": False,
                "skip_min_seqlen_q": False,
                "sink": False,
                "dbias": False,
                "store_randval": False,
                "deterministic": False,
                "kv_memory_layout": "vectorized",
                "kv_lookup_table": "sglang",
                "page_size": 1,
            },
            "algorithm": {
                "pipeline": "qr"
                if "fp8" in prec
                else ("qr_async" if dq >= 64 else "qr"),
                "tile": list(tile),
                "wave": [2, 1, 1, 2, 1, 1, 1, 1, 1]
                if "fp8" in prec
                else [4, 1, 1, 4, 1, 1, 1, 1, 1],
                "warp": [32, 32, 32, 32, 32, 32, 16, 16, 16]
                if "fp8" in prec
                else [32, 32, 16, 32, 32, 16, 16, 16, 16],
                "padding": [True, True, True, True],
                "block_per_cu": 1,
                "num_wave_groups": 1,
                "max_splits_log2": 0,
                "max_seq_len_q": 0,
            },
        }
    )


def bwd_codegen_jsons(key: tuple, arch: str) -> list:
    """Produce 3 JSON strings for bwd stages: dot_do_o, dq_dk_dv, convert_dq."""
    prec, dq, dv, mask, bias, lse, drop, is_varlen = key
    mode = "group" if is_varlen else "batch"
    cfg = BWD_CONFIGS.get(dq, BWD_CONFIGS[128])
    bwd_tile = cfg["tile"]
    bwd_wave = cfg["wave"]
    bwd_warp = cfg["warp"]

    base_sig = {
        "data_type": prec,
        "mode": mode,
        "vlayout": "r",
        "hdim_q": dq,
        "hdim_v": dv,
        "mask": mask,
        "bias": bias,
        "lse": True,
        "dropout": drop,
        "qscale": "no",
        "rope": "none",
        "logits": False,
        "paged_kv": False,
        "fp8_static_quant": False,
        "skip_min_seqlen_q": False,
        "sink": False,
        "dbias": False,
        "store_randval": False,
        "deterministic": False,
        "kv_memory_layout": "vectorized",
        "kv_lookup_table": "sglang",
        "page_size": 1,
    }
    base_alg = {
        "pipeline": "bwd",
        "padding": [True, True, True, True],
        "block_per_cu": 1,
        "num_wave_groups": 1,
        "max_splits_log2": 0,
        "max_seq_len_q": 0,
        "use_trload": False,
    }

    dot_bm0 = max(bwd_tile[0], 64)
    dot_json = json.dumps(
        {
            "arch": arch,
            "signature": {**base_sig, "family": "bwd_dot_do_o"},
            "algorithm": {
                **base_alg,
                "tile": [dot_bm0, 0, 0, 0, 0, dv],
                "wave": [1, 1, 1, 1, 1, 1, 1, 1, 1],
                "warp": [16, 16, 16, 16, 16, 16, 16, 16, 16],
            },
        }
    )

    dqdkdv_json = json.dumps(
        {
            "arch": arch,
            "signature": {**base_sig, "family": "bwd_dq_dk_dv"},
            "algorithm": {
                **base_alg,
                "tile": bwd_tile,
                "wave": bwd_wave,
                "warp": bwd_warp + bwd_warp[:3],
            },
        }
    )

    cvt_bm0 = max(bwd_tile[0], 64)
    cvt_json = json.dumps(
        {
            "arch": arch,
            "signature": {**base_sig, "family": "bwd_convert_dq"},
            "algorithm": {
                **base_alg,
                "tile": [cvt_bm0, 0, 0, 0, 0, dq],
                "wave": [1, 1, 1, 1, 1, 1, 1, 1, 1],
                "warp": [16, 16, 16, 16, 16, 16, 16, 16, 16],
            },
        }
    )

    return [dot_json, dqdkdv_json, cvt_json]


# =========================================================================
# Phase 1 -- JIT build (no GPU, pure hipcc subprocesses)
# =========================================================================


def _jit_one(key: tuple, out_dir: Path, arch: str) -> Tuple[bool, str, float]:
    """JIT-compile a single kernel config. Runs hipcc only, never touches GPU."""
    t0 = time.perf_counter()
    out_dir.mkdir(parents=True, exist_ok=True)

    codegen_dir = DISPATCHER_DIR / "codegen"
    ctypes_src = DISPATCHER_DIR / "bindings" / "ctypes" / "fmha_ctypes_lib.cpp"
    static_lib = DISPATCHER_DIR / "build" / "libck_tile_dispatcher.a"
    if not static_lib.exists():
        return (False, "libck_tile_dispatcher.a not found", time.perf_counter() - t0)

    hipcc = "hipcc"
    cfg_json = config_to_codegen_json(key, arch)

    # 1. codegen
    r = subprocess.run(
        [
            sys.executable,
            str(codegen_dir / "fmha" / "generate_fallback.py"),
            "--output-dir",
            str(out_dir),
            "--gpu-target",
            arch,
            "--config-json",
            cfg_json,
        ],
        capture_output=True,
        text=True,
        cwd=str(codegen_dir),
    )
    if r.returncode != 0:
        return (False, f"codegen: {r.stderr[:200]}", time.perf_counter() - t0)

    dispatch_hdr = out_dir / "fmha_python_dispatch.hpp"
    if not dispatch_hdr.exists():
        return (False, "no dispatch header", time.perf_counter() - t0)

    sys.path.insert(0, str(PYTHON_DIR))
    from fmha_utils import fmha_compile_flags  # noqa: E402

    inc = [
        f"-I{out_dir}",
        f"-I{out_dir / 'dispatcher_wrappers'}",
    ]
    # fmha_compile_flags provides hipcc + all standard flags; strip hipcc (element 0)
    base_flags = fmha_compile_flags(arch, family="fwd")[1:]

    # 2. compile kernel .cpp files
    kernel_objs = []
    for cpp in sorted(out_dir.glob("fmha_*.cpp")):
        obj = cpp.with_suffix(".o")
        r = subprocess.run(
            [hipcc, "-c", *base_flags, *inc, str(cpp), "-o", str(obj)],
            capture_output=True,
            text=True,
        )
        if r.returncode != 0:
            return (False, f"kernel: {r.stderr[:200]}", time.perf_counter() - t0)
        kernel_objs.append(str(obj))

    # 3. compile ctypes lib
    ctypes_obj = out_dir / "fmha_ctypes_lib.o"
    r = subprocess.run(
        [
            hipcc,
            "-c",
            *base_flags,
            *inc,
            f"-include{dispatch_hdr}",
            f'-DGFX_ARCH="{arch}"',
            str(ctypes_src),
            "-o",
            str(ctypes_obj),
        ],
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        return (False, f"ctypes: {r.stderr[:200]}", time.perf_counter() - t0)

    # 4. link .so
    name = config_name(key)
    so_path = out_dir / f"libdispatcher_fmha_{name}.so"
    r = subprocess.run(
        [
            hipcc,
            "-shared",
            "-fPIC",
            str(ctypes_obj),
            *kernel_objs,
            str(static_lib),
            "-lamdhip64",
            "-o",
            str(so_path),
        ],
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        return (False, f"link: {r.stderr[:200]}", time.perf_counter() - t0)

    return (True, str(so_path), time.perf_counter() - t0)


def _jit_one_bwd(key: tuple, out_dir: Path, arch: str) -> Tuple[bool, str, float]:
    """JIT-compile all 3 bwd stages into one .so."""
    t0 = time.perf_counter()
    out_dir.mkdir(parents=True, exist_ok=True)

    codegen_dir = DISPATCHER_DIR / "codegen"
    ctypes_src = DISPATCHER_DIR / "bindings" / "ctypes" / "fmha_ctypes_lib.cpp"
    static_lib = DISPATCHER_DIR / "build" / "libck_tile_dispatcher.a"
    if not static_lib.exists():
        return (False, "libck_tile_dispatcher.a not found", time.perf_counter() - t0)

    hipcc = "hipcc"
    jsons = bwd_codegen_jsons(key, arch)

    # 1. codegen all 3 stages into the same dir
    for stage_json in jsons:
        r = subprocess.run(
            [
                sys.executable,
                str(codegen_dir / "fmha" / "codegen.py"),
                "--output-dir",
                str(out_dir),
                "--gpu-target",
                arch,
                "--config-json",
                stage_json,
            ],
            capture_output=True,
            text=True,
            cwd=str(codegen_dir),
        )
        if r.returncode != 0:
            return (False, f"codegen: {r.stderr[:200]}", time.perf_counter() - t0)

    # 1b. generate dispatch header combining all wrappers
    wrapper_dir = out_dir / "dispatcher_wrappers"
    if not wrapper_dir.exists():
        return (False, "no wrappers dir", time.perf_counter() - t0)

    sys.path.insert(0, str(codegen_dir))
    sys.path.insert(0, str(codegen_dir / "fmha"))
    from generate_fallback import generate_dispatch_header

    generate_dispatch_header(out_dir, wrapper_dir)

    dispatch_hdr = out_dir / "fmha_python_dispatch.hpp"
    from fmha_utils import fmha_compile_flags  # noqa: E402

    inc = [
        f"-I{out_dir}",
        f"-I{wrapper_dir}",
    ]
    base_flags = fmha_compile_flags(arch, family="bwd")[1:]

    # 2. compile all kernel .cpp files
    kernel_objs = []
    for cpp in sorted(out_dir.glob("fmha_*.cpp")):
        obj = cpp.with_suffix(".o")
        r = subprocess.run(
            [hipcc, "-c", *base_flags, *inc, str(cpp), "-o", str(obj)],
            capture_output=True,
            text=True,
        )
        if r.returncode != 0:
            return (
                False,
                f"kernel({cpp.name}): {r.stderr[:200]}",
                time.perf_counter() - t0,
            )
        kernel_objs.append(str(obj))

    # 3. compile ctypes lib
    ctypes_obj = out_dir / "fmha_ctypes_lib.o"
    r = subprocess.run(
        [
            hipcc,
            "-c",
            *base_flags,
            *inc,
            f"-include{dispatch_hdr}",
            f'-DGFX_ARCH="{arch}"',
            str(ctypes_src),
            "-o",
            str(ctypes_obj),
        ],
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        return (False, f"ctypes: {r.stderr[:200]}", time.perf_counter() - t0)

    # 4. link .so
    name = config_name(key)
    so_path = out_dir / f"libdispatcher_fmha_bwd_{name}.so"
    r = subprocess.run(
        [
            hipcc,
            "-shared",
            "-fPIC",
            str(ctypes_obj),
            *kernel_objs,
            str(static_lib),
            "-lamdhip64",
            "-o",
            str(so_path),
        ],
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        return (False, f"link: {r.stderr[:200]}", time.perf_counter() - t0)

    return (True, str(so_path), time.perf_counter() - t0)


# =========================================================================
# Phase 2 -- GPU tests (sequential, each in its own subprocess)
# =========================================================================


def find_ck_exe() -> Optional[str]:
    for p in [
        "/tmp/ck_fmha_full/bin/tile_example_fmha_fwd",
        "/tmp/ck_fmha_build/bin/tile_example_fmha_fwd",
    ]:
        if os.path.exists(p):
            return p
    return None


def run_ck_test(exe: str, case: TestCase) -> Tuple[bool, str]:
    cmd = [exe] + to_ck_cli_args(case)
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        return (r.returncode == 0, "")
    except subprocess.TimeoutExpired:
        return (False, "timeout")
    except Exception as e:
        return (False, str(e)[:60])


MASK_INT = {"0": 0, "1": 1, "2": 2}
BIAS_INT = {"n": 0, "e": 1, "a": 2}


def run_dispatcher_test(
    so_path: str, case: TestCase, key: tuple, arch: str
) -> Tuple[bool, str]:
    """Run one test in an isolated subprocess -- never touches our process's HIP."""
    dq = case.hdim_q
    dv = case.effective_hdim_v()
    nk = case.effective_nhead_k()
    traits_dq = key[1]  # tile-rounded hdim for kernel matching
    traits_dv = key[2]

    if case.seqlen_k <= 0 or case.seqlen_q <= 0:
        return (True, "edge-case-ok")

    mi = MASK_INT.get(case.mask, 1 if case.mask.startswith(("t:", "b:")) else 0)
    bi = BIAS_INT.get(case.bias, 0)
    scale = 1.0 / (dq**0.5)

    # Build a tiny runner script executed in a fresh process
    runner = f"""\
import ctypes, numpy as np, sys
lib = ctypes.CDLL("{so_path}")
lib.fmha_dispatcher_initialize.argtypes = [ctypes.c_char_p]
lib.fmha_dispatcher_initialize.restype = ctypes.c_int
lib.fmha_dispatcher_run_fwd.argtypes = [
    ctypes.c_void_p,ctypes.c_void_p,ctypes.c_void_p,ctypes.c_void_p,
    ctypes.c_int,ctypes.c_int,ctypes.c_int,ctypes.c_int,ctypes.c_int,
    ctypes.c_int,ctypes.c_int,ctypes.c_float,
    ctypes.c_int,ctypes.c_int,ctypes.c_int,ctypes.c_int,
    ctypes.c_int,ctypes.c_int,ctypes.c_int,
    ctypes.c_int,
    ctypes.c_char_p,ctypes.c_int,
    ctypes.c_int,ctypes.c_int,
    ctypes.c_int,ctypes.c_int,ctypes.c_int,
    ctypes.POINTER(ctypes.c_float)]
lib.fmha_dispatcher_run_fwd.restype = ctypes.c_int
lib.fmha_dispatcher_cleanup.argtypes = []
lib.fmha_dispatcher_cleanup.restype = None
rc = lib.fmha_dispatcher_initialize(b"{arch}")
if rc != 0: print("INIT_FAIL"); sys.exit(1)
np.random.seed(42)
grp={case.mode}
perm={case.perm}
if grp:
  Q=np.ascontiguousarray((np.random.randn({case.batch}*{case.seqlen_q},{case.nhead_q},{dq})*0.3).astype(np.float16))
  K=np.ascontiguousarray((np.random.randn({case.batch}*{case.seqlen_k},{nk},{dq})*0.3).astype(np.float16))
  V=np.ascontiguousarray((np.random.randn({case.batch}*{case.seqlen_k},{nk},{dv})*0.3).astype(np.float16))
  O=np.ascontiguousarray(np.zeros(({case.batch}*{case.seqlen_q},{case.nhead_q},{dv}),dtype=np.float16))
elif perm==1:
  Q=np.ascontiguousarray((np.random.randn({case.batch},{case.nhead_q},{case.seqlen_q},{dq})*0.3).astype(np.float16))
  K=np.ascontiguousarray((np.random.randn({case.batch},{nk},{case.seqlen_k},{dq})*0.3).astype(np.float16))
  V=np.ascontiguousarray((np.random.randn({case.batch},{nk},{case.seqlen_k},{dv})*0.3).astype(np.float16))
  O=np.ascontiguousarray(np.zeros(({case.batch},{case.nhead_q},{case.seqlen_q},{dv}),dtype=np.float16))
else:
  Q=np.ascontiguousarray((np.random.randn({case.batch},{case.seqlen_q},{case.nhead_q},{dq})*0.3).astype(np.float16))
  K=np.ascontiguousarray((np.random.randn({case.batch},{case.seqlen_k},{nk},{dq})*0.3).astype(np.float16))
  V=np.ascontiguousarray((np.random.randn({case.batch},{case.seqlen_k},{nk},{dv})*0.3).astype(np.float16))
  O=np.ascontiguousarray(np.zeros(({case.batch},{case.seqlen_q},{case.nhead_q},{dv}),dtype=np.float16))
t=ctypes.c_float(0.0)
rc=lib.fmha_dispatcher_run_fwd(Q.ctypes.data,K.ctypes.data,V.ctypes.data,O.ctypes.data,\
{case.batch},{case.nhead_q},{nk},{case.seqlen_q},{case.seqlen_k},{dq},{dv},\
{scale},{mi},{bi},{case.lse},{int(case.p_drop > 0)},{traits_dq},{traits_dv},1,{case.perm},b"{case.prec}",{case.mode},\
{-1 if mi == 0 else -1},{-1 if mi == 0 else 0},0,0,0,ctypes.byref(t))
lib.fmha_dispatcher_cleanup()
if rc!=0: print(f"RC{{rc}}"); sys.exit(1)
nz=int(np.count_nonzero(O))
if nz==0: print("ZEROS"); sys.exit(1)
print(f"OK {{t.value:.3f}}ms nz={{nz}}")
"""
    try:
        r = subprocess.run(
            [sys.executable, "-c", runner],
            capture_output=True,
            text=True,
            timeout=30,
            env={**os.environ, "HIP_VISIBLE_DEVICES": "0"},
        )
        out = r.stdout.strip()
        err = r.stderr.strip()
        if r.returncode == 0 and out.startswith("OK"):
            return (True, out)
        msg = out
        if err:
            msg = msg + " ERR:" + err[:80] if msg else err[:120]
        return (False, msg[:160])
    except subprocess.TimeoutExpired:
        return (False, "timeout")


# =========================================================================
# Main
# =========================================================================


def _run_phase(
    label: str,
    cases,
    configs,
    jit_fn,
    test_fn,
    ck_exe,
    ck_bwd_exe,
    args,
    jit_root,
    is_bwd=False,
):
    """Run JIT + test for a set of cases. Returns (jit_time, test_time, stats_dict)."""
    case_key_map: Dict[int, tuple] = {}
    for i, c in enumerate(cases):
        case_key_map[i] = config_key(c)

    lib_for: Dict[tuple, Optional[str]] = {}
    jit_stats = Counter()
    jit_t0 = time.perf_counter()

    if not args.skip_jit:
        print(f"\n--- {label} JIT ({len(configs)} cfgs, {args.workers} workers) ---")
        futures = {}
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            for key in configs:
                name = ("bwd_" if is_bwd else "") + config_name(key)
                out = jit_root / name
                futures[pool.submit(jit_fn, key, out, args.arch)] = (key, name, out)
            done = 0
            for f in as_completed(futures):
                key, name, out = futures[f]
                ok, msg, elapsed = f.result()
                done += 1
                if ok:
                    lib_for[key] = msg
                    jit_stats["ok"] += 1
                else:
                    lib_for[key] = None
                    jit_stats["fail"] += 1
                if done % max(1, len(configs) // 20) == 0 or done <= 3 or not ok:
                    tag = "OK" if ok else f"FAIL({msg[:50]})"
                    print(f"  [{done}/{len(configs)}] {name}  {elapsed:.1f}s  {tag}")
    else:
        for key in configs:
            name = ("bwd_" if is_bwd else "") + config_name(key)
            out = jit_root / name
            sos = sorted(out.glob("libdispatcher_fmha_*.so")) if out.exists() else []
            lib_for[key] = str(sos[0]) if sos else None
            jit_stats["ok" if sos else "missing"] += 1

    jit_elapsed = time.perf_counter() - jit_t0
    print(f"  JIT done: {dict(jit_stats)}  ({jit_elapsed:.0f}s)")

    ck_cnt = Counter()
    disp_cnt = Counter()
    par_cnt = Counter()
    failures = []
    test_t0 = time.perf_counter()
    exe = ck_bwd_exe if is_bwd else ck_exe

    print(f"\n--- {label} tests: {len(cases)} cases (sequential) ---")
    for i, case in enumerate(cases):
        if (i + 1) % 50 == 0 or i == 0:
            el = time.perf_counter() - test_t0
            rate = (i + 1) / max(el, 0.01)
            print(f"  [{i + 1}/{len(cases)}] {el:.0f}s  ({rate:.1f}/s)")

        ck_ok = run_ck_test(exe, case)[0] if exe else None
        key = case_key_map.get(i)
        so = lib_for.get(key) if key else None
        if so:
            d_ok, d_msg = test_fn(so, case, key, args.arch)
        else:
            d_ok, d_msg = None, "no-lib"

        ck_cnt["pass" if ck_ok else ("fail" if ck_ok is False else "skip")] += 1
        disp_cnt["pass" if d_ok else ("fail" if d_ok is False else "skip")] += 1
        if ck_ok is not None and d_ok is not None:
            if ck_ok == d_ok:
                par_cnt["match"] += 1
            else:
                par_cnt["mismatch"] += 1
                failures.append(
                    dict(
                        idx=i,
                        dir=label,
                        ck=ck_ok,
                        disp=d_ok,
                        msg=d_msg,
                        hq=case.hdim_q,
                        hv=case.effective_hdim_v(),
                        mask=case.mask,
                        bias=case.bias,
                        nq=case.nhead_q,
                        nk=case.effective_nhead_k(),
                        sq=case.seqlen_q,
                        sk=case.seqlen_k,
                    )
                )
        else:
            par_cnt["n/a"] += 1
        if d_ok is False:
            dv = case.effective_hdim_v()
            nk = case.effective_nhead_k()
            print(
                f"    FAIL[{i}] h={case.hdim_q}x{dv} m={case.mask} b={case.bias}"
                f" nq={case.nhead_q} nk={nk} -> {d_msg[:80]}"
            )

    test_elapsed = time.perf_counter() - test_t0
    return (
        jit_elapsed,
        test_elapsed,
        dict(
            jit=dict(jit_stats),
            ck=dict(ck_cnt),
            dispatcher=dict(disp_cnt),
            parity=dict(par_cnt),
            failures=failures[:100],
        ),
    )


def find_ck_bwd_exe() -> Optional[str]:
    for p in [
        "/tmp/ck_fmha_full/bin/tile_example_fmha_bwd",
        "/tmp/ck_fmha_build/bin/tile_example_fmha_bwd",
    ]:
        if os.path.exists(p):
            return p
    return None


def run_dispatcher_bwd_test(
    so_path: str, case: TestCase, key: tuple, arch: str
) -> Tuple[bool, str]:
    """Backward test stub -- validates kernel loads and produces nonzero grads."""
    if case.seqlen_k <= 0 or case.seqlen_q <= 0:
        return (True, "edge-case-ok")

    # For now, just verify the bwd .so loads and initializes (kernel selection).
    # Full GPU bwd execution requires run_bwd ABI updates matching fwd.
    runner = f"""\
import ctypes, sys
lib = ctypes.CDLL("{so_path}")
lib.fmha_dispatcher_initialize.argtypes = [ctypes.c_char_p]
lib.fmha_dispatcher_initialize.restype = ctypes.c_int
lib.fmha_dispatcher_kernel_count.argtypes = []
lib.fmha_dispatcher_kernel_count.restype = ctypes.c_int
lib.fmha_dispatcher_cleanup.argtypes = []
lib.fmha_dispatcher_cleanup.restype = None
rc = lib.fmha_dispatcher_initialize(b"{arch}")
if rc != 0: print("INIT_FAIL"); sys.exit(1)
n = lib.fmha_dispatcher_kernel_count()
lib.fmha_dispatcher_cleanup()
if n < 3: print(f"KERNELS={{n}}"); sys.exit(1)
print(f"OK kernels={{n}}")
"""
    try:
        r = subprocess.run(
            [sys.executable, "-c", runner],
            capture_output=True,
            text=True,
            timeout=15,
            env={**os.environ, "HIP_VISIBLE_DEVICES": "0"},
        )
        out = r.stdout.strip()
        err = r.stderr.strip()
        if r.returncode == 0 and out.startswith("OK"):
            return (True, out)
        msg = out
        if err:
            msg = msg + " ERR:" + err[:80] if msg else err[:120]
        return (False, msg[:160])
    except subprocess.TimeoutExpired:
        return (False, "timeout")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-cases", type=int, default=0, help="0 = all")
    parser.add_argument("--max-configs", type=int, default=0)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--arch", default="gfx950")
    parser.add_argument("--skip-jit", action="store_true")
    parser.add_argument("--skip-ck", action="store_true")
    parser.add_argument("--fwd-only", action="store_true")
    parser.add_argument("--bwd-only", action="store_true")
    parser.add_argument("--report", default="parity_report.json")
    args = parser.parse_args()

    ck_exe = find_ck_exe() if not args.skip_ck else None
    ck_bwd_exe = find_ck_bwd_exe() if not args.skip_ck else None

    print("=" * 80)
    print("FMHA Full Parity Test (fwd + bwd)")
    print("=" * 80)
    print(f"  CK fwd exe:   {ck_exe or 'N/A'}")
    print(f"  CK bwd exe:   {ck_bwd_exe or 'N/A'}")
    print(f"  GPU arch:      {args.arch}")
    print(f"  JIT workers:   {args.workers}")

    jit_root = Path("/tmp/fmha_parity_jit")
    jit_root.mkdir(parents=True, exist_ok=True)

    all_results = {}
    total_jit = 0.0
    total_test = 0.0

    # ---- Forward ----
    if not args.bwd_only:
        fwd_cases = generate_fwd_fp16_bf16_matrix()
        if args.max_cases > 0:
            fwd_cases = fwd_cases[: args.max_cases]
        fwd_configs = {}
        for c in fwd_cases:
            k = config_key(c)
            fwd_configs[k] = True
        if args.max_configs > 0:
            fwd_configs = dict(list(fwd_configs.items())[: args.max_configs])
        print(f"\n  FWD: {len(fwd_cases)} cases, {len(fwd_configs)} configs")

        jt, tt, stats = _run_phase(
            "FWD",
            fwd_cases,
            fwd_configs,
            _jit_one,
            run_dispatcher_test,
            ck_exe,
            ck_bwd_exe,
            args,
            jit_root,
        )
        all_results["fwd"] = stats
        total_jit += jt
        total_test += tt

    # ---- Backward ----
    if not args.fwd_only:
        bwd_cases = generate_bwd_matrix()
        if args.max_cases > 0:
            bwd_cases = bwd_cases[: args.max_cases]
        bwd_configs = {}
        for c in bwd_cases:
            k = config_key(c)
            bwd_configs[k] = True
        if args.max_configs > 0:
            bwd_configs = dict(list(bwd_configs.items())[: args.max_configs])
        print(f"\n  BWD: {len(bwd_cases)} cases, {len(bwd_configs)} configs")

        jt, tt, stats = _run_phase(
            "BWD",
            bwd_cases,
            bwd_configs,
            _jit_one_bwd,
            run_dispatcher_bwd_test,
            ck_exe,
            ck_bwd_exe,
            args,
            jit_root,
            is_bwd=True,
        )
        all_results["bwd"] = stats
        total_jit += jt
        total_test += tt

    # ---- Padding edge cases ----
    if not args.bwd_only:
        pad_cases = generate_padding_matrix()
        pad_configs = {}
        for c in pad_cases:
            k = config_key(c)
            pad_configs[k] = True
        print(f"\n  PAD: {len(pad_cases)} cases, {len(pad_configs)} configs")
        jt, tt, stats = _run_phase(
            "PAD",
            pad_cases,
            pad_configs,
            _jit_one,
            run_dispatcher_test,
            ck_exe,
            ck_bwd_exe,
            args,
            jit_root,
        )
        all_results["padding"] = stats
        total_jit += jt
        total_test += tt

    # ---- FP8 ----
    if not args.bwd_only:
        fp8_cases = generate_fp8_matrix()
        fp8_configs = {}
        for c in fp8_cases:
            k = config_key(c)
            fp8_configs[k] = True
        print(f"\n  FP8: {len(fp8_cases)} cases, {len(fp8_configs)} configs")
        jt, tt, stats = _run_phase(
            "FP8",
            fp8_cases,
            fp8_configs,
            _jit_one,
            run_dispatcher_test,
            ck_exe,
            ck_bwd_exe,
            args,
            jit_root,
        )
        all_results["fp8"] = stats
        total_jit += jt
        total_test += tt

    # ---- SplitKV ----
    if not args.bwd_only:
        skv_cases = generate_splitkv_matrix()
        if args.max_cases > 0:
            skv_cases = skv_cases[: args.max_cases]
        skv_configs = {}
        for c in skv_cases:
            k = config_key(c)
            skv_configs[k] = True
        print(f"\n  SKV: {len(skv_cases)} cases, {len(skv_configs)} configs")
        jt, tt, stats = _run_phase(
            "SKV",
            skv_cases,
            skv_configs,
            _jit_one,
            run_dispatcher_test,
            ck_exe,
            ck_bwd_exe,
            args,
            jit_root,
        )
        all_results["splitkv"] = stats
        total_jit += jt
        total_test += tt

    # ---- Report ----
    print(f"\n{'=' * 80}")
    print("FMHA Full Parity Report")
    print(f"{'=' * 80}")
    print(f"  JIT total:   {total_jit:.0f}s")
    print(f"  Test total:  {total_test:.0f}s")
    for direction, stats in all_results.items():
        d = stats["dispatcher"]
        p = stats["parity"]
        print(f"\n  [{direction.upper()}]")
        print(f"    CK:         {stats['ck']}")
        print(
            f"    Dispatcher:  {d.get('pass', 0)} pass, {d.get('fail', 0)} fail,"
            f" {d.get('skip', 0)} skip"
        )
        print(
            f"    Parity:      {p.get('match', 0)} match, {p.get('mismatch', 0)} mismatch"
        )
        if stats.get("failures"):
            print("    Failures[0:5]:")
            for f in stats["failures"][:5]:
                print(
                    f"      [{f['idx']}] ck={f['ck']} disp={f['disp']}"
                    f" h={f['hq']}x{f['hv']} -> {f['msg'][:50]}"
                )
    print(f"{'=' * 80}")

    with open(args.report, "w") as fp:
        json.dump(
            dict(jit_time_s=total_jit, test_time_s=total_test, results=all_results),
            fp,
            indent=2,
        )
    print(f"\nSaved {args.report}")

    total_fail = sum(
        r["dispatcher"].get("fail", 0) + r["dispatcher"].get("skip", 0)
        for r in all_results.values()
    )
    return 1 if total_fail > 0 else 0


if __name__ == "__main__":
    sys.exit(main())