diff --git a/include/llama.h b/include/llama.h index 05eb3b8fd..4573775dd 100644 --- a/include/llama.h +++ b/include/llama.h @@ -47,10 +47,10 @@ #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' #define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN -#define LLAMA_SESSION_VERSION 9 +#define LLAMA_SESSION_VERSION 10 #define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ -#define LLAMA_STATE_SEQ_VERSION 3 +#define LLAMA_STATE_SEQ_VERSION 4 #define LLAMA_SERVER_MAGIC 0x6c6d7376u // 'lmsv' #define LLAMA_SERVER_VERSION 1 diff --git a/src/llama.cpp b/src/llama.cpp index 4987911c8..0e66a7eed 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -9277,6 +9277,29 @@ struct llama_data_write { } } } + + // DSA lightning-indexer key cache (kv_self.kr_l): one row per cell, like K/V. + const uint32_t dsa_indexer_state = !kv_self.kr_l.empty() ? 1 : 0; + write(&dsa_indexer_state, sizeof(dsa_indexer_state)); + + if (dsa_indexer_state != 0) { + for (uint32_t il = 0; il < n_layer; ++il) { + const bool has_kr_cache = need_kv && il < kv_self.kr_l.size() && kv_self.kr_l[il] != nullptr; + if (!has_kr_cache) continue; + + const int32_t kr_type_i = has_kr_cache ? (int32_t) kv_self.kr_l[il]->type : -1; + write(&kr_type_i, sizeof(kr_type_i)); + + const uint64_t kr_size_row = has_kr_cache ? ggml_row_size(kv_self.kr_l[il]->type, kv_self.kr_l[il]->ne[0]) : 0; + write(&kr_size_row, sizeof(kr_size_row)); + + for (const auto & range : cell_ranges) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * kr_size_row; + write_tensor_data(kv_self.kr_l[il], range.first * kr_size_row, buf_size, il); + } + } + } } void write_kv_cache(const struct llama_context * ctx, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) { @@ -9806,6 +9829,48 @@ struct llama_data_read { } } } + + // DSA lightning-indexer key cache (kv_self.kr_l). + uint32_t dsa_indexer_state_ref = 0; + read_to(&dsa_indexer_state_ref, sizeof(dsa_indexer_state_ref)); + + const bool has_dsa_indexer = !kv_self.kr_l.empty(); + if ((dsa_indexer_state_ref != 0) != has_dsa_indexer) { + LLAMA_LOG_ERROR("%s: incompatible DSA indexer key cache presence (file=%u, model=%d)\n", + __func__, dsa_indexer_state_ref, has_dsa_indexer ? 1 : 0); + return false; + } + + if (dsa_indexer_state_ref != 0) { + for (uint32_t il = 0; il < n_layer; ++il) { + const bool has_kr_cache = need_kv && il < kv_self.kr_l.size() && kv_self.kr_l[il] != nullptr; + if (!has_kr_cache) continue; + + int32_t kr_type_i_ref; + read_to(&kr_type_i_ref, sizeof(kr_type_i_ref)); + const int32_t kr_type_i = (int32_t) kv_self.kr_l[il]->type; + if (kr_type_i != kr_type_i_ref) { + LLAMA_LOG_ERROR("%s: mismatched DSA indexer key type (%d != %d, layer %d)\n", __func__, kr_type_i, kr_type_i_ref, il); + return false; + } + + uint64_t kr_size_row_ref; + read_to(&kr_size_row_ref, sizeof(kr_size_row_ref)); + const size_t kr_size_row = ggml_row_size(kv_self.kr_l[il]->type, kv_self.kr_l[il]->ne[0]); + if (kr_size_row != kr_size_row_ref) { + LLAMA_LOG_ERROR("%s: mismatched DSA indexer key row size (%zu != %zu, layer %d)\n", __func__, kr_size_row, (size_t) kr_size_row_ref, il); + return false; + } + + if (cell_count) { + if (kv_self.kr_l[il]->extra) { + read_kv_cache_data_split(ctx, kv_self.kr_l[il], read(cell_count * kr_size_row), kv_self.head, kr_size_row, cell_count, il); + } else { + ggml_backend_tensor_set(kv_self.kr_l[il], read(cell_count * kr_size_row), kv_self.head * kr_size_row, cell_count * kr_size_row); + } + } + } + } return true; }