From a009e04e8a15b20feb1f5623a1ca6a377995bfe0 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 12:23:23 +0300 Subject: [PATCH 1/7] metal : cleanup --- ggml/src/ggml-metal/ggml-metal.metal | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index deb616105ae9..82e29d5ad7cd 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2531,7 +2531,7 @@ kernel void kernel_rwkv_wkv7_f32( constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]]; constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]]; -constant short FC_gated_delta_net_K [[function_constant(FC_GATED_DELTA_NET + 2)]]; +constant short FC_gated_delta_net_K [[function_constant(FC_GATED_DELTA_NET + 2)]]; #if 1 template @@ -2549,6 +2549,7 @@ kernel void kernel_gated_delta_net_impl( uint3 ntg[[threads_per_threadgroup]]) { #define S_v FC_gated_delta_net_ne20 #define G FC_gated_delta_net_ne30 +#define K FC_gated_delta_net_K const uint tx = tpitg.x; const uint ty = tpitg.y; @@ -2562,8 +2563,6 @@ kernel void kernel_gated_delta_net_impl( const float scale = 1.0f / sqrt((float)S_v); - const uint K = FC_gated_delta_net_K; - // input state layout (D, K, n_seqs): per-seq stride is K*H*D; we read slot 0. // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous const uint state_in_base = (i23*K*args.ne21 + i21)*S_v*S_v + i20*S_v; @@ -2666,6 +2665,7 @@ kernel void kernel_gated_delta_net_impl( #undef S_v #undef G +#undef K } typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t; From 5122abb537b3369804941552d0c167cf23c9bff5 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 12:57:58 +0300 Subject: [PATCH 2/7] llama : fix faulty bitwise check in recurrent memory --- src/llama-memory-recurrent.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 084c5d9ea4f3..852501b5c60e 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -743,7 +743,7 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq cell_ranges.emplace_back(cell_range_begin, size); } - if (flags % LLAMA_STATE_SEQ_FLAGS_ON_DEVICE && cell_ranges.size() > 1) { + if ((flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) && cell_ranges.size() > 1) { GGML_ABORT("cannot save/load multiple ranges of cells to/from device memory\n"); } From 1e2a77ae065419a443983f44df864c5cbbf719c4 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 12:59:57 +0300 Subject: [PATCH 3/7] server : disable RS-based MTP in combination with other spec types --- src/llama-memory-recurrent.cpp | 9 +++++++++ tools/server/server-context.cpp | 29 +++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 852501b5c60e..64c7e726fdca 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -719,6 +719,15 @@ size_t llama_memory_recurrent::size_s_bytes() const { void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); + // [TAG_RS_STATE_ROLLBACK_SUPPORT] + if (n_rs_seq != 0) { + for (uint32_t i = 0; i < rs_idx.size(); ++i) { + if (rs_idx[i] != 0) { + GGML_ABORT("recurrent state read/write is not supported with partial rollback"); + } + } + } + std::vector> cell_ranges; // ranges, from inclusive, to exclusive uint32_t cell_count = 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index df6ab623d775..a643d872f7d8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -750,6 +750,35 @@ struct server_context_impl { add_bos_token = llama_vocab_get_add_bos(vocab); + // [TAG_RS_STATE_ROLLBACK_SUPPORT] + // TODO: ngram speculative methods require checkpointing in addition to partial RS rollback + // currently this is not supported. so we enable partial rollback only when MTP runs on its own + // in this case, we don't have to do checkpoints to perform speculative decoding + if (llama_n_rs_seq(ctx_tgt) > 0) { + bool print_warning = true; + + auto & types = params_base.speculative.types; + + // filter all non-MTP speculative types + for (int i = 0; i < (int) params_base.speculative.types.size(); i++) { + if (types[i] == COMMON_SPECULATIVE_TYPE_NONE) { + continue; + } + if (types[i] != COMMON_SPECULATIVE_TYPE_DRAFT_MTP) { + if (print_warning && !types.empty()) { + SRV_WRN("%s", "partial recurrent rollback is enabled, only MTP speculative decoding is supported\n"); + } + + SRV_WRN("removing speculative type %d (%s) from speculative types\n", + types[i], common_speculative_type_to_str(types[i]).c_str()); + types.erase(types.begin() + i); + + print_warning = false; + i--; + } + } + } + if (params_base.speculative.has_dft()) { // TODO speculative: move to common/speculative.cpp? const auto & params_spec = params_base.speculative.draft; From 931754806812f68c1cfd4aaf5ba535a12120738f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 13:06:54 +0300 Subject: [PATCH 4/7] spec : add TODOs --- common/common.h | 2 +- common/speculative.cpp | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/common/common.h b/common/common.h index 222f3177b243..dda4973bb7c7 100644 --- a/common/common.h +++ b/common/common.h @@ -303,7 +303,7 @@ struct common_params_speculative_draft { int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding float p_split = 0.1f; // speculative decoding split probability - float p_min = 0.75f; // minimum speculative decoding probability (greedy) + float p_min = 0.75f; // minimum speculative decoding probability (greedy) // TODO: change default to 0.0f common_params_model mparams; diff --git a/common/speculative.cpp b/common/speculative.cpp index 0a78bc1136c7..a56a705e2e8d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -424,7 +424,7 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { for (auto & s : smpls) { common_params_sampling sparams; sparams.no_perf = false; - sparams.top_k = 1; + sparams.top_k = 1; // TODO: re-enable top_k == 10 and utilize `p_min` spec param sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams)); } @@ -1494,6 +1494,10 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, u GGML_ASSERT(impl); + // TODO: currently only the implementation that generated the draft is used to accept it + // however, some implementations (such as MTP) need to also "see" the accepted tokens + // extend `common_speculative_impl::accept()` with an extra argument `bool is_other` to + // inform the implementation if the accepted tokens are from another implementation { common_time_meas tm(impl->t_accept_us, !impl->gen_perf); if (n_accepted > 0) { From 1271f826c23402d023d9b98269d76c6b04ef60de Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 13:10:33 +0300 Subject: [PATCH 5/7] cont : fix comment --- tools/server/server-context.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a643d872f7d8..21e623d227d1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -752,8 +752,7 @@ struct server_context_impl { // [TAG_RS_STATE_ROLLBACK_SUPPORT] // TODO: ngram speculative methods require checkpointing in addition to partial RS rollback - // currently this is not supported. so we enable partial rollback only when MTP runs on its own - // in this case, we don't have to do checkpoints to perform speculative decoding + // currently this is not supported. so we remove the non-MTP types from the list if (llama_n_rs_seq(ctx_tgt) > 0) { bool print_warning = true; From baabc98cf438cb9d11d6ecd67d56d8b7da9a544b Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 13:27:00 +0300 Subject: [PATCH 6/7] cont : update comment --- common/speculative.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index a56a705e2e8d..3488b9393c5a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1497,7 +1497,8 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, u // TODO: currently only the implementation that generated the draft is used to accept it // however, some implementations (such as MTP) need to also "see" the accepted tokens // extend `common_speculative_impl::accept()` with an extra argument `bool is_other` to - // inform the implementation if the accepted tokens are from another implementation + // inform the implementation if the accepted tokens are from another implementation and + // pass the accepted tokens to all remaining implementations using `is_other == true` { common_time_meas tm(impl->t_accept_us, !impl->gen_perf); if (n_accepted > 0) { From 5c11caea63a01393b73af82caef90b46928d6c09 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 15 May 2026 13:47:39 +0300 Subject: [PATCH 7/7] common : fix logic for ngram + mtp compat --- common/common.cpp | 24 ++++++++++++++++++++++++ tools/server/server-context.cpp | 28 ---------------------------- 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 9f597b4feb80..8b6d182f5495 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -7,6 +7,7 @@ #include "log.h" #include "llama.h" #include "sampling.h" +#include "speculative.h" #include "unicode.h" #include @@ -1247,6 +1248,29 @@ common_init_result::common_init_result(common_params & params) : cparams.n_samplers = pimpl->samplers_seq_config.size(); } + // [TAG_RS_STATE_ROLLBACK_SUPPORT] + // TODO: ngram speculative methods require checkpointing in addition to partial RS rollback + // currently this is not supported. so we disable the partial rollback + if (cparams.n_rs_seq > 0 && (llama_model_is_recurrent(model) || llama_model_is_hybrid(model))) { + auto & types = params.speculative.types; + + for (int i = 0; i < (int) types.size(); i++) { + if (types[i] == COMMON_SPECULATIVE_TYPE_NONE) { + continue; + } + if (types[i] == COMMON_SPECULATIVE_TYPE_DRAFT_MTP) { + continue; + } + + cparams.n_rs_seq = 0; + + LOG_WRN("%s: recurrent state rollback is not compatible with '%s' - disabling rollback support\n", __func__, + common_speculative_type_to_str(types[i]).c_str()); + + break; + } + } + llama_context * lctx = llama_init_from_model(model, cparams); if (lctx == NULL) { LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 21e623d227d1..df6ab623d775 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -750,34 +750,6 @@ struct server_context_impl { add_bos_token = llama_vocab_get_add_bos(vocab); - // [TAG_RS_STATE_ROLLBACK_SUPPORT] - // TODO: ngram speculative methods require checkpointing in addition to partial RS rollback - // currently this is not supported. so we remove the non-MTP types from the list - if (llama_n_rs_seq(ctx_tgt) > 0) { - bool print_warning = true; - - auto & types = params_base.speculative.types; - - // filter all non-MTP speculative types - for (int i = 0; i < (int) params_base.speculative.types.size(); i++) { - if (types[i] == COMMON_SPECULATIVE_TYPE_NONE) { - continue; - } - if (types[i] != COMMON_SPECULATIVE_TYPE_DRAFT_MTP) { - if (print_warning && !types.empty()) { - SRV_WRN("%s", "partial recurrent rollback is enabled, only MTP speculative decoding is supported\n"); - } - - SRV_WRN("removing speculative type %d (%s) from speculative types\n", - types[i], common_speculative_type_to_str(types[i]).c_str()); - types.erase(types.begin() + i); - - print_warning = false; - i--; - } - } - } - if (params_base.speculative.has_dft()) { // TODO speculative: move to common/speculative.cpp? const auto & params_spec = params_base.speculative.draft;