tqcq/llama.cpp
0
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-06-26 19:55:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

cont : kv-cells cp/set for non-cont slots

Browse Source 
ggml-ci
This commit is contained in:
Georgi Gerganov
2025-06-21 15:26:01 +03:00
parent f875d6cb72
commit 39d0b1e8df
3 changed files with 64 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/llama-kv-cache-unified.cpp
View File

@ -388,7 +388,8 @@ llama_kv_cache_unified::slot_info_vec_t llama_kv_cache_unified::prepare(const st
struct state {
uint32_t head_old; // old position of the head, before placing the ubatch
uint32_t head_new; // new position of the head, after placing the ubatch
slot_info sinfo; // slot info for the ubatch
llama_kv_cells_unified cells; // copy of the old cells, before placing the ubatch
};
@ -409,13 +410,8 @@ llama_kv_cache_unified::slot_info_vec_t llama_kv_cache_unified::prepare(const st
// remeber the position that we found
res.push_back(sinfo_new);
// TODO: temporary
if (supports_set_rows) {
GGML_ASSERT(sinfo_new.is_cont());
}
// store the old state of the cells in the recovery stack
states.push_back({head, sinfo_new.head(), cells.cp(sinfo_new.head(), ubatch.n_tokens)});
states.push_back({head, sinfo_new, cells.cp(sinfo_new.idxs)});
// now emplace the ubatch
apply_ubatch(sinfo_new, ubatch);
@ -423,7 +419,7 @@ llama_kv_cache_unified::slot_info_vec_t llama_kv_cache_unified::prepare(const st
// iterate backwards and restore the cells to their original state
for (auto it = states.rbegin(); it != states.rend(); ++it) {
cells.set(it->head_new, it->cells);
cells.set(it->sinfo.idxs, it->cells);
head = it->head_old;
}

14
src/llama-kv-cache-unified.h
View File

@ -49,20 +49,6 @@ public:
return idxs.empty();
}
// TODO: tmp until kv cells support non-cont slots
bool is_cont() const {
bool res = true;
for (uint32_t i = 1; i < idxs.size(); ++i) {
if (idxs[i] != idxs[i - 1] + 1) {
res = false;
break;
}
}
return res;
}
// TODO: implement
//std::vector<idx_vec_t> seq_idxs;
};

70
src/llama-kv-cells.h
View File

@ -105,10 +105,29 @@ public:
res.resize(n);
for (uint32_t j = 0; j < n; ++j) {
res.pos[j] = pos[i + j];
res.seq[j] = seq[i + j];
const auto idx = i + j;
assert(shift[i + j] == 0);
res.pos[j] = pos[idx];
res.seq[j] = seq[idx];
assert(shift[idx] == 0);
}
return res;
}
llama_kv_cells_unified cp(const std::vector<uint32_t> & idxs) const {
llama_kv_cells_unified res;
res.resize(idxs.size());
for (uint32_t j = 0; j < idxs.size(); ++j) {
const auto idx = idxs[j];
res.pos[j] = pos[idx];
res.seq[j] = seq[idx];
assert(shift[idx] == 0);
}
return res;
@ -119,26 +138,57 @@ public:
assert(i + other.pos.size() <= pos.size());
for (uint32_t j = 0; j < other.pos.size(); ++j) {
if (pos[i + j] == -1 && other.pos[j] != -1) {
const auto idx = i + j;
if (pos[idx] == -1 && other.pos[j] != -1) {
used.insert(i + j);
}
if (pos[i + j] != -1 && other.pos[j] == -1) {
if (pos[idx] != -1 && other.pos[j] == -1) {
used.erase(i + j);
}
if (pos[i + j] != -1) {
if (pos[idx] != -1) {
seq_pos_rm(i + j);
}
pos[i + j] = other.pos[j];
seq[i + j] = other.seq[j];
pos[idx] = other.pos[j];
seq[idx] = other.seq[j];
if (pos[i + j] != -1) {
if (pos[idx] != -1) {
seq_pos_add(i + j);
}
assert(shift[i + j] == 0);
assert(shift[idx] == 0);
}
}
void set(const std::vector<uint32_t> & idxs, const llama_kv_cells_unified & other) {
assert(idxs.size() == other.pos.size());
for (uint32_t j = 0; j < other.pos.size(); ++j) {
const auto idx = idxs[j];
if (pos[idx] == -1 && other.pos[j] != -1) {
used.insert(idx);
}
if (pos[idx] != -1 && other.pos[j] == -1) {
used.erase(idx);
}
if (pos[idx] != -1) {
seq_pos_rm(idx);
}
pos[idx] = other.pos[j];
seq[idx] = other.seq[j];
if (pos[idx] != -1) {
seq_pos_add(idx);
}
assert(shift[idx] == 0);
}
}