diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 31832aa144..8fca8e1bc0 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2945,93 +2945,6 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } -ggml_tensor * llm_graph_context::build_attn( - llm_graph_input_attn_kv * inp, - ggml_tensor * wo, - ggml_tensor * wo_b, - ggml_tensor * wo_s, - ggml_tensor * q_cur, - ggml_tensor * k_cur, - ggml_tensor * v_cur, - ggml_tensor * kq_b, - ggml_tensor * sinks, - ggml_tensor * v_mla, - ggml_tensor * top_k, - float kq_scale, - int il) const { - GGML_ASSERT(v_mla == nullptr); - GGML_ASSERT(inp->self_k_rot == nullptr && inp->self_v_rot == nullptr); - - // these nodes are added to the graph together so that they are not reordered - // by doing so, the number of splits in the graph is reduced - // expand k later to enable rope fusion which directly writes into k-v cache - ggml_build_forward_expand(gf, q_cur); - ggml_build_forward_expand(gf, v_cur); - ggml_build_forward_expand(gf, k_cur); - - const auto * mctx_cur = inp->mctx; - - // store to KV cache - { - const auto & k_idxs = inp->get_k_idxs(); - const auto & v_idxs = inp->get_v_idxs(); - - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); - } - - ggml_tensor * kq_mask = build_attn_mask_top_k(inp->get_kq_mask(), top_k); - - 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, kq_b, kq_mask, sinks, v_mla, kq_scale, il); - cb(cur, "kqv_out", il); - - if (wo) { - cur = build_lora_mm(wo, cur, wo_s); - } - - if (wo_b) { - cur = ggml_add(ctx0, cur, wo_b); - } - - return cur; -} - -// Restrict a KQ mask to the positions named by top_k. Lifted verbatim from the DSA -// build_attn below so non-MLA sparse architectures can share it; node sequence unchanged. -ggml_tensor * llm_graph_context::build_attn_mask_top_k( - ggml_tensor * kq_mask, - ggml_tensor * top_k) const { - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); - - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - - // combine with the original kq mask - return ggml_add(ctx0, kq_mask_top_k, kq_mask); -} - ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_dsa * inp, ggml_tensor * wo, @@ -3064,7 +2977,31 @@ ggml_tensor * llm_graph_context::build_attn( const auto & kq_mask = inp->get_kq_mask_mla(); - ggml_tensor * kq_mask_top_k = build_attn_mask_top_k(kq_mask, top_k); + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); diff --git a/src/llama-graph.h b/src/llama-graph.h index deb0bd444d..b388e028cb 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1206,27 +1206,6 @@ struct llm_graph_context { float kq_scale, int il) const; - // as above, but attending only to positions named by top_k; for sparse attention - // architectures that are plain GQA rather than MLA - ggml_tensor * build_attn( - llm_graph_input_attn_kv * inp, - ggml_tensor * wo, - ggml_tensor * wo_b, - ggml_tensor * wo_s, - ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] - ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] - ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] - ggml_tensor * kq_b, - ggml_tensor * sinks, // [n_head_q] - ggml_tensor * v_mla, - ggml_tensor * top_k, // [n_top_k, n_tokens] - float kq_scale, - int il) const; - - ggml_tensor * build_attn_mask_top_k( - ggml_tensor * kq_mask, - ggml_tensor * top_k) const; - llm_graph_input_attn_k * build_attn_inp_k() const; ggml_tensor * build_attn( diff --git a/src/models/models.h b/src/models/models.h index b952790e8c..62e3fb47a7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2314,6 +2314,18 @@ struct llama_model_qwen4exp : public llama_model_base { int * sections, int il); + // dense self-attention restricted to the cells named by top_k. arch-local rather + // than a build_attn overload so that no shared attention path changes: the sparse + // MLA architectures keep their own copy of the same mask construction. + ggml_tensor * build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + // QSA: token indices this layer's queries may attend to, or nullptr for dense ggml_tensor * build_qsa_top_k( const llama_kv_cache_context * mctx_idx, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 275b1a3922..d2557390eb 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -514,13 +514,85 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); - // build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. + // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); cb(top_k, "indexer_top_k", il); return top_k; } +// Dense GQA self-attention restricted to the cells named by top_k. +// +// This is the plain-KV counterpart of the MLA sparse path in llm_graph_context::build_attn +// for llm_graph_input_attn_k_dsa, and the mask construction below is a copy of that one. +// It is kept here rather than factored into a shared helper so that the attention path used +// by every other architecture is untouched by this arch. +ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + GGML_ASSERT(inp->self_k_rot == nullptr && inp->self_v_rot == nullptr); + + // these nodes are added to the graph together so that they are not reordered + // by doing so, the number of splits in the graph is reduced + // expand k later to enable rope fusion which directly writes into k-v cache + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + + ggml_tensor * kq_mask = inp->get_kq_mask(); + + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + + 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, kq_mask_top_k, nullptr, nullptr, kq_scale, il); + cb(cur, "kqv_out", il); + + return cur; +} + ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( llm_graph_input_attn_kv * inp, const llama_kv_cache_context * mctx_idx, @@ -592,9 +664,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; if (top_k) { - cur = build_attn(inp, - nullptr, nullptr, nullptr, - Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, top_k, kq_scale, il); + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); } else { cur = build_attn(inp, nullptr, nullptr, nullptr,