Commit d454a204e for llama.cpp
commit d454a204ef2b63af357744809c05af897611a72e
Author: Aman Gupta <amangupta052@gmail.com>
Date: Thu Oct 22 01:02:22 2026 +0800
llama: create new fused op for QSA mask (#30640)
* llama: create new fused op for QSA scoring
* enroll glm5.3-next too
* tests : initialize the remaining tensors in test_kpool_mask
The harness sentinels were left uninitialized, so leftover memory could read as NaN and fail the NaN check.
* ggml: generalize the k-pool mask op into ggml_top_k_mask
The k-pool mask op only covered the pooled QSA selection of qwen4exp
and glm5-next, while the DSA models and the DeepSeek V4 CSA build the
same top-k mask by hand with fill, set_rows and add.
ggml_top_k_mask(mask, top_k, score, blocks, extra) keeps the mask
cells of the live picks and the extra cells, and sets the rest to
-inf. Every row tensor is 4D in the layout of top_k, and the mask is
broadcast to those rows, so a selection can also be made per head.
score, blocks and extra are optional: without score every pick is
live, without blocks a pick is a cell, and extra holds the cells a
row keeps without a pick.
llm_graph_context::build_fused_top_k_mask views the row tensors in
the layout of kq_mask and registers the fused node. The k-pool
models, the DSA attention of build_attn and the DeepSeek V4 CSA mask
use it when the probe enables the op, and keep their current path on
backends without it.
---------
Co-authored-by: Pascal <admin@serveurperso.com>
diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h
index 482bd3666..ecf0a3a6d 100644
--- a/ggml/include/ggml-rpc.h
+++ b/ggml/include/ggml-rpc.h
@@ -8,10 +8,10 @@ extern "C" {
#define RPC_PROTO_MAJOR_VERSION 8
#define RPC_PROTO_MINOR_VERSION 0
-#define RPC_PROTO_PATCH_VERSION 0
+#define RPC_PROTO_PATCH_VERSION 1
#ifdef __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
#endif
#define GGML_RPC_MAX_SERVERS 16
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index c37ba9471..3c2e6734b 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -585,6 +585,7 @@ extern "C" {
GGML_OP_DSV4_HC_COMB,
GGML_OP_DSV4_HC_PRE,
GGML_OP_DSV4_HC_POST,
+ GGML_OP_TOP_K_MASK,
GGML_OP_UNARY,
@@ -2732,6 +2733,26 @@ extern "C" {
struct ggml_tensor * post,
struct ggml_tensor * comb);
+ // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf
+ //
+ // mask: [n_kv, m1, m2, m3] f16 or f32, broadcast to the rows of top_k
+ // top_k: [n_top, ne1, ne2, ne3] i32, picks of each row
+ // score: [n_pick, ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live
+ // blocks: [n_blk, n_pick] i32, can be NULL: the member cells of each pick, else a pick is a cell
+ // extra: [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick
+ // res: [n_kv, ne1, ne2, ne3] same type as mask
+ //
+ // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf
+ // picks and cells out of range are skipped
+ //
+ GGML_API struct ggml_tensor * ggml_top_k_mask(
+ struct ggml_context * ctx,
+ struct ggml_tensor * mask,
+ struct ggml_tensor * top_k,
+ struct ggml_tensor * score,
+ struct ggml_tensor * blocks,
+ struct ggml_tensor * extra);
+
// custom operators
typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp
index edbe4da41..17f871c89 100644
--- a/ggml/src/ggml-backend-meta.cpp
+++ b/ggml/src/ggml-backend-meta.cpp
@@ -1055,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
} break;
case GGML_OP_DSV4_HC_COMB:
case GGML_OP_DSV4_HC_PRE:
- case GGML_OP_DSV4_HC_POST: {
+ case GGML_OP_DSV4_HC_POST:
+ case GGML_OP_TOP_K_MASK: {
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
} break;
case GGML_OP_UNARY: {
diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
index ab82e3d6c..4c33369ee 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.c
+++ b/ggml/src/ggml-cpu/ggml-cpu.c
@@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
{
ggml_compute_forward_dsv4_hc_post(params, tensor);
} break;
+ case GGML_OP_TOP_K_MASK:
+ {
+ ggml_compute_forward_top_k_mask(params, tensor);
+ } break;
case GGML_OP_MAP_CUSTOM1:
{
ggml_compute_forward_map_custom1(params, tensor);
@@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
case GGML_OP_DSV4_HC_COMB:
case GGML_OP_DSV4_HC_PRE:
case GGML_OP_DSV4_HC_POST:
+ case GGML_OP_TOP_K_MASK:
{
n_tasks = n_threads;
} break;
diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
index 7dff1d11b..02d5efffa 100644
--- a/ggml/src/ggml-cpu/ops.cpp
+++ b/ggml/src/ggml-cpu/ops.cpp
@@ -11736,6 +11736,92 @@ void ggml_compute_forward_dsv4_hc_post(
}
}
+// ggml_compute_forward_top_k_mask
+
+void ggml_compute_forward_top_k_mask(
+ const ggml_compute_params * params,
+ 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];
+
+ GGML_ASSERT(dst->type == mask->type);
+ GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
+
+ const int64_t n_kv = mask->ne[0];
+ const int64_t ne1 = dst->ne[1];
+ const int64_t ne2 = dst->ne[2];
+ const int64_t n_rows = ggml_nrows(dst);
+ const int64_t n_top = top_k->ne[0];
+ const int64_t n_blk = blocks ? blocks->ne[0] : 1;
+ const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+ const int64_t n_extra = extra ? extra->ne[0] : 0;
+
+ const size_t ts = ggml_type_size(dst->type);
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ for (int64_t ir = ith; ir < n_rows; ir += nth) {
+ const int64_t i1 = ir % ne1;
+ const int64_t i2 = (ir / ne1) % ne2;
+ const int64_t i3 = ir / (ne1*ne2);
+
+ const char * mask_row = (const char *) mask->data
+ + (i1 % mask->ne[1])*mask->nb[1] + (i2 % mask->ne[2])*mask->nb[2] + (i3 % mask->ne[3])*mask->nb[3];
+ char * dst_row = (char *) dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
+
+ if (dst->type == GGML_TYPE_F16) {
+ const ggml_fp16_t ninf = GGML_CPU_FP32_TO_FP16(-INFINITY);
+ for (int64_t c = 0; c < n_kv; ++c) {
+ ((ggml_fp16_t *) dst_row)[c] = ninf;
+ }
+ } else {
+ for (int64_t c = 0; c < n_kv; ++c) {
+ ((float *) dst_row)[c] = -INFINITY;
+ }
+ }
+
+ auto keep = [&](int32_t c) {
+ if (c >= 0 && c < n_kv) {
+ memcpy(dst_row + c*ts, mask_row + c*ts, ts);
+ }
+ };
+
+ const char * top_k_row = (const char *) top_k->data + i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3];
+
+ for (int64_t j = 0; j < n_top; ++j) {
+ const int32_t p = *(const int32_t *) (top_k_row + j*top_k->nb[0]);
+ if (p < 0 || p >= n_pick) {
+ continue;
+ }
+ if (score) {
+ const float s = *(const float *) ((const char *) score->data
+ + p*score->nb[0] + i1*score->nb[1] + i2*score->nb[2] + i3*score->nb[3]);
+ if (!(s > -INFINITY)) {
+ continue;
+ }
+ }
+ if (!blocks) {
+ keep(p);
+ continue;
+ }
+ for (int64_t k = 0; k < n_blk; ++k) {
+ keep(*(const int32_t *) ((const char *) blocks->data + k*blocks->nb[0] + p*blocks->nb[1]));
+ }
+ }
+
+ if (extra) {
+ const char * extra_row = (const char *) extra->data + i1*extra->nb[1] + i2*extra->nb[2] + i3*extra->nb[3];
+ for (int64_t j = 0; j < n_extra; ++j) {
+ keep(*(const int32_t *) (extra_row + j*extra->nb[0]));
+ }
+ }
+ }
+}
+
// ggml_compute_forward_rwkv_wkv7
static void ggml_compute_forward_rwkv_wkv7_f32(
diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h
index 2728b08b6..00a2ea7dc 100644
--- a/ggml/src/ggml-cpu/ops.h
+++ b/ggml/src/ggml-cpu/ops.h
@@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p
void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index f6ff9f893..e43807392 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -70,6 +70,7 @@
#include "ggml-cuda/cumsum.cuh"
#include "ggml-cuda/fill.cuh"
#include "ggml-cuda/lightning-indexer.cuh"
+#include "ggml-cuda/top-k-mask.cuh"
#include "ggml.h"
#include <algorithm>
@@ -2570,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_LIGHTNING_INDEXER:
ggml_cuda_lightning_indexer(ctx, dst);
break;
+ case GGML_OP_TOP_K_MASK:
+ ggml_cuda_op_top_k_mask(ctx, dst);
+ break;
default:
return false;
}
@@ -6003,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return true;
case GGML_OP_LIGHTNING_INDEXER:
return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+ case GGML_OP_TOP_K_MASK:
+ return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;
default:
return false;
diff --git a/ggml/src/ggml-cuda/top-k-mask.cu b/ggml/src/ggml-cuda/top-k-mask.cu
new file mode 100644
index 000000000..9ff995240
--- /dev/null
+++ b/ggml/src/ggml-cuda/top-k-mask.cu
@@ -0,0 +1,141 @@
+#include "top-k-mask.cuh"
+#include "convert.cuh"
+
+#define CUDA_TOP_K_MASK_BLOCK_SIZE 512
+
+// one block per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells
+template <typename T>
+static __global__ void top_k_mask_kernel(
+ const T * mask, const int32_t * top_k, const float * score, const int32_t * blocks, const int32_t * extra, T * dst,
+ const int64_t n_kv, const int64_t ne1, const int64_t ne2,
+ const int64_t m1, const int64_t m2, const int64_t m3,
+ const int64_t n_top, const int64_t n_pick, const int64_t n_blk, const int64_t n_extra,
+ const int64_t sm1, const int64_t sm2, const int64_t sm3,
+ const int64_t sd1, const int64_t sd2, const int64_t sd3,
+ const int64_t sk0, const int64_t sk1, const int64_t sk2, const int64_t sk3,
+ const int64_t ss0, const int64_t ss1, const int64_t ss2, const int64_t ss3,
+ const int64_t sb0, const int64_t sb1,
+ const int64_t se0, const int64_t se1, const int64_t se2, const int64_t se3,
+ const bool vec) {
+ ggml_cuda_pdl_lc();
+
+ const int64_t ir = blockIdx.x;
+ const int64_t i1 = ir % ne1;
+ const int64_t i2 = (ir / ne1) % ne2;
+ const int64_t i3 = ir / (ne1*ne2);
+
+ const T * mask_row = mask + (i1 % m1)*sm1 + (i2 % m2)*sm2 + (i3 % m3)*sm3;
+ T * dst_row = dst + i1*sd1 + i2*sd2 + i3*sd3;
+ const int32_t * top_k_row = top_k + i1*sk1 + i2*sk2 + i3*sk3;
+
+ ggml_cuda_pdl_sync();
+
+ const T ninf = ggml_cuda_cast<T>(-INFINITY);
+ if (vec) {
+ constexpr int nv = sizeof(int4)/sizeof(T);
+ T fill[nv];
+#pragma unroll
+ for (int j = 0; j < nv; ++j) {
+ fill[j] = ninf;
+ }
+ const int4 fill4 = *(const int4 *) fill;
+ for (int64_t i = threadIdx.x; i < n_kv/nv; i += blockDim.x) {
+ ((int4 *) dst_row)[i] = fill4;
+ }
+ } else {
+ for (int64_t i = threadIdx.x; i < n_kv; i += blockDim.x) {
+ dst_row[i] = ninf;
+ }
+ }
+
+ __syncthreads();
+
+ // one thread per pick or extra cell; cells kept twice get the same value
+ for (int64_t i = threadIdx.x; i < n_top + n_extra; i += blockDim.x) {
+ if (i >= n_top) {
+ const int32_t c = extra[(i - n_top)*se0 + i1*se1 + i2*se2 + i3*se3];
+ if (c >= 0 && c < n_kv) {
+ dst_row[c] = mask_row[c];
+ }
+ continue;
+ }
+
+ const int32_t p = top_k_row[i*sk0];
+ if (p < 0 || p >= n_pick) {
+ continue;
+ }
+ if (score && !(score[p*ss0 + i1*ss1 + i2*ss2 + i3*ss3] > -INFINITY)) {
+ continue;
+ }
+ if (!blocks) {
+ if (p < n_kv) {
+ dst_row[p] = mask_row[p];
+ }
+ continue;
+ }
+ for (int64_t k = 0; k < n_blk; ++k) {
+ const int32_t c = blocks[k*sb0 + p*sb1];
+ if (c >= 0 && c < n_kv) {
+ dst_row[c] = mask_row[c];
+ }
+ }
+ }
+}
+
+template <typename T>
+static void top_k_mask_cuda(ggml_backend_cuda_context & ctx, 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];
+
+ const int64_t n_kv = mask->ne[0];
+
+ GGML_ASSERT(dst->nb[0] == sizeof(T) && mask->nb[0] == sizeof(T));
+ GGML_ASSERT(dst->nb[1] % sizeof(T) == 0 && mask->nb[1] % sizeof(T) == 0);
+
+ // 16-byte stores need every dst row aligned
+ const bool vec = (n_kv*sizeof(T)) % sizeof(int4) == 0 &&
+ dst->nb[1] % sizeof(int4) == 0 && dst->nb[2] % sizeof(int4) == 0 && dst->nb[3] % sizeof(int4) == 0 &&
+ (uintptr_t) dst->data % sizeof(int4) == 0;
+
+ const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+
+ const size_t si = sizeof(int32_t);
+ const size_t sf = sizeof(float);
+
+ const dim3 block_dims(CUDA_TOP_K_MASK_BLOCK_SIZE, 1, 1);
+ const dim3 grid_dims(ggml_nrows(dst), 1, 1);
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+ ggml_cuda_kernel_launch(top_k_mask_kernel<T>, launch_params,
+ (const T *) mask->data, (const int32_t *) top_k->data,
+ score ? (const float *) score->data : nullptr,
+ blocks ? (const int32_t *) blocks->data : nullptr,
+ extra ? (const int32_t *) extra->data : nullptr,
+ (T *) dst->data,
+ n_kv, dst->ne[1], dst->ne[2],
+ mask->ne[1], mask->ne[2], mask->ne[3],
+ top_k->ne[0], n_pick, blocks ? blocks->ne[0] : 1, extra ? extra->ne[0] : 0,
+ mask->nb[1]/sizeof(T), mask->nb[2]/sizeof(T), mask->nb[3]/sizeof(T),
+ dst->nb[1]/sizeof(T), dst->nb[2]/sizeof(T), dst->nb[3]/sizeof(T),
+ top_k->nb[0]/si, top_k->nb[1]/si, top_k->nb[2]/si, top_k->nb[3]/si,
+ score ? score->nb[0]/sf : 0, score ? score->nb[1]/sf : 0, score ? score->nb[2]/sf : 0, score ? score->nb[3]/sf : 0,
+ blocks ? blocks->nb[0]/si : 0, blocks ? blocks->nb[1]/si : 0,
+ extra ? extra->nb[0]/si : 0, extra ? extra->nb[1]/si : 0, extra ? extra->nb[2]/si : 0, extra ? extra->nb[3]/si : 0,
+ vec);
+}
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ switch (dst->type) {
+ case GGML_TYPE_F16:
+ top_k_mask_cuda<half>(ctx, dst);
+ break;
+ case GGML_TYPE_F32:
+ top_k_mask_cuda<float>(ctx, dst);
+ break;
+ default:
+ GGML_ABORT("unsupported type");
+ }
+}
diff --git a/ggml/src/ggml-cuda/top-k-mask.cuh b/ggml/src/ggml-cuda/top-k-mask.cuh
new file mode 100644
index 000000000..930d6667b
--- /dev/null
+++ b/ggml/src/ggml-cuda/top-k-mask.cuh
@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 01dc8d05c..004970db0 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -1095,6 +1095,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"DSV4_HC_COMB",
"DSV4_HC_PRE",
"DSV4_HC_POST",
+ "TOP_K_MASK",
"UNARY",
@@ -1112,7 +1113,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"GLU",
};
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
@@ -1210,6 +1211,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"dsv4_hc_comb(mixes, scale, base)",
"dsv4_hc_pre(x, weights)",
"dsv4_hc_post(x, residual, post, comb)",
+ "top_k_mask(mask, top_k, score, blocks, extra)",
"unary(x)",
@@ -1227,7 +1229,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"glu(x)",
};
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
@@ -6630,6 +6632,48 @@ struct ggml_tensor * ggml_dsv4_hc_post(
return result;
}
+// ggml_top_k_mask
+
+struct ggml_tensor * ggml_top_k_mask(
+ struct ggml_context * ctx,
+ struct ggml_tensor * mask,
+ struct ggml_tensor * top_k,
+ struct ggml_tensor * score,
+ struct ggml_tensor * blocks,
+ struct ggml_tensor * extra) {
+ GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
+ GGML_ASSERT(mask->nb[0] == ggml_type_size(mask->type));
+ GGML_ASSERT(top_k->type == GGML_TYPE_I32);
+ GGML_ASSERT(top_k->ne[1] % mask->ne[1] == 0 && top_k->ne[2] % mask->ne[2] == 0 && top_k->ne[3] % mask->ne[3] == 0);
+
+ if (score) {
+ GGML_ASSERT(score->type == GGML_TYPE_F32);
+ GGML_ASSERT(score->ne[1] == top_k->ne[1] && score->ne[2] == top_k->ne[2] && score->ne[3] == top_k->ne[3]);
+ }
+
+ if (blocks) {
+ GGML_ASSERT(blocks->type == GGML_TYPE_I32);
+ GGML_ASSERT(blocks->ne[2] == 1 && blocks->ne[3] == 1);
+ GGML_ASSERT(!score || blocks->ne[1] == score->ne[0]);
+ }
+
+ if (extra) {
+ GGML_ASSERT(extra->type == GGML_TYPE_I32);
+ GGML_ASSERT(extra->ne[1] == top_k->ne[1] && extra->ne[2] == top_k->ne[2] && extra->ne[3] == top_k->ne[3]);
+ }
+
+ struct ggml_tensor * result = ggml_new_tensor_4d(ctx, mask->type, mask->ne[0], top_k->ne[1], top_k->ne[2], top_k->ne[3]);
+
+ result->op = GGML_OP_TOP_K_MASK;
+ result->src[0] = mask;
+ result->src[1] = top_k;
+ result->src[2] = score;
+ result->src[3] = blocks;
+ result->src[4] = extra;
+
+ return result;
+}
+
////////////////////////////////////////////////////////////////////////////////
struct ggml_hash_set ggml_hash_set_new(size_t size) {
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 9ee64b566..ca84fc405 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -165,6 +165,12 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {
/*.n_tokens_per_seq =*/ 1,
};
+static const llm_fused_op_probe llm_fused_op_top_k_mask_probe = {
+ /*.op =*/ LLM_FUSED_OP_TOP_K_MASK,
+ /*.name =*/ "fused top-k mask",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
llama_context::llama_context(
const llama_model & model,
llama_context_params params) :
@@ -328,6 +334,10 @@ llama_context::llama_context(
cparams.fused_dsv4_hc_post = true;
cparams.auto_fhc = true;
+ // only the indexer models build this mask, skip the probe for the others
+ cparams.fused_top_k_mask = true;
+ cparams.auto_ftkm = hparams.indexer_top_k > 0;
+
// with causal attention, the batch size is limited by the context size
cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
@@ -671,6 +681,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);
cparams.auto_fhc = false;
}
+
+ if (cparams.auto_ftkm) {
+ LLAMA_LOG_INFO("%s: resolving fused top-k mask support:\n", func);
+ resolve(llm_fused_op_top_k_mask_probe, cparams.fused_top_k_mask);
+ cparams.auto_ftkm = false;
+ }
}
static int llama_graph_n_input_tensors(ggml_cgraph * gf) {
diff --git a/src/llama-cparams.h b/src/llama-cparams.h
index f4eb181a7..0c260f7b6 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -48,6 +48,8 @@ struct llama_cparams {
bool fused_dsv4_hc_comb;
bool fused_dsv4_hc_post;
bool auto_fhc;
+ bool fused_top_k_mask; // use fused top-k sparse attention mask
+ bool auto_ftkm;
bool no_perf;
bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
bool op_offload;
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 59caffe87..8c31b5130 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3375,6 +3375,27 @@ ggml_tensor * llm_graph_context::build_attn(
return cur;
}
+ggml_tensor * llm_graph_context::build_fused_top_k_mask(
+ ggml_tensor * kq_mask,
+ ggml_tensor * top_k,
+ ggml_tensor * score,
+ ggml_tensor * blocks,
+ ggml_tensor * extra,
+ int il) const {
+ const int64_t m1 = kq_mask->ne[1];
+ const int64_t m2 = kq_mask->ne[2];
+ const int64_t m3 = kq_mask->ne[3];
+
+ auto rows = [&](ggml_tensor * t) {
+ return t ? ggml_reshape_4d(ctx0, t, t->ne[0], m1, m2, m3) : nullptr;
+ };
+
+ ggml_tensor * cur = ggml_top_k_mask(ctx0, kq_mask, rows(top_k), rows(score), blocks, rows(extra));
+ res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il});
+
+ return cur;
+}
+
ggml_tensor * llm_graph_context::build_attn(
llm_graph_input_attn_k_dsa * inp,
ggml_tensor * wo,
@@ -3407,31 +3428,36 @@ ggml_tensor * llm_graph_context::build_attn(
const auto & kq_mask = inp->get_kq_mask_mla();
- // prepare new kq mask - starts filled with -INFINITY
- ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
+ ggml_tensor * kq_mask_top_k = nullptr;
+ if (cparams.fused_top_k_mask) {
+ kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+ } else {
+ // prepare new kq mask - starts filled with -INFINITY
+ ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
- // reshape KQ mask into tensor with rows of size 1:
- // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
- kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
+ // reshape KQ mask into tensor with rows of size 1:
+ // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
+ kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
- // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
- ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
+ // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
+ ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
- // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
- // this will be our source of zero values for unmasking top k mask elements
- ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
- zeros = ggml_fill(ctx0, zeros, 0.0f);
+ // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
+ // this will be our source of zero values for unmasking top k mask elements
+ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
+ zeros = ggml_fill(ctx0, zeros, 0.0f);
- // modify KQ mask by unmasking elements that are in top_k indices
- // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
- ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
+ // modify KQ mask by unmasking elements that are in top_k indices
+ // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
+ kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
- // reshape to restore the original shape of KQ mask:
- // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
- kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
+ // reshape to restore the original shape of KQ mask:
+ // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
+ kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
- // combine with the original kq mask
- kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+ // combine with the original kq mask
+ kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+ }
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
diff --git a/src/llama-graph.h b/src/llama-graph.h
index abe5133c9..9c5f3aea9 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -53,6 +53,7 @@ enum llm_fused_op {
LLM_FUSED_OP_DSV4_HC_PRE,
LLM_FUSED_OP_DSV4_HC_COMB,
LLM_FUSED_OP_DSV4_HC_POST,
+ LLM_FUSED_OP_TOP_K_MASK,
};
enum llm_ffn_op_type : int {
@@ -1335,6 +1336,15 @@ struct llm_graph_context {
float kq_scale,
int il) const;
+ // ggml_top_k_mask with top_k, score and extra viewed in the row layout of kq_mask, see ggml.h
+ ggml_tensor * build_fused_top_k_mask(
+ ggml_tensor * kq_mask, // [n_kv, n_tokens/n_stream, 1, n_stream]
+ ggml_tensor * top_k, // [n_top, n_tokens]
+ ggml_tensor * score, // [n_pick, n_tokens], can be null
+ ggml_tensor * blocks, // [n_blk, n_pick], can be null
+ ggml_tensor * extra, // [n_extra, n_tokens], can be null
+ int il) const;
+
llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
llm_graph_input_dsv4 * build_inp_dsv4() const;
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index af04dddd3..184585a9e 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -683,6 +683,13 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask(
GGML_ASSERT(kq_mask);
GGML_ASSERT(top_k);
+ if (cparams.fused_top_k_mask) {
+ ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+ cb(kq_mask_top_k, name, il);
+
+ return kq_mask_top_k;
+ }
+
ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3],
kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 90870bd96..5056c901b 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -323,13 +323,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
inp->n_sel = (uint32_t) n_sel;
- // the scatter maps the dead slots of this mask to dump rows
- inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
- ggml_set_input(inp->sel_mask);
- // set_input_kpool always fills the mask, so it stays allocated in every graph
- ggml_build_forward_expand(gf, inp->sel_mask);
- // one view for all layers: each view of a host input is a separate upload
- inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+ // the scatter maps the dead slots of this mask to dump rows, the fused mask op reads liveness from the scores instead
+ if (!cparams.fused_top_k_mask) {
+ inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
+ ggml_set_input(inp->sel_mask);
+ // set_input_kpool always fills the mask, so it stays allocated in every graph
+ ggml_build_forward_expand(gf, inp->sel_mask);
+ // one view for all layers: each view of a host input is a separate upload
+ inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+ }
}
inp->n_new = n_new;
@@ -973,6 +975,17 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED
+ // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+ if (cparams.fused_top_k_mask) {
+ ggml_tensor * tail = hparams.indexer_kpool_select_tail ? inp_kpool->tail_idxs : nullptr;
+ GGML_ASSERT(kpool*n_top_pool + (tail ? tail->ne[0] : 0) == (int64_t) inp_kpool->n_sel);
+
+ ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, il);
+ cb(sel, "indexer_sel", il);
+
+ return sel;
+ }
+
// The selection mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
ggml_tensor * sel_score = ggml_get_rows(ctx0,
ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens]
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index 22e715984..ec0b438fa 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -813,6 +813,16 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], unordered
cb(top_k, "indexer_top_k", il);
+ // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+ if (cparams.fused_top_k_mask) {
+ GGML_ASSERT(kpool*n_top_pool + inp_kpool->tail_idxs->ne[0] == inp_kpool->n_sel);
+
+ ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs, il);
+ cb(sel, "indexer_sel", il);
+
+ return sel;
+ }
+
// the top blocks, then the incomplete tail with n_kv for missing cells
ggml_tensor * sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs,
ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens]
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 58f3275da..9f2d1f54d 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9128,6 +9128,148 @@ struct test_lightning_indexer : public test_case {
}
};
+// GGML_OP_TOP_K_MASK
+struct test_top_k_mask : public test_case {
+ const ggml_type type;
+ const int64_t n_kv;
+ const int64_t n_blk; // 0: no blocks, a pick is a cell
+ const int64_t n_pick; // number of blocks, ignored without blocks
+ const int64_t n_top;
+ const int64_t n_extra; // 0: no extra cells
+ const bool score; // false: every pick is live
+ const int64_t nb; // rows per stream
+ const int64_t nh; // rows of top_k per mask row, broadcast over ne2
+ const int64_t ns; // streams
+ const bool real; // blocks of consecutive cells, live picks, as in a real cache
+
+ std::string vars() override {
+ return VARS_TO_STR11(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real);
+ }
+
+ // the result is a copy of the mask or -inf, the harness checks the -inf pattern, the finite values must match exactly
+ double max_err() override {
+ return 0.0;
+ }
+
+ double err(const float * a, const float * b, size_t n) override {
+ double result = 0.0;
+ for (size_t i = 0; i < n; ++i) {
+ if (std::isfinite(a[i]) && std::isfinite(b[i])) {
+ result = std::max(result, (double) fabsf(a[i] - b[i]));
+ }
+ }
+ return result;
+ }
+
+ size_t op_size(ggml_tensor * t) override {
+ const size_t ts = ggml_type_size(type);
+ const size_t nc = std::max<int64_t>(n_blk, 1);
+ return ggml_nbytes(t) + nb*nh*ns*(n_top*(sizeof(int32_t) + (score ? sizeof(float) : 0) + nc*(sizeof(int32_t) + 2*ts)) + n_extra*(sizeof(int32_t) + 2*ts));
+ }
+
+ test_top_k_mask(ggml_type type = GGML_TYPE_F16, int64_t n_kv = 256, int64_t n_blk = 4, int64_t n_pick = 64,
+ int64_t n_top = 16, int64_t n_extra = 3, bool score = true, int64_t nb = 32, int64_t nh = 1, int64_t ns = 1, bool real = false)
+ : type(type), n_kv(n_kv), n_blk(n_blk), n_pick(n_blk ? n_pick : n_kv), n_top(n_top), n_extra(n_extra), score(score),
+ nb(nb), nh(nh), ns(ns), real(real) {}
+
+ ggml_tensor * build_graph(ggml_context * ctx) override {
+ ggml_tensor * m = ggml_new_tensor_4d(ctx, type, n_kv, nb, 1, ns);
+ ggml_set_name(m, "m");
+
+ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top, nb, nh, ns);
+ ggml_set_name(top_k, "top_k");
+
+ ggml_tensor * s = nullptr;
+ if (score) {
+ s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_pick, nb, nh, ns);
+ ggml_set_name(s, "score");
+ }
+
+ ggml_tensor * blocks = nullptr;
+ if (n_blk > 0) {
+ blocks = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_blk, n_pick);
+ ggml_set_name(blocks, "blocks");
+ }
+
+ ggml_tensor * extra = nullptr;
+ if (n_extra > 0) {
+ extra = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_extra, nb, nh, ns);
+ ggml_set_name(extra, "extra");
+ }
+
+ ggml_tensor * out = ggml_top_k_mask(ctx, m, top_k, s, blocks, extra);
+ ggml_set_name(out, "out");
+
+ return out;
+ }
+
+ void initialize_tensors(ggml_context * ctx) override {
+ std::mt19937 gen(0x4B50);
+
+ // cells and picks sometimes fall out of range, the op skips them
+ auto init_cells = [&](ggml_tensor * t) {
+ std::vector<int32_t> data(ggml_nelements(t));
+ if (real) {
+ // block p holds cells p*n_blk..p*n_blk + n_blk - 1, the extra cells are the last cells
+ const bool is_blocks = strcmp(t->name, "blocks") == 0;
+ for (int64_t i = 0; i < (int64_t) data.size(); ++i) {
+ const int64_t c = is_blocks ? i : n_kv - 1 - i % t->ne[0];
+ data[i] = c < n_kv ? (int32_t) c : (int32_t) n_kv;
+ }
+ } else {
+ std::uniform_int_distribution<int32_t> cell(0, n_kv - 1);
+ for (auto & c : data) {
+ const uint32_t r = gen() % 16;
+ c = r == 0 ? (int32_t) n_kv : r == 1 ? -1 : cell(gen);
+ }
+ }
+ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+ };
+
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ if (strcmp(t->name, "m") == 0) {
+ if (type == GGML_TYPE_F16) {
+ init_tensor_kq_mask(t);
+ } else {
+ std::vector<float> data(ggml_nelements(t));
+ std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+ for (auto & v : data) {
+ v = gen() % 4 == 0 ? -INFINITY : dis(gen);
+ }
+ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+ }
+ } else if (strcmp(t->name, "score") == 0) {
+ // -inf marks the picks a row cannot see
+ std::vector<float> data(ggml_nelements(t));
+ std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+ for (auto & v : data) {
+ v = !real && gen() % 4 == 0 ? -INFINITY : dis(gen);
+ }
+ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+ } else if (strcmp(t->name, "top_k") == 0) {
+ std::vector<int32_t> data(ggml_nelements(t));
+ std::vector<int32_t> order(n_pick);
+ for (int64_t r = 0; r < ggml_nrows(t); ++r) {
+ for (int64_t i = 0; i < n_pick; ++i) {
+ order[i] = i;
+ }
+ std::shuffle(order.begin(), order.end(), gen);
+ for (int64_t j = 0; j < n_top; ++j) {
+ const bool out = !real && j == 0 && r % 5 == 0;
+ data[r*n_top + j] = out ? (r % 2 ? -1 : (int32_t) n_pick) : order[j % n_pick];
+ }
+ }
+ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+ } else if (strcmp(t->name, "blocks") == 0 || strcmp(t->name, "extra") == 0) {
+ init_cells(t);
+ } else {
+ // the harness sentinels and the output, the callback checks the sentinels for NaN
+ init_tensor_uniform(t);
+ }
+ }
+ }
+};
+
// Deserializable generic test case
struct input_tensor {
ggml_type type;
@@ -9912,6 +10054,25 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true));
test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true));
+ for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+ // blocks with extra cells and scores (k-pool)
+ test_cases.emplace_back(new test_top_k_mask(type));
+ test_cases.emplace_back(new test_top_k_mask(type, 1001, 4, 70, 16, 3, true, 7));
+ test_cases.emplace_back(new test_top_k_mask(type, 256, 4, 64, 64, 0, true, 33));
+ test_cases.emplace_back(new test_top_k_mask(type, 512, 8, 16, 32, 7, true, 5, 1, 4));
+ test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 1));
+ test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 64, 1, 2));
+ // cells without scores (DSA)
+ test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 16, 0, false, 32));
+ test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 64, 0, false, 7, 1, 3));
+ test_cases.emplace_back(new test_top_k_mask(type, 4096, 0, 0, 2048, 0, false, 5));
+ // cells with scores and extra cells
+ test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 32, 5, true, 9));
+ // blocks picked per head, the mask broadcast over the heads
+ test_cases.emplace_back(new test_top_k_mask(type, 512, 16, 32, 4, 0, true, 6, 4, 1));
+ test_cases.emplace_back(new test_top_k_mask(type, 1024, 8, 128, 16, 8, false, 3, 2, 2));
+ }
+
// glu ops
for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {
for (int v : {0, 1}) {
@@ -12745,6 +12906,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
}
}
+ // top_k_mask, Flash-Next QSA shapes: 512 pools of 4 + a 3-cell tail, and DSA shapes: 2048 cells
+ for (int64_t n_kv : { 18432, 131072 }) {
+ for (int64_t nb : { 1, 2048 }) {
+ test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 4, n_kv/4, 512, 3, true, nb, 1, 1, true));
+ test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 0, 0, 2048, 0, false, nb, 1, 1, true));
+ }
+ }
+
// launch-overhead isolation: single L2_NORM launch vs batched siblings at the GDN
// production shape (strided qkv views) -- perf-mode only, the eval list has its own
// 2/4-wide coverage