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

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"