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
This commit is contained in:
dev
@@ -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"
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user