Commit 27b7a8a00 for llama.cpp
commit 27b7a8a009de2f5f75d8111b6171381c920df764
Author: Aman Gupta <amangupta052@gmail.com>
Date: Tue Oct 27 21:57:45 2026 +0800
llama : build the k-pool sparse selection independently of the kq_mask shape (#31243)
* llama : build the k-pool sparse selection independently of the kq_mask shape
qwen4exp (QSA) and glm5-next pick cache cells by index and no longer read the kq_mask values, but they still shaped the selection like the dense kq_mask. Build the selection as [n_kv, n_tokens] over the cache cells instead, and split it into the streams of the K view right before the attention.
This lets the dense mask change its layout (e.g. per-sequence row lists) without touching the sparse layers.
* improve comment
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* llama : drop the unused build_fused_top_k_mask args, review nits
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index a49134e74..a2d1c7def 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3398,20 +3398,10 @@ ggml_tensor * llm_graph_context::build_attn(
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,
- ggml_tensor * mask,
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];
+ top_k = ggml_reshape_4d(ctx0, top_k, top_k->ne[0], kq_mask->ne[1], kq_mask->ne[2], 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, mask ? mask : kq_mask, rows(top_k), rows(score), blocks, rows(extra));
+ ggml_tensor * cur = ggml_top_k_mask(ctx0, kq_mask, top_k, nullptr, nullptr, nullptr);
res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il});
return cur;
@@ -3456,7 +3446,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, nullptr, il);
+ kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, 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 f56a305e9..c1a7d36ee 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1336,15 +1336,10 @@ 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
- // mask replaces kq_mask as the source of the kept values, it broadcasts to the rows
+ // ggml_top_k_mask with top_k 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
- 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 4d58b2fd2..da51549e4 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -686,7 +686,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, nullptr, il);
+ ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, 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 46019d6aa..d0917ed8a 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -887,7 +887,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer(
// Scores pools of kpool consecutive tokens, expands the selected pools and the incomplete tail into an additive mask
ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
- ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer,
+ ggml_tensor * cur, ggml_tensor * qr, const llama_layer & layer,
const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il) {
const auto * mctx_lid = mctx_hyb->get_idx();
@@ -919,6 +919,9 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
auto kpool_cache = mctx_hyb->get_kpool_access(ctx0, il, n_embd_indexer);
const int64_t n_kv = mctx_lid->get_n_kv();
+ // flash attention requires an f16 mask
+ const ggml_type type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
// Pool the entries completed by this ubatch. The last one is a dummy when the ubatch completes none.
ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_idxs);
rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new);
@@ -977,8 +980,9 @@ 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 * keep = ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, kq_mask->type, n_kv), 0.0f);
- ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, keep, il);
+ ggml_tensor * keep = ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, type_mask, n_kv), 0.0f);
+ ggml_tensor * sel = ggml_top_k_mask(ctx0, keep, top_k, score, inp_kpool->pool_idxs, tail); // [n_kv, n_tokens]
+ res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, sel, il});
cb(sel, "indexer_sel", il);
return sel;
@@ -1008,12 +1012,12 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
ggml_build_forward_expand(gf, sel_idx);
- ggml_tensor * mask_all = ggml_new_tensor_4d(ctx0, kq_mask->type, n_kv + n_sel, 1, 1, 1);
+ ggml_tensor * mask_all = ggml_new_tensor_4d(ctx0, type_mask, n_kv + n_sel, 1, 1, 1);
mask_all = ggml_fill(ctx0, mask_all, -INFINITY);
mask_all = ggml_repeat_4d(ctx0, mask_all, n_kv + n_sel, n_tokens, 1, 1);
mask_all = ggml_reshape_3d(ctx0, mask_all, 1, n_kv + n_sel, n_tokens);
- ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, kq_mask->type, n_sel, 1, 1, 1);
+ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, type_mask, n_sel, 1, 1, 1);
zeros = ggml_fill(ctx0, zeros, 0.0f);
zeros = ggml_repeat_4d(ctx0, zeros, n_sel, n_tokens, 1, 1);
zeros = ggml_reshape_3d(ctx0, zeros, 1, n_sel, n_tokens);
@@ -1029,11 +1033,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1));
- // Fold causal visibility before shared-indexer reuse.
- GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1]*kq_mask->ne[2]*kq_mask->ne[3] == n_tokens);
- 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_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0);
sel = ggml_cont(ctx0, sel);
cb(sel, "indexer_sel", il);
@@ -1075,11 +1075,17 @@ llama_model_glm5_next::graph::build_dsa_layer(
q_absorbed = ggml_permute(ctx0, q_absorbed, 0, 2, 1, 3);
cb(q_absorbed, "q_absorbed", il);
- ggml_tensor * kq_mask = inp_attn->get_kq_mask();
+ ggml_tensor * k = mctx_mla->get_k(ctx0, il);
+ ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * sel = prev_sel;
if (il >= (int) hparams.n_layer() || hparams.is_indexer_full(il)) { // the NextN block always has a full indexer
- sel = build_kpool_select(cur, qr, kq_mask, layer, mctx_hyb, inp_kpool, il);
+ sel = build_kpool_select(cur, qr, layer, mctx_hyb, inp_kpool, il);
+
+ // the selection replaces the kq_mask: split it into the streams of the K view
+ const int64_t n_stream = k->ne[3];
+ GGML_ASSERT(sel->ne[0] == k->ne[2] && sel->ne[1] % n_stream == 0);
+ sel = ggml_reshape_4d(ctx0, sel, sel->ne[0], sel->ne[1]/n_stream, 1, n_stream);
} else {
GGML_ASSERT(sel != nullptr && "shared indexer layer must follow a full indexer layer");
}
@@ -1088,14 +1094,8 @@ llama_model_glm5_next::graph::build_dsa_layer(
ggml_build_forward_expand(gf, kv_cmpr);
ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, kv_cmpr, inp_attn->get_k_idxs(), il));
- // the selection replaces the kq_mask and has the same shape
- // with no layers reading it, the kq_mask isn't allocated
- GGML_ASSERT(ggml_are_same_shape(sel, kq_mask));
cb(sel, "kq_mask_dsa", il);
- ggml_tensor * k = mctx_mla->get_k(ctx0, il);
- ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
-
ggml_tensor * out = build_attn_mha(q_absorbed, k, v, nullptr, sel, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
cb(out, "kqv_out", il);
diff --git a/src/models/models.h b/src/models/models.h
index 5d9b6077a..8f2048a31 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -2792,13 +2792,12 @@ struct llama_model_qwen4exp : public llama_model_base {
// the QSA layers share one set of k-pool inputs, see llama_memory_hybrid_idx
llm_graph_input_kpool * build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb);
- // QSA: the additive mask [n_kv, n_tokens] of the top blocks and the tail, kq_mask included
+ // QSA: the additive mask [n_kv, n_tokens] of the top blocks and the tail
ggml_tensor * build_qsa_sel(
const llama_memory_hybrid_idx_context * mctx_hyb,
llm_graph_input_kpool * inp_kpool,
ggml_tensor * cur,
ggml_tensor * inp_pos,
- ggml_tensor * kq_mask,
int * sections,
int il);
@@ -3026,7 +3025,7 @@ struct llama_model_glm5_next : public llama_model_base {
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
- ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer,
+ ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, const llama_layer & layer,
const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il);
// attention output (cur) and the selection to carry to the next shared layer (sel)
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index f5a81f8f6..c3e2b010d 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -818,7 +818,6 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
llm_graph_input_kpool * inp_kpool,
ggml_tensor * cur,
ggml_tensor * inp_pos,
- ggml_tensor * kq_mask,
int * sections,
int il) {
const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
@@ -888,13 +887,17 @@ 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);
+ // flash attention requires an f16 mask
+ const ggml_type type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
// 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 * 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);
+ ggml_tensor * keep = ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, type_mask, inp_kpool->n_kv), 0.0f);
+ ggml_tensor * sel = ggml_top_k_mask(ctx0, keep, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs); // [n_kv, n_tokens]
+ res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, sel, il});
cb(sel, "indexer_sel", il);
return sel;
@@ -916,12 +919,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
// seeding from sel_idx ties the scatter storage lifetime to this layer
const int64_t n_kv = inp_kpool->n_kv;
- ggml_tensor * mask_all = ggml_new_tensor_4d(ctx0, kq_mask->type, n_kv + n_sel, 1, 1, 1);
+ ggml_tensor * mask_all = ggml_new_tensor_4d(ctx0, type_mask, n_kv + n_sel, 1, 1, 1);
mask_all = ggml_fill(ctx0, mask_all, -INFINITY);
mask_all = ggml_repeat_4d(ctx0, mask_all, n_kv + n_sel, n_tokens, 1, 1);
mask_all = ggml_reshape_3d(ctx0, mask_all, 1, n_kv + n_sel, n_tokens);
- ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, kq_mask->type, n_sel, 1, 1, 1);
+ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, type_mask, n_sel, 1, 1, 1);
zeros = ggml_fill(ctx0, zeros, 0.0f);
zeros = ggml_repeat_4d(ctx0, zeros, n_sel, n_tokens, 1, 1);
zeros = ggml_reshape_3d(ctx0, zeros, 1, n_sel, n_tokens);
@@ -946,10 +949,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1));
- GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1]*kq_mask->ne[2]*kq_mask->ne[3] == n_tokens);
- 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_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0);
sel = ggml_cont(ctx0, sel);
cb(sel, "indexer_sel", il);
@@ -995,16 +995,16 @@ 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 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);
-
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
+ // the selection replaces the kq_mask: split it into the streams of the K view
+ const int64_t n_stream = k->ne[3];
+ GGML_ASSERT(sel->ne[0] == k->ne[2] && sel->ne[1] % n_stream == 0);
+ sel = ggml_reshape_4d(ctx0, sel, sel->ne[0], sel->ne[1]/n_stream, 1, n_stream);
+ cb(sel, "kq_mask_qsa", il);
+
ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, sel, nullptr, nullptr, n_sel, kq_scale, il);
cb(cur, "kqv_out", il);
@@ -1030,7 +1030,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
// indexer reads the same block input as q/k/v; no cache or no ratio means dense
const bool qsa = inp_kpool != nullptr && hparams.dsv4_compress_ratios[il] > 0;
- ggml_tensor * sel = qsa ? build_qsa_sel(mctx_hyb, inp_kpool, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr;
+ ggml_tensor * sel = qsa ? build_qsa_sel(mctx_hyb, inp_kpool, cur, inp_pos, sections, il) : nullptr;
// Qwen3Next uses a single Q projection that outputs query + gate
ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ]