feat(EX): wire xllm CUB topk_softmax kernel into MoE routing

Upstream: xllm/kernels/cuda/moe/moe_topk_softmax_kernels.cuh (Apache 2.0)
Adapted: CHECK→TORCH_CHECK, include path fix, cuda/functional guard, pybind11

Call chain now:
  qwen3_5.py:_pure_pytorch_experts()
    → _ex_moe_topk_softmax (fused CUB kernel, 1 launch)
    → fallback: torch.softmax + torch.topk (3 launches)

Files:
  ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh — xllm kernel (adapted)
  ex_engine/csrc/moe/device_utils.cuh — xllm device utils
  ex_engine/csrc/moe/moe_topk_softmax_ext.cu — pybind11 wrapper
  ex_engine/python/moe_topk.py — JIT loader (same pattern as flash_qla_sm70)
  qwen3_5.py — import + use in _pure_pytorch_experts()
  patch_ops.sh — deploy kernel sources for JIT

qwen3_6_scripts/patch_ops.sh

@@ -183,19 +183,24 @@ EX_ENGINE_SRC="/workspace/ex_engine"
 if [ -d "$EX_ENGINE_SRC/python" ]; then
     # Deploy into vllm's model dir so qwen3_5.py can import it
     EX_DST="$VLLM/model_executor/models/ex_engine"
-    mkdir -p "$EX_DST/python"
-    mkdir -p "$EX_DST/csrc"
+    mkdir -p "$EX_DST/python" "$EX_DST/csrc"
     cp "$EX_ENGINE_SRC/python/"*.py "$EX_DST/python/" 2>/dev/null || true
-    # ix_moe_bridge.cpp needs to be next to the python module for JIT compile
+    # ix_moe_bridge.cpp for JIT compile
     cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/csrc/" 2>/dev/null || true
     cp "$EX_ENGINE_SRC/csrc/ix_moe_bridge.cpp" "$EX_DST/python/" 2>/dev/null || true
-    # Also make ex_engine importable from Python path
     touch "$EX_DST/__init__.py"
     touch "$EX_DST/python/__init__.py"
-    # Copy built .so files if they exist
+    # Copy built .so files
     if [ -d "$EX_ENGINE_SRC/build" ]; then
         cp "$EX_ENGINE_SRC/build/"*.so "$EX_DST/" 2>/dev/null || true
     fi
+    # Deploy MoE CUDA kernel sources for JIT compilation
+    if [ -d "$EX_ENGINE_SRC/csrc/moe" ]; then
+        mkdir -p "$EX_DST/csrc/moe"
+        cp "$EX_ENGINE_SRC/csrc/moe/"*.cu "$EX_DST/csrc/moe/" 2>/dev/null || true
+        cp "$EX_ENGINE_SRC/csrc/moe/"*.cuh "$EX_DST/csrc/moe/" 2>/dev/null || true
+        echo "[patch_ops] MoE CUDA kernel sources deployed for JIT"
+    fi
     echo "[patch_ops] EX Engine deployed to $EX_DST"
     ls -la "$EX_DST/csrc/" 2>/dev/null || true
     if [ -n "$VLLM2" ]; then
@@ -207,6 +212,19 @@ else
     echo "[patch_ops] WARNING: EX Engine not found — MoE uses slow PyTorch fallback"
 fi
 
+# Also deploy ex_engine Python package to system path for direct import
+EX_PY_DST="/usr/local/corex/lib/python3/dist-packages/ex_engine"
+if [ -d "$EX_ENGINE_SRC/python" ]; then
+    mkdir -p "$EX_PY_DST"
+    cp "$EX_ENGINE_SRC/python/"*.py "$EX_PY_DST/" 2>/dev/null || true
+    if [ -d "$EX_ENGINE_SRC/csrc/moe" ]; then
+        mkdir -p "$EX_PY_DST/../ex_engine/csrc/moe"
+        cp "$EX_ENGINE_SRC/csrc/moe/"*.cu "$EX_PY_DST/../ex_engine/csrc/moe/" 2>/dev/null || true
+        cp "$EX_ENGINE_SRC/csrc/moe/"*.cuh "$EX_PY_DST/../ex_engine/csrc/moe/" 2>/dev/null || true
+    fi
+    echo "[patch_ops] EX Engine Python package deployed to $EX_PY_DST"
+fi
+
 echo "[patch_ops] DONE — EX Engine + SM70 GDN kernel + serving layer deployed"
 echo "[patch_ops] Deployed: qwen3_5.py, flash_qla_sm70, ex_engine factors, paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, serving layer"
 echo "[patch_ops] EX factors replace: vllm_moe_topk_softmax (2304 calls/token), gdn_chunk_fwd (NaN fix)"

qwen3_6_scripts/qwen3_5.py

@@ -127,6 +127,21 @@ try:
 except ImportError:
     pass
 
+# EX Engine: fused MoE topk_softmax CUDA kernel (xllm CUB-based)
+_ex_moe_topk_softmax = None
+_ex_moe_topk_available = False
+try:
+    from ex_engine.python.moe_topk import moe_topk_softmax as _ex_moe_topk_softmax
+    _ex_moe_topk_available = True
+    logger.info("EX Engine MoE topk_softmax kernel available")
+except ImportError:
+    try:
+        from vllm.model_executor.models.ex_engine.moe_topk import moe_topk_softmax as _ex_moe_topk_softmax
+        _ex_moe_topk_available = True
+        logger.info("EX Engine MoE topk_softmax kernel available (vllm path)")
+    except ImportError:
+        pass
+
 
 # ---------------------------------------------------------------------------
 # ixformer-accelerated ops (drop-in replacements for torch ops)
@@ -1065,9 +1080,21 @@ class Qwen3_5MoeSparseBlock(nn.Module):
         Output is partial (pre-all-reduce), same contract as FusedMoE
         with reduce_results=False.
         """
-        # Routing: fused topk+softmax via ixformer C++ bridge (if available)
-        # Falls back to PyTorch softmax → topk → renormalize
-        if _ix_bridge_available:
+        # Routing: fused topk+softmax dispatch chain
+        # Tier 1: EX Engine CUB kernel → Tier 2: ix_bridge → Tier 3: PyTorch
+        if _ex_moe_topk_available:
+            T_tok = router_logits.shape[0]
+            topk_weights = torch.empty(T_tok, self.top_k, dtype=torch.float32,
+                                       device=router_logits.device)
+            topk_ids = torch.empty(T_tok, self.top_k, dtype=torch.int32,
+                                   device=router_logits.device)
+            token_expert_indices = torch.empty(T_tok, self.top_k, dtype=torch.int32,
+                                               device=router_logits.device)
+            _ex_moe_topk_softmax(topk_weights, topk_ids, token_expert_indices,
+                                 router_logits.float(), True)
+            topk_ids = topk_ids.to(torch.long)
+            topk_weights = topk_weights.to(hidden_states.dtype)
+        elif _ix_bridge_available:
             topk_weights, topk_ids = _ix_topk_softmax(
                 router_logits, self.top_k, renormalize=True)
             topk_weights = topk_weights.to(hidden_states.dtype)