feat(CRITICAL): ix_bridge call chain + upstream xllm/ds_vllm sync

3 changes that close the MoE performance gap:

1. _custom_ops.py: Add ix_bridge as Priority 0 for topk_softmax
   - Before: tries our .cu kernel (fails) → PyTorch fallback (1-3 TPS)
   - After: tries ix_bridge → ixformer::infer::topk_softmax() → FAST
   - Call chain: _custom_ops.topk_softmax() → ix_bridge.topk_softmax()
     → ix_moe_bridge.so → ixformer::infer::topk_softmax()

2. Dockerfile: Add ix_moe_bridge.cpp precompile step
   - This was the missing link: code existed but was never compiled
   - Uses torch.utils.cpp_extension.load() to link against libixformer.so

3. upstream_ref sync from GitHub (cloned, not rewritten):
   - xLLM-AI/xllm: ILU kernels + CUDA MoE + GDN fp32 state mgmt
   - Deep-Spark/vllm: latest MoE kernel sources
This commit is contained in:
project6-dev
2026-08-11 01:27:33 +00:00
parent 26e6cb4019
commit 56146f8130
16 changed files with 5755 additions and 1 deletions

View File

@@ -974,9 +974,44 @@ def invoke_fused_moe_kernel(
_moe_topk_ext = None
_moe_topk_init_done = False
_ix_bridge_mod = None
_ix_bridge_init_done = False
def _init_ix_bridge():
"""Try to load ix_bridge which calls ixformer::infer::topk_softmax() via C++ pybind."""
global _ix_bridge_mod, _ix_bridge_init_done
_ix_bridge_init_done = True
try:
from ex_engine.python.ix_bridge import is_available, topk_softmax as _ix_ts
if is_available():
_ix_bridge_mod = True
logger.info("topk_softmax: ix_bridge → ixformer::infer::topk_softmax() LOADED")
return
except Exception as e:
logger.info("topk_softmax: ix_bridge unavailable (%s)", e)
# Also try direct import from workspace
try:
import sys
for p in ['/workspace/ex_engine/python', '/workspace/ex_engine',
'/usr/local/corex/lib/python3/dist-packages/ex_engine/python']:
if p not in sys.path:
sys.path.insert(0, p)
from ix_bridge import is_available, topk_softmax as _ix_ts
if is_available():
_ix_bridge_mod = True
logger.info("topk_softmax: ix_bridge (direct) → ixformer::infer LOADED")
return
except Exception as e:
logger.info("topk_softmax: ix_bridge direct import failed (%s)", e)
def _init_moe_topk():
global _moe_topk_ext, _moe_topk_init_done
_moe_topk_init_done = True
# 0. Try ix_bridge first (calls ixformer C++ SDK directly)
_init_ix_bridge()
if _ix_bridge_mod:
return # ix_bridge loaded, no need for CUDA kernel
# 1. Try import precompiled module (torch cache from Docker build)
try:
import moe_topk_softmax_v3 as ext
@@ -1038,6 +1073,21 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
if not _moe_topk_init_done:
_init_moe_topk()
# Priority 0: ix_bridge → ixformer::infer::topk_softmax() (fastest, uses SDK)
if _ix_bridge_mod:
try:
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk
gating = gating_output if isinstance(gating_output, torch.Tensor) else gating_output
topk_k = topk_weights.shape[1]
weights, ids = _ix_topk(gating, topk_k, renormalize=False)
topk_weights.copy_(weights.to(topk_weights.dtype))
topk_ids.copy_(ids.to(topk_ids.dtype))
# token_expert_indicies not produced by ix_bridge, fill with topk_ids
token_expert_indicies.copy_(ids.to(token_expert_indicies.dtype))
return
except Exception as e:
logger.warning("topk_softmax ix_bridge failed (%s), trying CUDA kernel", e)
# Priority 1: Our CUDA kernel (fused warp-shuffle, ~5x faster than PyTorch)
if _moe_topk_ext is not None:
try: