|
| 1 | +#include "ggml.h" |
| 2 | +#include "common.h" |
| 3 | +#include "llama.h" |
| 4 | +#include "log.h" |
| 5 | +#include "ngram-cache.h" |
| 6 | + |
| 7 | +#include <cmath> |
| 8 | +#include <cstdint> |
| 9 | +#include <cstdio> |
| 10 | +#include <fstream> |
| 11 | +#include <string> |
| 12 | +#include <vector> |
| 13 | +#include <unordered_map> |
| 14 | + |
| 15 | +int main(int argc, char ** argv){ |
| 16 | + gpt_params params; |
| 17 | + |
| 18 | + if (!gpt_params_parse(argc, argv, params)) { |
| 19 | + return 1; |
| 20 | + } |
| 21 | + |
| 22 | + const int n_draft = params.n_draft; |
| 23 | + |
| 24 | + // init llama.cpp |
| 25 | + llama_backend_init(); |
| 26 | + llama_numa_init(params.numa); |
| 27 | + |
| 28 | + llama_model * model = NULL; |
| 29 | + llama_context * ctx = NULL; |
| 30 | + |
| 31 | + // load the model |
| 32 | + std::tie(model, ctx) = llama_init_from_gpt_params(params); |
| 33 | + llama_set_rng_seed(ctx, params.seed); |
| 34 | + GGML_ASSERT(llama_n_vocab(model) < (1 << 16)); |
| 35 | + |
| 36 | + // tokenize the prompt |
| 37 | + const bool add_bos = llama_should_add_bos_token(model); |
| 38 | + LOG("add_bos tgt: %d\n", add_bos); |
| 39 | + |
| 40 | + std::vector<llama_token> inp; |
| 41 | + inp = ::llama_tokenize(ctx, params.prompt, add_bos, true); |
| 42 | + |
| 43 | + llama_ngram_cache ngram_cache_context; |
| 44 | + llama_ngram_cache ngram_cache_dynamic; |
| 45 | + llama_ngram_cache ngram_cache_static; |
| 46 | + int64_t t_draft_flat_us = 0; |
| 47 | + int64_t t_draft_us = 0; |
| 48 | + |
| 49 | + { |
| 50 | + const int64_t t_start_draft_us = ggml_time_us(); |
| 51 | + |
| 52 | + if (!params.lookup_cache_static.empty()) { |
| 53 | + try { |
| 54 | + ngram_cache_static = llama_ngram_cache_load(params.lookup_cache_static); |
| 55 | + } catch (std::system_error const &) { |
| 56 | + fprintf(stderr, "error: failed to open static lookup cache: %s", params.lookup_cache_static.c_str()); |
| 57 | + exit(1); |
| 58 | + } |
| 59 | + } |
| 60 | + |
| 61 | + if (!params.lookup_cache_dynamic.empty()) { |
| 62 | + try { |
| 63 | + ngram_cache_dynamic = llama_ngram_cache_load(params.lookup_cache_dynamic); |
| 64 | + } catch (std::system_error const &) {} // if the file does not exist it will simply be created at the end of the program |
| 65 | + } |
| 66 | + |
| 67 | + t_draft_flat_us += ggml_time_us() - t_start_draft_us; |
| 68 | + } |
| 69 | + |
| 70 | + const int n_input = inp.size(); |
| 71 | + const int n_ctx = params.n_ctx; |
| 72 | + |
| 73 | + int n_drafted = 0; |
| 74 | + int n_accept = 0; |
| 75 | + |
| 76 | + const int64_t t_start_ms = ggml_time_ms(); |
| 77 | + |
| 78 | + // Iterate over input tokens in chunks of size n_ctx. |
| 79 | + // Each chunk is treated as if a sequential generation but with pre-determined tokens to ensure reproducibility. |
| 80 | + for (int i_start = 0; i_start + n_ctx < n_input; i_start += n_ctx) { |
| 81 | + const std::vector<llama_token> inp_slice(inp.begin() + i_start, inp.begin() + i_start + n_ctx); |
| 82 | + std::vector<llama_token> pseudo_output; |
| 83 | + pseudo_output.push_back(inp_slice[0]); |
| 84 | + |
| 85 | + while ((int) pseudo_output.size() < n_ctx) { |
| 86 | + // Simulate drafting and decoding from draft: |
| 87 | + std::vector<llama_token> draft; |
| 88 | + draft.push_back(pseudo_output.back()); |
| 89 | + |
| 90 | + { |
| 91 | + const int64_t t_start_draft_us = ggml_time_us(); |
| 92 | + llama_ngram_cache_draft(pseudo_output, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static); |
| 93 | + t_draft_us += ggml_time_us() - t_start_draft_us; |
| 94 | + } |
| 95 | + |
| 96 | + n_drafted += draft.size() - 1; |
| 97 | + |
| 98 | + for (size_t j = 1; j < draft.size() && (int) pseudo_output.size() < n_ctx; ++j) { |
| 99 | + const llama_token ground_truth = inp_slice[pseudo_output.size()]; |
| 100 | + const llama_token drafted = draft[j]; |
| 101 | + |
| 102 | + if (ground_truth != drafted) { |
| 103 | + break; |
| 104 | + } |
| 105 | + |
| 106 | + ++n_accept; |
| 107 | + pseudo_output.push_back(ground_truth); |
| 108 | + |
| 109 | + { |
| 110 | + const int64_t t_start_draft_us = ggml_time_us(); |
| 111 | + llama_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false); |
| 112 | + t_draft_us += ggml_time_us() - t_start_draft_us; |
| 113 | + } |
| 114 | + } |
| 115 | + |
| 116 | + // After each simulated batch decoding simulate the sampling of a single token: |
| 117 | + if ((int) pseudo_output.size() < n_ctx) { |
| 118 | + pseudo_output.push_back(inp_slice[pseudo_output.size()]); |
| 119 | + { |
| 120 | + const int64_t t_start_draft_us = ggml_time_us(); |
| 121 | + llama_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false); |
| 122 | + t_draft_us += ggml_time_us() - t_start_draft_us; |
| 123 | + } |
| 124 | + } |
| 125 | + |
| 126 | + draft.erase(draft.begin()); |
| 127 | + |
| 128 | + } |
| 129 | + if (i_start > 0 && i_start / 100000 != (i_start - n_ctx) / 100000) { |
| 130 | + const int64_t t_now_ms = ggml_time_ms(); |
| 131 | + const int64_t eta_ms = (n_input - i_start) * (t_now_ms - t_start_ms) / i_start; |
| 132 | + const int64_t eta_min = eta_ms / (60*1000); |
| 133 | + const int64_t eta_s = (eta_ms - 60*1000*eta_min) / 1000; |
| 134 | + |
| 135 | + LOG_TEE("%d/%d done, ETA: %02ld:%02ld\n", i_start, n_input, eta_min, eta_s); |
| 136 | + } |
| 137 | + |
| 138 | + // After each chunk, update the dynamic ngram cache with the context ngram cache: |
| 139 | + llama_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context); |
| 140 | + ngram_cache_context.clear(); |
| 141 | + } |
| 142 | + |
| 143 | + LOG_TEE("\n"); |
| 144 | + |
| 145 | + LOG_TEE("\n"); |
| 146 | + LOG_TEE("n_draft = %d\n", n_draft); |
| 147 | + LOG_TEE("n_predict = %d\n", n_input - n_input % n_ctx); |
| 148 | + LOG_TEE("n_drafted = %d\n", n_drafted); |
| 149 | + LOG_TEE("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3); |
| 150 | + LOG_TEE("t_draft = %.2f ms, %.2f us per token, %.2f tokens per second\n", |
| 151 | + t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us)); |
| 152 | + LOG_TEE("n_accept = %d\n", n_accept); |
| 153 | + LOG_TEE("accept = %.3f%%\n", 100.0f * n_accept / n_drafted); |
| 154 | + |
| 155 | + llama_free(ctx); |
| 156 | + llama_free_model(model); |
| 157 | + |
| 158 | + llama_backend_free(); |
| 159 | + |
| 160 | + fprintf(stderr, "\n\n"); |
| 161 | + |
| 162 | + return 0; |
| 163 | +} |
0 commit comments