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"]
|
|
|
|
_available: dict[str, bool] = {}
|
|
_modules: dict[str, object] = {}
|
|
|
|
for _name in KERNEL_NAMES:
|
|
try:
|
|
_mod = importlib.import_module(f".{_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)
|