LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Fused indexer top_k (CUDA) (#2103)

Browse Source 
* WIP: indexer_topk on CUDA

* Forgot these

* WIP

* WIP

* This seems to work

* Minor

* Fix bug. Fix suggested by @sayap using GLM-5.2
This commit is contained in:
Kawrakow
2026-07-12 08:08:04 +03:00
committed by GitHub
parent 3c6cbf6e2a
commit e913807605
8 changed files with 276 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda.cu
+5
View File
@@ -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 <algorithm>
#include <array>
@@ -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__)
ggml/src/ggml-cuda/argsort.cu
+1 -1
View File
@@ -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,
ggml/src/ggml-cuda/argsort.cuh
+9
View File
@@ -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);
ggml/src/ggml-cuda/indexer_topk.cu
+216
View File
@@ -0,0 +1,216 @@
#include "indexer_topk.cuh"
#include "mmq.cuh"
#include "quantize.cuh"
#include "convert.cuh"
#include "argsort.cuh"
template <typename kq_t, typename mask_t>
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 <typename kq_t, typename mask_t>
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<int>(k_max_rows, q->ne[2]);
int nstep = (q->ne[2] + max_rows - 1)/max_rows;
ggml_cuda_pool_alloc<half> kq(ctx.pool(), int64_t(n_kv)*q->ne[1]*max_rows);
ggml_cuda_pool_alloc<float> score(ctx.pool(), int64_t(n_kv)*max_rows);
ggml_cuda_pool_alloc<int> sorted(ctx.pool(), int64_t(n_kv)*max_rows);
ggml_cuda_pool_alloc<half> 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<<<grid, k_block_size, 0, ctx.stream()>>>(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<<<grid, k_block_size, 0, ctx.stream()>>>(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<<<nrows, k_block_size, 0, ctx.stream()>>>(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<float> kq(ctx.pool(), int64_t(n_kv)*max_rows*n_head);
ggml_cuda_pool_alloc<float> score(ctx.pool(), int64_t(n_kv)*max_rows);
ggml_cuda_pool_alloc<int> sorted(ctx.pool(), int64_t(n_kv)*max_rows);
ggml_cuda_pool_alloc<float> k_f32(ctx.pool());
ggml_cuda_pool_alloc<char> 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<<<nrows, 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<<<nrows, 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], nrows, GGML_SORT_ORDER_DESC, ctx.stream());
CUDA_CHECK(cudaGetLastError());
k_copy_topk<<<nrows, k_block_size, 0, ctx.stream()>>>(sorted.get(), (int *)((char *)dst->data + first*dst->nb[1]),
k->ne[1], dst->ne[0]);
CUDA_CHECK(cudaGetLastError());
}
}
ggml/src/ggml-cuda/indexer_topk.cuh
+8
View File
@@ -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);
ggml/src/ggml-cuda/mmq.cu
+32 -27
View File
@@ -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<GGML_TYPE_Q4_0>(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);
ggml/src/ggml-cuda/mmq.cuh
+2
View File
@@ -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);
src/graphs/build_deepseek2.cpp
+3 -1
View File
@@ -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) {