[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/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

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/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/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/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))