kv-cache : improve find_slot impl
This commit is contained in:
parent
2ac5be3a58
commit
a70293bc25
|
|
@ -624,14 +624,15 @@ llama_kv_cache_unified::slot_info llama_kv_cache_unified::find_slot(const llama_
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
uint32_t n_found = 0;
|
|
||||||
uint32_t n_tested = 0;
|
uint32_t n_tested = 0;
|
||||||
|
|
||||||
const uint32_t n_test = cont ? n_tokens : 1;
|
const uint32_t n_test = cont ? n_tokens : 1;
|
||||||
|
|
||||||
slot_info res;
|
slot_info res;
|
||||||
|
|
||||||
res.idxs.resize(n_tokens);
|
auto & idxs = res.idxs;
|
||||||
|
|
||||||
|
idxs.reserve(n_tokens);
|
||||||
|
|
||||||
while (true) {
|
while (true) {
|
||||||
if (head_cur + n_test > cells.size()) {
|
if (head_cur + n_test > cells.size()) {
|
||||||
|
|
@ -677,20 +678,18 @@ llama_kv_cache_unified::slot_info llama_kv_cache_unified::find_slot(const llama_
|
||||||
n_tested++;
|
n_tested++;
|
||||||
|
|
||||||
if (can_use) {
|
if (can_use) {
|
||||||
res.idxs[n_found] = idx;
|
idxs.push_back(idx);
|
||||||
|
|
||||||
n_found++;
|
|
||||||
} else {
|
} else {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if (n_found == n_tokens) {
|
if (idxs.size() == n_tokens) {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
if (cont) {
|
if (cont) {
|
||||||
n_found = 0;
|
idxs.clear();
|
||||||
}
|
}
|
||||||
|
|
||||||
if (n_tested >= cells.size()) {
|
if (n_tested >= cells.size()) {
|
||||||
|
|
@ -700,7 +699,7 @@ llama_kv_cache_unified::slot_info llama_kv_cache_unified::find_slot(const llama_
|
||||||
}
|
}
|
||||||
|
|
||||||
// we didn't find a suitable slot - return empty result
|
// we didn't find a suitable slot - return empty result
|
||||||
if (n_found < n_tokens) {
|
if (idxs.size() < n_tokens) {
|
||||||
res.clear();
|
res.clear();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue