feat: support HEAD_DIM=32 and split extension into loader/ops

- add case 32 to decode/prefill dispatch switch
- fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory
- restructure decode dispatch to #ifndef/#else/#endif matching prefill
- split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API

astrai/extension/__init__.py

@@ -1,91 +1,19 @@
-import importlib
-import logging
+"""CUDA attention kernel wrappers with torch fallback.
 
-import torch
-import torch.nn.functional as F
+Public API:
+    - ``gqa_decode_attn`` — single-query decode attention
+    - ``gqa_prefill_attn`` — multi-query prefill attention
 
-logger = logging.getLogger(__name__)
+Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``)
+when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
+"""
 
-_available: dict[str, bool] = {}
-_modules: dict[str, object] = {}
+from astrai.extension.loader import KERNEL_NAMES, is_available
+from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
 
-for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
-    try:
-        _mod = importlib.import_module(f".{_name}", package=__package__)
-        _available[_name] = True
-        _modules[_name] = _mod
-    except ImportError:
-        _available[_name] = False
-        _modules[_name] = None
-
-
-def _expand_kv_heads(
-    k: torch.Tensor, v: torch.Tensor, q_head: int
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Expand K/V heads to match Q heads for GQA fallback."""
-    kv_head = k.size(1)
-    if kv_head == q_head:
-        return k, v
-    group = q_head // kv_head
-    k = k.repeat_interleave(group, dim=1)
-    v = v.repeat_interleave(group, dim=1)
-    return k, v
-
-
-def _torch_fallback(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None,
-    is_causal: bool,
-    scale: float | None,
-) -> torch.Tensor:
-    k, v = _expand_kv_heads(k, v, q.size(1))
-    attn_mask = mask[:, None, None, :] if mask is not None else None
-    return F.scaled_dot_product_attention(
-        q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
-    )
-
-
-def gqa_decode_attn(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    is_causal: bool = False,
-    causal_offset: int = 0,
-    scale: float | None = None,
-) -> torch.Tensor:
-    if _available["gqa_decode_attn"]:
-        return _modules["gqa_decode_attn"].gqa_decode_attn(
-            q,
-            k,
-            v,
-            mask=mask,
-            is_causal=is_causal,
-            causal_offset=causal_offset,
-            scale=scale,
-        )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
-
-
-def gqa_prefill_attn(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    is_causal: bool = False,
-    causal_offset: int = 0,
-    scale: float | None = None,
-) -> torch.Tensor:
-    if _available["gqa_prefill_attn"]:
-        return _modules["gqa_prefill_attn"].gqa_prefill_attn(
-            q,
-            k,
-            v,
-            mask=mask,
-            is_causal=is_causal,
-            causal_offset=causal_offset,
-            scale=scale,
-        )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
+__all__ = [
+    "gqa_decode_attn",
+    "gqa_prefill_attn",
+    "is_available",
+    "KERNEL_NAMES",
+]