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