Repository navigation
fix(coordinator): run the prompt-token parity check on conditional-decode hits and stop buffering when nothing is rewritten - #3889
Open
shimib wants to merge 2 commits into
Conversation
Contributor
|
@shimib This PR has merge conflicts with gh pr checkout 3889
git fetch origin main
git rebase origin/main
# resolve conflicts
git push --force-with-lease |
…code hits and stop buffering when nothing is rewritten Follow-ups to llm-d#3751 from the post-merge review on that PR. A served conditional-decode hit is a KV-reuse case like any other, but the step streamed the worker's response straight to the client, so a renderer drift there went unseen. The parity setup moves out of the decode step into a shared helper, and conditional-decode runs it on its probe: the check is installed on the writer, a streamed probe is asked for include_usage on a cloned body so a 412 continues the pipeline with the client's body untouched, and the step takes the same prompt_token_parity parameter. With warn as the default, every decode response with rendered tokens went through the usage writer, which dropped Content-Length unconditionally and held a non-streaming body until the end before forwarding it, purely to read usage. The writer now has a passthrough mode for the observe-only case: it forwards every write as it arrives, keeps Content-Length, and reads usage on the side (complete SSE lines as they come, a bounded copy of a JSON body at the flush). The cached-token rewrite and the injected-frame drop are unchanged. ids_forwarded now reads the leg's body instead of the route: the generate route, completions with the IDs as the prompt, or kv_transfer_params .prompt_token_ids on the chat-shaped routes. A step that starts forwarding IDs (llm-d#3741) needs no change here. A stream_options that is present but not an object is forwarded as the client sent it and the check is skipped, instead of being replaced by a valid object. Tests: missing usage is a no-op, the completions route carries its own label values, a malformed stream_options is left alone, observe-only keeps Content-Length, and the conditional-decode hit, miss, streaming and mode cases. Docs: the api_type values are APIType.String() (vllm_generate, not generate), the check needs a render step, and conditional-decode gained the parameter in the architecture table and the sample config. The parameter key moves to the const block in params.go. The field-name constants stay local until llm-d#3756 lands, which moves neighbouring names into pkg/common/request. Part of llm-d#3748. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Shimi Bandiel <shimib@google.com>
shimib
force-pushed
the
feat/prompt-token-parity-followup
branch
from
November 3, 2026 15:56
6b7a28d to
5b20b3f
Compare
The usage-wrapper keys (choices, response, message), the Messages cache split fields, the stream_options keys, and prompt_token_ids described the same request and response contract as the Field constants in pkg/common/request but lived as unexported copies next to their users. Move them beside their siblings so there is one name per field, and name prompt with the rest. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Shimi Bandiel <shimib@google.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
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.
What type of PR is this?
/kind bug
/kind cleanup
What this PR does / why we need it:
Follow-ups to #3751 from @roytman's post-merge review there.
parity.go) and conditional-decode runs it on its probe: the check rides the response writer, a streamed probe is asked forinclude_usageon the cloned body so a 412 continues the pipeline with the client's body untouched, and the step takes the sameprompt_token_parityparameter.warnas the default, every decode response with rendered tokens went through the usage writer, which droppedContent-Lengthunconditionally and held a non-streaming body until the end. The writer now has a passthrough mode for the case where nothing can be rewritten: writes are forwarded as they arrive,Content-Lengthis kept, and usage is read on the side (complete SSE lines as they come; a copy of a JSON body, capped at 32 MiB, at the flush). The cached-token rewrite and the injected-frame drop are unchanged.ids_forwardedreads the body, not the route: the generate route, completions with the IDs as the prompt, orkv_transfer_params.prompt_token_idson the chat-shaped routes. feat(coordinator): send rendered token ids to decode in kv_transfer_params #3741 needs no change here when it lands.stream_options(present but not an object) is forwarded as the client sent it and the check is skipped, instead of being replaced by a valid object.api_typelabel values areAPIType.String()(vllm_generate, notgenerate); the check needs arenderstep; conditional-decode's parameter in the architecture table and the sample config.ParamPromptTokenParitymoves to the const block inparams.go.pkg/common/request(second commit, rebased on feat(coordinator): add mm_payload render modes (tensors, none) #3756): the usage wrappers (choices,response,message), the Messages cache-split fields, thestream_optionskeys,prompt_token_idsandpromptjoin theFieldconstants there instead of unexported copies next to their users.Not in this PR, on purpose: an automated e2e for a render/worker mismatch is not possible in the kind suite today: the simulator serves both the render and the worker role with one tokenizer, so the two counts cannot differ; the manual GPU run on #3751 stands, and I'll note the gap on #3748.
Which issue(s) this PR fixes:
Part of #3748.
Release note:
Validation:
go test ./pkg/common/usage/ ./pkg/coordinator/...: passes. New tests: conditional-decode hit mismatch, 412 miss (nothing written, no injection leak), streaming hit inject-and-drop, off/error/unknown modes; decode missing-usage no-op, completions label values, malformedstream_optionsleft alone, observe-only keepsContent-Length;idsForwardedInandstreamOptionstables; usage writer passthrough (eager forwarding andContent-Lengthon JSON and SSE, rewrite mode still drops it, observe cap,ReadFrom).golangci-lint run --config=./.golangci.yml --new-from-rev=upstream/mainat v2.14.0: 0 issues.BenchmarkResponseWriterChunkedJSONunchanged (128/149 allocs/op).Real environment (coordinator binary from this branch in a pod, gateway pointed at a vLLM worker, render at a CPU
vllm launch rendertier, Qwen3-VL-8B on A100; vLLM serves everyPrefer: if-availableprobe as a hit, so the conditional-decode path ran on every request):include_usage), completions, Responses (batch, stream), Messages, image by URL: all 200,prompt_token_mismatch_totalabsent--served-model-name Qwen/Qwen3-VL-8B-Instructenable_thinking=false, 13 vs 9), logged withstep=conditional-decode; client got a 200step=decode, as on #3751In every run a non-streaming chat response arrived with its
Content-Length(no chunked encoding) and its first byte with the headers; before this change the observe-only writer stripped the header and held the body until the end. Injected usage frames were dropped and client-requested ones kept on streamed replies.