f32ef97013
moe_topk_softmax_v3.cu: BI-V100 verified (2026-08-10) - 64 experts, topk=8, warp shuffle, zero shared memory - renormalize: sum=1.0 ✓, no NaN ✓, no duplicate ids ✓ - 881 token batch ✓ - Compiler: corex clang/16, --cuda-gpu-arch=ivcore10 - Stream: c10::cuda::getCurrentCUDAStream() corex_moe.py: loads CUDA kernel, NO Python fallback - Searches pre-compiled .so → JIT compile from source → error - MoE pipeline: CUDA topk → cublas expert GEMM → ixformer silu_and_mul precompile_moe_topk.py: Docker build-time compilation + verification Key finding from real machine probing: ixformer::infer::topk_softmax is DECLARED in ixformer.h but NOT IMPLEMENTED in any .so in the base image (nm -D scan: zero hits). Must compile our own kernel.
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()
|