test_utils.py 797 Bytes
Newer Older
1
2
3
4
5
6
7
8
"""
Tests for miscellaneous utilities
"""

import pytest
import torch

from vllm.platforms import current_platform
9
from .utils import torch_version
10

11
12
if torch_version.startswith("2.4"):
    from tests.kernels.utils import opcheck
13

14
15
16
17
18

# def test_convert_fp8_opcheck():
#     data = torch.randn((256, 256), dtype=torch.float32, device="cuda")
#     result = torch.empty_like(data, dtype=torch.float8_e4m3fn)
#     opcheck(torch.ops._C_cache_ops.convert_fp8, (result, data, 1.0, "fp8"))
19
20
21
22
23
24
25
26
27


@pytest.mark.skipif(not current_platform.is_cuda(),
                    reason="Only supported for CUDA")
def test_cuda_utils_opcheck():
    opcheck(torch.ops._C_cuda_utils.get_device_attribute, (0, 0))
    opcheck(
        torch.ops._C_cuda_utils.
        get_max_shared_memory_per_block_device_attribute, (0, ))