388 lines
15 KiB
Plaintext
388 lines
15 KiB
Plaintext
#include "dsa_attn.cuh"
|
|
|
|
static inline bool v_is_k_view(const ggml_tensor * K, const ggml_tensor * V) {
|
|
if (!V || !V->data) return false;
|
|
auto k_data = (const char *)K->data;
|
|
auto v_data = (const char *)V->data;
|
|
auto k_row_size = ggml_row_size(K->type, K->ne[0]);
|
|
auto v_row_size = ggml_row_size(V->type, V->ne[0]);
|
|
return v_data >= k_data && v_data + v_row_size <= k_data + k_row_size;
|
|
}
|
|
|
|
static __global__ void k_prepare_mask(int nidx, const int * __restrict__ idx, const half * __restrict__ m_in,
|
|
half * __restrict__ m_out, size_t stride_idx, size_t stride_m) {
|
|
int row = blockIdx.x;
|
|
int col = blockIdx.y*blockDim.x + threadIdx.x;
|
|
idx += row*stride_idx;
|
|
int ii = idx[col];
|
|
m_out[row*nidx + col] = ii >= 0 ? m_in[row*stride_m + ii] : __float2half(-INFINITY);
|
|
}
|
|
|
|
static __global__ void k_prepare_one_batch_kv(int nk, int ncol, const int * idx, const char * k_in,
|
|
half * k_out, size_t stride_k, size_t stride_idx) {
|
|
int row = blockIdx.y;
|
|
int col = blockIdx.x;
|
|
int i = idx[row*stride_idx + col];
|
|
if (i < 0) {
|
|
i = 0;
|
|
}
|
|
auto k_row = (const half *)(k_in + stride_k * i);
|
|
k_out += (row*ncol + col)*nk;
|
|
for (int j = threadIdx.x; j < nk; j += blockDim.x) {
|
|
k_out[j] = k_row[j];
|
|
}
|
|
}
|
|
|
|
static __global__ void k_prepare_one_batch_kv_q8_0(int nk, int ncol, const int * idx, const char * k_in,
|
|
half * k_out, size_t stride_k, size_t stride_idx) {
|
|
int row = blockIdx.y;
|
|
int col = blockIdx.x;
|
|
int i = idx[row*stride_idx + col];
|
|
if (i < 0) {
|
|
i = 0;
|
|
}
|
|
auto k_row = (const block_q8_0 *)(k_in + stride_k * i);
|
|
k_out += (row*ncol + col)*nk;
|
|
for (int j = threadIdx.x; j < nk; j += blockDim.x) {
|
|
k_out[j] = k_row[j/32].d * (half)k_row[j/32].qs[j%32];
|
|
}
|
|
}
|
|
|
|
static __global__ void k_prepare_one_batch_q(int ne0, int ne1, size_t nb1, size_t nb2,
|
|
const float * q_in, half * q_out) {
|
|
int i0 = blockIdx.x*blockDim.x + threadIdx.x;
|
|
if (i0 >= ne0) {
|
|
return;
|
|
}
|
|
int i1 = blockIdx.y;
|
|
int i2 = blockIdx.z;
|
|
q_out[i0 + (i2 + i1*ne1)*ne0] = __float2half(q_in[i0 + i1*nb1 + i2*nb2]);
|
|
}
|
|
|
|
static __global__ void k_copy_dst(int nelem, int ncols, const float * kqv32, const float * inv_sum, float * dst) {
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= nelem) {
|
|
return;
|
|
}
|
|
dst[i] = kqv32[i] * inv_sum[i / ncols];
|
|
}
|
|
|
|
template <int ncols_template, int block_size_template>
|
|
static __global__ void soft_max_f16_simple(half * x, const half * mask, const float * sinks, const int ncols_par, const int nrows_y,
|
|
const float scale, float * store_inv_sum) {
|
|
const int ncols = ncols_template == 0 ? ncols_par : ncols_template;
|
|
|
|
const int tid = threadIdx.x;
|
|
const int rowx = blockIdx.x;
|
|
const int rowy = rowx / nrows_y;
|
|
|
|
const int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
|
|
|
|
const int warp_id = threadIdx.x / WARP_SIZE;
|
|
const int lane_id = threadIdx.x % WARP_SIZE;
|
|
|
|
extern __shared__ float data_soft_max_f32[];
|
|
float * buf_iw = data_soft_max_f32; // shared memory buffer for inter-warp communication
|
|
// shared memory buffer to cache values between iterations:
|
|
float * vals = buf_iw + WARP_SIZE;
|
|
|
|
float max_val = sinks ? sinks[rowx % nrows_y] : -INFINITY;
|
|
|
|
#pragma unroll
|
|
for (int col0 = 0; col0 < ncols; col0 += block_size) {
|
|
const int col = col0 + tid;
|
|
|
|
if (ncols_template == 0 && col >= ncols) {
|
|
break;
|
|
}
|
|
|
|
const int64_t ix = (int64_t)rowx*ncols + col;
|
|
const int64_t iy = (int64_t)rowy*ncols + col;
|
|
|
|
const float val = scale*__half2float(x[ix]) + __half2float(mask[iy]);
|
|
|
|
vals[col] = val;
|
|
max_val = max(max_val, val);
|
|
}
|
|
|
|
// find the max value in the block
|
|
max_val = warp_reduce_max(max_val);
|
|
if (block_size > WARP_SIZE) {
|
|
if (warp_id == 0) {
|
|
buf_iw[lane_id] = -INFINITY;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (lane_id == 0) {
|
|
buf_iw[warp_id] = max_val;
|
|
}
|
|
__syncthreads();
|
|
|
|
max_val = buf_iw[lane_id];
|
|
max_val = warp_reduce_max(max_val);
|
|
}
|
|
|
|
float tmp = 0.0f; // partial sum
|
|
|
|
if (store_inv_sum) {
|
|
#pragma unroll
|
|
for (int col0 = 0; col0 < ncols; col0 += block_size) {
|
|
const int col = col0 + tid;
|
|
|
|
if (ncols_template == 0 && col >= ncols) {
|
|
break;
|
|
}
|
|
|
|
const float val = expf(vals[col] - max_val);
|
|
tmp += val;
|
|
const int64_t ix = (int64_t)rowx*ncols + col;
|
|
x[ix] = __float2half(val);
|
|
}
|
|
} else {
|
|
#pragma unroll
|
|
for (int col0 = 0; col0 < ncols; col0 += block_size) {
|
|
const int col = col0 + tid;
|
|
|
|
if (ncols_template == 0 && col >= ncols) {
|
|
break;
|
|
}
|
|
|
|
const float val = expf(vals[col] - max_val);
|
|
tmp += val;
|
|
vals[col] = val;
|
|
}
|
|
}
|
|
|
|
// find the sum of exps in the block
|
|
tmp = warp_reduce_sum(tmp);
|
|
if (block_size > WARP_SIZE) {
|
|
__syncthreads();
|
|
if (warp_id == 0) {
|
|
buf_iw[lane_id] = 0.0f;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (lane_id == 0) {
|
|
buf_iw[warp_id] = tmp;
|
|
}
|
|
__syncthreads();
|
|
|
|
tmp = buf_iw[lane_id];
|
|
tmp = warp_reduce_sum(tmp);
|
|
}
|
|
|
|
if (sinks) {
|
|
tmp += expf(sinks[rowx % nrows_y] - max_val);
|
|
}
|
|
|
|
const float inv_sum = 1.0f / tmp;
|
|
|
|
if (store_inv_sum) {
|
|
if (tid == 0) {
|
|
store_inv_sum[rowx] = inv_sum;
|
|
}
|
|
return;
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int col0 = 0; col0 < ncols; col0 += block_size) {
|
|
const int col = col0 + tid;
|
|
|
|
if (ncols_template == 0 && col >= ncols) {
|
|
return;
|
|
}
|
|
|
|
const int64_t ix = (int64_t)rowx*ncols + col;
|
|
x[ix] = __float2half(vals[col] * inv_sum);
|
|
}
|
|
}
|
|
|
|
#define CUDA_SOFT_MAX_BLOCK_SIZE 1024
|
|
|
|
// nrows_y is Q->ne[2]
|
|
// nrows_x is Q->ne[2] * nrows
|
|
static void soft_max_f16_cuda_simple(half * x, const half * mask, const float * sinks, const int ncols_x, const int nrows_x,
|
|
const int nrows_y, const float scale, float * store_inv_sum, cudaStream_t stream) {
|
|
int nth = WARP_SIZE;
|
|
while (nth < ncols_x && nth < CUDA_SOFT_MAX_BLOCK_SIZE) nth *= 2;
|
|
const dim3 block_dims(nth, 1, 1);
|
|
const dim3 block_nums(nrows_x, 1, 1);
|
|
const size_t shmem = (GGML_PAD(ncols_x, WARP_SIZE) + WARP_SIZE)*sizeof(float);
|
|
static_assert(CUDA_SOFT_MAX_BLOCK_SIZE == 1024, "These values need to be adjusted.");
|
|
|
|
GGML_ASSERT(shmem < ggml_cuda_info().devices[ggml_cuda_get_device()].smpb);
|
|
|
|
switch (ncols_x) {
|
|
case 32:
|
|
soft_max_f16_simple<32, 32><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 64:
|
|
soft_max_f16_simple<64, 64><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 128:
|
|
soft_max_f16_simple<128, 128><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 256:
|
|
soft_max_f16_simple<256, 256><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 512:
|
|
soft_max_f16_simple<512, 512><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 1024:
|
|
soft_max_f16_simple<1024, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 2048:
|
|
soft_max_f16_simple<2048, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
case 4096:
|
|
soft_max_f16_simple<4096, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
default:
|
|
soft_max_f16_simple<0, 0><<<block_nums, block_dims, shmem, stream>>>(x, mask, sinks, ncols_x, nrows_y, scale, store_inv_sum);
|
|
break;
|
|
}
|
|
}
|
|
|
|
bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
|
if (!dst) return false;
|
|
|
|
constexpr int k_max_rows = 32;
|
|
|
|
const ggml_tensor * Q = dst->src[0];
|
|
const ggml_tensor * K = dst->src[1];
|
|
const ggml_tensor * V = dst->src[2];
|
|
const ggml_tensor * mask = dst->src[3];
|
|
const ggml_tensor * sink = dst->src[4];
|
|
const ggml_tensor * indexer = dst->src[5];
|
|
|
|
if (!Q || !K || !V || !mask || !indexer) return false;
|
|
|
|
if (indexer->ne[0] % 256 != 0) return false; // lazyness to add checks and handle tails in case of not multiple of 256
|
|
// But are there DSA variants where top_k is not a multiple of 256?
|
|
//if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency
|
|
if (Q->ne[1] <= 16) {
|
|
if (indexer->ne[0] >= K->ne[1]) return false;
|
|
} else {
|
|
if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency
|
|
}
|
|
if (K->ne[2] > 1 || K->ne[3] > 1 || mask->ne[2] > 1 || mask->ne[3] > 1 || Q->ne[3] > 1) return false;
|
|
if ((K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_Q8_0) ||
|
|
(V->type != GGML_TYPE_F16 && V->type != GGML_TYPE_Q8_0) || mask->type != GGML_TYPE_F16 || Q->type != GGML_TYPE_F32) return false;
|
|
if (K->ne[0] != Q->ne[0]) return false;
|
|
|
|
//printf("%s(%s)\n", __func__, dst->name);
|
|
|
|
float scale;
|
|
memcpy(&scale, dst->op_params, sizeof(float));
|
|
|
|
const half alpha = 1.0f;
|
|
const half beta = 0.0f;
|
|
float alpha_32 = 1.0f, beta_32 = 0.0f;
|
|
|
|
int max_rows = std::min<int>(Q->ne[1], k_max_rows);
|
|
bool is_k_view = v_is_k_view(K, V);
|
|
auto mask_size = indexer->ne[0]*Q->ne[1]; // mask is relatively small, so we can do it once for the whole calculation
|
|
auto k_cache_size = indexer->ne[0]*K->ne[0]*max_rows;
|
|
auto v_cache_size = indexer->ne[0]*V->ne[0]*max_rows;
|
|
auto q_size = Q->ne[0]*Q->ne[2]*max_rows;
|
|
auto kq_size = indexer->ne[0]*Q->ne[2]*max_rows;
|
|
auto kqv_size = V->ne[0]*Q->ne[2]*max_rows;
|
|
ggml_cuda_pool_alloc<half> q16(ctx.pool(), q_size);
|
|
ggml_cuda_pool_alloc<half> kq16(ctx.pool(), kq_size);
|
|
ggml_cuda_pool_alloc<float> kqv32(ctx.pool(), kqv_size);
|
|
ggml_cuda_pool_alloc<half> mask16(ctx.pool(), mask_size);
|
|
ggml_cuda_pool_alloc<half> k16(ctx.pool(), k_cache_size);
|
|
ggml_cuda_pool_alloc<float> inv_sum(ctx.pool(), max_rows);
|
|
ggml_cuda_pool_alloc<half> v16(ctx.pool());
|
|
size_t v_offset = 0;
|
|
if (is_k_view) {
|
|
v_offset = (const half *)V->data - (const half *)K->data;
|
|
} else {
|
|
v16.alloc(v_cache_size);
|
|
}
|
|
auto stride_idx = indexer->nb[1]/sizeof(int);
|
|
{
|
|
dim3 grid(Q->ne[1], indexer->ne[0]/256, 1);
|
|
k_prepare_mask<<<grid, 256, 0, ctx.stream()>>>(indexer->ne[0], (const int * )indexer->data,
|
|
(const half *)mask->data, mask16.get(), stride_idx, mask->nb[1]/sizeof(half));
|
|
}
|
|
|
|
int nstep = (Q->ne[1] + max_rows - 1)/max_rows;
|
|
|
|
for (int istep = 0; istep < nstep; ++istep) {
|
|
int first = istep*max_rows;
|
|
int last = std::min<int>(first + max_rows, Q->ne[1]);
|
|
int nrows = last - first;
|
|
{
|
|
dim3 grid(indexer->ne[0], nrows, 1);
|
|
if (K->type == GGML_TYPE_F16) {
|
|
k_prepare_one_batch_kv<<<grid, 256, 0, ctx.stream()>>>(K->ne[0], indexer->ne[0],
|
|
(const int *)indexer->data + stride_idx*first,
|
|
(const char *)K->data, k16.get(), K->nb[1], stride_idx);
|
|
} else {
|
|
k_prepare_one_batch_kv_q8_0<<<grid, 256, 0, ctx.stream()>>>(K->ne[0], indexer->ne[0],
|
|
(const int *)indexer->data + stride_idx*first,
|
|
(const char *)K->data, k16.get(), K->nb[1], stride_idx);
|
|
}
|
|
if (!is_k_view) {
|
|
if (V->type == GGML_TYPE_F16) {
|
|
k_prepare_one_batch_kv<<<grid, 256, 0, ctx.stream()>>>(V->ne[0], indexer->ne[0],
|
|
(const int *)indexer->data + stride_idx*first,
|
|
(const char *)V->data, v16.get(), V->nb[1], stride_idx);
|
|
} else {
|
|
k_prepare_one_batch_kv_q8_0<<<grid, 256, 0, ctx.stream()>>>(V->ne[0], indexer->ne[0],
|
|
(const int *)indexer->data + stride_idx*first,
|
|
(const char *)V->data, v16.get(), V->nb[1], stride_idx);
|
|
}
|
|
}
|
|
}
|
|
{
|
|
int nblock = (Q->ne[0] + 255)/256;
|
|
dim3 grid(nblock, nrows, Q->ne[2]);
|
|
k_prepare_one_batch_q<<<grid, 256, 0, ctx.stream()>>>(Q->ne[0], Q->ne[2],
|
|
Q->nb[1]/sizeof(float), Q->nb[2]/sizeof(float),
|
|
(const float *)((const char *)Q->data + first*Q->nb[1]), q16.get());
|
|
}
|
|
|
|
CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), ctx.stream()));
|
|
CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
|
|
indexer->ne[0], Q->ne[2], Q->ne[0],
|
|
&alpha, k16.get(), K->ne[0], K->ne[0]*indexer->ne[0],
|
|
q16.get(), Q->ne[0], Q->ne[0]*Q->ne[2],
|
|
&beta, kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2], nrows));
|
|
|
|
soft_max_f16_cuda_simple(kq16.get(), mask16.get() + first*indexer->ne[0],
|
|
sink ? (const float *)sink->data : nullptr,
|
|
indexer->ne[0], Q->ne[2]*nrows,
|
|
Q->ne[2], scale, inv_sum.get(), ctx.stream());
|
|
CUDA_CHECK(cudaGetLastError());
|
|
|
|
if (is_k_view) {
|
|
CUBLAS_CHECK(cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
|
|
V->ne[0], Q->ne[2], indexer->ne[0],
|
|
&alpha_32, k16.get() + v_offset, CUDA_R_16F, K->ne[0], K->ne[0]*indexer->ne[0],
|
|
kq16.get(), CUDA_R_16F, indexer->ne[0], indexer->ne[0]*Q->ne[2],
|
|
&beta_32, kqv32.get(), CUDA_R_32F, V->ne[0], V->ne[0]*Q->ne[2], nrows,
|
|
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
|
} else {
|
|
CUBLAS_CHECK(cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
|
|
V->ne[0], Q->ne[2], indexer->ne[0],
|
|
&alpha_32, v16.get(), CUDA_R_16F, V->ne[0], V->ne[0]*indexer->ne[0],
|
|
kq16.get(), CUDA_R_16F, indexer->ne[0], indexer->ne[0]*Q->ne[2],
|
|
&beta_32, kqv32.get(), CUDA_R_32F, V->ne[0], V->ne[0]*Q->ne[2], nrows,
|
|
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
|
}
|
|
|
|
{
|
|
int nelem = V->ne[0]*Q->ne[2]*nrows;
|
|
int nblock = (nelem + 255)/256;
|
|
k_copy_dst<<<nblock, 256, 0, ctx.stream()>>>(nelem, dst->ne[0], kqv32.get(), inv_sum.get(),
|
|
(float *)((char *)dst->data + dst->nb[2]*first));
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return true;
|
|
}
|