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
+1
View File
@@ -16,6 +16,7 @@ from astrai.dataset import (
Store,
StoreFactory,
)
from astrai.extension import available
from astrai.factory import BaseFactory
from astrai.inference import (
GenerationRequest,
+13
View File
@@ -0,0 +1,13 @@
import importlib
import logging
logger = logging.getLogger(__name__)
available: dict[str, bool] = {}
for _name in ["gqa_decode_attn"]:
try:
importlib.import_module(f".{_name}", package=__package__)
available[_name] = True
except ImportError:
available[_name] = False