Commit ae93877c8 for llama.cpp
commit ae93877c8f3bf9f39f0ee4ffb3ef1697c416a7c3
Author: Yonghye Kwon <developer.0hye@gmail.com>
Date: Fri Nov 13 22:02:36 2026 +0900
vulkan: add TOP_K_MASK (#31544)
* vulkan: add TOP_K_MASK
Assisted-by: Claude Opus 5.5
* vulkan: fix the TOP_K_MASK barrier and use the common offset helper
Assisted-by: Claude Opus 5.5
* vulkan: declare ggml_vk_top_k_mask in the common header
Assisted-by: Claude Opus 5.5
diff --git a/docs/ops.md b/docs/ops.md
index 3c2bcd239..c2ae5cbfb 100644
--- a/docs/ops.md
+++ b/docs/ops.md
@@ -126,7 +126,7 @@ Legend:
| TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-| TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
| UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
diff --git a/docs/ops/Vulkan.csv b/docs/ops/Vulkan.csv
index 65ac8bf33..dbd6412ef 100644
--- a/docs/ops/Vulkan.csv
+++ b/docs/ops/Vulkan.csv
@@ -19439,3 +19439,31 @@
"Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=q4_1","support","0","no","Vulkan"
"Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=q4_0","support","0","no","Vulkan"
"Vulkan0","LIGHTNING_INDEXER","hsk=128,nh=64,kv=256,nb=512,ns=4,nm=1,type_K=iq4_nl","support","0","no","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=4,n_pick=64,n_top=16,n_extra=3,score=1,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=4,n_pick=64,n_top=64,n_extra=0,score=1,nb=33,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=1,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=64,nh=1,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=256,n_blk=0,n_pick=256,n_top=16,n_extra=0,score=0,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=0,n_pick=1001,n_top=64,n_extra=0,score=0,nb=7,nh=1,ns=3,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=4096,n_blk=0,n_pick=4096,n_top=2048,n_extra=0,score=0,nb=5,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=0,n_pick=512,n_top=32,n_extra=5,score=1,nb=9,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=16,n_pick=32,n_top=4,n_extra=0,score=1,nb=6,nh=4,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1024,n_blk=8,n_pick=128,n_top=16,n_extra=8,score=0,nb=3,nh=2,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f16,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=4,n_pick=64,n_top=16,n_extra=3,score=1,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=4,n_pick=64,n_top=64,n_extra=0,score=1,nb=33,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=1,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=4,n_pick=1024,n_top=512,n_extra=3,score=1,nb=64,nh=1,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=256,n_blk=0,n_pick=256,n_top=16,n_extra=0,score=0,nb=32,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=0,n_pick=1001,n_top=64,n_extra=0,score=0,nb=7,nh=1,ns=3,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=4096,n_blk=0,n_pick=4096,n_top=2048,n_extra=0,score=0,nb=5,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=0,n_pick=512,n_top=32,n_extra=5,score=1,nb=9,nh=1,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=16,n_pick=32,n_top=4,n_extra=0,score=1,nb=6,nh=4,ns=1,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1024,n_blk=8,n_pick=128,n_top=16,n_extra=8,score=0,nb=3,nh=2,ns=2,real=0,bcast=0","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=1001,n_blk=4,n_pick=70,n_top=16,n_extra=3,score=1,nb=7,nh=1,ns=1,real=0,bcast=1","support","1","yes","Vulkan"
+"Vulkan0","TOP_K_MASK","type=f32,n_kv=512,n_blk=8,n_pick=16,n_top=32,n_extra=7,score=1,nb=5,nh=1,ns=4,real=0,bcast=1","support","1","yes","Vulkan"
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-common.h b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
index 4e05fb818..b74062931 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
@@ -109,6 +109,7 @@ void ggml_vk_quantize_q8_1(ggml_backend_vk_context * ctx, vk_context& subctx, co
void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst);
void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst);
void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst, const ggml_tensor * gate_scale_in = nullptr);
+void ggml_vk_top_k_mask(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst);
void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx);
bool ggml_vk_use_mul_mat_vec_id(const struct ggml_cgraph * cgraph, int node_idx);
void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx);
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
index 1ef00477b..caa852e28 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
@@ -219,6 +219,29 @@ struct vk_op_dsv4_hc_post_push_constants {
float gate_scale_out;
};
+struct vk_op_top_k_mask_push_constants {
+ uint32_t n_kv;
+ uint32_t n_top;
+ uint32_t n_pick;
+ uint32_t n_blk;
+ uint32_t n_extra;
+ uint32_t has_score;
+ uint32_t ninf;
+
+ uint32_t ne1; uint32_t ne2;
+
+ uint32_t m1; uint32_t m2; uint32_t m3;
+ uint32_t nbm1; uint32_t nbm2; uint32_t nbm3;
+ uint32_t nbd1; uint32_t nbd2; uint32_t nbd3;
+
+ uint32_t m_offset;
+ uint32_t k_offset;
+ uint32_t s_offset;
+ uint32_t b_offset;
+ uint32_t e_offset;
+ uint32_t d_offset;
+};
+
struct vk_op_count_experts_push_constants {
uint32_t ne00;
uint32_t ne01;
@@ -972,6 +995,18 @@ template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context *
p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
}
+// the fifth source, extra, comes from dst
+template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_top_k_mask_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
+ const ggml_tensor * extra = dst->src[4];
+
+ p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
+ p.k_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
+ p.s_offset = src2 ? get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type) : 0;
+ p.b_offset = src3 ? get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type) : 0;
+ p.e_offset = extra ? get_misalign_bytes(ctx, extra) / ggml_type_size(extra->type) : 0;
+ p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
+}
+
template <typename T> size_t push_constant_size(const T &t) {
static_assert(std::is_class<T>::value, "T must be a struct/class");
GGML_UNUSED(t);
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index 805963e94..e72afc6ae 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -977,6 +977,8 @@ struct vk_device_struct {
vk_pipeline pipeline_dsv4_hc_pre_gated_f32;
vk_pipeline pipeline_dsv4_hc_post_f32;
vk_pipeline pipeline_dsv4_hc_post_nocomb_f32;
+ vk_pipeline pipeline_top_k_mask_f16;
+ vk_pipeline pipeline_top_k_mask_f32;
std::map<vk_solve_tri_pipeline_state, vk_pipeline> pipeline_solve_tri_f32;
vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16;
vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16;
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 89761dc1b..785d6fa69 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -3823,6 +3823,10 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 1 }, 1);
ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_nocomb_f32,"dsv4_hc_post_nocomb_f32",dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 0 }, 1);
+ const uint32_t top_k_mask_wg = std::min({256u, device->properties.limits.maxComputeWorkGroupInvocations, device->properties.limits.maxComputeWorkGroupSize[0]});
+ ggml_vk_create_pipeline(device, device->pipeline_top_k_mask_f16, "top_k_mask_f16", top_k_mask_f16_len, top_k_mask_f16_data, "main", 6, sizeof(vk_op_top_k_mask_push_constants), {1, 1, 1}, {top_k_mask_wg}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_top_k_mask_f32, "top_k_mask_f32", top_k_mask_f32_len, top_k_mask_f32_data, "main", 6, sizeof(vk_op_top_k_mask_push_constants), {1, 1, 1}, {top_k_mask_wg}, 1);
+
for (auto &s : device->pipeline_solve_tri_f32) {
const vk_solve_tri_pipeline_state &state = s.first;
@@ -7625,6 +7629,46 @@ void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, con
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 });
}
+void ggml_vk_top_k_mask(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+ const ggml_tensor * mask = dst->src[0];
+ const ggml_tensor * top_k = dst->src[1];
+ const ggml_tensor * score = dst->src[2];
+ const ggml_tensor * blocks = dst->src[3];
+ const ggml_tensor * extra = dst->src[4];
+ VK_LOG_DEBUG("ggml_vk_top_k_mask(" << mask << ", " << top_k << ", " << score << ", " << blocks << ", " << extra << ", " << dst << ")");
+
+ vk_pipeline pipeline = dst->type == GGML_TYPE_F16 ? ctx->device->pipeline_top_k_mask_f16 : ctx->device->pipeline_top_k_mask_f32;
+ GGML_ASSERT(pipeline != nullptr);
+
+ ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+
+ const vk_subbuffer m_buf = ggml_vk_tensor_subbuffer(ctx, mask, true);
+ const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, top_k, true);
+ const vk_subbuffer s_buf = score ? ggml_vk_tensor_subbuffer(ctx, score, true) : k_buf;
+ const vk_subbuffer b_buf = blocks ? ggml_vk_tensor_subbuffer(ctx, blocks, true) : k_buf;
+ const vk_subbuffer e_buf = extra ? ggml_vk_tensor_subbuffer(ctx, extra, true) : k_buf;
+ const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
+
+ vk_op_top_k_mask_push_constants pc = {
+ (uint32_t)mask->ne[0],
+ (uint32_t)top_k->ne[0],
+ (uint32_t)(blocks ? blocks->ne[1] : score ? score->ne[0] : mask->ne[0]),
+ blocks ? (uint32_t)blocks->ne[0] : 0,
+ extra ? (uint32_t)extra->ne[0] : 0,
+ score ? 1u : 0u,
+ dst->type == GGML_TYPE_F16 ? 0xFC00u : 0xFF800000u,
+ (uint32_t)dst->ne[1], (uint32_t)dst->ne[2],
+ (uint32_t)mask->ne[1], (uint32_t)mask->ne[2], (uint32_t)mask->ne[3],
+ ggml_vk_nb_elem(mask, 1), ggml_vk_nb_elem(mask, 2), ggml_vk_nb_elem(mask, 3),
+ ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), ggml_vk_nb_elem(dst, 3),
+ 0, 0, 0, 0, 0, 0,
+ };
+ init_pushconst_tensor_offsets(ctx, pc, mask, top_k, score, blocks, dst);
+
+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { m_buf, k_buf, s_buf, b_buf, e_buf, d_buf }, pc,
+ { (uint32_t)dst->ne[1], (uint32_t)dst->ne[2], (uint32_t)dst->ne[3] });
+}
+
void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) {
ggml_tensor * dst = cgraph->nodes[node_idx];
ggml_tensor * src0 = dst->src[0];
@@ -13211,6 +13255,11 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in
break;
+ case GGML_OP_TOP_K_MASK:
+ ggml_vk_top_k_mask(ctx, compute_ctx, node);
+
+ break;
+
case GGML_OP_GATED_DELTA_NET:
ggml_vk_gated_delta_net(ctx, compute_ctx, node);
@@ -16488,6 +16537,30 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
}
return true;
}
+ case GGML_OP_TOP_K_MASK:
+ {
+ const ggml_tensor * mask = op->src[0];
+ if ((op->type != GGML_TYPE_F16 && op->type != GGML_TYPE_F32) || mask->type != op->type) {
+ return false;
+ }
+
+ // one workgroup per row
+ const auto & max_wg = device->properties.limits.maxComputeWorkGroupCount;
+ if (op->ne[1] > max_wg[0] || op->ne[2] > max_wg[1] || op->ne[3] > max_wg[2]) {
+ return false;
+ }
+
+ // the shader expects contiguous rows in mask and dst, and contiguous tensors in the other inputs
+ if (!ggml_is_contiguous_rows(op) || !ggml_is_contiguous_rows(mask)) {
+ return false;
+ }
+ for (int i = 1; i <= 4; ++i) {
+ if (op->src[i] && !ggml_is_contiguous(op->src[i])) {
+ return false;
+ }
+ }
+ return true;
+ }
case GGML_OP_GATED_DELTA_NET:
{
const uint32_t S_v = op->src[2]->ne[0];
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/top_k_mask.comp b/ggml/src/ggml-vulkan/vulkan-shaders/top_k_mask.comp
new file mode 100644
index 000000000..0e2c5b160
--- /dev/null
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/top_k_mask.comp
@@ -0,0 +1,95 @@
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_KHR_memory_scope_semantics : enable
+
+// One workgroup per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells.
+// The op only moves bits, so f16 needs no float16 arithmetic and the score check does not depend on float controls.
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 256;
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter
+{
+ uint n_kv;
+ uint n_top;
+ uint n_pick;
+ uint n_blk; // 0: a pick is a cell
+ uint n_extra;
+ uint has_score;
+ uint ninf; // bits of -inf: a 16-bit literal would need shaderInt16
+
+ uint ne1; uint ne2;
+
+ uint m1; uint m2; uint m3;
+ uint nbm1; uint nbm2; uint nbm3;
+ uint nbd1; uint nbd2; uint nbd3;
+
+ uint m_offset;
+ uint k_offset;
+ uint s_offset;
+ uint b_offset;
+ uint e_offset;
+ uint d_offset;
+};
+
+layout(binding = 0, std430) readonly buffer M { MASK_TYPE data_m[]; };
+layout(binding = 1, std430) readonly buffer K { int data_k[]; };
+layout(binding = 2, std430) readonly buffer S { uint data_s[]; };
+layout(binding = 3, std430) readonly buffer B { int data_b[]; };
+layout(binding = 4, std430) readonly buffer E { int data_e[]; };
+// coherent: the fill and the copy of a cell can come from different invocations
+layout(binding = 5, std430) coherent writeonly buffer D { MASK_TYPE data_d[]; };
+
+void keep(const int c, const uint m_row, const uint d_row) {
+ if (c >= 0 && uint(c) < n_kv) {
+ data_d[d_row + uint(c)] = data_m[m_row + uint(c)];
+ }
+}
+
+// score > -inf
+bool is_live(const uint s) {
+ return s != 0xFF800000u && (s & 0x7FFFFFFFu) <= 0x7F800000u;
+}
+
+void main() {
+ const uint tid = gl_LocalInvocationIndex;
+
+ const uint i1 = gl_WorkGroupID.x;
+ const uint i2 = gl_WorkGroupID.y;
+ const uint i3 = gl_WorkGroupID.z;
+
+ const uint row = i1 + ne1*(i2 + ne2*i3);
+ const uint m_row = m_offset + (i1 % m1)*nbm1 + (i2 % m2)*nbm2 + (i3 % m3)*nbm3;
+ const uint d_row = d_offset + i1*nbd1 + i2*nbd2 + i3*nbd3;
+
+ for (uint i = tid; i < n_kv; i += BLOCK_SIZE) {
+ data_d[d_row + i] = MASK_TYPE(ninf);
+ }
+
+ controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup, gl_StorageSemanticsBuffer, gl_SemanticsAcquireRelease);
+
+ // one invocation per pick or extra cell; cells kept twice get the same value
+ for (uint i = tid; i < n_top + n_extra; i += BLOCK_SIZE) {
+ if (i >= n_top) {
+ keep(data_e[e_offset + row*n_extra + (i - n_top)], m_row, d_row);
+ continue;
+ }
+
+ const int p = data_k[k_offset + row*n_top + i];
+ if (p < 0 || uint(p) >= n_pick) {
+ continue;
+ }
+ if (has_score != 0 && !is_live(data_s[s_offset + row*n_pick + uint(p)])) {
+ continue;
+ }
+ if (n_blk == 0) {
+ keep(p, m_row, d_row);
+ continue;
+ }
+ for (uint k = 0; k < n_blk; ++k) {
+ keep(data_b[b_offset + uint(p)*n_blk + k], m_row, d_row);
+ }
+ }
+}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
index 7ccccfa1d..d26357498 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
@@ -1120,6 +1120,8 @@ void process_shaders() {
string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {});
string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {});
string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {});
+ string_to_spv("top_k_mask_f16", "top_k_mask.comp", {{"MASK_TYPE", "uint16_t"}});
+ string_to_spv("top_k_mask_f32", "top_k_mask.comp", {{"MASK_TYPE", "uint"}});
string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));