Merge pull request #254 from HKUSTDial:optim_triton_version

Refactor test utilities and add CUDA tensor operation tests

Author: LoserCheems

tests/test_dense_base_backward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_backward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_dense_base_backward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_backward_base_case(
+        kind="dense",
+        batch_size=2,
+        seqlen_q=80,
+        seqlen_k=96,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )

tests/test_dense_base_forward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_forward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_dense_base_forward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_forward_base_case(
+        kind="dense",
+        batch_size=2,
+        seqlen_q=96,
+        seqlen_k=128,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )

tests/test_dense_varlen_backward.py

@@ -0,0 +1,22 @@
+import pytest
+import torch
+
+from test_utils import run_backward_varlen_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_dense_varlen_backward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_backward_varlen_case(
+        kind="dense",
+        lens_q=[19, 27, 41],
+        lens_k=[25, 31, 43],
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )

tests/test_dense_varlen_forward.py

@@ -0,0 +1,22 @@
+import pytest
+import torch
+
+from test_utils import run_forward_varlen_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_dense_varlen_forward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_forward_varlen_case(
+        kind="dense",
+        lens_q=[17, 33, 29],
+        lens_k=[23, 37, 31],
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )

tests/test_gated_base_backward.py

@@ -0,0 +1,24 @@
+import pytest
+import torch
+
+from test_utils import run_backward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_gated_base_backward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_backward_base_case(
+        kind="gated",
+        batch_size=2,
+        seqlen_q=80,
+        seqlen_k=96,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+        is_logsigmoid_gate=True,
+    )

tests/test_gated_base_forward.py

@@ -0,0 +1,24 @@
+import pytest
+import torch
+
+from test_utils import run_forward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_gated_base_forward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_forward_base_case(
+        kind="gated",
+        batch_size=2,
+        seqlen_q=96,
+        seqlen_k=128,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+        is_logsigmoid_gate=True,
+    )

tests/test_gated_varlen_backward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_backward_varlen_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_gated_varlen_backward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_backward_varlen_case(
+        kind="gated",
+        lens_q=[19, 27, 41],
+        lens_k=[25, 31, 43],
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+        is_logsigmoid_gate=True,
+    )

tests/test_gated_varlen_forward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_forward_varlen_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_gated_varlen_forward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_forward_varlen_case(
+        kind="gated",
+        lens_q=[17, 33, 29],
+        lens_k=[23, 37, 31],
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+        is_logsigmoid_gate=True,
+    )

tests/test_sparse_base_backward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_backward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_sparse_base_backward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_backward_base_case(
+        kind="sparse",
+        batch_size=2,
+        seqlen_q=80,
+        seqlen_k=96,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )

tests/test_sparse_base_forward.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from test_utils import run_forward_base_case, set_seed
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required"
+)
+
+
+@pytest.mark.parametrize("is_causal", [False, True])
+def test_sparse_base_forward_correctness(is_causal: bool) -> None:
+    set_seed(0)
+    run_forward_base_case(
+        kind="sparse",
+        batch_size=2,
+        seqlen_q=96,
+        seqlen_k=128,
+        num_heads_q=8,
+        num_heads_kv=4,
+        head_dim=64,
+        is_causal=is_causal,
+    )