cdcf115037
BI-V100 real device results: moe_group_gemm: err=0.000015 PASS moe_decode_cutlass: NaN=False PASS cutlass grouped: 4.77ms vs torch.mm loop: 9.38ms → 1.97x speedup Fix: gemm_grouped.cu ldb=K (not N) for ColumnMajor B view Link: -lcuinfer from /usr/local/corex-3.2.3/lib64/libcuinfer.so.7
57 lines
1.5 KiB
Bash
57 lines
1.5 KiB
Bash
#!/usr/bin/env bash
|
|
# build_cuinfer_gemm.sh — Compile cuinfer GEMM wrapper
|
|
#
|
|
# Links: libcuinfer.so (from /usr/local/corex/lib64/)
|
|
# Output: cuinfer_gemm_wrapper.so
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
SRC="${SCRIPT_DIR}/cuinfer_gemm_wrapper.cu"
|
|
HDR="${SCRIPT_DIR}/cuinfer_handle.h"
|
|
|
|
echo "[cuinfer_gemm] Building cuinfer_gemm_wrapper.so"
|
|
|
|
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
|
CUINFER_LIB=""
|
|
for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib"; do
|
|
if [[ -f "${d}/libcuinfer.so" ]]; then
|
|
CUINFER_LIB="${d}"
|
|
break
|
|
fi
|
|
done
|
|
|
|
python3 << PYEOF
|
|
import os, sys, shutil
|
|
|
|
src = "${SRC}"
|
|
hdr_dir = "${SCRIPT_DIR}"
|
|
cuinfer_lib = "${CUINFER_LIB}"
|
|
|
|
ldflags = []
|
|
if cuinfer_lib:
|
|
ldflags = [f"-L{cuinfer_lib}", "-lcuinfer", f"-Wl,-rpath,{cuinfer_lib}"]
|
|
|
|
try:
|
|
from torch.utils.cpp_extension import load
|
|
mod = load(
|
|
name="cuinfer_gemm_wrapper",
|
|
sources=[src],
|
|
extra_include_paths=[hdr_dir],
|
|
extra_cflags=["-O2", "-std=c++17"],
|
|
extra_cuda_cflags=["-O2"],
|
|
extra_ldflags=ldflags,
|
|
verbose=True,
|
|
)
|
|
print("[cuinfer_gemm] ✓ OK")
|
|
|
|
import importlib
|
|
spec = importlib.util.find_spec("cuinfer_gemm_wrapper")
|
|
if spec and spec.origin:
|
|
shutil.copy2(spec.origin, os.path.join(hdr_dir, "cuinfer_gemm_wrapper.so"))
|
|
print(f"[cuinfer_gemm] ✓ Saved")
|
|
|
|
except Exception as e:
|
|
print(f"[cuinfer_gemm] ERROR: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
PYEOF
|