@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 509 lines 19.1 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 """Compile-guard tests for the ROCm RDNA3 W4A16 kernels (dense + MoE).
5
6 Verifies that the gfx1100 compilation and dispatch guards are hermetic:
7 - On gfx1100: all ops exist, dispatch selects RDNA3 kernels.
8 - On CDNA (gfx942/gfx950) or other non-gfx1100: ops must NOT exist,
9 dispatch must fall through to Triton/Marlin, and no RDNA3 code
10 path is reachable.
11
12 The negative (non-gfx1100) tests verify at three layers:
13 1. Compile-level: on non-gfx1100 hardware, the RDNA3 ops are absent
14 from the compiled _rocm_C extension — real binary verification.
15 2. Static source analysis: parses CMakeLists.txt and torch_bindings.cpp
16 to verify that all RDNA3 .cu files and op registrations are inside
17 gfx1100-only guards.
18 3. Runtime mock: patches on_gfx1100() to False and verifies that the
19 Python dispatch chain rejects the RDNA3 path.
20
21 Run `pytest tests/kernels/quantization/test_rdna3_compile_guards.py`.
22 """
23
24 from pathlib import Path
25 from unittest.mock import patch
26
27 import pytest
28 import regex as re
29 import torch
30
31 import vllm
32 from vllm.platforms import current_platform
33
34 if not current_platform.is_rocm():
35 pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
36
37 from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
38 kInt4Static,
39 kInt4Static32,
40 kInt4Static32Asym,
41 kInt4StaticAsym,
42 kInt8Static,
43 )
44 from vllm.platforms.rocm import on_gfx1100 # noqa: E402
45
46 gfx1100_only = pytest.mark.skipif(
47 not on_gfx1100(),
48 reason="Requires gfx1100 hardware",
49 )
50
51 not_gfx1100 = pytest.mark.skipif(
52 on_gfx1100(),
53 reason="This test verifies non-gfx1100 builds — skip on gfx1100",
54 )
55
56 RDNA3_OPS = ["gptq_gemm_rdna3", "gptq_gemm_rdna3_wmma", "moe_gptq_gemm_rdna3"]
57 RDNA3_CU_FILES = [
58 "q_gemm_rdna3.cu",
59 "q_gemm_rdna3_wmma.cu",
60 "moe_q_gemm_rdna3.cu",
61 ]
62
63
64 def _find_repo_root() -> Path | None:
65 """Walk up from this file to find the repo root (has CMakeLists.txt)."""
66 for parent in [Path(__file__).resolve(), *Path(__file__).resolve().parents]:
67 if (parent / "CMakeLists.txt").exists() and (parent / "csrc").is_dir():
68 return parent
69 return None
70
71
72 REPO_ROOT = _find_repo_root()
73
74 # Directory of the *installed* vllm python package. The .py guard checks read
75 # from here so they verify the code that is actually imported at runtime — this
76 # works even on CI images that ship the wheel instead of the python source tree
77 # (where only csrc/ + CMakeLists.txt are checked out for building).
78 VLLM_PKG_DIR: Path | None = (
79 Path(vllm.__file__).parent if getattr(vllm, "__file__", None) else None
80 )
81
82 needs_source = pytest.mark.skipif(
83 REPO_ROOT is None,
84 reason="C/CMake source tree not available (installed package only)",
85 )
86
87
88 def _read_source_or_skip(*relparts: str) -> str:
89 """Read a C/CMake source file from the repo tree, or skip if absent.
90
91 Used for csrc/ and CMakeLists.txt — these only exist in a source checkout,
92 not in the installed wheel.
93 """
94 assert REPO_ROOT is not None # callers are gated by @needs_source
95 path = REPO_ROOT.joinpath(*relparts)
96 if not path.exists():
97 pytest.skip(f"{path} not present in this source tree")
98 return path.read_text()
99
100
101 def _read_pkg_source_or_skip(*relparts: str) -> str:
102 """Read a python source file from the installed vllm package.
103
104 Reflects the code actually loaded at runtime, so these guard checks run in
105 CI against the wheel — no source checkout required. Only skips for an
106 exotic install layout (namespace/zipimport) where __file__ is unavailable.
107 """
108 if VLLM_PKG_DIR is None:
109 pytest.skip("vllm package directory not resolvable (zip/namespace?)")
110 assert VLLM_PKG_DIR is not None # narrow for mypy (skip above is NoReturn)
111 path = VLLM_PKG_DIR.joinpath(*relparts)
112 if not path.exists():
113 pytest.skip(f"{path} not present in installed vllm package")
114 return path.read_text()
115
116
117 # ============================================================================
118 # Part A: POSITIVE — on gfx1100, ops exist and dispatch works
119 # ============================================================================
120
121
122 @gfx1100_only
123 @pytest.mark.parametrize("op_name", RDNA3_OPS)
124 def test_op_registered_on_gfx1100(op_name):
125 """On gfx1100, all RDNA3 ops must be registered in _rocm_C."""
126 assert hasattr(torch.ops, "_rocm_C"), "_rocm_C module not loaded"
127 assert hasattr(torch.ops._rocm_C, op_name), (
128 f"_rocm_C.{op_name} not registered — "
129 "check CMakeLists.txt VLLM_ROCM_HAS_GFX1100 "
130 "and torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100"
131 )
132
133
134 @gfx1100_only
135 def test_all_ops_present_or_all_absent():
136 """The 3 RDNA3 ops are behind the same #ifdef — all present or all absent.
137
138 Catches someone accidentally moving an op outside the guard.
139 """
140 has_rocm_c = hasattr(torch.ops, "_rocm_C")
141 if not has_rocm_c:
142 pytest.skip("_rocm_C not loaded")
143
144 present = {op: hasattr(torch.ops._rocm_C, op) for op in RDNA3_OPS}
145 values = set(present.values())
146 assert len(values) == 1, (
147 f"Guard inconsistency — some RDNA3 ops registered, others not: "
148 f"{present}. Check torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100 block."
149 )
150
151
152 # ============================================================================
153 # Part B: NEGATIVE — compile-level verification on non-gfx1100
154 # ============================================================================
155
156
157 @not_gfx1100
158 @pytest.mark.parametrize("op_name", RDNA3_OPS)
159 def test_op_absent_on_non_gfx1100(op_name):
160 """On non-gfx1100 (CDNA), RDNA3 ops must NOT exist in _rocm_C.
161
162 This is the real compile-level check: the binary was built without
163 gfx1100 support, so the ops should not have been compiled or registered.
164 """
165 if not hasattr(torch.ops, "_rocm_C"):
166 return
167 assert not hasattr(torch.ops._rocm_C, op_name), (
168 f"_rocm_C.{op_name} is registered on non-gfx1100 hardware — "
169 "compile guard is broken: check CMakeLists.txt "
170 "VLLM_ROCM_HAS_GFX1100 and torch_bindings.cpp #ifdef"
171 )
172
173
174 @not_gfx1100
175 def test_rocm_moe_not_supported_on_non_gfx1100():
176 """The RDNA3 MoE experts must not be selectable on non-gfx1100 hardware."""
177 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
178 Rdna3WNA16Experts,
179 )
180
181 assert Rdna3WNA16Experts._supports_current_device() is False, (
182 "Rdna3WNA16Experts reported support on non-gfx1100 — dispatch guard is broken"
183 )
184
185
186 @not_gfx1100
187 def test_dense_kernel_rejects_on_non_gfx1100():
188 """RDNA3W4A16LinearKernel.can_implement must reject on non-gfx1100."""
189 from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
190 MPLinearLayerConfig,
191 )
192 from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
193 RDNA3W4A16LinearKernel,
194 )
195 from vllm.scalar_type import scalar_types
196
197 config = MPLinearLayerConfig(
198 full_weight_shape=(1024, 256),
199 partition_weight_shape=(1024, 256),
200 weight_type=scalar_types.uint4b8,
201 act_type=torch.float16,
202 group_size=128,
203 zero_points=False,
204 )
205 ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
206 assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}"
207
208
209 # ============================================================================
210 # Part C: Static source analysis (build-level guards)
211 # ============================================================================
212
213
214 @needs_source
215 class TestCMakeGuards:
216 """Verify CMakeLists.txt only compiles RDNA3 .cu files for gfx1100."""
217
218 @staticmethod
219 def _read_cmake():
220 return _read_source_or_skip("CMakeLists.txt")
221
222 def test_rdna3_cu_files_inside_gfx1100_conditional(self):
223 """All RDNA3 .cu files must be listed inside the
224 ``if(VLLM_GPU_ARCHES MATCHES "gfx1100")`` block, not unconditionally.
225 """
226 cmake = self._read_cmake()
227 for cu_file in RDNA3_CU_FILES:
228 assert cu_file in cmake, f"{cu_file} not found in CMakeLists.txt"
229
230 lines = cmake.splitlines()
231 in_gfx1100_block = False
232 for line in lines:
233 if 'VLLM_GPU_ARCHES MATCHES "gfx1100"' in line:
234 in_gfx1100_block = True
235 if in_gfx1100_block and "endif()" in line:
236 in_gfx1100_block = False
237 if cu_file in line:
238 assert in_gfx1100_block, (
239 f"{cu_file} is listed OUTSIDE the gfx1100 "
240 f"conditional in CMakeLists.txt — CDNA builds "
241 f"would compile RDNA3 code. Line: {line.strip()}"
242 )
243
244 def test_compile_definition_only_for_gfx1100(self):
245 """VLLM_ROCM_GFX1100 compile definition must be conditional."""
246 cmake = self._read_cmake()
247 lines = cmake.splitlines()
248 in_gfx1100_block = False
249 for line in lines:
250 if "VLLM_ROCM_HAS_GFX1100)" in line:
251 in_gfx1100_block = True
252 if in_gfx1100_block and "endif()" in line:
253 in_gfx1100_block = False
254 if "VLLM_ROCM_GFX1100" in line and "target_compile_definitions" in line:
255 assert in_gfx1100_block, (
256 "VLLM_ROCM_GFX1100 compile definition is set outside "
257 "the VLLM_ROCM_HAS_GFX1100 conditional — CDNA builds "
258 f"would define it. Line: {line.strip()}"
259 )
260
261
262 @needs_source
263 class TestTorchBindingsGuards:
264 """Verify torch_bindings.cpp gates all RDNA3 ops behind #ifdef."""
265
266 @staticmethod
267 def _read_bindings():
268 return _read_source_or_skip("csrc", "rocm", "torch_bindings.cpp")
269
270 def test_all_rdna3_ops_inside_ifdef(self):
271 """Every rdna3 op def/impl must be between #ifdef VLLM_ROCM_GFX1100
272 and #endif. If any is outside, a CDNA build would try to register
273 the op and link a symbol that doesn't exist.
274 """
275 src = self._read_bindings()
276 lines = src.splitlines()
277
278 inside_guard = False
279 rdna3_lines_outside = []
280
281 for i, line in enumerate(lines, 1):
282 if "#ifdef VLLM_ROCM_GFX1100" in line:
283 inside_guard = True
284 elif line.strip() == "#endif" and inside_guard:
285 inside_guard = False
286
287 if (
288 "rdna3" in line.lower()
289 and not line.strip().startswith("//")
290 and not inside_guard
291 ):
292 rdna3_lines_outside.append((i, line.strip()))
293
294 assert not rdna3_lines_outside, (
295 "RDNA3 op references found OUTSIDE #ifdef VLLM_ROCM_GFX1100 "
296 "in torch_bindings.cpp — these would break CDNA builds:\n"
297 + "\n".join(f" L{n}: {s}" for n, s in rdna3_lines_outside)
298 )
299
300 def test_no_unconditional_rdna3_includes(self):
301 """No #include of RDNA3-specific headers outside the guard."""
302 src = self._read_bindings()
303 lines = src.splitlines()
304
305 inside_guard = False
306 for i, line in enumerate(lines, 1):
307 if "#ifdef VLLM_ROCM_GFX1100" in line:
308 inside_guard = True
309 elif line.strip() == "#endif" and inside_guard:
310 inside_guard = False
311
312 if "#include" in line and "rdna3" in line.lower():
313 assert inside_guard, (
314 f"L{i}: RDNA3 include outside gfx1100 guard: {line.strip()}"
315 )
316
317
318 class TestCustomOpsGuards:
319 """Verify _custom_ops.py gates register_fake behind hasattr checks."""
320
321 @staticmethod
322 def _read_custom_ops():
323 return _read_pkg_source_or_skip("_custom_ops.py")
324
325 def test_register_fake_guarded_by_hasattr(self):
326 """Every register_fake for an RDNA3 op must be preceded by a hasattr
327 check — otherwise it would crash on import on CDNA where the ops
328 don't exist.
329 """
330 src = self._read_custom_ops()
331 for op in RDNA3_OPS:
332 pattern = rf'register_fake\(\s*"_rocm_C::{op}"\s*\)'
333 match = re.search(pattern, src)
334 if match is None:
335 continue
336 preceding = src[: match.start()]
337 last_hasattr = preceding.rfind(f'hasattr(torch.ops._rocm_C, "{op}")')
338 assert last_hasattr != -1, (
339 f'register_fake("_rocm_C::{op}") is not preceded by a '
340 f"hasattr check — would crash on CDNA import"
341 )
342 gap = preceding[last_hasattr:].count("\n")
343 assert gap <= 5, (
344 f"hasattr guard for {op} is {gap} lines before "
345 f"register_fake — suspiciously far; verify it's the "
346 f"actual guard and not a coincidence"
347 )
348
349 def test_no_toplevel_rocm_c_import(self):
350 """No top-level ``from vllm._rocm_C import`` — would crash on CDNA."""
351 src = self._read_custom_ops()
352 for line in src.splitlines():
353 stripped = line.strip()
354 if stripped.startswith("#") or stripped.startswith("//"):
355 continue
356 assert "from vllm._rocm_C import" not in stripped, (
357 f"Top-level import of _rocm_C in _custom_ops.py would "
358 f"crash on CDNA: {stripped}"
359 )
360
361
362 # ============================================================================
363 # Part D: Runtime mock (simulate CDNA on gfx1100 hardware)
364 # ============================================================================
365
366
367 class TestMoEDispatchMocked:
368 """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
369
370 def test_kernel_unavailable_when_mocked_cdna(self):
371 """The device gate must reject when not on gfx1100."""
372 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
373 rdna3_moe_kernel_available,
374 )
375
376 with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
377 assert rdna3_moe_kernel_available() is False
378
379 @pytest.mark.parametrize(
380 "weight_key",
381 [kInt8Static, kInt4StaticAsym, kInt4Static32Asym, None],
382 )
383 def test_quant_scheme_rejects_non_symmetric_int4(self, weight_key):
384 """Only symmetric int4 weight-only schemes reach the kernel."""
385 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
386 Rdna3WNA16Experts,
387 )
388
389 assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is False
390
391 @pytest.mark.parametrize("weight_key", [kInt4Static, kInt4Static32])
392 def test_quant_scheme_accepts_symmetric_int4(self, weight_key):
393 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
394 Rdna3WNA16Experts,
395 )
396
397 assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is True
398
399 def test_kernel_unavailable_when_op_missing(self):
400 """The device gate returns False when the C++ op doesn't exist."""
401 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
402 rdna3_moe_kernel_available,
403 )
404
405 fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
406 with patch.object(torch, "ops", create=True) as mock_ops:
407 mock_ops._rocm_C = fake_rocm_c
408 assert rdna3_moe_kernel_available() is False
409
410 def test_kernel_unavailable_when_rocm_c_absent(self):
411 """The device gate returns False when _rocm_C doesn't exist at all."""
412 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
413 rdna3_moe_kernel_available,
414 )
415
416 fake_ops = type("FakeOps", (), {})()
417 with patch.object(torch, "ops", fake_ops):
418 assert rdna3_moe_kernel_available() is False
419
420
421 class TestDenseKernelSelectionMocked:
422 """Mock on_gfx1100() and verify dense RDNA3 kernel is not selected."""
423
424 @gfx1100_only
425 def test_can_implement_rejects_when_mocked_cdna(self):
426 """RDNA3W4A16LinearKernel.can_implement must reject on mocked CDNA."""
427 from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
428 MPLinearLayerConfig,
429 )
430 from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
431 RDNA3W4A16LinearKernel,
432 )
433 from vllm.scalar_type import scalar_types
434
435 config = MPLinearLayerConfig(
436 full_weight_shape=(1024, 256),
437 partition_weight_shape=(1024, 256),
438 weight_type=scalar_types.uint4b8,
439 act_type=torch.float16,
440 group_size=128,
441 zero_points=False,
442 )
443 ok, _ = RDNA3W4A16LinearKernel.can_implement(config)
444 assert ok is True
445
446 with (
447 patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
448 patch("vllm.platforms.rocm._ON_GFX1100", False),
449 ):
450 ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
451 assert ok is False, f"RDNA3 kernel accepted on simulated CDNA: {reason}"
452
453 @gfx1100_only
454 def test_chooser_skips_rdna3_when_mocked_cdna(self):
455 """choose_mp_linear_kernel must NOT return RDNA3 on mocked CDNA."""
456 from vllm.model_executor.kernels.linear import (
457 choose_mp_linear_kernel,
458 )
459 from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
460 MPLinearLayerConfig,
461 )
462 from vllm.scalar_type import scalar_types
463
464 config = MPLinearLayerConfig(
465 full_weight_shape=(1024, 256),
466 partition_weight_shape=(1024, 256),
467 weight_type=scalar_types.uint4b8,
468 act_type=torch.float16,
469 group_size=128,
470 zero_points=False,
471 )
472 with (
473 patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
474 patch("vllm.platforms.rocm._ON_GFX1100", False),
475 ):
476 chosen = choose_mp_linear_kernel(config)
477 assert chosen.__name__ != "RDNA3W4A16LinearKernel", (
478 "RDNA3 kernel was selected on simulated CDNA — "
479 "choose_mp_linear_kernel guard is broken"
480 )
481
482
483 class TestWNA16OracleWiring:
484 """The RDNA3 backend must reach the kernel only through the oracle."""
485
486 def test_backend_maps_to_rdna3_experts(self):
487 from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
488 Rdna3WNA16Experts,
489 )
490 from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
491 WNA16MoEBackend,
492 backend_to_kernel_cls,
493 map_wna16_backend,
494 )
495
496 assert backend_to_kernel_cls(WNA16MoEBackend.RDNA3) == [Rdna3WNA16Experts]
497 assert map_wna16_backend("rdna3") == WNA16MoEBackend.RDNA3
498
499 def test_backend_is_offered_before_the_triton_fallback(self):
500 """Priority order: the native kernel outranks Triton when supported."""
501 from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
502 WNA16MoEBackend,
503 _get_priority_backends,
504 )
505
506 backends = _get_priority_backends()
507 assert backends.index(WNA16MoEBackend.RDNA3) < backends.index(
508 WNA16MoEBackend.TRITON
509 )