[Bugfix] Fix tokenizer config for higgs-audio-v3 test error (#4918)

Signed-off-by: herotai214 <herotai214@gmail.com>

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -19,6 +19,10 @@ 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)
@@ -55,6 +59,23 @@ 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,
@@ -68,6 +89,7 @@ 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"):
@@ -87,7 +109,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is not 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
@@ -96,6 +118,13 @@ 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 == [
@@ -155,7 +184,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    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"]
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,6 +50,48 @@ 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,6 +22,7 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
+        "vae_use_tiling": True,
     },
 )
 
@@ -32,6 +33,7 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
+        "vae_use_tiling": True,
     },
 )
 
@@ -45,10 +47,6 @@ 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."""
@@ -74,7 +72,6 @@ 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`` sits below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sit 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=2)
+            acks = await llm_engine.sleep(stage_ids=[0], level=1)
             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=2)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             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 2
+            # Sleep Level 1
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=2)
+            await diffusion_engine.sleep(stage_ids=[0], level=1)
             # 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 2 Deep Sleep (Full Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
             reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,6 +407,47 @@ 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
@@ -444,7 +485,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=2)
+        acks = await engine.sleep(level=1)
         statuses = [get_ack_info(ack, "status") for ack in acks]
         assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
 
@@ -500,7 +541,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=2)
+        acks = await engine.sleep(stage_ids=[0, 1], level=1)
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -339,6 +339,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -7,6 +7,7 @@ without requiring the actual checkpoint or GPU.
 """
 
 import asyncio
+import json
 import time
 from collections import OrderedDict, defaultdict
 from types import SimpleNamespace
@@ -1309,7 +1310,10 @@ class TestVoiceCloneReferenceCache:
         parent = tmp_path / "OmniVoice"
         subdir = parent / "audio_tokenizer"
         subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text("{}", encoding="utf-8")
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "higgs_audio_v2_tokenizer"}),
+            encoding="utf-8",
+        )
 
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
@@ -1317,6 +1321,32 @@ 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,6 +480,52 @@ 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

@@ -81,6 +81,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         top_k=None,
         top_p=None,
         generator=None,
+        generators=None,
     ):
         self.calls.append(
             {
@@ -90,6 +91,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "top_k": top_k,
                 "top_p": top_p,
                 "generator": generator,
+                "generators": generators,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -306,6 +308,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "top_k": 9,
             "top_p": 0.55,
             "generator": runner.talker_mtp.calls[0]["generator"],
+            "generators": None,
         }
     ]
     assert runner.talker_mtp.calls[0]["generator"] is not None
@@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     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)."""
+    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=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.requests["r2"].sampling_params = SimpleNamespace(
+        seed=22,
+        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,
+    )
+    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)
+
+    # 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
+    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]
+
+    # The per-request generator stream persists across steps...
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+    assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
+    assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
+
+    # ...and is evicted once its request finishes.
+    del runner.requests["r2"]
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+    assert set(runner._talker_mtp_generators) == {"r1"}
+    assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
+
+
 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,6 +6,7 @@ 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
@@ -1488,5 +1489,12 @@ 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,6 +512,21 @@ 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,
@@ -838,23 +853,40 @@ 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_parallel_state._PCP = _SP
+        # 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",
+        )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    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:
+    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:
         # 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_model_parallel_group(
+        vllm_parallel_state._DP = init_vllm_model_parallel_group(
             group_ranks=rank_generator.get_ranks("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="data",
+            group_name="dp",
         )
 
     global _FS, _EXPERT_PARALLEL_GROUP_RANKS
@@ -869,11 +901,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_model_parallel_group(
+        vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=ep_group_ranks,
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="expert",
+            group_name="ep",
         )
         _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -3,6 +3,8 @@
 
 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
 
@@ -24,6 +26,30 @@ 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.
 
@@ -62,6 +88,7 @@ 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/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2235,13 +2235,8 @@ class HunyuanImage3Model(nn.Module):
             # processed with quantization, LoRA, fine-tuning, etc.
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
+            # KV-cache scales are renamed via maybe_remap_kv_scale_name below;
+            # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810).
 
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:

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 OmniACK, OmniSleepTask, OmniWakeTask
+from vllm_omni.diffusion.data import CuMemTag, 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,7 +141,8 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
-        self._is_sleeping: bool = False
+        self._sleeping_tags: set[str] = set()
+        self._level2_sleeping: bool = False
         self.final_output_task: asyncio.Task | None = None
         self.event_resolver = AsyncEventResolver(orchestrator=self)
         self.config_path = self.engine.config_path
@@ -208,10 +209,13 @@ 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
@@ -221,6 +225,11 @@ 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
@@ -306,6 +315,14 @@ 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", [])
@@ -922,11 +939,32 @@ class AsyncOmni(EngineClient, OmniBase):
                 if ack is not None:
                     await self.event_resolver.resolve(ack)
                     final_acks.append(ack)
-        self._is_sleeping = True
+        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
         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
@@ -940,7 +978,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=tags, task_id=task_id)
+        task = OmniWakeTask(tags=requested_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:
@@ -951,7 +989,14 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_acks.append(ack)
         current_omni_platform.synchronize()
         await asyncio.sleep(0.1)
-        self._is_sleeping = False
+
+        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
         logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
         return final_acks
 
@@ -961,7 +1006,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 self._is_sleeping
+        return bool(getattr(self, "_sleeping_tags", None))
 
     async def add_lora(self, lora_request: LoRARequest) -> bool:
         """Load a new LoRA adapter into all stages.

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -14,7 +14,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import torch
 import torch.nn as nn
@@ -742,6 +742,27 @@ class CodePredictorWrapper(nn.Module):
     #  Forward -- re-prefill + inline sampling
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _multinomial(
+        probs: torch.Tensor,
+        generator: torch.Generator | None,
+        generators: Sequence[torch.Generator | None] | None,
+    ) -> torch.Tensor:
+        """Sample one code per row, optionally with per-row generators.
+
+        Per-row generators keep explicitly-seeded requests deterministic in a
+        multi-row batch: each row consumes draws only from its own generator,
+        so the transformer forward can stay batched (#4883).
+        """
+        if generators is None:
+            return torch.multinomial(probs, num_samples=1, generator=generator)
+        return torch.cat(
+            [
+                torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator)
+                for row, row_generator in enumerate(generators)
+            ]
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -753,9 +774,12 @@ class CodePredictorWrapper(nn.Module):
         top_k: int = 50,
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -846,7 +870,7 @@ class CodePredictorWrapper(nn.Module):
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
                 probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = torch.multinomial(probs, num_samples=1, generator=generator)
+                code = self._multinomial(probs, generator, generators)
             else:
                 # "per_call" mode: temperature-scaled + top-k
                 if use_sampling:
@@ -855,7 +879,7 @@ class CodePredictorWrapper(nn.Module):
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
                     probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = torch.multinomial(probs, num_samples=1, generator=generator)
+                    code = self._multinomial(probs, generator, generators)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -22,6 +22,7 @@ generation.
 
 from __future__ import annotations
 
+import json
 import os
 import re
 from typing import Any
@@ -170,14 +171,25 @@ _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))
-    if os.path.isfile(os.path.join(expanded, "config.json")):
+    config_path = os.path.join(expanded, "config.json")
+    if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path):
         return expanded
     nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR)
-    if os.path.isfile(os.path.join(nested, "config.json")):
+    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):
         return nested
     return None
 
@@ -201,7 +213,9 @@ 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):
-        return os.path.dirname(cached_config)
+        candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config))
+        if candidate is not None:
+            return candidate
 
     from huggingface_hub.constants import HF_HUB_CACHE
 
@@ -240,7 +254,9 @@ def _load_audio_tokenizer():
             _K2_OMNIVOICE_REPO,
             allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
         )
-        audio_tokenizer_dir = os.path.join(repo_path, _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")
     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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
 import numpy as np
@@ -318,6 +318,12 @@ 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"))
 
@@ -1056,6 +1062,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_k: int | None = None,
         top_p: float | None = None,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1095,6 +1102,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_k=top_k,
             top_p=top_p,
             generator=generator,
+            generators=generators,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

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

vllm_omni/utils/mm_outputs.py

@@ -106,7 +106,13 @@ 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
+    element: object,
+    idx: int,
+    start: int,
+    end: int,
+    pass_lists_through: bool = False,
+    seq_len: int | None = None,
+    scheduled_seq_len: int | None = None,
 ):
     """Build an mm payload element corresponding to one request index
     from an element containing 0 or more CPU tensors.
@@ -120,21 +126,39 @@ 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 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.
+        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.
     """
-    # 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:
+    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)
+    ):
         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)
+            sk: to_payload_element(
+                sv,
+                idx,
+                start,
+                end,
+                pass_lists_through=pass_lists_through,
+                seq_len=seq_len,
+                scheduled_seq_len=scheduled_seq_len,
+            )
             for sk, sv in element.items()
         }
     elif isinstance(element, list):

vllm_omni/worker/gpu_ar_model_runner.py

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

vllm_omni/worker/gpu_model_runner.py

@@ -1826,7 +1826,33 @@ class OmniGPUModelRunner(GPUModelRunner):
                 seed = extra_args.get("tts_local_seed")
             return int(seed) if seed is not None else None
 
-        if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
+        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.
@@ -1849,28 +1875,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[:decode_batch_size].copy_(saved_text)
             return
 
-        generator = None
-        if decode_req_ids:
-            first_req_id = decode_req_ids[0]
-            seed = _explicit_talker_seed(first_req_id)
-            if seed is not None:
-                generators = getattr(self, "_talker_mtp_generators", None)
-                if generators is None:
-                    generators = {}
-                    self._talker_mtp_generators = generators
-                generator = generators.get(first_req_id)
-                if generator is None or generator.device != req_input_ids.device:
-                    generator = torch.Generator(device=req_input_ids.device)
-                    generator.manual_seed(int(seed))
-                    generators[first_req_id] = generator
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
             "temperature": subtalker_params.get("temperature"),
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        if generator is not None:
-            talker_kwargs["generator"] = generator
+        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 getattr(self.model, "talker_mtp_accepts_req_infos", False):
             talker_kwargs["req_ids"] = decode_req_ids
             talker_kwargs["req_infos"] = [