From 949079535567992e17d9595c0e2765061eb04b49 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 30 Aug 2026 20:36:18 +0200 Subject: [PATCH] qwen4exp: attend the selected cells instead of masking the window QSA names about n_top_k cells per query, but build_attn_qsa turned that selection into a mask over the whole window and handed the full K and V to flash attention. On CUDA the only skip is a trailing cut, and the selection always keeps the current tail, so it never fires: the sparsity saved no work at all. The selected cells are now gathered into a window of their own, one per query, and the queries ride the stream axis of the attention so each one carries the window its own selection named. The mask comes from gathering the attention mask at those cells, which leaves the same values the mask path would put there. The scan reads the window once for all the queries of a stream while the gather moves the selected cells twice, so the two meet at 2*n_tps*width == n_kv. Below that margin the scan still wins, which keeps prompt processing on the existing path; only generation and short chunks take the gather. Flash attention is required, since the gather reads the value side as rows. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs: 132k context generation 56.7 -> 61.1 t/s Nothing at 55k, and the gain grows with depth, as the scan cost follows n_kv while the gather cost does not. Prompt processing is unchanged. Retrieval from a 132k token context is unaffected. --- src/models/models.h | 16 +++++++ src/models/qwen4exp.cpp | 95 ++++++++++++++++++++++++++++++++++++++--- 2 files changed, 104 insertions(+), 7 deletions(-) diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af..7180051bf60 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2320,6 +2320,22 @@ struct llama_model_qwen4exp : public llama_model_base { float kq_scale, int il); + // attention over the whole window, masked down to the selection + ggml_tensor * build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + + // attention over the selected cells alone, gathered into a window per query + ggml_tensor * build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, // so the layers sharing a ratio share one input set std::map qsa_inps; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 8f0e47b1fef..67fd63f2e59 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -712,6 +712,37 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_kv = mctx_cur->get_n_kv(); + + // the scan reads the window once for all the queries of a stream, the gather moves the + // selected cells twice, so the two meet at 2*n_tps*width == n_kv. the margin below keeps + // the win clear and the windows small enough to stay in the compute buffer of a decode + // graph. flash attention keeps the value side as rows, which is what the gather reads + ggml_tensor * cur = cparams.flash_attn && 4*n_tps*width < n_kv + ? build_qsa_gather(inp, q_cur, top_k, kq_scale, il) + : build_qsa_scan (inp, q_cur, top_k, kq_scale, il); + cb(cur, "kqv_out", il); + + // the rotation is its own inverse, so undo it on the value side of the output + if (inp->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + } + + return cur; +} + +// The window stays whole and the mask hides every cell the selection leaves out. +// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; + ggml_tensor * kq_mask = inp->get_kq_mask(); // prepare new kq mask - starts filled with -INFINITY @@ -744,15 +775,65 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( 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, 0, kq_scale, il); - cb(cur, "kqv_out", il); + return build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il); +} - // the rotation is its own inverse, so undo it on the value side of the output - if (inp->self_v_rot) { - cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); - } +// The selected cells are gathered into a window of their own, one per query, so the key and +// value traffic follows the budget instead of the whole cache. The queries then ride the stream +// axis of the attention: each one carries the window its own selection named. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; - return cur; + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_stream = top_k->ne[3]; + const int64_t n_q = n_tps*n_stream; + + ggml_tensor * k_all = mctx_cur->get_k(ctx0, il); + ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); + + const int64_t n_kv = k_all->ne[2]; + + // a cell holds its heads back to back, so one row of the gather is one whole cell + ggml_tensor * k_cells = ggml_view_3d(ctx0, k_all, k_all->ne[0]*k_all->ne[1], n_kv, n_stream, + k_all->nb[2], k_all->nb[3], 0); + ggml_tensor * v_cells = ggml_view_3d(ctx0, v_all, v_all->ne[0]*v_all->ne[1], n_kv, n_stream, + v_all->nb[2], v_all->nb[3], 0); + + // a cell index names a cell of its own stream, and the ubatch lays the queries of a stream + // out contiguously, so the flat index of a window is the index of its query + ggml_tensor * idx_stream = ggml_reshape_2d(ctx0, top_k, width*n_tps, n_stream); + + ggml_tensor * k_sel = ggml_get_rows(ctx0, k_cells, idx_stream); + ggml_tensor * v_sel = ggml_get_rows(ctx0, v_cells, idx_stream); + + k_sel = ggml_reshape_4d(ctx0, k_sel, k_all->ne[0], k_all->ne[1], width, n_q); + v_sel = ggml_reshape_4d(ctx0, v_sel, v_all->ne[0], v_all->ne[1], width, n_q); + cb(k_sel, "qsa_k_sel", il); + cb(v_sel, "qsa_v_sel", il); + + // gathering the attention mask at the selected cells leaves the same values the scan path + // would put there, so the window carries the reach of its query + ggml_tensor * kq_mask = inp->get_kq_mask(); + + GGML_ASSERT(kq_mask->nb[3] == kq_mask->nb[1]*n_tps); + + ggml_tensor * mask_cells = ggml_view_3d(ctx0, kq_mask, 1, n_kv, n_q, + kq_mask->nb[0], kq_mask->nb[1], 0); + + ggml_tensor * idx_query = ggml_reshape_3d(ctx0, top_k, width, n_q, 1); + + ggml_tensor * mask = ggml_get_rows(ctx0, mask_cells, idx_query); + + mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mask, width, 1, 1, n_q), GGML_TYPE_F16); + cb(mask, "qsa_mask_sel", il); + + return build_attn_mha(q_cur, k_sel, v_sel, nullptr, mask, nullptr, nullptr, 0, kq_scale, il); } ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(