"""Minimal smoke test for flash-attn-4-sm120-sncbl.

Verifies:
- Package imports cleanly.
- `flash_attn_func` runs on CUDA and produces output matching a naive
  PyTorch reference within bf16 numerical tolerance.
- Runs on SM120 (our target arch); skips on other GPUs.

Not a correctness suite — the 5 bundled PRs each have their own test
file upstream (tests/cute/test_*.py). This is just "does the package
load and compute anything on SM120".
"""

from __future__ import annotations

import pytest
import torch


def _get_sm() -> tuple[int, int]:
    if not torch.cuda.is_available():
        pytest.skip("CUDA not available")
    return torch.cuda.get_device_capability(0)


def _ref_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool) -> torch.Tensor:
    # q/k/v: (B, S, H, D)
    scale = 1.0 / (q.shape[-1] ** 0.5)
    qf, kf, vf = q.float(), k.float(), v.float()
    # (B, H, S_q, S_k)
    scores = torch.einsum("bshd,bthd->bhst", qf, kf) * scale
    if causal:
        s_q, s_k = q.shape[1], k.shape[1]
        mask = torch.triu(
            torch.full((s_q, s_k), float("-inf"), device=q.device),
            diagonal=s_k - s_q + 1,
        )
        scores = scores + mask
    probs = torch.softmax(scores, dim=-1)
    out = torch.einsum("bhst,bthd->bshd", probs, vf)
    return out.to(q.dtype)


def test_import():
    from flash_attn_4_sm120_sncbl import flash_attn_func, flash_attn_varlen_func  # noqa: F401


@pytest.mark.parametrize("causal", [False, True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_forward_correctness(causal: bool, dtype: torch.dtype):
    major, _ = _get_sm()
    if major != 12:
        pytest.skip(f"SM120 smoke test; current SM is {major}.x")
    from flash_attn_4_sm120_sncbl import flash_attn_func

    torch.manual_seed(0)
    B, S_q, S_k, H, D = 1, 128, 128, 4, 64
    q = torch.randn(B, S_q, H, D, device="cuda", dtype=dtype)
    k = torch.randn(B, S_k, H, D, device="cuda", dtype=dtype)
    v = torch.randn(B, S_k, H, D, device="cuda", dtype=dtype)

    out, *_ = flash_attn_func(q, k, v, causal=causal)
    ref = _ref_attention(q, k, v, causal=causal)

    max_diff = (out.float() - ref.float()).abs().max().item()
    mean_diff = (out.float() - ref.float()).abs().mean().item()
    # bf16 accum tol; matches the tolerance used by upstream FA tests.
    assert max_diff < 1e-1, f"max_diff={max_diff:.4f} mean={mean_diff:.5f}"


if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-v"]))