Split the special ops test cases (#2600)

Author: tengqm

tests/test_apply_rotary_pos_emb.py

@@ -0,0 +1,130 @@
+import random
+import time
+from typing import Optional
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+from . import conftest as cfg
+
+random.seed(time.time() // 100)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.cohere.modeling_cohere.rotate_half
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py
+def rotate_interleave(x):
+    """Rotates interleave the hidden dims of the input."""
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def _torch_apply_rotary_pos_emb(
+    q,
+    k,
+    cos,
+    sin,
+    position_ids: Optional[torch.Tensor] = None,
+    rotary_interleaved: bool = False,
+):
+    q = q.float()
+    k = k.float()
+    if position_ids is None:
+        cos = cos[None, : q.size(-3), None, :]
+        sin = sin[None, : q.size(-3), None, :]
+    else:
+        cos = cos[position_ids].unsqueeze(-2)  # [bs, seq_len, 1, dim/2]
+        sin = sin[position_ids].unsqueeze(-2)  # [bs, seq_len, 1, dim/2]
+    if rotary_interleaved:
+        cos = torch.repeat_interleave(cos, 2, dim=-1)  # [bs, seq_len, 1, dim]
+        sin = torch.repeat_interleave(sin, 2, dim=-1)  # [bs, seq_len, 1, dim]
+        rotate_fn = rotate_interleave
+    else:
+        cos = torch.cat([cos, cos], dim=-1)  # [bs, seq_len, 1, dim]
+        sin = torch.cat([sin, sin], dim=-1)  # [bs, seq_len, 1, dim]
+        rotate_fn = rotate_half
+
+    q_embed = (q * cos) + (rotate_fn(q) * sin)
+    k_embed = (k * cos) + (rotate_fn(k) * sin)
+
+    return q_embed, k_embed
+
+
+def _get_rope_cos_sin(max_seq_len, dim, dtype, base=10000, device=flag_gems.device):
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+    t = torch.arange(max_seq_len, device=device, dtype=inv_freq.dtype)
+    freqs = torch.outer(t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+
+    return cos, sin
+
+
+@pytest.mark.apply_rotary_pos_emb
+@pytest.mark.parametrize("batch_size", [2] if cfg.TO_CPU else [4, 8])
+@pytest.mark.parametrize("max_seq_len", [16] if cfg.TO_CPU else [512, 2048])
+@pytest.mark.parametrize("q_heads,k_heads", [(8, 1), (6, 2), (1, 1), (8, 8)])
+@pytest.mark.parametrize("head_dim", [8] if cfg.TO_CPU else [64, 96, 128, 256])
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+@pytest.mark.parametrize("rotary_interleaved", [True, False])
+@pytest.mark.parametrize("has_pos_id", [True, False])
+def test_apply_rotary_pos_emb(
+    batch_size,
+    max_seq_len,
+    q_heads,
+    k_heads,
+    head_dim,
+    dtype,
+    has_pos_id,
+    rotary_interleaved,
+):
+    seq_len = torch.randint(1, max_seq_len, (1,)).item()
+    q = torch.randn(
+        (batch_size, seq_len, q_heads, head_dim), dtype=dtype, device=flag_gems.device
+    )
+    k = torch.randn(
+        (batch_size, seq_len, k_heads, head_dim), dtype=dtype, device=flag_gems.device
+    )
+
+    position_ids = torch.randint(
+        0, max_seq_len, (batch_size, seq_len), device=flag_gems.device
+    )
+    cos, sin = _get_rope_cos_sin(max_seq_len, head_dim, dtype, device=flag_gems.device)
+
+    ref_q = utils.to_reference(q, True)
+    ref_k = utils.to_reference(k, True)
+    ref_cos = utils.to_reference(cos, True)
+    ref_sin = utils.to_reference(sin, True)
+    ref_position_ids = utils.to_reference(position_ids)
+
+    q_embed_ref, k_embed_ref = _torch_apply_rotary_pos_emb(
+        q=ref_q,
+        k=ref_k,
+        cos=ref_cos,
+        sin=ref_sin,
+        position_ids=ref_position_ids if has_pos_id else None,
+        rotary_interleaved=rotary_interleaved,
+    )
+    q_embed_out, k_embed_out = flag_gems.apply_rotary_pos_emb(
+        q=q,
+        k=k,
+        cos=cos,
+        sin=sin,
+        position_ids=position_ids if has_pos_id else None,
+        rotary_interleaved=rotary_interleaved,
+    )
+
+    utils.gems_assert_close(q_embed_out, q_embed_ref, dtype)
+    utils.gems_assert_close(k_embed_out, k_embed_ref, dtype)

tests/test_arange.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+from . import conftest as cfg
+
+
+@pytest.mark.arange
+@pytest.mark.parametrize("start", utils.ARANGE_START)
+@pytest.mark.parametrize("step", [1, 2, 5])
+@pytest.mark.parametrize("end", [128, 256, 1024])
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.ALL_INT_DTYPES + [None])
+@pytest.mark.parametrize("device", [flag_gems.device, None])
+@pytest.mark.parametrize(
+    "pin_memory", [False, None]
+)  # Since triton only target to GPU, pin_memory only used in CPU tensors.
+def test_arange(start, step, end, dtype, device, pin_memory):
+    ref_out = torch.arange(
+        start,
+        end,
+        step,
+        dtype=dtype,
+        device="cpu" if cfg.TO_CPU else device,
+        pin_memory=pin_memory,
+    )
+    with flag_gems.use_gems():
+        res_out = torch.arange(
+            start, end, step, dtype=dtype, device=device, pin_memory=pin_memory
+        )
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_assert_async.py

@@ -0,0 +1,38 @@
+import pytest
+import torch
+
+import flag_gems
+
+
+@pytest.mark.assert_async
+@pytest.mark.parametrize(
+    "shape, value, expected_err, match_str",
+    [
+        ((), 1.0, None, None),
+        ((2,), 1.0, RuntimeError, "is ambiguous"),
+        ((1,), 1.0, None, None),
+    ],
+)
+def test_assert_async(shape, value, expected_err, match_str):
+    msg = "Assertion failed!"
+    inp_pt = torch.full(shape, value, device=flag_gems.device)
+    inp_triton = inp_pt.clone()
+    if expected_err:
+        with flag_gems.use_gems():
+            with pytest.raises(expected_err, match=match_str):
+                flag_gems._assert_async(inp_triton, msg)
+                if value == 0:
+                    torch.cuda.synchronize()
+    else:
+        with flag_gems.use_gems():
+            flag_gems._assert_async(inp_triton, msg)
+            torch.cuda.synchronize()
+
+    if expected_err:
+        with pytest.raises(expected_err, match=match_str):
+            torch._assert_async(inp_pt, msg)
+            if value == 0:
+                torch.cuda.synchronize()
+    else:
+        torch._assert_async(inp_pt, msg)
+        torch.cuda.synchronize()

tests/test_cat.py

@@ -0,0 +1,86 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+CAT_SHAPES = [
+    [(1, 32), (8, 32)],
+    [(16, 128), (32, 128)],
+    [(1024, 1024), (1024, 1024)],
+    [(1, 1024, 256), (8, 1024, 256), (16, 1024, 256)],
+    [(16, 320, 15), (32, 320, 15), (64, 320, 15)],
+    [(16, 128, 64, 64), (16, 128, 64, 64), (24, 128, 64, 64), (32, 128, 64, 64)],
+]
+
+
+def gen_cat_shapes_dim(shapes):
+    results = []
+    for tensor_shapes in shapes:
+        assert all(
+            [len(s) == len(tensor_shapes[0]) for s in tensor_shapes]
+        ), "All tensor rank must agree."
+
+        assert all(
+            [s[-1] == tensor_shapes[0][-1] for s in tensor_shapes]
+        ), "All tensor must have same shape except cat dim."
+
+        rank = len(tensor_shapes[0])
+        results.append([tensor_shapes, 0])
+        for dim in range(1, rank):
+            results.append(
+                [[(s[dim], *s[1:dim], s[0], *s[dim + 1 :]) for s in tensor_shapes], dim]
+            )
+            results.append(
+                [
+                    [(s[dim], *s[1:dim], s[0], *s[dim + 1 :]) for s in tensor_shapes],
+                    dim - rank,
+                ]
+            )
+    return results
+
+
+@pytest.mark.cat
+@pytest.mark.parametrize("shape, dim", gen_cat_shapes_dim(CAT_SHAPES))
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.INT_DTYPES)
+def test_cat(shape, dim, dtype):
+    if dtype in utils.FLOAT_DTYPES:
+        inp = [torch.randn(s, dtype=dtype, device=flag_gems.device) for s in shape]
+    else:
+        inp = [
+            torch.randint(low=0, high=0x7FFF, size=s, dtype=dtype, device="cpu").to(
+                flag_gems.device
+            )
+            for s in shape
+        ]
+    ref_inp = [utils.to_reference(_) for _ in inp]
+    ref_out = torch.cat(ref_inp, dim)
+
+    with flag_gems.use_gems():
+        res_out = torch.cat(inp, dim)
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.cat
+@pytest.mark.parametrize(
+    "shape, dim",
+    [
+        (((0, 3), (2, 3)), 0),
+        (((0, 3), (0, 3)), 0),
+        (((0,), (0,)), 0),
+        (((0,), (1, 3)), -1),
+        (((0,), (1, 2, 3)), -2),
+        (((0,), (1, 1, 2, 3)), -3),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float32])
+def test_accuracy_cat_empty_tensor(shape, dim, dtype):
+    inp = [torch.randn(s, dtype=dtype, device=flag_gems.device) for s in shape]
+    ref_inp = [utils.to_reference(_) for _ in inp]
+    ref_out = torch.cat(ref_inp, dim)
+
+    with flag_gems.use_gems():
+        res_out = torch.cat(inp, dim)
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_conj_physical.py

@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.conj_physical
+@pytest.mark.parametrize("shape", [(256,), (32, 64), (2, 3, 4)])
+@pytest.mark.parametrize("is_complex", [True, False])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])
+def test_conj_physical(shape, is_complex, dtype):
+    device = flag_gems.device
+
+    if is_complex:
+        real = torch.randn(shape, dtype=torch.float32, device=device)
+        imag = torch.randn(shape, dtype=torch.float32, device=device)
+        input = torch.complex(real, imag)
+        out_dtype = input.dtype
+    else:
+        input = torch.randn(shape, dtype=dtype, device=device)
+        out_dtype = dtype
+
+    ref_input = utils.to_reference(input, True)
+    ref_out = torch.conj_physical(ref_input)
+    with flag_gems.use_gems():
+        res_out = torch.conj_physical(input)
+
+    utils.gems_assert_close(res_out, ref_out, out_dtype, reduce_dim=1)

tests/test_contiguous.py

@@ -0,0 +1,35 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.contiguous
+@pytest.mark.parametrize("shape", utils.SPECIAL_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.ALL_INT_DTYPES)
+def test_accuracy_contiguous(shape, dtype):
+    if shape[0] <= 2:
+        return
+
+    if dtype in utils.FLOAT_DTYPES:
+        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    else:
+        inp = torch.randint(
+            low=-10000, high=10000, size=shape, dtype=dtype, device="cpu"
+        ).to(flag_gems.device)
+
+    inp = inp[::2]
+    assert inp.is_contiguous() is False
+
+    ref_inp = utils.to_reference(inp)
+    ref_out = ref_inp.contiguous()
+    with flag_gems.use_gems():
+        res_out = inp.contiguous()
+
+    assert res_out.is_contiguous() is True
+    assert res_out.is_contiguous() is True
+    assert res_out.stride() == ref_out.stride()
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_diag.py

@@ -0,0 +1,36 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.diag
+@pytest.mark.parametrize("shape", utils.UT_SHAPES_1D + utils.UT_SHAPES_2D)
+@pytest.mark.parametrize("diagonal", [-2, -1, 0, 1, 2])
+@pytest.mark.parametrize(
+    "dtype", utils.FLOAT_DTYPES + utils.INT_DTYPES + utils.BOOL_TYPES
+)
+def test_diag(shape, diagonal, dtype):
+    if flag_gems.vendor_name == "kunlunxin":
+        torch.manual_seed(0)
+        torch.cuda.manual_seed_all(0)
+
+    if dtype in utils.FLOAT_DTYPES:
+        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    elif dtype in utils.BOOL_TYPES:
+        inp = torch.randint(0, 2, size=shape, dtype=dtype, device="cpu").to(
+            flag_gems.device
+        )
+    else:
+        inp = torch.randint(0, 0x7FFF, size=shape, dtype=dtype, device="cpu").to(
+            flag_gems.device
+        )
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.diag(ref_inp, diagonal)
+    with flag_gems.use_gems():
+        res_out = torch.diag(inp, diagonal)
+
+    utils.gems_assert_equal(res_out, ref_out)