| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for Helion utility functions.""" |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.kernels.helion.utils import canonicalize_gpu_name |
| 8 | |
| 9 | |
| 10 | @pytest.mark.parametrize( |
| 11 | "driver_reported_name,expected", |
| 12 | [ |
| 13 | ("NVIDIA H200", "nvidia_h200"), |
| 14 | ("NVIDIA A100-SXM4-80GB", "nvidia_a100"), |
| 15 | ("NVIDIA H100 80GB HBM3", "nvidia_h100"), |
| 16 | ("NVIDIA H100 PCIe", "nvidia_h100"), |
| 17 | ("NVIDIA H100 SXM5", "nvidia_h100"), |
| 18 | ("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"), |
| 19 | ("AMD Instinct MI300X", "amd_instinct_mi300x"), |
| 20 | ("AMD Instinct MI250X / MI250", "amd_instinct_mi250x_mi250"), |
| 21 | ("Tesla V100-SXM2-32GB", "tesla_v100"), |
| 22 | ], |
| 23 | ) |
| 24 | def test_canonicalize_gpu_name(driver_reported_name, expected): |
| 25 | """Test GPU name canonicalization.""" |
| 26 | assert canonicalize_gpu_name(driver_reported_name) == expected |
| 27 | |
| 28 | |
| 29 | @pytest.mark.parametrize("invalid_name", ["", " ", "\t", "\n"]) |
| 30 | def test_canonicalize_gpu_name_rejects_empty(invalid_name): |
| 31 | """Test that empty or whitespace-only names are rejected.""" |
| 32 | with pytest.raises(ValueError, match="cannot be empty"): |
| 33 | canonicalize_gpu_name(invalid_name) |