refactor: clarify seeded Qwen3-TTS MTP generator dispatch

Signed-off-by: Sy03 <1370724210@qq.com>

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -19,10 +19,6 @@ class _FakeGroup:
         self.group_ranks = group_ranks
         self.parallel_mode = parallel_mode
         self.device_group = object()
-        self.device_communicator = kwargs.get("device_communicator")
-        reduce_scatter = kwargs.get("reduce_scatter")
-        if reduce_scatter is not None:
-            self.reduce_scatter = reduce_scatter
         self.ulysses_group = kwargs.get("ulysses_group")
         self.ring_group = kwargs.get("ring_group")
         self.local_group = next(group for group in group_ranks if local_rank in group)
@@ -59,23 +55,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         created_groups.append(group)
         return group
 
-    def fake_init_vllm_model_parallel_group(
-        group_ranks,
-        local_rank,
-        backend,
-        group_name,
-    ):
-        del backend
-        group = _FakeGroup(
-            [list(ranks) for ranks in group_ranks],
-            local_rank,
-            f"vllm_{group_name}",
-            device_communicator=object(),
-            reduce_scatter=lambda tensor, **kwargs: tensor,
-        )
-        created_groups.append(group)
-        return group
-
     fake_world_group = SimpleNamespace(
         rank_in_group=local_rank,
         local_rank=local_rank,
@@ -89,7 +68,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
-    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
     for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
@@ -109,7 +87,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
     assert parallel_state.vllm_parallel_state._PCP.world_size == 2
     assert parallel_state._DP.world_size == 2
     assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
@@ -118,13 +96,6 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._TP.world_size == 2
     assert parallel_state._PP.world_size == 2
 
-    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
-    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
-    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
-    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
-    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
-    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
-
     assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
     assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
     assert parallel_state.vllm_parallel_state._EP.local_group == [
@@ -184,9 +155,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
-    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,48 +50,6 @@ class TestSetForwardContextNumTokens:
         mock_get.assert_not_called()
 
 
-class TestSetForwardContextDPMetadata:
-    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
-
-    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = None
-        fake_group = mocker.MagicMock()
-        fake_group.world_size = 2
-        fake_group.cpu_group = object()
-
-        def _all_gather_object(output, obj, group):
-            assert group is fake_group.cpu_group
-            output[:] = [obj, obj + 1]
-
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
-        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is not None
-        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
-
-    def test_keeps_existing_dp_metadata(self, mocker):
-        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
-
-        existing_metadata = object()
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.dp_metadata = existing_metadata
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
-
-        hunyuan_moe._set_forward_context_dp_metadata(1024)
-
-        assert mock_ctx.dp_metadata is existing_metadata
-        mock_all_gather.assert_not_called()
-
-
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -22,7 +22,6 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
-        "vae_use_tiling": True,
     },
 )
 
@@ -33,7 +32,6 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
-        "vae_use_tiling": True,
     },
 )
 
@@ -47,6 +45,10 @@ pytestmark = [
     ),
 ]
 
+_SKIP_ISSUE_3236 = pytest.mark.skip(
+    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
+)
+
 
 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -72,6 +74,7 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     container.close()
 
 
+@_SKIP_ISSUE_3236
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
     """End-to-end test for MagiHuman generating video and audio."""

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
         windows stitch into a coherent PCM stream. The byte-count gate is the same as
         the sync paths; per-chunk audio content is verified offline against Whisper.
 
-        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
         check still catches catastrophic codec failure (silence, white noise, sample
         scramble all give HNR << 0) while allowing for the sliding-window codec's
         slightly-noisier-than-sync output. The default 1.0 dB threshold has only

tests/entrypoints/test_omni_sleep_mode.py

@@ -179,7 +179,7 @@ class TestOmniSleepMode:
         try:
             device_id = 0
             used_before = get_device_global_memory_used_gib(device_id)
-            acks = await llm_engine.sleep(stage_ids=[0], level=1)
+            acks = await llm_engine.sleep(stage_ids=[0], level=2)
             await asyncio.sleep(1.5)
             used_after = get_device_global_memory_used_gib(device_id)
             drop_gib = used_before - used_after
@@ -209,7 +209,7 @@ class TestOmniSleepMode:
         """Diffusion Worker stage signal loop"""
         try:
             logger.info("Starting Diffusion Worker Handshake Test")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
 
             def _get_status(ack):
                 return ack.status if hasattr(ack, "status") else ack.get("status")
@@ -264,9 +264,9 @@ class TestOmniSleepMode:
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
-            # Sleep Level 1
+            # Sleep Level 2
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=1)
+            await diffusion_engine.sleep(stage_ids=[0], level=2)
             # Wake-up
             logger.info("Waking up (Reloading Weights)...")
             await diffusion_engine.wake_up(stage_ids=[0])
@@ -353,8 +353,8 @@ class TestOmniSleepMode:
             logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
 
             # Sleep
-            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
 
             reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
             reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,47 +407,6 @@ class TestOmniSleepMode:
             diffusion_engine.shutdown()
             await asyncio.sleep(1)
 
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
-        NotImplementedError instead of silently producing corrupted output."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=2)
-            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
-                await llm_engine.wake_up(stage_ids=[0])
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
-        is still asleep after wake_up(tags=["weights"]), instead of crashing with
-        CUDA illegal memory access."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
-            with pytest.raises(RuntimeError, match="partially or fully asleep"):
-                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
-                    pass
-        finally:
-            llm_engine.shutdown()
-
-    @pytest.mark.asyncio
-    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
-        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
-        safe no-op instead of raising a cumem CUDA invalid argument error."""
-        try:
-            await llm_engine.sleep(stage_ids=[0], level=1)
-            first_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert len(first_acks) > 0, "First wake_up() should return ACKs"
-            second_acks = await llm_engine.wake_up(stage_ids=[0])
-            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
-        finally:
-            llm_engine.shutdown()
-
 
 # ---------------------------------------------------------------------------
 # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
@@ -485,7 +444,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int):
         async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
             pass
 
-        acks = await engine.sleep(level=1)
+        acks = await engine.sleep(level=2)
         statuses = [get_ack_info(ack, "status") for ack in acks]
         assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
 
@@ -541,7 +500,7 @@ async def test_multistage_sleep_h100(tp_size: int):
         async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
             pass
 
-        acks = await engine.sleep(stage_ids=[0, 1], level=1)
+        acks = await engine.sleep(stage_ids=[0, 1], level=2)
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -7,7 +7,6 @@ without requiring the actual checkpoint or GPU.
 """
 
 import asyncio
-import json
 import time
 from collections import OrderedDict, defaultdict
 from types import SimpleNamespace
@@ -1310,10 +1309,7 @@ class TestVoiceCloneReferenceCache:
         parent = tmp_path / "OmniVoice"
         subdir = parent / "audio_tokenizer"
         subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text(
-            json.dumps({"model_type": "higgs_audio_v2_tokenizer"}),
-            encoding="utf-8",
-        )
+        (subdir / "config.json").write_text("{}", encoding="utf-8")
 
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
@@ -1321,32 +1317,6 @@ class TestVoiceCloneReferenceCache:
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
 
-    def test_audio_tokenizer_dir_rejects_omnivoice_config(self, tmp_path, monkeypatch):
-        import huggingface_hub
-        import huggingface_hub.constants as hub_constants
-
-        from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok
-
-        parent = tmp_path / "OmniVoice"
-        subdir = parent / "audio_tokenizer"
-        subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text(
-            json.dumps({"model_type": "omnivoice"}),
-            encoding="utf-8",
-        )
-
-        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
-        monkeypatch.setenv("HIGGS_AUDIO_V2_TOKENIZER_PATH", "")
-        monkeypatch.setattr(
-            huggingface_hub,
-            "try_to_load_from_cache",
-            lambda **_: str(subdir / "config.json"),
-        )
-        monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path / "empty_cache"))
-
-        assert tok._normalize_audio_tokenizer_dir(str(parent)) is None
-        assert tok._resolve_audio_tokenizer_dir() is None
-
     def test_higgs_v3_ref_code_cache_returns_clone(self):
         from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 

tests/worker/test_gpu_ar_model_runner.py

@@ -480,52 +480,6 @@ def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_onl
     assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3])
 
 
-def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch):
-    """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
-    runner = _make_async_output_runner(engine_output_type="latent")
-    runner.model.omni_pooler_payload_include_hidden = False
-    runner._async_chunk = False
-    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
-
-    monkeypatch.setattr(
-        GPUARModelRunner,
-        "_resolve_pooler_payload_req_ids",
-        lambda self, req_ids: ("latent", req_ids),
-    )
-    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
-    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
-    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
-
-    layers = torch.arange(96, dtype=torch.long).reshape(3, 32)
-    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
-        runner,
-        scheduler_output=SimpleNamespace(
-            total_num_scheduled_tokens=5,
-            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
-        ),
-        hidden_states=torch.randn(3, 4),
-        staged_hidden_states_cpu=None,
-        multimodal_outputs={"hidden_states": {"layers": {0: layers}}},
-        req_ids_output_copy=["r1", "r2", "r3"],
-        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
-        valid_sampled_token_ids=[[101], [102], [103]],
-        logprobs_lists=None,
-        prompt_logprobs_dict={},
-        num_nans_in_logits=None,
-        kv_connector_output=None,
-        ec_connector_output=None,
-        cudagraph_stats=None,
-        kv_extracted_req_ids=None,
-        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
-        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
-    )
-
-    assert output.inter_stage_outputs is not None
-    assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1])
-    assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2])
-    assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3])
-
-
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     runner = _make_async_output_runner()
     runner.model.omni_pooler_payload_include_hidden = False

tests/worker/test_omni_gpu_model_runner.py

@@ -171,6 +171,12 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     return runner
 
 
+def _enable_talker_seed(runner, *, accepts_generators=False):
+    runner.model.talker_mtp_seed_extra_arg = "tts_local_seed"
+    if accepts_generators:
+        runner.model.talker_mtp_accepts_generators = True
+
+
 def _make_runner_for_mimo(req_id="r_mimo"):
     """Create a minimal runner with MiMoAudio-like model and request state."""
     runner = object.__new__(OmniGPUModelRunner)
@@ -255,6 +261,7 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
 
     def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
@@ -269,6 +276,34 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
     assert runner.talker_mtp.calls[0]["generator"] is None
 
 
+def test_talker_mtp_forward_ignores_tts_local_seed_without_model_seed_key(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=42,
+        extra_args={"tts_local_seed": 42},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model.talker_mtp_accepts_generators = True
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert runner.talker_mtp.calls[0]["generators"] is None
+
+
 def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
@@ -280,6 +315,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
         extra_args={"tts_local_seed": 42},
     )
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(
             subtalker_sampling_params={
@@ -314,7 +350,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
     assert runner.talker_mtp.calls[0]["generator"] is not None
 
 
-def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
+def test_talker_mtp_forward_keeps_unmarked_models_scalar_for_explicit_seed(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
     monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
@@ -329,6 +365,7 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
         extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
 
     def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
@@ -348,13 +385,13 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
 
     assert [call["batch_size"] for call in runner.talker_mtp.calls] == [1, 1]
     assert all(call["generator"] is not None for call in runner.talker_mtp.calls)
+    assert all(call["generators"] is None for call in runner.talker_mtp.calls)
     assert runner.talker_mtp.calls[0]["generator"] is not runner.talker_mtp.calls[1]["generator"]
     assert torch.equal(runner.talker_mtp_input_ids.gpu, saved_input_ids)
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
-def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
-    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+def test_talker_mtp_forward_keeps_explicit_seeded_requests_batched(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
     monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
@@ -369,10 +406,7 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
         extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
-    runner.model = SimpleNamespace(
-        talker_mtp_output_key=("codes", "audio"),
-        talker_mtp_accepts_per_row_generators=True,
-    )
+    _enable_talker_seed(runner, accepts_generators=True)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
 
     def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
@@ -386,8 +420,9 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
 
     # One batched call with distinct per-row generators, not two scalar calls.
     assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
-    row_generators = runner.talker_mtp.calls[0]["generators"]
     assert runner.talker_mtp.calls[0]["generator"] is None
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert row_generators is not None
     assert len(row_generators) == 2
     assert all(generator is not None for generator in row_generators)
     assert row_generators[0] is not row_generators[1]
@@ -404,6 +439,62 @@ def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypa
     assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
 
 
+def test_talker_mtp_forward_batches_mixed_explicit_seed_requests(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("seeded", "unseeded"), hidden_size=4)
+    runner.requests["seeded"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner, accepts_generators=True)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["seeded", "unseeded"], inputs_embeds)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    generators = runner.talker_mtp.calls[0]["generators"]
+    assert generators is not None
+    assert len(generators) == 2
+    assert generators[0] is not None
+    assert generators[1] is None
+
+
+def test_talker_mtp_forward_does_not_pass_generators_without_explicit_seed(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.talker_mtp = CaptureTalkerMTP()
+    _enable_talker_seed(runner, accepts_generators=True)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert runner.talker_mtp.calls[0]["generators"] is None
+
+
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""

vllm_omni/diffusion/data.py

@@ -6,7 +6,6 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
-from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -1489,12 +1488,5 @@ class OmniWakeTask:
     tags: list[str] | None = None
 
 
-class CuMemTag(str, Enum):
-    """Tags representing specific CuMem allocations for sleep/wake state tracking."""
-
-    WEIGHTS = "weights"
-    KV_CACHE = "kv_cache"
-
-
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}

vllm_omni/diffusion/distributed/parallel_state.py

@@ -512,21 +512,6 @@ def init_model_parallel_group(
         )
 
 
-def init_vllm_model_parallel_group(
-    group_ranks: list[list[int]],
-    local_rank: int,
-    backend: str,
-    group_name: str,
-) -> vllm_parallel_state.GroupCoordinator:
-    return vllm_parallel_state.init_model_parallel_group(
-        group_ranks=group_ranks,
-        local_rank=local_rank,
-        backend=backend,
-        group_name=group_name,
-        use_device_communicator=True,
-    )
-
-
 def init_dit_group(
     dit_parallel_size: int,
     backend: str,
@@ -853,40 +838,23 @@ def initialize_model_parallel(
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
         # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
-        # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator
-        # and reduce_scatter(), which the diffusion SP coordinator intentionally
-        # does not own. Keep the rank layout but build a vLLM coordinator.
-        vllm_parallel_state._PCP = init_vllm_model_parallel_group(
-            group_ranks=sp_group_ranks,
-            local_rank=get_world_group().local_rank,
-            backend=backend,
-            group_name="pcp",
-        )
+        vllm_parallel_state._PCP = _SP
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    tp_group_ranks = rank_generator.get_ranks("tp")
-    if use_moe_parallel_mapping:
-        vllm_parallel_state._TP = init_vllm_model_parallel_group(
-            group_ranks=tp_group_ranks,
-            local_rank=get_world_group().local_rank,
-            backend=backend,
-            group_name="tp",
-        )
-    else:
-        vllm_parallel_state._TP = init_model_parallel_group(
-            group_ranks=tp_group_ranks,
-            local_rank=get_world_group().local_rank,
-            backend=backend,
-            parallel_mode="tensor",
-        )
-    if use_moe_parallel_mapping:
+    vllm_parallel_state._TP = init_model_parallel_group(
+        group_ranks=rank_generator.get_ranks("tp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="tensor",
+    )
+    if use_moe_parallel_mapping and cfg_parallel_size > 1:
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
-        vllm_parallel_state._DP = init_vllm_model_parallel_group(
+        vllm_parallel_state._DP = init_model_parallel_group(
             group_ranks=rank_generator.get_ranks("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
-            group_name="dp",
+            parallel_mode="data",
         )
 
     global _FS, _EXPERT_PARALLEL_GROUP_RANKS
@@ -901,11 +869,11 @@ def initialize_model_parallel(
     _EXPERT_PARALLEL_GROUP_RANKS = None
     if use_moe_parallel_mapping:
         ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
-        vllm_parallel_state._EP = init_vllm_model_parallel_group(
+        vllm_parallel_state._EP = init_model_parallel_group(
             group_ranks=ep_group_ranks,
             local_rank=get_world_group().local_rank,
             backend=backend,
-            group_name="ep",
+            parallel_mode="expert",
         )
         _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -3,8 +3,6 @@
 
 from typing import Any
 
-import torch
-import vllm.distributed.parallel_state as _vllm_ps
 import vllm.forward_context as _vllm_fc
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
@@ -26,30 +24,6 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
         forward_context.in_profile_run = False
 
 
-def _set_forward_context_dp_metadata(num_tokens: int) -> None:
-    """Populate vLLM MoE DP token metadata for diffusion-only DP mappings."""
-    if not _vllm_fc.is_forward_context_available():
-        return
-    forward_context = _vllm_fc.get_forward_context()
-    if getattr(forward_context, "dp_metadata", None) is not None:
-        return
-
-    dp_group = getattr(_vllm_ps, "_DP", None)
-    if dp_group is None or getattr(dp_group, "world_size", 1) <= 1:
-        return
-
-    gathered_num_tokens: list[int | None] = [None] * dp_group.world_size
-    torch.distributed.all_gather_object(
-        gathered_num_tokens,
-        int(num_tokens),
-        group=dp_group.cpu_group,
-    )
-    if any(count is None for count in gathered_num_tokens):
-        raise RuntimeError(f"Failed to gather MoE DP token counts: {gathered_num_tokens}")
-
-    forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64))
-
-
 class HunyuanFusedMoEDefault:
     """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
 
@@ -88,7 +62,6 @@ class HunyuanFusedMoEDefault:
                 hidden_states = args[0]
             if hidden_states is not None:
                 _set_forward_context_num_tokens(hidden_states.shape[0])
-                _set_forward_context_dp_metadata(hidden_states.shape[0])
 
         moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
 

vllm_omni/entrypoints/async_omni.py

@@ -26,7 +26,7 @@ from vllm.tasks import SupportedTask
 from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask
+from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import (
@@ -141,8 +141,7 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
-        self._sleeping_tags: set[str] = set()
-        self._level2_sleeping: bool = False
+        self._is_sleeping: bool = False
         self.final_output_task: asyncio.Task | None = None
         self.event_resolver = AsyncEventResolver(orchestrator=self)
         self.config_path = self.engine.config_path
@@ -209,13 +208,10 @@ class AsyncOmni(EngineClient, OmniBase):
 
     def get_diffusion_od_config(self) -> Any | None:
         """Return the diffusion-stage config when the pipeline has one."""
-        saw_diffusion_stage = False
         for stage_client in self.engine.stage_clients:
             if getattr(stage_client, "stage_type", None) != "diffusion":
                 continue
 
-            saw_diffusion_stage = True
-
             od_config = getattr(stage_client, "od_config", None)
             if od_config is not None:
                 return od_config
@@ -225,11 +221,6 @@ class AsyncOmni(EngineClient, OmniBase):
             if od_config is not None:
                 return od_config
 
-        # Out-of-process diffusion clients don't carry od_config (it lives in the
-        # worker); fall back to the engine's model_class_name resolution.
-        if saw_diffusion_stage:
-            return self.engine.get_diffusion_od_config()
-
         return None
 
     @property
@@ -315,14 +306,6 @@ class AsyncOmni(EngineClient, OmniBase):
 
         logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
-        _sleeping_tags = getattr(self, "_sleeping_tags", None)
-        if _sleeping_tags:
-            raise RuntimeError(
-                f"Generation rejected: Engine is partially or fully asleep. "
-                f"Currently sleeping tags: {list(_sleeping_tags)}. "
-                f"Please perform a full wake_up before generating."
-            )
-
         # Reject diffusion list-prompt early with a clear API error.
         if isinstance(prompt, list) and any(
             getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
@@ -939,32 +922,11 @@ class AsyncOmni(EngineClient, OmniBase):
                 if ack is not None:
                     await self.event_resolver.resolve(ack)
                     final_acks.append(ack)
-        if not hasattr(self, "_sleeping_tags"):
-            self._sleeping_tags = set()
-        self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
-        if level == 2:
-            self._level2_sleeping = True
+        self._is_sleeping = True
         return final_acks
 
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
         self._final_output_handler()
-
-        if getattr(self, "_level2_sleeping", False):
-            raise NotImplementedError(
-                "wake_up() after sleep(level=2) is not yet implemented: weights were "
-                "discarded from GPU and reloading from disk is not yet supported. "
-                "Use sleep(level=1) instead, which offloads weights to CPU RAM "
-                "and supports fast DMA restore."
-            )
-        _current_tags = getattr(self, "_sleeping_tags", set())
-        if tags is None:
-            requested_tags = list(_current_tags)
-        else:
-            requested_tags = [t for t in tags if t in _current_tags]
-        if not requested_tags:
-            logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
-            return []
-
         if stage_ids is None:
             stage_ids = list(range(len(self.engine.stage_clients)))
         total_workers = 0
@@ -978,7 +940,7 @@ class AsyncOmni(EngineClient, OmniBase):
         task_id = str(uuid.uuid4())
         self.event_resolver.watch_task(task_id, expected_count=total_workers)
         logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
+        task = OmniWakeTask(tags=tags, task_id=task_id)
         rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
         final_acks = []
         for stage_res in rpc_results:
@@ -989,14 +951,7 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_acks.append(ack)
         current_omni_platform.synchronize()
         await asyncio.sleep(0.1)
-
-        for t in requested_tags:
-            if hasattr(self, "_sleeping_tags"):
-                self._sleeping_tags.discard(t)
-        # Only clear the level-2 flag once all tags are warm, in case partial
-        # wake support (e.g. tags=["kv_cache"] only) is added in the future.
-        if not getattr(self, "_sleeping_tags", None):
-            self._level2_sleeping = False
+        self._is_sleeping = False
         logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
         return final_acks
 
@@ -1006,7 +961,7 @@ class AsyncOmni(EngineClient, OmniBase):
         TODO(AsyncOmni): query the orchestrator once all stage backends expose
         a real sleeping-state RPC. For now we track the requested state locally.
         """
-        return bool(getattr(self, "_sleeping_tags", None))
+        return self._is_sleeping
 
     async def add_lora(self, lora_request: LoRARequest) -> bool:
         """Load a new LoRA adapter into all stages.

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -22,7 +22,6 @@ generation.
 
 from __future__ import annotations
 
-import json
 import os
 import re
 from typing import Any
@@ -171,25 +170,14 @@ _AUDIO_TOKENIZER_PATH_ENVS = (
 )
 
 
-def _is_higgs_audio_tokenizer_config(config_path: str) -> bool:
-    try:
-        with open(config_path, encoding="utf-8") as f:
-            config = json.load(f)
-    except (OSError, json.JSONDecodeError):
-        return False
-    return config.get("model_type") == "higgs_audio_v2_tokenizer"
-
-
 def _normalize_audio_tokenizer_dir(path: str) -> str | None:
     if not path:
         return None
     expanded = os.path.abspath(os.path.expanduser(path))
-    config_path = os.path.join(expanded, "config.json")
-    if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path):
+    if os.path.isfile(os.path.join(expanded, "config.json")):
         return expanded
     nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR)
-    nested_config_path = os.path.join(nested, "config.json")
-    if os.path.isfile(nested_config_path) and _is_higgs_audio_tokenizer_config(nested_config_path):
+    if os.path.isfile(os.path.join(nested, "config.json")):
         return nested
     return None
 
@@ -213,9 +201,7 @@ def _resolve_audio_tokenizer_dir() -> str | None:
         filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json",
     )
     if isinstance(cached_config, str) and os.path.isfile(cached_config):
-        candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config))
-        if candidate is not None:
-            return candidate
+        return os.path.dirname(cached_config)
 
     from huggingface_hub.constants import HF_HUB_CACHE
 
@@ -254,9 +240,7 @@ def _load_audio_tokenizer():
             _K2_OMNIVOICE_REPO,
             allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
         )
-        audio_tokenizer_dir = _normalize_audio_tokenizer_dir(repo_path)
-        if audio_tokenizer_dir is None:
-            raise RuntimeError(f"Downloaded {_K2_OMNIVOICE_REPO} does not contain a valid Higgs audio tokenizer")
+        audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
     device = current_omni_platform.get_torch_device()
     model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
     return model.eval()

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -266,6 +266,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
     """vLLM-AR talker: step-wise layer-0 codec decoding.
     Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`."""
 
+    talker_mtp_accepts_generators = True
+    talker_mtp_seed_extra_arg = "tts_local_seed"
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             # Talker backbone (Qwen3 decoder-only).
@@ -318,13 +321,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
-        # talker_mtp samples with per-row generators, so explicitly-seeded
-        # requests stay batched instead of one scalar forward per row (#4883).
-        # Only valid while talker_mtp receives the unpadded active batch (this
-        # talker is not graph-wrapped); a padded batch would need the runner to
-        # pad the generators list as well.
-        self.talker_mtp_accepts_per_row_generators = True
-
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -902,8 +902,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         ) = self.execute_model_state
         # Clear ephemeral state.
         self.execute_model_state = None
-        hidden_seq_len = int(hidden_states.shape[0])
-        scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
+        seq_len = hidden_states.shape[0]
 
         # Apply structured output bitmasks if present.
         if grammar_output is not None:
@@ -1142,8 +1141,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                                 start=start,
                                 end=end,
                                 pass_lists_through=False,
-                                seq_len=hidden_seq_len,
-                                scheduled_seq_len=scheduled_seq_len,
+                                seq_len=seq_len,
                             )
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))

vllm_omni/utils/mm_outputs.py

@@ -106,13 +106,7 @@ def _to_cpu(value):
 
 
 def to_payload_element(
-    element: object,
-    idx: int,
-    start: int,
-    end: int,
-    pass_lists_through: bool = False,
-    seq_len: int | None = None,
-    scheduled_seq_len: int | None = None,
+    element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None
 ):
     """Build an mm payload element corresponding to one request index
     from an element containing 0 or more CPU tensors.
@@ -126,39 +120,21 @@ def to_payload_element(
             passthrough data; this should be False in normal cases, but True
             if we need to avoid splitting nonempty lists prior to calling
             postprocess, which is the case for prefix cache.
-        seq_len: Optional hidden-aligned batch length (``hidden_states.shape[0]``).
-            When a tensor's first dimension equals this value, it is sliced
-            per request using ``start:end``.
-        scheduled_seq_len: Optional scheduler-aligned batch length
-            (``scheduler_output.total_num_scheduled_tokens``). Some full-payload
-            mm tensors (e.g. batched ``codes.audio`` with tail-only hidden states)
-            are laid out by scheduled tokens instead of the hidden tail shape.
-            When omitted, ``seq_len`` is reused for backward compatibility.
+        seq_len: Optional sequence length (i.e., dim 0 of hidden states).
+            When set, a tensor whose first dimension equals seq_len is
+            sliced per request. The prefix cache passthrough also passes
+            the total scheduled token count here so 1D (seq_len,) metadata
+            that is intentionally not cached is still split per request.
     """
-    if scheduled_seq_len is None:
-        scheduled_seq_len = seq_len
-
-    # Cached per-token tensors are merged elsewhere; here a first dim equal to
-    # either the hidden-aligned or scheduler-aligned batch length means a
-    # per-request slice is required.
-    if isinstance(element, torch.Tensor) and (
-        (seq_len is not None and element.shape[0] == seq_len)
-        or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len)
-    ):
+    # Cached per-token tensors are merged elsewhere; here a first dim
+    # equal to seq_len means a per-request slice is required.
+    if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len:
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)
     # and running a model without prefix caching.
     elif isinstance(element, dict):
         return {
-            sk: to_payload_element(
-                sv,
-                idx,
-                start,
-                end,
-                pass_lists_through=pass_lists_through,
-                seq_len=seq_len,
-                scheduled_seq_len=scheduled_seq_len,
-            )
+            sk: to_payload_element(sv, idx, start, end, pass_lists_through=pass_lists_through, seq_len=seq_len)
             for sk, sv in element.items()
         }
     elif isinstance(element, list):

vllm_omni/worker/gpu_ar_model_runner.py

@@ -839,8 +839,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         end: int,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        hidden_seq_len: int,
-        scheduled_seq_len: int,
+        seq_len: int,
     ) -> dict[str, object]:
         if combined_multimodal_outputs:
             return self._build_combined_prefix_cache_mm_payload(
@@ -877,8 +876,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 start=start,
                 end=end,
                 pass_lists_through=False,
-                seq_len=hidden_seq_len,
-                scheduled_seq_len=scheduled_seq_len,
+                seq_len=seq_len,
             )
         return mm_payload
 
@@ -896,8 +894,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         mm_cpu: dict[str, object] | None,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        hidden_seq_len: int,
-        scheduled_seq_len: int,
+        seq_len: int,
     ) -> dict[str, object]:
         payload: dict[str, object] = {}
         if not audio_sparse_output:
@@ -923,8 +920,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             end=end,
             audio_sparse_output=audio_sparse_output,
             sparse_mm_index=sparse_mm_index,
-            hidden_seq_len=hidden_seq_len,
-            scheduled_seq_len=scheduled_seq_len,
+            seq_len=seq_len,
         )
         payload.update(mm_payload)
         return payload
@@ -1735,8 +1731,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # The actual multimodal wire transport uses multimodal_outputs instead.
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
-            hidden_seq_len = int(hidden_states.shape[0])
-            scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
+            mm_seq_len = int(scheduler_output.total_num_scheduled_tokens)
             mm_cpu = None
             if self.omni_prefix_cache is not None:
                 (
@@ -1792,8 +1787,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         mm_cpu=mm_cpu,
                         audio_sparse_output=audio_sparse_output,
                         sparse_mm_index=sparse_mm_index,
-                        hidden_seq_len=hidden_seq_len,
-                        scheduled_seq_len=scheduled_seq_len,
+                        seq_len=mm_seq_len,
                     )
                     pooler_output.append(flatten_payload(payload))
 

vllm_omni/worker/gpu_model_runner.py

@@ -84,6 +84,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         self.omni_prefix_cache = None
         self._sampled_token_ids_cpu_override = None
         self._omni_query_start_loc_model_kwarg = False
+        self._talker_mtp_generators: dict[str, torch.Generator] = {}
 
     def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
         override_fn = self._sampled_token_ids_cpu_override
@@ -1785,6 +1786,84 @@ class OmniGPUModelRunner(GPUModelRunner):
             ec_connector_output,
         )
 
+    def _explicit_talker_seed(self, req_id: str) -> int | None:
+        seed_extra_arg = getattr(self.model, "talker_mtp_seed_extra_arg", None)
+        if not seed_extra_arg:
+            return None
+
+        sampling_params = getattr(self.requests[req_id], "sampling_params", None)
+        extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+        if not isinstance(extra_args, dict):
+            return None
+
+        seed = extra_args.get(seed_extra_arg)
+        return int(seed) if seed is not None else None
+
+    def _resolve_talker_mtp_generators(
+        self,
+        decode_req_ids: list[str],
+        device: torch.device,
+    ) -> list[torch.Generator | None] | None:
+        """Resolve per-row generators aligned to ``decode_req_ids``.
+
+        ``row_generators[i]`` belongs to the batched input row for
+        ``decode_req_ids[i]``. Explicitly seeded rows use cached per-request
+        generators, while unseeded rows stay as ``None`` to keep global RNG
+        semantics. Return ``None`` when the whole batch is unseeded so callers
+        can use the normal batched path without sampling kwargs.
+        """
+        generator_cache = getattr(self, "_talker_mtp_generators", None)
+        if generator_cache is None:
+            generator_cache = {}
+            self._talker_mtp_generators = generator_cache
+
+        row_generators: list[torch.Generator | None] = []
+        any_seeded = False
+        for req_id in decode_req_ids:
+            seed = self._explicit_talker_seed(req_id)
+            if seed is None:
+                row_generators.append(None)
+                continue
+
+            generator = generator_cache.get(req_id)
+            if generator is None or generator.device != device:
+                generator = torch.Generator(device=device)
+                generator.manual_seed(seed)
+                generator_cache[req_id] = generator
+            row_generators.append(generator)
+            any_seeded = True
+
+        for stale_id in [rid for rid in generator_cache if rid not in self.requests]:
+            del generator_cache[stale_id]
+
+        return row_generators if any_seeded else None
+
+    def _talker_mtp_forward_row_by_row(
+        self,
+        decode_req_ids: list[str],
+        inputs_embeds: torch.Tensor,
+        start_offsets: list[int] | None = None,
+    ) -> None:
+        """Seeded fallback for models without batched per-row generators."""
+        decode_batch_size = len(decode_req_ids)
+        saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
+        saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
+        saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
+        saved_text = self.text_step.gpu[:decode_batch_size].clone()
+        try:
+            for row, req_id in enumerate(decode_req_ids):
+                self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
+                self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
+                self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
+                self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
+                row_offsets = None if start_offsets is None else [start_offsets[row]]
+                self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
+        finally:
+            self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
+            self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
+            self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
+            self.text_step.gpu[:decode_batch_size].copy_(saved_text)
+
     def _talker_mtp_forward(
         self,
         decode_req_ids: list[str],
@@ -1818,62 +1897,22 @@ class OmniGPUModelRunner(GPUModelRunner):
         if not isinstance(subtalker_params, dict):
             subtalker_params = {}
 
-        def _explicit_talker_seed(req_id: str) -> int | None:
-            sampling_params = getattr(self.requests[req_id], "sampling_params", None)
-            extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
-            seed = None
-            if isinstance(extra_args, dict):
-                seed = extra_args.get("tts_local_seed")
-            return int(seed) if seed is not None else None
-
-        def _row_generator(req_id: str) -> torch.Generator | None:
-            seed = _explicit_talker_seed(req_id)
-            if seed is None:
-                return None
-            cache = getattr(self, "_talker_mtp_generators", None)
-            if cache is None:
-                cache = {}
-                self._talker_mtp_generators = cache
-            generator = cache.get(req_id)
-            if generator is None or generator.device != req_input_ids.device:
-                generator = torch.Generator(device=req_input_ids.device)
-                generator.manual_seed(seed)
-                cache[req_id] = generator
-            return generator
-
-        row_generators = [_row_generator(req_id) for req_id in decode_req_ids]
-        cache = getattr(self, "_talker_mtp_generators", None)
-        if cache:
-            # Generators live as long as their request; drop finished ones.
-            for stale_id in [rid for rid in cache if rid not in self.requests]:
-                del cache[stale_id]
-
-        if (
-            decode_batch_size > 1
-            and any(generator is not None for generator in row_generators)
-            and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False)
-        ):
-            # A torch.Generator is a single stream. Using one generator for a
-            # multi-row batch would make explicitly-seeded requests depend on
-            # other rows in the same scheduler step, so keep that path scalar.
-            saved_input_ids = self.talker_mtp_input_ids.gpu[:decode_batch_size].clone()
-            saved_embeds = self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].clone()
-            saved_hidden = self.last_talker_hidden.gpu[:decode_batch_size].clone()
-            saved_text = self.text_step.gpu[:decode_batch_size].clone()
-            try:
-                for row, req_id in enumerate(decode_req_ids):
-                    self.talker_mtp_input_ids.gpu[:1].copy_(saved_input_ids[row : row + 1])
-                    self.talker_mtp_inputs_embeds.gpu[:1].copy_(saved_embeds[row : row + 1])
-                    self.last_talker_hidden.gpu[:1].copy_(saved_hidden[row : row + 1])
-                    self.text_step.gpu[:1].copy_(saved_text[row : row + 1])
-                    row_offsets = None if start_offsets is None else [start_offsets[row]]
-                    self._talker_mtp_forward([req_id], inputs_embeds, row_offsets)
-            finally:
-                self.talker_mtp_input_ids.gpu[:decode_batch_size].copy_(saved_input_ids)
-                self.talker_mtp_inputs_embeds.gpu[:decode_batch_size].copy_(saved_embeds)
-                self.last_talker_hidden.gpu[:decode_batch_size].copy_(saved_hidden)
-                self.text_step.gpu[:decode_batch_size].copy_(saved_text)
-            return
+        row_generators = self._resolve_talker_mtp_generators(
+            decode_req_ids,
+            req_input_ids.device,
+        )
+        generator = None
+        generators = None
+        if row_generators is not None:
+            if decode_batch_size == 1:
+                generator = row_generators[0]
+            elif getattr(self.model, "talker_mtp_accepts_generators", False):
+                # Unseeded rows intentionally keep the default global RNG semantics,
+                # represented by None, even when batched with explicitly seeded rows.
+                generators = row_generators
+            else:
+                self._talker_mtp_forward_row_by_row(decode_req_ids, inputs_embeds, start_offsets)
+                return
 
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
@@ -1881,11 +1920,10 @@ class OmniGPUModelRunner(GPUModelRunner):
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        if decode_batch_size == 1:
-            if row_generators[0] is not None:
-                talker_kwargs["generator"] = row_generators[0]
-        elif any(generator is not None for generator in row_generators):
-            talker_kwargs["generators"] = row_generators
+        if generator is not None:
+            talker_kwargs["generator"] = generator
+        if generators is not None:
+            talker_kwargs["generators"] = generators
         if getattr(self.model, "talker_mtp_accepts_req_infos", False):
             talker_kwargs["req_ids"] = decode_req_ids
             talker_kwargs["req_infos"] = [