[dongteng]

Follow-up to #7438, addressing the native MRv2 receive-path gap identified in the follow-up review.

Contributor intent (required human answers)

1. What outcome do you want to achieve?

Extend the atomic async-chunk protection from #7438 to native MRv2 receivers, so token-budget pressure does not split a received chunk or desynchronize its data and metadata.

2. What is the highest-leverage variable for that outcome?

The scheduler must recognize the async-chunk receive contract before applying token-budget limits, independently of whether reception uses the legacy adapter or the native coordinator.

3. How will this work move that variable?

Recognize native async receivers and use their coordinator readiness state. Reuse the existing whole-chunk admission and oversized-chunk rejection paths, while preserving the native in-flight fence and ordinary prompt/sender-only behavior. Add regression coverage for both receive implementations.

4. Whose customer or user needs does this address, and what are those needs?

Deployments receiving async chunks through native MRv2, including Qwen3-TTS Code2Wav stages, need received chunks and their metadata to remain consistent when multiple requests compete for the scheduling budget.

Purpose

Before this change, _requires_atomic_chunk() only recognized chunk_transfer_adapter.receives_chunks. Native MRv2 receivers use input_coordinator with no legacy adapter, so they bypassed the existing whole-chunk guards and oversized-chunk rejection.

This PR:

• Recognizes both legacy adapter receivers and native downstream async-chunk coordinators.
• Exposes the native coordinator's public receives_chunks property and shares that receive predicate between chunk processing and generation scheduling, without reading coordinator-private fields in the generation scheduler.
• Uses the appropriate receiver's requests_with_ready_chunks for oversized-chunk rejection.
• Preserves the existing native in-flight fence and the behavior of sender-only, stage-0, full-payload, and ordinary prompt paths.
• Parameterizes legacy/native regression coverage for waiting/running budget exhaustion, exact fit, intact next-step admission, oversized rejection, and terminal error delivery.
• Adds native cleanup coverage across waiting, KV-holding waiting, and running queues, checking exactly one terminal ERROR and the existing resource-release path.
• Defers native oversized rejection while a previous execution is in flight, preserving readiness and delivering its valid audio before terminal ERROR and cleanup. Regression coverage checks this ordering and that another ready stream can still run.
• Updates the async-chunk design contract in the pushed review-fix commit.

For example, a 1168-token received chunk with only 128 tokens remaining must wait intact, rather than scheduling a 128-token prefix. A ready chunk exceeding the entire scheduling budget must terminate with one ERROR and cleanup instead of waiting indefinitely.

This addresses missing native protection; it does not claim that #7438 introduced a new MRv2 regression. Production changes are limited to the generation scheduler and scheduling coordinator, with no orchestrator, model, runner, or CI configuration changes. After correcting two test-fixture omissions, the focused and broader related Linux CPU regressions passed. The latest related result is 585 passed, 16 warnings in 5.41s on changes subsequently committed and pushed as cdcc185aa; applicable server hooks and DCO sign-off passed with the documented skips. No current-head CI result is claimed.

Test Plan

• vLLM Version: 0.31.0.
• vLLM-Omni Commit: pushed review-fix commit cdcc185aa; original implementation d48822304; documentation cleanup 4c747b147; development base b60a7e9e8. The contributor's successful push advanced the existing fork branch from 4c747b147 to cdcc185aa.
• Environment: Linux, Python 3.12; Omni editable installation points to the follow-up worktree.
• Regression scope: scheduler tests and the two related orchestrator suites. The historical 571-test result below was obtained on changes subsequently committed as d48822304; the latest 585-test result is from review-fix working-tree changes subsequently committed as cdcc185aa. No additional pytest run after commit is claimed.
• Supplemental documentation cleanup: formatting-only changes to async_chunk.md have passed local Markdownlint and were committed and pushed as 4c747b147. The Python implementation and tests were not changed by this cleanup; pytest was not rerun after this documentation-only commit.

python -m pytest tests/core/sched/ \
  tests/engine/test_orchestrator.py \
  tests/engine/test_orchestrator_first_audio.py \
  -m "core_model and cpu" -q -ra

Test Result

• Latest review-fix related regression: 585 passed, 16 warnings in 5.41s, from the uploaded working tree; see the review-follow-up section for commands and the fixture corrections. This includes the previously reported 101 focused tests; the counts are not added together.
• Original targeted regression: 571 passed, 16 warnings in 5.61s. This predates the 2026-10-19 review revisions and does not validate the newly added in-flight/oversized combination.
• python -m pip check: No broken requirements found.
• For the implementation change, applicable pre-commit checks passed, including Ruff, typos, mypy-3.10, test marks, SPDX, forbidden imports, and the new torch.cuda call-site check. Commit sign-off passed for both commits.

The implementation-change pre-commit command was:

SKIP=actionlint,markdownlint-cli2 python -m pre_commit run --files \
  docs/design/feature/async_chunk.md \
  tests/core/sched/test_generation_scheduler_atomic_chunk.py \
  tests/core/sched/test_generation_scheduler_restore.py \
  tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py \
  tests/core/sched/test_omni_scheduler_mixin_shared.py \
  vllm_omni/core/sched/omni_generation_scheduler.py

actionlint was skipped after tool initialization was network-blocked; this PR changes no workflow files. The original server Markdownlint hook could not initialize because its Node.js environment could not be downloaded. A subsequent Windows check used the repository-pinned markdownlint-cli2 v0.23.2 (markdownlint v0.41.1) with Node v22.17.0 and the unchanged repository .markdownlint.yaml configuration. It found 107 formatting issues in the document; after formatting-only cleanup, the same check exited successfully with 0 issues:

Finding: docs/design/feature/async_chunk.md
Linting: 1 file
Summary: 0 issues in 0 files

This supplemental result applies to the cleaned-up document committed and pushed as 4c747b147. It does not change the historical skipped-hook result on d48822304, and does not claim that the document before the formatting cleanup passed. The server hook remained skipped for the documented initialization limitation. GitHub pre-commit currently skips Markdownlint, so its status is not used as proof of this check.

Broader repository test attempts and outstanding limitations

The following CPU-tagged partitions were attempted in this worktree before the final test-dependency installation and targeted rerun. They have not all been rerun in the final environment and are not presented as a complete successful repository-wide validation for this head. Tag selection did not prevent model-resource access or GPU cases in these suites.

python -m pytest tests/model_executor -m "core_model and cpu" -q -ra
python -m pytest tests/diffusion/models -m "core_model and cpu" -q -ra
python -m pytest tests/entrypoints tests/engine tests/diffusion \
  --ignore=tests/diffusion/models -m "core_model and cpu" -q -ra
python -m pytest tests/ \
  --ignore=tests/diffusion --ignore=tests/model_executor \
  --ignore=tests/entrypoints --ignore=tests/engine \
  -m "core_model and cpu" -q -ra

Partition	Actual result
Model executor	4913 passed, 17 failed, 6 errors, 21 skipped, 1759 deselected, 72 warnings
Diffusion models	4652 passed, 16 skipped, 494 deselected, 29 warnings; no failures/errors
Entrypoints + engine + remaining diffusion	8516 passed, 20 failed, 2 errors, 18 skipped, 1827 deselected, 1 xfailed, 98 warnings
Remaining tests	Collection interrupted: 28 errors, 9 skipped, 1125 deselected, 17 warnings; test bodies did not run

• Model executor: 12 failures and 6 setup errors occurred while required Qwen/OmniVoice resources were unavailable over the network. Two failures lacked cosyvoice2; three lacked voxcpm. The 21 skips comprise 14 missing-extra cases (flashcosyvoice, cosyvoice2, librosa, voxcpm), 3 unavailable tokenizer/remote-code resource cases, and 4 legacy/ROCm/upstream-API exclusions. A socket-cleanup FileNotFoundError appeared after the pytest summary and has not been independently reproduced.
• Diffusion models: the 16 skips comprise 4 optional-package/source cases, 9 unavailable model/tokenizer/checkpoint cases, 1 unconfigured multi-GPU SP case, and 2 unavailable upstream-API cases.
• Entrypoints/engine/remaining diffusion: 12 failures and 2 setup errors occurred with unavailable Hugging Face configuration/pipeline resources; 8 failures occurred during four-rank CUDA device selection with invalid device ordinal. The 18 skips comprise 7 optional-dependency cases, 2 VAD/HF-resource cases, and 9 unsupported-platform cases. The xfail is an existing images/edits stop-token case.
• Remaining tests: 27 collection errors were missing torchmetrics, mistune, openpyxl, or ray; one was a network-blocked CosyVoice3 snapshot download. The 9 skips cover unavailable NeMo/openpi packages, MiMo resources, and NPU dependencies. The four missing test packages were subsequently installed, but this partition was not rerun.

No unmodified-baseline comparison has been completed for these broader failures; they are not all asserted to be unrelated to this PR. Complete repository-wide validation remains outstanding. No native MRv2 GPU end-to-end reproduction or Buildkite result is claimed for this follow-up head; Buildkite admission remains maintainer-owned.

Review follow-up validation (2026-10-19)

The local patch addresses both public receive-predicate coupling and oversized rejection during a native in-flight execution. Restore fixtures now use real scheduling coordinators. New CPU regressions cover the shared public receive gate, in-flight deferral without blocking another stream, and prior-audio delivery followed by exactly one ERROR through the real cleanup path.

Local Ruff 0.14.10 code and formatting checks passed for the seven affected Python files. Isolated probes of the actual production methods reproduced premature rejection before the fix, then verified deferral while in flight and rejection after the fence clears. Those probes use minimal stand-ins; they are not repository pytest or GPU validation. The first Linux run of the three focused files below returned 3 failed, 98 passed, 14 warnings in 1.68s. All three failures were in the in-flight cleanup cases: the KV cache mock's estimate_cached_tokens() returned a Mock rather than the integer required by the real prefill statistics finalizer. The fixture now explicitly returns 0 cached tokens, preserving the real statistics/output/cleanup path and all assertions; production code is unchanged by this fixture correction. After uploading that correction, the contributor reran all three focused files on Linux: 101 passed, 14 warnings in 1.50s. The 14 warnings are Torch JIT deprecations. This focused run preceded the related-suite fixture correction below; the final review-fix changes were subsequently committed as cdcc185aa. This Windows environment has no vLLM installation; the reported pytest results are from the contributor's server.

Focused CPU run, from the follow-up worktree with the existing matching vLLM/Omni installation and pytest-mock (no GPU or model weights required):

python -m pytest tests/core/sched/test_generation_scheduler_atomic_chunk.py \
  tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py \
  tests/core/sched/test_omni_scheduling_coordinator.py \
  -m "core_model and cpu" --run-level=core_model -q -ra

The first related-suite run returned 4 failed, 581 passed, 16 warnings in 5.84s. All four failures used the shared legacy scheduler fixture in test_generation_scheduler_restore.py, which bypasses the constructor with __new__ and omitted _native_data_plane. The fixture now explicitly defaults that flag to False, consistent with its legacy adapter; native tests still explicitly override it to True. This was a one-line test-fixture correction with no production change or weakened assertions. Local Ruff and syntax checks passed. After uploading the corrected fixture, the contributor reran the entire related suite on Linux: 585 passed, 16 warnings in 5.41s. The warnings are 14 Torch JIT, one deprecated API entrypoint, and one audioop deprecation.

Successful related-suite rerun with the CI marker/level selection:

python -m pytest tests/core/sched/ \
  tests/engine/test_orchestrator.py \
  tests/engine/test_orchestrator_first_audio.py \
  -m "core_model and cpu" --run-level=core_model -q -ra

Both initial failed runs remain recorded; the corrected focused and entire related runs passed as stated above. The 101 focused tests are included in the 585-test related suite and are not counted again. Local read-only review found no new confirmed correctness issue; server status confirmed the expected eight-file scope and git diff --check reported no whitespace errors. The review-fix commit delta is 158 insertions and 32 deletions, not the full PR delta from the development base.

The current review-fix pre-commit rerun covered all eight changed files (the six original files plus vllm_omni/core/sched/omni_scheduling_coordinator.py and tests/core/sched/test_omni_scheduling_coordinator.py), using SKIP=actionlint,markdownlint-cli2. The result was 13 Passed, 6 Skipped, with no automatic file modifications reported. Ruff check/format, typos, mypy-3.10, CI test marks, SPDX, forbidden imports, and the new torch.cuda call-site check passed. The two explicit environment-limited skips retain their documented status; the other four skips had no applicable files. The commit-time rerun also passed its applicable checks, and Sign-off Commit passed when creating cdcc185aa; no automatic file modifications were reported.

Before committing, staging verification passed: the index contained only the expected eight modified files, with 158 insertions and 32 deletions, and git diff --cached --check reported no whitespace errors. Commit cdcc185aa was then created successfully with the same file count and change statistics. Post-commit status/log confirmed a clean worktree at that commit, and the contributor subsequently pushed it successfully to the existing fork branch. No full repository-wide, GPU end-to-end, or Buildkite pass is claimed for this revision; current-head CI results remain to be checked.

AI assistance

OpenAI Codex helped

[chatgpt-codex-connector]

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-18T14:44:17.947112Z	4c747b1	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]

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.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/ar_runtime.md, docs/design/module/cache_management.md, docs/design/module/input_output_modality_contracts.md.

Module owners: @Gaohan123 @alex-jw-brooks @amy-why-3459

Routing: @Gaohan123 via module of the changed files, CODEOWNERS; @alex-jw-brooks via module of the changed files; @amy-why-3459 via module of the changed files

@dongteng, 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 6a5ba26800b5 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[dongteng]

@amy-why-3459 This is the follow-up to #7438 for native MRv2 atomic-chunk scheduling, with regression coverage and validation results included. Thanks!

[amy-why-3459]

Thanks for the follow-up. Extending the atomic-chunk contract to native receivers looks right to me: native update_request_metadata replaces prompt_token_ids and resets num_computed_tokens the same way the legacy commit does, so len(prompt_token_ids) is a valid full-chunk length on both paths, and readiness survives a budget deferral because _clear_chunk_ready only discards scheduled ids. The parametrized legacy/native coverage is a nice improvement.

Two comments inline (private-attr coupling of the receive predicate, and oversized rejection vs. the native in-flight fence). Minor: the PR description still says the async_chunk.md formatting cleanup is not committed, but it is already in 4c747b1, so the test-plan section may want a refresh.

[amy-why-3459]

This reaches into coordinator._async_chunk / coordinator._stage_id from the scheduler, and the tests now have to fake those private fields on SimpleNamespace coordinators. The same predicate already gates OmniSchedulingCoordinator.process_pending_chunks (if self._stage_id == 0 or not self._async_chunk: return). Could we expose it once as a public property on the coordinator (e.g. receives_chunks, mirroring OmniChunkTransferAdapter.receives_chunks) and use it in both places? Then the generation scheduler can treat both receive paths uniformly (owner.receives_chunks / owner.requests_with_ready_chunks), and a future change to the receiver rule (e.g. the stage_receives_chunks gate in the mixin init) can't drift between the two call sites.

[amy-why-3459]

On the native path, readiness can be set while the previous chunk of the same request is still executing (num_in_flight_tokens > 0). That is the case TestNativeInFlightFence covers for a fitting chunk. If the newly ready chunk is oversized, this loop finishes the request right away, while the previous chunk is still in flight. That chunk's output is then dropped (the request is gone from self.requests when it returns), so the client loses the last valid audio segment before the ERROR. Since the ready flag is preserved across steps anyway, would it be safer to skip native requests with num_in_flight_tokens > 0 here and reject on the next step once the fence clears? That would also match the doc line about not acting on a ready chunk while a previous execution still owns the request state. A native in-flight + oversized test case would pin whichever behavior you choose.

[amy-why-3459]

Re-checked cdcc185. Both earlier points are addressed: _requires_atomic_chunk() and process_pending_chunks() now share the public OmniSchedulingCoordinator.receives_chunks predicate, and _reject_oversized_ready_chunks() skips native requests with num_in_flight_tokens > 0, so the previous step's output is delivered before the oversized chunk is rejected on a later step. The regression that drives the real update_from_output path before rejection covers this well. No further comments on this update.

[amy-why-3459]

@dongteng Could you please rebase this PR onto the latest base branch and resolve the merge conflicts? Thanks!

[dongteng]

@dongteng Could you please rebase this PR onto the latest base branch and resolve the merge conflicts? Thanks!

@amy-why-3459 Sure,I’ll handle it.

[amy-why-3459]

Re-checked the rebased tip af5a95e5 (vs the previously reviewed cdcc185a content).

The production changes that closed the earlier threads are still present after the rebase: OmniSchedulingCoordinator.receives_chunks is the shared receive predicate for _requires_atomic_chunk() / process_pending_chunks(), and _reject_oversized_ready_chunks() continues to skip native requests with num_in_flight_tokens > 0 so the prior step’s output can land before an oversized rejection. The new tip commit itself is test-fixture alignment for vLLM 0.32 (receive_path= construction and _kv_fetch_stages = None); no further production comments from me on this head.

[dongteng]

Thanks @amy-why-3459! Could you add the ready label to trigger Buildkite when you get a chance?

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 03a4c3af2b29 refers to an earlier head; the pull request now points at 6a5ba26800b5.

[dongteng]

@amy-why-3459 I checked, and the Buildkite failure is caused by an upstream AMD nightly job-count mismatch, unrelated to this PR’s changes. Could you please update the branch with the latest main and rerun Buildkite? Thanks!