ViperEkura
2bc4d2b8a8
- loader.py: lazy/cached import; is_available defers the actual load; get_module raises on unavailable
- ops/{attention,rotary,fp8}: use get_module instead of touching private _modules or their own _mod() cache
- package-data: ship astrai.extension.lib *.so in built wheels (non-editable installs previously lost every kernel)
84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
"""Dynamic discovery and loading of compiled CUDA kernel modules.
|
|
|
|
Each kernel is built by the CMake build in ``csrc/CMakeLists.txt`` into a
|
|
``.so`` placed in ``astrai/extension/lib/`` — the module name equals the
|
|
``.so`` name equals the pybind name (e.g. ``attn_decode``, defined via
|
|
``TORCH_EXTENSION_NAME``). ``KERNEL_NAMES`` is discovered automatically from
|
|
the ``.so`` files present, so adding a kernel to the CMake ``KERNELS``
|
|
registry needs no change here.
|
|
|
|
Loading is **lazy and centralized**: module names are discovered eagerly
|
|
(cheap glob), but each ``.so`` is imported on first use via the single
|
|
``get_module`` accessor, then cached. The wrapper modules (``ops/*.py``) never
|
|
touch the internals or keep their own caches — they call ``get_module(name)``
|
|
(or ``is_available(name)`` when a torch fallback is acceptable). A kernel that
|
|
failed to build (or is running on a CPU-only machine) is ``None`` in the cache,
|
|
so ``is_available`` returns ``False`` and ``get_module`` raises a clear error.
|
|
"""
|
|
|
|
import glob
|
|
import importlib
|
|
import logging
|
|
import os
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_LIB_DIR = os.path.join(os.path.dirname(__file__), "lib")
|
|
|
|
|
|
def _discover_kernel_names() -> list[str]:
|
|
"""Return the module names of the compiled kernel ``.so`` files in lib/."""
|
|
names: list[str] = []
|
|
for path in glob.glob(os.path.join(_LIB_DIR, "*.so")):
|
|
# strip the "<soabi>.so" suffix, e.g. attn_decode.cpython-312-...so
|
|
names.append(os.path.basename(path).split(".", 1)[0])
|
|
return sorted(names)
|
|
|
|
|
|
KERNEL_NAMES = _discover_kernel_names()
|
|
|
|
_available: dict[str, bool] = {}
|
|
_modules: dict[str, object] = {}
|
|
|
|
|
|
def _try_load(name: str) -> object:
|
|
"""Import and cache the ``name`` kernel module (lazy, one attempt).
|
|
|
|
Returns the module, or ``None`` if it is unavailable. Cached so each
|
|
``.so`` is imported at most once per process.
|
|
"""
|
|
if name not in _modules:
|
|
try:
|
|
_modules[name] = importlib.import_module(
|
|
f".lib.{name}", package=__package__
|
|
)
|
|
_available[name] = True
|
|
except ImportError:
|
|
logger.warning("kernel '%s' failed to import; marking unavailable", name)
|
|
_modules[name] = None
|
|
_available[name] = False
|
|
return _modules[name]
|
|
|
|
|
|
def is_available(name: str) -> bool:
|
|
"""Return ``True`` if the compiled kernel ``name`` could be loaded."""
|
|
if name not in _available:
|
|
_try_load(name)
|
|
return _available.get(name, False)
|
|
|
|
|
|
def get_module(name: str) -> object:
|
|
"""Return the loaded kernel module for ``name``, importing it on first use.
|
|
|
|
Raises ``RuntimeError`` if the kernel is unavailable (not built, or failed
|
|
to import) — callers that can tolerate a torch fallback should check
|
|
``is_available(name)`` first instead.
|
|
"""
|
|
mod = _try_load(name)
|
|
if mod is None:
|
|
raise RuntimeError(
|
|
f"CUDA kernel '{name}' is not available. "
|
|
f"Build with CSRC_KERNELS=true (or use the torch-native fallback)."
|
|
)
|
|
return mod
|