39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
|
|
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
|
|
return VllmConfig(
|
|
device_config=DeviceConfig(device="cpu"),
|
|
compilation_config=CompilationConfig(backend=backend, mode=mode),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"mode",
|
|
[CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
|
|
)
|
|
def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
|
|
"""Regression for #4964: inductor-active configs must not switch to native-only."""
|
|
priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
|
|
assert priority.rms_norm == ["vllm_c", "native"]
|
|
assert priority.fused_add_rms_norm == ["vllm_c", "native"]
|
|
|
|
|
|
def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
import vllm.envs as envs
|
|
|
|
monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
|
|
priority = CudaOmniPlatform.get_default_ir_op_priority(
|
|
_vllm_config(mode=CompilationMode.VLLM_COMPILE),
|
|
)
|
|
assert priority.rms_norm == ["oink", "vllm_c", "native"]
|
|
assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]
|