82 lines
2.6 KiB
Plaintext
82 lines
2.6 KiB
Plaintext
#pragma once
|
|
#include "gqa_common.cuh"
|
|
|
|
__global__ void gqa_prefill_attn_kernel(GQAParams p) {
|
|
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 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;
|
|
|
|
float m = -FLT_MAX, l = 0.0f, acc[8] = {0};
|
|
|
|
int tiles = (p.kv_len + Bc - 1) / Bc;
|
|
int tt = blockDim.x * blockDim.y;
|
|
|
|
for (int ti = 0; ti < tiles; ti++) {
|
|
int kv0 = ti * Bc;
|
|
int tlen = min(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];
|
|
}
|
|
__syncthreads();
|
|
|
|
int lim = tlen;
|
|
if (p.is_causal && q_row < p.q_len) {
|
|
int ep = q_row + p.causal_offset + 1;
|
|
if (kv0 >= ep)
|
|
lim = 0;
|
|
else if (kv0 + tlen > ep)
|
|
lim = ep - kv0;
|
|
}
|
|
|
|
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);
|
|
|
|
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);
|
|
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;
|
|
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;
|
|
float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
|
|
for (int i = 0; i < dpw; i++)
|
|
p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
|
|
}
|
|
}
|