refactor: deduplicate kernel code with shared MMA and entry-point helpers
This commit is contained in:
parent
88f8dca2c2
commit
29b0423c4e
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
#include "gqa_decode_attn.cuh"
|
||||
#include <torch/extension.h>
|
||||
#include "gqa_entry_utils.cuh"
|
||||
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "gqa_decode_attn_mma.cuh"
|
||||
|
|
@ -83,36 +83,10 @@ torch::Tensor gqa_decode_attn(
|
|||
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);
|
||||
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<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
}
|
||||
|
||||
auto O = torch::empty_like(q);
|
||||
p.o = (bf16*)O.data_ptr();
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<KD, NC8>(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<NC8, DN8>(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<DN8, KT2>(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<KD, NC8>(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<NC8, DN8>(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<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
#pragma once
|
||||
#include <torch/extension.h>
|
||||
#include "gqa_common.cuh"
|
||||
|
||||
inline void gqa_pack_params(
|
||||
torch::Tensor q,
|
||||
torch::Tensor k,
|
||||
torch::Tensor v,
|
||||
c10::optional<torch::Tensor> mask,
|
||||
bool is_causal,
|
||||
int64_t causal_offset,
|
||||
c10::optional<double> 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<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
}
|
||||
}
|
||||
|
|
@ -104,3 +104,131 @@ template <int N>
|
|||
__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 <int KD, int NC8>
|
||||
__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 <int NC8, int DN8>
|
||||
__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 <int DN8, int KT2>
|
||||
__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]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
#include "gqa_prefill_attn.cuh"
|
||||
#include <torch/extension.h>
|
||||
#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<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);
|
||||
|
||||
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<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
}
|
||||
|
||||
auto O = torch::empty_like(q);
|
||||
p.o = (bf16*)O.data_ptr();
|
||||
|
|
|
|||
|
|
@ -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<KD, NC8>(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<NC8, DN8>(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<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
|
||||
} // if active (warp-level causal skip)
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue