nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

speculative : fix out-of-bounds read in ngram-map on prompt shrink (#23936)

Browse Source 
* speculative : fix out-of-bounds read in ngram-map on prompt shrink

* speculative : fix ngram-map cleanup cutoff after prompt shrink
This commit is contained in:
o7si
2026-07-07 15:25:04 +08:00
committed by GitHub
parent a8cfdbb9e4
commit defa95c306
1 changed files with 15 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
common/ngram-map.cpp
View File

@@ -125,6 +125,16 @@ void common_ngram_map_begin(
LOG_DBG("%s: begin, idx_last_draft=%zu, new begin=%zu, #keys=%zu\n", __func__, LOG_DBG("%s: begin, idx_last_draft=%zu, new begin=%zu, #keys=%zu\n", __func__,
map.idx_last_check, size_begin, map.keys.size()); map.idx_last_check, size_begin, map.keys.size());
size_t idx_begin_cleanup = map.size_last_begin;
if (idx_begin_cleanup > size_begin) {
if (size_begin > (size_t) map.size_key + map.size_value) {
idx_begin_cleanup = size_begin - map.size_key - map.size_value;
} else {
idx_begin_cleanup = 0;
}
LOG_INF("%s: shrink cleanup begin: %zu -> %zu\n", __func__, map.size_last_begin, idx_begin_cleanup);
}
size_t count_map_entries_upd = 0; size_t count_map_entries_upd = 0;
if (!map.key_map.empty() && size_begin < map.idx_last_check) { if (!map.key_map.empty() && size_begin < map.idx_last_check) {
if (map.show_key_map_stats) { if (map.show_key_map_stats) {
@@ -150,27 +160,23 @@ void common_ngram_map_begin(
// Update the map from hash to key index (clear outdated entries). // Update the map from hash to key index (clear outdated entries).
for (size_t i = 0; i < map.key_map.size(); ++i) { for (size_t i = 0; i < map.key_map.size(); ++i) {
uint32_t key_idx = map.key_map[i]; uint32_t key_idx = map.key_map[i];
if (key_idx >= map.size_last_begin) { if (key_idx != 0 && key_idx >= idx_begin_cleanup) {
map.key_map[i] = 0; map.key_map[i] = 0;
count_map_entries_upd++; count_map_entries_upd++;
} }
} }
map.key_map_last_idx = (map.size_last_begin > 0) ? map.size_last_begin - 1 : 0; map.key_map_last_idx = (idx_begin_cleanup > 0) ? (uint32_t) (idx_begin_cleanup - 1) : 0;
} }
if (size_begin < map.idx_last_check && !map.keys.empty()) { if (size_begin < map.idx_last_check && !map.keys.empty()) {
// The next token generation will start at index size_begin.
// The tokens between map.size_last_begin and size_begin are no longer valid.
//
// Refresh map: Remove all entries with index >= map.size_last_begin.
size_t count_keys = map.keys.size(); size_t count_keys = map.keys.size();
size_t count_keys_del = 0; size_t count_keys_del = 0;
size_t count_values_del = 0; size_t count_values_del = 0;
for (int32_t i = map.keys.size() - 1; i >= 0; --i) { for (int32_t i = map.keys.size() - 1; i >= 0; --i) {
common_ngram_map_key & key = map.keys[i]; common_ngram_map_key & key = map.keys[i];
if (key.key_idx >= map.size_last_begin) { if (key.key_idx >= idx_begin_cleanup) {
// Delete the key. // Delete the key.
LOG_DBG("%s: delete key %d at index %zu (>= size_last_begin=%zu)\n", __func__, i, key.key_idx, map.size_last_begin); LOG_DBG("%s: delete key %d at index %zu (>= idx_begin_cleanup=%zu)\n", __func__, i, key.key_idx, idx_begin_cleanup);
map.keys.erase(map.keys.begin() + i); map.keys.erase(map.keys.begin() + i);
count_keys_del++; count_keys_del++;
continue; continue;
@@ -182,7 +188,7 @@ void common_ngram_map_begin(
// Check the indices of the values. // Check the indices of the values.
for (int16_t j = COMMON_NGRAM_MAX_VALUES - 1; j >= 0; --j) { for (int16_t j = COMMON_NGRAM_MAX_VALUES - 1; j >= 0; --j) {
common_ngram_map_value & value = key.values[j]; common_ngram_map_value & value = key.values[j];
if (value.value_idx >= map.size_last_begin) { if (value.value_idx != 0 && value.value_idx >= idx_begin_cleanup) {
// Delete the value. // Delete the value.
count_values_del++; count_values_del++;