vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix restoring zero cell count

Browse source
This commit is contained in:
Jan Boon 2024-03-30 02:35:33 +08:00
parent a71ec3db7b
commit bf1d4932f8
1 changed files with 18 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama.cpp
View file

@ -15283,7 +15283,7 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
} }
// Allocate the new cells for the slot // Allocate the new cells for the slot
{ if (cell_count) {
llama_batch batch = llama_batch_init(cell_count, 0, 1); llama_batch batch = llama_batch_init(cell_count, 0, 1);
batch.n_tokens = cell_count; batch.n_tokens = cell_count;
for (uint32_t i = 0; i < cell_count; ++i) { for (uint32_t i = 0; i < cell_count; ++i) {
@ -15327,10 +15327,12 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
return 0; return 0;
} }
if (cell_count) {
// Read and set the keys for the whole cell range // Read and set the keys for the whole cell range
ggml_backend_tensor_set(kv_self.k_l[il], inp, kv_head * k_size_row, cell_count * k_size_row); ggml_backend_tensor_set(kv_self.k_l[il], inp, kv_head * k_size_row, cell_count * k_size_row);
inp += cell_count * k_size_row; inp += cell_count * k_size_row;
} }
}
// For each layer, read the values for each cell (transposed) // For each layer, read the values for each cell (transposed)
for (int il = 0; il < (int)n_layer; ++il) { for (int il = 0; il < (int)n_layer; ++il) {
@ -15339,6 +15341,7 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
memcpy(&v_size_el_ref, inp, sizeof(v_size_el_ref)); memcpy(&v_size_el_ref, inp, sizeof(v_size_el_ref));
inp += sizeof(v_size_el_ref); inp += sizeof(v_size_el_ref);
if (cell_count) {
const size_t v_size_el = ggml_type_size(kv_self.v_l[il]->type); const size_t v_size_el = ggml_type_size(kv_self.v_l[il]->type);
if (v_size_el != v_size_el_ref) { if (v_size_el != v_size_el_ref) {
llama_kv_cache_seq_rm(kv_self, dest_seq_id, -1, -1); llama_kv_cache_seq_rm(kv_self, dest_seq_id, -1, -1);
@ -15352,6 +15355,7 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src,
inp += cell_count * v_size_el; inp += cell_count * v_size_el;
} }
} }
}
const size_t nread = inp - src; const size_t nread = inp - src;
return nread; return nread;