fix: connect ex_engine to Docker build pipeline

Path breaks fixed: 1. Dockerfile: COPY entire ex_engine/ instead of 5 individual files 2. patch_ops.sh: EX_ENGINE_DIR fallback to /workspace/ex_engine 3. build_moe_bridge.sh: search csrc/ under both SCRIPT_DIR and SCRIPT_DIR/ex_engine New build steps added to patch_ops.sh: - Build gemm_grouped.so (CUTLASS Cu10 grouped GEMM, verified on device) - Build corex_batched_gemm.so (CUTLASS batched GEMM for decode) - Build ix_moe_bridge.so (7-step fused MoE pipeline) - Deploy all ex_engine/python/*.py to VLLM_ROOT/ex_engine/python/ Previously broken: ix_ops.py, patch_vllm_ops.py, ix_startup_patch.py were never deployed because EX_ENGINE_DIR resolved to nonexistent path. xllm_activation/norm/rope/cache.so were copied but had no Python caller.
2026-08-17 06:56:11 +00:00
parent 461addf428
commit ee516bd206
4 changed files with 247 additions and 25 deletions
--- a/8
+++ b/8
@@ -4,12 +4,8 @@ WORKDIR /workspace/
 # Copy all our engine patches
 COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
 COPY ./computility-run.yaml /workspace/computility-run.yaml
-# Copy ex_engine source for MoE bridge compilation
-COPY ./ex_engine/csrc/moe_ops_impl.cu /workspace/qwen3_6_scripts/ex_engine_src/csrc/moe_ops_impl.cu
-COPY ./ex_engine/csrc/ix_full_bridge_v2.cpp /workspace/qwen3_6_scripts/ex_engine_src/csrc/ix_full_bridge_v2.cpp
-COPY ./ex_engine/build_moe_bridge.sh /workspace/qwen3_6_scripts/ex_engine_src/build_moe_bridge.sh
-COPY ./ex_engine/python/moe_dispatch.py /workspace/qwen3_6_scripts/ex_engine_src/python/moe_dispatch.py
-COPY ./ex_engine/python/patch_moe_hot_path.py /workspace/qwen3_6_scripts/ex_engine_src/python/patch_moe_hot_path.py
+# Copy entire ex_engine — python dispatch, csrc, build scripts, headers
+COPY ./ex_engine /workspace/ex_engine
 # Make patch script executable and run it
 RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
    bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
--- a/build_moe_bridge.sh
+++ b/build_moe_bridge.sh
@@ -18,17 +18,26 @@ echo "[moe_bridge] Building ix_moe_bridge.so"
 echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"

 # --- Locate sources ---
-MOE_CU="${SCRIPT_DIR}/ex_engine/csrc/moe_ops_impl.cu"
-BRIDGE_CPP="${SCRIPT_DIR}/ex_engine/csrc/ix_full_bridge_v2.cpp"
+# Support both layouts:
+#   1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
+#   2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
+MOE_CU=""
+BRIDGE_CPP=""
+for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
+    [[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
+    [[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
+done

-if [[ ! -f "$MOE_CU" ]]; then
-    echo "[moe_bridge] ERROR: $MOE_CU not found" >&2
+if [[ -z "$MOE_CU" ]]; then
+    echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
    exit 1
 fi
-if [[ ! -f "$BRIDGE_CPP" ]]; then
-    echo "[moe_bridge] ERROR: $BRIDGE_CPP not found" >&2
+if [[ -z "$BRIDGE_CPP" ]]; then
+    echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
    exit 1
 fi
+echo "[moe_bridge] MOE_CU: ${MOE_CU}"
+echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"

 # --- Locate libraries ---
 COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
--- a/ex_engine/build_moe_bridge.sh
+++ b/ex_engine/build_moe_bridge.sh
@@ -0,0 +1,165 @@
+#!/usr/bin/env bash
+# build_moe_bridge.sh — Compile MoE ops + bridge into ix_moe_bridge.so
+#
+# Links against:
+#   libcuinfer.so   (cuinferCustomGemm, cuinferTopK — confirmed in symbol dump)
+#   libixformer.so  (silu_and_mul, rms_norm, flash_attn, etc — confirmed)
+#
+# Real device compiler: corex clang/16, NOT nvcc
+# Reference: ex_engine/build_ix_bridge.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+VLLM_ROOT="${1:-}"
+
+echo "[moe_bridge] Building ix_moe_bridge.so"
+echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
+
+# --- Locate sources ---
+# Support both layouts:
+#   1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
+#   2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
+MOE_CU=""
+BRIDGE_CPP=""
+for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
+    [[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
+    [[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
+done
+
+if [[ -z "$MOE_CU" ]]; then
+    echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+if [[ -z "$BRIDGE_CPP" ]]; then
+    echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+echo "[moe_bridge] MOE_CU: ${MOE_CU}"
+echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"
+
+# --- Locate libraries ---
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
+
+# Find libcuinfer.so
+CUINFER_SO=""
+for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib" "/usr/lib64" "/usr/lib"; do
+    if [[ -f "${d}/libcuinfer.so" ]]; then
+        CUINFER_SO="${d}/libcuinfer.so"
+        break
+    fi
+done
+
+# Find libixformer.so and ixformer Python package
+IX_LIB_DIR=""
+IX_SO_FILES=()
+for d in \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \
+    "$(python3 -c 'import ixformer, os; print(os.path.dirname(ixformer.__file__))' 2>/dev/null || echo '')"; do
+    if [[ -d "$d" ]]; then
+        IX_LIB_DIR="$d"
+        while IFS= read -r so; do
+            IX_SO_FILES+=("$so")
+        done < <(find "$d" -name "*.so" -type f 2>/dev/null)
+        break
+    fi
+done
+
+echo "[moe_bridge] COREX_ROOT: ${COREX_ROOT}"
+echo "[moe_bridge] cuinfer: ${CUINFER_SO:-NOT FOUND}"
+echo "[moe_bridge] ixformer dir: ${IX_LIB_DIR:-NOT FOUND}"
+echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}"
+
+# --- Build via torch.utils.cpp_extension ---
+mkdir -p "${SCRIPT_DIR}/prebuilt"
+
+python3 << 'PYEOF'
+import os, sys, glob, shutil
+
+script_dir = os.environ.get("SCRIPT_DIR", ".")
+vllm_root = os.environ.get("VLLM_ROOT", "")
+
+moe_cu = os.path.join(script_dir, "ex_engine", "csrc", "moe_ops_impl.cu")
+bridge_cpp = os.path.join(script_dir, "ex_engine", "csrc", "ix_full_bridge_v2.cpp")
+
+# Collect linker flags
+extra_ldflags = []
+rpath_dirs = set()
+
+corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
+for search_dir in [
+    os.path.join(corex_root, "lib64"),
+    os.path.join(corex_root, "lib"),
+]:
+    if os.path.isdir(search_dir):
+        rpath_dirs.add(search_dir)
+        for so in glob.glob(os.path.join(search_dir, "libcuinfer*.so*")):
+            extra_ldflags.append(so)
+
+# ixformer .so files
+try:
+    import ixformer
+    ix_dir = os.path.dirname(ixformer.__file__)
+    rpath_dirs.add(ix_dir)
+    for so in glob.glob(os.path.join(ix_dir, "*.so")):
+        extra_ldflags.append(so)
+    for so in glob.glob(os.path.join(ix_dir, "lib*.so")):
+        if so not in extra_ldflags:
+            extra_ldflags.append(so)
+except ImportError:
+    # Search common paths
+    for d in [
+        os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
+        os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
+    ]:
+        if os.path.isdir(d):
+            rpath_dirs.add(d)
+            for so in glob.glob(os.path.join(d, "*.so")):
+                extra_ldflags.append(so)
+
+for d in rpath_dirs:
+    extra_ldflags.append(f"-Wl,-rpath,{d}")
+
+print(f"[moe_bridge] Linking against {len(extra_ldflags)} items")
+for f in extra_ldflags[:10]:
+    print(f"  {f}")
+
+try:
+    from torch.utils.cpp_extension import load
+
+    mod = load(
+        name="ix_moe_bridge",
+        sources=[moe_cu, bridge_cpp],
+        extra_include_paths=[os.path.join(script_dir, "csrc")],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_cuda_cflags=["-O2", ],
+        extra_ldflags=extra_ldflags,
+        verbose=True,
+    )
+    print("[moe_bridge] ✓ Compilation successful")
+
+    # Find and copy the built .so
+    import importlib
+    spec = importlib.util.find_spec("ix_moe_bridge")
+    if spec and spec.origin:
+        dst = os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so")
+        shutil.copy2(spec.origin, dst)
+        print(f"[moe_bridge] ✓ Saved to {dst}")
+
+        if vllm_root:
+            vllm_dst = os.path.join(vllm_root, "ex_engine", "ix_moe_bridge.so")
+            os.makedirs(os.path.dirname(vllm_dst), exist_ok=True)
+            shutil.copy2(spec.origin, vllm_dst)
+            print(f"[moe_bridge] ✓ Deployed to {vllm_dst}")
+    else:
+        print("[moe_bridge] ⚠ Could not locate compiled .so via importlib")
+
+except Exception as e:
+    print(f"[moe_bridge] ERROR: {e}", file=sys.stderr)
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+PYEOF
+
+echo "[moe_bridge] Done"
--- a/qwen3_6_scripts/patch_ops.sh
+++ b/qwen3_6_scripts/patch_ops.sh
@@ -202,8 +202,9 @@ fi
 # --- Deploy ix_bridge Python integration layer --------------------------------
 build_stage "deploying ix_bridge operator replacements"
 EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")"
-if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR" ]; then
-    EX_ENGINE_DIR="$(cd "$(dirname "$0")" && pwd)/../ex_engine"
+if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR/python" ]; then
+    # Dockerfile puts ex_engine at /workspace/ex_engine
+    EX_ENGINE_DIR="/workspace/ex_engine"
 fi

 if [ -d "$EX_ENGINE_DIR/python" ]; then
@@ -332,22 +333,73 @@ if b"max_completion_tokens" not in installed:
    raise SystemExit("protocol.py missing max_completion_tokens field")
 PY

+build_stage "building CUTLASS grouped GEMM (gemm_grouped.so)"
+if [[ -f "${EX_ENGINE_DIR}/build_gemm_grouped.sh" ]]; then
+    bash "${EX_ENGINE_DIR}/build_gemm_grouped.sh" 2>&1 || {
+        echo "[WARN] gemm_grouped build failed — will use torch.mm fallback"
+    }
+    # Deploy compiled .so if it exists
+    for so in "${EX_ENGINE_DIR}"/gemm_grouped.so "${EX_ENGINE_DIR}"/csrc/gemm_grouped.so; do
+        if [[ -f "$so" ]]; then
+            cp "$so" "${VLLM_ROOT}/gemm_grouped.so"
+            echo "[patch_ops] deployed gemm_grouped.so → ${VLLM_ROOT}/"
+            break
+        fi
+    done
+fi
+
+build_stage "building CUTLASS batched GEMM (corex_batched_gemm.so)"
+if [[ -f "${EX_ENGINE_DIR}/xllm_kernels/cuda/corex_batched_gemm_kernel.cu" ]]; then
+    python3 << PYEOF
+import os, sys, shutil
+try:
+    from torch.utils.cpp_extension import load
+    ex = "${EX_ENGINE_DIR}"
+    cutlass_inc = ""
+    for d in ["/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass/include",
+              "/usr/local/corex/include/cutlass", "/usr/include/cutlass"]:
+        if os.path.isdir(d):
+            cutlass_inc = d
+            break
+    if not cutlass_inc:
+        print("[batched_gemm] No cutlass headers — skip"); sys.exit(0)
+    mod = load(
+        name="corex_batched_gemm",
+        sources=[
+            os.path.join(ex, "xllm_kernels/cuda/corex_batched_gemm_kernel.cu"),
+            os.path.join(ex, "xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp"),
+        ],
+        extra_include_paths=[cutlass_inc],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_cuda_cflags=["-O2", f"-I{cutlass_inc}"],
+        extra_ldflags=["/usr/local/corex/lib64/libcuinfer.so", "-Wl,-rpath,/usr/local/corex/lib64"],
+        verbose=False,
+    )
+    print("[batched_gemm] ✓ Compiled")
+    import importlib
+    spec = importlib.util.find_spec("corex_batched_gemm")
+    if spec and spec.origin:
+        shutil.copy2(spec.origin, "${VLLM_ROOT}/corex_batched_gemm.so")
+        print("[batched_gemm] ✓ Deployed to ${VLLM_ROOT}/")
+except Exception as e:
+    print(f"[batched_gemm] WARN: {e}")
+PYEOF
+fi
+
 build_stage "building MoE bridge (ix_moe_bridge.so)"
-if [[ -f "./ex_engine_src/build_moe_bridge.sh" ]]; then
-    bash ./ex_engine_src/build_moe_bridge.sh "${VLLM_ROOT}" 2>&1 || {
+if [[ -f "${EX_ENGINE_DIR}/csrc/ix_moe_bridge.cpp" ]]; then
+    SCRIPT_DIR="${EX_ENGINE_DIR}" bash "${EX_ENGINE_DIR}/build_moe_bridge.sh" "${VLLM_ROOT}" 2>&1 || {
        echo "[WARN] MoE bridge build failed — will use Python fallback"
    }
 fi

-build_stage "deploying MoE dispatch modules"
-EX_DIR="${VLLM_ROOT}/ex_engine/python"
-mkdir -p "${EX_DIR}"
-for pyfile in moe_dispatch.py patch_moe_hot_path.py; do
-    if [[ -f "./ex_engine_src/python/${pyfile}" ]]; then
-        cp "./ex_engine_src/python/${pyfile}" "${EX_DIR}/${pyfile}"
-        echo "  ✓ ${pyfile}"
-    fi
-done
+build_stage "deploying all ex_engine Python modules"
+EX_PY_DIR="${VLLM_ROOT}/ex_engine/python"
+mkdir -p "${EX_PY_DIR}"
+if [[ -d "${EX_ENGINE_DIR}/python" ]]; then
+    cp "${EX_ENGINE_DIR}/python/"*.py "${EX_PY_DIR}/" 2>/dev/null
+    echo "[patch_ops] deployed $(ls -1 "${EX_PY_DIR}"/*.py 2>/dev/null | wc -l) Python modules → ${EX_PY_DIR}/"
+fi

 build_stage "compiling submission Python sources"
 find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile