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);