35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
|
"""
|
||
|
|
Precompile moe_topk_softmax_v3.cu → .so during Docker build.
|
||
|
|
Same pattern as precompile_gdn.py.
|
||
|
|
|
||
|
|
Run: python3 ex_engine/precompile_moe_topk.py
|
||
|
|
"""
|
||
|
|
import os, sys
|
||
|
|
|
||
|
|
def main():
|
||
|
|
cu_path = os.path.join(os.path.dirname(__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} ...")
|
||
|
|
from torch.utils.cpp_extension import load
|
||
|
|
ext = load(
|
||
|
|
name="moe_topk_softmax_v3",
|
||
|
|
sources=[cu_path],
|
||
|
|
extra_cuda_cflags=["-O3"],
|
||
|
|
verbose=True,
|
||
|
|
)
|
||
|
|
print("[MOE] ✓ moe_topk_softmax_v3.so compiled successfully")
|
||
|
|
|
||
|
|
# Verify
|
||
|
|
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 __name__ == "__main__":
|
||
|
|
main()
|