Skip to content

[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
Niko1221:mainfrom
macintog:fix/budget-cut-in-flight-tokens
Open

macintog wants to merge 1 commit into
Niko1221:mainfrom
macintog:fix/budget-cut-in-flight-tokens

Conversation

@macintog

Copy link
Copy Markdown

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 the
answer. 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_tokens
81,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:

Thinking cap Tokens generated past the cut Wraps Answer request
61,440 0 39 reused, read 19
61,440 1 26 reused, read 18
61,440 2 to 5 25 re-read 61,463 tokens, mean 30.8 s
24,576 / 32,768 / 49,152 0 or 1 14 reused, read 18 or 19
24,576 / 32,768 / 49,152 2 or 3 4 re-read 24,599 / 32,791 / 49,175 tokens

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):

  1. Service stops reading at the first clean point at or past the budget: thought >= budget, no tag held back and
    no character split across tokens (serve/server.py ~5341). Closing the generator runs StrataEngine.generate's
    finally (~2280), which sends STOP and drains to DONE, discarding the T lines it reads.
  2. The engine checks for STOP only after a verify window (src/program/generate.cpp ~15198), and one window can
    commit several tokens (~15145).
  3. At DONE the engine's live session is every generated token but the last (live.swap(consumed), ~15235).
  4. The answer request's prompt is prompt + seg + wrap-up (serve/server.py ~5474), where seg is the thinking the
    server 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 that
    matches 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's
log 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

  • The solo StrataEngine.generate drain records the T ids it reads after STOP (the tokens in flight) in
    self.overrun, reset per request.
  • The budget-wrap branch appends them to seg before the wrap-up and passes them through the parser as reasoning,
    so the client sees what the model saw. They count in completion_tokens.
  • They are left out, which keeps the old behavior, when they hold a stop id, a < (a tag could start) or a split
    character (U+FFFD), or when carrying them would leave the answer no room within max_tokens.
  • The thinking can now run past the budget by those tokens: 1 to 3 in our patched runs. Without the change, the engine
    ran up to 5 past the cap on the RTX server and 6 on the Mac.
  • The server logs thinking budget reached (N tokens, K more in flight) when it carries K tokens.
  • Unchanged: the forced-call (opens) branch, which has the same pattern but is not exercised here;
    reasoning_close_retry; and generate_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:
overrun is only read in the wrap branch, and the existing budget tests (off by default, wrap-up, accounting) pass
unchanged.

Validation

Passed:

  • Regression check, python -m unittest serve.test_server.ThinkingBudget serve.test_server.ReasoningCloseRetry,
    macOS 27.2, Python 3.14.9:

    • this branch (7dd5cf18 plus this commit): 27 tests OK;
    • 7dd5cf18 with 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 is
      kept (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 2
    with MTP, --mmap-experts), this commit's change on v0.2.3.1 (7ddc17a; without the no-room condition, which a
    1024-token reply never reaches), only serve/ swapped between arms. One 6,262-token code-review prompt, 20
    requests per arm, reasoning_budget_tokens 256, max_tokens 1024, reasoning_effort xhigh. Requests sent
    temperature 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.

    Arm Wraps Tokens past the cut: 1 / 3 Answer requests that re-read Answer request read time
    7ddc17a 20 19 / 1 1 (the cap + 3 one: 6257 reused + 279 read in 587 ms) 228 ms median otherwise
    patched 20 18 / 2 0 (both cap + 3: 6520 reused + 19 read in 240 / 233 ms) 236 ms median

    The patched server logged 1 more in flight 18 times and 3 more in flight twice. Every reply finished length
    at 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_tokens 81,920,
    so a 61,440 cap): 61442 generated, logged 2 more in flight, then the answer request 68624 reused + 19 read in 238 ms. The same server before the change: 61446 generated, then 5421 reused + 61463 read in 105559 ms.

Failed: none observed.

Not run:

  • The full test suite on main: some of its tests fake engine crashes, which raise macOS crash dialogs here.
  • The patched server on CUDA. The change is in the Python server, which every backend shares, and touches no backend
    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).
  • A sampled (non-greedy) run of the patched server.

Limits

  • The thinking can exceed reasoning_budget_tokens by the tokens in flight.
  • The forced-call (opens) branch has the same pattern and is not changed.
  • At the small budget the trigger is rare (1 or 2 in 20), so the A/B rests on 2 patched triggering cases against 1
    unpatched; the in-service case adds one at the full cap.
  • The A/B prompt reviews a private change, so it isn't attached; the trigger depends on where the cut falls in a
    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_effort xhigh, budget
65,536, max_tokens 81,920, top_p 0.95, top_k 20 and temperature 1.0 (0.6 from 2026-10-21); the smaller caps come
from other requests with smaller budgets. Its engine log's strata serve: prompt … lines from
2026-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:

import collections, re, sys
LINE = re.compile(r"prompt (\d+) tokens = (\d+) reused \+ (\d+) read in (\d+) ms.*?, (\d+) generated")
rows = [tuple(map(int, m.groups())) for m in map(LINE.search, open(sys.argv[1])) if m]
table = collections.Counter()
for i, (p, reused, read, ms, gen) in enumerate(rows[:-1]):
    cap = next((c for c in (24576, 32768, 49152, 61440) if c <= gen <= c + 8), None)
    p2, reused2, read2, ms2, _ = rows[i + 1]
    if cap is None or reused or not p + cap <= p2 <= p + cap + 30 or reused2 == 0:
        continue
    table[cap, gen - cap, "reused" if read2 <= 40 else "re-read"] += 1
for key, n in sorted(table.items()):
    print(*key, n)

A/B, the triggering runs (engine log, the thinking request, then the answer request):

7ddc17a:  prompt 6262 tokens = 6257 reused + 5 read in 73 ms (68.8 tok/s), 259 generated in 4205 ms (61.6 tok/s), drafts accepted 145 of 179, 1 checkpoints
          prompt 6536 tokens = 6257 reused + 279 read in 587 ms (475.6 tok/s), 750 generated in 13183 ms (56.9 tok/s), drafts accepted 356 of 499, 1 checkpoints
patched:  prompt 6262 tokens = 6257 reused + 5 read in 64 ms (78.3 tok/s), 259 generated in 4295 ms (60.3 tok/s), drafts accepted 150 of 178, 1 checkpoints
          prompt 6539 tokens = 6520 reused + 19 read in 240 ms (79.2 tok/s), 747 generated in 13357 ms (55.9 tok/s), drafts accepted 354 of 467, 1 checkpoints
patched:  prompt 6262 tokens = 6257 reused + 5 read in 72 ms (69.1 tok/s), 259 generated in 4160 ms (62.3 tok/s), drafts accepted 137 of 161, 1 checkpoints
          prompt 6539 tokens = 6520 reused + 19 read in 233 ms (81.6 tok/s), 747 generated in 13297 ms (56.2 tok/s), drafts accepted 350 of 485, 1 checkpoints

In service on the Mac (6,539 = 6,262 + 256 + 3 + 18 above; 68,643 = 7,183 + 61,440 + 2 + 18 here):

before:  prompt 5426 tokens = 0 reused + 5426 read in 26249 ms (206.7 tok/s), 61446 generated in 978342 ms (62.8 tok/s), drafts accepted 39623 of 42860, 1 checkpoints
         prompt 66884 tokens = 5421 reused + 61463 read in 105559 ms (582.3 tok/s), 20462 generated in 298263 ms (68.6 tok/s), drafts accepted 14382 of 15204, 4 checkpoints
after:   [strata] thinking budget reached (61440 tokens, 2 more in flight): wrapping up the thinking
         prompt 7183 tokens = 0 reused + 7183 read in 10138 ms (708.5 tok/s), 61442 generated in 927779 ms (66.2 tok/s), drafts accepted 40321 of 43793, 1 checkpoints
         prompt 68643 tokens = 68624 reused + 19 read in 238 ms (80.0 tok/s), 521 generated in 9963 ms (52.3 tok/s), drafts accepted 213 of 333, 1 checkpoints

…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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant