feat: split kernel defs from bindings, add prefill tiled kernel and pure C tests
- Split .cuh/.cu for gqa_decode_attn and gqa_prefill_attn - gqa_prefill_attn: tiled shared-memory K/V, fused load, compute-opt, mask support - Add pure C tests under csrc/tests/ for fast nvcc-only iteration - Update .gitignore for build artifacts
This commit is contained in:
parent
579b8c3129
commit
bcdd93e0eb
|
|
@ -11,6 +11,7 @@
|
|||
!csrc/**/*.py
|
||||
|
||||
!csrc/**/*.cu
|
||||
!csrc/**/*.cuh
|
||||
|
||||
!scripts/**/*.sh
|
||||
|
||||
|
|
@ -32,4 +33,3 @@
|
|||
|
||||
# Allow build files
|
||||
!/setup.py
|
||||
!/AGENTS.md
|
||||
|
|
|
|||
|
|
@ -1,86 +1,8 @@
|
|||
// per-KV-head block, K shared in smem, each thread handles hd/32 elements
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <cfloat>
|
||||
// torch binding for gqa_decode_attn
|
||||
// kernel defined in gqa_decode_attn.cuh
|
||||
#include "gqa_decode_attn.cuh"
|
||||
#include <torch/extension.h>
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
constexpr int CHUNK = 64;
|
||||
|
||||
__inline__ __device__ float warp_reduce_sum(float val) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void gqa_decode_attn_kernel(
|
||||
const bf16* __restrict__ q_ptr,
|
||||
const bf16* __restrict__ k_ptr,
|
||||
const bf16* __restrict__ v_ptr,
|
||||
const bool* __restrict__ mask_ptr,
|
||||
bf16* __restrict__ out_ptr,
|
||||
int B, int n_heads, int n_kv_heads, int seq_len, int hd
|
||||
) {
|
||||
int batch = blockIdx.x / n_kv_heads;
|
||||
int kv_head = blockIdx.x % n_kv_heads;
|
||||
int group_size = blockDim.y;
|
||||
int q_head = kv_head * group_size + threadIdx.y;
|
||||
int lane = threadIdx.x;
|
||||
int hd_per_thread = hd / 32;
|
||||
|
||||
float q_reg[8];
|
||||
int q_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
q_reg[i] = __bfloat162float(q_ptr[q_off + i]);
|
||||
|
||||
int kv_base = ((batch * n_kv_heads + kv_head) * seq_len) * hd;
|
||||
int mask_base = batch * seq_len;
|
||||
|
||||
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
|
||||
float scale = rsqrtf((float)hd);
|
||||
|
||||
extern __shared__ __align__(16) bf16 k_smem[];
|
||||
|
||||
for (int chunk_start = 0; chunk_start < seq_len; chunk_start += CHUNK) {
|
||||
int this_chunk = min(CHUNK, seq_len - chunk_start);
|
||||
|
||||
int total = this_chunk * hd;
|
||||
for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
|
||||
k_smem[i] = k_ptr[kv_base + chunk_start * hd + i];
|
||||
__syncthreads();
|
||||
|
||||
for (int s = 0; s < this_chunk; s++) {
|
||||
float partial = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
partial += q_reg[i] * __bfloat162float(k_smem[s * hd + lane * hd_per_thread + i]);
|
||||
partial = warp_reduce_sum(partial) * scale;
|
||||
|
||||
if (!mask_ptr[mask_base + chunk_start + s]) partial = -FLT_MAX;
|
||||
|
||||
float new_m = fmaxf(m, partial);
|
||||
float alpha = expf(m - new_m);
|
||||
float beta = expf(partial - new_m);
|
||||
d = d * alpha + beta;
|
||||
|
||||
int v_off = kv_base + (chunk_start + s) * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta;
|
||||
m = new_m;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
int out_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d);
|
||||
}
|
||||
|
||||
torch::Tensor gqa_decode_attn(
|
||||
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor mask
|
||||
) {
|
||||
|
|
@ -97,7 +19,7 @@ torch::Tensor gqa_decode_attn(
|
|||
int group_size = n_heads / n_kv;
|
||||
auto out = torch::empty_like(q);
|
||||
|
||||
size_t smem = CHUNK * hd * sizeof(bf16); // K chunk
|
||||
size_t smem = DC_CHUNK * hd * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(B * n_kv);
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,78 @@
|
|||
// gqa_decode_attn.cuh — header-only decode kernel
|
||||
#pragma once
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cfloat>
|
||||
#include <algorithm>
|
||||
using std::min;
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
constexpr int DC_CHUNK = 64;
|
||||
|
||||
__device__ inline float warp_reduce_sum(float val) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void gqa_decode_attn_kernel(
|
||||
const bf16* __restrict__ q_ptr,
|
||||
const bf16* __restrict__ k_ptr,
|
||||
const bf16* __restrict__ v_ptr,
|
||||
const bool* __restrict__ mask_ptr,
|
||||
bf16* __restrict__ out_ptr,
|
||||
int B, int n_heads, int n_kv_heads, int seq_len, int hd
|
||||
) {
|
||||
int batch = blockIdx.x / n_kv_heads;
|
||||
int kv_head = blockIdx.x % n_kv_heads;
|
||||
int group_size = blockDim.y;
|
||||
int q_head = kv_head * group_size + threadIdx.y;
|
||||
int lane = threadIdx.x;
|
||||
int hd_per_thread = hd / 32;
|
||||
|
||||
float q_reg[8];
|
||||
int q_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
q_reg[i] = __bfloat162float(q_ptr[q_off + i]);
|
||||
|
||||
int kv_base = ((batch * n_kv_heads + kv_head) * seq_len) * hd;
|
||||
int mask_base = batch * seq_len;
|
||||
|
||||
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
|
||||
float scale = rsqrtf((float)hd);
|
||||
|
||||
extern __shared__ __align__(16) bf16 k_smem[];
|
||||
|
||||
for (int chunk_start = 0; chunk_start < seq_len; chunk_start += DC_CHUNK) {
|
||||
int this_chunk = min(DC_CHUNK, seq_len - chunk_start);
|
||||
|
||||
int total = this_chunk * hd;
|
||||
for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
|
||||
k_smem[i] = k_ptr[kv_base + chunk_start * hd + i];
|
||||
__syncthreads();
|
||||
|
||||
for (int s = 0; s < this_chunk; s++) {
|
||||
float partial = 0.0f;
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
partial += q_reg[i] * __bfloat162float(k_smem[s * hd + lane * hd_per_thread + i]);
|
||||
partial = warp_reduce_sum(partial) * scale;
|
||||
|
||||
if (!mask_ptr[mask_base + chunk_start + s]) partial = -FLT_MAX;
|
||||
|
||||
float new_m = fmaxf(m, partial);
|
||||
float alpha = expf(m - new_m);
|
||||
float beta = expf(partial - new_m);
|
||||
d = d * alpha + beta;
|
||||
|
||||
int v_off = kv_base + (chunk_start + s) * hd + lane * hd_per_thread;
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta;
|
||||
m = new_m;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
int out_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread;
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d);
|
||||
}
|
||||
|
|
@ -1,79 +1,8 @@
|
|||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <cfloat>
|
||||
// torch binding for gqa_prefill_attn
|
||||
// kernel defined in gqa_prefill_attn.cuh
|
||||
#include "gqa_prefill_attn.cuh"
|
||||
#include <torch/extension.h>
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
__inline__ __device__ float warp_reduce_sum(float val) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void gqa_prefill_attn_kernel(
|
||||
const bf16* __restrict__ q_ptr,
|
||||
const bf16* __restrict__ k_ptr,
|
||||
const bf16* __restrict__ v_ptr,
|
||||
const bool* __restrict__ mask_ptr,
|
||||
bf16* __restrict__ out_ptr,
|
||||
int B, int n_heads, int n_kv_heads, int q_len, int kv_len, int hd,
|
||||
int use_mask, int is_causal, int causal_offset
|
||||
) {
|
||||
int flat_id = blockIdx.x;
|
||||
int pos = flat_id % q_len;
|
||||
flat_id /= q_len;
|
||||
int q_head = flat_id % n_heads;
|
||||
int batch = flat_id / n_heads;
|
||||
int kv_head = q_head / (n_heads / n_kv_heads);
|
||||
int lane = threadIdx.x;
|
||||
int hd_per_thread = hd / 32;
|
||||
|
||||
// each thread handles hd/32 elements of Q
|
||||
float q_reg[8];
|
||||
int q_off = ((batch * n_heads + q_head) * q_len + pos) * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
q_reg[i] = __bfloat162float(q_ptr[q_off + i]);
|
||||
|
||||
int kv_base = ((batch * n_kv_heads + kv_head) * kv_len) * hd;
|
||||
int limit = is_causal ? min(pos + causal_offset + 1, kv_len) : kv_len;
|
||||
|
||||
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
|
||||
float scale = rsqrtf((float)hd);
|
||||
|
||||
int mask_stride = q_len * kv_len;
|
||||
int mask_off = batch * mask_stride + pos * kv_len;
|
||||
|
||||
for (int s = 0; s < limit; s++) {
|
||||
float partial = 0.0f;
|
||||
int k_off = kv_base + s * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
partial += q_reg[i] * __bfloat162float(k_ptr[k_off + i]);
|
||||
partial = warp_reduce_sum(partial) * scale;
|
||||
|
||||
if (use_mask && !mask_ptr[mask_off + s]) partial = -FLT_MAX;
|
||||
|
||||
float new_m = fmaxf(m, partial);
|
||||
float alpha = expf(m - new_m);
|
||||
float beta = expf(partial - new_m);
|
||||
d = d * alpha + beta;
|
||||
|
||||
int v_off = kv_base + s * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta;
|
||||
m = new_m;
|
||||
}
|
||||
|
||||
int out_off = ((batch * n_heads + q_head) * q_len + pos) * hd + lane * hd_per_thread;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < hd_per_thread; i++)
|
||||
out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d);
|
||||
}
|
||||
|
||||
torch::Tensor gqa_prefill_attn(
|
||||
torch::Tensor q, torch::Tensor k, torch::Tensor v,
|
||||
c10::optional<torch::Tensor> mask,
|
||||
|
|
@ -85,34 +14,34 @@ torch::Tensor gqa_prefill_attn(
|
|||
TORCH_CHECK(k.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(v.dtype() == torch::kBFloat16);
|
||||
|
||||
int B = q.size(0), n_heads = q.size(1), q_len = q.size(2), hd = q.size(3);
|
||||
int n_kv = k.size(1), kv_len = k.size(2);
|
||||
TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32");
|
||||
int B = q.size(0), Hq = q.size(1), q_len = q.size(2), D = q.size(3);
|
||||
int Hk = k.size(1), kv_len = k.size(2);
|
||||
TORCH_CHECK(D % 32 == 0, "head_dim must be multiple of 32");
|
||||
|
||||
bool use_mask = mask.has_value();
|
||||
const bool* mask_ptr = nullptr;
|
||||
if (use_mask) {
|
||||
TORCH_CHECK(mask.value().dtype() == torch::kBool);
|
||||
TORCH_CHECK(mask.value().dim() == 2);
|
||||
TORCH_CHECK(mask.value().size(0) == B);
|
||||
TORCH_CHECK(mask.value().size(1) == kv_len);
|
||||
mask_ptr = mask.value().data_ptr<bool>();
|
||||
}
|
||||
|
||||
auto out = torch::empty_like(q);
|
||||
auto O = torch::empty_like(q);
|
||||
|
||||
dim3 block(32);
|
||||
dim3 grid(B * n_heads * q_len);
|
||||
dim3 grid((q_len + Br - 1) / Br, Hq, B);
|
||||
dim3 block(32, Br, 1);
|
||||
size_t smem = 2 * Bc * D * sizeof(bf16);
|
||||
|
||||
gqa_prefill_attn_kernel<<<grid, block>>>(
|
||||
reinterpret_cast<const bf16*>(q.data_ptr()),
|
||||
reinterpret_cast<const bf16*>(k.data_ptr()),
|
||||
reinterpret_cast<const bf16*>(v.data_ptr()),
|
||||
mask_ptr,
|
||||
reinterpret_cast<bf16*>(out.data_ptr()),
|
||||
B, n_heads, n_kv, q_len, kv_len, hd,
|
||||
(int)use_mask, (int)is_causal, (int)causal_offset
|
||||
gqa_prefill_attn_kernel<<<grid, block, smem>>>(
|
||||
(const bf16*)q.data_ptr(), (const bf16*)k.data_ptr(),
|
||||
(const bf16*)v.data_ptr(), mask_ptr, (bf16*)O.data_ptr(),
|
||||
B, Hq, Hk, q_len, kv_len, D, (int)is_causal, (int)causal_offset, (int)use_mask
|
||||
);
|
||||
return out;
|
||||
return O;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill attention (naive)");
|
||||
m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill v3 (compute-opt)");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,100 @@
|
|||
// gqa_prefill_attn.cuh — header-only kernel definition (no torch dependency)
|
||||
#pragma once
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cfloat>
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
static constexpr int Br = 32;
|
||||
static constexpr int Bc = 64;
|
||||
|
||||
__device__ inline float warp_sum(float v) {
|
||||
for (int off = 16; off > 0; off >>= 1)
|
||||
v += __shfl_xor_sync(0xffffffff, v, off);
|
||||
return v;
|
||||
}
|
||||
|
||||
__global__ void gqa_prefill_attn_kernel(
|
||||
const bf16* __restrict__ Q, const bf16* __restrict__ K,
|
||||
const bf16* __restrict__ V, const bool* __restrict__ mask,
|
||||
bf16* __restrict__ O,
|
||||
int B, int Hq, int Hk, int q_len, int kv_len, int D,
|
||||
int is_causal, int causal_offset, int use_mask
|
||||
) {
|
||||
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 = D >> 5;
|
||||
|
||||
int kv_head = q_head / (Hq / Hk);
|
||||
|
||||
float qs[8] = {0};
|
||||
if (q_row < q_len) {
|
||||
float sc = rsqrtf((float)D);
|
||||
int q_off = (((batch * Hq + q_head) * q_len + q_row) * D) + d_part * dpw;
|
||||
for (int i = 0; i < dpw; i++)
|
||||
qs[i] = __bfloat162float(Q[q_off + i]) * sc;
|
||||
}
|
||||
|
||||
int kv_base = ((batch * Hk + kv_head) * kv_len) * D;
|
||||
|
||||
extern __shared__ __align__(16) bf16 smem[];
|
||||
bf16* sK = smem;
|
||||
bf16* sV = smem + Bc * D;
|
||||
|
||||
float m = -FLT_MAX, l = 0.0f, acc[8] = {0};
|
||||
|
||||
int tiles = (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, kv_len - kv0);
|
||||
|
||||
for (int i = threadIdx.y * blockDim.x + threadIdx.x;
|
||||
i < tlen * D; i += tt) {
|
||||
int r = i / D, c = i % D, idx = r * D + c;
|
||||
int g_off = kv_base + (kv0 + r) * D + c;
|
||||
sK[idx] = K[g_off];
|
||||
sV[idx] = V[g_off];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int lim = tlen;
|
||||
if (is_causal && q_row < q_len) {
|
||||
int ep = q_row + 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 * D + d_part * dpw + i]);
|
||||
dot = warp_sum(dot);
|
||||
|
||||
if (use_mask && !mask[batch * q_len * kv_len + q_row * 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 * D + d_part * dpw + i]) * be;
|
||||
m = nm;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (q_row < q_len) {
|
||||
int o_off = (((batch * Hq + q_head) * q_len + q_row) * D) + d_part * dpw;
|
||||
float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
|
||||
for (int i = 0; i < dpw; i++)
|
||||
O[o_off + i] = __float2bfloat16(acc[i] * rl);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,120 @@
|
|||
// Pure-C test for decode kernel
|
||||
// compile: nvcc -I csrc csrc/tests/gqa_decode_test.cu -o test && ./test
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <sys/time.h>
|
||||
#include "../kernels/gqa_decode_attn.cuh"
|
||||
|
||||
static double now_ms() {
|
||||
struct timeval tv;
|
||||
gettimeofday(&tv, NULL);
|
||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
||||
}
|
||||
|
||||
static void cpu_decode(const float* Q, const float* K, const float* V,
|
||||
const bool* mask, float* O,
|
||||
int B, int Hq, int Hk, int seq_len, int D) {
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
int n_rep = Hq / Hk;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int h = 0; h < Hq; h++) {
|
||||
int kv_h = h / n_rep;
|
||||
float mv = -INFINITY, sv = 0.0f;
|
||||
float accum[256] = {0};
|
||||
for (int s = 0; s < seq_len; s++) {
|
||||
if (!mask[b * seq_len + s]) continue;
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < D; d++)
|
||||
dot += Q[((b * Hq + h) * 1 + 0) * D + d]
|
||||
* K[((b * Hk + kv_h) * seq_len + s) * D + d];
|
||||
dot *= scale;
|
||||
float nm = fmaxf(mv, dot);
|
||||
float al = expf(mv - nm);
|
||||
float be = expf(dot - nm);
|
||||
sv = sv * al + be;
|
||||
for (int d = 0; d < D; d++)
|
||||
accum[d] = accum[d] * al
|
||||
+ V[((b * Hk + kv_h) * seq_len + s) * D + d] * be;
|
||||
mv = nm;
|
||||
}
|
||||
float inv = 1.0f / sv;
|
||||
for (int d = 0; d < D; d++)
|
||||
O[((b * Hq + h) * 1 + 0) * D + d] = accum[d] * inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static bf16 f2bf(float x) { return __float2bfloat16(x); }
|
||||
static float bf2f(bf16 x) { return __bfloat162float(x); }
|
||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
||||
|
||||
int main() {
|
||||
const int configs[][5] = {
|
||||
{1, 2, 1, 64, 32}, // B,Hq,Hk,seq_len,D
|
||||
{1, 32, 4, 512, 128},
|
||||
{1, 32, 4, 1024, 128},
|
||||
};
|
||||
int n_cfgs = sizeof(configs) / sizeof(configs[0]);
|
||||
|
||||
for (int ci = 0; ci < n_cfgs; ci++) {
|
||||
int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
|
||||
int sl = configs[ci][3], D = configs[ci][4], gs = Hq / Hk;
|
||||
printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d ===\n", B,Hq,Hk,sl,D,gs);
|
||||
|
||||
size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
|
||||
float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
|
||||
for (size_t i=0;i<nQ;i++) hQ[i]=randf();
|
||||
for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
|
||||
|
||||
bool* hMask=new bool[B*sl];
|
||||
for (int i=0;i<B*sl;i++) hMask[i]=true;
|
||||
|
||||
bf16 *dQ,*dK,*dV,*dO,*tmp;
|
||||
bool* dMask;
|
||||
cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
|
||||
cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
|
||||
cudaMalloc(&dMask,B*sl);
|
||||
|
||||
tmp=new bf16[max(nQ,nKV)];
|
||||
for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
|
||||
cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
|
||||
cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
|
||||
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
|
||||
|
||||
size_t smem=DC_CHUNK*D*sizeof(bf16);
|
||||
dim3 block(32, gs);
|
||||
dim3 grid(B*Hk);
|
||||
printf("grid=(%d,1,1) block=(%d,%d,1) smem=%zu\n",
|
||||
grid.x, block.x, block.y, smem);
|
||||
|
||||
double t0=now_ms();
|
||||
gqa_decode_attn_kernel<<<grid,block,smem>>>(dQ,dK,dV,dMask,dO,
|
||||
B,Hq,Hk,sl,D);
|
||||
cudaDeviceSynchronize();
|
||||
double kms=now_ms()-t0;
|
||||
cudaError_t err=cudaGetLastError();
|
||||
if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
|
||||
|
||||
bf16* hOut=new bf16[nQ];
|
||||
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
|
||||
|
||||
float* ref=new float[nQ];
|
||||
cpu_decode(hQ,hK,hV,hMask,ref,B,Hq,Hk,sl,D);
|
||||
|
||||
float max_err=0;
|
||||
for (size_t i=0;i<nQ;i++){
|
||||
float d=fabsf(bf2f(hOut[i])-ref[i]);
|
||||
if(d>max_err) max_err=d;
|
||||
}
|
||||
printf("kernel: %.3f ms max_err: %.6e\n\n",kms,max_err);
|
||||
|
||||
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
|
||||
delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
|
||||
}
|
||||
printf("All tests passed!\n");
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
// Pure-C test: compile with nvcc -I csrc csrc/tests/gqa_prefill_test.cu -o test && ./test
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <sys/time.h>
|
||||
#include "../kernels/gqa_prefill_attn.cuh"
|
||||
|
||||
static double now_ms() {
|
||||
struct timeval tv;
|
||||
gettimeofday(&tv, NULL);
|
||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
||||
}
|
||||
|
||||
static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
|
||||
int B, int Hq, int Hk, int q_len, int kv_len, int D,
|
||||
int is_causal, int causal_off) {
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
int n_rep = Hq / Hk;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int h = 0; h < Hq; h++) {
|
||||
for (int qi = 0; qi < q_len; qi++) {
|
||||
int kv_h = h / n_rep;
|
||||
float mv = -INFINITY, sv = 0.0f;
|
||||
float accum[256] = {0};
|
||||
int lim = is_causal ? min(kv_len, qi + causal_off + 1) : kv_len;
|
||||
for (int kj = 0; kj < lim; kj++) {
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < D; d++)
|
||||
dot += Q[((b*Hq + h)*q_len + qi)*D + d]
|
||||
* K[((b*Hk + kv_h)*kv_len + kj)*D + d];
|
||||
dot *= scale;
|
||||
float nm = fmaxf(mv, dot);
|
||||
float al = expf(mv - nm);
|
||||
float be = expf(dot - nm);
|
||||
sv = sv * al + be;
|
||||
for (int d = 0; d < D; d++)
|
||||
accum[d] = accum[d] * al
|
||||
+ V[((b*Hk + kv_h)*kv_len + kj)*D + d] * be;
|
||||
mv = nm;
|
||||
}
|
||||
float inv = 1.0f / sv;
|
||||
for (int d = 0; d < D; d++)
|
||||
O[((b*Hq + h)*q_len + qi)*D + d] = accum[d] * inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
|
||||
static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
|
||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
||||
|
||||
int main() {
|
||||
const int configs[][7] = {
|
||||
{1,2,1,64,128,32,0}, // tiny: B,Hq,Hk,q,kv,D,causal
|
||||
{1,32,4,512,512,128,0}, // standard
|
||||
{1,32,4,128,256,128,0}, // medium
|
||||
};
|
||||
int n_configs = sizeof(configs) / sizeof(configs[0]);
|
||||
|
||||
for (int ci = 0; ci < n_configs; ci++) {
|
||||
int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
|
||||
int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
|
||||
int causal=configs[ci][6];
|
||||
printf("=== B=%d Hq=%d Hk=%d q=%d kv=%d D=%d causal=%d ===\n",
|
||||
B,Hq,Hk,ql,kl,D,causal);
|
||||
|
||||
size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*D;
|
||||
float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
|
||||
for (size_t i=0;i<nQ;i++) hQ[i]=randf();
|
||||
for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
|
||||
|
||||
bf16 *dQ,*dK,*dV,*dO,*tmp;
|
||||
cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
|
||||
cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
|
||||
tmp=new bf16[max(nQ,nKV)];
|
||||
for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
|
||||
cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
|
||||
cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
|
||||
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
|
||||
dim3 grid((ql+Br-1)/Br, Hq, B);
|
||||
dim3 block(32, Br, 1);
|
||||
size_t smem=2*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>>>(dQ,dK,dV,nullptr,dO,
|
||||
B,Hq,Hk,ql,kl,D,causal,0,0);
|
||||
cudaDeviceSynchronize();
|
||||
double kms=now_ms()-t0;
|
||||
cudaError_t err=cudaGetLastError();
|
||||
if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
|
||||
|
||||
bf16* hOut=new bf16[nQ];
|
||||
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
|
||||
|
||||
float* ref=new float[nQ];
|
||||
cpu_attention(hQ,hK,hV,ref,B,Hq,Hk,ql,kl,D,causal,0);
|
||||
|
||||
float max_err=0;
|
||||
for (size_t i=0;i<nQ;i++) {
|
||||
float d=fabsf(bf2f(hOut[i])-ref[i]);
|
||||
if(d>max_err) max_err=d;
|
||||
}
|
||||
printf("kernel: %.3f ms max_err: %.6e\n\n",kms,max_err);
|
||||
|
||||
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
|
||||
delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
|
||||
}
|
||||
printf("All tests passed!\n");
|
||||
return 0;
|
||||
}
|
||||
Loading…
Reference in New Issue