[fix] merge conflicts

tests/test_infer/test_ops/cuda/test_kv_cache_memcpy.py

@@ -1,8 +1,10 @@
 import pytest
 import torch
+import torch.nn.functional as F
 
 from colossalai.kernel.kernel_loader import InferenceOpsLoader
 from colossalai.utils import get_current_device
+from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2
 from tests.test_infer.test_ops.triton.test_kvcache_copy import prepare_data
 
 inference_ops = InferenceOpsLoader().load()
@@ -10,12 +12,7 @@ inference_ops = InferenceOpsLoader().load()
 HEAD_DIM = 4
 
 
-@pytest.mark.parametrize("bsz", [4, 7, 32])
-@pytest.mark.parametrize("block_size", [16, 32, 64])
-@pytest.mark.parametrize("max_num_blocks_per_seq", [8, 32])
-@pytest.mark.parametrize("num_kv_heads", [16])
-@pytest.mark.parametrize("same_context_len", [True, False])
-def test_copy_kv_to_caches(
+def run_decode_copy_kv_to_caches(
     bsz: int,
     block_size: int,
     max_num_blocks_per_seq: int,
@@ -61,5 +58,65 @@ def test_copy_kv_to_caches(
     assert torch.equal(v_target, v_source)
 
 
+def run_context_copy_kv_to_cache(
+    bsz: int,
+    block_size: int,
+    max_num_blocks_per_seq: int,
+    num_kv_heads: int,
+    same_context_len: bool,
+):
+    torch.manual_seed(123)
+
+    assert isinstance(num_kv_heads, int) and num_kv_heads > 0, "Invalid number of kv heads."
+    max_seq_len = max_num_blocks_per_seq * block_size
+    dtype = torch.float16
+    device = get_current_device()
+
+    if same_context_len:
+        context_lengths = torch.tensor([max_seq_len for _ in range(bsz)], dtype=torch.int32, device=device)
+    else:
+        context_lengths = torch.randint(low=1, high=max_seq_len, size=(bsz,), dtype=torch.int32, device=device)
+
+    num_tokens = torch.sum(context_lengths).item()
+
+    max_seq_len_in_batch = context_lengths.max()
+    cu_seqlens = F.pad(torch.cumsum(context_lengths, dim=0, dtype=torch.torch.int32), (1, 0))
+
+    kv_size = (num_tokens, num_kv_heads, HEAD_DIM)
+    key = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+    value = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+
+    k_cache_ref, v_cache_ref, block_tables = generate_caches_and_block_tables_v2(
+        key, value, context_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
+    )
+
+    block_tables = block_tables.to(device=device)
+    k_cache = torch.zeros_like(k_cache_ref)
+    v_cache = torch.zeros_like(v_cache_ref)
+
+    inference_ops.context_kv_cache_memcpy(
+        key, value, k_cache, v_cache, context_lengths, cu_seqlens, block_tables, max_seq_len_in_batch
+    )
+
+    assert torch.equal(k_cache, k_cache_ref)
+    assert torch.equal(v_cache, v_cache_ref)
+
+
+@pytest.mark.parametrize("bsz", [4, 7, 32])
+@pytest.mark.parametrize("block_size", [16, 32, 64])
+@pytest.mark.parametrize("max_num_blocks_per_seq", [8, 32])
+@pytest.mark.parametrize("num_kv_heads", [16])
+@pytest.mark.parametrize("same_context_len", [True, False])
+def test_kv_cache_memcopy(
+    bsz: int,
+    block_size: int,
+    max_num_blocks_per_seq: int,
+    num_kv_heads: int,
+    same_context_len: bool,
+):
+    run_context_copy_kv_to_cache(bsz, block_size, max_num_blocks_per_seq, num_kv_heads, same_context_len)
+    run_decode_copy_kv_to_caches(bsz, block_size, max_num_blocks_per_seq, num_kv_heads, same_context_len)
+
+
 if __name__ == "__main__":
-    test_copy_kv_to_caches(4, 32, 8, 16, True)
+    test_kv_cache_memcopy(4, 32, 8, 16, True)

tests/test_infer/test_ops/cuda/test_rotary_embdding_unpad.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pytest
 import torch
 from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
@@ -10,11 +11,18 @@ from tests.test_infer.test_ops.triton.kernel_utils import mock_alloc_block_table
 from tests.test_infer.test_ops.triton.test_rotary_embdding_unpad import torch_rotary_emb
 
 
+def numpy_allclose(x, y, rtol, atol):
+    x_numpy = x.detach().cpu().numpy()
+    y_numpy = y.detach().cpu().numpy()
+
+    np.testing.assert_allclose(x_numpy, y_numpy, rtol=rtol, atol=atol)
+
+
 @pytest.mark.parametrize("BATCH_SIZE", [4])
 @pytest.mark.parametrize("SEQ_LEN", [64])
 @pytest.mark.parametrize("H", [32])
 @pytest.mark.parametrize("D", [64])
-@pytest.mark.parametrize("dtype", [torch.float16])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
 def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
     torch.manual_seed(10)
     TOTAL_TOKENS = BATCH_SIZE * SEQ_LEN
@@ -54,17 +62,36 @@ def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
 
     kv_seq_lengths = past_kv_seq_lengths + 1
     block_tables = block_tables.to(device="cuda")
-    q_ref = torch_rotary_emb(new_q, cos[:BATCH_SIZE], sin[:BATCH_SIZE])
-    k_ref = torch_rotary_emb(new_k, cos[:BATCH_SIZE], sin[:BATCH_SIZE])
 
     new_q_copy = new_q.clone()
     new_k_copy = new_k.clone()
 
+    if dtype == torch.float16:
+        rtol = 1e-3
+        atol = 1e-3
+
+        new_q_fp16 = new_q.clone()
+        new_k_fp16 = new_k.clone()
+
+        high_precision_cos = cos[:BATCH_SIZE].to(torch.float32)
+        high_precision_sin = sin[:BATCH_SIZE].to(torch.float32)
+        high_precision_q = new_q.to(torch.float32)
+        high_precision_k = new_k.to(torch.float32)
+        q_ref = torch_rotary_emb(high_precision_q, high_precision_cos, high_precision_sin).to(torch.float16)
+        k_ref = torch_rotary_emb(high_precision_k, high_precision_cos, high_precision_sin).to(torch.float16)
+
+    else:
+        rtol = 1e-5
+        atol = 1e-7
+
+        q_ref = torch_rotary_emb(new_q, cos[:BATCH_SIZE], sin[:BATCH_SIZE])
+        k_ref = torch_rotary_emb(new_k, cos[:BATCH_SIZE], sin[:BATCH_SIZE])
+
     inference_ops.rotary_embedding_and_cache_copy(
-        new_q, new_k, new_v, cos, sin, k_cache, v_cache, kv_seq_lengths, block_tables
+        new_q, new_k, new_v, cos, sin, k_cache, v_cache, kv_seq_lengths, block_tables, True
     )
 
-    inference_ops.rotary_embedding(new_q_copy, new_k_copy, cos, sin)
+    inference_ops.rotary_embedding(new_q_copy, new_k_copy, cos, sin, True)
 
     past_kv_seq_len = kv_seq_lengths - 1
     target_block_ids = block_tables[range(0, block_tables.size(0)), past_kv_seq_len // block_size]
@@ -74,18 +101,26 @@ def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
     v_target = v_cache[target_block_ids, :, offsets_in_block, :].squeeze()
     v_source = new_v.squeeze()
 
-    assert torch.allclose(new_q, q_ref, atol=1e-6, rtol=1e-6)
-    assert torch.allclose(k_target, k_ref, atol=1e-6, rtol=1e-6)
+    numpy_allclose(new_q, q_ref, rtol=rtol, atol=atol)
+    numpy_allclose(k_target, k_ref, rtol=rtol, atol=atol)
 
-    assert torch.allclose(new_q_copy, q_ref, atol=1e-6, rtol=1e-6)
-    assert torch.allclose(new_k_copy, k_ref, atol=1e-6, rtol=1e-6)
+    numpy_allclose(new_q_copy, q_ref, rtol=rtol, atol=atol)
+    numpy_allclose(new_k_copy, k_ref, rtol=rtol, atol=atol)
 
     assert k_target.shape == k_source.shape
-    assert torch.allclose(k_target, k_source, atol=1e-6, rtol=1e-6)
+    numpy_allclose(k_target, k_source, rtol=rtol, atol=atol)
 
     assert v_target.shape == v_source.shape
     assert torch.equal(v_target, v_source)
 
+    if dtype == torch.float16:
+        # After testing cuda fp16 high_precision, it was found to have higher precision than torch fp16. Therefore, the threshold here has been relaxed to pass the test.
+        rtol = 1e-3
+        atol = 1e-1
+        inference_ops.rotary_embedding(new_q_fp16, new_k_fp16, cos, sin, False)
+        numpy_allclose(new_q_copy, new_q_fp16, rtol=rtol, atol=atol)
+        numpy_allclose(new_k_copy, new_k_fp16, rtol=rtol, atol=atol)
+
 
 if __name__ == "__main__":
-    test_rotary_emb(16, 512, 4, 128, torch.float16)
+    test_rotary_emb(16, 64, 4, 128, torch.float16)