[Refactor] Retire the generic deferred MTP path in MRV2 by sphinxkkkbc · Pull Request #9346 · vllm-project/vllm-omni

sphinxkkkbc · 2026-10-25T04:25:44Z

PLEASE FILL IN THE PR DESCRIPTION HERE.

Contributor intent (required human answers)

1. What outcome do you want to achieve?

Implement #9173 by retiring the generic legacy deferred MTP path in MRV2.

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

The coexistence of deferred and current-step MTP scheduling, together with MTP state distributed across components.

3. How will this work move that variable?

Consolidate MTP execution state in MTPState, use post_sample() for the generic current-step path, and remove generic deferred MTP dispatch from preprocess.

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

Contributors who need a simpler, more maintainable MTP execution lifecycle.

Purpose

Implements #9173. The motivation is described in the issue; this PR focuses on retiring the generic legacy deferred MTP path, only changes MRV2. Model adaptations are limited to models integrated with MRV2; V1 execution remains unchanged.

What's Changed

Renamed EagerMTPState to MTPState in mtp_model_state.py and consolidates MTP-specific buffers, request generators, execution helpers, and CUDA Graph management there. It replaces run_eager_mtp() with model_state.post_sample(), removes generic deferred MTP dispatch from preprocess, and separates decode-input composition metadata from the rows recorded for current-step MTP.

Most helpers retain their existing behavior, and model-side mtp / talker_mtp hook names and signatures remain unchanged. Preprocess continues to consume cached frame embeddings and compose the following backbone input.

Model Side Change

Qwen3 Omni and Qwen3 TTS use current-step MTP, with first-audio decoding and delivery configured separately.

VoxCPM2 migrates its callback to post_sample() while retaining its per-patch audio and stop handling.

MOSS Local now always uses MossLocalModelState, and mrv2_gpu_slot_state is removed. Configurations that previously omitted this option or set it to false now migrate from generic OmniModelState to the dedicated state. Model-owned deferred/eager scheduling remains controlled by mrv2_eager_mtp. The default CUDA YAML already selected the dedicated state, so its state selection is unchanged. Existing NPU deployments remain on MRV1; custom MRV2 Local deployments on NPU have not been validated.

vLLM Version: v0.32.0 (5df63a7203)

vLLM-Omni Commit: 1dd816263fb2e5009448620374d943cecd0fdc38

I haven't benchmarked performance, but I'd be happy to do so if reviewers would find it useful. This PR primarily refactors MTP scheduling and state ownership rather than introducing performance-oriented changes to the execution path.

Test Plan

Added tests verify that current-step MTP is enabled by default for models declaring the capability, independently of first-audio configuration, and that generic preprocess no longer dispatches deferred MTP. An additional boundary test checks explicit rejection of unsupported multi-token sampling.

Existing tests now access resources through MTPState, with obsolete fields and duplicate initialization removed. MOSS factory and deployment tests reflect unconditional dedicated-state selection, while existing coverage for frame publication, chunked prefill, EOS handling, request reordering, slot reuse, RNG, graph dispatch, and output ownership is retained.

CPU-focused coverage:

pytest -v \
  tests/worker_v2/test_omni_model_state.py \
  tests/worker_v2/test_omni_ar_model_runner.py \
  tests/worker_v2/test_omni_gpu_model_runner.py \
  tests/worker_v2/test_omni_ar_model_runner_cuda.py \
  tests/worker_v2/test_eager_frame_kernels.py \
  tests/worker/test_talker_mtp_graph_budget.py \
  tests/model_executor/models/moss_tts/test_stage0_first_audio.py \
  tests/model_executor/models/moss_tts/test_local_model_state.py \
  tests/model_executor/models/moss_tts/test_local_eager_mtp.py \
  tests/model_executor/models/voxcpm2/test_per_patch_audio.py \
  tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py \
  tests/config/test_config_factory.py

Test Result

849 passed, no failed test.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.

Note

To avoid agent spam, vLLM-Omni limits contributors without write access to 5 open PRs. If this blocks well-intentioned critical work, please contact a committer to request bypass list access.
(anything written below this line will be removed by GitHub Actions)

^{Need help on this PR? Tag @codesmith-bot with what you need.}

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

chatgpt-codex-connector · 2026-10-25T07:45:51Z

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-25T07:45:45.478710Z We couldn’t complete this review. Comment `@codex review` to try again.	`c427588`	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 · 2026-10-25T07:45:54Z

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 · 2026-10-25T07:48:50Z

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

Module owners: @tzhouam @fake0fan @Gaohan123

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

@sphinxkkkbc, 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.

amy-why-3459

Thanks for consolidating MTP into MTPState / post_sample and retiring the generic deferred path — the ownership split and the updated CPU coverage look directionally right.

A few concrete points below on the runner gate, the Qwen3-TTS enablement expansion, and the Moss Local always-on state choice. Not blocking on style; mainly want to make sure the retired deferred path does not leave a silent crash or config landmine.

amy-why-3459 · 2026-10-25T07:51:54Z

-        if multimodal_outputs and run_eager_mtp is not None:
-            run_eager_mtp(
+        post_sample = getattr(self.model_state, "post_sample", None)
+        if post_sample is not None:


Previously this was if multimodal_outputs and run_eager_mtp is not None. Dropping the multimodal_outputs check means post_sample always runs when the method exists.

OmniModelState.post_sample is a no-op without _has_mtp, and MTPState.post_sample early-returns when recorded is empty — those cases are fine. But if preprocess recorded MTP rows (_has_mtp + non-empty recorded) and multimodal_outputs is still None/empty here, MTPState.post_sample will hit multimodal_outputs.get(...) (AttributeError) or the "requires retained codes.audio" RuntimeError.

Worth either restoring a soft multimodal_outputs guard at this call site, or hardening MTPState.post_sample to tolerate a missing payload (clear early return / explicit error) so a missing retain path cannot turn into an opaque crash. VoxCPM2's override also now always runs its per-patch end tracking; that looks intentional, but please confirm it is OK when multimodal_outputs is empty.

Thanks for noting this. Added the empty-MM guard to OmniModelState.post_sample(), so the default implementation only dispatches built-in MTP when MM output is available. The runner still calls the lifecycle hook unconditionally. VoxCPM2 keeps its own override and existing per-step patch decoding and stop handling; it does not go through this guard.

amy-why-3459 · 2026-10-25T07:51:54Z

        # its reference context independently of the two-stage codec profile.
        self.stream_ref_context_frames = stream_ref_context_frames(vllm_config) if self.stream_decode else 0
-        self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) or self.stream_decode
+        self.has_talker_mtp = True


Behavior change vs. the previous mtp_eager_frames = talker_first_audio_enabled(...) or self.stream_decode.

has_talker_mtp = True always flips eager_frames_active in OmniModelState, so current-step MTP + post_sample becomes the only path. make_omni_output still only allocates the token-major codes.audio / codec_frame_valid placeholders under async_chunk and eager_frames_active (_make_eager_omni_output). Without async-chunk (or without that placeholder path), preprocess can record MTP rows while forward returns {} / no retained tensors, and post_sample then fails the retained-output check.

Please confirm every supported Qwen3-TTS MRV2 deploy path either (a) always has async_chunk (or otherwise retains those tensors), or (b) gate has_talker_mtp / eager placeholders so the retired deferred path cannot leave a hard failure for previously-working configs. A short note in the PR / model docs would help too.

[P1] Preserve request isolation in synchronous Qwen3-TTS

With MRV2, async_chunk=False, and prefix caching disabled, make_omni_output() still gathers the retained prefill-sized codes.audio buffers. Current-step MTP fills the outgoing tensor but does not replace those buffers. For example, after prefill spans of 3 and 4 tokens, the next two-request decode produces a [7, Q] codec tensor for a two-token batch. _slice_pooler_value() sees that the leading dimension does not match the current token count and clones the whole tensor into each request, mixing their codec frames.

Please allocate outputs matching the current token spans whenever current-step MTP is active, including synchronous mode and its metadata, and add a two-request prefill-to-decode regression test.

This is based on source tracing at c427588; no native/GPU reproduction was run.

Fixed the non-async output assembly issue raised by @amy-why-3459 and detailed by @hsliuustc0106. Qwen3 TTS now allocates codes and validity from the current step’s token spans, rather than stale request-buffer codes, while preserving synchronous reference-code metadata. post_sample() then fills the current frames.

Also fixed the corresponding missing-validity container in Qwen3 Omni’s non-async path. CPU regression tests reproduce the previous failures and now pass for async/non-async and mixed prefill/decode batches, including request isolation.

amy-why-3459 · 2026-10-25T07:51:54Z

-
-            state_cls = MossLocalModelState
-        return state_cls(vllm_config, self, encoder_cache, device)
+        return MossLocalModelState(vllm_config, self, encoder_cache, device)


Unconditional MossLocalModelState (and dropping the mrv2_gpu_slot_state gate for first-frame decoder below) is a silent migration for any deploy that omitted / set false mrv2_gpu_slot_state and previously stayed on generic OmniModelState + deferred MTP.

Local state also rejects mtp_inputs in preprocess, so those configs cannot fall back. Please call this out in the PR description / release notes, and confirm the default CUDA yaml (already Local) is the only supported MRV2 Local path going forward — especially for Ascend/NPU where UVA is skipped and Local takes the non-UVA gather path.

Agreed—this changes state selection for configurations that omitted or explicitly disabled mrv2_gpu_slot_state. I’ll call out that migration in the PR description.

MRV2 Local now always selects MossLocalModelState; its model-owned deferred/eager scheduling remains controlled separately by mrv2_eager_mtp. The default CUDA YAML already selected the dedicated state, so its state selection is unchanged. The repository’s NPU configuration still uses MRV1; this PR does not establish support for a custom NPU MRV2 Local deployment. Removing the first-frame decoder’s slot-state gate follows the unconditional state selection; the other first-audio eligibility checks remain.

amy-why-3459 · 2026-10-25T07:51:54Z

@@ -233,49 +232,6 @@ def test_async_output_preserves_fixed_token_scores_on_host(monkeypatch) -> None:
    torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req-0"], scores)


test_chunked_fixed_token_scores_are_copied_from_cuda (CUDA L4 hardware test for chunked fixed-token prompt logprobs → host) looks deleted here without a replacement in this PR. It does not appear tied to the MTP refactor. Was this intentional (moved / covered elsewhere), or an accidental drop?

This test was moved unchanged to tests/worker_v2/test_omni_ar_model_runner_cuda.py due to pre-commit requirement, rather than removed. The separate file avoids combining the original file’s module-level CPU mark with accelerator marks.

vllm-omni-review-bot · 2026-10-25T07:53:37Z

Omni ReviewBot triage note

Automated triage of commit b51d7a8d191c produced:

Priority: high. Prompt maintainer attention is suggested.

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

…ssion test Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

amy-why-3459

Thanks for the follow-up on tip 1dd81626 — prior points look addressed from the code side.

Empty-MM guard in OmniModelState.post_sample (_has_mtp and multimodal_outputs) looks right, and the new regression covers None / {}.
Qwen3 TTS / Omni non-async now size codes / validity from the current step's token spans (Omni via _make_eager_talker_output; TTS via current_codes), which closes the request-isolation hole. The new sync/async regression tests look solid.
Confirmed test_chunked_fixed_token_scores_are_copied_from_cuda lives in tests/worker_v2/test_omni_ar_model_runner_cuda.py on this tip (moved for pre-commit, not dropped).
Moss Local always-on state callout is present in the PR description — thanks.

One tiny doc nit: the description still pins vLLM-Omni Commit to c4275885…; worth bumping to this tip when you next edit the body. No new substantive blockers on this delta from my side.

sphinxkkkbc · 2026-10-27T03:30:40Z

@amy-why-3459 Thanks for checking! I've updated the PR description with the latest SHA, test results, and the MOSS Local state migration details. Full-model inference and performance remain untested.

Brings in vllm-project#9507, which only appends batched decode rows to the MTP rows when the model state runs current-step MTP. Since vllm-project#9346 every batched row was treated as deferred MTP, so Chatterbox's non-MTP decode hooks failed the first Model Runner V2 decode step; with main's guard they run as they did before vllm-project#9346, and Chatterbox needs no change of its own. Signed-off-by: Will Rice <wrice20@gmail.com>

sphinxkkkbc added 3 commits October 25, 2026 11:03

refactor

c44e864

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

cleanup

52abd96

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

fix pre-commit pytest mark conflict

c427588

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

hsliuustc0106 added core related to core module: cache, scheduler, engine, worker, modelrunner refactor refactoring for better code scalability and quality labels Oct 25, 2026 — with Cursor

sphinxkkkbc marked this pull request as ready for review October 25, 2026 07:45

sphinxkkkbc requested review from Gaohan123, NickCao, Sy0307, alex-jw-brooks, amy-why-3459, gcanlin, hsliuustc0106, linyueqian, lishunyang12, tzhouam, yenuo26 and ywang96 as code owners October 25, 2026 07:45

amy-why-3459 reviewed Oct 25, 2026

View reviewed changes

hsliuustc0106 added the high priority high priority issue, needs to be done asap label Oct 25, 2026

resolve comments, add guard for non-async+current_step mtp with regre…

1dd8162

…ssion test Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

amy-why-3459 reviewed Oct 27, 2026

View reviewed changes

Merge branch 'main' into refactor/unify_mtp_contract

b51d7a8

amy-why-3459 added the ready label to trigger buildkite CI label Oct 27, 2026

amy-why-3459 enabled auto-merge (squash) October 27, 2026 03:39

amy-why-3459 approved these changes Oct 27, 2026

View reviewed changes

amy-why-3459 merged commit 74d3e26 into vllm-project:main Oct 27, 2026
8 of 11 checks passed

sphinxkkkbc mentioned this pull request Oct 27, 2026

[RFC]: Make current-step MTP the MRV2 default via model_state.post_sample #9173

Closed

1 task done

tzhouam mentioned this pull request Oct 27, 2026

[cherry-pick][release/v0.32.0] Backport 18 bugfixes for milestone issues (batch 7) #9470

Merged

DarkMatter9527 mentioned this pull request Oct 28, 2026

[Hardware][Ascend] Support VoxCPM2 on the NPU Model Runner V2 #9507

Merged

sphinxkkkbc mentioned this pull request Oct 28, 2026

[Bugfix] Guard MTP input collection in batched MRV2 preprocess #9540

Open

twu3202 mentioned this pull request Oct 30, 2026

[Model] Run the Fish DAC stage on Model Runner V2 #9625

Open

		@@ -233,49 +232,6 @@ def test_async_output_preserves_fixed_token_scores_on_host(monkeypatch) -> None:
		torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req-0"], scores)

Conversation

sphinxkkkbc commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Contributor intent (required human answers)

1. What outcome do you want to achieve?

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

3. How will this work move that variable?

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

Purpose

What's Changed

Model Side Change

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector Bot commented Oct 25, 2026

Codex Review Summary

Uh oh!

chatgpt-codex-connector Bot commented Oct 25, 2026

Uh oh!

vllm-omni-review-bot commented Oct 25, 2026

Uh oh!

amy-why-3459 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vllm-omni-review-bot commented Oct 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Omni ReviewBot triage note

Uh oh!

amy-why-3459 left a comment

Choose a reason for hiding this comment

Uh oh!

sphinxkkkbc commented Oct 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

sphinxkkkbc commented Oct 25, 2026 •

edited

Loading

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

edited

Loading