qwen36-mtp-rys_delta.patch

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;
-}
\ 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<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);