From 02b0156b3ce2773c7ddf60a524da64aab83689f1 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 4 Sep 2026 05:13:28 +0000 Subject: [PATCH 1/3] server: on a full KV cache, terminate one slot instead of every slot When a decode cannot find KV space and the retry ladder bottoms out at n_batch == 1, the server sends an error to every processing slot and clears every context. One conversation overflowing therefore kills all of them, which on a shared unified cache is the common case rather than a corner one. Implements the TODO already sitting above that handler: terminate only the largest active sequence and continue with the rest. The blocker it notes, removing the tokens from the current batch, is tractable because off splits the batch cleanly -- everything at index >= off is undecoded by construction of the loop in update_slots -- so the victim's undecoded entries can be dropped and the batch re-rendered. Leaving the survivors running is safe: ret == 1 comes from FAILED_PREPARE in llama_context::decode, before output_reserve, before the sampling transaction and before any compute, so no cells were assigned and no survivor's KV was touched. render() rebuilds the batch from tokens, so surviving entries after a removed one shift down; i_batch and spec_i_batch are remapped accordingly. Without that the slots would sample from stale indices with no error raised at all. With one slot left there is nobody to spare and the previous behaviour, erroring that slot, is correct and is what happens. --- tools/server/server-context.cpp | 95 +++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..167705e1e4ec 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2644,6 +2644,89 @@ struct server_context_impl { } } + // Terminate ONE slot to make room, rather than every conversation on the server. + // + // This is the TODO sitting above the context-exceeded handler: "try to terminate only + // the largest active slot/sequence and continue with the rest / need to remove the + // tokens from the current batch too". The second sentence is the hard part, and it is + // tractable here only because `off` splits the batch cleanly: everything at index + // `>= off` has not been decoded yet, by construction of the loop in update_slots. + // + // Safe to leave the survivors running because a `ret == 1` decode committed nothing. + // That return comes from LLAMA_MEMORY_STATUS_FAILED_PREPARE in llama_context::decode, + // before output_reserve, before the sampling transaction opens and before any compute, + // so no cells were assigned and no survivor's KV was touched by the failed call. + // + // Returns false when there is nobody to spare, and then the caller's original + // behaviour -- error everyone, which is now just the one slot -- is the correct one. + bool terminate_one_slot(int32_t off, const std::string & reason) { + server_slot * victim = nullptr; + int n_processing = 0; + iterate(slots, [&](server_slot & slot) { + if (!slot.is_processing()) { + return; + } + n_processing++; + // Largest sequence, as the TODO says: on a crash path the figure that matters + // is cells freed per conversation lost. (Scheduling preemption elsewhere is + // better served by newest-first; this is not that, and the two should not be + // made to agree for the sake of it.) + if (victim == nullptr || slot.prompt.n_tokens() > victim->prompt.n_tokens()) { + victim = &slot; + } + }); + + if (victim == nullptr || n_processing <= 1) { + return false; + } + + const int32_t victim_id = victim->id; + + SRV_WRN("terminating slot %d (%d tokens) to free space, %d other slot(s) continue\n", + victim_id, (int) victim->prompt.n_tokens(), n_processing - 1); + + send_error(*victim, reason); + victim->release(); + // seq_rm inside frees this sequence's cells, which is the entire point of the + // exercise; the tokens it already decoded below `off` go with them. + victim->prompt_clear(); + victim->i_batch = -1; + victim->spec_i_batch.clear(); + + // Drop the victim's UNDECODED entries and remember where every survivor moved to. + std::vector remap(batch.tokens.size(), -1); + std::vector kept; + kept.reserve(batch.tokens.size()); + for (int32_t i = 0; i < (int32_t) batch.tokens.size(); i++) { + if (i >= off && batch.tokens[i].id_slot == victim_id) { + continue; + } + remap[i] = (int32_t) kept.size(); + kept.push_back(batch.tokens[i]); + } + batch.tokens = std::move(kept); + batch.batch_rendered = false; + batch.render(); + + // render() rebuilds the batch wholesale from `tokens`, so every surviving entry + // after a removed one shifts DOWN. `i_batch` and `spec_i_batch` are absolute + // indices into that batch and must move with it. Skipping this does not throw: + // post_decode goes on to sample from whatever logits now sit at the stale index, + // silently, which is a worse outcome than the crash being fixed here. + iterate(slots, [&](server_slot & slot) { + if (slot.i_batch >= 0 && slot.i_batch < (int32_t) remap.size()) { + slot.i_batch = remap[slot.i_batch]; + } + for (auto & i : slot.spec_i_batch) { + if (i >= 0 && i < (int32_t) remap.size()) { + i = remap[i]; + } + } + }); + + return true; + } + // @ngxson : for debugging only int64_t t_pre_decode = 0; int64_t t_decode = 0; @@ -3593,6 +3676,18 @@ struct server_context_impl { if (!err.empty()) { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); + // Out of KV space, which is a capacity problem with one sequence too + // many in it, not a fault in every conversation on the server. Free + // the largest and carry on with the rest. Only for ret == 1: an + // invalid batch or a compute error says nothing about capacity, and + // dropping a slot would not address either. + if (ret == 1 && n_batch == 1 && terminate_one_slot(off, err)) { + // A whole sequence just went, so the halving that got us down to + // 1 no longer describes what will fit. + n_batch = llama_n_batch(ctx_tgt); + return false; // retry, with the batch minus that slot + } + for (auto & slot : slots) { if (slot.is_processing()) { send_error(slot, err); From 5760e9235e8bb5a8fe40f1db6cae8f1f4f78feb0 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 4 Sep 2026 05:15:01 +0000 Subject: [PATCH 2/3] server: compact the embedding array alongside the tokens In embd mode server_batch::embd is a flat n_embd-per-token array running parallel to tokens, and render() passes embd.data() to the batch directly. Dropping entries from tokens without dropping the matching spans from embd misaligns every embedding after the first removal, with nothing raised. --- tools/server/server-context.cpp | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 167705e1e4ec..b3f80006a338 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2697,14 +2697,33 @@ struct server_context_impl { std::vector remap(batch.tokens.size(), -1); std::vector kept; kept.reserve(batch.tokens.size()); + // In embd mode `embd` is a flat n_embd-per-token array running parallel to + // `tokens`, and render() hands embd.data() straight to the batch. Compacting one + // without the other misaligns every embedding after the first removal, silently. + std::vector kept_embd; + if (batch.has_embd) { + kept_embd.reserve(batch.embd.size()); + } for (int32_t i = 0; i < (int32_t) batch.tokens.size(); i++) { if (i >= off && batch.tokens[i].id_slot == victim_id) { continue; } remap[i] = (int32_t) kept.size(); kept.push_back(batch.tokens[i]); + if (batch.has_embd) { + const size_t stride = (size_t) batch.n_embd; + const size_t begin = (size_t) i * stride; + if (begin + stride <= batch.embd.size()) { + kept_embd.insert(kept_embd.end(), + batch.embd.begin() + begin, + batch.embd.begin() + begin + stride); + } + } } batch.tokens = std::move(kept); + if (batch.has_embd) { + batch.embd = std::move(kept_embd); + } batch.batch_rendered = false; batch.render(); From 2948afdf41738d6b579085bb4ccc2113bafde2cb Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 4 Sep 2026 05:16:03 +0000 Subject: [PATCH 3/3] server: note why slot_batched is safe to leave pointing at a terminated slot --- tools/server/server-context.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b3f80006a338..16bf72eb8b83 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2743,6 +2743,12 @@ struct server_context_impl { } }); + // `batch.slot_batched` may be the victim, and is intentionally left alone. It is + // not dereferenced inside the decode loop -- the lora and embedding settings it + // supplies were applied once before the loop -- and `slots` is a fixed vector, so + // release() does not destroy the object and the pointer does not dangle. Every + // survivor passed can_batch_with() against it, so those settings remain correct + // for the tokens that are left. return true; }