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:
2026-08-04 21:27:22 +08:00
parent cdf9145ecf
commit 49a9c6b3d2
6 changed files with 214 additions and 201 deletions
+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;