Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/inference/modeling/models/nopadding_llama.py b/colossalai/inference/modeling/models/nopadding_llama.py
--- a/colossalai/inference/modeling/models/nopadding_llama.py
+++ b/colossalai/inference/modeling/models/nopadding_llama.py
@@ -101,12 +101,22 @@ def llama_model_forward(
         use_cuda_kernel = False
 
     hidden_states = self.embed_tokens(input_tokens_ids)
-    if use_cuda_kernel and inputmetadata != torch.float32 and use_flash_attn2:
-        cu_seqlens = F.pad(torch.cumsum(sequence_lengths, dim=0, dtype=torch.torch.int32), (1, 0))
+    if use_cuda_kernel:
+        if inputmetadata != torch.float32 and use_flash_attn2:
+            cu_seqlens = F.pad(torch.cumsum(sequence_lengths, dim=0, dtype=torch.torch.int32), (1, 0))
+
+        hidden_dim = self._cos_cached.size(-1)
+        total_length = hidden_states.size(0)
+        cos = torch.empty((total_length, hidden_dim), dtype=self._cos_cached.dtype, device=self._cos_cached.device)
+        sin = torch.empty((total_length, hidden_dim), dtype=self._sin_cached.dtype, device=self._sin_cached.device)
+        inference_ops.get_cos_and_sin(
+            self._cos_cached, self._sin_cached, cos, sin, sequence_lengths, kv_seq_len, inputmetadata.is_prompts
+        )
+        cos_sin = (cos, sin)
+
     else:
         cu_seqlens = None
-
-    cos_sin = get_xine_cache(sequence_lengths, self._cos_cached, self._sin_cached, inputmetadata.is_prompts)
+        cos_sin = get_xine_cache(sequence_lengths, self._cos_cached, self._sin_cached, inputmetadata.is_prompts)
 
     sm_scale = 1.0 / (inputmetadata.head_dim**0.5)
 
diff --git a/extensions/inference/inference_ops_cuda.py b/extensions/inference/inference_ops_cuda.py
--- a/extensions/inference/inference_ops_cuda.py
+++ b/extensions/inference/inference_ops_cuda.py
@@ -16,6 +16,7 @@ def sources_files(self):
                 "cuda/fused_rotary_emb_and_cache_kernel.cu",
                 "cuda/activation_kernel.cu",
                 "cuda/rms_layernorm_kernel.cu",
+                "cuda/get_cos_and_sin_kernel.cu",
             ]
         ]
         return ret


Test patch:
diff --git a/tests/test_infer/test_ops/cuda/test_get_cos_and_sin.py b/tests/test_infer/test_ops/cuda/test_get_cos_and_sin.py
new file mode 100644
--- /dev/null
+++ b/tests/test_infer/test_ops/cuda/test_get_cos_and_sin.py
@@ -0,0 +1,53 @@
+import numpy as np
+import pytest
+import torch
+
+from colossalai.kernel.kernel_loader import InferenceOpsLoader
+from tests.test_infer.test_ops.triton.test_xine_copy import get_cos_sin
+
+inference_ops = InferenceOpsLoader().load()
+
+
+def numpy_equal(x, y):
+    x_numpy = x.detach().cpu().numpy()
+    y_numpy = y.detach().cpu().numpy()
+
+    np.testing.assert_equal(x_numpy, y_numpy)
+
+
+@pytest.mark.parametrize("BATCH_SIZE", [4])
+@pytest.mark.parametrize("MAX_SEQ_LEN", [64])
+@pytest.mark.parametrize("HEAD_DIM", [64])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+def test_get_cos_and_sin(BATCH_SIZE, MAX_SEQ_LEN, HEAD_DIM, dtype):
+    MAX_TOTAL_TOKENS = BATCH_SIZE * MAX_SEQ_LEN
+    cos_cache = torch.randn((MAX_TOTAL_TOKENS, HEAD_DIM), dtype=dtype, device="cuda")
+    sin_cache = torch.randn((MAX_TOTAL_TOKENS, HEAD_DIM), dtype=dtype, device="cuda")
+    lengths = torch.randint(2, MAX_SEQ_LEN, (BATCH_SIZE,), device="cuda").to(torch.int32)
+
+    max_seq_len_in_batch = lengths.max()
+
+    # prefill
+    cos_ref, sin_ref = get_cos_sin(lengths, cos_cache, sin_cache, is_prompts=True, dtype=dtype)
+
+    cos = torch.zeros_like(cos_ref)
+    sin = torch.zeros_like(sin_ref)
+
+    inference_ops.get_cos_and_sin(cos_cache, sin_cache, cos, sin, lengths, max_seq_len_in_batch, True)
+
+    numpy_equal(cos, cos_ref)
+    numpy_equal(sin, sin_ref)
+
+    # decoding
+    ncos_ref, nsin_ref = get_cos_sin(lengths, cos_cache, sin_cache, is_prompts=False, dtype=dtype)
+
+    cos = torch.zeros_like(ncos_ref)
+    sin = torch.zeros_like(nsin_ref)
+
+    inference_ops.get_cos_and_sin(cos_cache, sin_cache, cos, sin, lengths, max_seq_len_in_batch, False)
+    numpy_equal(cos, ncos_ref)
+    numpy_equal(sin, nsin_ref)
+
+
+if __name__ == "__main__":
+    test_get_cos_and_sin(16, 4096, 256, torch.float16)