From e4acba83d669494c7dab729ff79e3d24b9a5964d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 4 Sep 2026 01:29:43 +0000 Subject: [PATCH 1/4] server: keep a speculative block inside one sub-batch, and shift its indices Fixes #24840. When a decode fails for want of KV cache, update_slots halves n_batch and retries the same offset with a smaller view. The non-speculative path accounts for that (`tok_idx = slot.i_batch - off`), but the speculative path passed slot.spec_i_batch to common_sampler_sample_and_accept_n unshifted, so with a non-zero offset it addressed the wrong logits. Rather than sample wrongly, post_decode threw, which aborts every slot in flight: speculative batch index 4 is not inside the current sub-batch [0, 4) Three parts, because shifting alone is not enough. A slot's spec_i_batch entries are contiguous and their logits must all come from ONE decode, so a view that cuts through a block leaves half of them in a decode that has already happened. The view loop now ends a view just before a block it would otherwise bisect, which costs one extra decode call and keeps every block whole. post_decode then treats a block that is not in this view as ordinary, because after the change it belongs entirely to another view and will be sampled when that view is decoded. It still throws for a block that is split anyway, which needs n_batch below n_draft + 1 and so is only reachable at the bottom of the retry ladder. The speculative sampling loop skips a slot whose block is not in this view; without that the shift below it produces negative indices. Finally the indices are shifted into view-local space before sampling, which is the part the issue describes. Reached in practice whenever several long conversations share one cache under --parallel N --kv-unified: the KV-full retry is common there, and every occurrence killed every chat on the server rather than one. --- tools/server/server-context.cpp | 85 ++++++++++++++++++++++++++++++--- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..c82ba621518e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2752,7 +2752,39 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); for (int32_t off = 0; off < batch.size(); off = off_next) { - const int32_t n_tokens = std::min(n_batch, batch.size() - off); + int32_t n_tokens = std::min(n_batch, batch.size() - off); + + // Do not let the view end in the middle of a slot's speculative block. + // + // A slot's spec_i_batch entries are contiguous (see the push_back loop where + // they are built), and common_sampler_sample_and_accept_n needs the logits + // for all of them from ONE decode. A view that cuts through a block leaves + // half its logits in a decode that has already happened, which is why + // post_decode had no option but to abort the slots. Ending the view just + // before the block instead costs one extra decode call and keeps the block + // whole. See https://github.com/ggml-org/llama.cpp/issues/24840 + if (n_tokens < batch.size() - off) { + const int32_t view_end = off + n_tokens; + int32_t cut = view_end; + iterate(slots, [&](server_slot & slot) { + if (slot.spec_i_batch.empty()) { + return; + } + const int32_t first = slot.spec_i_batch.front(); + const int32_t last = slot.spec_i_batch.back(); + // straddles the end of the view: stop short of it + if (first > off && first < view_end && last >= view_end) { + cut = std::min(cut, first); + } + }); + n_tokens = cut - off; + // A block longer than the whole view cannot be kept whole by cutting. + // That needs n_batch below n_draft + 1, which only the retry ladder + // reaches; fall back to the original view and let post_decode report it. + if (n_tokens <= 0) { + n_tokens = std::min(n_batch, batch.size() - off); + } + } try { scoped_timer t(t_decode, n_decode); // TODO @ngxson : maybe handle n_batch == 1 here instead of inside decode() @@ -3671,14 +3703,36 @@ struct server_context_impl { return idx >= off && idx < off + n_batch_tokens; }; - // TODO @ngxson : it's tricky to make sub-batch compatible with common_sampler_sample_and_accept_n, - // so for now we will throw an error in this case: https://github.com/ggml-org/llama.cpp/issues/24840 - iterate(slots, [&](server_slot & slot) { - for (auto & i : slot.spec_i_batch) { + // A slot whose speculative block is not in THIS view is not an error: the batch + // was split, and the block belongs to another view, which will sample it when it + // is decoded. The loop that builds the views keeps each block whole, so a block + // is either entirely inside or entirely outside. + // + // Aborting every slot here instead is https://github.com/ggml-org/llama.cpp/issues/24840, + // reached whenever a KV-full decode halves n_batch, which is exactly when several + // long conversations share one cache. + auto spec_is_inside_view = [&](const server_slot & slot) { + for (const auto & i : slot.spec_i_batch) { if (!is_inside_view(i)) { - throw std::runtime_error(string_format("speculative batch index %d is not inside the current sub-batch [%d, %d)", i, off, off + n_batch_tokens)); + return false; } } + return true; + }; + iterate(slots, [&](server_slot & slot) { + if (slot.spec_i_batch.empty() || spec_is_inside_view(slot)) { + return; + } + const int32_t first = slot.spec_i_batch.front(); + const int32_t last = slot.spec_i_batch.back(); + if (first >= off && last < off + n_batch_tokens) { + return; + } + if (first < off + n_batch_tokens && last >= off + n_batch_tokens && first >= off) { + // Split anyway: only reachable when the retry ladder drove n_batch below + // one block. Report it rather than sample from logits that do not exist. + throw std::runtime_error(string_format("speculative batch index %d is not inside the current sub-batch [%d, %d)", last, off, off + n_batch_tokens)); + } }); auto accept_special_token = [&](server_slot & slot, llama_token token) { @@ -3785,6 +3839,13 @@ struct server_context_impl { return; } + // Its logits are in a different view of this batch, so it is sampled when + // that view is decoded, not now. Without this the shift below produces + // negative indices and reads whatever is at the front of the wrong view. + if (!spec_is_inside_view(slot)) { + return; + } + // save the original draft size const size_t n_draft = slot.spec_draft.size(); @@ -3795,7 +3856,17 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + // Shifted according to the current sub-batch, exactly as tok_idx is for + // the non-speculative path above. These are indices into the batch that + // was just decoded, and that batch is the VIEW, not the whole batch. + // Passing the absolute indices is the second half of #24840: with off + // non-zero they address the wrong logits. + std::vector spec_i_view; + spec_i_view.reserve(slot.spec_i_batch.size()); + for (const auto & i : slot.spec_i_batch) { + spec_i_view.push_back(i - off); + } + auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, spec_i_view, slot.spec_draft); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); From 69e6dfba1e764b4e3a5f6a93ca3d0ac7bab7b562 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 4 Sep 2026 01:43:20 +0000 Subject: [PATCH 2/4] server: never halve n_batch below one speculative block Second half of #24840, and the half that matters in production. Keeping a block inside one sub-batch only works while a sub-batch can hold one. The KV-full retry halves n_batch without a floor, so it walks 2048, 1024, ... 4, 2, 1, and a view of 1 or 2 cannot serve a 3-index block however it is positioned. post_decode is then left with nothing to do but abort every slot on the server. Of 78 occurrences recorded in production logs at --spec-type draft-mtp --spec-draft-n-max 2, thirty were exactly that: a block starting at the view's own offset and reaching past its end, against views of 1 and 2. The other 48 are blocks that straddle or sit beyond a wider view, which the previous commit handles. Together they cover all 78. Halving past one block buys no memory worth having, since the difference is a couple of cells, and costs every conversation in flight. Note this is measured against the shape of the failures, not against a run: the CPU reproduction in the workspace uses ngram-simple, which ignores --spec-draft-n-max and drafts blocks of 49, so it exercises a regime no view-fitting fix can serve. Validating this properly needs an MTP draft model on a GPU. --- tools/server/server-context.cpp | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c82ba621518e..f42ed3f5c57b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3644,6 +3644,22 @@ struct server_context_impl { // retry with half the batch size to try to find a free slot in the KV cache if (!try_clear_idle_slots()) { n_batch /= 2; + + // ... but never below one speculative block. A slot's spec_i_batch + // entries all need logits from ONE decode, so a view narrower than a + // block cannot serve it however the view is positioned, and post_decode + // is left with nothing to do but abort every slot on the server. + // + // Of 78 occurrences recorded in production logs, 30 were exactly this: + // a block starting at the view's own offset and reaching past its end, + // with views of 1 and 2 against a 3-index block. Halving past that point + // buys no memory worth having -- the difference is a couple of cells -- + // and costs every conversation in flight. + int32_t n_spec_min = 1; + iterate(slots, [&](server_slot & slot) { + n_spec_min = std::max(n_spec_min, (int32_t) slot.spec_i_batch.size()); + }); + n_batch = std::max(n_batch, n_spec_min); } SRV_WRN("failed to find free space in the KV cache, retrying with smaller batch size, off = %d, n_batch = %d, ret = %d\n", off, n_batch, ret); @@ -3731,7 +3747,7 @@ struct server_context_impl { if (first < off + n_batch_tokens && last >= off + n_batch_tokens && first >= off) { // Split anyway: only reachable when the retry ladder drove n_batch below // one block. Report it rather than sample from logits that do not exist. - throw std::runtime_error(string_format("speculative batch index %d is not inside the current sub-batch [%d, %d)", last, off, off + n_batch_tokens)); + throw std::runtime_error(string_format("speculative block [%d, %d] straddles the current sub-batch [%d, %d)", first, last, off, off + n_batch_tokens)); } }); From 8d6118d448c9498fea48bb50109cda4988c9188f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 4 Sep 2026 03:05:38 +0000 Subject: [PATCH 3/4] server: do not let the speculative floor trap the retry ladder Bug in the previous commit, mine. Flooring n_batch at one speculative block clamped it there permanently, and the ladder reaching n_batch == 1 is precisely what terminates this retry: decode() reports "Context size has been exceeded" only for n_batch == 1 && ret == 1. A cache that genuinely cannot fit anything would therefore retry forever instead of reporting, which is worse than the crash the floor was added to prevent. Pause AT the floor once, then carry on halving. A block still gets one whole-view attempt, and the terminating case is still reachable. --- tools/server/server-context.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f42ed3f5c57b..86bc85064e03 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3643,6 +3643,7 @@ struct server_context_impl { // retry with half the batch size to try to find a free slot in the KV cache if (!try_clear_idle_slots()) { + const int32_t n_batch_prev = n_batch; n_batch /= 2; // ... but never below one speculative block. A slot's spec_i_batch @@ -3659,7 +3660,14 @@ struct server_context_impl { iterate(slots, [&](server_slot & slot) { n_spec_min = std::max(n_spec_min, (int32_t) slot.spec_i_batch.size()); }); - n_batch = std::max(n_batch, n_spec_min); + // Pause AT the floor, once, rather than clamping to it forever. The + // ladder reaching n_batch == 1 is what terminates this retry: decode() + // reports "Context size has been exceeded" only for n_batch == 1, so a + // hard clamp above 1 turns a cache that genuinely cannot fit anything + // into an infinite retry loop. Halving from the floor continues past it. + if (n_batch < n_spec_min && n_batch_prev > n_spec_min) { + n_batch = n_spec_min; + } } SRV_WRN("failed to find free space in the KV cache, retrying with smaller batch size, off = %d, n_batch = %d, ret = %d\n", off, n_batch, ret); From 945cdfc915bdbada7ccf3bf2317401c000016856 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 4 Sep 2026 20:19:15 +0000 Subject: [PATCH 4/4] server: drop the draft, not the conversation, when a block cannot fit the view The previous commit stopped the retry ladder from cutting through a speculative block, and reports the case it cannot fix by throwing. That throw reaches abort_all_slots, which ends every conversation on the server because one of them was drafting into a full cache: the same shape as the KV-full send_error path. spec_i_batch[0] is the index of the token that was actually sampled last round rather than a draft, and in this case it is inside the current view, so the slot can still sample its next token the ordinary way. Give up the prediction and keep the conversation: the cost is this one step's speedup. Only reachable when the ladder has driven n_batch below one block, which needs a cache full enough that the view is narrower than the draft. --- tools/server/server-context.cpp | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 86bc85064e03..802a7ffac2d3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3754,8 +3754,27 @@ struct server_context_impl { } if (first < off + n_batch_tokens && last >= off + n_batch_tokens && first >= off) { // Split anyway: only reachable when the retry ladder drove n_batch below - // one block. Report it rather than sample from logits that do not exist. - throw std::runtime_error(string_format("speculative block [%d, %d] straddles the current sub-batch [%d, %d)", first, last, off, off + n_batch_tokens)); + // one block, i.e. a cache so full that the view is narrower than the draft. + // + // Give up the DRAFT rather than the conversation. spec_i_batch[0] is the + // index of the token that was actually sampled last round, not a draft + // (see the push_back loop in server_slot::add_to_batch), and it is inside + // this view, so the slot can still sample its next token the ordinary way. + // The drafts behind it are a prediction; losing them costs this one step + // its speedup. + // + // Throwing here instead calls abort_all_slots, which ends EVERY + // conversation on the server because one of them was drafting into a full + // cache. Measured on the branch's own base: four chats, four errors, all of + // them this. It is the same failure shape as the KV-full send_error path + // that unslothai/llama.cpp#183 narrows. + SRV_WRN("speculative block [%d, %d] does not fit the current sub-batch [%d, %d); " + "dropping the draft for slot %d and sampling one token\n", + first, last, off, off + n_batch_tokens, slot.id); + slot.spec_draft.clear(); + slot.spec_i_batch.clear(); + slot.i_batch = first; + return; } });