ViperEkura
9f83d982ec
- CUDAExtension module names changed from astrai.extension.<name> to astrai.extension.lib.<name> - Compiled .so files now land in astrai/extension/lib/ instead of alongside Python source - loader.py imports from .lib.<name> subpackage - Add astrai/extension/lib/__init__.py to make lib a proper package - Separates compiled artifacts from Python source for cleaner directory structure
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
"""Dynamic discovery and loading of compiled CUDA kernel modules.
|
|
|
|
Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed
|
|
in this package directory. On import we try to load each one; kernels that
|
|
failed to build (or are running on a CPU-only machine) are marked unavailable
|
|
so the wrapper functions can fall back to ``torch`` SDPA.
|
|
"""
|
|
|
|
import importlib
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode", "rotary_emb"]
|
|
|
|
_available: dict[str, bool] = {}
|
|
_modules: dict[str, object] = {}
|
|
|
|
for _name in KERNEL_NAMES:
|
|
try:
|
|
_mod = importlib.import_module(f".lib.{_name}", package=__package__)
|
|
_available[_name] = True
|
|
_modules[_name] = _mod
|
|
except ImportError:
|
|
_available[_name] = False
|
|
_modules[_name] = None
|
|
|
|
|
|
def is_available(name: str) -> bool:
|
|
"""Return ``True`` if the compiled kernel ``name`` was loaded."""
|
|
return _available.get(name, False)
|
|
|
|
|
|
def get_module(name: str) -> object:
|
|
"""Return the loaded kernel module for ``name``, or ``None`` if unavailable."""
|
|
return _modules.get(name)
|