diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 415dff33e9..7e8b80a3a0 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 a49f5432d2..f56a305e93 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 184585a9e8..11699909ff 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 eb2e5bd598..3955609907 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 fc9acd43e3..97e71d3d21 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 0fc5c0b5f0..5716a2b010 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> 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