diff --git a/csrc/kernels/gqa_common.cuh b/csrc/kernels/gqa_common.cuh index 26bab88..981d601 100644 --- a/csrc/kernels/gqa_common.cuh +++ b/csrc/kernels/gqa_common.cuh @@ -7,15 +7,6 @@ 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; diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index f7261d7..f88aa40 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -1,5 +1,5 @@ #include "gqa_decode_attn.cuh" -#include +#include "gqa_entry_utils.cuh" #ifndef ASTRAI_NO_MMA #include "gqa_decode_attn_mma.cuh" @@ -76,43 +76,17 @@ static void dispatch_decode(GQAParams& p) { torch::Tensor gqa_decode_attn( torch::Tensor q, - torch::Tensor k, + torch::Tensor k, torch::Tensor v, c10::optional mask, - bool is_causal = false, + 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); - TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1"); - 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); + gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); + TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1"); 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; - } auto O = torch::empty_like(q); p.o = (bf16*)O.data_ptr(); diff --git a/csrc/kernels/gqa_decode_attn.cuh b/csrc/kernels/gqa_decode_attn.cuh index 65e5079..3a9826e 100644 --- a/csrc/kernels/gqa_decode_attn.cuh +++ b/csrc/kernels/gqa_decode_attn.cuh @@ -1,6 +1,14 @@ #pragma once #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(GQAParams p) { int batch = blockIdx.x / p.kv_head; int kv_head = blockIdx.x % p.kv_head; diff --git a/csrc/kernels/gqa_decode_attn_mma.cuh b/csrc/kernels/gqa_decode_attn_mma.cuh index 41a24b0..84daf85 100644 --- a/csrc/kernels/gqa_decode_attn_mma.cuh +++ b/csrc/kernels/gqa_decode_attn_mma.cuh @@ -109,92 +109,16 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { } __syncwarp(); - // S = Q @ K^T (Q already pre-scaled, so Sacc includes scale) + // S = Q @ K^T + online softmax + O += P @ V (shared MMA functions) float Sacc[NC8][4]; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f; - int krow_l = n8 * 8 + (lane & 7); - int kcol_h = (lane & 8) ? 8 : 0; -#pragma unroll - for (int kt = 0; kt < KD; kt++) { - unsigned b[2]; - ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]); - mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); - } - } + mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); - // ---- online softmax (Q pre-scaled → no per-tile scale multiply) ---- - float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - int cc = kv0 + n8 * 8 + 2 * tid4; - bool bc0 = (cc >= p.kv_len) || - (has_mask && !p.mask[mask_base + cc]); - bool bc1 = (cc + 1 >= p.kv_len) || - (has_mask && !p.mask[mask_base + cc + 1]); - bool cz = p.is_causal; - int off = p.causal_offset; - bool bad0 = bc0 || (cz && cc > off); - bool bad1 = bc1 || (cz && (cc + 1) > off); - float s0 = bad0 ? -FLT_MAX : Sacc[n8][0]; - float s1 = bad1 ? -FLT_MAX : Sacc[n8][1]; - float s2 = bad0 ? -FLT_MAX : Sacc[n8][2]; - float s3 = bad1 ? -FLT_MAX : Sacc[n8][3]; - Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3; - rmax0 = fmaxf(rmax0, fmaxf(s0, s1)); - rmax1 = fmaxf(rmax1, fmaxf(s2, s3)); - } - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1)); - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2)); + int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + mma_softmax_tile(kv0, maxc, maxc, + mask_base, p.mask, has_mask, + Sacc, Oacc, m0, m1, l0, l1, lane); - float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1); - float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0); - float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1); - - float rsum0 = 0.0f, rsum1 = 0.0f; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0); - float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0); - float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1); - float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1); - Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3; - rsum0 += p0 + p1; - rsum1 += p2 + p3; - } - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1); - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2); - l0 = l0 * corr0 + rsum0; - l1 = l1 * corr1 + rsum1; - m0 = nm0; m1 = nm1; - -#pragma unroll - for (int j = 0; j < DN8; j++) { - Oacc[j][0] *= corr0; Oacc[j][1] *= corr0; - Oacc[j][2] *= corr1; Oacc[j][3] *= corr1; - } - - // O += P @ V -#pragma unroll - for (int kt2 = 0; kt2 < KT2; kt2++) { - unsigned Pa[4]; - Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]); - Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]); - Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]); - Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]); - int vrow_l = kt2 * 16 + (lane & 15); -#pragma unroll - for (int dn8 = 0; dn8 < DN8; dn8++) { - unsigned b[2]; - ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); - mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); - } - } + mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); __syncwarp(); // sK/sV reused next tile } @@ -322,86 +246,14 @@ __global__ void gqa_decode_attn_mma_splitk_kernel(GQAParams p, __syncwarp(); float Sacc[NC8][4]; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f; - int krow_l = n8 * 8 + (lane & 7); - int kcol_h = (lane & 8) ? 8 : 0; -#pragma unroll - for (int kt = 0; kt < KD; kt++) { - unsigned b[2]; - ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]); - mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); - } - } + mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); - float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - int cc = kv0 + n8 * 8 + 2 * tid4; - bool bc0 = (cc >= p.kv_len) || (has_mask && !p.mask[mask_base + cc]); - bool bc1 = (cc + 1 >= p.kv_len) || (has_mask && !p.mask[mask_base + cc + 1]); - bool cz = p.is_causal; - int off = p.causal_offset; - bool bad0 = bc0 || (cz && cc > off); - bool bad1 = bc1 || (cz && (cc + 1) > off); - float s0 = bad0 ? -FLT_MAX : Sacc[n8][0]; - float s1 = bad1 ? -FLT_MAX : Sacc[n8][1]; - float s2 = bad0 ? -FLT_MAX : Sacc[n8][2]; - float s3 = bad1 ? -FLT_MAX : Sacc[n8][3]; - Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3; - rmax0 = fmaxf(rmax0, fmaxf(s0, s1)); - rmax1 = fmaxf(rmax1, fmaxf(s2, s3)); - } - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1)); - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2)); + int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + mma_softmax_tile(kv0, maxc, maxc, + mask_base, p.mask, has_mask, + Sacc, Oacc, m0, m1, l0, l1, lane); - float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1); - float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0); - float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1); - - float rsum0 = 0.0f, rsum1 = 0.0f; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0); - float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0); - float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1); - float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1); - Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3; - rsum0 += p0 + p1; - rsum1 += p2 + p3; - } - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1); - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2); - l0 = l0 * corr0 + rsum0; - l1 = l1 * corr1 + rsum1; - m0 = nm0; m1 = nm1; - -#pragma unroll - for (int j = 0; j < DN8; j++) { - Oacc[j][0] *= corr0; Oacc[j][1] *= corr0; - Oacc[j][2] *= corr1; Oacc[j][3] *= corr1; - } - -#pragma unroll - for (int kt2 = 0; kt2 < KT2; kt2++) { - unsigned Pa[4]; - Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]); - Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]); - Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]); - Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]); - int vrow_l = kt2 * 16 + (lane & 15); -#pragma unroll - for (int dn8 = 0; dn8 < DN8; dn8++) { - unsigned b[2]; - ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); - mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); - } - } + mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); __syncwarp(); } diff --git a/csrc/kernels/gqa_entry_utils.cuh b/csrc/kernels/gqa_entry_utils.cuh new file mode 100644 index 0000000..19b0605 --- /dev/null +++ b/csrc/kernels/gqa_entry_utils.cuh @@ -0,0 +1,42 @@ +#pragma once +#include +#include "gqa_common.cuh" + +inline void gqa_pack_params( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + c10::optional mask, + bool is_causal, + int64_t causal_offset, + c10::optional scale, + GQAParams& p +) { + 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); + + p.batch = (int)q.size(0); + p.q_head = (int)q.size(1); + p.kv_head = (int)k.size(1); + p.q_len = (int)q.size(2); + p.kv_len = (int)k.size(2); + p.head_dim = (int)q.size(3); + p.use_mask = mask.has_value() ? 1 : 0; + p.is_causal = is_causal ? 1 : 0; + 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; + } +} diff --git a/csrc/kernels/gqa_mma_utils.cuh b/csrc/kernels/gqa_mma_utils.cuh index fa4d65c..dae9310 100644 --- a/csrc/kernels/gqa_mma_utils.cuh +++ b/csrc/kernels/gqa_mma_utils.cuh @@ -104,3 +104,131 @@ template __device__ __forceinline__ void cp_async_wait_group() { asm volatile("cp.async.wait_group %0;" :: "n"(N)); } + +// --------------------------------------------------------------------------- +// Shared MMA compute functions — used by both decode and prefill MMA kernels. +// Extracted because S=Q@K^T, online softmax, and P@V are structurally identical +// between the two kernels; only the per-row causal/mask bounds differ. +// --------------------------------------------------------------------------- + +// S = Q @ K^T (Qa pre-loaded and pre-scaled by the caller). +// LD and SWIZ_MASK are constexpr in the calling kernel — passing them as +// runtime ints lets the compiler fold them while keeping the signature clean. +template +__device__ inline void mma_compute_scores( + const unsigned Qa[KD][4], + const bf16* __restrict__ sK, + int LD, int SWIZ_MASK, int lane, + float Sacc[NC8][4]) +{ + #pragma unroll + for (int n8 = 0; n8 < NC8; n8++) { + Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f; + int krow_l = n8 * 8 + (lane & 7); + int kcol_h = (lane & 8) ? 8 : 0; + #pragma unroll + for (int kt = 0; kt < KD; kt++) { + unsigned b[2]; + ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]); + mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); + } + } +} + +// Online softmax + Oacc rescale for one K/V tile. maxc0/maxc1 are the per-row +// KV column bounds — prefill passes per-query-row causal limits while decode +// passes the same value for both rows (q_len==1). Sacc is consumed in place +// (replaced by P = exp(S - nm) for the subsequent P@V step). +template +__device__ inline void mma_softmax_tile( + int kv0, + int maxc0, int maxc1, + int mask_base, + const bool* __restrict__ mask, + bool has_mask, + float Sacc[NC8][4], + float Oacc[DN8][4], + float& m0, float& m1, + float& l0, float& l1, + int lane) +{ + int tid4 = lane & 3; + + float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; + #pragma unroll + for (int n8 = 0; n8 < NC8; n8++) { + int cc = kv0 + n8 * 8 + 2 * tid4; + int c1 = cc + 1; + bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base + cc]); + bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base + c1]); + bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base + cc]); + bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base + c1]); + float s0 = b0 ? -FLT_MAX : Sacc[n8][0]; + float s1 = b1 ? -FLT_MAX : Sacc[n8][1]; + float s2 = b2 ? -FLT_MAX : Sacc[n8][2]; + float s3 = b3 ? -FLT_MAX : Sacc[n8][3]; + Sacc[n8][0] = s0; Sacc[n8][1] = s1; + Sacc[n8][2] = s2; Sacc[n8][3] = s3; + rmax0 = fmaxf(rmax0, fmaxf(s0, s1)); + rmax1 = fmaxf(rmax1, fmaxf(s2, s3)); + } + rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1)); + rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2)); + rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1)); + rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2)); + + float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1); + float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0); + float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1); + + float rsum0 = 0.0f, rsum1 = 0.0f; + #pragma unroll + for (int n8 = 0; n8 < NC8; n8++) { + float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0); + float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0); + float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1); + float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1); + Sacc[n8][0] = p0; Sacc[n8][1] = p1; + Sacc[n8][2] = p2; Sacc[n8][3] = p3; + rsum0 += p0 + p1; + rsum1 += p2 + p3; + } + rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1); + rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2); + rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1); + rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2); + l0 = l0 * corr0 + rsum0; + l1 = l1 * corr1 + rsum1; + m0 = nm0; m1 = nm1; + + #pragma unroll + for (int j = 0; j < DN8; j++) { + Oacc[j][0] *= corr0; Oacc[j][1] *= corr0; + Oacc[j][2] *= corr1; Oacc[j][3] *= corr1; + } +} + +// O += P @ V (Sacc must contain P = attention weights after softmax). +template +__device__ inline void mma_pv_accumulate( + float Sacc[][4], + const bf16* __restrict__ sV, + int LD, int SWIZ_MASK, int lane, + float Oacc[DN8][4]) +{ + #pragma unroll + for (int kt2 = 0; kt2 < KT2; kt2++) { + unsigned Pa[4]; + Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]); + Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]); + Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]); + Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]); + int vrow_l = kt2 * 16 + (lane & 15); + #pragma unroll + for (int dn8 = 0; dn8 < DN8; dn8++) { + unsigned b[2]; + ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); + mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); + } + } +} diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index e233a43..8d2f335 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -1,5 +1,5 @@ #include "gqa_prefill_attn.cuh" -#include +#include "gqa_entry_utils.cuh" #ifndef ASTRAI_NO_MMA #include "gqa_prefill_attn_mma.cuh" @@ -41,35 +41,9 @@ torch::Tensor gqa_prefill_attn( 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); - 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); + gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16"); - 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; - } auto O = torch::empty_like(q); p.o = (bf16*)O.data_ptr(); diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh index 87d1f8e..b9bbbd6 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/csrc/kernels/gqa_prefill_attn_mma.cuh @@ -163,101 +163,19 @@ void gqa_prefill_attn_mma_kernel(GQAParams p) { // Warp-level causal skip if (!use_skip || kv0 <= max_kv) { - // S = Q @ K^T → Sacc[n8][0..3] (n8: 8 kv cols each) + // S = Q @ K^T + online softmax + O += P @ V (shared MMA functions) float Sacc[NC8][4]; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f; - int krow_l = n8 * 8 + (lane & 7); - int kcol_h = (lane & 8) ? 8 : 0; -#pragma unroll - for (int kt = 0; kt < KD; kt++) { - unsigned b[2]; - ldmatrix_x2(b, &bK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]); - mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); - } - } + mma_compute_scores(Qa, bK, LD, SWIZ_MASK, lane, Sacc); - // ---- online softmax (in registers) ---- - // Q is pre-scaled, so Sacc already includes the attention scale. int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1) : p.kv_len; int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1) : p.kv_len; - float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - int cc = kv0 + n8 * 8 + 2 * tid4; - int c1 = cc + 1; - bool b0 = (cc >= maxc0) || (has_mask && !p.mask[mb + cc]); - bool b1 = (c1 >= maxc0) || (has_mask && !p.mask[mb + c1]); - bool b2 = (cc >= maxc1) || (has_mask && !p.mask[mb + cc]); - bool b3 = (c1 >= maxc1) || (has_mask && !p.mask[mb + c1]); - float s0 = b0 ? -FLT_MAX : Sacc[n8][0]; - float s1 = b1 ? -FLT_MAX : Sacc[n8][1]; - float s2 = b2 ? -FLT_MAX : Sacc[n8][2]; - float s3 = b3 ? -FLT_MAX : Sacc[n8][3]; - Sacc[n8][0] = s0; Sacc[n8][1] = s1; - Sacc[n8][2] = s2; Sacc[n8][3] = s3; - rmax0 = fmaxf(rmax0, fmaxf(s0, s1)); - rmax1 = fmaxf(rmax1, fmaxf(s2, s3)); - } - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1)); - rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1)); - rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2)); + mma_softmax_tile(kv0, maxc0, maxc1, + mb, p.mask, has_mask, + Sacc, Oacc, m0, m1, l0, l1, lane); - float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1); - float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0); - float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1); - - float rsum0 = 0.0f, rsum1 = 0.0f; -#pragma unroll - for (int n8 = 0; n8 < NC8; n8++) { - float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f - : __expf(Sacc[n8][0] - nm0); - float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f - : __expf(Sacc[n8][1] - nm0); - float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f - : __expf(Sacc[n8][2] - nm1); - float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f - : __expf(Sacc[n8][3] - nm1); - Sacc[n8][0] = p0; Sacc[n8][1] = p1; - Sacc[n8][2] = p2; Sacc[n8][3] = p3; - rsum0 += p0 + p1; - rsum1 += p2 + p3; - } - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1); - rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1); - rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2); - l0 = l0 * corr0 + rsum0; - l1 = l1 * corr1 + rsum1; - m0 = nm0; m1 = nm1; - - // rescale O accumulator by per-row correction -#pragma unroll - for (int j = 0; j < DN8; j++) { - Oacc[j][0] *= corr0; Oacc[j][1] *= corr0; - Oacc[j][2] *= corr1; Oacc[j][3] *= corr1; - } - - // O += P @ V -#pragma unroll - for (int kt2 = 0; kt2 < KT2; kt2++) { - unsigned Pa[4]; - Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]); - Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]); - Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]); - Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]); - int vrow_l = kt2 * 16 + (lane & 15); -#pragma unroll - for (int dn8 = 0; dn8 < DN8; dn8++) { - unsigned b[2]; - ldmatrix_x2_trans(b, &bV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); - mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); - } - } + mma_pv_accumulate(Sacc, bV, LD, SWIZ_MASK, lane, Oacc); } // if active (warp-level causal skip) }