refactor: separate extension ops and backends

astrai/extension/__init__.py

@@ -15,25 +15,25 @@ Each wrapper calls its compiled CUDA kernel directly. Fallback to torch
 SDPA is handled by the attention backend, not the wrapper functions.
 """
 
-from astrai.extension.attention_backend import (
+from astrai.extension.backend import (
     ATTN_BACKEND,
     AttentionBackend,
     AttentionBackendFactory,
     CudaBackend,
     FlashAttnBackend,
     TorchNativeBackend,
+    apply_rotary_emb,
     attention,
     attn_backend,
     get_backend,
 )
-from astrai.extension.attention_ops import (
+from astrai.extension.loader import KERNEL_NAMES, is_available
+from astrai.extension.ops import (
     TensorLayout,
     attn_decode,
     attn_paged_decode,
     attn_prefill,
 )
-from astrai.extension.loader import KERNEL_NAMES, is_available
-from astrai.extension.rotary_backend import apply_rotary_emb
 
 __all__ = [
     "ATTN_BACKEND",

astrai/extension/backend/__init__.py

@@ -0,0 +1,27 @@
+"""Backend selection, fallbacks, and execution policies."""
+
+from astrai.extension.backend.attention import (
+    ATTN_BACKEND,
+    AttentionBackend,
+    AttentionBackendFactory,
+    CudaBackend,
+    FlashAttnBackend,
+    TorchNativeBackend,
+    attention,
+    attn_backend,
+    get_backend,
+)
+from astrai.extension.backend.rotary import apply_rotary_emb
+
+__all__ = [
+    "ATTN_BACKEND",
+    "AttentionBackend",
+    "AttentionBackendFactory",
+    "CudaBackend",
+    "FlashAttnBackend",
+    "TorchNativeBackend",
+    "apply_rotary_emb",
+    "attention",
+    "attn_backend",
+    "get_backend",
+]

astrai/extension/backend/attention.py

@@ -32,7 +32,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 import contextvars
 import enum
 import functools
-import importlib
 import os
 import threading
 from abc import ABC, abstractmethod
@@ -43,13 +42,18 @@ import torch
 import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension.attention_ops import (
+from astrai.extension.loader import is_available
+from astrai.extension.ops.attention import (
     attn_paged_decode,
     attn_paged_prefill,
 )
-from astrai.extension.loader import is_available
 from astrai.factory import BaseFactory
 
+try:
+    import flash_attn as _flash_attn
+except Exception:
+    _flash_attn = None
+
 if TYPE_CHECKING:
     from astrai.inference.cache import KVCache
 
@@ -67,7 +71,7 @@ _current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
 def flash_attn_available() -> bool:
     if not torch.cuda.is_available():
         return False
-    fa = _get_flash_attn()
+    fa = _flash_attn
     if fa is None:
         return False
 
@@ -90,14 +94,6 @@ def flash_attn_available() -> bool:
         return False
 
 
-@functools.lru_cache(maxsize=1)
-def _get_flash_attn():
-    try:
-        return importlib.import_module("flash_attn")
-    except Exception:
-        return None
-
-
 class ATTN_BACKEND(enum.Enum):
     """Backend selector enum, mirroring ``torch.nn.attention.SDPBackend``."""
 
@@ -145,7 +141,7 @@ def _backend_supports(
         if q.dtype not in (torch.float16, torch.bfloat16):
             return False
         if fwd is not None:
-            return q.ndim == 3 and hasattr(_get_flash_attn(), "flash_attn_varlen_func")
+            return q.ndim == 3 and hasattr(_flash_attn, "flash_attn_varlen_func")
         if attn_mask is None or is_causal:
             return True
         return attn_mask.dim() == 4
@@ -680,7 +676,7 @@ class FlashAttnBackend(AttentionBackend):
                 "FlashAttnBackend does not support a custom attention mask; "
                 "use a causal mask or select TorchNativeBackend."
             )
-        fa = _get_flash_attn()
+        fa = _flash_attn
         if fa is None:
             raise RuntimeError(
                 "FlashAttnBackend requires the optional 'flash-attn' package. "
@@ -702,7 +698,7 @@ class FlashAttnBackend(AttentionBackend):
         kv_cache: "KVCache",
         layer_id: int,
     ) -> Tensor:
-        fa = _get_flash_attn()
+        fa = _flash_attn
         if fa is None or not hasattr(fa, "flash_attn_varlen_func"):
             raise RuntimeError("packed inference requires flash_attn_varlen_func")
         kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k

astrai/extension/backend/rotary.py

@@ -11,6 +11,7 @@ import torch
 from torch import Tensor
 
 from astrai.extension.loader import is_available
+from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary
 
 _cache = {"available": None}
 
@@ -48,7 +49,5 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
         and x.is_cuda
         and x.dtype == torch.bfloat16
     ):
-        from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
-
         return _cuda_rotary(x, freqs_cis)
     return _torch_apply(x, freqs_cis)

astrai/extension/fp8.py

@@ -1,8 +1,8 @@
 """FP8 training: scaling state and aten::linear dispatch.
 
-Layered (see also ``fp8_ops.py`` for the CUDA interface adapter):
+Layered (see also ``ops/fp8.py`` for the CUDA interface adapter):
 
-1. Kernel interface: "fp8_ops" — the only module touching the pybind.
+1. Kernel interface: ``ops.fp8`` - the only module touching the pybind.
 2. Training state (this module): per-tensor scales, amax history, delayed
    scaling, and the ``fp8_autocast`` context (TE-style, like
    ``torch.autocast``).
@@ -25,7 +25,7 @@ from contextlib import contextmanager
 import torch
 from torch.library import Library
 
-from astrai.extension.fp8_ops import (
+from astrai.extension.ops.fp8 import (
     linear_backward_scaled,
     linear_forward_scaled,
 )

astrai/extension/ops/__init__.py

@@ -0,0 +1,19 @@
+"""Stateless wrappers around compiled extension kernels."""
+
+from astrai.extension.ops.attention import (
+    TensorLayout,
+    attn_decode,
+    attn_paged_decode,
+    attn_paged_prefill,
+    attn_prefill,
+)
+from astrai.extension.ops.rotary import rotary_emb
+
+__all__ = [
+    "TensorLayout",
+    "attn_decode",
+    "attn_paged_decode",
+    "attn_paged_prefill",
+    "attn_prefill",
+    "rotary_emb",
+]

astrai/extension/ops/attention.py

@@ -1,4 +1,4 @@
-"""Attention kernel wrapper functions — one entry point per compiled kernel.
+"""Attention kernel wrapper functions - one entry point per compiled kernel.
 
 Each wrapper calls its CUDA kernel directly. If the kernel is not
 available, raises ``RuntimeError``. Fallback to torch SDPA is the

astrai/extension/ops/rotary.py

@@ -2,7 +2,7 @@
 
 Calls the compiled CUDA kernel directly. If the kernel is not available,
 raises ``RuntimeError``. Fallback to torch complex multiply is the
-responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
+responsibility of ``astrai.extension.backend.rotary.apply_rotary_emb``.
 
 Layout: x is packed [tokens, n_heads, head_dim] or dense
 [batch, seq_len, n_heads, head_dim]. ``freqs_cis`` has matching token axes.

astrai/inference/runtime/executor.py

@@ -7,7 +7,7 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
-from astrai.extension.attention_backend import (
+from astrai.extension.backend.attention import (
     CudaBackend,
     get_backend,
 )

astrai/model/components/__init__.py

@@ -1,4 +1,4 @@
-from astrai.extension.rotary_backend import apply_rotary_emb
+from astrai.extension.backend.rotary import apply_rotary_emb
 from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding

astrai/model/components/attention.py

@@ -5,8 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
-from astrai.extension import attention
-from astrai.extension.rotary_backend import apply_rotary_emb
+from astrai.extension.backend import apply_rotary_emb, attention
 from astrai.factory import BaseFactory
 from astrai.inference.cache import KVCache
 from astrai.model.components.linear import Linear