diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index f0b0536..e233a43 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -8,10 +8,17 @@ template static void dispatch_prefill(GQAParams& p) { #ifndef ASTRAI_NO_MMA - constexpr int WARPS = 4, BC = 16, BR = 16; - // Double-buffered K/V doubles smem, so BC is halved to 16 to keep 3+ - // blocks/SM. Register-hint MIN_BLOCKS tuned per HEAD_DIM's smem footprint. - constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2; + constexpr int WARPS = 4, BR = 16; + // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier + + // 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). + 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. diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh index 8266886..87d1f8e 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/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; diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu index 6d2a753..2576eb9 100644 --- a/csrc/tests/gqa_prefill_test.cu +++ b/csrc/tests/gqa_prefill_test.cu @@ -25,8 +25,10 @@ static double now_ms() { template static void launch_prefill(GQAParams& p) { #ifndef ASTRAI_NO_MMA - constexpr int WARPS = 4, BC = 16, BR = 16; - constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2; + 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); gqa_prefill_attn_mma_kernel<<>>(p);