fix: correct gqa_decode_attn reduction + add gqa_prefill_attn

- gqa_decode_attn: rewrite to per-KV-head, K in smem
- gqa_prefill_attn: new kernel for Q_len > 1 with GQA

astrai/extension/__init__.py

@@ -5,7 +5,7 @@ logger = logging.getLogger(__name__)
 
 available: dict[str, bool] = {}
 
-for _name in ["gqa_decode_attn"]:
+for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
     try:
         importlib.import_module(f".{_name}", package=__package__)
         available[_name] = True

csrc/build.py

@@ -27,3 +27,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
 
 
 register("gqa_decode_attn")
+register("gqa_prefill_attn")

csrc/kernels/gqa_decode_attn.cu

@@ -1,3 +1,4 @@
+// per-KV-head block, K shared in smem, each thread handles hd/32 elements
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
 #include <cmath>
@@ -6,6 +7,8 @@
 
 using bf16 = __nv_bfloat16;
 
+constexpr int CHUNK = 64;
+
 __inline__ __device__ float warp_reduce_sum(float val) {
     for (int offset = 16; offset > 0; offset >>= 1)
         val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
@@ -13,47 +16,69 @@ __inline__ __device__ float warp_reduce_sum(float val) {
 }
 
 __global__ void gqa_decode_attn_kernel(
-    const bf16* q_ptr, const bf16* k_ptr, const bf16* v_ptr,
+    const bf16* __restrict__ q_ptr,
-    const bool*  mask_ptr, bf16* out_ptr,
+    const bf16* __restrict__ k_ptr,
+    const bf16* __restrict__ v_ptr,
+    const bool*  __restrict__ mask_ptr,
+    bf16* __restrict__ out_ptr,
     int B, int n_heads, int n_kv_heads, int seq_len, int hd
 ) {
-    int batch = blockIdx.x / n_heads;
+    int batch = blockIdx.x / n_kv_heads;
-    int q_head = blockIdx.x % n_heads;
+    int kv_head = blockIdx.x % n_kv_heads;
-    int kv_head = q_head / (n_heads / n_kv_heads);
+    int group_size = blockDim.y;
-    int tid = threadIdx.x;
+    int q_head = kv_head * group_size + threadIdx.y;
+    int lane = threadIdx.x;
+    int hd_per_thread = hd / 32;
+
+    float q_reg[8];
+    int q_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        q_reg[i] = __bfloat162float(q_ptr[q_off + i]);
 
-    float q_val = __bfloat162float(
-        q_ptr[((batch * n_heads + q_head) * 1) * hd + tid]);
     int kv_base = ((batch * n_kv_heads + kv_head) * seq_len) * hd;
     int mask_base = batch * seq_len;
 
-    float m = -FLT_MAX, d = 0.0f, acc = 0.0f;
+    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
-    __shared__ float smem[2];
+    float scale = rsqrtf((float)hd);
-    float scale = 1.0f / sqrtf((float)hd);
 
-    for (int s = 0; s < seq_len; s++) {
+    extern __shared__ __align__(16) bf16 k_smem[];
-        int off = kv_base + s * hd + tid;
-        float partial = q_val * __bfloat162float(k_ptr[off]);
-        partial = warp_reduce_sum(partial) * scale;
 
-        if (tid % 32 == 0) smem[tid / 32] = partial;
+    for (int chunk_start = 0; chunk_start < seq_len; chunk_start += CHUNK) {
-        __syncthreads();
+        int this_chunk = min(CHUNK, seq_len - chunk_start);
-        if (tid == 0) smem[0] = smem[0] + smem[1];
+
+        int total = this_chunk * hd;
+        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
+            k_smem[i] = k_ptr[kv_base + chunk_start * hd + i];
         __syncthreads();
 
-        float score = smem[0];
+        for (int s = 0; s < this_chunk; s++) {
-        if (!mask_ptr[mask_base + s]) score = -FLT_MAX;
+            float partial = 0.0f;
+            #pragma unroll
+            for (int i = 0; i < hd_per_thread; i++)
+                partial += q_reg[i] * __bfloat162float(k_smem[s * hd + lane * hd_per_thread + i]);
+            partial = warp_reduce_sum(partial) * scale;
 
-        float new_m = fmaxf(m, score);
+            if (!mask_ptr[mask_base + chunk_start + s]) partial = -FLT_MAX;
-        float alpha = expf(m - new_m);
+
-        float beta  = expf(score - new_m);
+            float new_m = fmaxf(m, partial);
-        d = d * alpha + beta;
+            float alpha = expf(m - new_m);
-        acc = acc * alpha + __bfloat162float(v_ptr[off]) * beta;
+            float beta  = expf(partial - new_m);
-        m = new_m;
+            d = d * alpha + beta;
+
+            int v_off = kv_base + (chunk_start + s) * hd + lane * hd_per_thread;
+            #pragma unroll
+            for (int i = 0; i < hd_per_thread; i++)
+                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta;
+            m = new_m;
+        }
+        __syncthreads();
     }
 
-    int out_off = ((batch * n_heads + q_head) * 1) * hd + tid;
+    int out_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
-    out_ptr[out_off] = __float2bfloat16(acc / d);
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d);
 }
 
 torch::Tensor gqa_decode_attn(
@@ -68,9 +93,15 @@ torch::Tensor gqa_decode_attn(
 
     int B = q.size(0), n_heads = q.size(1), n_kv = k.size(1);
     int seq_len = k.size(2), hd = q.size(3);
+    TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32");
+    int group_size = n_heads / n_kv;
     auto out = torch::empty_like(q);
 
-    gqa_decode_attn_kernel<<<dim3(B * n_heads), dim3(hd)>>>(
+    size_t smem = CHUNK * hd * sizeof(bf16); // K chunk
+    dim3 block(32, group_size);
+    dim3 grid(B * n_kv);
+
+    gqa_decode_attn_kernel<<<grid, block, smem>>>(
         reinterpret_cast<const bf16*>(q.data_ptr()),
         reinterpret_cast<const bf16*>(k.data_ptr()),
         reinterpret_cast<const bf16*>(v.data_ptr()),
@@ -82,5 +113,5 @@ torch::Tensor gqa_decode_attn(
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode attention (fused)");
+    m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode v2 (per-KV-head, shared K)");
 }

csrc/kernels/gqa_prefill_attn.cu

@@ -0,0 +1,118 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cmath>
+#include <cfloat>
+#include <torch/extension.h>
+
+using bf16 = __nv_bfloat16;
+
+__inline__ __device__ float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+__global__ void gqa_prefill_attn_kernel(
+    const bf16* __restrict__ q_ptr,
+    const bf16* __restrict__ k_ptr,
+    const bf16* __restrict__ v_ptr,
+    const bool*  __restrict__ mask_ptr,
+    bf16* __restrict__ out_ptr,
+    int B, int n_heads, int n_kv_heads, int q_len, int kv_len, int hd,
+    int use_mask, int is_causal, int causal_offset
+) {
+    int flat_id = blockIdx.x;
+    int pos = flat_id % q_len;
+    flat_id /= q_len;
+    int q_head = flat_id % n_heads;
+    int batch = flat_id / n_heads;
+    int kv_head = q_head / (n_heads / n_kv_heads);
+    int lane = threadIdx.x;
+    int hd_per_thread = hd / 32;
+
+    // each thread handles hd/32 elements of Q
+    float q_reg[8];
+    int q_off = ((batch * n_heads + q_head) * q_len + pos) * hd + lane * hd_per_thread;
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        q_reg[i] = __bfloat162float(q_ptr[q_off + i]);
+
+    int kv_base = ((batch * n_kv_heads + kv_head) * kv_len) * hd;
+    int limit = is_causal ? min(pos + causal_offset + 1, kv_len) : kv_len;
+
+    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
+    float scale = rsqrtf((float)hd);
+
+    int mask_stride = q_len * kv_len;
+    int mask_off = batch * mask_stride + pos * kv_len;
+
+    for (int s = 0; s < limit; s++) {
+        float partial = 0.0f;
+        int k_off = kv_base + s * hd + lane * hd_per_thread;
+        #pragma unroll
+        for (int i = 0; i < hd_per_thread; i++)
+            partial += q_reg[i] * __bfloat162float(k_ptr[k_off + i]);
+        partial = warp_reduce_sum(partial) * scale;
+
+        if (use_mask && !mask_ptr[mask_off + s]) 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 + s * hd + lane * hd_per_thread;
+        #pragma unroll
+        for (int i = 0; i < hd_per_thread; i++)
+            acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta;
+        m = new_m;
+    }
+
+    int out_off = ((batch * n_heads + q_head) * q_len + pos) * hd + lane * hd_per_thread;
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d);
+}
+
+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);
+
+    int B = q.size(0), n_heads = q.size(1), q_len = q.size(2), hd = q.size(3);
+    int n_kv = k.size(1), kv_len = k.size(2);
+    TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32");
+
+    bool use_mask = mask.has_value();
+    const bool* mask_ptr = nullptr;
+    if (use_mask) {
+        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        mask_ptr = mask.value().data_ptr<bool>();
+    }
+
+    auto out = torch::empty_like(q);
+
+    dim3 block(32);
+    dim3 grid(B * n_heads * q_len);
+
+    gqa_prefill_attn_kernel<<<grid, block>>>(
+        reinterpret_cast<const bf16*>(q.data_ptr()),
+        reinterpret_cast<const bf16*>(k.data_ptr()),
+        reinterpret_cast<const bf16*>(v.data_ptr()),
+        mask_ptr,
+        reinterpret_cast<bf16*>(out.data_ptr()),
+        B, n_heads, n_kv, q_len, kv_len, hd,
+        (int)use_mask, (int)is_causal, (int)causal_offset
+    );
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill attention (naive)");
+}