refactor: unify kernel module loading and packaging

- 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)

astrai/extension/loader.py

@@ -5,9 +5,15 @@ Each kernel is built by the CMake build in ``csrc/CMakeLists.txt`` into a
 ``.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.
+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
@@ -34,21 +40,44 @@ KERNEL_NAMES = _discover_kernel_names()
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}
 
-for _name in KERNEL_NAMES:
+
+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:
-        _mod = importlib.import_module(f".lib.{_name}", package=__package__)
-        _available[_name] = True
-        _modules[_name] = _mod
+            _modules[name] = importlib.import_module(
+                f".lib.{name}", package=__package__
+            )
+            _available[name] = True
         except ImportError:
-        _available[_name] = False
-        _modules[_name] = None
+            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`` was loaded."""
+    """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``, or ``None`` if unavailable."""
-    return _modules.get(name)
+    """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

astrai/extension/ops/attention.py

@@ -17,7 +17,7 @@ from typing import Optional
 
 import torch
 
-from astrai.extension.loader import _available, _modules
+from astrai.extension.loader import get_module
 
 
 class TensorLayout(enum.IntEnum):
@@ -30,14 +30,6 @@ class TensorLayout(enum.IntEnum):
     BLHD = 1  # [batch, seq_len, n_heads, head_dim]
 
 
-def _check_available(name: str):
-    if not _available.get(name):
-        raise RuntimeError(
-            f"CUDA kernel '{name}' is not available. "
-            f"Build with CSRC_KERNELS=true or use a torch-native backend."
-        )
-
-
 def attn_decode(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -57,9 +49,9 @@ def attn_decode(
     Returns:
         [batch, 1, n_heads, head_dim] (blhd, bf16)
     """
-    _check_available("attn_decode")
+    mod = get_module("attn_decode")
     causal_offset = (k.size(1) - 1) if is_causal else -1
-    return _modules["attn_decode"].attn_decode(
+    return mod.attn_decode(
         q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD
     )
 
@@ -83,9 +75,9 @@ def attn_prefill(
     Returns:
         [batch, q_len, n_heads, head_dim] (blhd, bf16)
     """
-    _check_available("attn_prefill")
+    mod = get_module("attn_prefill")
     causal_offset = (k.size(1) - q.size(1)) if is_causal else -1
-    return _modules["attn_prefill"].attn_prefill(
+    return mod.attn_prefill(
         q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD
     )
 
@@ -129,9 +121,9 @@ def attn_paged_decode(
     Returns:
         [batch, n_heads, head_dim] (bf16, 3D)
     """
-    _check_available("attn_paged_decode")
+    mod = get_module("attn_paged_decode")
     causal_offset = 0 if is_causal else -1
-    return _modules["attn_paged_decode"].attn_paged_decode(
+    return mod.attn_paged_decode(
         q,
         k_cache,
         v_cache,
@@ -183,9 +175,9 @@ def attn_paged_prefill(
     Returns:
         [total_q, n_heads, head_dim] (bf16, 3D)
     """
-    _check_available("attn_paged_prefill")
+    mod = get_module("attn_paged_prefill")
     causal_offset = 0 if is_causal else -1
-    return _modules["attn_paged_prefill"].attn_paged_prefill(
+    return mod.attn_paged_prefill(
         q,
         k_cache,
         v_cache,

astrai/extension/ops/fp8.py

@@ -19,27 +19,11 @@ this module is stateless.
 import torch
 from torch.library import custom_op
 
-from astrai.extension.loader import get_module, is_available
+from astrai.extension.loader import get_module
 
 # fmt string -> kernel int (0 = E4M3, 1 = E5M2)
 _FMT_TO_INT = {"e4m3": 0, "e5m2": 1}
 
-# The pybind module is loaded once at first use and cached: the loader
-# resolves modules at import time and never reloads them, so every call
-# after the first is a single None check.
-_MOD: object | None = None
-
-
-def _mod() -> object:
-    global _MOD
-    if _MOD is None:
-        if not is_available("fp8_ops"):
-            raise RuntimeError(
-                "CUDA kernel 'fp8_ops' is not available. Build with CSRC_KERNELS=true."
-            )
-        _MOD = get_module("fp8_ops")
-    return _MOD
-
 
 def _fmt_int(fmt: str) -> int:
     try:
@@ -72,7 +56,7 @@ def _fp8_quantize_fake(x, scale, fmt):
 def _fp8_quantize_cuda(x, scale, fmt):
     if x.dtype != torch.bfloat16:
         raise TypeError(f"fp8 quantize requires bf16 input, got {x.dtype}")
-    return _mod().quantize_bf16(x, scale, int(fmt))
+    return get_module("fp8_ops").quantize_bf16(x, scale, int(fmt))
 
 
 @fp8_quantize.register_kernel("cpu")
@@ -110,7 +94,7 @@ def _fp8_gemm_cuda(a, b, sa, sb, out_dtype=0, out_scale=None):
         raise TypeError(
             f"fp8 GEMM requires matching fp8 inputs, got {a.dtype}/{b.dtype}"
         )
-    return _mod().mm_fp8(a, b, sa, sb, int(out_dtype), out_scale)
+    return get_module("fp8_ops").mm_fp8(a, b, sa, sb, int(out_dtype), out_scale)
 
 
 @fp8_gemm.register_kernel("cpu")
@@ -165,7 +149,7 @@ def linear_forward_fp8(x, w, bias, sx, sw, fmt: str = "e4m3"):
         raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")
     if bias is None:
         bias = torch.empty(0, device=x.device, dtype=x.dtype)
-    return _mod().linear_forward_fp8(x, w, bias, sx, sw, _fmt_int(fmt))
+    return get_module("fp8_ops").linear_forward_fp8(x, w, bias, sx, sw, _fmt_int(fmt))
 
 
 def linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt: str = "e5m2"):
@@ -183,4 +167,6 @@ def linear_backward_fp8(g, x, w, masks, sg, sw, sx, fmt: str = "e5m2"):
         raise TypeError(
             f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
         )
-    return _mod().linear_backward_fp8(g, x, w, list(masks), sg, sw, sx, _fmt_int(fmt))
+    return get_module("fp8_ops").linear_backward_fp8(
+        g, x, w, list(masks), sg, sw, sx, _fmt_int(fmt)
+    )

astrai/extension/ops/rotary.py

@@ -10,15 +10,7 @@ Layout: x is packed [tokens, n_heads, head_dim] or dense
 
 import torch
 
-from astrai.extension.loader import _available, _modules
-
-
-def _check_available():
-    if not _available.get("rotary_emb"):
-        raise RuntimeError(
-            "CUDA kernel 'rotary_emb' is not available. "
-            "Build with CSRC_KERNELS=true or use the torch fallback."
-        )
+from astrai.extension.loader import get_module
 
 
 def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
@@ -31,9 +23,9 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     Returns:
         Tensor with the same shape as ``x``.
     """
-    _check_available()
+    mod = get_module("rotary_emb")
     if not x.is_contiguous():
         x = x.contiguous()
     if not freqs_cis.is_contiguous():
         freqs_cis = freqs_cis.contiguous()
-    return _modules["rotary_emb"].rotary_emb(x, freqs_cis)
+    return mod.rotary_emb(x, freqs_cis)

pyproject.toml

@@ -37,6 +37,9 @@ flash = ["flash-attn>=2.6"]
 [tool.setuptools.packages.find]
 where = ["."]
 
+[tool.setuptools.package-data]
+"astrai.extension.lib" = ["*.so"]
+
 [tool.setuptools.dynamic]
 version = { attr = "astrai.__version__" }