happyz
/
llama.cpp
mirror of https://github.com/ggerganov/llama.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

kv-cache : bounds-check when accessing slot_info indices

This commit is contained in:
Georgi Gerganov 2025-07-02 13:39:10 +03:00
parent a70293bc25
commit f3da97e61b
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -717,7 +717,7 @@ void llama_kv_cache_unified::apply_ubatch(const slot_info & sinfo, const llama_u
assert(ubatch.n_tokens == sinfo.idxs.size()); assert(ubatch.n_tokens == sinfo.idxs.size());
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
const auto idx = sinfo.idxs[i]; const auto idx = sinfo.idxs.at(i);
if (!cells.is_empty(idx)) { if (!cells.is_empty(idx)) {
assert(cells.seq_count(idx) == 1); assert(cells.seq_count(idx) == 1);
@ -915,7 +915,7 @@ void llama_kv_cache_unified::set_input_k_idxs(ggml_tensor * dst, const llama_uba
int64_t * data = (int64_t *) dst->data; int64_t * data = (int64_t *) dst->data;
for (int64_t i = 0; i < n_tokens; ++i) { for (int64_t i = 0; i < n_tokens; ++i) {
data[i] = sinfo.idxs[i]; data[i] = sinfo.idxs.at(i);
} }
} }
@ -930,7 +930,7 @@ void llama_kv_cache_unified::set_input_v_idxs(ggml_tensor * dst, const llama_uba
int64_t * data = (int64_t *) dst->data; int64_t * data = (int64_t *) dst->data;
for (int64_t i = 0; i < n_tokens; ++i) { for (int64_t i = 0; i < n_tokens; ++i) {
data[i] = sinfo.idxs[i]; data[i] = sinfo.idxs.at(i);
} }
} }

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

@ -42,7 +42,7 @@ public:
idx_vec_t idxs; idx_vec_t idxs;
uint32_t head() const { uint32_t head() const {
return idxs[0]; return idxs.at(0);
} }
bool empty() const { bool empty() const {