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..1d749ac7 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}" @@ -66,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 = [] diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index ac007739..1c24b056 100755 --- 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 diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 130894f7..343e904a 100644 --- 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}") \ No newline at end of file