From 579b8c3129dab26231b349acbbae934071544628 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 6 Jul 2026 13:44:23 +0800 Subject: [PATCH] fix: correct gqa_decode_attn reduction + add gqa_prefill_attn - gqa_decode_attn: rewrite to per-KV-head, K in smem - gqa_prefill_attn: new kernel for Q_len > 1 with GQA --- astrai/extension/__init__.py | 2 +- csrc/build.py | 1 + csrc/kernels/gqa_decode_attn.cu | 91 ++++++++++++++++-------- csrc/kernels/gqa_prefill_attn.cu | 118 +++++++++++++++++++++++++++++++ 4 files changed, 181 insertions(+), 31 deletions(-) create mode 100644 csrc/kernels/gqa_prefill_attn.cu diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index a861be1..d787fa9 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -5,7 +5,7 @@ logger = logging.getLogger(__name__) available: dict[str, bool] = {} -for _name in ["gqa_decode_attn"]: +for _name in ["gqa_decode_attn", "gqa_prefill_attn"]: try: importlib.import_module(f".{_name}", package=__package__) available[_name] = True diff --git a/csrc/build.py b/csrc/build.py index d8310fe..7e86212 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -27,3 +27,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs): register("gqa_decode_attn") +register("gqa_prefill_attn") diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index c13a3b4..2318b24 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -1,3 +1,4 @@ +// per-KV-head block, K shared in smem, each thread handles hd/32 elements #include #include #include @@ -6,6 +7,8 @@ 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); @@ -13,47 +16,69 @@ __inline__ __device__ float warp_reduce_sum(float val) { } __global__ void gqa_decode_attn_kernel( - const bf16* q_ptr, const bf16* k_ptr, const bf16* v_ptr, - const bool* mask_ptr, bf16* out_ptr, + 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_heads; - int q_head = blockIdx.x % n_heads; - int kv_head = q_head / (n_heads / n_kv_heads); - int tid = threadIdx.x; + 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]); - float q_val = __bfloat162float( - q_ptr[((batch * n_heads + q_head) * 1) * hd + tid]); 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 = 0.0f; - __shared__ float smem[2]; - float scale = 1.0f / sqrtf((float)hd); + float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f}; + float scale = rsqrtf((float)hd); - for (int s = 0; s < seq_len; s++) { - int off = kv_base + s * hd + tid; - float partial = q_val * __bfloat162float(k_ptr[off]); - partial = warp_reduce_sum(partial) * scale; + extern __shared__ __align__(16) bf16 k_smem[]; - if (tid % 32 == 0) smem[tid / 32] = partial; - __syncthreads(); - if (tid == 0) smem[0] = smem[0] + smem[1]; + 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(); - float score = smem[0]; - if (!mask_ptr[mask_base + s]) score = -FLT_MAX; + 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; - float new_m = fmaxf(m, score); - float alpha = expf(m - new_m); - float beta = expf(score - new_m); - d = d * alpha + beta; - acc = acc * alpha + __bfloat162float(v_ptr[off]) * beta; - m = new_m; + 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 + tid; - out_ptr[out_off] = __float2bfloat16(acc / d); + 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( @@ -68,9 +93,15 @@ torch::Tensor gqa_decode_attn( 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); - gqa_decode_attn_kernel<<>>( + size_t smem = CHUNK * hd * sizeof(bf16); // K chunk + dim3 block(32, group_size); + dim3 grid(B * n_kv); + + gqa_decode_attn_kernel<<>>( reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), reinterpret_cast(v.data_ptr()), @@ -82,5 +113,5 @@ torch::Tensor gqa_decode_attn( } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode attention (fused)"); + m.def("gqa_decode_attn", &gqa_decode_attn, "GQA decode v2 (per-KV-head, shared K)"); } diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu new file mode 100644 index 0000000..11c0064 --- /dev/null +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -0,0 +1,118 @@ +#include +#include +#include +#include +#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, + 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()); + TORCH_CHECK(q.dtype() == torch::kBFloat16); + 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"); + + bool use_mask = mask.has_value(); + const bool* mask_ptr = nullptr; + if (use_mask) { + TORCH_CHECK(mask.value().dtype() == torch::kBool); + mask_ptr = mask.value().data_ptr(); + } + + auto out = torch::empty_like(q); + + dim3 block(32); + dim3 grid(B * n_heads * q_len); + + 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 + ); + return out; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill attention (naive)"); +}