perf: reduce MMA kernel registers, switch to static smem

- Move Qa[KD][4] into tile loop (reload from sQ per tile)
  cutting ~32 resident registers for HEAD_DIM=128
- Replace extern __shared__ with static template-sized smem
  (no cudaFuncSetAttribute or dynamic allocation needed)
- Add __launch_bounds__ with MIN_BLOCKS param, dispatch by HEAD_DIM
  (hd=128→4, hd=64→6, hd=32→6)
- Remove dynamic smem from scalar kernel and C test
- Result: hd=128 168→128 regs, 25%→33% occupancy

csrc/kernels/gqa_prefill_attn.cu

@@ -8,20 +8,19 @@
 template <int HEAD_DIM>
 static void dispatch_prefill(GQAParams& p) {
 #ifndef ASTRAI_NO_MMA
-    constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM;
+    constexpr int WARPS = 4, BC = 32, BR = 16;
+    // Target higher occupancy for smaller HEAD_DIM (fewer Oacc/Qa registers)
+    constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 4 : 2;
     dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
     dim3 block(WARPS * 32, 1, 1);
-    // sK + sV (each BC*LD) + shared sQ staging (BR*LD)
+    // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed.
-    int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
+    // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params.
-    cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC>,
+    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
-                         cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-    gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block, smem>>>(p);
 #else
     constexpr int G = 8, ROWS = 32, P_BC = 32;
     dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
     dim3 block(G, ROWS, 1);
-    size_t smem = 2 * P_BC * HEAD_DIM * sizeof(bf16);
+    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
-    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block, smem>>>(p);
 #endif
 }
 

csrc/kernels/gqa_prefill_attn.cuh

@@ -43,9 +43,8 @@ __global__ void gqa_prefill_attn_kernel_t(GQAParams p) {
 
     int kv_head = q_head / (p.q_head / p.kv_head);
 
-    extern __shared__ __align__(16) bf16 smem[];
+    __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
-    bf16* sK = smem;
+    __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
-    bf16* sV = sK + P_BC * HEAD_DIM;
 
     float qreg[DPT];
     if (q_row < p.q_len) {

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -2,24 +2,29 @@
 #include "gqa_common.cuh"
 #include "gqa_mma_utils.cuh"
 
-// Tensor-core prefill, register-resident flash attention (raw mma.sync PTX).
+// 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
-// cores via mma.sync.m16n8k16 (f32 accumulate). Q stays resident in registers;
+// cores via mma.sync.m16n8k16 (f32 accumulate). Q is staged in static shared
-// S, O, and the online-softmax stats (m, l) live in registers too — nothing is
+// memory (sQ) and reloaded per tile via ldmatrix — this avoids keeping KD*4
-// staged through shared memory except the cooperatively-loaded K/V tiles. The
+// fragment registers resident across the tile loop, cutting ~32 regs for
-// mma fragment layout is used directly: the S accumulator (f32) maps element-
+// HEAD_DIM=128 and enabling 4 blocks/SM (33% occupancy, up from 25%). S, O,
+// and the online-softmax stats (m, l) live in registers. Shared memory is
+// statically sized via template parameters — no dynamic 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> with BC a multiple of 16.
+// <HEAD_DIM, WARPS, BC, MIN_BLOCKS> with BC a multiple of 16.
 //
-// Optimizations: shared sQ staging (single area, serialized per-warp load)
+// Optimizations: shared sQ staging (single area, serialized per-warp load) with
-// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→
+// per-tile reload → cuts registers; pre-scale Q by attention scale during Q
-// shared for K/V; scalar fallback only for the last partial tile; causal tile
+// load; cp.async global→shared for K/V; scalar fallback only for the last
-// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col)
+// partial tile; causal tile skipping (block-level early break + warp-level
-// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM).
+// skip); XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without
+// LD padding (LD=HEAD_DIM).
 
-template <int HEAD_DIM, int WARPS, int BC>
+template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
-__global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
+__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
+void gqa_prefill_attn_mma_kernel(GQAParams p) {
     constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
     constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
@@ -39,15 +44,17 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     const int kv_head = q_head / (p.q_head / p.kv_head);
     const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
 
-    extern __shared__ __align__(16) bf16 smem[];
+    // Static shared memory — sized by template parameters at compile time.
-    bf16* sK = smem;                       // [BC][LD]
+    // No extern __shared__ / cudaFuncSetAttribute needed.
-    bf16* sV = sK + BC * LD;               // [BC][LD]
+    __shared__ __align__(16) bf16 sK[BC * LD];
-    bf16* sQ = sV + BC * LD;               // shared staging [BR][LD]
+    __shared__ __align__(16) bf16 sV[BC * LD];
+    __shared__ __align__(16) bf16 sQ[BR * LD];
 
-    // Q resident A-fragments (loaded once per warp via shared staging).
+    // Load Q into sQ with pre-scaling (staged per-warp to avoid smem conflicts).
     // Pre-scale by attention scale so softmax doesn't need to multiply later.
+    // Q fragments are NOT kept resident — reloaded from sQ each tile via
+    // ldmatrix to cut ~KD*4 registers.
     const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
-    unsigned Qa[KD][4];
     bf16 scale_bf16 = __float2bfloat16(p.scale);
     int qrow_l = (lane & 7) + (lane & 8);       // 0..15
     int qcol_l = (lane & 16) ? 8 : 0;
@@ -61,9 +68,6 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
                 sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
             }
             __syncwarp();
-        #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)]);
         }
         __syncthreads();  // prevent next warp from overwriting sQ prematurely
     }
@@ -124,6 +128,14 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
         // Warp-level causal skip
         if (!use_skip || kv0 <= max_kv) {
 
+        // Reload Q fragments from sQ each tile — saves KD*4 resident registers.
+        // ldmatrix.sync is warp-cooperative; all 32 lanes execute it together
+        // (the if condition is uniform per warp).
+        unsigned Qa[KD][4];
+#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)]);
+
         // S = Q @ K^T  → Sacc[n8][0..3]   (n8: 8 kv cols each)
         float Sacc[NC8][4];
 #pragma unroll
@@ -243,4 +255,4 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
                 __float2bfloat16(Oacc[dn8][3] * rl1);
         }
     }
-}
+}

csrc/tests/gqa_prefill_test.cu

@@ -97,14 +97,13 @@ int main() {
         constexpr int G=8, ROWS=32, P_BC=32;
         dim3 grid((ql+ROWS-1)/ROWS, Hq, B);
         dim3 block(G, ROWS, 1);
-        size_t smem=2*P_BC*D*sizeof(bf16);
+        printf("grid=(%d,%d,%d) block=(%d,%d,%d)\n",
-        printf("grid=(%d,%d,%d) block=(%d,%d,%d) smem=%zu\n",
+               grid.x,grid.y,grid.z, block.x,block.y,block.z);
-               grid.x,grid.y,grid.z, block.x,block.y,block.z, smem);
 
         double t0=now_ms();
         switch (D) {
-            case 64:  gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            case 64:  gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block>>>(p); break;
-            case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block>>>(p); break;
             default: printf("unsupported D=%d\n",D); return 1;
         }
         cudaDeviceSynchronize();