main
py 24 lines 872 Bytes
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for miscellaneous utilities."""
4
5 import torch
6
7 from tests.kernels.utils import opcheck
8
9
10 def test_convert_fp8_opcheck():
11 data = torch.randn((256, 256), dtype=torch.float32, device="cuda")
12 result = torch.empty_like(data, dtype=torch.float8_e4m3fn)
13 opcheck(torch.ops._C_cache_ops.convert_fp8, (result, data, 1.0, "fp8"))
14
15
16 # TODO: Add this back, currently fails with
17 # csrc/cuda_utils_kernels.cu:15 'invalid argument'
18 # @pytest.mark.skipif(not current_platform.is_cuda(),
19 # reason="Only supported for CUDA")
20 # def test_cuda_utils_opcheck():
21 # opcheck(torch.ops._C_cuda_utils.get_device_attribute, (0, 0))
22 # opcheck(
23 # torch.ops._C_cuda_utils.
24 # get_max_shared_memory_per_block_device_attribute, (0, ))