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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user