ViperEkura
41dcf0feb9
- PagedAttentionParams uses flat KV pool + req_to_token + kv_indptr/qo_indptr instead of page_table - MMA split-KV decode and split-Q prefill kernels with indirect ragged-batch addressing - Prefill kernel accepts 4D mask (causal-aware); decode kernel supports 2D mask - CudaBackend is inference-only: kv_cache=None raises, no torch fallback - benchmark.py: required --ckpt, --backend/--compare options - Parallel build isolates build-temp/build-lib per subprocess - Standalone test covers decode/prefill with mask, 27 cases pass
43 lines
1.2 KiB
Python
43 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",
|
|
]
|
|
|
|
_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)
|