wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

qwen4exp : build the QSA selection without the kq_mask (#30965)

Browse Source 
* qwen4exp : build the QSA selection without the kq_mask

* clean

* lang

* qwen4exp : reshape the fused QSA selection to the kq_mask streams

ggml_top_k_mask with a single keep row returns [n_kv, n_tokens], while
build_attn_qsa expects the kq_mask layout [n_kv, n_tokens/n_stream, 1, n_stream].
With several streams the shape assert failed at graph reserve.

* test the QSA optimization

* reword

* update build_fused_top_k_mask instead of handwiring

* simplify

* keep il at the end for consistency
This commit is contained in:
Alexey Kopytko authored and GitHub committed 2026-10-23 19:49:08 +02:00
1 parent da2ee1e7a6
commit 5c31142682
6 files changed
+28 -14
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/llama-graph.cpp
+7 -3
View File
@@ -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);