Skip to content

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
llm-d:mainfrom
shimib:feat/prompt-token-parity-followup
Open

shimib wants to merge 2 commits into
llm-d:mainfrom
shimib:feat/prompt-token-parity-followup

Conversation

@shimib

@shimib shimib commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

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.

  • Conditional-decode hits are checked. A served cache hit is a KV-reuse case, but the step streamed the worker's response straight to the client, so renderer drift there went unseen. The parity setup moves into a shared helper (parity.go) and conditional-decode runs it on its probe: the check rides the response writer, a streamed probe is asked for include_usage on the cloned body so a 412 continues the pipeline with the client's body untouched, and the step takes the same prompt_token_parity parameter.
  • Observe-only is a passthrough. 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. The writer now has a passthrough mode for the case where nothing can be rewritten: writes are forwarded as they arrive, Content-Length is 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_forwarded reads the body, not the route: the generate route, completions with the IDs as the prompt, or kv_transfer_params.prompt_token_ids on the chat-shaped routes. feat(coordinator): send rendered token ids to decode in kv_transfer_params #3741 needs no change here when it lands.
  • Malformed 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.
  • Docs: the api_type label values are APIType.String() (vllm_generate, not generate); the check needs a render step; conditional-decode's parameter in the architecture table and the sample config. ParamPromptTokenParity moves to the const block in params.go.
  • Field names live in 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, the stream_options keys, prompt_token_ids and prompt join the Field constants 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:

The coordinator's prompt-token parity check now also covers served `conditional-decode` cache hits, through the same `prompt_token_parity` parameter on that step. When nothing is rewritten the decode response is forwarded as it arrives, keeping `Content-Length`. The `ids_forwarded` label of `prompt_token_mismatch_total` follows the request body (including `kv_transfer_params.prompt_token_ids`), and the documented `api_type` value for the generate route is `vllm_generate`.

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, malformed stream_options left alone, observe-only keeps Content-Length; idsForwardedIn and streamOptions tables; usage writer passthrough (eager forwarding and Content-Length on JSON and SSE, rewrite mode still drops it, observe cap, ReadFrom).
  • golangci-lint run --config=./.golangci.yml --new-from-rev=upstream/main at v2.14.0: 0 issues. BenchmarkResponseWriterChunkedJSON unchanged (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 render tier, Qwen3-VL-8B on A100; vLLM serves every Prefer: if-available probe as a hit, so the conditional-decode path ran on every request):

scenario pipeline render tier result
match render → conditional-decode → decode Qwen3-VL-8B chat (batch, stream, stream + include_usage), completions, Responses (batch, stream), Messages, image by URL: all 200, prompt_token_mismatch_total absent
drift same Qwen3-8B template under --served-model-name Qwen/Qwen3-VL-8B-Instruct one mismatch (enable_thinking=false, 13 vs 9), logged with step=conditional-decode; client got a 200
match / drift render → decode as above same results with step=decode, as on #3751

In 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.

@shimib
shimib requested review from a team and roytman as code owners November 3, 2026 00:02
@github-actions github-actions Bot added size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files. kind/bug Categorizes issue or PR as related to a bug. kind/cleanup area/coordinator area/docs needs-rebase PR has merge conflicts and needs a rebase labels Nov 3, 2026
@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown
Contributor

@shimib This PR has merge conflicts with main. Please rebase on main and push to resolve them:

gh pr checkout 3889
git fetch origin main
git rebase origin/main
# resolve conflicts
git push --force-with-lease

@roytman roytman self-assigned this Nov 3, 2026
…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
shimib force-pushed the feat/prompt-token-parity-followup branch from 6b7a28d to 5b20b3f Compare November 3, 2026 15:56
@github-actions github-actions Bot removed the needs-rebase PR has merge conflicts and needs a rebase label Nov 3, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/coordinator area/docs kind/bug Categorizes issue or PR as related to a bug. kind/cleanup size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants