.gitignore

@@ -7,13 +7,8 @@
 # Allow specific file types and root files
 !astrai/**/*.py
 !scripts/**/*.py
-!tests/**/*.py
-!csrc/**/*.py
-
-!csrc/**/*.cu
-!csrc/**/*.cuh
-
 !scripts/**/*.sh
+!tests/**/*.py
 
 # Allow GitHub files
 !/.github/**
@@ -28,8 +23,3 @@
 !/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/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/release/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,9 +59,8 @@ End-to-end walkthrough in 5 steps:
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .                                          # pure PyTorch (no CUDA kernels)
+pip install -e .
-# CSRC_KERNELS=true pip install -e . --no-build-isolation  # optional: fused CUDA kernels
+# pip install -e ".[dev]"    # optional: dev dependencies (pytest, ruff)
-# 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/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/release/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,7 @@
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .                                          # 纯 PyTorch（不含 CUDA 内核）
+pip install -e .
-# CSRC_KERNELS=true pip install -e . --no-build-isolation  # 可选：融合 CUDA 内核加速
 # pip install -e ".[dev]"    # 可选：开发依赖（pytest, ruff）
 ```
 

astrai/extension/__init__.py

@@ -1,91 +0,0 @@
-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

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

csrc/build.py

@@ -1,36 +0,0 @@
-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

@@ -1,35 +0,0 @@
-#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

@@ -1,66 +0,0 @@
-#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

@@ -1,59 +0,0 @@
-#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

@@ -1,96 +0,0 @@
-#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

@@ -1,137 +0,0 @@
-#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

@@ -1,244 +0,0 @@
-#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

@@ -1,124 +0,0 @@
-// 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

@@ -1,127 +0,0 @@
-// 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,13 +1,12 @@
 """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 ContiguousCache, PageCache
+from astrai.inference import KVCache
 from astrai.model.transformer import AutoRegressiveLM
 
 
@@ -25,14 +24,41 @@ class GenerationBenchmark:
         config: AutoRegressiveLMConfig,
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
-        cache_type: str = "contiguous",
+        page_size: int = 128,
     ):
         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(
@@ -42,12 +68,8 @@ class GenerationBenchmark:
         num_trials: int = 10,
     ) -> BenchmarkResult:
         for _ in range(3):
-            prompt_ids = torch.randint(
+            prompt_ids, _ = self._prepare_inputs(
-                0,
+                batch_size, prompt_length, prompt_length
-                self.config.vocab_size,
-                (batch_size, prompt_length),
-                device=self.device,
-                dtype=torch.long,
             )
             _ = self.model(prompt_ids)
         torch.cuda.synchronize()
@@ -56,15 +78,12 @@ class GenerationBenchmark:
         total_tokens = batch_size * prompt_length * num_trials
 
         for trial in range(num_trials):
-            prompt_ids = torch.randint(
+            prompt_ids, _ = self._prepare_inputs(
-                0,
+                batch_size, prompt_length, prompt_length
-                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()
@@ -88,7 +107,6 @@ class GenerationBenchmark:
                 "prompt_length": prompt_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
-                "cache": "none",
             },
         )
 
@@ -102,56 +120,29 @@ 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 = torch.randint(
+            prompt_ids, gen_ids = self._prepare_inputs(
-                0,
-                self.config.vocab_size,
-                (batch_size, prompt_length),
-                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,
-            )
-
-            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,
+                prompt_length,
-                    self.config.n_kv_heads,
+                prompt_length + gen_length,
-                    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)]
+            n_pages = (prompt_length + gen_length + page_size - 1) // page_size
-            for tid in task_ids:
+            total = n_pages * batch_size
-                cache.task_alloc(tid, [0] * max_seq)
+            pages = []
-                for p in range(max_seq):
+            for _ in range(total):
-                    cache.task_extend(tid, p)
+                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,
+            )
 
-            cv = cache.bind_tasks(task_ids, prompt_length, self.device)
+            cv = self._page_cache.bind(page_table, total_len=prompt_length)
             _ = self.model(
                 prompt_ids,
                 paged_cache=cv,
@@ -161,35 +152,37 @@ 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()
 
+            start.record()
+            current_pos = prompt_length
             for i in range(gen_length):
-                pos = prompt_length + i
+                input_token = gen_ids[:, i : i + 1]
-                cv = cache.bind_tasks(task_ids, pos + 1, self.device)
+                cv = self._page_cache.bind(page_table, total_len=current_pos + 1)
                 _ = self.model(
-                    gen_ids[:, i : i + 1],
+                    input_token,
                     paged_cache=cv,
                     position_ids=torch.full(
                         (batch_size, 1),
-                        pos,
+                        current_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)"
@@ -206,7 +199,6 @@ class GenerationBenchmark:
                 "gen_length": gen_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
-                "cache": self.cache_type,
             },
         )
 
@@ -224,42 +216,6 @@ 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,
@@ -271,29 +227,23 @@ if __name__ == "__main__":
         norm_eps=1e-5,
     )
 
-    benchmark = GenerationBenchmark(
+    benchmark = GenerationBenchmark(config)
-        config, device=args.device, dtype=dtype_map[args.dtype], cache_type=args.cache
-    )
 
     print("=" * 80)
-    print(
+    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
-        f"Running AutoRegressiveLM Benchmark (device={args.device}, dtype={args.dtype})"
-    )
     print("=" * 80)
 
-    if not args.decode_only:
     prefill_result = benchmark.run_prefill_benchmark(
-            batch_size=args.batch_size,
+        batch_size=4,
-            prompt_length=args.prompt_length,
+        prompt_length=512,
-            num_trials=args.num_trials,
+        num_trials=5,
     )
     print_benchmark_result(prefill_result)
 
-    if not args.prefill_only:
     gen_result = benchmark.run_decoding_benchmark(
-            batch_size=args.batch_size,
+        batch_size=4,
-            prompt_length=args.prompt_length,
+        prompt_length=512,
-            gen_length=args.gen_length,
+        gen_length=128,
-            num_trials=args.num_trials,
+        num_trials=5,
     )
     print_benchmark_result(gen_result)

setup.py

@@ -1,58 +0,0 @@
-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)