Skip to content

[Bugfix][release/v0.32.0] Reject unsupported pre-tokenized diffusion prompt ids - #9335

Merged
tzhouam merged 3 commits into
vllm-project:release/v0.32.0from
fusheng-ji:fix/diffusion-pre-tokenized-prompt-contract-0.32.0
Oct 27, 2026
Merged

tzhouam merged 3 commits into
vllm-project:release/v0.32.0from
fusheng-ji:fix/diffusion-pre-tokenized-prompt-contract-0.32.0

Conversation

@fusheng-ji

@fusheng-ji fusheng-ji commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

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_ids raise OmniClientError(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_ids used for multi-stage hand-offs.

Purpose

Backport of #8448, reviewed at 9fda1a28a839, onto release/v0.32.0 at 6497aaf8e85f. 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, declared True and undeclared None bypass rejection; explicit False rejects populated ID fields. This preserves admission for verl-style rollout custom pipelines. Native prompt_token_ids and 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 on prompt_utils.py. The release-branch merge (6f65f8a, release parent ab216c0) preserves that feature patch; the current PR diff contains no unrelated main-branch changes. The subsequent commit 25a78e4 adds only the review-requested hardening for unsized ID values and its regression coverage. The model-adding guide and OmniCustomPrompt docstring 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 under vllm_omni/ declares supports_pre_tokenized_prompt_ids = True. The concrete test custom pipeline QwenImagePipelineWithLogProbForTest declares 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.

CUDA_VISIBLE_DEVICES='' VLLM_TARGET_DEVICE=cpu python -m pytest -q \
  tests/diffusion/test_diffusion_engine.py \
  -m 'core_model and cpu' --run-level=core_model

CUDA_VISIBLE_DEVICES='' VLLM_TARGET_DEVICE=cpu python -m pytest -q \
  tests/diffusion/test_diffusion_engine_cleanup.py \
  tests/diffusion/test_diffusion_request.py \
  tests/diffusion/test_inline_stage_diffusion_client.py \
  tests/diffusion/test_stage_diffusion_client.py \
  tests/diffusion/test_multiproc_engine_concurrency.py \
  tests/diffusion/test_diffusion_engine_rpc_routing.py \
  tests/diffusion/test_request_cancellation.py \
  tests/diffusion/test_stage_payload_kv_coexistence.py \
  tests/diffusion/test_diffusion_scheduler.py \
  -m 'core_model and cpu' --run-level=core_model

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-Image and the Qwen/Qwen2-1.5B-Instruct tokenizer. VLLM_TARGET_DEVICE=cuda, VLLM_WORKER_MULTIPROC_METHOD=spawn, and PYTHONPATH points to the pinned release checkout.

python -m pytest -s -v tests/e2e/features/custom_pipeline/ \
  -m core_model --run-level=core_model --junitxml=custom-pipeline.xml

python -m pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py \
  --run-level=core_model --junitxml=verl.xml

Test Result

  • Original backport (33216af): engine suite 105 passed, 1 deselected; nine adjacent suites 356 passed.
  • Scalar-input follow-up (25a78e4): engine CPU suite 141 passed, 1 deselected. All 24 new rejection cases failed with TypeError before the fix; the expanded suite covers both ID fields across all four admission paths and preserves custom-pipeline bypass.
  • Follow-up checks on the two changed Python files: Ruff lint/format, typos, whitespace/debug, SPDX, test marks, forbidden imports, CUDA-call ratchet and git diff --check pass. Mypy reports one pre-existing diagnostic at tests/diffusion/test_diffusion_engine.py:961, reproduced unchanged against the original 6f65f8a file.
  • Admission-only before/after probe using the exact release-base engine, the production scheduler and the real QwenImage pipeline capability: the release base queues unsupported positive IDs, negative IDs, a zero-valued tensor and a scalar-zero tensor; the backport raises 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.
  • Original backport checks (33216af): Ruff 0.14.10 lint/format, typos, whitespace/debug checks, SPDX headers, test marks, forbidden imports, CUDA-call ratchet and git diff --check pass. 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.
  • GPU E2E rerun on the original backport commit 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.
  • The initial shared-IPC attempt had 12 custom-pipeline passes, 2 MP worker-startup failures and 1 VeRL setup error, all involving missing multiprocessing semaphores (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/shm and /tmp mounts; 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.

…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>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ⚠️ Failed 2026-10-26T09:45:04.810261Z
We couldn’t complete this review. Comment @codex review to try again.
6f65f8a Draft marked ready
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@amy-why-3459 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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.

@vllm-omni-review-bot

Copy link
Copy Markdown

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.

@vllm-omni-review-bot

vllm-omni-review-bot commented Oct 26, 2026 •

Copy link
Copy Markdown

Omni ReviewBot triage note

Automated triage of commit 25a78e4ed4dd produced:

  • Priority: high. Prompt maintainer attention is suggested.

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 amy-why-3459 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks for re-checking and confirming the fix!

@BLANKETusers BLANKETusers 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.

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_ids and 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

  1. 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.
  2. 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, so release reviewers know the allowlist.
  3. 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:

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.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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")

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.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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)

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.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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.

@tzhouam tzhouam added the ready label to trigger buildkite CI label Oct 27, 2026
@tzhouam
tzhouam enabled auto-merge (squash) October 27, 2026 07:52
@tzhouam
tzhouam disabled auto-merge October 27, 2026 07:53
@tzhouam
tzhouam merged commit 04653fc into vllm-project:release/v0.32.0 Oct 27, 2026
9 of 11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working diffusion codes related to diffusion models ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants