diff --git a/tests/kernels/test_uva.py b/tests/kernels/test_uva.py new file mode 100644 index 0000000000000..1f8a174a0f066 --- /dev/null +++ b/tests/kernels/test_uva.py @@ -0,0 +1,61 @@ +import pytest +import torch + +from vllm.utils import get_cuda_view_from_cpu_tensor, is_uva_available + +CUDA_DEVICES = [ + f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2) +] + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("device", CUDA_DEVICES) +def test_cpu_write(device): + torch.set_default_device(device) + cpu_tensor = torch.zeros(10, + 10, + device="cpu", + pin_memory=True, + dtype=torch.int32) + cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor) + assert cuda_view.device.type == "cuda" + + assert cuda_view[0, 0] == 0 + assert cuda_view[2, 3] == 0 + assert cuda_view[4, 5] == 0 + + cpu_tensor[0, 0] = 1 + cpu_tensor[2, 3] = 2 + cpu_tensor[4, 5] = -1 + + cuda_view.mul_(2) + assert cuda_view[0, 0] == 2 + assert cuda_view[2, 3] == 4 + assert cuda_view[4, 5] == -2 + + +@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") +@pytest.mark.parametrize("device", CUDA_DEVICES) +def test_gpu_write(device): + torch.set_default_device(device) + cpu_tensor = torch.zeros(10, + 10, + device="cpu", + pin_memory=True, + dtype=torch.int32) + cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor) + assert cuda_view.device.type == "cuda" + + assert cuda_view[0, 0] == 0 + assert cuda_view[2, 3] == 0 + assert cuda_view[4, 5] == 0 + + cuda_view[0, 0] = 1 + cuda_view[2, 3] = 2 + cuda_view[4, 5] = -1 + cuda_view.mul_(2) + + torch.cuda.synchronize() + assert cpu_tensor[0, 0] == 2 + assert cpu_tensor[2, 3] == 4 + assert cpu_tensor[4, 5] == -2