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

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

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: