feat(backends): preserve token ids across chat turns - #1592
Draft
noaakl wants to merge 14 commits into
Draft
Conversation
jakelorocco
reviewed
Aug 27, 2026
Contributor
There was a problem hiding this comment.
Hello; thank you for putting together the draft PR. I added some initial comments but we will also take a look as a team to provide some more indepth feedback.
I think it might also be helpful if you could provide a minimum viable example to reproduce both the problematic and correct behavior that you are describing. I don't think this needs to be done with Mellea, but if you can just highlight exactly what is getting lost with an example / example tokens, that would be very helpful for myself.
I think the biggest potential issue here is that granite switch (and our adapters) actually run through a function called _generate_from_intrinsic. That function utilizes some lower level transformations to modify the input / output of any given request. We need to understand how this relates to multi-turn conversations and the switch based adapters that this is being implemented for.
Also; our io.yamls tend to re-write the context so I would be interested to hear how this works across multiple turns.
Comment on lines
+85
to
+87
| # This lives here rather than in its own module because `OpenAIBackend` is the | ||
| # only consumer: a local-tokenizer backend has no use for the route, and a | ||
| # multi-provider proxy cannot rely on it. |
Contributor
There was a problem hiding this comment.
Can you please expand on why a local-tokenizer backend doesn't need this functionality? or in other words, why we wouldn't want this functionality on the LocalHFBackend?
Comment on lines
+115
to
+120
| Do NOT pass the retained ids as `prev_ids`. They are what the server actually | ||
| saw, and they diverge from a fresh re-render exactly when `encode(decode(ids))` | ||
| loses a token -- which is the case retaining ids exists to survive. Comparing | ||
| against them would make this raise on the one conversation the feature is for. | ||
| The retained ids are the prefix the caller SPLICES onto this delta, not the | ||
| thing it compares against. |
Contributor
There was a problem hiding this comment.
It might be nice to include an example of what these control tokens are or a link to the documentation; something like:
encode(decode([0, 1, 2, 3, 4, 5, 6])) -> [1, 2, 3, 4, 5]
Comment on lines
+1755
to
+1771
| # Placed here, AFTER extra_body is merged, so the chat_template_kwargs handed | ||
| # to /tokenize are the ones this request would actually have sent -- including | ||
| # `adapter_name` arriving via user extra_body and `enable_thinking` set above. | ||
| # Dispatching earlier would tokenize under a different template than the turn | ||
| # is generated under, and the delta would describe the wrong render. | ||
| if isinstance(ctx, ChatContext) and ctx.retains_token_ids: | ||
| return await self._generate_via_token_ids( | ||
| ctx, | ||
| conversation, | ||
| (extra_params.get("extra_body") or {}).get("chat_template_kwargs"), | ||
| action=action, | ||
| linearized_context=linearized_context, | ||
| _format=_format, | ||
| model_options=model_opts, | ||
| has_tools=use_tools, | ||
| ) | ||
|
|
Contributor
There was a problem hiding this comment.
I think this would need to fire under the generate_from_intrinsic path in order to get the proper formatting and tokens required.
Comment on lines
+68
to
+83
| retain_token_ids (bool): Opt into id-preserving history. When `True`, a | ||
| backend that supports it sends the exact token ids already sent plus | ||
| only the new turn's, instead of re-rendering the conversation from | ||
| text. Re-rendering drops the control tokens a chat template inserted | ||
| and cannot reproduce ids exactly (`encode(decode(ids))` is not the | ||
| identity), both of which break a server's prefix cache. Defaults to | ||
| `False`, so behaviour is unchanged unless asked for. | ||
| sent_token_ids (tuple[int, ...]): Ids the server has already seen, | ||
| verbatim. Empty until a backend records a turn. A tuple so a caller | ||
| cannot mutate the context's state through it. | ||
| sent_model_id (str | None): Model those ids were produced by. Ids are not | ||
| portable across vocabularies, so a backend can refuse rather than | ||
| reinterpret a prefix produced by a different model. | ||
| sent_message_count (int): How many chat messages `sent_token_ids` covers. A | ||
| backend needs this to re-render exactly the already-sent side of the | ||
| conversation, which is what the new turn's ids are subtracted against. |
Contributor
There was a problem hiding this comment.
What is the reason for the context to handle this instead of just having a flag on the backend to attempt to tokenize and try to attach tokens to a cblock?
Contributor
planetf1
requested changes
Sep 4, 2026
Contributor
There was a problem hiding this comment.
The adapter-function refactor tracked by #1144 should merge in the next few days. This PR will need rebasing on main afterwards: both changes modify the intrinsic request path, and the resolution needs to retain both the pre-tokenized completion route and the adapter lifecycle handling.
Can we cover this at three levels?
- Unit tests for prompt construction and fallback: exact ids passed to
/v1/completions, parser shape, retained-prefix bookkeeping, and client-side retention observability. - Integration tests using the in-memory tracing and metrics exporters: the post event must carry the pre event’s generation id, a non-negative duration, and the retention signal.
- A real vLLM e2e test for the intended outcome: compare a normal re-rendering control with retained ids over the same multi-turn adapter conversation, and assert the retained-id arm produces more server prefix-cache reuse. Include an adapter-to-base transition and
THINKING=True; cache-hit evidence alone is not enough.
The e2e case should use the existing e2e, openai, and vllm markers, GPU gating, and slow if it exceeds one minute.
Could we open and link a follow-up for client-visible retention observability? Users without access to vLLM metrics need to see whether Mellea used retained ids and how large the retained prefix was. This should report client-side retention rather than claim a server cache hit. The follow-up should also document supported servers and fallback behaviour.
| output._call.action = action | ||
| output._call.context = linearized_context | ||
| if isinstance(action, Component): | ||
| output.parsed_repr = action._parse(output) |
Contributor
There was a problem hiding this comment.
_generate_from_raw() stores a single completion choice in output.raw.response, but Message._parse() expects the chat-completion shape: response["choices"][0]. Re-parsing the restored Message therefore raises KeyError: "choices" on the first retaining mfuncs.chat turn. Can we either retain a chat-shaped response here or parse the completion result directly? A regression test through mfuncs.chat would cover this.
|
|
||
| glog = output._generate_log | ||
| await invoke_hook( | ||
| HookType.GENERATION_POST_CALL, |
Contributor
There was a problem hiding this comment.
This post-call hook runs before Backend.generate_from_context() attaches the generation id to the returned thunk. The token-id path therefore emits a post event with generation_id=None, so the tracing plugin cannot close the span opened by the pre event. The raw completion path also has no start time, causing the latency metric to record -0.001 s. Can we emit the post event after the wrapper assigns the id and initialise timing on this path? A focused tracing/metrics test would keep both behaviours covered.
A chat template writes adapter control tokens into the rendered prompt, and
they exist only in the token ids that render produced. Re-rendering a
conversation from `messages` on a later turn drops them, and `encode(decode(ids))`
is not the identity, so the re-derived prefix stops matching what the server
cached. Every KV block from the first divergence onward is lost, and each turn's
history is reinterpreted under the base model rather than the adapter that
produced it.
Adds an opt-in policy that keeps the ids instead of re-deriving them:
ctx = ChatContext(retain_token_ids=True)
`ChatContext` gains the policy plus the state it needs -- `sent_token_ids`,
`sent_model_id`, `sent_message_count` -- all propagated to descendant nodes and
cleared on a root reset, since ids are per-conversation. `PreTokenizedCBlock`
carries vocabulary ids that bypass the formatter entirely; it has no string form,
because there is no text whose re-encoding is guaranteed to reproduce them.
On `OpenAIBackend`, a retaining context routes to `/v1/completions` with
`prompt=[ids]` rather than posting messages, since the chat endpoint re-renders
and re-tokenizes server-side and would silently revert the policy. The new
turn's ids come from subtracting two fresh `/tokenize` renders -- the already-sent
messages, and the whole conversation -- and are spliced onto the retained
prefix. The retained ids are never compared against a re-render: they differ
from one exactly when a token fails to round-trip, which is the case this
policy exists to survive.
The tokenizer API is reached at the server root, not under `/v1`, which is where
vLLM serves it. Combinations the completions endpoint cannot honour are refused
rather than silently degraded: tool calling, streaming, a string-valued
reasoning level, and ids produced by a different model. A history that shrank --
a compactor dropping turns, or the token-budget truncation `view_for_generation`
applies once a model_id is bound -- is refused too, because the already-sent
side can no longer be identified.
Not implemented, deliberately: the checkpoint guards this policy needs to be
fully safe (switch_type == "multi", aLoRA-only placement, a chat_template_features
capability gate, and the bf16 control-token ceiling). All four require reading
the served checkpoint's config.json, and mellea exposes no control-token ids
today. A constant with no caller would read as a guard that exists.
Known gaps, both needing a live vLLM server to settle: the turn terminator is
appended with no overlap check against the emitted ids, which would double the
EOS token if vLLM reports it; and `return_token_ids` requires vLLM 0.10.2+,
below which no ids are reported and the policy silently never retains.
Assisted-by: Claude Code
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: noaa <noaa.kless@ibm.com>
The prefix that must stay byte-identical for a server cache hit is the one sent to the model, not the one stored as chat history, so retention cannot live only on the chat path. `_generate_from_intrinsic` now reuses a retained prefix via `_reuse_intrinsic_prefix_ids`, sending exact ids to `/v1/completions` and adapting the text-shaped reply back to a `ChatCompletion` so the existing result processor runs unchanged. It reuses without committing: the io.yaml rewriter replaces the conversation, so a rewritten request must never become the next canonical prefix. `sent_message_count` alone cannot identify a reusable prefix once a rewriter is involved, since an edited historical turn or dropped oldest turns leave the count intact. `ChatContext` gains `sent_prompt_digest`, a per-message fingerprint over role/content/tool_calls, canonicalized so the same turn fingerprints identically whether the chat serializer or the intrinsic path shaped it. Reuse is refused on a mismatch, so the prefix is proven unchanged rather than assumed. The digest is over text, which keeps it immune to the `encode(decode(ids))` non-identity this policy exists to survive. Requests the completions transport cannot honour fall back to the chat endpoint rather than degrade: tools, logprobs (score adapters read a different shape from that endpoint), a string reasoning level, and server-rendered documents. Also documents why this is not on `LocalHFBackend` -- Granite Switch activation is OpenAI-only until generative-computing#1018, so no local path injects the control tokens this preserves -- and records the concrete divergence in `derive_delta`: for `granite-switch-4.1-3b-preview` an adapter control token substitutes for the role marker (`100356` in place of `100264`), same length, so a length check misses it. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
…xtend `retain_token_ids` is an optimization, so a prefix that can no longer be extended should cost the cache, not the turn. `DeltaNotDerivable` was propagated deliberately, on the reasoning that falling back to messages would serve the re-rendering policy under a context that promised otherwise. But the conditions that raise it are ordinary: a caller edits an earlier message, a compactor drops the oldest turn, documents arrive mid-conversation. Each of those failed the generation outright rather than making it slightly slower. `_generate_from_context` now catches it and does not return, so execution falls through to the chat send already directly below -- correct output via a normal text render, with a warning recording that the prefix cache was re-primed and earlier control tokens dropped from it. Caught at the dispatch site rather than inside `_generate_via_token_ids_inner`, which has no access to the tools, extra_params, reasoning_params or backend_specific a chat request needs; there the fallback path is the next statement. Safe because `_build_prompt_ids` is the only raiser and runs before anything is sent, so the fall-through is a clean first attempt rather than a retry. Cheap because the model, shrink and digest guards all precede `/tokenize`: a digest mismatch is refused with no round trips at all. `TokenizeUnavailable` is deliberately still propagated. It means no usable `/tokenize` route exists, so retention can never work against that server, and swallowing it would leave `retain_token_ids` permanently inert with no signal -- the silent degradation this policy exists to make visible. A changed prefix is a per-turn condition; a missing route is misconfiguration. The retained ids are not cleared on fallback either, so a later turn that lines up with the prefix again resumes reuse rather than one divergent turn forfeiting the cache for the rest of the conversation. Assisted-by: Claude Code Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
Rationale that existed only to answer review questions -- why the id-space logic is not on `LocalHFBackend`, and why the retained state lives on `ChatContext` rather than on the backend -- belongs in the PR discussion, not in the source. Docstrings and comments describe what the code does; the argument for a choice already made reads as a conversation the reader was not part of. The `derive_delta` divergence example stays: its two causes are what make the "do not pass retained ids as `prev_ids`" rule comprehensible, so they document the contract rather than the review. Also corrects the `Raises:` entry on `_generate_via_token_ids`, which named `_generate_from_context` as the caller catching `DeltaNotDerivable`. The catch is in `_generate_from_chat_context_standard`. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
Every refusal message advised abandoning `retain_token_ids`, which was true when an unextendable prefix ended the request. It no longer is: the chat dispatch catches `DeltaNotDerivable` and re-renders the turn, so the output is unaffected and only the prefix-cache hit is lost. The advice now described a failure that does not happen. The two messages in `_build_prompt_ids` say what becomes of the turn, since every caller that reaches them falls back. The two in `derive_delta` state only that the ids cannot be extended -- it is reachable directly, where no fallback is guaranteed. Diagnostics (id counts, divergence index, likely causes) are unchanged. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
The token-id fast path generates via `/v1/completions`, whose reply is text-shaped, while it runs inside a chat turn where every consumer of `raw.response` dispatches on provider and expects chat shape. `Message._parse` reads `response["choices"][0]["message"]` and would raise KeyError, and `_retained_ids` read `token_ids` at the top level rather than on the choice. Add `_completion_choice_as_chat_response` to rebuild a completions choice into the exact shape a vLLM chat reply has under `return_token_ids` (message content plus `token_ids` on the choice), normalize `raw.response` before parsing, and read `token_ids` off the choice. The transport swap is now invisible downstream. Assisted-by: Claude Code Signed-off-by: noaa <noaa.kless@ibm.com>
The OpenAI token-id retention path materializes the reply eagerly (it derives the retained ids from it), so it returns an already-computed thunk. Such a thunk short-circuits `astream()`, so the `generation_post_call` hook astream normally fires never runs. The path fired it manually instead, but from inside `_generate_from_context` -- before the public `generate_from_context` wrapper assigns `_call.generation_id`. The hook therefore carried `generation_id=None`, so GenerationTracingPlugin could not close the open PRE span, and the path never set `_gen.start`, so LatencyMetricsPlugin recorded -0.001s. Add an opt-in `_CallInfo.fire_post_call_on_return` flag: a backend that returns an already-computed thunk sets it, and the wrapper fires the post-call once `generation_id` is assigned -- with the correct id and a real latency. The token-id path now stamps `_gen.start` before the request and sets the flag instead of firing the hook itself. The flag is off by default, so every other backend (and the computed-thunk DummyBackend) is unchanged. Tests for this change are committed separately. Assisted-by: Claude Code Signed-off-by: noaa <noaa.kless@ibm.com>
… path Add and extend unit/integration/e2e tests for the token-id history feature: - test_openai_token_id_postcall_unit: parametrized _retained_ids coverage of both vLLM id shapes (plain ints and "token_id:NNNN" strings) and every reject path; keep the achat-boundary KeyError regression test. - test_token_id_retention_telemetry: drive the deferred post-call through the real BackendTracingPlugin/LatencyMetricsPlugin with in-memory OTel exporters, asserting the generation span closes and duration is non-negative. - test_openai_token_id_e2e: add a THINKING=True exact-prefix reuse turn. - test_hook_call_sites: cover the wrapper firing post-call for a flagged already-computed thunk and skipping it for an unflagged one. All e2e verified against live vLLM (plain Granite + Granite Switch for the intrinsic path). No overlap: each test maps to one distinct behavior. Assisted-by: Claude Code Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
noaakl
force-pushed
the
feat/token-id-history
branch
from
September 7, 2026 12:14
2560052 to
7433066
Compare
Each of these let the prompt actually sent diverge from the conversation the caller described, with no error and no symptom beyond a fallen prefix-cache hit rate. `derive_delta` cannot catch any of them: it compares two fresh renders and never looks at the retained ids. - Template kwargs are recorded on the context (`ChatContext.sent_template_kwargs`) and the already-sent side is re-rendered under the kwargs it was actually sent with. A kwarg introduced mid-conversation appears in BOTH renders of the subtraction and cancels out, so `documents=[...]` supplied on turn 3 would reach neither the reused prefix nor the delta -- a RAG adapter activated against an empty context with every other guard passing. Now refused by name. The chat path also declines reuse when `documents` are present, matching the intrinsic path: they are rendered server-side by the chat template and are not a /tokenize parameter, so pre-tokenized ids omit them entirely. - A control-token ceiling of 188, the count over which the coded switch recovers a write address exactly in bf16. Past it two control tokens key one codeword and the memory head returns the mean of their expert ids -- an arbitrary adapter, no error in the output. A served model exposes no `adapter_token_ids`, so the ids are learned from a /tokenize diff (adapter render vs plain, positional: the control token substitutes for the role marker, so lengths match), cached per adapter, and probed only after every guard that can refuse without a round trip. Over the ceiling the prefix is dropped and the transcript re-rendered, counted by `token_id_reprefills`; a full render that is itself over raises, since re-baselining cannot reduce it. - The turn terminator is probed under the turn's own chat template kwargs and cached per kwargs rather than once globally. Granite 4.2 closes an assistant turn differently depending on `enable_thinking`, so a terminator derived under the template defaults was spliced into a sequence closed the other way -- one wrong id mid-conversation. The cache also moves to the instance: a terminator is a property of one server's chat template. Tests: 15 new cases in test_openai_token_id_guards_unit, each watched failing first. The documents gate is additionally verified by disabling it and confirming the test catches it. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
planetf1
reviewed
Sep 8, 2026
planetf1
left a comment
planetf1
left a comment
Contributor
There was a problem hiding this comment.
The exact-prefix e2e coverage is useful, but it demonstrates cache eligibility rather than cache reuse: it has no non-retaining control or server cache-metric assertion. I think that remains worth covering, alongside the client-visible retention observability follow-up.
| output._meta["retained_token_ids"] = retained | ||
| output._meta["retained_model_id"] = self._model_id | ||
| # Covers every rendered message plus the assistant turn just produced. | ||
| output._meta["retained_message_count"] = len(conversation) + 1 |
Contributor
There was a problem hiding this comment.
retained_message_count covers the generated assistant message (len(conversation) + 1), but the digest recorded below is only over conversation. On the next turn, _build_prompt_ids() verifies only the digest-length prefix, so editing that retained assistant message still passes the guard and reuses ids containing the original assistant response.
I reproduced this with a retained [user, assistant] prefix: changing only the assistant content returned the old retained prefix without raising DeltaNotDerivable. Could we make the digest cover every message claimed by sent_message_count, with a regression test for an edited assistant turn?
| # raise KeyError. Normalized FIRST, so the `_parse` below and | ||
| # `_retained_ids` further down both see one shape. | ||
| if isinstance(output.raw.response, dict): | ||
| output.raw.response = _completion_choice_as_chat_response( |
Contributor
There was a problem hiding this comment.
This transport conversion is not fully invisible downstream. _generate_from_raw() stores one CompletionChoice; _completion_choice_as_chat_response() rebuilds a response with only choices and usage. That drops top-level response metadata such as the provider response id/model and choice fields such as logprobs.
The normal chat post_processing() path that populates GenerationMetadata.response_id, response_model, and finish reasons does not run afterwards. Could we preserve or explicitly map this metadata when adapting the response shape, with a retaining-turn regression test?
| "_sent_prompt_digest", | ||
| ) | ||
|
|
||
| # Class-level defaults: `_rebuild_chat_context` builds nodes via `__new__` |
Contributor
There was a problem hiding this comment.
Manual compaction rebuilds a ChatContext through this helper, but _configure() restores only the compactor, token limit, and model id. The new nodes therefore use the class default retain_token_ids=False.
A caller who manually compacts a retaining context silently loses the policy for later turns. Could the rebuilt context retain the policy while clearing the now-invalid retained ids/count/digest? A compacted ChatContext(retain_token_ids=True) regression test would cover it.
| f"/tokenize replied without a 'tokens' list (got keys " | ||
| f"{sorted(payload)}), so its ids cannot be trusted." | ||
| ) | ||
| return [int(t) for t in tokens] |
Contributor
There was a problem hiding this comment.
This accepts coercible values rather than exact token ids: True, 1.5, or "12" are converted by int() and can then be sent through the exact-id path. A malformed or incompatible /tokenize response should fail as TokenizeUnavailable, rather than quietly constructing a different prompt.
This is a defensive edge case, but it seems worth validating that every entry is a non-boolean int before retaining it.
…robed ones The token-id retention ceiling guard undercounted control tokens. `_control_count` only recognizes ids in `_control_token_id_set`, which the per-adapter `/tokenize` probe populated LAZILY -- one id per adapter actually invoked. A control token for a registered-but-never-invoked adapter was therefore counted as zero, so a prompt genuinely over MAX_RETAINED_CONTROL_TOKENS could pass the guard and be sent, silently misrouting to the mean of two experts. Seed the full set from metadata instead. A composed Granite Switch model's `adapter_index.json` records every adapter's control-token id; carry it onto `Identity.control_token_id` (which survives both the deprecated shim and the `_discover_embedded_adapters` composed-Adapter rebuild) and union across all registered adapters up front. The `/tokenize` probe stays as the fallback only for an `adapter_name` passed as a raw template kwarg with no registered metadata. Verified the id source: `adapter_index.json`, `config.json` (adapter_token_ids), and the tokenizer (control-token string -> id) agree for all 12 adapters of gs_4.1_3b, so the metadata id is exactly what the served model receives. - core: `Identity` gains `control_token_id: int | None = None`. - adapters: `EmbeddedIntrinsicAdapter` parses `control_token.id` from the index and sets it on its `Identity`. - openai: `_seed_control_tokens_from_adapters()` unions every registered adapter's id; `_learn_control_tokens` seeds first and probes only an unregistered adapter. - tests: 5 cases incl. an un-invoked adapter's token now counted (0 -> 2). Assisted-by: Claude Code Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
`_rebuild_chat_context` builds every node with `__new__` and re-applies configuration by hand, so a field it does not enumerate falls back to the class default. `_retain_token_ids` was not enumerated: a caller who compacted a retaining context silently lost the policy for every later turn, with no error and no symptom beyond a fallen prefix-cache hit rate. The helper now takes the policy and `_configure` applies it, and both call sites (`WindowCompactor.compact`, `LLMSummarizeCompactor.compact`) pass the source context's value. The retained ids, count, digest and template kwargs are deliberately NOT carried: compaction has just dropped turns, so ids covering them describe a conversation that no longer exists. They stay at the class defaults, the same policy-versus-state split `_make_root` makes. Tests: two cases in a new TestCompactionPreservesContextPolicy -- one asserting the policy survives a rebuild and still propagates to nodes appended afterwards, one pinning that the now-invalid retained state does not. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
Each of these let the prompt actually sent differ from the conversation the caller
described, or dropped something the caller asked for, with no error and no symptom
beyond a fallen prefix-cache hit rate.
- The retained digest now covers every message `sent_message_count` claims. The
count included the assistant turn just produced (`len(conversation) + 1`) while
the digest fingerprinted only `conversation`, so the reply itself was reused on
trust: editing it left the guard passing and spliced ids carrying the ORIGINAL
text. Both facts are now derived from ONE list, whose assistant entry is
serialized through the same `to_chat_messages` -> `message_to_openai_message`
pipeline the next turn renders history with, so the fingerprint recorded here is
the one computed then. `_build_prompt_ids` additionally refuses a digest whose
length does not equal the count -- fewer leaves the newest messages unproven,
more reaches past the retained boundary -- and compares over `retained_count`
rather than the digest's own length.
- `_prompt_digest` fingerprints every field on the message instead of `role`,
`content` and `tool_calls`. These dicts are what goes on the wire, so a field the
chat template renders but the digest ignores is a prompt change the guard cannot
see; `reasoning_content` and `tool_call_id` were both invisible. Empty values are
dropped so absent-versus-`None` is still not a difference, which is what kept the
two serializers' output comparable in the first place.
- Response-side metadata survives the completions transport. `_generate_from_raw`
stores the per-choice dump, which carries no `id` or `model`; the chat
`post_processing()` that fills `mot.generation` never runs on this path, so every
retained turn -- and every batch completion -- reported `None` for `response_id`,
`response_model` and `finish_reasons`. They are now set from the enclosing
completion, with the finish reason taken from that thunk's own choice rather than
every choice in the batch, and `_completion_choice_as_chat_response` carries the
top-level identifiers plus the choice's `logprobs` through the shape adaptation.
- `/tokenize` ids are validated, not coerced. `int()` turned `True` into 1, `2.9`
into 2 and `"12"` into 12, and `int("abc")` raised a bare `ValueError` that
escaped the callers catching `TokenizeUnavailable` to fall back. These ids become
the prompt, and both sides of `derive_delta` come through this reader, so a
consistent corruption cancels out of the subtraction and reaches the server with
every other guard passing.
Also declines id reuse when `logprobs` are requested, from either the model-option
or the `extra_body` channel. The two endpoints report logprobs in incompatible
shapes and this path adapts only the reply's envelope, so reuse handed consumers a
payload they cannot read inside a reply labelled a chat completion.
`_reuse_intrinsic_prefix_ids` already declined on those grounds; the chat path now
matches it.
Tests: 13 new cases across the three token-id unit modules, each watched failing
first.
Assisted-by: Claude Code
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: noaa <noaa.kless@ibm.com>
Every existing assertion in this module compares ids on the CLIENT: it shows the
prompt was eligible for a cache hit, since its leading tokens are byte-identical to
what the server already saw. Only the server's own counters show that the hit
happened.
Two cases, both reading `vllm:prefix_cache_{queries,hits}_total` from `/metrics` as
a delta around turn 2, since turn 1 is what populates the cache:
- a retaining conversation hits a majority of the blocks it queries;
- a non-retaining conversation over the same two turns hits no more than the
retaining one. Asserted as `>=` rather than `>` deliberately -- on a model whose
text round-trips exactly the two are legitimately equal, and a strict `>` would
fail on precisely the servers where the policy is redundant rather than wrong.
Both skip when the server exports neither counter, so a build that reports only the
v0 hit-rate gauge does not fail them.
Assisted-by: Claude Code
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: noaa <noaa.kless@ibm.com>
A turn carrying `documents` used to forgo id reuse entirely, on the grounds that
`/tokenize` has no `documents` field and the ids would omit them. The field name is
missing there, but the render is not: vLLM binds `documents` as a chat-template
VARIABLE, merging it into the template kwargs before rendering
(`ChatCompletionRequest.build_chat_params`), and `/tokenize` accepts arbitrary template
kwargs. Passing it as `chat_template_kwargs={"documents": [...]}` therefore produces the
same prompt the generation endpoint would, so the prefix is reusable like any other.
This is the workload where the prefix cache is worth the most -- documents make prompts
long, and a RAG conversation re-sends them on every turn.
`_template_kwargs_with_documents` does the conversion, with the top-level field winning
over a hand-set `chat_template_kwargs["documents"]` so the precedence matches the
server's own `merge_kwargs` ordering. An empty list is treated as absent: it adds nothing
to the render, and recording the key would look like drift against a prefix that has it
unset.
The turn terminator is probed without `documents` (as it already was without
`adapter_name`): how a template closes an assistant turn does not depend on the system
block, and keying its cache on the documents would spend two `/tokenize` round trips per
new document set.
Documents that appear MID-conversation are still refused, by the existing template-kwargs
guard rather than by a gate of their own: they re-render the already-sent region, so they
land on both sides of the subtraction and cancel out of the delta -- reaching neither the
reused prefix nor the new turn. The refusal message already advises supplying
`documents=[...]` from the first turn.
Only the token-id path is touched. Every edit is inside the `ctx.retains_token_ids`
branch or in code only that branch calls; the chat send is unchanged, and a context
without `retain_token_ids` never reaches any of it.
Tests: 4 unit cases (reuse on a documents turn, the same for an intrinsic, the terminator
probe not keying on documents, and drift when they arrive late), each watched failing
first, plus an e2e that asserts turn 2 of a documents conversation both splices the exact
prefix and is HIT by the server's prefix cache -- the only check that can catch a
`/tokenize` render diverging from the chat template's, since the assembled prompt is
never returned. The e2e needs VLLM_TEST_BASE_URL and has not been run.
Assisted-by: Claude Code
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: noaa <noaa.kless@ibm.com>
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
3 participants
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.
A chat template writes adapter control tokens into the rendered prompt, and they exist only in the token ids that render produced. Re-rendering a conversation from
messageson a later turn drops them, andencode(decode(ids))is not the identity, so the re-derived prefix stops matching what the server cached. Every KV block from the first divergence onward is lost, and each turn's history is reinterpreted under the base model rather than the adapter that produced it.Adds an opt-in policy that keeps the ids instead of re-deriving them:
ChatContextgains the policy plus the state it needs --sent_token_ids,sent_model_id,sent_message_count-- all propagated to descendant nodes and cleared on a root reset, since ids are per-conversation.PreTokenizedCBlockcarries vocabulary ids that bypass the formatter entirely; it has no string form, because there is no text whose re-encoding is guaranteed to reproduce them.On
OpenAIBackend, a retaining context routes to/v1/completionswithprompt=[ids]rather than posting messages, since the chat endpoint re-renders and re-tokenizes server-side and would silently revert the policy. The new turn's ids come from subtracting two fresh/tokenizerenders -- the already-sent messages, and the whole conversation -- and are spliced onto the retained prefix. The retained ids are never compared against a re-render: they differ from one exactly when a token fails to round-trip, which is the case this policy exists to survive.The tokenizer API is reached at the server root, not under
/v1, which is where vLLM serves it. Combinations the completions endpoint cannot honour are refused rather than silently degraded: tool calling, streaming, a string-valued reasoning level, and ids produced by a different model. A history that shrank -- a compactor dropping turns, or the token-budget truncationview_for_generationapplies once a model_id is bound -- is refused too, because the already-sent side can no longer be identified.Known gaps, both needing a live vLLM server to settle:
return_token_idsrequires vLLM 0.10.2+, below which no ids are reported and the policy silently never retains.Assisted-by: Claude Code
Pull Request
Issue
Fixes #
Description
Testing
Attribution
Adding a new component, requirement, sampling strategy, or tool?
If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.
NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.