furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

GLM-DSA: much better PP long context performance (CUDA)

Browse Source
This commit is contained in:
Kawrakow
2026-07-11 15:38:14 +00:00
parent ff27465bac
commit fdada9003e
3 changed files with 322 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/dsa_attn.cu
+312
View File
@@ -0,0 +1,312 @@
#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;
m_out[row*nidx + col] = m_in[row*stride_m + idx[col]];
}
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];
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_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, const half * kqv16, float * dst) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= nelem) {
return;
}
dst[i] = __half2float(kqv16[i]);
}
template <int ncols_template, int block_size_template>
static __global__ void soft_max_f16_simple(half * x, const half * mask, const int ncols_par, const int nrows_y, const float scale) {
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; // broadcast the mask in the row dimension
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 = -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
#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);
}
const float inv_sum = 1.0f / tmp;
#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
static void soft_max_f16_cuda_simple(half * x, const half * mask, const int ncols_x, const int nrows_x,
const int nrows_y, const float scale, 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, ncols_x, nrows_y, scale);
break;
case 64:
soft_max_f16_simple<64, 64><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 128:
soft_max_f16_simple<128, 128><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 256:
soft_max_f16_simple<256, 256><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 512:
soft_max_f16_simple<512, 512><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 1024:
soft_max_f16_simple<1024, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 2048:
soft_max_f16_simple<2048, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
case 4096:
soft_max_f16_simple<4096, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
break;
default:
soft_max_f16_simple<0, 0><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
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 (sink) return false; // We do not support sinks at this point
if (!Q || !K || !V || !mask || !indexer) return false;
if (indexer->ne[0] % 256 != 0) return false; // lazyness to add checks and handle tailes 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 (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 || V->type != GGML_TYPE_F16 || mask->type != GGML_TYPE_F16 || Q->type != GGML_TYPE_F32) return false;
if (K->ne[0] != Q->ne[0]) return false;
float scale;
memcpy(&scale, dst->op_params, sizeof(float));
const half alpha = 1.0f;
const half beta = 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<half> kqv16(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<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);
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);
if (!is_k_view) {
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);
}
}
{
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(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], indexer->ne[0], Q->ne[2]*nrows,
Q->ne[2], scale, ctx.stream());
CUDA_CHECK(cudaGetLastError());
if (is_k_view) {
CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
V->ne[0], Q->ne[2], indexer->ne[0],
&alpha, k16.get() + v_offset, K->ne[0], K->ne[0]*indexer->ne[0],
kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2],
&beta, kqv16.get(), V->ne[0], V->ne[0]*Q->ne[2], nrows));
} else {
CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
V->ne[0], Q->ne[2], indexer->ne[0],
&alpha, v16.get(), V->ne[0], V->ne[0]*indexer->ne[0],
kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2],
&beta, kqv16.get(), V->ne[0], V->ne[0]*Q->ne[2], nrows));
}
{
int nelem = V->ne[0]*Q->ne[2]*nrows;
int nblock = (nelem + 255)/256;
k_copy_dst<<<nblock, 256, 0, ctx.stream()>>>(nelem, kqv16.get(),
(float *)((char *)dst->data + dst->nb[2]*first));
}
}
return true;
}
ggml/src/ggml-cuda/dsa_attn.cuh
+3
View File
@@ -0,0 +1,3 @@
#include "common.cuh"
bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
ggml/src/ggml-cuda/fattn.cu
+7 -1
View File
@@ -13,7 +13,7 @@
#include "fattn-mma-f16-interface.cuh"
#include "fattn-new-mma.cuh"
#include "fattn.cuh"
#include "convert.cuh"
#include "dsa_attn.cuh"
#include <cstdint>
@@ -43,6 +43,12 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
const int32_t precision = KQV->op_params[3];
const int32_t n_swa = KQV->op_params[4];
if (dst->src[5]) {
if (ggml_cuda_dsa_attn_ext(ctx, dst)) {
return;
}
}
ggml_tensor local_dst, Kl, Vl, Ml;
if (n_swa > 0) {
int ntokens = std::max(FATTN_KQ_STRIDE, int(Q->ne[1]));