Skip to content

Commit 31c6ecf

Browse files
ggerganovPetros Sideris
authored and
baramofme
committed
spec : parallel drafting support (ggml-org#22838)
* spec : refactor * spec : drop support for incompatible vocabs * spec : update common_speculative_init() * cont : pass seq_id * cont : dedup ctx_seq_rm_type * server : sketch the ctx_dft decode loop * server : draft prompt cache and checkpoints * server : improve ctx names * server, spec : transition to unified spec context * cont : sync main and drft contexts * cont : async drft eval when possible * cont : handle non-ckpt models * cont : pass correct n_past for drafting * cont : process images throught the draft context * spec : handle draft running out of context * server : fix mtmd draft processing * server : fix URL for draft model * server : add comment * server : clean-up + dry * speculative-simple : update * spec : fix n_past type * server : fix slot ctx_drft ptr * tools : update readme * naming : improve consistency * spec : refactor for multi-sequence speculative context * cont : prepare params * cont : prepare params * spec : support parallel drafts * server : support parallel drafting * llama : reuse device buffers when possible * server, spec : clean-up * cont : clean-up * cont : minor * spec : reset `drafting` flag at the end * spec : introduce `common_speculative_process()` * spec : allow for multiple spec types (chain of speculators) * replace old type field of type common_speculative_type in the common_params_speculative struct with a vector to allow multiple types to be specified * introduce common_get_enabled_speculative_impls(const std::vector<enum common_speculative_type>) to figure out which implementations the user has enabled * introduce common_speculative_type_from_names(const std::vector<std::string> & names) to parse the already user provided spec types * all speculators run sequentially, best one wins (we verify its drafted tokens) * maximize expected accepted tokens for current round by calculating the product between the probability of accepting current token (n_acc_tokens / n_gen_drafts) and the draft's length --------- Co-authored-by: Petros Sideris <petros.sideris@nokia.com>
1 parent 0f53c83 commit 31c6ecf

7 files changed

Lines changed: 66 additions & 118 deletions

File tree

‎common/arg.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3554,7 +3554,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
35543554
).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL"));
35553555
add_opt(common_arg(
35563556
{"--spec-type"}, common_speculative_all_types_str(),
3557-
string_format("comma-separated list of types of speculative decoding to use (default: %s)\n",
3557+
string_format("type of speculative decoding to use when no draft model is provided (default: %s)\n",
35583558
common_speculative_type_name_str(params.speculative.types).c_str()),
35593559
[](common_params & params, const std::string & value) {
35603560
const auto enabled_types = string_split<std::string>(value, ',');
@@ -4047,10 +4047,10 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
40474047
{"--spec-default"},
40484048
string_format("enable default speculative decoding config"),
40494049
[](common_params & params) {
4050-
params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MOD;
4051-
params.speculative.ngram_size_n = 24;
4052-
params.speculative.n_min = 48;
4053-
params.speculative.n_max = 64;
4050+
params.speculative.types = { COMMON_SPECULATIVE_TYPE_NGRAM_MOD };
4051+
params.speculative.ngram_mod.n_match = 24;
4052+
params.speculative.ngram_mod.n_min = 48;
4053+
params.speculative.ngram_mod.n_max = 64;
40544054
}
40554055
).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
40564056

‎common/common.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1455,7 +1455,7 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) {
14551455

14561456
// try to remove the last tokens
14571457
if (!llama_memory_seq_rm(mem, 0, 1, -1)) {
1458-
LOG_TRC("%s: the context does not support partial sequence removal\n", __func__);
1458+
LOG_WRN("%s: the context does not support partial sequence removal\n", __func__);
14591459
res = COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
14601460
goto done;
14611461
}

‎common/speculative.cpp‎

Lines changed: 40 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121

2222
const std::map<std::string, common_speculative_type> common_speculative_type_from_name_map = {
2323
{"none", COMMON_SPECULATIVE_TYPE_NONE},
24-
{"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE},
25-
{"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
24+
{"draft", COMMON_SPECULATIVE_TYPE_DRAFT},
25+
{"eagle3", COMMON_SPECULATIVE_TYPE_EAGLE3},
2626
{"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
2727
{"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
2828
{"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
@@ -145,15 +145,15 @@ struct common_speculative_impl {
145145
virtual void accept(llama_seq_id seq_id, uint16_t n_accepted) = 0;
146146
};
147147

148-
struct common_speculative_impl_draft_simple : public common_speculative_impl {
148+
struct common_speculative_state_draft : public common_speculative_impl {
149149
common_params_speculative_draft params;
150150

151151
llama_batch batch;
152152

153153
std::vector<common_sampler_ptr> smpls;
154154

155-
common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq)
156-
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq)
155+
common_speculative_state_draft(const common_params_speculative & params, uint32_t n_seq)
156+
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT, n_seq)
157157
, params(params.draft)
158158
{
159159
auto * ctx_dft = this->params.ctx_dft;
@@ -206,7 +206,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
206206
}
207207
}
208208

209-
~common_speculative_impl_draft_simple() override {
209+
~common_speculative_state_draft() override {
210210
llama_batch_free(batch);
211211
}
212212

@@ -340,11 +340,11 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
340340
}
341341
};
342342

343-
struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
343+
struct common_speculative_state_eagle3 : public common_speculative_impl {
344344
//common_params_speculative_eagle3 params;
345345

346-
common_speculative_impl_draft_eagle3(const common_params_speculative & /*params*/, uint32_t n_seq)
347-
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) {}
346+
common_speculative_state_eagle3(const common_params_speculative & /*params*/, uint32_t n_seq)
347+
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_EAGLE3, n_seq) {}
348348

349349
void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
350350
// noop
@@ -365,13 +365,13 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
365365
};
366366

367367
// state of self-speculation (simple implementation, not ngram-map)
368-
struct common_speculative_impl_ngram_simple : public common_speculative_impl {
368+
struct common_speculative_state_ngram_simple : public common_speculative_impl {
369369
common_params_speculative_ngram_map params;
370370

371371
// shared across all sequences
372372
common_ngram_simple_config config;
373373

374-
common_speculative_impl_ngram_simple(
374+
common_speculative_state_ngram_simple(
375375
const common_params_speculative & params, uint32_t n_seq,
376376
common_ngram_simple_config config)
377377
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq)
@@ -405,13 +405,13 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
405405
}
406406
};
407407

408-
struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
408+
struct common_speculative_state_ngram_map_k : public common_speculative_impl {
409409
common_params_speculative_ngram_map params;
410410

411411
// n_seq configs
412412
std::vector<common_ngram_map> config;
413413

414-
common_speculative_impl_ngram_map_k(
414+
common_speculative_state_ngram_map_k(
415415
const common_params_speculative & params,
416416
const common_ngram_map & config,
417417
uint32_t n_seq)
@@ -453,7 +453,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
453453
}
454454
};
455455

456-
struct common_speculative_impl_ngram_mod : public common_speculative_impl {
456+
struct common_speculative_state_ngram_mod : public common_speculative_impl {
457457
common_params_speculative_ngram_mod params;
458458

459459
// shared across all sequences
@@ -475,7 +475,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
475475

476476
std::vector<seq_info> sinfos;
477477

478-
common_speculative_impl_ngram_mod(
478+
common_speculative_state_ngram_mod(
479479
const common_params_speculative & params,
480480
uint32_t n_seq)
481481
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq)
@@ -611,8 +611,8 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
611611
}
612612

613613
mod.reset();
614-
n_low = 0;
615-
i_last = 0;
614+
sinfo.n_low = 0;
615+
sinfo.i_last = 0;
616616
}
617617
} else {
618618
sinfo.n_low = 0;
@@ -621,7 +621,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
621621
}
622622
};
623623

624-
struct common_speculative_impl_ngram_cache : public common_speculative_impl {
624+
struct common_speculative_state_ngram_cache : public common_speculative_impl {
625625
common_params_speculative_ngram_cache params;
626626

627627
uint16_t n_draft;
@@ -639,7 +639,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
639639

640640
std::vector<seq_info> sinfos;
641641

642-
common_speculative_impl_ngram_cache(
642+
common_speculative_state_ngram_cache(
643643
const common_params_speculative & params,
644644
uint32_t n_seq,
645645
uint16_t n_draft,
@@ -775,7 +775,7 @@ static common_ngram_map get_common_ngram_map(
775775
return common_ngram_map(size_key, size_value, key_only, min_hits);
776776
}
777777

778-
static common_speculative_impl_ngram_cache create_state_ngram_cache(
778+
static common_speculative_state_ngram_cache create_state_ngram_cache(
779779
const common_speculative_config & config,
780780
uint32_t n_seq,
781781
const std::string & path_static,
@@ -786,7 +786,7 @@ static common_speculative_impl_ngram_cache create_state_ngram_cache(
786786
bool save_static = false;
787787
bool save_dynamic = false;
788788

789-
common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
789+
common_speculative_state_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
790790

791791
return state;
792792
}
@@ -818,8 +818,8 @@ const char * common_speculative_all_types_str() {
818818
std::string common_speculative_type_to_str(common_speculative_type type) {
819819
switch (type) {
820820
case COMMON_SPECULATIVE_TYPE_NONE: return "none";
821-
case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: return "draft-simple";
822-
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3";
821+
case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft";
822+
case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3";
823823
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple";
824824
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k";
825825
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
@@ -872,9 +872,9 @@ common_speculative * common_speculative_init(common_params_speculative & params,
872872
{
873873
uint32_t enabled_configs = common_get_enabled_speculative_configs(params.types);
874874

875-
bool has_draft_model_path = !params.draft.mparams.path.empty();
875+
bool has_draft = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT));
876+
bool has_draft_model = !params.draft.mparams.path.empty();
876877

877-
bool has_draft_simple = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE));
878878
// bool has_mtp = false; // TODO: add MTP here
879879
bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3
880880

@@ -906,18 +906,18 @@ common_speculative * common_speculative_init(common_params_speculative & params,
906906
if (has_ngram_cache) {
907907
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, params));
908908
}
909-
if (has_draft_simple) {
910-
if (!has_draft_model_path) {
909+
if (has_draft) {
910+
if (!has_draft_model) {
911911
LOG_WRN("%s: draft model is not specified - cannot use 'draft' type\n", __func__);
912-
has_draft_simple = false;
912+
has_draft = false;
913913
}
914-
} else if (has_draft_model_path) {
914+
} else if (has_draft_model) {
915915
LOG_WRN("%s: draft model is specified but 'draft' speculative type is not explicitly enabled - enabling it\n", __func__);
916-
has_draft_simple = true;
916+
has_draft = true;
917917
}
918918

919-
if (has_draft_simple) {
920-
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, params));
919+
if (has_draft) {
920+
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT, params));
921921
}
922922
// TODO: add MTP here
923923
if (has_draft_eagle3) {
@@ -932,12 +932,12 @@ common_speculative * common_speculative_init(common_params_speculative & params,
932932
switch (config.type) {
933933
case COMMON_SPECULATIVE_TYPE_NONE:
934934
break;
935-
case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: {
936-
impls.push_back(std::make_unique<common_speculative_impl_draft_simple>(config.params, n_seq));
935+
case COMMON_SPECULATIVE_TYPE_DRAFT: {
936+
impls.push_back(std::make_unique<common_speculative_state_draft>(config.params, n_seq));
937937
break;
938938
}
939-
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: {
940-
impls.push_back(std::make_unique<common_speculative_impl_draft_eagle3>(config.params, n_seq));
939+
case COMMON_SPECULATIVE_TYPE_EAGLE3: {
940+
impls.push_back(std::make_unique<common_speculative_state_eagle3>(config.params, n_seq));
941941
break;
942942
}
943943
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
@@ -950,7 +950,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
950950
/* .size_ngram = */ ngram_size_key,
951951
/* .size_mgram = */ mgram_size_value
952952
};
953-
auto state = std::make_unique<common_speculative_impl_ngram_simple>(
953+
auto state = std::make_unique<common_speculative_state_ngram_simple>(
954954
/* .params = */ config.params,
955955
/* .n_seq = */ n_seq,
956956
/* .state = */ config_simple
@@ -961,21 +961,21 @@ common_speculative * common_speculative_init(common_params_speculative & params,
961961
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
962962
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: {
963963
impls.push_back(
964-
std::make_unique<common_speculative_impl_ngram_map_k>(
964+
std::make_unique<common_speculative_state_ngram_map_k>(
965965
config.params, get_common_ngram_map(config.type, config.params.ngram_map_k), n_seq));
966966
break;
967967
}
968968
case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: {
969969
impls.push_back(
970-
std::make_unique<common_speculative_impl_ngram_mod>(config.params, n_seq));
970+
std::make_unique<common_speculative_state_ngram_mod>(config.params, n_seq));
971971
break;
972972
}
973973
case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
974974
auto state = create_state_ngram_cache(
975975
config, n_seq,
976976
params.ngram_cache.lookup_cache_static,
977977
params.ngram_cache.lookup_cache_dynamic);
978-
impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
978+
impls.push_back(std::make_unique<common_speculative_state_ngram_cache>(state));
979979
break;
980980
}
981981
default:

‎tools/cli/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@
199199
| `--spec-draft-device, -devd, --device-draft <dev1,dev2,..>` | comma-separated list of devices to use for offloading the draft model (none = don't offload)<br/>use --list-devices to see a list of available devices |
200200
| `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) |
201201
| `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)<br/>(env: LLAMA_ARG_SPEC_DRAFT_MODEL) |
202-
| `--spec-type none,draft-simple,draft-eagle3,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
202+
| `--spec-type [none\|ngram-cache\|ngram-simple\|ngram-map-k\|ngram-map-k4v\|ngram-mod]` | type of speculative decoding to use when no draft model is provided (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
203203
| `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) |
204204
| `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) |
205205
| `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) |

‎tools/server/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ For the full list of features, please refer to [server's changelog](https://gith
248248
| `--spec-draft-device, -devd, --device-draft <dev1,dev2,..>` | comma-separated list of devices to use for offloading the draft model (none = don't offload)<br/>use --list-devices to see a list of available devices |
249249
| `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) |
250250
| `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)<br/>(env: LLAMA_ARG_SPEC_DRAFT_MODEL) |
251-
| `--spec-type none,draft-simple,draft-eagle3,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
251+
| `--spec-type [none\|ngram-cache\|ngram-simple\|ngram-map-k\|ngram-map-k4v\|ngram-mod]` | type of speculative decoding to use when no draft model is provided (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
252252
| `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) |
253253
| `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) |
254254
| `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) |

0 commit comments

Comments
 (0)