Skip to content

[Bugfix][Frontend] Require a KV connector for decode-side prompt id reuse - #62333

Open
he-yufeng wants to merge 1 commit into
vllm-project:mainfrom
he-yufeng:fix/kv-reuse-requires-connector
Open

he-yufeng wants to merge 1 commit into
vllm-project:mainfrom
he-yufeng:fix/kv-reuse-requires-connector

Conversation

@he-yufeng

Copy link
Copy Markdown
Contributor

Fixes #62307.

Why

kv_transfer_params["prompt_token_ids"] is the disaggregated-serving reuse path: the decode stage skips templating and tokenization because the prefill stage's ids arrive alongside the transfer coordination (do_remote_prefill / do_remote_decode). The reuse funnel _reused_prompt_token_ids checked the ids' shape (#55771) and vocabulary (#59555), but never checked that the server actually runs a KV transfer connector. On a plain server, any client could send harmless-looking messages plus attacker-chosen ids and have the model run on the ids instead, with the chat template (and any system prompt it carries) silently dropped. First affected release: 0.27.0 (#48145).

What changed

One gate in vllm/renderers/online_renderer.py::_reused_prompt_token_ids: after the existing echo / non-text-content bail-outs (which neutralize the ids and still render messages), the ids are now rejected with a 400 unless renderer.config.kv_transfer_config is set. The connector is the capability boundary the reuse contract already assumes, and both entry funnels (preprocess_chat and the Harmony branch) pass through this helper, so chat completions and responses are covered by the same check.

Not done deliberately: restricting reuse to /v1/chat/completions only. The docs example uses chat completions, but disaggregated deployments front the same engine with the responses API too, and the connector gate is what removes the deceptive surface; an endpoint ban would break those deployments without adding safety. Happy to add it if maintainers want the stricter reading.

Tests

  • New test_chat_kv_transfer_prompt_token_ids_rejected_without_kv_connector: well-formed ids on a connector-less server now raise VLLMValidationError naming --kv-transfer-config; goes red if the gate is removed (mutation-checked).
  • New test_chat_kv_transfer_prompt_token_ids_reused_with_kv_connector: with a connector present, well-formed ids skip templating and tokenization (conversation == [], engine input carries the ids, remaining kv_transfer_params preserved).
  • The two existing render-side shape/vocabulary tests now give their mock engine a connector sentinel, since the gate runs before the shape checks; the Harmony-path tests use a MagicMock renderer whose config satisfies the gate unchanged.
  • Local run of tests/entrypoints/openai/chat_completion/test_serving_chat.py: 208 passed; the 10 errors are TestGPTOSS* fixture setup failures that reproduce identically on pristine upstream/main on this machine (missing harmony deps), unrelated to this change. ruff check and ruff format clean on both touched files.

…euse

Fixes vllm-project#62307.

kv_transfer_params["prompt_token_ids"] is the disaggregated-serving
reuse path: the decode stage skips templating and tokenization because
the prefill stage's ids arrive with the transfer coordination. But the
renderer honored the key on every server, so on a plain server a client
could supply ids that silently replace the rendered messages, skipping
the chat template, truncate_prompt_tokens, and tokenization.

The reuse gate now rejects the key with a 400 unless the server runs a
KV transfer connector (--kv-transfer-config), which is the capability
the contract already assumes. echo and non-text-content requests still
ignore the ids and render messages as before. The shape and vocabulary
checks added in vllm-project#55771/vllm-project#59555 keep their behavior once the gate passes.

The two existing render-side shape tests now give their mock engine a
connector sentinel, since the gate runs first; a new pair pins the gate
(no connector -> 400 even for well-formed ids; connector -> ids skip
templating and tokenization).

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@hickeyma hickeyma left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the bug in #62307 but I don't think kv_transfer_config is the right approach. llm-d-router#3706 forwards prompt_token_ids from the EPP on aggregated deployments too. Those have no KV connector which means every forwarded request would get a 400.

Two changes would cover both cases:

  • Without a connector, ignore the IDs and render messages instead of returning a 400. This is similar that what we do for echo and non-text content. That closes the hole and router requests still work.
  • Add an opt-in flag for deployments that want the IDs honored without a connector. --enable-prompt-token-ids was already suggested on #55770.

#64721 sends prompt_features through this same helper, so whatever we pick needs to cover that as well.

cc @DarkLight1337 @praveingk

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: kv_transfer_params.prompt_token_ids replaces the rendered chat prompt on servers with no KV connector

2 participants