KBaba7/llama.cpp
0
1#include "ngram-cache.h"2#include "common.h"3#include "log.h"4 5#include <cinttypes>6#include <cstdint>7#include <cstdio>8#include <fstream>9#include <thread>10 11void common_ngram_cache_update(common_ngram_cache & ngram_cache, int ngram_min, int ngram_max,12 std::vector<llama_token> & inp, int nnew, bool print_progress) {13 const int64_t t_start_ms = ggml_time_ms();14 const int64_t inp_size = inp.size();15 16 const int64_t n_todo = inp_size * (ngram_max - ngram_min + 1);17 int64_t n_done = 0;18 19 for (int64_t ngram_size = ngram_min; ngram_size <= ngram_max; ++ngram_size) {20 const int64_t i_start = std::max(inp_size - nnew, ngram_size);21 for (int64_t i = i_start; i < inp_size; ++i) {22 const int64_t ngram_start = i - ngram_size;23 common_ngram ngram(&inp[ngram_start], ngram_size);24 const llama_token token = inp[i];25 26 common_ngram_cache::iterator part_it = ngram_cache.find(ngram);27 if (part_it == ngram_cache.end()) {28 common_ngram_cache_part part;29 part.emplace(token, 1);30 ngram_cache.emplace(ngram, part);31 } else {32 common_ngram_cache_part::iterator token_count_it = part_it->second.find(token);33 if (token_count_it == part_it->second.end()) {34 part_it->second.emplace(token, 1);35 } else {36 token_count_it->second++;37 }38 }39 ++n_done;40 41 if (print_progress && n_done % 10000000 == 0) {42 const int64_t t_now_ms = ggml_time_ms();43 const int64_t eta_ms = (inp_size*(ngram_max-ngram_min+1) - n_done) * (t_now_ms - t_start_ms) / n_done;44 const int64_t eta_min = eta_ms / (60*1000);45 const int64_t eta_s = (eta_ms - 60*1000*eta_min) / 1000;46 47 fprintf(stderr, "%s: %" PRId64 "/%" PRId64 " done, ETA: %02" PRId64 ":%02" PRId64 "\n", __func__, n_done, n_todo, eta_min, eta_s);48 }49 }50 }51}52 53// Helper function to get a token from the combined, speculative sequence of inp and draft.54static llama_token get_token(const std::vector<llama_token> & inp, const std::vector<llama_token> & draft, const size_t i) {55 return i < inp.size() ? inp[i] : draft[1 + i - inp.size()];56}57 58// If sample size or percentage are below these thresholds the draft is aborted early:59constexpr int draft_min_sample_size_lax[LLAMA_NGRAM_MAX] = { 2, 2, 1, 1};60constexpr int draft_min_percent_lax[LLAMA_NGRAM_MAX] = {66, 50, 50, 50};61constexpr int draft_min_sample_size_strict[LLAMA_NGRAM_MAX] = { 4, 3, 2, 2};62constexpr int draft_min_percent_strict[LLAMA_NGRAM_MAX] = {75, 66, 66, 66};63 64// Helper function that tries to draft a token from only the static ngram cache:65static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram ngram_static) {66 common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static);67 if (part_static_it == nc_static.end()) {68 return LLAMA_TOKEN_NULL;69 }70 const common_ngram_cache_part part_static = part_static_it->second;71 72 int max_count_static = 0;73 int sum_count_static = 0;74 llama_token max_token = LLAMA_TOKEN_NULL;75 76 for (std::pair<llama_token, int> token_count_static : part_static) {77 const llama_token token = token_count_static.first;78 const int32_t count_static = token_count_static.second;79 80 if (count_static > max_count_static) {81 max_token = token;82 max_count_static = count_static;83 }84 sum_count_static += count_static;85 }86 87 if (sum_count_static < draft_min_sample_size_lax[LLAMA_NGRAM_STATIC-1]) {88 return LLAMA_TOKEN_NULL;89 }90 if (100*max_count_static < draft_min_percent_lax[LLAMA_NGRAM_STATIC-1]*sum_count_static) {91 return LLAMA_TOKEN_NULL;92 }93 return max_token;94}95 96// Try to draft a token from primary cache (context/dynamic), validate with static cache:97static llama_token try_draft(98 common_ngram_cache & nc_primary, const std::vector<common_ngram> & ngrams_primary, common_ngram_cache_part & part_static,99 const int * min_sample_size, const int * min_percent) {100 101 llama_token drafted_token = LLAMA_TOKEN_NULL;102 103 for (int i = ngrams_primary.size()-1; i >= 0 && drafted_token == LLAMA_TOKEN_NULL; --i) {104 const common_ngram ngram_primary = ngrams_primary[i];105 106 common_ngram_cache::iterator part_primary_it = nc_primary.find(ngram_primary);107 if (part_primary_it == nc_primary.end()) {108 continue;109 }110 const common_ngram_cache_part part_primary = part_primary_it->second;111 112 int max_count_primary = 0;113 int max_count_static = 0;114 int sum_count_primary = 0;115 llama_token max_token = LLAMA_TOKEN_NULL;116 117 for (std::pair<llama_token, int> token_count_primary : part_primary) {118 const llama_token token = token_count_primary.first;119 120 common_ngram_cache_part::iterator token_count_static_it = part_static.find(token);121 122 const int32_t count_primary = token_count_primary.second;123 const int32_t count_static = token_count_static_it != part_static.end() ? 100*token_count_static_it->second : 1;124 125 if (count_primary*count_static > max_count_primary*max_count_static) {126 max_token = token;127 max_count_primary = count_primary;128 max_count_static = count_static;129 }130 sum_count_primary += count_primary;131 }132 133 if (sum_count_primary < min_sample_size[i]) {134 continue;135 }136 if (100*max_count_primary < min_percent[i]*sum_count_primary) {137 continue;;138 }139 drafted_token = max_token;140 }141 142 return drafted_token;143}144 145void common_ngram_cache_draft(146 std::vector<llama_token> & inp, std::vector<llama_token> & draft, int n_draft, int ngram_min, int ngram_max,147 common_ngram_cache & nc_context, common_ngram_cache & nc_dynamic, common_ngram_cache & nc_static148) {149 GGML_ASSERT(draft.size() == 1);150 const int inp_size = inp.size();151 152 if (inp_size < LLAMA_NGRAM_STATIC) {153 return;154 }155 156 while ((int) draft.size()-1 < n_draft) {157 llama_token drafted_token = LLAMA_TOKEN_NULL;158 159 const int ngram_start_static = inp_size-LLAMA_NGRAM_STATIC + draft.size()-1;160 common_ngram ngram_static;161 for (int j = ngram_start_static; j < ngram_start_static + LLAMA_NGRAM_STATIC; ++j) {162 ngram_static.tokens[j-ngram_start_static] = get_token(inp, draft, j);163 }164 common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static);165 common_ngram_cache_part part_static;166 if (part_static_it != nc_static.end()) {167 part_static = part_static_it->second;168 }169 170 // cd = context + dynamic171 std::vector<common_ngram> ngrams_cd;172 for (int ngram_size_cd = ngram_min; ngram_size_cd <= ngram_max; ++ngram_size_cd) {173 const int ngram_start_cd = inp_size-ngram_size_cd + draft.size()-1;174 common_ngram ngram_cd;175 for (int j = ngram_start_cd; j < ngram_start_cd + ngram_size_cd; ++j) {176 ngram_cd.tokens[j-ngram_start_cd] = get_token(inp, draft, j);177 }178 ngrams_cd.push_back(ngram_cd);179 }180 if (drafted_token == LLAMA_TOKEN_NULL) {181 drafted_token = try_draft(nc_context, ngrams_cd, part_static, draft_min_sample_size_lax, draft_min_percent_lax);182 }183 if (drafted_token == LLAMA_TOKEN_NULL) {184 drafted_token = try_draft(nc_dynamic, ngrams_cd, part_static, draft_min_sample_size_strict, draft_min_percent_strict);185 }186 if (drafted_token == LLAMA_TOKEN_NULL) {187 drafted_token = try_draft(nc_static, ngram_static);188 }189 190 if (drafted_token == LLAMA_TOKEN_NULL) {191 break;192 }193 194 LOG(" - draft candidate: token=%d\n", drafted_token);195 draft.push_back(drafted_token);196 }197}198 199void common_ngram_cache_save(common_ngram_cache & ngram_cache, std::string & filename) {200 std::ofstream file_out(filename, std::ios::binary);201 for (std::pair<common_ngram, common_ngram_cache_part> item : ngram_cache) {202 const common_ngram ngram = item.first;203 common_ngram_cache_part token_counts = item.second;204 GGML_ASSERT(!token_counts.empty());205 const int32_t ntokens = token_counts.size();206 GGML_ASSERT(ntokens > 0);207 208 file_out.write(reinterpret_cast<const char *>(&ngram), sizeof(common_ngram));209 file_out.write(reinterpret_cast<const char *>(&ntokens), sizeof(int32_t));210 for (std::pair<llama_token, int32_t> item2 : token_counts) {211 const llama_token token = item2.first;212 const int32_t count = item2.second;213 GGML_ASSERT(count > 0);214 215 file_out.write(reinterpret_cast<const char *>(&token), sizeof(llama_token));216 file_out.write(reinterpret_cast<const char *>(&count), sizeof(int32_t));217 }218 }219 220}221 222common_ngram_cache common_ngram_cache_load(std::string & filename) {223 std::ifstream hashmap_file(filename, std::ios::binary);224 if (!hashmap_file) {225 throw std::ifstream::failure("Unable to open file " + filename);226 }227 common_ngram_cache ngram_cache;228 229 common_ngram ngram;230 int32_t ntokens;231 llama_token token;232 int32_t count;233 234 char * ngramc = reinterpret_cast<char*>(&ngram);235 char * ntokensc = reinterpret_cast<char*>(&ntokens);236 char * tokenc = reinterpret_cast<char*>(&token);237 char * countc = reinterpret_cast<char*>(&count);238 while(hashmap_file.read(ngramc, sizeof(common_ngram))) {239 GGML_ASSERT(!hashmap_file.eof());240 GGML_ASSERT(hashmap_file.read(ntokensc, sizeof(int32_t)));241 GGML_ASSERT(ntokens > 0);242 common_ngram_cache_part token_counts;243 244 for (int i = 0; i < ntokens; ++i) {245 GGML_ASSERT(!hashmap_file.eof());246 GGML_ASSERT(hashmap_file.read(tokenc, sizeof(llama_token)));247 GGML_ASSERT(!hashmap_file.eof());248 GGML_ASSERT(hashmap_file.read(countc, sizeof(int32_t)));249 GGML_ASSERT(count > 0);250 token_counts.emplace(token, count);251 }252 253 ngram_cache.emplace(ngram, token_counts);254 }255 GGML_ASSERT(hashmap_file.eof());256 257 return ngram_cache;258}259 260void common_ngram_cache_merge(common_ngram_cache & ngram_cache_target, common_ngram_cache & ngram_cache_add) {261 for (std::pair<common_ngram, common_ngram_cache_part> ngram_part : ngram_cache_add) {262 const common_ngram ngram = ngram_part.first;263 common_ngram_cache_part part = ngram_part.second;264 265 common_ngram_cache::iterator part_merged_it = ngram_cache_target.find(ngram);266 if (part_merged_it == ngram_cache_target.end()) {267 ngram_cache_target.emplace(ngram, part);268 continue;269 }270 271 for (std::pair<llama_token, int32_t> token_count : part) {272 const llama_token token = token_count.first;273 const int32_t count = token_count.second;274 GGML_ASSERT(count > 0);275 276 common_ngram_cache_part::iterator token_count_merged_it = part_merged_it->second.find(token);277 if (token_count_merged_it == part_merged_it->second.end()) {278 part_merged_it->second.emplace(token, count);279 continue;280 }281 282 token_count_merged_it->second += count;283 }284 }285}286