ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: correct prefill mask index, unify GQA kernel interface 

- Fix mask indexing: batch*q_len*kv_len -> batch*kv_len
- Add csrc/kernels/gqa_common.cuh with shared GQAParams struct
- Unify decode/prefill Python API: both accept (q,k,v,mask=None,...)
- Decode now supports optional mask, is_causal, causal_offset, scale
- Rename struct fields: B->batch, Hq->q_head, Hk->kv_head, D->head_dim
- Use py::arg() for correct None/defaults handling in pybind11
- Update pure C tests and build instructions (-arch=sm_89)
This commit is contained in:
ViperEkura 2026-07-06 16:52:19 +08:00
parent bcdd93e0eb
commit 11fa807cfc
7 changed files with 191 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
csrc/kernels/gqa_common.cuh Normal file
View File

@ -0,0 +1,35 @@
#pragma once
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <algorithm>
using bf16 = __nv_bfloat16;
using std::min;
constexpr int DC_CHUNK = 64;
constexpr int Br = 32, Bc = 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;
}
struct GQAParams {
int batch;
int q_head;
int kv_head;
int q_len;
int kv_len;
int head_dim;
int use_mask;
int is_causal;
int causal_offset;
float scale;
const bf16* __restrict__ q;
const bf16* __restrict__ k;
const bf16* __restrict__ v;
const bool* __restrict__ mask;
bf16* __restrict__ o;
};

71
csrc/kernels/gqa_decode_attn.cu
View File

@ -1,39 +1,66 @@
// torch binding for gqa_decode_attn
// kernel defined in gqa_decode_attn.cuh
#include "gqa_decode_attn.cuh" #include "gqa_decode_attn.cuh"
#include <torch/extension.h> #include <torch/extension.h>
torch::Tensor gqa_decode_attn( torch::Tensor gqa_decode_attn(
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor mask torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
c10::optional<torch::Tensor> mask,
bool is_causal = false,
int64_t causal_offset = 0,
c10::optional<double> scale = c10::nullopt
) { ) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && mask.is_cuda()); TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
TORCH_CHECK(q.dtype() == torch::kBFloat16); TORCH_CHECK(q.dtype() == torch::kBFloat16);
TORCH_CHECK(k.dtype() == torch::kBFloat16); TORCH_CHECK(k.dtype() == torch::kBFloat16);
TORCH_CHECK(v.dtype() == torch::kBFloat16); TORCH_CHECK(v.dtype() == torch::kBFloat16);
TORCH_CHECK(mask.dtype() == torch::kBool);
TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1"); TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1");
int B = q.size(0), n_heads = q.size(1), n_kv = k.size(1); GQAParams p;
int seq_len = k.size(2), hd = q.size(3); p.batch = q.size(0);
TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32"); p.q_head = q.size(1);
int group_size = n_heads / n_kv; p.kv_head = k.size(1);
auto out = torch::empty_like(q); p.q_len = 1;
p.kv_len = k.size(2);
p.head_dim = q.size(3);
TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
p.use_mask = mask.has_value();
p.is_causal = (int)is_causal;
p.causal_offset = (int)causal_offset;
p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
p.q = (const bf16*)q.data_ptr();
p.k = (const bf16*)k.data_ptr();
p.v = (const bf16*)v.data_ptr();
if (p.use_mask) {
TORCH_CHECK(mask.value().dtype() == torch::kBool);
TORCH_CHECK(mask.value().dim() == 2);
TORCH_CHECK(mask.value().size(0) == p.batch);
TORCH_CHECK(mask.value().size(1) == p.kv_len);
p.mask = mask.value().data_ptr<bool>();
} else {
p.mask = nullptr;
}
size_t smem = DC_CHUNK * hd * sizeof(bf16); auto O = torch::empty_like(q);
p.o = (bf16*)O.data_ptr();
int group_size = p.q_head / p.kv_head;
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
dim3 block(32, group_size); dim3 block(32, group_size);
dim3 grid(B * n_kv); dim3 grid(p.batch * p.kv_head);
gqa_decode_attn_kernel<<<grid, block, smem>>>( gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
reinterpret_cast<const bf16*>(q.data_ptr()), return O;
reinterpret_cast<const bf16*>(k.data_ptr()),
reinterpret_cast<const bf16*>(v.data_ptr()),
mask.data_ptr<bool>(),
reinterpret_cast<bf16*>(out.data_ptr()),
B, n_heads, n_kv, seq_len, hd
);
return out;
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode v2 (per-KV-head, shared K)"); m.def("gqa_decode_attn", &gqa_decode_attn,
py::arg("q"),
py::arg("k"),
py::arg("v"),
py::arg("mask") = py::none(),
py::arg("is_causal") = false,
py::arg("causal_offset") = 0,
py::arg("scale") = py::none(),
"GQA decode (per-KV-head, shared K)");
} }

65
csrc/kernels/gqa_decode_attn.cuh
View File

@ -1,78 +1,59 @@
// gqa_decode_attn.cuh — header-only decode kernel
#pragma once #pragma once
#include <cuda_bf16.h> #include "gqa_common.cuh"
#include <cuda_runtime.h>
#include <cfloat>
#include <algorithm>
using std::min;
using bf16 = __nv_bfloat16;
constexpr int DC_CHUNK = 64; __global__ void gqa_decode_attn_kernel(GQAParams p) {
int batch = blockIdx.x / p.kv_head;
__device__ inline float warp_reduce_sum(float val) { int kv_head = blockIdx.x % p.kv_head;
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 group_size = blockDim.y;
int q_head = kv_head * group_size + threadIdx.y; int q_head = kv_head * group_size + threadIdx.y;
int lane = threadIdx.x; int lane = threadIdx.x;
int hd_per_thread = hd / 32; int hd_per_thread = p.head_dim / 32;
float q_reg[8]; float q_reg[8];
int q_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread; int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
for (int i = 0; i < hd_per_thread; i++) for (int i = 0; i < hd_per_thread; i++)
q_reg[i] = __bfloat162float(q_ptr[q_off + i]); q_reg[i] = __bfloat162float(p.q[q_off + i]);
int kv_base = ((batch * n_kv_heads + kv_head) * seq_len) * hd; int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
int mask_base = batch * seq_len; int mask_base = batch * p.kv_len;
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f}; float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
float scale = rsqrtf((float)hd);
extern __shared__ __align__(16) bf16 k_smem[]; extern __shared__ __align__(16) bf16 k_smem[];
for (int chunk_start = 0; chunk_start < seq_len; chunk_start += DC_CHUNK) { for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) {
int this_chunk = min(DC_CHUNK, seq_len - chunk_start); int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
int total = this_chunk * hd; int total = this_chunk * p.head_dim;
for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) 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]; k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i];
__syncthreads(); __syncthreads();
for (int s = 0; s < this_chunk; s++) { for (int s = 0; s < this_chunk; s++) {
float partial = 0.0f; float partial = 0.0f;
for (int i = 0; i < hd_per_thread; i++) for (int i = 0; i < hd_per_thread; i++)
partial += q_reg[i] * __bfloat162float(k_smem[s * hd + lane * hd_per_thread + i]); partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
partial = warp_reduce_sum(partial) * scale; partial = warp_reduce_sum(partial) * p.scale;
if (!mask_ptr[mask_base + chunk_start + s]) partial = -FLT_MAX; if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
partial = -FLT_MAX;
if (p.is_causal && (chunk_start + s) > p.causal_offset)
partial = -FLT_MAX;
float new_m = fmaxf(m, partial); float new_m = fmaxf(m, partial);
float alpha = expf(m - new_m); float alpha = expf(m - new_m);
float beta = expf(partial - new_m); float beta = expf(partial - new_m);
d = d * alpha + beta; d = d * alpha + beta;
int v_off = kv_base + (chunk_start + s) * hd + lane * hd_per_thread; int v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread;
for (int i = 0; i < hd_per_thread; i++) for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(v_ptr[v_off + i]) * beta; acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta;
m = new_m; m = new_m;
} }
__syncthreads(); __syncthreads();
} }
int out_off = ((batch * n_heads + q_head) * 1) * hd + lane * hd_per_thread; int out_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
for (int i = 0; i < hd_per_thread; i++) for (int i = 0; i < hd_per_thread; i++)
out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d); p.o[out_off + i] = __float2bfloat16(acc_reg[i] / d);
} }

61
csrc/kernels/gqa_prefill_attn.cu
View File

@ -1,12 +1,13 @@
// torch binding for gqa_prefill_attn
// kernel defined in gqa_prefill_attn.cuh
#include "gqa_prefill_attn.cuh" #include "gqa_prefill_attn.cuh"
#include <torch/extension.h> #include <torch/extension.h>
torch::Tensor gqa_prefill_attn( torch::Tensor gqa_prefill_attn(
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
c10::optional<torch::Tensor> mask, c10::optional<torch::Tensor> mask,
bool is_causal = false, int64_t causal_offset = 0, bool is_causal = false,
int64_t causal_offset = 0,
c10::optional<double> scale = c10::nullopt c10::optional<double> scale = c10::nullopt
) { ) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda()); TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
@ -14,34 +15,50 @@ torch::Tensor gqa_prefill_attn(
TORCH_CHECK(k.dtype() == torch::kBFloat16); TORCH_CHECK(k.dtype() == torch::kBFloat16);
TORCH_CHECK(v.dtype() == torch::kBFloat16); TORCH_CHECK(v.dtype() == torch::kBFloat16);
int B = q.size(0), Hq = q.size(1), q_len = q.size(2), D = q.size(3); GQAParams p;
int Hk = k.size(1), kv_len = k.size(2); p.batch = q.size(0);
TORCH_CHECK(D % 32 == 0, "head_dim must be multiple of 32"); p.q_head = q.size(1);
p.kv_head = k.size(1);
bool use_mask = mask.has_value(); p.q_len = q.size(2);
const bool* mask_ptr = nullptr; p.kv_len = k.size(2);
if (use_mask) { p.head_dim = q.size(3);
TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
p.use_mask = mask.has_value();
p.is_causal = (int)is_causal;
p.causal_offset = (int)causal_offset;
p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
p.q = (const bf16*)q.data_ptr();
p.k = (const bf16*)k.data_ptr();
p.v = (const bf16*)v.data_ptr();
if (p.use_mask) {
TORCH_CHECK(mask.value().dtype() == torch::kBool); TORCH_CHECK(mask.value().dtype() == torch::kBool);
TORCH_CHECK(mask.value().dim() == 2); TORCH_CHECK(mask.value().dim() == 2);
TORCH_CHECK(mask.value().size(0) == B); TORCH_CHECK(mask.value().size(0) == p.batch);
TORCH_CHECK(mask.value().size(1) == kv_len); TORCH_CHECK(mask.value().size(1) == p.kv_len);
mask_ptr = mask.value().data_ptr<bool>(); p.mask = mask.value().data_ptr<bool>();
} else {
p.mask = nullptr;
} }
auto O = torch::empty_like(q); auto O = torch::empty_like(q);
p.o = (bf16*)O.data_ptr();
dim3 grid((q_len + Br - 1) / Br, Hq, B); dim3 grid((p.q_len + Br - 1) / Br, p.q_head, p.batch);
dim3 block(32, Br, 1); dim3 block(32, Br, 1);
size_t smem = 2 * Bc * D * sizeof(bf16); size_t smem = 2 * Bc * p.head_dim * sizeof(bf16);
gqa_prefill_attn_kernel<<<grid, block, smem>>>( gqa_prefill_attn_kernel<<<grid, block, smem>>>(p);
(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 O; return O;
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill v3 (compute-opt)"); m.def("gqa_prefill_attn", &gqa_prefill_attn,
py::arg("q"),
py::arg("k"),
py::arg("v"),
py::arg("mask") = py::none(),
py::arg("is_causal") = false,
py::arg("causal_offset") = 0,
py::arg("scale") = py::none(),
"GQA prefill (tiled, K+V smem)");
} }

77
csrc/kernels/gqa_prefill_attn.cuh
View File

@ -1,82 +1,63 @@
// gqa_prefill_attn.cuh — header-only kernel definition (no torch dependency)
#pragma once #pragma once
#include <cuda_bf16.h> #include "gqa_common.cuh"
#include <cuda_runtime.h>
#include <cfloat>
using bf16 = __nv_bfloat16; __global__ void gqa_prefill_attn_kernel(GQAParams p) {
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_tile = blockIdx.x;
int q_head = blockIdx.y; int q_head = blockIdx.y;
int batch = blockIdx.z; int batch = blockIdx.z;
int q_row = q_tile * Br + threadIdx.y; int q_row = q_tile * Br + threadIdx.y;
int d_part = threadIdx.x; int d_part = threadIdx.x;
int dpw = D >> 5; int dpw = p.head_dim >> 5;
int kv_head = q_head / (Hq / Hk); int kv_head = q_head / (p.q_head / p.kv_head);
float qs[8] = {0}; float qs[8] = {0};
if (q_row < q_len) { if (q_row < p.q_len) {
float sc = rsqrtf((float)D); int q_off = (((batch * p.q_head + q_head) * p.q_len + q_row) * p.head_dim) + d_part * dpw;
int q_off = (((batch * Hq + q_head) * q_len + q_row) * D) + d_part * dpw;
for (int i = 0; i < dpw; i++) for (int i = 0; i < dpw; i++)
qs[i] = __bfloat162float(Q[q_off + i]) * sc; qs[i] = __bfloat162float(p.q[q_off + i]) * p.scale;
} }
int kv_base = ((batch * Hk + kv_head) * kv_len) * D; int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
extern __shared__ __align__(16) bf16 smem[]; extern __shared__ __align__(16) bf16 smem[];
bf16* sK = smem; bf16* sK = smem;
bf16* sV = smem + Bc * D; bf16* sV = smem + Bc * p.head_dim;
float m = -FLT_MAX, l = 0.0f, acc[8] = {0}; float m = -FLT_MAX, l = 0.0f, acc[8] = {0};
int tiles = (kv_len + Bc - 1) / Bc; int tiles = (p.kv_len + Bc - 1) / Bc;
int tt = blockDim.x * blockDim.y; int tt = blockDim.x * blockDim.y;
for (int ti = 0; ti < tiles; ti++) { for (int ti = 0; ti < tiles; ti++) {
int kv0 = ti * Bc; int kv0 = ti * Bc;
int tlen = min(Bc, kv_len - kv0); int tlen = min(Bc, p.kv_len - kv0);
for (int i = threadIdx.y * blockDim.x + threadIdx.x; for (int i = threadIdx.y * blockDim.x + threadIdx.x;
i < tlen * D; i += tt) { i < tlen * p.head_dim; i += tt) {
int r = i / D, c = i % D, idx = r * D + c; int r = i / p.head_dim, c = i % p.head_dim, idx = r * p.head_dim + c;
int g_off = kv_base + (kv0 + r) * D + c; int g_off = kv_base + (kv0 + r) * p.head_dim + c;
sK[idx] = K[g_off]; sK[idx] = p.k[g_off];
sV[idx] = V[g_off]; sV[idx] = p.v[g_off];
} }
__syncthreads(); __syncthreads();
int lim = tlen; int lim = tlen;
if (is_causal && q_row < q_len) { if (p.is_causal && q_row < p.q_len) {
int ep = q_row + causal_offset + 1; int ep = q_row + p.causal_offset + 1;
if (kv0 >= ep) lim = 0; if (kv0 >= ep)
else if (kv0 + tlen > ep) lim = ep - kv0; lim = 0;
else if (kv0 + tlen > ep)
lim = ep - kv0;
} }
for (int s = 0; s < lim; s++) { for (int s = 0; s < lim; s++) {
float dot = 0.0f; float dot = 0.0f;
for (int i = 0; i < dpw; i++) for (int i = 0; i < dpw; i++)
dot += qs[i] * __bfloat162float(sK[s * D + d_part * dpw + i]); dot += qs[i] * __bfloat162float(sK[s * p.head_dim + d_part * dpw + i]);
dot = warp_sum(dot); dot = warp_reduce_sum(dot);
if (use_mask && !mask[batch * q_len * kv_len + q_row * kv_len + kv0 + s]) if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s])
dot = -FLT_MAX; dot = -FLT_MAX;
float nm = fmaxf(m, dot); float nm = fmaxf(m, dot);
@ -85,16 +66,16 @@ __global__ void gqa_prefill_attn_kernel(
l = l * al + be; l = l * al + be;
for (int i = 0; i < dpw; i++) for (int i = 0; i < dpw; i++)
acc[i] = acc[i] * al + __bfloat162float(sV[s * D + d_part * dpw + i]) * be; acc[i] = acc[i] * al + __bfloat162float(sV[s * p.head_dim + d_part * dpw + i]) * be;
m = nm; m = nm;
} }
__syncthreads(); __syncthreads();
} }
if (q_row < q_len) { if (q_row < p.q_len) {
int o_off = (((batch * Hq + q_head) * q_len + q_row) * D) + d_part * dpw; 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; float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
for (int i = 0; i < dpw; i++) for (int i = 0; i < dpw; i++)
O[o_off + i] = __float2bfloat16(acc[i] * rl); p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
} }
} }

12
csrc/tests/gqa_decode_test.cu
View File

@ -1,5 +1,4 @@
// Pure-C test for decode kernel // Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
// compile: nvcc -I csrc csrc/tests/gqa_decode_test.cu -o test && ./test
#include <cstdio> #include <cstdio>
#include <cstdlib> #include <cstdlib>
#include <cmath> #include <cmath>
@ -85,6 +84,12 @@ int main() {
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice); cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice); cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
GQAParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
p.use_mask=1; p.is_causal=0; p.causal_offset=0;
p.scale=1.0f/sqrtf((float)D);
p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
size_t smem=DC_CHUNK*D*sizeof(bf16); size_t smem=DC_CHUNK*D*sizeof(bf16);
dim3 block(32, gs); dim3 block(32, gs);
dim3 grid(B*Hk); dim3 grid(B*Hk);
@ -92,8 +97,7 @@ int main() {
grid.x, block.x, block.y, smem); grid.x, block.x, block.y, smem);
double t0=now_ms(); double t0=now_ms();
gqa_decode_attn_kernel<<<grid,block,smem>>>(dQ,dK,dV,dMask,dO, gqa_decode_attn_kernel<<<grid,block,smem>>>(p);
B,Hq,Hk,sl,D);
cudaDeviceSynchronize(); cudaDeviceSynchronize();
double kms=now_ms()-t0; double kms=now_ms()-t0;
cudaError_t err=cudaGetLastError(); cudaError_t err=cudaGetLastError();

11
csrc/tests/gqa_prefill_test.cu
View File

@ -1,4 +1,4 @@
// Pure-C test: compile with nvcc -I csrc csrc/tests/gqa_prefill_test.cu -o test && ./test // Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
#include <cstdio> #include <cstdio>
#include <cstdlib> #include <cstdlib>
#include <cmath> #include <cmath>
@ -81,6 +81,12 @@ int main() {
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]); for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice); cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
GQAParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
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 grid((ql+Br-1)/Br, Hq, B);
dim3 block(32, Br, 1); dim3 block(32, Br, 1);
size_t smem=2*Bc*D*sizeof(bf16); size_t smem=2*Bc*D*sizeof(bf16);
@ -88,8 +94,7 @@ int main() {
grid.x,grid.y,grid.z, block.x,block.y,block.z, smem); grid.x,grid.y,grid.z, block.x,block.y,block.z, smem);
double t0=now_ms(); double t0=now_ms();
gqa_prefill_attn_kernel<<<grid,block,smem>>>(dQ,dK,dV,nullptr,dO, gqa_prefill_attn_kernel<<<grid,block,smem>>>(p);
B,Hq,Hk,ql,kl,D,causal,0,0);
cudaDeviceSynchronize(); cudaDeviceSynchronize();
double kms=now_ms()-t0; double kms=now_ms()-t0;
cudaError_t err=cudaGetLastError(); cudaError_t err=cudaGetLastError();