From e7268bf6a99688f893abf58295f49454cb2b4eef Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 14 Oct 2026 08:45:55 +0300 Subject: [PATCH] llama : build the k-pool selection in the causal mask layout (#30480) * llama : build the k-pool selection in the causal mask layout - the glm5-next indexer folded the causal mask in through a [n_kv, n_tokens] view: one reshape of the mask ahead of the add, and one of the result ahead of every layer reading the selection, including the shared-indexer layers - build the selection directly in the mask layout, as qwen4exp already does, and drop the now-degenerate reshape in its attention as well - assert the selection width matches the input it is taken from Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next * cont : remove comment --- src/models/glm5-next.cpp | 15 +++++++++------ src/models/qwen4exp.cpp | 8 ++++---- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 32199a260..a1b722950 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -988,6 +988,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( } } const int64_t n_sel = sel_idx->ne[0]; + GGML_ASSERT(n_sel == (int64_t) inp_kpool->n_sel); ggml_build_forward_expand(gf, sel_idx); @@ -1011,11 +1012,13 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( sel_idx = ggml_cast(ctx0, idx_f, GGML_TYPE_I32); ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1)); - sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0); // 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); - sel = ggml_add(ctx0, sel, ggml_reshape_2d(ctx0, kq_mask, n_kv, 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_add(ctx0, sel, kq_mask); cb(sel, "indexer_sel", il); return sel; @@ -1070,14 +1073,14 @@ ggml_tensor * 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. - ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); - cb(mask, "kq_mask_dsa", il); + // The scatter selection already includes the causal mask, and it comes in the mask layout + 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, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il); + 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); out = ggml_mul_mat(ctx0, layer.wo, out); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index f68a8622c..e055abdf5 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -905,16 +905,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 mask already carries the causal mask + // the selection already carries the causal mask, and it comes in the mask layout ggml_tensor * kq_mask = inp->get_kq_mask(); - ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); - cb(mask, "kq_mask_qsa", il); + 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); - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, mask, nullptr, nullptr, n_sel, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, sel, nullptr, nullptr, n_sel, kq_scale, il); cb(cur, "kqv_out", il); // the rotation is its own inverse, so undo it on the value side of the output