nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

common : fix state save in common_prompt_batch_decode (#23468)

Browse Source 
* common : fix state save in common_prompt_batch_decode

This commit addresses a bug in common_prompt_batch_decode that affects
the session state store/restore in completion.cpp and
save-load-state.cpp.

The motivation for this is that currently the code is saving n-1 tokens
in both the session_tokens and in the KV cache. Then when loading the
session tokens, and if the prompt matches, it would replay the last
saved token (n-1) into the next position, effectively replaying the
same token in the wrong position.

The fix is to store all n tokens in session_tokens, while the memory
state only reflects n-1 processed tokens as the saving happens before
the last token is decoded in common_prompt_batch_decode.

I ran both completion.cpp and save-load-state.cpp with a transformer, a
recurrent, and a hybrid model.

Resolves: https://github.com/ggml-org/llama.cpp/issues/23400

Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com>
This commit is contained in:
Daniel Bevenius
2026-06-02 15:44:15 +02:00
committed by GitHub
parent 60130d18f9
commit 0b7154066e
4 changed files with 31 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tools/completion/completion.cpp
View File

@@ -373,16 +373,10 @@ int llama_completion(int argc, char ** argv) {
__func__, n_match, embd_inp.size());
}
if (session_tokens.size() == n_match) {
// [TAG_CONTEXT_STATE_LOGITS]
// in this case, we are going to reuse the logits from the session
// if we ever decide to remove the logits from the session, we need to handle this somehow
// ref: https://github.com/ggml-org/llama.cpp/pull/18862#issuecomment-3756330941
}
// remove any "future" tokens that we might have inherited from the previous session
if (session_tokens.size() > n_match) {
if (!llama_memory_seq_rm(mem, -1, n_match, -1)) {
llama_pos pos = n_match > 0 ? (llama_pos)(n_match - 1) : 0;
if (!llama_memory_seq_rm(mem, -1, pos, -1)) {
LOG_WRN("%s: unable to reuse common prefix (for example, when the memory is recurrent)\n", __func__);
llama_memory_clear(mem, true);
session_tokens.clear();
@@ -398,7 +392,7 @@ int llama_completion(int argc, char ** argv) {
// Logits are not stored as part of the session state so we need to
// "replay" the last token to get logits for sampling.
if (!session_tokens.empty() && n_match > 0 && n_match == session_tokens.size()) {
if (!common_replay_last_token(ctx, session_tokens.back(), n_match)) {
if (!common_replay_last_token(ctx, session_tokens.back(), n_match - 1)) {
return 1;
}
@@ -695,12 +689,14 @@ int llama_completion(int argc, char ** argv) {
if (!embd.empty()) {
const bool is_last_batch = (n_consumed >= (int) embd_inp.size());
const bool save_now = session_do_save && is_last_batch;
if (!common_prompt_batch_decode(ctx, embd, n_past, params.n_batch, path_session, save_now)) {
session_tokens.insert(session_tokens.end(), embd.begin(), embd.end());
if (!common_prompt_batch_decode(ctx, session_tokens, embd.size(), n_past, params.n_batch, path_session, save_now)) {
return 1;
}
session_tokens.insert(session_tokens.end(), embd.begin(), embd.end());
n_session_consumed = session_tokens.size();
session_do_save = false;
n_session_consumed += embd.size();
if (save_now) {
session_do_save = false;
}
LOG_DBG("n_past = %d\n", n_past);
@@ -991,7 +987,10 @@ int llama_completion(int argc, char ** argv) {
if (!path_session.empty() && params.prompt_cache_all && !params.prompt_cache_ro) {
LOG("\n%s: saving final output to session file '%s'\n", __func__, path_session.c_str());
session_tokens.insert(session_tokens.end(), embd.begin(), embd.end());
llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size());
LOG_INF("saved final session to %s, n_tokens = %ld\n", path_session.data(), session_tokens.size());
}
LOG("\n\n");