build: migrate CUDA kernel build to CMake
Replace torch CUDAExtension/ParallelBuildExtension with a CMake-based build. Each kernel compiles as an independent pybind11 module in parallel via cmake --build -j, outputting to astrai/extension/lib. - Add csrc/CMakeLists.txt (5 kernel targets, torch/pybind11 linking) - setup.py: _CMakeBuildExt invokes cmake; auto-detect CUDA arch via torch - Remove csrc/build.py (REGISTRY/build flags now in CMakeLists) - Fix rel-err eps in attn_test.cu (1e-8 -> 1e-4, bf16 scale) - Update docs/developer/cuda_kernels.md build section - .gitignore: allow csrc/CMakeLists.txt
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
cmake_minimum_required(VERSION 3.18)
|
||||
project(astrai_kernels LANGUAGES CUDA CXX)
|
||||
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
set(CMAKE_CUDA_STANDARD 17)
|
||||
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
|
||||
if(NOT DEFINED TORCH_HOME)
|
||||
set(TORCH_HOME "$ENV{TORCH_HOME}")
|
||||
endif()
|
||||
if(NOT TORCH_HOME)
|
||||
message(FATAL_ERROR "TORCH_HOME must point at the torch install dir (site-packages/torch)")
|
||||
endif()
|
||||
|
||||
if(NOT DEFINED PYTHON_INCLUDE_DIR)
|
||||
set(PYTHON_INCLUDE_DIR "/usr/include/python${PYTHON_VERSION_MAJOR}.${PYTHON_VERSION_MINOR}")
|
||||
endif()
|
||||
|
||||
if(NOT DEFINED ASTRAI_CUDA_ARCH)
|
||||
if(DEFINED ENV{ASTRAI_CUDA_ARCH})
|
||||
set(ASTRAI_CUDA_ARCH "$ENV{ASTRAI_CUDA_ARCH}")
|
||||
else()
|
||||
set(ASTRAI_CUDA_ARCH 80)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
set(TORCH_LIB_DIR "${TORCH_HOME}/lib")
|
||||
set(CUDA_LIB_DIR "/usr/local/cuda/lib64")
|
||||
|
||||
set(CXX_FLAGS -O3 -funroll-loops)
|
||||
set(NVCC_FLAGS -O3
|
||||
--expt-relaxed-constexpr
|
||||
--use_fast_math
|
||||
"--ptxas-options=-O3,-v"
|
||||
--extra-device-vectorization
|
||||
--threads=16)
|
||||
|
||||
set(TORCH_LIBS
|
||||
"${TORCH_LIB_DIR}/libtorch_python.so"
|
||||
"${TORCH_LIB_DIR}/libtorch_cuda.so"
|
||||
"${TORCH_LIB_DIR}/libc10_cuda.so"
|
||||
"${TORCH_LIB_DIR}/libtorch_cpu.so"
|
||||
"${TORCH_LIB_DIR}/libtorch.so"
|
||||
"${TORCH_LIB_DIR}/libc10.so"
|
||||
CUDA::cudart)
|
||||
|
||||
set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
|
||||
|
||||
set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb)
|
||||
|
||||
foreach(name ${KERNELS})
|
||||
add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu")
|
||||
|
||||
target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name})
|
||||
|
||||
target_include_directories(${name} PRIVATE
|
||||
"${TORCH_HOME}/include"
|
||||
"${TORCH_HOME}/include/torch/csrc/api/include"
|
||||
"${PYTHON_INCLUDE_DIR}")
|
||||
|
||||
target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
|
||||
target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
|
||||
|
||||
target_compile_options(${name} PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:${CXX_FLAGS}>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:${NVCC_FLAGS}>)
|
||||
|
||||
set_target_properties(${name} PROPERTIES
|
||||
PREFIX ""
|
||||
SUFFIX ".${PY_SOABI}.so"
|
||||
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../astrai/extension/lib")
|
||||
endforeach()
|
||||
@@ -1,76 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def cuda_toolkit_version() -> tuple[int, int] | None:
|
||||
"""Return ``(major, minor)`` of the nvcc on PATH, or ``None``.
|
||||
|
||||
Used by ``setup.py`` to detect nvcc/torch CUDA version mismatches
|
||||
(e.g. nvcc 13.0 with a cu128 torch wheel) which cause cryptic ABI errors.
|
||||
"""
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
nvcc = shutil.which("nvcc")
|
||||
if nvcc is None:
|
||||
return None
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
[nvcc, "--version"], stderr=subprocess.STDOUT, text=True
|
||||
)
|
||||
for line in out.splitlines():
|
||||
if "release" in line:
|
||||
ver = line.split("release")[1].split(",")[0].strip()
|
||||
major, minor = ver.split(".")
|
||||
return (int(major), int(minor))
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _arch_flags() -> list[str]:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
cap = torch.cuda.get_device_capability()
|
||||
else:
|
||||
cap = (8, 0)
|
||||
ver = f"{cap[0]}{cap[1]}"
|
||||
flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
|
||||
# tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
|
||||
# kernel dispatch at build time via this define rather than at runtime.
|
||||
if cap[0] < 8:
|
||||
flags.append("-DASTRAI_NO_MMA")
|
||||
return flags
|
||||
|
||||
|
||||
_kernels_dir = Path("csrc/kernels")
|
||||
REGISTRY: dict[str, dict] = {}
|
||||
|
||||
CXX_FLAGS = ["-O3", "-funroll-loops"]
|
||||
NVCC_FLAGS = [
|
||||
"-O3",
|
||||
"--expt-relaxed-constexpr",
|
||||
"--use_fast_math",
|
||||
"--ptxas-options=-O3,-v",
|
||||
"--extra-device-vectorization",
|
||||
"--threads=16",
|
||||
]
|
||||
|
||||
|
||||
def register(name: str, sources: list[str] | None = None, **kwargs):
|
||||
if sources is None:
|
||||
sources = [str(_kernels_dir / f"{name}.cu")]
|
||||
REGISTRY[name] = {
|
||||
"sources": sources,
|
||||
"cxx_flags": [*CXX_FLAGS],
|
||||
"nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
|
||||
"extra_link_args": kwargs.pop("extra_link_args", []),
|
||||
**kwargs,
|
||||
}
|
||||
|
||||
|
||||
register("attn_decode")
|
||||
register("attn_prefill")
|
||||
register("attn_paged_decode")
|
||||
register("attn_paged_prefill")
|
||||
register("rotary_emb")
|
||||
@@ -83,7 +83,7 @@ static int run_decode_test(int B, int Hq, int Hk, int sl, int D, int causal) {
|
||||
for (size_t i=0;i<nQ;i++){
|
||||
float err=fabsf(bf2f(hOut[i])-ref[i]);
|
||||
if(err>max_abs_err) max_abs_err=err;
|
||||
float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
|
||||
float rel=err/fmaxf(fabsf(ref[i]), 1e-4f);
|
||||
if(rel>max_rel_err) max_rel_err=rel;
|
||||
}
|
||||
const float atol=0.01f, rtol=0.01f;
|
||||
@@ -206,7 +206,7 @@ static int run_prefill_test(int B, int Hq, int Hk, int ql, int kl, int D, int ca
|
||||
for (size_t i=0;i<nQ;i++) {
|
||||
float err=fabsf(bf2f(hOut[i])-ref[i]);
|
||||
if(err>max_abs_err) max_abs_err=err;
|
||||
float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
|
||||
float rel=err/fmaxf(fabsf(ref[i]), 1e-4f);
|
||||
if(rel>max_rel_err) max_rel_err=rel;
|
||||
}
|
||||
const float atol=0.01f, rtol=0.01f;
|
||||
|
||||
Reference in New Issue
Block a user