feat(MoE): wire full ix_fused_moe_forward as Tier 0 dispatch

ix_bridge.py: expose all 6 ixformer::infer functions + fused_moe_forward()
qwen3_5.py: 4-tier MoE dispatch (fused C++ → CUB topk → ix topk → PyTorch)
patch_ops.sh: deploy ix_moe_bridge.cpp to 4 search paths for JIT

ex_engine/python/ix_bridge.py

@@ -1,71 +1,179 @@
 """
-ix_bridge.py — Load ix_moe_bridge C++ extension at runtime.
+ix_bridge.py — Full ixformer MoE pipeline bridge.
 
-Calls ixformer::infer::topk_softmax() via C++ torch extension,
-bypassing the missing Python binding in ixformer.functions.
+Loads ix_moe_bridge.so via JIT and exposes both individual ops and the full
+fused MoE forward pass that replaces the Python for-loop in qwen3_5.py.
 
-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)
+Pipeline (mirrors xllm/core/layers/ilu/fused_moe.cpp):
+  topk_softmax → moe_gen_idx → moe_expand_input → group_gemm(w13)
+  → silu_and_mul → group_gemm(w2) → moe_combine_result
+
+All 6 ixformer::infer C++ functions are called through ix_moe_bridge.cpp
+which forward-declares them and links against the base image SDK.
 """
 
 import os
 import logging
 import torch
+from typing import Tuple, Optional, List
 
 logger = logging.getLogger("ex_engine.ix_bridge")
 
 _ix_bridge = None
+_ix_bridge_loaded = False  # True after attempt, even if failed
 _ix_bridge_available = False
 
+
+def _find_cpp_source():
+    """Find ix_moe_bridge.cpp in multiple locations."""
+    candidates = []
+    # 1. Relative to this file: ex_engine/csrc/
+    here = os.path.dirname(os.path.abspath(__file__))
+    candidates.append(os.path.join(here, "..", "csrc", "ix_moe_bridge.cpp"))
+    # 2. Deployed path inside vllm model dir
+    candidates.append(os.path.join(here, "ix_moe_bridge.cpp"))
+    # 3. /workspace paths
+    candidates.append("/workspace/ex_engine/csrc/ix_moe_bridge.cpp")
+    candidates.append("/workspace/qwen3_6_scripts/ix_moe_bridge.cpp")
+
+    for c in candidates:
+        p = os.path.normpath(c)
+        if os.path.exists(p):
+            return p
+    return None
+
+
 def _load_bridge():
-    """JIT-compile and load ix_moe_bridge.so"""
-    global _ix_bridge, _ix_bridge_available
-    if _ix_bridge is not None:
+    """JIT-compile and load ix_moe_bridge.so — called once."""
+    global _ix_bridge, _ix_bridge_loaded, _ix_bridge_available
+    if _ix_bridge_loaded:
         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
+    _ix_bridge_loaded = True
+
+    cpp_file = _find_cpp_source()
+    if cpp_file is None:
+        logger.warning("ix_moe_bridge.cpp not found in any search path")
         return False
-    
+
     try:
         from torch.utils.cpp_extension import load
-        logger.info("JIT-compiling ix_moe_bridge.cpp ...")
+        logger.info("JIT-compiling ix_moe_bridge.cpp from %s ...", cpp_file)
         _ix_bridge = load(
             name="ix_moe_bridge",
             sources=[cpp_file],
-            extra_cflags=["-O2"],
+            extra_cflags=["-O2", "-std=c++17"],
             verbose=False,
         )
         _ix_bridge_available = True
-        logger.info("ix_moe_bridge loaded successfully: %s", dir(_ix_bridge))
+        fns = [x for x in dir(_ix_bridge) if not x.startswith("_")]
+        logger.info("ix_moe_bridge loaded: %s", fns)
         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):
+
+def is_available() -> bool:
+    """Check if bridge is available (lazy-load on first call)."""
+    if not _ix_bridge_loaded:
+        _load_bridge()
+    return _ix_bridge_available
+
+
+# =========================================================================
+# Individual ops (thin wrappers with type safety)
+# =========================================================================
+
+def topk_softmax(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    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.
+    Fused topk+softmax via ixformer::infer::topk_softmax.
+    Returns: (topk_weights [T, K] fp32, topk_ids [T, K] int32)
     """
-    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."
-            )
-    
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available — JIT compile failed")
     return _ix_bridge.topk_softmax(gating_output, topk, renormalize)
+
+
+def moe_gen_idx(
+    expert_id: torch.Tensor,
+    expert_num: int,
+) -> List[torch.Tensor]:
+    """
+    Build expert permutation maps.
+    Returns: [src_dst, dst_src, expert_sizes, cumsum]
+    """
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.moe_gen_idx(expert_id, expert_num)
+
+
+def moe_expand_input(
+    input: torch.Tensor,
+    gather_index: torch.Tensor,
+    combine_idx: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    """Gather tokens by expert assignment."""
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.moe_expand_input(input, gather_index, combine_idx, topk)
+
+
+def group_gemm(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    token_count: torch.Tensor,
+    output_n: int,
+) -> torch.Tensor:
+    """Batched expert GEMM via ixformer."""
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.group_gemm(inputs, weights, token_count, output_n)
+
+
+def silu_and_mul(input: torch.Tensor) -> torch.Tensor:
+    """Fused SiLU gate activation."""
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.silu_and_mul(input)
+
+
+def moe_combine_result(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+) -> torch.Tensor:
+    """Weighted reduce for MoE output."""
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.moe_combine_result(input, weight)
+
+
+# =========================================================================
+# Full fused MoE forward — replaces _pure_pytorch_experts() entirely
+# =========================================================================
+
+def fused_moe_forward(
+    hidden_states: torch.Tensor,    # (T, H)
+    router_logits: torch.Tensor,    # (T, E)
+    w13: torch.Tensor,              # (E, 2*I, H) gate_up
+    w2: torch.Tensor,               # (E, H, I) down
+    topk: int,
+    num_experts: int,
+    renormalize: bool = True,
+) -> torch.Tensor:
+    """
+    Full fused MoE forward via ixformer C++ pipeline.
+
+    Pipeline: topk → gen_idx → expand → gemm1(w13) → silu → gemm2(w2) → combine
+
+    Returns: (T, H) — partial output, needs all-reduce after.
+    """
+    if not is_available():
+        raise RuntimeError("ix_moe_bridge not available")
+    return _ix_bridge.fused_moe_forward(
+        hidden_states, router_logits, w13, w2, topk, num_experts, renormalize
+    )