main
py 60 lines 1.95 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 """Check that DeepSeek V4 and the JIT monitor leave the HIP symbols alone.
5
6 TileLang ships a libhip_stub.so that takes over `hipFree` in the global symbol
7 table once it is loaded, so neither the model import nor the JIT monitor may
8 pull TileLang in on ROCm. Both claims are about process-global state, which is
9 why this runs in an interpreter of its own.
10 """
11
12 import ctypes
13 import sys
14
15 from vllm.model_executor.layers import mhc # noqa: F401
16 from vllm.models import deepseek_v4 # noqa: F401
17 from vllm.utils import jit_monitor
18
19
20 class DlInfo(ctypes.Structure):
21 _fields_ = [
22 ("dli_fname", ctypes.c_char_p),
23 ("dli_fbase", ctypes.c_void_p),
24 ("dli_sname", ctypes.c_char_p),
25 ("dli_saddr", ctypes.c_void_p),
26 ]
27
28
29 def _must_not_run() -> None:
30 raise AssertionError("TileLang JIT monitor must not run on ROCm")
31
32
33 jit_monitor._active = False
34 jit_monitor._setup_triton_autotuning_print = lambda: None
35 jit_monitor._setup_triton_jit_hook = lambda: None
36 jit_monitor._setup_cutedsl_jit_hook = lambda: None
37 jit_monitor._setup_tilelang_jit_hook = _must_not_run
38 jit_monitor.activate()
39
40 imported = [
41 name for name in sys.modules if name == "tilelang" or name.startswith("tilelang.")
42 ]
43 assert not imported, f"TileLang was imported: {imported}"
44
45 libdl = ctypes.CDLL("libdl.so.2")
46 dlsym = libdl.dlsym
47 dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
48 dlsym.restype = ctypes.c_void_p
49 dladdr = libdl.dladdr
50 dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(DlInfo)]
51 dladdr.restype = ctypes.c_int
52
53 address = dlsym(None, b"hipFree")
54 assert address is not None, "hipFree is not available in the global symbol table"
55 info = DlInfo()
56 assert dladdr(address, ctypes.byref(info))
57 source = info.dli_fname.decode() if info.dli_fname else "<unknown>"
58 assert "libhip_stub.so" not in source, source
59 assert "libamdhip64.so" in source, source
60 print("OK")