Repository navigation
[Engine] serve: keep the thinking tokens generated past a budget cut, so the answer doesn't re-read the thinking - #3971
Open
macintog wants to merge 1 commit into
Conversation
…past the cut At reasoning_budget_tokens the request loop stops reading at exactly the budget and closes the engine generator, which sends STOP and drains to DONE, discarding the T lines still in flight. The engine stops only between verify windows, so it has often generated (and committed to its state) a few tokens past the cut. The continuation prompt (prompt + read tokens + wrap-up) then diverges from the engine's live state a few tokens back, and the engine falls back to the checkpoint at the end of the prompt: it re-reads the whole thinking (61,463 tokens at a 61,440 cap: ~31 s on an RTX 3080 Ti, ~105 s on an M4 Max) in about a quarter of budget wraps. The drain now keeps those tokens (`overrun`), and the budget wrap appends them to the thinking before the wrap-up, so the continuation extends what the engine holds. They are left out (the old behavior) when they hold a stop token, a '<' or a split character, or when carrying them would leave the answer no room within max_tokens. The thinking can run past the budget by those few tokens. The forced-call (`opens`) path and the batched path are unchanged. Tests: ThinkingBudget gains an engine that leaves tokens in flight; the new test fails without the fix. Authored-by-model: claude-opus-5-5@high (Claude Code subagent)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem and result
When a request's thinking reaches
reasoning_budget_tokens, the server stops the engine, appends a short wrap-up("I have thought about this long enough; time to give my answer." and
</think>), and asks the engine for theanswer. The engine still holds the prompt and the thinking, so that answer request should read only the wrap-up and
one more token: 18 or 19 tokens, in about 0.2 s.
In about a quarter of these budget wraps it reads the whole thinking again before answering. With
max_tokens81,920 and a budget of 65,536 the server caps the thinking at 61,440 tokens (#984), so that is 61,463 tokens read per
request: 30.8 s on average on an RTX 3080 Ti, and 105.6 s in the one case we logged on an M4 Max.
It happens because the engine doesn't stop exactly when the server stops reading. It finishes the step it is on,
which can generate a few more thinking tokens, and the server discards them. The answer request then starts from
the thinking without those tokens, which no longer matches what the engine holds, so the engine goes back to its
last checkpoint, near the end of the prompt, and reads the whole thinking again.
This PR keeps those tokens and puts them in the answer request, before the wrap-up, so the answer request matches
what the engine holds (except in the few cases, listed under Approach, where it leaves them out as before). On the M4 Max, the unpatched server took 105.6 s to re-read 61,463 tokens on such a
wrap; with this change, a wrap with two tokens past the cut read 19 tokens in 0.24 s.
How often, and the deciding factor. On our RTX 3080 Ti server, 29 of 108 wraps re-read the thinking. Each of the
29 had generated two or more tokens past the cut, and none of the other 79 had:
One token past the cut does no harm: the engine never commits the last token it generates, so its state still
matches. Two or more broke the match in every case we logged.
In the code (at
7dd5cf18):Servicestops reading at the first clean point at or past the budget:thought >= budget, no tag held back andno character split across tokens (
serve/server.py~5341). Closing the generator runsStrataEngine.generate'sfinally(~2280), which sendsSTOPand drains toDONE, discarding theTlines it reads.STOPonly after a verify window (src/program/generate.cpp~15198), and one window cancommit several tokens (~15145).
DONEthe engine's live session is every generated token but the last (live.swap(consumed), ~15235).prompt + seg + wrap-up(serve/server.py~5474), wheresegis the thinking theserver read. When the live session holds a committed token past the cut that differs from the wrap-up's first
token, it is not a prefix of that prompt,
starts_with(live, …)fails (~12590), and the newest checkpoint thatmatches is the prompt's last one, 5 tokens before its end (hence
reused= prompt − 5 in the logs).Not the cause: re-tokenization (the answer request is built from token ids),
reasoning_close_retry(not triggered),--prompt-cache-every(it counts prompt tokens; decoding takes no checkpoints), the conversation cache (the Mac'slog shows it off).
Related discussion
None (a focused fix). Related issues that cover other parts of the budget: #123, #984, #1053, and closed #1537 (a
live-session checkpoint for a later request; it does not cover an answer request that leaves the live session a few
tokens before its end).
Approach
StrataEngine.generatedrain records theTids it reads afterSTOP(the tokens in flight) inself.overrun, reset per request.segbefore the wrap-up and passes them through the parser as reasoning,so the client sees what the model saw. They count in
completion_tokens.<(a tag could start) or a splitcharacter (U+FFFD), or when carrying them would leave the answer no room within
max_tokens.ran up to 5 past the cap on the RTX server and 6 on the Mac.
thinking budget reached (N tokens, K more in flight)when it carries K tokens.opens) branch, which has the same pattern but is not exercised here;reasoning_close_retry; andgenerate_batched(--batch), whose drain is separate.This changes the default path for every budget wrap that finds tokens in flight, not only for those that used to
re-read: a wrap with one token in flight (the common case: 26 of 108 RTX wraps, 18 of 20 patched A/B runs) now
carries that token in its thinking and reads 19 tokens instead of 18. Requests without a budget wrap are unchanged:
overrunis only read in the wrap branch, and the existing budget tests (off by default, wrap-up, accounting) passunchanged.
Validation
Passed:
Regression check,
python -m unittest serve.test_server.ThinkingBudget serve.test_server.ReasoningCloseRetry,macOS 27.2, Python 3.14.9:
7dd5cf18plus this commit): 27 tests OK;7dd5cf18with only this commit's test changes: 1 failure,test_the_continuation_keeps_the_tokens_generated_past_the_budget. The two tests that check the old behavior iskept (a
<in flight, no room to answer) pass there too.A/B on an M4 Max, 128 GB (Qwen3.8-Flash-Next Unsloth UD-Q4_K_XL, engine 0.1.48-metal built from source,
--spec 2with MTP,
--mmap-experts), this commit's change onv0.2.3.1(7ddc17a; without the no-room condition, which a1024-token reply never reaches), only
serve/swapped between arms. One 6,262-token code-review prompt, 20requests per arm,
reasoning_budget_tokens256,max_tokens1024,reasoning_effortxhigh. Requests senttemperature 0.6, top_p 0.95, top_k 20, min_p 0, but the Metal engine at this version doesn't apply request
sampling ([Engine] The Mac (Metal) engine ignores a request's sampling (temperature, top_p, top_k, min_p, seed, penalties) and drops logit_bias #3006), so both arms decoded greedily.
7ddc17a6257 reused + 279 read in 587 ms)6520 reused + 19 read in 240 / 233 ms)The patched server logged
1 more in flight18 times and3 more in flighttwice. Every reply finishedlengthat 1024 tokens, with no think tags in content or reasoning. Memory pressure stayed normal, with no swap growth.
In service on the same Mac since 2026-10-23 (the same change on
v0.2.3.1, budget 65,536,max_tokens81,920,so a 61,440 cap):
61442 generated, logged2 more in flight, then the answer request68624 reused + 19 read in 238 ms. The same server before the change:61446 generated, then5421 reused + 61463 read in 105559 ms.Failed: none observed.
Not run:
main: some of its tests fake engine crashes, which raise macOS crash dialogs here.source, but its behavior on the RTX server is unverified.
--batch(generate_batched), HIP and SYCL hosts, and DeepSeek-V4.1 (its own stop token and tags).Limits
reasoning_budget_tokensby the tokens in flight.opens) branch has the same pattern and is not changed.unpatched; the in-service case adds one at the full cap.
verify window, not on the prompt.
Evidence: how the RTX counts were made, and the A/B and in-service log lines
The RTX server: RTX 3080 Ti, Windows, Qwen3.8-Flash-Next IQ3_S, int8 KV, MTP, running that week's Strata
release (0.2.3.1 with engine 0.1.48 at its end). Most requests were code reviews at
reasoning_effortxhigh, budget65,536,
max_tokens81,920, top_p 0.95, top_k 20 and temperature 1.0 (0.6 from 2026-10-21); the smaller caps comefrom other requests with smaller budgets. Its engine log's
strata serve: prompt …lines from2026-10-19 to 2026-10-23, in order. A wrap is a thinking request that read its prompt fresh and generated between the cap
and cap + 8, followed by a request whose prompt is that prompt plus the cap plus up to 30 tokens. 31 such pairs,
logged while the server ran without its prompt cache (every answer request
0 reused), are left out. The pairing:A/B, the triggering runs (engine log, the thinking request, then the answer request):
In service on the Mac (6,539 = 6,262 + 256 + 3 + 18 above; 68,643 = 7,183 + 61,440 + 2 + 18 here):