server: on a full KV cache, terminate one slot instead of every slot - #183
server: on a full KV cache, terminate one slot instead of every slot#183danielhanchen wants to merge 3 commits into
Conversation
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.
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.
|
You have reached your Codex usage limits for security reviews. Please try again later. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2948afdf41
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| batch.batch_rendered = false; | ||
| batch.render(); |
There was a problem hiding this comment.
When KV exhaustion occurs for a batch with has_embd, the initial server_batch::render() has already set batch.token to nullptr. This second render() calls common_batch_add() for every retained entry, and that helper writes through batch.token, so the recovery path dereferences null and crashes the server instead of terminating one slot. Restore tokens_ptr before rebuilding the batch, as server_batch::clear() does.
Useful? React with 👍 / 👎.
1 participant
With
--parallel N --kv-unified, one overflowing request kills every request that happensto be decoding at the time.
The bug
--kv-unifiedgives N slots one shared pool of N cells while telling each slot it has allof them. That is the point: it lets a single long chat use the whole context. It also means
the pool can run out while several slots are mid-answer.
When it does, the server calls
send_erroron every processing slot. A user whose chathad nothing to do with the overflow loses their answer because somebody else's grew.
What this changes
Terminate ONE slot and let the rest carry on. The victim is the slot holding the most
cells, so the eviction frees as much as possible per turn lost, and one slot always makes
progress rather than N slots repeatedly failing together.
Three commits:
requires. Removing one slot's entries from the middle of a batch moves the entries after
it, so
i_batchandspec_i_batchare remapped for the slots that shift down.embdalongsidetokens. This one is mine: the first version compacted thetoken array and left the embedding array, which is a flat
n_embd-per-token buffer, soa multimodal batch would have been fed embeddings misaligned to their tokens. Silently.
slot_batchedis safe to leave pointing at a terminated slot.The remap is checked by brute force over 20000 randomly shaped batches, against the
invariant that entries below the removal point never move.
Notes
server-context.cpponly, and only inside the branch that already handled this case byfailing everyone. Compiles clean against the repo's own
compile_commands.json.Unlike the sibling fix in #182, I have not been able to put a number on
this one: reaching it requires the pool to overflow with more than one slot decoding, and on
the configuration I have, the speculative index bug in that PR fires first. The change is
argued from the code path rather than from a measurement, which is why it is a separate PR.
There is also a case the two PRs would compose on. #182 leaves one situation throwing: a
speculative block that begins at the view offset and reaches past its end cannot be cut, and
it reports rather than sampling logits that were never computed. Terminating that one slot,
which is what this PR makes possible, is a better answer than either throwing or killing
everybody. I have not written that change here; it belongs after both of these land.
No automated CI runs on this PR, which is expected rather than missing: the only
pull_request-triggered workflow in this fork is thepr-set.jsonlint, gated on paths thischange does not touch. Prebuild checks arrive when a PR is pinned into
scripts/unsloth/pr-set.json, which also ships it, so I have left that to a maintainer.