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
This commit is contained in:
parent
d7da51569f
commit
579b8c3129
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -27,3 +27,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
|
|||
|
||||
|
||||
register("gqa_decode_attn")
|
||||
register("gqa_prefill_attn")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
// per-KV-head block, K shared in smem, each thread handles hd/32 elements
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
|
@ -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<<<dim3(B * n_heads), dim3(hd)>>>(
|
||||
size_t smem = CHUNK * hd * sizeof(bf16); // K chunk
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(B * n_kv);
|
||||
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(
|
||||
reinterpret_cast<const bf16*>(q.data_ptr()),
|
||||
reinterpret_cast<const bf16*>(k.data_ptr()),
|
||||
reinterpret_cast<const bf16*>(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)");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,118 @@
|
|||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <cfloat>
|
||||
#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,
|
||||
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());
|
||||
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<bool>();
|
||||
}
|
||||
|
||||
auto out = torch::empty_like(q);
|
||||
|
||||
dim3 block(32);
|
||||
dim3 grid(B * n_heads * q_len);
|
||||
|
||||
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
|
||||
);
|
||||
return out;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("gqa_prefill_attn", &gqa_prefill_attn, "GQA prefill attention (naive)");
|
||||
}
|
||||
Loading…
Reference in New Issue