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

Move embedding management to speculative (#1825)

Browse Source 
* refactor speculative decoding with companion context and draft result structures

* feat: add common speculative feature handling in server context

* refactor: move embedings outside server

* feat: harden draft input hidden state in llama context

* remove unused functions

* refactor: streamline speculative feature handling and remove unused code

* remove redundant code

* remove more unused variables

* refactor: implement speculative feature handling
This commit is contained in:
Samuel Oliveira Alves
2026-05-20 17:42:48 +03:00
committed by GitHub
parent dd67a9fb24
commit 11a1fea9e2
10 changed files with 1043 additions and 416 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative.cpp
+621 -75
View File
@@ -20,7 +20,6 @@
#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5
void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx);
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx);
const std::vector<enum common_speculative_type> common_speculative_types = {
COMMON_SPECULATIVE_TYPE_NONE,
@@ -168,7 +167,30 @@ struct common_speculative_state {
virtual void accept(uint16_t n_accepted) = 0;
};
static void mtp_invalidate_cached_draft(const llama_context * ctx);
struct common_speculative_state_mtp;
static common_speculative_state_mtp * common_speculative_get_mtp_state(common_speculative * spec);
static const common_speculative_state_mtp * common_speculative_get_mtp_state(const common_speculative * spec);
static void mtp_invalidate_cached_drafts(common_speculative_state_mtp & state);
static std::vector<llama_token> mtp_speculative_gen_draft(
common_speculative_state_mtp & state,
struct common_sampler * smpl,
struct llama_context * ctx,
int n_draft,
float p_min,
llama_token id_last,
llama_pos n_past,
llama_seq_id seq_id,
bool constant_draft_positions = false);
static int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch & batch, bool is_prompt_warmup);
struct mtp_last_embd {
std::vector<float> embd;
float prob = 0.0f;
int last_id = -1;
};
struct common_speculative_state_mtp : public common_speculative_state {
llama_context * ctx_tgt;
@@ -176,6 +198,9 @@ struct common_speculative_state_mtp : public common_speculative_state {
common_sampler * smpl;
// For Gemma 4 external MTP assistant: draft positions are held constant
bool constant_draft_positions = false;
int n_embd = 0;
std::unordered_map<llama_seq_id, std::vector<float>> target_hidden_by_seq;
std::unordered_map<llama_seq_id, mtp_last_embd> draft_cache_by_seq;
common_speculative_state_mtp(
enum common_speculative_type type,
@@ -193,6 +218,7 @@ struct common_speculative_state_mtp : public common_speculative_state {
};
smpl = common_sampler_init(llama_get_model(ctx_mtp), sparams);
llama_set_mtp_target_context(ctx_mtp, ctx_tgt);
n_embd = llama_mtp_state_n_embd(ctx_mtp);
LOG_INF("%s: MTP context ready (n_ctx=%d, constant_draft_positions=%s)\n", __func__,
llama_n_ctx(ctx_mtp), constant_draft_positions ? "true" : "false");
@@ -207,7 +233,8 @@ struct common_speculative_state_mtp : public common_speculative_state {
void begin(const llama_tokens & prompt) override {
GGML_UNUSED(prompt);
mtp_invalidate_cached_draft(ctx_mtp);
target_hidden_by_seq.clear();
draft_cache_by_seq.clear();
}
void draft(
@@ -245,7 +272,20 @@ struct common_speculative_state_mtp : public common_speculative_state {
llama_context * ctx = ctx_mtp;
const auto hidden_it = target_hidden_by_seq.find(seq_id);
if (hidden_it == target_hidden_by_seq.end() || (int) hidden_it->second.size() != n_embd) {
LOG_WRN("%s: missing target hidden state for seq_id %d\n", __func__, (int) seq_id);
result.clear();
return;
}
if (!llama_set_draft_input_hidden_state_copy(ctx, hidden_it->second.data(), hidden_it->second.size())) {
result.clear();
return;
}
result = mtp_speculative_gen_draft(
*this,
smpl,
ctx,
params.n_max,
@@ -1320,33 +1360,6 @@ void common_speculative_begin(common_speculative * spec, const llama_tokens & pr
}
}
struct mtp_last_embd {
std::vector<float> embd;
float prob;
int last_id = -1;
};
// Hopefully never called concurrently from multiple threads
static mtp_last_embd & mtp_get_last_embd(const llama_context * ctx) {
static std::unordered_map<const llama_context *, mtp_last_embd> map;
auto & last = map[ctx];
if (last.embd.empty()) {
auto n_embd = llama_mtp_state_n_embd(ctx);
last.embd.resize(n_embd);
}
return last;
}
static void mtp_invalidate_cached_draft(const llama_context * ctx) {
if (ctx == nullptr) {
return;
}
auto & last = mtp_get_last_embd(ctx);
last.last_id = -1;
last.prob = 0.0f;
}
llama_tokens common_speculative_draft(
common_speculative * spec,
common_params_speculative & params,
@@ -1437,12 +1450,390 @@ void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) {
}
if (impl->type != COMMON_SPECULATIVE_TYPE_MTP) {
if (auto * ctx_mtp = common_speculative_get_mtp_ctx(spec); ctx_mtp != nullptr) {
mtp_invalidate_cached_draft(ctx_mtp);
if (auto * mtp_state = common_speculative_get_mtp_state(spec); mtp_state != nullptr) {
mtp_invalidate_cached_drafts(*mtp_state);
}
}
}
static bool common_speculative_has_type(const common_speculative * spec, common_speculative_type type) {
if (spec == nullptr) {
return false;
}
return std::any_of(spec->configs.begin(), spec->configs.end(), [type](const common_speculative_config & config) {
return config.type == type;
});
}
static int common_speculative_ctx_mtp_n_embd(llama_context * ctx) {
return ctx ? (int) llama_mtp_state_n_embd(ctx) : 0;
}
static bool common_speculative_batch_token_has_seq_id(
const llama_batch & batch,
int token_index,
llama_seq_id seq_id) {
if (batch.n_seq_id == nullptr || batch.seq_id == nullptr || batch.n_seq_id[token_index] <= 0 || batch.seq_id[token_index] == nullptr) {
return false;
}
for (int i = 0; i < batch.n_seq_id[token_index]; ++i) {
if (batch.seq_id[token_index][i] == seq_id) {
return true;
}
}
return false;
}
static bool common_speculative_batch_is_exact_single_seq(
const llama_batch & batch,
llama_seq_id seq_id) {
if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
return false;
}
for (int i = 0; i < batch.n_tokens; ++i) {
if (batch.n_seq_id[i] != 1 || batch.seq_id[i] == nullptr || batch.seq_id[i][0] != seq_id) {
return false;
}
}
return true;
}
static int common_speculative_copy_seq_batch(
const llama_batch & batch,
llama_seq_id seq_id,
llama_batch & seq_batch) {
if (batch.token == nullptr || batch.pos == nullptr) {
return -1;
}
if (batch.n_tokens < 1) {
return 0;
}
std::vector<int> token_indices;
token_indices.reserve(batch.n_tokens);
for (int i = 0; i < batch.n_tokens; ++i) {
if (common_speculative_batch_token_has_seq_id(batch, i, seq_id)) {
token_indices.push_back(i);
}
}
if (token_indices.empty()) {
return 0;
}
seq_batch = llama_batch_init((int) token_indices.size(), 0, 1);
for (const int i : token_indices) {
common_batch_add(seq_batch, batch.token[i], batch.pos[i], { seq_id }, batch.logits != nullptr && batch.logits[i]);
}
return (int) token_indices.size();
}
static bool common_speculative_feature_view_copy_batch_rows(
const common_speculative_feature_view & view,
const llama_batch & batch,
llama_seq_id seq_id,
std::vector<float> * hidden_rows) {
if (hidden_rows == nullptr || view.kind != COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE || view.width <= 0 || batch.n_tokens <= 0 || batch.pos == nullptr) {
return false;
}
std::unordered_map<llama_pos, const float *> rows_by_pos;
rows_by_pos.reserve(view.rows.size());
for (const auto & row : view.rows) {
if (row.seq_id == seq_id && row.data != nullptr) {
rows_by_pos[row.pos] = row.data;
}
}
hidden_rows->clear();
hidden_rows->reserve((size_t) batch.n_tokens * view.width);
for (int i = 0; i < batch.n_tokens; ++i) {
auto it = rows_by_pos.find(batch.pos[i]);
if (it == rows_by_pos.end()) {
hidden_rows->clear();
return false;
}
hidden_rows->insert(hidden_rows->end(), it->second, it->second + view.width);
}
return hidden_rows->size() == (size_t) batch.n_tokens * view.width;
}
static bool common_speculative_capture_target_features(
common_speculative * spec,
const common_speculative_feature_view & features);
static bool common_speculative_feature_view_from_hidden_rows(
const std::vector<float> & hidden_rows,
int32_t width,
llama_seq_id seq_id,
llama_pos pos_base,
common_speculative_feature_view & view) {
view = {};
view.kind = COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE;
view.width = width;
if (width <= 0 || hidden_rows.empty() || hidden_rows.size() % (size_t) width != 0) {
return false;
}
const size_t n_rows = hidden_rows.size() / (size_t) width;
view.rows.reserve(n_rows);
for (size_t i = 0; i < n_rows; ++i) {
view.rows.push_back({
/* .seq_id = */ seq_id,
/* .pos = */ pos_base + (llama_pos) i,
/* .data = */ hidden_rows.data() + i * (size_t) width,
});
}
return true;
}
static bool common_speculative_collect_target_batch_features(
const common_speculative * spec,
llama_context * ctx,
const llama_batch & batch,
common_speculative_feature_view & features) {
features = {};
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP)) {
return true;
}
if (!llama_spec_get_hidden_feature_view(ctx, batch, features)) {
return false;
}
return true;
}
static bool common_speculative_collect_target_seq_batch_features(
const common_speculative * spec,
llama_context * ctx,
const llama_batch & batch,
llama_seq_id seq_id,
common_speculative_feature_view & features) {
features = {};
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP)) {
return true;
}
if (!llama_spec_get_hidden_feature_view_for_seq(ctx, batch, seq_id, features)) {
return false;
}
return true;
}
bool common_speculative_capture_output_hidden(
common_speculative * spec,
llama_context * ctx,
int32_t output_index,
llama_seq_id seq_id,
llama_pos pos) {
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP)) {
return true;
}
common_speculative_feature_view features;
if (!llama_spec_get_hidden_feature_view_from_output_index(ctx, output_index, seq_id, pos, features)) {
return false;
}
return common_speculative_capture_target_features(spec, features);
}
bool common_speculative_ensure_sequence_hidden(
common_speculative * spec,
llama_context * ctx,
llama_seq_id seq_id,
llama_pos pos) {
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || common_speculative_has_sequence_hidden(spec, seq_id)) {
return true;
}
return common_speculative_capture_output_hidden(spec, ctx, -1, seq_id, pos);
}
int32_t common_speculative_on_target_seq_batch(
common_speculative * spec,
llama_context * ctx_tgt,
const llama_batch & batch,
llama_seq_id seq_id,
bool is_prompt_warmup) {
llama_context * ctx_mtp = common_speculative_get_companion_ctx(spec);
ctx_mtp = ctx_mtp ? ctx_mtp : ctx_tgt;
if (ctx_tgt == nullptr || ctx_mtp == nullptr || batch.n_tokens <= 0) {
return 0;
}
const int n_embd_src = common_speculative_ctx_mtp_n_embd(ctx_tgt);
const int n_embd_dst = common_speculative_ctx_mtp_n_embd(ctx_mtp);
if (n_embd_src <= 0 || n_embd_dst <= 0) {
return -1;
}
if (n_embd_src != n_embd_dst) {
LOG_ERR("MTP warmup hidden state width mismatch: n_embd_src = %d, n_embd_dst = %d\n", n_embd_src, n_embd_dst);
return -1;
}
common_speculative_feature_view feature_view;
const llama_batch * batch_for_spec = &batch;
llama_batch seq_batch = {};
const bool needs_seq_split = is_prompt_warmup && !common_speculative_batch_is_exact_single_seq(batch, seq_id);
if (needs_seq_split) {
const int n_seq_tokens = common_speculative_copy_seq_batch(batch, seq_id, seq_batch);
if (n_seq_tokens <= 0) {
return n_seq_tokens < 0 ? -1 : 0;
}
if (!common_speculative_collect_target_seq_batch_features(spec, ctx_tgt, batch, seq_id, feature_view)) {
llama_batch_free(seq_batch);
return -1;
}
batch_for_spec = &seq_batch;
} else {
if (!common_speculative_collect_target_batch_features(spec, ctx_tgt, batch, feature_view)) {
return -1;
}
}
const int32_t ret = common_speculative_on_target_batch(spec, *batch_for_spec, feature_view, is_prompt_warmup);
if (needs_seq_split) {
llama_batch_free(seq_batch);
}
return ret;
}
bool common_speculative_copy_output_hidden_rows(
const common_speculative * spec,
llama_context * ctx,
const std::vector<int32_t> & output_indices,
std::vector<float> & hidden_rows) {
hidden_rows.clear();
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP)) {
return true;
}
return llama_spec_copy_hidden_rows_from_output_indices(ctx, output_indices, hidden_rows);
}
static bool common_speculative_build_commit_tokens(
common_speculative_type spec_type_used,
llama_token sampled_before,
const std::vector<llama_token> & ids,
std::vector<llama_token> & commit_tokens) {
commit_tokens.clear();
if (ids.empty()) {
return true;
}
if (spec_type_used == COMMON_SPECULATIVE_TYPE_MTP) {
commit_tokens = ids;
return true;
}
commit_tokens.reserve(ids.size());
commit_tokens.push_back(sampled_before);
if (ids.size() > 1) {
commit_tokens.insert(commit_tokens.end(), ids.begin(), ids.end() - 1);
}
return commit_tokens.size() == ids.size();
}
static bool common_speculative_apply_hidden_rows(
common_speculative * spec,
llama_seq_id seq_id,
llama_pos pos_base,
const std::vector<llama_token> & ids,
const std::vector<float> & hidden_rows) {
auto * mtp_state = common_speculative_get_mtp_state(spec);
if (mtp_state == nullptr || ids.empty()) {
return true;
}
const size_t expected_floats = ids.size() * (size_t) mtp_state->n_embd;
if (mtp_state->n_embd <= 0 || hidden_rows.size() != expected_floats) {
return false;
}
llama_batch accepted_batch = llama_batch_init(ids.size(), 0, 1);
for (size_t i = 0; i < ids.size(); ++i) {
common_batch_add(accepted_batch, ids[i], pos_base + (llama_pos) i, { seq_id }, true);
}
common_speculative_feature_view feature_view;
const bool have_feature_view = common_speculative_feature_view_from_hidden_rows(
hidden_rows, mtp_state->n_embd, seq_id, pos_base, feature_view);
const int32_t ret = have_feature_view
? common_speculative_on_target_batch(spec, accepted_batch, feature_view, false)
: -1;
llama_batch_free(accepted_batch);
return ret == 0;
}
bool common_speculative_commit_accepted_hidden_rows(
common_speculative * spec,
common_speculative_type spec_type_used,
llama_seq_id seq_id,
llama_pos pos_base,
llama_token sampled_before,
const std::vector<llama_token> & ids,
const std::vector<float> & hidden_rows) {
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || ids.empty()) {
return true;
}
std::vector<llama_token> commit_tokens;
if (!common_speculative_build_commit_tokens(spec_type_used, sampled_before, ids, commit_tokens)) {
return false;
}
return common_speculative_apply_hidden_rows(spec, seq_id, pos_base, commit_tokens, hidden_rows);
}
bool common_speculative_commit_accepted_output(
common_speculative * spec,
llama_context * ctx,
common_speculative_type spec_type_used,
llama_seq_id seq_id,
llama_pos pos_base,
llama_token sampled_before,
const std::vector<llama_token> & ids,
const std::vector<int32_t> & output_indices) {
if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || ids.empty()) {
return true;
}
std::vector<float> hidden_rows;
if (!common_speculative_copy_output_hidden_rows(spec, ctx, output_indices, hidden_rows)) {
return false;
}
return common_speculative_commit_accepted_hidden_rows(
spec,
spec_type_used,
seq_id,
pos_base,
sampled_before,
ids,
hidden_rows);
}
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;
@@ -1485,20 +1876,199 @@ void common_speculative_print_stats(const common_speculative * spec, double slot
// MTP
// ----------------------------------------------------------------------------
llama_context * common_speculative_get_mtp_ctx(common_speculative * spec) {
if (!spec) return nullptr;
static common_speculative_state_mtp * common_speculative_get_mtp_state(common_speculative * spec) {
if (!spec) {
return nullptr;
}
for (auto & impl : spec->impls) {
if (impl->type == COMMON_SPECULATIVE_TYPE_MTP) {
auto * mtp_state = dynamic_cast<common_speculative_state_mtp *>(impl.get());
if (mtp_state) {
return mtp_state->ctx_mtp;
}
if (impl->type != COMMON_SPECULATIVE_TYPE_MTP) {
continue;
}
if (auto * mtp_state = dynamic_cast<common_speculative_state_mtp *>(impl.get())) {
return mtp_state;
}
}
return nullptr;
}
static const common_speculative_state_mtp * common_speculative_get_mtp_state(const common_speculative * spec) {
return common_speculative_get_mtp_state(const_cast<common_speculative *>(spec));
}
static mtp_last_embd & mtp_get_last_embd(common_speculative_state_mtp & state, llama_seq_id seq_id) {
auto & last = state.draft_cache_by_seq[seq_id];
if ((int) last.embd.size() != state.n_embd) {
last.embd.resize(state.n_embd);
}
return last;
}
static void mtp_invalidate_cached_draft(common_speculative_state_mtp & state, llama_seq_id seq_id) {
auto it = state.draft_cache_by_seq.find(seq_id);
if (it == state.draft_cache_by_seq.end()) {
return;
}
it->second.last_id = -1;
it->second.prob = 0.0f;
}
static void mtp_invalidate_cached_drafts(common_speculative_state_mtp & state) {
for (auto & entry : state.draft_cache_by_seq) {
entry.second.last_id = -1;
entry.second.prob = 0.0f;
}
}
static void mtp_store_target_hidden(
common_speculative_state_mtp & state,
llama_seq_id seq_id,
const float * hidden,
int32_t width) {
if (hidden == nullptr || width <= 0) {
return;
}
auto & stored = state.target_hidden_by_seq[seq_id];
stored.assign(hidden, hidden + width);
}
static void mtp_clear_target_hidden(common_speculative_state_mtp & state, llama_seq_id seq_id) {
state.target_hidden_by_seq.erase(seq_id);
state.draft_cache_by_seq.erase(seq_id);
}
static bool common_speculative_capture_target_features(common_speculative * spec, const common_speculative_feature_view & features) {
auto * mtp_state = common_speculative_get_mtp_state(spec);
if (mtp_state == nullptr || features.kind != COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE || features.width <= 0) {
return false;
}
bool captured = false;
for (const auto & row : features.rows) {
if (row.data == nullptr) {
continue;
}
mtp_store_target_hidden(*mtp_state, row.seq_id, row.data, features.width);
mtp_invalidate_cached_draft(*mtp_state, row.seq_id);
captured = true;
}
return captured;
}
bool common_speculative_has_sequence_hidden(const common_speculative * spec, llama_seq_id seq_id) {
const auto * mtp_state = common_speculative_get_mtp_state(spec);
if (mtp_state == nullptr) {
return false;
}
auto it = mtp_state->target_hidden_by_seq.find(seq_id);
return it != mtp_state->target_hidden_by_seq.end() && !it->second.empty();
}
void common_speculative_clear_sequence_hidden(common_speculative * spec, llama_seq_id seq_id) {
auto * mtp_state = common_speculative_get_mtp_state(spec);
if (mtp_state == nullptr) {
return;
}
mtp_clear_target_hidden(*mtp_state, seq_id);
}
llama_context * common_speculative_get_companion_ctx(common_speculative * spec) {
if (auto * mtp_state = common_speculative_get_mtp_state(spec); mtp_state != nullptr) {
return mtp_state->ctx_mtp;
}
return nullptr;
}
static int32_t mtp_accept_batch(
common_speculative_state_mtp & state,
const llama_batch & accepted_batch,
llama_seq_id seq_id,
const float * hidden_rows) {
if (accepted_batch.n_tokens == 0 || hidden_rows == nullptr) {
return 0;
}
const size_t hidden_rows_floats = (size_t) accepted_batch.n_tokens * state.n_embd;
if (!llama_set_draft_input_hidden_state_copy(state.ctx_mtp, hidden_rows, hidden_rows_floats)) {
return -1;
}
if (mtp_update_kv_cache(state.ctx_mtp, accepted_batch, false) != 0) {
return -1;
}
auto & last = mtp_get_last_embd(state, seq_id);
const float * embd = llama_get_embeddings_ith(state.ctx_mtp, accepted_batch.n_tokens - 1);
if (embd != nullptr) {
std::memcpy(last.embd.data(), embd, last.embd.size() * sizeof(float));
if (!llama_set_draft_input_hidden_state_copy(state.ctx_mtp, last.embd.data(), last.embd.size())) {
return -1;
}
last.last_id = common_sampler_sample_speculative(nullptr, state.ctx_mtp, accepted_batch.n_tokens - 1, &last.prob);
}
return 0;
}
int32_t common_speculative_on_target_batch(
common_speculative * spec,
const llama_batch & batch,
const common_speculative_feature_view & features,
bool is_prompt_warmup) {
auto * mtp_state = common_speculative_get_mtp_state(spec);
if (mtp_state == nullptr) {
return 0;
}
if (features.kind != COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE || features.width <= 0 || batch.n_tokens <= 0) {
return 0;
}
if (batch.n_seq_id == nullptr || batch.seq_id == nullptr || batch.n_seq_id[0] <= 0 || batch.seq_id[0] == nullptr) {
return -1;
}
const llama_seq_id seq_id = batch.seq_id[0][0];
for (int i = 0; i < batch.n_tokens; ++i) {
if (batch.n_seq_id[i] != 1 || batch.seq_id[i] == nullptr || batch.seq_id[i][0] != seq_id) {
return -1;
}
}
std::vector<float> hidden_rows_storage;
if (!common_speculative_feature_view_copy_batch_rows(features, batch, seq_id, &hidden_rows_storage)) {
return -1;
}
const float * first_hidden = hidden_rows_storage.data();
const float * last_hidden = hidden_rows_storage.data() + (size_t) (batch.n_tokens - 1) * features.width;
mtp_store_target_hidden(*mtp_state, seq_id, last_hidden, features.width);
if (mtp_state->constant_draft_positions) {
mtp_invalidate_cached_draft(*mtp_state, seq_id);
return 0;
}
if (is_prompt_warmup) {
if (!llama_set_draft_input_hidden_state_copy(mtp_state->ctx_mtp, hidden_rows_storage.data(), hidden_rows_storage.size())) {
return -1;
}
const int32_t ret = mtp_update_kv_cache(mtp_state->ctx_mtp, batch, true);
mtp_invalidate_cached_draft(*mtp_state, seq_id);
return ret;
}
return mtp_accept_batch(*mtp_state, batch, seq_id, first_hidden);
}
common_speculative_type common_speculative_current_type(const common_speculative * spec) {
if (spec == nullptr || spec->curr_impl == nullptr) {
return COMMON_SPECULATIVE_TYPE_NONE;
@@ -1513,13 +2083,14 @@ void common_speculative_context_shift(
llama_pos kv_keep,
llama_pos kv_discard,
llama_pos kv_past) {
if (auto * ctx_mtp = common_speculative_get_mtp_ctx(spec); ctx_mtp != nullptr) {
if (auto * ctx_mtp = common_speculative_get_companion_ctx(spec); ctx_mtp != nullptr) {
llama_kv_cache_seq_rm (ctx_mtp, seq_id, kv_keep, kv_keep + kv_discard);
llama_kv_cache_seq_add(ctx_mtp, seq_id, kv_keep + kv_discard, kv_past, -kv_discard);
}
}
std::vector<llama_token> mtp_speculative_gen_draft(
common_speculative_state_mtp & state,
struct common_sampler * smpl,
struct llama_context * ctx,
int n_draft,
@@ -1535,7 +2106,7 @@ std::vector<llama_token> mtp_speculative_gen_draft(
if (!smpl) return drafts;
if (n_draft <= 0) {
mtp_invalidate_cached_draft(ctx);
mtp_invalidate_cached_draft(state, seq_id);
return drafts;
}
@@ -1551,7 +2122,7 @@ std::vector<llama_token> mtp_speculative_gen_draft(
llama_pos current_n_past = n_past;
const int n_embd = llama_mtp_state_n_embd(ctx);
auto & last = mtp_get_last_embd(ctx);
auto & last = mtp_get_last_embd(state, seq_id);
int i0 = 0;
if (last.last_id >= 0) {
if (last.prob < p_min) {
@@ -1561,7 +2132,11 @@ std::vector<llama_token> mtp_speculative_gen_draft(
last.last_id = -1;
drafts.push_back(current_input_id);
current_n_past++;
llama_set_draft_input_hidden_state(ctx, last.embd.data());
if (!llama_set_draft_input_hidden_state_copy(ctx, last.embd.data(), last.embd.size())) {
llama_batch_free(mtp_batch);
llama_set_mtp_op_type(ctx, MTP_OP_NONE);
return drafts;
}
i0 = 1;
}
@@ -1591,7 +2166,9 @@ std::vector<llama_token> mtp_speculative_gen_draft(
// Keep a stable copy because later decode steps reuse ctx->embd storage.
memcpy(last.embd.data(), emb, n_embd * sizeof(float));
llama_set_draft_input_hidden_state(ctx, last.embd.data());
if (!llama_set_draft_input_hidden_state_copy(ctx, last.embd.data(), last.embd.size())) {
break;
}
current_input_id = id_next;
current_n_past++;
@@ -1655,34 +2232,3 @@ int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch
llama_set_mtp_op_type(ctx, MTP_OP_NONE);
return ret;
}
void mtp_accept_tokens(
struct llama_context * ctx,
const std::vector<llama_token> & ids,
int32_t n_past_base,
llama_seq_id seq_id) {
if (ids.empty()) {
return;
}
llama_batch accepted_batch = llama_batch_init(ids.size(), 0, 1);
for (size_t i = 0; i < ids.size(); ++i) {
common_batch_add(accepted_batch, ids[i], n_past_base + i, { seq_id }, true);
}
if (mtp_update_kv_cache(ctx, accepted_batch, false) != 0) {
LOG_ERR("failed to update MTP KV cache for accepted tokens\n");
llama_batch_free(accepted_batch);
return;
}
auto & last = mtp_get_last_embd(ctx);
auto embd = llama_get_embeddings_ith(ctx, ids.size() - 1);
if (embd) {
std::memcpy(last.embd.data(), embd, last.embd.size()*sizeof(float));
llama_set_draft_input_hidden_state(ctx, last.embd.data());
last.last_id = common_sampler_sample_speculative(nullptr, ctx, ids.size() - 1, &last.prob);
}
llama_batch_free(accepted_batch);
}
common/speculative.h
+65 -22
View File
@@ -1,11 +1,19 @@
#pragma once
#include "llama.h"
#include "llama-spec-features.h"
#include "common.h"
#include "spec-tuner.h"
struct common_speculative;
using common_speculative_feature_kind = llama_spec_feature_kind;
using common_speculative_feature_row_view = llama_spec_feature_row_view;
using common_speculative_feature_view = llama_spec_feature_view;
static constexpr common_speculative_feature_kind COMMON_SPECULATIVE_FEATURE_NONE = LLAMA_SPEC_FEATURE_NONE;
static constexpr common_speculative_feature_kind COMMON_SPECULATIVE_FEATURE_HIDDEN_STATE = LLAMA_SPEC_FEATURE_HIDDEN_STATE;
// comma separated list of all types
std::string common_speculative_type_name_str();
@@ -41,11 +49,66 @@ llama_tokens common_speculative_draft(
// informs the speculative decoder that n_accepted tokens were accepted by the target model
void common_speculative_accept(common_speculative * spec, uint16_t n_accepted);
bool common_speculative_ensure_sequence_hidden(
common_speculative * spec,
llama_context * ctx,
llama_seq_id seq_id,
llama_pos pos);
bool common_speculative_capture_output_hidden(
common_speculative * spec,
llama_context * ctx,
int32_t output_index,
llama_seq_id seq_id,
llama_pos pos);
bool common_speculative_copy_output_hidden_rows(
const common_speculative * spec,
llama_context * ctx,
const std::vector<int32_t> & output_indices,
std::vector<float> & hidden_rows);
bool common_speculative_commit_accepted_hidden_rows(
common_speculative * spec,
common_speculative_type spec_type_used,
llama_seq_id seq_id,
llama_pos pos_base,
llama_token sampled_before,
const std::vector<llama_token> & ids,
const std::vector<float> & hidden_rows);
bool common_speculative_commit_accepted_output(
common_speculative * spec,
llama_context * ctx,
common_speculative_type spec_type_used,
llama_seq_id seq_id,
llama_pos pos_base,
llama_token sampled_before,
const std::vector<llama_token> & ids,
const std::vector<int32_t> & output_indices);
bool common_speculative_has_sequence_hidden(const common_speculative * spec, llama_seq_id seq_id);
void common_speculative_clear_sequence_hidden(common_speculative * spec, llama_seq_id seq_id);
llama_context * common_speculative_get_companion_ctx(common_speculative * spec);
int32_t common_speculative_on_target_seq_batch(
common_speculative * spec,
llama_context * ctx,
const llama_batch & batch,
llama_seq_id seq_id,
bool is_prompt_warmup);
int32_t common_speculative_on_target_batch(
common_speculative * spec,
const llama_batch & batch,
const common_speculative_feature_view & features,
bool is_prompt_warmup);
// print statistics about the speculative decoding
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);
// get the MTP context from the speculative object (nullptr if not MTP type)
llama_context * common_speculative_get_mtp_ctx(common_speculative * spec);
common_speculative_type common_speculative_current_type(const common_speculative * spec);
// Context shift for MTP to match how server handle main model
@@ -55,23 +118,3 @@ void common_speculative_context_shift(
llama_pos kv_keep,
llama_pos kv_discard,
llama_pos kv_past);
// Generates speculative draft tokens using the Multi-Token Prediction (MTP) architecture.
std::vector<llama_token> mtp_speculative_gen_draft(
struct common_sampler * smpl,
struct llama_context * ctx,
int n_draft,
float p_min,
llama_token id_last,
llama_pos n_past,
llama_seq_id seq_id,
bool constant_draft_positions = false);
int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, bool is_prompt_warmup);
void mtp_accept_tokens(
struct llama_context * ctx,
const std::vector<llama_token> & ids,
int32_t n_past_base,
llama_seq_id seq_id
);
examples/imatrix/imatrix.cpp
+6 -3
View File
@@ -7,6 +7,7 @@
#include "common.h"
#include "llama.h"
#include "llama-spec-features.h"
#include <cmath>
#include <cstdio>
@@ -98,8 +99,6 @@ static bool add_and_check_nans(int n, const float * x, float * y, int * counts)
return add_and_check_nans_scalar(n, x, y, counts);
}
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx);
void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx);
static llama_model * ik_load_model_from_params(const gpt_params & params, const llama_model_params & mparams) {
@@ -927,7 +926,11 @@ static bool compute_draft_imatrix_batch(
}
llama_set_mtp_op_type(ctx_dft, MTP_OP_DRAFT_GEN);
llama_set_draft_input_hidden_state(ctx_dft, hidden);
if (!llama_set_draft_input_hidden_state_copy(ctx_dft, hidden, (size_t) batch_size * n_embd_dft)) {
llama_set_mtp_op_type(ctx_dft, MTP_OP_NONE);
fprintf(stderr, "%s: failed to stage paired draft hidden snapshot\n", __func__);
return false;
}
const int ret = llama_decode(ctx_dft, llama_batch_get_one(draft_tokens + batch_start, batch_size, batch_pos, 0));
llama_set_mtp_op_type(ctx_dft, MTP_OP_NONE);
examples/server/server-context.cpp
+79 -298
View File
@@ -17,8 +17,6 @@
#include <regex>
#include <exception>
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx);
static void server_prompt_checkpoint_update(server_prompt_checkpoint & ckpt, llama_context * ctx, int id, int64_t n_tokens, llama_pos pos_min = -1, llama_pos pos_max = -1, int32_t offset = 0) {
if (pos_min == -1) {
pos_min = llama_kv_cache_seq_pos_min(ctx, id);
@@ -52,162 +50,20 @@ static bool params_use_gemma4_external_mtp(const gpt_params & params_base) {
llama_model_is_gemma4_mtp_assistant(params_base.speculative.model_dft);
}
static llama_context * get_slot_mtp_ctx(server_slot & slot, llama_context * ctx) {
llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec);
return mtp_ctx ? mtp_ctx : ctx;
}
static int get_ctx_mtp_n_embd(llama_context * ctx) {
return ctx ? (int) llama_mtp_state_n_embd(ctx) : 0;
}
static int get_slot_mtp_n_embd(server_slot & slot, llama_context * ctx) {
return get_ctx_mtp_n_embd(get_slot_mtp_ctx(slot, ctx));
}
static void cache_slot_mtp_hidden(server_slot & slot, const float * hidden, int n_embd) {
if (hidden == nullptr || n_embd <= 0) {
return;
}
slot.mtp_hidden_state.assign(hidden, hidden + n_embd);
}
static void sync_slot_mtp_hidden(server_slot & slot, llama_context * ctx) {
if (!slot.has_mtp || !slot.spec || slot.mtp_hidden_state.empty()) {
return;
}
const int n_embd = get_slot_mtp_n_embd(slot, ctx);
if (n_embd <= 0 || slot.mtp_hidden_state.size() < (size_t) n_embd) {
return;
}
const int n_hidden = slot.mtp_hidden_state.size() / n_embd;
llama_set_draft_input_hidden_state(get_slot_mtp_ctx(slot, ctx), slot.mtp_hidden_state.data() + (n_hidden - 1) * n_embd);
}
static void cache_and_sync_slot_mtp_hidden(server_slot & slot, llama_context * ctx, const float * hidden, int n_embd) {
cache_slot_mtp_hidden(slot, hidden, n_embd);
sync_slot_mtp_hidden(slot, ctx);
}
static void cache_and_sync_slot_mtp_hidden_from_rows(server_slot & slot, llama_context * ctx, const std::vector<float> & rows, int n_embd) {
if (rows.empty() || n_embd <= 0) {
return;
}
const size_t n_rows = rows.size() / n_embd;
if (n_rows == 0) {
return;
}
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,
const std::vector<float> & mtp_hidden_state,
const std::vector<llama_token> & ids,
int32_t mtp_n_past_base,
int n_embd) {
if (!slot.has_mtp || mtp_hidden_state.empty() || n_embd <= 0) {
return;
}
llama_context * mtp_ctx = get_slot_mtp_ctx(slot, ctx);
if (slot.use_gemma4_external_mtp) {
cache_and_sync_slot_mtp_hidden_from_rows(slot, ctx, mtp_hidden_state, n_embd);
return;
}
slot.mtp_hidden_state = mtp_hidden_state;
llama_set_draft_input_hidden_state(mtp_ctx, slot.mtp_hidden_state.data());
mtp_accept_tokens(mtp_ctx, ids, mtp_n_past_base, slot.id);
}
static void set_external_mtp_hidden(server_slot & slot, llama_context * ctx, const float * hidden, int n_embd) {
if (!slot.has_mtp || !slot.spec || hidden == nullptr || n_embd <= 0) {
return;
}
cache_and_sync_slot_mtp_hidden(slot, ctx, hidden, n_embd);
}
struct server_mtp_warmup {
llama_context * ctx_tgt;
server_slot * slot;
};
static int32_t server_mtp_warmup_batch(
llama_context * ctx_tgt,
llama_context * ctx_mtp,
const llama_batch * batch,
server_slot & slot) {
if (!ctx_tgt || !ctx_mtp || !batch || batch->n_tokens <= 0) {
return 0;
}
const float * emb = llama_get_embeddings(ctx_tgt);
const int n_embd_src = get_ctx_mtp_n_embd(ctx_tgt);
const int n_embd_dst = get_ctx_mtp_n_embd(ctx_mtp);
if (emb == nullptr || n_embd_src <= 0 || n_embd_dst <= 0) {
return -1;
}
if (n_embd_src != n_embd_dst) {
LOG_ERROR("MTP warmup hidden state width mismatch", {
{"n_embd_src", n_embd_src},
{"n_embd_dst", n_embd_dst},
});
return -1;
}
const float * last_hidden = emb + (batch->n_tokens - 1) * n_embd_src;
if (slot.use_gemma4_external_mtp) {
cache_and_sync_slot_mtp_hidden(slot, ctx_tgt, last_hidden, n_embd_dst);
return 0;
}
cache_slot_mtp_hidden(slot, last_hidden, n_embd_dst);
llama_set_draft_input_hidden_state(ctx_mtp, emb);
return mtp_update_kv_cache(ctx_mtp, *batch, true);
}
static int32_t server_mtp_media_warmup_callback(void * user_data, const llama_batch * batch) {
auto * data = static_cast<server_mtp_warmup *>(user_data);
if (data == nullptr || data->slot == nullptr) {
return 0;
}
return server_mtp_warmup_batch(data->ctx_tgt, get_slot_mtp_ctx(*data->slot, data->ctx_tgt), batch, *data->slot);
return batch != nullptr
? common_speculative_on_target_seq_batch(data->slot->spec, data->ctx_tgt, *batch, data->slot->id, true)
: 0;
}
static bool server_response_needs_chat_parse(oaicompat_type oaicompat) {
@@ -592,7 +448,6 @@ void server_context::init() {
params_base.speculative.cparams_dft.embeddings = true;
slot.has_mtp = true;
slot.use_gemma4_external_mtp = has_external_mtp;
slot.params.speculative.cparams_dft = params_base.speculative.cparams_dft;
slot.batch_spec = llama_batch_init(slot.params.speculative.get_max_stage_n_max() + 1, 0, 1);
@@ -755,6 +610,7 @@ void server_slot::reset() {
n_past_prompt = 0;
n_sent_text = 0;
drafted.clear();
drafted_spec_type = COMMON_SPECULATIVE_TYPE_NONE;
i_batch_dft.clear();
spec_ckpt.clear();
n_sent_token_probs = 0;
@@ -773,7 +629,9 @@ void server_slot::reset() {
checkpoint_pos = 0;
image_just_processed = false;
do_checkpoint = false;
mtp_hidden_state.clear();
if (spec != nullptr) {
common_speculative_clear_sequence_hidden(spec, id);
}
positional_bans.clear();
ban_phrases.clear();
@@ -3652,19 +3510,19 @@ void server_context::add_sampled_tokens() {
const llama_pos draft_base_pos = slot.has_mtp ? slot.cache_tokens.pos_next() : -1;
if (slot.has_mtp) {
if (!slot.mtp_hidden_state.empty()) {
sync_slot_mtp_hidden(slot, ctx);
} else {
if (!common_speculative_ensure_sequence_hidden(slot.spec, ctx, slot.id, draft_base_pos - 1)) {
LOG_ERROR("MTP hidden state is empty during speculation", {});
const float* emb_neg1 = llama_get_embeddings_ith(ctx, -1);
if (emb_neg1) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
cache_and_sync_slot_mtp_hidden(slot, ctx, emb_neg1, n_embd);
}
}
}
llama_tokens draft = common_speculative_draft(slot.spec, params_spec, cached_text_tokens, slot.sampled, draft_base_pos, slot.id);
llama_tokens draft = common_speculative_draft(
slot.spec,
params_spec,
cached_text_tokens,
slot.sampled,
draft_base_pos,
slot.id);
slot.drafted_spec_type = common_speculative_current_type(slot.spec);
const int n_draft_max = slot.get_n_draft_max();
@@ -3689,6 +3547,7 @@ void server_context::add_sampled_tokens() {
// fallback to normal decoding
slot.i_batch = slot.i_batch_dft[0];
slot.drafted.clear();
slot.drafted_spec_type = COMMON_SPECULATIVE_TYPE_NONE;
slot.i_batch_dft.clear();
} else {
// keep track of total number of drafted tokens tested
@@ -3705,6 +3564,7 @@ void server_context::add_sampled_tokens() {
}
else {
// no speculative decoding
slot.drafted_spec_type = COMMON_SPECULATIVE_TYPE_NONE;
slot.i_batch = batch.n_tokens;
common_batch_add(batch, slot.sampled, slot.cache_tokens.pos_next(), { slot.id }, true);
@@ -4222,10 +4082,8 @@ void server_context::extend_context(const int32_t n_tokens) {
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,
llama_token sampled_before,
const std::vector<llama_token> & ids, int n_draft,
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;
@@ -4240,35 +4098,17 @@ 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()) {
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());
}
}
if (!common_speculative_commit_accepted_hidden_rows(
slot.spec,
spec_type_used,
slot.id,
mtp_n_past_base,
sampled_before,
ids,
mtp_hidden_state_pre)) {
common_speculative_clear_sequence_hidden(slot.spec, slot.id);
} else if (spec_type_used != COMMON_SPECULATIVE_TYPE_MTP) {
SLT_DBG(slot, "%s", "synced MTP target hidden state from accepted-prefix rows after per-step restore");
}
}
@@ -4303,31 +4143,23 @@ static void restore_speculative_checkpoint(
SLT_ERR(slot, "failed to re-decode accepted tokens after checkpoint restore: %d\n", ret);
}
if (slot.has_mtp) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
const int n_accepted = (int)ids.size();
slot.mtp_hidden_state.resize(n_accepted * n_embd);
for (int j = 0; j < n_accepted; j++) {
const float * emb_j = llama_get_embeddings_ith(ctx, j);
if (emb_j) {
memcpy(slot.mtp_hidden_state.data() + j * n_embd, emb_j, n_embd * sizeof(float));
}
std::vector<int32_t> redecoded_indices(n_accepted);
for (int j = 0; j < n_accepted; ++j) {
redecoded_indices[j] = j;
}
if (slot.use_gemma4_external_mtp) {
cache_and_sync_slot_mtp_hidden_from_rows(slot, ctx, slot.mtp_hidden_state, n_embd);
} else {
llama_context * mtp_ctx = get_slot_mtp_ctx(slot, ctx);
llama_set_draft_input_hidden_state(mtp_ctx, slot.mtp_hidden_state.data());
mtp_accept_tokens(mtp_ctx, ids, slot.spec_ckpt.n_past, slot.id);
if (n_accepted > 1) {
memmove(slot.mtp_hidden_state.data(),
slot.mtp_hidden_state.data() + (n_accepted - 1) * n_embd,
n_embd * sizeof(float));
}
if (!common_speculative_commit_accepted_output(
slot.spec,
ctx,
spec_type_used,
slot.id,
slot.spec_ckpt.n_past,
sampled_before,
ids,
redecoded_indices)) {
common_speculative_clear_sequence_hidden(slot.spec, slot.id);
}
slot.mtp_hidden_state.resize(n_embd);
}
for (llama_token id : ids) {
@@ -4350,9 +4182,8 @@ void server_context::speculative_decoding_accept() {
}
const llama_token sampled_before = slot.sampled;
const common_speculative_type spec_type_used = common_speculative_current_type(slot.spec);
const common_speculative_type spec_type_used = slot.drafted_spec_type;
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) {
@@ -4381,48 +4212,28 @@ void server_context::speculative_decoding_accept() {
continue;
}
const bool any_rejected = (ids.size() - 1) < n_draft;
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;
std::vector<int32_t> accepted_output_indices;
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);
const int n_embd = get_ctx_mtp_n_embd(ctx);
if (!ids.empty()) {
mtp_hidden_state_pre.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_hidden_state_pre.data() + i * n_embd, emb_i, n_embd * sizeof(float));
}
}
accepted_output_indices.assign(slot.i_batch_dft.begin(), slot.i_batch_dft.begin() + ids.size());
}
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) {
mtp_hidden_state_pre.resize(n_embd);
memcpy(mtp_hidden_state_pre.data(), emb0, n_embd * sizeof(float));
if (any_rejected && slot.spec_ckpt.valid && !accepted_output_indices.empty()) {
if (!common_speculative_copy_output_hidden_rows(slot.spec, ctx, accepted_output_indices, mtp_hidden_state_pre)) {
mtp_hidden_state_pre.clear();
}
}
}
slot.i_batch_dft.clear();
slot.drafted.clear();
slot.drafted_spec_type = COMMON_SPECULATIVE_TYPE_NONE;
slot.n_past += ids.size();
slot.n_decoded += ids.size();
@@ -4446,40 +4257,22 @@ void server_context::speculative_decoding_accept() {
slot.n_past = slot.cache_tokens.n_tokens();
// 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, spec_type_used, ids, n_draft, mtp_commit_tokens, mtp_commit_states, mtp_hidden_state_seed, mtp_hidden_state_pre, mtp_n_past_base);
restore_speculative_checkpoint(slot, ctx, model, spec_type_used, sampled_before, ids, n_draft, mtp_hidden_state_pre, mtp_n_past_base);
} else {
if (slot.has_mtp && !mtp_hidden_state_pre.empty()) {
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());
}
}
if (slot.has_mtp && !accepted_output_indices.empty()) {
if (!common_speculative_commit_accepted_output(
slot.spec,
ctx,
spec_type_used,
slot.id,
mtp_n_past_base,
sampled_before,
ids,
accepted_output_indices)) {
common_speculative_clear_sequence_hidden(slot.spec, slot.id);
} else if (spec_type_used != COMMON_SPECULATIVE_TYPE_MTP) {
SLT_DBG(slot, "%s", "synced MTP target hidden state from accepted-prefix rows");
}
}
llama_kv_cache_seq_rm(ctx, slot.id, slot.cache_tokens.pos_next(slot.n_past), -1);
@@ -4858,26 +4651,22 @@ void server_context::process_batch_tokens(int32_t & n_batch) {
continue; // continue loop of n_batch
}
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)) {
bool has_tokens_for_slot = (batch_view.n_tokens > 0 && batch_view.n_seq_id[0] > 0 && batch_view.seq_id[0][0] == slot.id);
if (has_tokens_for_slot) {
mtp_warmup_slot = &slot;
break;
}
}
for (auto & slot : slots) {
if (!slot.spec || !slot.has_mtp) {
continue;
}
}
if (mtp_warmup_slot && mtp_warmup_slot->spec && mtp_warmup_slot->has_mtp) {
llama_context * mtp_ctx = get_slot_mtp_ctx(*mtp_warmup_slot, ctx);
if (server_mtp_warmup_batch(ctx, mtp_ctx, &batch_view, *mtp_warmup_slot) != 0) {
LOG_ERROR("%s\n", "failed to warm up MTP state from prompt batch");
if ((slot.state != SLOT_STATE_PROCESSING || slot.n_decoded != 0) &&
(slot.state != SLOT_STATE_IDLE || slot.command != SLOT_COMMAND_LOAD_PROMPT)) {
continue;
}
if (common_speculative_on_target_seq_batch(slot.spec, ctx, batch_view, slot.id, true) != 0) {
LOG_ERROR("failed to warm up MTP state from prompt batch for slot %d\n", slot.id);
}
}
}
for (auto& slot : slots) {
bool is_active_slot = (slot.state == SLOT_STATE_PROCESSING);
@@ -4929,15 +4718,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) {
const int tok_idx = slot.i_batch - i;
if (slot.has_mtp && slot.n_decoded == 0) {
const float* emb_i = llama_get_embeddings_ith(ctx, tok_idx);
if (emb_i) {
const int n_embd = get_ctx_mtp_n_embd(ctx);
if (slot.use_gemma4_external_mtp) {
set_external_mtp_hidden(slot, ctx, emb_i, n_embd);
} else {
cache_slot_mtp_hidden(slot, emb_i, n_embd);
}
}
(void) common_speculative_capture_output_hidden(slot.spec, ctx, tok_idx, slot.id, slot.n_past);
}
apply_server_biases(slot);
examples/server/server-context.h
+1 -2
View File
@@ -136,6 +136,7 @@ struct server_slot {
// sampling
llama_token sampled; // in speculative mode, this is the last accepted token
llama_tokens drafted;
common_speculative_type drafted_spec_type = COMMON_SPECULATIVE_TYPE_NONE;
json json_schema;
@@ -171,8 +172,6 @@ struct server_slot {
decltype(ctx_sampling->elb_states) elb_prev_states;
bool has_mtp = false;
bool use_gemma4_external_mtp = false;
std::vector<float> mtp_hidden_state;
// saves recurrent state before a speculative batch so it can be restored on rejection
server_speculative_checkpoint spec_ckpt;
src/CMakeLists.txt
+1
View File
@@ -40,6 +40,7 @@ endif()
add_library(llama
../include/llama.h
llama.cpp
llama-spec-features.cpp
llama-vocab.cpp
llama-grammar.cpp
llama-sampling.cpp
src/llama-context.h
+6
View File
@@ -4,6 +4,8 @@
#include "llama-cparams.h"
#include "llama-sampling.h"
#include "llama-spec-features.h"
struct llama_model;
#include <vector>
@@ -245,6 +247,7 @@ struct llama_context {
std::vector<int32_t> output_ids; // map batch token positions to ids of the logits and embd buffers
size_t output_size = 0; // capacity (of tokens positions) for the output buffers
int32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch
int32_t n_outputs_embd = 0; // number of embedding rows produced for the current logical batch
bool logits_all = false;
@@ -272,6 +275,8 @@ struct llama_context {
void * abort_callback_data = nullptr;
const float * draft_input_hidden_state = nullptr;
size_t draft_input_hidden_state_n_floats = 0;
std::vector<float> draft_input_hidden_state_owned;
// input tensors
struct ggml_tensor * inp_tokens; // I32 [n_batch]
@@ -315,3 +320,4 @@ struct llama_context {
void set_mtp_op_type(llama_mtp_op_type value);
};
src/llama-spec-features.cpp
+182
View File
@@ -0,0 +1,182 @@
#include "llama-spec-features.h"
#include <random>
#include "llama-model.h"
#include "llama-context.h"
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) {
if (ctx == nullptr) {
return 0;
}
const auto & hparams = ctx->model.hparams;
if (ctx->cparams.mtp && ctx->model.arch == LLM_ARCH_GEMMA4_MTP && hparams.mtp_backbone_n_embd > 0) {
return hparams.mtp_backbone_n_embd;
}
return hparams.n_embd;
}
bool llama_set_draft_input_hidden_state_copy(
struct llama_context * ctx,
const float * hidden_state,
size_t n_floats) {
if (ctx == nullptr || hidden_state == nullptr || n_floats == 0) {
return false;
}
ctx->draft_input_hidden_state_owned.assign(hidden_state, hidden_state + n_floats);
ctx->draft_input_hidden_state = ctx->draft_input_hidden_state_owned.data();
ctx->draft_input_hidden_state_n_floats = n_floats;
return true;
}
static bool llama_spec_prepare_hidden_feature_view(
struct llama_context * ctx,
int32_t n_rows,
llama_spec_feature_view & view) {
view.kind = LLAMA_SPEC_FEATURE_HIDDEN_STATE;
view.width = 0;
view.rows.clear();
if (ctx == nullptr || n_rows < 0) {
return false;
}
llama_synchronize(ctx);
if (ctx->embd == nullptr) {
return false;
}
view.width = (int32_t) llama_mtp_state_n_embd(ctx);
if (view.width <= 0 || ctx->n_outputs_embd < n_rows) {
view.width = 0;
return false;
}
view.rows.reserve(n_rows);
return true;
}
bool llama_spec_get_hidden_feature_view(
struct llama_context * ctx,
const llama_batch & batch,
llama_spec_feature_view & view) {
if (batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
return false;
}
if (!llama_spec_prepare_hidden_feature_view(ctx, batch.n_tokens, view)) {
return false;
}
for (int32_t i = 0; i < batch.n_tokens; ++i) {
if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) {
view.rows.clear();
return false;
}
view.rows.push_back({
/* .seq_id = */ batch.seq_id[i][0],
/* .pos = */ batch.pos[i],
/* .data = */ ctx->embd + (size_t) i * view.width,
});
}
return true;
}
bool llama_spec_get_hidden_feature_view_for_seq(
struct llama_context * ctx,
const llama_batch & batch,
llama_seq_id seq_id,
llama_spec_feature_view & view) {
if (batch.n_tokens <= 0 || batch.pos == nullptr || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
return false;
}
if (!llama_spec_prepare_hidden_feature_view(ctx, batch.n_tokens, view)) {
return false;
}
for (int32_t i = 0; i < batch.n_tokens; ++i) {
if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) {
view.rows.clear();
return false;
}
for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) {
if (batch.seq_id[i][j] != seq_id) {
continue;
}
view.rows.push_back({
/* .seq_id = */ seq_id,
/* .pos = */ batch.pos[i],
/* .data = */ ctx->embd + (size_t) i * view.width,
});
break;
}
}
return !view.rows.empty();
}
bool llama_spec_get_hidden_feature_view_from_output_index(
struct llama_context * ctx,
int32_t output_index,
llama_seq_id seq_id,
llama_pos pos,
llama_spec_feature_view & view) {
if (!llama_spec_prepare_hidden_feature_view(ctx, 1, view)) {
return false;
}
if (output_index < 0) {
output_index += ctx->n_outputs_embd;
}
if (output_index < 0 || output_index >= ctx->n_outputs_embd) {
view.rows.clear();
return false;
}
view.rows.push_back({
/* .seq_id = */ seq_id,
/* .pos = */ pos,
/* .data = */ ctx->embd + (size_t) output_index * view.width,
});
return true;
}
bool llama_spec_copy_hidden_rows_from_output_indices(
struct llama_context * ctx,
const std::vector<int32_t> & output_indices,
std::vector<float> & hidden_rows) {
hidden_rows.clear();
if (output_indices.empty()) {
return false;
}
llama_spec_feature_view view;
if (!llama_spec_prepare_hidden_feature_view(ctx, (int32_t) output_indices.size(), view)) {
return false;
}
hidden_rows.reserve((size_t) output_indices.size() * view.width);
for (int32_t output_index : output_indices) {
if (output_index < 0) {
output_index += ctx->n_outputs_embd;
}
if (output_index < 0 || output_index >= ctx->n_outputs_embd) {
hidden_rows.clear();
return false;
}
const float * row = ctx->embd + (size_t) output_index * view.width;
hidden_rows.insert(hidden_rows.end(), row, row + view.width);
}
return hidden_rows.size() == (size_t) output_indices.size() * view.width;
}
src/llama-spec-features.h
+54
View File
@@ -0,0 +1,54 @@
#pragma once
#include "llama.h"
#include <vector>
struct llama_context;
enum llama_spec_feature_kind {
LLAMA_SPEC_FEATURE_NONE,
LLAMA_SPEC_FEATURE_HIDDEN_STATE,
};
struct llama_spec_feature_row_view {
llama_seq_id seq_id = 0;
llama_pos pos = -1;
const float * data = nullptr;
};
struct llama_spec_feature_view {
llama_spec_feature_kind kind = LLAMA_SPEC_FEATURE_NONE;
int32_t width = 0;
std::vector<llama_spec_feature_row_view> rows;
};
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx);
bool llama_set_draft_input_hidden_state_copy(
struct llama_context * ctx,
const float * hidden_state,
size_t n_floats);
bool llama_spec_get_hidden_feature_view(
struct llama_context * ctx,
const llama_batch & batch,
llama_spec_feature_view & view);
bool llama_spec_get_hidden_feature_view_for_seq(
struct llama_context * ctx,
const llama_batch & batch,
llama_seq_id seq_id,
llama_spec_feature_view & view);
bool llama_spec_get_hidden_feature_view_from_output_index(
struct llama_context * ctx,
int32_t output_index,
llama_seq_id seq_id,
llama_pos pos,
llama_spec_feature_view & view);
bool llama_spec_copy_hidden_rows_from_output_indices(
struct llama_context * ctx,
const std::vector<int32_t> & output_indices,
std::vector<float> & hidden_rows);
src/llama.cpp
+28 -16
View File
@@ -17,6 +17,7 @@
#include "llama-cparams.h"
#include "llama-hparams.h"
#include "llama-context.h"
#include "llama-spec-features.h"
#include "llama-quantize.h"
#include "unicode.h"
@@ -25,7 +26,6 @@
#include "ggml-alloc.h"
#include "ggml-backend.h"
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx);
void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx);
// TODO: fix these includes
@@ -4650,11 +4650,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) {
// Make sure enough space is available for outputs.
// Returns max number of outputs for which space was reserved.
static uint32_t llama_output_embd_width(const llama_context & lctx) {
const auto & hparams = lctx.model.hparams;
if (lctx.cparams.mtp && lctx.model.arch == LLM_ARCH_GEMMA4_MTP && hparams.mtp_backbone_n_embd > 0) {
return hparams.mtp_backbone_n_embd;
}
return hparams.n_embd;
return llama_mtp_state_n_embd(&lctx);
}
static bool llama_context_has_mtp_outputs(const llama_context & lctx) {
@@ -4775,12 +4771,19 @@ static void llama_graph_compute(
static bool prepare_mtp_graph_inputs(struct llama_context & lctx) {
ggml_tensor * dst = lctx.inp_mtp_states;
const float * src = lctx.draft_input_hidden_state;
const size_t expected_floats = ggml_nbytes(dst) / sizeof(float);
if (!src) {
LLAMA_LOG_ERROR("%s: Source hidden state is null\n", __func__);
return false;
}
if (lctx.draft_input_hidden_state_n_floats != expected_floats) {
LLAMA_LOG_ERROR("%s: Source hidden state size mismatch (have %zu floats, need %zu)\n",
__func__, lctx.draft_input_hidden_state_n_floats, expected_floats);
return false;
}
ggml_backend_tensor_set(dst, src, 0, ggml_nbytes(dst));
return true;
}
@@ -5262,6 +5265,7 @@ static int llama_decode_internal(
// set to total number of outputs in the batch, for use in llama_get_logits_ith
lctx.n_outputs = n_outputs;
lctx.n_outputs_embd = n_outputs_embd;
// wait for the computation to finish (automatically done when obtaining the model output)
//llama_synchronize(&lctx);
@@ -7918,7 +7922,8 @@ struct llama_data_write {
}
void write_embeddings(const struct llama_context * ctx) {
const uint64_t embeddings_size = std::min((uint64_t) ctx->embd_size, (uint64_t) ctx->n_outputs * ctx->model.hparams.n_embd);
const uint64_t row_width = llama_output_embd_width(*ctx);
const uint64_t embeddings_size = std::min((uint64_t) ctx->embd_size, (uint64_t) ctx->n_outputs_embd * row_width);
write(&embeddings_size, sizeof(embeddings_size));
@@ -8214,6 +8219,13 @@ struct llama_data_read {
throw std::runtime_error("embeddings buffer too small");
}
const uint64_t row_width = llama_output_embd_width(*ctx);
if (row_width == 0 || (embeddings_size % row_width) != 0) {
throw std::runtime_error("invalid embeddings payload size");
}
ctx->n_outputs_embd = embeddings_size / row_width;
if (embeddings_size) {
read_to(ctx->embd, embeddings_size * sizeof(float));
}
@@ -9320,9 +9332,9 @@ float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i) {
}
if (i < 0) {
j = ctx->n_outputs + i;
j = ctx->n_outputs_embd + i;
if (j < 0) {
throw std::runtime_error(format("negative index out of range [0, %d)", ctx->n_outputs));
throw std::runtime_error(format("negative index out of range [0, %d)", ctx->n_outputs_embd));
}
} else if ((size_t) i >= ctx->output_ids.size()) {
throw std::runtime_error(format("out of range [0, %lu)", ctx->output_ids.size()));
@@ -9333,12 +9345,12 @@ float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i) {
if (j < 0) {
throw std::runtime_error(format("batch.logits[%d] != true", i));
}
if (j >= ctx->n_outputs) {
if (j >= ctx->n_outputs_embd) {
// This should not happen
throw std::runtime_error(format("corrupt output buffer (j=%d, n_outputs=%d)", j, ctx->n_outputs));
throw std::runtime_error(format("corrupt output buffer (j=%d, n_outputs_embd=%d)", j, ctx->n_outputs_embd));
}
return ctx->embd + j*ctx->model.hparams.n_embd;
return ctx->embd + (size_t) j * llama_output_embd_width(*ctx);
} catch (const std::exception & err) {
LLAMA_LOG_ERROR("%s: invalid embeddings id %d, reason: %s\n", __func__, i, err.what());
#ifndef NDEBUG
@@ -10566,11 +10578,11 @@ void llama_set_offload_policy(struct llama_context * lctx, int op, bool on_or_of
}
void llama_set_draft_input_hidden_state(struct llama_context * ctx, const float * hidden_state) {
ctx->draft_input_hidden_state_owned.clear();
ctx->draft_input_hidden_state = hidden_state;
}
uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) {
return llama_output_embd_width(*ctx);
ctx->draft_input_hidden_state_n_floats = ctx->inp_mtp_states
? ggml_nbytes(ctx->inp_mtp_states) / sizeof(float)
: 0;
}
void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx) {