LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Allow dual speculative decoding (#1789)

Browse Source 
* wip: test logic to use multiple specs

* feat: introduce composite speculative decoding stages

* handle MTP context and draft invalidation

* fix: allow gemma mtp for speculative stages

* fix: normalize spec stage keys

* refactor: remove enable_mtp flag and improve speculative stage handling

* fix: update cached text tokens handling for stage chains

* feat: implement sync for external MTP after non-MTP accept
This commit is contained in:
Samuel Oliveira Alves
2026-05-15 10:10:40 +03:00
committed by GitHub
parent 53cd4d0ff0
commit f4f4b3ff26
9 changed files with 1044 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/server/server-context.cpp
+344 -78
View File
@@ -104,6 +104,32 @@ static void cache_and_sync_slot_mtp_hidden_from_rows(server_slot & slot, llama_c
cache_and_sync_slot_mtp_hidden(slot, ctx, rows.data() + (n_rows - 1) * n_embd, n_embd);
}
static const float * mtp_hidden_last_row(const std::vector<float> & rows, int n_embd) {
if (n_embd <= 0 || rows.size() < (size_t) n_embd) {
return nullptr;
}
const size_t n_rows = rows.size() / n_embd;
if (n_rows == 0) {
return nullptr;
}
return rows.data() + (n_rows - 1) * n_embd;
}
static bool sync_external_mtp_after_non_mtp_accept(
server_slot & slot,
llama_context * ctx,
const std::vector<float> & mtp_commit_states,
int n_embd) {
if (!slot.use_gemma4_external_mtp || mtp_commit_states.empty() || n_embd <= 0) {
return false;
}
cache_and_sync_slot_mtp_hidden_from_rows(slot, ctx, mtp_commit_states, n_embd);
return true;
}
static void apply_slot_mtp_accept(
server_slot & slot,
llama_context * ctx,
@@ -183,6 +209,12 @@ static int32_t server_mtp_media_warmup_callback(void * user_data, const llama_ba
return server_mtp_warmup_batch(data->ctx_tgt, get_slot_mtp_ctx(*data->slot, data->ctx_tgt), batch, *data->slot);
}
static bool server_response_needs_chat_parse(oaicompat_type oaicompat) {
return oaicompat == OAICOMPAT_TYPE_CHAT ||
oaicompat == OAICOMPAT_TYPE_ANTHROPIC ||
oaicompat == OAICOMPAT_TYPE_RESP;
}
void server_speculative_checkpoint::clear() {
valid = false;
per_step_enabled = false;
@@ -229,6 +261,104 @@ static bool save_speculative_checkpoint(server_slot & slot, llama_model * model,
return true;
}
static void server_remove_speculative_stage(common_params_speculative & spec, common_speculative_type type) {
spec.stages.erase(std::remove_if(spec.stages.begin(), spec.stages.end(), [type](const common_speculative_stage_params & stage) {
return stage.type == type;
}), spec.stages.end());
if (spec.type == type) {
spec.type = COMMON_SPECULATIVE_TYPE_NONE;
const auto resolved = spec.get_resolved_stages();
spec.type = resolved.empty() ? COMMON_SPECULATIVE_TYPE_NONE : resolved.front().type;
}
}
static bool server_speculative_has_mtp(const common_params_speculative & spec) {
return spec.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
}
static bool server_speculative_same_stage_types(
const common_params_speculative & lhs,
const common_params_speculative & rhs) {
const auto lhs_stages = lhs.get_resolved_stages();
const auto rhs_stages = rhs.get_resolved_stages();
if (lhs_stages.size() != rhs_stages.size()) {
return false;
}
for (size_t i = 0; i < lhs_stages.size(); ++i) {
if (lhs_stages[i].type != rhs_stages[i].type) {
return false;
}
}
return true;
}
static void server_reject_dead_speculative_request_overrides(const json & data) {
if (json_value_ptr(data, "speculative.type") != nullptr) {
throw std::runtime_error("Error: speculative.type request override is not supported; keep the startup stage types and use speculative.stages or speculative.n_max/n_min/p_min");
}
if (json_value_ptr(data, "speculative.ngram_size_n") != nullptr ||
json_value_ptr(data, "speculative.ngram_size_m") != nullptr ||
json_value_ptr(data, "speculative.ngram_min_hits") != nullptr ||
json_value_ptr(data, "speculative.suffix_min_match_len") != nullptr ||
json_value_ptr(data, "speculative.suffix_max_depth") != nullptr) {
throw std::runtime_error("Error: structural speculative overrides are startup-only; per-request overrides only support speculative.n_max, speculative.n_min, speculative.p_min, and speculative.stages");
}
}
static common_speculative_stage_params server_parse_speculative_stage_json(const json & stage_json) {
if (!stage_json.is_object()) {
throw std::runtime_error("Error: speculative.stages entries must be objects");
}
if (!stage_json.contains("type") || !stage_json["type"].is_string()) {
throw std::runtime_error("Error: speculative.stages entries must include a string 'type'");
}
common_speculative_stage_params stage;
stage.type = common_speculative_type_from_name(stage_json["type"].get<std::string>());
if (stage.type == COMMON_SPECULATIVE_TYPE_COUNT) {
throw std::runtime_error("Error: unknown speculative stage type in speculative.stages");
}
for (const auto & item : stage_json.items()) {
if (item.key() == "type") {
continue;
}
if (item.key() == "n_max") {
stage.n_max = item.value().get<int32_t>();
if (stage.n_max < 0) {
throw std::runtime_error("Error: speculative.stages[].n_max must be >= 0");
}
continue;
}
if (item.key() == "n_min") {
stage.n_min = item.value().get<int32_t>();
if (stage.n_min < 0) {
throw std::runtime_error("Error: speculative.stages[].n_min must be >= 0");
}
continue;
}
if (item.key() == "p_min") {
stage.p_min = item.value().get<float>();
if (stage.p_min < 0.0f) {
throw std::runtime_error("Error: speculative.stages[].p_min must be >= 0");
}
continue;
}
throw std::runtime_error("Error: per-request speculative.stages only support type, n_max, n_min, and p_min; structural stage overrides are startup-only");
}
return stage;
}
server_context::~server_context() {
if (ctx) {
llama_free(ctx);
@@ -328,9 +458,13 @@ bool server_context::load_model(const gpt_params& params_) {
LOG_ERROR("%s\n", "err: speculative decode is not supported by multimodal");
return false;
}
if (params_base.speculative.type != COMMON_SPECULATIVE_TYPE_NONE &&
params_base.speculative.type != COMMON_SPECULATIVE_TYPE_MTP) {
const auto spec_stages = params_base.speculative.get_resolved_stages();
const bool multimodal_spec_supported = spec_stages.empty() ||
(spec_stages.size() == 1 && spec_stages.front().type == COMMON_SPECULATIVE_TYPE_MTP);
if (!multimodal_spec_supported) {
params_base.speculative.type = COMMON_SPECULATIVE_TYPE_NONE;
params_base.speculative.stages.clear();
params_base.has_mtp = false;
SRV_WRN("%s\n", "speculative decoding is not supported by multimodal, it will be disabled");
}
}
@@ -380,9 +514,12 @@ bool server_context::load_model(const gpt_params& params_) {
params_base.speculative.cparams_dft = cparams_dft;
}
else if (params_base.has_mtp && llama_model_n_nextn_layer(model) == 0) {
LOG_WARNING("WARNING: -mtp flag provided, but model has 0 NextN layers. MTP will be disabled.\n", {});
if (server_speculative_has_mtp(params_base.speculative) &&
llama_model_n_nextn_layer(model) == 0 &&
!params_use_gemma4_external_mtp(params_base)) {
LOG_WARNING("WARNING: MTP speculative stage requested, but model has 0 NextN layers. MTP will be disabled.\n", {});
params_base.has_mtp = false;
server_remove_speculative_stage(params_base.speculative, COMMON_SPECULATIVE_TYPE_MTP);
}
return true;
}
@@ -433,13 +570,14 @@ void server_context::init() {
slot.ga_n = ga_n;
slot.ga_w = ga_w;
slot.params.speculative = params_base.speculative;
slot.sparams = params_base.sparams;
if (params_base.has_mtp) {
const bool wants_mtp_stage = server_speculative_has_mtp(params_base.speculative);
if (wants_mtp_stage) {
const bool has_external_mtp = params_use_gemma4_external_mtp(params_base);
if (llama_model_n_nextn_layer(model) > 0 || has_external_mtp) {
params_base.speculative.type = COMMON_SPECULATIVE_TYPE_MTP;
params_base.pooling_type = LLAMA_POOLING_TYPE_NONE;
if (!has_external_mtp) {
@@ -452,25 +590,24 @@ void server_context::init() {
slot.has_mtp = true;
slot.use_gemma4_external_mtp = has_external_mtp;
slot.params.speculative.type = COMMON_SPECULATIVE_TYPE_MTP;
slot.params.speculative.n_min = 0;
slot.params.speculative.cparams_dft = params_base.speculative.cparams_dft;
slot.batch_spec = llama_batch_init(slot.params.speculative.n_max + 1, 0, 1);
slot.batch_spec = llama_batch_init(slot.params.speculative.get_max_stage_n_max() + 1, 0, 1);
SLT_DBG(slot, "batch_spec contains %d tokens\n", slot.batch_spec.n_tokens);
SRV_INF("%s\n", "MTP needs embeddings on decode, enabling");
llama_set_embeddings(ctx, true);
}
else {
SRV_WRN("%s\n", "MTP enabled via flag, but model has 0 NextN layers. Disabling speculative.");
params_base.speculative.type = COMMON_SPECULATIVE_TYPE_NONE;
SRV_WRN("%s\n", "MTP speculative stage requested, but model has 0 NextN layers. Removing MTP from the configured stage chain.");
params_base.has_mtp = false;
server_remove_speculative_stage(params_base.speculative, COMMON_SPECULATIVE_TYPE_MTP);
slot.params.speculative = params_base.speculative;
slot.has_mtp = false;
}
}
const bool requested_spec = params_base.speculative.type != COMMON_SPECULATIVE_TYPE_NONE ||
params_base.speculative.has_dft();
const bool requested_spec = !params_base.speculative.get_resolved_stages().empty();
bool can_spec = true;
if (!params_base.dry_run) {
@@ -633,6 +770,7 @@ void server_slot::reset() {
checkpoint_pos = 0;
image_just_processed = false;
do_checkpoint = false;
mtp_hidden_state.clear();
positional_bans.clear();
ban_phrases.clear();
@@ -711,7 +849,7 @@ int server_slot::get_n_draft_max() const {
}
// determine the max draft that fits the current slot state
int n_draft_max = params.speculative.n_max;
int n_draft_max = params.speculative.get_max_stage_n_max();
// note: slot.prompt is not yet expanded with the `id` token sampled above
// also, need to leave space for 1 extra token to allow context shifts
@@ -723,8 +861,9 @@ int server_slot::get_n_draft_max() const {
SLT_DBG(*this, "max possible draft: %d\n", n_draft_max);
if (n_draft_max < params.speculative.n_min) {
SLT_DBG(*this, "the max possible draft is too small: %d < %d - skipping speculative decoding\n", n_draft_max, params.speculative.n_min);
const int min_usable_draft = params.speculative.get_min_usable_stage_n_min();
if (n_draft_max < min_usable_draft) {
SLT_DBG(*this, "the max possible draft is too small: %d < %d - skipping speculative decoding\n", n_draft_max, min_usable_draft);
n_draft_max = 0;
}
return n_draft_max;
@@ -1239,40 +1378,83 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task)
slot.params.post_sampling_probs = json_value(data, "post_sampling_probs", defaults.post_sampling_probs);
// speculative decoding parameters
slot.params.speculative.n_max = json_value(data, "speculative.n_max", params_base.speculative.n_max);
slot.params.speculative.n_min = json_value(data, "speculative.n_min", params_base.speculative.n_min);
slot.params.speculative.p_min = json_value(data, "speculative.p_min", params_base.speculative.p_min);
try {
slot.params.speculative = defaults.speculative;
slot.params.speculative.n_max = json_value(data, "speculative.n_max", params_base.speculative.n_max);
slot.params.speculative.n_min = json_value(data, "speculative.n_min", params_base.speculative.n_min);
slot.params.speculative.p_min = json_value(data, "speculative.p_min", params_base.speculative.p_min);
slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min);
slot.params.speculative.n_min = std::max(slot.params.speculative.n_min, 0);
slot.params.speculative.n_max = std::max(slot.params.speculative.n_max, 0);
server_reject_dead_speculative_request_overrides(data);
slot.params.speculative.type = common_speculative_type_from_name(json_value(data, "speculative.type", common_speculative_type_to_str(defaults.speculative.type)));
const json stages = json_value(data, "speculative.stages", json());
if (!stages.is_null()) {
if (!stages.is_array()) {
throw std::runtime_error("Error: speculative.stages must be an array");
}
// Clamp speculative parameters
slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min);
slot.params.speculative.n_min = std::max(slot.params.speculative.n_min, 0);
slot.params.speculative.n_max = std::max(slot.params.speculative.n_max, 0);
const auto default_stages = defaults.speculative.get_resolved_stages();
if (stages.size() != default_stages.size()) {
throw std::runtime_error("Error: speculative.stages must provide the same number of stages configured at server startup");
}
if (slot.can_speculate() &&
llama_model_has_recurrent(model) &&
slot.params.speculative.n_max > params_base.speculative.n_max) {
send_error(task,
"Error: speculative.n_max=" + std::to_string(slot.params.speculative.n_max) +
" exceeds the recurrent speculative startup limit of " + std::to_string(params_base.speculative.n_max) +
"; restart the server with a higher --draft-max to reserve checkpoint capacity",
ERROR_TYPE_INVALID_REQUEST);
slot.params.speculative.stages = default_stages;
for (size_t i = 0; i < stages.size(); ++i) {
const auto stage_override = server_parse_speculative_stage_json(stages[i]);
if (stage_override.type != default_stages[i].type) {
throw std::runtime_error("Error: speculative.stages must preserve the stage types configured at server startup");
}
if (stage_override.has_n_max_override()) {
slot.params.speculative.stages[i].n_max = stage_override.n_max;
}
if (stage_override.has_n_min_override()) {
slot.params.speculative.stages[i].n_min = stage_override.n_min;
}
if (stage_override.has_p_min_override()) {
slot.params.speculative.stages[i].p_min = stage_override.p_min;
}
}
const auto resolved = slot.params.speculative.get_resolved_stages();
slot.params.speculative.type = resolved.empty() ? COMMON_SPECULATIVE_TYPE_NONE : resolved.front().type;
}
slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min);
slot.params.speculative.n_min = std::max(slot.params.speculative.n_min, 0);
slot.params.speculative.n_max = std::max(slot.params.speculative.n_max, 0);
if (slot.can_speculate() &&
llama_model_has_recurrent(model) &&
slot.params.speculative.n_max > params_base.speculative.n_max) {
send_error(task,
"Error: speculative.n_max=" + std::to_string(slot.params.speculative.n_max) +
" exceeds the recurrent speculative startup limit of " + std::to_string(params_base.speculative.n_max) +
"; restart the server with a higher --draft-max to reserve checkpoint capacity",
ERROR_TYPE_INVALID_REQUEST);
return false;
}
if (!server_speculative_same_stage_types(slot.params.speculative, defaults.speculative)) {
throw std::runtime_error("Error: per-request speculative stages must match the server startup stage types; only stage parameter overrides are supported");
}
if (slot.params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && !slot.has_mtp) {
throw std::runtime_error("Error: MTP speculative stage requested, but the server was not started with MTP support");
}
if (slot.params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_DRAFT) && !params_base.speculative.has_dft()) {
throw std::runtime_error("Error: draft speculative stage requested, but no draft model is loaded");
}
std::string spec_error;
if (!common_speculative_validate_chain(slot.params.speculative, &spec_error)) {
throw std::runtime_error("Error: invalid speculative request configuration: " + spec_error);
}
} catch (const std::exception & e) {
send_error(task, e.what(), ERROR_TYPE_INVALID_REQUEST);
return false;
}
slot.params.speculative.ngram_size_n = json_value(data, "speculative.ngram_size_n", defaults.speculative.ngram_size_n);
slot.params.speculative.ngram_size_m = json_value(data, "speculative.ngram_size_m", defaults.speculative.ngram_size_m);
slot.params.speculative.ngram_min_hits = json_value(data, "speculative.ngram_m_hits", defaults.speculative.ngram_min_hits);
slot.params.speculative.ngram_size_n = std::max(std::min(1, (int)slot.params.speculative.ngram_size_n), 1024);
slot.params.speculative.ngram_size_m = std::max(std::min(1, (int)slot.params.speculative.ngram_size_m), 1024);
slot.params.speculative.ngram_min_hits = std::max(std::min(1, (int)slot.params.speculative.ngram_min_hits), 1024);
if (slot.sparams.penalty_last_n < -1) {
throw std::runtime_error("Error: repeat_last_n must be >= -1");
@@ -2378,31 +2560,33 @@ void server_context::send_partial_response(server_slot& slot, completion_token_o
{"id_slot", slot.id},
{"multimodal", false}
};
slot.update_chat_msg(true, res->oaicompat_msg_diffs);
if (server_response_needs_chat_parse(slot.params.oaicompat)) {
slot.update_chat_msg(true, res->oaicompat_msg_diffs);
res->anthropic_has_reasoning = !slot.chat_msg.reasoning_content.empty();
res->anthropic_has_reasoning = !slot.chat_msg.reasoning_content.empty();
res->anthropic_thinking_block_started = slot.anthropic_thinking_block_started;
res->anthropic_text_block_started = slot.anthropic_text_block_started;
res->anthropic_thinking_block_started = slot.anthropic_thinking_block_started;
res->anthropic_text_block_started = slot.anthropic_text_block_started;
res->oai_resp_thinking_block_started = slot.oai_resp_thinking_block_started;
res->oai_resp_text_block_started = slot.oai_resp_text_block_started;
res->oai_resp_thinking_block_started = slot.oai_resp_thinking_block_started;
res->oai_resp_text_block_started = slot.oai_resp_text_block_started;
for (const auto& diff : res->oaicompat_msg_diffs) {
if (!diff.reasoning_content_delta.empty() && !slot.anthropic_thinking_block_started) {
slot.anthropic_thinking_block_started = true;
}
if (!diff.content_delta.empty() && !slot.anthropic_text_block_started) {
slot.anthropic_text_block_started = true;
}
if (!diff.reasoning_content_delta.empty() && !slot.oai_resp_thinking_block_started) {
slot.oai_resp_thinking_block_started = true;
}
if (!diff.content_delta.empty() && !slot.oai_resp_text_block_started) {
slot.oai_resp_text_block_started = true;
}
if (!diff.tool_call_delta.name.empty()) {
slot.oai_resp_fc_id = diff.tool_call_delta.id;
for (const auto& diff : res->oaicompat_msg_diffs) {
if (!diff.reasoning_content_delta.empty() && !slot.anthropic_thinking_block_started) {
slot.anthropic_thinking_block_started = true;
}
if (!diff.content_delta.empty() && !slot.anthropic_text_block_started) {
slot.anthropic_text_block_started = true;
}
if (!diff.reasoning_content_delta.empty() && !slot.oai_resp_thinking_block_started) {
slot.oai_resp_thinking_block_started = true;
}
if (!diff.content_delta.empty() && !slot.oai_resp_text_block_started) {
slot.oai_resp_text_block_started = true;
}
if (!diff.tool_call_delta.name.empty()) {
slot.oai_resp_fc_id = diff.tool_call_delta.id;
}
}
}
@@ -2439,7 +2623,9 @@ void server_context::send_final_response(server_slot& slot) {
res->post_sampling_probs = slot.params.post_sampling_probs;
res->oaicompat = slot.params.oaicompat;
res->oaicompat_cmpl_id = slot.params.oaicompat_cmpl_id;
res->oaicompat_msg = slot.update_chat_msg(false, res->oaicompat_msg_diffs);
if (server_response_needs_chat_parse(slot.params.oaicompat)) {
res->oaicompat_msg = slot.update_chat_msg(false, res->oaicompat_msg_diffs);
}
res->oai_resp_id = slot.oai_resp_id;
res->oai_resp_reasoning_id = slot.oai_resp_reasoning_id;
res->oai_resp_message_id = slot.oai_resp_message_id;
@@ -3452,7 +3638,7 @@ void server_context::add_sampled_tokens() {
}
static const llama_tokens empty_prompt;
const llama_tokens & cached_text_tokens = slot.has_mtp
const llama_tokens & cached_text_tokens = slot.has_mtp && !slot.params.speculative.has_composite_stage_chain()
? empty_prompt
: slot.cache_tokens.get_text_tokens();
@@ -3491,14 +3677,14 @@ void server_context::add_sampled_tokens() {
common_batch_add(batch, slot.sampled, slot.cache_tokens.pos_next(), { slot.id }, true);
slot.cache_tokens.push_back(slot.sampled);
if (slot.params.speculative.n_min > (int)draft.size()) {
SLT_DBG(slot, "ignoring small draft: %d < %d\n", (int)draft.size(), slot.params.speculative.n_min);
const int min_usable_draft = slot.params.speculative.get_min_usable_stage_n_min();
if (min_usable_draft > (int)draft.size()) {
SLT_DBG(slot, "ignoring small draft: %d < %d\n", (int)draft.size(), min_usable_draft);
// fallback to normal decoding
slot.i_batch = slot.i_batch_dft[0];
slot.drafted.clear();
slot.i_batch_dft.clear();
}
else {
} else {
// keep track of total number of drafted tokens tested
slot.n_draft_total += draft.size();
@@ -3586,6 +3772,7 @@ void server_context::apply_checkpoint(server_slot & slot) {
slot.n_past_prompt = 0;
slot.n_past_se = 0;
slot.ga_i = 0;
slot.cache_tokens.keep_first(0);
pos_next = 0;
common_sampler_reset(slot.ctx_sampling);
}
@@ -4028,7 +4215,11 @@ void server_context::extend_context(const int32_t n_tokens) {
// Restore recurrent state and re-decode accepted tokens after speculative-decode rejection.
static void restore_speculative_checkpoint(
server_slot & slot, llama_context * ctx, llama_model * model,
common_speculative_type spec_type_used,
const std::vector<llama_token> & ids, int n_draft,
const std::vector<llama_token> & mtp_commit_tokens,
const std::vector<float> & mtp_commit_states,
const std::vector<float> & mtp_hidden_state_seed,
const std::vector<float> & mtp_hidden_state_pre, int32_t mtp_n_past_base) {
if (slot.spec_ckpt.per_step_enabled) {
const int step = (int)ids.size() - 1;
@@ -4043,8 +4234,36 @@ static void restore_speculative_checkpoint(
// Update MTP KV cache and hidden state using embeddings collected before checkpoint restore.
if (slot.has_mtp && !mtp_hidden_state_pre.empty()) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
apply_slot_mtp_accept(slot, ctx, mtp_hidden_state_pre, ids, mtp_n_past_base, n_embd);
llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec);
llama_context * mtp_target = mtp_ctx ? mtp_ctx : ctx;
if (spec_type_used == COMMON_SPECULATIVE_TYPE_MTP) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
apply_slot_mtp_accept(slot, ctx, mtp_hidden_state_pre, ids, mtp_n_past_base, n_embd);
} else if (!mtp_commit_tokens.empty() && !mtp_commit_states.empty()) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
if (sync_external_mtp_after_non_mtp_accept(slot, ctx, mtp_commit_states, n_embd)) {
SLT_DBG(slot, "%s", "synced external MTP hidden state from accepted-prefix rows after per-step restore");
} else {
const float * seed_hidden = mtp_hidden_last_row(mtp_hidden_state_seed, n_embd);
if (seed_hidden == nullptr) {
SLT_WRN(slot, "%s", "missing MTP seed hidden state for accepted-prefix replay after per-step restore");
slot.mtp_hidden_state.clear();
} else {
llama_batch accepted_batch = llama_batch_init(mtp_commit_tokens.size(), 0, 1);
for (size_t i = 0; i < mtp_commit_tokens.size(); ++i) {
common_batch_add(accepted_batch, mtp_commit_tokens[i], mtp_n_past_base + i, { slot.id }, true);
}
llama_set_draft_input_hidden_state(mtp_target, seed_hidden);
mtp_update_kv_cache(mtp_target, accepted_batch, false);
llama_batch_free(accepted_batch);
slot.mtp_hidden_state.assign(mtp_commit_states.end() - n_embd, mtp_commit_states.end());
}
}
}
}
SLT_DBG(slot, "per-step restore: step=%d (rejected %d drafts)\n",
@@ -4124,7 +4343,10 @@ void server_context::speculative_decoding_accept() {
continue;
}
const llama_token sampled_before = slot.sampled;
const common_speculative_type spec_type_used = common_speculative_current_type(slot.spec);
size_t n_draft = slot.drafted.size();
const std::vector<float> mtp_hidden_state_seed = slot.has_mtp ? slot.mtp_hidden_state : std::vector<float>{};
slot.ctx_sampling->to_generated_text = &slot.generated_text;
if (n_draft > 0) {
@@ -4155,6 +4377,8 @@ void server_context::speculative_decoding_accept() {
int32_t mtp_n_past_base = 0;
std::vector<float> mtp_hidden_state_pre;
std::vector<llama_token> mtp_commit_tokens;
std::vector<float> mtp_commit_states;
if (slot.has_mtp) {
const int32_t n_pre_spec_tokens = slot.cache_tokens.n_tokens() - (int32_t)(slot.drafted.size() + 1);
mtp_n_past_base = slot.cache_tokens.pos_next(n_pre_spec_tokens);
@@ -4168,6 +4392,20 @@ void server_context::speculative_decoding_accept() {
memcpy(mtp_hidden_state_pre.data() + i * n_embd, emb_i, n_embd * sizeof(float));
}
}
if (spec_type_used != COMMON_SPECULATIVE_TYPE_MTP) {
mtp_commit_tokens.reserve(ids.size());
mtp_commit_tokens.push_back(sampled_before);
mtp_commit_tokens.insert(mtp_commit_tokens.end(), ids.begin(), ids.end() - 1);
mtp_commit_states.resize(ids.size() * n_embd);
for (size_t i = 0; i < ids.size(); ++i) {
const float * emb_i = llama_get_embeddings_ith(ctx, slot.i_batch_dft[i]);
if (emb_i) {
memcpy(mtp_commit_states.data() + i * n_embd, emb_i, n_embd * sizeof(float));
}
}
}
} else {
const float* emb0 = llama_get_embeddings_ith(ctx, 0);
if (emb0) {
@@ -4204,11 +4442,39 @@ void server_context::speculative_decoding_accept() {
// for recurrent/hybrid models: if any drafts were rejected, restore recurrent state
const bool any_rejected = (ids.size() - 1) < n_draft;
if (any_rejected && slot.spec_ckpt.valid) {
restore_speculative_checkpoint(slot, ctx, model, ids, n_draft, mtp_hidden_state_pre, mtp_n_past_base);
restore_speculative_checkpoint(slot, ctx, model, spec_type_used, ids, n_draft, mtp_commit_tokens, mtp_commit_states, mtp_hidden_state_seed, mtp_hidden_state_pre, mtp_n_past_base);
} else {
if (slot.has_mtp && !mtp_hidden_state_pre.empty()) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
apply_slot_mtp_accept(slot, ctx, mtp_hidden_state_pre, ids, mtp_n_past_base, n_embd);
llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec);
llama_context * mtp_target = mtp_ctx ? mtp_ctx : ctx;
if (spec_type_used == COMMON_SPECULATIVE_TYPE_MTP) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
apply_slot_mtp_accept(slot, ctx, mtp_hidden_state_pre, ids, mtp_n_past_base, n_embd);
} else if (!mtp_commit_tokens.empty() && !mtp_commit_states.empty()) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
if (sync_external_mtp_after_non_mtp_accept(slot, ctx, mtp_commit_states, n_embd)) {
SLT_DBG(slot, "%s", "synced external MTP hidden state from accepted-prefix rows");
} else {
const float * seed_hidden = mtp_hidden_last_row(mtp_hidden_state_seed, n_embd);
if (seed_hidden == nullptr) {
SLT_WRN(slot, "%s", "missing MTP seed hidden state for accepted-prefix replay");
slot.mtp_hidden_state.clear();
} else {
llama_batch accepted_batch = llama_batch_init(mtp_commit_tokens.size(), 0, 1);
for (size_t i = 0; i < mtp_commit_tokens.size(); ++i) {
common_batch_add(accepted_batch, mtp_commit_tokens[i], mtp_n_past_base + i, { slot.id }, true);
}
llama_set_draft_input_hidden_state(mtp_target, seed_hidden);
mtp_update_kv_cache(mtp_target, accepted_batch, false);
llama_batch_free(accepted_batch);
slot.mtp_hidden_state.assign(mtp_commit_states.end() - n_embd, mtp_commit_states.end());
}
}
}
}
llama_kv_cache_seq_rm(ctx, slot.id, slot.cache_tokens.pos_next(slot.n_past), -1);
discard_speculative_checkpoint(slot, ctx);
@@ -4586,8 +4852,8 @@ void server_context::process_batch_tokens(int32_t & n_batch) {
continue; // continue loop of n_batch
}
server_slot * mtp_warmup_slot = nullptr;
if (params_base.has_mtp) {
server_slot * mtp_warmup_slot = nullptr;
if (server_speculative_has_mtp(params_base.speculative)) {
for (auto& slot : slots) {
if ((slot.state == SLOT_STATE_PROCESSING && slot.n_decoded == 0) ||
(slot.state == SLOT_STATE_IDLE && slot.command == SLOT_COMMAND_LOAD_PROMPT)) {
@@ -4632,7 +4898,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) {
if (slot.n_decoded == 0 && slot.can_speculate()) {
static const llama_tokens empty_prompt;
const llama_tokens & spec_prompt = slot.has_mtp
const llama_tokens & spec_prompt = slot.has_mtp && !slot.params.speculative.has_composite_stage_chain()
? empty_prompt
: slot.cache_tokens.get_text_tokens();
common_speculative_begin(slot.spec, spec_prompt);