furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-20 09:45:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

include DSA indexer state in kv/slot serializer (#2146)

Browse Source
This commit is contained in:
usrlocalben
2026-07-18 09:01:26 +03:00
committed by GitHub
parent fbcc743c70
commit 65891dcdb3
2 changed files with 67 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/llama.h
+2 -2
View File
@@ -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
src/llama.cpp
+65
View File
@@ -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;
}