KBaba7/llama.cpp

1#include "arg.h"2#include "ggml.h"3#include "common.h"4#include "ngram-cache.h"5#include "sampling.h"6#include "log.h"7#include "llama.h"8 9#include <cstdint>10#include <cstdio>11#include <fstream>12#include <string>13#include <vector>14 15int main(int argc, char ** argv){16    common_params params;17 18    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) {19        return 1;20    }21 22    common_init();23 24    // max. number of additional tokens to draft if match is found25    const int n_draft = params.speculative.n_max;26 27    const bool dump_kv_cache = params.dump_kv_cache;28 29    // init llama.cpp30    llama_backend_init();31    llama_numa_init(params.numa);32 33    // load the model34    common_init_result llama_init = common_init_from_params(params);35 36    llama_model * model = llama_init.model.get();37    llama_context * ctx = llama_init.context.get();38 39    const llama_vocab * vocab = llama_model_get_vocab(model);40 41    // tokenize the prompt42    std::vector<llama_token> inp;43    inp = common_tokenize(ctx, params.prompt, true, true);44 45    common_ngram_cache ngram_cache_context;46    common_ngram_cache ngram_cache_dynamic;47    common_ngram_cache ngram_cache_static;48    int64_t t_draft_flat_us = 0;49    int64_t t_draft_us = 0;50 51    {52        // Fill up context ngram cache with tokens from user input:53        const int64_t t_start_draft_us = ggml_time_us();54        common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, inp.size(), false);55 56        if (!params.lookup_cache_static.empty()) {57            try {58                ngram_cache_static = common_ngram_cache_load(params.lookup_cache_static);59            } catch (std::ifstream::failure const &) {60                LOG_ERR("failed to open static lookup cache: %s", params.lookup_cache_static.c_str());61                exit(1);62            }63        }64 65        if (!params.lookup_cache_dynamic.empty()) {66            try {67                ngram_cache_dynamic = common_ngram_cache_load(params.lookup_cache_dynamic);68            } catch (std::ifstream::failure const &) {} // if the file does not exist it will simply be created at the end of the program69        }70 71        t_draft_flat_us += ggml_time_us() - t_start_draft_us;72    }73 74    const int max_context_size     = llama_n_ctx(ctx);75    const int max_tokens_list_size = max_context_size - 4;76 77    if ((int) inp.size() > max_tokens_list_size) {78        LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size);79        return 1;80    }81 82    LOG("\n\n");83 84    for (auto id : inp) {85        LOG("%s", common_token_to_piece(ctx, id).c_str());86    }87 88    fflush(stderr);89 90    const int n_input = inp.size();91 92    const auto t_enc_start = ggml_time_us();93 94    llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1));95    llama_decode(ctx, llama_batch_get_one(&inp.back(),           1));96 97    const auto t_enc_end = ggml_time_us();98 99    int n_predict = 0;100    int n_drafted = 0;101    int n_accept  = 0;102 103    int n_past = inp.size();104 105    bool has_eos = false;106 107    struct common_sampler * smpl = common_sampler_init(model, params.sampling);108 109    std::vector<llama_token> draft;110 111    llama_batch batch_tgt = llama_batch_init(params.n_ctx, 0, 1);112 113    // debug114    struct llama_kv_cache_view kvc_view = llama_kv_cache_view_init(ctx, 1);115 116    const auto t_dec_start = ggml_time_us();117 118    while (true) {119        // debug120        if (dump_kv_cache) {121            llama_kv_cache_view_update(ctx, &kvc_view);122            common_kv_cache_dump_view_seqs(kvc_view, 40);123        }124 125        // print current draft sequence126        LOG_DBG("drafted %s\n", string_from(ctx, draft).c_str());127 128        int i_dft = 0;129        while (true) {130            // sample from the target model131            llama_token id = common_sampler_sample(smpl, ctx, i_dft);132 133            common_sampler_accept(smpl, id, true);134 135            const std::string token_str = common_token_to_piece(ctx, id);136 137            if (!params.use_color) {138                LOG("%s", token_str.c_str());139            }140 141            if (llama_vocab_is_eog(vocab, id)) {142                has_eos = true;143            }144 145            ++n_predict;146 147            // check if the target token matches the draft148            if (i_dft < (int) draft.size() && id == draft[i_dft]) {149                LOG_DBG("the sampled target token matches the %dth drafted token (%d, '%s') - accepted\n", i_dft, id, token_str.c_str());150                ++n_accept;151                ++n_past;152                ++i_dft;153                inp.push_back(id);154                {155                    // Update context ngram cache with the newly accepted token:156                    const int64_t t_start_draft_us = ggml_time_us();157                    common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false);158                    t_draft_us += ggml_time_us() - t_start_draft_us;159                }160 161                if (params.use_color) {162                    // color accepted draft token163                    LOG("\033[34m%s\033[0m", token_str.c_str());164                    fflush(stdout);165                }166                continue;167            }168 169            if (params.use_color) {170                LOG("%s", token_str.c_str());171            }172            fflush(stdout);173 174 175            LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());176 177            draft.clear();178            draft.push_back(id);179            inp.push_back(id);180            {181                // Update context ngram cache with the newly accepted token:182                const int64_t t_start_draft_us = ggml_time_us();183                common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false);184                t_draft_us += ggml_time_us() - t_start_draft_us;185            }186            break;187        }188 189        if ((params.n_predict > 0 && n_predict > params.n_predict) || has_eos) {190            break;191        }192 193        // KV cache management194        // clean the cache of draft tokens that weren't accepted195        llama_kv_cache_seq_rm(ctx, 0, n_past, -1);196 197        common_batch_clear(batch_tgt);198        common_batch_add(batch_tgt, draft[0], n_past, { 0 }, true);199 200        // Draft already contains a single token sampled from the model:201        GGML_ASSERT(draft.size() == 1);202        GGML_ASSERT(draft[0] == inp.back());203        const int64_t t_start_draft_us = ggml_time_us();204 205        common_ngram_cache_draft(inp, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static);206 207        for (size_t i = 1; i < draft.size(); ++i) {208            common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);209        }210 211        t_draft_us += ggml_time_us() - t_start_draft_us;212        n_drafted += draft.size() - 1;213 214        llama_decode(ctx, batch_tgt);215        ++n_past;216 217        draft.erase(draft.begin());218    }219 220    auto t_dec_end = ggml_time_us();221 222    // Update dynamic ngram cache with context ngram cache and save it to disk:223    common_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context);224    common_ngram_cache_save(ngram_cache_dynamic, params.lookup_cache_dynamic);225 226    LOG("\n\n");227 228    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));229    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));230 231    LOG_INF("\n");232    LOG_INF("n_draft      = %d\n", n_draft);233    LOG_INF("n_predict    = %d\n", n_predict);234    LOG_INF("n_drafted    = %d\n", n_drafted);235    LOG_INF("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3);236    LOG_INF("t_draft      = %.2f ms, %.2f us per token, %.2f tokens per second\n",237            t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us));238    LOG_INF("n_accept     = %d\n", n_accept);239    LOG_INF("accept       = %.3f%%\n", 100.0f * n_accept / n_drafted);240 241    LOG_INF("\ntarget:\n\n");242    common_perf_print(ctx, smpl);243 244    common_sampler_free(smpl);245 246    llama_batch_free(batch_tgt);247 248    llama_backend_free();249 250    LOG("\n\n");251 252    return 0;253}254