[zsariboga]

Rebased on #1852's PromptEncoder (the encoder now on main), as asked. The branch is one commit on main (477ac13c). It no longer depends on #1455.

What it changes

Everything new is opt-in behind STRATA_REPLY_REUSE=1. With the variable unset, every prompt is a fresh encode and the server behaves exactly as main does.

• PromptEncoder.remember_reply(prompt_ids, gen_ids) (tools/strata_tokenizer.py, +48 lines). After a reply, the server keeps the prompt plus the reply's own ids as one encoder entry.
  • The entry's cuts are the reply's token boundaries. A cut never falls inside a UTF-8 character (the incremental decoder's state is checked) and never inside a special tag (_straddles).
  • In encode, a prompt that shares the text past the reply's start may cut at those boundaries without the max_special_len - 1 margin. The engine holds those ids, and the text from the cut on is encoded afresh.
  • The ids always decode to the prompt's text. This is the one place where prompt ids may differ from a full encode: they are the model's own pieces, which the engine already holds.
• Thinking budget stop (serve/server.py). The engine sends a few tokens after the budget's STOP before it sees it, and keeps them in its session. They are now collected while draining (StrataEngine.drained). With the switch, they join the thinking when they are plain text (no stop token, no <, no split character), so the wrap-up continuation extends the session instead of reading the thinking again.
• The previous version's STRATA_PROMPT_DUMP diagnosis and the raised encoder limits are gone. serve: tokenise the prompt from the last shared prefix, not from the start #1852's encoder has no such limits.

Why

In agent sessions the engine often read a whole reply again on the next turn, although the reply's state was still in its session. The model writes some text in pieces BPE would merge (Turkish letters, digits, code), so the next turn's fresh encode leaves the engine's ids at the first such piece. A reply has no checkpoint inside it, so the engine went back to the turn's start. One real case: 364 tokens reused and 1903 read, where about 885 could have been reused.

Measured (the previous version, 0.1.40.3)

Setup: RTX 5070 Ti Laptop 12 GB, Core Ultra 7 255HX, 64 GB DDR5-5600, Windows 11, Swift 1.5 abliterated IQ3_XXS, opencode.

• In an engine log of 1740 requests, ~400 s of 2836 s of prompt reading (~14%) re-read a reply the engine still held. About 150 s of that came from divergences inside ordinary replies.
• A tool turn after a budget wrap-up: reused 362 / 364 before, 969 / 971 after; read 1379-3635 before, 583 / 1770 after.

Not re-measured end to end on this rebase. The unit tests below check the encoder and the budget path. I run the same encoder change on 0.1.45 locally.

Tests

• serve/test_prompt_encoder.py ReplyReuse (4 new, with a small BPE vocabulary where 'ab' / 'abc' / ' a' merge):
  • the next turn keeps the reply's own ids up to where the text differs;
  • nothing is kept without its prompt;
  • a tag across the cut is encoded whole;
  • a split character is never a cut.
• serve/test_server.py ThinkingBudget (3 new):
  • drained tokens join the thinking with the switch;
  • drained tokens with a tag are left out;
  • drained tokens are dropped without the switch.
• python -m unittest discover -s serve -t . -p "test_*.py" (Windows): 855 tests, one failure. It is test_keepawake test_service_holds_for_the_whole_request_including_early_close, which also failed on unchanged main here and passes when run on its own (timing).

Python only: no engine, HIP or SYCL file is touched.

[Niko1221]

Thank you. #1455 is covered by #1852 in our branch, so this change needs a rewrite on that encoder. We have not measured it (it needs a long agent session). If you can rebase it on #1852's PromptEncoder we will measure it then.

[zsariboga]

Rebased on #1852's PromptEncoder as one commit on main (477ac13), and made opt-in (STRATA_REPLY_REUSE=1; unset, every prompt is a fresh encode as before). The description above is updated with what changed, the tests, and what was not re-measured on this rebase.