108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Precompile ix_moe_bridge.cpp during Docker build.
|
||
|
|
|
||
|
|
This bridges Python ↔ ixformer::infer C++ API (topk_softmax, group_gemm, etc).
|
||
|
|
Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
|
||
|
|
"""
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import glob
|
||
|
|
import logging
|
||
|
|
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
logger = logging.getLogger("precompile_ix_bridge")
|
||
|
|
|
||
|
|
def find_ixformer_libs():
|
||
|
|
"""Find libixformer.so and related libraries for linking."""
|
||
|
|
extra_ldflags = []
|
||
|
|
ixf_lib_dirs = set()
|
||
|
|
|
||
|
|
try:
|
||
|
|
import ixformer
|
||
|
|
ixf_dir = os.path.dirname(ixformer.__file__)
|
||
|
|
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
||
|
|
if "cpython" not in so:
|
||
|
|
extra_ldflags.append(so)
|
||
|
|
ixf_lib_dirs.add(os.path.dirname(so))
|
||
|
|
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
|
||
|
|
if so not in extra_ldflags:
|
||
|
|
extra_ldflags.append(so)
|
||
|
|
except ImportError:
|
||
|
|
logger.warning("ixformer not installed")
|
||
|
|
|
||
|
|
corex_lib = "/usr/local/corex/lib64"
|
||
|
|
if os.path.isdir(corex_lib):
|
||
|
|
for lib in ["libixformer.so", "libixattn.so", "libcublas.so"]:
|
||
|
|
p = os.path.join(corex_lib, lib)
|
||
|
|
if os.path.exists(p) and p not in extra_ldflags:
|
||
|
|
extra_ldflags.append(p)
|
||
|
|
ixf_lib_dirs.add(corex_lib)
|
||
|
|
|
||
|
|
for d in ixf_lib_dirs:
|
||
|
|
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
||
|
|
|
||
|
|
return extra_ldflags
|
||
|
|
|
||
|
|
def main():
|
||
|
|
# Find the .cpp source
|
||
|
|
search = [
|
||
|
|
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
|
||
|
|
os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
|
||
|
|
]
|
||
|
|
|
||
|
|
cpp_path = None
|
||
|
|
for p in search:
|
||
|
|
if os.path.isfile(p):
|
||
|
|
cpp_path = p
|
||
|
|
break
|
||
|
|
|
||
|
|
if not cpp_path:
|
||
|
|
logger.error("ix_moe_bridge.cpp not found in: %s", search)
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
logger.info("Compiling ix_moe_bridge from %s", cpp_path)
|
||
|
|
|
||
|
|
extra_ldflags = find_ixformer_libs()
|
||
|
|
logger.info("Link flags: %s", extra_ldflags)
|
||
|
|
|
||
|
|
if not extra_ldflags:
|
||
|
|
logger.error("No ixformer libraries found — cannot compile bridge")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
try:
|
||
|
|
from torch.utils.cpp_extension import load
|
||
|
|
mod = load(
|
||
|
|
name="ix_moe_bridge",
|
||
|
|
sources=[cpp_path],
|
||
|
|
extra_cflags=["-O2", "-std=c++17"],
|
||
|
|
extra_ldflags=extra_ldflags,
|
||
|
|
verbose=True,
|
||
|
|
)
|
||
|
|
fns = [x for x in dir(mod) if not x.startswith("_")]
|
||
|
|
logger.info("SUCCESS: ix_moe_bridge compiled with functions: %s", fns)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("FAILED to compile ix_moe_bridge: %s", e)
|
||
|
|
# Also try ix_full_bridge.cpp
|
||
|
|
full_path = cpp_path.replace("ix_moe_bridge", "ix_full_bridge")
|
||
|
|
if os.path.isfile(full_path):
|
||
|
|
logger.info("Trying ix_full_bridge.cpp instead...")
|
||
|
|
try:
|
||
|
|
mod = load(
|
||
|
|
name="ix_full_bridge",
|
||
|
|
sources=[full_path],
|
||
|
|
extra_cflags=["-O2", "-std=c++17"],
|
||
|
|
extra_ldflags=extra_ldflags,
|
||
|
|
verbose=True,
|
||
|
|
)
|
||
|
|
fns = [x for x in dir(mod) if not x.startswith("_")]
|
||
|
|
logger.info("SUCCESS: ix_full_bridge compiled with functions: %s", fns)
|
||
|
|
except Exception as e2:
|
||
|
|
logger.error("FAILED ix_full_bridge too: %s", e2)
|
||
|
|
sys.exit(1)
|
||
|
|
else:
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|