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
This commit is contained in:
project6
2026-08-10 07:50:28 +00:00
parent c0cc4e7dc9
commit 8b6f3fd242
3 changed files with 62 additions and 23 deletions

View File

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