| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm platform and extension smoke tests. |
| 4 | |
| 5 | These tests catch ROCm-specific regressions early: |
| 6 | |
| 7 | 1. **Extension import smoke tests**: Verify that ROCm native extensions load |
| 8 | without crashing. Build/linking issues surface here before runtime. |
| 9 | |
| 10 | 2. **GCN arch parsing tests**: Verify _capability_from_gcn_arch correctly |
| 11 | parses all known GPU architectures. Wrong parsing causes silent feature |
| 12 | gate failures or crashes from selecting unsupported kernels. |
| 13 | |
| 14 | 3. **AITER ops availability tests**: Verify rocm_aiter_ops accessor methods |
| 15 | work without crashing, regardless of hardware support. |
| 16 | |
| 17 | 4. **Env var propagation tests**: Verify ROCm-specific env vars correctly |
| 18 | gate their respective features (e.g., shuffle KV cache, skinny GEMM). |
| 19 | |
| 20 | Related coverage: |
| 21 | - AITER kernel numerics: ``tests/kernels/core/test_rocm_aiter_ops.py`` |
| 22 | - Attention backend selection: ``tests/v1/attention/test_rocm_attention_backends_selection.py`` |
| 23 | - GEMM dispatch: ``tests/model_executor/layers/test_rocm_unquantized_gemm.py`` |
| 24 | """ # noqa: E501 |
| 25 | |
| 26 | import importlib |
| 27 | |
| 28 | import pytest |
| 29 | |
| 30 | from vllm.platforms import current_platform |
| 31 | |
| 32 | pytestmark = pytest.mark.skipif( |
| 33 | not current_platform.is_rocm(), reason="ROCm-specific tests" |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | # --------------------------------------------------------------------------- |
| 38 | # Extension import smoke tests |
| 39 | # --------------------------------------------------------------------------- |
| 40 | |
| 41 | |
| 42 | def test_rocm_extension_imports(): |
| 43 | """ROCm native extension loads without error. |
| 44 | |
| 45 | Catches build/linking issues before runtime. If this fails, the ROCm |
| 46 | build is broken and nothing else will work. |
| 47 | """ |
| 48 | import vllm._rocm_C # noqa: F401 |
| 49 | |
| 50 | |
| 51 | def test_aiter_ops_module_imports(): |
| 52 | """AITER ops module imports without error. |
| 53 | |
| 54 | The _aiter_ops module registers custom ops and initializes AITER state. |
| 55 | Import failures here indicate missing dependencies or incompatible |
| 56 | AITER versions. |
| 57 | """ |
| 58 | import vllm._aiter_ops # noqa: F401 |
| 59 | |
| 60 | |
| 61 | def test_rocm_platform_module_imports(): |
| 62 | """ROCm platform module imports without error. |
| 63 | |
| 64 | The platform module queries hardware via amdsmi and sets up arch flags. |
| 65 | Import failures indicate missing amdsmi or HIP runtime issues. |
| 66 | """ |
| 67 | import vllm.platforms.rocm # noqa: F401 |
| 68 | |
| 69 | |
| 70 | # --------------------------------------------------------------------------- |
| 71 | # GCN arch parsing tests |
| 72 | # --------------------------------------------------------------------------- |
| 73 | |
| 74 | |
| 75 | @pytest.mark.parametrize( |
| 76 | ("gcn_arch", "expected_major", "expected_minor"), |
| 77 | [ |
| 78 | pytest.param("gfx90a", 9, 0, id="MI200-gfx90a"), |
| 79 | pytest.param("gfx942", 9, 4, id="MI300X-gfx942"), |
| 80 | pytest.param("gfx950", 9, 5, id="MI350-gfx950"), |
| 81 | pytest.param("gfx1100", 11, 0, id="RX7900-gfx1100"), |
| 82 | pytest.param("gfx1101", 11, 0, id="RX7800-gfx1101"), |
| 83 | pytest.param("gfx1151", 11, 5, id="StrixHalo-gfx1151"), |
| 84 | pytest.param("gfx1200", 12, 0, id="RDNA4-gfx1200"), |
| 85 | pytest.param("gfx1201", 12, 0, id="RX9070-gfx1201"), |
| 86 | pytest.param("gfx1250", 12, 5, id="CDNA5-gfx1250"), |
| 87 | ], |
| 88 | ) |
| 89 | def test_capability_from_gcn_arch_known_gpus(gcn_arch, expected_major, expected_minor): |
| 90 | """_capability_from_gcn_arch parses known GPU architectures correctly. |
| 91 | |
| 92 | This is critical for feature gating - wrong parsing causes: |
| 93 | - FP8 kernels selected on hardware that doesn't support them |
| 94 | - MX format support incorrectly enabled/disabled |
| 95 | - Custom allreduce used on unsupported hardware |
| 96 | """ |
| 97 | from vllm.platforms.rocm import _capability_from_gcn_arch |
| 98 | |
| 99 | result = _capability_from_gcn_arch(gcn_arch) |
| 100 | assert result is not None, f"Failed to parse {gcn_arch}" |
| 101 | assert result == (expected_major, expected_minor), ( |
| 102 | f"Wrong capability for {gcn_arch}: got {result}, " |
| 103 | f"expected ({expected_major}, {expected_minor})" |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | @pytest.mark.parametrize( |
| 108 | "gcn_arch", |
| 109 | [ |
| 110 | pytest.param("not_a_gfx_string", id="non-gfx-string"), |
| 111 | pytest.param("cuda_arch", id="cuda-style"), |
| 112 | pytest.param("", id="empty-string"), |
| 113 | ], |
| 114 | ) |
| 115 | def test_capability_from_gcn_arch_returns_none_for_non_gfx(gcn_arch): |
| 116 | """_capability_from_gcn_arch returns None for non-GFX strings. |
| 117 | |
| 118 | Non-GFX strings should return None so the caller can fall back to |
| 119 | torch.cuda.get_device_capability. |
| 120 | """ |
| 121 | from vllm.platforms.rocm import _capability_from_gcn_arch |
| 122 | |
| 123 | result = _capability_from_gcn_arch(gcn_arch) |
| 124 | assert result is None |
| 125 | |
| 126 | |
| 127 | @pytest.mark.parametrize( |
| 128 | "gcn_arch", |
| 129 | [ |
| 130 | pytest.param("gfx8", id="too-few-digits"), |
| 131 | pytest.param("gfx12345", id="too-many-digits"), |
| 132 | pytest.param("gfx700", id="unsupported-major-7"), |
| 133 | pytest.param("gfx1500", id="future-major-15"), |
| 134 | ], |
| 135 | ) |
| 136 | def test_capability_from_gcn_arch_raises_for_malformed(gcn_arch): |
| 137 | """_capability_from_gcn_arch raises ValueError for malformed GFX strings. |
| 138 | |
| 139 | Malformed strings that look like GFX but can't be parsed should raise |
| 140 | so the user gets a clear error message asking them to file an issue. |
| 141 | """ |
| 142 | from vllm.platforms.rocm import _capability_from_gcn_arch |
| 143 | |
| 144 | with pytest.raises(ValueError, match="GCN arch"): |
| 145 | _capability_from_gcn_arch(gcn_arch) |
| 146 | |
| 147 | |
| 148 | # --------------------------------------------------------------------------- |
| 149 | # AITER ops availability tests |
| 150 | # --------------------------------------------------------------------------- |
| 151 | |
| 152 | |
| 153 | def test_rocm_aiter_ops_accessor_methods_dont_crash(): |
| 154 | """rocm_aiter_ops accessor methods work without crashing. |
| 155 | |
| 156 | These methods gate kernel selection. Even if AITER isn't available, |
| 157 | they should return False gracefully rather than crashing. |
| 158 | """ |
| 159 | from vllm._aiter_ops import rocm_aiter_ops |
| 160 | |
| 161 | # These should all return bool without crashing |
| 162 | assert isinstance(rocm_aiter_ops.is_enabled(), (bool, type(None))) |
| 163 | assert isinstance(rocm_aiter_ops.is_linear_enabled(), (bool, type(None))) |
| 164 | assert isinstance(rocm_aiter_ops.is_fused_moe_enabled(), (bool, type(None))) |
| 165 | assert isinstance(rocm_aiter_ops.is_mla_enabled(), (bool, type(None))) |
| 166 | assert isinstance(rocm_aiter_ops.is_mha_enabled(), (bool, type(None))) |
| 167 | |
| 168 | |
| 169 | def test_rocm_aiter_ops_is_aiter_found(): |
| 170 | """is_aiter_found returns consistent results.""" |
| 171 | from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found |
| 172 | |
| 173 | # The cached global should match the function result |
| 174 | assert is_aiter_found() == IS_AITER_FOUND |
| 175 | |
| 176 | |
| 177 | def test_rocm_aiter_ops_refresh_env_variables_doesnt_crash(): |
| 178 | """refresh_env_variables works without crashing. |
| 179 | |
| 180 | This is called after monkeypatching env vars in tests. It should |
| 181 | never crash, even with unusual env var combinations. |
| 182 | """ |
| 183 | from vllm._aiter_ops import rocm_aiter_ops |
| 184 | |
| 185 | # Should not raise |
| 186 | rocm_aiter_ops.refresh_env_variables() |
| 187 | |
| 188 | |
| 189 | # --------------------------------------------------------------------------- |
| 190 | # Env var propagation tests |
| 191 | # --------------------------------------------------------------------------- |
| 192 | |
| 193 | |
| 194 | @pytest.mark.parametrize("enabled", [True, False]) |
| 195 | def test_shuffle_kv_cache_env_propagates_to_rocm_aiter_ops(enabled, monkeypatch): |
| 196 | """VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT propagates to AITER state.""" |
| 197 | import vllm.envs as envs |
| 198 | from vllm._aiter_ops import rocm_aiter_ops |
| 199 | |
| 200 | monkeypatch.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "1" if enabled else "0") |
| 201 | importlib.reload(envs) |
| 202 | rocm_aiter_ops.refresh_env_variables() |
| 203 | |
| 204 | assert envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT is enabled |
| 205 | assert rocm_aiter_ops.is_shuffle_kv_cache_enabled() is enabled |
| 206 | |
| 207 | |
| 208 | @pytest.mark.parametrize("enabled", [True, False]) |
| 209 | def test_skinny_gemm_env_controls_rocm_fp8_scaled_mm_support(enabled, monkeypatch): |
| 210 | """VLLM_ROCM_USE_SKINNY_GEMM gates ROCm FP8 scaled-mm kernel eligibility.""" |
| 211 | import vllm.envs as envs |
| 212 | from vllm.model_executor.kernels.linear.scaled_mm.rocm import ( |
| 213 | ROCmFP8ScaledMMLinearKernel, |
| 214 | ) |
| 215 | |
| 216 | monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "1" if enabled else "0") |
| 217 | importlib.reload(envs) |
| 218 | monkeypatch.setattr("vllm.platforms.rocm.on_mi3xx", lambda: True) |
| 219 | |
| 220 | supported, reason = ROCmFP8ScaledMMLinearKernel.is_supported() |
| 221 | |
| 222 | assert supported is enabled |
| 223 | if enabled: |
| 224 | assert reason is None |
| 225 | else: |
| 226 | assert reason == "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." |