Is it reasonable to add C APIs to get & set kv-cache?

[AsakusaRinne]

Greetings!

I'm openning this discussion to ask if it's reasonable to add a C API to get and set kv_cache directly. :)

Description

Currently, there are following C APIs related with kv-cache (only list part of them here).

1. llama_model_kv_override allows users to override the key-value pairs of the model meta data.
2. llama_kv_cache_seq_add and llama_kv_cache_seq_div allow users to manipulate the kv-cache directly.
3. llama_copy_state_data and llama_set_state_data allow users to get and set the state data, including embeddings, logits and kv-cache.

Since there is already an API to get and set the state data, including kv-cache, and the kv-cache does allow to be modified from outside, I wonder if it's possible to add the following three APIs.

// Get the kv_cache of specified sequence ids.
// seq_ids: the sequence ids of the kv-cache in the batch.
// If `seq_id == nullptr`, it will return the full kv-cache of all sequences.
LLAMA_API llama_kv_cache_partial llama_get_kv_cache(
            struct llama_context * ctx,
                         llama_seq_id*  seq_ids);

// Set the kv_cache to the specified sequence.
// src: the source buffer which stores the kv-cache to be set.
// Return true if the kv-cache is valid and the operation is successful.
LLAMA_API bool llama_set_kv_cache(
            struct llama_context * ctx,
                   const llama_kv_cache_partial * src);


struct llama_kv_cache_partial {
    // The data type of the kv-cache, in which `llama_model_kv_type` is an enum.
    llama_model_kv_type data_type;

    // Number of sequences in the kv-cache part.
    int32_t n_seq;

    // seq_ids for each sequence in this kv-cache part.
    llama_seq_id * seq_ids;

    // length of the keys for each sequence in byte.
    llama_seq_id * key_lengths;

    // length of the values for each sequence in byte.
    llama_seq_id * value_lengths;

    // key data for each sequence.
    uint8_t ** keys;

    // key data for each sequence.
    uint8_t ** values;
};

It's not a formal proposal because llama.cpp actually uses cells to store the kv-cache and padding is required when setting kv-cache. If it's at least possible to be implemented, I'd like to step further to improve the proposal and contribute.

I have searched the issues and discussions, but I didn't find related topics. I'll appreciate it if you could let me know if there is any similar discussion that I missed.

Application case

As for the application of this API, let's consider a server using batched decoding. At a moment, there are three users online, which are a, b and c. Therefore, their quiries will be put into one batch to generate the response.

Then, a and b log in again after being offline for some time. Besides, at this moment, d and e are also online.

a and b want to continue their previous sessions but no cache was found, while d and e already has some kv-cache in the memory. In this case, to get the best performance, saving kv-cache is a good option. The state data from llama_copy_state_data is for all the sequences, including a's, b's and c's. However, loading the state will remove the state of d and e, and what's more, will load an unexpected state from c. Thus, it sounds reasonable to add the API to get and set kv-cache directly.

Conclusion

As a conclusion, here are my questions.

1. Is it reasonable to add the C API to get and set kv_cache directly?
2. In the case I described above, is embeddings and logits also necessary to get and set, rather than only kv-cache?
3. If it's not reasonable to add such APIs, is there any way to deal with the case I described above?

Any suggestion will be appreciated!

Best regards,
Rinne

[AsakusaRinne]

Sorry, I didn't notice the newly added APIs, as below. I'm going to close this discussion now and try to work with these APIs. I'll reopen this if it proves that the APIs I described is indeed necessary.

    LLAMA_API size_t llama_state_seq_set_data(
            struct llama_context * ctx,
                   const uint8_t * src,
                    llama_seq_id   dest_seq_id);

    LLAMA_API size_t llama_state_seq_save_file(
            struct llama_context * ctx,
                      const char * filepath,
                    llama_seq_id   seq_id,
               const llama_token * tokens,
                          size_t   n_token_count);

    LLAMA_API size_t llama_state_seq_load_file(
            struct llama_context * ctx,
                      const char * filepath,
                    llama_seq_id   dest_seq_id,
                     llama_token * tokens_out,
                          size_t   n_token_capacity,
                          size_t * n_token_count_out);