From b85ca9e8991f1c3fa6c85539741d6159ab74f3b2 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Mon, 5 Oct 2026 20:20:29 +0200 Subject: [PATCH] llama : re-reserve the sched when the nextn extraction flags change The speculative MTP init enables NextN extraction on the target and draft contexts after both were created and their schedulers reserved. With unmasked extraction the trunk graph keeps every token through the last layer instead of cropping to the output rows, so the first decode reallocates to that batch's shape and the next, wider batch trips GGML_SCHED_DEBUG_REALLOC. Invalidate the reserve when the flags change so the next compute re-reserves with the new graph shape. Assisted-by: Claude --- src/llama-context.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 07b4c6148eb5..830a5b364d5f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1232,6 +1232,11 @@ void llama_context::set_embeddings(bool value) { void llama_context::set_embeddings_nextn(bool value, bool masked) { LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked); + if (cparams.embeddings_nextn != value || cparams.embeddings_nextn_masked != masked) { + // these flags change the graph shape + sched_need_reserve = true; + } + cparams.embeddings_nextn = value; cparams.embeddings_nextn_masked = masked; }