| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import multiprocessing |
| 4 | import types |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | |
| 10 | |
| 11 | def _test_oink_availability_impl( |
| 12 | device_capability: tuple[int, int], |
| 13 | has_rmsnorm: bool, |
| 14 | has_fused_add_rms_norm: bool, |
| 15 | expected_available: bool, |
| 16 | expected_fused: bool, |
| 17 | ) -> None: |
| 18 | """Test OINK support detection with mocked state.""" |
| 19 | import torch |
| 20 | |
| 21 | from vllm import platforms |
| 22 | |
| 23 | # Mock device capability (class method, override on class) |
| 24 | dc = platforms.interface.DeviceCapability(*device_capability) |
| 25 | platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc |
| 26 | |
| 27 | # Mock oink ops |
| 28 | oink_ops = types.SimpleNamespace() |
| 29 | if has_rmsnorm: |
| 30 | oink_ops.rmsnorm = lambda x, w, eps: x |
| 31 | if has_fused_add_rms_norm: |
| 32 | oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None |
| 33 | |
| 34 | torch.ops.oink = oink_ops |
| 35 | |
| 36 | # Now import vllm modules with mocks in place (fresh import with mocked platform) |
| 37 | import vllm.kernels.oink_ops # noqa: F401 |
| 38 | from vllm.ir.ops import fused_add_rms_norm, rms_norm |
| 39 | |
| 40 | # Verify support checks |
| 41 | assert rms_norm.impls["oink"].supported is expected_available |
| 42 | assert fused_add_rms_norm.impls["oink"].supported is expected_fused |
| 43 | |
| 44 | |
| 45 | @pytest.mark.parametrize( |
| 46 | "device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused", |
| 47 | [ |
| 48 | # Case 1: < SM100, ops not supported |
| 49 | ((9, 0), True, False, False, False), |
| 50 | # Case 2: CUDA available and SM100, rmsnorm op registered |
| 51 | ((10, 0), True, False, True, False), |
| 52 | # Case 3: SM100 with both rmsnorm and fused_add_rms_norm |
| 53 | ((10, 0), True, True, True, True), |
| 54 | ], |
| 55 | ) |
| 56 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") |
| 57 | def test_oink_availability_checks( |
| 58 | device_capability: tuple[int, int], |
| 59 | has_rmsnorm: bool, |
| 60 | has_fused_add_rms_norm: bool, |
| 61 | expected_available: bool, |
| 62 | expected_fused: bool, |
| 63 | ): |
| 64 | """Test OINK support detection with clean import state for each parameter set.""" |
| 65 | # Use spawn to run function in fresh process with clean imports |
| 66 | # TODO migrate to spawn utility: |
| 67 | # https://github.com/vllm-project/vllm/issues/41415 |
| 68 | ctx = multiprocessing.get_context("spawn") |
| 69 | process = ctx.Process( |
| 70 | target=_test_oink_availability_impl, |
| 71 | args=( |
| 72 | device_capability, |
| 73 | has_rmsnorm, |
| 74 | has_fused_add_rms_norm, |
| 75 | expected_available, |
| 76 | expected_fused, |
| 77 | ), |
| 78 | ) |
| 79 | process.start() |
| 80 | process.join() |
| 81 | |
| 82 | if process.exitcode != 0: |
| 83 | raise AssertionError( |
| 84 | f"Subprocess test failed with exit code {process.exitcode}" |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | def test_can_view_as_2d_stride_guard(): |
| 89 | # No global import |
| 90 | import torch |
| 91 | |
| 92 | # Import the helper from the kernels module. |
| 93 | from vllm.kernels.oink_ops import _can_view_as_2d |
| 94 | |
| 95 | x = torch.zeros((2, 3, 4)) |
| 96 | assert _can_view_as_2d(x) is True |
| 97 | |
| 98 | # Size-1 dims should be ignored by the viewability check. |
| 99 | # Create a tensor where stride(0) != stride(1) * size(1) due to padding, |
| 100 | # but view(-1, H) is still valid because dim 1 has size 1. |
| 101 | base = torch.zeros((2, 10, 4)) |
| 102 | x_singleton = base[:, :1, :] |
| 103 | x_singleton.view(-1, x_singleton.shape[-1]) |
| 104 | assert _can_view_as_2d(x_singleton) is True |
| 105 | |
| 106 | # Middle-dimension stride break: view(-1, hidden) should be invalid. |
| 107 | x2 = x[:, ::2, :] |
| 108 | with pytest.raises(RuntimeError): |
| 109 | x2.view(-1, x2.shape[-1]) |
| 110 | assert _can_view_as_2d(x2) is False |