platform test baseline4

2026-08-17 07:19:44 +00:00
parent be4d661191
commit 8211a45464
4 changed files with 183 additions and 51 deletions
--- a/qwen3_6_scripts/patch_ops.sh
+++ b/qwen3_6_scripts/patch_ops.sh
@@ -202,19 +202,36 @@ 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
-    # 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
@@ -229,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)
@@ -332,22 +349,83 @@ if b"max_completion_tokens" not in installed:
    raise SystemExit("protocol.py missing max_completion_tokens field")
 PY

-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 || {
-        echo "[WARN] MoE bridge build failed — will use Python fallback"
+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 "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 "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 || {
+        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 "compiling submission Python sources"
 find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
--- 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()

@@ -2796,4 +2831,4 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
                weight_loader(param, loaded_weight)
        _bi100_model_trace(
            f"MoE load_weights complete items={loaded_count} "
-            f"vision_items={vision_loaded_count}")
+            f"vision_items={vision_loaded_count}")