From ee516bd20693b7b046fc21e937128d4e0159fa2e Mon Sep 17 00:00:00 2001 From: dev Date: Mon, 17 Aug 2026 06:56:11 +0000 Subject: [PATCH] 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. --- Dockerfile | 8 +- build_moe_bridge.sh | 21 +++-- ex_engine/build_moe_bridge.sh | 165 ++++++++++++++++++++++++++++++++++ qwen3_6_scripts/patch_ops.sh | 78 +++++++++++++--- 4 files changed, 247 insertions(+), 25 deletions(-) create mode 100644 ex_engine/build_moe_bridge.sh diff --git a/Dockerfile b/Dockerfile index bc2b17e7..ec4625aa 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 ; \ diff --git a/build_moe_bridge.sh b/build_moe_bridge.sh index 32ff7ef5..6caf15a5 100644 --- 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}" diff --git a/ex_engine/build_moe_bridge.sh b/ex_engine/build_moe_bridge.sh new file mode 100644 index 00000000..6caf15a5 --- /dev/null +++ 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" \ No newline at end of file diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index ac007739..05ada242 100755 --- 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