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(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);