@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c0adee9231

[ROCm][CI] Add ROCm misc ops and env tests (#53279)

Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Divakar Verma committed Sep 1, 2026 at 15:03 UTC c0adee92317d436b88b621751122a72c7361229e
1 file changed +226
tests/kernels/core/test_rocm_misc_ops.py new
+226
@@ -0,0 +1,226 @@
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."