e04a3bace9
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
231 lines
8.1 KiB
Python
231 lines
8.1 KiB
Python
#!/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)
|