File size: 2,644 Bytes
8e9f35a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import importlib
import pytest
import torch
import torch.nn.functional as F
from kernels import get_kernel
kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)
def _inputs(tokens=256, heads=4):
torch.manual_seed(42)
q = torch.randn(
1,
tokens,
heads,
128,
device="cuda",
dtype=torch.bfloat16,
)
return q, torch.randn_like(q), torch.randn_like(q)
@pytest.mark.kernels_ci
def test_backend_dispatch_contract():
interface = importlib.import_module(f"{kernel.__name__}.interface")
assert interface._backend_for_arch((8, 0), cute_available=True) == "triton"
assert interface._backend_for_arch((8, 9), cute_available=True) == "triton"
assert interface._backend_for_arch((9, 0), cute_available=True) == "cute_sm90"
assert interface._backend_for_arch((10, 0), cute_available=True) == "cute_sm100"
assert interface._backend_for_arch((12, 0), cute_available=True) == "cute_sm120"
assert interface._backend_for_arch((9, 0), cute_available=False) == "triton"
assert interface._backend_for_arch((10, 0), cute_available=False) == "triton"
assert interface._backend_for_arch((12, 0), cute_available=False) == "triton"
@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_full_sink_matches_sdpa():
capability = torch.cuda.get_device_capability()
if capability[0] < 8:
pytest.skip("Sol-Attn requires compute capability 8.0 or newer")
q, k, v = _inputs()
expected = F.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
).transpose(1, 2)
actual = kernel.sol_attn(
q,
k,
v,
tau=1.0,
thresh_type="exact",
sink_start=0,
sink_tokens=q.shape[1],
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2)
@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_selected_backend_matches_triton_reference():
capability = torch.cuda.get_device_capability()
if capability[0] < 8:
pytest.skip("Sol-Attn requires compute capability 8.0 or newer")
triton_ref = importlib.import_module(f"{kernel.__name__}.triton_ref")
q, k, v = _inputs()
expected = triton_ref.sol_attn(
q,
k,
v,
tau=1.0,
thresh_type="exact",
)
actual = kernel.sol_attn(
q,
k,
v,
tau=1.0,
thresh_type="exact",
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2)
|