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:
@@ -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
|
||||
Reference in New Issue
Block a user