From b37189aae4d92cc4bcebf909b692b03f1c66174e Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 10 Aug 2026 18:46:02 +0200 Subject: [PATCH] Actually fix quantized indexer cache on CUDA (#2286) --- ggml/src/ggml-cuda/indexer_topk.cu | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/indexer_topk.cu b/ggml/src/ggml-cuda/indexer_topk.cu index ce89db37..7c8aa8ae 100644 --- a/ggml/src/ggml-cuda/indexer_topk.cu +++ b/ggml/src/ggml-cuda/indexer_topk.cu @@ -87,7 +87,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 ((k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q8_0) && q->type == GGML_TYPE_F32) { constexpr size_t k_max_buf_size = 1 << 28; size_t per_row = size_t(n_kv)*(q->ne[1]*sizeof(half) + sizeof(int) + sizeof(float)) + q->ne[0]*q->ne[1]*sizeof(half); int max_rows = (k_max_buf_size + per_row - 1)/per_row; @@ -98,6 +98,17 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds 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); + ggml_cuda_pool_alloc k_f16(ctx.pool()); + + auto k_data = (const half *)k->data; + if (k->type == GGML_TYPE_Q8_0) { + k_f16.alloc(k->ne[0]*k->ne[1]); + auto to_fp16_cuda = ggml_get_to_fp16_cuda(k->type); + GGML_ASSERT(to_fp16_cuda); + to_fp16_cuda(k->data, k_f16.get(), k->ne[0]*k->ne[1], 1, ctx.stream()); + CUDA_CHECK(cudaGetLastError()); + k_data = k_f16.get(); + } auto to_fp16_cuda = ggml_get_to_fp16_cuda(q->type); GGML_ASSERT(to_fp16_cuda); @@ -118,9 +129,9 @@ void ggml_cuda_op_indexer_topk(ggml_backend_cuda_context & ctx, ggml_tensor * ds 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], + &alpha, 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));