refactor: rename gqa_* to attn_*, split-KV for all decode paths

- Rename all csrc/kernels/gqa_*.cuh/cu to attn_*, with _split_q / _split_kv
  strategy suffix and optional _mma compute suffix
- Remove non-split MMA decode kernel, keep only split-KV path
- Convert scalar decode fallback to split-KV (o_part/ml_part + combine)
- Move combine kernel to attn_decode_split_kv.cuh (shared by both paths)
- Rename GQAParams to AttentionParams
- Update all C++ #include, PYBIND11, and Python extension references

astrai/extension/__init__.py

@@ -1,19 +1,19 @@
 """CUDA attention kernel wrappers with torch fallback.
 
 Public API:
-    - ``gqa_decode_attn`` — single-query decode attention
+    - ``attn_decode`` — single-query decode attention
-    - ``gqa_prefill_attn`` — multi-query prefill attention
+    - ``attn_prefill`` — multi-query prefill attention
 
-Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``)
+Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
 when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
 """
 
 from astrai.extension.loader import KERNEL_NAMES, is_available
-from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
+from astrai.extension.ops import attn_decode, attn_prefill
 
 __all__ = [
-    "gqa_decode_attn",
+    "attn_decode",
-    "gqa_prefill_attn",
+    "attn_prefill",
     "is_available",
     "KERNEL_NAMES",
 ]

astrai/extension/loader.py

@@ -11,7 +11,7 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-KERNEL_NAMES = ["gqa_decode_attn", "gqa_prefill_attn"]
+KERNEL_NAMES = ["attn_decode", "attn_prefill"]
 
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}

astrai/extension/ops.py

@@ -42,7 +42,7 @@ def _torch_fallback(
     )
 
 
-def gqa_decode_attn(
+def attn_decode(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
@@ -51,8 +51,8 @@ def gqa_decode_attn(
     causal_offset: int = 0,
     scale: float | None = None,
 ) -> torch.Tensor:
-    if _available["gqa_decode_attn"]:
+    if _available["attn_decode"]:
-        return _modules["gqa_decode_attn"].gqa_decode_attn(
+        return _modules["attn_decode"].attn_decode(
             q,
             k,
             v,
@@ -64,7 +64,7 @@ def gqa_decode_attn(
     return _torch_fallback(q, k, v, mask, is_causal, scale)
 
 
-def gqa_prefill_attn(
+def attn_prefill(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
@@ -73,8 +73,8 @@ def gqa_prefill_attn(
     causal_offset: int = 0,
     scale: float | None = None,
 ) -> torch.Tensor:
-    if _available["gqa_prefill_attn"]:
+    if _available["attn_prefill"]:
-        return _modules["gqa_prefill_attn"].gqa_prefill_attn(
+        return _modules["attn_prefill"].attn_prefill(
             q,
             k,
             v,

csrc/build.py

@@ -42,5 +42,5 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
     }
 
 
-register("gqa_decode_attn")
+register("attn_decode")
-register("gqa_prefill_attn")
+register("attn_prefill")

csrc/kernels/attn_common.cuh

@@ -7,7 +7,7 @@
 using bf16 = __nv_bfloat16;
 using std::min;
 
-struct GQAParams {
+struct AttentionParams {
     int batch;
     int q_head;
     int kv_head;
@@ -22,5 +22,9 @@ struct GQAParams {
     const bf16* __restrict__ k;
     const bf16* __restrict__ v;
     const bool* __restrict__ mask;
+    
     bf16* __restrict__ o;
+    float* __restrict__ o_part;
+    float* __restrict__ ml_part;
+    int num_splits;
 };

csrc/kernels/attn_decode.cu

@@ -1,70 +1,67 @@
-#include "gqa_decode_attn.cuh"
+#include "attn_decode_split_kv.cuh"
-#include "gqa_entry_utils.cuh"
+#include "attn_entry_utils.cuh"
 
 #ifndef ASTRAI_NO_MMA
-#include "gqa_decode_attn_mma.cuh"
+#include "attn_decode_split_kv_mma.cuh"
 #endif
 
-// Scalar fallback: one warp per query head, per (batch, kv_head) block.
-static void launch_scalar_decode(const GQAParams& p) {
-    int group_size = p.q_head / p.kv_head;
-    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-    gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, group_size), smem>>>(p);
-}
-
-#ifndef ASTRAI_NO_MMA
-// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask.
-static bool decode_use_mma(const GQAParams& p) {
-    int G = p.q_head / p.kv_head;
-    return !p.use_mask && G > 1 && G <= 16;
-}
-
 // Decode has only batch*kv_head independent tasks; without split-K the grid is
 // tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count
-// that fills the device (~2 blocks/SM), capped by the tile count and 32.
+// that fills the device (~2 blocks/SM), capped by the tile count, min work per
-static int decode_num_splits(const GQAParams& p, int tiles_total) {
+// split (at least 8 tiles), and 32.
+static int decode_num_splits(const AttentionParams& p, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
     int base_blocks = p.kv_head * p.batch;
     int desired = 2 * (sm_count > 0 ? sm_count : 64);
     int n = (desired + base_blocks - 1) / base_blocks;
-    return std::max(1, std::min(n, std::min(tiles_total, 32)));
+    int max_by_work = tiles_total / 8;
+    return std::max(1, std::min({n, tiles_total, 32, max_by_work}));
+}
+
+// Scalar fallback: one warp per query head, split-KV across grid.z.
+static void launch_scalar_decode(AttentionParams& p) {
+    int group_size = p.q_head / p.kv_head;
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    p.num_splits = decode_num_splits(p, chunks_total);
+
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
+
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+
+#ifndef ASTRAI_NO_MMA
+// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask.
+static bool decode_use_mma(const AttentionParams& p) {
+    int G = p.q_head / p.kv_head;
+    return !p.use_mask && G > 1 && G <= 16;
 }
 
 template <int HEAD_DIM, int BC>
-static void launch_mma_decode(GQAParams& p) {
+static void launch_mma_decode(AttentionParams& p) {
-    constexpr int BR = 16, LD = HEAD_DIM;  // XOR swizzle → no padding
-    int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    int num_splits = decode_num_splits(p, tiles_total);
+    p.num_splits = decode_num_splits(p, tiles_total);
 
-    // Enough (batch, kv_head) work to fill the SMs → single pass, direct write.
-    if (num_splits <= 1) {
-        cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
-                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-        gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
-            <<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
-        return;
-    }
-
-    // Split-K (FlashDecoding): partition kv across blocks, then reduce.
     auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, num_splits, p.head_dim}, fopt);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, num_splits, 2}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
 
-    cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
-                         cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
-    gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-        <<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
-            p, o_part.data_ptr<float>(), ml_part.data_ptr<float>(), num_splits);
-    gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
-        o_part.data_ptr<float>(), ml_part.data_ptr<float>(), p.o,
-        num_splits, p.head_dim);
 }
 #endif
 
 template <int HEAD_DIM>
-static void dispatch_decode(GQAParams& p) {
+static void dispatch_decode(AttentionParams& p) {
 #ifndef ASTRAI_NO_MMA
     if (decode_use_mma(p)) {
         launch_mma_decode<HEAD_DIM, 32>(p);
@@ -74,7 +71,7 @@ static void dispatch_decode(GQAParams& p) {
     launch_scalar_decode(p);
 }
 
-torch::Tensor gqa_decode_attn(
+torch::Tensor attn_decode(
     torch::Tensor q,
     torch::Tensor k,
     torch::Tensor v,
@@ -83,8 +80,8 @@ torch::Tensor gqa_decode_attn(
     int64_t causal_offset = 0,
     c10::optional<double> scale = c10::nullopt
 ) {
-    GQAParams p;
+    AttentionParams p;
-    gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
 
@@ -112,7 +109,7 @@ torch::Tensor gqa_decode_attn(
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("gqa_decode_attn", &gqa_decode_attn,
+    m.def("attn_decode", &attn_decode,
         py::arg("q"),
         py::arg("k"),
         py::arg("v"),

csrc/kernels/attn_decode_split_kv.cuh

@@ -1,5 +1,5 @@
 #pragma once
-#include "gqa_common.cuh"
+#include "attn_common.cuh"
 
 constexpr int DC_CHUNK = 64;
 
@@ -9,9 +9,10 @@ __device__ inline float warp_reduce_sum(float val) {
     return val;
 }
 
-__global__ void gqa_decode_attn_kernel(GQAParams p) {
+__global__ void attn_decode_split_kv_kernel(AttentionParams p) {
     int batch = blockIdx.x / p.kv_head;
     int kv_head = blockIdx.x % p.kv_head;
+    int split = blockIdx.z;
     int group_size = blockDim.y;
     int q_head = kv_head * group_size + threadIdx.y;
     int lane = threadIdx.x;
@@ -29,7 +30,14 @@ __global__ void gqa_decode_attn_kernel(GQAParams p) {
 
     extern __shared__ __align__(16) bf16 k_smem[];
 
-    for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) {
+    // Split-KV: each split processes a contiguous subset of chunks
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
+    int ch_begin = split * chunks_per_split;
+    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
+
+    for (int ci = ch_begin; ci < ch_end; ci++) {
+        int chunk_start = ci * DC_CHUNK;
         int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
 
         int total = this_chunk * p.head_dim;
@@ -61,7 +69,47 @@ __global__ void gqa_decode_attn_kernel(GQAParams p) {
         __syncthreads();
     }
 
-    int out_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    // ---- write UN-normalised partials for this split ----
-    for (int i = 0; i < hd_per_thread; i++)
+    size_t bh = (size_t)batch * p.q_head + q_head;
-        p.o[out_off + i] = __float2bfloat16(acc_reg[i] / d);
+    size_t slot = bh * p.num_splits + split;
+    int d0 = lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++) {
+        int dd = d0 + i;
+        p.o_part[slot * p.head_dim + dd] = acc_reg[i];
+    }
+    if (lane == 0) {
+        p.ml_part[slot * 2] = m;
+        p.ml_part[slot * 2 + 1] = d;
+    }
+}
+
+// Reduce split-K partials into the final bf16 output. One block per (batch,
+// q_head); each thread owns one head_dim element and folds across all splits
+// with a numerically-stable online rescale.
+__global__ void attn_decode_combine_kernel(AttentionParams p) {
+    int bh = blockIdx.x;
+    int d = threadIdx.x;
+    if (d >= p.head_dim) return;
+
+    size_t split_base = (size_t)bh * p.num_splits;
+
+    const float* mlp = p.ml_part + split_base * 2;
+    float mstar = -FLT_MAX;
+    for (int s = 0; s < p.num_splits; s++)
+        mstar = fmaxf(mstar, mlp[s * 2]);
+
+    float lstar = 0.0f;
+    for (int s = 0; s < p.num_splits; s++) {
+        float mi = mlp[s * 2];
+        if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar);
+    }
+
+    const float* op = p.o_part + split_base * p.head_dim;
+    float acc = 0.0f;
+    for (int s = 0; s < p.num_splits; s++) {
+        float mi = mlp[s * 2];
+        if (mi > -FLT_MAX) acc += op[s * p.head_dim + d] * __expf(mi - mstar);
+    }
+    float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f;
+    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -0,0 +1,164 @@
+#pragma once
+#include "attn_common.cuh"
+#include "attn_mma_utils.cuh"
+
+// Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
+//
+// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
+// on its own. But GQA gives us G = q_head / kv_head query heads that all share
+// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
+// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
+// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
+// not the flops). The KV sequence is partitioned across gridDim.z blocks so that
+// a decode with only batch*kv_head independent tasks can fill all SMs. Each
+// (batch, kv_head, split) block computes an UN-normalised partial (Oacc, m, l)
+// over its KV slice; the combine kernel below reduces across splits. Fixes the
+// "grid too small" bottleneck (0.04 waves/SM → many blocks) for long-context,
+// small-batch decode.
+//
+// Partial layout (float, contiguous):
+//   o_part : [batch, q_head, num_splits, HEAD_DIM]
+//   ml_part: [batch, q_head, num_splits, 2]  (m, l)
+//
+// Optimizations:
+//   - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
+//   - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
+//   - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply
+//   - single-buffer: keeps smem small for high occupancy
+template <int HEAD_DIM, int BC>
+__global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) {
+    constexpr int BR = 16;
+    constexpr int KD = HEAD_DIM / 16;
+    constexpr int NC8 = BC / 8;
+    constexpr int KT2 = BC / 16;
+    constexpr int DN8 = HEAD_DIM / 8;
+    constexpr int LD = HEAD_DIM;
+    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
+
+    const int lane = threadIdx.x;
+    const int gid = lane >> 2;
+    const int tid4 = lane & 3;
+
+    const int kv_head = blockIdx.x;
+    const int batch = blockIdx.y;
+    const int split = blockIdx.z;
+    const int G = p.q_head / p.kv_head;
+    const int q_head0 = kv_head * G;
+
+    __shared__ __align__(16) bf16 sK[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sV[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM];
+
+    bf16 scale_bf16 = __float2bfloat16(p.scale);
+    for (int i = lane; i < BR * HEAD_DIM; i += 32) {
+        int r = i / HEAD_DIM, d = i % HEAD_DIM;
+        bf16 val = __float2bfloat16(0.0f);
+        if (r < G) {
+            int qh = q_head0 + r;
+            val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];
+        }
+        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
+    }
+    __syncwarp();
+
+    unsigned Qa[KD][4];
+    int qrow_l = (lane & 7) + (lane & 8);
+    int qcol_l = (lane & 16) ? 8 : 0;
+#pragma unroll
+    for (int kt = 0; kt < KD; kt++)
+        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
+
+    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 mask_base = batch * p.kv_len;
+    const int tiles_total = (p.kv_len + BC - 1) / BC;
+    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
+    const int ti_begin = split * tiles_per_split;
+    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
+    const int has_mask = p.use_mask && p.mask;
+
+    for (int ti = ti_begin; ti < ti_end; ti++) {
+        int kv0 = ti * BC;
+
+        bool full_tile = (kv0 + BC <= p.kv_len);
+        if (full_tile) {
+            constexpr int VEC = 8;
+            int total = BC * HEAD_DIM;
+#pragma unroll
+            for (int i = lane * VEC; i < total; i += 32 * VEC) {
+                int r = i / HEAD_DIM, d = i % HEAD_DIM;
+                int kc = kv0 + r;
+                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
+                            &p.k[kv_base + kc * HEAD_DIM + d]);
+                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
+                            &p.v[kv_base + kc * HEAD_DIM + d]);
+            }
+            cp_async_commit();
+            cp_async_wait_all();
+        } else {
+            for (int i = lane; i < BC * HEAD_DIM; i += 32) {
+                int r = i / HEAD_DIM, d = i % HEAD_DIM;
+                int kc = kv0 + r;
+                bf16 z = __float2bfloat16(0.0f);
+                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
+                    (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
+                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
+                    (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
+            }
+        }
+        __syncwarp();
+
+        float Sacc[NC8][4];
+        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
+
+        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
+                                    mask_base, p.mask, has_mask,
+                                    Sacc, Oacc, m0, m1, l0, l1, lane);
+
+        mma_pv_accumulate<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
+        __syncwarp();
+    }
+
+    // ---- write UN-normalised partials for this split ----
+    auto split_slot = [&](int h) -> size_t {
+        size_t bh = (size_t)batch * p.q_head + h;
+        return bh * p.num_splits + split;
+    };
+#pragma unroll
+    for (int dn8 = 0; dn8 < DN8; dn8++) {
+        int d = dn8 * 8 + 2 * tid4;
+        int r0 = gid, r1 = gid + 8;
+        if (r0 < G) {
+            int h = q_head0 + r0;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][0];
+            op[d + 1] = Oacc[dn8][1];
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][2];
+            op[d + 1] = Oacc[dn8][3];
+        }
+    }
+    if (tid4 == 0) {
+        int r0 = gid, r1 = gid + 8;
+        if (r0 < G) {
+            int h = q_head0 + r0;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m0; mp[1] = l0;
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m1; mp[1] = l1;
+        }
+    }
+}
+

csrc/kernels/attn_entry_utils.cuh

@@ -1,8 +1,8 @@
 #pragma once
 #include <torch/extension.h>
-#include "gqa_common.cuh"
+#include "attn_common.cuh"
 
-inline void gqa_pack_params(
+inline void attn_pack_params(
     torch::Tensor q,
     torch::Tensor k,
     torch::Tensor v,
@@ -10,7 +10,7 @@ inline void gqa_pack_params(
     bool is_causal,
     int64_t causal_offset,
     c10::optional<double> scale,
-    GQAParams& p
+    AttentionParams& p
 ) {
     TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
     TORCH_CHECK(q.dtype() == torch::kBFloat16);

csrc/kernels/attn_prefill.cu

@@ -1,12 +1,12 @@
-#include "gqa_prefill_attn.cuh"
+#include "attn_prefill_split_q.cuh"
-#include "gqa_entry_utils.cuh"
+#include "attn_entry_utils.cuh"
 
 #ifndef ASTRAI_NO_MMA
-#include "gqa_prefill_attn_mma.cuh"
+#include "attn_prefill_split_q_mma.cuh"
 #endif
 
 template <int HEAD_DIM>
-static void dispatch_prefill(GQAParams& p) {
+static void dispatch_prefill(AttentionParams& p) {
 #ifndef ASTRAI_NO_MMA
     constexpr int WARPS = 4, BR = 16;
     // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
@@ -23,16 +23,16 @@ static void dispatch_prefill(GQAParams& p) {
     dim3 block(WARPS * 32, 1, 1);
     // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed.
     // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params.
-    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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);
-    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
 #endif
 }
 
-torch::Tensor gqa_prefill_attn(
+torch::Tensor attn_prefill(
     torch::Tensor q,
     torch::Tensor k,
     torch::Tensor v,
@@ -41,8 +41,8 @@ torch::Tensor gqa_prefill_attn(
     int64_t causal_offset = 0,
     c10::optional<double> scale = c10::nullopt
 ) {
-    GQAParams p;
+    AttentionParams p;
-    gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
 
     auto O = torch::empty_like(q);
@@ -69,7 +69,7 @@ torch::Tensor gqa_prefill_attn(
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("gqa_prefill_attn", &gqa_prefill_attn,
+    m.def("attn_prefill", &attn_prefill,
         py::arg("q"),
         py::arg("k"),
         py::arg("v"),

csrc/kernels/attn_prefill_split_q.cuh

@@ -1,5 +1,5 @@
 #pragma once
-#include "gqa_common.cuh"
+#include "attn_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
@@ -31,7 +31,7 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) {
 }
 
 template <int HEAD_DIM, int G, int ROWS, int P_BC>
-__global__ void gqa_prefill_attn_kernel_t(GQAParams p) {
+__global__ void attn_prefill_split_q_kernel_t(AttentionParams p) {
     constexpr int DPT = HEAD_DIM / G;
 
     int q_tile = blockIdx.x;

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -1,6 +1,6 @@
 #pragma once
-#include "gqa_common.cuh"
+#include "attn_common.cuh"
-#include "gqa_mma_utils.cuh"
+#include "attn_mma_utils.cuh"
 
 // Tensor-core prefill 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
@@ -38,7 +38,7 @@
 
 template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
 __global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
-void gqa_prefill_attn_mma_kernel(GQAParams p) {
+void attn_prefill_split_q_mma_kernel(AttentionParams 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)

csrc/kernels/gqa_decode_attn_mma.cuh

@@ -1,321 +0,0 @@
-#pragma once
-#include "gqa_common.cuh"
-#include "gqa_mma_utils.cuh"
-
-// Tensor-core decode via GQA head-packing with cp.async loads.
-//
-// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
-// on its own. But GQA gives us G = q_head / kv_head query heads that all share
-// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
-// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
-// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
-// not the flops). Fragment layout is identical to the prefill mma kernel; the
-// only differences are (1) the M rows come from different heads at position 0
-// instead of different sequence positions of one head, and (2) causal masking is
-// a single scalar bound shared by every row. One warp owns one (batch, kv_head);
-// requires G <= 16.
-//
-// Optimizations:
-//   - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
-//   - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
-//   - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply
-//   - single-buffer: keeps smem small for high occupancy
-
-template <int HEAD_DIM, int BC>
-__global__ void gqa_decode_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;       // XOR swizzle handles bank conflicts, zero waste
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
-
-    const int lane = threadIdx.x;  // single warp
-    const int gid = lane >> 2;     // 0..7 → rows gid, gid+8
-    const int tid4 = lane & 3;
-
-    const int kv_head = blockIdx.x;
-    const int batch = blockIdx.y;
-    const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head * G;
-
-    extern __shared__ __align__(16) bf16 smem[];
-    bf16* sK = smem;                  // [BC][LD]
-    bf16* sV = sK + BC * LD;          // [BC][LD]
-    bf16* sQ = sV + BC * LD;          // [BR][LD]
-
-    // ---- stage Q into shared (pre-scaled, swizzled) ----
-    bf16 scale_bf16 = __float2bfloat16(p.scale);
-    for (int i = lane; i < BR * HEAD_DIM; i += 32) {
-        int r = i / HEAD_DIM, d = i % HEAD_DIM;
-        bf16 val = __float2bfloat16(0.0f);
-        if (r < G) {
-            int qh = q_head0 + r;
-            val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];  // q_len == 1
-        }
-        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
-    }
-    __syncwarp();
-
-    // Q resident A-fragments
-    unsigned Qa[KD][4];
-    int qrow_l = (lane & 7) + (lane & 8);
-    int qcol_l = (lane & 16) ? 8 : 0;
-#pragma unroll
-    for (int kt = 0; kt < KD; kt++)
-        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
-
-    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 mask_base = batch * p.kv_len;
-    const int tiles = (p.kv_len + BC - 1) / BC;
-    const int has_mask = p.use_mask && p.mask;
-
-    for (int ti = 0; ti < tiles; ti++) {
-        int kv0 = ti * BC;
-
-        // ---- load K/V tile to shared (cp.async on full tiles) ----
-        bool full_tile = (kv0 + BC <= p.kv_len);
-        if (full_tile) {
-            constexpr int VEC = 8;  // 8 bf16 = 16 bytes per cp.async
-            int total = BC * HEAD_DIM;
-#pragma unroll
-            for (int i = lane * VEC; i < total; i += 32 * VEC) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.k[kv_base + kc * HEAD_DIM + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.v[kv_base + kc * HEAD_DIM + d]);
-            }
-            cp_async_commit();
-            cp_async_wait_all();
-        } else {
-            for (int i = lane; i < BC * HEAD_DIM; i += 32) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                bf16 z = __float2bfloat16(0.0f);
-                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
-                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
-            }
-        }
-        __syncwarp();
-
-        // S = Q @ K^T + online softmax + O += P @ V  (shared MMA functions)
-        float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
-
-        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
-        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    mask_base, p.mask, has_mask,
-                                    Sacc, Oacc, m0, m1, l0, l1, lane);
-
-        mma_pv_accumulate<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
-        __syncwarp();  // 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;
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int o_off = (batch * p.q_head + q_head0 + r0) * HEAD_DIM + d;
-            p.o[o_off] = __float2bfloat16(Oacc[dn8][0] * rl0);
-            p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
-        }
-        if (r1 < G) {
-            int o_off = (batch * p.q_head + q_head0 + r1) * HEAD_DIM + d;
-            p.o[o_off] = __float2bfloat16(Oacc[dn8][2] * rl1);
-            p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
-        }
-    }
-}
-
-// ---------------------------------------------------------------------------
-// Split-K (FlashDecoding) decode: identical math to the kernel above, but the
-// KV sequence is partitioned across gridDim.z blocks so that a decode with only
-// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head, split)
-// block computes an UN-normalised partial (Oacc, m, l) over its KV slice; the
-// combine kernel below reduces across splits. Fixes the "grid too small"
-// bottleneck (0.04 waves/SM → many blocks) for long-context, small-batch decode.
-//
-// Partial layout (float, contiguous):
-//   o_part : [batch, q_head, num_splits, HEAD_DIM]
-//   ml_part: [batch, q_head, num_splits, 2]  (m, l)
-template <int HEAD_DIM, int BC>
-__global__ void gqa_decode_attn_mma_splitk_kernel(GQAParams p,
-                                                  float* __restrict__ o_part,
-                                                  float* __restrict__ ml_part,
-                                                  int num_splits) {
-    constexpr int BR = 16;
-    constexpr int KD = HEAD_DIM / 16;
-    constexpr int NC8 = BC / 8;
-    constexpr int KT2 = BC / 16;
-    constexpr int DN8 = HEAD_DIM / 8;
-    constexpr int LD = HEAD_DIM;
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
-
-    const int lane = threadIdx.x;
-    const int gid = lane >> 2;
-    const int tid4 = lane & 3;
-
-    const int kv_head = blockIdx.x;
-    const int batch = blockIdx.y;
-    const int split = blockIdx.z;
-    const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head * G;
-
-    extern __shared__ __align__(16) bf16 smem[];
-    bf16* sK = smem;
-    bf16* sV = sK + BC * LD;
-    bf16* sQ = sV + BC * LD;
-
-    bf16 scale_bf16 = __float2bfloat16(p.scale);
-    for (int i = lane; i < BR * HEAD_DIM; i += 32) {
-        int r = i / HEAD_DIM, d = i % HEAD_DIM;
-        bf16 val = __float2bfloat16(0.0f);
-        if (r < G) {
-            int qh = q_head0 + r;
-            val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];
-        }
-        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
-    }
-    __syncwarp();
-
-    unsigned Qa[KD][4];
-    int qrow_l = (lane & 7) + (lane & 8);
-    int qcol_l = (lane & 16) ? 8 : 0;
-#pragma unroll
-    for (int kt = 0; kt < KD; kt++)
-        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
-
-    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 mask_base = batch * p.kv_len;
-    const int tiles_total = (p.kv_len + BC - 1) / BC;
-    const int tiles_per_split = (tiles_total + num_splits - 1) / num_splits;
-    const int ti_begin = split * tiles_per_split;
-    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
-    const int has_mask = p.use_mask && p.mask;
-
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        int kv0 = ti * BC;
-
-        bool full_tile = (kv0 + BC <= p.kv_len);
-        if (full_tile) {
-            constexpr int VEC = 8;
-            int total = BC * HEAD_DIM;
-#pragma unroll
-            for (int i = lane * VEC; i < total; i += 32 * VEC) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.k[kv_base + kc * HEAD_DIM + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.v[kv_base + kc * HEAD_DIM + d]);
-            }
-            cp_async_commit();
-            cp_async_wait_all();
-        } else {
-            for (int i = lane; i < BC * HEAD_DIM; i += 32) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                bf16 z = __float2bfloat16(0.0f);
-                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
-                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
-            }
-        }
-        __syncwarp();
-
-        float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
-
-        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
-        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
-                                    mask_base, p.mask, has_mask,
-                                    Sacc, Oacc, m0, m1, l0, l1, lane);
-
-        mma_pv_accumulate<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
-        __syncwarp();
-    }
-
-    // ---- write UN-normalised partials for this split ----
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int hh = q_head0 + r0;
-            float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
-            op[d] = Oacc[dn8][0];
-            op[d + 1] = Oacc[dn8][1];
-        }
-        if (r1 < G) {
-            int hh = q_head0 + r1;
-            float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
-            op[d] = Oacc[dn8][2];
-            op[d + 1] = Oacc[dn8][3];
-        }
-    }
-    if (tid4 == 0) {
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r0) * num_splits + split) * 2;
-            mp[0] = m0; mp[1] = l0;
-        }
-        if (r1 < G) {
-            float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r1) * num_splits + split) * 2;
-            mp[0] = m1; mp[1] = l1;
-        }
-    }
-}
-
-// Reduce split-K partials into the final bf16 output. One block per (batch,
-// q_head); each thread owns one head_dim element and folds across all splits
-// with a numerically-stable online rescale.
-__global__ void gqa_decode_combine_kernel(const float* __restrict__ o_part,
-                                          const float* __restrict__ ml_part,
-                                          bf16* __restrict__ out,
-                                          int num_splits, int head_dim) {
-    int bh = blockIdx.x;  // batch * q_head + head
-    int d = threadIdx.x;
-    if (d >= head_dim) return;
-
-    const float* mlp = ml_part + (size_t)bh * num_splits * 2;
-    float mstar = -FLT_MAX;
-    for (int s = 0; s < num_splits; s++)
-        mstar = fmaxf(mstar, mlp[s * 2]);
-
-    float lstar = 0.0f;
-    for (int s = 0; s < num_splits; s++) {
-        float mi = mlp[s * 2];
-        if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar);
-    }
-
-    const float* op = o_part + (size_t)bh * num_splits * head_dim;
-    float acc = 0.0f;
-    for (int s = 0; s < num_splits; s++) {
-        float mi = mlp[s * 2];
-        if (mi > -FLT_MAX) acc += op[s * head_dim + d] * __expf(mi - mstar);
-    }
-    float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f;
-    out[(size_t)bh * head_dim + d] = __float2bfloat16(acc * inv);
-}

csrc/tests/attn_decode_test.cu

@@ -2,16 +2,16 @@
 Pure-C test:
 nvcc -I csrc -arch=sm_89 -O3 \
     --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/gqa_decode_test.cu -o test && ./test
+    csrc/tests/attn_decode_test.cu -o test && ./test
 */
 
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
 #include <sys/time.h>
-#include "../kernels/gqa_decode_attn.cuh"
+#include "../kernels/attn_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
-#include "../kernels/gqa_decode_attn_mma.cuh"
+#include "../kernels/attn_decode_split_kv_mma.cuh"
 #endif
 
 static double now_ms() {
@@ -28,63 +28,58 @@ struct DecodeScratch {
     float* ml_part = nullptr;
 };
 
-// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
+static int decode_num_splits(const AttentionParams& p, int tiles_total) {
-// scalar fallback otherwise), mirroring dispatch_decode() in gqa_decode_attn.cu.
-#ifndef ASTRAI_NO_MMA
-static bool decode_use_mma(const GQAParams& p) {
-    int G = p.q_head / p.kv_head;
-    return !p.use_mask && G > 1 && G <= 16;
-}
-
-static int decode_num_splits(const GQAParams& p, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
     int base_blocks = p.kv_head * p.batch;
     int desired = 2 * (sm_count > 0 ? sm_count : 64);
     int n = (desired + base_blocks - 1) / base_blocks;
-    return max(1, min(n, min(tiles_total, 32)));
+    int max_by_work = tiles_total / 8;
+    return max(1, min(min(min(n, tiles_total), 32), max_by_work));
+}
+
+// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
+// scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu.
+#ifndef ASTRAI_NO_MMA
+static bool decode_use_mma(const AttentionParams& p) {
+    int G = p.q_head / p.kv_head;
+    return !p.use_mask && G > 1 && G <= 16;
 }
 
 template <int HEAD_DIM, int BC>
-static void launch_mma_decode(GQAParams& p, DecodeScratch& sc) {
+static void launch_mma_decode(AttentionParams& p, DecodeScratch& sc) {
-    constexpr int BR = 16, LD = HEAD_DIM;
-    int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    int num_splits = decode_num_splits(p, tiles_total);
+    p.num_splits = decode_num_splits(p, tiles_total);
+    p.o_part = sc.o_part;
+    p.ml_part = sc.ml_part;
 
-    if (num_splits <= 1) {
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
-        cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
+        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
-                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-        gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
-            <<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
-        return;
-    }
-    cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
-                         cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-    gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
-        <<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
-            p, sc.o_part, sc.ml_part, num_splits);
-    gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
-        sc.o_part, sc.ml_part, p.o, num_splits, p.head_dim);
 }
 #endif
 
-static void launch_scalar_decode(const GQAParams& p) {
+static void launch_scalar_decode(AttentionParams& p, DecodeScratch& sc) {
     int gs = p.q_head / p.kv_head;
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    p.num_splits = decode_num_splits(p, chunks_total);
+    p.o_part = sc.o_part;
+    p.ml_part = sc.ml_part;
+
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-    gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, gs), smem>>>(p);
+    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, gs), smem>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 
 template <int HEAD_DIM>
-static void dispatch_decode_t(GQAParams& p, DecodeScratch& sc) {
+static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) {
 #ifndef ASTRAI_NO_MMA
     if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
 #endif
-    (void)sc;
+    launch_scalar_decode(p, sc);
-    launch_scalar_decode(p);
 }
 
-static void dispatch_decode(GQAParams& p, DecodeScratch& sc) {
+static void dispatch_decode(AttentionParams& p, DecodeScratch& sc) {
     switch (p.head_dim) {
         case 32:  dispatch_decode_t<32>(p, sc);  break;
         case 64:  dispatch_decode_t<64>(p, sc);  break;
@@ -169,7 +164,7 @@ static void bench() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        GQAParams p;
+        AttentionParams 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=0; p.is_causal=0; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);
@@ -245,7 +240,7 @@ int main() {
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
         cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
 
-        GQAParams p;
+        AttentionParams 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=0; p.is_causal=0; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);

csrc/tests/attn_prefill_test.cu

@@ -2,16 +2,16 @@
 Pure-C test:
 nvcc -I csrc -arch=sm_89 -O3 \
     --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/gqa_prefill_test.cu -o test && ./test
+    csrc/tests/attn_prefill_test.cu -o test && ./test
 */
 
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
 #include <sys/time.h>
-#include "../kernels/gqa_prefill_attn.cuh"
+#include "../kernels/attn_prefill_split_q.cuh"
 #ifndef ASTRAI_NO_MMA
-#include "../kernels/gqa_prefill_attn_mma.cuh"
+#include "../kernels/attn_prefill_split_q_mma.cuh"
 #endif
 
 static double now_ms() {
@@ -21,9 +21,9 @@ static double now_ms() {
 }
 
 // Launch the production prefill path (tensor-core MMA on sm_80+, else the
-// scalar fallback), mirroring dispatch_prefill() in gqa_prefill_attn.cu.
+// scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
 template <int HEAD_DIM>
-static void launch_prefill(GQAParams& p) {
+static void launch_prefill(AttentionParams& p) {
 #ifndef ASTRAI_NO_MMA
     constexpr int WARPS = 4, BR = 16;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
@@ -31,16 +31,16 @@ static void launch_prefill(GQAParams& p) {
                              : (HEAD_DIM <= 128) ? 3 : 2;
     dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
     dim3 block(WARPS * 32, 1, 1);
-    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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);
-    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
 #endif
 }
 
-static void dispatch_prefill(GQAParams& p) {
+static void dispatch_prefill(AttentionParams& p) {
     switch (p.head_dim) {
         case 64:  launch_prefill<64>(p);  break;
         case 128: launch_prefill<128>(p); break;
@@ -123,7 +123,7 @@ static void bench() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        GQAParams p;
+        AttentionParams 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);
@@ -191,7 +191,7 @@ int main() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        GQAParams p;
+        AttentionParams 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);