fix: resolve all 8 deployment pipeline breaks

Breaks found and fixed: 1. Dockerfile: COPY 5 individual files → COPY entire ex_engine/ 2. patch_ops.sh EX_ENGINE_DIR: /workspace/ex_engine not found → added fallback 3. patch_ops.sh deploy: ix_ops.py to ex_engine/ (flat) → ex_engine/python/ (correct package) 4. ix_startup_patch.py: import from vllm.ex_engine.patch_vllm_ops → vllm.ex_engine.python.patch_vllm_ops 5. ix_moe_bridge.so: only deployed to ex_engine/ → also copy to model_executor/models/ and vllm root 6. ex_engine/__init__.py: missing re-exports → add imports so 'from vllm.ex_engine import x' works 7. gemm_grouped.so: compiled but never imported → add import + flag + prefill GEMM path in qwen3_5.py 8. build_moe_bridge.sh Python heredoc: SCRIPT_DIR not exported + wrong nested path → export + search both layouts Also added: - CUTLASS batched GEMM compile step (corex_batched_gemm.so for decode) - Full ex_engine/python/*.py deployment (was deploying only 2 of 19 files) - EX_ENGINE_INFRA_AUDIT.md documenting all findings
2026-08-17 07:01:41 +00:00
parent ee516bd206
commit 9ea0a1d4f4
5 changed files with 286 additions and 25 deletions
--- a/EX_ENGINE_INFRA_AUDIT.md
+++ b/EX_ENGINE_INFRA_AUDIT.md
@@ -0,0 +1,172 @@
+# 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在用) | ✓ 正常工作 |
--- a/build_moe_bridge.sh
+++ b/build_moe_bridge.sh
@@ -75,14 +75,28 @@ 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", "")

-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")
+# 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 = []
--- a/ex_engine/build_moe_bridge.sh
+++ b/ex_engine/build_moe_bridge.sh
@@ -75,14 +75,28 @@ 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", "")

-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")
+# 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 = []
--- a/qwen3_6_scripts/patch_ops.sh
+++ b/qwen3_6_scripts/patch_ops.sh
@@ -208,14 +208,30 @@ if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR/python" ]; then
 fi

 if [ -d "$EX_ENGINE_DIR/python" ]; then
-    # Create ex_engine package inside vllm
+    # Create ex_engine package inside vllm with correct Python package structure
+    mkdir -p "${VLLM_ROOT}/ex_engine/python"
    mkdir -p "${VLLM_ROOT}/ex_engine/csrc"
-    echo '"""ex_engine — Algorithm factor replacement for BI-V100."""' > "${VLLM_ROOT}/ex_engine/__init__.py"

-    # 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/"
+    # __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 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
@@ -230,7 +246,7 @@ import logging
 _logger = logging.getLogger("ix_startup_patch")
 def apply():
    try:
-        from vllm.ex_engine.patch_vllm_ops import apply_all_patches
+        from vllm.ex_engine.python.patch_vllm_ops import apply_all_patches
        n = apply_all_patches()
        if n > 0:
            _logger.info("ix_startup_patch: %d patches applied", n)
@@ -391,6 +407,16 @@ 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"
    }
+    # 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"
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -143,6 +143,11 @@ 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:
@@ -212,6 +217,11 @@ _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))
@@ -1884,21 +1894,46 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                expert_counts = torch.bincount(
                    flat_eids, minlength=w13.shape[0]).tolist()

-            start = 0
-            for eid, count in enumerate(expert_counts):
-                end = start + count
-                if count == 0:
+            # --- 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 = 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()