Skip to content

Commit eef008e

Browse files
ggerganovggerganov
committed
speculative : reuse grammar parser + better logs and comments
1 parent 57c111f commit eef008e

File tree

1 file changed

+41
-43
lines changed

1 file changed

+41
-43
lines changed

‎examples/speculative/speculative.cpp

Lines changed: 41 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -114,47 +114,34 @@ int main(int argc, char ** argv) {
114114
struct llama_grammar * grammar_dft = NULL;
115115
struct llama_grammar * grammar_tgt = NULL;
116116

117-
grammar_parser::parse_state parsed_grammar_dft;
118-
grammar_parser::parse_state parsed_grammar_tgt;
117+
grammar_parser::parse_state parsed_grammar;
119118

120119
std::vector<llama_grammar *> grammar_mem(n_draft, NULL);
121120

121+
// if requested - load the grammar, error checking is omitted for brevity
122122
if (!params.grammar.empty()) {
123-
// dft
124-
{
125-
parsed_grammar_dft = grammar_parser::parse(params.grammar.c_str());
126-
// will be empty (default) if there are parse errors
127-
if (parsed_grammar_dft.rules.empty()) {
128-
return 1;
129-
}
130-
131-
std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar_dft.c_rules());
132-
grammar_dft = llama_grammar_init(grammar_rules.data(), grammar_rules.size(), parsed_grammar_dft.symbol_ids.at("root"));
123+
parsed_grammar = grammar_parser::parse(params.grammar.c_str());
124+
// will be empty (default) if there are parse errors
125+
if (parsed_grammar.rules.empty()) {
126+
return 1;
133127
}
134128

135-
// tgt
136-
{
137-
parsed_grammar_tgt = grammar_parser::parse(params.grammar.c_str());
138-
// will be empty (default) if there are parse errors
139-
if (parsed_grammar_tgt.rules.empty()) {
140-
return 1;
141-
}
142-
143-
std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar_tgt.c_rules());
144-
grammar_tgt = llama_grammar_init(grammar_rules.data(), grammar_rules.size(), parsed_grammar_tgt.symbol_ids.at("root"));
145-
}
129+
std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar.c_rules());
130+
grammar_dft = llama_grammar_init(grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
131+
grammar_tgt = llama_grammar_init(grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
146132
}
147133

148134
const auto t_dec_start = ggml_time_us();
149135

150136
while (true) {
151137
LOG("drafted: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx_dft, drafted));
152138

153-
// sample from the drafted tokens if any
154139
int i_dft = 0;
155140
while (true) {
141+
// sample from the target model
156142
const llama_token id = llama_sample_token(ctx_tgt, NULL, grammar_tgt, params, last_tokens, candidates, i_dft);
157143

144+
// remember which tokens were sampled - used for repetition penalties during sampling
158145
last_tokens.erase(last_tokens.begin());
159146
last_tokens.push_back(id);
160147

@@ -170,8 +157,9 @@ int main(int argc, char ** argv) {
170157

171158
++n_predict;
172159

160+
// check if the draft matches the target
173161
if (i_dft < (int) drafted.size() && id == drafted[i_dft]) {
174-
LOG("drafted token %d accepted\n", id);
162+
LOG("the sampled target token matches the %dth drafted token (%d, '%s') - accepted\n", i_dft, id, token_str.c_str());
175163
++n_accept;
176164
++n_past_tgt;
177165
++n_past_dft;
@@ -180,25 +168,20 @@ int main(int argc, char ** argv) {
180168
continue;
181169
}
182170

171+
// the drafted token was rejected or we are out of drafted tokens
172+
183173
if (i_dft < (int) drafted.size()) {
184-
LOG("drafted token %d rejected\n", id);
174+
LOG("the %dth drafted token (%d, '%s') does not match the sampled target token (%d, '%s') - rejected\n",
175+
i_dft, drafted[i_dft], llama_token_to_piece(ctx_dft, drafted[i_dft]).c_str(), id, token_str.c_str());
185176

186177
if (grammar_mem[i_dft]) {
187178
grammar_dft = llama_grammar_copy(grammar_mem[i_dft]);
188-
LOG("restored grammar %d\n", i_dft);
179+
LOG("restored draft grammar state %d\n", i_dft);
189180
}
181+
} else {
182+
LOG("out of drafted tokens\n");
190183
}
191184

192-
for (auto & g : grammar_mem) {
193-
if (g) {
194-
llama_grammar_free(g);
195-
g = NULL;
196-
}
197-
}
198-
199-
LOG("i_dft = %d, drafted.size() = %d\n", i_dft, (int) drafted.size());
200-
201-
// the drafted token was rejected or we are out of drafted tokens
202185
llama_eval(ctx_dft, &id, 1, n_past_dft, params.n_threads);
203186
++n_past_dft;
204187

@@ -212,11 +195,20 @@ int main(int argc, char ** argv) {
212195
break;
213196
}
214197

198+
for (int i = 0; i < (int) grammar_mem.size(); ++i) {
199+
auto & g = grammar_mem[i];
200+
if (g) {
201+
LOG("freeing grammar state %d\n", i);
202+
llama_grammar_free(g);
203+
g = NULL;
204+
}
205+
}
206+
215207
if (n_predict > params.n_predict || has_eos) {
216208
break;
217209
}
218210

219-
// sample n_draft tokens from the draft model picking the best token
211+
// sample n_draft tokens from the draft model using greedy decoding
220212
int n_past_cur = n_past_dft;
221213
for (int i = 0; i < n_draft; ++i) {
222214
// remember the grammar state
@@ -244,11 +236,13 @@ int main(int argc, char ** argv) {
244236
LOG(" - draft candidate %3d: %6d (%8.3f) '%s'\n", i, cur_p.data[i].id, cur_p.data[i].p, llama_token_to_piece(ctx_dft, cur_p.data[i].id).c_str());
245237
}
246238

247-
// too low probability, stop drafting
239+
// TODO: better logic?
248240
if (cur_p.data[0].p < 2*cur_p.data[1].p) {
241+
LOG("stopping drafting, probability too low: %8.f < 2*%8.f\n", cur_p.data[0].p, cur_p.data[1].p);
249242
break;
250243
}
251244

245+
// drafted token
252246
const llama_token id = cur_p.data[0].id;
253247

254248
if (grammar_dft != NULL) {
@@ -258,17 +252,21 @@ int main(int argc, char ** argv) {
258252
drafted.push_back(id);
259253
++n_drafted;
260254

261-
if (i < n_draft - 1) {
262-
// evaluate the drafted token on the draft model
263-
llama_eval(ctx_dft, &drafted.back(), 1, n_past_cur, params.n_threads);
264-
++n_past_cur;
255+
// no need to evaluate the last drafted token, since we won't use the result
256+
if (i == n_draft - 1) {
257+
break;
265258
}
259+
260+
// evaluate the drafted token on the draft model
261+
llama_eval(ctx_dft, &drafted.back(), 1, n_past_cur, params.n_threads);
262+
++n_past_cur;
266263
}
267264

268265
// evaluate the target model on the drafted tokens
269266
llama_eval(ctx_tgt, drafted.data(), drafted.size(), n_past_tgt, params.n_threads);
270267
++n_past_tgt;
271268

269+
// the first token is always proposed by the traget model before the speculation loop
272270
drafted.erase(drafted.begin());
273271
}
274272

0 commit comments

Comments
 (0)