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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user