"""K04 flash_attn (aiter Triton _attn_fwd) — MLA prefill attention, head_dim_qk=192, head_dim_v=128.
Per-rank heads = 64/TP4 = 16 (matches trace 'fmha_fwd_hd192x128 ... qh16'). causal.
Optimizable: kernel_jit.py (_attn_fwd). Reference: installed-original flash_attn_func (golden, fwd deterministic)."""
import os
import sys
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
_TASK_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(_TASK_ROOT, "scripts"))
torch.cuda.init()
torch.empty(1, device="cuda")
import host as loc  # noqa: E402
from aiter.ops.triton.attention.mha import flash_attn_func as ref_attn  # noqa: E402
from _bench_common import make_argparser, run_modes  # noqa: E402

H = 16          # heads per rank (64 // TP4)
DQK = 192       # qk_nope(128)+qk_rope(64)
DV = 128        # v_head_dim
DTYPE = torch.bfloat16
SCALE = 1.0 / (DQK ** 0.5)
# (name, batch, seqlen)
SHAPES = [("s512", 1, 512), ("s1024", 1, 1024), ("s2048", 1, 2048)]


def build(b, s):
    torch.manual_seed(0)
    q = torch.randn(b, s, H, DQK, device="cuda", dtype=DTYPE)
    k = torch.randn(b, s, H, DQK, device="cuda", dtype=DTYPE)
    v = torch.randn(b, s, H, DV, device="cuda", dtype=DTYPE)
    return q, k, v


def main():
    args = make_argparser("k04 flash_attn prefill").parse_args()
    cases = []
    for name, b, s in SHAPES:
        q, k, v = build(b, s)

        def run(q=q, k=k, v=v):
            return loc.flash_attn_func(q, k, v, softmax_scale=SCALE, causal=True)

        def check(q=q, k=k, v=v):
            o = loc.flash_attn_func(q, k, v, softmax_scale=SCALE, causal=True)
            r = ref_attn(q, k, v, softmax_scale=SCALE, causal=True)
            o = o[0] if isinstance(o, tuple) else o
            r = r[0] if isinstance(r, tuple) else r
            return torch.allclose(o.float(), r.float(), atol=2e-2, rtol=2e-2)

        cases.append({"name": name, "run": run, "check": check})
    run_modes(args, cases)


if __name__ == "__main__":
    main()