diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index 93d03221..e4065bf8 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -14,6 +14,24 @@ ALLOW_MISMATCH_ENV = "FREETOKEN_ALLOW_CUDA_MISMATCH" _TRUE_VALUES = {"1", "true", "yes", "on"} +_AMDHIP64_VERSIONED_RE = re.compile(r"libamdhip64\.so\.(\d+(?:\.\d+)*)$") + + +def select_versioned_rocm_runtime(paths): + """Return the highest numeric libamdhip64 soname from ``paths``. + + Path/string lexical order is not a version order: for example ``7.9`` sorts + after ``7.14``. Keep this helper package-independent so both setup.py and + runtime JIT discovery can use the same selection contract. + """ + candidates = [] + for path in paths: + match = _AMDHIP64_VERSIONED_RE.fullmatch(os.path.basename(os.fspath(path))) + if match is None: + continue + version = tuple(int(part) for part in match.group(1).split(".")) + candidates.append((version, os.fspath(path), path)) + return max(candidates, key=lambda item: (item[0], item[1]))[2] if candidates else None def _is_rocm() -> bool: diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index eb7e71e1..3cb5feaf 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -97,9 +97,7 @@ #endif #ifndef cudaDevAttrCanUseHostPointerForRegisteredMem -// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). -// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. -#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeCanUseHostPointerForRegisteredMem #endif #ifndef cudaFuncSetAttribute diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index b59a588f..533b1222 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import importlib import os import pathlib @@ -64,6 +65,17 @@ def _hip_cflags(extra: List[str]) -> List[str]: return flags + [f"--offload-arch={arch}" for arch in arches] +def _select_versioned_rocm_runtime(paths): + """Load the standalone toolchain selector without importing kernel package state.""" + path = pathlib.Path(__file__).with_name("_toolchain.py") + spec = importlib.util.spec_from_file_location("_freetoken_toolchain_runtime", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load ROCm runtime selector from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.select_versioned_rocm_runtime(paths) + + @cache def _rocm_link_flags() -> List[str]: """Make ROCm's runtime library discoverable to JIT link commands. @@ -72,6 +84,11 @@ def _rocm_link_flags() -> List[str]: ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via an explicit linker search path without modifying the Python environment. + + The compatibility directory is keyed by the resolved runtime origin so a + long-lived user cache cannot retain a link to a previous ROCm SDK after the + environment or image changes. Different tensor-parallel ranks using the same + runtime still converge on the same cache path. """ candidates: list[pathlib.Path] = [] if os.getenv("ROCM_HOME"): @@ -93,18 +110,24 @@ def _rocm_link_flags() -> List[str]: unversioned = library_dir / "libamdhip64.so" link_dir = library_dir if not unversioned.exists(): - versioned = sorted(library_dir.glob("libamdhip64.so.*")) - if not versioned: + versioned = _select_versioned_rocm_runtime(library_dir.glob("libamdhip64.so.*")) + if versioned is None: continue - link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" + runtime_target = versioned.resolve() + cache_key = hashlib.sha256(str(runtime_target).encode("utf-8")).hexdigest()[:16] + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" / cache_key link_dir.mkdir(parents=True, exist_ok=True) compat_link = link_dir / "libamdhip64.so" - if not compat_link.exists() and not compat_link.is_symlink(): + if not compat_link.exists(): try: - compat_link.symlink_to(versioned[-1]) + compat_link.symlink_to(runtime_target) except FileExistsError: # Multiple tensor-parallel ranks may prepare the same cache. pass + if not compat_link.exists(): + raise RuntimeError( + f"ROCm runtime compatibility link is unavailable: {compat_link} -> {runtime_target}" + ) return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] diff --git a/setup.py b/setup.py index bac07c6f..98ec6d74 100644 --- a/setup.py +++ b/setup.py @@ -12,12 +12,16 @@ KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") -def _check_toolchain() -> None: +def _toolchain_module(): path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - module.check_nvcc_matches_torch() + return module + + +def _check_toolchain() -> None: + _toolchain_module().check_nvcc_matches_torch() def _is_rocm() -> bool: @@ -46,9 +50,11 @@ def _rocm_paths() -> tuple[list[str], list[str], str]: continue if (library_dir / "libamdhip64.so").exists(): return [str(include_dir)], [str(library_dir)], "amdhip64" - versioned = sorted(library_dir.glob("libamdhip64.so.*")) - if versioned: - return [str(include_dir)], [str(library_dir)], f":{versioned[-1].name}" + versioned = _toolchain_module().select_versioned_rocm_runtime( + library_dir.glob("libamdhip64.so.*") + ) + if versioned is not None: + return [str(include_dir)], [str(library_dir)], f":{versioned.name}" searched = ", ".join(str(path) for path in dict.fromkeys(candidates)) raise RuntimeError( diff --git a/tests/kernels/test_rocm_host_pointer_capability.py b/tests/kernels/test_rocm_host_pointer_capability.py new file mode 100644 index 00000000..2cf4a87e --- /dev/null +++ b/tests/kernels/test_rocm_host_pointer_capability.py @@ -0,0 +1,23 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +HIP_COMPAT = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include" / "freetoken" / "hip_compat.h" + + +def test_rocm_host_pointer_capability_is_not_aliased_to_uva(): + text = HIP_COMPAT.read_text() + + assert ( + "#define cudaDevAttrCanUseHostPointerForRegisteredMem " + "hipDeviceAttributeCanUseHostPointerForRegisteredMem" + ) in text + assert ( + "#define cudaDevAttrCanUseHostPointerForRegisteredMem " + "hipDeviceAttributeUnifiedAddressing" + ) not in text + + +if __name__ == "__main__": + test_rocm_host_pointer_capability_is_not_aliased_to_uva() + print("ROCM_HOST_POINTER_SOURCE_CONTRACT=PASS") diff --git a/tests/kernels/test_rocm_runtime_link_cache.py b/tests/kernels/test_rocm_runtime_link_cache.py new file mode 100644 index 00000000..fb1b47cb --- /dev/null +++ b/tests/kernels/test_rocm_runtime_link_cache.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path +import sys +import tempfile +from unittest import mock + + +ROOT = Path(__file__).parents[2] +PYTHON_ROOT = ROOT / "python" +UTILS = PYTHON_ROOT / "freetoken" / "kernel" / "utils.py" + +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + + +def _load_utils_module(): + spec = importlib.util.spec_from_file_location("_freetoken_kernel_utils_test", UTILS) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_versioned_rocm_runtime_link_cache_tracks_runtime_origin_and_numeric_version() -> None: + module = _load_utils_module() + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + home = root / "home" + first_root = root / "rocm-first" + second_root = root / "rocm-second" + first_old = first_root / "lib" / "libamdhip64.so.7.9" + first_new = first_root / "lib" / "libamdhip64.so.7.14" + first_debug = first_root / "lib" / "libamdhip64.so.debug" + second_lib = second_root / "lib" / "libamdhip64.so.7.15" + first_old.parent.mkdir(parents=True) + second_lib.parent.mkdir(parents=True) + home.mkdir() + first_old.write_bytes(b"first-old") + first_new.write_bytes(b"first-new") + first_debug.write_bytes(b"debug") + second_lib.write_bytes(b"second") + + with mock.patch.dict(os.environ, {"HOME": str(home), "ROCM_HOME": str(first_root)}, clear=False): + module._rocm_link_flags.cache_clear() + first_flags = module._rocm_link_flags() + + first_link_dir = Path(next(flag[2:] for flag in first_flags if flag.startswith("-L"))) + first_link = first_link_dir / "libamdhip64.so" + assert first_link.is_symlink() + assert first_link.resolve() == first_new.resolve() + + # Model a long-lived cache surviving a ROCm SDK/image change. A stale + # compat symlink must not pin JIT linking to the vanished runtime, and + # the first selection must be numeric (7.14 > 7.9), not lexical. + first_new.unlink() + + with mock.patch.dict(os.environ, {"HOME": str(home), "ROCM_HOME": str(second_root)}, clear=False): + module._rocm_link_flags.cache_clear() + second_flags = module._rocm_link_flags() + + second_link_dir = Path(next(flag[2:] for flag in second_flags if flag.startswith("-L"))) + second_link = second_link_dir / "libamdhip64.so" + assert second_link.is_symlink() + assert second_link.exists() + assert second_link.resolve() == second_lib.resolve() + assert second_link_dir != first_link_dir + + utils_text = UTILS.read_text() + assert "select_versioned_rocm_runtime" in utils_text + assert 'sorted(library_dir.glob("libamdhip64.so.*"))' not in utils_text + + +if __name__ == "__main__": + test_versioned_rocm_runtime_link_cache_tracks_runtime_origin_and_numeric_version() + print("ROCM_JIT_RUNTIME_RESOLUTION=PASS_NUMERIC_SELECTION_AND_CACHE_LIFETIME") diff --git a/tests/kernels/test_rocm_versioned_runtime_selection.py b/tests/kernels/test_rocm_versioned_runtime_selection.py new file mode 100644 index 00000000..6266766a --- /dev/null +++ b/tests/kernels/test_rocm_versioned_runtime_selection.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +TOOLCHAIN = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" +SETUP = ROOT / "setup.py" + + +def _load_toolchain(): + spec = importlib.util.spec_from_file_location("_freetoken_toolchain_test", TOOLCHAIN) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_rocm_versioned_runtime_selection_is_numeric_not_lexical() -> None: + module = _load_toolchain() + candidates = [ + Path("/sdk/lib/libamdhip64.so.7"), + Path("/sdk/lib/libamdhip64.so.7.9"), + Path("/sdk/lib/libamdhip64.so.7.14"), + Path("/sdk/lib/libamdhip64.so.debug"), + ] + selected = module.select_versioned_rocm_runtime(candidates) + assert selected is not None + assert selected.name == "libamdhip64.so.7.14" + + setup_text = SETUP.read_text() + assert "select_versioned_rocm_runtime" in setup_text + assert 'sorted(library_dir.glob("libamdhip64.so.*"))' not in setup_text + + +if __name__ == "__main__": + test_rocm_versioned_runtime_selection_is_numeric_not_lexical() + print("ROCM_VERSIONED_RUNTIME_SELECTION=PASS_NUMERIC_7_14_OVER_7_9")