diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a8939e481ce7..bf0913c366a6 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -248,6 +248,7 @@ llama_kv_cache::llama_kv_cache( : nullptr; if (k_idx) { ggml_format_name(k_idx, "cache_k_idx_l%d", il); + msa_strict_slots = (n_stream == n_seq_max); } std::vector k_idx_stream; @@ -1025,6 +1026,11 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, uint32_t head_cur = v_heads[seq_to_stream[seq_id]]; + // MSA block selection assumes slot == logical position (append-only streams), which Head-based placement can technically violate after tail trims + if (msa_strict_slots) { + head_cur = 0; + } + // if we have enough unused cells before the current head -> // better to start searching from the beginning of the cache, hoping to fill it if (head_cur > cells.get_used() + 2*n_tokens) { @@ -1142,6 +1148,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & const auto idx = sinfo.idxs[s][ii]; + if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) { + LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: " + "writing pos %d into cell %u (stream %u). The indexer cache " + "would desync and block selection would silently corrupt. " + "This is a bug, please report it with reproduction steps.\n", + __func__, ubatch.pos[i], idx, sinfo.strm[s]); + GGML_ABORT("MSA: slot != pos"); + } + if (!cells.is_empty(idx)) { assert(cells.seq_count(idx) == 1); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index e5afa20101be..d5a92f4405b5 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -263,6 +263,9 @@ class llama_kv_cache : public llama_memory_i { // env: LLAMA_KV_CACHE_DEBUG int debug = 0; + // set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq) + bool msa_strict_slots = false; + // this is the SWA type of the cache - not to be confused with the model SWA type const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;