Brunobkr/llama.cpp_AlgMor24_github

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