250 lines
9.4 KiB
Python
250 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
verify_xllm_moe_load.py — 在 BI-V100 真机验证 xllm_moe.so 能否加载
|
|
|
|
用法: python3 verify_xllm_moe_load.py
|
|
或 Docker 内: python3 /workspace/qwen3_6_scripts/verify_xllm_moe_load.py
|
|
|
|
验证内容:
|
|
1. 所有 prebuilt .so 的 dlopen 可行性
|
|
2. xllm_moe.so 的 moe_fused_topk 是否可调用
|
|
3. corex_moe_topk_softmax.so 的 moe_topk_softmax 是否可调用
|
|
4. ix_startup_patch.apply() 完整执行结果
|
|
5. _custom_ops.topk_softmax 实际走哪条路径
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import importlib
|
|
import importlib.util
|
|
import traceback
|
|
|
|
def section(title):
|
|
print(f"\n{'='*65}")
|
|
print(f" {title}")
|
|
print(f"{'='*65}")
|
|
|
|
def try_load_so(name, path):
|
|
"""Try to load a .so file, return (module, error_string)."""
|
|
if not os.path.isfile(path):
|
|
return None, f"FILE NOT FOUND: {path}"
|
|
try:
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
if spec is None:
|
|
return None, "spec_from_file_location returned None"
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
funcs = [x for x in dir(mod) if not x.startswith('_')]
|
|
return mod, f"OK ({len(funcs)} functions: {', '.join(funcs[:10])})"
|
|
except Exception as e:
|
|
return None, f"{type(e).__name__}: {e}"
|
|
|
|
# =========================================================================
|
|
section("1. Locate all .so files")
|
|
# =========================================================================
|
|
|
|
# Search paths (same as xllm_ops.py)
|
|
search_dirs = []
|
|
try:
|
|
import vllm
|
|
vllm_root = os.path.dirname(vllm.__file__)
|
|
search_dirs.append(vllm_root)
|
|
print(f" vllm root: {vllm_root}")
|
|
except ImportError:
|
|
print(" vllm not installed")
|
|
vllm_root = None
|
|
|
|
for d in [
|
|
"/workspace/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10",
|
|
"/workspace/qwen3_6_scripts/ex_engine/prebuilt",
|
|
os.path.join(os.path.dirname(__file__), "prebuilt", "corex-3.2.3-ivcore10"),
|
|
os.path.join(os.path.dirname(__file__), "qwen3_6_scripts", "prebuilt", "corex-3.2.3-ivcore10"),
|
|
]:
|
|
if os.path.isdir(d):
|
|
search_dirs.append(os.path.normpath(d))
|
|
|
|
search_dirs = list(dict.fromkeys(search_dirs)) # dedupe preserving order
|
|
print(f" search dirs: {search_dirs}")
|
|
|
|
# =========================================================================
|
|
section("2. Load test: all xllm_*.so and key bridge .so")
|
|
# =========================================================================
|
|
|
|
test_modules = [
|
|
"xllm_moe",
|
|
"xllm_norm",
|
|
"xllm_activation",
|
|
"xllm_cache",
|
|
"xllm_rope",
|
|
"ix_moe_bridge",
|
|
"ix_full_bridge",
|
|
"corex_moe_topk_softmax",
|
|
"corex_moe_index_combine",
|
|
"corex_moe_direct_routed",
|
|
"corex_moe_weight_gather",
|
|
"corex_moe_exact_reduce",
|
|
"gemm_grouped",
|
|
"corex_batched_gemm",
|
|
]
|
|
|
|
loaded_modules = {}
|
|
for name in test_modules:
|
|
found = False
|
|
for d in search_dirs:
|
|
path = os.path.join(d, f"{name}.so")
|
|
if os.path.isfile(path):
|
|
mod, status = try_load_so(name, path)
|
|
tag = "✓" if mod else "✗"
|
|
print(f" {tag} {name:35s} {status}")
|
|
if mod:
|
|
loaded_modules[name] = mod
|
|
found = True
|
|
break
|
|
if not found:
|
|
print(f" - {name:35s} NOT FOUND in any search dir")
|
|
|
|
# =========================================================================
|
|
section("3. Functional test: xllm_moe.moe_fused_topk")
|
|
# =========================================================================
|
|
|
|
if "xllm_moe" in loaded_modules:
|
|
mod = loaded_modules["xllm_moe"]
|
|
print(f" Exported functions: {[x for x in dir(mod) if not x.startswith('_')]}")
|
|
if hasattr(mod, "moe_fused_topk"):
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
# moe_fused_topk(gating_output, topk, renormalize, correction_bias, scoring_func)
|
|
gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
|
w, ids = mod.moe_fused_topk(gating, 2)
|
|
print(f" ✓ moe_fused_topk: weights shape={w.shape}, ids shape={ids.shape}")
|
|
print(f" weights={w[0].tolist()}, ids={ids[0].tolist()}")
|
|
else:
|
|
print(f" - no CUDA device, skipping functional test")
|
|
except Exception as e:
|
|
print(f" ✗ moe_fused_topk call FAILED: {e}")
|
|
traceback.print_exc()
|
|
else:
|
|
print(f" ✗ moe_fused_topk NOT in module attrs")
|
|
print(f" available: {[x for x in dir(mod) if not x.startswith('_')]}")
|
|
else:
|
|
print(f" ✗ xllm_moe.so not loaded, cannot test")
|
|
|
|
# =========================================================================
|
|
section("4. Functional test: corex_moe_topk_softmax.moe_topk_softmax")
|
|
# =========================================================================
|
|
|
|
if "corex_moe_topk_softmax" in loaded_modules:
|
|
mod = loaded_modules["corex_moe_topk_softmax"]
|
|
print(f" Exported functions: {[x for x in dir(mod) if not x.startswith('_')]}")
|
|
if hasattr(mod, "moe_topk_softmax"):
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
|
w, ids = mod.moe_topk_softmax(gating, 2, True)
|
|
print(f" ✓ moe_topk_softmax: weights shape={w.shape}, ids shape={ids.shape}")
|
|
print(f" weights={w[0].tolist()}, ids={ids[0].tolist()}")
|
|
else:
|
|
print(f" - no CUDA device")
|
|
except Exception as e:
|
|
print(f" ✗ moe_topk_softmax call FAILED: {e}")
|
|
traceback.print_exc()
|
|
else:
|
|
print(f" ✗ moe_topk_softmax NOT in module attrs")
|
|
else:
|
|
print(f" ✗ corex_moe_topk_softmax.so not loaded, cannot test")
|
|
|
|
# =========================================================================
|
|
section("5. ixformer.functions.vllm_moe_topk_softmax existence")
|
|
# =========================================================================
|
|
|
|
try:
|
|
import ixformer.functions as ixf_F
|
|
has_it = hasattr(ixf_F, "vllm_moe_topk_softmax")
|
|
print(f" ixformer.functions.vllm_moe_topk_softmax: {'✓ EXISTS' if has_it else '✗ MISSING (expected on BI-V100)'}")
|
|
if not has_it:
|
|
# Check what MoE-related functions DO exist
|
|
moe_funcs = [n for n in dir(ixf_F) if 'moe' in n.lower() or 'topk' in n.lower()]
|
|
print(f" MoE-related functions that DO exist: {moe_funcs if moe_funcs else 'NONE'}")
|
|
except ImportError as e:
|
|
print(f" ixformer not available: {e}")
|
|
|
|
# =========================================================================
|
|
section("6. ix_startup_patch.apply() test")
|
|
# =========================================================================
|
|
|
|
try:
|
|
from vllm import ix_startup_patch
|
|
print(f" ix_startup_patch imported from: {ix_startup_patch.__file__}")
|
|
n = ix_startup_patch.apply()
|
|
print(f" ✓ apply() returned: {n} patches applied")
|
|
except ImportError as e:
|
|
print(f" ✗ import failed: {e}")
|
|
except Exception as e:
|
|
print(f" ✗ apply() failed: {type(e).__name__}: {e}")
|
|
traceback.print_exc()
|
|
|
|
# =========================================================================
|
|
section("7. _custom_ops.topk_softmax dispatch test")
|
|
# =========================================================================
|
|
|
|
try:
|
|
import vllm._custom_ops as ops
|
|
print(f" _custom_ops from: {ops.__file__}")
|
|
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
|
tw = torch.empty(4, 2, device="cuda", dtype=torch.float32)
|
|
ti = torch.empty(4, 2, device="cuda", dtype=torch.int32)
|
|
tei = torch.empty(4, 2, device="cuda", dtype=torch.int32)
|
|
try:
|
|
ops.topk_softmax(tw, ti, tei, gating)
|
|
print(f" ✓ topk_softmax succeeded")
|
|
print(f" weights={tw[0].tolist()}, ids={ti[0].tolist()}")
|
|
except Exception as e:
|
|
print(f" ✗ topk_softmax FAILED: {e}")
|
|
|
|
# Check which path was taken by inspecting the function
|
|
import inspect
|
|
src = inspect.getsource(ops.topk_softmax)
|
|
if "hasattr(ixf_F" in src:
|
|
print(f" → using PATCHED _custom_ops.py (has fallback chain)")
|
|
elif "ixf_F.vllm_moe_topk_softmax" in src and "hasattr" not in src:
|
|
print(f" → using ORIGINAL _custom_ops.py (NO fallback, will crash)")
|
|
else:
|
|
print(f" → using UNKNOWN version of topk_softmax")
|
|
else:
|
|
print(f" - no CUDA device")
|
|
except Exception as e:
|
|
print(f" ✗ {e}")
|
|
traceback.print_exc()
|
|
|
|
# =========================================================================
|
|
section("8. Library compatibility check")
|
|
# =========================================================================
|
|
|
|
try:
|
|
import torch
|
|
print(f" PyTorch: {torch.__version__}")
|
|
print(f" CUDA available: {torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f" CUDA version: {torch.version.cuda}")
|
|
print(f" Device: {torch.cuda.get_device_name(0)}")
|
|
|
|
# Check if libcudart.so.10.2 exists (required by all xllm .so)
|
|
import ctypes
|
|
for lib in ["libcudart.so.10.2", "libcuinfer.so.7", "libc10.so",
|
|
"libtorch_python.so", "libtorch_cuda.so"]:
|
|
try:
|
|
ctypes.CDLL(lib)
|
|
print(f" ✓ {lib}")
|
|
except OSError as e:
|
|
print(f" ✗ {lib}: {e}")
|
|
except Exception as e:
|
|
print(f" error: {e}")
|
|
|
|
print(f"\n{'='*65}")
|
|
print(" DONE — paste this entire output back")
|
|
print(f"{'='*65}") |