feat(backends): preserve token ids across chat turns - #1592
Conversation
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>
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.
| # 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. |
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?
| 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. |
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]
| # 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, | ||
| ) | ||
|
|
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.
| 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. |
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?
2 participants
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.