Repository navigation
[Bugfix][release/v0.32.0] Reject unsupported pre-tokenized diffusion prompt ids - #9335
Conversation
…prompt ids Backport the seven-file feature patch from vllm-project#8448, approved at 9fda1a2, onto release/v0.32.0 at 6497aaf. Preserve tensor-aware admission checks, custom-pipeline capability compatibility, documentation and regression tests without bringing main-branch merges. Add the required SPDX header to prompt_utils.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…t-contract-0.32.0
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
amy-why-3459
left a comment
There was a problem hiding this comment.
Thanks for the clean release/v0.32.0 backport of #8448.
Covered on this tip (6f65f8a): the new pipeline_utils capability resolution (custom / Diffusers-adapter / registry), _validate_pre_tokenized_prompt at admission (before preprocess / scheduler), pre_tokenized_prompt_keys presence rules (numel() for tensors; empty containers admitted; prompt_token_ids / bare lists left alone), the tri-state custom-pipeline behavior, the RL-test pipeline flag, and the new CPU admission tests.
The contract looks right for #7855: unsupported non-empty prompt_ids / negative_prompt_ids become OmniClientError(400) instead of a silent empty-condition success, while undeclared custom pipelines stay unchecked for the verl-style rollout path.
One non-blocking note on the helper (inline). No blocking issue spotted in the admission path itself given the reported CPU suite and B200 custom-pipeline / VeRL E2E results.
| token_ids = prompt.get(key) | ||
| if token_ids is None: | ||
| continue | ||
| size = token_ids.numel() if isinstance(token_ids, torch.Tensor) else len(token_ids) |
There was a problem hiding this comment.
Minor hardening: for non-tensors this uses len(token_ids). A malformed scalar (e.g. prompt_ids: 101) would raise TypeError here and surface as a 500 instead of the same OmniClientError(400) path. Wrapping with a small try/except (or treating non-sized values as populated) would keep the failure mode consistent; optional for this backport if call sites are already typed.
There was a problem hiding this comment.
Thanks for catching this. Addressed in 25a78e4: unsized values now count as populated, so malformed scalars—including 0 and False—follow the existing OmniClientError(400) path for unsupported pipelines. Tensor handling, empty values, and custom-pipeline bypass remain unchanged.
Added regression coverage for both ID fields across all four admission paths and for custom-pipeline compatibility. The 24 new rejection cases failed before the fix; the CPU engine suite now reports 141 passed, 1 deselected.
|
This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/input_output_modality_contracts.md, docs/design/module/diffusion/diffusion_runtime.md. Module owners: @david6666666 @Isotr0py @princepride Routing: @david6666666 via module of the changed files, module named in the PR description, CODEOWNERS; @Isotr0py via module of the changed files, module named in the PR description; @princepride via module of the changed files, module named in the PR description @fusheng-ji, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
amy-why-3459
left a comment
There was a problem hiding this comment.
Re-checked the delta on tip 25a78e4 after the prior note on unsized id values.
pre_tokenized_prompt_keys now treats TypeError from len() / non-tensor sizing as populated (size = 1), so plain ints / bools still hit the admission-time OmniClientError path instead of slipping past as empty. The expanded parametrization in test_diffusion_engine.py (integer / integer-zero / boolean on reject, step-streaming, and custom-pipeline bypass) covers that contract cleanly.
No further issues on this tip; the earlier non-blocking note is addressed.
| try: | ||
| size = token_ids.numel() if isinstance(token_ids, torch.Tensor) else len(token_ids) | ||
| except TypeError: | ||
| # Unsized values must still reach the admission-time client error. |
There was a problem hiding this comment.
Thanks — wrapping numel() / len() in TypeError → size = 1 is the right shape here. Unsized scalars / bools stay "populated" so unsupported pipelines still fail at admission rather than silently admitting.
There was a problem hiding this comment.
Thanks for re-checking and confirming the fix!
BLANKETusers
left a comment
There was a problem hiding this comment.
Thanks @fusheng-ji — reviewing this as a diffusion-runtime admission contract backport of #8448 / #7855 onto release/v0.32.0.
Agreement
Accepting caller prompt_ids / negative_prompt_ids on pipelines that only tokenize text is a silent correctness bug (request "succeeds" with empty/wrong conditioning). Failing at every admission entrypoint with OmniClientError(400) before preprocess/schedule is the right place.
What looks good
- Capability resolved once via
supports_pre_tokenized_prompt_ids, with the tri-state custom-pipeline rule (True/False/ undeclared→None) so verl-style rollout custom pipelines stay compatible. - Presence uses
numel()/ sized length so zero-valued and scalar tensors still count as populated, while empty tensors remain absent. - Native
prompt_token_idsand bare token lists are intentionally out of scope — correct for multi-stage hand-offs. - Regression surface is strong: admission matrix across entrypoints + custom-pipeline / Diffusers-adapter cases, plus contributor B200 custom-pipeline and VeRL E2E.
Asks
- Please confirm the release-branch cherry-pick is limited to the approved #8448 patch (+ SPDX on
prompt_utils.py) with no unintended main-only behavior. - One line in the PR body listing which in-tree pipelines (if any) already declare
supports_pre_tokenized_prompt_ids = Trueon this release tip, so release reviewers know the allowlist. - After merge, please link the release cherry-pick note back to #8448 / #7855 for tracking.
Happy to support landing this on the release branch once those are confirmed.
| """ | ||
|
|
||
| supported = getattr(self, "_supports_pre_tokenized_prompt_ids", None) | ||
| if supported is None or supported: |
There was a problem hiding this comment.
Confirming the tri-state: None (undeclared custom pipeline) and True both skip rejection — only explicit False rejects. Please keep this documented in the PR body / release note so release reviewers know verl-style rollouts stay admitted.
There was a problem hiding this comment.
Updated the PR description to explicitly document this: True and undeclared None bypass rejection; explicit False rejects populated ID fields. This preserves compatibility for verl-style rollout custom pipelines.
| # ``OmniCustomPrompt`` keys that carry caller-tokenized ids. The native vLLM | ||
| # spelling ``prompt_token_ids`` is deliberately not listed: multi-stage models | ||
| # use it to hand an AR stage's tokens to their diffusion stage. | ||
| PRE_TOKENIZED_PROMPT_KEYS = ("prompt_ids", "negative_prompt_ids") |
There was a problem hiding this comment.
Good that prompt_token_ids is intentionally omitted — multi-stage hand-offs must not trip this admission gate. One line in the PR body listing which in-tree pipelines (if any) already declare supports_pre_tokenized_prompt_ids = True on this release tip would help reviewers.
There was a problem hiding this comment.
Added the inventory to the PR description. At 25a78e4, no production pipeline under vllm_omni/ declares supports_pre_tokenized_prompt_ids = True. The test custom pipeline QwenImagePipelineWithLogProbForTest declares it for the RL-rollout tests.
| custom_cls = resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None)) | ||
| if custom_cls is not None: | ||
| declared = getattr(custom_cls, flag, None) | ||
| return None if declared is None else bool(declared) |
There was a problem hiding this comment.
Tri-state on custom pipelines looks correct (undeclared→None, declared False→reject, True→admit). Please confirm the cherry-pick onto release/v0.32.0 is limited to the approved #8448 surface (+ SPDX on prompt_utils.py) with no unintended main-only behavior.
There was a problem hiding this comment.
Confirmed by comparing the patches: the initial backport matches the approved seven-file #8448 patch, plus the SPDX header on prompt_utils.py. The release-branch merge preserves that patch. The only subsequent change is the review-requested scalar hardening and regression coverage in 25a78e4; no unrelated main-branch changes are included. Updated the PR description accordingly.
Contributor intent
1. What outcome do you want to achieve?
Bring the admission fix in #8448 to
release/v0.32.0, as requested in the approval, so unsupported pre-tokenized diffusion prompts produce an actionable error instead of a misleading successful request.2. What is the highest-leverage variable for that outcome?
Whether engine admission checks that the selected pipeline actually supports the caller's pre-tokenized IDs before preprocessing or scheduling the request.
3. How will this work move that variable?
Backport the approved per-pipeline capability, admission check, documentation and regression tests. Unsupported non-empty
prompt_ids/negative_prompt_idsraiseOmniClientError(400)before entering the scheduler.4. Whose customer or user needs does this address, and what are those needs?
Diffusion API callers and RL rollout integrations need requests to preserve their intended conditioning or fail clearly. Custom pipelines that consume IDs must remain compatible, as must native
prompt_token_idsused for multi-stage hand-offs.Purpose
Backport of #8448, reviewed at
9fda1a28a839, ontorelease/v0.32.0at6497aaf8e85f. Related bug: #7855.The release engine accepts caller-tokenized IDs even when its pipeline only reads text. This change resolves the pipeline capability once and rejects unsupported IDs at every admission entrypoint. Tensor presence uses
numel(), so zero-valued and scalar tensors count as populated while empty tensors remain absent. For user-supplied custom pipelines, declaredTrueand undeclaredNonebypass rejection; explicitFalserejects populated ID fields. This preserves admission for verl-style rollout custom pipelines. Nativeprompt_token_idsand bare token lists retain their existing behavior.Scope audit: the initial backport (
33216af) matches the approved seven-file #8448 feature patch, plus the required SPDX header onprompt_utils.py. The release-branch merge (6f65f8a, release parentab216c0) preserves that feature patch; the current PR diff contains no unrelated main-branch changes. The subsequent commit25a78e4adds only the review-requested hardening for unsized ID values and its regression coverage. The model-adding guide andOmniCustomPromptdocstring describe the capability. The relevant diffusion runtime and input/output design pages are drafts; this backport retains the documentation scope of the approved fix.At this PR tip (
25a78e4), no production in-tree pipeline undervllm_omni/declaressupports_pre_tokenized_prompt_ids = True. The concrete test custom pipelineQwenImagePipelineWithLogProbForTestdeclares it for the RL-rollout tests.Post-merge tracking: add a release cherry-pick note linking to the merged #9335 in both #8448 and #7855. This follow-up remains pending while #9335 is open.
Test Plan
Environment: Python 3.12.14, torch 2.13.0, vLLM 0.32.0, Diffusers 0.40.0, Transformers 5.19.0; CPU checks with
CUDA_VISIBLE_DEVICES='' VLLM_TARGET_DEVICE=cpu.Original backport validation commit:
33216af86c970038d3bb0787745be526f69338f1.Scalar-input follow-up CPU validation commit:
25a78e4ed4dda85c0a860bd4d913f13d54a647f7.GPU environment: Python 3.12.14, torch 2.13.0 / CUDA 13.0, vLLM 0.32.0, Diffusers 0.40.0, Transformers 5.19.0, Ray 2.56.1; one NVIDIA B200, driver 580.126.20. The unmodified guards use
tiny-random/Qwen-Imageand theQwen/Qwen2-1.5B-Instructtokenizer.VLLM_TARGET_DEVICE=cuda,VLLM_WORKER_MULTIPROC_METHOD=spawn, andPYTHONPATHpoints to the pinned release checkout.Test Result
33216af): engine suite 105 passed, 1 deselected; nine adjacent suites 356 passed.25a78e4): engine CPU suite 141 passed, 1 deselected. All 24 new rejection cases failed withTypeErrorbefore the fix; the expanded suite covers both ID fields across all four admission paths and preserves custom-pipeline bypass.git diff --checkpass. Mypy reports one pre-existing diagnostic attests/diffusion/test_diffusion_engine.py:961, reproduced unchanged against the original6f65f8afile.OmniClientError(status_code=400)and queues none for all four. Empty-tensor and text controls are admitted on both. This probe does not run model inference.33216af): Ruff 0.14.10 lint/format, typos, whitespace/debug checks, SPDX headers, test marks, forbidden imports, CUDA-call ratchet andgit diff --checkpass. The full pre-commit run reports 27 mypy diagnostics and 29 Markdown issues; comparison with the exact release base has zero added or removed diagnostics after ignoring shifted line/column numbers. Unrelated Markdown autoformats were reverted. The final relevant-hook run skips those two verified baseline failures.33216af86c970038d3bb0787745be526f69338f1, using 1 NVIDIA B200: Custom Pipeline 14 passed (423.43s), VeRL 1 passed (294.41s); 0 failures, 0 errors, 0 skipped in both. VeRL returned valid images for all four concurrent tokenized requests, including the logprobs case. Both pytest processes exited 0 without timeout. All 5,081 tracked source-file hashes and HEAD remained unchanged before/after validation.SemLock._rebuild,ENOENT). This is consistent with the cluster's job epilog, which deletes shared-memory/temp files by username. The complete successful rerun used private user/mount/IPC namespaces with separate/dev/shmand/tmpmounts; no PR source or test was changed. These are contributor-run GPU results, not an official Buildkite result.AI assistance: Codex prepared the backport, ran its CPU regression and GPU E2E validation, implemented and CPU-tested the scalar-input follow-up, and helped draft the PR wording. The intent wording follows the original PR's purpose at the author's request.