ViperEkura
a5b238dd86
- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy - custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd - aten::linear and linear_backward dispatch on CUDA key, zero model changes - per-tensor scale or raw cast; single-GPU smoke loss matches bf16
44 lines
1.2 KiB
Python
44 lines
1.2 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",
|
|
"attn_paged_prefill",
|
|
"rotary_emb",
|
|
"fp8_mm",
|
|
]
|
|
|
|
_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)
|