Commit d78f40d9d for llama.cpp
commit d78f40d9d34d7a72ecdcd2e854f75ff86afc8ca3
Author: Alexey Kopytko <alexey@kopytko.com>
Date: Tue Oct 27 18:54:36 2026 +0900
glm5-next : QSA without kq_mask (+ DSA) (#31009)
* glm5-next : QSA without kq_mask
* DSA: dsa top-k mask change tested with layers 0,3 of DeepSeek-V3.2
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 7309705dc..a49134e74 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -634,7 +634,10 @@ bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
- mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
+ // leave unallocated when the top-k layers use the indexer mask
+ if (self_kq_mask_mla->buffer) {
+ mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
+ }
mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
@@ -1218,7 +1221,10 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_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. glm5-next DSA)
+ if (inp_attn->self_kq_mask->buffer) {
+ mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+ }
const int64_t n_rs = mctx->get_recr()->get_n_rs();
@@ -3441,7 +3447,12 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
}
- const auto & kq_mask = inp->get_kq_mask_mla();
+ // top_k indexes the cells, and the indexer mask of the same type is built the same way
+ // let's use it, so when no other layer reads the MLA mask is not allocated
+ ggml_tensor * kq_mask = inp->get_kq_mask_mla();
+ if (inp->get_kq_mask_lid()->type == kq_mask->type) {
+ kq_mask = inp->get_kq_mask_lid();
+ }
ggml_tensor * kq_mask_top_k = nullptr;
if (cparams.fused_top_k_mask) {
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index d1e9b503d..46019d6aa 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -971,12 +971,14 @@ 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
+ // 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_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, nullptr, il);
+ 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);
cb(sel, "indexer_sel", il);
return sel;
@@ -1032,7 +1034,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
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;
@@ -1086,7 +1088,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 scatter selection already includes the causal mask, and it comes in the mask layout
+ // 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);