ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

build: migrate CUDA kernel build to CMake

Browse Source 
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:
ViperEkura
2026-08-04 21:27:22 +08:00
parent cdf9145ecf
commit 49a9c6b3d2
6 changed files with 214 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -9,6 +9,7 @@
!scripts/**/*.py
!tests/**/*.py
!csrc/**/*.py
!csrc/CMakeLists.txt
!csrc/**/*.cu
!csrc/**/*.h
csrc/CMakeLists.txt
+74
View File
@@ -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()
csrc/build.py
-76
View File
@@ -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")
csrc/tests/attn_test.cu
+2 -2
View File
@@ -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;
docs/developer/cuda_kernels.md
+13 -4
View File
@@ -51,23 +51,32 @@ CSRC_KERNELS=true pip install -e . --no-build-isolation
# Rebuild after editing .cu/.cuh files
CSRC_KERNELS=true python setup.py build_ext --inplace
# Output: astrai/extension/lib/*.so
# Or invoke CMake directly
cmake -S csrc -B build/cmake \
-DTORCH_HOME=<site-packages>/torch \
-DPYTHON_INCLUDE_DIR=<python include> \
-DPY_SOABI=cpython-312-x86_64-linux-gnu
cmake --build build/cmake -j 16
```
### Architecture flags
`csrc/build.py` auto-detects the GPU compute capability and generates the appropriate `nvcc` gencode flag:
`setup.py` passes the GPU compute capability to CMake via `ASTRAI_CUDA_ARCH` (default `89`, i.e. sm_89 / L20):
- **sm_80+** (Ampere and later): enables tensor-core MMA path (`mma.sync.m16n8k16.bf16`)
- **Below sm_80**: adds `-DASTRAI_NO_MMA` to disable the MMA path at compile time
### Build configuration
`csrc/CMakeLists.txt` defines the CUDA extension build:
```
NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
--ptxas-options=-O3,-v --extra-device-vectorization --threads=8
--ptxas-options=-O3,-v --extra-device-vectorization --threads=16
```
The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 5). Each entry maps a kernel name to its source files and build flags.
Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all five kernel targets in parallel via `cmake --build -j N`.
## Attention Backend
@@ -146,7 +155,7 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
```
csrc/
├── build.py # Build system: REGISTRY, _arch_flags, nvcc flags
├── CMakeLists.txt # CMake build: 5 kernel targets, torch/pybind11 linking
├── kernels/
│ ├── attn_common.h # Shared attention params (AttentionParams, PagedAttentionParams)
│ ├── attn_decode.cu # Basic decode kernel (registered)
setup.py
+107 -102
View File
@@ -19,8 +19,6 @@ def _should_build():
if force == "false":
return False
try:
import shutil
import torch
return shutil.which("nvcc") is not None and torch.cuda.is_available()
@@ -28,125 +26,132 @@ def _should_build():
return False
ext_modules = []
cmdclass = {}
if _should_build():
def _torch_prefix():
"""Return the torch install dir (site-packages/torch) used for headers/libs."""
try:
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
from csrc.build import REGISTRY, cuda_toolkit_version
return str(Path(torch.__file__).parent.resolve())
except Exception:
return os.environ.get("TORCH_HOME", "")
# Preflight: warn if nvcc major version != torch's bundled CUDA major version.
# A mismatch (e.g. nvcc 13.0 + cu128 torch) causes cryptic ABI/header errors.
nvcc_ver = cuda_toolkit_version()
torch_cuda = torch.version.cuda
def _python_include():
import sysconfig
return sysconfig.get_path("include")
def _python_soabi():
import sysconfig
ext = sysconfig.get_config_var("EXT_SUFFIX").lstrip(".")
return ext[: -len(".so")]
class _CMakeBuildExt(_build_ext):
def run(self):
src = Path(__file__).parent
build_dir = src / "build" / "cmake"
torch_home = _torch_prefix()
if not torch_home:
raise RuntimeError(
"torch not found; cannot build kernels. "
"Activate the environment or set TORCH_HOME."
)
nvcc_ver = _cuda_toolkit_version()
torch_cuda = _torch_cuda_version()
if nvcc_ver is not None and torch_cuda is not None:
torch_major = int(torch_cuda.split(".")[0])
if nvcc_ver[0] != torch_major:
if nvcc_ver[0] != int(torch_cuda.split(".")[0]):
warnings.warn(
f"CUDA version mismatch: nvcc is {nvcc_ver[0]}.{nvcc_ver[1]} "
f"but torch was built with CUDA {torch_cuda}. "
f"This may cause compilation errors. "
f"Install a matching torch wheel: "
f"pip install torch --index-url "
f"https://download.pytorch.org/whl/cu{nvcc_ver[0]}{nvcc_ver[1]}",
f"Install a matching torch wheel.",
stacklevel=2,
)
_torch_lib = torch.utils.cpp_extension.library_paths()[0]
cmake = shutil.which("cmake")
if cmake is None:
raise RuntimeError("cmake not found on PATH; install it to build kernels")
for name, info in REGISTRY.items():
ext_modules.append(
CUDAExtension(
f"astrai.extension.lib.{name}",
info["sources"],
extra_compile_args={
"cxx": info["cxx_flags"],
"nvcc": info["nvcc_flags"],
},
extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
)
)
# Parallel build — each extension is an independent ninja project, so we
# can compile them concurrently. BuildExtension compiles them serially by
# default; this subclass dispatches each extension to a subprocess.
# Set BUILD_PARALLEL=N to override (default: min(n_exts, 4)).
_single_ext = os.environ.get("ASTRAI_BUILD_SINGLE_EXT", "")
class ParallelBuildExtension(BuildExtension):
def build_extensions(self):
if _single_ext:
self.extensions = [e for e in self.extensions if e.name == _single_ext]
if not self.extensions:
return
super().build_extensions()
return
n = len(self.extensions)
max_workers = int(os.environ.get("BUILD_PARALLEL", 8))
if max_workers <= 1 or n <= 1:
super().build_extensions()
return
# Each subprocess gets its own build-temp / build-lib so the
# ninja files (build.ninja, .ninja_log) never race. The built
# .so files are then collected into the parent's build_lib so the
# normal setuptools copy steps (inplace / editable wheel) work.
names = [e.name for e in self.extensions]
env = {**os.environ, "BUILD_PARALLEL": "1"}
base = os.path.join("build", "parallel")
os.makedirs(base, exist_ok=True)
procs = {}
for i in range(0, len(names), max_workers):
batch = names[i : i + max_workers]
for name in batch:
e = {**env, "ASTRAI_BUILD_SINGLE_EXT": name}
tag = name.replace(".", "_")
subdir = os.path.join(base, tag)
cmd = [
sys.executable,
__file__,
"build_ext",
"--build-temp",
os.path.join(subdir, "temp"),
"--build-lib",
os.path.join(subdir, "lib"),
parallel = os.environ.get("BUILD_PARALLEL", "16")
cfg = [
cmake,
"-S",
str(src / "csrc"),
"-B",
str(build_dir),
f"-DTORCH_HOME={torch_home}",
f"-DPYTHON_INCLUDE_DIR={_python_include()}",
f"-DPY_SOABI={_python_soabi()}",
]
procs[name] = subprocess.Popen(
cmd, env=e, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
)
for name in batch:
out, _ = procs[name].communicate()
if procs[name].returncode != 0:
sys.stdout.write(out.decode())
raise RuntimeError(
f"parallel build failed for {name} "
f"(exit {procs[name].returncode})"
)
self._collect_extensions(
os.path.join(base, name.replace(".", "_"), "lib")
arch = os.environ.get("ASTRAI_CUDA_ARCH")
if not arch:
arch = _detect_cuda_arch()
if arch:
cfg.append(f"-DASTRAI_CUDA_ARCH={arch}")
subprocess.run(cfg, check=True)
subprocess.run(
[cmake, "--build", str(build_dir), "-j", parallel], check=True
)
def _collect_extensions(self, sub_lib):
src = os.path.join(sub_lib, "astrai", "extension", "lib")
if not os.path.isdir(src):
return
dst = os.path.join(self.build_lib, "astrai", "extension", "lib")
os.makedirs(dst, exist_ok=True)
for f in os.listdir(src):
if f.endswith(".so"):
shutil.copy2(os.path.join(src, f), os.path.join(dst, f))
cmdclass["build_ext"] = ParallelBuildExtension
def _cuda_toolkit_version():
import shutil
import subprocess
if not cmdclass:
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()
return tuple(int(x) for x in ver.split("."))
except Exception:
pass
return None
class _NullBuildExt(_build_ext):
def _detect_cuda_arch():
"""Detect real GPU compute capability via torch (nvidia-smi may be spoofed).
Returns something like ``"89"`` or ``"103"``, or ``None`` if unavailable.
"""
try:
import torch
if torch.cuda.is_available():
major, minor = torch.cuda.get_device_capability()
return f"{major}{minor}"
except Exception:
pass
return None
def _torch_cuda_version():
try:
import torch
return torch.version.cuda
except Exception:
return None
class _NullBuildExt(_build_ext):
def build_extensions(self):
pass
cmdclass = {}
if _should_build():
cmdclass["build_ext"] = _CMakeBuildExt
else:
cmdclass["build_ext"] = _NullBuildExt
setup(ext_modules=ext_modules, cmdclass=cmdclass)
setup(ext_modules=[], cmdclass=cmdclass)