diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index f815cd7665..3330b66262 100644 --- a/src/llama-kv-cache.cpp +++ b/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 diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index f00ec2137b..b211ca36f9 100644 --- a/src/llama-kv-cache.h +++ b/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); }; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index e7f6e565dc..38e3e61c57 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/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 { diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index a17e184209..8c867b56eb 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/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;