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
+29
View File
@@ -0,0 +1,29 @@
from pathlib import Path
def _arch_flag():
import torch
if torch.cuda.is_available():
cap = torch.cuda.get_device_capability()
ver = f"{cap[0]}{cap[1]}"
return f"-gencode=arch=compute_{ver},code=sm_{ver}"
return "-gencode=arch=compute_80,code=sm_80"
_kernels_dir = Path("csrc/kernels")
REGISTRY: dict[str, dict] = {}
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,
"nvcc_flags": ["-O3", "--expt-relaxed-constexpr", _arch_flag()],
"extra_link_args": kwargs.pop("extra_link_args", []),
**kwargs,
}
register("gqa_decode_attn")