From a90169b45d5403ea0cc95caee36993fc1a38d207 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 8 Jul 2026 13:16:23 +0000 Subject: [PATCH 1/8] WIP: indexer_topk on CUDA --- ggml/src/ggml-cuda.cu | 5 +++ ggml/src/ggml-cuda/argsort.cu | 2 +- ggml/src/ggml-cuda/argsort.cuh | 9 ++++++ ggml/src/ggml-cuda/mmq.cu | 59 ++++++++++++++++++---------------- ggml/src/ggml-cuda/mmq.cuh | 2 ++ 5 files changed, 49 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index e006e661cf..53d358d799 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -57,6 +57,7 @@ #include "ggml-cuda/tri.cuh" #include "ggml-cuda/delta-net.cuh" #include "ggml-cuda/blend.cuh" +#include "ggml-cuda/indexer_topk.cuh" #include #include @@ -4126,6 +4127,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_FLASH_ATTN_EXT: ggml_cuda_flash_attn_ext(ctx, dst); break; + case GGML_OP_INDEXER_TOPK: + ggml_cuda_op_indexer_topk(ctx, dst); + break; default: return false; } @@ -5028,6 +5032,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons op->src[1]->ne[0] == op->src[0]->ne[1] && op->src[3]->ne[0] == op->src[0]->ne[2]; case GGML_OP_DELTA_NET: + case GGML_OP_INDEXER_TOPK: return true; case GGML_OP_FLASH_ATTN_EXT: #if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 2ef26b9a5c..0aaf5861d1 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -466,7 +466,7 @@ static __global__ void init_indices(int * indices, const int ncols, const int nr } } -static void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, +void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, const int ncols, diff --git a/ggml/src/ggml-cuda/argsort.cuh b/ggml/src/ggml-cuda/argsort.cuh index 331f373bfe..7248d98eea 100644 --- a/ggml/src/ggml-cuda/argsort.cuh +++ b/ggml/src/ggml-cuda/argsort.cuh @@ -17,3 +17,12 @@ void cuda_bailingmoev2_experts(ggml_backend_cuda_context & ctx, ggml_tensor * ds void cuda_glm45moe_experts(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * topk); void cuda_openai_experts(ggml_backend_cuda_context & ctx, ggml_tensor * topk, ggml_tensor * softmax); + +void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, + const float * x, + int * dst, + const int ncols, + const int nrows, + ggml_sort_order order, + cudaStream_t stream); + diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 5991547efa..65b87303c6 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -7,33 +7,9 @@ #include "mmq.cuh" -void ggml_cuda_op_mul_mat_q( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, - const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, - const int64_t src1_padded_row_size, cudaStream_t stream) { - - const int64_t ne00 = src0->ne[0]; - const int64_t nb01 = ggml_row_size(src0->type, ne00); - - const int64_t ne10 = src1->ne[0]; - const int64_t ne11 = src1->ne[1]; - GGML_ASSERT(ne10 % QK8_1 == 0); - - const int64_t ne0 = dst->ne[0]; - - const int64_t row_diff = row_high - row_low; - - int id = ggml_cuda_get_device(); - const int compute_capability = ggml_cuda_info().devices[id].cc; - - // the main device has a larger memory buffer to hold the results from all GPUs - // nrows_dst == nrows of the matrix that the kernel writes into - const int64_t nrows_dst = id == ctx.device ? ne0 : row_diff; - - const mmq_args args = {src0_dd_i, src1_ddq_i, dst_dd_i, ne00, row_diff, nb01, src1_padded_row_size, src1_ncols, ne11, nrows_dst}; - - switch (src0->type) { +void ggml_cuda_op_mul_mat_q(ggml_backend_cuda_context & ctx, enum ggml_type type, const mmq_args & args) { + auto stream = ctx.stream(); + switch (type) { case GGML_TYPE_Q4_0: mul_mat_q_case(ctx, args, stream); break; @@ -164,6 +140,35 @@ void ggml_cuda_op_mul_mat_q( GGML_ABORT("fatal error"); break; } +} + +void ggml_cuda_op_mul_mat_q( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream) { + + const int64_t ne00 = src0->ne[0]; + const int64_t nb01 = ggml_row_size(src0->type, ne00); + + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + GGML_ASSERT(ne10 % QK8_1 == 0); + + const int64_t ne0 = dst->ne[0]; + + const int64_t row_diff = row_high - row_low; + + int id = ggml_cuda_get_device(); + const int compute_capability = ggml_cuda_info().devices[id].cc; + + // the main device has a larger memory buffer to hold the results from all GPUs + // nrows_dst == nrows of the matrix that the kernel writes into + const int64_t nrows_dst = id == ctx.device ? ne0 : row_diff; + + const mmq_args args = {src0_dd_i, src1_ddq_i, dst_dd_i, ne00, row_diff, nb01, src1_padded_row_size, src1_ncols, ne11, nrows_dst}; + + ggml_cuda_op_mul_mat_q(ctx, src0->type, args); GGML_UNUSED(src1); GGML_UNUSED(dst); diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index ee10d9ed64..eeb5dbcce4 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -4314,4 +4314,6 @@ void ggml_cuda_op_mul_mat_q( const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, cudaStream_t stream); +void ggml_cuda_op_mul_mat_q(ggml_backend_cuda_context & ctx, enum ggml_type type, const mmq_args & args); + bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11); From 6c279d1a2d37fa2776a530c498b3d296933fc255 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 8 Jul 2026 14:58:28 +0000 Subject: [PATCH 2/8] Forgot these --- ggml/src/ggml-cuda/indexer_topk.cu | 226 ++++++++++++++++++++++++++++ ggml/src/ggml-cuda/indexer_topk.cuh | 8 + 2 files changed, 234 insertions(+) create mode 100644 ggml/src/ggml-cuda/indexer_topk.cu create mode 100644 ggml/src/ggml-cuda/indexer_topk.cuh diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu new file mode 100644 index 0000000000..373dc7236f --- /dev/null +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -0,0 +1,226 @@ +#include "indexer_topk.cuh" +#include "mmq.cuh" +#include "quantize.cuh" +#include "convert.cuh" +#include "argsort.cuh" + +template +static __global__ void k_fused_relu_mul_sum_rows(const float * __restrict__ kq, const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, const int ncols, const int nhead, size_t nbm) { + const int row = blockIdx.x; + const int col = threadIdx.x; + + int64_t step = ncols*nhead; + auto this_w = w + blockIdx.x*nhead; + auto this_m = (const mask_t *)((const char *)m + nbm*row); + + for (int i = col; i < ncols; i += blockDim.x) { + float sum = (float)this_m[i]; + auto this_kq = kq + blockIdx.x * step; + for (int head = 0; head < nhead; ++head) { + float relu = this_kq[i]; + relu = relu > 0.0f ? relu : 0.0f; + sum += relu * this_w[head]; + this_kq += ncols; + } + dst[ncols*row + i] = sum; + } +} + +static __global__ void k_copy_topk(const int * __restrict__ sorted, int * dst, const int ncols, const int n_top_k) { + const int row = blockIdx.x; + const int col = threadIdx.x; + sorted += int64_t(ncols)*row; + dst += int64_t(n_top_k)*row; + for (int i = col; i < n_top_k; i += blockDim.x) { + dst[i] = sorted[i]; + } +} + +template +static __global__ void k_fused_gemm_relu_mul_sum_rows(const half * __restrict__ k, const float * __restrict__ q, + const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, + int head_size, int n_kv, int n_head, size_t nbm) { + const int row_k = blockIdx.x; + const int row_q = blockIdx.y; + const int col = threadIdx.x; + + k += size_t(row_k) * head_size; + q += size_t(row_q) * head_size * n_head; + w += size_t(row_q) * n_head; + auto this_m = (const mask_t *)((const char *)m + nbm*row_q); + + auto k2 = (const half2 *)k; + auto q2 = (const float2 *)q; + float score = (float)this_m[row_k]; + for (int head = 0; head < n_head; ++head) { + half2 sum = {}; + for (int i = col; i < head_size/2; i += blockDim.x) { + auto qh = __float22half2_rn(q2[i]); + sum += k2[i] * qh; + } + float sumf = (float)(sum.x + sum.y); + sumf = warp_reduce_sum(sumf); + sumf = sumf > 0.0f ? sumf : 0.0f; + score += sumf * w[head]; + q2 += head_size/2; + } + if (col == 0) { + dst[size_t(n_kv)*row_q + row_k] = score; + } +} + +//template +//static __global__ void k_fused_gemm_relu_mul_sum_rows(const half * __restrict__ k, const float * __restrict__ q, +// const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, +// int head_size, int n_kv, int n_head, size_t nbm) { +// const int row_k = blockIdx.x; +// const int row_q = blockIdx.y; +// const int col = threadIdx.x % WARP_SIZE; +// const int head = threadIdx.x / WARP_SIZE; +// +// k += size_t(row_k) * head_size; +// q += size_t(row_q) * head_size * n_head + head * head_size; +// w += size_t(row_q) * n_head; +// auto this_m = (const mask_t *)((const char *)m + nbm*row_q); +// +// __shared__ float s[32]; +// +// auto k2 = (const half2 *)k; +// auto q2 = (const float2 *)q; +// half2 sum = {}; +// for (int i = col; i < head_size/2; i += WARP_SIZE) { +// auto qh = __float22half2_rn(q2[i]); +// sum += k2[i] * qh; +// } +// float sumf = (float)(sum.x + sum.y); +// sumf = warp_reduce_sum(sumf); +// sumf = sumf > 0.0f ? sumf * w[head] : 0.0f; +// if (col == 0) { +// s[head] = sumf; +// } +// __syncthreads(); +// if (head == 0) { +// sumf = s[col]; +// sumf = warp_reduce_sum(sumf); +// if (col == 0) { +// dst[size_t(n_kv)*row_q + row_k] = sumf + (float)this_m[row_k]; +// } +// } +//} + +void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + auto op = ggml_unary_op(dst->op_params[0]); + GGML_ASSERT(op == GGML_UNARY_OP_RELU); + auto k = dst->src[0]; + auto q = dst->src[1]; + auto w = dst->src[2]; + auto m = dst->src[3]; + int n_top_k = dst->ne[0]; + int n_kv = k->ne[1]; + int n_head = q->ne[1]; + GGML_ASSERT(k->type == GGML_TYPE_F16 || ggml_is_quantized(k->type)); + GGML_ASSERT(k->ne[2] == 1 || k->ne[3] == 1); + GGML_ASSERT(k->ne[1] > n_top_k); + GGML_ASSERT(k->ne[1] == m->ne[0]); + GGML_ASSERT(k->ne[0] == q->ne[0]); + GGML_ASSERT(q->ne[2] == m->ne[1]); + GGML_ASSERT(q->ne[1] == w->ne[0]); + GGML_ASSERT(q->ne[2] == w->ne[1]); + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(w->type == GGML_TYPE_F32); + GGML_ASSERT(m->type == GGML_TYPE_F32 || m->type == GGML_TYPE_F16); + GGML_ASSERT(dst->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(w)); + + constexpr int k_block_size = 256; + + if (k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32) { + printf("%s: using alternative\n", __func__); + ggml_cuda_pool_alloc score (ctx.pool(), q->ne[2] * n_kv); + ggml_cuda_pool_alloc sorted(ctx.pool(), q->ne[2] * n_kv); + dim3 grid(n_kv, q->ne[2], 1); + if (m->type == GGML_TYPE_F16) { + k_fused_gemm_relu_mul_sum_rows<<>>((const half *)k->data, (const float *)q->data, + (const float *)w->data, (const half *)m->data, score.get(), k->ne[0], k->ne[1], q->ne[1], m->nb[1]); + } else { + k_fused_gemm_relu_mul_sum_rows<<>>((const half *)k->data, (const float *)q->data, + (const float *)w->data, (const float *)m->data, score.get(), k->ne[0], k->ne[1], q->ne[1], m->nb[1]); + } + CUDA_CHECK(cudaGetLastError()); + + argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], q->ne[2], GGML_SORT_ORDER_DESC, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + + k_copy_topk<<ne[2], k_block_size, 0, ctx.stream()>>>(sorted.get(), (int *)dst->data, k->ne[1], dst->ne[0]); + CUDA_CHECK(cudaGetLastError()); + + return; + } + + constexpr int64_t k_max_work_buffer_elements = 1 << 28; + + int max_rows = k_max_work_buffer_elements / n_kv / n_head; + if (max_rows < 1) max_rows = 1; + if (max_rows > q->ne[2]) max_rows = q->ne[2]; + + int nstep = (q->ne[2] + max_rows - 1)/max_rows; + + ggml_cuda_pool_alloc kq(ctx.pool(), int64_t(n_kv)*max_rows*n_head); + ggml_cuda_pool_alloc score(ctx.pool(), int64_t(n_kv)*max_rows); + ggml_cuda_pool_alloc sorted(ctx.pool(), int64_t(n_kv)*max_rows); + ggml_cuda_pool_alloc k_f32(ctx.pool()); + ggml_cuda_pool_alloc q_converted(ctx.pool()); + auto q_padded = GGML_PAD(q->ne[0], MATRIX_ROW_PADDING); + if (ggml_is_quantized(k->type)) { + auto nbytes_q = q->ne[1] * max_rows * sizeof(block_q8_1)/QK8_1; + nbytes_q += get_mmq_x_max_host(ggml_cuda_info().devices[ctx.device].cc)*sizeof(block_q8_1_mmq); + q_converted.alloc(nbytes_q); + } else { + k_f32.alloc(k->ne[0]*k->ne[1]); + auto to_fp32_cuda = ggml_get_to_fp32_cuda(k->type); + to_fp32_cuda(k->data, k_f32.get(), k->ne[1]*k->ne[0], 1, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + } + + for (int istep = 0; istep < nstep; ++istep) { + int first = istep*max_rows; + int last = std::min(first + max_rows, int(q->ne[2])); + int nrows = last - first; + auto q_data = (const char *)q->data + istep*max_rows*q->nb[2]; + auto m_data = (const char *)m->data + istep*max_rows*m->nb[1]; + if (ggml_is_quantized(k->type)) { + quantize_mmq_q8_1_cuda((const float *)q_data, q_converted.get(), q->ne[0], nrows, 1, q_padded, k->type, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + mmq_args args{(const char *)k->data, q_converted.get(), kq.get(), + k->ne[0], k->ne[1], int64_t(k->nb[1]), + q_padded, q->ne[1]*nrows, q->ne[1]*nrows, k->ne[1]}; + ggml_cuda_op_mul_mat_q(ctx, k->type, args); + CUDA_CHECK(cudaGetLastError()); + } else { + // I wonder if it makes sense to use CUBLAS. If we did simple dot products we could fuse the + // relu, mul, sum_rows all in one kernel, avoiding the k*q intermediate result. + const float alpha = 1.0f; + const float beta = 0.0f; + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasSgemm(ctx.cublas_handle(ctx.device), CUBLAS_OP_T, CUBLAS_OP_N, + k->ne[1], q->ne[1]*nrows, q->ne[0], + &alpha, k_f32.get(), k->ne[0], + (const float *)q_data, q->ne[0], + &beta, kq.get(), k->ne[1])); + } + // sum_row1(relu(kq)*w) + if (m->type == GGML_TYPE_F32) { + k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const float *)m_data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + } else { + k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const half *)m_data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + } + CUDA_CHECK(cudaGetLastError()); + + argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], nrows, GGML_SORT_ORDER_DESC, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + + k_copy_topk<<>>(sorted.get(), (int *)((char *)dst->data + first*dst->nb[1]), k->ne[1], dst->ne[0]); + CUDA_CHECK(cudaGetLastError()); + } + +} diff --git a/ggml/src/ggml-cuda/indexer_topk.cuh b/ggml/src/ggml-cuda/indexer_topk.cuh new file mode 100644 index 0000000000..4f919e9ae0 --- /dev/null +++ b/ggml/src/ggml-cuda/indexer_topk.cuh @@ -0,0 +1,8 @@ +// +// Copyright (C) 2024 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// +#include "common.cuh" + +void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * dst); From b2b1602fe434a678e5988af3b6c89f4433fbb042 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 8 Jul 2026 15:33:23 +0000 Subject: [PATCH 3/8] WIP --- ggml/src/ggml-cuda/indexer_topk.cu | 46 +++++++++++++++++++++++++++--- 1 file changed, 42 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index 373dc7236f..50f95ec559 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -4,8 +4,8 @@ #include "convert.cuh" #include "argsort.cuh" -template -static __global__ void k_fused_relu_mul_sum_rows(const float * __restrict__ kq, const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, const int ncols, const int nhead, size_t nbm) { +template +static __global__ void k_fused_relu_mul_sum_rows(const kq_t * __restrict__ kq, const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, const int ncols, const int nhead, size_t nbm) { const int row = blockIdx.x; const int col = threadIdx.x; @@ -17,7 +17,7 @@ static __global__ void k_fused_relu_mul_sum_rows(const float * __restrict__ kq, float sum = (float)this_m[i]; auto this_kq = kq + blockIdx.x * step; for (int head = 0; head < nhead; ++head) { - float relu = this_kq[i]; + float relu = (float)this_kq[i]; relu = relu > 0.0f ? relu : 0.0f; sum += relu * this_w[head]; this_kq += ncols; @@ -134,7 +134,7 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds constexpr int k_block_size = 256; - if (k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32) { + if (false && k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32) { printf("%s: using alternative\n", __func__); ggml_cuda_pool_alloc score (ctx.pool(), q->ne[2] * n_kv); ggml_cuda_pool_alloc sorted(ctx.pool(), q->ne[2] * n_kv); @@ -157,6 +157,44 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds return; } + if (k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32 && q->ne[2] <= 16) { + ggml_cuda_pool_alloc kq(ctx.pool(), int64_t(n_kv)*q->ne[2]*q->ne[1]); + ggml_cuda_pool_alloc score(ctx.pool(), int64_t(n_kv)*q->ne[2]); + ggml_cuda_pool_alloc sorted(ctx.pool(), int64_t(n_kv)*q->ne[2]); + ggml_cuda_pool_alloc q_f16(ctx.pool(), q->ne[0]*q->ne[1]*q->ne[2]); + + auto to_fp16_cuda = ggml_get_to_fp16_cuda(q->type); + to_fp16_cuda((const float *)q->data, q_f16.get(), q->ne[0]*q->ne[1]*q->ne[2], 1, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + + const half alpha = 1.0f; + const half beta = 0.0f; + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasGemmEx(ctx.cublas_handle(ctx.device), CUBLAS_OP_T, CUBLAS_OP_N, + k->ne[1], q->ne[1]*q->ne[2], q->ne[0], + &alpha, (const half *)k->data, CUDA_R_16F, k->ne[0], + q_f16.get(), CUDA_R_16F, q->ne[0], + &beta, kq.get(), CUDA_R_16F, k->ne[1], + CUBLAS_COMPUTE_16F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + + if (m->type == GGML_TYPE_F32) { + k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const float *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + } else { + k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const half *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + } + CUDA_CHECK(cudaGetLastError()); + + argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], q->ne[2], GGML_SORT_ORDER_DESC, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + + k_copy_topk<<ne[2], k_block_size, 0, ctx.stream()>>>(sorted.get(), (int *)dst->data, k->ne[1], dst->ne[0]); + CUDA_CHECK(cudaGetLastError()); + + return; + + } + constexpr int64_t k_max_work_buffer_elements = 1 << 28; int max_rows = k_max_work_buffer_elements / n_kv / n_head; From 98471e74e4fd4f88bbee5654cf625f600db997a5 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Thu, 9 Jul 2026 13:13:34 +0000 Subject: [PATCH 4/8] WIP --- ggml/src/ggml-cuda/indexer_topk.cu | 36 ++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index 50f95ec559..86bda056ee 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -26,6 +26,29 @@ static __global__ void k_fused_relu_mul_sum_rows(const kq_t * __restrict__ kq, c } } +template +static __global__ void k_fused_relu_mul_sum_rows_2(const kq_t * __restrict__ kq, const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, const int ncols, const int nhead, size_t nbm) { + const int row = blockIdx.x; + const int col = blockIdx.y*blockDim.x + threadIdx.x; + if (col >= ncols) { + return; + } + + int64_t step = ncols*nhead; + auto this_w = w + blockIdx.x*nhead; + auto this_m = (const mask_t *)((const char *)m + nbm*row); + + float sum = (float)this_m[col]; + auto this_kq = kq + row * step; + for (int head = 0; head < nhead; ++head) { + float relu = (float)this_kq[col]; + relu = relu > 0.0f ? relu : 0.0f; + sum += relu * this_w[head]; + this_kq += ncols; + } + dst[ncols*row + col] = sum; +} + static __global__ void k_copy_topk(const int * __restrict__ sorted, int * dst, const int ncols, const int n_top_k) { const int row = blockIdx.x; const int col = threadIdx.x; @@ -178,11 +201,20 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + int nblocks = (k->ne[1] + k_block_size - 1)/k_block_size; + dim3 grid(q->ne[2], nblocks, 1); if (m->type == GGML_TYPE_F32) { - k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const float *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + k_fused_relu_mul_sum_rows_2<<>>(kq.get(), (const float *)w->data, (const float *)m->data, + score.get(), k->ne[1], q->ne[1], m->nb[1]); } else { - k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const half *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + k_fused_relu_mul_sum_rows_2<<>>(kq.get(), (const float *)w->data, (const half *)m->data, + score.get(), k->ne[1], q->ne[1], m->nb[1]); } + //if (m->type == GGML_TYPE_F32) { + // k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const float *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + //} else { + // k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const half *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + //} CUDA_CHECK(cudaGetLastError()); argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], q->ne[2], GGML_SORT_ORDER_DESC, ctx.stream()); From d0c0955d9c40a55686fb544e68eff2017350aa7b Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Thu, 9 Jul 2026 14:58:35 +0000 Subject: [PATCH 5/8] This seems to work --- ggml/src/ggml-cuda/indexer_topk.cu | 174 ++++++++--------------------- 1 file changed, 47 insertions(+), 127 deletions(-) diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index 86bda056ee..1cb5e70f3b 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -59,78 +59,6 @@ static __global__ void k_copy_topk(const int * __restrict__ sorted, int * dst, c } } -template -static __global__ void k_fused_gemm_relu_mul_sum_rows(const half * __restrict__ k, const float * __restrict__ q, - const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, - int head_size, int n_kv, int n_head, size_t nbm) { - const int row_k = blockIdx.x; - const int row_q = blockIdx.y; - const int col = threadIdx.x; - - k += size_t(row_k) * head_size; - q += size_t(row_q) * head_size * n_head; - w += size_t(row_q) * n_head; - auto this_m = (const mask_t *)((const char *)m + nbm*row_q); - - auto k2 = (const half2 *)k; - auto q2 = (const float2 *)q; - float score = (float)this_m[row_k]; - for (int head = 0; head < n_head; ++head) { - half2 sum = {}; - for (int i = col; i < head_size/2; i += blockDim.x) { - auto qh = __float22half2_rn(q2[i]); - sum += k2[i] * qh; - } - float sumf = (float)(sum.x + sum.y); - sumf = warp_reduce_sum(sumf); - sumf = sumf > 0.0f ? sumf : 0.0f; - score += sumf * w[head]; - q2 += head_size/2; - } - if (col == 0) { - dst[size_t(n_kv)*row_q + row_k] = score; - } -} - -//template -//static __global__ void k_fused_gemm_relu_mul_sum_rows(const half * __restrict__ k, const float * __restrict__ q, -// const float * __restrict__ w, const mask_t * __restrict__ m, float * __restrict__ dst, -// int head_size, int n_kv, int n_head, size_t nbm) { -// const int row_k = blockIdx.x; -// const int row_q = blockIdx.y; -// const int col = threadIdx.x % WARP_SIZE; -// const int head = threadIdx.x / WARP_SIZE; -// -// k += size_t(row_k) * head_size; -// q += size_t(row_q) * head_size * n_head + head * head_size; -// w += size_t(row_q) * n_head; -// auto this_m = (const mask_t *)((const char *)m + nbm*row_q); -// -// __shared__ float s[32]; -// -// auto k2 = (const half2 *)k; -// auto q2 = (const float2 *)q; -// half2 sum = {}; -// for (int i = col; i < head_size/2; i += WARP_SIZE) { -// auto qh = __float22half2_rn(q2[i]); -// sum += k2[i] * qh; -// } -// float sumf = (float)(sum.x + sum.y); -// sumf = warp_reduce_sum(sumf); -// sumf = sumf > 0.0f ? sumf * w[head] : 0.0f; -// if (col == 0) { -// s[head] = sumf; -// } -// __syncthreads(); -// if (head == 0) { -// sumf = s[col]; -// sumf = warp_reduce_sum(sumf); -// if (col == 0) { -// dst[size_t(n_kv)*row_q + row_k] = sumf + (float)this_m[row_k]; -// } -// } -//} - void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { auto op = ggml_unary_op(dst->op_params[0]); GGML_ASSERT(op == GGML_UNARY_OP_RELU); @@ -157,71 +85,61 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds constexpr int k_block_size = 256; - if (false && k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32) { - printf("%s: using alternative\n", __func__); - ggml_cuda_pool_alloc score (ctx.pool(), q->ne[2] * n_kv); - ggml_cuda_pool_alloc sorted(ctx.pool(), q->ne[2] * n_kv); - dim3 grid(n_kv, q->ne[2], 1); - if (m->type == GGML_TYPE_F16) { - k_fused_gemm_relu_mul_sum_rows<<>>((const half *)k->data, (const float *)q->data, - (const float *)w->data, (const half *)m->data, score.get(), k->ne[0], k->ne[1], q->ne[1], m->nb[1]); - } else { - k_fused_gemm_relu_mul_sum_rows<<>>((const half *)k->data, (const float *)q->data, - (const float *)w->data, (const float *)m->data, score.get(), k->ne[0], k->ne[1], q->ne[1], m->nb[1]); - } - CUDA_CHECK(cudaGetLastError()); - - argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], q->ne[2], GGML_SORT_ORDER_DESC, ctx.stream()); - CUDA_CHECK(cudaGetLastError()); - - k_copy_topk<<ne[2], k_block_size, 0, ctx.stream()>>>(sorted.get(), (int *)dst->data, k->ne[1], dst->ne[0]); - CUDA_CHECK(cudaGetLastError()); + if (k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32) { + constexpr int k_max_rows = 16; + int max_rows = std::min(k_max_rows, q->ne[2]); + int nstep = (q->ne[2] + max_rows - 1)/max_rows; - return; - } - - if (k->type == GGML_TYPE_F16 && q->type == GGML_TYPE_F32 && q->ne[2] <= 16) { - ggml_cuda_pool_alloc kq(ctx.pool(), int64_t(n_kv)*q->ne[2]*q->ne[1]); - ggml_cuda_pool_alloc score(ctx.pool(), int64_t(n_kv)*q->ne[2]); - ggml_cuda_pool_alloc sorted(ctx.pool(), int64_t(n_kv)*q->ne[2]); - ggml_cuda_pool_alloc q_f16(ctx.pool(), q->ne[0]*q->ne[1]*q->ne[2]); + ggml_cuda_pool_alloc kq(ctx.pool(), int64_t(n_kv)*q->ne[1]*max_rows); + ggml_cuda_pool_alloc score(ctx.pool(), int64_t(n_kv)*max_rows); + ggml_cuda_pool_alloc sorted(ctx.pool(), int64_t(n_kv)*max_rows); + ggml_cuda_pool_alloc q_f16(ctx.pool(), q->ne[0]*q->ne[1]*max_rows); auto to_fp16_cuda = ggml_get_to_fp16_cuda(q->type); - to_fp16_cuda((const float *)q->data, q_f16.get(), q->ne[0]*q->ne[1]*q->ne[2], 1, ctx.stream()); - CUDA_CHECK(cudaGetLastError()); + GGML_ASSERT(to_fp16_cuda); const half alpha = 1.0f; const half beta = 0.0f; - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); - CUBLAS_CHECK(cublasGemmEx(ctx.cublas_handle(ctx.device), CUBLAS_OP_T, CUBLAS_OP_N, - k->ne[1], q->ne[1]*q->ne[2], q->ne[0], + + for (int istep = 0; istep < nstep; ++istep) { + int first_row = max_rows*istep; + int last_row = std::min(first_row + k_max_rows, int(q->ne[2])); + int nrows = last_row - first_row; + + to_fp16_cuda((const float *)q->data + q->ne[0]*q->ne[1]*first_row, q_f16.get(), q->ne[0]*q->ne[1]*nrows, 1, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(ctx.device), ctx.stream())); + CUBLAS_CHECK(cublasGemmEx(ctx.cublas_handle(ctx.device), CUBLAS_OP_T, CUBLAS_OP_N, + k->ne[1], q->ne[1]*nrows, q->ne[0], &alpha, (const half *)k->data, CUDA_R_16F, k->ne[0], q_f16.get(), CUDA_R_16F, q->ne[0], &beta, kq.get(), CUDA_R_16F, k->ne[1], CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - int nblocks = (k->ne[1] + k_block_size - 1)/k_block_size; - dim3 grid(q->ne[2], nblocks, 1); - if (m->type == GGML_TYPE_F32) { - k_fused_relu_mul_sum_rows_2<<>>(kq.get(), (const float *)w->data, (const float *)m->data, - score.get(), k->ne[1], q->ne[1], m->nb[1]); - } else { - k_fused_relu_mul_sum_rows_2<<>>(kq.get(), (const float *)w->data, (const half *)m->data, - score.get(), k->ne[1], q->ne[1], m->nb[1]); - } - //if (m->type == GGML_TYPE_F32) { - // k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const float *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); - //} else { - // k_fused_relu_mul_sum_rows<<ne[2], k_block_size, 0, ctx.stream()>>>(kq.get(), (const float *)w->data, (const half *)m->data, score.get(), k->ne[1], q->ne[1], m->nb[1]); - //} - CUDA_CHECK(cudaGetLastError()); + int nblocks = (k->ne[1] + k_block_size - 1)/k_block_size; + dim3 grid(nrows, nblocks, 1); + if (m->type == GGML_TYPE_F32) { + k_fused_relu_mul_sum_rows_2<<>>(kq.get(), + (const float *)w->data + first_row*q->ne[1], + (const float *)((const char *)m->data + first_row*m->nb[1]), + score.get(), k->ne[1], q->ne[1], m->nb[1]); + } else { + k_fused_relu_mul_sum_rows_2<<>>(kq.get(), + (const float *)w->data + first_row*q->ne[1], + (const half *)((const char *)m->data + first_row*m->nb[1]), + score.get(), k->ne[1], q->ne[1], m->nb[1]); + } + CUDA_CHECK(cudaGetLastError()); - argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], q->ne[2], GGML_SORT_ORDER_DESC, ctx.stream()); - CUDA_CHECK(cudaGetLastError()); + argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], max_rows, GGML_SORT_ORDER_DESC, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); - k_copy_topk<<ne[2], k_block_size, 0, ctx.stream()>>>(sorted.get(), (int *)dst->data, k->ne[1], dst->ne[0]); - CUDA_CHECK(cudaGetLastError()); + k_copy_topk<<>>(sorted.get(), + (int *)((char *)dst->data + first_row*dst->nb[1]), k->ne[1], dst->ne[0]); + CUDA_CHECK(cudaGetLastError()); + } return; @@ -278,18 +196,20 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds (const float *)q_data, q->ne[0], &beta, kq.get(), k->ne[1])); } - // sum_row1(relu(kq)*w) if (m->type == GGML_TYPE_F32) { - k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const float *)m_data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const float *)m_data, + score.get(), k->ne[1], q->ne[1], m->nb[1]); } else { - k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const half *)m_data, score.get(), k->ne[1], q->ne[1], m->nb[1]); + k_fused_relu_mul_sum_rows<<>>(kq.get(), (const float *)w->data, (const half *)m_data, + score.get(), k->ne[1], q->ne[1], m->nb[1]); } CUDA_CHECK(cudaGetLastError()); argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], nrows, GGML_SORT_ORDER_DESC, ctx.stream()); CUDA_CHECK(cudaGetLastError()); - k_copy_topk<<>>(sorted.get(), (int *)((char *)dst->data + first*dst->nb[1]), k->ne[1], dst->ne[0]); + k_copy_topk<<>>(sorted.get(), (int *)((char *)dst->data + first*dst->nb[1]), + k->ne[1], dst->ne[0]); CUDA_CHECK(cudaGetLastError()); } From a9f7234cce9ee1924317049f7a1f70dc6bd4fa37 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Thu, 9 Jul 2026 15:09:40 +0000 Subject: [PATCH 6/8] Minor --- src/graphs/build_deepseek2.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 4dab202eee..77c8a64f6c 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -496,7 +496,9 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer( cb(indexer_score, "dsa_indexer_score_sink", il); ggml_build_forward_expand(gf, indexer_score); } - return ggml_indexer_topk(ctx0, indexer_k_b, indexer_q, indexer_weights, indexer_score, GGML_UNARY_OP_RELU, n_top_k); + auto topk = ggml_indexer_topk(ctx0, indexer_k_b, indexer_q, indexer_weights, indexer_score, GGML_UNARY_OP_RELU, n_top_k); + ggml_build_forward_expand(gf, topk); + return topk; } if (indexer_q->ne[2] <= 8) { From ff27465bacce31a749063b0852588a657b19945c Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sat, 11 Jul 2026 06:52:20 +0000 Subject: [PATCH 7/8] Fix bug. Fix suggested by @sayap using GLM-5.2 --- ggml/src/ggml-cuda/indexer_topk.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index 1cb5e70f3b..f03ec23a5e 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -133,10 +133,10 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds } CUDA_CHECK(cudaGetLastError()); - argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], max_rows, GGML_SORT_ORDER_DESC, ctx.stream()); + argsort_f32_i32_cuda_cub(ctx.pool(), score.get(), sorted.get(), k->ne[1], nrows, GGML_SORT_ORDER_DESC, ctx.stream()); CUDA_CHECK(cudaGetLastError()); - k_copy_topk<<>>(sorted.get(), + k_copy_topk<<>>(sorted.get(), (int *)((char *)dst->data + first_row*dst->nb[1]), k->ne[1], dst->ne[0]); CUDA_CHECK(cudaGetLastError()); } From fdada9003ed6b9573eb7fa881dfdd1dbf155e982 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sat, 11 Jul 2026 15:38:14 +0000 Subject: [PATCH 8/8] GLM-DSA: much better PP long context performance (CUDA) --- ggml/src/ggml-cuda/dsa_attn.cu | 312 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/dsa_attn.cuh | 3 + ggml/src/ggml-cuda/fattn.cu | 8 +- 3 files changed, 322 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-cuda/dsa_attn.cu create mode 100644 ggml/src/ggml-cuda/dsa_attn.cuh diff --git a/ggml/src/ggml-cuda/dsa_attn.cu b/ggml/src/ggml-cuda/dsa_attn.cu new file mode 100644 index 0000000000..ba4639601e --- /dev/null +++ b/ggml/src/ggml-cuda/dsa_attn.cu @@ -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 +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><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 64: + soft_max_f16_simple<64, 64><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 128: + soft_max_f16_simple<128, 128><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 256: + soft_max_f16_simple<256, 256><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 512: + soft_max_f16_simple<512, 512><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 1024: + soft_max_f16_simple<1024, 1024><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 2048: + soft_max_f16_simple<2048, 1024><<>>(x, mask, ncols_x, nrows_y, scale); + break; + case 4096: + soft_max_f16_simple<4096, 1024><<>>(x, mask, ncols_x, nrows_y, scale); + break; + default: + soft_max_f16_simple<0, 0><<>>(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(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 q16(ctx.pool(), q_size); + ggml_cuda_pool_alloc kq16(ctx.pool(), kq_size); + ggml_cuda_pool_alloc kqv16(ctx.pool(), kqv_size); + ggml_cuda_pool_alloc mask16(ctx.pool(), mask_size); + ggml_cuda_pool_alloc k16(ctx.pool(), k_cache_size); + ggml_cuda_pool_alloc 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<<>>(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(first + max_rows, Q->ne[1]); + int nrows = last - first; + { + dim3 grid(indexer->ne[0], nrows, 1); + k_prepare_one_batch_kv<<>>(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<<>>(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<<>>(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<<>>(nelem, kqv16.get(), + (float *)((char *)dst->data + dst->nb[2]*first)); + + } + + } + + return true; +} diff --git a/ggml/src/ggml-cuda/dsa_attn.cuh b/ggml/src/ggml-cuda/dsa_attn.cuh new file mode 100644 index 0000000000..7fa2e8df01 --- /dev/null +++ b/ggml/src/ggml-cuda/dsa_attn.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 02b35c4d0c..cc9666be24 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -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 @@ -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]));