perf: enlarge prefill KV tile to BC=32 for D<=128

- Kernel is latency-bound (25% occupancy), not compute/bandwidth-bound
- BC=16 wasted a cp.async wait + barrier + loop overhead per tiny tile
- Double KV tile to BC=32 for D<=128; D=256 stays 16 (64KB > 48KB smem cap)
- Retune MIN_BLOCKS per head_dim (32->6, 64->4, 128->3, 256->2)
- Result: ~6-8% faster on L20, 0.93-1.20x vs torch SDPA, correctness unchanged

csrc/kernels/gqa_prefill_attn.cu

@@ -8,10 +8,17 @@
 template <int HEAD_DIM>
 static void dispatch_prefill(GQAParams& p) {
 #ifndef ASTRAI_NO_MMA
-    constexpr int WARPS = 4, BC = 16, BR = 16;
+    constexpr int WARPS = 4, BR = 16;
-    // Double-buffered K/V doubles smem, so BC is halved to 16 to keep 3+
+    // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
-    // blocks/SM. Register-hint MIN_BLOCKS tuned per HEAD_DIM's smem footprint.
+    // loop overhead over more tensor-core work (this kernel is latency-bound,
-    constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2;
+    // 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).
+    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.

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -26,6 +26,9 @@
 // (because it runs before the next prefetch) guards the buffer being refilled,
 // so no second barrier is needed. Predicated cp.async (cp_async_16_pred)
 // zero-fills rows past kv_len, unifying full and partial tiles on one path.
+// BC=32 (D<=128) amortizes the per-tile wait+barrier+loop overhead over more
+// tensor-core work — this kernel is latency-bound (low occupancy from high
+// 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;

csrc/tests/gqa_prefill_test.cu

@@ -25,8 +25,10 @@ static double now_ms() {
 template <int HEAD_DIM>
 static void launch_prefill(GQAParams& p) {
 #ifndef ASTRAI_NO_MMA
-    constexpr int WARPS = 4, BC = 16, BR = 16;
+    constexpr int WARPS = 4, BR = 16;
-    constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2;
+    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);
     gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);