c5b7c069a4
Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import ctypes
|
|
import importlib
|
|
import importlib.util
|
|
import sys
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils import import_utils
|
|
|
|
|
|
class _PassConfigKey:
|
|
TL_DISABLE_WARP_SPECIALIZED = "disable_warp_specialized"
|
|
TL_DISABLE_TMA_LOWER = "disable_tma_lower"
|
|
TL_PTXAS_REGISTER_USAGE_LEVEL = "ptxas_register_usage_level"
|
|
|
|
|
|
def _install_tilelang_stub(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> dict[str, int]:
|
|
calls = {"jit_decorate": 0, "compiled_call": 0}
|
|
|
|
tilelang: Any = ModuleType("tilelang")
|
|
|
|
def jit(**kwargs: Any) -> Any:
|
|
def decorate(func: Any) -> Any:
|
|
calls["jit_decorate"] += 1
|
|
|
|
def compiled(*args: Any, **kw: Any) -> Any:
|
|
calls["compiled_call"] += 1
|
|
return func.__name__
|
|
|
|
return compiled
|
|
|
|
return decorate
|
|
|
|
tilelang.PassConfigKey = _PassConfigKey
|
|
tilelang.jit = jit
|
|
|
|
monkeypatch.setattr(import_utils, "has_tilelang", lambda: True)
|
|
monkeypatch.setitem(sys.modules, "tilelang", tilelang)
|
|
monkeypatch.setitem(
|
|
sys.modules, "tilelang.language", ModuleType("tilelang.language")
|
|
)
|
|
monkeypatch.delitem(sys.modules, "vllm.tilelang_utils", raising=False)
|
|
|
|
return calls
|
|
|
|
|
|
def test_tilelang_jit_decorator_is_lazy_only_on_rocm(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
if not (current_platform.is_cuda() or current_platform.is_rocm()):
|
|
pytest.skip("Test requires CUDA or ROCm")
|
|
|
|
calls = _install_tilelang_stub(monkeypatch)
|
|
module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels"
|
|
monkeypatch.delitem(sys.modules, module_name, raising=False)
|
|
module = importlib.import_module(module_name)
|
|
|
|
if current_platform.is_rocm():
|
|
assert calls["jit_decorate"] == 0
|
|
else:
|
|
assert calls["jit_decorate"] > 0
|
|
|
|
decorated_calls = calls["jit_decorate"]
|
|
assert module.mhc_post_tilelang() == "mhc_post_tilelang"
|
|
if current_platform.is_rocm():
|
|
assert calls["jit_decorate"] == 1
|
|
else:
|
|
assert calls["jit_decorate"] == decorated_calls
|
|
assert calls["compiled_call"] == 1
|
|
|
|
|
|
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
|
|
def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
if importlib.util.find_spec("tilelang") is None:
|
|
pytest.skip("Test requires TileLang to be installed")
|
|
|
|
class DlInfo(ctypes.Structure):
|
|
_fields_ = [
|
|
("dli_fname", ctypes.c_char_p),
|
|
("dli_fbase", ctypes.c_void_p),
|
|
("dli_sname", ctypes.c_char_p),
|
|
("dli_saddr", ctypes.c_void_p),
|
|
]
|
|
|
|
libdl = ctypes.CDLL("libdl.so.2")
|
|
dlsym = libdl.dlsym
|
|
dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
|
|
dlsym.restype = ctypes.c_void_p
|
|
dladdr = libdl.dladdr
|
|
dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(DlInfo)]
|
|
dladdr.restype = ctypes.c_int
|
|
|
|
from vllm.model_executor.layers import mhc # noqa: F401
|
|
from vllm.models import deepseek_v4 # noqa: F401
|
|
from vllm.utils import jit_monitor
|
|
|
|
monkeypatch.setattr(jit_monitor, "_active", False)
|
|
monkeypatch.setattr(jit_monitor, "_setup_triton_autotuning_print", lambda: None)
|
|
monkeypatch.setattr(jit_monitor, "_setup_triton_jit_hook", lambda: None)
|
|
monkeypatch.setattr(jit_monitor, "_setup_cutedsl_jit_hook", lambda: None)
|
|
monkeypatch.setattr(
|
|
jit_monitor,
|
|
"_setup_tilelang_jit_hook",
|
|
lambda: pytest.fail("TileLang JIT monitor must not run on ROCm"),
|
|
)
|
|
jit_monitor.activate()
|
|
|
|
assert not any(
|
|
name == "tilelang" or name.startswith("tilelang.") for name in sys.modules
|
|
)
|
|
|
|
address = dlsym(None, b"hipFree")
|
|
assert address is not None, "hipFree is not available in the global symbol table"
|
|
info = DlInfo()
|
|
assert dladdr(address, ctypes.byref(info))
|
|
source = info.dli_fname.decode() if info.dli_fname else "<unknown>"
|
|
assert "libhip_stub.so" not in source, source
|
|
assert "libamdhip64.so" in source, source
|