fix(bridge): link against libixformer.so for silu_and_mul symbol
- ix_bridge.py: auto-discover ixformer .so files, pass as extra_ldflags - ix_moe_bridge.cpp: fix mangled header from bad sed, add #include <optional> - verify_single_gpu.py: also pass extra_ldflags during JIT compile The undefined symbol _ZN8ixformer5infer12silu_and_mulERN2at6TensorES3_ lives in libixformer.so — need to explicitly link it.
This commit is contained in:
@@ -1,13 +1,8 @@
|
||||
// ix_moe_bridge.cpp — Full MoE pipeline bridge to ixformer C++ API
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
//
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
// Exposes ALL 6 MoE functions from ixformer::infer (ixformer.h):
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
// 1. topk_softmax — fused routing
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
// 2. moe_compute_token_index_api — permutation maps (src_dst, dst_src)
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
// 3. moe_expand_input — gather tokens by expert
|
||||
// 4. moe_w16a16_group_gemm — batched expert GEMM
|
||||
// 5. silu_and_mul — fused activation
|
||||
@@ -20,6 +15,9 @@ static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
#include <torch/extension.h>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
#include <optional>
|
||||
|
||||
static const std::optional<torch::Tensor> kNoneTensor = {};
|
||||
|
||||
// Forward-declare ixformer C++ API (from base image SDK)
|
||||
namespace ixformer {
|
||||
|
||||
@@ -51,6 +51,39 @@ def _load_bridge():
|
||||
_loaded = True
|
||||
|
||||
from torch.utils.cpp_extension import load
|
||||
import glob
|
||||
|
||||
# Find ixformer .so libraries to link against
|
||||
extra_ldflags = []
|
||||
ixf_lib_dirs = set()
|
||||
try:
|
||||
import ixformer
|
||||
ixf_dir = os.path.dirname(ixformer.__file__)
|
||||
# Link against all .so in the ixformer package
|
||||
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
||||
if "cpython" not in so: # skip the Python extension .so
|
||||
extra_ldflags.append(so)
|
||||
ixf_lib_dirs.add(os.path.dirname(so))
|
||||
# Also try the _C and _ixformer_torch extensions
|
||||
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
|
||||
extra_ldflags.append(so)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Also check /usr/local/corex/lib64 for libixattn etc
|
||||
corex_lib = "/usr/local/corex/lib64"
|
||||
if os.path.isdir(corex_lib):
|
||||
for lib in ["libixattn.so", "libixformer.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)
|
||||
|
||||
# Add rpath so the .so can find its dependencies at runtime
|
||||
for d in ixf_lib_dirs:
|
||||
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
||||
|
||||
logger.info("ix_bridge extra_ldflags: %s", extra_ldflags)
|
||||
|
||||
for cpp_name in _CPP_NAMES:
|
||||
cpp_path = _find_cpp(cpp_name)
|
||||
@@ -63,6 +96,7 @@ def _load_bridge():
|
||||
name=mod_name,
|
||||
sources=[cpp_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=extra_ldflags,
|
||||
verbose=False,
|
||||
)
|
||||
_available = True
|
||||
|
||||
@@ -52,12 +52,42 @@ def step0_compile_bridge():
|
||||
print(f" Source: {cpp_path}")
|
||||
|
||||
from torch.utils.cpp_extension import load
|
||||
import glob
|
||||
|
||||
# Find ixformer .so to link against
|
||||
extra_ldflags = []
|
||||
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)
|
||||
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
|
||||
extra_ldflags.append(so)
|
||||
except ImportError:
|
||||
pass
|
||||
corex_lib = "/usr/local/corex/lib64"
|
||||
if os.path.isdir(corex_lib):
|
||||
for lib in ["libixattn.so", "libixformer.so"]:
|
||||
p = os.path.join(corex_lib, lib)
|
||||
if os.path.exists(p) and p not in extra_ldflags:
|
||||
extra_ldflags.append(p)
|
||||
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
|
||||
try:
|
||||
import ixformer
|
||||
extra_ldflags.append(f"-Wl,-rpath,{os.path.dirname(ixformer.__file__)}")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
print(f" Link libs: {[os.path.basename(x) for x in extra_ldflags if not x.startswith('-')]}")
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
bridge = load(
|
||||
name="ix_full_bridge",
|
||||
sources=[cpp_path],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
extra_ldflags=extra_ldflags,
|
||||
verbose=True,
|
||||
)
|
||||
dt = time.time() - t0
|
||||
|
||||
Reference in New Issue
Block a user