基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/fused_moe/scripts/build_fused_moe_i8_tn_pybind.sh

@@ -18,9 +18,39 @@ mkdir -p "$BUILD_DIR"
 PYTHON_INCLUDE="$("${PYTHON_BIN}" -c "import sysconfig; print(sysconfig.get_path('include'))")"
 PYTHON_LIB="$("${PYTHON_BIN}" -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))")"
 PYTHON_LDLIB="$("${PYTHON_BIN}" -c "import sysconfig; print(sysconfig.get_config_var('LDLIBRARY'))")"
+PYTHON_SHARED_LIB="$("${PYTHON_BIN}" -c "import os, sys, sysconfig
+libdir = sysconfig.get_config_var('LIBDIR') or ''
+ldlibrary = sysconfig.get_config_var('LDLIBRARY') or ''
+library = sysconfig.get_config_var('LIBRARY') or ''
+ldversion = sysconfig.get_config_var('LDVERSION') or f'{sys.version_info.major}.{sys.version_info.minor}'
+version = f'{sys.version_info.major}.{sys.version_info.minor}'
+names = [ldlibrary, library, f'libpython{ldversion}.so', f'libpython{ldversion}.so.1.0', f'libpython{version}.so', f'libpython{version}.so.1.0']
+seen = set()
+for name in names:
+    if not name or name in seen:
+        continue
+    seen.add(name)
+    candidates = [name]
+    if name.endswith('.a'):
+        candidates = [name[:-2] + '.so', name[:-2] + '.so.1.0'] + candidates
+    for candidate in candidates:
+        path = os.path.join(libdir, candidate)
+        if os.path.isfile(path) and '.so' in os.path.basename(path):
+            print(path)
+            raise SystemExit(0)
+print('')")"
 
 PYBIND11_INCLUDE="$ROOT_DIR/third_party/pybind11/include"
 
+if [[ -z "$PYTHON_SHARED_LIB" ]]; then
+    echo "[ERROR] Cannot resolve Python shared library from PYTHON_BIN=$PYTHON_BIN" >&2
+    echo "        LIBDIR=$PYTHON_LIB LDLIBRARY=$PYTHON_LDLIB" >&2
+    exit 1
+fi
+
+PYTHON_LIB="$(dirname "$PYTHON_SHARED_LIB")"
+PYTHON_LDLIB="$(basename "$PYTHON_SHARED_LIB")"
+
 "$MACA_PATH/mxgpu_llvm/bin/mxcc" \
     -std=c++17 \
     -O2 \
@@ -42,9 +72,9 @@ g++ \
     -fPIC \
     "$BUILD_DIR/fused_moe_i8_tn_pybind.o" \
     -L"$MACA_PATH/lib" \
-    -L"/opt/conda/lib" \
-    -Wl,-rpath,"/opt/conda/lib" \
-    -l:"libpython3.12.so" \
+    -L"$PYTHON_LIB" \
+    -Wl,-rpath,"$PYTHON_LIB" \
+    -l:"$PYTHON_LDLIB" \
     -lmcruntime \
     -lmccompiler \
     -o "$SO"