From a54dbda3bbe8aa9d1c1125986cc1e0fa0e7ab433 Mon Sep 17 00:00:00 2001
From: project6-dev <dev@project6.ai>
Date: Mon, 10 Aug 2026 06:29:11 +0000
Subject: [PATCH] fix(bridge): link against libixformer.so for silu_and_mul
 symbol
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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.
---
 ex_engine/csrc/ix_moe_bridge.cpp |  8 +++-----
 ex_engine/python/ix_bridge.py    | 34 ++++++++++++++++++++++++++++++++
 verify_single_gpu.py             | 30 ++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/ex_engine/csrc/ix_moe_bridge.cpp b/ex_engine/csrc/ix_moe_bridge.cpp
index 7a3d5f5b..6e294984 100644
--- a/ex_engine/csrc/ix_moe_bridge.cpp
+++ b/ex_engine/csrc/ix_moe_bridge.cpp
@@ -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 {
diff --git a/ex_engine/python/ix_bridge.py b/ex_engine/python/ix_bridge.py
index 9b55d7c8..84a6ab89 100644
--- a/ex_engine/python/ix_bridge.py
+++ b/ex_engine/python/ix_bridge.py
@@ -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
diff --git a/verify_single_gpu.py b/verify_single_gpu.py
index 9105eee7..08fbd504 100644
--- a/verify_single_gpu.py
+++ b/verify_single_gpu.py
@@ -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