diff --git a/common/chat.cpp b/common/chat.cpp index ed1c0e54..7f3008ca 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -28,9 +28,59 @@ #include #include #include +#include +#include using json = nlohmann::ordered_json; +static std::string trim_copy(const std::string & s) { + size_t start = 0; + while (start < s.size() && std::isspace(static_cast(s[start]))) { + start++; + } + size_t end = s.size(); + while (end > start && std::isspace(static_cast(s[end - 1]))) { + end--; + } + return s.substr(start, end - start); +} + +static std::string canonicalize_tool_call_arguments(const std::string & arguments, bool is_partial) { + std::string out = trim_copy(arguments); + if (!is_partial && !out.empty()) { + try { + const auto parsed = nlohmann::json::parse(out); + out = parsed.dump(); + } catch (const std::exception &) { + } + } + return out; +} + +static void dedupe_tool_calls(std::vector & tool_calls, bool is_partial, const common_chat_parser_params & params) { + if (!params.parse_tool_calls || tool_calls.size() < 2) { + return; + } + + std::unordered_set seen; + seen.reserve(tool_calls.size()); + + std::vector out; + out.reserve(tool_calls.size()); + + for (auto & tc : tool_calls) { + std::string key = tc.name; + key.push_back('\x1f'); + key += canonicalize_tool_call_arguments(tc.arguments, is_partial); + + if (seen.insert(key).second) { + out.push_back(std::move(tc)); + } + } + + tool_calls = std::move(out); +} + static std::string format_time(const std::chrono::system_clock::time_point & now, const std::string & format) { auto time = std::chrono::system_clock::to_time_t(now); auto local_time = *std::localtime(&time); @@ -2499,6 +2549,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars } mapper->from_ast(ctx.ast, result); + dedupe_tool_calls(msg.tool_calls, is_partial, params); + if (ctx.is_debug()) { fprintf(stderr, "\nAST for partial parse (fail):\n%s\n", ctx.ast.dump().c_str()); fflush(stderr); @@ -2519,6 +2571,7 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars mapper = std::make_unique(msg); } mapper->from_ast(ctx.ast, result); + dedupe_tool_calls(msg.tool_calls, is_partial, params); if (ctx.is_debug()) { fprintf(stderr, "\nAST for %s parse:\n%s\n", is_partial ? "partial" : "full", ctx.ast.dump().c_str()); diff --git a/common/common.cpp b/common/common.cpp index d308fe9d..9b580ae1 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1054,6 +1054,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.speculative.autotune = true; return true; } + if (arg == "--mtp-adaptive") { + params.speculative.mtp_adaptive = true; + return true; + } if (arg == "--chunks") { CHECK_ARG params.n_chunks = std::stoi(argv[i]); @@ -2745,6 +2749,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", "-hft, --hf-token TOKEN", "Hugging Face access token (default: value from HF_TOKEN environment variable)" }); options.push_back({ "*", "-mtp, --multi-token-prediction", "whether to use multi-token-prediction (if supported) (default: %s)", params.has_mtp ? "true" : "false" }); options.push_back({ "*", "-no-mtp, --no-multi-token-prediction", "whether to use multi-token-prediction (if supported) (default: %s)", !params.has_mtp ? "true" : "false" }); + options.push_back({ "*", "--mtp-adaptive", "server-side MTP gate: calibrate no-MTP speed, then disable MTP on poor recent speed or acceptance" }); options.push_back({ "*", "--draft-max, --draft, --draft-n N", "number of tokens to draft for speculative decoding (default: %d)", params.speculative.n_max }); options.push_back({ "*", "--draft-min, --draft-n-min N", "minimum number of draft tokens to use for speculative decoding" }); diff --git a/common/common.h b/common/common.h index 734d93de..423ca846 100644 --- a/common/common.h +++ b/common/common.h @@ -210,6 +210,7 @@ struct common_params_speculative { std::string cache_type_v = ""; // KV cache data type for V for the draft model bool autotune = false; // automatically optimize speculative params for max tokens/sec + bool mtp_adaptive = false; // server-side MTP runtime gate bool has_dft() const { return !model.empty() || !params.empty(); diff --git a/common/speculative.cpp b/common/speculative.cpp index d63edd74..70151855 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -1353,6 +1354,11 @@ void common_speculative_context_shift( } } +static bool mtp_draft_gpu_argmax_enabled() { + const char * value = std::getenv("LLAMA_MTP_DRAFT_GPU_ARGMAX"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + std::vector mtp_speculative_gen_draft( struct common_sampler * smpl, struct llama_context * ctx, @@ -1370,6 +1376,9 @@ std::vector mtp_speculative_gen_draft( common_sampler_reset(smpl); llama_batch mtp_batch = llama_batch_init(1, 0, 1); + const bool use_gpu_argmax = p_min <= 0.0f && mtp_draft_gpu_argmax_enabled(); + + llama_set_mtp_draft_gpu_argmax(ctx, use_gpu_argmax); llama_set_mtp_op_type(ctx, MTP_OP_DRAFT_GEN); llama_token current_input_id = id_last; @@ -1383,9 +1392,17 @@ std::vector mtp_speculative_gen_draft( break; } - float prob; - llama_token id_next = common_sampler_sample_speculative(smpl, ctx, 0, &prob); + float prob = 1.0f; + llama_token id_next = LLAMA_TOKEN_NULL; + if (use_gpu_argmax) { + id_next = llama_get_mtp_draft_argmax_ith(ctx, 0); + } else { + id_next = common_sampler_sample_speculative(smpl, ctx, 0, p_min > 0.0f ? &prob : nullptr); + } + if (id_next == LLAMA_TOKEN_NULL) { + break; + } drafts.push_back(id_next); const float * emb = llama_get_embeddings_ith(ctx, 0); @@ -1403,10 +1420,8 @@ std::vector mtp_speculative_gen_draft( llama_batch_free(mtp_batch); llama_set_mtp_op_type(ctx, MTP_OP_NONE); - // Purge the metadata for the draft tokens. - // This prevents cache state corruption where two cells map to the same logical position. - if (!drafts.empty()) { - llama_kv_cache_seq_rm(ctx, seq_id, n_past, current_n_past); + if (current_n_past > n_past + 1) { + llama_kv_cache_seq_rm(ctx, seq_id, n_past + 1, current_n_past); } return drafts; @@ -1436,7 +1451,7 @@ void mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, b } for (int i = 0; i < mtp_batch.n_tokens; ++i) { - mtp_batch.logits[i] = true; + mtp_batch.logits[i] = false; } llama_decode(ctx, mtp_batch); llama_set_mtp_op_type(ctx, MTP_OP_NONE); @@ -1452,8 +1467,12 @@ void mtp_accept_tokens( return; } - llama_batch accepted_batch = llama_batch_init(ids.size(), 0, 1); - for (size_t i = 0; i < ids.size(); ++i) { + if (ids.size() == 1) { + return; + } + + llama_batch accepted_batch = llama_batch_init(ids.size() - 1, 0, 1); + for (size_t i = 1; i < ids.size(); ++i) { common_batch_add(accepted_batch, ids[i], n_past_base + i, { seq_id }, true); } diff --git a/examples/imatrix/imatrix.cpp b/examples/imatrix/imatrix.cpp index 8bc44587..6be1eb98 100644 --- a/examples/imatrix/imatrix.cpp +++ b/examples/imatrix/imatrix.cpp @@ -638,6 +638,37 @@ static void process_logits( } } +static bool run_mtp_imatrix_warmup(llama_context * ctx, llama_token * tokens, const float * hidden_states, int32_t n_tokens, llama_pos pos_0) { + const llama_model * model = llama_get_model(ctx); + if (llama_model_n_nextn_layer(model) <= 0) { + return true; + } + + llama_batch mtp_batch = llama_batch_init(n_tokens, 0, 1); + mtp_batch.n_tokens = n_tokens; + for (int32_t i = 0; i < n_tokens; ++i) { + mtp_batch.token[i] = tokens[i]; + mtp_batch.pos[i] = pos_0 + i; + mtp_batch.n_seq_id[i] = 1; + mtp_batch.seq_id[i][0] = 0; + mtp_batch.logits[i] = 1; + } + + llama_set_draft_input_hidden_state(ctx, hidden_states); + llama_set_mtp_op_type(ctx, MTP_OP_WARMUP); + const int ret = llama_decode(ctx, mtp_batch); + llama_set_mtp_op_type(ctx, MTP_OP_NONE); + llama_set_draft_input_hidden_state(ctx, nullptr); + llama_batch_free(mtp_batch); + + if (ret != 0) { + fprintf(stderr, "%s: failed to eval MTP warmup batch\n", __func__); + return false; + } + + return true; +} + static bool compute_imatrix(llama_context * ctx, const gpt_params & params) { const bool add_bos = llama_should_add_bos_token(llama_get_model(ctx)); GGML_ASSERT(llama_add_eos_token(llama_get_model(ctx)) != 1); @@ -680,12 +711,17 @@ static bool compute_imatrix(llama_context * ctx, const gpt_params & params) { const int n_chunk = params.n_chunks < 0 ? n_chunk_max : std::min(params.n_chunks, n_chunk_max); const int n_vocab = llama_n_vocab(llama_get_model(ctx)); const int n_batch = params.n_batch; + const bool collect_mtp = params.has_mtp && llama_model_n_nextn_layer(llama_get_model(ctx)) > 0; + const int n_embd = collect_mtp ? llama_model_n_embd(llama_get_model(ctx)) : 0; int count = 0; double nll = 0.0; double nll2 = 0.0; fprintf(stderr, "%s: computing over %d chunks with batch_size %d\n", __func__, n_chunk, n_batch); + if (collect_mtp) { + fprintf(stderr, "%s: MTP warmup collection enabled\n", __func__); + } std::vector workers(std::thread::hardware_concurrency() - 1); @@ -701,6 +737,18 @@ static bool compute_imatrix(llama_context * ctx, const gpt_params & params) { const int end = start + n_ctx; std::vector logits; + if (params.compute_ppl && collect_mtp) { + logits.reserve((size_t)n_ctx * n_vocab); + } + std::vector mtp_hidden_states; + std::vector mtp_tokens; + if (collect_mtp) { + mtp_hidden_states.resize((size_t)n_ctx * n_embd); + mtp_tokens.assign(tokens.begin() + start, tokens.begin() + end); + if (add_bos) { + mtp_tokens[0] = llama_token_bos(llama_get_model(ctx)); + } + } const auto t_start = std::chrono::high_resolution_clock::now(); @@ -725,12 +773,36 @@ static bool compute_imatrix(llama_context * ctx, const gpt_params & params) { return false; } + if (params.compute_ppl && (num_batches > 1 || collect_mtp)) { + const auto * batch_logits = llama_get_logits(ctx); + logits.insert(logits.end(), batch_logits, batch_logits + batch_size * n_vocab); + } + + if (collect_mtp) { + float * hidden_dst = mtp_hidden_states.data() + (size_t)j * n_batch * n_embd; + for (int k = 0; k < batch_size; ++k) { + const float * emb = llama_get_embeddings_ith(ctx, k); + if (!emb) { + fprintf(stderr, "%s: failed to read main-model hidden state for token %d\n", __func__, k); + return false; + } + std::memcpy(hidden_dst + (size_t)k * n_embd, emb, (size_t)n_embd * sizeof(float)); + } + } + // restore the original token in case it was set to BOS tokens[batch_start] = token_org; + } - if (params.compute_ppl && num_batches > 1) { - const auto * batch_logits = llama_get_logits(ctx); - logits.insert(logits.end(), batch_logits, batch_logits + batch_size * n_vocab); + if (collect_mtp) { + llama_kv_cache_clear(ctx); + const int mtp_batch = std::max(1, params.n_ubatch); + for (int mtp_start = 0; mtp_start < n_ctx; mtp_start += mtp_batch) { + const int mtp_size = std::min(n_ctx - mtp_start, mtp_batch); + if (!run_mtp_imatrix_warmup(ctx, mtp_tokens.data() + mtp_start, + mtp_hidden_states.data() + (size_t)mtp_start * n_embd, mtp_size, mtp_start)) { + return false; + } } } @@ -749,7 +821,7 @@ static bool compute_imatrix(llama_context * ctx, const gpt_params & params) { if (params.compute_ppl) { const int first = n_ctx/2; - const auto all_logits = num_batches > 1 ? logits.data() : llama_get_logits(ctx); + const auto all_logits = !logits.empty() ? logits.data() : llama_get_logits(ctx); process_logits(n_vocab, all_logits + first*n_vocab, tokens.data() + start + first, n_ctx - 1 - first, workers, nll, nll2, logit_history.data() + start + first, prob_history.data() + start + first); count += n_ctx - first - 1; diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index b38d13a0..347213f3 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -22,6 +22,107 @@ static void log_text(const gpt_params & params_base, const std::string & text) { } } +static constexpr int32_t MTP_ADAPTIVE_BASELINE_TOKENS = 4; +static constexpr int32_t MTP_ADAPTIVE_MIN_WINDOWS = 4; +static constexpr int32_t MTP_ADAPTIVE_MAX_BAD_WINDOWS = 2; +static constexpr double MTP_ADAPTIVE_EMA_ALPHA = 0.25; +static constexpr double MTP_ADAPTIVE_MIN_ACCEPT = 0.25; +static constexpr double MTP_ADAPTIVE_MIN_TPS_RATIO = 0.98; + +static void mtp_adaptive_update_ema(double & ema, double value) { + if (ema <= 0.0) { + ema = value; + } else { + ema = MTP_ADAPTIVE_EMA_ALPHA * value + (1.0 - MTP_ADAPTIVE_EMA_ALPHA) * ema; + } +} + +static bool mtp_adaptive_enabled(const server_slot & slot) { + return slot.has_mtp && slot.params.speculative.mtp_adaptive; +} + +static void mtp_adaptive_disable(server_slot & slot, const char * reason, double value, double threshold) { + if (slot.mtp_adaptive_disabled) { + return; + } + + slot.mtp_adaptive_disabled = true; + slot.mtp_adaptive_step_start_us = 0; + slot.mtp_adaptive_no_mtp_step_start_us = 0; + + SLT_WRN(slot, + "adaptive MTP disabled: %s (value %.3f, threshold %.3f, no_mtp_tps %.2f, mtp_tps %.2f, accept %.3f)\n", + reason, value, threshold, slot.mtp_adaptive_no_mtp_tps, slot.mtp_adaptive_mtp_tps, + slot.mtp_adaptive_accept); +} + +static void mtp_adaptive_note_no_mtp(server_slot & slot, int64_t t_now_us) { + if (!mtp_adaptive_enabled(slot) || slot.mtp_adaptive_disabled || slot.mtp_adaptive_no_mtp_step_start_us <= 0) { + return; + } + + const int64_t elapsed_us = t_now_us - slot.mtp_adaptive_no_mtp_step_start_us; + slot.mtp_adaptive_no_mtp_step_start_us = 0; + if (elapsed_us <= 100) { + return; + } + + mtp_adaptive_update_ema(slot.mtp_adaptive_no_mtp_tps, 1e6 / (double) elapsed_us); + slot.mtp_adaptive_baseline_seen++; + + if (slot.mtp_adaptive_baseline_seen == MTP_ADAPTIVE_BASELINE_TOKENS) { + SLT_DBG(slot, "adaptive MTP no-MTP baseline ready: %.2f tok/s over %d tokens\n", + slot.mtp_adaptive_no_mtp_tps, slot.mtp_adaptive_baseline_seen); + } +} + +static void mtp_adaptive_note_mtp(server_slot & slot, size_t n_draft, size_t n_output, int64_t t_now_us) { + if (!mtp_adaptive_enabled(slot) || slot.mtp_adaptive_disabled || n_draft == 0) { + return; + } + + const size_t n_accepted = n_output > 0 ? n_output - 1 : 0; + const double acceptance = (double) n_accepted / (double) n_draft; + mtp_adaptive_update_ema(slot.mtp_adaptive_accept, acceptance); + slot.mtp_adaptive_windows++; + + if (slot.mtp_adaptive_step_start_us > 0) { + const int64_t elapsed_us = t_now_us - slot.mtp_adaptive_step_start_us; + slot.mtp_adaptive_step_start_us = 0; + if (elapsed_us > 100 && n_output > 0) { + mtp_adaptive_update_ema(slot.mtp_adaptive_mtp_tps, (double) n_output * 1e6 / (double) elapsed_us); + } + } + + if (slot.mtp_adaptive_windows < MTP_ADAPTIVE_MIN_WINDOWS) { + return; + } + + if (slot.mtp_adaptive_accept < MTP_ADAPTIVE_MIN_ACCEPT) { + slot.mtp_adaptive_low_accept++; + } else { + slot.mtp_adaptive_low_accept = 0; + } + + if (slot.mtp_adaptive_low_accept >= MTP_ADAPTIVE_MAX_BAD_WINDOWS) { + mtp_adaptive_disable(slot, "low acceptance", slot.mtp_adaptive_accept, MTP_ADAPTIVE_MIN_ACCEPT); + return; + } + + if (slot.mtp_adaptive_baseline_seen >= MTP_ADAPTIVE_BASELINE_TOKENS && + slot.mtp_adaptive_no_mtp_tps > 0.0 && slot.mtp_adaptive_mtp_tps > 0.0 && + slot.mtp_adaptive_mtp_tps < slot.mtp_adaptive_no_mtp_tps * MTP_ADAPTIVE_MIN_TPS_RATIO) { + slot.mtp_adaptive_slow_windows++; + } else { + slot.mtp_adaptive_slow_windows = 0; + } + + if (slot.mtp_adaptive_slow_windows >= MTP_ADAPTIVE_MAX_BAD_WINDOWS) { + mtp_adaptive_disable(slot, "slower than no-MTP", slot.mtp_adaptive_mtp_tps, + slot.mtp_adaptive_no_mtp_tps * MTP_ADAPTIVE_MIN_TPS_RATIO); + } +} + void server_speculative_checkpoint::clear() { valid = false; per_step_enabled = false; @@ -456,6 +557,16 @@ void server_slot::reset() { // Reset speculative decoding stats n_draft_total = 0; n_draft_accepted = 0; + mtp_adaptive_disabled = false; + mtp_adaptive_baseline_seen = 0; + mtp_adaptive_windows = 0; + mtp_adaptive_low_accept = 0; + mtp_adaptive_slow_windows = 0; + mtp_adaptive_step_start_us = 0; + mtp_adaptive_no_mtp_step_start_us = 0; + mtp_adaptive_no_mtp_tps = 0.0; + mtp_adaptive_mtp_tps = 0.0; + mtp_adaptive_accept = 0.0; chat_msg = {}; json_schema = json(); generated_tool_call_ids.clear(); @@ -510,13 +621,19 @@ void server_slot::add_token_string(const completion_token_output& token) { } bool server_slot::can_speculate() const { - return (!!spec || has_mtp); + return !mtp_adaptive_disabled && (!!spec || has_mtp); } int server_slot::get_n_draft_max() const { if (!can_speculate()) { return 0; } + if (has_mtp && params.speculative.mtp_adaptive && + mtp_adaptive_baseline_seen < MTP_ADAPTIVE_BASELINE_TOKENS) { + SLT_DBG(*this, "adaptive MTP collecting no-MTP baseline: %d/%d\n", + mtp_adaptive_baseline_seen, MTP_ADAPTIVE_BASELINE_TOKENS); + return 0; + } // determine the max draft that fits the current slot state int n_draft_max = params.speculative.n_max; @@ -1049,6 +1166,7 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) slot.params.speculative.n_max = json_value(data, "speculative.n_max", params_base.speculative.n_max); slot.params.speculative.n_min = json_value(data, "speculative.n_min", params_base.speculative.n_min); slot.params.speculative.p_min = json_value(data, "speculative.p_min", params_base.speculative.p_min); + slot.params.speculative.mtp_adaptive = json_value(data, "speculative.mtp_adaptive", defaults.speculative.mtp_adaptive); slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min); slot.params.speculative.n_min = std::max(slot.params.speculative.n_min, 0); @@ -1608,6 +1726,10 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) bool do_checkpoint = params_base.ctx_checkpoints_n > 0; // make checkpoints only for completion tasks do_checkpoint = do_checkpoint && task.type == SERVER_TASK_TYPE_COMPLETION; + if (do_checkpoint && llama_model_is_split_mode_graph(llama_get_model(slot.ctx))) { + LLAMA_LOG_WARN("%s: disabling recurrent checkpoints for split-mode graph; partial sequence snapshots are unstable on this path\n", __func__); + do_checkpoint = false; + } // make a checkpoint of the parts of the memory that cannot be rolled back. // checkpoints are created only if: // - the model architecture is marked as recurrent or hybrid @@ -3155,6 +3277,10 @@ void server_context::add_sampled_tokens() { } } + if (mtp_adaptive_enabled(slot)) { + slot.mtp_adaptive_step_start_us = ggml_time_us(); + } + llama_tokens draft = common_speculative_draft(slot.spec, params_spec, cached_text_tokens, slot.sampled); const int n_draft_max = slot.get_n_draft_max(); @@ -3180,6 +3306,7 @@ void server_context::add_sampled_tokens() { slot.i_batch = slot.i_batch_dft[0]; slot.drafted.clear(); slot.i_batch_dft.clear(); + slot.mtp_adaptive_step_start_us = 0; } else { // keep track of total number of drafted tokens tested @@ -3197,6 +3324,9 @@ void server_context::add_sampled_tokens() { else { // no speculative decoding slot.i_batch = batch.n_tokens; + if (mtp_adaptive_enabled(slot) && !slot.mtp_adaptive_disabled) { + slot.mtp_adaptive_no_mtp_step_start_us = ggml_time_us(); + } common_batch_add(batch, slot.sampled, slot.cache_tokens.pos_next(), { slot.id }, true); @@ -3349,6 +3479,10 @@ bool server_context::create_checkpoint(server_slot & slot) { } void server_context::batch_pending_prompt(const int32_t n_ubatch, const int32_t n_batch, int32_t & batch_type) { + const bool serialize_recurrent_graph_prompts = + llama_model_has_recurrent(llama_get_model(ctx)) && + llama_model_is_split_mode_graph(llama_get_model(ctx)); + if (params_base.cont_batching || batch.n_tokens == 0) { for (auto& slot : slots) { // this slot still has a prompt to be processed @@ -3688,6 +3822,10 @@ void server_context::batch_pending_prompt(const int32_t n_ubatch, const int32_t } } + if (serialize_recurrent_graph_prompts && batch.n_tokens > 0) { + break; + } + if (batch.n_tokens >= n_batch) { break; } @@ -3770,11 +3908,22 @@ static void restore_speculative_checkpoint( common_batch_add(re_batch, ids[j], slot.spec_ckpt.n_past + 1 + j, { slot.id }, j == n_re - 2); } + const int n_embd = slot.has_mtp ? llama_model_n_embd(llama_get_model(ctx)) : 0; + const size_t mtp_hidden_state_needed = (size_t)n_re * (size_t)n_embd; + const bool can_reuse_mtp_hidden_state = + slot.has_mtp && + n_embd > 0 && + mtp_hidden_state_pre.size() >= mtp_hidden_state_needed; + const bool need_redecode_logits = + slot.sparams.n_probs > 0 && !slot.params.post_sampling_probs; + const bool need_redecode_mtp_hidden_state = + slot.has_mtp && !can_reuse_mtp_hidden_state; + if (slot.has_mtp) { for (int j = 0; j < re_batch.n_tokens; j++) { - re_batch.logits[j] = true; + re_batch.logits[j] = need_redecode_mtp_hidden_state || need_redecode_logits; } - llama_set_embeddings(ctx, true); + llama_set_embeddings(ctx, need_redecode_mtp_hidden_state); } const int ret = llama_decode(ctx, re_batch); @@ -3782,14 +3931,18 @@ 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 = llama_model_n_embd(llama_get_model(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)); + if (can_reuse_mtp_hidden_state) { + slot.mtp_hidden_state.assign( + mtp_hidden_state_pre.begin(), + mtp_hidden_state_pre.begin() + mtp_hidden_state_needed); + } else { + 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)); + } } } @@ -3811,8 +3964,8 @@ static void restore_speculative_checkpoint( } llama_batch_free(re_batch); - SLT_DBG(slot, "spec checkpoint restored: re-decoded %d tokens (rejected %d drafts)\n", - n_re, (int)(n_draft - (ids.size() - 1))); + SLT_DBG(slot, "spec checkpoint restored: re-decoded %d tokens (rejected %d drafts, reused_mtp_hidden=%d)\n", + n_re, (int)(n_draft - (ids.size() - 1)), can_reuse_mtp_hidden_state ? 1 : 0); } } @@ -3862,6 +4015,7 @@ void server_context::speculative_decoding_accept() { slot.n_decoded += ids.size(); const int64_t t_current = ggml_time_us(); slot.t_token_generation = std::max(1, t_current - slot.t_start_generation) / 1e3; + mtp_adaptive_note_mtp(slot, n_draft, ids.size(), t_current); // update how many tokens out of those tested were accepted slot.n_draft_accepted += ids.size() - 1; @@ -4363,6 +4517,7 @@ void server_context::process_batch_tokens(int32_t & n_batch) { } slot.t_token_generation = std::max(1, t_current - slot.t_start_generation) / 1e3; + mtp_adaptive_note_no_mtp(slot, ggml_time_us()); result.tok = id; result.prob = 1.0f; // TODO: set it here instead of doing inside populate_token_probs diff --git a/examples/server/server-context.h b/examples/server/server-context.h index 074787b5..63d6b8bb 100644 --- a/examples/server/server-context.h +++ b/examples/server/server-context.h @@ -169,6 +169,16 @@ struct server_slot { bool has_mtp = false; std::vector mtp_hidden_state; + bool mtp_adaptive_disabled = false; + int32_t mtp_adaptive_baseline_seen = 0; + int32_t mtp_adaptive_windows = 0; + int32_t mtp_adaptive_low_accept = 0; + int32_t mtp_adaptive_slow_windows = 0; + int64_t mtp_adaptive_step_start_us = 0; + int64_t mtp_adaptive_no_mtp_step_start_us = 0; + double mtp_adaptive_no_mtp_tps = 0.0; + double mtp_adaptive_mtp_tps = 0.0; + double mtp_adaptive_accept = 0.0; // saves recurrent state before a speculative batch so it can be restored on rejection server_speculative_checkpoint spec_ckpt; diff --git a/examples/server/server.cpp b/examples/server/server.cpp index feaf1b4e..8151e1e5 100644 --- a/examples/server/server.cpp +++ b/examples/server/server.cpp @@ -1086,6 +1086,21 @@ int main(int argc, char ** argv) { const std::string oaicompat_model_name = requested_model_name.empty() ? fallback_model_name : requested_model_name; + + const auto infer_id_slot_from_model = [](const std::string & model) -> int { + const auto ends_with = [](const std::string & s, const std::string & suffix) -> bool { + return s.size() >= suffix.size() && s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0; + }; + + if (ends_with(model, "-slot0") || ends_with(model, "-s0")) { + return 0; + } + if (ends_with(model, "-slot1") || ends_with(model, "-s1")) { + return 1; + } + return -1; + }; + const int inferred_id_slot = infer_id_slot_from_model(oaicompat_model_name); for (size_t i = 0; i < inputs.size(); i++) { server_task task = server_task(type); @@ -1099,6 +1114,10 @@ int main(int argc, char ** argv) { // ctx_server.params, // data); task.id_slot = json_value(data, "id_slot", -1); + if (task.id_slot < 0 && inferred_id_slot >= 0 && inferred_id_slot < ctx_server.params_base.n_parallel) { + task.id_slot = inferred_id_slot; + task.data["id_slot"] = task.id_slot; + } // OAI-compat task.params.oaicompat = oaicompat; @@ -1254,18 +1273,33 @@ int main(int argc, char ** argv) { }; const auto handle_models = [¶ms, &model_meta](const httplib::Request & req, httplib::Response & res) { + (void) req; + + json data = json::array(); + + const auto add_model = [&](const std::string & id, const json & extra_meta = json::object()) { + json meta = model_meta; + for (const auto & kv : extra_meta.items()) { + meta[kv.key()] = kv.value(); + } + data.push_back({ + {"id", id}, + {"object", "model"}, + {"created", std::time(0)}, + {"owned_by", "llamacpp"}, + {"meta", meta}, + {"max_model_len", params.n_ctx}, + }); + }; + + add_model(params.model_alias); + for (int32_t i = 0; i < params.n_parallel; ++i) { + add_model(params.model_alias + std::string("-slot") + std::to_string(i), {{"slot_pinned", i}}); + } + json models = { {"object", "list"}, - {"data", { - { - {"id", params.model_alias}, - {"object", "model"}, - {"created", std::time(0)}, - {"owned_by", "llamacpp"}, - {"meta", model_meta}, - {"max_model_len", params.n_ctx}, //vllm specs - }, - }} + {"data", data}, }; res.set_content(models.dump(), "application/json; charset=utf-8"); diff --git a/include/llama.h b/include/llama.h index ac0a275b..5eb54e46 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1562,6 +1562,10 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns( LLAMA_API void llama_set_draft_input_hidden_state(struct llama_context * ctx, const float * hidden_state); + LLAMA_API void llama_set_mtp_draft_gpu_argmax(struct llama_context * ctx, bool enabled); + + LLAMA_API llama_token llama_get_mtp_draft_argmax_ith(struct llama_context * ctx, int32_t i); + #ifdef __cplusplus } #endif diff --git a/src/graphs/build_qwen35.cpp b/src/graphs/build_qwen35.cpp index fb19d679..8e3d939b 100644 --- a/src/graphs/build_qwen35.cpp +++ b/src/graphs/build_qwen35.cpp @@ -153,7 +153,11 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp( struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); - struct ggml_tensor * inp_out_ids = (n_outputs < n_tokens) ? build_inp_out_ids() : nullptr; + const bool mtp_cache_update_only = + cparams.mtp_op_type == MTP_OP_WARMUP || + cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED; + + struct ggml_tensor * inp_out_ids = (!mtp_cache_update_only && n_outputs < n_tokens) ? build_inp_out_ids() : nullptr; ggml_tensor * token_emb = build_inp_embd_mtp(model.tok_embd); @@ -210,11 +214,21 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp( cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "ffn_out", il); + if (mtp_cache_update_only) { + cb(cur, "result_mtp_cache_update", -1); + return cur; + } + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "result_norm", -1); cur = build_output(lctx, ctx0, cur, model.output, nullptr, cb); cb(cur, "result_output", -1); + if (lctx.mtp_draft_gpu_argmax && cparams.mtp_op_type == MTP_OP_DRAFT_GEN) { + cur = ggml_argmax(ctx0, cur); + cb(cur, "result_mtp_argmax", -1); + } + return cur; -} \ No newline at end of file +} diff --git a/src/llama-context.h b/src/llama-context.h index 7b6e56cf..d4f8ae19 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -264,6 +264,8 @@ struct llama_context { void * abort_callback_data = nullptr; const float * draft_input_hidden_state = nullptr; + bool mtp_draft_gpu_argmax = false; + std::vector mtp_draft_argmax; // input tensors struct ggml_tensor * inp_tokens; // I32 [n_batch] @@ -289,6 +291,7 @@ struct llama_context { struct Prev; std::unique_ptr prev; + std::unique_ptr prev_mtp; void reset_scheduler(); bool can_reuse_graph(const llama_batch & u_batch); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 7053952c..cae099cb 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -36,6 +36,20 @@ static inline const char * llm_expert_gating_func_name(llm_expert_gating_func_ty } } +static bool llm_detect_qwen35_recurrent_layer(const llama_model_loader & ml, uint32_t il, uint32_t fallback_interval) { + const std::string ssm_name = "blk." + std::to_string(il) + ".ssm_conv1d.weight"; + if (ml.get_tensor_meta(ssm_name.c_str()) != nullptr) { + return true; + } + + const std::string attn_q_name = "blk." + std::to_string(il) + ".attn_q.weight"; + if (ml.get_tensor_meta(attn_q_name.c_str()) != nullptr) { + return false; + } + + return ((il + 1) % fallback_interval != 0); +} + void llm_load_hparams( llama_model_loader & ml, @@ -507,7 +521,7 @@ void llm_load_hparams( uint32_t full_attn_interval = 4; ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); for (uint32_t i = 0; i < hparams.n_layer; ++i) { - hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); + hparams.recurrent_layer_arr[i] = llm_detect_qwen35_recurrent_layer(ml, i, full_attn_interval); } } @@ -546,7 +560,7 @@ void llm_load_hparams( const uint32_t n_main_layers = hparams.n_layer - hparams.nextn_predict_layers; for (uint32_t i = 0; i < hparams.n_layer; ++i) { if (i < n_main_layers) { - hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); + hparams.recurrent_layer_arr[i] = llm_detect_qwen35_recurrent_layer(ml, i, full_attn_interval); } else { hparams.recurrent_layer_arr[i] = false; } @@ -562,6 +576,10 @@ void llm_load_hparams( model.type = hparams.n_embd == 2560 ? e_model::MODEL_4B : e_model::MODEL_9B; break; case 64: // without MTP layer case 65: // with MTP layer (64 main + 1 MTP) + case 67: // RYS 11-14 without MTP layer + case 68: // RYS 11-14 with MTP layer + case 69: // RYS 15-20 without MTP layer + case 70: // RYS 15-20 with MTP layer model.type = e_model::MODEL_27B; break; default: model.type = e_model::MODEL_UNKNOWN; } diff --git a/src/llama.cpp b/src/llama.cpp index f7b55bbf..b8661462 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -548,36 +548,63 @@ struct llama_context::Prev { int all_seq_id; int n_outputs; int n_kv; + int n_tokens; llama_mtp_op_type mtp_op_type; + bool mtp_draft_gpu_argmax; ggml_cgraph * graph; }; void llama_context::reset_scheduler() { ggml_backend_sched_reset(sched); prev.reset(); + prev_mtp.reset(); } bool llama_context::can_reuse_graph(const llama_batch & u_batch) { - if (!prev || !prev->graph) return false; - if (u_batch.n_tokens > 1) return false; - if (u_batch.embd) return false; if (!cparams.graph_reuse) return false; - return u_batch.all_seq_id == prev->all_seq_id && + auto the_prev = cparams.mtp_op_type == MTP_OP_NONE ? prev.get() : prev_mtp.get(); + if (!the_prev || !the_prev->graph) return false; + //if (u_batch.n_tokens > 1) return false; + if (u_batch.embd) return false; + return u_batch.all_seq_id == the_prev->all_seq_id && kv_self.head > 0 && - kv_self.n == prev->n_kv && - n_outputs == prev->n_outputs && - cparams.mtp_op_type == prev->mtp_op_type && + kv_self.n == the_prev->n_kv && + n_outputs == the_prev->n_outputs && + u_batch.n_tokens == the_prev->n_tokens && + cparams.mtp_op_type == the_prev->mtp_op_type && + mtp_draft_gpu_argmax == the_prev->mtp_draft_gpu_argmax && update_cache_copies(); } +/* +static void why_not_reuse_previous(const llama_batch & u_batch, const llama_context & ctx, const llama_context::Prev * the_prev) { + if (!the_prev) { printf(" previous is null\n"); return; } + if (!the_prev->graph) { printf(" previous graph is null\n"); return; } + if (!ctx.cparams.graph_reuse) { printf(" graph_reuse is false\n"); return; } + if (u_batch.embd) { printf(" ubatch.embd is not null\n"); return; } + if (u_batch.all_seq_id != the_prev->all_seq_id) { printf(" all_seq_id is not the same\n"); return; } + if (ctx.kv_self.head == 0) { printf(" kv_self.head = 0\n"); return; } + if (ctx.kv_self.n != the_prev->n_kv) { printf(" kv_self.n is not the same\n"); return; } + if (ctx.n_outputs != the_prev->n_outputs) { printf(" n_outputs is not the same\n"); return; } + if (u_batch.n_tokens != the_prev->n_tokens) { printf(" n_tokens is not the same\n"); return; } + if (ctx.cparams.mtp_op_type != the_prev->mtp_op_type) { printf(" mtp_op_type is not the same\n"); return; } + printf(" update_cache_copies() must have failed\n"); +} +*/ + bool llama_context::update_cache_copies() { - const int n_layer = model.mtp ? model.hparams.n_layer - : model.hparams.n_layer - model.hparams.nextn_predict_layers; //cache_copies.size()/2; + const int n_layer = model.mtp && cparams.mtp_op_type != MTP_OP_NONE ? + model.hparams.n_layer : model.hparams.n_layer - model.hparams.nextn_predict_layers; //cache_copies.size()/2; auto layer_has_attention_kv = [&](int il) { return !model.hparams.is_recurrent(il); }; - if ((int)kv_self.k_l.size() != n_layer) return false; - if (!(kv_self.v_l.empty() || (int)kv_self.v_l.size() == n_layer)) return false; + + if ((int)kv_self.k_l.size() < n_layer) { + return false; + } + if (!kv_self.v_l.empty() && (int)kv_self.v_l.size() < n_layer) { + return false; + } for (int il = 0; il < n_layer; ++il) { if (!layer_has_attention_kv(il) || kv_self.k_l[il] == nullptr) { continue; @@ -594,7 +621,9 @@ bool llama_context::update_cache_copies() { for (int id = 0; id < kl->n_device; ++id) { if (!kl->splits[id]) continue; auto& c = cache_copies[2*model.splits.size()*il + 2*id + 0]; - if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kl->splits[id]) return false; + if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kl->splits[id]) { + return false; + } c.cpy->view_offs = kv_self.head*c.step; c.cpy->src[1]->data = (char *)kl->splits[id]->data + c.cpy->view_offs; c.cpy->data = c.cpy->src[1]->data; @@ -603,29 +632,26 @@ bool llama_context::update_cache_copies() { for (int id = 0; id < vl->n_device; ++id) { if (!vl->splits[id]) continue; auto& c = cache_copies[2*model.splits.size()*il + 2*id + 1]; - if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != vl->splits[id]) return false; + if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != vl->splits[id]) { + return false; + } c.cpy->view_offs = kv_self.head*c.step; c.cpy->src[1]->data = (char *)vl->splits[id]->data + c.cpy->view_offs; c.cpy->data = c.cpy->src[1]->data; } } else { - for (int il = 0; il < n_layer; ++il) { - if (!layer_has_attention_kv(il) || kv_self.k_l[il] == nullptr) { - continue; - } - auto& c = cache_copies[2*il+0]; - if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kv_self.k_l[il]) return false; - c.cpy->view_offs = kv_self.head*c.step; - c.cpy->src[1]->data = (char *)kv_self.k_l[il]->data + c.cpy->view_offs; - c.cpy->data = c.cpy->src[1]->data; + auto& c = cache_copies[2*il+0]; + if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kv_self.k_l[il]) { + return false; } - if (kv_self.v_l.empty()) return true; - for (int il = 0; il < n_layer; ++il) { - if (!layer_has_attention_kv(il) || kv_self.v_l[il] == nullptr) { - continue; - } + c.cpy->view_offs = kv_self.head*c.step; + c.cpy->src[1]->data = (char *)kv_self.k_l[il]->data + c.cpy->view_offs; + c.cpy->data = c.cpy->src[1]->data; + if (!kv_self.v_l.empty() && kv_self.v_l[il]) { auto& c = cache_copies[2*il+1]; - if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kv_self.v_l[il]) return false; + if (!c.cpy || c.cpy->op != GGML_OP_CPY || c.cpy->view_src != kv_self.v_l[il]) { + return false; + } c.cpy->view_offs = kv_self.head*c.step; c.cpy->src[1]->data = (char *)kv_self.v_l[il]->data + c.cpy->view_offs; c.cpy->data = c.cpy->src[1]->data; @@ -1516,7 +1542,7 @@ bool llama_kv_cache::per_step_alloc(int max_tokens) { } bool llama_kv_cache::per_step_restore(int step) { - if (ckpt.per_step_ssm.empty() || step < 0) { + if (ckpt.per_step_ssm.empty() || step < 0 || step >= ckpt.per_step_max_allocated) { return false; } @@ -1582,7 +1608,7 @@ bool llama_kv_cache::per_step_restore(int step) { n_restored++; } - return true; + return n_restored > 0; } static void llama_kv_cache_clear(struct llama_kv_cache & cache) { @@ -4084,6 +4110,27 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { const auto & cparams = lctx.cparams; const auto & hparams = lctx.model.hparams; + const bool has_mtp = lctx.model.hparams.nextn_predict_layers > 0 && lctx.cparams.mtp; + const bool mtp_cache_update_only = + has_mtp && + (cparams.mtp_op_type == MTP_OP_WARMUP || + cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED); + if (mtp_cache_update_only && n_outputs == 0) { + lctx.mtp_draft_argmax.clear(); + if (lctx.output_ids.empty()) { + lctx.output_ids.resize(cparams.n_batch); + } + + lctx.logits = nullptr; + lctx.embd = nullptr; + lctx.output_size = 0; + lctx.logits_size = 0; + lctx.embd_size = 0; + std::fill(lctx.output_ids.begin(), lctx.output_ids.end(), -1); + lctx.n_outputs = 0; + + return 0; + } const size_t n_outputs_max = std::max(n_outputs, (size_t) cparams.n_seq_max); @@ -4091,9 +4138,11 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { const auto n_vocab = hparams.n_vocab; const auto n_embd = hparams.n_embd; + const bool mtp_draft_gpu_argmax = + has_mtp && cparams.mtp_op_type == MTP_OP_DRAFT_GEN && lctx.mtp_draft_gpu_argmax; + // TODO: use a per-batch flag for logits presence instead - const bool has_mtp = lctx.model.hparams.nextn_predict_layers > 0 && lctx.cparams.mtp; - const bool has_logits = !cparams.embeddings || has_mtp; + const bool has_logits = !mtp_draft_gpu_argmax && (!cparams.embeddings || has_mtp); const bool has_embd = lctx.is_encoding || (cparams.embeddings && (cparams.pooling_type == LLAMA_POOLING_TYPE_NONE)) || has_mtp; const size_t logits_size = has_logits ? n_vocab*n_outputs_max : 0; @@ -4140,6 +4189,12 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { // set all ids as invalid (negative) std::fill(lctx.output_ids.begin(), lctx.output_ids.end(), -1); + if (mtp_draft_gpu_argmax) { + lctx.mtp_draft_argmax.assign(n_outputs_max, LLAMA_TOKEN_NULL); + } else { + lctx.mtp_draft_argmax.clear(); + } + if (has_mtp) { // MTP uses a large output footprint, clear only the active region. const size_t clear_size = (logits_size + embd_size) * sizeof(float); @@ -4254,9 +4309,16 @@ static int llama_decode_internal( // this indicates we are doing pooled embedding, so we ignore batch.logits and output all tokens const bool embd_pooled = cparams.embeddings && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE; const bool has_mtp = cparams.mtp && hparams.nextn_predict_layers > 0; - + const bool mtp_cache_update_only = + has_mtp && + (cparams.mtp_op_type == MTP_OP_WARMUP || + cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED); + const bool mtp_draft_gpu_argmax = + has_mtp && cparams.mtp_op_type == MTP_OP_DRAFT_GEN && lctx.mtp_draft_gpu_argmax; // count outputs - if (batch_all.logits && !embd_pooled) { + if (mtp_cache_update_only) { + n_outputs = 0; + } else if (batch_all.logits && !embd_pooled) { for (uint32_t i = 0; i < n_tokens_all; ++i) { n_outputs += batch_all.logits[i] != 0; } @@ -4268,7 +4330,7 @@ static int llama_decode_internal( } // reserve output buffer - n_outputs_embd = has_mtp ? n_tokens_all : n_outputs; + n_outputs_embd = mtp_cache_update_only ? 0 : (has_mtp ? n_tokens_all : n_outputs); if (llama_output_reserve(lctx, std::max(n_outputs, n_outputs_embd)) < std::max(n_outputs, n_outputs_embd)) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %zu outputs\n", __func__, std::max(n_outputs, n_outputs_embd)); return -2; @@ -4357,7 +4419,9 @@ static int llama_decode_internal( { int32_t n_outputs_new = 0; - if (u_batch.logits && !embd_pooled) { + if (mtp_cache_update_only) { + n_outputs_new = 0; + } else if (u_batch.logits && !embd_pooled) { for (uint32_t i = 0; i < n_tokens; i++) { n_outputs_new += u_batch.logits[i] != 0; } @@ -4438,21 +4502,15 @@ static int llama_decode_internal( printf("prelude(...): %d us\n", int(tim2-tim1)); #endif - - //if (n_tokens_all == 1) { - // printf("================= %s\n", __func__); - // printf(" all_pos_0 = %d, all_pos_1 = %d, all_seq_id = %d\n", batch_all.all_pos_0, batch_all.all_pos_1, batch_all.all_seq_id); - // printf(" embd = %p, logits = %p, token = %p\n", (const void *)batch_all.embd, (const void *)batch_all.logits, (const void *)batch_all.token); - // printf(" n_outputs = %d, kv_self.n = %d\n", n_outputs, kv_self.n); - //} - //printf("kv_self.n = %5d, kv_self.used = %5d, kv_self.head = %5d\n", kv_self.n, kv_self.used, kv_self.head); - #if IK_PRINT_TIMING tim1 = ggml_time_us(); #endif + auto & prev = cparams.mtp_op_type == MTP_OP_NONE ? lctx.prev : lctx.prev_mtp; ggml_cgraph * gf = nullptr; if (!lctx.can_reuse_graph(u_batch)) { - lctx.reset_scheduler(); + //lctx.reset_scheduler(); + ggml_backend_sched_reset(lctx.sched); + prev.reset(); ggml_backend_sched_set_eval_callback(lctx.sched, lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data); #if IK_PRINT_TIMING tim2 = ggml_time_us(); @@ -4476,14 +4534,15 @@ static int llama_decode_internal( tim2 = ggml_time_us(); printf("sched_alloc_graph(...): %d us\n", int(tim2-tim1)); #endif - if (u_batch.n_tokens == 1 && u_batch.embd == nullptr && lctx.cparams.graph_reuse) { - lctx.prev = std::make_unique(llama_context::Prev{ + //if (u_batch.n_tokens == 1 && u_batch.embd == nullptr && lctx.cparams.graph_reuse) { + if (u_batch.embd == nullptr && lctx.cparams.graph_reuse) { + prev = std::make_unique(llama_context::Prev{ (int)u_batch.all_seq_id, (int)lctx.n_outputs, (int)lctx.kv_self.n, - cparams.mtp_op_type, gf}); + (int)u_batch.n_tokens, cparams.mtp_op_type, lctx.mtp_draft_gpu_argmax, gf}); } } else { - //printf("Reusing graph\n"); - gf = lctx.prev->graph; + //printf("Reusing graph with n_kv = %d, n_tokens = %d\n", (int)prev->n_kv, (int)prev->n_tokens); + gf = prev->graph; } if (cparams.mtp_op_type != MTP_OP_NONE) { @@ -4495,6 +4554,7 @@ static int llama_decode_internal( // the output is always the last tensor in the graph struct ggml_tensor * res = gf->nodes[gf->n_nodes - 1]; struct ggml_tensor * embd = nullptr; + struct ggml_tensor * mtp_argmax = nullptr; if (lctx.n_outputs == 0) { // no output @@ -4505,6 +4565,9 @@ static int llama_decode_internal( const bool use_qwen_mtp_embd = has_mtp && lctx.model.arch == LLM_ARCH_QWEN35; if (cparams.embeddings || has_mtp) { for (int i = gf->n_nodes - 1; i >= 0; --i) { + if (mtp_draft_gpu_argmax && strcmp(gf->nodes[i]->name, "result_mtp_argmax") == 0) { + mtp_argmax = gf->nodes[i]; + } if (use_qwen_mtp_embd && strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) { // Qwen 3.5 uses raw hidden state before the final shared-head normalization. embd = gf->nodes[i]; @@ -4526,6 +4589,10 @@ static int llama_decode_internal( GGML_ASSERT(strcmp(res->name, "result_output") == 0 && "missing result_output tensor"); } } + if (mtp_draft_gpu_argmax) { + GGML_ASSERT(mtp_argmax != nullptr && "missing MTP draft argmax tensor"); + res = nullptr; + } } // LLAMA_LOG_INFO("graph build time: %.3f ms (%d nodes, %d leafs)\n", (ggml_time_us() - t_start_us)/1000.0, gf->n_nodes, gf->n_leafs); #if IK_PRINT_TIMING == 1 @@ -4566,12 +4633,26 @@ static int llama_decode_internal( // ggml_graph_dump_dot(gf, NULL, "llama.dot"); //} + if (mtp_argmax) { + ggml_backend_t backend_argmax = ggml_backend_sched_get_tensor_backend(lctx.sched, mtp_argmax); + GGML_ASSERT(backend_argmax != nullptr); + + const int32_t n_outputs_new = lctx.n_outputs; + if (n_outputs_new) { + GGML_ASSERT(n_outputs_prev + n_outputs_new <= n_outputs); + if (lctx.mtp_draft_argmax.size() < n_outputs) { + lctx.mtp_draft_argmax.resize(n_outputs, LLAMA_TOKEN_NULL); + } + llama_token * argmax_out = lctx.mtp_draft_argmax.data() + n_outputs_prev; + ggml_backend_tensor_get_async(backend_argmax, mtp_argmax, argmax_out, 0, n_outputs_new*sizeof(llama_token)); + } + } + // extract logits if (res) { #if IK_PRINT_TIMING tim1 = ggml_time_us(); #endif - // Do not process logits if MTP is only updating the KV cache. if (cparams.mtp_op_type != MTP_OP_WARMUP && cparams.mtp_op_type != MTP_OP_UPDATE_ACCEPTED) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(lctx.sched, res); @@ -4609,7 +4690,7 @@ static int llama_decode_internal( } // extract embeddings - if (embd && (cparams.mtp_op_type == MTP_OP_NONE || cparams.mtp_op_type == MTP_OP_DRAFT_GEN)) { + if (embd && (cparams.mtp_op_type == MTP_OP_NONE || cparams.mtp_op_type == MTP_OP_DRAFT_GEN)) { #if IK_PRINT_TIMING tim1 = ggml_time_us(); #endif @@ -6983,7 +7064,7 @@ bool llama_spec_ckpt_restore(struct llama_context * ctx, llama_seq_id seq_id, return false; } const llama_pos accepted_pos = n_past + accepted_step; - if (seq_id >= 0 && (uint32_t)seq_id < kv.size) { + if (kv.recurrent && seq_id >= 0 && (uint32_t)seq_id < kv.size) { kv.cells[seq_id].pos = accepted_pos; } llama_kv_cache_seq_rm(kv, seq_id, accepted_pos + 1, -1); @@ -6991,7 +7072,9 @@ bool llama_spec_ckpt_restore(struct llama_context * ctx, llama_seq_id seq_id, } case LLAMA_SPEC_CKPT_GPU_FALLBACK: - kv.checkpoint_restore(); + if (!kv.checkpoint_restore()) { + return false; + } llama_kv_cache_seq_rm(kv, seq_id, n_past, -1); return false; @@ -8494,6 +8577,47 @@ void llama_set_mtp_op_type(llama_context * ctx, llama_mtp_op_type mtp_op_type) { ctx->set_mtp_op_type(mtp_op_type); } +void llama_set_mtp_draft_gpu_argmax(struct llama_context * ctx, bool enabled) { + if (ctx->mtp_draft_gpu_argmax != enabled) { + ctx->mtp_draft_gpu_argmax = enabled; + ctx->prev_mtp.reset(); + } +} + +llama_token llama_get_mtp_draft_argmax_ith(struct llama_context * ctx, int32_t i) { + int32_t j = -1; + llama_synchronize(ctx); + + try { + if (ctx->mtp_draft_argmax.empty()) { + throw std::runtime_error("no MTP draft argmax output"); + } + + if (i < 0) { + j = ctx->n_outputs + i; + if (j < 0) { + throw std::runtime_error(format("negative index out of range [0, %d)", ctx->n_outputs)); + } + } else if ((size_t) i >= ctx->output_ids.size()) { + throw std::runtime_error(format("out of range [0, %lu)", ctx->output_ids.size())); + } else { + j = ctx->output_ids[i]; + } + + if (j < 0) { + throw std::runtime_error(format("batch.logits[%d] != true", i)); + } + if (j >= ctx->n_outputs || (size_t) j >= ctx->mtp_draft_argmax.size()) { + throw std::runtime_error(format("corrupt MTP argmax buffer (j=%d, n_outputs=%d)", j, ctx->n_outputs)); + } + + return ctx->mtp_draft_argmax[j]; + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: invalid MTP argmax id %d, reason: %s\n", __func__, i, err.what()); + return LLAMA_TOKEN_NULL; + } +} + void llama_synchronize(struct llama_context * ctx) { ggml_backend_sched_synchronize(ctx->sched);