#!/usr/bin/env python3
"""CUDA smoke check for the local Triton + SageAttention installation."""

from __future__ import annotations

import json

import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from sageattention import sageattn
from torch.nn.attention import SDPBackend, sdpa_kernel


@triton.jit
def _add_kernel(x_ptr, y_ptr, out_ptr, count: tl.constexpr):
    offsets = tl.arange(0, count)
    x = tl.load(x_ptr + offsets)
    y = tl.load(y_ptr + offsets)
    tl.store(out_ptr + offsets, x + y)


def main() -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable")

    left = torch.rand(128, device="cuda", dtype=torch.float32)
    right = torch.rand_like(left)
    triton_result = torch.empty_like(left)
    _add_kernel[(1,)](left, right, triton_result, count=left.numel())
    torch.cuda.synchronize()
    triton_max_error = float((triton_result - (left + right)).abs().max())
    if triton_max_error != 0.0:
        raise RuntimeError(f"Triton vector-add mismatch: {triton_max_error}")

    shape = (1, 8, 512, 128)
    query = torch.randn(shape, device="cuda", dtype=torch.float16)
    key = torch.randn_like(query)
    value = torch.randn_like(query)
    torch.backends.cuda.enable_math_sdp(True)
    with sdpa_kernel(SDPBackend.MATH):
        reference = F.scaled_dot_product_attention(query, key, value)
    sage = sageattn(query, key, value)
    torch.cuda.synchronize()

    absolute_error = (sage.float() - reference.float()).abs()
    mean_absolute_error = float(absolute_error.mean())
    max_absolute_error = float(absolute_error.max())
    if mean_absolute_error >= 0.1:
        raise RuntimeError(
            f"SageAttention mean absolute error is too high: {mean_absolute_error}"
        )

    print(
        json.dumps(
            {
                "status": "ok",
                "torch": torch.__version__,
                "torch_cuda": torch.version.cuda,
                "triton": triton.__version__,
                "gpu": torch.cuda.get_device_name(0),
                "compute_capability": list(torch.cuda.get_device_capability(0)),
                "triton_vector_add_max_error": triton_max_error,
                "sage_shape": list(shape),
                "sage_mean_absolute_error": mean_absolute_error,
                "sage_max_absolute_error": max_absolute_error,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()
