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:
claude
2026-08-11 09:49:25 +00:00
parent f4d4219280
commit 3a5cc2a589
2 changed files with 10 additions and 3 deletions

View File

@@ -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