kv-cache : SWA checkpoints store only non-masked cells (#23981)

This commit is contained in:
Georgi Gerganov 2026-06-02 11:06:29 +03:00 committed by GitHub
parent f7a0777a5c
commit 2365315955
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1876,7 +1876,19 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
uint32_t cell_range_begin = cells.size();
for (uint32_t i = 0; i < cells.size(); ++i) {
if (!cells.is_empty(i) && (seq_id == -1 || cells.seq_has(i, seq_id))) {
bool add_cell = true;
add_cell = add_cell && !cells.is_empty(i);
add_cell = add_cell && (seq_id == -1 || cells.seq_has(i, seq_id));
// check the cell is not SWA-masked
if (add_cell && seq_id != -1) {
const bool is_masked = llama_hparams::is_masked_swa(n_swa, swa_type, cells.pos_get(i), cells.seq_pos_max(seq_id));
add_cell = !is_masked;
}
if (add_cell) {
++cell_count;
if (cell_range_begin == cells.size()) {
cell_range_begin = i;
@ -2129,7 +2141,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
sinfo = find_slot(ubatch, false);
if (sinfo.empty()) {
LLAMA_LOG_ERROR("%s: failed to find available cells in kv cache\n", __func__);
LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count);
return false;
}