From 85dc77146081969575b80f9f38dbcf19c7bdeb8c Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 10 Jul 2026 00:39:47 +0800 Subject: [PATCH] perf: reduce MMA kernel registers, switch to static smem MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 | 15 ++++--- csrc/kernels/gqa_prefill_attn.cuh | 5 +-- csrc/kernels/gqa_prefill_attn_mma.cuh | 58 ++++++++++++++++----------- csrc/tests/gqa_prefill_test.cu | 9 ++--- 4 files changed, 48 insertions(+), 39 deletions(-) diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index 5aa1f9c..51d474e 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -8,20 +8,19 @@ template 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) - int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16); - cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_prefill_attn_mma_kernel<<>>(p); + // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed. + // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params. + gqa_prefill_attn_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); dim3 block(G, ROWS, 1); - size_t smem = 2 * P_BC * HEAD_DIM * sizeof(bf16); - gqa_prefill_attn_kernel_t<<>>(p); + gqa_prefill_attn_kernel_t<<>>(p); #endif } diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/gqa_prefill_attn.cuh index d0ca014..c360e9b 100644 --- a/csrc/kernels/gqa_prefill_attn.cuh +++ b/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[]; - bf16* sK = smem; - bf16* sV = sK + P_BC * HEAD_DIM; + __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM]; + __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM]; float qreg[DPT]; if (q_row < p.q_len) { diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh index 474fe9a..5b1fd82 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/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; -// S, O, and the online-softmax stats (m, l) live in registers too — nothing is -// staged through shared memory except the cooperatively-loaded K/V tiles. The -// mma fragment layout is used directly: the S accumulator (f32) maps element- +// cores via mma.sync.m16n8k16 (f32 accumulate). Q is staged in static shared +// memory (sQ) and reloaded per tile via ldmatrix — this avoids keeping KD*4 +// fragment registers resident across the tile loop, cutting ~32 regs for +// 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 -// with BC a multiple of 16. +// with BC a multiple of 16. // -// Optimizations: shared sQ staging (single area, serialized per-warp load) -// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→ -// shared for K/V; scalar fallback only for the last partial tile; causal tile -// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col) -// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM). +// Optimizations: shared sQ staging (single area, serialized per-warp load) with +// per-tile reload → cuts registers; pre-scale Q by attention scale during Q +// load; cp.async global→shared for K/V; scalar fallback only for the last +// partial tile; causal tile skipping (block-level early break + warp-level +// skip); XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without +// LD padding (LD=HEAD_DIM). -template -__global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { +template +__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[]; - bf16* sK = smem; // [BC][LD] - bf16* sV = sK + BC * LD; // [BC][LD] - bf16* sQ = sV + BC * LD; // shared staging [BR][LD] + // Static shared memory — sized by template parameters at compile time. + // No extern __shared__ / cudaFuncSetAttribute needed. + __shared__ __align__(16) bf16 sK[BC * 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); } } -} \ No newline at end of file +} diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu index 19ba520..3acebec 100644 --- a/csrc/tests/gqa_prefill_test.cu +++ b/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) smem=%zu\n", - grid.x,grid.y,grid.z, block.x,block.y,block.z, smem); + printf("grid=(%d,%d,%d) block=(%d,%d,%d)\n", + grid.x,grid.y,grid.z, block.x,block.y,block.z); double t0=now_ms(); switch (D) { - case 64: gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<>>(p); break; - case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<>>(p); break; + case 64: gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<>>(p); break; + case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<>>(p); break; default: printf("unsupported D=%d\n",D); return 1; } cudaDeviceSynchronize();