refactor: remove ineffective __launch_bounds__ from prefill kernel

- Remove MIN_BLOCKS template param and __launch_bounds__ attribute
- Profiling shows smem (not registers) is the occupancy bottleneck for D>=64, making the hint a no-op
- D=64 sees 2-4% speedup, D=128 unchanged (smem-capped at 1 block/SM)
- Update comment blocks in kernel header and both dispatch sites
- Verified correctness via standalone CUDA test (max_err ~1e-4)

csrc/kernels/attn_prefill.cu

@@ -15,15 +15,11 @@ static void dispatch_prefill(AttentionParams<bf16>& p) {
     // 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.
     // 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);
+    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
+// thread group. Templated on <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
-// 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).
 //
 // 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
@@ -40,9 +35,8 @@ using bf16 = __nv_bfloat16;
 // 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>
+template <int HEAD_DIM, int WARPS, int BC>
-__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
+__global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
-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_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);