dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

build: add MoE topk kernel precompile to patch_ops.sh + Dockerfile pipeline

Browse Source 
patch_ops.sh: step 7 precompiles moe_topk_softmax_v3.cu during Docker build
corex_moe.py: expanded .so/.cu search paths for both pre-compiled and JIT scenarios

Docker build flow:
  1. COPY ex_engine/ → /workspace/ex_engine/
  2. patch_ops.sh deploys corex_moe.py + corex_gdn.py to vllm models dir
  3. patch_ops.sh runs precompile_moe_topk.py → .so cached
  4. At runtime, corex_moe.py loads cached .so (no JIT delay)

Competition submission ready.
This commit is contained in:
project6-dev
2026-08-10 04:26:33 +00:00
parent f32ef97013
commit 44d36e6ccc
2 changed files with 24 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
ex_engine/python/corex_moe.py
View File

@@ -52,8 +52,16 @@ def _load_topk_kernel():
# JIT compile from source
cu_search = [
"/workspace/ex_engine/csrc/moe_topk_softmax_v3.cu",
"/workspace/qwen3_6_scripts/../ex_engine/csrc/moe_topk_softmax_v3.cu",
# Deployed by patch_ops.sh into vllm models dir
os.path.join(os.path.dirname(os.path.abspath(__file__)), "moe_topk_softmax_v3.cu"),
]
# Also search in vllm model_executor/models/
try:
import vllm
vllm_models = os.path.join(os.path.dirname(vllm.__file__), "model_executor", "models")
cu_search.append(os.path.join(vllm_models, "moe_topk_softmax_v3.cu"))
except Exception:
pass
for cu_path in cu_search:
if os.path.isfile(cu_path):
try:

16
qwen3_6_scripts/patch_ops.sh
View File

@@ -241,7 +241,21 @@ if [ -d "$EX_ENGINE_SRC/python" ]; then
echo "[patch_ops] EX Engine Python package deployed to $EX_PY_DST"
fi
echo "[patch_ops] DONE — EX Engine + SM70 GDN kernel + serving layer deployed"
# 7. Precompile MoE topk_softmax CUDA kernel (.cu → .so)
# This replaces the missing ixf_F.vllm_moe_topk_softmax with our own CUDA kernel
MOE_TOPK_CU="/workspace/ex_engine/csrc/moe_topk_softmax_v3.cu"
if [ -f "$MOE_TOPK_CU" ]; then
echo "[patch_ops] Precompiling moe_topk_softmax_v3.cu ..."
python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 || \
echo "[patch_ops] WARNING: MoE topk precompile failed — will JIT at runtime"
# Also deploy .cu source to vllm for JIT fallback
cp "$MOE_TOPK_CU" "$VLLM/model_executor/models/" 2>/dev/null || true
if [ -n "$VLLM2" ]; then
cp "$MOE_TOPK_CU" "$VLLM2/model_executor/models/" 2>/dev/null || true
fi
fi
echo "[patch_ops] DONE — EX Engine + SM70 GDN kernel + MoE topk kernel + serving layer deployed"
echo "[patch_ops] Deployed: qwen3_5.py, flash_qla_sm70, ex_engine factors, paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, serving layer"
echo "[patch_ops] EX factors replace: vllm_moe_topk_softmax (2304 calls/token), gdn_chunk_fwd (NaN fix)"
echo "[patch_ops] NOT deployed (base image native): model_runner.py, _custom_ops.py, sampler.py, logits_processor.py, arg_utils.py"