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
+"""CUDA attention kernel wrappers with torch fallback.
-import logging
 
-import torch
+Public API:
-import torch.nn.functional as F
+    - ``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] = {}
+from astrai.extension.loader import KERNEL_NAMES, is_available
-_modules: dict[str, object] = {}
+from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
 
-for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
+__all__ = [
-    try:
+    "gqa_decode_attn",
-        _mod = importlib.import_module(f".{_name}", package=__package__)
+    "gqa_prefill_attn",
-        _available[_name] = True
+    "is_available",
-        _modules[_name] = _mod
+    "KERNEL_NAMES",
-    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)

astrai/extension/loader.py

@@ -0,0 +1,36 @@
+"""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 = ["gqa_decode_attn", "gqa_prefill_attn"]
+
+_available: dict[str, bool] = {}
+_modules: dict[str, object] = {}
+
+for _name in KERNEL_NAMES:
+    try:
+        _mod = importlib.import_module(f".{_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)

astrai/extension/ops.py

@@ -0,0 +1,86 @@
+"""GQA attention wrapper functions — one entry point per compiled kernel.
+
+Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
+available, otherwise falls back to ``torch`` SDPA.
+
+Add new kernel wrappers here; split into per-variant files only if this file
+grows large.
+"""
+
+import torch
+import torch.nn.functional as F
+
+from astrai.extension.loader import _available, _modules
+
+
+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:
+    """Reference attention via ``scaled_dot_product_attention``."""
+    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)

csrc/kernels/gqa_decode_attn.cu

@@ -21,13 +21,20 @@ static void dispatch_decode(GQAParams& p) {
         gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(p);
         return;
     }
-#endif
     // scalar fallback (per-KV-head, one warp per query head)
     int group_size = p.q_head / p.kv_head;
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
     dim3 block(32, group_size);
     dim3 grid(p.batch * p.kv_head);
     gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
+#else
+    // scalar fallback (per-KV-head, one warp per query head)
+    int group_size = p.q_head / p.kv_head;
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+    dim3 block(32, group_size);
+    dim3 grid(p.batch * p.kv_head);
+    gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
+#endif
 }
 
 torch::Tensor gqa_decode_attn(
@@ -74,6 +81,9 @@ torch::Tensor gqa_decode_attn(
     p.o = (bf16*)O.data_ptr();
 
     switch (p.head_dim) {
+        case 32:
+            dispatch_decode<32>(p);
+            break;
         case 64:
             dispatch_decode<64>(p);
             break;
@@ -85,7 +95,7 @@ torch::Tensor gqa_decode_attn(
             break;
         default:
             TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim,
-                        " (supported: 64, 128, 256)");
+                        " (supported: 32, 64, 128, 256)");
     }
     return O;
 }

csrc/kernels/gqa_mma_utils.cuh

@@ -58,10 +58,12 @@ __device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
 
 // XOR swizzle for shared-memory column at 8-bf16 chunk granularity.
 // Eliminates ldmatrix bank conflicts without LD padding: consecutive rows
-// land in distinct bank groups. swiz_col(d, r) = ((d>>3)^(r&7))<<3 | (d&7).
+// land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7).
-// Works for any d; aligned (d%8==0) simplifies to d ^ ((r&7)<<3).
+// mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for
-__device__ __forceinline__ int swiz_col(int d, int r) {
+// HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps
-    return ((d >> 3) ^ (r & 7)) << 3 | (d & 7);
+// existing HEAD_DIM>=64 call sites working unchanged.
+__device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
+    return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
 }
 
 // cp.async: copy 16 bytes (8 bf16) from global to shared memory directly,

csrc/kernels/gqa_prefill_attn.cu

@@ -68,6 +68,9 @@ torch::Tensor gqa_prefill_attn(
     p.o = (bf16*)O.data_ptr();
 
     switch (p.head_dim) {
+        case 32:
+            dispatch_prefill<32>(p);
+            break;
         case 64:
             dispatch_prefill<64>(p);
             break;
@@ -79,7 +82,7 @@ torch::Tensor gqa_prefill_attn(
             break;
         default:
             TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
-                        " (supported: 64,128,256)");
+                        " (supported: 32,64,128,256)");
     }
     return O;
 }

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -30,6 +30,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     constexpr int KT2 = BC / 16;       // P k-tiles (K=16 each)
     constexpr int DN8 = HEAD_DIM / 8;  // O n-tiles (N=8 each)
     constexpr int LD = HEAD_DIM;   // XOR swizzle (swiz_col) handles bank conflicts
+    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);  // chunk bits, stay within LD
 
     const int warp = threadIdx.x / 32;
     const int lane = threadIdx.x % 32;
@@ -61,12 +62,12 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
                 int qr = qrow0 + r;
                 bf16 qv = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d]
                                          : __float2bfloat16(0.0f);
-                sQ[r * LD + swiz_col(d, r)] = __hmul(qv, scale_bf16);
+                sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
             }
             __syncwarp();
         #pragma unroll
             for (int kt = 0; kt < KD; kt++)
-                ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l)]);
+                ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
         }
         __syncthreads();  // prevent next warp from overwriting sQ prematurely
     }
@@ -106,8 +107,8 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
                 int r = i / HEAD_DIM;
                 int d = i % HEAD_DIM;
                 int kc = kv0 + r;
-                cp_async_16(&sK[r * LD + swiz_col(d, r)], &p.k[kv_base + kc * HEAD_DIM + d]);
+                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.k[kv_base + kc * HEAD_DIM + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r)], &p.v[kv_base + kc * HEAD_DIM + d]);
+                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.v[kv_base + kc * HEAD_DIM + d]);
             }
             cp_async_commit();
             cp_async_wait_all();
@@ -116,9 +117,9 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
                 int r = i / HEAD_DIM, d = i % HEAD_DIM;
                 int kc = kv0 + r;
                 bf16 z = __float2bfloat16(0.0f);
-                sK[r * LD + swiz_col(d, r)] = (kc < p.kv_len)
+                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len)
                                      ? p.k[kv_base + kc * HEAD_DIM + d] : z;
-                sV[r * LD + swiz_col(d, r)] = (kc < p.kv_len)
+                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len)
                                      ? p.v[kv_base + kc * HEAD_DIM + d] : z;
             }
         }
@@ -137,7 +138,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
 #pragma unroll
             for (int kt = 0; kt < KD; kt++) {
                 unsigned b[2];
-                ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l)]);
+                ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
                 mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
             }
         }
@@ -218,7 +219,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
 #pragma unroll
             for (int dn8 = 0; dn8 < DN8; dn8++) {
                 unsigned b[2];
-                ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l)]);
+                ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
                 mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
             }
         }