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, ex: 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.path.empty()) {
28 LOG_ERR("%s: --model-draft is required\n", __func__);
29 return 1;
30 }
31
32 // init llama.cpp
33 llama_backend_init();
34 llama_numa_init(numa: 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 model
43 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: model_tgt);
49
50 // load the draft model
51 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 params.tensor_buft_overrides = params.speculative.tensor_buft_overrides;
63
64 common_init_result llama_init_dft = common_init_from_params(params);
65
66 //model_dft = llama_init_dft.model.get();
67 ctx_dft = llama_init_dft.context.get();
68
69 if (!common_speculative_are_compatible(ctx_tgt, ctx_dft)) {
70 LOG_INF("the draft model '%s' is not compatible with the target model '%s'. tokens will be translated between the draft and target models.\n", params.speculative.model.path.c_str(), params.model.path.c_str());
71 }
72
73 // Tokenize the prompt
74 std::vector<llama_token> inp;
75 inp = common_tokenize(ctx: ctx_tgt, text: params.prompt, add_special: true, parse_special: true);
76
77 if (llama_n_ctx(ctx: ctx_tgt) < (uint32_t) inp.size()) {
78 LOG_ERR("%s: the prompt exceeds the context size (%d tokens, ctx %d)\n", __func__, (int) inp.size(), llama_n_ctx(ctx_tgt));
79
80 return 1;
81 }
82
83 if (llama_n_batch(ctx: ctx_tgt) < (uint32_t) inp.size()) {
84 LOG_ERR("%s: the prompt exceeds the batch size (%d tokens, batch %d)\n", __func__, (int) inp.size(), llama_n_batch(ctx_tgt));
85
86 return 1;
87 }
88
89 LOG("\n\n");
90
91 for (auto id : inp) {
92 LOG("%s", common_token_to_piece(ctx_tgt, id).c_str());
93 }
94
95 // how many tokens to draft each time
96 int n_draft = params.speculative.n_max;
97 int n_draft_min = params.speculative.n_min;
98
99 float p_min = params.speculative.p_min;
100
101 int n_predict = 0;
102 int n_drafted = 0;
103 int n_accept = 0;
104
105 // used to determine end of generation
106 bool has_eos = false;
107
108 // ================================================
109 // everything until here is standard initialization
110 // the relevant stuff for speculative decoding starts here
111
112 const auto t_enc_start = ggml_time_us();
113
114 // target model sampling context
115 struct common_sampler * smpl = common_sampler_init(model: model_tgt, params: params.sampling);
116
117 // eval the prompt
118 llama_decode(ctx: ctx_tgt, batch: llama_batch_get_one(tokens: inp.data(), n_tokens: inp.size() - 1));
119
120 // note: keep the last token separate!
121 llama_token id_last = inp.back();
122
123 // all tokens currently in the target context
124 llama_tokens prompt_tgt(inp.begin(), inp.end() - 1);
125 prompt_tgt.reserve(n: llama_n_ctx(ctx: ctx_tgt));
126
127 int n_past = inp.size() - 1;
128
129 // init the speculator
130 struct common_speculative_params params_spec;
131 params_spec.n_draft = n_draft;
132 params_spec.n_reuse = llama_n_ctx(ctx: ctx_dft) - n_draft;
133 params_spec.p_min = p_min;
134
135 struct common_speculative * spec = common_speculative_init(ctx_tgt, ctx_dft);
136 for (auto &pair : params.speculative.replacements) {
137 common_speculative_add_replacement_tgt_dft(spec, source: pair.first.c_str(), dest: pair.second.c_str());
138 }
139
140 llama_batch batch_tgt = llama_batch_init(n_tokens: llama_n_batch(ctx: ctx_tgt), embd: 0, n_seq_max: 1);
141
142 const auto t_enc_end = ggml_time_us();
143
144 const auto t_dec_start = ggml_time_us();
145
146 while (true) {
147 // optionally, generate draft tokens that can be appended to the target batch
148 //
149 // this is the most important part of the speculation. the more probable tokens that are provided here
150 // the better the performance will be. in theory, this computation can be performed asynchronously and even
151 // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens
152 // from a cache or lookup tables.
153 //
154 llama_tokens draft = common_speculative_gen_draft(spec, params: params_spec, prompt: prompt_tgt, id_last);
155
156 //LOG_DBG("draft: %s\n", string_from(ctx_dft, draft).c_str());
157
158 // always have a token to evaluate from before - id_last
159 common_batch_clear(batch&: batch_tgt);
160 common_batch_add (batch&: batch_tgt, id: id_last, pos: n_past++, seq_ids: { 0 }, logits: true);
161
162 // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]
163 {
164 // do not waste time on small drafts
165 if (draft.size() < (size_t) n_draft_min) {
166 draft.clear();
167 }
168
169 for (size_t i = 0; i < draft.size(); ++i) {
170 common_batch_add(batch&: batch_tgt, id: draft[i], pos: n_past + i, seq_ids: { 0 }, logits: true);
171 }
172
173 //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());
174
175 llama_decode(ctx: ctx_tgt, batch: batch_tgt);
176 }
177
178 // sample from the full target batch and return the accepted tokens based on the target sampler
179 //
180 // for each token to be accepted, the sampler would have to sample that same token
181 // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the
182 // available logits from the batch and sample the next token until we run out of logits or the sampler
183 // disagrees with the draft
184 //
185 const auto ids = common_sampler_sample_and_accept_n(gsmpl: smpl, ctx: ctx_tgt, draft);
186
187 //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());
188
189 GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token
190
191 n_past += ids.size() - 1;
192 n_drafted += draft.size(); // note: we ignore the discarded small drafts
193 n_accept += ids.size() - 1;
194 n_predict += ids.size();
195
196 // process the accepted tokens and update contexts
197 //
198 // this is the standard token post-processing that we normally do
199 // in this case, we do it for a group of accepted tokens at once
200 //
201 for (size_t i = 0; i < ids.size(); ++i) {
202 prompt_tgt.push_back(x: id_last);
203
204 id_last = ids[i];
205
206 if (llama_vocab_is_eog(vocab, token: id_last)) {
207 has_eos = true;
208 break;
209 }
210
211 const std::string token_str = common_token_to_piece(ctx: ctx_tgt, token: id_last);
212
213 if (params.use_color && i + 1 < ids.size()) {
214 LOG("\u001b[%dm%s\u001b[37m", (36 - 0 % 6), token_str.c_str());
215 } else {
216 LOG("%s", token_str.c_str());
217 }
218 }
219
220 LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last);
221
222 {
223 LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);
224
225 llama_memory_seq_rm(mem: llama_get_memory(ctx: ctx_tgt), seq_id: 0, p0: n_past, p1: -1);
226 }
227
228 if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {
229 break;
230 }
231 }
232
233 auto t_dec_end = ggml_time_us();
234
235 const int n_input = inp.size();
236
237 LOG("\n\n");
238
239 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));
240 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));
241
242 LOG_INF("\n");
243 LOG_INF("n_draft = %d\n", n_draft);
244 LOG_INF("n_predict = %d\n", n_predict);
245 LOG_INF("n_drafted = %d\n", n_drafted);
246 LOG_INF("n_accept = %d\n", n_accept);
247 LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted);
248
249 LOG_INF("\n");
250 LOG_INF("draft:\n\n");
251
252 llama_perf_context_print(ctx: ctx_dft);
253
254 LOG_INF("\n");
255 LOG_INF("target:\n\n");
256 common_perf_print(ctx: ctx_tgt, gsmpl: smpl);
257
258 common_sampler_free(gsmpl: smpl);
259 common_speculative_free(spec);
260
261 llama_backend_free();
262
263 LOG("\n\n");
264
265 return 0;
266}
267