server: shift the speculative batch indices on the sub-batch retry - #182
server: shift the speculative batch indices on the sub-batch retry#182danielhanchen wants to merge 4 commits into
Conversation
…indices Fixes ggml-org#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.
Second half of ggml-org#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.
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.
|
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. |
… 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.
|
Added What was leftThe commit before it stops the retry ladder from cutting through a speculative block, and reports the case it cannot fix by throwing. That throw reaches What it does now
MeasurementFour concurrent chats, 3000-token prompts,
The solid result is the third row and the What remains in every row is genuine cache exhaustion. Four chats each holding a 3000-token prompt plus a long answer do not fit in 16384 cells, and a bare llama-server has no admission control to say so, which is what #183 narrows from "every slot" to "one slot". With a scheduler in front that reserves the drafts, the same four chats complete 4 of 4 with zero context-exceeded errors. |
![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: 945cdfc915
ℹ️ 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".
| slot.spec_draft.clear(); | ||
| slot.spec_i_batch.clear(); | ||
| slot.i_batch = first; |
There was a problem hiding this comment.
When KV pressure reduces the view below the speculative block size, clearing these vectors does not actually discard the draft: handle_last_sampled_token() has already appended every draft token to slot.prompt.tokens, and the successful partial decode may already have written some of them to both target and draft caches; the outer loop will also decode the remaining batch entries afterward. The ordinary sample taken from first is therefore subsequently positioned after unverified draft tokens, silently making them part of the conversation context and producing incorrect continuations. This fallback must remove the drafts from the prompt and cache/batch state (or fail safely) before sampling normally.
Useful? React with 👍 / 👎.
1 participant
--spec-type draft-mtpwith--parallel N --kv-unifiedthrowsspeculative batch index %d is not inside the current sub-batchand kills the request,whenever the KV cache is full enough that llama-server has to fall back to a smaller batch.
The bug
When a decode fails to fit, the server retries the same batch in smaller pieces. On that
path it already does two things for
slot.i_batch: skips a slot whose index lies outsidethe current view, and shifts the index by the view's offset, with the comment "shifted
according to the current sub-batch". It does neither for
slot.spec_i_batch.That asymmetry is the whole defect. The speculative indices stay absolute while the view
they are read against becomes relative, so the bounds check fires on indices that were
correct when they were recorded.
The check is also the only thing standing between this and a wrong answer: with it removed
the same indices would sample from another slot's logits.
What this changes
Three commits, smallest first:
spec_i_batchentries are contiguous, so a block that would straddle the end of a viewcuts the view short instead, and the surviving indices are shifted by the offset exactly
as
i_batchis.n_batchbelow one speculative block. The retry ladder halves the batcheach time; without a floor it can end up smaller than a single block, which no amount of
cutting can fit.
reaching past its end cannot be cut, so the floor must not stop the ladder from
descending, or the retry never terminates.
The error message that remains is deliberately different: a block that genuinely cannot be
placed now says
speculative block [%d, %d] straddles the current sub-batch [%d, %d), sothe two situations are distinguishable in a log.
Measurement
Two independent runs, because they answer different questions.
The branch, against its own base. Both arms built from this tree, differing only by these
three commits, run as a bare llama-server with no caller in front of it. Four concurrent
chats,
-c 16384 --parallel 4 --kv-unified --spec-type draft-mtp --spec-draft-n-max 2,3000-token prompts, driven until the pool overflows:
is not inside the current sub-batchstraddles the current sub-batchThe same change on the base Unsloth Studio actually ships (b10715-mix, 92cedc867),
interleaved control and treatment over two rounds with the order reversed, each arm confirmed
from
/proc/PID/mapsof the live server: 18 occurrences of the original error on control,0 on treatment.
What this does and does not fix
It removes the bounds violation this issue is about. It is NOT a complete fix, and the table
above says so: two of the four chats still fail on this branch, with a different error.
The remaining case is a speculative block that begins exactly at the view's own offset and
reaches past its end,
[4, 5]against[4, 5). It cannot be cut, because cutting the viewbefore it leaves an empty view, and the retry ladder has already driven
n_batchbelow oneblock by the time it appears. The code reports it rather than sampling from logits that were
never computed, which is the safe half of the choice.
Sampling only the part of the block inside the view would look like an obvious fix and is
not one: the drafted tokens beyond the view are already in the batch at positions this
sub-batch does not decode, so discarding them mid-flight risks leaving stale cells, and the
failure mode of getting that wrong is a wrong answer rather than an error. I have not
convinced myself of a correct version, so I have not written one.
What the change buys, then, is narrower than "the bug is gone" and still worth having:
can be told apart in a log rather than presenting as one bounds error;
handles are the common ones;
does, none of these are reached at all. The value there is that the residual overshoot
degrades to
Context size has been exceeded, which a caller can see coming.The natural next step is to make the unfittable case terminate that one slot through the
ordinary context-exceeded path instead of throwing, which composes with
#183. That is a separate change and is not in this PR.
Notes
server-context.cpponly. Builds clean, and the measurements above are from that build. Noeffect without
--spec-type: every new branch is inside a block guarded on a non-emptyspec_i_batch.There is no automated CI on this PR, and that is expected rather than a gap: the only
pull_request-triggered workflow in this fork is thepr-set.jsonlint, which is gated onpaths this change does not touch. Prebuild checks arrive when a PR is pinned into
scripts/unsloth/pr-set.json, which also ships it into the nightly builds, so I have leftthat decision to a maintainer.