diff --git a/csrc/kernels/gqa_common.cuh b/csrc/kernels/gqa_common.cuh new file mode 100644 index 0000000..26bab88 --- /dev/null +++ b/csrc/kernels/gqa_common.cuh @@ -0,0 +1,35 @@ +#pragma once +#include +#include +#include +#include + +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; +}; diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index 0e4f35b..bda3e93 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -1,39 +1,66 @@ -// torch binding for gqa_decode_attn -// kernel defined in gqa_decode_attn.cuh #include "gqa_decode_attn.cuh" #include 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 mask, + bool is_causal = false, + int64_t causal_offset = 0, + c10::optional 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(k.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"); - int B = q.size(0), n_heads = q.size(1), n_kv = k.size(1); - int seq_len = k.size(2), hd = q.size(3); - TORCH_CHECK(hd % 32 == 0, "head_dim must be multiple of 32"); - int group_size = n_heads / n_kv; - auto out = torch::empty_like(q); + GQAParams p; + p.batch = q.size(0); + p.q_head = q.size(1); + p.kv_head = k.size(1); + 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(); + } 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 grid(B * n_kv); + dim3 grid(p.batch * p.kv_head); - gqa_decode_attn_kernel<<>>( - reinterpret_cast(q.data_ptr()), - reinterpret_cast(k.data_ptr()), - reinterpret_cast(v.data_ptr()), - mask.data_ptr(), - reinterpret_cast(out.data_ptr()), - B, n_heads, n_kv, seq_len, hd - ); - return out; + gqa_decode_attn_kernel<<>>(p); + return O; } 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)"); } diff --git a/csrc/kernels/gqa_decode_attn.cuh b/csrc/kernels/gqa_decode_attn.cuh index 7119fa9..65e5079 100644 --- a/csrc/kernels/gqa_decode_attn.cuh +++ b/csrc/kernels/gqa_decode_attn.cuh @@ -1,78 +1,59 @@ -// gqa_decode_attn.cuh — header-only decode kernel #pragma once -#include -#include -#include -#include -using std::min; -using bf16 = __nv_bfloat16; +#include "gqa_common.cuh" -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; +__global__ void gqa_decode_attn_kernel(GQAParams p) { + int batch = blockIdx.x / p.kv_head; + int kv_head = blockIdx.x % p.kv_head; int group_size = blockDim.y; int q_head = kv_head * group_size + threadIdx.y; int lane = threadIdx.x; - int hd_per_thread = hd / 32; + int hd_per_thread = p.head_dim / 32; 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++) - 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 mask_base = batch * seq_len; + int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim; + int mask_base = batch * p.kv_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); + for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) { + 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) - 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(); 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; + partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]); + 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 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; + int v_off = kv_base + (chunk_start + s) * p.head_dim + 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; + acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta; m = new_m; } __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++) - out_ptr[out_off + i] = __float2bfloat16(acc_reg[i] / d); + p.o[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 25187be..2b8b4e4 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -1,12 +1,13 @@ -// torch binding for gqa_prefill_attn -// kernel defined in gqa_prefill_attn.cuh #include "gqa_prefill_attn.cuh" #include 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 mask, - bool is_causal = false, int64_t causal_offset = 0, + bool is_causal = false, + int64_t causal_offset = 0, c10::optional scale = c10::nullopt ) { 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(v.dtype() == torch::kBFloat16); - 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) { + GQAParams p; + p.batch = q.size(0); + p.q_head = q.size(1); + p.kv_head = k.size(1); + p.q_len = q.size(2); + 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) == B); - TORCH_CHECK(mask.value().size(1) == kv_len); - mask_ptr = mask.value().data_ptr(); + TORCH_CHECK(mask.value().size(0) == p.batch); + TORCH_CHECK(mask.value().size(1) == p.kv_len); + p.mask = mask.value().data_ptr(); + } else { + p.mask = nullptr; } 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); - size_t smem = 2 * Bc * D * sizeof(bf16); + size_t smem = 2 * Bc * p.head_dim * sizeof(bf16); - 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 - ); + gqa_prefill_attn_kernel<<>>(p); return O; } 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)"); } diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/gqa_prefill_attn.cuh index 9523470..2567089 100644 --- a/csrc/kernels/gqa_prefill_attn.cuh +++ b/csrc/kernels/gqa_prefill_attn.cuh @@ -1,82 +1,63 @@ -// gqa_prefill_attn.cuh — header-only kernel definition (no torch dependency) #pragma once -#include -#include -#include +#include "gqa_common.cuh" -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 -) { +__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 = 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}; - 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; + 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(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[]; 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}; - int tiles = (kv_len + Bc - 1) / Bc; + 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, kv_len - kv0); + int tlen = min(Bc, p.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]; + 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 (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; + 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 * D + d_part * dpw + i]); - dot = warp_sum(dot); + dot += qs[i] * __bfloat162float(sK[s * p.head_dim + d_part * dpw + i]); + 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; float nm = fmaxf(m, dot); @@ -85,16 +66,16 @@ __global__ void gqa_prefill_attn_kernel( 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; + acc[i] = acc[i] * al + __bfloat162float(sV[s * p.head_dim + 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; + 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++) - O[o_off + i] = __float2bfloat16(acc[i] * rl); + p.o[o_off + i] = __float2bfloat16(acc[i] * rl); } } diff --git a/csrc/tests/gqa_decode_test.cu b/csrc/tests/gqa_decode_test.cu index 47aa95a..e4756ff 100644 --- a/csrc/tests/gqa_decode_test.cu +++ b/csrc/tests/gqa_decode_test.cu @@ -1,5 +1,4 @@ -// Pure-C test for decode kernel -// compile: nvcc -I csrc csrc/tests/gqa_decode_test.cu -o test && ./test +// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test #include #include #include @@ -85,6 +84,12 @@ int main() { cudaMemcpy(dV,tmp,nKV*2,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); dim3 block(32, gs); dim3 grid(B*Hk); @@ -92,8 +97,7 @@ int main() { grid.x, block.x, block.y, smem); double t0=now_ms(); - gqa_decode_attn_kernel<<>>(dQ,dK,dV,dMask,dO, - B,Hq,Hk,sl,D); + gqa_decode_attn_kernel<<>>(p); cudaDeviceSynchronize(); double kms=now_ms()-t0; cudaError_t err=cudaGetLastError(); diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu index c98a212..27811f1 100644 --- a/csrc/tests/gqa_prefill_test.cu +++ b/csrc/tests/gqa_prefill_test.cu @@ -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 #include #include @@ -81,6 +81,12 @@ int main() { for (size_t i=0;i>>(dQ,dK,dV,nullptr,dO, - B,Hq,Hk,ql,kl,D,causal,0,0); + gqa_prefill_attn_kernel<<>>(p); cudaDeviceSynchronize(); double kms=now_ms()-t0; cudaError_t err=cudaGetLastError();