mirror of https://github.com/hpcaitech/ColossalAI
aibig-modeldata-parallelismdeep-learningdistributed-computingfoundation-modelsheterogeneous-traininghpcinferencelarge-scalemodel-parallelismpipeline-parallelism
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
39 lines
1.1 KiB
39 lines
1.1 KiB
import pytest |
|
import torch |
|
from packaging import version |
|
|
|
try: |
|
pass |
|
|
|
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()
|
|
|