#!/usr/bin/env python3
import argparse

import torch
import torch.nn.functional as F

from liger_kernel.ops.geglu import LigerGELUMulFunction
from liger_kernel.ops.rms_norm import LigerRMSNormFunction
from liger_kernel.ops.swiglu import LigerSiLUMulFunction

EPS = 1e-6
TOL = 0.1
BATCH = 8
HIDDEN = 4096


def ref_rmsnorm(x, weight, eps):
    xf = x.float()
    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype)


def rmsnorm_case(dev):
    x = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
    weight = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) * 0.1
    return (x, weight, EPS, 0, "llama", True, None), ref_rmsnorm(x, weight, EPS)


def silu_case(dev):
    a = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
    b = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
    return (a, b), F.silu(a) * b


def gelu_case(dev):
    a = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
    b = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
    return (a, b), F.gelu(a, approximate="tanh") * b


CASES = [
    ("LigerRMSNorm", LigerRMSNormFunction, rmsnorm_case),
    ("LigerSiLUMul", LigerSiLUMulFunction, silu_case),
    ("LigerGELUMul", LigerGELUMulFunction, gelu_case),
]


def run_case(name, fn, build, devices):
    print(f"\n===== {name} =====", flush=True)

    # Launch on device 0 first, then keep device 0 active while passing
    # device 1 tensors to reproduce the mismatch.
    for dev in devices:
        torch.manual_seed(0)
        args, ref = build(dev)
        torch.cuda.set_device(devices[0])

        try:
            with torch.no_grad():
                out = fn.apply(*args)
            finite = bool(torch.isfinite(out).all().item())
            max_abs = float((out.float() - ref.float()).abs().max().item())
            ok = finite and max_abs < TOL
            print(
                f"  input={dev} active=cuda:{torch.cuda.current_device()} "
                f"finite={finite} max|diff|={max_abs:.6f} -> {'OK' if ok else 'CORRUPT'}",
                flush=True,
            )
        except Exception as exc:
            print(
                f"  input={dev} active=cuda:{torch.cuda.current_device()} "
                f"CRASH {type(exc).__name__}: {str(exc).splitlines()[0][:120]}",
                flush=True,
            )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--devices", nargs=2, default=["cuda:0", "cuda:1"])
    args = parser.parse_args()

    if torch.cuda.device_count() < 2:
        raise SystemExit(f"need >=2 CUDA devices, found {torch.cuda.device_count()}")

    print(f"liger-kernel multi-GPU repro on {args.devices}", flush=True)
    for name, fn, build in CASES:
        run_case(name, fn, build, args.devices)


if __name__ == "__main__":
    main()