fix: fail-fast on ix_bridge failure + probe script for real machine
1. ix_bridge.py: RuntimeError instead of silent PyTorch fallback If JIT compile fails, crash immediately with diagnostic message. 0 score with no error log is worse than a visible crash. 2. qwen3_5.py: explicit WARNING log on import failure (not silent) Shows exact error so we can diagnose from docker log. 3. probe_ixformer_symbols.py: definitive test for real machine - Finds all ixformer .so files - nm/objdump for topk_softmax C++ symbol - Checks Python bindings - Attempts JIT compile + link (the real test) - Prints PASS/FAIL with next-step instructions Run on real machine: python3 probe_ixformer_symbols.py
This commit is contained in:
EX Engine
@@ -57,21 +57,15 @@ def topk_softmax(gating_output: torch.Tensor, topk: int, renormalize: bool = Tru
|
||||
"""
|
||||
Fused topk+softmax via ixformer C++ API.
|
||||
|
||||
Args:
|
||||
gating_output: (num_tokens, num_experts) router logits
|
||||
topk: number of experts to select
|
||||
renormalize: whether to renormalize weights
|
||||
|
||||
Returns:
|
||||
(topk_weights, topk_indices) — both (num_tokens, topk)
|
||||
FAIL FAST: if bridge not available, raises RuntimeError immediately.
|
||||
No silent fallback — 0 score with no error log is worse than a crash.
|
||||
"""
|
||||
if not _ix_bridge_available:
|
||||
if not _load_bridge():
|
||||
# Fallback to pure PyTorch
|
||||
probs = torch.softmax(gating_output.float(), dim=-1)
|
||||
topk_w, topk_ids = torch.topk(probs, topk, dim=-1)
|
||||
if renormalize:
|
||||
topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True)
|
||||
return topk_w, topk_ids.to(torch.int32)
|
||||
raise RuntimeError(
|
||||
"ix_moe_bridge: FATAL — ixformer C++ topk_softmax not available. "
|
||||
"JIT compile failed. Run probe_ixformer_symbols.py on real machine "
|
||||
"to diagnose. Cannot fall back silently — would produce 0 score."
|
||||
)
|
||||
|
||||
return _ix_bridge.topk_softmax(gating_output, topk, renormalize)
|
||||
|
||||
230
probe_ixformer_symbols.py
Normal file
230
probe_ixformer_symbols.py
Normal file
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
probe_ixformer_symbols.py — 在真机上跑，探测 ixformer C++ 符号表
|
||||
|
||||
用法: python3 probe_ixformer_symbols.py
|
||||
|
||||
输出:
|
||||
1. ixformer 所有 .so 文件路径
|
||||
2. 每个 .so 里包含 topk_softmax / moe / gdn / attention 的符号
|
||||
3. 结论：ix_moe_bridge.cpp 能不能链接成功
|
||||
"""
|
||||
|
||||
import subprocess, sys, os, glob
|
||||
|
||||
def find_ixformer_so():
|
||||
"""找到 ixformer 的所有 .so 文件"""
|
||||
paths = []
|
||||
# 方法1: 从 Python import 路径找
|
||||
try:
|
||||
import ixformer
|
||||
pkg_dir = os.path.dirname(ixformer.__file__)
|
||||
paths.extend(glob.glob(os.path.join(pkg_dir, "**/*.so"), recursive=True))
|
||||
paths.extend(glob.glob(os.path.join(pkg_dir, "**/*.so.*"), recursive=True))
|
||||
print(f"[1] ixformer package dir: {pkg_dir}")
|
||||
except ImportError:
|
||||
print("[1] ixformer not importable")
|
||||
|
||||
# 方法2: 搜索常见路径
|
||||
for base in ["/usr/local/corex/lib64", "/usr/local/corex/lib",
|
||||
"/usr/local/lib", "/usr/lib"]:
|
||||
paths.extend(glob.glob(os.path.join(base, "**/libixformer*"), recursive=True))
|
||||
paths.extend(glob.glob(os.path.join(base, "**/*ixformer*.so"), recursive=True))
|
||||
paths.extend(glob.glob(os.path.join(base, "**/libixattn*"), recursive=True))
|
||||
paths.extend(glob.glob(os.path.join(base, "**/libixinfer*"), recursive=True))
|
||||
|
||||
# 方法3: 从 torch 找已加载的 .so
|
||||
try:
|
||||
import torch
|
||||
# ixformer 的 C++ 后端可能是 _ixformer_torch.so 或 _C.so
|
||||
try:
|
||||
import ixformer._ixformer_torch as ixt
|
||||
if hasattr(ixt, '__file__') and ixt.__file__:
|
||||
paths.append(ixt.__file__)
|
||||
print(f"[2] _ixformer_torch: {ixt.__file__}")
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
import ixformer._C as ic
|
||||
if hasattr(ic, '__file__') and ic.__file__:
|
||||
paths.append(ic.__file__)
|
||||
print(f"[2] _C: {ic.__file__}")
|
||||
except:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
||||
return list(set(paths))
|
||||
|
||||
def nm_grep(so_path, patterns):
|
||||
"""用 nm 查符号，grep 匹配"""
|
||||
results = []
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["nm", "-D", "--demangle", so_path],
|
||||
capture_output=True, text=True, timeout=10)
|
||||
for line in out.stdout.splitlines():
|
||||
for p in patterns:
|
||||
if p.lower() in line.lower():
|
||||
results.append(line.strip())
|
||||
except Exception as e:
|
||||
# nm 可能不存在，用 objdump
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["objdump", "-T", so_path],
|
||||
capture_output=True, text=True, timeout=10)
|
||||
for line in out.stdout.splitlines():
|
||||
for p in patterns:
|
||||
if p.lower() in line.lower():
|
||||
results.append(line.strip())
|
||||
except Exception as e2:
|
||||
results.append(f"ERROR: nm/objdump failed: {e}, {e2}")
|
||||
return results
|
||||
|
||||
def check_python_binding():
|
||||
"""检查 Python 层面有没有 topk_softmax"""
|
||||
print("\n=== Python Binding Check ===")
|
||||
try:
|
||||
import ixformer.functions as ixf
|
||||
attrs = dir(ixf)
|
||||
moe_attrs = [a for a in attrs if 'moe' in a.lower() or 'topk' in a.lower()
|
||||
or 'softmax' in a.lower() or 'expert' in a.lower()]
|
||||
print(f" ixformer.functions MoE-related: {moe_attrs}")
|
||||
if not moe_attrs:
|
||||
print(f" ixformer.functions ALL ({len(attrs)}): {attrs}")
|
||||
except Exception as e:
|
||||
print(f" ixformer.functions: {e}")
|
||||
|
||||
try:
|
||||
import ixformer
|
||||
# 搜索所有子模块
|
||||
for attr_name in dir(ixformer):
|
||||
obj = getattr(ixformer, attr_name)
|
||||
if hasattr(obj, 'topk_softmax'):
|
||||
print(f" FOUND: ixformer.{attr_name}.topk_softmax")
|
||||
if hasattr(obj, 'moe_topk_softmax'):
|
||||
print(f" FOUND: ixformer.{attr_name}.moe_topk_softmax")
|
||||
except:
|
||||
pass
|
||||
|
||||
def check_torch_ops():
|
||||
"""检查 torch.ops 注册"""
|
||||
print("\n=== torch.ops Check ===")
|
||||
try:
|
||||
import torch
|
||||
# 检查是否有 ixformer 注册的 ops
|
||||
for ns in ['ixformer', '_ixformer', 'ixf', '_C']:
|
||||
try:
|
||||
ns_obj = getattr(torch.ops, ns, None)
|
||||
if ns_obj:
|
||||
ops = [x for x in dir(ns_obj) if 'topk' in x.lower() or 'moe' in x.lower()]
|
||||
if ops:
|
||||
print(f" torch.ops.{ns} MoE ops: {ops}")
|
||||
else:
|
||||
print(f" torch.ops.{ns} exists but no MoE ops: {dir(ns_obj)[:10]}...")
|
||||
except:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
||||
def try_jit_compile():
|
||||
"""尝试 JIT 编译 ix_moe_bridge.cpp 看链接是否成功"""
|
||||
print("\n=== JIT Compile Test ===")
|
||||
test_cpp = "/tmp/ix_probe_test.cpp"
|
||||
with open(test_cpp, "w") as f:
|
||||
f.write("""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// Forward-declare — this is what ix_moe_bridge.cpp needs
|
||||
namespace ixformer { namespace infer {
|
||||
void topk_softmax(torch::Tensor&, torch::Tensor&, torch::Tensor&,
|
||||
torch::Tensor&, bool);
|
||||
}}
|
||||
|
||||
void test_link() {
|
||||
auto a = torch::empty({1,1});
|
||||
auto b = torch::empty({1,1});
|
||||
auto c = torch::empty({1,1});
|
||||
auto d = torch::empty({1,1});
|
||||
ixformer::infer::topk_softmax(a, b, c, d, false);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("test_link", &test_link);
|
||||
}
|
||||
""")
|
||||
try:
|
||||
from torch.utils.cpp_extension import load
|
||||
ext = load(name="ix_probe_test", sources=[test_cpp],
|
||||
extra_cflags=["-O0"], verbose=True)
|
||||
print(" JIT COMPILE + LINK: SUCCESS ✓")
|
||||
print(" ixformer::infer::topk_softmax symbol resolved!")
|
||||
return True
|
||||
except Exception as e:
|
||||
err = str(e)
|
||||
if "undefined reference" in err or "undefined symbol" in err:
|
||||
print(f" JIT LINK FAILED: symbol not found in .so")
|
||||
print(f" Error: {err[:500]}")
|
||||
else:
|
||||
print(f" JIT COMPILE FAILED: {err[:500]}")
|
||||
return False
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 70)
|
||||
print("ixformer Symbol Probe")
|
||||
print("=" * 70)
|
||||
|
||||
# Step 1: Find .so files
|
||||
print("\n=== .so Files ===")
|
||||
so_files = find_ixformer_so()
|
||||
if not so_files:
|
||||
print(" No ixformer .so files found!")
|
||||
for f in sorted(set(so_files)):
|
||||
size = os.path.getsize(f) if os.path.exists(f) else 0
|
||||
print(f" {f} ({size/1024/1024:.1f} MB)")
|
||||
|
||||
# Step 2: Search for symbols
|
||||
patterns = ["topk_softmax", "moe_topk", "topk_gating",
|
||||
"moe_compute_token", "moe_expand", "moe_output_reduce",
|
||||
"moe_w16a16", "group_gemm"]
|
||||
print("\n=== Symbol Search (MoE-related) ===")
|
||||
found_any = False
|
||||
for f in sorted(set(so_files)):
|
||||
results = nm_grep(f, patterns)
|
||||
if results:
|
||||
found_any = True
|
||||
print(f"\n {os.path.basename(f)}:")
|
||||
for r in results[:20]:
|
||||
print(f" {r}")
|
||||
if not found_any:
|
||||
print(" No MoE symbols found in any .so")
|
||||
# Also search for ANY ixformer::infer symbols
|
||||
print("\n=== Symbol Search (ixformer::infer namespace) ===")
|
||||
for f in sorted(set(so_files)):
|
||||
results = nm_grep(f, ["ixformer", "infer"])
|
||||
if results:
|
||||
print(f"\n {os.path.basename(f)} ({len(results)} matches):")
|
||||
for r in results[:30]:
|
||||
print(f" {r}")
|
||||
|
||||
# Step 3: Python binding
|
||||
check_python_binding()
|
||||
|
||||
# Step 4: torch.ops
|
||||
check_torch_ops()
|
||||
|
||||
# Step 5: JIT compile test (the definitive answer)
|
||||
jit_ok = try_jit_compile()
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 70)
|
||||
if jit_ok:
|
||||
print("RESULT: ix_moe_bridge.cpp CAN link to ixformer::infer::topk_softmax")
|
||||
print("ACTION: proceed with C++ bridge approach")
|
||||
else:
|
||||
print("RESULT: ix_moe_bridge.cpp CANNOT link to ixformer C++ API")
|
||||
print("ACTION: need alternative — options:")
|
||||
print(" A) Build topk_softmax kernel from upstream_ref/xllm CUDA source")
|
||||
print(" B) Build from upstream_ref/ds_vllm/csrc/moe/topk_softmax_kernels.cu")
|
||||
print(" C) Keep PyTorch path but add explicit error logging")
|
||||
print("=" * 70)
|
||||
@@ -74,6 +74,7 @@ except ImportError:
|
||||
# ix_bridge: C++ bridge to ixformer::infer::topk_softmax (bypasses missing Python binding)
|
||||
_ix_bridge_module = None
|
||||
_ix_bridge_available = False
|
||||
_ix_topk_softmax = None
|
||||
try:
|
||||
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax
|
||||
_ix_bridge_available = True
|
||||
@@ -88,8 +89,11 @@ except ImportError:
|
||||
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk_softmax
|
||||
_ix_bridge_available = True
|
||||
logger.info("ix_bridge: ixformer C++ topk_softmax available (deployed path)")
|
||||
except ImportError:
|
||||
logger.info("ix_bridge: not available, MoE uses PyTorch topk")
|
||||
except ImportError as e:
|
||||
# NOT silent: log the exact error so we can diagnose from docker logs
|
||||
logger.warning(
|
||||
"ix_bridge: IMPORT FAILED (%s). MoE will use PyTorch topk. "
|
||||
"This is 3x slower. Run probe_ixformer_symbols.py to diagnose.", e)
|
||||
_corex_gdn_available = False
|
||||
_corex_moe_available = False
|
||||
|
||||
|
||||
Reference in New Issue
Block a user