KBaba7/llama.cpp

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 <cstdio>9#include <cstring>10#include <string>11#include <vector>12 13int main(int argc, char ** argv) {14    common_params params;15 16    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) {17        return 1;18    }19 20    if (params.n_predict < -1) {21        LOG_ERR("%s: --n-predict must be >= -1\n", __func__);22        return 1;23    }24 25    common_init();26 27    if (params.speculative.model.empty()) {28        LOG_ERR("%s: --model-draft is required\n", __func__);29        return 1;30    }31 32    // init llama.cpp33    llama_backend_init();34    llama_numa_init(params.numa);35 36    llama_model * model_tgt = NULL;37    //llama_model * model_dft = NULL;38 39    llama_context * ctx_tgt = NULL;40    llama_context * ctx_dft = NULL;41 42    // load the target model43    common_init_result llama_init_tgt = common_init_from_params(params);44 45    model_tgt = llama_init_tgt.model.get();46    ctx_tgt   = llama_init_tgt.context.get();47 48    const llama_vocab * vocab = llama_model_get_vocab(model_tgt);49 50    // load the draft model51    params.devices      = params.speculative.devices;52    params.model        = params.speculative.model;53    params.n_ctx        = params.speculative.n_ctx;54    params.n_batch      = params.speculative.n_ctx > 0 ? params.speculative.n_ctx : params.n_batch;55    params.n_gpu_layers = params.speculative.n_gpu_layers;56 57    if (params.speculative.cpuparams.n_threads > 0) {58        params.cpuparams.n_threads = params.speculative.cpuparams.n_threads;59    }60 61    params.cpuparams_batch.n_threads = params.speculative.cpuparams_batch.n_threads;62    common_init_result llama_init_dft = common_init_from_params(params);63 64    //model_dft = llama_init_dft.model.get();65    ctx_dft   = llama_init_dft.context.get();66 67    if (!common_speculative_are_compatible(ctx_tgt, ctx_dft)) {68        return 1;69    }70 71    // Tokenize the prompt72    std::vector<llama_token> inp;73    inp = common_tokenize(ctx_tgt, params.prompt, true, true);74 75    if (llama_n_ctx(ctx_tgt) < (uint32_t) inp.size()) {76        LOG_ERR("%s: the prompt exceeds the context size (%d tokens, ctx %d)\n", __func__, (int) inp.size(), llama_n_ctx(ctx_tgt));77 78        return 1;79    }80 81    if (llama_n_batch(ctx_tgt) < (uint32_t) inp.size()) {82        LOG_ERR("%s: the prompt exceeds the batch size (%d tokens, batch %d)\n", __func__, (int) inp.size(), llama_n_batch(ctx_tgt));83 84        return 1;85    }86 87    LOG("\n\n");88 89    for (auto id : inp) {90        LOG("%s", common_token_to_piece(ctx_tgt, id).c_str());91    }92 93    // how many tokens to draft each time94    int n_draft     = params.speculative.n_max;95    int n_draft_min = params.speculative.n_min;96 97    float p_min = params.speculative.p_min;98 99    int n_predict = 0;100    int n_drafted = 0;101    int n_accept  = 0;102 103    // used to determine end of generation104    bool has_eos = false;105 106    // ================================================107    // everything until here is standard initialization108    // the relevant stuff for speculative decoding starts here109 110    const auto t_enc_start = ggml_time_us();111 112    // target model sampling context113    struct common_sampler * smpl = common_sampler_init(model_tgt, params.sampling);114 115    // eval the prompt116    llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1));117 118    // note: keep the last token separate!119    llama_token id_last = inp.back();120 121    // all tokens currently in the target context122    llama_tokens prompt_tgt(inp.begin(), inp.end() - 1);123    prompt_tgt.reserve(llama_n_ctx(ctx_tgt));124 125    int n_past = inp.size() - 1;126 127    // init the speculator128    struct common_speculative_params params_spec;129    params_spec.n_draft = n_draft;130    params_spec.n_reuse = llama_n_ctx(ctx_dft) - n_draft;131    params_spec.p_min   = p_min;132 133    struct common_speculative * spec = common_speculative_init(ctx_dft);134 135    llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);136 137    const auto t_enc_end = ggml_time_us();138 139    const auto t_dec_start = ggml_time_us();140 141    while (true) {142        // optionally, generate draft tokens that can be appended to the target batch143        //144        // this is the most important part of the speculation. the more probable tokens that are provided here145        // the better the performance will be. in theory, this computation can be performed asynchronously and even146        // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens147        // from a cache or lookup tables.148        //149        llama_tokens draft = common_speculative_gen_draft(spec, params_spec, prompt_tgt, id_last);150 151        //LOG_DBG("draft: %s\n", string_from(ctx_dft, draft).c_str());152 153        // always have a token to evaluate from before - id_last154        common_batch_clear(batch_tgt);155        common_batch_add  (batch_tgt, id_last, n_past++, { 0 }, true);156 157        // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]158        {159            // do not waste time on small drafts160            if (draft.size() < (size_t) n_draft_min) {161                draft.clear();162            }163 164            for (size_t i = 0; i < draft.size(); ++i) {165                common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);166            }167 168            //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());169 170            llama_decode(ctx_tgt, batch_tgt);171        }172 173        // sample from the full target batch and return the accepted tokens based on the target sampler174        //175        // for each token to be accepted, the sampler would have to sample that same token176        // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the177        // available logits from the batch and sample the next token until we run out of logits or the sampler178        // disagrees with the draft179        //180        const auto ids = common_sampler_sample_and_accept_n(smpl, ctx_tgt, draft);181 182        //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());183 184        GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token185 186        n_past    += ids.size() - 1;187        n_drafted += draft.size(); // note: we ignore the discarded small drafts188        n_accept  += ids.size() - 1;189        n_predict += ids.size();190 191        // process the accepted tokens and update contexts192        //193        // this is the standard token post-processing that we normally do194        // in this case, we do it for a group of accepted tokens at once195        //196        for (size_t i = 0; i < ids.size(); ++i) {197            prompt_tgt.push_back(id_last);198 199            id_last = ids[i];200 201            if (llama_vocab_is_eog(vocab, id_last)) {202                has_eos = true;203                break;204            }205 206            const std::string token_str = common_token_to_piece(ctx_tgt, id_last);207 208            if (params.use_color && i + 1 < ids.size()) {209                LOG("\u001b[%dm%s\u001b[37m", (36 - 0 % 6), token_str.c_str());210            } else {211                LOG("%s", token_str.c_str());212            }213        }214 215        LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last);216 217        {218            LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);219 220            llama_kv_cache_seq_rm(ctx_tgt, 0, n_past, -1);221        }222 223        if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {224            break;225        }226    }227 228    auto t_dec_end = ggml_time_us();229 230    const int n_input = inp.size();231 232    LOG("\n\n");233 234    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));235    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));236 237    LOG_INF("\n");238    LOG_INF("n_draft   = %d\n", n_draft);239    LOG_INF("n_predict = %d\n", n_predict);240    LOG_INF("n_drafted = %d\n", n_drafted);241    LOG_INF("n_accept  = %d\n", n_accept);242    LOG_INF("accept    = %.3f%%\n", 100.0f * n_accept / n_drafted);243 244    LOG_INF("\n");245    LOG_INF("draft:\n\n");246 247    llama_perf_context_print(ctx_dft);248 249    LOG_INF("\n");250    LOG_INF("target:\n\n");251    common_perf_print(ctx_tgt, smpl);252 253    common_sampler_free(smpl);254    common_speculative_free(spec);255 256    llama_backend_free();257 258    LOG("\n\n");259 260    return 0;261}262