fix(rocm): correct runtime capability and resolution semantics

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:

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

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}"]
 

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(

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")

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")

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")