KBaba7/llama.cpp
0
1#include "arg.h"2#include "common.h"3#include "sampling.h"4#include "log.h"5#include "llama.h"6 7#include <algorithm>8#include <cstdio>9#include <cstring>10#include <random>11#include <set>12#include <string>13#include <vector>14 15#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 12816#define SPEC_VOCAB_CHECK_START_TOKEN_ID 517 18struct seq_draft {19 bool active = false;20 bool drafting = false;21 bool skip = false;22 23 int i_batch_dft = 0;24 std::vector<int> i_batch_tgt;25 26 std::vector<llama_token> tokens;27 std::vector<std::vector<llama_token_data>> dists;28 29 struct common_sampler * smpl = nullptr;30};31 32int main(int argc, char ** argv) {33 common_params params;34 35 // needed to get candidate probs even for temp <= 0.036 params.sampling.n_probs = 128;37 38 if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) {39 return 1;40 }41 42 if (params.n_predict < -1) {43 LOG_ERR("%s: --n-predict must be >= -1\n", __func__);44 return 1;45 }46 47 common_init();48 49 if (params.speculative.model.empty()) {50 LOG_ERR("%s: --model-draft is required\n", __func__);51 return 1;52 }53 54 // max number of parallel drafting sequences (i.e. tree branches)55 const int n_seq_dft = params.n_parallel;56 57 // probability threshold for splitting a draft branch (only for n_seq_dft > 1)58 const float p_draft_split = params.speculative.p_split;59 60 std::default_random_engine rng(params.sampling.seed == LLAMA_DEFAULT_SEED ? std::random_device()() : params.sampling.seed);61 std::uniform_real_distribution<> u_dist;62 63 // init llama.cpp64 llama_backend_init();65 llama_numa_init(params.numa);66 67 llama_model * model_tgt = NULL;68 llama_model * model_dft = NULL;69 70 llama_context * ctx_tgt = NULL;71 llama_context * ctx_dft = NULL;72 73 // load the target model74 common_init_result llama_init_tgt = common_init_from_params(params);75 76 model_tgt = llama_init_tgt.model.get();77 ctx_tgt = llama_init_tgt.context.get();78 79 // load the draft model80 params.devices = params.speculative.devices;81 params.model = params.speculative.model;82 params.n_gpu_layers = params.speculative.n_gpu_layers;83 if (params.speculative.cpuparams.n_threads > 0) {84 params.cpuparams.n_threads = params.speculative.cpuparams.n_threads;85 }86 87 params.cpuparams_batch.n_threads = params.speculative.cpuparams_batch.n_threads;88 common_init_result llama_init_dft = common_init_from_params(params);89 90 model_dft = llama_init_dft.model.get();91 ctx_dft = llama_init_dft.context.get();92 93 const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt);94 const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft);95 96 const bool vocab_type_tgt = llama_vocab_type(vocab_tgt);97 LOG_DBG("vocab_type tgt: %d\n", vocab_type_tgt);98 99 const bool vocab_type_dft = llama_vocab_type(vocab_dft);100 LOG_DBG("vocab_type dft: %d\n", vocab_type_dft);101 102 if (vocab_type_tgt != vocab_type_dft) {103 LOG_ERR("%s: draft model vocab type must match target model to use speculation but ", __func__);104 LOG_ERR("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt);105 return 1;106 }107 108 if (109 llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||110 llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||111 llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) ||112 llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft)113 ) {114 LOG_ERR("%s: draft model special tokens must match target model to use speculation\n", __func__);115 return 1;116 }117 118 {119 const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt);120 const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft);121 const int vocab_diff = n_vocab_tgt > n_vocab_dft122 ? n_vocab_tgt - n_vocab_dft123 : n_vocab_dft - n_vocab_tgt;124 125 if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {126 LOG_ERR("%s: draft model vocab must closely match target model to use speculation but ", __func__);127 LOG_ERR("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",128 n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);129 return 1;130 }131 132 for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) {133 const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i);134 const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);135 if (std::strcmp(token_text_tgt, token_text_dft) != 0) {136 LOG_ERR("%s: draft model vocab must match target model to use speculation but ", __func__);137 LOG_ERR("token %d content differs - target '%s', draft '%s'\n", i,138 common_token_to_piece(ctx_tgt, i).c_str(),139 common_token_to_piece(ctx_dft, i).c_str());140 return 1;141 }142 }143 }144 145 146 // Tokenize the prompt147 std::vector<llama_token> inp;148 inp = common_tokenize(ctx_tgt, params.prompt, true, true);149 150 const int max_context_size = llama_n_ctx(ctx_tgt);151 const int max_tokens_list_size = max_context_size - 4;152 153 if ((int) inp.size() > max_tokens_list_size) {154 LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size);155 return 1;156 }157 158 LOG("\n\n");159 160 for (auto id : inp) {161 LOG("%s", common_token_to_piece(ctx_tgt, id).c_str());162 }163 164 const int n_input = inp.size();165 166 const auto t_enc_start = ggml_time_us();167 168 // eval the prompt with both models169 llama_decode(ctx_tgt, llama_batch_get_one( inp.data(), n_input - 1));170 llama_decode(ctx_tgt, llama_batch_get_one(&inp.back(), 1));171 llama_decode(ctx_dft, llama_batch_get_one( inp.data(), n_input));172 173 const auto t_enc_end = ggml_time_us();174 175 // the 2 models should have the same vocab176 //GGML_ASSERT(n_vocab == llama_vocab_n_tokens(model_dft));177 178 // how many tokens to draft each time179 int n_draft = params.speculative.n_max;180 181 int n_predict = 0;182 int n_drafted = 0;183 int n_accept = 0;184 185 int n_past_tgt = inp.size();186 int n_past_dft = inp.size();187 188 // used to determine end of generation189 bool has_eos = false;190 191 // target model sampling context (reuse the llama_context's sampling instance)192 struct common_sampler * smpl = common_sampler_init(model_tgt, params.sampling);193 194 // draft sequence data195 std::vector<seq_draft> drafts(n_seq_dft);196 197 for (int s = 0; s < n_seq_dft; ++s) {198 // allocate llama_sampler for each draft sequence199 drafts[s].smpl = common_sampler_init(model_dft, params.sampling);200 }201 202 llama_batch batch_dft = llama_batch_init(llama_n_batch(ctx_dft), 0, 1);203 llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, n_seq_dft);204 205 const auto t_dec_start = ggml_time_us();206 207 // sample from the last token of the prompt208 drafts[0].i_batch_tgt.resize(1);209 drafts[0].i_batch_tgt[0] = 0;210 211 while (true) {212 std::set<int> active_seqs = {};213 214 // print current draft sequences215 for (int s = 0; s < n_seq_dft; ++s) {216 if (!drafts[s].active) {217 continue;218 }219 220 active_seqs.insert(s);221 const auto & tokens = drafts[s].tokens;222 223 LOG_DBG("draft %d: %s\n", s, string_from(ctx_dft, tokens).c_str());224 }225 226 int i_dft = 0;227 int s_keep = 0;228 229 llama_token token_id;230 std::string token_str;231 232 // loop until we fail to accept a drafted token or we run out of drafted tokens233 while (true) {234 235 // check if the target token matches any of the drafts236 // for stochastic sampling, attempt to match the token with the drafted tokens237 {238 bool accept = false;239 if (params.sampling.temp > 0) {240 // stochastic verification241 common_sampler_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft], true);242 243 auto & dist_tgt = *common_sampler_get_candidates(smpl);244 245 float p_tgt = 0.0f;246 float p_dft = 0.0f;247 248 while (active_seqs.size() > 0) {249 // randomly select a sequence to verify from active sequences250 std::uniform_int_distribution<unsigned int> u_int_dist(0, active_seqs.size() - 1);251 int s = *std::next(active_seqs.begin(), u_int_dist(rng));252 if (i_dft >= (int) drafts[s].tokens.size()) {253 drafts[s].active = false;254 active_seqs.erase(s);255 continue;256 }257 if (accept) {258 // if we already accepted a token, we can skip the rest259 if (drafts[s].tokens[i_dft] != drafts[s_keep].tokens[i_dft]) {260 drafts[s].active = false;261 active_seqs.erase(s);262 }263 continue;264 }265 266 LOG_DBG("verifying sequence #%d at pos #%d from %d active sequence(s)\n", s, i_dft, (int) active_seqs.size());267 float r = u_dist(rng);268 llama_token_data_array dist_dft = { drafts[s].dists[i_dft].data() , drafts[s].dists[i_dft].size(), LLAMA_TOKEN_NULL, true };269 270 //GGML_ASSERT(dist_tgt.size <= dist_dft.size);271 272 // acquire the token probabilities assigned by the draft and target models273 for (size_t i = 0; i < dist_tgt.size; i++) {274 if (dist_tgt.data[i].id == drafts[s].tokens[i_dft]) {275 p_tgt = dist_tgt.data[i].p;276 break;277 }278 }279 for (size_t i = 0; i < dist_dft.size; i++) {280 if (dist_dft.data[i].id == drafts[s].tokens[i_dft]) {281 p_dft = dist_dft.data[i].p;282 break;283 }284 }285 LOG_DBG("r = %f, p_dft = %f, p_tgt = %f\n", r, p_dft, p_tgt);286 if (r <= p_tgt / p_dft) {287 s_keep = s;288 accept = true;289 token_id = drafts[s].tokens[i_dft];290 token_str = common_token_to_piece(ctx_tgt, token_id);291 common_sampler_accept(smpl, token_id, true);292 293 LOG_DBG("draft token %d of sequence %d (%d, '%s') accepted\n", i_dft, s, token_id, token_str.c_str());294 break;295 } else {296 LOG_DBG("draft token %d of sequence %d (%d, '%s') rejected\n", i_dft, s, drafts[s].tokens[i_dft], common_token_to_piece(ctx_tgt, drafts[s].tokens[i_dft]).c_str());297 drafts[s].active = false;298 299 // calculate residual probability300 GGML_ASSERT(dist_tgt.sorted);301 GGML_ASSERT(dist_dft.sorted);302 303 // sort dist by id304 std::sort(dist_tgt.data, dist_tgt.data + dist_tgt.size, [](const llama_token_data &a, const llama_token_data &b) {305 return a.id < b.id;306 });307 std::sort(dist_dft.data, dist_dft.data + dist_dft.size, [](const llama_token_data &a, const llama_token_data &b) {308 return a.id < b.id;309 });310 311 float sum_probs = 0.0f;312 313 for (size_t i = 0; i < dist_tgt.size; i++) {314 if (i < dist_dft.size) {315 dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p - dist_dft.data[i].p);316 } else {317 dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p);318 }319 320 sum_probs += dist_tgt.data[i].p;321 }322 323 for (size_t i = 0; i < dist_tgt.size; i++) {324 dist_tgt.data[i].p /= sum_probs;325 }326 327 // sort dist_tgt by p desc328 std::sort(dist_tgt.data, dist_tgt.data + dist_tgt.size, [](const llama_token_data &a, const llama_token_data &b) {329 return a.p > b.p;330 });331 }332 333 active_seqs.erase(s);334 for(int i = 0; i < n_seq_dft; i++) {335 if (i == s) {336 continue;337 }338 if (drafts[i].tokens[i_dft] == drafts[s].tokens[i_dft]) {339 // synchronize active status for sequences with the same drafted token340 drafts[i].active = drafts[i].active && accept;341 if (!drafts[i].active) {342 active_seqs.erase(s);343 }344 }345 }346 }347 348 if (!accept) {349 // all drafted tokens were rejected350 // sample from the target model351 LOG_DBG("all drafted tokens were rejected, sampling from residual distribution\n");352 std::vector<float> probs(dist_tgt.size);353 for (size_t i = 0; i < dist_tgt.size; ++i) {354 probs[i] = dist_tgt.data[i].p;355 }356 357 std::discrete_distribution<> dist(probs.begin(), probs.end());358 359 const int idx = dist(rng);360 361 token_id = dist_tgt.data[idx].id;362 common_sampler_accept(smpl, token_id, true);363 token_str = common_token_to_piece(ctx_tgt, token_id);364 }365 } else {366 // greedy verification367 368 // sample from the target model369 LOG_DBG("sampling target: s_keep = %3d, i_dft = %3d, i_batch_tgt = %3d\n", s_keep, i_dft, drafts[s_keep].i_batch_tgt[i_dft]);370 token_id = common_sampler_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft]);371 372 common_sampler_accept(smpl, token_id, true);373 374 token_str = common_token_to_piece(ctx_tgt, token_id);375 376 for (int s = 0; s < n_seq_dft; ++s) {377 if (!drafts[s].active) {378 continue;379 }380 381 if (i_dft < (int) drafts[s].tokens.size() && token_id == drafts[s].tokens[i_dft]) {382 LOG_DBG("the sampled target token matches the %dth drafted token of sequence %d (%d, '%s') - accepted\n", i_dft, s, token_id, token_str.c_str());383 384 s_keep = s;385 accept = true;386 } else {387 drafts[s].active = false;388 }389 }390 }391 392 if (llama_vocab_is_eog(vocab_tgt, token_id)) {393 has_eos = true;394 }395 ++n_predict;396 397 if (accept) {398 ++n_accept;399 ++n_past_tgt;400 ++n_past_dft;401 ++i_dft;402 if (params.use_color) {403 // Color token according to its origin sequence404 LOG("\u001b[%dm%s\u001b[37m", (36 - s_keep % 6), token_str.c_str());405 } else {406 LOG("%s", token_str.c_str());407 }408 continue;409 } else {410 LOG("%s", token_str.c_str());411 break;412 }413 }414 }415 416 {417 LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", token_id, token_str.c_str());418 419 // TODO: simplify420 {421 LOG_DBG("keeping sequence %d, n_past_tgt = %d, n_past_dft = %d\n", s_keep, n_past_tgt, n_past_dft);422 423 llama_kv_cache_seq_keep(ctx_dft, s_keep);424 llama_kv_cache_seq_cp (ctx_dft, s_keep, 0, -1, -1);425 llama_kv_cache_seq_keep(ctx_dft, 0);426 427 llama_kv_cache_seq_rm (ctx_tgt, s_keep, n_past_tgt, -1);428 llama_kv_cache_seq_keep(ctx_tgt, s_keep);429 llama_kv_cache_seq_cp (ctx_tgt, s_keep, 0, -1, -1);430 llama_kv_cache_seq_keep(ctx_tgt, 0);431 }432 433 for (int s = 0; s < n_seq_dft; ++s) {434 drafts[s].active = false;435 drafts[s].tokens.clear();436 drafts[s].i_batch_tgt.clear();437 drafts[s].dists.clear();438 }439 // note: will be erased after the speculation phase440 drafts[0].tokens.push_back(token_id);441 drafts[0].dists.push_back(std::vector<llama_token_data>());442 drafts[0].i_batch_tgt.push_back(0);443 444 common_batch_clear(batch_dft);445 common_batch_add (batch_dft, token_id, n_past_dft, { 0 }, true);446 447 llama_kv_cache_seq_rm(ctx_dft, 0, n_past_dft, -1);448 // LOG_DBG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str());449 llama_decode(ctx_dft, batch_dft);450 451 ++n_past_dft;452 }453 454 if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {455 break;456 }457 458 if (drafts[0].smpl) {459 common_sampler_free(drafts[0].smpl);460 }461 drafts[0].smpl = common_sampler_clone(smpl);462 463 int n_seq_cur = 1;464 int n_past_cur = n_past_dft;465 466 for (int s = 0; s < n_seq_dft; ++s) {467 drafts[s].active = false;468 drafts[s].drafting = false;469 }470 drafts[0].active = true;471 drafts[0].drafting = true;472 drafts[0].i_batch_dft = 0;473 474 common_batch_clear(batch_tgt);475 common_batch_add (batch_tgt, drafts[0].tokens[0], n_past_tgt, { 0 }, true);476 477 // sample n_draft tokens from the draft model using tree-based sampling478 for (int i = 0; i < n_draft; ++i) {479 batch_dft.n_tokens = 0;480 481 for (int s = 0; s < n_seq_dft; ++s) {482 drafts[s].skip = false;483 }484 485 for (int s = 0; s < n_seq_dft; ++s) {486 if (!drafts[s].drafting || drafts[s].skip) {487 continue;488 }489 490 common_sampler_sample(drafts[s].smpl, ctx_dft, drafts[s].i_batch_dft, true);491 492 const auto * cur_p = common_sampler_get_candidates(drafts[s].smpl);493 494 for (int k = 0; k < std::min(n_seq_dft + 3, (int) cur_p->size); ++k) {495 LOG_DBG(" - draft candidate %3d for seq %3d, pos %3d: %6d (%8.3f) '%s'\n",496 k, s, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());497 }498 499 std::vector<int> sa(1, s);500 501 // attempt to split the branch if the probability is high enough502 for (int f = 1; f < 8; ++f) {503 if (n_seq_cur < n_seq_dft && cur_p->data[f].p > p_draft_split) {504 LOG_DBG("splitting seq %3d into %3d\n", s, n_seq_cur);505 506 llama_kv_cache_seq_rm(ctx_dft, n_seq_cur, -1, -1);507 llama_kv_cache_seq_cp(ctx_dft, s, n_seq_cur, -1, -1);508 509 // all previous tokens from this branch are now also part of the new branch510 for (int t = 0; t < batch_tgt.n_tokens; ++t) {511 for (int p = 0; p < batch_tgt.n_seq_id[t]; ++p) {512 if (batch_tgt.seq_id[t][p] == s) {513 batch_tgt.seq_id[t][batch_tgt.n_seq_id[t]] = n_seq_cur;514 batch_tgt.n_seq_id[t]++;515 break;516 }517 }518 }519 520 // copy the draft state521 drafts[n_seq_cur].active = true;522 drafts[n_seq_cur].drafting = true;523 drafts[n_seq_cur].skip = true;524 525 drafts[n_seq_cur].tokens = drafts[s].tokens;526 drafts[n_seq_cur].dists = drafts[s].dists;527 drafts[n_seq_cur].i_batch_dft = drafts[s].i_batch_dft;528 drafts[n_seq_cur].i_batch_tgt = drafts[s].i_batch_tgt;529 530 if (drafts[n_seq_cur].smpl) {531 common_sampler_free(drafts[n_seq_cur].smpl);532 }533 drafts[n_seq_cur].smpl = common_sampler_clone(drafts[s].smpl);534 535 sa.push_back(n_seq_cur);536 537 n_seq_cur++;538 } else {539 break;540 }541 }542 543 // add drafted token for each sequence544 for (int is = 0; is < (int) sa.size(); ++is) {545 const llama_token id = cur_p->data[is].id;546 547 const int s = sa[is];548 549 common_sampler_accept(drafts[s].smpl, id, true);550 551 drafts[s].tokens.push_back(id);552 // save cur_p.data into drafts[s].dists553 drafts[s].dists.push_back({cur_p->data, cur_p->data + cur_p->size});554 555 // add unique drafted tokens to the target batch556 drafts[s].i_batch_tgt.push_back(batch_tgt.n_tokens);557 558 common_batch_add(batch_tgt, id, n_past_tgt + i + 1, { s }, true);559 560 // add the token to the batch for batched decoding with the draft model561 drafts[s].i_batch_dft = batch_dft.n_tokens;562 563 common_batch_add(batch_dft, id, n_past_cur, { s }, true);564 565 if (batch_tgt.n_tokens > n_draft) {566 drafts[s].drafting = false;567 }568 }569 }570 571 // no sequence is drafting anymore572 if (batch_dft.n_tokens == 0) {573 break;574 }575 576 // evaluate the drafted tokens on the draft model577 llama_decode(ctx_dft, batch_dft);578 ++n_past_cur;579 ++n_drafted;580 581 if (batch_tgt.n_tokens > n_draft) {582 break;583 }584 }585 586 // evaluate the target model on the drafted tokens587 {588 llama_kv_cache_seq_keep(ctx_tgt, 0);589 for (int s = 1; s < n_seq_dft; ++s) {590 llama_kv_cache_seq_cp(ctx_tgt, 0, s, -1, -1);591 }592 593 // LOG_DBG("target batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_tgt, batch_tgt).c_str());594 llama_decode(ctx_tgt, batch_tgt);595 ++n_past_tgt;596 }597 598 // the first token is always proposed by the target model before the speculation loop so we erase it here599 for (int s = 0; s < n_seq_dft; ++s) {600 if (!drafts[s].active) {601 continue;602 }603 604 drafts[s].tokens.erase(drafts[s].tokens.begin());605 drafts[s].dists.erase(drafts[s].dists.begin());606 }607 }608 609 auto t_dec_end = ggml_time_us();610 611 LOG("\n\n");612 613 LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input, (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f));614 LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict / ((t_dec_end - t_dec_start) / 1e6f));615 616 LOG_INF("\n");617 LOG_INF("n_draft = %d\n", n_draft);618 LOG_INF("n_predict = %d\n", n_predict);619 LOG_INF("n_drafted = %d\n", n_drafted);620 LOG_INF("n_accept = %d\n", n_accept);621 LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted);622 623 LOG_INF("\n");624 LOG_INF("draft:\n\n");625 // TODO: print sampling/grammar timings for all drafts626 llama_perf_context_print(ctx_dft);627 628 LOG_INF("\n");629 LOG_INF("target:\n\n");630 common_perf_print(ctx_tgt, smpl);631 632 common_sampler_free(smpl);633 for (int s = 0; s < n_seq_dft; ++s) {634 common_sampler_free(drafts[s].smpl);635 }636 637 llama_batch_free(batch_dft);638 639 llama_backend_free();640 641 LOG("\n\n");642 643 return 0;644}645