fix(MoE): robust CUDA kernel loading + no-GPU precompile

1. precompile_moe_topk.py: skip GPU verification during Docker build
   (torch.cuda.is_available() check — .so compilation doesn't need GPU)

2. _custom_ops.py topk_softmax init: 3-tier loading
   - import precompiled module (torch cache)
   - scan known .so paths (torch_extensions cache dirs)
   - JIT compile from .cu source
   - PyTorch fallback with WARNING (not silent — must know if CUDA failed)

3. patch_ops.sh: report .so location after precompile for debugging

ex_engine/precompile_moe_topk.py

@@ -1,34 +1,52 @@
 """
 Precompile moe_topk_softmax_v3.cu → .so during Docker build.
-Same pattern as precompile_gdn.py.
+Build-only — does NOT require GPU. Verification deferred to runtime.
 
-Run: python3 ex_engine/precompile_moe_topk.py
+The .so will be cached by torch and loaded at runtime via:
+  import moe_topk_softmax_v3
 """
 import os, sys
 
 def main():
-    cu_path = os.path.join(os.path.dirname(__file__), "csrc", "moe_topk_softmax_v3.cu")
+    cu_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "csrc", "moe_topk_softmax_v3.cu")
     if not os.path.isfile(cu_path):
         print(f"[MOE] ERROR: {cu_path} not found")
         sys.exit(1)
 
     print(f"[MOE] Compiling {cu_path} ...")
+
+    # Detect corex compiler (BI-V100 Docker image)
+    corex_clang = "/usr/local/corex/bin/clang++"
+    use_corex = os.path.isfile(corex_clang)
+
     from torch.utils.cpp_extension import load
+
+    extra_cuda_cflags = ["-O3"]
+    extra_ldflags = []
+
+    if use_corex:
+        print(f"[MOE] Using corex clang at {corex_clang}")
+        # corex torch extension picks up CUDA_HOME automatically
+        # No special flags needed — torch.utils.cpp_extension handles ivcore10
+
     ext = load(
         name="moe_topk_softmax_v3",
         sources=[cu_path],
-        extra_cuda_cflags=["-O3"],
+        extra_cuda_cflags=extra_cuda_cflags,
+        extra_ldflags=extra_ldflags,
         verbose=True,
     )
-    print("[MOE] ✓ moe_topk_softmax_v3.so compiled successfully")
+    print("[MOE] ✓ moe_topk_softmax_v3.so compiled")
 
-    # Verify
+    # Optional GPU verification — skip if no GPU (Docker build)
     import torch
-    gating = torch.randn(4, 64, device='cuda', dtype=torch.float16)
-    w, ids, _ = ext.moe_topk_softmax(gating, 8, True)
-    assert not w.isnan().any(), "NaN in topk weights!"
-    assert torch.allclose(w.sum(dim=-1), torch.ones(4, device='cuda'), atol=1e-3)
-    print("[MOE] ✓ Runtime verification passed")
+    if torch.cuda.is_available():
+        gating = torch.randn(4, 64, device='cuda', dtype=torch.float16)
+        w, ids, _ = ext.moe_topk_softmax(gating, 8, True)
+        assert not w.isnan().any(), "NaN in topk weights!"
+        print("[MOE] ✓ GPU verification passed")
+    else:
+        print("[MOE] No GPU — skipping runtime verification (will verify at first inference)")
 
 if __name__ == "__main__":
     main()

qwen3_6_scripts/_custom_ops.py

@@ -828,30 +828,47 @@ def invoke_fused_moe_kernel(
 
 
 # ---------- topk_softmax: CUDA kernel → PyTorch fallback ----------
-# Our moe_topk_softmax_v3.cu is a fused warp-shuffle CUDA kernel
-# specialized for 64 experts. It's precompiled during Docker build.
-# If loading fails (no GPU at build time), falls back to PyTorch.
+# moe_topk_softmax_v3.cu: fused warp-shuffle kernel, 64 experts, zero SMEM.
+# Precompiled during Docker build → .so cached by torch.
+# If not found, JIT from .cu source. PyTorch last resort.
 _moe_topk_ext = None
 _moe_topk_init_done = False
 
 def _init_moe_topk():
     global _moe_topk_ext, _moe_topk_init_done
     _moe_topk_init_done = True
-    # Try loading precompiled .so first
+    # 1. Try import precompiled module (torch cache from Docker build)
     try:
         import moe_topk_softmax_v3 as ext
         _moe_topk_ext = ext
-        logger.info("topk_softmax: loaded CUDA kernel (moe_topk_softmax_v3)")
+        logger.info("topk_softmax: loaded precompiled CUDA kernel")
         return
     except ImportError:
         pass
-    # Try JIT compile from source
-    import os, glob
+    # 2. Try loading from known .so paths
+    import glob
+    so_patterns = [
+        "/workspace/ex_engine/build/moe_topk_softmax_v3*.so",
+        "/root/.cache/torch_extensions/*/moe_topk_softmax_v3/*.so",
+        "/tmp/torch_extensions/*/moe_topk_softmax_v3/*.so",
+    ]
+    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
+                _moe_topk_ext = ext
+                logger.info("topk_softmax: loaded CUDA kernel from %s", so_path)
+                return
+            except Exception:
+                pass
+    # 3. JIT compile from .cu source
+    import os
     search_paths = [
         "/workspace/ex_engine/csrc/moe_topk_softmax_v3.cu",
-        os.path.join(os.path.dirname(__file__), "moe_topk_softmax_v3.cu"),
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "moe_topk_softmax_v3.cu"),
     ]
-    # Also search vllm model dir where patch_ops.sh copies it
     for base in ["/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models",
                  "/usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/models"]:
         search_paths.append(os.path.join(base, "moe_topk_softmax_v3.cu"))
@@ -869,8 +886,9 @@ def _init_moe_topk():
                 logger.info("topk_softmax: JIT compiled CUDA kernel from %s", cu_path)
                 return
             except Exception as e:
-                logger.warning("topk_softmax: JIT compile failed (%s), trying next", e)
-    logger.info("topk_softmax: no CUDA kernel available, using PyTorch fallback")
+                logger.warning("topk_softmax: JIT compile failed (%s)", e)
+                break  # Don't retry same source with different paths
+    logger.warning("topk_softmax: CUDA kernel unavailable — PyTorch fallback (SLOW)")
 
 
 def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,

qwen3_6_scripts/patch_ops.sh

@@ -267,7 +267,10 @@ if [ -f "$MOE_TOPK_CU" ]; then
     echo "[patch_ops] Precompiling moe_topk_softmax_v3.cu ..."
     python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 || \
         echo "[patch_ops] WARNING: MoE topk precompile failed — will JIT at runtime"
-    # Also deploy .cu source to vllm for JIT fallback
+    # Find and report the compiled .so location
+    echo "[patch_ops] Searching for compiled .so ..."
+    find /root/.cache/torch_extensions /tmp/torch_extensions -name "*.so" -path "*moe_topk*" 2>/dev/null | head -3
+    # Also deploy .cu source to vllm dir for runtime JIT fallback
     cp "$MOE_TOPK_CU" "$VLLM/model_executor/models/" 2>/dev/null || true
     if [ -n "$VLLM2" ]; then
         cp "$MOE_TOPK_CU" "$VLLM2/model_executor/models/" 2>/dev/null || true