Saving state

[HarewVlad]

Hi!
In llama_copy_state_data i see this:

// v is not contiguous, copy row by row
tmp_buf.resize(elt_size*kv_head);
for (int ir = 0; ir < (int) n_embd_v_gqa; ++ir) {
    ggml_backend_tensor_get(kv_self.v_l[il], tmp_buf.data(), ir*elt_size*n_ctx, tmp_buf.size());
    data_ctx->write(tmp_buf.data(), tmp_buf.size());
}

Can it be converted to something like this to reduce time to get V cache?

tmp_buf.resize(n_embd_v_gqa*elt_size*n_ctx);
ggml_backend_tensor_get(kv_self.v_l[il], tmp_buf.data(), 0, tmp_buf.size());
data_ctx->write(tmp_buf.data(), tmp_buf.size());

[ggerganov]

Yes, it's also more correct since kv_head does not necessarily correspond to the size of the KV cache. It works for main since the generation is (almost) always single-sequence.

It will use more data though since we will be storing the full cache regardless if it has been used or not. There should be some logic to skip empty cells. Overall, saving the state could be improved