Upload qwen36-mtp-rys_delta.patch with huggingface_hub

qwen36-mtp-rys_delta.patch

@@ -0,0 +1,1292 @@
+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 <utility>
+ #include <vector>
+ #include <fstream>
++#include <cctype>
++#include <unordered_set>
+ 
+ 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<unsigned char>(s[start]))) {
++        start++;
++    }
++    size_t end = s.size();
++    while (end > start && std::isspace(static_cast<unsigned char>(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<common_chat_tool_call> & tool_calls, bool is_partial, const common_chat_parser_params & params) {
++    if (!params.parse_tool_calls || tool_calls.size() < 2) {
++        return;
++    }
++
++    std::unordered_set<std::string> seen;
++    seen.reserve(tool_calls.size());
++
++    std::vector<common_chat_tool_call> 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<common_chat_peg_mapper>(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 <algorithm>
+ #include <cstring>
++#include <cstdlib>
+ #include <iomanip>
+ #include <map>
+ 
+@@ -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<llama_token> mtp_speculative_gen_draft(
+     struct common_sampler * smpl,
+     struct llama_context * ctx,
+@@ -1370,6 +1376,9 @@ std::vector<llama_token> 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<llama_token> 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<llama_token> 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<std::thread> 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<float> logits;
++        if (params.compute_ppl && collect_mtp) {
++            logits.reserve((size_t)n_ctx * n_vocab);
++        }
++        std::vector<float> mtp_hidden_states;
++        std::vector<llama_token> 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<int>(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<int64_t>(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<int64_t>(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<float> 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 = [&params, &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;
+-}
+
++}
+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<llama_token> 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> prev;
++    std::unique_ptr<Prev> 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<size_t>(n_outputs, n_outputs_embd)) < std::max<size_t>(n_outputs, n_outputs_embd)) {
+         LLAMA_LOG_ERROR("%s: could not reserve space for batch with %zu outputs\n", __func__, std::max<size_t>(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>(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>(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);
+