From af13241f6a06a2acb3a04063bcf648830bda07af Mon Sep 17 00:00:00 2001 From: Chrono Date: Thu, 2 Jul 2026 23:11:24 +0200 Subject: [PATCH 1/7] Log indexer cache size on launch --- src/llama-kv-cache.cpp | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 5f950c2c347e..53189642e0fd 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -327,11 +327,22 @@ llama_kv_cache::llama_kv_cache( { const size_t memory_size_k = size_k_bytes(); const size_t memory_size_v = size_v_bytes(); - - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__, - (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, + const size_t memory_size_k_idx = size_k_idx_bytes(); + + const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx; + + if (memory_size_k_idx > 0) { + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB, K_idx (%s): %7.2f MiB\n", __func__, + (float)memory_size_total / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, + ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), + ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f), + ggml_type_name(GGML_TYPE_F32), (float)memory_size_k_idx / (1024.0f * 1024.0f)); + } else { + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__, + (float)memory_size_total / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f)); + } } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] @@ -1898,6 +1909,18 @@ size_t llama_kv_cache::size_v_bytes() const { return size_v_bytes; } +size_t llama_kv_cache::size_k_idx_bytes() const { + size_t size_k_idx_bytes = 0; + + for (const auto & layer : layers) { + if (layer.k_idx) { + size_k_idx_bytes += ggml_nbytes(layer.k_idx); + } + } + + return size; +} + ggml_tensor * llama_kv_cache::build_rope_shift( const llama_cparams & cparams, ggml_context * ctx, From 453adcf187588d75391898741c9c806cf6e8b113 Mon Sep 17 00:00:00 2001 From: Chrono Date: Thu, 2 Jul 2026 23:43:11 +0200 Subject: [PATCH 2/7] Log indexer cache size on launch and disallow ctx shift --- src/llama-kv-cache.cpp | 12 +++++++++++- src/llama-kv-cache.h | 1 + 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 53189642e0fd..c601488cc0b4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -889,6 +889,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co if (layer.v_stream[ssrc]) { ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } + if (layer.k_idx_stream[ssrc]) { + GGML_ASSERT(layer.k_idx_stream[sdst]); + ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]); + } } } } @@ -1219,6 +1223,12 @@ bool llama_kv_cache::get_can_shift() const { if (hparams.n_pos_per_embd() > 1) { return false; } + //shifting would leave k_idx stale + for (const auto & layer : layers) { + if (layer.k_idx) { + return false; + } + } return true; } @@ -1918,7 +1928,7 @@ size_t llama_kv_cache::size_k_idx_bytes() const { } } - return size; + return size_k_idx_bytes; } ggml_tensor * llama_kv_cache::build_rope_shift( diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c0ecc193ff48..e5afa20101be 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -295,6 +295,7 @@ class llama_kv_cache : public llama_memory_i { size_t size_k_bytes() const; size_t size_v_bytes() const; + size_t size_k_idx_bytes() const; ggml_tensor * build_rope_shift( const llama_cparams & cparams, From d7687615ee0f1f5b891508c0385fe2da38fee6ae Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 3 Jul 2026 00:12:58 +0200 Subject: [PATCH 3/7] Add prompt caching support --- src/llama-kv-cache.cpp | 114 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index c601488cc0b4..39ecd9f7eeaa 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2238,6 +2238,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t } } + if (size_k_idx_bytes() > 0) { + const uint32_t has_k_idx_u32 = 1; + io.write(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + for (const auto & layer : layers) { + const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0; + io.write(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[cr.strm]); + + const int32_t k_idx_type_i = (int32_t) layer.k_idx->type; + io.write(&k_idx_type_i, sizeof(k_idx_type_i)); + + const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + io.write(&k_idx_size_row, sizeof(k_idx_size_row)); + + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * k_idx_size_row; + const size_t offset = range.first * k_idx_size_row; + + io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size); + } + } + } + if (!v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2486,6 +2516,90 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } } + if (size_k_idx_bytes() > 0) { + uint32_t has_k_idx_u32 = 0; + io.read_to(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + if (has_k_idx_u32 != 1) { + LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__); + return false; + } + + size_t dst_stream_idx = 0; + while (dst_stream_idx < sinfo.strm.size() && (uint32_t) sinfo.strm[dst_stream_idx] != strm) { + ++dst_stream_idx; + } + + if (dst_stream_idx == sinfo.strm.size()) { + LLAMA_LOG_ERROR("%s: failed to find destination stream %u in slot info\n", __func__, strm); + return false; + } + + const auto & dst_idxs = sinfo.idxs[dst_stream_idx]; + + if (dst_idxs.size() != cell_count) { + LLAMA_LOG_ERROR( + "%s: invalid k_idx slot count: got %zu, expected %u\n", + __func__, dst_idxs.size(), cell_count); + return false; + } + + std::vector tmp_buf; + + for (auto & layer : layers) { + uint32_t layer_has_k_idx = 0; + io.read_to(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0; + + if (layer_has_k_idx != expected_layer_has_k_idx) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx state for layer: got %u, expected %u\n", + __func__, layer_has_k_idx, expected_layer_has_k_idx); + return false; + } + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[strm]); + + int32_t k_idx_type_i = -1; + io.read_to(&k_idx_type_i, sizeof(k_idx_type_i)); + + if (k_idx_type_i != (int32_t) layer.k_idx->type) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx type: got %d, expected %d\n", + __func__, k_idx_type_i, (int32_t) layer.k_idx->type); + return false; + } + + uint64_t k_idx_size_row = 0; + io.read_to(&k_idx_size_row, sizeof(k_idx_size_row)); + + const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + + if (k_idx_size_row != expected_k_idx_size_row) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx row size: got %" PRIu64 ", expected %" PRIu64 "\n", + __func__, k_idx_size_row, expected_k_idx_size_row); + return false; + } + + tmp_buf.resize((size_t) k_idx_size_row); + + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_to(tmp_buf.data(), (size_t) k_idx_size_row); + + ggml_backend_tensor_set( + layer.k_idx_stream[strm], + tmp_buf.data(), + (size_t) dst_idxs[i] * (size_t) k_idx_size_row, + (size_t) k_idx_size_row); + } + } + } if (!this->v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; From 946a13139c03a8132e4cfb4d3419f32bb9ce8635 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 3 Jul 2026 01:05:02 +0200 Subject: [PATCH 4/7] Fix compile issue --- src/llama-kv-cache.cpp | 42 ++++++++++++++---------------------------- 1 file changed, 14 insertions(+), 28 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 39ecd9f7eeaa..dd4f69e06948 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -328,9 +328,8 @@ llama_kv_cache::llama_kv_cache( const size_t memory_size_k = size_k_bytes(); const size_t memory_size_v = size_v_bytes(); const size_t memory_size_k_idx = size_k_idx_bytes(); - const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx; - + if (memory_size_k_idx > 0) { LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB, K_idx (%s): %7.2f MiB\n", __func__, (float)memory_size_total / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, @@ -1223,7 +1222,7 @@ bool llama_kv_cache::get_can_shift() const { if (hparams.n_pos_per_embd() > 1) { return false; } - //shifting would leave k_idx stale + // shifting would leave k_idx stale for (const auto & layer : layers) { if (layer.k_idx) { return false; @@ -2518,7 +2517,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 if (size_k_idx_bytes() > 0) { uint32_t has_k_idx_u32 = 0; - io.read_to(&has_k_idx_u32, sizeof(has_k_idx_u32)); + io.read(&has_k_idx_u32, sizeof(has_k_idx_u32)); if (has_k_idx_u32 != 1) { LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__); @@ -2535,20 +2534,9 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - const auto & dst_idxs = sinfo.idxs[dst_stream_idx]; - - if (dst_idxs.size() != cell_count) { - LLAMA_LOG_ERROR( - "%s: invalid k_idx slot count: got %zu, expected %u\n", - __func__, dst_idxs.size(), cell_count); - return false; - } - - std::vector tmp_buf; - for (auto & layer : layers) { uint32_t layer_has_k_idx = 0; - io.read_to(&layer_has_k_idx, sizeof(layer_has_k_idx)); + io.read(&layer_has_k_idx, sizeof(layer_has_k_idx)); const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0; @@ -2566,7 +2554,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 GGML_ASSERT(layer.k_idx_stream[strm]); int32_t k_idx_type_i = -1; - io.read_to(&k_idx_type_i, sizeof(k_idx_type_i)); + io.read(&k_idx_type_i, sizeof(k_idx_type_i)); if (k_idx_type_i != (int32_t) layer.k_idx->type) { LLAMA_LOG_ERROR( @@ -2576,7 +2564,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } uint64_t k_idx_size_row = 0; - io.read_to(&k_idx_size_row, sizeof(k_idx_size_row)); + io.read(&k_idx_size_row, sizeof(k_idx_size_row)); const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); @@ -2587,16 +2575,14 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - tmp_buf.resize((size_t) k_idx_size_row); - - for (uint32_t i = 0; i < cell_count; ++i) { - io.read_to(tmp_buf.data(), (size_t) k_idx_size_row); - - ggml_backend_tensor_set( - layer.k_idx_stream[strm], - tmp_buf.data(), - (size_t) dst_idxs[i] * (size_t) k_idx_size_row, - (size_t) k_idx_size_row); + if (cell_count) { + if (sinfo.is_contiguous()) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row); + } else { + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row); + } + } } } } From 0dee698e8e5d33957645e2e159e50a7811377883 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 3 Jul 2026 01:26:53 +0200 Subject: [PATCH 5/7] Remove stale code --- src/llama-kv-cache.cpp | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index dd4f69e06948..4463abd35a64 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2524,17 +2524,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } - size_t dst_stream_idx = 0; - while (dst_stream_idx < sinfo.strm.size() && (uint32_t) sinfo.strm[dst_stream_idx] != strm) { - ++dst_stream_idx; - } - - if (dst_stream_idx == sinfo.strm.size()) { - LLAMA_LOG_ERROR("%s: failed to find destination stream %u in slot info\n", __func__, strm); - return false; - } - - for (auto & layer : layers) { + for (const auto & layer : layers) { uint32_t layer_has_k_idx = 0; io.read(&layer_has_k_idx, sizeof(layer_has_k_idx)); @@ -2586,6 +2576,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } } } + if (!this->v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; From 0c128f100ea195d68cb1ad0da31a71727d0d3804 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 3 Jul 2026 01:39:58 +0200 Subject: [PATCH 6/7] Remove stale comment --- src/llama-kv-cache.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 4463abd35a64..110fadf6eecb 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -259,7 +259,6 @@ llama_kv_cache::llama_kv_cache( } const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il); - // Indexer keys share the SAME buffer as k/v for this layer. ggml_tensor * k_idx = n_embd_k_idx > 0 ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream) : nullptr; From ece3dae8bbde929458d15ddcd2bc1f5ea5326d7e Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 3 Jul 2026 01:51:06 +0200 Subject: [PATCH 7/7] fix compile bug --- src/llama-kv-cache.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 110fadf6eecb..7e50fc4d9b8c 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2559,8 +2559,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 if (k_idx_size_row != expected_k_idx_size_row) { LLAMA_LOG_ERROR( - "%s: mismatched k_idx row size: got %" PRIu64 ", expected %" PRIu64 "\n", - __func__, k_idx_size_row, expected_k_idx_size_row); + "%s: mismatched k_idx row size: got %zu, expected %zu\n", + __func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row); return false; }