From fe215a8ccdce6b844d2a3a3bbde08ae76a6284bf Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Thu, 3 Sep 2026 12:35:31 -0500 Subject: [PATCH] qwen4exp: gather selected cells for depth-constant TG attention (#2404) * qwen4exp: gather selected cells for depth-constant TG attention * qwen4exp: skip the index pad when the width is already aligned --- src/graphs/build_qwen4exp.cpp | 87 +++++++++++++++++++++++++++++++++-- src/llama-build-context.cpp | 9 ++-- src/llama-build-context.h | 2 +- 3 files changed, 90 insertions(+), 8 deletions(-) diff --git a/src/graphs/build_qwen4exp.cpp b/src/graphs/build_qwen4exp.cpp index 99a38134..30232065 100644 --- a/src/graphs/build_qwen4exp.cpp +++ b/src/graphs/build_qwen4exp.cpp @@ -253,7 +253,8 @@ static ggml_tensor * qwen4exp_qsa_mask( ggml_tensor * inp_pos, ggml_tensor * KQ_mask, int il, - const llm_build_cb & cb) { + const llm_build_cb & cb, + ggml_tensor ** top_k_out) { const llama_hparams & hparams = bctx.hparams; const llama_model & model = bctx.model; const llama_kv_cache & kv_self = bctx.kv_self; @@ -396,6 +397,7 @@ static ggml_tensor * qwen4exp_qsa_mask( ggml_build_forward_expand(gf, fused); ggml_tensor * mask = ggml_indexer_mask(ctx0, KQ_mask, fused); cb(mask, "qsa_mask", il); + *top_k_out = fused; return mask; } @@ -422,10 +424,82 @@ static ggml_tensor * qwen4exp_qsa_mask( ggml_tensor * mask = ggml_indexer_mask(ctx0, KQ_mask, top_k); cb(mask, "qsa_mask", il); + *top_k_out = top_k; return mask; } +static bool qwen4exp_qsa_gather( + llm_build_context & bctx, + ggml_context * ctx0, + ggml_tensor * top_k, + ggml_tensor * KQ_mask, + int il, + const llm_build_cb & cb, + ggml_tensor ** k_out, + ggml_tensor ** v_out, + ggml_tensor ** mask_out) { + *k_out = *v_out = nullptr; + + const llama_hparams & hparams = bctx.hparams; + const llama_kv_cache & kv_self = bctx.kv_self; + + ggml_tensor * k_cache = kv_self.k_l[il]; + ggml_tensor * v_cache = kv_self.v_l[il]; + + if (bctx.n_tokens != 1 || !bctx.cparams.flash_attn || k_cache->extra || v_cache->extra) { + return false; + } + + const int32_t n_head_kv = hparams.n_head_kv(il); + const int32_t n_embd_head_k = hparams.n_embd_head_k(il); + const int32_t n_embd_head_v = hparams.n_embd_head_v(il); + // K cache is [n_embd_head_k, n_head_kv*cells]; one cell is n_head_kv rows + const size_t k_cell = k_cache->nb[1]*n_head_kv; + const size_t v_cell = ggml_row_size(v_cache->type, hparams.n_embd_v_gqa(il)); + + const int32_t n_cells = kv_self.rows(il); + const int32_t width = top_k->ne[0]; + const int32_t n_pad = GGML_PAD(width, llama_kv_cache::get_padding(bctx.cparams.flash_attn)); + + // not worth the extra copy until the cache is a few times the selection width + if (3*n_pad > bctx.n_kv) { + return false; + } + + // pad the indices to the FA granularity; extra entries hit cell 0 and get masked. No I32 PAD, but it only copies + ggml_tensor * idx = ggml_reshape_1d(ctx0, top_k, width); + if (n_pad != width) { + idx = ggml_reshape_4d_ext(ctx0, idx, GGML_TYPE_F32, width, 1, 1, 1); + idx = ggml_pad(ctx0, idx, n_pad - width, 0, 0, 0); + idx = ggml_reshape_1d(ctx0, ggml_reshape_4d_ext(ctx0, idx, GGML_TYPE_I32, n_pad, 1, 1, 1), n_pad); + } + + ggml_tensor * k = ggml_reshape_4d_ext(ctx0, k_cache, GGML_TYPE_F32, k_cell/sizeof(float), n_cells, 1, 1); + k = ggml_get_rows(ctx0, k, idx); + k = ggml_reshape_4d_ext(ctx0, k, k_cache->type, n_embd_head_k, n_head_kv, n_pad, 1); + *k_out = ggml_view_3d(ctx0, k, n_embd_head_k, n_pad, n_head_kv, k->nb[2], k->nb[1], 0); + cb(*k_out, "qsa_k_sel", il); + + ggml_tensor * v = ggml_reshape_4d_ext(ctx0, v_cache, GGML_TYPE_F32, v_cell/sizeof(float), n_cells, 1, 1); + v = ggml_get_rows(ctx0, v, idx); + v = ggml_reshape_4d_ext(ctx0, v, v_cache->type, n_embd_head_v, n_head_kv, n_pad, 1); + *v_out = ggml_view_3d(ctx0, v, n_embd_head_v, n_pad, n_head_kv, v->nb[2], v->nb[1], 0); + cb(*v_out, "qsa_v_sel", il); + + // gather the mask too: a selected cell can belong to another sequence or lie past its end + ggml_tensor * m = ggml_view_2d(ctx0, KQ_mask, 1, bctx.n_kv, ggml_element_size(KQ_mask), 0); + m = ggml_get_rows(ctx0, m, ggml_reshape_1d(ctx0, top_k, width)); + ggml_tensor * tail = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_pad*KQ_mask->ne[1] - width); + tail = ggml_fill_inplace(ctx0, tail, -INFINITY); + m = ggml_concat(ctx0, ggml_reshape_1d(ctx0, m, width), tail, 0); + m = ggml_cast(ctx0, m, GGML_TYPE_F16); + *mask_out = ggml_reshape_2d(ctx0, m, n_pad, KQ_mask->ne[1]); + cb(*mask_out, "qsa_mask_sel", il); + + return true; +} + ggml_cgraph * llm_build_context::build_qwen4exp() { ggml_cgraph * gf = new_graph_custom(); @@ -530,13 +604,20 @@ ggml_cgraph * llm_build_context::build_qwen4exp() { } else { // the indexer reads the same block input as q/k/v, and returns the causal mask // itself when the layer carries no compression ratio + ggml_tensor * top_k = nullptr; ggml_tensor * mask = hparams.is_qsa(il) - ? qwen4exp_qsa_mask(*this, ctx0, lctx, gf, cur, inp_pos, KQ_mask, il, cb) + ? qwen4exp_qsa_mask(*this, ctx0, lctx, gf, cur, inp_pos, KQ_mask, il, cb, &top_k) : KQ_mask; + ggml_tensor * k_sel = nullptr; + ggml_tensor * v_sel = nullptr; + const bool gathered = top_k && + qwen4exp_qsa_gather(*this, ctx0, top_k, KQ_mask, il, cb, &k_sel, &v_sel, &mask); + cur = build_std_attention(gf, nullptr, cur, inp_pos, nullptr, nullptr, mask, nullptr, nullptr, KQ_scale, 0.0f, 0, il, true, false, - /* add_input */ false, /* is_norm */ false, /* is_multi */ true); + /* add_input */ false, /* is_norm */ false, /* is_multi */ true, + nullptr, -1, 0.0f, nullptr, gathered ? &k_sel : nullptr, gathered ? &v_sel : nullptr); } res_hc = qwen4exp_hc_combine(ctx0, hparams, res_hc, cur, inject, n_embd, il, cb); diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 61373eca..a9b9300e 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -3075,7 +3075,8 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens ggml_tensor * input, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors_in, ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il, bool do_rope, bool add_graph_split, bool add_input, bool is_norm, bool is_multi, - ggml_tensor * post_norm, int kv_il, float post_norm_eps, post_norm_data * pnd) { + ggml_tensor * post_norm, int kv_il, float post_norm_eps, post_norm_data * pnd, + ggml_tensor ** k_view, ggml_tensor ** v_view) { float freq_base_l = n_swa > 0 ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base; float freq_scale_l = n_swa > 0 ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train; @@ -3463,7 +3464,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cur = llm_build_kv(ctx0, lctx, kv_self, gf, nullptr, nullptr, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il, - nullptr, nullptr, swa_head); + k_view, v_view, swa_head); cb(cur, "wqkv", il); auto gate = llm_build_lora_mm(lctx, ctx0, wqkv_gate, input_normed); if (model.arch == LLM_ARCH_LAGUNA) { @@ -3505,7 +3506,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens if (gate) { cur = llm_build_kv(ctx0, lctx, kv_self, gf, nullptr, nullptr, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il, - nullptr, nullptr, swa_head); + k_view, v_view, swa_head); if (false && cur->ne[1] == 1) { // we need to add GGML_UNARY_OP_SIGMOID to the ops supported by ggml_fused_mul_unary cur = ggml_fused_mul_unary(ctx0, cur, gate, GGML_UNARY_OP_SIGMOID); } else { @@ -3523,7 +3524,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cur = llm_build_kv(ctx0, lctx, kv_self, gf, model.layers[il].wo, model.layers[il].bo, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il, - nullptr, nullptr, swa_head); + k_view, v_view, swa_head); } } diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 9c3ce942..6c4ef9b0 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -614,7 +614,7 @@ llm_expert_gating_func_type gating_op, ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il, bool do_rope = true, bool add_graph_split = false, bool add_input = false, bool is_norm = false, bool is_multi = false, ggml_tensor * post_norm = nullptr, int kv_il = -1, float post_norm_eps = 0.0f, - post_norm_data * pnd = nullptr); + post_norm_data * pnd = nullptr, ggml_tensor ** k_view = nullptr, ggml_tensor ** v_view = nullptr); static ggml_tensor * build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur, ggml_tensor * output, const llm_build_cb & cb);