Skip to content

kv-cache: add LLAMA_STATE_SEQ_FLAGS_APPEND for incremental state restore #21221

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
sammshen wants to merge 1 commit into
base: master
Choose a base branch
from
+9 −4
Open

kv-cache: add LLAMA_STATE_SEQ_FLAGS_APPEND for incremental state restore #21221

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
3 changes: 3 additions & 0 deletions include/llama.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -872,6 +872,9 @@ extern "C" {
// work only with partial states, such as SWA KV cache or recurrent cache (e.g. Mamba)
#define LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY 1

// restore without clearing existing sequence data (append to existing KV cache entries)
#define LLAMA_STATE_SEQ_FLAGS_APPEND 2

typedef uint32_t llama_state_seq_flags;

LLAMA_API size_t llama_state_seq_get_size_ext(
Expand Down
8 changes: 5 additions & 3 deletions src/llama-kv-cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1722,7 +1722,7 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
slot_info sinfo;

bool res = true;
res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id, flags);
res = res && state_read_data(io, strm, cell_count, sinfo);

if (!res) {
Expand Down Expand Up @@ -1868,13 +1868,15 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
}
}

bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id) {
bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id, llama_state_seq_flags flags) {
auto & cells = v_cells[strm];
auto & head = v_heads[strm];

if (dest_seq_id != -1) {
// single sequence
seq_rm(dest_seq_id, -1, -1);
if (!(flags & LLAMA_STATE_SEQ_FLAGS_APPEND)) {
seq_rm(dest_seq_id, -1, -1);
}

llama_batch_allocr balloc(hparams.n_pos_per_embd());

Expand Down
2 changes: 1 addition & 1 deletion src/llama-kv-cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,7 @@ class llama_kv_cache : public llama_memory_i {
void state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id = -1) const;
void state_write_data(llama_io_write_i & io, const cell_ranges_t & cr) const;

bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1);
bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1, llama_state_seq_flags flags = 0);

@sammshen sammshen Mar 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

default no change

bool state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo);
};

Expand Down