[armaanamatya]

Purpose

When stage 0 has more than one replica, AsyncOmniEngine._ensure_stage_replica_mm_uuids scopes each multimodal cache key to the chosen replica, but it writes the scoped UUIDs back into the caller's prompt dict (prompt["multi_modal_uuids"] = scoped_uuids). On a later call it treats any UUIDs already in the dict as user-supplied and uses them without hashing the media.

So when a request is built from a dict that already went through the engine (reused, retried, or shallow-copied from a template), its new media gets keyed under the old media's hash. Two requests with different media that land on the same replica share one processor-cache key. The frontend sender cache sees a hit and omits the new tensor, and the replica then fills the request from its receiver cache, silently serving the features cached for the first request.

Repro on main (calling the helper directly): shallow copies of a used template with three different videos, routed to the same replica, all get stage0:rep1:stage0:rep0:<hash of the first video>.

Fix: scope a shallow copy of the prompt and pass that copy to process_inputs, leaving the caller's dict untouched. _scope_stage0_multimodal_cache_to_replica now returns (replica_id, prompt). The only reader of multi_modal_uuids is the stage-0 renderer (vllm_omni/inputs/preprocess.py), which gets the copy, so downstream stages that use original_prompt are unaffected.

Behaviour change to note in the changelog: a retry built from a prompt dict that already went through the engine and lands on a different replica is now scoped from the original UUIDs, so it gets stage0:rep1:<hash> like a first-time request on that replica, instead of the doubly prefixed stage0:rep1:stage0:rep0:<hash> the old write-back produced.

Possibly related to #8039 (asked the reporter there). #7295 is a separate failure: it is on the rendered chat path, which this scoping does not cover, so this PR does not fix it.

cc @tzhouam @fake0fan @bjf-frz @vllm-omni-review-bot

Test Plan

• Two new CPU tests in tests/engine/test_async_omni_engine_input.py:
  • test_mm_uuid_scoping_leaves_caller_prompt_untouched: user UUIDs are still replica-prefixed, and the caller's dict keeps its original value.
  • test_mm_uuid_scoping_keys_reused_prompt_by_its_own_media: a used template plus three copies with new media, over two replicas, get four distinct keys with a single stage0: prefix each.
• pytest tests/engine/test_async_omni_engine_input.py
• ruff check / ruff format --check on the changed files

vLLM Version: 0.30.0

vLLM-Omni Commit: 75570ee

Test Result

• With the fix: 23 passed.
• With the source change reverted: both new tests fail, the other 21 pass.
• ruff: clean.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/engine_orchestration.md, docs/design/module/entrypoints.md, docs/design/module/observability.md.

Module owners: @fake0fan @NickCao @tzhouam

Routing: @fake0fan via module of the changed files, CODEOWNERS; @NickCao via module of the changed files, CODEOWNERS; @tzhouam via module of the changed files, CODEOWNERS

@armaanamatya, 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]

Omni ReviewBot triage note

Automated triage of commit 52d0d91c2e38 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[armaanamatya]

Self-review:

• The only caller of _scope_stage0_multimodal_cache_to_replica is _build_add_request_message, which now passes the scoped copy to process_inputs. The only reader of multi_modal_uuids is the stage-0 renderer (vllm_omni/inputs/preprocess.py), so original_prompt and downstream stages are unaffected.
• The copy is shallow, so multi_modal_data and additional_information are still shared with the caller, as before. Only the multi_modal_uuids key differs.
• User-supplied UUIDs are still replica-prefixed. The single-replica, diffusion, and no-serviceable-replica paths return the prompt unchanged, same as before.
• Ran tests/engine and tests/entrypoints/test_serve.py locally with -m "not slow" on CPU: 1354 passed, 2 skipped. test_async_omni_engine_abort.py::test_abort failed only because it needs a GPU (L4), which was hidden for that run. With the source change reverted, both new tests fail.

@tzhouam @fake0fan could one of you add the ready label so L2 CI runs? Thanks.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 95e4810d3700 refers to an earlier head; the pull request now points at 52d0d91c2e38.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict under experiment vllm-omni-strict-5050-20260829.

[vllm-omni-review-bot]

Omni ReviewBot review

PR description

Stage-0 multimodal cache scoping used to write replica-prefixed UUIDs back into the caller's prompt dict. A later request built from that same dict (retry, template, or shallow copy) then treated those engine keys as user UUIDs, skipped hashing the new media, and could silently reuse another request's cached features. The engine now scopes a shallow copy and sends only that copy into process_inputs, so original_prompt and the caller keep their original UUID view while each request is keyed by its own media.

Change flow

flowchart TD
    Caller["[EXISTING] Caller prompt dict"]:::existing
    Build["[CHANGED] _build_add_request_message"]:::changed
    Scope["[CHANGED] _scope_stage0_multimodal_cache_to_replica"]:::changed
    Ensure["[CHANGED] _ensure_stage_replica_mm_uuids"]:::changed
    Stage0["[EXISTING] process_inputs / stage-0 renderer"]:::existing
    Downstream["[EXISTING] original_prompt downstream"]:::existing
    Tests["[NEW] test_mm_uuid_scoping_*"]:::new
    Caller --> Build
    Build --> Scope
    Scope --> Ensure
    Ensure -->|"scoped copy, no write-back"| Stage0
    Caller -->|"unscoped original"| Downstream
    Tests -->|"pins reused-template keys"| Ensure
    classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
    classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
    classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
    classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

No actionable findings.

[vllm-omni-review-bot]

Omni ReviewBot review

PR description

When stage 0 has more than one replica, the engine used to write replica-scoped multimodal UUIDs back into the caller’s prompt dict. A later request built from that same dict then treated those values as user-supplied UUIDs, so new media reused the old processor-cache key and could silently receive the first request’s features. This change scopes UUIDs on a shallow copy and passes that copy to process_inputs, leaving the caller dict and original_prompt unscoped. User-supplied UUIDs are still replica-prefixed on the copy; single-replica, diffusion, and no-serviceable-replica paths still return the original prompt.

Change flow

flowchart TD
  caller["[EXISTING] Caller prompt dict"]:::existing
  build["[CHANGED] AsyncOmniEngine._build_add_request_message"]:::changed
  scope["[CHANGED] _scope_stage0_multimodal_cache_to_replica / _ensure_stage_replica_mm_uuids"]:::changed
  copy["[NEW] Shallow prompt copy with replica-scoped UUIDs"]:::new
  proc["[EXISTING] input_processor.process_inputs"]:::existing
  orig["[EXISTING] original_prompt for downstream stages"]:::existing
  tests["[NEW] test_mm_uuid_scoping_* reuse and immutability tests"]:::new
  caller --> build
  build --> scope
  scope --> copy
  copy --> proc
  caller --> orig
  copy --> tests
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

CI at 52d0d91c2e38 (2026-09-26T20:41:34.571311+00:00): required check(s) blocking: buildkite/vllm-omni (missing), and -5 more.

No actionable findings.

[armaanamatya]

Update from #8039: the reporter checked their deployment and it runs one stage-0 replica per engine, so this is not what they are hitting. They did confirm the write-back into the caller's dict is a real bug independent of their case ("#8062 looks like the right fix").

No change to the PR. It is still mergeable on current main (27a6321a); neither touched file has moved since the base. @tzhouam @fake0fan @NickCao PTAL.

[armaanamatya]

Re-validated against current main (27a6321a): with this PR's commit applied on top, tests/engine/test_async_omni_engine_input.py gives 23 passed (CPU, vLLM 0.30.0). All GitHub checks are green and it is mergeable. The bot reviewed both heads with no findings.

@fake0fan @NickCao @tzhouam could one of you review and add ready so Buildkite runs? It is a two-file fix.

@vllm-omni-review-bot please review.

[hsliuustc0106]

Approving — this fixes the stale-scoping-key bug verified during the #8039 triage (write-back into the caller's dict at the old L155, then read as user-supplied UUIDs on the next call, keying new media under the old media's hash).

What I checked:

• Signature change is safe: both private helpers have exactly one caller each on main and both call sites are updated here.
• Copy semantics are correct: dict(prompt) + wholesale replacement of multi_modal_uuids never aliases into the caller's nested lists, so the caller's dict is provably untouched (first new test asserts exactly this).
• The regression tests encode both failure modes directly: a reused template with swapped media now produces 4 distinct keys, single stage0: prefix (no double-scoping), and the round-robin replica split is asserted. This is precisely the CPU repro from the #8039 thread.
• One behavioral improvement worth naming in the changelog: a retry that lands on a different replica now re-scopes from the original UUIDs instead of reusing stale stage0:repN:... keys — the old write-back was wrong for that path too.

Non-blocking: the Any return type on _ensure_stage_replica_mm_uuids is fine for a private helper; a dict-narrowing can come with future typing work.

For the record: this fixes a real bug on main but does not explain #8039's tensor-level contamination (the reporter's prompts carry uuids=None), so that issue stays open independently of this merge.

[armaanamatya]

@hsliuustc0106 thanks for the review. I added the retry-on-a-different-replica behaviour change to the description so it is there for the changelog: such a retry is now scoped from the original UUIDs and gets stage0:rep1:<hash> like a first-time request on that replica, instead of the doubly prefixed key the old write-back produced. No code change, so the approval stands as is. The Any return on _ensure_stage_replica_mm_uuids can tighten to dict with the typing pass you mentioned.

Merge is blocked only on the Buildkite gate, which needs the ready label. Could you add it so it runs on 52d0d91c? I have not rebased: the branch still merges cleanly against main, and I would rather not invalidate the approval with a new push.

On #8039, agreed: this closes only the stale-key thread, and the issue stays open for thread 2.