refactor: extension dispatch layer with CUDA/torch fallback

- Add gqa_decode_attn/gqa_prefill_attn dispatch functions
- Internal _available/__modules with underscore prefix
- CUDA kernel path with F.scaled_dot_product_attention fallback
- GQA head expansion in fallback path

.gitignore

@@ -7,8 +7,13 @@
 # Allow specific file types and root files
 !astrai/**/*.py
 !scripts/**/*.py
-!scripts/**/*.sh
 !tests/**/*.py
+!csrc/**/*.py
+
+!csrc/**/*.cu
+!csrc/**/*.cuh
+
+!scripts/**/*.sh
 
 # Allow GitHub files
 !/.github/**
@@ -23,3 +28,8 @@
 !/LICENSE
 !/pyproject.toml
 !/README.md
+# Allow extension modules (only source .py)
+!/astrai/extension/**/*.py
+
+# Allow build files
+!/setup.py

README.md

@@ -9,7 +9,7 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
@@ -59,8 +59,9 @@ End-to-end walkthrough in 5 steps:
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
+pip install -e .                                          # pure PyTorch (no CUDA kernels)
-# pip install -e ".[dev]"    # optional: dev dependencies (pytest, ruff)
+# CSRC_KERNELS=true pip install -e . --no-build-isolation  # optional: fused CUDA kernels
+# pip install -e ".[dev]"                                  # dev dependencies (pytest, ruff)
 ```
 
 **2. Download model**

assets/docs/README-zh-CN.md

@@ -15,7 +15,7 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
@@ -65,8 +65,9 @@
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
+pip install -e .                                          # 纯 PyTorch（不含 CUDA 内核）
-# pip install -e ".[dev]"    # 可选：开发依赖（pytest, ruff）
+# CSRC_KERNELS=true pip install -e . --no-build-isolation  # 可选：融合 CUDA 内核加速
+# pip install -e ".[dev]"                                  # 可选：开发依赖（pytest, ruff）
 ```
 
 **2. 下载模型**

astrai/extension/__init__.py

@@ -0,0 +1,91 @@
+import importlib
+import logging
+
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+_available: dict[str, bool] = {}
+_modules: dict[str, object] = {}
+
+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)

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,36 @@
+from pathlib import Path
+
+
+def _arch_flags() -> list[str]:
+    import torch
+
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability()
+    else:
+        cap = (8, 0)
+    ver = f"{cap[0]}{cap[1]}"
+    flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
+    # tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
+    # kernel dispatch at build time via this define rather than at runtime.
+    if cap[0] < 8:
+        flags.append("-DASTRAI_NO_MMA")
+    return flags
+
+
+_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_flags()],
+        "extra_link_args": kwargs.pop("extra_link_args", []),
+        **kwargs,
+    }
+
+
+register("gqa_decode_attn")
+register("gqa_prefill_attn")

csrc/kernels/gqa_common.cuh

@@ -0,0 +1,35 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cfloat>
+#include <algorithm>
+
+using bf16 = __nv_bfloat16;
+using std::min;
+
+constexpr int DC_CHUNK = 64;
+constexpr int Br = 32, Bc = 64;
+
+__device__ inline float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+struct GQAParams {
+    int batch;
+    int q_head;
+    int kv_head;
+    int q_len;
+    int kv_len;
+    int head_dim;
+    int use_mask;
+    int is_causal;
+    int causal_offset;
+    float scale;
+    const bf16* __restrict__ q;
+    const bf16* __restrict__ k;
+    const bf16* __restrict__ v;
+    const bool* __restrict__ mask;
+    bf16* __restrict__ o;
+};

csrc/kernels/gqa_decode_attn.cu

@@ -0,0 +1,66 @@
+#include "gqa_decode_attn.cuh"
+#include <torch/extension.h>
+
+torch::Tensor gqa_decode_attn(
+    torch::Tensor q,
+    torch::Tensor k, 
+    torch::Tensor v,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal = false, 
+    int64_t causal_offset = 0,
+    c10::optional<double> scale = c10::nullopt
+) {
+    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.dtype() == torch::kBFloat16);
+    TORCH_CHECK(v.dtype() == torch::kBFloat16);
+    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1");
+
+    GQAParams p;
+    p.batch = q.size(0);
+    p.q_head = q.size(1);
+    p.kv_head = k.size(1);
+    p.q_len = 1;
+    p.kv_len = k.size(2);
+    p.head_dim = q.size(3);
+    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
+    p.use_mask = mask.has_value();
+    p.is_causal = (int)is_causal;
+    p.causal_offset = (int)causal_offset;
+    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.q = (const bf16*)q.data_ptr();
+    p.k = (const bf16*)k.data_ptr();
+    p.v = (const bf16*)v.data_ptr();
+    if (p.use_mask) {
+        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        TORCH_CHECK(mask.value().dim() == 2);
+        TORCH_CHECK(mask.value().size(0) == p.batch);
+        TORCH_CHECK(mask.value().size(1) == p.kv_len);
+        p.mask = mask.value().data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+    }
+
+    auto O = torch::empty_like(q);
+    p.o = (bf16*)O.data_ptr();
+
+    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);
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_decode_attn", &gqa_decode_attn,
+        py::arg("q"),
+        py::arg("k"),
+        py::arg("v"),
+        py::arg("mask") = py::none(),
+        py::arg("is_causal") = false,
+        py::arg("causal_offset") = 0,
+        py::arg("scale") = py::none(),
+        "GQA decode (per-KV-head, shared K)");
+}

csrc/kernels/gqa_decode_attn.cuh

@@ -0,0 +1,59 @@
+#pragma once
+#include "gqa_common.cuh"
+
+__global__ void gqa_decode_attn_kernel(GQAParams p) {
+    int batch = blockIdx.x / p.kv_head;
+    int kv_head = blockIdx.x % p.kv_head;
+    int group_size = blockDim.y;
+    int q_head = kv_head * group_size + threadIdx.y;
+    int lane = threadIdx.x;
+    int hd_per_thread = p.head_dim / 32;
+
+    float q_reg[8];
+    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++)
+        q_reg[i] = __bfloat162float(p.q[q_off + i]);
+
+    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
+    int mask_base = batch * p.kv_len;
+
+    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
+
+    extern __shared__ __align__(16) bf16 k_smem[];
+
+    for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) {
+        int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
+
+        int total = this_chunk * p.head_dim;
+        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
+            k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i];
+        __syncthreads();
+
+        for (int s = 0; s < this_chunk; s++) {
+            float partial = 0.0f;
+            for (int i = 0; i < hd_per_thread; i++)
+                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
+            partial = warp_reduce_sum(partial) * p.scale;
+
+            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
+                partial = -FLT_MAX;
+            if (p.is_causal && (chunk_start + s) > p.causal_offset)
+                partial = -FLT_MAX;
+
+            float new_m = fmaxf(m, partial);
+            float alpha = expf(m - new_m);
+            float beta  = expf(partial - new_m);
+            d = d * alpha + beta;
+
+            int v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread;
+            for (int i = 0; i < hd_per_thread; i++)
+                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta;
+            m = new_m;
+        }
+        __syncthreads();
+    }
+
+    int out_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++)
+        p.o[out_off + i] = __float2bfloat16(acc_reg[i] / d);
+}

csrc/kernels/gqa_prefill_attn.cu

@@ -0,0 +1,96 @@
+#include "gqa_prefill_attn.cuh"
+#include <torch/extension.h>
+
+#ifndef ASTRAI_NO_MMA
+#include "gqa_prefill_attn_mma.cuh"
+#endif
+
+template <int HEAD_DIM>
+static void dispatch_prefill(GQAParams& p) {
+#ifndef ASTRAI_NO_MMA
+    constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM + 8;
+    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
+    dim3 block(WARPS * 32, 1, 1);
+    int smem = (2 * BC * LD + WARPS * BR * LD) * (int)sizeof(bf16);
+    cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC>,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block, smem>>>(p);
+#else
+    constexpr int G = 8, ROWS = 32, P_BC = 32;
+    dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
+    dim3 block(G, ROWS, 1);
+    size_t smem = 2 * P_BC * HEAD_DIM * sizeof(bf16);
+    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block, smem>>>(p);
+#endif
+}
+
+torch::Tensor gqa_prefill_attn(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal = false,
+    int64_t causal_offset = 0,
+    c10::optional<double> scale = c10::nullopt
+) {
+    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.dtype() == torch::kBFloat16);
+    TORCH_CHECK(v.dtype() == torch::kBFloat16);
+
+    GQAParams p;
+    p.batch = q.size(0);
+    p.q_head = q.size(1);
+    p.kv_head = k.size(1);
+    p.q_len = q.size(2);
+    p.kv_len = k.size(2);
+    p.head_dim = q.size(3);
+    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
+    p.use_mask = mask.has_value();
+    p.is_causal = (int)is_causal;
+    p.causal_offset = (int)causal_offset;
+    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.q = (const bf16*)q.data_ptr();
+    p.k = (const bf16*)k.data_ptr();
+    p.v = (const bf16*)v.data_ptr();
+    if (p.use_mask) {
+        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        TORCH_CHECK(mask.value().dim() == 2);
+        TORCH_CHECK(mask.value().size(0) == p.batch);
+        TORCH_CHECK(mask.value().size(1) == p.kv_len);
+        p.mask = mask.value().data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+    }
+
+    auto O = torch::empty_like(q);
+    p.o = (bf16*)O.data_ptr();
+
+    switch (p.head_dim) {
+        case 64:
+            dispatch_prefill<64>(p);
+            break;
+        case 128:
+            dispatch_prefill<128>(p);
+            break;
+        case 256:
+            dispatch_prefill<256>(p);
+            break;
+        default:
+            TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
+                        " (supported: 64,128,256)");
+    }
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_prefill_attn", &gqa_prefill_attn,
+        py::arg("q"),
+        py::arg("k"),
+        py::arg("v"),
+        py::arg("mask") = py::none(),
+        py::arg("is_causal") = false,
+        py::arg("causal_offset") = 0,
+        py::arg("scale") = py::none(),
+        "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
+}

csrc/kernels/gqa_prefill_attn.cuh

@@ -0,0 +1,137 @@
+#pragma once
+#include "gqa_common.cuh"
+
+// v9: group-split register blocking. G threads cooperate on one query row,
+// each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps
+// occupancy high; the S dot product is reduced across the G-lane group with a
+// short shuffle chain (log2(G) shuffles) instead of a full 32-lane warp reduce.
+// Online (per-kv) softmax — cheap because acc[] is only HEAD_DIM/G long.
+// Templated on <HEAD_DIM, G, ROWS, P_BC>. Block = (G, ROWS). G power-of-two,
+// G*ROWS a multiple of 32 with groups warp-aligned.
+
+template <int G>
+__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
+#pragma unroll
+    for (int o = G / 2; o > 0; o >>= 1)
+        v += __shfl_xor_sync(mask, v, o);
+    return v;
+}
+
+// load 8 contiguous bf16 from (16-byte aligned) smem as one float4, unpack to
+// 8 floats — cuts shared-load instructions 8x vs scalar bf16 loads.
+__device__ __forceinline__ void ld8(const bf16* p, float* o) {
+    float4 raw = *reinterpret_cast<const float4*>(p);
+    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&raw);
+#pragma unroll
+    for (int j = 0; j < 4; j++) {
+        float2 f = __bfloat1622float2(h[j]);
+        o[2 * j] = f.x;
+        o[2 * j + 1] = f.y;
+    }
+}
+
+template <int HEAD_DIM, int G, int ROWS, int P_BC>
+__global__ void gqa_prefill_attn_kernel_t(GQAParams p) {
+    constexpr int DPT = HEAD_DIM / G;
+
+    int q_tile = blockIdx.x;
+    int q_head = blockIdx.y;
+    int batch  = blockIdx.z;
+    int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
+    int row    = threadIdx.y;  // 0..ROWS-1
+    int q_row  = q_tile * ROWS + row;
+
+    int kv_head = q_head / (p.q_head / p.kv_head);
+
+    extern __shared__ __align__(16) bf16 smem[];
+    bf16* sK = smem;
+    bf16* sV = sK + P_BC * HEAD_DIM;
+
+    float qreg[DPT];
+    if (q_row < p.q_len) {
+        int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+#pragma unroll
+        for (int i = 0; i < DPT; i++)
+            qreg[i] = __bfloat162float(p.q[q_off + i]) * p.scale;
+    }
+
+    float m = -FLT_MAX, l = 0.0f;
+    float acc[DPT];
+#pragma unroll
+    for (int i = 0; i < DPT; i++)
+        acc[i] = 0.0f;
+
+    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
+    int tiles   = (p.kv_len + P_BC - 1) / P_BC;
+    int tt      = G * ROWS;
+    int lid     = row * G + gpos;
+
+    // per-group shuffle mask: only the G lanes of this row's group participate,
+    // so causal masking (differing loop bounds across rows in a warp) is safe.
+    int lane_in_warp = lid & 31;
+    unsigned gmask = (G == 32) ? 0xFFFFFFFFu
+                               : (((1u << G) - 1u) << (lane_in_warp & ~(G - 1)));
+
+    for (int ti = 0; ti < tiles; ti++) {
+        int kv0  = ti * P_BC;
+        int tlen = min(P_BC, p.kv_len - kv0);
+
+        for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
+            int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM);
+            sK[i] = p.k[gidx];
+            sV[i] = p.v[gidx];
+        }
+        __syncthreads();
+
+        int lim = tlen;
+        if (p.is_causal && q_row < p.q_len) {
+            int ep = q_row + p.causal_offset + 1;
+            if (kv0 >= ep)
+                lim = 0;
+            else if (kv0 + tlen > ep)
+                lim = ep - kv0;
+        }
+
+        for (int s = 0; s < lim; s++) {
+            const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
+            float part = 0.0f;
+#pragma unroll
+            for (int i = 0; i < DPT; i += 8) {
+                float k8[8];
+                ld8(kr + i, k8);
+#pragma unroll
+                for (int j = 0; j < 8; j++)
+                    part = fmaf(qreg[i + j], k8[j], part);
+            }
+            float dot = group_reduce_sum<G>(part, gmask);
+
+            if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s])
+                dot = -FLT_MAX;
+
+            float nm = fmaxf(m, dot);
+            float al = __expf(m - nm);
+            float be = __expf(dot - nm);
+            l = l * al + be;
+
+            const bf16* vr = sV + s * HEAD_DIM + gpos * DPT;
+#pragma unroll
+            for (int i = 0; i < DPT; i += 8) {
+                float v8[8];
+                ld8(vr + i, v8);
+#pragma unroll
+                for (int j = 0; j < 8; j++)
+                    acc[i + j] = fmaf(v8[j], be, acc[i + j] * al);
+            }
+            m = nm;
+        }
+        __syncthreads();
+    }
+
+    if (q_row < p.q_len) {
+        int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+        float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
+#pragma unroll
+        for (int i = 0; i < DPT; i++)
+            p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
+    }
+}

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -0,0 +1,244 @@
+#pragma once
+#include "gqa_common.cuh"
+
+// Tensor-core prefill, register-resident flash attention (raw mma.sync PTX).
+// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
+// cores via mma.sync.m16n8k16 (f32 accumulate). Q stays resident in registers;
+// S, O, and the online-softmax stats (m, l) live in registers too — nothing is
+// staged through shared memory except the cooperatively-loaded K/V tiles. The
+// mma fragment layout is used directly: the S accumulator (f32) maps element-
+// for-element onto the P matrix_a (bf16) operand, so softmax needs no shuffle
+// repack; row reductions fold across the 4-lane thread group. Templated on
+// <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
+
+// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
+// (only compiled when ASTRAI_HAS_MMA is set, i.e. built for sm_80+)
+__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
+                                          const unsigned* b, const float* c) {
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
+        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+}
+
+// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
+__device__ __forceinline__ unsigned ld2(const bf16* p) {
+    return *reinterpret_cast<const unsigned*>(p);
+}
+__device__ __forceinline__ unsigned pk2(float a, float b) {
+    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
+    return *reinterpret_cast<unsigned*>(&v);
+}
+// pack two (non-contiguous) bf16 into one .b32
+__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
+    __nv_bfloat162 v;
+    v.x = a;
+    v.y = b;
+    return *reinterpret_cast<unsigned*>(&v);
+}
+
+// ldmatrix: cooperatively load mma fragments from smem (one instruction per
+// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
+// scalar per-thread fragment packing, cutting shared-load instructions and bank
+// conflicts. Each lane supplies the shared address of one 8-wide row.
+__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}
+
+template <int HEAD_DIM, int WARPS, int BC>
+__global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
+    constexpr int BR = 16;
+    constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
+    constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
+    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 + 8;   // padded smem row stride (kills ldmatrix
+                                       // bank conflicts: consecutive rows land
+                                       // in distinct banks instead of colliding)
+
+    const int warp = threadIdx.x / 32;
+    const int lane = threadIdx.x % 32;
+    const int gid = lane >> 2;   // 0..7  → rows gid, gid+8
+    const int tid4 = lane & 3;   // 0..3
+    const int nthreads = WARPS * 32;
+
+    const int q_head = blockIdx.y;
+    const int batch = blockIdx.z;
+    const int kv_head = q_head / (p.q_head / p.kv_head);
+    const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
+
+    extern __shared__ __align__(16) bf16 smem[];
+    bf16* sK = smem;                                          // [BC][LD]
+    bf16* sV = sK + BC * LD;                                  // [BC][LD]
+    bf16* sQ = sV + BC * LD + warp * (BR * LD);               // per-warp [BR][LD]
+
+    // stage Q into smem (zero-padded past q_len)
+    const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+    for (int i = lane; i < BR * HEAD_DIM; i += 32) {
+        int r = i / HEAD_DIM, d = i % HEAD_DIM;
+        int qr = qrow0 + r;
+        sQ[r * LD + d] = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d] : __float2bfloat16(0.0f);
+    }
+    __syncwarp();
+
+    // Q resident A-fragments: Qa[kt][0..3] (loaded once via ldmatrix.x4)
+    unsigned Qa[KD][4];
+    int qrow_l = (lane & 7) + (lane & 8);       // 0..15
+    int qcol_l = (lane & 16) ? 8 : 0;
+#pragma unroll
+    for (int kt = 0; kt < KD; kt++)
+        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + kt * 16 + qcol_l]);
+
+    float Oacc[DN8][4];
+#pragma unroll
+    for (int j = 0; j < DN8; j++)
+        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
+    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
+
+    const int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
+    const int tiles = (p.kv_len + BC - 1) / BC;
+    const int qr0 = qrow0 + gid;      // row for c0/c1
+    const int qr1 = qrow0 + gid + 8;  // row for c2/c3
+
+    for (int ti = 0; ti < tiles; ti++) {
+        int kv0 = ti * BC;
+
+        for (int i = threadIdx.x; i < BC * HEAD_DIM; i += nthreads) {
+            int r = i / HEAD_DIM, d = i % HEAD_DIM;
+            int kc = kv0 + r;
+            bf16 z = __float2bfloat16(0.0f);
+            sK[r * LD + d] = (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
+            sV[r * LD + d] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
+        }
+        __syncthreads();
+
+        // S = Q @ K^T  → Sacc[n8][0..3]   (n8: 8 kv cols each)
+        float Sacc[NC8][4];
+#pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++) {
+            Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
+            int kv = kv0 + n8 * 8 + gid;
+            int krow_l = n8 * 8 + (lane & 7);        // kv within tile
+            int kcol_h = (lane & 8) ? 8 : 0;         // which k-half
+#pragma unroll
+            for (int kt = 0; kt < KD; kt++) {
+                unsigned b[2];
+                ldmatrix_x2(b, &sK[krow_l * LD + kt * 16 + kcol_h]);
+                mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
+            }
+            (void)kv;
+        }
+
+        // ---- online softmax (in registers) ----
+        // scale + mask, then per-row (gid, gid+8) max over held cols
+        float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
+#pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++) {
+            int cc = kv0 + n8 * 8 + 2 * tid4;  // col for c0/c2
+            bool bc0 = (cc >= p.kv_len) ||
+                       (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc]);
+            bool bc1 = (cc + 1 >= p.kv_len) ||
+                       (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc + 1]);
+            bool cz = p.is_causal;
+            int off = p.causal_offset;
+            bool bad0 = bc0 || (cz && cc > qr0 + off);
+            bool bad1 = bc1 || (cz && (cc + 1) > qr0 + off);
+            bool bad2 = bc0 || (cz && cc > qr1 + off);
+            bool bad3 = bc1 || (cz && (cc + 1) > qr1 + off);
+            float s0 = bad0 ? -FLT_MAX : Sacc[n8][0] * p.scale;
+            float s1 = bad1 ? -FLT_MAX : Sacc[n8][1] * p.scale;
+            float s2 = bad2 ? -FLT_MAX : Sacc[n8][2] * p.scale;
+            float s3 = bad3 ? -FLT_MAX : Sacc[n8][3] * p.scale;
+            Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
+            rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
+            rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
+        }
+        // reduce max across the 4-lane group (tid4)
+        rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
+        rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
+        rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
+        rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
+
+        float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
+        float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0);
+        float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1);
+
+        float rsum0 = 0.0f, rsum1 = 0.0f;
+#pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++) {
+            float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
+            float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
+            float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
+            float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
+            Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
+            rsum0 += p0 + p1;
+            rsum1 += p2 + p3;
+        }
+        rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
+        rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
+        rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
+        rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
+        l0 = l0 * corr0 + rsum0;
+        l1 = l1 * corr1 + rsum1;
+        m0 = nm0; m1 = nm1;
+
+        // rescale O accumulator by per-row correction
+#pragma unroll
+        for (int j = 0; j < DN8; j++) {
+            Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
+            Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
+        }
+
+        // O += P @ V
+#pragma unroll
+        for (int kt2 = 0; kt2 < KT2; kt2++) {
+            unsigned Pa[4];
+            Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
+            Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
+            Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
+            Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
+            int vrow_l = kt2 * 16 + (lane & 15);  // kv within tile (0..15)
+#pragma unroll
+            for (int dn8 = 0; dn8 < DN8; dn8++) {
+                unsigned b[2];
+                ldmatrix_x2_trans(b, &sV[vrow_l * LD + dn8 * 8]);
+                mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
+            }
+        }
+        __syncthreads();  // sK/sV reused next tile
+    }
+
+    // ---- write output ----
+    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
+    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
+    const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+#pragma unroll
+    for (int dn8 = 0; dn8 < DN8; dn8++) {
+        int d = dn8 * 8 + 2 * tid4;
+        if (qr0 < p.q_len) {
+            p.o[o_base + qr0 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][0] * rl0);
+            p.o[o_base + qr0 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
+        }
+        if (qr1 < p.q_len) {
+            p.o[o_base + qr1 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][2] * rl1);
+            p.o[o_base + qr1 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
+        }
+    }
+}

csrc/tests/gqa_decode_test.cu

@@ -0,0 +1,124 @@
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
+#include <cstdio>
+#include <cstdlib>
+#include <cmath>
+#include <sys/time.h>
+#include "../kernels/gqa_decode_attn.cuh"
+
+static double now_ms() {
+    struct timeval tv;
+    gettimeofday(&tv, NULL);
+    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
+}
+
+static void cpu_decode(const float* Q, const float* K, const float* V,
+                        const bool* mask, float* O,
+                        int B, int Hq, int Hk, int seq_len, int D) {
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hk;
+    for (int b = 0; b < B; b++) {
+        for (int h = 0; h < Hq; h++) {
+            int kv_h = h / n_rep;
+            float mv = -INFINITY, sv = 0.0f;
+            float accum[256] = {0};
+            for (int s = 0; s < seq_len; s++) {
+                if (!mask[b * seq_len + s]) continue;
+                float dot = 0.0f;
+                for (int d = 0; d < D; d++)
+                    dot += Q[((b * Hq + h) * 1 + 0) * D + d]
+                         * K[((b * Hk + kv_h) * seq_len + s) * D + d];
+                dot *= scale;
+                float nm = fmaxf(mv, dot);
+                float al = expf(mv - nm);
+                float be = expf(dot - nm);
+                sv = sv * al + be;
+                for (int d = 0; d < D; d++)
+                    accum[d] = accum[d] * al
+                             + V[((b * Hk + kv_h) * seq_len + s) * D + d] * be;
+                mv = nm;
+            }
+            float inv = 1.0f / sv;
+            for (int d = 0; d < D; d++)
+                O[((b * Hq + h) * 1 + 0) * D + d] = accum[d] * inv;
+        }
+    }
+}
+
+static bf16 f2bf(float x) { return __float2bfloat16(x); }
+static float bf2f(bf16 x) { return __bfloat162float(x); }
+static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
+
+int main() {
+    const int configs[][5] = {
+        {1, 2, 1, 64, 32},    // B,Hq,Hk,seq_len,D
+        {1, 32, 4, 512, 128},
+        {1, 32, 4, 1024, 128},
+    };
+    int n_cfgs = sizeof(configs) / sizeof(configs[0]);
+
+    for (int ci = 0; ci < n_cfgs; ci++) {
+        int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
+        int sl = configs[ci][3], D = configs[ci][4], gs = Hq / Hk;
+        printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d ===\n", B,Hq,Hk,sl,D,gs);
+
+        size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
+        float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
+        for (size_t i=0;i<nQ;i++) hQ[i]=randf();
+        for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+        bool* hMask=new bool[B*sl];
+        for (int i=0;i<B*sl;i++) hMask[i]=true;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        bool* dMask;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        cudaMalloc(&dMask,B*sl);
+
+        tmp=new bf16[max(nQ,nKV)];
+        for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+        cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
+
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+        p.use_mask=1; p.is_causal=0; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
+
+        size_t smem=DC_CHUNK*D*sizeof(bf16);
+        dim3 block(32, gs);
+        dim3 grid(B*Hk);
+        printf("grid=(%d,1,1) block=(%d,%d,1) smem=%zu\n",
+               grid.x, block.x, block.y, smem);
+
+        double t0=now_ms();
+        gqa_decode_attn_kernel<<<grid,block,smem>>>(p);
+        cudaDeviceSynchronize();
+        double kms=now_ms()-t0;
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+        bf16* hOut=new bf16[nQ];
+        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+        float* ref=new float[nQ];
+        cpu_decode(hQ,hK,hV,hMask,ref,B,Hq,Hk,sl,D);
+
+        float max_err=0;
+        for (size_t i=0;i<nQ;i++){
+            float d=fabsf(bf2f(hOut[i])-ref[i]);
+            if(d>max_err) max_err=d;
+        }
+        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
+        delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
+    }
+    printf("All tests passed!\n");
+    return 0;
+}

csrc/tests/gqa_prefill_test.cu

@@ -0,0 +1,127 @@
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
+#include <cstdio>
+#include <cstdlib>
+#include <cmath>
+#include <sys/time.h>
+#include "../kernels/gqa_prefill_attn.cuh"
+
+static double now_ms() {
+    struct timeval tv;
+    gettimeofday(&tv, NULL);
+    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
+}
+
+static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
+                          int B, int Hq, int Hk, int q_len, int kv_len, int D,
+                          int is_causal, int causal_off) {
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hk;
+    for (int b = 0; b < B; b++) {
+        for (int h = 0; h < Hq; h++) {
+            for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
+                float mv = -INFINITY, sv = 0.0f;
+                float accum[256] = {0};
+                int lim = is_causal ? min(kv_len, qi + causal_off + 1) : kv_len;
+                for (int kj = 0; kj < lim; kj++) {
+                    float dot = 0.0f;
+                    for (int d = 0; d < D; d++)
+                        dot += Q[((b*Hq + h)*q_len + qi)*D + d]
+                             * K[((b*Hk + kv_h)*kv_len + kj)*D + d];
+                    dot *= scale;
+                    float nm = fmaxf(mv, dot);
+                    float al = expf(mv - nm);
+                    float be = expf(dot - nm);
+                    sv = sv * al + be;
+                    for (int d = 0; d < D; d++)
+                        accum[d] = accum[d] * al
+                                 + V[((b*Hk + kv_h)*kv_len + kj)*D + d] * be;
+                    mv = nm;
+                }
+                float inv = 1.0f / sv;
+                for (int d = 0; d < D; d++)
+                    O[((b*Hq + h)*q_len + qi)*D + d] = accum[d] * inv;
+            }
+        }
+    }
+}
+
+static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
+static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
+static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
+
+int main() {
+    const int configs[][7] = {
+        {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
+        {1,32,4,512,512,128,0},  // standard
+        {1,32,4,128,256,128,0},  // medium
+        {1,4,2,256,256,128,1},   // causal
+    };
+    int n_configs = sizeof(configs) / sizeof(configs[0]);
+
+    for (int ci = 0; ci < n_configs; ci++) {
+        int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
+        int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
+        int causal=configs[ci][6];
+        printf("=== B=%d Hq=%d Hk=%d q=%d kv=%d D=%d causal=%d ===\n",
+               B,Hq,Hk,ql,kl,D,causal);
+
+        size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*D;
+        float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
+        for (size_t i=0;i<nQ;i++) hQ[i]=randf();
+        for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        tmp=new bf16[max(nQ,nKV)];
+        for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+
+        constexpr int G=8, ROWS=32, P_BC=32;
+        dim3 grid((ql+ROWS-1)/ROWS, Hq, B);
+        dim3 block(G, ROWS, 1);
+        size_t smem=2*P_BC*D*sizeof(bf16);
+        printf("grid=(%d,%d,%d) block=(%d,%d,%d) smem=%zu\n",
+               grid.x,grid.y,grid.z, block.x,block.y,block.z, smem);
+
+        double t0=now_ms();
+        switch (D) {
+            case 64:  gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            default: printf("unsupported D=%d\n",D); return 1;
+        }
+        cudaDeviceSynchronize();
+        double kms=now_ms()-t0;
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+        bf16* hOut=new bf16[nQ];
+        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+        float* ref=new float[nQ];
+        cpu_attention(hQ,hK,hV,ref,B,Hq,Hk,ql,kl,D,causal,0);
+
+        float max_err=0;
+        for (size_t i=0;i<nQ;i++) {
+            float d=fabsf(bf2f(hOut[i])-ref[i]);
+            if(d>max_err) max_err=d;
+        }
+        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
+    }
+    printf("All tests passed!\n");
+    return 0;
+}

scripts/tools/benchmark.py

@@ -1,12 +1,13 @@
 """Benchmark AutoRegressiveLM with KVCache"""
 
+import argparse
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
 
 from astrai.config import AutoRegressiveLMConfig
-from astrai.inference import KVCache
+from astrai.inference import ContiguousCache, PageCache
 from astrai.model.transformer import AutoRegressiveLM
 
 
@@ -24,41 +25,14 @@ class GenerationBenchmark:
         config: AutoRegressiveLMConfig,
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
-        page_size: int = 128,
+        cache_type: str = "contiguous",
     ):
         self.config = config
         self.device = device
         self.dtype = dtype
+        self.cache_type = cache_type
         self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype)
         self.model.eval()
-        head_dim = config.dim // config.n_heads
-        n_pages = (config.max_len * 4 + page_size - 1) // page_size
-        self._page_cache = KVCache(
-            config.n_layers,
-            n_pages,
-            page_size,
-            config.n_kv_heads,
-            head_dim,
-            device,
-            dtype,
-        )
-
-    def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
-        prompt_ids = torch.randint(
-            low=0,
-            high=self.config.vocab_size,
-            size=(batch_size, prompt_length),
-            device=self.device,
-            dtype=torch.long,
-        )
-        gen_ids = torch.randint(
-            low=0,
-            high=self.config.vocab_size,
-            size=(batch_size, total_length - prompt_length),
-            device=self.device,
-            dtype=torch.long,
-        )
-        return prompt_ids, gen_ids
 
     @torch.inference_mode()
     def run_prefill_benchmark(
@@ -68,8 +42,12 @@ class GenerationBenchmark:
         num_trials: int = 10,
     ) -> BenchmarkResult:
         for _ in range(3):
-            prompt_ids, _ = self._prepare_inputs(
+            prompt_ids = torch.randint(
-                batch_size, prompt_length, prompt_length
+                0,
+                self.config.vocab_size,
+                (batch_size, prompt_length),
+                device=self.device,
+                dtype=torch.long,
             )
             _ = self.model(prompt_ids)
         torch.cuda.synchronize()
@@ -78,12 +56,15 @@ class GenerationBenchmark:
         total_tokens = batch_size * prompt_length * num_trials
 
         for trial in range(num_trials):
-            prompt_ids, _ = self._prepare_inputs(
+            prompt_ids = torch.randint(
-                batch_size, prompt_length, prompt_length
+                0,
+                self.config.vocab_size,
+                (batch_size, prompt_length),
+                device=self.device,
+                dtype=torch.long,
             )
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
-
             start.record()
             _ = self.model(prompt_ids)
             end.record()
@@ -107,6 +88,7 @@ class GenerationBenchmark:
                 "prompt_length": prompt_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
+                "cache": "none",
             },
         )
 
@@ -120,29 +102,56 @@ class GenerationBenchmark:
     ) -> BenchmarkResult:
         total_time = 0.0
         total_tokens = batch_size * gen_length * num_trials
-        page_size = self._page_cache.page_size
 
         for trial in range(num_trials):
-            prompt_ids, gen_ids = self._prepare_inputs(
+            prompt_ids = torch.randint(
-                batch_size,
+                0,
-                prompt_length,
+                self.config.vocab_size,
-                prompt_length + gen_length,
+                (batch_size, prompt_length),
-            )
-
-            n_pages = (prompt_length + gen_length + page_size - 1) // page_size
-            total = n_pages * batch_size
-            pages = []
-            for _ in range(total):
-                p = self._page_cache._pool.alloc()
-                assert p >= 0, "OOM"
-                pages.append(p)
-            page_table = torch.tensor(
-                [pages[i * n_pages : (i + 1) * n_pages] for i in range(batch_size)],
-                dtype=torch.long,
                 device=self.device,
+                dtype=torch.long,
+            )
+            gen_ids = torch.randint(
+                0,
+                self.config.vocab_size,
+                (batch_size, gen_length),
+                device=self.device,
+                dtype=torch.long,
             )
 
-            cv = self._page_cache.bind(page_table, total_len=prompt_length)
+            head_dim = self.config.dim // self.config.n_heads
+            max_seq = prompt_length + gen_length
+
+            if self.cache_type == "contiguous":
+                cache = ContiguousCache(
+                    self.config.n_layers,
+                    batch_size,
+                    max_seq,
+                    self.config.n_kv_heads,
+                    head_dim,
+                    self.device,
+                    self.dtype,
+                )
+            else:
+                page_size = 128
+                n_pages = (max_seq + page_size - 1) // page_size * batch_size
+                cache = PageCache(
+                    self.config.n_layers,
+                    n_pages,
+                    page_size,
+                    self.config.n_kv_heads,
+                    head_dim,
+                    self.device,
+                    self.dtype,
+                )
+
+            task_ids = [f"b{i}" for i in range(batch_size)]
+            for tid in task_ids:
+                cache.task_alloc(tid, [0] * max_seq)
+                for p in range(max_seq):
+                    cache.task_extend(tid, p)
+
+            cv = cache.bind_tasks(task_ids, prompt_length, self.device)
             _ = self.model(
                 prompt_ids,
                 paged_cache=cv,
@@ -152,37 +161,35 @@ class GenerationBenchmark:
                 .unsqueeze(0)
                 .expand(batch_size, -1),
             )
-
             torch.cuda.synchronize()
 
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
-
             start.record()
-            current_pos = prompt_length
+
             for i in range(gen_length):
-                input_token = gen_ids[:, i : i + 1]
+                pos = prompt_length + i
-                cv = self._page_cache.bind(page_table, total_len=current_pos + 1)
+                cv = cache.bind_tasks(task_ids, pos + 1, self.device)
                 _ = self.model(
-                    input_token,
+                    gen_ids[:, i : i + 1],
                     paged_cache=cv,
                     position_ids=torch.full(
                         (batch_size, 1),
-                        current_pos,
+                        pos,
                         dtype=torch.long,
                         device=self.device,
                     ),
                 )
-                current_pos += 1
+
             end.record()
             torch.cuda.synchronize()
 
+            for tid in task_ids:
+                cache.task_free(tid)
+
             trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
-            for idx in pages:
-                self._page_cache._pool.free(idx)
-
             print(
                 f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
                 f"({gen_length / trial_time:.1f} tok/s)"
@@ -199,6 +206,7 @@ class GenerationBenchmark:
                 "gen_length": gen_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
+                "cache": self.cache_type,
             },
         )
 
@@ -216,6 +224,42 @@ def print_benchmark_result(result: BenchmarkResult):
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="AutoRegressiveLM benchmark")
+    parser.add_argument(
+        "--device", type=str, default="cuda", help="Device (default: cuda)"
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["bfloat16", "float16", "float32"],
+        help="Dtype",
+    )
+    parser.add_argument(
+        "--cache",
+        type=str,
+        default="contiguous",
+        choices=["contiguous", "paged"],
+        help="KV cache type",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
+    parser.add_argument("--prompt_length", type=int, default=512, help="Prompt length")
+    parser.add_argument("--gen_length", type=int, default=128, help="Generation length")
+    parser.add_argument("--num_trials", type=int, default=5, help="Number of trials")
+    parser.add_argument(
+        "--prefill_only", action="store_true", help="Run prefill benchmark only"
+    )
+    parser.add_argument(
+        "--decode_only", action="store_true", help="Run decoding benchmark only"
+    )
+    args = parser.parse_args()
+
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+
     config = AutoRegressiveLMConfig(
         vocab_size=10000,
         dim=1536,
@@ -227,23 +271,29 @@ if __name__ == "__main__":
         norm_eps=1e-5,
     )
 
-    benchmark = GenerationBenchmark(config)
+    benchmark = GenerationBenchmark(
+        config, device=args.device, dtype=dtype_map[args.dtype], cache_type=args.cache
+    )
 
     print("=" * 80)
-    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
+    print(
+        f"Running AutoRegressiveLM Benchmark (device={args.device}, dtype={args.dtype})"
+    )
     print("=" * 80)
 
-    prefill_result = benchmark.run_prefill_benchmark(
+    if not args.decode_only:
-        batch_size=4,
+        prefill_result = benchmark.run_prefill_benchmark(
-        prompt_length=512,
+            batch_size=args.batch_size,
-        num_trials=5,
+            prompt_length=args.prompt_length,
-    )
+            num_trials=args.num_trials,
-    print_benchmark_result(prefill_result)
+        )
+        print_benchmark_result(prefill_result)
 
-    gen_result = benchmark.run_decoding_benchmark(
+    if not args.prefill_only:
-        batch_size=4,
+        gen_result = benchmark.run_decoding_benchmark(
-        prompt_length=512,
+            batch_size=args.batch_size,
-        gen_length=128,
+            prompt_length=args.prompt_length,
-        num_trials=5,
+            gen_length=args.gen_length,
-    )
+            num_trials=args.num_trials,
-    print_benchmark_result(gen_result)
+        )
+        print_benchmark_result(gen_result)

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)