0478628f17
Tested on real machine (cc-b2042074, BI-V100, IX-ML 3.2.3): _moe_C.topk_softmax() → SUCCESS, correct output Two fixes proven on hardware: 1. cuda_compat.h: WARP_SIZE=64 (BI-V100 warp is 64, not 32) 2. topk_softmax_kernels.cu: cub/block/block_reduce.cuh instead of cub/cub.cuh (cub.cuh pulls radix_sort which has WARP_SIZE conflict) Key finding: ixformer SDK on this base image does NOT have topk_softmax. The ixformer::infer namespace from xllm's ixformer.h is for newer SDK. We MUST compile our own _moe_C kernel — which now works. Build flags (clang 16, ivcore10): CUDA: -O3 -cl-fast-relaxed-math (NOT --use_fast_math) C++: -O2 -std=c++17 Dockerfile simplified: 3 steps (was 6) _custom_ops.py: _moe_C as Priority 0, in-place vllm API
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Precompile _moe_C extension: topk_softmax + moe_align_block_size.
|
|
|
|
Proven on real BI-V100 hardware:
|
|
- WARP_SIZE=64 (not 32)
|
|
- cub/block/block_reduce.cuh (not cub/cub.cuh which pulls radix_sort)
|
|
- -cl-fast-relaxed-math (not --use_fast_math which is nvcc-only)
|
|
"""
|
|
import os, sys, logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("precompile_moe")
|
|
|
|
def main():
|
|
import torch
|
|
from torch.utils.cpp_extension import load
|
|
|
|
base = os.path.dirname(os.path.abspath(__file__))
|
|
v055 = os.path.join(base, "csrc", "moe_v055")
|
|
|
|
sources = [
|
|
os.path.join(v055, "topk_softmax_kernels.cu"),
|
|
os.path.join(v055, "moe_align_block_size_kernels.cu"),
|
|
os.path.join(v055, "moe_pybind.cpp"),
|
|
]
|
|
for s in sources:
|
|
if not os.path.exists(s):
|
|
logger.error("MISSING: %s", s)
|
|
sys.exit(1)
|
|
|
|
include_paths = [
|
|
v055,
|
|
os.path.join(base, "csrc", "moe"),
|
|
os.path.join(base, "csrc"),
|
|
"/usr/local/corex/include",
|
|
]
|
|
|
|
logger.info("Sources: %s", sources)
|
|
logger.info("Compiling _moe_C...")
|
|
|
|
try:
|
|
mod = load(
|
|
name="_moe_C",
|
|
sources=sources,
|
|
extra_include_paths=include_paths,
|
|
extra_cuda_cflags=["-O3", "-cl-fast-relaxed-math"],
|
|
extra_cflags=["-O2", "-std=c++17"],
|
|
verbose=True,
|
|
)
|
|
fns = [x for x in dir(mod) if not x.startswith("_")]
|
|
logger.info("SUCCESS: _moe_C functions: %s", fns)
|
|
except Exception as e:
|
|
logger.error("FAILED: %s", e)
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|