Commit 5c3114268 for llama.cpp
commit 5c3114268229074e703be21ffddd4d8fb880c6b2
Author: Alexey Kopytko <alexey@kopytko.com>
Date: Sat Oct 24 02:49:08 2026 +0900
qwen4exp : build the QSA selection without the kq_mask (#30965)
* qwen4exp : build the QSA selection without the kq_mask
* clean
* lang
* qwen4exp : reshape the fused QSA selection to the kq_mask streams
ggml_top_k_mask with a single keep row returns [n_kv, n_tokens], while
build_attn_qsa expects the kq_mask layout [n_kv, n_tokens/n_stream, 1, n_stream].
With several streams the shape assert failed at graph reserve.
* test the QSA optimization
* reword
* update build_fused_top_k_mask instead of handwiring
* simplify
* keep il at the end for consistency
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 415dff33e..7e8b80a3a 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -1164,7 +1164,10 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
- mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+ // leave the mask unallocated if all attention layers have their own (e.g. qwen4exp QSA)
+ if (inp_attn->self_kq_mask->buffer) {
+ mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+ }
if (inp_attn->self_k_rot) {
mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot);
@@ -3386,6 +3389,7 @@ ggml_tensor * llm_graph_context::build_fused_top_k_mask(
ggml_tensor * score,
ggml_tensor * blocks,
ggml_tensor * extra,
+ ggml_tensor * mask,
int il) const {
const int64_t m1 = kq_mask->ne[1];
const int64_t m2 = kq_mask->ne[2];
@@ -3395,7 +3399,7 @@ ggml_tensor * llm_graph_context::build_fused_top_k_mask(
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));
+ ggml_tensor * cur = ggml_top_k_mask(ctx0, mask ? mask : kq_mask, rows(top_k), rows(score), blocks, rows(extra));
res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il});
return cur;
@@ -3435,7 +3439,7 @@ ggml_tensor * llm_graph_context::build_attn(
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);
+ kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, nullptr, il);
} else {
// prepare new kq mask - starts filled with -INFINITY
ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
diff --git a/src/llama-graph.h b/src/llama-graph.h
index a49f5432d..f56a305e9 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1337,12 +1337,14 @@ struct llm_graph_context {
int il) const;
// ggml_top_k_mask with top_k, score and extra viewed in the row layout of kq_mask, see ggml.h
+ // mask replaces kq_mask as the source of the kept values, it broadcasts to the rows
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
+ ggml_tensor * mask, // [n_kv], kq_mask if null
int il) const;
llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index 184585a9e..11699909f 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -684,7 +684,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_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);
+ ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, nullptr, il);
cb(kq_mask_top_k, name, il);
return kq_mask_top_k;
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index eb2e5bd59..395560990 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -975,7 +975,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
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);
+ ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, nullptr, il);
cb(sel, "indexer_sel", il);
return sel;
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index fc9acd43e..97e71d3d2 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -831,11 +831,13 @@ 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
+ // keep the cells of the live picked pools and the tail, -inf elsewhere
+ // all retained cells are in the token's sequence, so the kq_mask is not needed
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);
+ ggml_tensor * keep = ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, kq_mask->type, inp_kpool->n_kv), 0.0f);
+ ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs, keep, il);
cb(sel, "indexer_sel", il);
return sel;
@@ -891,7 +893,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
const size_t row = sel->nb[2];
sel = ggml_view_4d(ctx0, sel, n_kv, kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3],
row, row*kq_mask->ne[1], row*kq_mask->ne[1]*kq_mask->ne[2], 0);
- sel = ggml_add(ctx0, sel, kq_mask);
+ sel = ggml_cont(ctx0, sel);
cb(sel, "indexer_sel", il);
return sel;
@@ -936,8 +938,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
- // the selection already carries the causal mask, and it comes in the mask layout
- ggml_tensor * kq_mask = inp->get_kq_mask();
+ // the selection replaces the kq_mask and has the same shape
+ // it is not allocated without dense layers reading the kq_mask
+ const ggml_tensor * kq_mask = inp->get_kq_mask();
GGML_ASSERT(ggml_are_same_shape(sel, kq_mask));
cb(sel, "kq_mask_qsa", il);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 0fc5c0b5f..5716a2b01 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9141,9 +9141,10 @@ struct test_top_k_mask : public test_case {
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
+ const bool bcast; // one mask row for all rows and streams
std::string vars() override {
- return VARS_TO_STR11(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real);
+ return VARS_TO_STR12(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real, bcast);
}
// the result is a copy of the mask or -inf, the harness checks the -inf pattern, the finite values must match exactly
@@ -9168,12 +9169,13 @@ struct test_top_k_mask : public test_case {
}
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)
+ 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,
+ bool bcast = 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) {}
+ nb(nb), nh(nh), ns(ns), real(real), bcast(bcast) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
- ggml_tensor * m = ggml_new_tensor_4d(ctx, type, n_kv, nb, 1, ns);
+ ggml_tensor * m = bcast ? ggml_new_tensor_1d(ctx, type, n_kv) : 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);
@@ -10071,6 +10073,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
// 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));
+ // one mask row for all rows and streams (QSA)
+ test_cases.emplace_back(new test_top_k_mask(type, 1001, 4, 70, 16, 3, true, 7, 1, 1, false, true));
+ test_cases.emplace_back(new test_top_k_mask(type, 512, 8, 16, 32, 7, true, 5, 1, 4, false, true));
}
// glu ops