From 3a5cc2a589de0b7cc601a70f9b11cd3884466abb Mon Sep 17 00:00:00 2001
From: claude <claude@project6.dev>
Date: Tue, 11 Aug 2026 09:49:25 +0000
Subject: [PATCH] =?UTF-8?q?fix(CRITICAL):=20moe=5Ftopk=5Fsoftmax=E7=94=A8i?=
 =?UTF-8?q?mportlib=E5=8A=A0=E8=BD=BD+patch=5Fops=E7=BC=96=E8=AF=91?=
 =?UTF-8?q?=E9=83=A8=E7=BD=B2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- _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
---
 qwen3_6_scripts/_custom_ops.py | 8 +++++---
 qwen3_6_scripts/patch_ops.sh   | 5 +++++
 2 files changed, 10 insertions(+), 3 deletions(-)

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"