From 2760be885386a4f60294bd7ebf89d75e4ae501bd Mon Sep 17 00:00:00 2001 From: bingoo <1575938147@qq.com> Date: Tue, 19 May 2026 15:54:41 +0800 Subject: [PATCH] fix: find and load libcudart dynamically via /proc/self/maps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add find_loaded_library() to search /proc/self/maps for a loaded .so - Use the found cudart path first before falling back to hardcoded path - Fixes loading cudart.so issues when CUDA_LIB_PATH does not match - Regression: norm/rmsnorm_silu PASS, dit_layernorm 35 PASS, moe smoke PASS Refs: MISMATCH_EXPERIMENT ยง21 --- flashinfer/jit/__init__.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 8378e0ab74..8f76883867 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -17,6 +17,7 @@ import ctypes import functools import os +from typing import Optional # Re-export from . import cubin_loader @@ -99,8 +100,38 @@ ) +def find_loaded_library(lib_name) -> Optional[str]: + """ + According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, + the file contains the memory maps of the process, which includes the + shared libraries loaded by the process. We can use this file to find the path of the + a loaded library. + """ + found = False + with open("/proc/self/maps") as f: + for line in f: + if lib_name in line: + found = True + break + if not found: + # the library is not loaded in the current process + return None + # if lib_name is libcudart, we need to match a line with: + # address /path/to/libcudart-hash.so.11.0 + start = line.index("/") + path = line[start:].strip() + filename = path.split("/")[-1] + assert filename.rpartition(".so")[0].startswith(lib_name), ( + f"Unexpected filename: {filename} for library {lib_name}" + ) + return path + + cuda_lib_path = os.environ.get( "CUDA_LIB_PATH", "/usr/local/cuda/targets/x86_64-linux/lib/" ) -if os.path.exists(f"{cuda_lib_path}/libcudart.so.12"): +process_cudart_path = find_loaded_library("libcudart") +if process_cudart_path is not None: + ctypes.CDLL(process_cudart_path, mode=ctypes.RTLD_GLOBAL) +elif os.path.exists(f"{cuda_lib_path}/libcudart.so.12"): ctypes.CDLL(f"{cuda_lib_path}/libcudart.so.12", mode=ctypes.RTLD_GLOBAL)