Felipe97/llama-cpp-compiled
01.1k
1#include "arg.h"2#include "common.h"3#include "sampling.h"4#include "speculative.h"5#include "log.h"6#include "llama.h"7 8#include <algorithm>9#include <clocale>10#include <cstdio>11#include <cstring>12#include <cinttypes>13#include <string>14#include <vector>15#include <utility>16 17int main(int argc, char ** argv) {18 std::setlocale(LC_NUMERIC, "C");19 20 common_params params;21 22 common_init();23 24 if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) {25 return 1;26 }27 28 if (params.n_predict < -1) {29 LOG_ERR("%s: --n-predict must be >= -1\n", __func__);30 return 1;31 }32 33 const auto output_limits = common_speculative_get_output_limits(34 params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative));35 params.n_outputs_max = output_limits.total;36 params.n_outputs_max_per_seq = output_limits.per_seq;37 38 // init llama.cpp39 llama_backend_init();40 llama_numa_init(params.numa);41 42 llama_model * model_tgt = NULL;43 44 llama_context * ctx_tgt = NULL;45 46 // load the target model47 auto llama_init_tgt = common_init_from_params(params);48 49 model_tgt = llama_init_tgt->model();50 ctx_tgt = llama_init_tgt->context();51 52 const llama_vocab * vocab = llama_model_get_vocab(model_tgt);53 54 // load the draft model (if any) - this also creates the MTP draft context when MTP speculation is enabled55 common_speculative_init_result_ptr spec_init;56 57 {58 common_params params_dft = common_base_params_to_speculative(params);59 60 spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt);61 62 params.speculative.draft.ctx_tgt = ctx_tgt;63 params.speculative.draft.ctx_dft = spec_init->context();64 }65 66 llama_context * ctx_dft = params.speculative.draft.ctx_dft;67 68 // check if the context supports partial sequence removal69 const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;70 const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;71 72 if (use_ckpt_tgt) {73 LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n");74 }75 76 // Tokenize the prompt77 std::vector<llama_token> inp;78 inp = common_tokenize(ctx_tgt, params.prompt, true, true);79 80 if (llama_n_ctx(ctx_tgt) < (uint32_t) inp.size()) {81 LOG_ERR("%s: the prompt exceeds the context size (%d tokens, ctx %d)\n", __func__, (int) inp.size(), llama_n_ctx(ctx_tgt));82 83 return 1;84 }85 86 if (llama_n_batch(ctx_tgt) < (uint32_t) inp.size()) {87 LOG_ERR("%s: the prompt exceeds the batch size (%d tokens, batch %d)\n", __func__, (int) inp.size(), llama_n_batch(ctx_tgt));88 89 return 1;90 }91 92 LOG("\n\n");93 94 for (auto id : inp) {95 LOG("%s", common_token_to_piece(ctx_tgt, id).c_str());96 }97 98 int n_predict = 0;99 int n_drafted = 0;100 int n_accept = 0;101 102 // used to determine end of generation103 bool has_eos = false;104 105 llama_seq_id seq_id = 0;106 107 // ================================================108 // everything until here is standard initialization109 // the relevant stuff for speculative decoding starts here110 111 const auto t_enc_start = ggml_time_us();112 113 // target model sampling context114 common_sampler_ptr smpl(common_sampler_init(model_tgt, params.sampling));115 116 // init the speculator117 const auto & params_spec = params.speculative;118 119 struct common_speculative * spec = common_speculative_init(params.speculative, 1);120 121 if (spec == nullptr) {122 LOG_ERR("%s", "failed to initialize speculative decoding\n");123 return 1;124 }125 126 // eval the prompt on the target and feed it to the speculative implementation(s)127 {128 llama_batch batch_prompt = llama_batch_init(inp.size(), 0, 1);129 for (size_t i = 0; i < inp.size() - 1; ++i) {130 common_batch_add(batch_prompt, inp[i], i, { seq_id }, false);131 }132 133 llama_decode(ctx_tgt, batch_prompt);134 135 if (!common_speculative_process(spec, batch_prompt)) {136 LOG_ERR("%s", "failed to process speculative prompt\n");137 return 1;138 }139 }140 141 // note: keep the last token separate!142 llama_token id_last = inp.back();143 144 // all tokens currently in the target context145 llama_tokens prompt_tgt(inp.begin(), inp.end() - 1);146 prompt_tgt.reserve(llama_n_ctx(ctx_tgt));147 148 int n_past = inp.size() - 1;149 150 common_speculative_begin(spec, seq_id, prompt_tgt);151 152 llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);153 154 llama_tokens draft;155 156 common_prompt_checkpoint ckpt;157 158 const auto t_enc_end = ggml_time_us();159 160 const auto t_dec_start = ggml_time_us();161 162 while (true) {163 // generate or reuse draft tokens164 //165 // this is the most important part of the speculation. the more probable tokens that are provided here166 // the better the performance will be. in theory, this computation can be performed asynchronously and even167 // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens168 // from a cache or lookup tables.169 //170 if (draft.empty()) {171 ckpt.update_pos(172 prompt_tgt.size(),173 llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), seq_id),174 llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id));175 176 if (use_ckpt_dft) {177 ckpt.update_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);178 }179 180 // determine the max draft that fits the remaining context and generation budget181 int n_draft_max = (int) llama_n_ctx(ctx_tgt) - n_past - 2;182 if (params.n_predict >= 0) {183 n_draft_max = std::min(n_draft_max, params.n_predict - n_predict - 1);184 }185 n_draft_max = std::max(n_draft_max, 0);186 187 // generate a new draft188 common_speculative_get_draft_params(spec, seq_id) = {189 /* .drafting = */ true,190 /* .n_max = */ n_draft_max,191 /* .pos0 = */ n_past,192 /* .id_last = */ id_last,193 /* .prompt = */ &prompt_tgt,194 /* .result = */ &draft, // output195 };196 common_speculative_draft(spec);197 198 // save a checkpoint of the target context before evaluating the draft199 // this allows us to restore the state if partial draft acceptance occurs200 if (!draft.empty()) {201 if (use_ckpt_tgt) {202 ckpt.update_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);203 }204 }205 206 // reset the draft context to the checkpoint before verification207 if (ctx_dft) {208 if (use_ckpt_dft) {209 ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);210 }211 212 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);213 }214 } else {215 // we have a previous (partial) draft to reuse from checkpoint restoration216 if (use_ckpt_tgt) {217 GGML_ASSERT(!ckpt.empty());218 }219 }220 221 // always have a token to evaluate from before - id_last222 common_batch_clear(batch_tgt);223 common_batch_add (batch_tgt, id_last, n_past++, { seq_id }, true);224 225 // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]226 {227 for (size_t i = 0; i < draft.size(); ++i) {228 common_batch_add(batch_tgt, draft[i], n_past + i, { seq_id }, true);229 }230 231 //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());232 233 llama_decode(ctx_tgt, batch_tgt);234 }235 236 // feed the batch to the speculative implementation(s) - this drives the draft model, MTP, Eagle3, etc.237 if (!common_speculative_process(spec, batch_tgt)) {238 LOG_ERR("%s", "failed to process speculative batch\n");239 break;240 }241 242 // only save the sampler sampler state if we use checkpoints243 common_sampler_ptr smpl_save;244 if (use_ckpt_tgt) {245 smpl_save.reset(common_sampler_clone(smpl.get()));246 }247 248 // save the size of the draft being verified249 const size_t n_draft = draft.size();250 251 // sample from the full target batch and return the accepted tokens based on the target sampler252 //253 // for each token to be accepted, the sampler would have to sample that same token254 // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the255 // available logits from the batch and sample the next token until we run out of logits or the sampler256 // disagrees with the draft257 //258 auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft);259 260 //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());261 262 GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token263 264 // check for partial draft acceptance:265 // if the context doesn't support partial sequence removal, restore the checkpoint266 // and make the accepted tokens the new partial draft for the next iteration267 if (use_ckpt_tgt && ids.size() - 1 < n_draft) {268 LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft);269 270 draft = std::move(ids);271 272 {273 ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);274 275 llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1);276 }277 278 if (ctx_dft) {279 ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);280 281 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1);282 }283 284 prompt_tgt.resize(ckpt.n_tokens);285 smpl = std::move(smpl_save);286 287 n_past = (int) prompt_tgt.size();288 289 continue;290 }291 292 common_speculative_accept(spec, seq_id, ids.size() - 1);293 294 // full acceptance: consume the draft and commit accepted tokens295 n_past += ids.size() - 1;296 n_drafted += n_draft; // note: we ignore the discarded small drafts297 n_accept += ids.size() - 1;298 n_predict += ids.size();299 300 // process the accepted tokens and update contexts301 //302 // this is the standard token post-processing that we normally do303 // in this case, we do it for a group of accepted tokens at once304 //305 for (size_t i = 0; i < ids.size(); ++i) {306 prompt_tgt.push_back(id_last);307 308 id_last = ids[i];309 310 if (llama_vocab_is_eog(vocab, id_last)) {311 has_eos = true;312 break;313 }314 315 const std::string token_str = common_token_to_piece(ctx_tgt, id_last);316 317 if (params.use_color && i + 1 < ids.size()) {318 LOG("\u001b[%dm%s\u001b[37m", (36 - 0 % 6), token_str.c_str());319 } else {320 LOG("%s", token_str.c_str());321 }322 }323 324 LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last);325 326 // clear the draft since it has been consumed327 draft.clear();328 329 {330 LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);331 332 llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1);333 334 if (ctx_dft) {335 llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, n_past, -1);336 }337 }338 339 if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {340 break;341 }342 }343 344 auto t_dec_end = ggml_time_us();345 346 const int n_input = inp.size();347 348 LOG("\n\n");349 350 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));351 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));352 353 LOG_INF("\n");354 LOG_INF("n_draft = %d\n", params_spec.draft.n_max);355 LOG_INF("n_predict = %d\n", n_predict);356 LOG_INF("n_drafted = %d\n", n_drafted);357 LOG_INF("n_accept = %d\n", n_accept);358 LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted);359 360 LOG_INF("\n");361 LOG_INF("draft:\n\n");362 common_speculative_print_stats(spec);363 364 LOG_INF("\n");365 LOG_INF("target:\n\n");366 common_perf_print(ctx_tgt, smpl.get());367 368 llama_batch_free(batch_tgt);369 370 common_speculative_free(spec);371 372 llama_backend_free();373 374 LOG("\n\n");375 376 return 0;377}378