From c88c916636611afb93f0ea6aa17a8b0edb6c5760 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 25 Aug 2026 19:09:43 +0000 Subject: [PATCH] llama: QSA sparse attention for qwen4exp The full-attention layers of this model do not attend to everything. An indexer scores one mean-pooled key per block of compress_ratio tokens and keeps a budget of the best blocks, plus the tail of tokens that do not yet form a complete block. Below indexer_top_k + compress_ratio - 1 cached tokens every block fits in the budget, so the result is exactly dense. What is reused rather than rebuilt: - the mask machinery. build_attn's DSA overload already turns a list of token indices into a KQ mask via ggml_set_rows, so that block is lifted out verbatim into build_attn_mask_top_k and shared with a new overload on llm_graph_input_attn_kv. DSA's node sequence is unchanged; the new overload exists because llama_kv_cache_dsa assumes MLA and cannot be dropped into a hybrid model. - the indexer key cache, which is the optional third cache added to llama_memory_hybrid in the previous commit. It holds raw keys, because pooling happens before the norm and the rotation. The graph expands block scores rather than block indices: giving every token of a block its block's score needs only a gather, where expanding indices would need an integer multiply-add that ggml has no op for. Since the budget is a whole number of blocks and a block's members tie exactly, the cut still lands on a block boundary. Everything that depends on cache layout is computed host-side in set_input_qsa. Blocks are cuts of the position line rather than of the cell array, so nothing assumes the cache is contiguous. Measured on the tiny fixture against vLLM, comparing the selected token indices directly rather than the logits: below the budget selection identical, and 1024-token logits are bit-identical to the pre-QSA dense path above the budget mean jaccard 0.975 The direct index comparison is what made this correct. The reference rectifies each head's dot product before summing over heads, which an earlier reading of it had missed; on logits alone the resulting port looked fine, because on a randomly initialised fixture the known-correct dense path already disagrees with vLLM by more than the bug did. Comparing the indices showed 0.794, and fixing the ReLU moved it to 0.975. --- conversion/qwen4exp.py | 3 +- src/llama-graph.cpp | 114 +++++++++++++++++++++------ src/llama-graph.h | 21 +++++ src/llama-kv-cache.cpp | 98 ++++++++++++++++++++++++ src/llama-kv-cache.h | 13 ++++ src/llama-model.cpp | 13 +++- src/models/models.h | 10 +++ src/models/qwen4exp.cpp | 165 +++++++++++++++++++++++++++++++++++++++- 8 files changed, 406 insertions(+), 31 deletions(-) diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 85510ce648..0f80a69661 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -127,7 +127,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): ] # Gemma zero-centred gammas the inherited norm.weight rule misses - if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight")): + if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight", + ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")): return [(self.map_tensor_name(name), data_torch + 1)] if name.endswith(".ple.conv1d.weight"): diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0..4b52ba5ea3 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2945,6 +2945,94 @@ 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 out of the +// DSA build_attn below so the sparse-attention architectures that do not use MLA +// can share it; the node sequence is 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, @@ -2977,31 +3065,7 @@ ggml_tensor * llm_graph_context::build_attn( const auto & kq_mask = inp->get_kq_mask_mla(); - // 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 * kq_mask_top_k = build_attn_mask_top_k(kq_mask, top_k); 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 b388e028cb..24b6100a5a 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1206,6 +1206,27 @@ struct llm_graph_context { float kq_scale, int il) const; + // as above, but attending only to the positions named by top_k. Used by + // architectures whose sparse attention is 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/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a7531..2dd77ed3e3 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1787,6 +1787,94 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch } } +void llama_kv_cache::set_input_qsa( + ggml_tensor * cell_blk, + ggml_tensor * blk_cells, + ggml_tensor * blk_pos, + ggml_tensor * bias, + const llama_ubatch * ubatch, + uint32_t ratio) const { + GGML_ASSERT(n_stream == 1 && "TODO: support multiple streams"); + GGML_ASSERT(ratio > 0); + + const auto & cells = v_cells[0]; + + GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); + + const int64_t n_kv = cell_blk->ne[0]; + const int64_t n_blocks = blk_pos->ne[0]/4; + const int64_t n_tokens = ubatch->n_tokens; + const int64_t r = ratio; + + int32_t * dst_cell_blk = (int32_t *) cell_blk->data; + int32_t * dst_blk_cells = (int32_t *) blk_cells->data; + int32_t * dst_blk_pos = (int32_t *) blk_pos->data; + float * dst_bias = (float *) bias->data; + + // block b is positions [b*ratio, (b+1)*ratio), so its first token sits at + // b*ratio. The three mrope sections all carry that position: exact for text, + // an approximation for the interleaved t/h/w positions of image tokens. + for (int64_t s = 0; s < 4; ++s) { + for (int64_t b = 0; b < n_blocks; ++b) { + dst_blk_pos[s*n_blocks + b] = (int32_t) (b*r); + } + } + + // a block that is not completely populated cannot be pooled. Those cells are + // exactly the tail of the sequence, which the bias below forces in whatever + // score they carry, so they are pointed at block 0 only to keep the gather + // in range. -1 marks a cell with no usable block at all. + std::vector blk_of(n_kv, -1); + std::vector filled(n_blocks, 0); + + std::fill(dst_blk_cells, dst_blk_cells + r*n_blocks, 0); + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + + const llama_pos p = cells.pos_get(j); + const int64_t b = p/r; + + if (b >= n_blocks) { + continue; + } + + blk_of[j] = (int32_t) b; + dst_blk_cells[b*r + (p%r)] = (int32_t) j; + filled[b]++; + } + + for (int64_t j = 0; j < n_kv; ++j) { + if (blk_of[j] >= 0 && filled[blk_of[j]] < r) { + blk_of[j] = -1; + } + dst_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + } + + for (int64_t i = 0; i < n_tokens; ++i) { + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + const llama_pos q = ubatch->pos[i]; + + // everything from here on is inside an incomplete block and is always + // attended to, which is what makes the whole selection land on block + // boundaries the way the reference implementation does + const llama_pos tail_start = (q + 1)/r*r; + + for (int64_t j = 0; j < n_kv; ++j) { + float v = -INFINITY; + + if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { + // finite, so it can never meet a -inf and produce a nan + v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } + + dst_bias[i*n_kv + j] = v; + } + } +} + void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const { GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); @@ -2645,6 +2733,16 @@ void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama kv->set_input_pos_bucket(dst, ubatch); } +void llama_kv_cache_context::set_input_qsa( + ggml_tensor * cell_blk, + ggml_tensor * blk_cells, + ggml_tensor * blk_pos, + ggml_tensor * bias, + const llama_ubatch * ubatch, + uint32_t ratio) const { + kv->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio); +} + void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { kv->set_input_k_rot(dst); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f9..8511dc4186 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -216,6 +216,16 @@ public: void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // block-compressed sparse attention (qwen4exp QSA), computed over this + // cache's cells. Blocks are cuts of the *position* line, not of the cell + // array, so nothing here assumes the cache is laid out contiguously: + // cell_blk I32 [n_kv] block each cell belongs to + // blk_cells I32 [ratio*n_blocks] the cells making up each block + // blk_pos I32 [4*n_blocks] mrope position rows of each block's first token + // bias F32 [n_kv, n_tokens] -inf where invisible, large where always visible + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio) const; + void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; @@ -398,6 +408,9 @@ public: void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio) const; + void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a5e82501ed..dad7c0b81b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,6 +2434,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + // left null for every architecture but the sparse-attention + // ones, which is what keeps the indexer cache from existing + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; @@ -2451,6 +2454,13 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + + if (arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0) { + // QSA runs on the dense-attention layers only + filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { @@ -2491,7 +2501,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx)); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; diff --git a/src/models/models.h b/src/models/models.h index d66c2c4c12..6f024896d2 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2308,6 +2308,16 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * build_layer_attn( llm_graph_input_attn_kv * inp_attn, + const llama_kv_cache_context * mctx_idx, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il); + + // QSA: the token indices this layer's queries may attend to, or nullptr + // to attend densely + ggml_tensor * build_qsa_top_k( + const llama_kv_cache_context * mctx_idx, ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 4d054da7f6..6a5247e7c0 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -1,4 +1,5 @@ #include "models.h" +#include "llama-memory-hybrid.h" #include "llama-memory-recurrent.h" void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { @@ -265,6 +266,14 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa auto * inp = build_inp_mem_hybrid(); + // present only when the GGUF carries indexer tensors, so a model without + // them still builds a dense graph + const llama_kv_cache_context * mctx_idx = inp->mctx->get_idx(); + if (mctx_idx) { + GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() && + "the indexer cache must track the attention cache cell for cell"); + } + ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); @@ -294,7 +303,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa if (hparams.is_recr(il)) { cur = build_layer_attn_linear(inp->get_recr(), cur, il); } else { - cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il); + cur = build_layer_attn(inp->get_attn(), mctx_idx, cur, inp_pos, sections, il); } res_hc = build_hc_combine(res_hc, cur, inject, il); @@ -363,8 +372,145 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( return ggml_mul(ctx0, normalized, gated); } +// QSA attends to a budget of whole blocks of `compress_ratio` tokens, chosen by +// scoring one mean-pooled indexer key per block, plus the tail of tokens that do +// not yet make up a complete block, which is always visible. Below +// indexer_top_k + compress_ratio - 1 cached tokens every block fits in the +// budget and the result is exactly dense attention. +// +// Everything that depends on the cache layout - which cells make up a block, +// which blocks a query may see - is computed host-side in set_input, so the +// graph only gathers, pools and scores. +class llm_graph_input_qsa : public llm_graph_input_i { +public: + llm_graph_input_qsa(const llama_kv_cache_context * mctx, uint32_t ratio) : + mctx(mctx), ratio(ratio) {} + virtual ~llm_graph_input_qsa() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio); + } + + ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] + ggml_tensor * cell_blk = nullptr; // I32 [n_kv] + ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks] + ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks] + ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens] + + const llama_kv_cache_context * mctx; + const uint32_t ratio; +}; + +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( + const llama_kv_cache_context * mctx_idx, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int64_t idx_dim = hparams.indexer_head_size; + const int64_t n_idx_h = hparams.indexer_n_head; + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_idx->get_n_kv(); + const int64_t n_blocks = (n_kv + r - 1)/r; + + GGML_ASSERT(r > 0); + + auto qsa = std::make_unique(mctx_idx, (uint32_t) r); + + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + qsa->cell_blk = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv); + qsa->blk_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r*n_blocks); + qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks); + qsa->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); + + ggml_set_input(qsa->cell_blk); + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + ggml_set_input(qsa->bias); + + llm_graph_input_qsa * inp = qsa.get(); + res->add_input(std::move(qsa)); + + // the cached indexer keys are raw: pooling happens before the norm and the + // rotation, so neither may be applied on the way in + ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); + k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); + cb(k_raw, "indexer_k_raw", il); + + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); + + // one key head, so the per-cell rows are contiguous and this view is dense + ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); + k_all = ggml_view_2d(ctx0, k_all, idx_dim, n_kv, k_all->nb[2], 0); + + ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); + members = ggml_reshape_3d(ctx0, members, idx_dim, r, n_blocks); + + // mean over the block's members. compress_ratio is small, so summing the + // slices costs less than transposing to reach ggml_sum_rows + ggml_tensor * pooled = nullptr; + for (int64_t i = 0; i < r; ++i) { + ggml_tensor * slice = ggml_cont(ctx0, + ggml_view_2d(ctx0, members, idx_dim, n_blocks, members->nb[2], i*members->nb[1])); + pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; + } + pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); + cb(pooled, "indexer_k_pooled", il); + + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks); + pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + pooled = ggml_reshape_2d(ctx0, pooled, idx_dim, n_blocks); + cb(pooled, "indexer_k", il); + + ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); + q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens); + q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il); + q = ggml_rope_multi(ctx0, q, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q, "indexer_q", il); + + // Each head's dot product is rectified before the heads are summed, as in + // DeepSeek's lightning indexer, but with no learned per-head weight. The + // reference then divides by a constant, which cannot reorder anything, so + // that is left out. + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, + ggml_reshape_2d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tokens)); + score = ggml_reshape_3d(ctx0, score, n_blocks, n_idx_h, n_tokens); + score = ggml_relu(ctx0, score); + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + score = ggml_sum_rows(ctx0, score); + score = ggml_reshape_2d(ctx0, score, n_blocks, n_tokens); + cb(score, "indexer_score", il); + + // Give every token of a block its block's score, rather than expanding the + // selected block indices: that would need an integer multiply-add, which + // ggml has no op for. Because the budget is a whole number of blocks and the + // members of a block tie exactly, the cut still lands on a block boundary. + // get_rows gathers rows, so the scores are transposed for the gather. + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_transpose(ctx0, score)), inp->cell_blk); + expanded = ggml_cont(ctx0, ggml_transpose(ctx0, expanded)); + expanded = ggml_add(ctx0, expanded, inp->bias); + cb(expanded, "indexer_score_tokens", il); + + // the reference returns indexer_top_k + compress_ratio - 1 tokens: a whole + // budget of blocks plus the incomplete tail + const int64_t width = 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)); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( llm_graph_input_attn_kv * inp, + const llama_kv_cache_context * mctx_idx, ggml_tensor * cur, ggml_tensor * inp_pos, int * sections, @@ -372,6 +518,11 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // The indexer reads the same block input as q/k/v. Without an indexer cache + // this falls back to dense attention, which is what the model computes + // anyway for anything shorter than the budget. + ggml_tensor * top_k = mctx_idx ? build_qsa_top_k(mctx_idx, cur, inp_pos, sections, il) : nullptr; + // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate @@ -427,9 +578,15 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // Attention computation const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - cur = build_attn(inp, - nullptr, nullptr, nullptr, - Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + if (top_k) { + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, top_k, kq_scale, il); + } else { + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + } cb(cur, "attn_pregate", il); ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate);