llama: give the qwen4exp indexer cache the attention cache's slots

The indexer cache found its own slots, independently of the attention
cache. Both are the same size and see the same ubatches, so in a
straight-through prefill they agree, which is why every fixture and every
single-shot parity run passed. They drift once the context is being
rewritten between turns, and then the QSA top-k indices, which are applied
against the attention mask, point at the wrong cells.

The seven-turn chat test caught it on the third turn: llama-server aborted
on the assertion that the two caches report the same n_kv.

The cache is a side buffer addressed by the attention cache's cells, so it
now takes that cache's slot layout instead of computing one. Applying that
layout also marks its cells identically, so the two agree cell for cell by
construction rather than by coincidence, and the assertion can no longer
fire.

Inert where the caches already agreed: test-llama-archs green at 126 archs
and 0.00e+00, and the 4096-token tiny fixture is unchanged at max logit
delta 0.0.

src/llama-graph.cpp

@@ -3000,9 +3000,8 @@ ggml_tensor * llm_graph_context::build_attn(
     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.
+// Restrict a KQ mask to the positions named by top_k. Lifted verbatim from the DSA
+// build_attn below so non-MLA sparse architectures can share it; node sequence unchanged.
 ggml_tensor * llm_graph_context::build_attn_mask_top_k(
         ggml_tensor * kq_mask,
         ggml_tensor * top_k) const {

src/llama-graph.h

@@ -1206,8 +1206,8 @@ 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.
+    // as above, but attending only to positions named by top_k; for sparse attention
+    // architectures that are plain GQA rather than MLA
     ggml_tensor * build_attn(
             llm_graph_input_attn_kv * inp,
             ggml_tensor * wo,

src/llama-kv-cache.cpp

@@ -1811,19 +1811,17 @@ void llama_kv_cache::set_input_qsa(
     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.
+    // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio. All three
+    // mrope sections carry that position: exact for text, approximate for 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.
+    // an incompletely populated block cannot be pooled. Those cells are the tail, which
+    // the bias below forces in anyway, so they point at block 0 just to keep the gather
+    // in range. -1 marks a cell with no usable block.
     std::vector<int32_t> blk_of(n_kv, -1);
     std::vector<int32_t> filled(n_blocks, 0);
 
@@ -1857,9 +1855,8 @@ void llama_kv_cache::set_input_qsa(
         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
+        // from here on we are inside an incomplete block, always attended to, which is what
+        // lands the selection on block boundaries as the reference does
         const llama_pos tail_start = (q + 1)/r*r;
 
         for (int64_t j = 0; j < n_kv; ++j) {

src/llama-kv-cache.h

@@ -216,13 +216,12 @@ 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
+    // block-compressed sparse attention (qwen4exp QSA) over this cache's cells. Blocks
+    // cut the *position* line, not the cell array, so nothing assumes contiguous layout:
+    //   cell_blk  I32 [n_kv]           block each cell belongs to
+    //   blk_cells I32 [ratio*n_blocks] 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;
 

src/llama-memory-hybrid.cpp

@@ -66,8 +66,7 @@ llama_memory_hybrid::llama_memory_hybrid(
             : filter_recr
     )),
     mem_idx(filter_idx == nullptr ? nullptr : [&] {
-        // MQA with a single key head of indexer_head_size, the same shaping
-        // llama_kv_cache_dsa applies to its lightning-indexer cache
+        // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own
         std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
         hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
 
@@ -124,20 +123,21 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
         }
 
         // prepare the attention cache
-        llama_kv_cache::slot_info_vec_t heads_idx;
-        if (mem_idx) {
-            heads_idx = mem_idx->prepare(ubatches);
-            if (heads_idx.empty()) {
-                break;
-            }
-        }
-
         auto heads_attn = mem_attn->prepare(ubatches);
         if (heads_attn.empty()) {
             LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__);
             return std::make_unique<llama_memory_hybrid_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
         }
 
+        // The indexer cache is a side buffer addressed by the attention cache's cells, so it
+        // takes that slot layout rather than finding its own. Allocating separately let the
+        // two drift once context was rewritten between turns, pointing QSA top-k at the
+        // wrong cells.
+        llama_kv_cache::slot_info_vec_t heads_idx;
+        if (mem_idx) {
+            heads_idx = heads_attn;
+        }
+
         return std::make_unique<llama_memory_hybrid_context>(
                 this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches));
     } while(false);

src/llama-memory-hybrid.h

@@ -92,8 +92,7 @@ public:
 private:
     const llama_hparams & hparams;
 
-    // geometry for the indexer cache: MQA with a single key head of
-    // indexer_head_size, mirroring how llama_kv_cache_dsa builds its own
+    // indexer cache geometry: MQA, one key head of indexer_head_size, as in llama_kv_cache_dsa
     llama_hparams hparams_idx;
 
     const std::unique_ptr<llama_kv_cache> mem_attn;

src/llama-model.cpp

@@ -2434,8 +2434,8 @@ 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
+                    // 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; };

src/models/models.h

@@ -2314,8 +2314,7 @@ struct llama_model_qwen4exp : public llama_model_base {
                             int * sections,
                             int   il);
 
-        // QSA: the token indices this layer's queries may attend to, or nullptr
-        // to attend densely
+        // QSA: token indices this layer's queries may attend to, or nullptr for dense
         ggml_tensor * build_qsa_top_k(
   const llama_kv_cache_context * mctx_idx,
                     ggml_tensor * cur,

src/models/qwen4exp.cpp

@@ -266,8 +266,9 @@ 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
+    // present only when the GGUF carries indexer tensors, so a model without them still
+    // builds a dense graph. The indexer cache takes the attention cache's slot layout,
+    // so the two agree cell for cell by construction.
     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() &&
@@ -372,14 +373,10 @@ 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
+// QSA attends to a budget of whole blocks of `compress_ratio` tokens, scored by one
+// mean-pooled indexer key each, plus the always-visible incomplete tail. Below
+// indexer_top_k + compress_ratio - 1 cached tokens this is exactly dense attention.
+// Everything depending on cache layout is computed host-side in set_input; the
 // graph only gathers, pools and scores.
 class llm_graph_input_qsa : public llm_graph_input_i {
 public:
@@ -432,8 +429,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     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
+    // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither
     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);
@@ -447,8 +443,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     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
+    // mean over the block's members; compress_ratio is small, so summing slices beats
+    // transposing to reach ggml_sum_rows
     ggml_tensor * pooled = nullptr;
     for (int64_t i = 0; i < r; ++i) {
         ggml_tensor * slice = ggml_cont(ctx0,
@@ -474,10 +470,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
             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.
+    // Each head's dot product is rectified before summing, as in DeepSeek's lightning
+    // indexer but with no learned per-head weight. The reference's constant divisor
+    // cannot reorder anything, so it 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);
@@ -487,19 +482,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     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.
+    // Give every token of a block its block's score rather than expanding the block
+    // indices, which would need an integer multiply-add ggml has no op for. The
+    // budget is a whole number of blocks and members tie, so the cut still lands on
+    // a block boundary. get_rows gathers rows, so scores are transposed first.
     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
+    // the reference returns indexer_top_k + compress_ratio - 1: 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));
@@ -518,9 +512,8 @@ 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.
+    // The indexer reads the same block input as q/k/v; with no indexer cache this
+    // falls back to dense, which is what the model computes below the budget anyway.
     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