feat: add optional CUDA kernel system (csrc/) + fused GQA decode attention

Structure:
  csrc/               -- .cu sources + build.py registry
  astrai/extension/   -- compiled .so + __init__.py (import dispatcher)
  setup.py            -- CUDAExtension from csrc/build.py REGISTRY

Control: CSRC_KERNELS=true|false env var at install time.
Fallback: astrai.extension.available dict for runtime detection.
This commit is contained in:
2026-07-06 12:09:58 +08:00
parent 2579658e15
commit e8e228d035
7 changed files with 201 additions and 2 deletions
+58
View File
@@ -0,0 +1,58 @@
import os
import sys
from pathlib import Path
from setuptools import setup
from setuptools.command.build_ext import build_ext as _build_ext
sys.path.insert(0, str(Path(__file__).parent))
os.makedirs("astrai/extension", exist_ok=True)
def _should_build():
force = os.environ.get("CSRC_KERNELS", "").strip().lower()
if force == "true":
return True
if force == "false":
return False
try:
import shutil
import torch
return shutil.which("nvcc") is not None and torch.cuda.is_available()
except Exception:
return False
ext_modules = []
cmdclass = {}
if _should_build():
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
from csrc.build import REGISTRY
_torch_lib = torch.utils.cpp_extension.library_paths()[0]
for name, info in REGISTRY.items():
ext_modules.append(
CUDAExtension(
f"astrai.extension.{name}",
info["sources"],
extra_compile_args={"cxx": ["-O3"], "nvcc": info["nvcc_flags"]},
extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
)
)
cmdclass["build_ext"] = BuildExtension
if not cmdclass:
class _NullBuildExt(_build_ext):
def build_extensions(self):
pass
cmdclass["build_ext"] = _NullBuildExt
setup(ext_modules=ext_modules, cmdclass=cmdclass)