diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index 9962a147..0d4b1c77 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -1054,9 +1054,11 @@ def _init_moe_topk(): for pattern in so_patterns: for so_path in glob.glob(pattern): try: - torch.ops.load_library(so_path) - # After load_library, the pybind module should be importable - import moe_topk_softmax_v3 as ext + import importlib.util + spec = importlib.util.spec_from_file_location( + "moe_topk_softmax_v3", so_path) + ext = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ext) _moe_topk_ext = ext logger.info("topk_softmax: loaded CUDA kernel from %s", so_path) return diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 310d7dca..6c277622 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -179,6 +179,11 @@ fi build_stage "installing hash-pinned CoreX 3.2.3 extensions" bash ./install_prebuilt_corex.sh "${VLLM_ROOT}" +build_stage "compiling moe_topk_softmax CUDA kernel" +cd /workspace && bash ex_engine/build_moe_topk.sh 2>&1 || echo "[WARN] moe_topk build failed (non-fatal)" +# Deploy to workspace search path (_custom_ops.py looks in /workspace/ex_engine/build/) +cd "${OLDPWD}" + build_stage "installing BI100 runtime modules" cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py" cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"