fix(CRITICAL): moe_topk_softmax用importlib加载+patch_ops编译部署
- _custom_ops.py: 改用importlib.util加载.so替代torch.ops.load_library - patch_ops.sh: docker build时编译moe_topk_softmax_v3 CUDA kernel - 真机验证: weights sum=1.0, GPU test passed
This commit is contained in:
@@ -1054,9 +1054,11 @@ def _init_moe_topk():
|
|||||||
for pattern in so_patterns:
|
for pattern in so_patterns:
|
||||||
for so_path in glob.glob(pattern):
|
for so_path in glob.glob(pattern):
|
||||||
try:
|
try:
|
||||||
torch.ops.load_library(so_path)
|
import importlib.util
|
||||||
# After load_library, the pybind module should be importable
|
spec = importlib.util.spec_from_file_location(
|
||||||
import moe_topk_softmax_v3 as ext
|
"moe_topk_softmax_v3", so_path)
|
||||||
|
ext = importlib.util.module_from_spec(spec)
|
||||||
|
spec.loader.exec_module(ext)
|
||||||
_moe_topk_ext = ext
|
_moe_topk_ext = ext
|
||||||
logger.info("topk_softmax: loaded CUDA kernel from %s", so_path)
|
logger.info("topk_softmax: loaded CUDA kernel from %s", so_path)
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -179,6 +179,11 @@ fi
|
|||||||
build_stage "installing hash-pinned CoreX 3.2.3 extensions"
|
build_stage "installing hash-pinned CoreX 3.2.3 extensions"
|
||||||
bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
|
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"
|
build_stage "installing BI100 runtime modules"
|
||||||
cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
|
cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
|
||||||
cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
|
cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
|
||||||
|
|||||||
Reference in New Issue
Block a user