Commit 276378515 for llama.cpp
commit 276378515ceebee5d14bcc8536603dd5c87491a2
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Tue Oct 27 12:13:23 2026 +0200
metal : add GGML_OP_TOP_K_MASK (#31230)
* metal : add GGML_OP_TOP_K_MASK
Port the fused top-k mask from the CUDA backend: a grid-stride kernel fills the output with -inf, then one
threadgroup per row copies the mask values of the kept cells over it. Function constants drop the score,
blocks and extra paths that a given graph does not use.
Assisted-by: dsh:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB)
* metal : clean up top-k mask kernel formatting
Assisted-by: dsh:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB)
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index f3f5d4d77..25a8a3ef6 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -507,6 +507,47 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexe
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask(
+ ggml_metal_library_t lib,
+ const ggml_tensor * op,
+ bool fill,
+ bool has_score,
+ bool has_blocks,
+ bool has_extra,
+ bool vec) {
+ GGML_ASSERT(op->op == GGML_OP_TOP_K_MASK);
+
+ char base[256];
+ char name[256];
+
+ snprintf(base, 256, "kernel_top_k_mask%s_%s", fill ? "_fill" : "", ggml_type_name(op->type));
+
+ ggml_metal_pipeline_with_params res;
+
+ ggml_metal_cv_t cv = ggml_metal_cv_init();
+
+ if (fill) {
+ snprintf(name, 256, "%s_vec=%d", base, vec ? 1 : 0);
+
+ ggml_metal_cv_set_bool(cv, vec, FC_TOP_K_MASK + 3);
+ } else {
+ snprintf(name, 256, "%s_sc=%d_bl=%d_ex=%d", base, has_score ? 1 : 0, has_blocks ? 1 : 0, has_extra ? 1 : 0);
+
+ ggml_metal_cv_set_bool(cv, has_score, FC_TOP_K_MASK + 0);
+ ggml_metal_cv_set_bool(cv, has_blocks, FC_TOP_K_MASK + 1);
+ ggml_metal_cv_set_bool(cv, has_extra, FC_TOP_K_MASK + 2);
+ }
+
+ res = ggml_metal_library_get_pipeline(lib, name);
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+ }
+
+ ggml_metal_cv_free(cv);
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, const ggml_tensor * op) {
char name[256];
const char * base = nullptr;
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index 6d77c57ff..dd7cdb958 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -126,6 +126,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_ad
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask (ggml_metal_library_t lib, const struct ggml_tensor * op, bool fill, bool has_score, bool has_blocks, bool has_extra, bool vec);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t nc, bool use_silu);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs, int32_t nc, bool use_silu);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 31ffe0585..451f65299 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1835,6 +1835,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
ggml_is_contiguous_rows(op->src[1]) &&
ggml_is_contiguous_rows(op->src[2]) &&
(op->src[3] == NULL || ggml_is_contiguous_rows(op->src[3]));
+ 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->src[0]->type) &&
+ op->nb[0] == ggml_type_size(op->type) &&
+ op->src[1]->type == GGML_TYPE_I32 &&
+ (op->src[2] == NULL || op->src[2]->type == GGML_TYPE_F32) &&
+ (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_I32) &&
+ (op->src[4] == NULL || op->src[4]->type == GGML_TYPE_I32);
case GGML_OP_SSM_SCAN:
return has_simdgroup_reduction;
case GGML_OP_SSM_CONV:
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 9db266f44..6e566084a 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -125,6 +125,7 @@
#define FC_FLASH_ATTN_EXT_TENSOR 2200
#define FC_LIGHTNING_INDEXER 2200
#define FC_MUL_MV_MMA 2300
+#define FC_TOP_K_MASK 2400
// op-specific constants
#define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -1295,6 +1296,38 @@ typedef struct {
uint64_t nbm3;
} ggml_metal_kargs_lightning_indexer;
+typedef struct {
+ int32_t n_kv;
+ int32_t m1;
+ int32_t m2;
+ int32_t m3;
+ int32_t n_top;
+ int32_t n_pick;
+ int32_t n_blk;
+ int32_t n_extra;
+ int64_t n_elems;
+ uint64_t nb1;
+ uint64_t nb2;
+ uint64_t nb3;
+ uint64_t nbm1;
+ uint64_t nbm2;
+ uint64_t nbm3;
+ uint64_t nbk0;
+ uint64_t nbk1;
+ uint64_t nbk2;
+ uint64_t nbk3;
+ uint64_t nbs0;
+ uint64_t nbs1;
+ uint64_t nbs2;
+ uint64_t nbs3;
+ uint64_t nbb0;
+ uint64_t nbb1;
+ uint64_t nbe0;
+ uint64_t nbe1;
+ uint64_t nbe2;
+ uint64_t nbe3;
+} ggml_metal_kargs_top_k_mask;
+
typedef struct {
int32_t n_tokens;
int32_t n_iter;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 0cebbe5d6..14712c3e8 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -345,6 +345,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
{
n_fuse = ggml_metal_op_lightning_indexer(ctx, idx);
} break;
+ case GGML_OP_TOP_K_MASK:
+ {
+ n_fuse = ggml_metal_op_top_k_mask(ctx, idx);
+ } break;
case GGML_OP_DSV4_HC_COMB:
case GGML_OP_DSV4_HC_PRE:
case GGML_OP_DSV4_HC_POST:
@@ -1424,6 +1428,124 @@ int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) {
return 1;
}
+int ggml_metal_op_top_k_mask(ggml_metal_op_t ctx, int idx) {
+ ggml_tensor * op = ctx->node(idx);
+
+ ggml_metal_library_t lib = ctx->lib;
+ ggml_metal_encoder_t enc = ctx->enc;
+
+ const ggml_tensor * mask = op->src[0];
+ const ggml_tensor * top_k = op->src[1];
+ const ggml_tensor * score = op->src[2];
+ const ggml_tensor * blocks = op->src[3];
+ const ggml_tensor * extra = op->src[4];
+
+ GGML_ASSERT(op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32);
+ GGML_ASSERT(mask->type == op->type);
+ GGML_ASSERT(top_k->type == GGML_TYPE_I32);
+ GGML_ASSERT(!score || score->type == GGML_TYPE_F32);
+ GGML_ASSERT(!blocks || blocks->type == GGML_TYPE_I32);
+ GGML_ASSERT(!extra || extra->type == GGML_TYPE_I32);
+
+ const int32_t n_kv = mask->ne[0];
+ const int32_t n_top = (int32_t) top_k->ne[0];
+ const int32_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+ const int32_t n_blk = blocks ? blocks->ne[0] : 0;
+ const int32_t n_extra = extra ? extra->ne[0] : 0;
+
+ const ggml_metal_buffer_id bid = ggml_metal_get_buffer_id(op);
+
+ // the vector fill stores 4 elements at a time
+ const int32_t ts = ggml_type_size(op->type);
+ const int64_t n_elems = ggml_nelements(op);
+
+ const bool has_score = score != NULL;
+ const bool has_blocks = blocks != NULL;
+ const bool has_extra = extra != NULL;
+
+ const bool vec = n_elems % 4 == 0 && bid.offs % (4*ts) == 0;
+
+ ggml_metal_kargs_top_k_mask args = {
+ /*.n_kv =*/ n_kv,
+ /*.m1 =*/ (int32_t) mask->ne[1],
+ /*.m2 =*/ (int32_t) mask->ne[2],
+ /*.m3 =*/ (int32_t) mask->ne[3],
+ /*.n_top =*/ n_top,
+ /*.n_pick =*/ n_pick,
+ /*.n_blk =*/ n_blk,
+ /*.n_extra =*/ n_extra,
+ /*.n_elems =*/ n_elems,
+ /*.nb1 =*/ op->nb[1],
+ /*.nb2 =*/ op->nb[2],
+ /*.nb3 =*/ op->nb[3],
+ /*.nbm1 =*/ mask->nb[1],
+ /*.nbm2 =*/ mask->nb[2],
+ /*.nbm3 =*/ mask->nb[3],
+ /*.nbk0 =*/ top_k->nb[0],
+ /*.nbk1 =*/ top_k->nb[1],
+ /*.nbk2 =*/ top_k->nb[2],
+ /*.nbk3 =*/ top_k->nb[3],
+ /*.nbs0 =*/ score ? score->nb[0] : 0,
+ /*.nbs1 =*/ score ? score->nb[1] : 0,
+ /*.nbs2 =*/ score ? score->nb[2] : 0,
+ /*.nbs3 =*/ score ? score->nb[3] : 0,
+ /*.nbb0 =*/ blocks ? blocks->nb[0] : 0,
+ /*.nbb1 =*/ blocks ? blocks->nb[1] : 0,
+ /*.nbe0 =*/ extra ? extra->nb[0] : 0,
+ /*.nbe1 =*/ extra ? extra->nb[1] : 0,
+ /*.nbe2 =*/ extra ? extra->nb[2] : 0,
+ /*.nbe3 =*/ extra ? extra->nb[3] : 0,
+ };
+
+ // the fill covers the whole dst, the scatter overwrites the kept cells with the mask values
+ {
+ auto pipeline_fill = ggml_metal_library_get_pipeline_top_k_mask(lib, op, true, false, false, false, vec);
+
+ const int64_t n_fill = n_elems/(vec ? 4 : 1);
+ const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_fill);
+
+ int nth = 32; // SIMD width
+
+ while (nth < nth_max && nth < n_fill) {
+ nth *= 2;
+ }
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline_fill);
+ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer(enc, bid, 1);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, (n_fill + nth - 1)/nth, 1, 1, nth, 1, 1);
+
+ ggml_metal_op_concurrency_reset(ctx);
+ }
+
+ {
+ auto pipeline = ggml_metal_library_get_pipeline_top_k_mask(lib, op, false, has_score, has_blocks, has_extra, false);
+
+ const int32_t n_max = n_top + n_extra;
+ const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
+
+ int nth = 32; // SIMD width
+
+ while (nth < nth_max && nth < n_max) {
+ nth *= 2;
+ }
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline);
+ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mask), 1);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(top_k), 2);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(score ? score : mask), 3);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(blocks ? blocks : mask), 4);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(extra ? extra : mask), 5);
+ ggml_metal_encoder_set_buffer(enc, bid, 6);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], op->ne[2], op->ne[3], nth, 1, 1);
+ }
+
+ return 1;
+}
+
int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h
index e14a2dd5f..0a77d8d69 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.h
+++ b/ggml/src/ggml-metal/ggml-metal-ops.h
@@ -60,6 +60,7 @@ int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_top_k_mask (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx);
diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal
index 90d3fbe21..5dfd4a60c 100644
--- a/ggml/src/ggml-metal/kernels/misc.metal
+++ b/ggml/src/ggml-metal/kernels/misc.metal
@@ -906,3 +906,99 @@ kernel void kernel_dsv4_hc_post_f32(
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
}
}
+
+constant bool FC_top_k_mask_has_score [[function_constant(FC_TOP_K_MASK + 0)]];
+constant bool FC_top_k_mask_has_blocks [[function_constant(FC_TOP_K_MASK + 1)]];
+constant bool FC_top_k_mask_has_extra [[function_constant(FC_TOP_K_MASK + 2)]];
+constant bool FC_top_k_mask_vec [[function_constant(FC_TOP_K_MASK + 3)]];
+
+// the fill runs before the scatter, both are fully parallel over the whole dst and over the picks of one row
+template <typename T, typename T4>
+kernel void kernel_top_k_mask_fill(
+ constant ggml_metal_kargs_top_k_mask & args,
+ device char * dst,
+ uint tid[[thread_position_in_grid]],
+ uint ntot[[threads_per_grid]]) {
+ const T ninf = T(-INFINITY);
+
+ if (FC_top_k_mask_vec) {
+ device T4 * dst_4 = (device T4 *) dst;
+ for (int64_t i = tid; i < args.n_elems/4; i += ntot) {
+ dst_4[i] = T4(ninf);
+ }
+ } else {
+ device T * dst_t = (device T *) dst;
+ for (int64_t i = tid; i < args.n_elems; i += ntot) {
+ dst_t[i] = ninf;
+ }
+ }
+}
+
+typedef decltype(kernel_top_k_mask_fill<half, half4>) kernel_top_k_mask_fill_h_t;
+typedef decltype(kernel_top_k_mask_fill<float, float4>) kernel_top_k_mask_fill_f_t;
+
+template [[host_name("kernel_top_k_mask_fill_f16")]] kernel kernel_top_k_mask_fill_h_t kernel_top_k_mask_fill<half, half4>;
+template [[host_name("kernel_top_k_mask_fill_f32")]] kernel kernel_top_k_mask_fill_f_t kernel_top_k_mask_fill<float, float4>;
+
+// one threadgroup per row: copy the mask cells of the live picks and of the extra cells over the fill
+template <typename T, typename T4>
+kernel void kernel_top_k_mask(
+ constant ggml_metal_kargs_top_k_mask & args,
+ device const char * mask,
+ device const char * top_k,
+ device const char * score,
+ device const char * blocks,
+ device const char * extra,
+ device char * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ uint3 tpitg[[thread_position_in_threadgroup]],
+ uint3 ntg_max[[threads_per_threadgroup]]) {
+ const int i1 = tgpig.x;
+ const int i2 = tgpig.y;
+ const int i3 = tgpig.z;
+
+ const int tiitg = tpitg.x;
+ const int ntg = ntg_max.x;
+
+ device const char * mask_row = mask + (i1 % args.m1)*args.nbm1 + (i2 % args.m2)*args.nbm2 + (i3 % args.m3)*args.nbm3;
+ device char * dst_row = dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3;
+ device const char * top_k_row = top_k + i1*args.nbk1 + i2*args.nbk2 + i3*args.nbk3;
+
+ device const T * mask_row_t = (device const T *) mask_row;
+ device T * dst_row_t = (device T *) dst_row;
+
+ for (int i = tiitg; i < args.n_top; i += ntg) {
+ const int32_t p = *(device const int32_t *) (top_k_row + i*args.nbk0);
+ if (p < 0 || p >= args.n_pick) {
+ continue;
+ }
+ if (FC_top_k_mask_has_score && !(*(device const float *) (score + p*args.nbs0 + i1*args.nbs1 + i2*args.nbs2 + i3*args.nbs3) > -INFINITY)) {
+ continue;
+ }
+ if (FC_top_k_mask_has_blocks) {
+ FOR_UNROLL (int k = 0; k < args.n_blk; ++k) {
+ const int32_t c = *(device const int32_t *) (blocks + k*args.nbb0 + p*args.nbb1);
+ if (c >= 0 && c < args.n_kv) {
+ dst_row_t[c] = mask_row_t[c];
+ }
+ }
+ } else if (p < args.n_kv) {
+ dst_row_t[p] = mask_row_t[p];
+ }
+ }
+
+ if (FC_top_k_mask_has_extra) {
+ for (int i = tiitg; i < args.n_extra; i += ntg) {
+ const int32_t c = *(device const int32_t *) (extra + i*args.nbe0 + i1*args.nbe1 + i2*args.nbe2 + i3*args.nbe3);
+ if (c >= 0 && c < args.n_kv) {
+ dst_row_t[c] = mask_row_t[c];
+ }
+ }
+ }
+}
+
+typedef decltype(kernel_top_k_mask<half, half4>) kernel_top_k_mask_h_t;
+typedef decltype(kernel_top_k_mask<float, float4>) kernel_top_k_mask_f_t;
+
+template [[host_name("kernel_top_k_mask_f16")]] kernel kernel_top_k_mask_h_t kernel_top_k_mask<half, half4>;
+template [[host_name("kernel_top_k_mask_f32")]] kernel kernel_top_k_mask_f_t kernel_top_k_mask<float, float4>;