fix(jit): isolate AITER extensions from HIP interposers

aiter/jit/core.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
+import ctypes
 import functools
 import importlib
+import importlib.util
 import json
 import logging
 import multiprocessing
@@ -603,13 +605,35 @@ def check_numa():
 __mds = {}
 
 
+_RTLD_DEEPBIND = (
+    0 if os.getenv("AITER_DISABLE_DEEPBIND") == "1" else getattr(os, "RTLD_DEEPBIND", 0)
+)
+_deep_handles = {}
+
+
+def _deep_import(mod_path: str) -> types.ModuleType:
+    """Import an extension deep-bound.
+
+    importlib takes no dlopen flags, so open the extension here first; glibc
+    dedupes by (st_dev, st_ino), so the import reuses this already-bound
+    handle. sys.setdlopenflags() is avoided on purpose: it is process-wide, so
+    a concurrent unrelated dlopen would inherit the flag. Handles are kept
+    because dropping a CDLL can dlclose the module out from under us.
+    """
+    spec = importlib.util.find_spec(mod_path)
+    origin = getattr(spec, "origin", None) or ""
+    if _RTLD_DEEPBIND and origin.endswith(".so") and origin not in _deep_handles:
+        _deep_handles[origin] = ctypes.CDLL(origin, mode=os.RTLD_NOW | _RTLD_DEEPBIND)
+    return importlib.import_module(mod_path)
+
+
 @torch_compile_guard()
 def get_module_custom_op(md_name: str) -> None:
     if md_name not in __mds:
         if "AITER_JIT_DIR" in os.environ:
-            __mds[md_name] = importlib.import_module(md_name)
+            __mds[md_name] = _deep_import(md_name)
         else:
-            __mds[md_name] = importlib.import_module(f"{__package__}.{md_name}")
+            __mds[md_name] = _deep_import(f"{__package__}.{md_name}")
         logger.info(f"import [{md_name}] under {__mds[md_name].__file__}")
 
 
@@ -1302,7 +1326,7 @@ def _ensure_loaded():
                 d_args.get("third_party", []),
                 flags_extra_hip_per_source=d_args.get("flags_extra_hip_per_source", {}),
             )
-        lib = ctypes.CDLL(so_path)
+        lib = ctypes.CDLL(so_path, mode=os.RTLD_LAZY | _RTLD_DEEPBIND)
         c_func = getattr(lib, fc_name)
 
         def _opt_sym(name, argtypes=(), restype=None):

csrc/cpp_itfs/utils.h

@@ -20,6 +20,18 @@ namespace aiter{
 
 #define DIVIDE_ROUND_UP(a, b) (((a) + (b)-1) / (b))
 
+__inline__ int aiter_dlopen_mode(int resolution_mode)
+{
+#ifdef RTLD_DEEPBIND
+    const char* disable_deepbind = std::getenv("AITER_DISABLE_DEEPBIND");
+    if(disable_deepbind == nullptr || std::string(disable_deepbind) != "1")
+    {
+        return resolution_mode | RTLD_DEEPBIND;
+    }
+#endif
+    return resolution_mode;
+}
+
 static std::once_flag init_libs_lru_cache, init_func_names_lru_cache, init_root_dir_flag;
 
 template<typename K, typename V>
@@ -92,7 +104,7 @@ class SharedLibrary {
 
 public:
     SharedLibrary(std::string& path) {
-        handle = dlopen(path.c_str(), RTLD_LAZY);
+        handle = dlopen(path.c_str(), aiter_dlopen_mode(RTLD_LAZY));
         if (!handle) {
             throw std::runtime_error(dlerror());
         }

csrc/cpp_itfs/utils.py

@@ -280,7 +280,14 @@ def final_func():
 def run_lib(func_name, folder=None):
     if folder is None:
         folder = func_name
-    lib = ctypes.CDLL(f"{BUILD_DIR}/{folder}/lib.so", os.RTLD_LAZY)
+    # Same HIP interposition guard as aiter.jit.core._RTLD_DEEPBIND. Recomputed
+    # rather than imported so this module stays standalone.
+    deepbind = (
+        0
+        if os.getenv("AITER_DISABLE_DEEPBIND") == "1"
+        else getattr(os, "RTLD_DEEPBIND", 0)
+    )
+    lib = ctypes.CDLL(f"{BUILD_DIR}/{folder}/lib.so", mode=os.RTLD_LAZY | deepbind)
     return getattr(lib, func_name)
 
 

op_tests/test_jit_loader.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import os
+import shutil
+import subprocess
+import sys
+import types
+
+import pytest
+
+from aiter.jit import core
+
+pytestmark = pytest.mark.skipif(
+    getattr(os, "RTLD_DEEPBIND", 0) == 0, reason="RTLD_DEEPBIND is a glibc extension"
+)
+
+# Stand-in for TileLang's libhip_stub.so, built with plain gcc so this does not
+# depend on TileLang being installed. It reproduces the real defect: the
+# exported symbol carries the R0600 name but resolves the unversioned one, so
+# the legacy implementation writes a differently laid out struct into an R0600
+# destination. Forwarding to the real runtime, rather than scribbling a fixed
+# pattern, keeps the write inside the caller's buffer -- so an unprotected
+# AITER fails the way it does in production, with hipErrorInvalidConfiguration,
+# instead of by SIGSEGV.
+_STUB_C = """
+#include <dlfcn.h>
+#include <stddef.h>
+typedef int (*legacy_fn)(void *, int);
+int hipGetDevicePropertiesR0600(void *prop, int device) {
+    static legacy_fn legacy = NULL;
+    if (legacy == NULL) {
+        const char *names[] = {
+            "libamdhip64.so", "libamdhip64.so.7",
+            "libamdhip64.so.6", "libamdhip64.so.5"
+        };
+        void *h = NULL;
+        for (size_t i = 0; i < sizeof(names) / sizeof(names[0]); ++i) {
+            h = dlopen(names[i], RTLD_NOW | RTLD_LOCAL);
+            if (h != NULL) break;
+        }
+        if (h == NULL) return 1;
+        legacy = (legacy_fn)dlsym(h, "hipGetDeviceProperties");
+        if (legacy == NULL) return 1;
+    }
+    return legacy(prop, device);
+}
+"""
+
+_INTERPOSE_SCRIPT = """
+import ctypes, math, sys, torch
+ctypes.CDLL(sys.argv[1], mode=ctypes.RTLD_GLOBAL)   # before AITER loads anything
+
+import aiter
+n, e, k = 16384, 128, 8
+gate = torch.randn(n, e, dtype=torch.float32, device="cuda")
+w = torch.empty(n, k, dtype=torch.float32, device="cuda")
+i = torch.empty(n, k, dtype=torch.int32, device="cuda")
+t = torch.empty(n, k, dtype=torch.int32, device="cuda")
+
+aiter.topk_softmax(w, i, t, gate, False)            # pybind extension path
+torch.cuda.synchronize()
+want = torch.topk(torch.softmax(gate, -1), k, dim=-1)[0]
+assert torch.allclose(w, want, atol=1e-5), "topk_softmax corrupted by interposed HIP"
+
+from aiter.ops.moe_op import topk_softmax_asm       # standalone ctypes path
+w_asm = torch.empty_like(w)
+i_asm = torch.empty_like(i)
+t_asm = torch.empty_like(t)
+topk_softmax_asm(w_asm, i_asm, t_asm, gate, False)
+torch.cuda.synchronize()
+assert torch.allclose(w_asm, want, atol=1e-5), "ctypes topk corrupted by interposed HIP"
+
+from aiter.fused_moe import moe_sorting             # sizes its grid from device props
+moe_sorting(i, w, e, 1024, torch.bfloat16)
+torch.cuda.synchronize()
+
+from aiter.ops.mha import fmha_v3_varlen_fwd        # Kimi-K3 failure path
+q = torch.randn((128, 8, 128), device="cuda", dtype=torch.bfloat16)
+cu = torch.tensor([0, 128], device="cuda", dtype=torch.int32)
+fmha_v3_varlen_fwd(
+    q, q, q, cu, cu, 128, 128, 0, 0.0, 1.0 / math.sqrt(128), 0.0,
+    False, False, -1, -1, False, False, 1,
+)
+torch.cuda.synchronize()
+print("OK")
+"""
+
+
+def test_deep_import_leaves_process_dlopen_flags_alone():
+    """A process-wide flag would be inherited by unrelated concurrent dlopens."""
+    handles = dict(core._deep_handles)
+    before = sys.getdlopenflags()
+    core._deep_import("json")  # pure Python: nothing to pre-open
+    assert sys.getdlopenflags() == before
+    assert core._deep_handles == handles
+
+
+def test_deep_import_predlopens_extension(monkeypatch):
+    origin = "/tmp/module_deepbind_test.so"
+    handle = object()
+    module = object()
+    seen = {}
+
+    monkeypatch.setattr(core, "_RTLD_DEEPBIND", os.RTLD_DEEPBIND)
+    monkeypatch.setattr(
+        core.importlib.util,
+        "find_spec",
+        lambda name: types.SimpleNamespace(origin=origin),
+    )
+
+    def fake_cdll(path, mode):
+        seen["path"] = path
+        seen["mode"] = mode
+        return handle
+
+    monkeypatch.setattr(core.ctypes, "CDLL", fake_cdll)
+    monkeypatch.setattr(core.importlib, "import_module", lambda name: module)
+    core._deep_handles.pop(origin, None)
+    try:
+        assert core._deep_import("module_deepbind_test") is module
+        assert seen == {"path": origin, "mode": os.RTLD_NOW | os.RTLD_DEEPBIND}
+        assert core._deep_handles[origin] is handle
+    finally:
+        core._deep_handles.pop(origin, None)
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [("0", os.RTLD_LAZY | os.RTLD_DEEPBIND), ("1", os.RTLD_LAZY)],
+)
+def test_generated_python_loader_mode(monkeypatch, value, expected):
+    """The standalone generated-library loader follows the same opt-out."""
+    from csrc.cpp_itfs import utils
+
+    symbol = object()
+    seen = {}
+
+    def fake_cdll(path, mode):
+        seen.update(path=path, mode=mode)
+        return types.SimpleNamespace(test_symbol=symbol)
+
+    monkeypatch.setenv("AITER_DISABLE_DEEPBIND", value)
+    monkeypatch.setattr(utils, "BUILD_DIR", "/tmp/aiter-loader-test")
+    monkeypatch.setattr(utils.ctypes, "CDLL", fake_cdll)
+    utils.run_lib.cache_clear()
+    try:
+        assert utils.run_lib("test_symbol", "test_folder") is symbol
+        assert seen == {
+            "path": "/tmp/aiter-loader-test/test_folder/lib.so",
+            "mode": expected,
+        }
+    finally:
+        utils.run_lib.cache_clear()
+
+
+def test_core_loader_optout_is_applied_at_import():
+    """The documented process-start opt-out must affect the actual mode flag."""
+    script = "from aiter.jit import core; print(core._RTLD_DEEPBIND)"
+    for value, expected in (("0", os.RTLD_DEEPBIND), ("1", 0)):
+        env = os.environ.copy()
+        env["AITER_DISABLE_DEEPBIND"] = value
+        proc = subprocess.run(
+            [sys.executable, "-c", script],
+            env=env,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        assert proc.returncode == 0, proc.stderr
+        assert int(proc.stdout.strip().splitlines()[-1]) == expected
+
+
+def test_aiter_survives_a_global_hip_interposer(tmp_path):
+    """AITER must stay correct with an interposer loaded first.
+
+    Runs in a subprocess: loaded objects and their global ordering cannot be
+    reset in-process, so the failing order has to be built from scratch.
+    """
+    torch = pytest.importorskip("torch")
+    if not torch.cuda.is_available():
+        pytest.skip("no GPU")
+    cc = shutil.which("gcc") or shutil.which("cc")
+    if cc is None:
+        pytest.skip("no C compiler for the interposer")
+
+    src = tmp_path / "interposer.c"
+    src.write_text(_STUB_C)
+    stub = tmp_path / "libinterposer.so"
+    build = subprocess.run(
+        [cc, "-shared", "-fPIC", "-o", str(stub), str(src), "-ldl"],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if build.returncode != 0:
+        pytest.skip(f"could not build interposer: {build.stderr}")
+
+    script = tmp_path / "interpose.py"
+    script.write_text(_INTERPOSE_SCRIPT)
+    env = os.environ.copy()
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    env["PYTHONPATH"] = os.pathsep.join(
+        path for path in (repo_root, env.get("PYTHONPATH")) if path
+    )
+    proc = subprocess.run(
+        [sys.executable, str(script), str(stub)],
+        env=env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert proc.returncode == 0, f"stdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
+    assert "OK" in proc.stdout
+
+    # Prove that the interposer is effective, so the protected run above
+    # cannot pass merely because the test setup failed to reproduce the bug.
+    unprotected_env = env.copy()
+    unprotected_env["AITER_DISABLE_DEEPBIND"] = "1"
+    unprotected = subprocess.run(
+        [sys.executable, str(script), str(stub)],
+        env=unprotected_env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert unprotected.returncode != 0, "interposer did not reproduce the bug"
+    failure = f"{unprotected.stdout}\n{unprotected.stderr}".lower()
+    assert "invalid argument" in failure or "invalid configuration" in failure
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))