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
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""
|
|
ix_bridge.py — Load ix_moe_bridge C++ extension at runtime.
|
|
|
|
Calls ixformer::infer::topk_softmax() via C++ torch extension,
|
|
bypassing the missing Python binding in ixformer.functions.
|
|
|
|
Build: JIT-compiled on first import via torch.utils.cpp_extension.load()
|
|
(same mechanism as flash_qla_sm70 GDN kernel — proven to work on BI-V100)
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import torch
|
|
|
|
logger = logging.getLogger("ex_engine.ix_bridge")
|
|
|
|
_ix_bridge = None
|
|
_ix_bridge_available = False
|
|
|
|
def _load_bridge():
|
|
"""JIT-compile and load ix_moe_bridge.so"""
|
|
global _ix_bridge, _ix_bridge_available
|
|
if _ix_bridge is not None:
|
|
return _ix_bridge_available
|
|
|
|
csrc_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "csrc")
|
|
cpp_file = os.path.join(csrc_dir, "ix_moe_bridge.cpp")
|
|
|
|
if not os.path.exists(cpp_file):
|
|
# Try deployed path (inside vllm model dir)
|
|
alt_dir = os.path.dirname(os.path.abspath(__file__))
|
|
cpp_file = os.path.join(alt_dir, "ix_moe_bridge.cpp")
|
|
|
|
if not os.path.exists(cpp_file):
|
|
logger.warning("ix_moe_bridge.cpp not found at %s", cpp_file)
|
|
_ix_bridge_available = False
|
|
return False
|
|
|
|
try:
|
|
from torch.utils.cpp_extension import load
|
|
logger.info("JIT-compiling ix_moe_bridge.cpp ...")
|
|
_ix_bridge = load(
|
|
name="ix_moe_bridge",
|
|
sources=[cpp_file],
|
|
extra_cflags=["-O2"],
|
|
verbose=False,
|
|
)
|
|
_ix_bridge_available = True
|
|
logger.info("ix_moe_bridge loaded successfully: %s", dir(_ix_bridge))
|
|
return True
|
|
except Exception as e:
|
|
logger.warning("ix_moe_bridge JIT compile failed: %s", e)
|
|
_ix_bridge_available = False
|
|
return False
|
|
|
|
def topk_softmax(gating_output: torch.Tensor, topk: int, renormalize: bool = True):
|
|
"""
|
|
Fused topk+softmax via ixformer C++ API.
|
|
|
|
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():
|
|
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)
|