From beed2f78ac42cf16710b763e6f3ba20665c6d233 Mon Sep 17 00:00:00 2001 From: Abdel Darwish <168705979+abdel-darwish-27@users.noreply.github.com> Date: Wed, 2 Sep 2026 07:48:52 +1000 Subject: [PATCH] qwen4exp : gather-based sparse attention for QSA decode The indexer used to pick the top 2048 KV cache entries, but those selections were turned into a mask over the full KV cache, so attention still ran across the entire context with unselected positions masked out. This meant the attention cost continued to grow with context length even though only ~2k tokens were actually being attended to. This patch instead gathers the selected keys and values into a compact buffer and runs regular dense attention over the gathered set. The attention mask is derived from the existing per-cell bias values, so this does not require any new model inputs. Prompt processing and batched inference are unchanged and continue to use the existing masked path. QWEN4EXP_QSA_GATHER=0 disables the gather path at runtime. On dual RTX A6000, IQ4_XS, q8_0 KV cache: 130k context decode 15.7 up to 23.6 tok/s (+50%), 62k +19%, 31k +6%. --- src/llama-graph.cpp | 6 ++- src/models/models.h | 7 ++- src/models/qwen4exp.cpp | 108 +++++++++++++++++++++++++++++++++++++--- 3 files changed, 112 insertions(+), 9 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8ea441f441d7..9c8b5f56d82b 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1090,7 +1090,11 @@ 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); + // qwen4exp's QSA gather graphs never reference the mask, so it has no buffer; + // the same guard the other attention inputs carry + if (inp_attn->self_kq_mask && 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); diff --git a/src/models/models.h b/src/models/models.h index 93a6b34945de..1985b9917a16 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2380,8 +2380,10 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * qsa_bias, float kq_scale, - int il); + int il, + bool gather = false); // 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 @@ -2394,7 +2396,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il); + int il, + bool gather = false); ggml_tensor * build_layer_attn_linear( llm_graph_input_rs * inp, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 1484c9b07bda..996a0c290c03 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -528,7 +528,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il) { + int il, + bool gather) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); const int64_t idx_dim = hparams.indexer_head_size; @@ -549,7 +550,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there - const bool blk_bias = kq_mask != nullptr && + // the gather path needs the per-cell bias: it carries the full visibility of every + // cell (foreign, future, unpooled, tail), so the gathered cells' bias values can + // serve directly as the attention mask over the gathered set + const bool blk_bias = !gather && kq_mask != nullptr && kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && cparams.causal_attn && !hparams.use_alibi; @@ -662,7 +666,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail - const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + // the gather path pads the width to a multiple of 256 so the gathered K/V satisfy flash + // attention's padding without extra ops; surplus cells arrive with a -inf bias and are + // masked out of the gathered attention + const int64_t width = gather + ? std::min(n_kv, GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256)) + : std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); @@ -681,8 +690,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * qsa_bias, float kq_scale, - int il) { + int il, + bool gather) { // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer // has already scored with its own query in build_qsa_top_k, so top_k is unaffected. if (inp->self_k_rot) { @@ -712,6 +723,70 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } + // decode-time sparse gather: instead of masking the full cache (which makes flash + // attention scan all n_kv cells to use only indexer_top_k of them), copy the selected + // cells out of the cache and attend densely over just those. this bounds the per-token + // attention cost by the top-k width instead of the context depth. + // restricted to single-token-per-stream ubatches (decode), where every stream carries + // exactly one top-k list, so the gathered K/V stay one tensor per stream. + if (gather) { + ggml_tensor * kf = mctx_cur->get_k(ctx0, il); // [hd_k, n_head_kv, n_kv, ns] + ggml_tensor * vf = mctx_cur->get_v(ctx0, il); // [hd_v, n_head_kv, n_kv, ns] + + const int64_t hd_k = kf->ne[0]; + const int64_t hd_v = vf->ne[0]; + const int64_t n_h_kv = kf->ne[1]; + const int64_t n_kv = kf->ne[2]; + const int64_t ns = kf->ne[3]; + const int64_t n_topk = top_k->ne[0]; + + GGML_ASSERT(top_k->ne[1] == 1 && "QSA gather requires single-token-per-stream ubatches"); + GGML_ASSERT(qsa_bias != nullptr && qsa_bias->ne[0] == n_kv && "QSA gather requires the per-cell bias"); + + // the heads of a cell are contiguous in the cache, so a cell can be gathered as one row + GGML_ASSERT(kf->nb[2] == ggml_row_size(kf->type, hd_k*n_h_kv)); + GGML_ASSERT(vf->nb[2] == ggml_row_size(vf->type, hd_v*n_h_kv)); + + ggml_tensor * k_cells = ggml_view_4d(ctx0, kf, hd_k*n_h_kv, n_kv, 1, ns, + kf->nb[2], kf->nb[2]*n_kv, kf->nb[3], 0); + ggml_tensor * v_cells = ggml_view_4d(ctx0, vf, hd_v*n_h_kv, n_kv, 1, ns, + vf->nb[2], vf->nb[2]*n_kv, vf->nb[3], 0); + + // top_k [n_topk, 1, 1, ns] -> the index layout ggml_get_rows expects: [n_topk, 1, ns, 1] + ggml_tensor * idx = ggml_reshape_4d(ctx0, top_k, n_topk, 1, ns, 1); + + // get_rows dequantizes the cells to F32; build_attn_mha casts to F16 for flash attention + ggml_tensor * k_g = ggml_get_rows(ctx0, k_cells, idx); // F32 [hd_k*n_h_kv, n_topk, 1, ns] + ggml_tensor * v_g = ggml_get_rows(ctx0, v_cells, idx); // F32 [hd_v*n_h_kv, n_topk, 1, ns] + + k_g = ggml_reshape_4d(ctx0, k_g, hd_k, n_h_kv, n_topk, ns); + v_g = ggml_reshape_4d(ctx0, v_g, hd_v, n_h_kv, n_topk, ns); + cb(k_g, "qsa_k_gathered", il); + cb(v_g, "qsa_v_gathered", il); + + // the gathered cells' bias values already encode visibility: 0 for a selected block + // member, 1e9 for the always-visible tail, -inf for anything the padded top-k width + // pulled in that the query must not see. clamping to (-inf, 0] turns that into the + // attention mask over the gathered set. + ggml_tensor * b1 = ggml_view_4d(ctx0, qsa_bias, 1, n_kv, 1, ns, + qsa_bias->nb[0], qsa_bias->nb[1], qsa_bias->nb[3], 0); + ggml_tensor * m_g = ggml_get_rows(ctx0, b1, idx); // F32 [1, n_topk, 1, ns] + m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns); + m_g = ggml_clamp(ctx0, m_g, -INFINITY, 0.0f); + m_g = ggml_cast(ctx0, m_g, GGML_TYPE_F16); // FA wants contiguous F16 + cb(m_g, "qsa_mask_gathered", il); + + ggml_tensor * cur = build_attn_mha(q_cur, k_g, v_g, nullptr, m_g, nullptr, nullptr, 0, 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; + } + ggml_tensor * kq_mask = inp->get_kq_mask(); // prepare new kq mask - starts filled with -INFINITY @@ -771,7 +846,26 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // indexer reads the same block input as q/k/v; no cache or no ratio means dense const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + // gather-based QSA decode: worth it once the cache is meaningfully deeper than the + // top-k width; below that the masked path costs about the same. QWEN4EXP_QSA_GATHER=0 + // disables it (A/B lever, and an escape hatch). + static const bool gather_enabled = [] { + const char * e = getenv("QWEN4EXP_QSA_GATHER"); + return e == nullptr || atoi(e) != 0; + }(); + + bool gather = false; + if (qsa && gather_enabled) { + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_hyb->get_idx()->get_n_kv(); + const int64_t width = GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256); + + const int64_t n_stream = mctx_hyb->get_n_stream(); + + gather = n_tokens == n_stream && n_kv >= 4*width; + } + + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il, gather) : nullptr; // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] @@ -824,7 +918,9 @@ 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_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); + ggml_tensor * qsa_bias = gather ? qsa_inps.at((uint32_t) hparams.dsv4_compress_ratios[il])->bias : nullptr; + + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, qsa_bias, kq_scale, il, gather); } else { cur = build_attn(inp, nullptr, nullptr, nullptr,