rm core tests for now

tests/test_core.py

@@ -1,73 +0,0 @@
-import math
-import pytest
-import torch
-
-import sage_attention as sa
-
-
-cuda_available = torch.cuda.is_available()
-
-
-def current_sm():
-    if not cuda_available:
-        return None
-    major, minor = torch.cuda.get_device_capability(0)
-    return f"sm{major}{minor}"
-
-
-@pytest.mark.skipif(not cuda_available, reason="CUDA is required")
-@pytest.mark.parametrize("tensor_layout", ["HND", "NHD"])
-@pytest.mark.parametrize("head_dim", [64, 128])
-@pytest.mark.parametrize("return_lse", [False, True])
-def test_sageattn_runs_and_shapes(tensor_layout, head_dim, return_lse):
-    device = "cuda"
-    dtype = torch.float16
-
-    # Small, nontrivial shapes; pad path will be exercised for head_dim=64
-    if tensor_layout == "HND":
-        q = torch.randn(2, 6, 129, head_dim, dtype=dtype, device=device)
-        k = torch.randn(2, 3, 257, head_dim, dtype=dtype, device=device)
-        v = torch.randn(2, 3, 257, head_dim, dtype=dtype, device=device)
-        expected_o_shape = (2, 6, 129, head_dim)
-        expected_lse_shape = (2, 6, 129)
-    else:
-        q = torch.randn(2, 129, 6, head_dim, dtype=dtype, device=device)
-        k = torch.randn(2, 257, 3, head_dim, dtype=dtype, device=device)
-        v = torch.randn(2, 257, 3, head_dim, dtype=dtype, device=device)
-        expected_o_shape = (2, 129, 6, head_dim)
-        expected_lse_shape = (2, 6, 129)
-
-    sm = current_sm()
-
-    # Some backends may not be compiled on this GPU; skip gracefully if unsupported
-    try:
-        out = sa.sageattn(
-            q, k, v, tensor_layout=tensor_layout, is_causal=False, return_lse=return_lse
-        )
-    except ValueError as e:
-        if "Unsupported CUDA architecture" in str(e):
-            pytest.skip(f"Unsupported arch for this build: {sm}")
-        raise
-
-    if return_lse:
-        o, lse = out
-        assert lse.shape == expected_lse_shape and torch.isfinite(lse).all()
-    else:
-        o = out
-
-    assert o.shape == expected_o_shape
-    assert o.dtype == dtype
-    assert o.device.type == "cuda"
-
-
-@pytest.mark.skipif(not cuda_available, reason="CUDA is required")
-def test_sageattn_raises_on_unsupported_head_dim():
-    device = "cuda"
-    dtype = torch.float16
-    # head_dim > 128 should raise
-    q = torch.randn(1, 2, 8, 192, dtype=dtype, device=device)
-    k = torch.randn(1, 1, 8, 192, dtype=dtype, device=device)
-    v = torch.randn(1, 1, 8, 192, dtype=dtype, device=device)
-
-    with pytest.raises(ValueError):
-        sa.sageattn(q, k, v)