Felipe97/llama-cpp-compiled

1#include "common.h"2#include "log.h"3#include "ngram-map.h"4 5#include <cinttypes>6#include <cstdint>7#include <cstdio>8#include <sstream>9 10// prime number used for LCG hash function (32 bit), it is near (sqrt(5) - 1)/2 * 2^32.11#define LCG_FACTOR 2654435761UL12 13// Compute the LCG hash of a n-gram of size len at offset start.14static uint32_t common_ngram_map_hash(const llama_tokens & tokens, size_t start, size_t len) {15    uint32_t hash = 0;16    for (size_t i = 0; i < len; ++i) {17        hash = hash * LCG_FACTOR + tokens[start + i];18    }19    return hash;20}21 22// Print the values of a sublist of `llama_tokens & inp` to a string in the form [v0, v1, v2, ...].23static std::string common_tokens_to_str(const llama_tokens & inp, size_t start, size_t length) {24    std::ostringstream oss;25    oss << '[';26    for (size_t i = 0; i < length; ++i) {27        if (i > 0) {28            oss << ", ";29        }30        oss << inp[start + i];31    }32    oss << ']';33    return oss.str();34}35 36 37// n-gram simple38//39 40/**41 * Perform speculative generation using the model's own token history.42 * Searches for a matching pattern in the token history and returns draft tokens.43 *44 * @param state     Current state of this implementation45 * @param tokens    Token history to search in46 * @param sampled   Last sampled token47 * @return Vector of draft tokens, empty if no matching pattern is found48 */49llama_tokens common_ngram_simple_draft(50        const common_ngram_simple_config & config,51        const llama_tokens & tokens, llama_token sampled) {52 53    // Simple implementation of self-speculative decoding without a draft model.54    //55    const size_t cur_len = tokens.size();56 57    const size_t n_draft_min = config.size_ngram; // size of n-gram to lookup in token history58    const size_t n_draft_max = config.size_mgram; // the m-gram following the found n-gram is used for draft59 60    // vector for tokens we want to verify.61    // return empty vector if there is no match.62    llama_tokens draft_tokens;63 64    // We need at least n_draft_min + n_draft_max + 1 tokens.65    if (cur_len <= static_cast<size_t>(n_draft_min + n_draft_max + 1)) {66        return draft_tokens;67    }68 69    // pattern search70    llama_tokens pattern;71    pattern.reserve(n_draft_min);72    for (size_t j = cur_len - n_draft_min + 1; j < cur_len; ++j) {73        pattern.push_back(tokens[j]);74    }75    pattern.push_back(sampled); // add the last token to the pattern76 77    size_t match_pos = 0; // we ignore position 0, position 0 == no match78                          // search backwards, but skip the current match (we are currently there)79    for (size_t j = cur_len - n_draft_min - 1; j > 0; --j) {80        bool match = true;81        for (size_t k = 0; k < pattern.size(); ++k) {82            if (tokens[j + k] != pattern[k]) {83                match = false;84                break;85            }86        }87        if (match) {88            match_pos = j;89            break;90        }91    }92    if (match_pos == 0) {93        return draft_tokens;94    }95 96    const size_t copy_max = std::min(97            n_draft_max,98            cur_len - (match_pos + n_draft_min)99            );100    if (copy_max < n_draft_min) {101        return draft_tokens;102    }103    LOG_DBG("%s: #tokens = %zu: found matching pattern at pos %zu, length %zu, draft length %zu\n",104            __func__, cur_len,105            match_pos, pattern.size(), copy_max);106 107    draft_tokens.reserve(copy_max);108    for (size_t j = 0; j < copy_max; ++j) {109        draft_tokens.push_back(tokens[match_pos + n_draft_min + j]);110    }111    return draft_tokens;112}113 114 115// n-gram map116//117 118// maximum number of counted values of a ngram map value.119#define COMMON_NGRAM_MAX_VALUE_COUNT 16380120 121void common_ngram_map_begin(122    common_ngram_map & map, const llama_tokens & tokens) {123    size_t size_begin = tokens.size();124 125    LOG_DBG("%s: begin, idx_last_draft=%zu, new begin=%zu, #keys=%zu\n", __func__,126            map.idx_last_check, size_begin, map.keys.size());127 128    size_t idx_begin_cleanup = map.size_last_begin;129    if (idx_begin_cleanup > size_begin) {130        if (size_begin > (size_t) map.size_key + map.size_value) {131            idx_begin_cleanup = size_begin - map.size_key - map.size_value;132        } else {133            idx_begin_cleanup = 0;134        }135        LOG_INF("%s: shrink cleanup begin: %zu -> %zu\n", __func__, map.size_last_begin, idx_begin_cleanup);136    }137 138    size_t count_map_entries_upd = 0;139    if (!map.key_map.empty() && size_begin < map.idx_last_check) {140        if (map.show_key_map_stats) {141            // Print statistics of hash map map_key.142            size_t count_nonzero = 0;143            uint32_t min_idx = UINT32_MAX;144            uint32_t max_idx = 0;145            for (size_t i = 0; i < map.key_map.size(); ++i) {146                uint32_t key_idx = map.key_map[i];147                if (key_idx != 0) {148                    ++count_nonzero;149                    if (key_idx < min_idx) min_idx = key_idx;150                    if (key_idx > max_idx) max_idx = key_idx;151                }152            }153            if (count_nonzero == 0) {154                min_idx = 0;155            }156            LOG_INF("%s: key_map stats: entries=%zu, min_idx=%u, max_idx=%u, key_map_last_idx=%u\n",157                    __func__, count_nonzero, min_idx, max_idx, map.key_map_last_idx);158        }159 160        // Update the map from hash to key index (clear outdated entries).161        for (size_t i = 0; i < map.key_map.size(); ++i) {162            uint32_t key_idx = map.key_map[i];163            if (key_idx != 0 && key_idx >= idx_begin_cleanup) {164                map.key_map[i] = 0;165                count_map_entries_upd++;166            }167        }168        map.key_map_last_idx = (idx_begin_cleanup > 0) ? (uint32_t) (idx_begin_cleanup - 1) : 0;169    }170 171    if (size_begin < map.idx_last_check && !map.keys.empty()) {172        size_t count_keys = map.keys.size();173        size_t count_keys_del = 0;174        size_t count_values_del = 0;175        for (int32_t i = map.keys.size() - 1; i >= 0; --i) {176            common_ngram_map_key & key = map.keys[i];177            if (key.key_idx >= idx_begin_cleanup) {178                // Delete the key.179                LOG_DBG("%s: delete key %d at index %zu (>= idx_begin_cleanup=%zu)\n", __func__, i, key.key_idx, idx_begin_cleanup);180                map.keys.erase(map.keys.begin() + i);181                count_keys_del++;182                continue;183            }184            if (map.key_only) {185                continue;186            }187 188            // Check the indices of the values.189            for (int16_t j = COMMON_NGRAM_MAX_VALUES - 1; j >= 0; --j) {190                common_ngram_map_value & value = key.values[j];191                if (value.value_idx != 0 && value.value_idx >= idx_begin_cleanup) {192                    // Delete the value.193                    count_values_del++;194 195                    // Move all values after this value to the left.196                    for (uint16_t k = j; k < COMMON_NGRAM_MAX_VALUES - 1; ++k) {197                        key.values[k] = key.values[k + 1];198                    }199                    // Clear the last value.200                    key.values[COMMON_NGRAM_MAX_VALUES - 1].value_idx = 0;201                    key.values[COMMON_NGRAM_MAX_VALUES - 1].value_num = 0;202                }203            }204            if (key.values[0].value_idx == 0) {205                // No values left, delete the key.206                LOG_DBG("%s: delete key %d at index %zu (no values left)\n", __func__, i, key.key_idx);207                map.keys.erase(map.keys.begin() + i);208                count_keys_del++;209            }210        }211 212        LOG_INF("%s: refresh map: idx_last_draft=%zu, new begin=%zu, #keys_checked=%zu, #keys_del=%zu, #values_del=%zu, #hashes_upd=%zu\n", __func__,213                map.idx_last_check, size_begin,214                count_keys, count_keys_del, count_values_del, count_map_entries_upd);215    }216 217    map.idx_last_check = size_begin;218    map.size_last_begin = size_begin;219}220 221void common_ngram_map_draft(common_ngram_map & map,222        const llama_tokens & inp, llama_token sampled,223        llama_tokens & draft) {224    // reset last key and value.225    map.last_draft_created   = false;226    map.last_draft_key_idx   = 0;227    map.last_draft_value_idx = 0;228 229    const size_t cur_len = inp.size();230    const uint16_t n = map.size_key;231    const uint16_t m = map.size_value;232    if (cur_len < static_cast<size_t>(2 * n + m)) {233        return;234    }235    if (cur_len >= static_cast<size_t>(UINT32_MAX)) {236        // key_map uses uint32_t instead of size_t.237        GGML_ABORT("%s: cur_len exceeds UINT32_MAX: %zu", __func__, cur_len);238    }239 240    if (map.idx_last_check > cur_len) {241        // Should not happen because of common_ngram_map_begin().242        GGML_ABORT("%s: map.idx_last_check > cur_len: %zu > %zu", __func__, map.idx_last_check, cur_len);243    }244    map.idx_last_check = cur_len;245 246    // search pattern, the key n-gram247    std::vector<llama_token> key_tokens;248    key_tokens.reserve(n);249    for (size_t j = cur_len - n + 1; j < cur_len; ++j) {250        key_tokens.push_back(inp[j]);251    }252    key_tokens.push_back(sampled);253 254    // search for the key in the map255    size_t match_pos = 0;256    if (map.size_last_begin > cur_len) {257        GGML_ABORT("%s: map.size_last_begin > cur_len: %zu > %zu", __func__, map.size_last_begin, cur_len);258    }259    if (!map.key_map.empty()) {260        // Search for the key in the map key_map from hash of ngrams to index of ngram.261        uint32_t idx_hash = (common_ngram_map_hash(key_tokens, 0, n) % map.key_map.size());262        uint32_t idx_key = map.key_map[idx_hash];263        if (idx_key != 0 && idx_key < cur_len - n - m - 1) {264            // Check if the key matches the key at idx_key (because of possible collisions).265            bool match = true;266            for (size_t k = 0; k < n; ++k) {267                if (inp[idx_key + k] != key_tokens[k]) {268                    match = false;269                    break;270                }271            }272            LOG_DBG("%s: key hash %x -> idx_key %d: match %d\n", __func__, idx_hash, idx_key, match ? 1 : 0);273            if (match) {274                match_pos = idx_key;275            }276        }277    }278    if (match_pos == 0 && map.size_last_begin > (size_t) (n + m + 1)) {279        // Search for the key in [1, map.size_last_begin - n - m -1], descending.280        for (size_t j = map.size_last_begin - n - m - 1; j > map.key_map_last_idx; --j) {281            // Check if the key matches the key.282            bool match = true;283            for (size_t k = 0; k < n; ++k) {284                if (inp[j + k] != key_tokens[k]) {285                    match = false;286                    break;287                }288            }289            if (match) {290               match_pos = j;291               break;292            }293        }294    }295    if (match_pos == 0) {296        // In case of a reasoning chat, the part after size_last_begin may be deleted/reordered later.297        //298        // Search in [size_last_begin, cur_len - n - m - 1], descending.299        for (size_t j = cur_len - n - m - 1; j > map.size_last_begin && j > map.key_map_last_idx; --j) {300            bool match = true;301            for (size_t k = 0; k < n; ++k) {302                if (inp[j + k] != key_tokens[k]) {303                    match = false;304                    break;305                }306            }307            if (match) {308               match_pos = j;309               break;310            }311        }312    }313    if (match_pos > 0) {314        LOG_DBG("%s: cur_len = %zu, n = %d, m = %d, sz_tkns = %zu, sampled = %d, match_pos = %zu\n", __func__,315            cur_len, n, m, key_tokens.size(), sampled, match_pos);316    }317 318    if (!map.key_map.empty()) {319        // Add hashes of new ngrams in key_map.320        //321        // Use the same order as above.322        if (map.size_last_begin > (size_t) (n + m + 1)) {323            for (size_t j = map.size_last_begin - n - m - 1; j > map.key_map_last_idx; --j) {324                // compute hash and store index of ngram at idx j in the map.325                uint32_t idx_hash = (common_ngram_map_hash(inp, j, n) % map.key_map.size());326                if (map.key_map[idx_hash] == 0) {327                    map.key_map[idx_hash] = j; // collisions may occur328                }329            }330        }331 332        for (size_t j = cur_len - n - m - 1; j > map.size_last_begin && j > map.key_map_last_idx; --j) {333            // compute hash and store index of ngram at idx j in the map.334            uint32_t idx_hash = (common_ngram_map_hash(inp, j, n) % map.key_map.size());335            if (map.key_map[idx_hash] == 0) {336                map.key_map[idx_hash] = j;337            }338        }339        map.key_map_last_idx = std::max(static_cast<uint32_t>(cur_len - n - m - 1), map.key_map_last_idx);340    }341 342    if (match_pos == 0) {343        return;344    }345 346    // We have a match, now we look for the statistics of the key.347    size_t key_offset = map.keys.size(); // offset in the map348    // We iterate through the std::vector<common_ngram_map_key> map->keys.349    for (size_t i = 0; i < map.keys.size(); ++i) {350        bool match = true;351        for (size_t j = 0; j < n; ++j) {352            if (inp[map.keys[i].key_idx + j] != key_tokens[j]) {353                match = false;354                break;355            }356        }357        if (match) {358            key_offset = i;359            break;360        }361    }362    if (key_offset == map.keys.size()) {363        // We create a new key-entry, it will get offset key_offset.364        common_ngram_map_key new_key;365        new_key.key_idx = match_pos;366        new_key.stat_idx = 0;367        new_key.key_num = 0;368        for (int i = 0; i < COMMON_NGRAM_MAX_VALUES; ++i) {369            new_key.values[i].value_num = 0;370            new_key.values[i].n_accepted = m;371        }372        map.keys.push_back(new_key);373    }374 375    // our key n-gram:376    common_ngram_map_key & curr_key = map.keys[key_offset];377 378    // update number of key hits379    curr_key.key_num = (uint16_t) std::min((int) map.keys[key_offset].key_num + 1,380            (int) COMMON_NGRAM_MAX_VALUE_COUNT);381 382    if (map.key_only) {383        // simple mode:384        // Fill in the draft with the m tokens following the key.385        // We work with value values[0] only.386        int n_draft_tokens = std::min((int) m, (int) curr_key.values[0].n_accepted);387 388        for (int i = 0; i < n_draft_tokens; ++i) {389            draft.push_back(inp[match_pos + n + i]);390        }391 392        LOG_DBG("%s: key_idx = %zu, key_offset = %zu, key_num = %d, draft.size = %zu\n", __func__,393                curr_key.key_idx, key_offset, curr_key.key_num, draft.size());394 395        map.last_draft_created   = true;396        map.last_draft_key_idx   = key_offset;397        map.last_draft_value_idx = 0; // value 0 is used for simple mode398        return;399    }400 401    if (curr_key.key_num < map.min_hits) {402        // not enough hits to consider this a good draft403        LOG_DBG("%s: key_offset = %zu, key_num = %d, min_hits = %d, no draft\n", __func__,404                key_offset, curr_key.key_num, map.min_hits);405        return;406    }407 408    // complex mode: examine the different m-grams after this key n-gram.409    //410 411    // determine all (max COMMON_NGRAM_MAX_VALUES) m-grams after the key n-gram.412    for (size_t i = curr_key.stat_idx; i <= match_pos; ++i) {413        // begins the key n-gram at index i?414        bool match_key = true;415        for (size_t k = 0; k < n; ++k) {416            if (inp[i + k] != key_tokens[k]) {417                match_key = false;418                break;419            }420        }421        if (!match_key) {422            continue;423        }424 425        // Do we haven a existing value m-gram or a new one after the key at index i?426        size_t idx_begin_value_key = i + n;427        int idx_value = -1;428        for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) {429            size_t idx_begin_value_v = curr_key.values[v].value_idx;430            if (idx_begin_value_v == 0) {431                // We found an empty value slot => we found a new value m-gram after the key n-gram.432                curr_key.values[v].value_idx = idx_begin_value_key;433                curr_key.values[v].value_num = 0;434                curr_key.values[v].n_accepted = m;435                idx_value = v;436                break;437            }438            bool match = true;439            for (size_t j = 0; j < m; ++j) {440                if (inp[idx_begin_value_key + j] != inp[idx_begin_value_v + j]) {441                    match = false;442                    break;443                }444            }445            if (match) {446                // We found an existing value m-gram after the key n-gram.447                idx_value = v;448                break;449            }450        }451        if (idx_value >= 0) {452            // We found a value m-gram of the key n-gram.453            curr_key.values[idx_value].value_num = (uint16_t) std::min((int) curr_key.values[idx_value].value_num + 1,454                    (int) COMMON_NGRAM_MAX_VALUE_COUNT);455        }456    }457    // the statistics are updated up to match_pos.458    curr_key.stat_idx = match_pos;459 460    // Do we have a value we could use for the draft?461    uint16_t max_occur = 0;462    int slot_max = 0;463    for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) {464        uint16_t curr_occur = curr_key.values[v].value_num;465        if (curr_occur > max_occur) {466            max_occur = curr_occur;467            slot_max = v;468        }469    }470    // What is sum of the other occurrences?471    uint32_t sum_occur = 0;472    for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) {473        if (v == slot_max) {474            continue;475        }476        uint16_t curr_occur = curr_key.values[v].value_num;477        sum_occur += curr_occur;478    }479 480    LOG_DBG("%s: key_offset = %zu, max_occur = %d, sum_occur = %d, slot_max = %d [%zu/%d, %zu/%d, %zu/%d, %zu/%d]\n", __func__,481            key_offset,482            max_occur, sum_occur, slot_max,483            curr_key.values[0].value_idx, curr_key.values[0].value_num,484            curr_key.values[1].value_idx, curr_key.values[1].value_num,485            curr_key.values[2].value_idx, curr_key.values[2].value_num,486            curr_key.values[3].value_idx, curr_key.values[3].value_num487        );488    // Print the tokens of the four values (if idx != 0), use LOG_INF489    for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) {490        if (curr_key.values[v].value_idx != 0) {491            LOG_DBG("%s: value[%d] = %s\n", __func__, v, common_tokens_to_str(inp, curr_key.values[v].value_idx, m).c_str());492        }493    }494 495    if (sum_occur > 0 && max_occur < 2 * sum_occur) {496        // The most frequent value is not much more frequent than the other values.497        // We do not use the draft.498        return;499    }500 501    // We use the most frequent value values[slot_max] for the draft.502    // Fill in the draft with the m tokens following the key.503    int n_draft_tokens = std::min((int) m, (int) curr_key.values[slot_max].n_accepted);504 505    for (int i = 0; i < n_draft_tokens; ++i) {506        draft.push_back(inp[match_pos + n + i]);507    }508 509    LOG_DBG("%s: key_offset = %zu, slot_max = %d, key_num = %d, draft.size = %zu\n", __func__,510            key_offset, slot_max,511            curr_key.key_num, draft.size());512 513    map.last_draft_created   = true;514    map.last_draft_key_idx   = key_offset;515    map.last_draft_value_idx = slot_max; // value used for draft generation.516}517 518void common_ngram_map_accept(common_ngram_map & map, uint16_t n_accepted) {519    if (!map.last_draft_created) {520        return;521    }522 523    // find the key and its chosen value.524    const size_t key_idx = map.last_draft_key_idx;525    const size_t val_idx = map.last_draft_value_idx;526 527    // find key corresponding to key_idx.528    common_ngram_map_key & curr_key = map.keys[key_idx];529    // find value corresponding to val_idx.530    struct common_ngram_map_value & curr_value = curr_key.values[val_idx]; // value used for draft generation.531 532    // update the value statistics533    LOG_DBG("common_ngram_map_send_accepted: n_accepted = %d, prev value_num = %d\n",534            n_accepted, curr_value.n_accepted);535    curr_value.n_accepted = n_accepted;536}537