Skip to content

memory: clear rollback markers for every shared sequence #365

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bri-prism wants to merge 1 commit into
base: seqcp-rollback-mtp-depth
Choose a base branch
from
Open

memory: clear rollback markers for every shared sequence #365

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/llama-memory-recurrent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1268,8 +1268,12 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const {
const llama_seq_id seq = *mem->cells[cell_idx].seq_id.begin();
if (seq >= 0 && (size_t) seq < mem->rs_idx.size()) {
idx = mem->rs_idx[seq];
// reset rollback idx
mem->rs_idx[seq] = 0;
}
// The restored state advances every sequence sharing this cell.
for (const llama_seq_id id : mem->cells[cell_idx].seq_id) {
if (id >= 0 && (size_t) id < mem->rs_idx.size()) {
mem->rs_idx[id] = 0;
}
}
}
return (int32_t)(idx * mem->size) + src0;
Expand Down
28 changes: 19 additions & 9 deletions tests/test-recurrent-state-rollback.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,10 +207,11 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode
return true;
}

// Fork a sequence with seq_cp while its partial rollback is pending, then decode one token on the fork.
// The reference rolls back seq 1 directly. The fork must match it bitwise.
// The tiny test weights make the recurrent part too small to show in the logits, so compare the state bytes.
static bool test_seq_cp_pending_rollback(const common_params & params, llama_model * model, const int n_vocab) {
// Tiny test weights can hide recurrent errors in the logits, so compare state bytes after a fork and optional shared continuation.
static bool test_seq_cp_pending_rollback(const common_params & params,
llama_model * model,
const int n_vocab,
bool shared) {
constexpr uint32_t n_prompt = 12;
constexpr uint32_t n_rollback = 3;
constexpr llama_pos p0 = n_prompt - n_rollback;
Expand All @@ -221,7 +222,7 @@ static bool test_seq_cp_pending_rollback(const common_params & params, llama_mod
cparams.n_ctx = 256;
cparams.n_batch = 256;
cparams.n_ubatch = 256;
cparams.kv_unified = false;
cparams.kv_unified = shared;

// fork: seq 0 -> rollback -> seq_cp to seq 1, ref: seq 1 -> rollback
std::vector<uint8_t> state[2];
Expand All @@ -235,7 +236,7 @@ static bool test_seq_cp_pending_rollback(const common_params & params, llama_mod
llama_memory_t mem = llama_get_memory(ctx);
const llama_seq_id src = fork ? 0 : 1;

llama_batch batch = llama_batch_init(n_prompt, 0, 1);
llama_batch batch = llama_batch_init(n_prompt, 0, 2);
for (llama_pos pos = 0; pos < (llama_pos) n_prompt; ++pos) {
common_batch_add(batch, (llama_token) ((5 * pos + 3) % n_vocab), pos, { src }, false);
}
Expand All @@ -247,8 +248,15 @@ static bool test_seq_cp_pending_rollback(const common_params & params, llama_mod
}

common_batch_clear(batch);
common_batch_add(batch, (llama_token) (p0 % n_vocab), p0, { 1 }, true);
const std::vector<llama_seq_id> ids =
shared && fork ? std::vector<llama_seq_id>{ 0, 1 } : std::vector<llama_seq_id>{ 1 };
common_batch_add(batch, (llama_token) (p0 % n_vocab), p0, ids, true);
ok = ok && llama_decode(ctx, batch) == 0;
if (shared) {
common_batch_clear(batch);
common_batch_add(batch, (llama_token) ((p0 + 1) % n_vocab), p0 + 1, { 1 }, true);
ok = ok && llama_decode(ctx, batch) == 0;
}
llama_batch_free(batch);

if (ok) {
Expand All @@ -269,7 +277,8 @@ static bool test_seq_cp_pending_rollback(const common_params & params, llama_mod
return false;
}

fprintf(stderr, "%s : forked seq matched the rolled-back state (%zu bytes)\n", __func__, state[0].size());
fprintf(stderr, "%s : forked seq matched the rolled-back state (shared=%d, %zu bytes)\n", __func__, shared,
state[0].size());
return true;
}

Expand Down Expand Up @@ -463,7 +472,8 @@ int main(int argc, char ** argv) {
return 1;
}

if (!test_seq_cp_pending_rollback(params, model, n_vocab)) {
if (!test_seq_cp_pending_rollback(params, model, n_vocab, false) ||
!test_seq_cp_pending_rollback(params, model, n_vocab, true)) {
return 1;
}

Expand Down