mirror of https://github.com/hpcaitech/ColossalAI
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import pytest
|
|
import torch
|
|
from packaging import version
|
|
from torch import nn
|
|
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from colossalai.kernel.triton.copy_kv_cache_dest import copy_kv_cache_to_dest
|
|
HAS_TRITON = True
|
|
except ImportError:
|
|
HAS_TRITON = False
|
|
print("please install triton from https://github.com/openai/triton")
|
|
|
|
TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse('11.4')
|
|
|
|
|
|
@pytest.mark.skipif(not TRITON_CUDA_SUPPORT or not HAS_TRITON,
|
|
reason="triton requires cuda version to be higher than 11.4")
|
|
def test_kv_cache_copy_op():
|
|
|
|
B_NTX = 32 * 2048
|
|
head_num = 8
|
|
head_dim = 64
|
|
|
|
cache = torch.randn((B_NTX, head_num, head_dim), device="cuda", dtype=torch.float16)
|
|
dest_index = torch.arange(0, B_NTX, device="cuda", dtype=torch.int32)
|
|
|
|
dest_data = torch.ones((B_NTX, head_num, head_dim), device="cuda", dtype=torch.float16)
|
|
|
|
copy_kv_cache_to_dest(cache, dest_index, dest_data)
|
|
|
|
assert torch.allclose(cache.cpu(), dest_data.cpu(), rtol=1e-3,
|
|
atol=1e-3), "copy_kv_cache_to_dest outputs from triton and torch are not matched"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_kv_cache_copy_op()
|