From be4d66119198340981856eda47223d2d76a47108 Mon Sep 17 00:00:00 2001 From: dev Date: Mon, 17 Aug 2026 07:04:50 +0000 Subject: [PATCH] =?UTF-8?q?revert:=20undo=202=20premature=20pushes=20(ee51?= =?UTF-8?q?6bd2,=209ea0a1d4)=20=E2=80=94=20code=20needs=20review=20first?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dockerfile | 8 +- EX_ENGINE_INFRA_AUDIT.md | 172 -------------------------------- build_moe_bridge.sh | 39 ++------ ex_engine/build_moe_bridge.sh | 179 ---------------------------------- qwen3_6_scripts/patch_ops.sh | 118 ++++------------------ qwen3_6_scripts/qwen3_5.py | 63 +++--------- 6 files changed, 48 insertions(+), 531 deletions(-) delete mode 100644 EX_ENGINE_INFRA_AUDIT.md delete mode 100644 ex_engine/build_moe_bridge.sh diff --git a/Dockerfile b/Dockerfile index ec4625aa..bc2b17e7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,8 +4,12 @@ WORKDIR /workspace/ # Copy all our engine patches COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts COPY ./computility-run.yaml /workspace/computility-run.yaml -# Copy entire ex_engine — python dispatch, csrc, build scripts, headers -COPY ./ex_engine /workspace/ex_engine +# 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 # 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/EX_ENGINE_INFRA_AUDIT.md b/EX_ENGINE_INFRA_AUDIT.md deleted file mode 100644 index eedb4af2..00000000 --- a/EX_ENGINE_INFRA_AUDIT.md +++ /dev/null @@ -1,172 +0,0 @@ -# EX_ENGINE 基建盘点 - -日期: 2026-08-17 -基于: commit 512f384a (CUTLASS Cu10 grouped GEMM 真机验证) + 后续 revert - ---- - -## 一、Docker容器里实际运行的状态 - -### 进了Docker且工作正常的 (prebuilt .so) - -| .so 文件 | qwen3_5.py flag | 状态 | -|---|---|---| -| corex_gdn_causal_conv.so | _USE_COREX_GDN_CAUSAL_CONV | ✓ 工作 | -| corex_gdn_gated_norm.so | _USE_COREX_GDN_GATED_NORM | ✓ 工作 | -| corex_gdn_beta_decay.so | _USE_COREX_GDN_BETA_DECAY | ✓ 工作 | -| corex_gdn_qk_map.so | _USE_COREX_GDN_QK_MAP | ✓ 工作 | -| corex_gdn_packed_decode.so | _USE_COREX_GDN_PACKED_DECODE | ✓ 工作 | -| corex_gdn_chunk_recurrent.so | _HAS_COREX_GDN_CHUNK | ✓ 工作 | -| corex_moe_topk_softmax.so | _USE_COREX_MOE_TOPK_SOFTMAX | ✓ 工作 | -| corex_moe_direct_routed.so | _USE_COREX_MOE_DIRECT_ROUTED | ✓ 工作 | -| corex_moe_weight_gather.so | _USE_COREX_MOE_WEIGHT_GATHER | ✓ 工作 | -| corex_moe_exact_reduce.so | _USE_COREX_MOE_EXACT_REDUCE | ✓ 工作 | -| corex_moe_index_combine.so | _USE_COREX_MOE_INDEX_COMBINE | ✓ 工作 | -| corex_attn_head_rms_norm.so | _USE_COREX_ATTN_HEAD_RMS_NORM | ✓ 工作 | -| corex_block_major_kv_transfer.so | (block_major_kv_cache.py用) | ✓ 工作 | -| corex_paged_kv_gather.so | (paged_attn.py用) | ✓ 工作 | -| corex_fused_paged_prefill.so | (corex_fa2用) | ✓ 工作 | -| xllm_moe.so | _USE_XLLM_MOE | ✓ 工作 | -| xllm_activation.so | (patch_vllm_ops用) | ? 见下 | -| xllm_norm.so | (patch_vllm_ops用) | ? 见下 | -| xllm_rope.so | (patch_vllm_ops用) | ? 见下 | -| xllm_cache.so | (patch_vllm_ops用) | ? 见下 | -| ix_full_bridge.so | (ix_ops.py用) | ? 见下 | - -### 进了Docker但断裂的 - -| 文件 | 问题 | -|---|---| -| ix_full_bridge.so (v1) | 已cp到$VLLM_ROOT/, 但python wrapper ix_ops.py没部署 | -| xllm_activation/norm/rope/cache.so | 已cp到$VLLM_ROOT/, 但patch_vllm_ops.py没部署，没hook | -| ix_fused_moe.py | 已cp到models/, 但找不到ix_moe_bridge.so → _HAS_IX_FUSED_MOE=False | - -### 没进Docker的关键文件 - -| 文件 | 功能 | 行数 | -|---|---|---| -| ex_engine/python/ix_ops.py | ix_full_bridge.so的Python wrapper | 343 | -| ex_engine/python/ix_ops_dispatch.py | 统一op dispatch (bridge→ixformer→raise) | 407 | -| ex_engine/python/patch_vllm_ops.py | monkey-patch vllm的silu/rms_norm/rope/cache | 201 | -| ex_engine/python/corex_moe.py | corex MoE pipeline wrapper | 237 | -| ex_engine/python/corex_gdn.py | corex GDN ops wrapper | 256 | -| ex_engine/python/corex_fa2.py | corex FlashAttn dispatch | 279 | -| ex_engine/python/corex_fa2_dispatch.py | FA2 3-mode dispatch | 231 | -| ex_engine/python/fused_moe_ilu.py | 7-step MoE pipeline (Python) | 205 | -| ex_engine/python/gemm_dispatch.py | GEMM dispatch (cutlass/cuinfer/torch) | 180 | -| ex_engine/csrc/gemm_grouped.cu | ✓ 真机验证的CUTLASS grouped GEMM | 188 | -| ex_engine/csrc/gemm_grouped_bind.cpp | pybind11 binding | 182 | -| ex_engine/build_gemm_grouped.sh | 编译脚本 | — | -| ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu | CUTLASS batched GEMM | 67 | -| ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp | binding | 129 | - ---- - -## 二、两个路径断裂的根因 - -### 断裂1: patch_ops.sh 找不到 ex_engine/python/ - -patch_ops.sh 第204行: -```bash -EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")" -``` - -Docker容器里的目录结构: -``` -/workspace/ -├── qwen3_6_scripts/ ← patch_ops.sh 在这里 -│ ├── patch_ops.sh -│ ├── ex_engine_src/ ← Dockerfile COPY进来的（只有5个文件） -│ │ ├── csrc/moe_ops_impl.cu -│ │ ├── csrc/ix_full_bridge_v2.cpp -│ │ ├── build_moe_bridge.sh -│ │ └── python/moe_dispatch.py, patch_moe_hot_path.py -│ └── prebuilt/corex-3.2.3-ivcore10/*.so -└── (没有 ex_engine/ 目录) -``` - -`$(dirname "$0")/../ex_engine` = `/workspace/ex_engine` → **不存在** - -结果: ix_ops.py, patch_vllm_ops.py, ix_startup_patch.py 全部没部署。 -xllm_activation/norm/rope/cache.so 虽然被cp到$VLLM_ROOT/但没有Python层调用它们。 - -### 断裂2: build_moe_bridge.sh 内部路径错误 - -build_moe_bridge.sh 第18-19行: -```bash -MOE_CU="${SCRIPT_DIR}/ex_engine/csrc/moe_ops_impl.cu" -BRIDGE_CPP="${SCRIPT_DIR}/ex_engine/csrc/ix_full_bridge_v2.cpp" -``` - -SCRIPT_DIR = `/workspace/qwen3_6_scripts/ex_engine_src` -实际路径 = `${SCRIPT_DIR}/csrc/moe_ops_impl.cu`（少了 `ex_engine/` 一层） - -结果: ix_moe_bridge.so 编译失败 → _HAS_IX_FUSED_MOE=False → 7-step fused MoE pipeline 未启用 - ---- - -## 三、ex_engine/ 文件去重审计 - -### 重复实现的功能（同一功能多个文件） - -**MoE topk softmax (5个文件做同一件事)**: -1. `xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh` (866行) ← xllm上游原版 -2. `csrc/moe/moe_topk_softmax_kernels.cuh` (855行) ← 几乎相同的拷贝 -3. `csrc/factor_moe_topk_softmax.cu` (260行) ← 独立提取版 -4. `csrc/moe/moe_topk_softmax_ext.cu` (55行) ← 另一个入口 -5. `csrc/moe_topk_softmax_v3.cu` (143行) ← 又一个版本 -6. prebuilt `corex_moe_topk_softmax.so` ← 已编译可用 -7. prebuilt `xllm_moe.so` ← 也包含此功能 - -**MoE combine (3个文件)**: -1. `xllm_kernels/cuda/moe/moe_combine.cu` (105行) ← xllm上游原版 -2. files_5 的 `factor_moe_combine.cu` ← 重写 -3. prebuilt `xllm_moe.so` ← 已编译可用 - -**MoE compute_index (2个文件)**: -1. `xllm_kernels/cuda/moe/moe_compute_index.cu` (156行) ← xllm上游原版 -2. files_5 的 `factor_moe_compute_index.cu` ← 重写 - -**MoE Python pipeline (3个文件)**: -1. `python/fused_moe_ilu.py` (205行) -2. `python/moe_dispatch.py` (171行) -3. files_5 的 `moe_pipeline.py` ← 重写 - -**Attention dispatch (2个文件)**: -1. `python/corex_fa2_dispatch.py` (231行) -2. files_5 的 `attn_dispatch.py` ← 重写 - -**C++ bridge (3个文件)**: -1. `csrc/ix_full_bridge.cpp` (90行) ← v1, 对应prebuilt ix_full_bridge.so -2. `csrc/ix_full_bridge_v2.cpp` (387行) ← v2, 没编译 -3. `csrc/ix_moe_bridge.cpp` (261行) ← MoE专用, 没编译 - -**ILU层 fused_moe (2处)**: -1. `xllm_layers/ilu/fused_moe.cpp` (797行) ← xllm上游 -2. `csrc/ilu_layer_fused_moe.cpp` (797行) ← 拷贝 - -### 上游cat但未修改的文件 - -| 目录 | 文件数 | 来源 | -|---|---|---| -| xllm_kernels/ilu/*.cpp | 7 | xllm上游ILU kernel接口 | -| xllm_layers/ilu/*.cpp | 2 | xllm上游ILU layer | -| xllm_layers/common/*.cpp | 5 | xllm上游common layer | -| xllm_layers/npu_torch/*.cpp | 10 | xllm上游NPU实现（不适用BI-V100） | -| xllm_layers/mlu/*.cpp | 4 | xllm上游MLU实现（不适用BI-V100） | -| xllm_models/*.h | 6 | xllm上游model定义 | -| moe/*.py | 14 | ds_vllm上游MoE模块 | -| fla_kernels/ | 6 | FLA库GDN kernel（Triton，BI-V100不能跑） | - ---- - -## 四、真机验证状态 - -| 组件 | commit | 真机结果 | -|---|---|---| -| gemm_grouped.cu (CUTLASS Cu10 TN) | cdcf1150 | ✓ err=0.000015 PASS, 1.97x vs torch.mm | -| corex_batched_gemm_kernel.cu | (sub 655用过) | ✓ decode单token 2.462ms | -| 16个prebuilt .so | (sub 694在用) | ✓ 正常加载 | -| ix_moe_bridge.so | 未编译 | ✗ 路径断裂 | -| ix_full_bridge_v2.so | 未编译 | ✗ 未进Docker | -| xllm_moe.so 的 fused_topk | (sub 694在用) | ✓ 正常工作 | diff --git a/build_moe_bridge.sh b/build_moe_bridge.sh index 1d749ac7..32ff7ef5 100644 --- a/build_moe_bridge.sh +++ b/build_moe_bridge.sh @@ -18,26 +18,17 @@ 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 +MOE_CU="${SCRIPT_DIR}/ex_engine/csrc/moe_ops_impl.cu" +BRIDGE_CPP="${SCRIPT_DIR}/ex_engine/csrc/ix_full_bridge_v2.cpp" -if [[ -z "$MOE_CU" ]]; then - echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2 +if [[ ! -f "$MOE_CU" ]]; then + echo "[moe_bridge] ERROR: $MOE_CU not found" >&2 exit 1 fi -if [[ -z "$BRIDGE_CPP" ]]; then - echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2 +if [[ ! -f "$BRIDGE_CPP" ]]; then + echo "[moe_bridge] ERROR: $BRIDGE_CPP not found" >&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}" @@ -75,28 +66,14 @@ echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}" # --- Build via torch.utils.cpp_extension --- mkdir -p "${SCRIPT_DIR}/prebuilt" -export SCRIPT_DIR VLLM_ROOT python3 << 'PYEOF' import os, sys, glob, shutil script_dir = os.environ.get("SCRIPT_DIR", ".") vllm_root = os.environ.get("VLLM_ROOT", "") -# Find source files — try direct csrc/ first, then ex_engine/csrc/ -moe_cu = "" -bridge_cpp = "" -for base in [script_dir, os.path.join(script_dir, "ex_engine")]: - candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu") - candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp") - if os.path.isfile(candidate_cu): - moe_cu = candidate_cu - if os.path.isfile(candidate_cpp): - bridge_cpp = candidate_cpp -if not moe_cu or not bridge_cpp: - print(f"[moe_bridge] ERROR: sources not found under {script_dir}") - sys.exit(1) -print(f"[moe_bridge] MOE_CU: {moe_cu}") -print(f"[moe_bridge] BRIDGE_CPP: {bridge_cpp}") +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 = [] diff --git a/ex_engine/build_moe_bridge.sh b/ex_engine/build_moe_bridge.sh deleted file mode 100644 index 1d749ac7..00000000 --- a/ex_engine/build_moe_bridge.sh +++ /dev/null @@ -1,179 +0,0 @@ -#!/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" - -export SCRIPT_DIR VLLM_ROOT -python3 << 'PYEOF' -import os, sys, glob, shutil - -script_dir = os.environ.get("SCRIPT_DIR", ".") -vllm_root = os.environ.get("VLLM_ROOT", "") - -# Find source files — try direct csrc/ first, then ex_engine/csrc/ -moe_cu = "" -bridge_cpp = "" -for base in [script_dir, os.path.join(script_dir, "ex_engine")]: - candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu") - candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp") - if os.path.isfile(candidate_cu): - moe_cu = candidate_cu - if os.path.isfile(candidate_cpp): - bridge_cpp = candidate_cpp -if not moe_cu or not bridge_cpp: - print(f"[moe_bridge] ERROR: sources not found under {script_dir}") - sys.exit(1) -print(f"[moe_bridge] MOE_CU: {moe_cu}") -print(f"[moe_bridge] BRIDGE_CPP: {bridge_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 1c24b056..ac007739 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -202,36 +202,19 @@ 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/python" ]; then - # Dockerfile puts ex_engine at /workspace/ex_engine - EX_ENGINE_DIR="/workspace/ex_engine" +if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR" ]; then + EX_ENGINE_DIR="$(cd "$(dirname "$0")" && pwd)/../ex_engine" fi if [ -d "$EX_ENGINE_DIR/python" ]; then - # Create ex_engine package inside vllm with correct Python package structure - mkdir -p "${VLLM_ROOT}/ex_engine/python" + # Create ex_engine package inside vllm mkdir -p "${VLLM_ROOT}/ex_engine/csrc" + echo '"""ex_engine — Algorithm factor replacement for BI-V100."""' > "${VLLM_ROOT}/ex_engine/__init__.py" - # __init__.py with re-exports so both import styles work: - # from ex_engine.python import ix_ops_dispatch (direct) - # from vllm.ex_engine import ix_ops_dispatch (via re-export) - cat > "${VLLM_ROOT}/ex_engine/__init__.py" << 'INIT_EOF' -"""ex_engine — Algorithm factor replacement for BI-V100.""" -# Re-export python subpackage members at top level for backward compat -# Allows: from vllm.ex_engine import ix_ops_dispatch -try: - from ex_engine.python.ix_ops_dispatch import * - from ex_engine.python import ix_ops_dispatch - from ex_engine.python import ix_ops - from ex_engine.python import patch_vllm_ops -except ImportError: - pass -INIT_EOF - echo '"""ex_engine.python — dispatch and bridge modules."""' > "${VLLM_ROOT}/ex_engine/python/__init__.py" - - # Deploy ALL Python modules - cp "$EX_ENGINE_DIR/python/"*.py "${VLLM_ROOT}/ex_engine/python/" - echo "[patch_ops] deployed $(ls -1 "${VLLM_ROOT}/ex_engine/python/"*.py | wc -l) modules → ${VLLM_ROOT}/ex_engine/python/" + # Deploy Python modules + cp "$EX_ENGINE_DIR/python/ix_ops.py" "${VLLM_ROOT}/ex_engine/ix_ops.py" + cp "$EX_ENGINE_DIR/python/patch_vllm_ops.py" "${VLLM_ROOT}/ex_engine/patch_vllm_ops.py" + echo "[patch_ops] deployed ix_ops.py + patch_vllm_ops.py → ${VLLM_ROOT}/ex_engine/" # Deploy bridge C++ source for JIT fallback for cpp in "$EX_ENGINE_DIR"/csrc/ix_full_bridge*.cpp "$EX_ENGINE_DIR"/csrc/ix_moe_bridge.cpp; do @@ -246,7 +229,7 @@ import logging _logger = logging.getLogger("ix_startup_patch") def apply(): try: - from vllm.ex_engine.python.patch_vllm_ops import apply_all_patches + from vllm.ex_engine.patch_vllm_ops import apply_all_patches n = apply_all_patches() if n > 0: _logger.info("ix_startup_patch: %d patches applied", n) @@ -349,83 +332,22 @@ 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_DIR}/csrc/ix_moe_bridge.cpp" ]]; then - SCRIPT_DIR="${EX_ENGINE_DIR}" bash "${EX_ENGINE_DIR}/build_moe_bridge.sh" "${VLLM_ROOT}" 2>&1 || { +if [[ -f "./ex_engine_src/build_moe_bridge.sh" ]]; then + bash ./ex_engine_src/build_moe_bridge.sh "${VLLM_ROOT}" 2>&1 || { echo "[WARN] MoE bridge build failed — will use Python fallback" } - # Deploy .so to all paths ix_fused_moe.py searches - for src in "${VLLM_ROOT}/ex_engine/ix_moe_bridge.so" \ - "${EX_ENGINE_DIR}/prebuilt/ix_moe_bridge.so"; do - if [[ -f "$src" ]]; then - cp "$src" "${VLLM_ROOT}/ix_moe_bridge.so" 2>/dev/null || true - cp "$src" "${VLLM_ROOT}/model_executor/models/ix_moe_bridge.so" 2>/dev/null || true - echo "[patch_ops] deployed ix_moe_bridge.so to vllm search paths" - break - fi - done fi -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 "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 "compiling submission Python sources" find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 15394e8e..130894f7 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -143,11 +143,6 @@ try: except ImportError: _corex_batched_gemm = None -try: - from vllm import gemm_grouped as _gemm_grouped -except ImportError: - _gemm_grouped = None - try: from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax except ImportError: @@ -217,11 +212,6 @@ _USE_COREX_MOE_DIRECT_ROUTED = ( _USE_COREX_BATCHED_GEMM = ( _corex_batched_gemm is not None and env_bool("BI100_MOE_BATCHED_GEMM", True)) -_USE_GEMM_GROUPED = ( - _gemm_grouped is not None - and env_bool("BI100_MOE_GEMM_GROUPED", True)) -if _USE_GEMM_GROUPED: - logger.info("gemm_grouped ENABLED — CUTLASS Cu10 grouped GEMM for MoE prefill") _USE_COREX_MOE_TOPK_SOFTMAX = ( _corex_moe_topk_softmax is not None and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True)) @@ -1894,46 +1884,21 @@ class Qwen3_5MoeSparseBlock(nn.Module): expert_counts = torch.bincount( flat_eids, minlength=w13.shape[0]).tolist() - # --- CUTLASS grouped GEMM path (replaces per-expert F.linear loop) --- - if _USE_GEMM_GROUPED and hidden_states.dtype == torch.float16: - # Sort tokens into expert order - sorted_hidden = hidden_states[sorted_tok_ids] # (T*topk, H) - expert_counts_t = torch.tensor( - expert_counts, dtype=torch.int32, - device=hidden_states.device) if not isinstance( - expert_counts, torch.Tensor) else expert_counts - - # Step 4: grouped GEMM w13 (gate_proj + up_proj) - gemm1_out = _gemm_grouped.moe_group_gemm( - sorted_hidden, w13, expert_counts_t) # (T*topk, 2*I) - gate, up = gemm1_out.chunk(2, dim=-1) - act_out = F.silu(gate) * up # (T*topk, I) - - # Step 6: grouped GEMM w2 (down_proj) - gemm2_out = _gemm_grouped.moe_group_gemm( - act_out, w2, expert_counts_t) # (T*topk, H) - - # Step 7: weighted combine back to token order - flat_weights = sorted_weights.unsqueeze(-1) # (T*topk, 1) - weighted = (gemm2_out * flat_weights).to(out.dtype) - out.index_add_(0, sorted_tok_ids, weighted) - else: - # Fallback: per-expert F.linear loop - start = 0 - for eid, count in enumerate(expert_counts): - end = start + count - if count == 0: - start = end - continue - tok_ids = sorted_tok_ids[start:end] - tokens = hidden_states[tok_ids] # (n, H) - gate_up = F.linear(tokens, w13[eid]) # (n, 2*I) - gate, up = gate_up.chunk(2, dim=-1) - act = F.silu(gate) * up # (n, I) - expert_out = F.linear(act, w2[eid]) # (n, H) - weights = sorted_weights[start:end].unsqueeze(-1) - out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) + start = 0 + for eid, count in enumerate(expert_counts): + end = start + count + if count == 0: start = end + continue + tok_ids = sorted_tok_ids[start:end] + tokens = hidden_states[tok_ids] # (n, H) + gate_up = F.linear(tokens, w13[eid]) # (n, 2*I) + gate, up = gate_up.chunk(2, dim=-1) + act = F.silu(gate) * up # (n, I) + expert_out = F.linear(act, w2[eid]) # (n, H) + weights = sorted_weights[start:end].unsqueeze(-1) + out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype)) + start = end return out # partial, all-reduce done in forward()