diff --git a/.gitignore b/.gitignore index 309443d..37ca78a 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ !csrc/**/*.py !csrc/**/*.cu +!csrc/**/*.cuh !scripts/**/*.sh @@ -32,4 +33,3 @@ # Allow build files !/setup.py -!/AGENTS.md diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index 2318b24..0e4f35b 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -1,86 +1,8 @@ -// per-KV-head block, K shared in smem, each thread handles hd/32 elements -#include -#include -#include -#include +// torch binding for gqa_decode_attn +// kernel defined in gqa_decode_attn.cuh +#include "gqa_decode_attn.cuh" #include -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); diff --git a/csrc/kernels/gqa_decode_attn.cuh b/csrc/kernels/gqa_decode_attn.cuh new file mode 100644 index 0000000..7119fa9 --- /dev/null +++ b/csrc/kernels/gqa_decode_attn.cuh @@ -0,0 +1,78 @@ +// gqa_decode_attn.cuh — header-only decode kernel +#pragma once +#include +#include +#include +#include +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); +} diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index 11c0064..25187be 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -1,79 +1,8 @@ -#include -#include -#include -#include +// torch binding for gqa_prefill_attn +// kernel defined in gqa_prefill_attn.cuh +#include "gqa_prefill_attn.cuh" #include -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 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(); } - 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<<>>( - reinterpret_cast(q.data_ptr()), - reinterpret_cast(k.data_ptr()), - reinterpret_cast(v.data_ptr()), - mask_ptr, - reinterpret_cast(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<<>>( + (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)"); } diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/gqa_prefill_attn.cuh new file mode 100644 index 0000000..9523470 --- /dev/null +++ b/csrc/kernels/gqa_prefill_attn.cuh @@ -0,0 +1,100 @@ +// gqa_prefill_attn.cuh — header-only kernel definition (no torch dependency) +#pragma once +#include +#include +#include + +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); + } +} diff --git a/csrc/tests/gqa_decode_test.cu b/csrc/tests/gqa_decode_test.cu new file mode 100644 index 0000000..47aa95a --- /dev/null +++ b/csrc/tests/gqa_decode_test.cu @@ -0,0 +1,120 @@ +// Pure-C test for decode kernel +// compile: nvcc -I csrc csrc/tests/gqa_decode_test.cu -o test && ./test +#include +#include +#include +#include +#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>>(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;imax_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; +} diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu new file mode 100644 index 0000000..c98a212 --- /dev/null +++ b/csrc/tests/gqa_prefill_test.cu @@ -0,0 +1,116 @@ +// Pure-C test: compile with nvcc -I csrc csrc/tests/gqa_prefill_test.cu -o test && ./test +#include +#include +#include +#include +#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>>(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;imax_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; +}