nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

CUDA: Fix non-F16 indexer top-k (#2272) 

* CUDA: Fix non-F16 indexer top-k

* CUDA: Handle strided F32 indexer keys

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing 2026-08-07 11:35:17 -05:00 committed by GitHub
parent 61b37f5f44
commit 40dffce685
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
ggml/src/ggml-cuda/indexer_topk.cu
View File

@ -70,7 +70,8 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
int n_kv = k->ne[1];
int n_head = q->ne[1];
//if (k->type != GGML_TYPE_F16 && !ggml_is_quantized(k->type)) printf("%s: K is %s?\n", __func__, ggml_type_name(k->type));
GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_F32 || ggml_is_quantized(k->type));
GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_BF16 ||
k->type == GGML_TYPE_F32 || 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]);
@ -159,16 +160,21 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
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());
const float * k_data = nullptr;
int k_ld = k->ne[0];
auto q_padded = GGML_PAD(q->ne[0], MATRIX_ROW_PADDING);
if (ggml_is_quantized(k->type)) {
auto nbytes_q = (size_t)(q_padded/QK8_1) * ((size_t) q->ne[1] * max_rows) * sizeof(block_q8_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 if (k->type == GGML_TYPE_F32) {
k_data = (const float *) k->data;
k_ld = k->nb[1]/sizeof(float);
} 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());
ggml_get_to_fp32_cuda(k->type)(k->data, k_f32.get(), k->ne[1]*k->ne[0], 1, ctx.stream());
CUDA_CHECK(cudaGetLastError());
k_data = k_f32.get();
}
for (int istep = 0; istep < nstep; ++istep) {
@ -177,6 +183,7 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
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];
auto w_data = (const float *)w->data + first*q->ne[1];
if (ggml_is_quantized(k->type)) {
quantize_mmq_q8_1_cuda((const float *)q_data, q_converted.get(), q->ne[0], q->ne[1]*nrows, 1, q_padded, k->type, ctx.stream());
CUDA_CHECK(cudaGetLastError());
@ -193,15 +200,15 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds
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],
&alpha, k_data, k_ld,
(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,
k_fused_relu_mul_sum_rows<<<nrows, k_block_size, 0, ctx.stream()>>>(kq.get(), 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,
k_fused_relu_mul_sum_rows<<<nrows, k_block_size, 0, ctx.stream()>>>(kq.get(), w_data, (const half *)m_data,
score.get(), k->ne[1], q->ne[1], m->nb[1]);
}
CUDA_CHECK(cudaGetLastError());