ViperEkura
16a55bb474
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/ - add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks) - split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu - fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads - update extension ops, loader, and kernel tests
55 lines
1.8 KiB
Python
55 lines
1.8 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. 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 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] = {}
|
|
|
|
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)
|