From 47ad34448e0a7b453da32a08090c32d1cb5e8b9e Mon Sep 17 00:00:00 2001 From: Yue Zhu Date: Tue, 21 Jul 2026 10:03:15 +0800 Subject: [PATCH 1/2] fix(hygon): initialize extensions before global preload The previous preload path opened Torch HIP libraries and flash_attn_2_cuda directly with ctypes. That bypassed Python extension initialization order and caused a reproducible double free during process teardown once the Hygon runtime context was active. Import Torch and FlashAttention through Python's module machinery first, then promote the initialized FlashAttention extension with RTLD_GLOBAL and retain its handle. This preserves symbol visibility for dlsym-based Hygon operators while keeping module ownership and teardown under Python. Keep the change scoped to the Hygon preload path; the MetaX HPCC loading behavior is unchanged. --- python/infinicore/_preload.py | 62 +++++++++++++++++++---------------- 1 file changed, 33 insertions(+), 29 deletions(-) diff --git a/python/infinicore/_preload.py b/python/infinicore/_preload.py index fda7fbb46..69121ff2f 100644 --- a/python/infinicore/_preload.py +++ b/python/infinicore/_preload.py @@ -1,11 +1,15 @@ import ctypes import glob +import importlib import importlib.util import os import sys from typing import Iterable, List +_PRELOADED_HANDLES: List[ctypes.CDLL] = [] + + def _candidate_prefixes(path: str) -> List[str]: """ Return HPCC install prefixes to search for libs. @@ -68,35 +72,35 @@ def preload_hpcc() -> None: def preload_torch_hip() -> None: """ - Best-effort preload of torch HIP runtime libs with RTLD_GLOBAL. + Best-effort import of the torch HIP runtime. - This helps external extensions resolve c10::hip symbols when they are - not recorded as direct DT_NEEDED dependencies. + Loading torch's shared libraries individually with ctypes bypasses + Python's extension initialization order and can corrupt HIP teardown. """ - spec = importlib.util.find_spec("torch") - if spec is None or not spec.origin: - return - torch_dir = os.path.dirname(spec.origin) - torch_libdir = os.path.join(torch_dir, "lib") - if not os.path.isdir(torch_libdir): - return - - libs = [ - "libtorch_global_deps.so", - "libc10.so", - "libc10_hip.so", - "libtorch_cpu.so", - "libtorch.so", - "libtorch_hip.so", - ] - for lib in libs: - full = os.path.join(torch_libdir, lib) - if os.path.exists(full): - try: - ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL) - except OSError: - # Best-effort preload, continue on errors. - pass + try: + importlib.import_module("torch") + except (ImportError, OSError): + pass + + +def _import_extension_global(module_name: str, path: str) -> None: + module = sys.modules.get(module_name) + if module is None: + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load extension module from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop(module_name, None) + raise + + module_path = getattr(module, "__file__", None) or path + _PRELOADED_HANDLES.append( + ctypes.CDLL(module_path, mode=ctypes.RTLD_GLOBAL) + ) def preload_flash_attn() -> None: @@ -145,9 +149,9 @@ def preload_flash_attn() -> None: if not os.path.exists(so_path): continue try: - ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL) + _import_extension_global("flash_attn_2_cuda", so_path) return - except OSError: + except (ImportError, OSError): continue From 3aa3cae8d7d1f4e6539789c865aa82b9912c10c1 Mon Sep 17 00:00:00 2001 From: Yue Zhu Date: Tue, 21 Jul 2026 10:03:26 +0800 Subject: [PATCH 2/2] fix(hygon): initialize the InfiniRT device context Register the Hygon backend in centralized context initialization. Without this call, Hygon devices are absent from the runtime table, cuda device aliases cannot resolve to HYGON, and the InfiniRT C++ graph bridge has no active Hygon runtime to capture. DeviceEnabled keeps the initialization compiled out for builds that do not include the Hygon backend. --- src/infinicore/context/context_impl.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index 9f6aa4a66..29a470679 100644 --- a/src/infinicore/context/context_impl.cc +++ b/src/infinicore/context/context_impl.cc @@ -91,6 +91,7 @@ ContextImpl::ContextImpl() { initializeRuntime(); initializeRuntime(); initializeRuntime(); + initializeRuntime(); initializeRuntime(); if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) {