Commit 84d9ce0fa for llama.cpp
commit 84d9ce0faddad002cd26a5a3fcd29f71fb426a1f
Author: lhez <lih@qti.qualcomm.com>
Date: Mon Nov 2 03:59:05 2026 -0800
opencl: add topk_mask (#31441)
diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt
index f905370a3..79fc37b60 100644
--- a/ggml/src/ggml-opencl/CMakeLists.txt
+++ b/ggml/src/ggml-opencl/CMakeLists.txt
@@ -67,6 +67,7 @@ set(GGML_OPENCL_KERNELS
argmax
argsort
top_k
+ top_k_mask
lightning_indexer
tri
fill
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index 6cfd0c0f9..97956bed1 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -873,6 +873,7 @@ struct ggml_backend_opencl_context {
ggml_cl_adreno_xmem_attn_state adreno_xmem_attn;
#endif
cl_kernel kernel_top_k_f32, kernel_top_k_merge_f32;
+ cl_kernel kernel_top_k_mask[2], kernel_top_k_mask_fill[2];
cl_kernel kernel_lightning_indexer;
cl_kernel kernel_get_rows_f32, kernel_get_rows_f16, kernel_get_rows_i32, kernel_get_rows_q4_0;
cl_kernel kernel_set_rows_f32_i64, kernel_set_rows_f32_i32, kernel_set_rows_f16_i64, kernel_set_rows_f16_i32;
@@ -2298,6 +2299,27 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
GGML_LOG_CONT(".");
}
+ // top_k_mask
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "top_k_mask.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("top_k_mask.cl");
+#endif
+ // Scores can contain NaN and -inf, so do not use the default fast-math options.
+ const std::string opts = "-cl-std=" + opencl_c_std +
+ (backend_ctx->adreno_use_large_buffer ? " -qcom-enable-large-buffer" : "");
+ for (int i = 0; i < 2; ++i) {
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), opts + " -DMASK_F16=" + std::to_string(i == 0));
+ CL_CHECK((backend_ctx->kernel_top_k_mask[i] = clCreateKernel(prog, "kernel_top_k_mask", &err), err));
+ CL_CHECK((backend_ctx->kernel_top_k_mask_fill[i] = clCreateKernel(prog, "kernel_top_k_mask_fill", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ }
+ GGML_LOG_CONT(".");
+ }
+
// get_rows
{
#ifdef GGML_OPENCL_EMBED_KERNELS
@@ -11800,6 +11822,14 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
op->src[2]->nb[0] == sizeof(float) && op->src[3]->nb[0] == sizeof(ggml_fp16_t);
case GGML_OP_TOP_K:
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_I32;
+ case GGML_OP_TOP_K_MASK:
+ return (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32) &&
+ op->src[0]->type == op->type &&
+ op->src[0]->nb[0] == ggml_type_size(op->type) && op->nb[0] == ggml_type_size(op->type) &&
+ op->src[1]->type == GGML_TYPE_I32 &&
+ (!op->src[2] || op->src[2]->type == GGML_TYPE_F32) &&
+ (!op->src[3] || op->src[3]->type == GGML_TYPE_I32) &&
+ (!op->src[4] || op->src[4]->type == GGML_TYPE_I32);
case GGML_OP_ARGSORT: {
load_cl_kernels_argsort(backend_ctx);
@@ -36968,6 +36998,119 @@ static void ggml_cl_top_k(ggml_backend_t backend, const ggml_tensor * src0, cons
}
}
+static void ggml_cl_top_k_mask(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ const ggml_tensor * score = dst->src[2];
+ const ggml_tensor * blocks = dst->src[3];
+ const ggml_tensor * extra = dst->src[4];
+ const int type = dst->type == GGML_TYPE_F16 ? 0 : 1;
+ cl_kernel kernel = backend_ctx->kernel_top_k_mask[type];
+
+ const ggml_tensor_extra_cl * extram = (const ggml_tensor_extra_cl *) src0->extra;
+ const ggml_tensor_extra_cl * extrak = (const ggml_tensor_extra_cl *) src1->extra;
+ const ggml_tensor_extra_cl * extras = score ? (const ggml_tensor_extra_cl *) score->extra : nullptr;
+ const ggml_tensor_extra_cl * extrab = blocks ? (const ggml_tensor_extra_cl *) blocks->extra : nullptr;
+ const ggml_tensor_extra_cl * extrae = extra ? (const ggml_tensor_extra_cl *) extra->extra : nullptr;
+ const ggml_tensor_extra_cl * extrad = (const ggml_tensor_extra_cl *) dst->extra;
+
+ const cl_mem buffer_s = extras ? extras->data_device : nullptr;
+ const cl_mem buffer_b = extrab ? extrab->data_device : nullptr;
+ const cl_mem buffer_e = extrae ? extrae->data_device : nullptr;
+
+ const cl_ulong offsetm = extram->offset + src0->view_offs;
+ const cl_ulong offsetk = extrak->offset + src1->view_offs;
+ const cl_ulong offsets = extras ? extras->offset + score->view_offs : 0;
+ const cl_ulong offsetb = extrab ? extrab->offset + blocks->view_offs : 0;
+ const cl_ulong offsete = extrae ? extrae->offset + extra->view_offs : 0;
+ const cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+ // single launch: fill == true, two launches: fill = false
+ // X2 benefits from a parallel fill for long single rows; mobile Adreno favors one launch.
+ const cl_int fill = !(backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E &&
+ ggml_nrows(dst) == 1 && dst->ne[0] >= 65536 && ggml_is_contiguous(dst));
+
+ if (!fill) {
+ GGML_ASSERT(ggml_is_contiguous(dst));
+
+ cl_kernel kernel_fill = backend_ctx->kernel_top_k_mask_fill[type];
+
+ const cl_ulong n = ggml_nelements(dst);
+ const size_t nv = 16/ggml_type_size(dst->type);
+ const size_t nth = MIN((size_t) 256, backend_ctx->get_kernel_workgroup_size(kernel_fill));
+
+ CL_CHECK(clSetKernelArg(kernel_fill, 0, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel_fill, 1, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel_fill, 2, sizeof(cl_ulong), &n));
+
+ size_t global[] = {((n + nv - 1)/nv + nth - 1)/nth*nth};
+ size_t local[] = {nth};
+ backend_ctx->enqueue_ndrange_kernel(kernel_fill, 1, global, local, dst);
+ }
+
+ GGML_TENSOR_LOCALS(cl_ulong, nem, src0, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbm, src0, nb);
+ GGML_TENSOR_LOCALS(cl_ulong, nek, src1, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbk, src1, nb);
+ GGML_TENSOR_LOCALS(cl_ulong, nes, score, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbs, score, nb);
+ GGML_TENSOR_LOCALS(cl_ulong, neb, blocks, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbb, blocks, nb);
+ GGML_TENSOR_LOCALS(cl_ulong, nee, extra, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbe, extra, nb);
+ GGML_TENSOR_LOCALS(cl_ulong, ned, dst, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nbd, dst, nb);
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extram->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offsetm));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extrak->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetk));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &buffer_s));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offsets));
+ CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &buffer_b));
+ CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offsetb));
+ CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_mem), &buffer_e));
+ CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &offsete));
+ CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nem0));
+ CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &nem1));
+ CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &nem2));
+ CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &nem3));
+ CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nbm1));
+ CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &nbm2));
+ CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &nbm3));
+ CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &nek0));
+ CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &nbk0));
+ CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &nbk1));
+ CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &nbk2));
+ CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &nbk3));
+ CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &nes0));
+ CL_CHECK(clSetKernelArg(kernel, 25, sizeof(cl_ulong), &nbs0));
+ CL_CHECK(clSetKernelArg(kernel, 26, sizeof(cl_ulong), &nbs1));
+ CL_CHECK(clSetKernelArg(kernel, 27, sizeof(cl_ulong), &nbs2));
+ CL_CHECK(clSetKernelArg(kernel, 28, sizeof(cl_ulong), &nbs3));
+ CL_CHECK(clSetKernelArg(kernel, 29, sizeof(cl_ulong), &neb0));
+ CL_CHECK(clSetKernelArg(kernel, 30, sizeof(cl_ulong), &neb1));
+ CL_CHECK(clSetKernelArg(kernel, 31, sizeof(cl_ulong), &nbb0));
+ CL_CHECK(clSetKernelArg(kernel, 32, sizeof(cl_ulong), &nbb1));
+ CL_CHECK(clSetKernelArg(kernel, 33, sizeof(cl_ulong), &nee0));
+ CL_CHECK(clSetKernelArg(kernel, 34, sizeof(cl_ulong), &nbe0));
+ CL_CHECK(clSetKernelArg(kernel, 35, sizeof(cl_ulong), &nbe1));
+ CL_CHECK(clSetKernelArg(kernel, 36, sizeof(cl_ulong), &nbe2));
+ CL_CHECK(clSetKernelArg(kernel, 37, sizeof(cl_ulong), &nbe3));
+ CL_CHECK(clSetKernelArg(kernel, 38, sizeof(cl_ulong), &ned1));
+ CL_CHECK(clSetKernelArg(kernel, 39, sizeof(cl_ulong), &ned2));
+ CL_CHECK(clSetKernelArg(kernel, 40, sizeof(cl_ulong), &nbd1));
+ CL_CHECK(clSetKernelArg(kernel, 41, sizeof(cl_ulong), &nbd2));
+ CL_CHECK(clSetKernelArg(kernel, 42, sizeof(cl_ulong), &nbd3));
+ CL_CHECK(clSetKernelArg(kernel, 43, sizeof(cl_int), &fill));
+
+ const size_t nth = MIN((size_t) 512, backend_ctx->get_kernel_workgroup_size(kernel));
+ size_t global[] = {nth, (size_t) ggml_nrows(dst)};
+ size_t local[] = {nth, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 2, global, local, dst);
+}
+
static void ggml_cl_argsort(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
@@ -37893,6 +38036,9 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
case GGML_OP_TOP_K:
func = ggml_cl_top_k;
break;
+ case GGML_OP_TOP_K_MASK:
+ func = ggml_cl_top_k_mask;
+ break;
case GGML_OP_ARGSORT:
if (!any_on_device) {
return false;
diff --git a/ggml/src/ggml-opencl/kernels/top_k_mask.cl b/ggml/src/ggml-opencl/kernels/top_k_mask.cl
new file mode 100644
index 000000000..00fc0fb13
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/top_k_mask.cl
@@ -0,0 +1,129 @@
+// Integer storage preserves the mask bits without FP16 conversion.
+#if MASK_F16
+#define MASK_T ushort
+#define MASK_V ushort8
+#define MASK_NV 8
+#define MASK_NINF 0xfc00
+#define MASK_STORE vstore8
+#else
+#define MASK_T uint
+#define MASK_V uint4
+#define MASK_NV 4
+#define MASK_NINF 0xff800000u
+#define MASK_STORE vstore4
+#endif
+
+kernel void kernel_top_k_mask_fill(
+ global char * dst,
+ ulong offsetd,
+ ulong n
+) {
+ global MASK_T * out = (global MASK_T *) (dst + offsetd);
+ const ulong c = get_global_id(0)*MASK_NV;
+ if (c + MASK_NV <= n) {
+ MASK_STORE((MASK_V) (MASK_NINF), 0, out + c);
+ } else {
+ for (ulong j = c; j < n; ++j) {
+ out[j] = MASK_NINF;
+ }
+ }
+}
+
+kernel void kernel_top_k_mask(
+ global const char * mask,
+ ulong offsetm,
+ global const char * top_k,
+ ulong offsetk,
+ global const char * score,
+ ulong offsets,
+ global const char * blocks,
+ ulong offsetb,
+ global const char * extra,
+ ulong offsete,
+ global char * dst,
+ ulong offsetd,
+ ulong nem0,
+ ulong nem1,
+ ulong nem2,
+ ulong nem3,
+ ulong nbm1,
+ ulong nbm2,
+ ulong nbm3,
+ ulong nek0,
+ ulong nbk0,
+ ulong nbk1,
+ ulong nbk2,
+ ulong nbk3,
+ ulong nes0,
+ ulong nbs0,
+ ulong nbs1,
+ ulong nbs2,
+ ulong nbs3,
+ ulong neb0,
+ ulong neb1,
+ ulong nbb0,
+ ulong nbb1,
+ ulong nee0,
+ ulong nbe0,
+ ulong nbe1,
+ ulong nbe2,
+ ulong nbe3,
+ ulong ned1,
+ ulong ned2,
+ ulong nbd1,
+ ulong nbd2,
+ ulong nbd3,
+ int fill
+) {
+ const ulong row = get_group_id(1);
+ const ulong i1 = row % ned1;
+ const ulong i2 = (row / ned1) % ned2;
+ const ulong i3 = row / (ned1*ned2);
+ const ulong tid = get_local_id(0);
+ const ulong nth = get_local_size(0);
+ const ulong n_pick = blocks ? neb1 : score ? nes0 : nem0;
+
+ global const MASK_T * mask_row = (global const MASK_T *) (mask + offsetm
+ + (i1 % nem1)*nbm1 + (i2 % nem2)*nbm2 + (i3 % nem3)*nbm3);
+ global MASK_T * dst_row = (global MASK_T *) (dst + offsetd + i1*nbd1 + i2*nbd2 + i3*nbd3);
+ global const char * top_k_row = top_k + offsetk + i1*nbk1 + i2*nbk2 + i3*nbk3;
+
+ if (fill) {
+ for (ulong c = tid; c < nem0/MASK_NV; c += nth) {
+ MASK_STORE((MASK_V) (MASK_NINF), c, dst_row);
+ }
+ for (ulong c = nem0/MASK_NV*MASK_NV + tid; c < nem0; c += nth) {
+ dst_row[c] = MASK_NINF;
+ }
+ barrier(CLK_GLOBAL_MEM_FENCE);
+ }
+
+ // Repeated cells copy the same mask value, as in the CUDA and Metal kernels.
+ for (ulong j = tid; j < nek0 + nee0; j += nth) {
+ if (j >= nek0) {
+ const int c = *(global const int *) (extra + offsete + (j - nek0)*nbe0 + i1*nbe1 + i2*nbe2 + i3*nbe3);
+ if (c >= 0 && (ulong) c < nem0) {
+ dst_row[c] = mask_row[c];
+ }
+ continue;
+ }
+
+ const int p = *(global const int *) (top_k_row + j*nbk0);
+ if (p < 0 || (ulong) p >= n_pick) {
+ continue;
+ }
+ if (score && !(*(global const float *) (score + offsets + p*nbs0 + i1*nbs1 + i2*nbs2 + i3*nbs3) > -INFINITY)) {
+ continue;
+ }
+ if (blocks) {
+ for (ulong k = 0; k < neb0; ++k) {
+ const int c = *(global const int *) (blocks + offsetb + k*nbb0 + p*nbb1);
+ if (c >= 0 && (ulong) c < nem0) {
+ dst_row[c] = mask_row[c];
+ }
+ }
+ } else if ((ulong) p < nem0) {
+ dst_row[p] = mask_row[p];
+ }
+ }
+}