Allow tuning of the best args for speculative decoding. (#1595)
* wip: build spec tuner for spefic args * wip: test different reward system * spec-tune: fix the reward to find best params given a good TPS * spec-tune: refactor logic for its own file * minor clean for comments and modules
This commit is contained in:
Samuel Oliveira Alves
GitHub
parent
0a6e4335f7
commit
3de81530c5
|
|
@ -80,6 +80,8 @@ add_library(${TARGET} STATIC
|
|||
peg-parser.cpp
|
||||
peg-parser.h
|
||||
speculative.cpp
|
||||
spec-tuner.cpp
|
||||
spec-tuner.h
|
||||
unicode.cpp
|
||||
unicode.h
|
||||
ngram-mod.cpp
|
||||
|
|
|
|||
|
|
@ -1021,6 +1021,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
|
|||
params.speculative.p_min = std::stof(argv[i]);
|
||||
return true;
|
||||
}
|
||||
if (arg == "--spec-autotune") {
|
||||
params.speculative.autotune = true;
|
||||
return true;
|
||||
}
|
||||
if (arg == "--chunks") {
|
||||
CHECK_ARG
|
||||
params.n_chunks = std::stoi(argv[i]);
|
||||
|
|
@ -2578,6 +2582,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
|
|||
options.push_back({ "*", "--spec-ngram-size-m N", "ngram size M for ngram-simple/ngram-map speculative decoding, length of draft m-gram (default: %d)\n", params.speculative.ngram_size_m });
|
||||
|
||||
options.push_back({ "*", "--spec-ngram-min-hits N", "minimum hits for ngram-map speculative decoding (default: %d)\n", params.speculative.ngram_min_hits });
|
||||
options.push_back({ "*", "--spec-autotune", "automatically tune speculative params to maximize tokens/sec" });
|
||||
|
||||
options.push_back({ "retrieval" });
|
||||
options.push_back({ "retrieval", " --context-file FNAME", "file to load context from (repeat to specify multiple files)" });
|
||||
|
|
|
|||
|
|
@ -200,6 +200,8 @@ struct common_params_speculative {
|
|||
std::string cache_type_k = ""; // KV cache data type for K for the draft model
|
||||
std::string cache_type_v = ""; // KV cache data type for V for the draft model
|
||||
|
||||
bool autotune = false; // automatically optimize speculative params for max tokens/sec
|
||||
|
||||
bool has_dft() const {
|
||||
return !model.empty() || !params.empty();
|
||||
//return !mparams_dft.path.empty() || !mparams_dft.hf_repo.empty();
|
||||
|
|
|
|||
|
|
@ -0,0 +1,356 @@
|
|||
#include "spec-tuner.h"
|
||||
|
||||
#include "ggml.h"
|
||||
#include "log.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <iomanip>
|
||||
#include <random>
|
||||
|
||||
int spec_tuner_coord::find_nearest_arm(float value) const {
|
||||
int idx = 0;
|
||||
float best_dist = 1e30f;
|
||||
for (int i = 0; i < (int)arms.size(); i++) {
|
||||
float dist = std::fabs(arms[i].value - value);
|
||||
if (dist < best_dist) {
|
||||
best_dist = dist;
|
||||
idx = i;
|
||||
}
|
||||
}
|
||||
return idx;
|
||||
}
|
||||
|
||||
int spec_tuner_coord::select_epsilon_greedy(double epsilon) const {
|
||||
static thread_local std::mt19937 rng(std::random_device{}());
|
||||
std::uniform_real_distribution<double> coin(0.0, 1.0);
|
||||
|
||||
if (coin(rng) < epsilon) {
|
||||
std::uniform_int_distribution<int> dist(0, (int)arms.size() - 1);
|
||||
return dist(rng);
|
||||
}
|
||||
return best_idx;
|
||||
}
|
||||
|
||||
void spec_tuner_coord::update(double reward) {
|
||||
auto & arm = arms[current_idx];
|
||||
arm.N += 1;
|
||||
arm.Q += (reward - arm.Q) / arm.N;
|
||||
|
||||
double best_Q = -1e30;
|
||||
for (int i = 0; i < (int)arms.size(); i++) {
|
||||
if (arms[i].N > 0 && arms[i].Q > best_Q) {
|
||||
best_Q = arms[i].Q;
|
||||
best_idx = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void spec_tuner_coord::reset_scores() {
|
||||
for (auto & arm : arms) {
|
||||
arm.Q = 0.0;
|
||||
arm.N = 0;
|
||||
}
|
||||
current_idx = user_idx;
|
||||
best_idx = user_idx;
|
||||
}
|
||||
|
||||
void spec_tuner_coord::build_grid_float(float lo, float hi, int n_points, float user_value) {
|
||||
arms.clear();
|
||||
for (int i = 0; i < n_points; i++) {
|
||||
float v = lo + (hi - lo) * i / std::max(1, n_points - 1);
|
||||
arms.push_back({v, 0.0, 0});
|
||||
}
|
||||
bool found = false;
|
||||
for (auto & a : arms) {
|
||||
if (std::fabs(a.value - user_value) < 1e-6f) { found = true; break; }
|
||||
}
|
||||
if (!found) {
|
||||
arms.push_back({user_value, 0.0, 0});
|
||||
std::sort(arms.begin(), arms.end(), [](const spec_tuner_arm & a, const spec_tuner_arm & b) {
|
||||
return a.value < b.value;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
void spec_tuner_coord::build_grid_int(int lo, int hi, int step, int user_value) {
|
||||
arms.clear();
|
||||
for (int v = lo; v <= hi; v += step) {
|
||||
arms.push_back({(float)v, 0.0, 0});
|
||||
}
|
||||
if (arms.empty() || (int)arms.back().value != hi) {
|
||||
arms.push_back({(float)hi, 0.0, 0});
|
||||
}
|
||||
bool found = false;
|
||||
for (auto & a : arms) {
|
||||
if ((int)a.value == user_value) { found = true; break; }
|
||||
}
|
||||
if (!found && user_value >= lo && user_value <= hi) {
|
||||
arms.push_back({(float)user_value, 0.0, 0});
|
||||
std::sort(arms.begin(), arms.end(), [](const spec_tuner_arm & a, const spec_tuner_arm & b) {
|
||||
return a.value < b.value;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
void spec_tuner::reset_exploration() {
|
||||
n_resets++;
|
||||
LOG_DBG("Autotune task change detected (n_low=%d) — resetting MAB (reset #%d)\n", n_low, n_resets);
|
||||
for (auto & coord : coords) {
|
||||
coord.reset_scores();
|
||||
}
|
||||
n_low = 0;
|
||||
cooldown = cooldown_max;
|
||||
step_ema = 0.0;
|
||||
n_calls = 0;
|
||||
}
|
||||
|
||||
void spec_tuner::write_best(common_params_speculative & params) const {
|
||||
for (const auto & coord : coords) {
|
||||
float val = coord.arms[coord.best_idx].value;
|
||||
if (coord.name == "n_max") params.n_max = (int32_t)val;
|
||||
else if (coord.name == "p_min") params.p_min = val;
|
||||
else if (coord.name == "n_min") params.n_min = (int32_t)val;
|
||||
else if (coord.name == "ngram_size_n") params.ngram_size_n = (uint16_t)val;
|
||||
else if (coord.name == "ngram_size_m") params.ngram_size_m = (uint16_t)val;
|
||||
else if (coord.name == "ngram_min_hits") params.ngram_min_hits = (uint16_t)val;
|
||||
}
|
||||
}
|
||||
|
||||
void spec_tuner::init(common_speculative_type type, const common_params_speculative & user_params) {
|
||||
enabled = true;
|
||||
spec_type = type;
|
||||
coords.clear();
|
||||
n_calls = 0;
|
||||
n_requests = 0;
|
||||
ema_tps = 0.0;
|
||||
step_ema = 0.0;
|
||||
n_low = 0;
|
||||
cooldown = 0;
|
||||
n_resets = 0;
|
||||
t_tuner_us = 0;
|
||||
last_n_drafted = 0;
|
||||
|
||||
// all types get n_max
|
||||
// For simplicity we will create a fixed grid of possible values
|
||||
{
|
||||
spec_tuner_coord coord;
|
||||
coord.name = "n_max";
|
||||
int hi = std::max(16, (int)user_params.n_max);
|
||||
coord.build_grid_int(1, hi, 1, user_params.n_max);
|
||||
coords.push_back(std::move(coord));
|
||||
}
|
||||
|
||||
if (type == COMMON_SPECULATIVE_TYPE_DRAFT) {
|
||||
{
|
||||
spec_tuner_coord coord;
|
||||
coord.name = "p_min";
|
||||
coord.build_grid_float(0.0f, 0.95f, 11, user_params.p_min);
|
||||
coords.push_back(std::move(coord));
|
||||
}
|
||||
{
|
||||
spec_tuner_coord coord;
|
||||
coord.name = "n_min";
|
||||
coord.build_grid_int(0, 6, 1, user_params.n_min);
|
||||
coords.push_back(std::move(coord));
|
||||
}
|
||||
}
|
||||
|
||||
// Ngram can change only n_max/n_min per call
|
||||
if (type == COMMON_SPECULATIVE_TYPE_NGRAM_MOD) {
|
||||
{
|
||||
spec_tuner_coord coord;
|
||||
coord.name = "n_min";
|
||||
int hi = std::max(0, std::min(4, (int)user_params.n_max - 1));
|
||||
coord.build_grid_int(0, hi, 1, user_params.n_min);
|
||||
coords.push_back(std::move(coord));
|
||||
}
|
||||
}
|
||||
|
||||
for (auto & coord : coords) {
|
||||
float user_val = 0.0f;
|
||||
if (coord.name == "n_max") user_val = (float)user_params.n_max;
|
||||
else if (coord.name == "p_min") user_val = user_params.p_min;
|
||||
else if (coord.name == "n_min") user_val = (float)user_params.n_min;
|
||||
else if (coord.name == "ngram_size_n") user_val = (float)user_params.ngram_size_n;
|
||||
else if (coord.name == "ngram_size_m") user_val = (float)user_params.ngram_size_m;
|
||||
else if (coord.name == "ngram_min_hits") user_val = (float)user_params.ngram_min_hits;
|
||||
|
||||
coord.user_idx = coord.find_nearest_arm(user_val);
|
||||
coord.best_idx = 0;
|
||||
coord.current_idx = 0;
|
||||
}
|
||||
|
||||
LOG_DBG("Autotune ε-greedy (ε=%.2f) per-draft-call, reward=per-step TPS\n", epsilon);
|
||||
for (const auto & coord : coords) {
|
||||
std::ostringstream oss;
|
||||
oss << " " << coord.name << ": [";
|
||||
for (size_t i = 0; i < coord.arms.size(); i++) {
|
||||
if (i > 0) oss << ", ";
|
||||
oss << coord.arms[i].value;
|
||||
}
|
||||
oss << "] (user=" << coord.arms[coord.user_idx].value << ")";
|
||||
LOG_DBG("%s\n", oss.str().c_str());
|
||||
}
|
||||
}
|
||||
|
||||
void spec_tuner::propose(common_params_speculative & params) {
|
||||
int64_t t_start = ggml_time_us();
|
||||
|
||||
// always select fresh arm for every draft call
|
||||
for (auto & coord : coords) {
|
||||
coord.current_idx = coord.select_epsilon_greedy(epsilon);
|
||||
|
||||
float val = coord.arms[coord.current_idx].value;
|
||||
if (coord.name == "n_max") params.n_max = (int32_t)val;
|
||||
else if (coord.name == "p_min") params.p_min = val;
|
||||
else if (coord.name == "n_min") params.n_min = (int32_t)val;
|
||||
else if (coord.name == "ngram_size_n") params.ngram_size_n = (uint16_t)val;
|
||||
else if (coord.name == "ngram_size_m") params.ngram_size_m = (uint16_t)val;
|
||||
else if (coord.name == "ngram_min_hits") params.ngram_min_hits = (uint16_t)val;
|
||||
}
|
||||
|
||||
enforce_constraints(params);
|
||||
t_tuner_us += (ggml_time_us() - t_start);
|
||||
}
|
||||
|
||||
void spec_tuner::enforce_constraints(common_params_speculative & params) {
|
||||
if (params.n_min < 0) params.n_min = 0;
|
||||
if (params.n_max < 1) params.n_max = 1;
|
||||
if (params.n_min > params.n_max) params.n_min = params.n_max;
|
||||
|
||||
if (params.p_min < 0.0f) params.p_min = 0.0f;
|
||||
if (params.p_min > 0.95f) params.p_min = 0.95f;
|
||||
|
||||
if (params.ngram_size_n < 1) params.ngram_size_n = 1;
|
||||
if (params.ngram_size_m < 1) params.ngram_size_m = 1;
|
||||
if (params.ngram_min_hits < 1) params.ngram_min_hits = 1;
|
||||
}
|
||||
|
||||
void spec_tuner::accept_feedback(int n_accepted, int n_drafted, double step_tps) {
|
||||
int64_t t_start = ggml_time_us();
|
||||
n_calls++;
|
||||
|
||||
// per-step TPS as reward: captures draft cost, verification cost, and acceptance benefit
|
||||
double reward = step_tps;
|
||||
|
||||
for (auto & coord : coords) {
|
||||
coord.update(reward);
|
||||
}
|
||||
|
||||
if (cooldown > 0) {
|
||||
cooldown--;
|
||||
if (step_ema <= 0.0) {
|
||||
step_ema = step_tps;
|
||||
} else {
|
||||
step_ema = step_ema_alpha * step_tps + (1.0 - step_ema_alpha) * step_ema;
|
||||
}
|
||||
} else if (step_ema <= 0.0) {
|
||||
step_ema = step_tps;
|
||||
} else {
|
||||
if (step_tps < step_ema * (1.0 - step_drop_pct)) {
|
||||
n_low++;
|
||||
if (n_low >= reset_after) {
|
||||
reset_exploration();
|
||||
t_tuner_us += (ggml_time_us() - t_start);
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
n_low = 0;
|
||||
}
|
||||
step_ema = step_ema_alpha * step_tps + (1.0 - step_ema_alpha) * step_ema;
|
||||
}
|
||||
|
||||
if (n_calls <= 5 || (n_calls % log_every == 0)) {
|
||||
std::ostringstream oss;
|
||||
oss << "Autotune call=" << n_calls
|
||||
<< " n_drafted=" << n_drafted
|
||||
<< " n_accepted=" << n_accepted
|
||||
<< " step_tps=" << std::fixed << std::setprecision(1) << step_tps
|
||||
<< " ema=" << std::fixed << std::setprecision(1) << step_ema;
|
||||
for (const auto & coord : coords) {
|
||||
bool is_int = (coord.name != "p_min");
|
||||
oss << " " << coord.name << "=";
|
||||
if (is_int) oss << (int)coord.arms[coord.current_idx].value;
|
||||
else oss << std::fixed << std::setprecision(2) << coord.arms[coord.current_idx].value;
|
||||
oss << "→best=";
|
||||
if (is_int) oss << (int)coord.arms[coord.best_idx].value;
|
||||
else oss << std::fixed << std::setprecision(2) << coord.arms[coord.best_idx].value;
|
||||
oss << "(Q=" << std::fixed << std::setprecision(1) << coord.arms[coord.best_idx].Q
|
||||
<< ",N=" << coord.arms[coord.best_idx].N << ")";
|
||||
}
|
||||
LOG_DBG("%s\n", oss.str().c_str());
|
||||
}
|
||||
|
||||
t_tuner_us += (ggml_time_us() - t_start);
|
||||
}
|
||||
|
||||
void spec_tuner::end_of_request(double slot_tps, int n_past, common_params_speculative & active_params) {
|
||||
int64_t t_start = ggml_time_us();
|
||||
n_requests++;
|
||||
|
||||
GGML_UNUSED(n_past);
|
||||
|
||||
if (ema_tps <= 0.0) {
|
||||
ema_tps = slot_tps;
|
||||
} else {
|
||||
ema_tps = ema_alpha * slot_tps + (1.0 - ema_alpha) * ema_tps;
|
||||
}
|
||||
|
||||
write_best(active_params);
|
||||
enforce_constraints(active_params);
|
||||
|
||||
t_tuner_us += (ggml_time_us() - t_start);
|
||||
print_best();
|
||||
}
|
||||
|
||||
void spec_tuner::print_best() const {
|
||||
{
|
||||
std::ostringstream oss;
|
||||
oss << "Autotune req=" << n_requests
|
||||
<< " calls=" << n_calls
|
||||
<< " tps=" << std::fixed << std::setprecision(2) << ema_tps;
|
||||
|
||||
if (n_resets > 0) oss << " resets=" << n_resets;
|
||||
if (n_low > 0) oss << " n_low=" << n_low;
|
||||
|
||||
oss << " best:";
|
||||
for (const auto & coord : coords) {
|
||||
bool is_int = (coord.name != "p_min");
|
||||
oss << " " << coord.name << "=";
|
||||
if (is_int) oss << (int)coord.arms[coord.best_idx].value;
|
||||
else oss << std::fixed << std::setprecision(2) << coord.arms[coord.best_idx].value;
|
||||
oss << "(Q=" << std::fixed << std::setprecision(2) << coord.arms[coord.best_idx].Q
|
||||
<< ",N=" << coord.arms[coord.best_idx].N << ")";
|
||||
}
|
||||
|
||||
if (!coords.empty()) {
|
||||
oss << " | n_max arms:";
|
||||
for (const auto & arm : coords[0].arms) {
|
||||
oss << " " << (int)arm.value << "(Q=" << std::fixed << std::setprecision(2) << arm.Q
|
||||
<< ",N=" << arm.N << ")";
|
||||
}
|
||||
}
|
||||
|
||||
oss << " tuner=" << std::fixed << std::setprecision(3) << t_tuner_us / 1000.0 << "ms";
|
||||
LOG_DBG("%s\n", oss.str().c_str());
|
||||
}
|
||||
|
||||
{
|
||||
std::ostringstream oss;
|
||||
oss << "Autotune reuse: ";
|
||||
for (const auto & coord : coords) {
|
||||
bool is_int = (coord.name != "p_min");
|
||||
if (coord.name == "n_max") oss << "--draft-max ";
|
||||
else if (coord.name == "p_min") oss << "--draft-p-min ";
|
||||
else if (coord.name == "n_min") oss << "--draft-min ";
|
||||
else if (coord.name == "ngram_size_n") oss << "--spec-ngram-size-n ";
|
||||
else if (coord.name == "ngram_size_m") oss << "--spec-ngram-size-m ";
|
||||
else if (coord.name == "ngram_min_hits") oss << "--spec-ngram-min-hits ";
|
||||
else oss << "--" << coord.name << " ";
|
||||
|
||||
if (is_int) oss << (int)coord.arms[coord.best_idx].value << " ";
|
||||
else oss << std::fixed << std::setprecision(2) << coord.arms[coord.best_idx].value << " ";
|
||||
}
|
||||
LOG_INF("%s\n", oss.str().c_str());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
#pragma once
|
||||
|
||||
#include "common.h"
|
||||
|
||||
struct spec_tuner_arm {
|
||||
float value;
|
||||
double Q = 0.0; // mean per-step Tokens-Per-Second (TPS)
|
||||
int N = 0;
|
||||
};
|
||||
|
||||
struct spec_tuner_coord {
|
||||
std::string name;
|
||||
std::vector<spec_tuner_arm> arms;
|
||||
int current_idx = 0;
|
||||
int best_idx = 0;
|
||||
int user_idx = 0;
|
||||
|
||||
int select_epsilon_greedy(double epsilon) const;
|
||||
|
||||
void update(double reward);
|
||||
|
||||
void reset_scores();
|
||||
|
||||
void build_grid_float(float lo, float hi, int n_points, float user_value);
|
||||
void build_grid_int(int lo, int hi, int step, int user_value);
|
||||
int find_nearest_arm(float value) const;
|
||||
};
|
||||
|
||||
struct spec_tuner {
|
||||
bool enabled = false;
|
||||
|
||||
double epsilon = 0.15; // 15% explore, 85% exploit
|
||||
|
||||
// task-change detection (per-call)
|
||||
// If tuner goes bad for 30 consecutive calls, reset the tuner.
|
||||
double step_ema = 0.0;
|
||||
double step_ema_alpha = 0.05;
|
||||
double step_drop_pct = 0.30;
|
||||
int n_low = 0;
|
||||
int reset_after = 30;
|
||||
int cooldown = 0;
|
||||
int cooldown_max = 50;
|
||||
int n_resets = 0;
|
||||
|
||||
int last_n_drafted = 0;
|
||||
uint64_t n_calls = 0;
|
||||
int log_every = 50;
|
||||
|
||||
// per-request tracking
|
||||
uint64_t n_requests = 0;
|
||||
int64_t t_tuner_us = 0;
|
||||
double ema_tps = 0.0;
|
||||
double ema_alpha = 0.3;
|
||||
|
||||
common_speculative_type spec_type = COMMON_SPECULATIVE_TYPE_NONE;
|
||||
std::vector<spec_tuner_coord> coords;
|
||||
|
||||
void init(common_speculative_type type, const common_params_speculative & user_params);
|
||||
void propose(common_params_speculative & params);
|
||||
void accept_feedback(int n_accepted, int n_drafted, double step_tps);
|
||||
void end_of_request(double slot_tps, int n_past, common_params_speculative & active_params);
|
||||
void enforce_constraints(common_params_speculative & params);
|
||||
void print_best() const;
|
||||
void reset_exploration();
|
||||
|
||||
void write_best(common_params_speculative & params) const;
|
||||
};
|
||||
|
|
@ -774,6 +774,9 @@ struct common_speculative_state_ngram_cache : public common_speculative_state {
|
|||
struct common_speculative {
|
||||
std::vector<std::unique_ptr<common_speculative_state>> impls; // list of implementations to use and their states
|
||||
common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats)
|
||||
std::unique_ptr<spec_tuner> tuner;
|
||||
int last_n_drafted = 0;
|
||||
int64_t t_step_start_us = 0;
|
||||
};
|
||||
|
||||
static common_ngram_map get_common_ngram_map(const common_speculative_config & config) {
|
||||
|
|
@ -1009,6 +1012,22 @@ common_speculative * common_speculative_init(
|
|||
/* .impls = */ std::move(impls)
|
||||
};
|
||||
|
||||
// initialize autotune if requested
|
||||
if (params.autotune && !result->impls.empty()) {
|
||||
auto actual_type = result->impls[0]->type;
|
||||
if (actual_type != COMMON_SPECULATIVE_TYPE_NONE &&
|
||||
actual_type != COMMON_SPECULATIVE_TYPE_EAGLE3) {
|
||||
result->tuner = std::make_unique<spec_tuner>();
|
||||
result->tuner->init(actual_type, params);
|
||||
LOG_DBG("Autotune initialized for %s, tuning %zu parameters\n",
|
||||
common_speculative_type_to_str(actual_type).c_str(),
|
||||
result->tuner->coords.size());
|
||||
} else {
|
||||
LOG_WRN("Autotune disabled — speculative type %s is not supported for autotuning\n",
|
||||
common_speculative_type_to_str(actual_type).c_str());
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
|
@ -1034,11 +1053,18 @@ void common_speculative_begin(common_speculative * spec, const llama_tokens & pr
|
|||
|
||||
llama_tokens common_speculative_draft(
|
||||
common_speculative * spec,
|
||||
const common_params_speculative & params,
|
||||
common_params_speculative & params,
|
||||
const llama_tokens & prompt_tgt, // specified in target model vocab
|
||||
llama_token id_last) {
|
||||
llama_tokens result;
|
||||
|
||||
spec->t_step_start_us = ggml_time_us();
|
||||
|
||||
// apply autotune proposal if enabled
|
||||
if (spec->tuner && spec->tuner->enabled) {
|
||||
spec->tuner->propose(params);
|
||||
}
|
||||
|
||||
spec->curr_impl = nullptr; // reset current implementation
|
||||
|
||||
for (auto & impl : spec->impls) {
|
||||
|
|
@ -1061,10 +1087,24 @@ llama_tokens common_speculative_draft(
|
|||
}
|
||||
}
|
||||
|
||||
// store draft count for tuner feedback
|
||||
if (spec->tuner && spec->tuner->enabled) {
|
||||
spec->last_n_drafted = (int)result.size();
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) {
|
||||
if (spec->tuner && spec->tuner->enabled && spec->t_step_start_us > 0) {
|
||||
int64_t step_time_us = ggml_time_us() - spec->t_step_start_us;
|
||||
double step_tps = (step_time_us > 100)
|
||||
? (n_accepted + 1.0) * 1e6 / (double)step_time_us
|
||||
: 0.0;
|
||||
spec->tuner->accept_feedback(n_accepted, spec->last_n_drafted, step_tps);
|
||||
spec->t_step_start_us = 0;
|
||||
}
|
||||
|
||||
if (n_accepted == 0) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -1085,7 +1125,7 @@ void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) {
|
|||
}
|
||||
}
|
||||
|
||||
void common_speculative_print_stats(const common_speculative * spec) {
|
||||
void common_speculative_print_stats(const common_speculative * spec, double slot_tps, int n_decoded, int n_past, common_params_speculative * active_params) {
|
||||
if (spec == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -1111,6 +1151,16 @@ void common_speculative_print_stats(const common_speculative * spec) {
|
|||
impl->n_acc_tokens,
|
||||
str_perf.c_str());
|
||||
}
|
||||
|
||||
if (spec->tuner && spec->tuner->enabled && slot_tps > 0.0 && n_decoded > 0) {
|
||||
auto * mutable_spec = const_cast<common_speculative *>(spec);
|
||||
if (active_params) {
|
||||
mutable_spec->tuner->end_of_request(slot_tps, n_past, *active_params);
|
||||
} else {
|
||||
common_params_speculative tmp_params;
|
||||
mutable_spec->tuner->end_of_request(slot_tps, n_past, tmp_params);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
#include "llama.h"
|
||||
#include "common.h"
|
||||
#include "spec-tuner.h"
|
||||
|
||||
struct common_speculative;
|
||||
|
||||
|
|
@ -30,7 +31,7 @@ void common_speculative_begin(common_speculative * spec, const llama_tokens & pr
|
|||
// sample up to n_draft tokens and add them to the batch using the draft model
|
||||
llama_tokens common_speculative_draft(
|
||||
common_speculative * spec,
|
||||
const common_params_speculative & params,
|
||||
common_params_speculative & params,
|
||||
const llama_tokens & prompt,
|
||||
llama_token id_last);
|
||||
|
||||
|
|
@ -38,7 +39,7 @@ llama_tokens common_speculative_draft(
|
|||
void common_speculative_accept(common_speculative * spec, uint16_t n_accepted);
|
||||
|
||||
// print statistics about the speculative decoding
|
||||
void common_speculative_print_stats(const common_speculative * spec);
|
||||
void common_speculative_print_stats(const common_speculative * spec, double slot_tps = 0.0, int n_decoded = 0, int n_past = 0, common_params_speculative * active_params = nullptr);
|
||||
|
||||
// Generates speculative draft tokens using the Multi-Token Prediction (MTP) architecture.
|
||||
std::vector<llama_token> mtp_speculative_gen_draft(
|
||||
|
|
|
|||
|
|
@ -598,7 +598,8 @@ void server_slot::print_timings() const {
|
|||
draft_ratio, n_draft_accepted, n_draft_total
|
||||
);
|
||||
}
|
||||
common_speculative_print_stats(spec);
|
||||
common_speculative_print_stats(spec, n_gen_second, n_decoded, n_past,
|
||||
const_cast<common_params_speculative *>(¶ms.speculative));
|
||||
}
|
||||
|
||||
void server_metrics::init() {
|
||||
|
|
@ -2825,8 +2826,8 @@ void server_context::add_sampled_tokens() {
|
|||
// generate draft tokens in speculative decoding mode
|
||||
// TODO: rework to have a single draft llama_context shared across all slots [TAG_SERVER_SPEC_REWORK]
|
||||
// perform the speculative drafting for all sequences at the same time in a single batch
|
||||
const int n_draft_max = slot.get_n_draft_max();
|
||||
if (n_draft_max > 0) {
|
||||
const int n_draft_max_pre = slot.get_n_draft_max();
|
||||
if (n_draft_max_pre > 0) {
|
||||
if (mctx) {
|
||||
// we should never reach this, as speculative is automatically disabled if mmproj is loaded
|
||||
GGML_ABORT("not supported by multimodal");
|
||||
|
|
@ -2834,7 +2835,7 @@ void server_context::add_sampled_tokens() {
|
|||
|
||||
const llama_tokens & cached_text_tokens = slot.cache_tokens.get_text_tokens();
|
||||
|
||||
const auto & params_spec = slot.params.speculative;
|
||||
auto & params_spec = slot.params.speculative;
|
||||
|
||||
if (slot.has_mtp) {
|
||||
if (!slot.mtp_hidden_state.empty()) {
|
||||
|
|
@ -2855,8 +2856,15 @@ void server_context::add_sampled_tokens() {
|
|||
|
||||
llama_tokens draft = common_speculative_draft(slot.spec, params_spec, cached_text_tokens, slot.sampled);
|
||||
|
||||
const int n_draft_max = slot.get_n_draft_max();
|
||||
|
||||
if (draft.size() > (size_t)n_draft_max) {
|
||||
SLT_WRN(slot, "draft size %d exceeds max %d, truncating\n", (int)draft.size(), n_draft_max);
|
||||
if (slot.params.speculative.autotune) {
|
||||
// expected near end-of-response when autotune shrinks n_max
|
||||
SLT_DBG(slot, "draft size %d exceeds max %d, truncating\n", (int)draft.size(), n_draft_max);
|
||||
} else {
|
||||
SLT_WRN(slot, "draft size %d exceeds max %d, truncating\n", (int)draft.size(), n_draft_max);
|
||||
}
|
||||
draft.resize(n_draft_max);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue