qwen4exp: keep the indexer cache in step across server slots

The QSA indexer keeps a side cache addressed by the cells of the attention
cache, so cell j has to hold the same token in both: the top-k indices it
produces are applied to the attention KQ mask. init_batch already hands the
indexer the attention cache's slot layout rather than letting it look for its
own, but the restore path did not. state_read called llama_kv_cache::state_read
on the two caches in turn and each ran its own find_slot over its own occupancy.
That agrees only for as long as nothing has already pushed the two caches apart,
which is the property a restore is supposed to re-establish rather than one it
can lean on.

The failure path was the worse half, and it is reachable from the public API
with nothing more than a short buffer. Truncating a good blob at 35 offsets and
feeding it to llama_state_seq_set_data left the two caches disagreeing at 5 of
them, and every one of 23 truncations of a whole-context blob did. Four of those
five land inside the attention section, so the attention cache drops the
sequence and the indexer keeps it; only the cut that lands in the indexer
section gives the opposite direction. llama_kv_cache::state_read cleans up its
own cache and rethrows, so whichever way it falls, nothing is left to bring the
two back together. The server papers over this by clearing the slot when a
prompt cache load fails; a caller of llama_state_seq_set_data that does not is
left with an indexer addressing cells that no longer mean what it thinks.

llama_kv_cache::state_read_sinfo reports the cells a restore landed in, or takes
a copy of them, and state_read_meta uses a supplied layout in place of find_slot
once it has checked that those cells are free here too. The indexer now adopts
the attention cache's restored layout by construction instead of reproducing it
by coincidence, and a layout that does not fit fails the read rather than being
applied over cells that already drifted. The hybrid restore is wrapped so that
any failure drops the sequence, or for a whole-context restore the context, from
all three caches at once, which is a state they do agree on.

src/llama-kv-cache.cpp

@@ -2038,6 +2038,15 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
 }
 
 void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+    state_read_sinfo(io, seq_id, flags, nullptr, nullptr);
+}
+
+void llama_kv_cache::state_read_sinfo(
+        llama_io_read_i & io,
+           llama_seq_id   seq_id,
+  llama_state_seq_flags   flags,
+      slot_info_vec_t *   sinfos_out,
+const slot_info_vec_t *   sinfos_in) {
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -2048,6 +2057,14 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
     // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
     GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
 
+    if (sinfos_out) {
+        sinfos_out->assign(n_stream, slot_info{});
+    }
+
+    if (sinfos_in && sinfos_in->size() != n_stream) {
+        throw std::runtime_error("failed to restore kv cache: mirrored slot layout has the wrong stream count");
+    }
+
     uint32_t n_stream_cur;
     io.read(&n_stream_cur, sizeof(n_stream_cur));
     if (n_stream_cur != n_stream) {
@@ -2059,6 +2076,10 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
         io.read(&cell_count, sizeof(cell_count));
 
         if (cell_count == 0) {
+            // a mirrored cache must be empty here as well, or the two no longer agree cell for cell
+            if (sinfos_in && !(*sinfos_in)[s].empty()) {
+                throw std::runtime_error("failed to restore kv cache: mirrored cache holds cells this one does not");
+            }
             continue;
         }
 
@@ -2067,7 +2088,7 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
         slot_info sinfo;
 
         bool res = true;
-        res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
+        res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id, sinfos_in ? &(*sinfos_in)[s] : nullptr);
 
         try {
             res = res && state_read_data(io, strm, cell_count, sinfo);
@@ -2083,6 +2104,10 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
             }
             throw std::runtime_error("failed to restore kv cache");
         }
+
+        if (sinfos_out) {
+            (*sinfos_out)[s] = sinfo;
+        }
     }
 }
 
@@ -2218,7 +2243,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
     }
 }
 
-bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id) {
+bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id, const slot_info * sinfo_in) {
     auto & cells = v_cells[strm];
     auto & head  = v_heads[strm];
 
@@ -2263,10 +2288,39 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
             ubatch.seq_id[i]   = &dest_seq_id;
         }
 
-        sinfo = find_slot(ubatch, false);
-        if (sinfo.empty()) {
-            LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__,  cell_count);
-            return false;
+        if (sinfo_in) {
+            // this cache mirrors another one, so it takes that cache's restored layout rather
+            // than searching for cells of its own
+            if (sinfo_in->empty() || sinfo_in->n_stream() != 1 || sinfo_in->idxs[0].size() != cell_count) {
+                LLAMA_LOG_ERROR("%s: mirrored slot layout holds %d cells, this cache restores %d\n", __func__,
+                        sinfo_in->empty() ? 0 : (int) sinfo_in->idxs[0].size(), cell_count);
+                return false;
+            }
+
+            sinfo = *sinfo_in;
+
+            // the layout is addressed by cell index, so it only means the same thing in both
+            // caches while their streams line up
+            sinfo.s0 = strm;
+            sinfo.s1 = strm;
+            sinfo.strm[0] = strm;
+
+            // seq_rm above freed exactly the cells this sequence held. anything else in the way
+            // is a cache that had already drifted, which this restore must not paper over
+            for (uint32_t i = 0; i < cell_count; ++i) {
+                const uint32_t idx = sinfo.idxs[0][i];
+
+                if (idx >= cells.size() || !cells.is_empty(idx)) {
+                    LLAMA_LOG_ERROR("%s: cell %u of the mirrored slot layout is not free\n", __func__, idx);
+                    return false;
+                }
+            }
+        } else {
+            sinfo = find_slot(ubatch, false);
+            if (sinfo.empty()) {
+                LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__,  cell_count);
+                return false;
+            }
         }
 
         // TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet

src/llama-kv-cache.h

@@ -168,6 +168,21 @@ public:
 
     const llama_kv_cells & get_cells(llama_seq_id seq_id) const;
 
+    // state_read, plus the cells the restored tokens were placed in.
+    // a cache that mirrors another one cell for cell (the qwen4exp indexer) cannot search for
+    // its own cells here: a second independent search only happens to agree with the first.
+    //   sinfos_out: if set, resized to n_stream and filled with the layout used; a stream that
+    //               carried no cells leaves an empty entry
+    //   sinfos_in : if set, the layout to use instead of searching for one. it must have one
+    //               entry per stream and the entry must match the cell count in the blob,
+    //               otherwise the read fails as it would on any other corrupt input
+    void state_read_sinfo(
+            llama_io_read_i & io,
+               llama_seq_id   seq_id,
+      llama_state_seq_flags   flags,
+          slot_info_vec_t *   sinfos_out,
+    const slot_info_vec_t *   sinfos_in);
+
     //
     // graph_build API
     //
@@ -320,7 +335,8 @@ private:
     void state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id = -1) const;
     void state_write_data(llama_io_write_i & io, const cell_ranges_t & cr) const;
 
-    bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count,       slot_info & sinfo, llama_seq_id dest_seq_id = -1);
+    // sinfo_in, when set, replaces the find_slot call: the cells are given by the caller
+    bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count,       slot_info & sinfo, llama_seq_id dest_seq_id = -1, const slot_info * sinfo_in = nullptr);
     bool state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo);
 };
 

src/llama-memory-hybrid-idx.cpp

@@ -476,19 +476,60 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se
 }
 
 void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
-    llama_memory_hybrid::state_read(io, seq_id, flags);
+    // note: this repeats llama_memory_hybrid::state_read because the indexer cache has to be
+    //       handed the cells the attention cache restored into, and because a restore that
+    //       fails halfway has to leave all three caches in the same state
 
-    // [TAG_HYBRID_IDX_STATE] must mirror the write order above.
-    // The indexer finds its own cells, which is safe because the two caches stay in lockstep:
-    // both state_read_meta calls run find_slot over the same occupancy and land on the same cells.
-    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
-        if (mem_idx) {
-            mem_idx->state_read(io, seq_id, flags);
+    // [TAG_HYBRID_IDX_SINFO]
+    // The indexer cache is addressed by the cells of the attention cache, so its restore adopts
+    // that layout instead of searching for cells of its own. Two independent find_slot calls
+    // agree only while nothing makes the two caches see different occupancy, and a restore is
+    // exactly the operation that can no longer promise that.
+    llama_kv_cache::slot_info_vec_t sinfos_attn;
+
+    try {
+        if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
+            get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr);
         }
+
+        get_mem_recr()->state_read(io, seq_id, flags);
+
+        // [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write
+        if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
+            if (mem_idx) {
+                mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
+            }
+        }
+
+        // [TAG_PLE_HISTORY] must mirror the write order above
+        ple_hist_state_read(io, seq_id);
+    } catch (...) {
+        // a half-restored context is the one state the indexer cache cannot be brought back from
+        // by itself: the attention cache holds the restored cells and the indexer the old ones.
+        // drop what was being restored from all of them, which is a state they do agree on.
+        state_drop(seq_id);
+
+        throw;
+    }
+}
+
+void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {
+    // dropped directly rather than through seq_rm, which the recurrent cache is allowed to
+    // refuse and which would then clear the other two caches and not it
+    if (seq_id < 0) {
+        clear(true);
+
+        return;
     }
 
-    // [TAG_PLE_HISTORY] must mirror the write order above
-    ple_hist_state_read(io, seq_id);
+    get_mem_attn()->seq_rm(seq_id, -1, -1);
+    get_mem_recr()->seq_rm(seq_id, -1, -1);
+
+    if (mem_idx) {
+        mem_idx->seq_rm(seq_id, -1, -1);
+    }
+
+    ple_hist.erase(seq_id);
 }
 
 llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const {

src/llama-memory-hybrid-idx.h

@@ -96,6 +96,11 @@ public:
     ple_history & ple_hist_get(llama_seq_id seq_id) const;
 
 private:
+    // forget seq_id (or, for seq_id < 0, everything) in every cache at once, so that a restore
+    // that failed partway cannot leave the indexer cache holding cells the attention cache does
+    // not. seq_id < 0 drops the whole context, as the caches themselves do on a failed restore.
+    void state_drop(llama_seq_id seq_id);
+
     // the indexer cache holds one key head per layer, so it needs its own hparams:
     // llama_kv_cache keeps a reference to what it is given
     llama_hparams hparams_idx;