diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index 3b0d745..eda117c 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -15,15 +15,11 @@ static void dispatch_prefill(AttentionParams& 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<<>>(p); + attn_prefill_split_q_mma_kernel<<>>(p); #else constexpr int G = 8, ROWS = 32, P_BC = 32; dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch); diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 617b6a1..10c67eb 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/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 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 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 @@ -40,9 +35,8 @@ using bf16 = __nv_bfloat16; // XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD // padding (LD=HEAD_DIM). -template -__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS) -void attn_prefill_split_q_mma_kernel(AttentionParams p) { +template +__global__ 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) diff --git a/csrc/tests/attn_prefill_test.cu b/csrc/tests/attn_prefill_test.cu index e67885e..6e315a5 100644 --- a/csrc/tests/attn_prefill_test.cu +++ b/csrc/tests/attn_prefill_test.cu @@ -18,11 +18,9 @@ static void launch_prefill(AttentionParams& 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<<>>(p); + attn_prefill_split_q_mma_kernel<<>>(p); #else constexpr int G = 8, ROWS = 32, P_BC = 32; dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);