nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 15 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
co-authored by fairydreaming
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
common/common.cpp
+13 -12
View File
@@ -1982,36 +1982,37 @@ bool common_replay_last_token(struct llama_context * ctx, llama_token last_token
bool common_prompt_batch_decode(
struct llama_context * ctx,
const std::vector<llama_token> & tokens,
const std::vector<llama_token> & all_tokens,
int n_new,
int & n_past,
int n_batch,
std::string_view state_path,
bool save_state) {
const int n_eval = tokens.size();
if (n_eval == 0) {
if (n_new == 0) {
return true;
}
const int offset = all_tokens.size() - n_new;
if (save_state && n_eval > 1) {
const int n_tokens_before_last = n_eval - 1;
if (save_state && n_new > 1) {
const int n_tokens_before_last = n_new - 1;
GGML_ASSERT(n_eval <= n_batch);
GGML_ASSERT(n_new <= n_batch);
// Decode all but the last token so we can save the memory state before decoding the last token.
// This is done so we can restore the session state later and replay the last token.
// Memory implementations in recurrent/hybrid models don't support removing tokens from their
// memory, so we can't just remove the last token from the memory and replay the last token which
// is the reason for this logic.
if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(tokens.data()), n_tokens_before_last))) {
if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(all_tokens.data() + offset), n_tokens_before_last))) {
LOG_ERR("%s : failed to eval\n", __func__);
return false;
}
n_past += n_tokens_before_last;
llama_state_save_file(ctx, state_path.data(), tokens.data(), n_tokens_before_last);
LOG_INF("saved session before last token to %s, n_tokens = %d\n", state_path.data(), n_tokens_before_last);
llama_state_save_file(ctx, state_path.data(), all_tokens.data(), all_tokens.size());
LOG_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size());
llama_token last_token = tokens.back();
llama_token last_token = all_tokens.back();
llama_batch batch = llama_batch_get_one(&last_token, 1);
int32_t pos = n_past;
batch.pos = &pos;
@@ -2022,11 +2023,11 @@ bool common_prompt_batch_decode(
}
n_past++;
} else {
if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(tokens.data()), n_eval))) {
if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(all_tokens.data() + offset), n_new))) {
LOG_ERR("%s : failed to eval\n", __func__);
return false;
}
n_past += n_eval;
n_past += n_new;
}
return true;
common/common.h
+2 -1
View File
@@ -930,7 +930,8 @@ void common_batch_add(
// tokens from memory, so this approach works across all model architectures.
bool common_prompt_batch_decode(
struct llama_context * ctx,
const std::vector<llama_token> & embd,
const std::vector<llama_token> & all_tokens,
int n_new,
int & n_past,
int n_batch,
std::string_view state_path,