feat(coordinator): check prompt-token parity between prefill and decode by shimib · Pull Request #3751 · llm-d/llm-d-router

shimib · 2026-10-29T20:57:58Z

What type of PR is this?

/kind feature

What this PR does / why we need it:

With forwarded prompt_token_ids (#3656, #3706), prefill runs on the token IDs the render step produced while decode may render the prompt itself: vLLM ignores the forwarded IDs when messages carry non-text content or echo is set, and vllm-project/vllm#62333 would add servers without a KV connector. If the two renders disagree (tokenizer, chat template, processor version, chat_template_kwargs), nothing fails: decode misses prefill's KV cache and recomputes the prompt. The request succeeds, TTFT jumps, and the only trace is a prefix-cache miss.

The decode step now compares the worker's usage.prompt_tokens (input_tokens on Responses and Messages) with len(TokenIDs), the count sent to prefill. A difference increments prompt_token_mismatch_total{model_name, api_type, ids_forwarded} and logs once per request. The response to the client is never changed.

New decode parameter prompt_token_parity: off, warn (default: count and log), error (also step_errors_total{error_code="prompt_token_mismatch"}). It is a post-response signal in every mode; the body is already on the wire when the usage block arrives, so error records rather than fails.
ids_forwarded is true for completions (the step sends the rendered IDs as the prompt) and the generate route, false for the routes that forward the client's text today. When the forwarded-IDs decode leg lands, a mismatch with true means the worker ignored the IDs; with false it means the renderers disagree.
Streamed chat and completions legs are asked for stream_options.include_usage when the client did not, so the final usage frame exists to compare, and that frame is dropped from the client's stream. A frame the client asked for is kept. stream_options reaches the step as raw JSON (the server decodes only the fields it inspects), which the check handles.
pkg/common/usage gains UsageBlock (top-level usage, response.usage for Responses stream events, message.usage for Messages message_start) and PromptTokens. The response writer takes Options{CachedTokens, OnUsage, DropUsageOnlyFrame} through NewResponseWriter; NewCachedTokensResponseWriter is unchanged. The cached-token rewrite is untouched: the existing tests pass as they were.
Docs: decode row in coordinator_architecture.md, the parameter in config/coordinator/coordinator.yaml, the metric and the new error code in metrics.coord.md.

Which issue(s) this PR fixes:

Fixes #3748. Part of #3656.

Release note:

The coordinator's `decode` step gains a `prompt_token_parity` parameter, `off`, `warn` (default) or `error`. It compares the decode worker's reported prompt token count with the rendered token count sent to prefill and counts a difference in `llm_d_coordinator_prompt_token_mismatch_total`, labelled by `model_name`, `api_type` and `ids_forwarded`; `error` also records a decode step error. The response is never changed. Streamed chat and completions legs request `stream_options.include_usage` when the client did not, and that trailing frame is dropped from the client's stream.

Validation:

Unit:

go test ./pkg/common/usage/ ./pkg/coordinator/...: passes. New: 6 tests in usage (block lookup across the four response shapes, observer on JSON and SSE, injected-frame drop, observer plus cached-token rewrite together) and 7 in steps (mismatch counts and keeps the response, match is silent, error mode records the step error, off and no-render skip, unknown mode rejected, streaming injects and drops the frame, client-requested frame kept for both body representations).
golangci-lint run --config=./.golangci.yml --new-from-rev=upstream/main at v2.14.0: 0 issues. gofmt, go vet, go mod tidy: clean.

Real environment (coordinator binary from this branch, gateway pointed directly at a vLLM worker, render at a CPU vllm launch render tier, Qwen3-VL-8B on A100):

scenario	render tier	decode worker	result
match	Qwen3-VL-8B	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, `warn`	Qwen3-8B template under `--served-model-name Qwen/Qwen3-VL-8B-Instruct`	Qwen3-VL-8B	exactly one mismatch: `chat_template_kwargs.enable_thinking=false` renders to 13 tokens on the Qwen3 template and 9 on the VL template; the other 7 requests match; the client received a 200 with the worker's count
drift, `error`	same	same	same mismatch plus `step_errors_total{error_code="prompt_token_mismatch",step="decode"} 1`

Also checked there: the injected usage frame does not reach the client, a client-requested one does, and the Responses stream has the same event count with and without the coordinator. The real-environment run is what found the raw-JSON stream_options case; the regression test for it is in this PR.

github-actions · 2026-10-30T18:15:05Z

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

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

With forwarded prompt_token_ids, prefill runs on the token IDs the render step produced while decode may render the prompt itself: vLLM ignores the forwarded IDs when messages carry non-text content or echo is set. If the two renders disagree nothing fails; decode misses prefill's KV cache and recomputes the prompt, and the only trace is a slower request. The decode step now compares the worker's usage.prompt_tokens (input_tokens on Responses and Messages) with len(TokenIDs), the count sent to prefill. A difference increments prompt_token_mismatch_total{model_name, api_type, ids_forwarded} and logs once per request; the response is never changed. A new decode parameter prompt_token_parity selects off, warn (default) or error, where error also records step_errors_total with the new prompt_token_mismatch code. Streamed chat and completions legs are asked for stream_options.include_usage when the client did not, so the final usage frame exists to compare, and that frame is dropped from the client's stream. The usage package gains UsageBlock and PromptTokens, and its response writer takes Options with an OnUsage observer and DropUsageOnlyFrame beside the existing cached-token rewrite. NewCachedTokensResponseWriter is unchanged. Part of llm-d#3656. Fixes llm-d#3748. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Shimi Bandiel <shimib@google.com>

ahg-g · 2026-10-30T20:09:46Z

+func PromptTokens(block any) (int, bool) {
+	fields, ok := block.(map[string]any)
+	if !ok {
+		return 0, false
+	}
+	for _, field := range []string{reqcommon.FieldPromptTokens, reqcommon.FieldInputTokens} {
+		if n, ok := intValue(fields[field]); ok {
+			return n, true
+		}
+	}
+	return 0, false
+}


Count the Messages prompt as input_tokens + cache_read_input_tokens + cache_creation_input_tokens, the way the EPP's Anthropic parser already does.

if n, ok := intValue(fields[reqcommon.FieldPromptTokens]); ok { return n, true } n, ok := intValue(fields[reqcommon.FieldInputTokens]) if !ok { return 0, false } read, _ := intValue(fields["cache_read_input_tokens"]) created, _ := intValue(fields["cache_creation_input_tokens"]) return n + read + created, true

Why: In the Messages schema input_tokens excludes cached tokens (anthropic.go applyInputTokens documents and sums this). When the decode worker reports the cache split, a correctly working P/D request (decode reads prefill's KV, so most of the prompt is cache_read_input_tokens) shows input_tokens far below len(TokenIDs) and is counted as a mismatch. The metric fires precisely when the KV transfer worked. Add a Messages case with cache_read_input_tokens to decode_parity_test.go.

Why it's safe: a worker that reports the full prompt in input_tokens with no cache fields sums to the same number as today.

Good catch, it would fire exactly when the handoff worked. Fixed in 45e6cc4: Messages now sums the three fields like the Anthropic parser; the other routes already include cached tokens in prompt_tokens. Test added.

ahg-g · 2026-10-30T20:09:46Z

+	options := streamOptions(reqCtx.Body)
+	if include, ok := options[fieldIncludeUsage].(bool); ok && include {
+		return false
+	}
+	merged := make(map[string]any, len(options)+1)
+	for k, v := range options {
+		merged[k] = v
+	}
+	merged[fieldIncludeUsage] = true
+	reqCtx.Body[reqcommon.FieldStreamOptions] = merged
+	return true


Skip the injection when the client set stream_options.continuous_usage_stats.

options := streamOptions(reqCtx.Body) if include, ok := options[fieldIncludeUsage].(bool); ok && include { return false } if continuous, _ := options["continuous_usage_stats"].(bool); continuous { return false }

Why: vLLM sends per-chunk usage only when include_usage && continuous_usage_stats, so a client that sent continuous_usage_stats: true alone gets no usage today. Setting include_usage turns on a usage block in every chunk; those chunks carry choices, so DropUsageOnlyFrame keeps them and the client's stream changes shape, against the "response is never changed" contract. It also makes every chunk pass the bytes.Contains(data, "usage") check, so the writer unmarshals each frame on the streaming path.

Agreed, verified in should_include_usage. Injection and check are now skipped when continuous_usage_stats is set. Test added.

…_usage_stats streams alone Review follow-ups from llm-d#3751. vLLM's Messages route subtracts cache_read_input_tokens and cache_creation_input_tokens from input_tokens, so a P/D request whose prompt was served from prefill's KV reported a short count and the parity check fired exactly when the transfer worked. usage.PromptTokens now sums the three fields on Messages, the way the EPP's Anthropic parser does; the chat and completions prompt_tokens already include cached tokens. vLLM reports usage on every chunk only when include_usage and continuous_usage_stats are both set. A client that sent continuous_usage_stats alone gets no usage today; injecting include_usage would put a usage block in each of its content chunks, which the writer cannot drop. The decode step now skips the injection and the check for that client. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Shimi Bandiel <shimib@google.com>

roytman · 2026-10-30T22:15:13Z

Oops, I have several comments here, but it is merged.
I will add them here, and @shimib , please see if they can be part of a follow-up PR.

roytman · 2026-10-30T22:23:10Z

A conditional-decode cache hit bypasses the parity check entirely. On a hit the conditional-decode step streams the response straight to the client and ends the pipeline, so no usage block is observed and include_usage is never injected on streamed hits. That path is a KV-reuse case where renderer drift is silent, which is what this metric is for; newParityCheck already skips only when no tokens were rendered, not when prefill did not run, so the comparison would be meaningful there too. Either cover it in this PR or file a follow-up issue and note the exclusion in the parameter docs.

roytman

@shimib, please take a look.
The most important comment is about conditional decoding.

roytman · 2026-10-30T22:23:50Z

+// and without a parity check the client gets the decode worker's bytes
+// untouched.
+func (s *DecodeStep) responseWriter(reqCtx *pipeline.RequestContext, parity *parityCheck, injectedUsage bool) (http.ResponseWriter, func() error) {
+	if reqCtx.PrefillCachedTokens == nil && parity == nil {


With prompt_token_parity defaulting to warn, this condition now routes every decode response that has rendered tokens through the rewriter, including responses where CachedTokens is nil and no frame was injected. In that mode no path can rewrite a byte, but the writer's writeHeader deletes Content-Length unconditionally (pre-existing behavior in cached_tokens.go, newly exposed by default here), so clients get chunked framing where they previously saw a bounded body, and a non-streaming body is buffered and unmarshalled into a map purely to observe usage.

Could the Del("Content-Length") in the writer be gated on r.cachedTokens != nil || r.dropUsageOnly, so a response that cannot change keeps its framing? The buffering itself is inherent to observing usage on the JSON path, but the default-on cost deserves a sentence in the PR description either way.

roytman · 2026-10-30T22:25:30Z

+// streamOptions returns the request's stream_options object. The server
+// decodes only the fields the pipeline inspects and keeps the rest as raw
+// JSON (decodeRequestBody), so the value is usually json.RawMessage here.
+func streamOptions(body map[string]any) map[string]any {


streamOptions returns nil both when stream_options is absent and when it is present but not an object (the server keeps it as raw JSON). injectIncludeUsage then overwrites the client's value with a valid object, so a request like {"stream":true,"stream_options":"x"} that the worker would reject with 400 is silently normalized and served, and its trailing usage frame is dropped. That contradicts the handler's contract that non-inspected fields are forwarded unchanged.

Could streamOptions distinguish absent (inject) from present-but-not-an-object (skip the check and leave the bytes untouched)?

roytman · 2026-10-30T22:26:20Z

+// usage reports a prompt token count different from the rendered one:
+// "off" skips the check, "warn" (default) counts and logs it, "error" also
+// records a step error. The response to the client is never changed.
+const ParamPromptTokenParity = "prompt_token_parity"


nit: the const block in params.go is the canonical home for step parameter keys (ParamKVConnector, ParamECConnector). Could ParamPromptTokenParity move there?

roytman · 2026-10-30T22:27:47Z

+// Responses stream events wrap the response object, and the Messages
+// message_start event wraps the message object.
+const (
+	fieldResponse = "response"


The canonical home for response-body field names is pkg/common/request, which already holds FieldUsage, FieldPromptTokens and FieldInputTokens that this package imports. fieldResponse, fieldMessage and fieldChoices describe the same class of contract, and choices and message also exist as unexported constants in pkg/sidecar/proxy, so landing this creates two names each that are free to diverge. Same point for the two cache-split constants below and for fieldIncludeUsage/fieldContinuousUsageStats in decode.go, whose parent key FieldStreamOptions already lives there.

Could these move to pkg/common/request/constants.go next to their siblings?

roytman · 2026-10-30T22:28:42Z

+	return step.(*DecodeStep), reqCtx, recorder, nil
+}
+
+func TestDecodeStep_PromptTokenParity_MismatchCountsAndKeepsResponse(t *testing.T) {


Three gaps against the issue's test list:

"a missing usage is a no-op" has no step-level test. The writer-level tests only cover UsageBlock returning false in isolation; a decode reply with no usage block (the older-worker case the issue names) should assert no counter increment and an unchanged body.

Every step-level test runs the chat route, so ids_forwarded="true" (completions and generate) and its api_type label values have no coverage. That boolean is the diagnostic distinction of the feature; a regression in it would be invisible. One table case over the completions path would close it.

Issue coordinator: prompt-token parity check between prefill and decode #3748 asks for an e2e with mismatched render and worker images asserting the counter increments and the request still succeeds. The PR adds none and does not say why, while claiming "Fixes coordinator: prompt-token parity check between prefill and decode #3748". Either add it or state the deferral in the description.

roytman · 2026-10-30T22:30:55Z

 *   **Note:** Not redundant with `execution_path_total` since a deferred probe does not indicate whether the subsequent path was `prefill-decode` or `encode-prefill-decode`.

+#### `prompt_token_mismatch_total` (Counter)
+*   **Labels:** `model_name`, `api_type` (`chat_completions`, `completions`, `responses`, `messages`, `generate`), `ids_forwarded` (`true` when the decode leg carried the rendered token IDs, `false` when it forwarded the client's text)


The generate route records vllm_generate (APITypeVLLMGenerate.String()), not generate, so the documented set does not include a value the metric emits. A PromQL equality filter written from this line matches no series for that route.

Suggested change

* **Labels:** `model_name`, `api_type` (`chat_completions`, `completions`, `responses`, `messages`, `generate`), `ids_forwarded` (`true` when the decode leg carried the rendered token IDs, `false` when it forwarded the client's text)

* **Labels:** `model_name`, `api_type` (`chat_completions`, `completions`, `responses`, `messages`, `vllm_generate`), `ids_forwarded` (`true` when the decode leg carried the rendered token IDs, `false` when it forwarded the client's text)

roytman · 2026-10-30T22:31:45Z

      # params:
      #   kv_connector: kv-nixl     # overrides pipeline.kv_connector
+      #
+      #   prompt_token_parity compares the prompt token count the decode


Could the comment note that the check requires a render step? Without one there are no rendered token IDs, newParityCheck returns nil, and the parameter silently does nothing. An operator enabling error on a pipeline without render would get no signal and no indication why.

…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 requested review from a team and roytman as code owners October 29, 2026 20:58

shimib requested review from ahg-g, elevran, hexfusion, liu-cong and vMaroon October 29, 2026 20:58

shimib force-pushed the feat/prompt-token-parity branch from 1802829 to edc2aa3 Compare October 30, 2026 18:26

github-actions Bot removed the needs-rebase PR has merge conflicts and needs a rebase label Oct 30, 2026

ahg-g reviewed Oct 30, 2026

View reviewed changes

github-actions Bot added size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files. and removed size/XL Denotes a PR that changes 500-999 lines, ignoring generated files. labels Oct 30, 2026

ahg-g approved these changes Oct 30, 2026

View reviewed changes

llm-d-router-release-notes Bot added this pull request to the merge queue Oct 30, 2026

Merged via the queue into llm-d:main with commit a528e87 Oct 30, 2026
56 checks passed

wguo7 mentioned this pull request Oct 30, 2026

fix(coordinator): rewrite cached tokens in streamed Responses usage #3596

Merged

2 tasks done

roytman reviewed Oct 30, 2026

View reviewed changes

This was referenced Oct 31, 2026

feat(coordinator): add mm_payload render modes (tensors, none) #3756

Merged

fix(coordinator): run the prompt-token parity check on conditional-decode hits and stop buffering when nothing is rewritten #3889

Open

liu-cong mentioned this pull request Nov 3, 2026

docs(review): add maintainer review playbook and rule bundle #3854

Draft

3 tasks done

This was referenced Nov 3, 2026

coordinator: prompt-token parity check between prefill and decode #3748

Open

[EPIC] Tokens-in/tokens-out path: render once, route by token IDs, decode with output_mode or derender #3656

Open

	* Labels: `model_name`, `api_type` (`chat_completions`, `completions`, `responses`, `messages`, `generate`), `ids_forwarded` (`true` when the decode leg carried the rendered token IDs, `false` when it forwarded the client's text)
	* Labels: `model_name`, `api_type` (`chat_completions`, `completions`, `responses`, `messages`, `vllm_generate`), `ids_forwarded` (`true` when the decode leg carried the rendered token IDs, `false` when it forwarded the client's text)

Conversation

shimib commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

roytman commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

roytman commented Oct 30, 2026

Uh oh!

roytman left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shimib commented Oct 29, 2026 •

edited

Loading

roytman commented Oct 30, 2026 •

edited

Loading