Super tiny add UT for copy_to_gpu_no_ce (#12270)
This commit is contained in:
@@ -369,7 +369,7 @@ def downcast_fp8(
|
||||
)
|
||||
|
||||
|
||||
def copy_to_gpu_no_ce(input: List[int], output: torch.Tensor):
|
||||
def copy_to_gpu_no_ce(input: torch.Tensor, output: torch.Tensor):
|
||||
torch.ops.sgl_kernel.copy_to_gpu_no_ce(input, output)
|
||||
|
||||
|
||||
|
||||
16
sgl-kernel/tests/test_copy.py
Normal file
16
sgl-kernel/tests/test_copy.py
Normal file
@@ -0,0 +1,16 @@
|
||||
import pytest
|
||||
import sgl_kernel
|
||||
import torch
|
||||
from sgl_kernel.elementwise import copy_to_gpu_no_ce
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", [64, 72])
|
||||
def test_copy_to_gpu_no_ce(size):
|
||||
tensor_cpu = torch.randint(0, 1000000, (size,), dtype=torch.int32, device="cpu")
|
||||
tensor_gpu = torch.empty_like(tensor_cpu, device="cuda")
|
||||
copy_to_gpu_no_ce(tensor_cpu, tensor_gpu)
|
||||
assert torch.all(tensor_cpu.cuda() == tensor_gpu)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
Reference in New Issue
Block a user