diff --git a/.gitignore b/.gitignore index bfa95b4..309443d 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file +!/README.md +# Allow extension modules (only source .py) +!/astrai/extension/**/*.py + +# Allow build files +!/setup.py +!/AGENTS.md diff --git a/astrai/__init__.py b/astrai/__init__.py index 49432c7..aac9fe0 100644 --- a/astrai/__init__.py +++ b/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, diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py new file mode 100644 index 0000000..a861be1 --- /dev/null +++ b/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 diff --git a/csrc/__init__.py b/csrc/__init__.py new file mode 100644 index 0000000..bcae4f2 --- /dev/null +++ b/csrc/__init__.py @@ -0,0 +1,2 @@ +# Source directory for CUDA kernels — build-time only. +# Compiled .so files live in astrAI/_ext/. diff --git a/csrc/build.py b/csrc/build.py new file mode 100644 index 0000000..d8310fe --- /dev/null +++ b/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") diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu new file mode 100644 index 0000000..c13a3b4 --- /dev/null +++ b/csrc/kernels/gqa_decode_attn.cu @@ -0,0 +1,86 @@ +#include +#include +#include +#include +#include + +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<<>>( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), + mask.data_ptr(), + reinterpret_cast(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)"); +} diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..21c1ec5 --- /dev/null +++ b/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)