Log analysis from dockerrizhi.txt (07-23 Sub168 run) reveals THREE corex modules, not two: 1. corex_gdn.py — GatedDeltaNet fused kernel (already implemented) 2. corex_moe.py — MoE routing + expert GEMM (already implemented) 3. corex_fa2.py — FlashAttention2 dispatch (NEW) corex_fa2.py handles 32/36 attention layers with three modes: :333 → FA2 packed prefill (B=2 Hq=4 Hkv=1 D=256 max_q=2048) :507 → FA2 paged chunked prefill (B=1 max_q=17 cache_blocks=2) :225 → FA2 paged decode (B=1 max_k=45455 partition=256) Wraps ixformer.contrib.vllm_flash_attn + ixf_F.vllm_single_query_cached_kv_attention. These .so files EXIST in the base image (libixattn.so). Also: wired corex_fa2 import into qwen3_5.py + deploy script.
84 lines
3.0 KiB
Bash
Executable File
84 lines
3.0 KiB
Bash
Executable File
#!/bin/bash
|
|
# deploy_corex_modules.sh — Deploy corex_gdn.py + corex_moe.py into vllm
|
|
#
|
|
# Competitor 168's Docker had these at:
|
|
# $VLLM/model_executor/models/corex_gdn.py
|
|
# $VLLM/model_executor/models/corex_moe.py
|
|
#
|
|
# Our qwen3_5.py already has import fallback for these (lines 117-125):
|
|
# from vllm.model_executor.models import corex_gdn as _corex_gdn_module
|
|
# from vllm.model_executor.models import corex_moe as _corex_moe_module
|
|
#
|
|
# This script copies our implementations there so the imports succeed.
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
|
SRC_DIR="${SCRIPT_DIR}/python"
|
|
|
|
# Find vllm install path
|
|
VLLM_MODELS=""
|
|
for candidate in \
|
|
/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models \
|
|
/usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/models \
|
|
/usr/local/lib/python3.10/site-packages/vllm/model_executor/models \
|
|
/workspace/vllm/model_executor/models; do
|
|
if [[ -d "$candidate" ]]; then
|
|
VLLM_MODELS="$candidate"
|
|
break
|
|
fi
|
|
done
|
|
|
|
if [[ -z "$VLLM_MODELS" ]]; then
|
|
# Try Python detection
|
|
VLLM_MODELS=$(python3 -c "
|
|
import os, vllm
|
|
print(os.path.join(os.path.dirname(vllm.__file__), 'model_executor', 'models'))
|
|
" 2>/dev/null || true)
|
|
fi
|
|
|
|
if [[ -z "$VLLM_MODELS" ]] || [[ ! -d "$VLLM_MODELS" ]]; then
|
|
echo "[COREX] ERROR: Cannot find vllm models directory"
|
|
exit 1
|
|
fi
|
|
|
|
echo "[COREX] Deploying to: $VLLM_MODELS"
|
|
|
|
# Deploy corex_gdn.py
|
|
if [[ ! -f "${VLLM_MODELS}/corex_gdn.py" ]]; then
|
|
cp "${SRC_DIR}/corex_gdn.py" "${VLLM_MODELS}/corex_gdn.py"
|
|
echo "[COREX] ✓ Deployed corex_gdn.py"
|
|
else
|
|
echo "[COREX] ✓ corex_gdn.py already exists (base image or prior deploy)"
|
|
fi
|
|
|
|
# Deploy corex_moe.py
|
|
if [[ ! -f "${VLLM_MODELS}/corex_moe.py" ]]; then
|
|
cp "${SRC_DIR}/corex_moe.py" "${VLLM_MODELS}/corex_moe.py"
|
|
echo "[COREX] ✓ Deployed corex_moe.py"
|
|
else
|
|
echo "[COREX] ✓ corex_moe.py already exists (base image or prior deploy)"
|
|
fi
|
|
|
|
# Deploy corex_fa2.py
|
|
if [[ ! -f "${VLLM_MODELS}/corex_fa2.py" ]]; then
|
|
cp "${SRC_DIR}/corex_fa2.py" "${VLLM_MODELS}/corex_fa2.py"
|
|
echo "[COREX] ✓ Deployed corex_fa2.py"
|
|
else
|
|
echo "[COREX] ✓ corex_fa2.py already exists (base image or prior deploy)"
|
|
fi
|
|
|
|
# Also deploy to ex_engine location (backup import path)
|
|
mkdir -p /workspace/ex_engine/python 2>/dev/null || true
|
|
cp "${SRC_DIR}/corex_gdn.py" /workspace/ex_engine/python/ 2>/dev/null || true
|
|
cp "${SRC_DIR}/corex_moe.py" /workspace/ex_engine/python/ 2>/dev/null || true
|
|
cp "${SRC_DIR}/corex_fa2.py" /workspace/ex_engine/python/ 2>/dev/null || true
|
|
|
|
echo "[COREX] Deploy complete"
|
|
echo "[COREX] Expected log on startup:"
|
|
echo " corex_gdn.py:NN → Loaded fused CoreX GDN decode operator ..."
|
|
echo " corex_gdn.py:NN → Using fused CoreX GDN prefill operator"
|
|
echo " corex_moe.py:NN → Using CoreX fused MoE prefill operator: tokens=N, kernel=expert-grouped-wmma"
|
|
echo " corex_fa2.py:NN → Using CoreX FA2 packed prefill: B=N Hq=4 Hkv=1 D=256 ..."
|
|
echo " corex_fa2.py:NN → Using CoreX paged decode: B=N Hq=4 Hkv=1 D=256 ..."
|