From e913807605f26e80feb2b5e9fe2b5c8e49ebd706 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sun, 12 Jul 2026 08:08:04 +0300 Subject: [PATCH] Fused indexer top_k (CUDA) (#2103) * WIP: indexer_topk on CUDA * Forgot these * WIP * WIP * This seems to work * Minor * Fix bug. Fix suggested by @sayap using GLM-5.2 --- ggml/src/ggml-cuda.cu | 5 + ggml/src/ggml-cuda/argsort.cu | 2 +- ggml/src/ggml-cuda/argsort.cuh | 9 ++ ggml/src/ggml-cuda/indexer_topk.cu | 216 ++++++++++++++++++++++++++++ ggml/src/ggml-cuda/indexer_topk.cuh | 8 ++ ggml/src/ggml-cuda/mmq.cu | 59 ++++---- ggml/src/ggml-cuda/mmq.cuh | 2 + src/graphs/build_deepseek2.cpp | 4 +- 8 files changed, 276 insertions(+), 29 deletions(-) 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.cu b/ggml/src/ggml-cuda.cu index dc2c86b1..b9a0d2f2 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 @@ -4129,6 +4130,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; } @@ -5032,6 +5036,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 2ef26b9a..0aaf5861 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 331f373b..7248d98e 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/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu new file mode 100644 index 00000000..f03ec23a --- /dev/null +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -0,0 +1,216 @@ +#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 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; + + 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 = (float)this_kq[i]; + relu = relu > 0.0f ? relu : 0.0f; + sum += relu * this_w[head]; + this_kq += ncols; + } + dst[ncols*row + i] = sum; + } +} + +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; + 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]; + } +} + +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) { + 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; + + 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); + GGML_ASSERT(to_fp16_cuda); + + const half alpha = 1.0f; + const half beta = 0.0f; + + 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(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], nrows, GGML_SORT_ORDER_DESC, ctx.stream()); + 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; + + } + + 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])); + } + 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 00000000..4f919e9a --- /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); diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 5991547e..65b87303 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 ee10d9ed..eeb5dbcc 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); diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 4dab202e..77c8a64f 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) {