Conversation
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?
There was a problem hiding this comment.
Because HF doesn't reuse anything between turns. Every call re-renders and re-tokenizes the whole conversation and prefills it from scratch, so there's no cached prefix for retained ids to line up with. Saving the history as ids would buy nothing.
There was a problem hiding this comment.
HF / Transformers does allow for caching. We should ensure that this approach is compatible with Mellea's caching mechanism and/or the default transformers caching mechanism that is utilized when passing in a cache object to the generate call.
Additionally, if we don't prefix cache here, is there a risk that the subsequent calls to the switch model result in different kv values since the control tokens would no longer be there? I remember seeing a discussion about this (or something similar) but I can't find it.
There was a problem hiding this comment.
You're right, but it looks like Mellea currently doesn't use the HF KV cache across turns. Adding support for that would be a larger, separate change.
If we don't prefix cache here, the subsequent calls to the switch model result in different kv values
since the control tokens would no longer be there. But the cost is only that we recompute the prefix
every turn instead of reusing it. So it is less efficient, not incorrect.
| 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]
There was a problem hiding this comment.
Two things get lost, both on granite-switch-4.1-3b-preview.
A reply stored as text can't give back the ids the server cached. Decoding loses which merge produced each token, so re-encoding gives the canonical split, and nothing constrains sampling to that split:
encode(decode([71, 4896])) -> [15339] # 'h' + 'ello' becomes 'hello'
encode(decode([383, 75, 385])) -> [15339] # 'he' + 'l' + 'lo' becomes 'hello'
The control token is gone when history is re-rendered. adapter_name applies to the turn being generated, so turn 2 renders turn 1 without it:
turn 1 sent [..., 198, 100356, 78191, 100265] 100356 = <|answerability|>
turn 2 render [..., 198, 100264, 78191, 100265, ...] 100264 = <|start_of_role|>
Because it substitutes for the role marker, the divergence starts at index 12. Everything before that is a cache hit. Everything after is recomputed, and turn 1 no longer carries the adapter that wrote it.
There was a problem hiding this comment.
Okay; I believe vLLM does prefix cache based off the words of the response not the token values, but that would only help us for the first case. Thank you for the explanation.
There was a problem hiding this comment.
I believe vLLM's prefix cache is actually based on token IDs rather than the decoded words/text. It hashes blocks containing the token IDs, so two different token sequences that decode to the same text would not be considered the same cached prefix.
| # 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?
There was a problem hiding this comment.
Reusing ids is only safe if nothing changed before the new turn, measured against what this conversation sent last time. A backend has nowhere to key that.
There was a problem hiding this comment.
What I'm saying is that for any given input / response, the tokens in that input / response should remain the same between calls. Yes, there's a risk that the context gets changed, but a core principle behind Mellea is that spans (cblocks, components, mots) should be able to maintain their kv / tokens even if contexts change (if desired by the end user).
If we can extract / save the tokens from the backend, maybe we should do that regardless of the context type.
There was a problem hiding this comment.
If the conversation changed since we saved the ids, we notice and send text instead, which then tokenizes into the right ids. The only cost is recomputing the prompt. So storing ids on spans can't produce a wrong answer, and saving them regardless of context type sounds right to me. Where they live may be a broader design question than this PR though, so worth settling separately.
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.
2560052 to
7433066
Compare
planetf1
left a comment
planetf1
left a comment
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.
Thanks. vLLM's prefix cache is keyed by the exact token ids of the prompt. Reuse it and turn 2 only prefills the new turn; miss it and the whole conversation is prefilled again. Today Mellea sends messages, so vLLM re-renders and re-tokenizes the history on every turn, and the ids it gets are not the ones it cached. This PR sends the ids we already know it has, so the cache actually hits. Two reasons text can't do that: 1. Re-render diff. MVE with msgs = [{"role": "user", "content": "What is 2+2?"}]
tok.apply_chat_template(msgs, add_generation_prompt=True, adapter_name="answerability")
tok.apply_chat_template(msgs, add_generation_prompt=True) # turn 2 re-renders it this wayThe control token replaces the role marker, so the length is unchanged. Two consequences: the cache is invalid from index 12 on, and turn 1's history is now read as base instead of the adapter that wrote it. 2. encode/decode. On On io.yaml rewriting: ids are reused only if the rewritten messages still start with exactly the messages the server already has. If not, we don't reuse: the turn goes out as normal chat messages, so the output is correct and only the cache hit is lost. |
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>
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>
…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>
All three are comment-only; no behavior changes. `derive_delta` named the substituting control token by id alone (`100356` in place of `100264`), which is unverifiable while reading. Both are now spelled out: `<|answerability|>` for `<|start_of_role|>` (confirmed against the `granite-switch-4.1-3b-preview` tokenizer's added_tokens). The `/tokenize` reader's comment had lost its subject mid-sentence -- "Same rule `PreTokenizedCBlock` applies on the way out (`core/base.py`): an int, and `bool` is not one" -- leaving the rule it states unreadable. The retention comment claimed `_prompt_digest` "projects only role/content/tool_calls", so a differing `replay_reasoning` decision could not perturb it. That was true of an earlier digest and is now wrong twice over: the digest fingerprints EVERY field (openai.py:213), and `reasoning_content` is included deliberately, precisely because a chat turn and `_generate_from_intrinsic` disagree about replaying it. What actually keeps a `Chat -> Intrinsic` reuse from being refused is the canonicalization -- absent and `None` fingerprint alike -- so the comment now says that instead of contradicting the guard directly below it. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
The rebase onto `main` brought generative-computing#1631, which renamed the latency histogram from `gen_ai.client.operation.duration` to `mellea.llm.request.duration` (`mellea/telemetry/metrics.py:663`). `_duration_points` filtered on the old name, so it returned nothing and `test_deferred_post_call_records_non_negative_duration` failed on the assertion that a metric was recorded at all -- not on the value it was written to guard. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
…nit__ `with_sent_token_ids` built its copy with `type(self)()`. The rebase onto `main` brought generative-computing#1582, whose point is that a `ChatContext` subclass may take required constructor arguments: `add()` and `_make_root()` therefore build nodes with `__new__` and copy `_propagated_fields` rather than calling the initializer. This method was the one remaining node factory still calling it, so a subclass like class Tagged(ChatContext): def __init__(self, tag: str, **kw): ... raised `TypeError: Tagged.__init__() missing 1 required positional argument: 'tag'` the moment a backend recorded a turn -- and a backend records on EVERY retained turn, so id retention was unusable for such a subclass while `add()` on the same context worked fine. Now built the same way as its neighbours: `type(self).__new__`, `Context.__init__`, then the `_propagated_fields` copy. The return type narrows from `ChatContext` to `Self`, which is what it was already returning at runtime. Test watched failing first (with the `TypeError` above) in test_recording_ids_keeps_a_subclass_with_required_ctor_args. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
565cf83 to
87e5e51
Compare
… suite Review ask: gate the vLLM e2e module on GPU resources. `require_gpu(min_vram_gb=12)` matches `test_openai_intrinsics.py:31`, which serves the same `granite-switch-4.1-3b-preview` over vLLM, plus the same `CICD` skip that module uses. The endpoint here is remote (`VLLM_TEST_BASE_URL`), so that env-var skip is what really decides whether these run; the resource marker keeps the declaration consistent for a runner that serves the model itself rather than pointing at someone else's. The `slow` marker is deliberately NOT added yet: the ask was "slow if >1 min" and these tests have never executed against a live endpoint, so there is no measured duration to justify it -- and `slow` is deselected by default (`pyproject.toml` addopts), so adding it on a guess would hide the module from the runs that would establish the number. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
`_retained_ids` appended the turn terminator unconditionally. With `return_token_ids` on, vLLM reports the stop token the model sampled, so the emitted ids can already end with the terminator or a prefix of it. Measured against a live vLLM serving granite-switch-4p1-3b: emitted ids ending `13, 100257` plus a `[100257, 198]` terminator produced `13, 100257, 100257, 198` where the chat template renders `13, 100257, 198`. The retained sequence then diverged from what the server had cached at that index -- reuse fell from 48 tokens to 32 -- and carried a stray end-of-turn token into every later turn's history. Append only the part the emitted ids do not already carry, matched as a prefix of the terminator rather than "ends with any terminator token": a reply ending in `198` does not mean the `100257` before it was emitted. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
A caller without access to the server's cache counters had no way to tell whether a turn was sent as retained ids or fell back to a chat render, nor how large the reused prefix was. `GenerationMetadata.token_id_retention` now carries `reused_prompt_tokens`, `new_prompt_tokens` and `prompt_tokens` for a turn that went out as ids, and stays `None` otherwise -- so its presence alone distinguishes reuse from a fallback. It is measured from the prompt actually built rather than from the context's state, so a turn whose prefix was dropped at the control-token ceiling honestly reports zero reuse. `TokenIdRetentionMetricsPlugin` emits `mellea.token_id_retention.turns` (tagged `reused=true|false`) and `.reused_prompt_tokens` from the `generation_post_call` hook, and the same values reach the generation span as `mellea.token_id_retention.*` attributes. Named for what it measures: this is what the client sent, not evidence that the server's prefix cache hit. Only the server's own counters can say the latter, and the docs say so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
… overlap Unit: the signal reports no reuse on a first turn, the prefix size on a splicing turn, zero on a re-baselined turn whose context still holds ids, and stays `None` on a turn that fell back to a chat send. Plus a parametrized case for the terminator overlap, pinning the rule as PREFIX overlap rather than "ends with any terminator token". Integration: the same values reach the in-memory OTel exporters -- the metric tagged `reused=true`, a no-reuse turn tagged `reused=false`, nothing at all for a non-retaining turn, and the span attributes present and absent to match. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
Run against a live vLLM serving granite-switch-4p1-3b with prefix caching on:
11/11, stable across repeat runs.
`test_retaining_ids_reuses_more_cache_than_re_rendering` is the outcome the
policy exists for, read from the server's own counters: 272 vs 240 reused tokens
on equal queries over the same five-turn adapter conversation. Three conditions
had to hold for the gap to be measurable, each established by measurement rather
than assumed, and each recorded in the docstring:
- every turn generated under an adapter. Control tokens accumulate one per turn
in the retained sequence (1, 2, 3 measured) while a re-render carries exactly
one, at the current generation prompt.
- long replies. vLLM caches whole 16-token blocks and never a trailing partial
one, so at MAX_NEW_TOKENS=16 the divergent control token lands in a block
that was never cached and both arms reused exactly 224 tokens. At 200 the
divergence falls inside complete blocks and the gap appears.
- greedy decoding, or the arms' transcripts differ in length and the comparison
is meaningless (an early version measured 376 vs 465 queried tokens purely
from reply-length variance).
The base-only comparison stays as a separate "never worse" test, since without
adapter turns both transports send the same ids and equal reuse is correct. It
also no longer compares hit RATES across arms that shared a transcript: the
first arm warmed the cache the second measured (94.12% vs 96.97%), and the
retained prompt's extra token moved the denominator without changing reuse.
Adds the adapter-to-base transition case, `slow`, and
`VLLM_TEST_ADAPTER_SOURCE`, which lets adapter metadata resolve from a local copy
while the request carries a served model name that only exists inside the
server's container.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: noaa <noaa.kless@ibm.com>
A `ChatContext(retain_token_ids=True)` was accepted by every backend but read by only one, so on any other backend the policy was silently inert: history was re-rendered from text on every turn, doing exactly the thing the flag was set to avoid. `Backend._supports_token_id_retention` defaults to `False` and `OpenAIBackend` sets it `True`. A retaining context on a backend that does not support it now warns once per instance rather than per turn. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
The docstring quality gate flagged it, and there is a real behaviour to state: `int(i)` over the incoming ids raises `TypeError` on `None` and `ValueError` on a non-numeric string, while a float is silently truncated. The backend that produces these ids validates them on the way in, so this is a backstop rather than the guard. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
d3fc51a to
cea3f3f
Compare
Rebased on main, and all three levels are now covered. |
…data Adapter discovery is client-side: it reads `adapter_index.json` and `io_configs/` to map adapter names to control-token ids. When the served model name is a path inside the server's container -- the normal shape for a locally-served checkpoint -- the client cannot read it and falls back to treating it as a Hub repo id, which raises `HFValidationError` on any path with more than one "/". Three cases hit that and FAILED where they should have skipped, so a run without `VLLM_TEST_ADAPTER_SOURCE` reported errors that said nothing about the code under test. They now skip with a reason naming the variable and what to point it at. Verified both ways against a live vLLM: 8 passed / 3 skipped without the variable, 11 passed with it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: noaa <noaa.kless@ibm.com>
|
Other than Jake's point (may need a followup issue to be opened?) all looks good to me and I'd be happy to approve once that thread is figured out. |
3 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.