From 11a1fea9e291f12ce2c803a9d7812c30ca806bcf Mon Sep 17 00:00:00 2001 From: Samuel Oliveira Alves <107287165+SamuelOliveirads@users.noreply.github.com> Date: Wed, 20 May 2026 11:42:48 -0300 Subject: [PATCH] Move embedding management to speculative (#1825) * 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 --- common/speculative.cpp | 696 +++++++++++++++++++++++++---- common/speculative.h | 87 +++- examples/imatrix/imatrix.cpp | 9 +- examples/server/server-context.cpp | 377 ++++------------ examples/server/server-context.h | 3 +- src/CMakeLists.txt | 1 + src/llama-context.h | 6 + src/llama-spec-features.cpp | 182 ++++++++ src/llama-spec-features.h | 54 +++ src/llama.cpp | 44 +- 10 files changed, 1043 insertions(+), 416 deletions(-) create mode 100644 src/llama-spec-features.cpp create mode 100644 src/llama-spec-features.h diff --git a/common/speculative.cpp b/common/speculative.cpp index bf78c33b..7c92e7cb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -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 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 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 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> target_hidden_by_seq; + std::unordered_map 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 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 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 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 * 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 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 & 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 & output_indices, + std::vector & 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 & ids, + std::vector & 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 & ids, + const std::vector & 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 & ids, + const std::vector & hidden_rows) { + if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || ids.empty()) { + return true; + } + + std::vector 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 & ids, + const std::vector & output_indices) { + if (!common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || ids.empty()) { + return true; + } + + std::vector 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(impl.get()); - if (mtp_state) { - return mtp_state->ctx_mtp; - } + if (impl->type != COMMON_SPECULATIVE_TYPE_MTP) { + continue; + } + + if (auto * mtp_state = dynamic_cast(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(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 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 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 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 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 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 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 & 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); -} diff --git a/common/speculative.h b/common/speculative.h index 981e42b5..06d4b580 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -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 & output_indices, + std::vector & 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 & ids, + const std::vector & 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 & ids, + const std::vector & 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 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 & ids, - int32_t n_past_base, - llama_seq_id seq_id -); diff --git a/examples/imatrix/imatrix.cpp b/examples/imatrix/imatrix.cpp index a59e3061..cff41565 100644 --- a/examples/imatrix/imatrix.cpp +++ b/examples/imatrix/imatrix.cpp @@ -7,6 +7,7 @@ #include "common.h" #include "llama.h" +#include "llama-spec-features.h" #include #include @@ -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); diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index c53a3134..a1292615 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -17,8 +17,6 @@ #include #include -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 & 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 & 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 & 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 & mtp_hidden_state, - const std::vector & 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(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 & ids, int n_draft, - const std::vector & mtp_commit_tokens, - const std::vector & mtp_commit_states, - const std::vector & mtp_hidden_state_seed, + llama_token sampled_before, + const std::vector & ids, int n_draft, const std::vector & 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 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 mtp_hidden_state_seed = slot.has_mtp ? slot.mtp_hidden_state : std::vector{}; 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 mtp_hidden_state_pre; - std::vector mtp_commit_tokens; - std::vector mtp_commit_states; + std::vector 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); diff --git a/examples/server/server-context.h b/examples/server/server-context.h index 02194ab9..9d643e02 100644 --- a/examples/server/server-context.h +++ b/examples/server/server-context.h @@ -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 mtp_hidden_state; // saves recurrent state before a speculative batch so it can be restored on rejection server_speculative_checkpoint spec_ckpt; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e53b9982..59f9b0a0 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -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 diff --git a/src/llama-context.h b/src/llama-context.h index a9e3c442..db0018d7 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -4,6 +4,8 @@ #include "llama-cparams.h" #include "llama-sampling.h" +#include "llama-spec-features.h" + struct llama_model; #include @@ -245,6 +247,7 @@ struct llama_context { std::vector 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 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); }; + diff --git a/src/llama-spec-features.cpp b/src/llama-spec-features.cpp new file mode 100644 index 00000000..5a32b848 --- /dev/null +++ b/src/llama-spec-features.cpp @@ -0,0 +1,182 @@ +#include "llama-spec-features.h" + +#include + +#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 & output_indices, + std::vector & 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; +} \ No newline at end of file diff --git a/src/llama-spec-features.h b/src/llama-spec-features.h new file mode 100644 index 00000000..7634197b --- /dev/null +++ b/src/llama-spec-features.h @@ -0,0 +1,54 @@ +#pragma once + +#include "llama.h" + +#include + +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 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 & output_indices, + std::vector & hidden_rows); \ No newline at end of file diff --git a/src/llama.cpp b/src/llama.cpp index bea371ee..3d773d16 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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) {