chore: bump version to 1.3.9

- Replace 4*NC8 per-element -FLT_MAX comparisons with 2 row-level pn guards; masked entries naturally underflow via expf(-FLT_MAX - nm) ≈ 0; pn only guards all-masked-row edge where nm == -FLT_MAX (exp(0)=1 not 0); ~1-3% speedup on prefill (verified via standalone CUDA bench); correctness verified: prefill 4/4, decode 3/3, paged decode 13/13
- Fix stale comments: remove false 'pre-scale Q' claim, correct occupancy numbers, remove phantom sQ from smem description

.github/workflows/release.yml

@@ -0,0 +1,71 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  build-pure:
+    name: Build pure-Python wheel
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build wheel (no CUDA)
+        run: |
+          pip wheel . --no-deps -w dist/
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: pure-wheel
+          path: dist/*.whl
+
+  build-cuda-linux:
+    name: Build CUDA wheel (Linux)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install torch (CUDA 12.8)
+        run: |
+          pip install torch --index-url https://download.pytorch.org/whl/cu128
+
+      - name: Setup CUDA
+        uses: Jimver/cuda-toolkit@v0.2.35
+        with:
+          cuda: "12.8.0"
+
+      - name: Build wheel (with CUDA kernels)
+        run: |
+          CSRC_KERNELS=true pip wheel . --no-deps --no-build-isolation -w dist/
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cuda-wheel-linux
+          path: dist/*.whl
+
+  release:
+    name: Attach wheels to release
+    needs: [build-pure, build-cuda-linux]
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: "*-wheel"
+          merge-multiple: true
+
+      - name: Create release & upload assets
+        uses: softprops/action-gh-release@v2
+        with:
+          files: ./*.whl
+          tag_name: ${{ github.ref_name }}
+          generate_release_notes: true

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.8"
+__version__ = "1.3.9"
 __author__ = "ViperEkura"
 
 from astrai.config import (

astrai/parallel/setup.py

@@ -1,14 +1,21 @@
 import os
+import socket
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from functools import wraps
-from typing import Callable
+from typing import Callable, Optional
 
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
 
+def find_free_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return str(s.getsockname()[1])
+
+
 def get_current_device():
     return os.environ["LOCAL_DEVICE"]
 
@@ -217,11 +224,13 @@ def spawn_parallel_fn(
     world_size: int,
     backend: str = "nccl",
     master_addr: str = "localhost",
-    master_port: str = "29500",
+    master_port: Optional[str] = None,
     device_type: str = "cuda",
     start_method: str = "spawn",
     **kwargs,
 ):
+    if master_port is None:
+        master_port = find_free_port()
     launcher = _detect_launcher()
     if launcher in ("torchelastic", "torchrun", "external"):
         strategy = TorchrunStrategy(

csrc/build.py

@@ -20,7 +20,7 @@ def _arch_flags() -> list[str]:
 _kernels_dir = Path("csrc/kernels")
 REGISTRY: dict[str, dict] = {}
 
-CXX_FLAGS = ["-O3", "-march=native", "-funroll-loops"]
+CXX_FLAGS = ["-O3", "-funroll-loops"]
 NVCC_FLAGS = [
     "-O3",
     "--expt-relaxed-constexpr",

csrc/kernels/attn_decode.cu

@@ -30,9 +30,12 @@ static void launch_scalar_decode(AttentionParams<bf16>& p) {
 }
 
 #ifndef ASTRAI_NO_MMA
-// MMA head-packing requires G <= 16 (sQ has BR=16 rows). sm_80+ tensor-core
+// MMA head-packing requires G <= 16 (BR=16 rows). sm_80+ tensor-core
 // + cp.async wins even at G=1 (decode is memory-bound, not compute-bound).
-template <int HEAD_DIM, int BC>
+// STAGES=2 (double-buffer) for D<=128 (smem 16 KB); STAGES=1 for D=256
+// (double-buffer would be 32 KB, near the 48 KB static cap — keep single
+// to preserve occupancy).
+template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
@@ -43,8 +46,7 @@ static void launch_mma_decode(AttentionParams<bf16>& p) {
     p.o_part = o_part.data_ptr<float>();
     p.ml_part = ml_part.data_ptr<float>();
 
-    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
-        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #endif

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -27,9 +27,18 @@ using bf16 = __nv_bfloat16;
 // 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>
+//   - Q loaded directly from global into mma A-operand registers (no sQ staging,
+//     no prologue syncwarp) — frees shared memory for double-buffering
+//   - Double-buffered KV (STAGES=2): next tile's cp.async overlaps current
+//     tile's MMA compute — hides global load latency / boosts bandwidth
+//     utilization for small-batch (low-occupancy) decode
+//   - Predicated cp.async (cp_async_16_pred) for full AND partial tiles on one
+//     uniform path — eliminates the scalar fallback branch
+//
+// Smem footprint (BC=32): STAGES=2 → 2*(sK+sV) = 2*2*32*HEAD_DIM*2 bytes.
+//   D=128: 16 KB (fits 48 KB static cap).  D=256: 32 KB (also fits).
+// STAGES=1 fallback (4/8 KB) for smem-constrained configs.
+template <int HEAD_DIM, int BC, int STAGES = 2>
 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;
@@ -38,6 +47,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     constexpr int DN8 = HEAD_DIM / 8;
     constexpr int LD = HEAD_DIM;
     constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
+    constexpr int VEC = 8;
+    constexpr int TOTAL = BC * HEAD_DIM;
 
     const int lane = threadIdx.x;
     const int gid = lane >> 2;
@@ -49,27 +60,31 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     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];
-
-    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)] = val;
-    }
-    __syncwarp();
+    // Double-buffered shared memory for K/V (no sQ needed — Q goes direct
+    // from global to registers).
+    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
+    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
+    // ---- Load Q directly from global into mma A-operand registers ----
+    // Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8
+    // cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9.
+    const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM;
+    const int qra = gid;
+    const int qrb = gid + 8;
+    const bool va = qra < G, vb = qrb < G;
     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)]);
+    for (int kt = 0; kt < KD; kt++) {
+        int c = kt * 16 + tid4 * 2;
+        const unsigned* pau = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qra * HEAD_DIM + c]);
+        const unsigned* pbu = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qrb * HEAD_DIM + c]);
+        Qa[kt][0] = va ? pau[0] : 0u;
+        Qa[kt][1] = vb ? pbu[0] : 0u;
+        Qa[kt][2] = va ? pau[4] : 0u;
+        Qa[kt][3] = vb ? pbu[4] : 0u;
+    }
 
     float Oacc[DN8][4];
 #pragma unroll
@@ -85,39 +100,48 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     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++) {
+    // ---- Load tile lambda: predicated cp.async, unified full/partial ----
+    auto load_tile = [&](int ti, int buf) {
         int kv0 = ti * BC;
-
-        bool full_tile = (kv0 + BC <= p.kv_len);
-        if (full_tile) {
-            constexpr int VEC = 8;
-            int total = BC * HEAD_DIM;
+        bf16* dK = sK + buf * BC * LD;
+        bf16* dV = sV + buf * BC * LD;
 #pragma unroll
-            for (int i = lane * VEC; i < total; i += 32 * VEC) {
+        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]);
+            bool valid = kc < p.kv_len;
+            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
+            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
+            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid);
         }
         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;
-            }
+    };
+
+    // ---- Prologue: issue first tile load ----
+    if (ti_begin < ti_end) {
+        load_tile(ti_begin, 0);
     }
+
+    for (int ti = ti_begin; ti < ti_end; ti++) {
+        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
+        int buf = (ti - ti_begin) & BUF_MASK;
+
+        // Wait for current tile, then issue next tile's prefetch (overlaps
+        // with this tile's compute). Single syncwarp covers both hazards.
+        // When STAGES==1, no prefetch — load happens at end of prior iter.
+        cp_async_wait_group<0>();
         __syncwarp();
+        if constexpr (STAGES > 1) {
+            if (ti + 1 < ti_end)
+                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
+        }
+
+        const bf16* bK = sK + buf * BC * LD;
+        const bf16* bV = sV + buf * BC * LD;
+        int kv0 = ti * BC;
 
         float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
+        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
 
         #pragma unroll
         for (int n8 = 0; n8 < NC8; n8++)
@@ -129,8 +153,13 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
                                     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);
+        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
         __syncwarp();
+
+        if constexpr (STAGES == 1) {
+            if (ti + 1 < ti_end)
+                load_tile(ti + 1, 0);
+        }
     }
 
     // ---- write UN-normalised partials for this split ----
@@ -169,4 +198,3 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         }
     }
 }
-

csrc/kernels/attn_mma_utils.cuh

@@ -114,7 +114,8 @@ __device__ __forceinline__ void cp_async_wait_group() {
 // between the two kernels; only the per-row causal/mask bounds differ.
 // ---------------------------------------------------------------------------
 
-// S = Q @ K^T  (Qa pre-loaded and pre-scaled by the caller).
+// S = Q @ K^T  (Qa pre-loaded by the caller; scale applied post-mma in the
+// caller to avoid bf16 precision loss).
 // LD and SWIZ_MASK are constexpr in the calling kernel — passing them as
 // runtime ints lets the compiler fold them while keeping the signature clean.
 template <int KD, int NC8>
@@ -138,10 +139,11 @@ __device__ inline void mma_compute_scores(
     }
 }
 
-// Online softmax + Oacc rescale for one K/V tile. maxc0/maxc1 are the per-row
-// KV column bounds — prefill passes per-query-row causal limits while decode
-// passes the same value for both rows (q_len==1). Sacc is consumed in place
-// (replaced by P = exp(S - nm) for the subsequent P@V step).
+// Online softmax + Oacc rescale for one K/V tile.
+//   maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits;
+//     decode: same value for both rows since q_len==1).
+// Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm),
+// rescales Oacc by exp(m_old - nm), and updates m/l — all in place.
 template <int NC8, int DN8>
 __device__ inline void mma_softmax_tile(
     int kv0,
@@ -157,6 +159,8 @@ __device__ inline void mma_softmax_tile(
 {
     int tid4 = lane & 3;
 
+    // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream
+    // without per-element sentinel checks. Compute tile-local row maxima.
     float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
     #pragma unroll
     for (int n8 = 0; n8 < NC8; n8++) {
@@ -175,22 +179,33 @@ __device__ inline void mma_softmax_tile(
         rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
         rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
     }
+    // Warp-reduce row maxima across the 4-lane thread group (xor 1, xor 2).
     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));
 
+    // nm = max(running max m, tile-local max rmax) — updated running maximum.
     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);
+    // corr rescales Oacc and l by exp(m_old - nm). When all-masked (m == nm ==
+    // -FLT_MAX), exp(0) = 1 — correct, no guard needed.
+    float corr0 = __expf(m0 - nm0);
+    float corr1 = __expf(m1 - nm1);
+    // pn guards only the all-masked-row edge: if nm == -FLT_MAX, exp(S - nm)
+    // gives 1 not 0 for masked entries. Two scalar masks replace 4*NC8
+    // per-element comparisons.
+    float pn0 = (nm0 == -FLT_MAX) ? 0.0f : 1.0f;
+    float pn1 = (nm1 == -FLT_MAX) ? 0.0f : 1.0f;
 
+    // P = exp(S - nm) for each element. Masked entries (Sacc = -FLT_MAX) give
+    // exp(-inf) ≈ 0 naturally; pn zero-fills the all-masked-row edge.
     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);
+        float p0 = pn0 * __expf(Sacc[n8][0] - nm0);
+        float p1 = pn0 * __expf(Sacc[n8][1] - nm0);
+        float p2 = pn1 * __expf(Sacc[n8][2] - nm1);
+        float p3 = pn1 * __expf(Sacc[n8][3] - nm1);
         Sacc[n8][0] = p0; Sacc[n8][1] = p1;
         Sacc[n8][2] = p2; Sacc[n8][3] = p3;
         rsum0 += p0 + p1;

csrc/kernels/attn_paged_decode.cu

@@ -25,15 +25,14 @@ static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
     p.ml_part = ml_part.data_ptr<float>();
 
     size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
-    paged_attn_decode_split_kv_kernel<<<
-        dim3(p.batch * p.kv_head, 1, p.num_splits),
-        dim3(32, group_size),
-        smem>>>(p);
+    dim3 grid = dim3(p.batch * p.kv_head, 1, p.num_splits);
+    dim3 block = dim3(32, group_size);
+    paged_attn_decode_split_kv_kernel<<<grid, block, smem>>>(p);
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 
 #ifndef ASTRAI_NO_MMA
-template <int HEAD_DIM, int BC>
+template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
@@ -44,8 +43,7 @@ static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
     p.o_part = o_part.data_ptr<float>();
     p.ml_part = ml_part.data_ptr<float>();
 
-    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
-        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #endif
@@ -104,13 +102,12 @@ torch::Tensor attn_paged_decode(
     p.batch = batch;
     p.q_head = q_head;
     p.kv_head = kv_head;
-    p.q_len = 1;
+    p.q_len = static_cast<int>(q.size(2));
     p.kv_len = static_cast<int>(kv_len);
     p.head_dim = head_dim;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
     p.is_causal = is_causal ? 1 : 0;
     p.causal_offset = static_cast<int>(causal_offset);
-    p.num_splits = 1;
     p.scale = scale_val;
     p.page_size = static_cast<int>(page_size);
     p.max_pages = max_pages;

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -11,7 +11,12 @@ using bf16 = __nv_bfloat16;
 // directly from the page pool through a page table, eliminating the gather
 // copy.  Each tile (BC=32) fits within a single page (page_size >= 32), so
 // the page-table lookup happens once per tile for cp.async.
-template <int HEAD_DIM, int BC>
+//
+// Optimizations mirror attn_decode_split_kv_mma_kernel:
+//   - Q loaded directly from global into mma A-operand registers (no sQ)
+//   - Double-buffered KV (STAGES=2) for D<=128, single-buffer for D=256
+//   - Predicated cp.async for unified full/partial tile path
+template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
     constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;
@@ -20,6 +25,8 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     constexpr int DN8 = HEAD_DIM / 8;
     constexpr int LD = HEAD_DIM;
     constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
+    constexpr int VEC = 8;
+    constexpr int TOTAL = BC * HEAD_DIM;
 
     const int lane = threadIdx.x;
     const int gid = lane >> 2;
@@ -31,31 +38,30 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     const int G = p.q_head / p.kv_head;
     const int q_head0 = kv_head_idx * 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];
-
-    // ---- load Q into registers via ldmatrix ----
-    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)] = val;
-    }
-    __syncwarp();
+    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
+    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
+    // ---- Load Q directly from global into mma A-operand registers ----
+    const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM;
+    const int qra = gid;
+    const int qrb = gid + 8;
+    const bool va = qra < G, vb = qrb < G;
     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)]);
+#pragma unroll
+    for (int kt = 0; kt < KD; kt++) {
+        int c = kt * 16 + tid4 * 2;
+        const unsigned* pau = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qra * HEAD_DIM + c]);
+        const unsigned* pbu = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qrb * HEAD_DIM + c]);
+        Qa[kt][0] = va ? pau[0] : 0u;
+        Qa[kt][1] = vb ? pbu[0] : 0u;
+        Qa[kt][2] = va ? pau[4] : 0u;
+        Qa[kt][3] = vb ? pbu[4] : 0u;
+    }
 
     float Oacc[DN8][4];
-    #pragma unroll
+#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;
@@ -72,57 +78,54 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     const int64_t pos_stride  = (int64_t)p.kv_head * HEAD_DIM;
     const int64_t head_off    = (int64_t)kv_head_idx * HEAD_DIM;
 
-    for (int ti = ti_begin; ti < ti_end; ti++) {
+    // ---- Load tile lambda: predicated cp.async, paged addressing ----
+    auto load_tile = [&](int ti, int buf) {
         int kv0 = ti * BC;
-
-        // phys_page is constant for the whole tile (BC <= page_size).
+        bf16* dK = sK + buf * BC * LD;
+        bf16* dV = sV + buf * BC * LD;
         int logical_page = kv0 / p.page_size;
         int phys_page = p.page_table[batch * p.max_pages + logical_page];
         bool page_valid = (phys_page >= 0);
-
-        bool full_tile = page_valid && (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) {
+#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;
+            bool valid = (kc < p.kv_len) && page_valid;
             int page_off = kc % p.page_size;
             int64_t gmem_base = (int64_t)phys_page * page_stride
                               + (int64_t)page_off * pos_stride
                               + head_off;
-                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.k_cache[gmem_base + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.v_cache[gmem_base + d]);
+            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
+            cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
+            cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
         }
         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);
-                if (kc < p.kv_len && page_valid) {
-                    int page_off = kc % p.page_size;
-                    int64_t gmem_base = (int64_t)phys_page * page_stride
-                                      + (int64_t)page_off * pos_stride
-                                      + head_off;
-                    sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.k_cache[gmem_base + d];
-                    sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.v_cache[gmem_base + d];
-                } else {
-                    sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = z;
-                    sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = z;
-                }
-            }
+    };
+
+    // ---- Prologue: issue first tile load ----
+    if (ti_begin < ti_end) {
+        load_tile(ti_begin, 0);
     }
+
+    for (int ti = ti_begin; ti < ti_end; ti++) {
+        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
+        int buf = (ti - ti_begin) & BUF_MASK;
+
+        cp_async_wait_group<0>();
         __syncwarp();
+        if constexpr (STAGES > 1) {
+            if (ti + 1 < ti_end)
+                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
+        }
+
+        const bf16* bK = sK + buf * BC * LD;
+        const bf16* bV = sV + buf * BC * LD;
+        int kv0 = ti * BC;
 
         float Sacc[NC8][4];
-        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
+        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
 
-        #pragma unroll
+#pragma unroll
         for (int n8 = 0; n8 < NC8; n8++)
             Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
             Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
@@ -132,8 +135,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
                                     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);
+        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
         __syncwarp();
+
+        if constexpr (STAGES == 1) {
+            if (ti + 1 < ti_end)
+                load_tile(ti + 1, 0);
+        }
     }
 
     // ---- write UN-normalised partials for this split ----
@@ -141,7 +149,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         size_t bh = (size_t)batch * p.q_head + h;
         return bh * p.num_splits + split;
     };
-    #pragma unroll
+#pragma unroll
     for (int dn8 = 0; dn8 < DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;
         int r0 = gid, r1 = gid + 8;

csrc/kernels/attn_prefill.cu

@@ -13,17 +13,15 @@ static void dispatch_prefill(AttentionParams<bf16>& p) {
     // loop overhead over more tensor-core work (this kernel is latency-bound,
     // not compute/bandwidth-bound), so BC=32 wins ~6-8% over BC=16 for
     // D<=128. D=256 stays at 16: BC=32 double-buffered would need 64KB smem,
-    // over the 48KB static cap. Both keep 3 blocks/SM (2 for D=256).
+    // over the 48KB static cap.
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
-    // Register-hint MIN_BLOCKS tuned per HEAD_DIM's (BC=32) smem+register
-    // footprint: the largest blocks/SM that avoids register spills.
-    constexpr int MIN_BLOCKS = (HEAD_DIM <= 32) ? 6 : (HEAD_DIM <= 64) ? 4
-                             : (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);
-    // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed.
-    // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params.
-    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
+    // Static shared memory — double-buffered K/V only (no sQ: Q goes direct
+    // to registers). 2*BC*LD bf16 each for sK and sV → 4*BC*HEAD_DIM*2 bytes.
+    // Occupancy is smem-capped: D=64→3 blocks/SM (16KB), D=128→1 (32KB),
+    // D=256→1 (32KB, BC=16).
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC><<<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);

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -16,12 +16,7 @@ using bf16 = __nv_bfloat16;
 // allocation. 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, MIN_BLOCKS> with BC a
-// multiple of 16.
-//
-// Occupancy: __launch_bounds__ forces the compiler to fit MIN_BLOCKS blocks/SM,
-// spilling to local memory as needed. MIN_BLOCKS is tuned per HEAD_DIM to the
-// double-buffered smem footprint (2*BC*LD for each of K/V).
+// thread group. Templated on <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
 //
 // Software pipeline: K/V are double-buffered and loaded via cp.async one tile
 // ahead, so the next tile streams from global memory while the current tile's
@@ -35,14 +30,14 @@ using bf16 = __nv_bfloat16;
 // register pressure), so fewer, larger tiles beat many tiny ones.
 //
 // Optimizations: load Q fragments directly from global in mma A-operand layout
-// (no sQ staging, no prologue barriers); pre-scale Q by attention scale during Q load; packed bf16x2 output stores;
+// (no sQ staging, no prologue barriers); post-multiply scale in float after
+// S=Q@K^T to avoid bf16 precision loss; packed bf16x2 output stores;
 // causal tile skipping (block-level prefetch bound + warp-level compute skip);
 // XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD
 // padding (LD=HEAD_DIM).
 
-template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
-__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
-void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
+template <int HEAD_DIM, int WARPS, int BC>
+__global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> 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/tests/attn_decode_test.cu

@@ -27,14 +27,14 @@ static bool decode_use_mma(const AttentionParams<bf16>& p) {
     return !p.use_mask && G > 1 && G <= 16;
 }
 
-template <int HEAD_DIM, int BC>
+template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
     p.o_part = sc.o_part;
     p.ml_part = sc.ml_part;
 
-    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES>
         <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
@@ -61,90 +61,64 @@ static void dispatch_decode_t(AttentionParams<bf16>& p, DecodeScratch& sc) {
 }
 
 static void dispatch_decode(AttentionParams<bf16>& 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;
-        case 128: dispatch_decode_t<128>(p, sc); break;
-        case 256: dispatch_decode_t<256>(p, sc); break;
-        default:  printf("bench: unsupported D=%d\n", p.head_dim);
-    }
+    dispatch_by_head_dim(p.head_dim, [&]<int D>() { dispatch_decode_t<D>(p, sc); });
 }
 
 // Warmed-up, CUDA-event timed sweep over the production decode MMA path.
-// Decode (q_len==1) is memory-bound: the two matmuls are GEMV-shaped, so we
-// report both effective K/V read bandwidth and the (small) attention FLOP/s.
-// FLOP/s = 2 matmuls (q@K^T, P@V), each 2*B*Hq*kv*D flops.
-// Bytes    = K + V read = 2 * B*Hk*kv*D * sizeof(bf16).
 static void bench() {
     const int cfgs[][5] = {
-        {1, 32, 4, 512, 128},    // B,Hq,Hk,seq,D
+        {1, 32, 4, 512, 128},    // B, Hq, Hk, kv_len, D
         {1, 32, 4, 1024, 128},
         {1, 32, 4, 2048, 128},
         {1, 32, 4, 4096, 128},
         {16, 32, 4, 2048, 128},
         {32, 32, 4, 1024, 128},
     };
-    int n = sizeof(cfgs)/sizeof(cfgs[0]);
     const int WARMUP = 10, ITERS = 100;
     printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
-    printf("---------------------------------------------------------------"
-           "----------------------------\n");
+    print_bench_header();
 
-    for (int ci = 0; ci < n; ci++) {
-        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
-        int sl=cfgs[ci][3], D=cfgs[ci][4];
-        size_t nQ=(size_t)B*Hq*D, nKV=(size_t)B*Hk*sl*D;
+    for (int ci = 0; ci < 6; ci++) {
+        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
+        int sl = cfgs[ci][3], D = cfgs[ci][4];
+        size_t nQ = (size_t)B * Hq * D;
+        size_t nKV = (size_t)B * Hk * sl * D;
 
-        bf16 *dQ,*dK,*dV,*dO,*tmp;
-        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
-        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
-        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+        bf16 *dQ, *dK, *dV, *dO;
+        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
+        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
+        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
+        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
+        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
+        delete[] tmp;
 
         AttentionParams<bf16> 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);
-        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+        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);
+        p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
 
         DecodeScratch sc;
         cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
         cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
 
-        for (int i=0;i<WARMUP;i++) dispatch_decode(p, sc);
-        cudaDeviceSynchronize();
-        cudaError_t err=cudaGetLastError();
-        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
-
-        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
-        cudaEventRecord(s);
-        for (int i=0;i<ITERS;i++) dispatch_decode(p, sc);
-        cudaEventRecord(e); cudaEventSynchronize(e);
-        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
-
-        double flops = 4.0*B*Hq*(double)sl*D;
-        double tflops = flops/(ms*1e-3)/1e12;
-        // HBM traffic: K + V read (B*Hk*sl*D each), bf16; Q/O negligible.
-        double bytes = 2.0 * (2.0*nKV);
-        double gbps = bytes/(ms*1e-3)/1e9;
+        auto launch = [&]() { dispatch_decode(p, sc); };
+        double flops = 4.0 * B * Hq * (double)sl * D;
+        double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
+        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
 
         char cfg[64];
         snprintf(cfg, sizeof(cfg),
                  "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
-                 B,Hq,Hk,1,sl,D,0);
-        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-               cfg, ms, gbps, tflops);
+                 B, Hq, Hk, 1, sl, D, 0);
+        print_bench_row(cfg, r);
 
-        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
-        cudaFree(sc.o_part);cudaFree(sc.ml_part);
-        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
+        cudaFree(sc.o_part); cudaFree(sc.ml_part);
     }
 }
 

csrc/tests/attn_paged_decode_test.cu

@@ -42,9 +42,10 @@ static void launch_paged_decode(PagedAttentionParams<bf16, float>& p) {
     int G_check = p.q_head / p.kv_head;
     bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
     if (use_mma) {
+        constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1;
         int tiles_total = (p.kv_len + 32 - 1) / 32;
         p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
-        paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32>
+        paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32, STAGES>
             <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     } else
 #endif
@@ -221,17 +222,16 @@ static const TestCase TESTS[] = {
 };
 
 static int dispatch_test(const TestCase& tc) {
-    switch (tc.head_dim) {
-        case 32:  return run_test<32>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
-        case 64:  return run_test<64>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
-        case 128: return run_test<128>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
-        case 256: return run_test<256>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
-        default:  return 1;
-    }
+    bool matched = false;
+    int r = 0;
+    dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
+        matched = true;
+        r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+    });
+    return matched ? r : 1;
 }
 
 // Warmed-up, CUDA-event timed sweep over paged decode configs.
-// Reports per-call latency and effective K/V read bandwidth.
 // Bytes = K + V read through page table (B*Hk*kv*D each), bf16.
 template <int HEAD_DIM>
 static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
@@ -281,43 +281,27 @@ static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
     pa.o_part = d_op; pa.ml_part = d_ml;
 
     const int WARMUP = 10, ITERS = 100;
-    for (int i = 0; i < WARMUP; i++) launch_paged_decode<HEAD_DIM>(pa);
-    cudaDeviceSynchronize();
-    CUDA_CHECK(cudaGetLastError());
-
-    cudaEvent_t s, e;
-    cudaEventCreate(&s); cudaEventCreate(&e);
-    cudaEventRecord(s);
-    for (int i = 0; i < ITERS; i++) launch_paged_decode<HEAD_DIM>(pa);
-    cudaEventRecord(e); cudaEventSynchronize(e);
-    float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= ITERS;
-
+    auto launch = [&]() { launch_paged_decode<HEAD_DIM>(pa); };
     double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
-    double tflops = flops / (ms * 1e-3) / 1e12;
     size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
-    double bytes = 2.0 * (2.0 * nKV);
-    double gbps = bytes / (ms * 1e-3) / 1e9;
+    double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
+    BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
 
     char cfg[64];
     snprintf(cfg, sizeof(cfg),
              "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
              B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
-    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-           cfg, ms, gbps, tflops);
+    print_bench_row(cfg, r);
 
     free(tmp);
     cudaFree(d_q); cudaFree(d_o);
     cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
     cudaFree(d_op); cudaFree(d_ml);
-    cudaEventDestroy(s); cudaEventDestroy(e);
 }
 
 static void bench() {
     printf("\n===== PAGED DECODE BENCH =====\n");
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
-    printf("---------------------------------------------------------------"
-           "----------------------------\n");
+    print_bench_header();
     bench_config<128>(1, 32, 4, 512, 128);
     bench_config<128>(1, 32, 4, 1024, 128);
     bench_config<128>(1, 32, 4, 2048, 128);

csrc/tests/attn_prefill_test.cu

@@ -18,11 +18,9 @@ static void launch_prefill(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     constexpr int WARPS = 4, BR = 16;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
-    constexpr int MIN_BLOCKS = (HEAD_DIM <= 32) ? 6 : (HEAD_DIM <= 64) ? 4
-                             : (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);
-    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC><<<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);

csrc/tests/test_utils.cuh

@@ -37,6 +37,70 @@ inline int compute_num_splits(int base_blocks, int tiles_total) {
         } \
     } while (0)
 
+struct BenchResult {
+    float ms;
+    double gbps;
+    double tflops;
+};
+
+template <typename Fn>
+BenchResult bench_kernel(Fn launch, int warmup, int iters,
+                         double flops, double bytes) {
+    for (int i = 0; i < warmup; i++) launch();
+    cudaDeviceSynchronize();
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("CUDA error before bench: %s\n", cudaGetErrorString(err));
+        return {0, 0, 0};
+    }
+
+    cudaEvent_t s, e;
+    cudaEventCreate(&s); cudaEventCreate(&e);
+    cudaEventRecord(s);
+    for (int i = 0; i < iters; i++) launch();
+    cudaEventRecord(e); cudaEventSynchronize(e);
+    float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= iters;
+    cudaEventDestroy(s); cudaEventDestroy(e);
+
+    return {ms, bytes / (ms * 1e-3) / 1e9, flops / (ms * 1e-3) / 1e12};
+}
+
+inline void print_bench_header() {
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+}
+
+inline void print_bench_row(const char* cfg, const BenchResult& r) {
+    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+           cfg, r.ms, r.gbps, r.tflops);
+}
+
+template <int... Ds>
+struct _HeadSwitch;
+
+template <int D>
+struct _HeadSwitch<D> {
+    template <typename Fn>
+    static void call(int hd, Fn&& fn) { if (hd == D) fn.template operator()<D>(); }
+};
+
+template <int D, int... Rest>
+struct _HeadSwitch<D, Rest...> {
+    template <typename Fn>
+    static void call(int hd, Fn&& fn) {
+        if (hd == D) fn.template operator()<D>();
+        else _HeadSwitch<Rest...>::call(hd, fn);
+    }
+};
+
+// Default set: 32, 64, 128, 256
+template <typename Fn>
+void dispatch_by_head_dim(int head_dim, Fn&& fn) {
+    _HeadSwitch<32, 64, 128, 256>::call(head_dim, fn);
+}
+
 // Generic CPU reference for multi-query / grouped-query attention.
 // Tensor shapes (all float*):
 //   Q : [B, Hq, q_len, D]