feat: add optional CUDA kernel system (csrc/) + fused GQA decode attention

Structure:
  csrc/               -- .cu sources + build.py registry
  astrai/extension/   -- compiled .so + __init__.py (import dispatcher)
  setup.py            -- CUDAExtension from csrc/build.py REGISTRY

Control: CSRC_KERNELS=true|false env var at install time.
Fallback: astrai.extension.available dict for runtime detection.

.gitignore

@@ -7,8 +7,12 @@
 # Allow specific file types and root files
 !astrai/**/*.py
 !scripts/**/*.py
-!scripts/**/*.sh
 !tests/**/*.py
+!csrc/**/*.py
+
+!csrc/**/*.cu
+
+!scripts/**/*.sh
 
 # Allow GitHub files
 !/.github/**
@@ -22,4 +26,10 @@
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
-!/README.md
+!/README.md
+# Allow extension modules (only source .py)
+!/astrai/extension/**/*.py
+
+# Allow build files
+!/setup.py
+!/AGENTS.md

astrai/__init__.py

@@ -16,6 +16,7 @@ from astrai.dataset import (
     Store,
     StoreFactory,
 )
+from astrai.extension import available
 from astrai.factory import BaseFactory
 from astrai.inference import (
     GenerationRequest,

astrai/extension/__init__.py

@@ -0,0 +1,13 @@
+import importlib
+import logging
+
+logger = logging.getLogger(__name__)
+
+available: dict[str, bool] = {}
+
+for _name in ["gqa_decode_attn"]:
+    try:
+        importlib.import_module(f".{_name}", package=__package__)
+        available[_name] = True
+    except ImportError:
+        available[_name] = False

csrc/__init__.py

@@ -0,0 +1,2 @@
+# Source directory for CUDA kernels — build-time only.
+# Compiled .so files live in astrAI/_ext/.

csrc/build.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+
+def _arch_flag():
+    import torch
+
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability()
+        ver = f"{cap[0]}{cap[1]}"
+        return f"-gencode=arch=compute_{ver},code=sm_{ver}"
+    return "-gencode=arch=compute_80,code=sm_80"
+
+
+_kernels_dir = Path("csrc/kernels")
+REGISTRY: dict[str, dict] = {}
+
+
+def register(name: str, sources: list[str] | None = None, **kwargs):
+    if sources is None:
+        sources = [str(_kernels_dir / f"{name}.cu")]
+    REGISTRY[name] = {
+        "sources": sources,
+        "nvcc_flags": ["-O3", "--expt-relaxed-constexpr", _arch_flag()],
+        "extra_link_args": kwargs.pop("extra_link_args", []),
+        **kwargs,
+    }
+
+
+register("gqa_decode_attn")

csrc/kernels/gqa_decode_attn.cu

@@ -0,0 +1,86 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cmath>
+#include <cfloat>
+#include <torch/extension.h>
+
+using bf16 = __nv_bfloat16;
+
+__inline__ __device__ float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+__global__ void gqa_decode_attn_kernel(
+    const bf16* q_ptr, const bf16* k_ptr, const bf16* v_ptr,
+    const bool*  mask_ptr, bf16* out_ptr,
+    int B, int n_heads, int n_kv_heads, int seq_len, int hd
+) {
+    int batch = blockIdx.x / n_heads;
+    int q_head = blockIdx.x % n_heads;
+    int kv_head = q_head / (n_heads / n_kv_heads);
+    int tid = threadIdx.x;
+
+    float q_val = __bfloat162float(
+        q_ptr[((batch * n_heads + q_head) * 1) * hd + tid]);
+    int kv_base = ((batch * n_kv_heads + kv_head) * seq_len) * hd;
+    int mask_base = batch * seq_len;
+
+    float m = -FLT_MAX, d = 0.0f, acc = 0.0f;
+    __shared__ float smem[2];
+    float scale = 1.0f / sqrtf((float)hd);
+
+    for (int s = 0; s < seq_len; s++) {
+        int off = kv_base + s * hd + tid;
+        float partial = q_val * __bfloat162float(k_ptr[off]);
+        partial = warp_reduce_sum(partial) * scale;
+
+        if (tid % 32 == 0) smem[tid / 32] = partial;
+        __syncthreads();
+        if (tid == 0) smem[0] = smem[0] + smem[1];
+        __syncthreads();
+
+        float score = smem[0];
+        if (!mask_ptr[mask_base + s]) score = -FLT_MAX;
+
+        float new_m = fmaxf(m, score);
+        float alpha = expf(m - new_m);
+        float beta  = expf(score - new_m);
+        d = d * alpha + beta;
+        acc = acc * alpha + __bfloat162float(v_ptr[off]) * beta;
+        m = new_m;
+    }
+
+    int out_off = ((batch * n_heads + q_head) * 1) * hd + tid;
+    out_ptr[out_off] = __float2bfloat16(acc / d);
+}
+
+torch::Tensor gqa_decode_attn(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor mask
+) {
+    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && mask.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.dtype() == torch::kBFloat16);
+    TORCH_CHECK(v.dtype() == torch::kBFloat16);
+    TORCH_CHECK(mask.dtype() == torch::kBool);
+    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1");
+
+    int B = q.size(0), n_heads = q.size(1), n_kv = k.size(1);
+    int seq_len = k.size(2), hd = q.size(3);
+    auto out = torch::empty_like(q);
+
+    gqa_decode_attn_kernel<<<dim3(B * n_heads), dim3(hd)>>>(
+        reinterpret_cast<const bf16*>(q.data_ptr()),
+        reinterpret_cast<const bf16*>(k.data_ptr()),
+        reinterpret_cast<const bf16*>(v.data_ptr()),
+        mask.data_ptr<bool>(),
+        reinterpret_cast<bf16*>(out.data_ptr()),
+        B, n_heads, n_kv, seq_len, hd
+    );
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode attention (fused)");
+}

setup.py

@@ -0,0 +1,58 @@
+import os
+import sys
+from pathlib import Path
+
+from setuptools import setup
+from setuptools.command.build_ext import build_ext as _build_ext
+
+sys.path.insert(0, str(Path(__file__).parent))
+os.makedirs("astrai/extension", exist_ok=True)
+
+
+def _should_build():
+    force = os.environ.get("CSRC_KERNELS", "").strip().lower()
+    if force == "true":
+        return True
+    if force == "false":
+        return False
+    try:
+        import shutil
+
+        import torch
+
+        return shutil.which("nvcc") is not None and torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+ext_modules = []
+cmdclass = {}
+
+if _should_build():
+    import torch
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+    from csrc.build import REGISTRY
+
+    _torch_lib = torch.utils.cpp_extension.library_paths()[0]
+
+    for name, info in REGISTRY.items():
+        ext_modules.append(
+            CUDAExtension(
+                f"astrai.extension.{name}",
+                info["sources"],
+                extra_compile_args={"cxx": ["-O3"], "nvcc": info["nvcc_flags"]},
+                extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
+            )
+        )
+    cmdclass["build_ext"] = BuildExtension
+
+if not cmdclass:
+
+    class _NullBuildExt(_build_ext):
+        def build_extensions(self):
+            pass
+
+    cmdclass["build_ext"] = _NullBuildExt
+
+setup(ext_modules=ext_modules, cmdclass=cmdclass)