diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index 2b8b4e4..bda8c9d 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -43,11 +43,25 @@ torch::Tensor gqa_prefill_attn( auto O = torch::empty_like(q); p.o = (bf16*)O.data_ptr(); - dim3 grid((p.q_len + Br - 1) / Br, p.q_head, p.batch); - dim3 block(32, Br, 1); - size_t smem = 2 * Bc * p.head_dim * sizeof(bf16); + 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 * p.head_dim * sizeof(bf16); - gqa_prefill_attn_kernel<<>>(p); + switch (p.head_dim) { + 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 256: + gqa_prefill_attn_kernel_t<256, G, ROWS, P_BC><<>>(p); + break; + default: + TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim, + " (supported: 64,128,256)"); + } return O; } diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/gqa_prefill_attn.cuh index 2567089..5aae0c9 100644 --- a/csrc/kernels/gqa_prefill_attn.cuh +++ b/csrc/kernels/gqa_prefill_attn.cuh @@ -1,44 +1,72 @@ #pragma once #include "gqa_common.cuh" -__global__ void gqa_prefill_attn_kernel(GQAParams p) { +// v9: group-split register blocking. G threads cooperate on one query row, +// each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps +// occupancy high; the S dot product is reduced across the G-lane group with a +// short shuffle chain (log2(G) shuffles) instead of a full 32-lane warp reduce. +// Online (per-kv) softmax — cheap because acc[] is only HEAD_DIM/G long. +// Templated on . Block = (G, ROWS). G power-of-two, +// G*ROWS a multiple of 32 with groups warp-aligned. + +template +__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) { +#pragma unroll + for (int o = G / 2; o > 0; o >>= 1) + v += __shfl_xor_sync(mask, v, o); + return v; +} + +template +__global__ void gqa_prefill_attn_kernel_t(GQAParams p) { + constexpr int DPT = HEAD_DIM / G; + int q_tile = blockIdx.x; int q_head = blockIdx.y; int batch = blockIdx.z; - int q_row = q_tile * Br + threadIdx.y; - int d_part = threadIdx.x; - int dpw = p.head_dim >> 5; + int gpos = threadIdx.x; // 0..G-1 (which d-chunk) + int row = threadIdx.y; // 0..ROWS-1 + int q_row = q_tile * ROWS + row; int kv_head = q_head / (p.q_head / p.kv_head); - float qs[8] = {0}; - if (q_row < p.q_len) { - int q_off = (((batch * p.q_head + q_head) * p.q_len + q_row) * p.head_dim) + d_part * dpw; - for (int i = 0; i < dpw; i++) - qs[i] = __bfloat162float(p.q[q_off + i]) * p.scale; - } - - int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim; - extern __shared__ __align__(16) bf16 smem[]; bf16* sK = smem; - bf16* sV = smem + Bc * p.head_dim; + bf16* sV = sK + P_BC * HEAD_DIM; - float m = -FLT_MAX, l = 0.0f, acc[8] = {0}; + float qreg[DPT]; + if (q_row < p.q_len) { + int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT; +#pragma unroll + for (int i = 0; i < DPT; i++) + qreg[i] = __bfloat162float(p.q[q_off + i]) * p.scale; + } - int tiles = (p.kv_len + Bc - 1) / Bc; - int tt = blockDim.x * blockDim.y; + float m = -FLT_MAX, l = 0.0f; + float acc[DPT]; +#pragma unroll + for (int i = 0; i < DPT; i++) + acc[i] = 0.0f; + + int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM; + int tiles = (p.kv_len + P_BC - 1) / P_BC; + int tt = G * ROWS; + int lid = row * G + gpos; + + // per-group shuffle mask: only the G lanes of this row's group participate, + // so causal masking (differing loop bounds across rows in a warp) is safe. + int lane_in_warp = lid & 31; + unsigned gmask = (G == 32) ? 0xFFFFFFFFu + : (((1u << G) - 1u) << (lane_in_warp & ~(G - 1))); for (int ti = 0; ti < tiles; ti++) { - int kv0 = ti * Bc; - int tlen = min(Bc, p.kv_len - kv0); + int kv0 = ti * P_BC; + int tlen = min(P_BC, p.kv_len - kv0); - for (int i = threadIdx.y * blockDim.x + threadIdx.x; - i < tlen * p.head_dim; i += tt) { - int r = i / p.head_dim, c = i % p.head_dim, idx = r * p.head_dim + c; - int g_off = kv_base + (kv0 + r) * p.head_dim + c; - sK[idx] = p.k[g_off]; - sV[idx] = p.v[g_off]; + for (int i = lid; i < tlen * HEAD_DIM; i += tt) { + int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM); + sK[i] = p.k[gidx]; + sV[i] = p.v[gidx]; } __syncthreads(); @@ -52,30 +80,35 @@ __global__ void gqa_prefill_attn_kernel(GQAParams p) { } for (int s = 0; s < lim; s++) { - float dot = 0.0f; - for (int i = 0; i < dpw; i++) - dot += qs[i] * __bfloat162float(sK[s * p.head_dim + d_part * dpw + i]); - dot = warp_reduce_sum(dot); + const bf16* kr = sK + s * HEAD_DIM + gpos * DPT; + float part = 0.0f; +#pragma unroll + for (int i = 0; i < DPT; i++) + part = fmaf(qreg[i], __bfloat162float(kr[i]), part); + float dot = group_reduce_sum(part, gmask); if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s]) dot = -FLT_MAX; float nm = fmaxf(m, dot); - float al = expf(m - nm); - float be = expf(dot - nm); + float al = __expf(m - nm); + float be = __expf(dot - nm); l = l * al + be; - for (int i = 0; i < dpw; i++) - acc[i] = acc[i] * al + __bfloat162float(sV[s * p.head_dim + d_part * dpw + i]) * be; + const bf16* vr = sV + s * HEAD_DIM + gpos * DPT; +#pragma unroll + for (int i = 0; i < DPT; i++) + acc[i] = fmaf(__bfloat162float(vr[i]), be, acc[i] * al); m = nm; } __syncthreads(); } if (q_row < p.q_len) { - int o_off = (((batch * p.q_head + q_head) * p.q_len + q_row) * p.head_dim) + d_part * dpw; + int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT; float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f; - for (int i = 0; i < dpw; i++) +#pragma unroll + for (int i = 0; i < DPT; i++) p.o[o_off + i] = __float2bfloat16(acc[i] * rl); } } diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu index 27811f1..9b94e86 100644 --- a/csrc/tests/gqa_prefill_test.cu +++ b/csrc/tests/gqa_prefill_test.cu @@ -87,14 +87,20 @@ int main() { p.scale=1.0f/sqrtf((float)D); p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO; - dim3 grid((ql+Br-1)/Br, Hq, B); - dim3 block(32, Br, 1); - size_t smem=2*Bc*D*sizeof(bf16); + 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); double t0=now_ms(); - gqa_prefill_attn_kernel<<>>(p); + switch (D) { + case 32: gqa_prefill_attn_kernel_t<32, 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(); double kms=now_ms()-t0; cudaError_t err=cudaGetLastError();