File size: 1,803 Bytes
54152e6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | #!/usr/bin/env bash
set -euo pipefail
RELEASE_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RUNTIME_ROOT="${RELEASE_ROOT}/runtime"
PYTHON_BIN="${PYTHON_BIN:-python}"
CUDA_TOOLKIT_ROOT="${CUDA_HOME:-/usr/local/cuda}"
NVCC="${CUDA_TOOLKIT_ROOT}/bin/nvcc"
if [[ ! -x "${NVCC}" ]]; then
echo "nvcc was not found at ${NVCC}" >&2
exit 1
fi
mkdir -p "${RUNTIME_ROOT}"
"${NVCC}" \
-std=c++17 -O3 -arch=sm_120a -shared -Xcompiler=-fPIC \
"${RUNTIME_ROOT}/nvfp4_linear.cu" \
-o "${RUNTIME_ROOT}/libmage_nvfp4_linear.so" \
-L"${CUDA_TOOLKIT_ROOT}/lib64" -lcublasLt -lcublas
PYTORCH_INCLUDE_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import include_paths
print(" ".join(f"-I{path}" for path in include_paths()))
PY
)"
PYTORCH_LIBRARY_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import library_paths
print(" ".join(f"-L{path}" for path in library_paths()))
PY
)"
PYTORCH_RPATH_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import library_paths
print(" ".join(f"-Wl,-rpath,{path}" for path in library_paths()))
PY
)"
PYTORCH_ABI_FLAG="$("${PYTHON_BIN}" - <<'PY'
import torch
print(f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}")
PY
)"
g++ \
-std=c++20 -O3 -shared -fPIC \
${PYTORCH_ABI_FLAG} \
${PYTORCH_INCLUDE_FLAGS} \
${PYTORCH_LIBRARY_FLAGS} \
${PYTORCH_RPATH_FLAGS} \
-I"${CUDA_TOOLKIT_ROOT}/include" \
-I"${RUNTIME_ROOT}" \
-o "${RUNTIME_ROOT}/libmage_nvfp4_torch_op.so" \
"${RUNTIME_ROOT}/sm120_linear_op.cpp" \
-L"${CUDA_TOOLKIT_ROOT}/lib64" \
-L"${RUNTIME_ROOT}" \
-Wl,-rpath,"${CUDA_TOOLKIT_ROOT}/lib64" \
-Wl,-rpath,'$ORIGIN' \
-ltorch -ltorch_cpu -ltorch_cuda -lc10 -lc10_cuda -lcudart \
-lmage_nvfp4_linear
echo "Built the packaged SM120 runtime in ${RUNTIME_ROOT}"
|