diff --git a/ex_engine/python/ix_bridge.py b/ex_engine/python/ix_bridge.py index f39c06c2..d0fe654b 100644 --- a/ex_engine/python/ix_bridge.py +++ b/ex_engine/python/ix_bridge.py @@ -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) diff --git a/probe_ixformer_symbols.py b/probe_ixformer_symbols.py new file mode 100644 index 00000000..1222ff60 --- /dev/null +++ b/probe_ixformer_symbols.py @@ -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 + +// 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) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 39abe28f..41acf746 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -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