perf: group-split register-blocking gqa_prefill kernel
- one query row per group of G=8 lanes, each owning HEAD_DIM/G dims of qreg[]/acc[] in registers - removes full 32-lane warp_reduce_sum; S dot reduces over only G lanes - templated on <HEAD_DIM,G,ROWS,P_BC>, block=(G,ROWS)=(8,32) - per-group shuffle mask so causal loop-bound divergence doesn't deadlock the shuffle - update pure-C test to the templated launch
This commit is contained in:
parent
11fa807cfc
commit
cc36530c73
|
|
@ -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<<<grid, block, smem>>>(p);
|
||||
switch (p.head_dim) {
|
||||
case 64:
|
||||
gqa_prefill_attn_kernel_t<64, G, ROWS, P_BC><<<grid, block, smem>>>(p);
|
||||
break;
|
||||
case 128:
|
||||
gqa_prefill_attn_kernel_t<128, G, ROWS, P_BC><<<grid, block, smem>>>(p);
|
||||
break;
|
||||
case 256:
|
||||
gqa_prefill_attn_kernel_t<256, G, ROWS, P_BC><<<grid, block, smem>>>(p);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
|
||||
" (supported: 64,128,256)");
|
||||
}
|
||||
return O;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <HEAD_DIM, G, ROWS, P_BC>. Block = (G, ROWS). G power-of-two,
|
||||
// G*ROWS a multiple of 32 with groups warp-aligned.
|
||||
|
||||
template <int G>
|
||||
__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 <int HEAD_DIM, int G, int ROWS, int P_BC>
|
||||
__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<G>(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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<<<grid,block,smem>>>(p);
|
||||
switch (D) {
|
||||
case 32: gqa_prefill_attn_kernel_t<32, G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
|
||||
case 64: gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
|
||||
case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
|
||||
default: printf("unsupported D=%d\n",D); return 1;
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
double kms=now_ms()-t0;
|
||||
cudaError_t err=cudaGetLastError();
|
||||
|
|
|
|||
Loading…
Reference in New Issue