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

@@ -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"):

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

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(

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<int32_t> blk_of(n_kv, -1);
+    std::vector<int32_t> 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);
 }

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;
 

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;

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,

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<llm_graph_input_qsa>(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<int64_t>(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);