1414
1515#include < algorithm>
1616#include < cassert>
17+ #include < cmath>
1718#include < cstring>
1819#include < iomanip>
1920#include < map>
@@ -138,6 +139,7 @@ struct common_speculative_impl {
138139 const common_speculative_type type;
139140
140141 uint32_t n_seq;
142+ int32_t n_max; // maximum draft length after implementation-specific limits
141143
142144 size_t n_call_begin = 0 ; // number of times this implementation was called for refresh.
143145 size_t n_call_draft = 0 ; // number of times this implementation was called for generation.
@@ -157,7 +159,7 @@ struct common_speculative_impl {
157159 int64_t t_draft_us = 0 ; // total time spent in generating drafts in this implementation in microseconds.
158160 int64_t t_accept_us = 0 ; // total time spent in accumulation of this implementation in microseconds.
159161
160- common_speculative_impl (common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {}
162+ common_speculative_impl (common_speculative_type type, uint32_t n_seq, int32_t n_max ) : type(type), n_seq(n_seq), n_max(n_max ) {}
161163
162164 virtual ~common_speculative_impl () = default ;
163165
@@ -182,7 +184,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
182184 std::vector<common_sampler_ptr> smpls;
183185
184186 common_speculative_impl_draft_simple (const common_params_speculative & params, uint32_t n_seq)
185- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE , n_seq)
187+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE , n_seq, params.draft.n_max )
186188 , params(params.draft)
187189 {
188190 auto * ctx_dft = this ->params .ctx_dft ;
@@ -452,7 +454,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
452454 std::vector<float > g_embd_buf;
453455
454456 common_speculative_impl_draft_eagle3 (const common_params_speculative & params, uint32_t n_seq)
455- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 , n_seq)
457+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 , n_seq, params.draft.n_max )
456458 , params(params.draft)
457459 {
458460 SPC_TRC (" %s" , " adding speculative implementation 'draft-eagle3'\n " );
@@ -937,7 +939,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
937939
938940 common_speculative_impl_draft_dflash (const common_params_speculative & params, uint32_t n_seq,
939941 common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH )
940- : common_speculative_impl(type, n_seq)
942+ : common_speculative_impl(type, n_seq, params.draft.n_max )
941943 , params(params.draft)
942944 , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK )
943945 {
@@ -983,6 +985,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
983985 this ->params .n_max = std::min (this ->params .n_max , n_draft_max);
984986 this ->params .n_min = std::min (this ->params .n_min , n_draft_max);
985987 }
988+ this ->n_max = this ->params .n_max ;
986989
987990 batch = llama_batch_init (llama_n_batch (ctx_dft), 0 , n_seq);
988991 batch_inject = llama_batch_init (llama_n_batch (ctx_dft), n_embd_dec, n_seq);
@@ -1315,7 +1318,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
13151318 std::vector<std::vector<float >> chain_h;
13161319
13171320 common_speculative_impl_draft_mtp (const common_params_speculative & params, uint32_t n_seq)
1318- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP , n_seq)
1321+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP , n_seq, params.draft.n_max )
13191322 , params(params.draft)
13201323 {
13211324 auto * ctx_tgt = this ->params .ctx_tgt ;
@@ -1382,6 +1385,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
13821385 c.reserve ((size_t ) (this ->params .n_max + 1 ) * n_embd);
13831386 }
13841387 }
1388+ this ->n_max = this ->params .n_max ;
13851389
13861390 pending_h.assign (n_seq, std::vector<float >(n_embd, 0 .0f ));
13871391
@@ -1726,7 +1730,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
17261730 common_speculative_impl_ngram_simple (
17271731 const common_params_speculative & params, uint32_t n_seq,
17281732 common_ngram_simple_config config)
1729- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE , n_seq)
1733+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE , n_seq, params.ngram_simple.size_m )
17301734 , params(params.ngram_simple)
17311735 , config(config)
17321736 {
@@ -1770,7 +1774,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
17701774 const common_ngram_map & config,
17711775 uint32_t n_seq)
17721776 : common_speculative_impl(config.key_only ? COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K
1773- : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V , n_seq)
1777+ : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V , n_seq, config.size_value )
17741778 {
17751779 for (uint32_t i = 0 ; i < n_seq; i++) {
17761780 this ->config .push_back (config);
@@ -1841,7 +1845,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
18411845 common_speculative_impl_ngram_mod (
18421846 const common_params_speculative & params,
18431847 uint32_t n_seq)
1844- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD , n_seq)
1848+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD , n_seq, params.ngram_mod.n_max )
18451849 , params(params.ngram_mod)
18461850 , mod(params.ngram_mod.n_match, 4 *1024 *1024 )
18471851 , verbose(std::getenv(" LLAMA_TRACE" ) != nullptr ) {
@@ -2017,7 +2021,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
20172021 const std::string & path_dynamic,
20182022 bool save_dynamic,
20192023 bool save_static)
2020- : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE , n_seq)
2024+ : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE , n_seq, n_draft )
20212025 , params(params.ngram_cache)
20222026 , n_draft(n_draft)
20232027 , save_dynamic(save_dynamic)
@@ -2138,6 +2142,8 @@ struct common_speculative {
21382142
21392143 // which implementaion was used for a given seq_id
21402144 std::vector<common_speculative_impl *> impl_last;
2145+
2146+ std::vector<double > synth_probs;
21412147};
21422148
21432149static common_ngram_map get_common_ngram_map (
@@ -2316,6 +2322,101 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
23162322 return n_max;
23172323}
23182324
2325+ int32_t common_speculative_n_max (const common_speculative * spec) {
2326+ int32_t n_max = 0 ;
2327+
2328+ if (spec == nullptr ) {
2329+ return n_max;
2330+ }
2331+
2332+ for (const auto & impl : spec->impls ) {
2333+ n_max = std::max (n_max, std::max (0 , impl->n_max ));
2334+ }
2335+
2336+ return n_max;
2337+ }
2338+
2339+ std::vector<double > common_speculative_synth_rates_resolve (const common_params_speculative * spec, int32_t n_max) {
2340+ const bool has_length = spec->synth_len != -1.0 ;
2341+ const bool has_rates = !spec->synth_rates .empty ();
2342+
2343+ if (!has_length && !has_rates) {
2344+ return {};
2345+ }
2346+ if (has_length && has_rates) {
2347+ throw std::invalid_argument (" synthetic acceptance length and rates are mutually exclusive" );
2348+ }
2349+
2350+ if (n_max <= 0 ) {
2351+ throw std::invalid_argument (" synthetic acceptance requires at least one speculative token" );
2352+ }
2353+
2354+ if (has_rates) {
2355+ const auto & rates = spec->synth_rates ;
2356+ if (rates.size () != (size_t ) n_max) {
2357+ throw std::invalid_argument (string_format (
2358+ " synthetic acceptance rates must contain %d values, got %zu" , n_max, rates.size ()));
2359+ }
2360+
2361+ for (size_t i = 0 ; i < rates.size (); ++i) {
2362+ if (!std::isfinite (rates[i]) || rates[i] < 0.0 || rates[i] > 1.0 ) {
2363+ throw std::invalid_argument (" synthetic acceptance rates must be finite and within [0, 1]" );
2364+ }
2365+ if (i > 0 && rates[i] > rates[i - 1 ]) {
2366+ throw std::invalid_argument (" synthetic acceptance rates must be monotonically non-increasing" );
2367+ }
2368+ }
2369+
2370+ return rates;
2371+ }
2372+
2373+ const double length = spec->synth_len ;
2374+ const double length_max = (double ) n_max + 1.0 ;
2375+ if (!std::isfinite (length) || length < 1.0 || length > length_max) {
2376+ throw std::invalid_argument (string_format (
2377+ " synthetic acceptance length must be finite and within [1, %.0f]" , length_max));
2378+ }
2379+
2380+ double p = 0.0 ;
2381+ if (length == length_max) {
2382+ p = 1.0 ;
2383+ } else if (length > 1.0 ) {
2384+ double p_min = 0.0 ;
2385+ double p_max = 1.0 ;
2386+ for (int i = 0 ; i < 32 ; ++i) {
2387+ const double p_mid = 0.5 * (p_min + p_max);
2388+ double sum = 0.0 ;
2389+ double term = p_mid;
2390+ for (int32_t j = 0 ; j < n_max; ++j) {
2391+ sum += term;
2392+ term *= p_mid;
2393+ }
2394+
2395+ if (sum < length - 1.0 ) {
2396+ p_min = p_mid;
2397+ } else {
2398+ p_max = p_mid;
2399+ }
2400+ }
2401+ p = 0.5 * (p_min + p_max);
2402+ }
2403+
2404+ std::vector<double > rates;
2405+ rates.reserve (n_max);
2406+ double rate = p;
2407+ for (int32_t i = 0 ; i < n_max; ++i) {
2408+ rates.push_back (rate);
2409+ rate *= p;
2410+ }
2411+
2412+ return rates;
2413+ }
2414+
2415+ const std::vector<double > & common_speculative_get_synth_probs (const common_speculative * spec) {
2416+ GGML_ASSERT (spec);
2417+ return spec->synth_probs ;
2418+ }
2419+
23192420common_params common_base_params_to_speculative (const common_params & params) {
23202421 const bool has_draft = params.speculative .has_dft ();
23212422
@@ -2568,13 +2669,39 @@ common_speculative * common_speculative_init(common_params_speculative & params,
25682669 return nullptr ;
25692670 }
25702671
2571- auto * result = new common_speculative {
2572- /* .dparams = */ common_speculative_draft_params_vec (n_seq),
2573- /* .impls = */ std::move (impls),
2574- /* .impl_last = */ std::vector<common_speculative_impl *>(n_seq, nullptr )
2575- };
2672+ common_speculative_ptr result (new common_speculative {
2673+ /* .dparams = */ common_speculative_draft_params_vec (n_seq),
2674+ /* .impls = */ std::move (impls),
2675+ /* .impl_last = */ std::vector<common_speculative_impl *>(n_seq, nullptr ),
2676+ /* .synth_probs = */ {},
2677+ });
25762678
2577- return result;
2679+ const int32_t n_max_configured = common_speculative_n_max (¶ms);
2680+ const int32_t n_max_effective = common_speculative_n_max (result.get ());
2681+ const auto rates = common_speculative_synth_rates_resolve (¶ms, n_max_effective);
2682+
2683+ std::vector<std::string> rates_str;
2684+ rates_str.reserve (rates.size ());
2685+ result->synth_probs .reserve (rates.size ());
2686+ double rate_prev = 1.0 ;
2687+ double acceptance_length = 1.0 ;
2688+ for (const double rate : rates) {
2689+ result->synth_probs .push_back (rate_prev > 0.0 ? rate / rate_prev : 0.0 );
2690+ rates_str.push_back (string_format (" %.6g" , rate));
2691+ rate_prev = rate;
2692+ acceptance_length += rate;
2693+ }
2694+ if (!result->synth_probs .empty ()) {
2695+ SPC_WRN (" %s" , " synthetic speculative acceptance is enabled for benchmarking; generated output is not valid\n " );
2696+ if (n_max_effective != n_max_configured) {
2697+ SPC_WRN (" synthetic acceptance draft limit was reduced from %d to %d by the initialized speculative implementations\n " ,
2698+ n_max_configured, n_max_effective);
2699+ }
2700+ SPC_INF (" synthetic acceptance: n_max = %zu, mean length = %.6f, rates = [%s]\n " ,
2701+ rates.size (), acceptance_length, string_join (rates_str, " , " ).c_str ());
2702+ }
2703+
2704+ return result.release ();
25782705}
25792706
25802707void common_speculative_free (common_speculative * spec) {
0 commit comments