wehub-resource-sync
112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import multiprocessing
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
|
|
def _test_oink_availability_impl(
|
|
device_capability: tuple[int, int],
|
|
has_rmsnorm: bool,
|
|
has_fused_add_rms_norm: bool,
|
|
expected_available: bool,
|
|
expected_fused: bool,
|
|
) -> None:
|
|
"""Test OINK support detection with mocked state."""
|
|
import torch
|
|
|
|
from vllm import platforms
|
|
|
|
# Mock device capability (class method, override on class)
|
|
dc = platforms.interface.DeviceCapability(*device_capability)
|
|
platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc
|
|
|
|
# Mock oink ops
|
|
oink_ops = types.SimpleNamespace()
|
|
if has_rmsnorm:
|
|
oink_ops.rmsnorm = lambda x, w, eps: x
|
|
if has_fused_add_rms_norm:
|
|
oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None
|
|
|
|
torch.ops.oink = oink_ops
|
|
|
|
# Now import vllm modules with mocks in place (fresh import with mocked platform)
|
|
import vllm.kernels.oink_ops # noqa: F401
|
|
from vllm.ir.ops import fused_add_rms_norm, rms_norm
|
|
|
|
# Verify support checks
|
|
assert rms_norm.impls["oink"].supported is expected_available
|
|
assert fused_add_rms_norm.impls["oink"].supported is expected_fused
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused",
|
|
[
|
|
# Case 1: < SM100, ops not supported
|
|
((9, 0), True, False, False, False),
|
|
# Case 2: CUDA available and SM100, rmsnorm op registered
|
|
((10, 0), True, False, True, False),
|
|
# Case 3: SM100 with both rmsnorm and fused_add_rms_norm
|
|
((10, 0), True, True, True, True),
|
|
],
|
|
)
|
|
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
|
|
def test_oink_availability_checks(
|
|
device_capability: tuple[int, int],
|
|
has_rmsnorm: bool,
|
|
has_fused_add_rms_norm: bool,
|
|
expected_available: bool,
|
|
expected_fused: bool,
|
|
):
|
|
"""Test OINK support detection with clean import state for each parameter set."""
|
|
|
|
# Use spawn to run function in fresh process with clean imports
|
|
# TODO migrate to spawn utility:
|
|
# https://github.com/vllm-project/vllm/issues/41415
|
|
ctx = multiprocessing.get_context("spawn")
|
|
process = ctx.Process(
|
|
target=_test_oink_availability_impl,
|
|
args=(
|
|
device_capability,
|
|
has_rmsnorm,
|
|
has_fused_add_rms_norm,
|
|
expected_available,
|
|
expected_fused,
|
|
),
|
|
)
|
|
process.start()
|
|
process.join()
|
|
|
|
if process.exitcode != 0:
|
|
raise AssertionError(
|
|
f"Subprocess test failed with exit code {process.exitcode}"
|
|
)
|
|
|
|
|
|
def test_can_view_as_2d_stride_guard():
|
|
# No global import
|
|
import torch
|
|
|
|
# Import the helper from the kernels module.
|
|
from vllm.kernels.oink_ops import _can_view_as_2d
|
|
|
|
x = torch.zeros((2, 3, 4))
|
|
assert _can_view_as_2d(x) is True
|
|
|
|
# Size-1 dims should be ignored by the viewability check.
|
|
# Create a tensor where stride(0) != stride(1) * size(1) due to padding,
|
|
# but view(-1, H) is still valid because dim 1 has size 1.
|
|
base = torch.zeros((2, 10, 4))
|
|
x_singleton = base[:, :1, :]
|
|
x_singleton.view(-1, x_singleton.shape[-1])
|
|
assert _can_view_as_2d(x_singleton) is True
|
|
|
|
# Middle-dimension stride break: view(-1, hidden) should be invalid.
|
|
x2 = x[:, ::2, :]
|
|
with pytest.raises(RuntimeError):
|
|
x2.view(-1, x2.shape[-1])
|
|
assert _can_view_as_2d(x2) is False
|