rebase: adapt to vLLM v0.29 APIs

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

docker/Dockerfile.ci

@@ -92,6 +92,10 @@ RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
 # plus `from deep_ep_cpp import EventHandle` -> ImportError, which together took
 # out 21 CI jobs on build #2947 that pass on main.
 #
+# Installing torch first means FlashInfer resolves against the final runtime.
+RUN uv pip install --system --force-reinstall \
+    "torch==2.13.0" "torchaudio==2.11.0" "torchvision==0.28.0" \
+    --torch-backend cu130
 
 # No --upgrade here: with torch already at its pinned version, --upgrade is what
 # would drag it forward again and reintroduce the mismatch this ordering fixes.

tests/diffusion/layers/test_rope_broadcast.py

@@ -16,7 +16,11 @@ import types
 import pytest
 import torch
 
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd
+from vllm_omni.diffusion.layers.rope import (
+    RotaryEmbedding,
+    RotaryEmbeddingWan,
+    apply_rotary_emb_mindiesd,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -239,6 +243,25 @@ class TestRotaryEmbeddingNativeShapeRegression:
         assert not torch.isnan(output).any()
 
 
+def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None:
+    """Source/CPU installs need not contain the wheel-vendored CUDA kernel."""
+    x = torch.randn(1, 5, 2, 8)
+
+    rope = RotaryEmbedding(is_neox_style=True)
+    rope.apply_rotary_emb_vllm_flash_attn = None
+    output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4))
+    assert output.shape == x.shape
+
+    wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+    wan_rope.apply_rotary_emb_vllm_flash_attn = None
+    output = wan_rope.forward_cuda(
+        x,
+        torch.randn(1, 5, 1, 4),
+        torch.randn(1, 5, 1, 4),
+    )
+    assert output.shape == x.shape
+
+
 def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
     """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
     captured: dict[str, object] = {}

tests/entrypoints/openai/test_batch_chat_completions.py

@@ -12,7 +12,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponseChoice,
     ChatMessage,
 )
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
 
 from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
 

tests/entrypoints/openai_api/test_audio_format.py

@@ -248,7 +248,7 @@ class TestResolveAudioFormat:
         assert result == "pcm"
 
     def test_invalid_format_returns_error(self, serving_chat):
-        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
         request = self._make_request({"format": "aac", "voice": "alloy"})
         result = serving_chat._resolve_audio_format(request)
@@ -256,7 +256,7 @@ class TestResolveAudioFormat:
         assert "aac" in result.error.message
 
     def test_all_supported_formats_accepted(self, serving_chat):
-        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
         for fmt in SUPPORTED_CHAT_AUDIO_FORMATS:
             request = self._make_request({"format": fmt, "voice": "alloy"})

tests/entrypoints/openai_api/test_modality_validation.py

@@ -13,7 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 

tests/entrypoints/openai_api/test_serving_chat_metrics.py

@@ -19,7 +19,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def test_omni_chat_completion_response_metrics():
     """Test OmniChatCompletionResponse metrics field works correctly."""
-    from vllm.entrypoints.openai.engine.protocol import UsageInfo
+    from vllm.entrypoints.serve.engine.protocol import UsageInfo
 
     from vllm_omni.entrypoints.openai.protocol.chat_completion import (
         OmniChatCompletionResponse,

tests/entrypoints/openai_api/test_serving_speech.py

@@ -25,7 +25,7 @@ from fastapi.responses import JSONResponse
 from fastapi.testclient import TestClient
 from pydantic import ValidationError
 from pytest_mock import MockerFixture
-from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
 
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler

tests/entrypoints/openai_api/test_video_stream_handler.py

@@ -281,7 +281,7 @@ class TestHandleQueryErrors:
     @pytest.mark.asyncio
     async def test_query_returns_error_response(self):
         """create_chat_completion returns ErrorResponse → error forwarded."""
-        from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+        from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse
 
         err = ErrorResponse(
             error=ErrorInfo(

tests/entrypoints/test_stream_finish_reason.py

@@ -411,7 +411,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
 @pytest.mark.asyncio
 async def test_audio_choice_error_response_is_not_iterated_as_choices():
     """Any ErrorResponse from the audio path is skipped, not iterated."""
-    from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+    from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 
     serving_chat = build_serving_chat()
     request = make_request(modalities=["text", "audio"])

tests/helpers/serving_chat.py

@@ -8,11 +8,11 @@ OmniRequestOutput objects, plus SSE stream collection/parsing utilities.
 import json
 from unittest.mock import MagicMock
 
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponseStreamChoice,
 )
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.outputs import CompletionOutput, RequestOutput
 

tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for MiniCPM-o 4.5 >30s audio preprocessing.
 
 Covers:
@@ -99,10 +99,10 @@ def _make_processor(
     processor.info = SimpleNamespace(audio_pattern="(<audio>./</audio>)")
     processor.data_parser = MiniCPMOMultiModalDataParser(target_sr=_SAMPLE_RATE)
 
-    def fake_base_call_hf_processor(prompts, mm_data, mm_kwargs, tok_kwargs, *, out_keys):
+    def fake_call_hf_processor_on_prompts(prompts, mm_data, mm_kwargs, *, out_keys):
         return {key: fake_hf_outputs[key] for key in out_keys}
 
-    processor._base_call_hf_processor = fake_base_call_hf_processor
+    processor._call_hf_processor_on_prompts = fake_call_hf_processor_on_prompts
     return processor
 
 
@@ -120,7 +120,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(45)]}, {})
 
         features = result["audio_features"]
         assert [feat.shape[-1] for feat in features] == chunk_lens
@@ -138,7 +138,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {})
 
         widths = [feat.shape[-1] for feat in result["audio_features"]]
         assert widths == [_CHUNK_FRAMES, 1200, _CHUNK_FRAMES, 800]
@@ -151,7 +151,7 @@ class TestProcessAudiosUnpadding:
         }
         processor = _make_processor(fake_outputs)
 
-        result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}, {})
+        result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {})
 
         features = result["audio_features"]
         assert len(features) == 1

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -266,8 +266,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     inputs = MagicMock()
     inputs.mm_data_items = mm_data_items
     inputs.hf_processor_mm_kwargs = {"use_audio_in_video": True}
-    inputs.tokenization_kwargs = {}
-    inputs.prompt = [1, 2, 3]
     inputs.get_mm_hashes = MagicMock(return_value={"video": ["v0"], "audio": ["a0"]})
 
     timing_ctx = MagicMock()
@@ -307,7 +305,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     )
     fake_self._paired_cache_keys = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys
     fake_self.info.parse_mm_data = MagicMock(return_value=SimpleNamespace(name="missing"))
-    fake_self._apply_hf_processor_main = MagicMock(return_value=([9], {"video_grid_thw": object()}, False))
+    fake_self._apply_hf_processor_main = MagicMock(return_value={"video_grid_thw": object()})
     fake_self._get_mm_fields_config = MagicMock(return_value={})
     fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]})
     fake_self._merge_mm_kwargs = capture_merge
@@ -316,13 +314,12 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
         return_value={"video": [object()], "audio": [object()]},
     ):
-        prompt_ids, mm_info, _ = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor(
+        mm_info = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor(
             fake_self,
             inputs,
             timing_ctx,
         )
 
-    assert prompt_ids == [9]
     # Both modalities were reprocessed as one unit after pair-miss expansion.
     assert fake_self.info.parse_mm_data.call_args.args[0] == {
         "video": ["video-0"],

tests/model_executor/models/test_weight_loader.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import WeightsMapper
+
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
+
+
+class _TinyModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.kept = nn.Linear(1, 1, bias=False)
+        self.skipped = nn.Linear(1, 1, bias=False)
+        self.rotary_embed = nn.Module()
+        self.rotary_embed.register_parameter("inv_freq", nn.Parameter(torch.zeros(1)))
+
+
+def test_auto_weights_loader_preserves_removed_skip_arguments() -> None:
+    model = _TinyModel()
+    model.skipped.weight.data.zero_()
+
+    loaded = AutoWeightsLoader(
+        model,
+        skip_prefixes=["skipped."],
+        skip_substrs=["rotary_embed.inv_freq"],
+    ).load_weights(
+        [
+            ("checkpoint.kept.weight", torch.ones(1, 1)),
+            ("skipped.weight", torch.full((1, 1), 2.0)),
+            ("rotary_embed.inv_freq", torch.ones(1)),
+        ],
+        mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""}),
+    )
+
+    assert loaded == {"kept.weight"}
+    torch.testing.assert_close(model.kept.weight, torch.ones(1, 1))
+    torch.testing.assert_close(model.skipped.weight, torch.zeros(1, 1))
+    torch.testing.assert_close(model.rotary_embed.inv_freq, torch.zeros(1))

tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by
 ``vllm_omni.patch``.
 
@@ -198,7 +198,7 @@ def test_reload_does_not_nest_wrapper():
         print("NESTED" if nested else "FLAT")
         """
     )
-    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120)
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
     if result.stdout.strip().startswith("SKIP:"):
         pytest.skip(result.stdout.strip())
     assert "FLAT" in result.stdout, (
@@ -265,9 +265,12 @@ def test_env_var_escape_hatch_disables_install(env_value, acceptable):
     )
     env = os.environ.copy()
     env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value
-    # 120s rather than 60s: cold CI containers can take 30-50s just to
-    # import vllm_omni → vllm → torch before the subprocess does any work.
-    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env)
+    # Generous timeout: cold CI containers can take 30-50s just to import
+    # vllm_omni → vllm → torch, and heavily loaded shared machines have
+    # been observed taking >2min for the same import under disk I/O
+    # contention. The subprocess does a single import + flag print, so 600s
+    # only bounds pathological hangs, not normal runtime.
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600, env=env)
     if result.stdout.strip().startswith("SKIP:"):
         pytest.skip("modelopt not available in subprocess")
     # Exact token match on the final stdout line — not a substring test, since

vllm_omni/benchmarks/patch/patch.py

@@ -2975,7 +2975,9 @@ async def benchmark(
             "duration": benchmark_duration,
             "completed": metrics.completed,
             "total_input_tokens": metrics.total_input,
+            "total_input_sequences": metrics.total_input_sequences,
             "request_throughput": metrics.request_throughput,
+            "input_sequence_throughput": metrics.input_sequence_throughput,
             "total_token_throughput": metrics.total_token_throughput,
             "input_lens": [output.prompt_len for output in outputs],
             "errors": [output.error for output in outputs],

vllm_omni/diffusion/layers/rope.py

@@ -1,3 +1,4 @@
+from importlib import import_module
 from importlib.util import find_spec
 
 import torch
@@ -98,6 +99,7 @@ class RotaryEmbedding(CustomOp):
         self.interleaved = not is_neox_style
         self.half_head_dim = half_head_dim
         self.apply_rotary_emb_flash_attn = None
+        self.apply_rotary_emb_vllm_flash_attn = None
         self.has_mindie = False
         # ``find_spec("flash_attn")`` is True as long as *any* package publishes
         # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships
@@ -111,6 +113,17 @@ class RotaryEmbedding(CustomOp):
                 self.apply_rotary_emb_flash_attn = apply_rotary
             except ImportError:
                 pass
+        # CUDA wheels vendor this kernel under vllm_flash_attn, while source
+        # and CPU-only installations may not contain the generated package.
+        # Resolve it lazily and retain the native implementation as a portable
+        # fallback for those environments.
+        if current_omni_platform.is_cuda():
+            try:
+                self.apply_rotary_emb_vllm_flash_attn = import_module(
+                    "vllm.vllm_flash_attn.layers.rotary"
+                ).apply_rotary_emb
+            except ImportError:
+                pass
         if find_spec("mindiesd") is not None:
             self.has_mindie = True
 
@@ -135,12 +148,18 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
-
         cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin)
 
+        if self.apply_rotary_emb_vllm_flash_attn is None:
+            return apply_rotary_emb_torch(
+                x,
+                cos,
+                sin,
+                interleaved=self.interleaved,
+            )
+
         x, squeezed = _ensure_batch_dim(x)
-        output = apply_rotary_emb(
+        output = self.apply_rotary_emb_vllm_flash_attn(
             x,
             cos,
             sin,
@@ -231,13 +250,14 @@ class RotaryEmbeddingWan(RotaryEmbedding):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+        if self.apply_rotary_emb_vllm_flash_attn is None:
+            return self.forward_native(x, cos, sin)
 
         if cos.dim() > 2:
             cos = cos.reshape(-1, cos.shape[-1])
             sin = sin.reshape(-1, sin.shape[-1])
 
-        return apply_rotary_emb(
+        return self.apply_rotary_emb_vllm_flash_attn(
             x,
             cos,
             sin,

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -17,7 +17,7 @@ from transformers.generation.configuration_utils import GenerationConfig
 from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin
 from transformers.utils.generic import ModelOutput
 from vllm.config.vllm import get_current_vllm_config
-from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+from vllm.model_executor.models.utils import WeightsMapper
 from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -33,6 +33,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
 from . import request_layout as request_layout_utils
 from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, TokenizerWrapper

vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py

@@ -24,7 +24,7 @@ from vllm.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from vllm.logger import init_logger
-from vllm.model_executor.layers.batch_invariant import linear_batch_invariant
+from vllm.model_executor.determinism.batch_invariant import linear_batch_invariant
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,

vllm_omni/entrypoints/async_omni.py

@@ -51,9 +51,9 @@ from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetric
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
-    from vllm.inputs.preprocess import InputPreprocessor
     from vllm.tokenizers import TokenizerLike
     from vllm.v1.engine import PauseMode
+    from vllm.v1.engine.input_processor import InputProcessor
 
     from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType
 
@@ -1714,7 +1714,7 @@ class AsyncOmni(EngineClient, OmniBase):
 
     # ==================== EngineClient Interface ====================
 
-    async def get_input_preprocessor(self) -> InputPreprocessor:
+    async def get_input_preprocessor(self) -> InputProcessor:
         """Get input preprocessor."""
         return self.input_processor
 
@@ -1755,10 +1755,12 @@ class AsyncOmni(EngineClient, OmniBase):
         """
         logger.debug("Weight update start requested (no-op in omni)")
 
-    async def finish_weight_update(self) -> None:
+    async def finish_weight_update(self, weight_version: str | None = None) -> None:
         """Finish the current weight update.
 
         Omni does not currently support weight transfer, so this is a no-op.
+        ``weight_version`` is accepted for upstream ``EngineClient`` protocol
+        compatibility (RLHF weight-transfer routers pass it positionally).
         """
         logger.debug("Weight update finish requested (no-op in omni)")
 

vllm_omni/entrypoints/cli/serve.py

@@ -18,7 +18,7 @@ from typing import Any, cast
 
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
-from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
+from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
 

vllm_omni/entrypoints/duplex/chat_fallback.py

@@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.duplex.protocol import DuplexSession

vllm_omni/entrypoints/openai/api_server.py

@@ -34,7 +34,10 @@ from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
-from vllm.entrypoints.launcher import serve_http, terminate_if_errored
+from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
+from vllm.entrypoints.launchers.cli_args import make_arg_parser
+from vllm.entrypoints.launchers.launcher import serve_http, terminate_if_errored
+from vllm.entrypoints.launchers.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
@@ -43,19 +46,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponse,
 )
-from vllm.entrypoints.openai.cli_args import make_arg_parser
 
 # yapf conflicts with isort for this block
 # yapf: disable
 # yapf: enable
 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
-from vllm.entrypoints.openai.engine.protocol import (
-    ErrorResponse,
-    ModelCard,
-    ModelList,
-    ModelPermission,
-    RequestResponseMetadata,
-)
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
@@ -64,14 +59,13 @@ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAISer
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
 from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
-
-# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
-# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
-# serve.exception_handling and re-exports it from the package root.
-try:
-    from vllm.entrypoints.serve import create_error_response
-except ImportError:
-    from vllm.entrypoints.serve.utils.error_response import create_error_response
+from vllm.entrypoints.serve import create_error_response
+from vllm.entrypoints.serve.engine.protocol import (
+    ErrorResponse,
+    ModelCard,
+    ModelList,
+    ModelPermission,
+)
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
@@ -85,7 +79,6 @@ from vllm.entrypoints.serve.utils.api_utils import (
 )
 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
-from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.speech_to_text.transcription.serving import (
     OpenAIServingTranscription,

vllm_omni/entrypoints/openai/batch_serving.py

@@ -11,7 +11,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponse,
     ChatCompletionResponseChoice,
 )
-from vllm.entrypoints.openai.engine.protocol import (
+from vllm.entrypoints.serve.engine.protocol import (
     ErrorResponse,
     UsageInfo,
 )

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -8,7 +8,7 @@ from typing import cast
 from uuid import uuid4
 
 import numpy as np
-from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.entrypoints.serve.engine.protocol import UsageInfo
 from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection
 from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone
 from vllm.logger import init_logger

vllm_omni/entrypoints/openai/serving_chat.py

@@ -54,8 +54,17 @@ except ImportError:
     soundfile = None
 
 
+from vllm.entrypoints.generate.base.protocol import (
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    FunctionCall,
+    FunctionDefinition,
+    RequestResponseMetadata,
+    ToolCall,
+)
 from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
-from vllm.entrypoints.launcher import terminate_if_errored
+from vllm.entrypoints.launchers.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
@@ -69,22 +78,15 @@ from vllm.entrypoints.openai.chat_completion.serving import (
     _get_mm_token_counts,
     _make_prompt_tokens_details,
 )
-from vllm.entrypoints.openai.engine.protocol import (
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ErrorInfo,
-    ErrorResponse,
-    FunctionCall,
-    FunctionDefinition,
-    RequestResponseMetadata,
-    ToolCall,
-    UsageInfo,
-)
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.entrypoints.serve.engine.protocol import (
+    ErrorInfo,
+    ErrorResponse,
+    UsageInfo,
+)
 from vllm.entrypoints.serve.engine.typing import ChatLikeRequest
 from vllm.entrypoints.serve.utils.api_utils import should_include_usage
 from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls

vllm_omni/entrypoints/openai/serving_speech.py

@@ -24,12 +24,11 @@ import soundfile as sf
 import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
+from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
 from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
-from vllm.entrypoints.launcher import terminate_if_errored
-from vllm.entrypoints.openai.engine.protocol import (
-    ErrorResponse,
-    RequestResponseMetadata,
-)
+from vllm.entrypoints.launchers.launcher import terminate_if_errored
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
+from vllm.inputs import tokens_input
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
 from vllm.utils import random_uuid

vllm_omni/entrypoints/openai/video_stream_session.py

@@ -20,7 +20,7 @@ from fastapi import WebSocket, WebSocketDisconnect
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter

vllm_omni/inputs/preprocess.py

@@ -1,9 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from typing import Any
 
 from typing_extensions import assert_never
-from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput
-from vllm.inputs.preprocess import InputPreprocessor
+from vllm.inputs import (
+    EmbedsInput,
+    MultiModalDataDict,
+    MultiModalInput,
+    MultiModalUUIDDict,
+    SingletonInput,
+    TextPrompt,
+    TokensInput,
+    TokensPrompt,
+    tokens_input,
+)
 from vllm.logger import init_logger
+from vllm.renderers import BaseRenderer, renderer_from_config
 from vllm.renderers.inputs import SingletonDictPrompt
 
 from vllm_omni.inputs.data import (
@@ -17,7 +30,175 @@ from vllm_omni.inputs.data import (
 logger = init_logger(__name__)
 
 
-class OmniInputPreprocessor(InputPreprocessor):
+class _InputPreprocessor:
+    """Compatibility shim for the upstream-removed ``InputPreprocessor``.
+
+    Upstream vLLM moved raw-prompt preprocessing into
+    :class:`vllm.renderers.BaseRenderer` and removed
+    ``vllm.inputs.preprocess.InputPreprocessor``. vLLM-Omni still exposes an
+    omni-specific preprocessor object for stage-0 input processing, so keep the
+    same thin renderer-delegating surface locally and adapt the few calls whose
+    renderer signature changed upstream.
+    """
+
+    def __init__(
+        self,
+        vllm_config: Any,
+        renderer: BaseRenderer | None = None,
+    ) -> None:
+        self.model_config = vllm_config.model_config
+        self.renderer = renderer or renderer_from_config(vllm_config)
+
+    @property
+    def tokenizer(self) -> Any | None:
+        return self.renderer.tokenizer
+
+    def _tokenize_prompt(
+        self,
+        prompt: str,
+        tokenization_kwargs: dict[str, Any] | None = None,
+    ) -> list[int]:
+        """Apply the model's tokenizer to a text prompt."""
+        tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {}))
+
+        tok_prompt = self.renderer._tokenize_singleton_prompt(
+            TextPrompt(prompt=prompt),
+            tok_params,
+        )
+
+        return tok_prompt["prompt_token_ids"]
+
+    def _process_multimodal(
+        self,
+        prompt: str | list[int],
+        mm_data: MultiModalDataDict,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        *,
+        mm_uuids: MultiModalUUIDDict | None = None,
+    ) -> MultiModalInput:
+        """Apply the model's multi-modal processor to a multi-modal prompt.
+
+        The upstream renderer no longer accepts raw ``str`` prompts or
+        ``tokenization_kwargs`` in ``_process_multimodal``; tokenize text
+        prompts here first to preserve the legacy ``InputPreprocessor`` API.
+        """
+        if isinstance(prompt, str):
+            prompt = self._tokenize_prompt(prompt, tokenization_kwargs)
+
+        return self.renderer._process_multimodal(
+            prompt,
+            mm_data,
+            mm_uuids=mm_uuids,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+
+    def _process_embeds(
+        self,
+        parsed_content: OmniEmbedsPrompt,
+    ) -> EmbedsInput:
+        return self.renderer._process_embeds(parsed_content)
+
+    def _truncate_inputs(
+        self,
+        inputs: list[int],
+        tokenization_kwargs: dict[str, Any] | None = None,
+    ) -> list[int]:
+        tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {}))
+
+        tok_prompt = self.renderer._tokenize_singleton_prompt(
+            TokensPrompt(prompt_token_ids=inputs),
+            tok_params,
+        )
+
+        return tok_prompt["prompt_token_ids"]
+
+    def _process_tokens(
+        self,
+        parsed_content: OmniTokensPrompt,
+        tokenization_kwargs: dict[str, Any] | None = None,
+    ) -> TokensInput | MultiModalInput:
+        prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
+
+        inputs: TokensInput | MultiModalInput
+        if multi_modal_data := parsed_content.get("multi_modal_data"):
+            inputs = self._process_multimodal(
+                prompt_token_ids,
+                multi_modal_data,
+                parsed_content.get("mm_processor_kwargs"),
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=parsed_content.get("multi_modal_uuids"),
+            )
+        else:
+            inputs = tokens_input(prompt_token_ids)
+
+        if prompt_text := parsed_content.get("prompt"):
+            inputs["prompt"] = prompt_text
+        if cache_salt := parsed_content.get("cache_salt"):
+            inputs["cache_salt"] = cache_salt
+
+        return inputs
+
+    def _process_text(
+        self,
+        parsed_content: OmniTextPrompt,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        *,
+        mm_uuids: MultiModalUUIDDict | None = None,
+    ) -> TokensInput | MultiModalInput:
+        prompt_text = parsed_content["prompt"]
+
+        inputs: TokensInput | MultiModalInput
+        if multi_modal_data := parsed_content.get("multi_modal_data"):
+            inputs = self._process_multimodal(
+                prompt_text,
+                multi_modal_data,
+                parsed_content.get("mm_processor_kwargs") or {},
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
+        else:
+            prompt_token_ids = self._tokenize_prompt(
+                prompt_text,
+                tokenization_kwargs=tokenization_kwargs,
+            )
+            inputs = tokens_input(prompt_token_ids)
+
+        inputs["prompt"] = prompt_text
+
+        if cache_salt := parsed_content.get("cache_salt"):
+            inputs["cache_salt"] = cache_salt
+
+        return inputs
+
+    def _prompt_to_llm_inputs(
+        self,
+        prompt: SingletonDictPrompt,
+        tokenization_kwargs: dict[str, Any] | None = None,
+        *,
+        mm_uuids: MultiModalUUIDDict | None = None,
+    ) -> SingletonInput:
+        """Extract the singleton inputs from a prompt."""
+        if "prompt_embeds" in prompt:
+            return self._process_embeds(prompt)  # type: ignore[arg-type]
+
+        if "prompt_token_ids" in prompt:
+            return self._process_tokens(
+                prompt,  # type: ignore[arg-type]
+                tokenization_kwargs=tokenization_kwargs,
+            )
+
+        if "prompt" in prompt:
+            return self._process_text(
+                prompt,  # type: ignore[arg-type]
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
+
+        assert_never(prompt)  # type: ignore[arg-type]
+
+
+class OmniInputPreprocessor(_InputPreprocessor):
     """Input preprocessor for omni models.
 
     Extends the base InputPreprocessor to handle omni-specific input

vllm_omni/model_executor/models/audex/audex_omni_processing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
 #   inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_features.py
@@ -222,6 +222,22 @@ class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]):
             tensor_type="pt",
         )
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()),
+            mm_data,
+            hf_processor_mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,

vllm_omni/model_executor/models/bagel/bagel.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from math import isqrt
 from typing import Any
 
@@ -20,7 +21,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.bagel import BagelForConditionalGeneration
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
 from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP
-from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
@@ -53,6 +53,7 @@ from vllm_omni.diffusion.models.bagel.bagel_transformer import (
     TimestepEmbedder,
 )
 from vllm_omni.diffusion.models.bagel.pipeline_bagel import default_ae_params
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
 
 class OmniBagelProcessor(BagelProcessor):
@@ -221,6 +222,20 @@ class OmniBagelDataParser(MultiModalDataParser):
 class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]):
     IMG2IMG_PLACEHOLDER = "<|fim_middle|>"
 
+    def apply(self, inputs, timing_ctx):
+        num_img2img = inputs.mm_data_items.get_all_counts().get("img2img", 0)
+        if num_img2img:
+            token_id = self.info.get_tokenizer().get_vocab().get(self.IMG2IMG_PLACEHOLDER)
+            if token_id is not None:
+                prompt_ids = list(inputs.prompt)
+                missing = num_img2img - prompt_ids.count(token_id)
+                if missing > 0:
+                    inputs = replace(
+                        inputs,
+                        prompt=[token_id] * missing + prompt_ids,
+                    )
+        return super().apply(inputs, timing_ctx)
+
     @staticmethod
     def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]:
         # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing.
@@ -296,6 +311,52 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
 
         return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        has_image = "images" in mm_data
+        has_img2img = "pixel_values_img2img" in mm_data
+        processed_data = BatchFeature()
+
+        if has_image:
+            image_kwargs = {**hf_processor_mm_kwargs, "is_img2img": False}
+            image_outputs = self.info.ctx.call_hf_processor(
+                processor,
+                {"text": prompt_text, "images": mm_data["images"]},
+                image_kwargs,
+            )
+            processed_data.update(image_outputs)
+
+        if has_img2img:
+            img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs)
+            img2img_kwargs["is_img2img"] = True
+            img2img_outputs = self.info.ctx.call_hf_processor(
+                processor,
+                {
+                    "text": prompt_text,
+                    "images": mm_data["pixel_values_img2img"],
+                },
+                img2img_kwargs,
+            )
+            pixel_values = img2img_outputs.pop("pixel_values", None)
+            if pixel_values is not None:
+                processed_data["pixel_values_img2img"] = pixel_values
+            for key, value in img2img_outputs.items():
+                processed_data.setdefault(key, value)
+
+        if not has_image and not has_img2img:
+            processed_data = BatchFeature(dict(passthrough_data))
+        else:
+            processed_data.update(passthrough_data)
+        return processed_data
+
     def _hf_processor_applies_updates(
         self,
         prompt_text: str,

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import partial
 from math import gcd
 from threading import Lock
@@ -109,6 +110,46 @@ class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo):
 
 
 class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]):
+    _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text"
+
+    def apply(self, inputs: ProcessorInputs, timing_ctx):
+        tokenizer = self.info.get_tokenizer()
+        prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False)
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        self._ensure_cached_runtime_components(model_dir, config)
+
+        text_token, text_token_len = extract_text_token(
+            prompt_text,
+            self.tokenizer,
+            config.allowed_special,
+        )
+        if inputs.mm_data_items.get_all_counts().get("audio", 0):
+            reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text")
+            if not isinstance(reference_text, str):
+                raise ValueError(f"prompt text is None : {reference_text}")
+            prompt_text_token, prompt_text_token_len = extract_text_token(
+                reference_text,
+                self.tokenizer,
+                config.allowed_special,
+            )
+            text_token, _ = concat_text_with_prompt_ids(
+                text_token,
+                text_token_len,
+                prompt_text_token,
+                prompt_text_token_len,
+            )
+
+        inputs = replace(
+            inputs,
+            prompt=text_token.reshape(-1).tolist(),
+            hf_processor_mm_kwargs={
+                **inputs.hf_processor_mm_kwargs,
+                self._OMNI_PROMPT_TEXT_KEY: prompt_text,
+            },
+        )
+        return super().apply(inputs, timing_ctx)
+
     def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None:
         cached_model_dir = getattr(self, "_cached_model_dir", None)
         if cached_model_dir == model_dir:
@@ -364,6 +405,24 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
 
         return ft
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        mm_kwargs = dict(hf_processor_mm_kwargs)
+        prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, ""))
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            prompt_text,
+            mm_data,
+            mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: "BatchFeature",

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 Tencent.
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
@@ -152,6 +155,22 @@ class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingIn
             tensor_type=None,
         )
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()),
+            mm_data,
+            hf_processor_mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,

vllm_omni/model_executor/models/gepard/gepard_talker.py

@@ -22,12 +22,13 @@ import torch.nn as nn
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM
-from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
+from vllm.model_executor.models.utils import maybe_prefix
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
 from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 # Adapted from
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm_image/modeling_glm_image.py
@@ -24,6 +24,7 @@
 import math
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import lru_cache
 from typing import Annotated, Literal
 
@@ -358,6 +359,24 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
             return self._apply_hf_processor(inputs, timing_ctx)
         return super()._cached_apply_hf_processor(inputs, timing_ctx)
 
+    def apply(self, inputs, timing_ctx):
+        """Ensure token prompts contain one source-image placeholder each."""
+        num_images = inputs.mm_data_items.get_all_counts().get("image", 0)
+        if num_images:
+            image_token_id = getattr(
+                self.info.get_hf_config(),
+                "image_token_id",
+                167855,
+            )
+            prompt_ids = list(inputs.prompt)
+            missing = num_images - prompt_ids.count(image_token_id)
+            if missing > 0:
+                inputs = replace(
+                    inputs,
+                    prompt=[image_token_id] * missing + prompt_ids,
+                )
+        return super().apply(inputs, timing_ctx)
+
     def _call_hf_processor(
         self,
         prompt: str,
@@ -827,6 +846,42 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
         # HF processor already expanded image placeholders in input_ids.
         return prompt_ids, mm_processed_data, True
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        """Process source images under vLLM's token-only prompt contract."""
+        mm_counts = mm_items.get_all_counts()
+        num_images = mm_counts.get("image", 0)
+        if num_images == 0:
+            return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_processor_mm_kwargs)})
+
+        valid_mm_items = mm_items.select({key for key, count in mm_counts.items() if count > 0})
+        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        images = processor_data.get("images")
+        if not images:
+            return BatchFeature(dict(passthrough_data))
+
+        processor = self.info.get_hf_processor()
+        image_inputs = processor.image_processor(
+            images=list(images),
+            return_tensors="pt",
+        )
+        image_grid_thw = image_inputs.get("image_grid_thw")
+        if image_grid_thw is not None:
+            source_grid_thw = image_grid_thw[:num_images]
+            image_inputs["image_grid_thw"] = source_grid_thw
+
+            target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype)
+            image_inputs["mrope_image_grid_thw"] = torch.cat(
+                [source_grid_thw, target_grid],
+                dim=0,
+            )
+
+        image_inputs.update(passthrough_data)
+        return image_inputs
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -865,6 +920,13 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
                 # to only include source image grids, so batching will work correctly
                 result["image_grid_thw"] = MultiModalFieldConfig.batched("image")
 
+                if "mrope_image_grid_thw" in hf_inputs:
+                    result["mrope_image_grid_thw"] = MultiModalFieldConfig.shared(
+                        "image",
+                        num_source_images,
+                        keep_on_cpu=True,
+                    )
+
         logger.debug(f"_get_mm_fields_config: result keys: {list(result.keys())}")
 
         return result

vllm_omni/model_executor/models/glm_tts/glm_tts.py

@@ -373,6 +373,55 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi
     handoff.
     """
 
+    _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text"
+
+    def apply(self, inputs: ProcessorInputs, timing_ctx):
+        source_tokenizer = self.info.get_tokenizer()
+        prompt_text = source_tokenizer.decode(
+            inputs.prompt,
+            skip_special_tokens=False,
+        )
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        self._ensure_cached_runtime_components(
+            model_dir,
+            config,
+            load_voice_clone=False,
+        )
+
+        normalized_text = _normalize_glm_tts_processor_text(
+            self.text_frontend,
+            prompt_text,
+        )
+        text_ids = self._encode_text(normalized_text)
+        prompt_parts = []
+        if inputs.mm_data_items.get_all_counts().get("audio", 0):
+            reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text")
+            if not isinstance(reference_text, str) or not reference_text.strip():
+                raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].")
+            normalized_reference = _normalize_glm_tts_processor_text(
+                self.text_frontend,
+                reference_text,
+                add_trailing_space=True,
+            )
+            prompt_parts.append(self._encode_text(normalized_reference))
+        prompt_parts.extend(
+            [
+                text_ids,
+                torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long),
+            ]
+        )
+        prompt_ids = torch.cat(prompt_parts, dim=1).reshape(-1).tolist()
+        inputs = replace(
+            inputs,
+            prompt=prompt_ids,
+            hf_processor_mm_kwargs={
+                **inputs.hf_processor_mm_kwargs,
+                self._OMNI_PROMPT_TEXT_KEY: prompt_text,
+            },
+        )
+        return super().apply(inputs, timing_ctx)
+
     def _ensure_cached_runtime_components(
         self,
         model_dir: str,
@@ -581,6 +630,24 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi
             }
         )
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        mm_kwargs = dict(hf_processor_mm_kwargs)
+        prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, ""))
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            prompt_text,
+            mm_data,
+            mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -1669,7 +1669,7 @@ class _BackboneWrapper(nn.Module):
         self.config = config
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        from vllm.model_executor.models.utils import AutoWeightsLoader
+        from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
         skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None
         loader = AutoWeightsLoader(self, skip_prefixes=skip)

vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py

@@ -0,0 +1,830 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Vendored from vllm/model_executor/models/hunyuan_v1.py at vllm@2cf0a6915c
+# (the upstream base vllm-omni origin/main was rebased against). Upstream
+# removed the native HunYuan V1 implementation in #53272 / #53615 (migrated
+# to the Transformers modeling backend). vllm-omni's HunyuanImage3 model
+# subclasses the removed HunYuanModel / HunYuanSparseMoeBlock / HunYuanMLP,
+# so those internals are preserved here to keep the omni diffusion model
+# working against current upstream.
+"""Inference-only HunYuan V1 internals vendored for vllm-omni HunyuanImage3."""
+
+import typing
+from collections.abc import Callable, Iterable
+from itertools import islice
+
+import regex as re
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_ep_group,
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEFactory,
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_layers,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backend import AttentionType
+
+
+def _is_moe(config: PretrainedConfig) -> bool:
+    num_experts = getattr(config, "num_experts", None)
+    if isinstance(num_experts, int):
+        return num_experts > 1
+    if isinstance(num_experts, list) and num_experts:
+        # Ensure all elements are integers before calling max.
+        if all(isinstance(e, int) for e in num_experts):
+            return max(num_experts) > 1
+        else:
+            return False
+    return False
+
+
+def _get_cla_factor(config: PretrainedConfig) -> int:
+    if not getattr(config, "use_cla", False):
+        return 1
+    return getattr(config, "cla_share_factor", 1)
+
+
+class HunYuanMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+        reduce_results: bool = True,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+            reduce_results=reduce_results,
+        )
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class HunYuanAttention(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        if hasattr(config, "head_dim") and config.head_dim:
+            self.head_dim = config.head_dim
+        elif hasattr(config, "attention_head_dim"):
+            self.head_dim = config.attention_head_dim
+        else:
+            self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.use_qk_norm = getattr(config, "use_qk_norm", False)
+        self.layer_id = layer_id
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=True,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+
+        if self.use_qk_norm:
+            self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        ori_k = k
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                q.view(-1, self.num_heads, self.head_dim),
+            )
+            k = self.key_layernorm(
+                k.view(-1, self.num_kv_heads, self.head_dim),
+            )
+
+        attn_output = self.attn(q, k, v)
+        # For o_proj
+        attn_output = attn_output.view(q.shape[0], -1)
+        output, _ = self.o_proj(attn_output)
+        return output, (ori_k, v)
+
+
+class HunYuanCrossAttention(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position_embeddings: int = 8192,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        if hasattr(config, "head_dim"):
+            self.head_dim = config.head_dim
+        elif hasattr(config, "attention_head_dim"):
+            self.head_dim = config.attention_head_dim
+        else:
+            self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.use_qk_norm = getattr(config, "use_qk_norm", False)
+        self.layer_id = layer_id
+
+        self.q_proj = ColumnParallelLinear(
+            hidden_size,
+            hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=True,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=AttentionType.ENCODER_DECODER,
+        )
+
+        if self.use_qk_norm:
+            self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        assert kv_states is not None
+        ori_k, v = kv_states  # use last layer kv,
+        k = ori_k
+        q, _ = self.q_proj(hidden_states)
+        k_tmp = torch.empty_like(k)  # Todo: reduant rotary embedding
+        q, _ = self.rotary_emb(positions, q, k_tmp)
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                q.view(-1, self.num_heads, self.head_dim),
+            )
+            k = self.key_layernorm(
+                k.view(-1, self.num_kv_heads, self.head_dim),
+            )
+
+        attn_output = self.attn(q, k, v)
+        # For o_proj
+        attn_output = attn_output.view(q.shape[0], -1)
+        output, _ = self.o_proj(attn_output)
+        return output, (ori_k, v)
+
+
+class HunYuanSparseMoeBlock(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        layer_id: int = -1,
+        prefix: str = "",
+        enable_eplb: bool = False,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        self.ep_group = get_ep_group().device_group
+        self.ep_size = self.ep_group.size()
+        self.n_routed_experts = config.num_experts
+
+        if self.tp_size > config.num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}."
+            )
+
+        # Get layer_id topk if config.moe_topk is a list
+        if isinstance(config.moe_topk, list):
+            assert layer_id >= 0
+            assert len(config.moe_topk) > layer_id
+            top_k = config.moe_topk[layer_id]
+        else:
+            top_k = config.moe_topk
+
+        # If it is moe, moe_intermediate_size is preferred
+        intermediate_size = config.intermediate_size
+        if config.moe_intermediate_size is not None:
+            intermediate_size = (
+                config.moe_intermediate_size
+                if isinstance(config.moe_intermediate_size, int)
+                else config.moe_intermediate_size[layer_id]
+            )
+
+        # Load balancing settings.
+        vllm_config = get_current_vllm_config()
+        eplb_config = vllm_config.parallel_config.eplb_config
+        self.enable_eplb = enable_eplb
+
+        self.n_logical_experts = self.n_routed_experts
+        self.n_redundant_experts = eplb_config.num_redundant_experts
+        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+        self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+        self.gate = ReplicatedLinear(
+            config.hidden_size,
+            config.num_experts,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.gate",
+        )
+        if config.use_mixed_mlp_moe > 0:
+            # Get layer_id num_shared_expert if config.num_shared_expert is
+            # a list.
+            if isinstance(config.num_shared_expert, list):
+                assert layer_id >= 0
+                assert len(config.num_shared_expert) > layer_id
+                num_shared_expert = config.num_shared_expert[layer_id]
+            else:
+                num_shared_expert = config.num_shared_expert
+
+            self.shared_mlp = HunYuanMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size * num_shared_expert,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_mlp",
+            )
+        else:
+            self.shared_mlp = None
+
+        self.experts = FusedMoEFactory(
+            shared_experts=self.shared_mlp,
+            num_experts=self.n_routed_experts,
+            top_k=top_k,
+            hidden_size=config.hidden_size,
+            intermediate_size=intermediate_size,
+            renormalize=top_k > 1,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts",
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # NOTE: hidden_states can have either 1D or 2D shape.
+        orig_shape = hidden_states.shape
+        hidden_dim = hidden_states.shape[-1]
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits)
+
+        return final_hidden_states.view(orig_shape)
+
+
+class HunYuanDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        layer_id: int = -1,
+        enable_eplb: bool = False,
+    ) -> None:
+        super().__init__()
+        assert layer_id >= 0
+        self.layer_id = layer_id
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = (
+            config.intermediate_size
+            if isinstance(config.intermediate_size, int)
+            else config.intermediate_size[layer_id]
+        )
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
+        cla_factor = _get_cla_factor(config)
+        attention_type = (
+            AttentionType.ENCODER_DECODER if layer_id >= 0 and layer_id % cla_factor != 0 else AttentionType.DECODER
+        )
+        if attention_type == AttentionType.DECODER:
+            self.self_attn = HunYuanAttention(
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
+                max_position_embeddings=max_position_embeddings,
+                quant_config=quant_config,
+                bias=attention_bias,
+                cache_config=cache_config,
+                prefix=f"{prefix}.self_attn",
+                layer_id=layer_id,
+            )
+        elif attention_type == AttentionType.ENCODER_DECODER:
+            self.self_attn = HunYuanCrossAttention(
+                config=config,
+                hidden_size=self.hidden_size,
+                num_heads=config.num_attention_heads,
+                num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
+                max_position_embeddings=max_position_embeddings,
+                quant_config=quant_config,
+                bias=attention_bias,
+                cache_config=cache_config,
+                prefix=f"{prefix}.self_attn",
+                layer_id=layer_id,
+            )
+        else:
+            raise RuntimeError(f"Unsupported attention type: {attention_type}")
+
+        if _is_moe(config):
+            self.mlp = HunYuanSparseMoeBlock(
+                config=config,
+                quant_config=quant_config,
+                layer_id=layer_id,
+                prefix=f"{prefix}.mlp",
+                enable_eplb=enable_eplb,
+            )
+        else:
+            self.mlp = HunYuanMLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                bias=getattr(config, "mlp_bias", False),
+                prefix=f"{prefix}.mlp",
+            )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        kv_states: tuple[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states, ori_kv_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_states=kv_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual, ori_kv_states
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (xd, seq_len) if xdrope is enabled for hunyuan-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class HunYuanModel(nn.Module, EagleModelMixin):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        eplb_config = vllm_config.parallel_config.eplb_config
+        enable_eplb = vllm_config.parallel_config.enable_eplb
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.vocab_size = config.vocab_size
+
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: HunYuanDecoderLayer(
+                config=config,
+                layer_id=int(prefix.split(".")[-1]),
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+                enable_eplb=enable_eplb,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        cla_factor = _get_cla_factor(self.config)
+        prev_kv_states = None
+        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        for i, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)):
+            hidden_states, residual, kv_states = layer(
+                positions,
+                hidden_states,
+                residual,
+                prev_kv_states,
+            )
+
+            if getattr(self.config, "use_cla", False) and i % cla_factor == 0:
+                prev_kv_states = kv_states
+            else:
+                prev_kv_states = None
+
+            self._maybe_add_hidden_state(aux_hidden_states, i + 1, hidden_states, residual)
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
+        return hidden_states
+
+    def _split_qkv_weight(self, qkv: torch.Tensor):
+        num_attention_heads = self.config.num_attention_heads
+        num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads)
+        num_key_value_groups = num_attention_heads // num_kv_heads
+        hidden_size = self.config.hidden_size
+
+        if hasattr(self.config, "head_dim"):
+            attention_head_dim = self.config.head_dim
+        elif hasattr(self.config, "attention_head_dim"):
+            attention_head_dim = self.config.attention_head_dim
+        else:
+            attention_head_dim = self.config.hidden_size // num_attention_heads
+
+        qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size)
+        q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1)
+        q = q.reshape(-1, hidden_size)
+        k = k.reshape(-1, hidden_size)
+        v = v.reshape(-1, hidden_size)
+        return torch.concat((q, k, v))
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        if _is_moe(self.config):
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            return fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.num_experts,
+                num_redundant_experts=self.num_redundant_experts,
+            )
+        else:
+            return []
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        cla_factor = _get_cla_factor(self.config)
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+
+        num_attention_heads = self.config.num_attention_heads
+        num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads)
+        split_params_mapping = [
+            (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),
+            (
+                ".qkv_proj",
+                ".qkv_proj",
+                num_attention_heads + num_kv_heads * 2,
+                [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)],
+                self._split_qkv_weight,
+            ),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "gate_proj_bias" in name:
+                name = name.replace("gate_proj_bias", "gate_proj.bias")
+            if "up_proj_bias" in name:
+                name = name.replace("up_proj_bias", "up_proj.bias")
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            # With tie_word_embeddings, we can skip lm_head.weight
+            # The weight might appear unnecessarily in the files if the model is
+            # processed with quantization, LoRA, fine-tuning, etc.
+            if self.config.tie_word_embeddings and "lm_head.weight" in name:
+                continue
+
+            is_found = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                if "mlp.experts" in name:
+                    continue
+                # cross layer only have q_proj, skip qkv pack
+                if weight_name == ".q_proj":
+                    match = re.search(r"layers\.\d+", name)
+                    if match:
+                        layer_id = int(match.group(0).split(".")[-1])
+                        if cla_factor > 1 and layer_id % cla_factor != 0:
+                            continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(name)
+                is_found = True
+                break
+            if is_found:
+                continue
+
+            for (
+                param_name,
+                weight_name,
+                den,
+                split_param,
+                func,
+            ) in split_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                assert loaded_weight.shape[0] % den == 0
+                units = loaded_weight.shape[0] // den
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                offset = 0
+                for shard_id, num in split_param:
+                    new_offset = offset + num * units
+                    if func:
+                        weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id)
+                    else:
+                        weight_loader(param, loaded_weight[offset:new_offset], shard_id)
+                    offset = new_offset
+
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                is_expert_weight = False
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    # this is an expert weight and should not be
+                    # attempted to load as other weights later
+                    is_expert_weight = True
+
+                    # Do not modify `name` since the loop may continue here
+                    # Instead, create a new variable
+                    name_mapped = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(name_mapped, self):
+                        continue
+                    param = params_dict[name_mapped]
+                    # We should ask the weight loader to return success or not
+                    # here since otherwise we may skip experts with other
+                    # available replicas.
+                    weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+                    success = weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                        return_success=True,
+                    )
+                    if success:
+                        name = name_mapped
+                        break
+                else:
+                    if is_expert_weight:
+                        # We've checked that this is an expert weight
+                        # However it's not mapped locally to this rank
+                        # So we simply skip it
+                        continue
+                    # Remapping the name of FP8 kv-scale.
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
+
+                    if is_pp_missing_parameter(name, self):
+                        continue
+
+                    if "mlp.gate.wg." in name:
+                        name = name.replace("wg.", "")
+
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import gc
 import math
 import typing
@@ -45,13 +45,6 @@ from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
-from vllm.model_executor.models.hunyuan_v1 import (
-    HunYuanMLP,
-    HunYuanModel,
-    HunYuanSparseMoeBlock,
-    _get_cla_factor,
-    _is_moe,
-)
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsLoRA,
@@ -61,7 +54,6 @@ from vllm.model_executor.models.interfaces import (
     _require_is_multimodal,
 )
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     PPMissingLayer,
     _merge_multimodal_embeddings,
     is_pp_missing_parameter,
@@ -91,8 +83,16 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import (
+    HunYuanMLP,
+    HunYuanModel,
+    HunYuanSparseMoeBlock,
+    _get_cla_factor,
+    _is_moe,
+)
 from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
 logger = init_logger(__name__)
 
@@ -1046,6 +1046,22 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
             batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long)
         return batch_feature
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()),
+            mm_data,
+            hf_processor_mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _hf_processor_applies_updates(
         self,
         prompt_text: str,

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2025 Xiaomi Corporation.
 import os
 from collections.abc import Iterable, Mapping, Sequence
@@ -34,7 +37,7 @@ from vllm.multimodal.processing import (
     PromptUpdate,
     PromptUpdateDetails,
 )
-from vllm.multimodal.utils import group_mm_kwargs_by_modality
+from vllm.multimodal.utils import group_and_batch_mm_kwargs
 from vllm.sequence import IntermediateTensors
 from vllm.utils.cache import LRUCache
 from vllm.utils.collection_utils import is_list_of
@@ -395,6 +398,22 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces
             tensor_type=None,
         )
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()),
+            dict(mm_data),
+            hf_processor_mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -484,7 +503,7 @@ class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProces
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_mimo_audio,
             )
         ]
@@ -626,7 +645,7 @@ class MiMoAudioForConditionalGeneration(
                 )
             mm_kwargs.append((mm_feature.modality, mm_item))
 
-        for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
+        for modality, num_items, mm_kwargs_group in group_and_batch_mm_kwargs(
             mm_kwargs,
             device=self.device,
             pin_memory=self.pin_memory,

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -420,6 +420,26 @@ class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingIn
             tok_kwargs=tok_kwargs,
         )
 
+    def _preprocess_hf_mm_data(
+        self,
+        mm_data: Mapping[str, object],
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data(
+            mm_data,
+            hf_processor_mm_kwargs,
+        )
+        mm_data = dict(mm_data)
+        audios = mm_data.pop("audios", [])
+        if audios:
+            mm_data["audio"] = audios
+            feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
+            hf_processor_mm_kwargs = {
+                **hf_processor_mm_kwargs,
+                "sampling_rate": feature_extractor.sampling_rate,
+            }
+        return mm_data, hf_processor_mm_kwargs
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,
@@ -480,7 +500,7 @@ class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingIn
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_mimo_audio,
             )
         ]

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -21,7 +21,6 @@ from transformers import AutoTokenizer, Qwen2Config, Qwen2Model
 from transformers.utils.hub import cached_file
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
@@ -30,6 +29,7 @@ from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
 from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig
 from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig
 
 from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 # Copyright 2024 ANT Group and the HuggingFace Inc. team.
 # Adapted from Ming repository modeling_bailingmm2.py and processing_bailingmm2.py
@@ -8,6 +9,7 @@
 
 import os
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import replace
 from typing import Annotated, Any
 
 import numpy as np
@@ -388,7 +390,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
             updates.append(
                 PromptReplacement(
                     modality="image",
-                    target=PLACEHOLDER_IMAGE_TOKEN_IN_TEXT,
+                    target=tokenizer.encode(
+                        PLACEHOLDER_IMAGE_TOKEN_IN_TEXT,
+                        add_special_tokens=False,
+                    ),
                     replacement=get_replacement_image,
                 )
             )
@@ -396,7 +401,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
             updates.append(
                 PromptReplacement(
                     modality="video",
-                    target=PLACEHOLDER_VIDEO_TOKEN_IN_TEXT,
+                    target=tokenizer.encode(
+                        PLACEHOLDER_VIDEO_TOKEN_IN_TEXT,
+                        add_special_tokens=False,
+                    ),
                     replacement=get_replacement_video,
                 )
             )
@@ -404,7 +412,10 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
             updates.append(
                 PromptReplacement(
                     modality="audio",
-                    target=PLACEHOLDER_AUDIO_TOKEN_IN_TEXT,
+                    target=tokenizer.encode(
+                        PLACEHOLDER_AUDIO_TOKEN_IN_TEXT,
+                        add_special_tokens=False,
+                    ),
                     replacement=get_replacement_audio,
                 )
             )
@@ -451,90 +462,71 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
 
         return config
 
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
+    def apply(self, inputs, timing_ctx):
+        """Add Ming image-generation query tokens before prompt updates."""
+        modalities = inputs.hf_processor_mm_kwargs.get("modalities") or []
+        if "image" in modalities or "img2img" in modalities:
+            from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
+                DEFAULT_NUM_QUERY_TOKENS,
+                IMAGE_PATCH_TOKEN,
+            )
+
+            tokenizer = self.info.get_tokenizer()
+            prompt_ids = list(inputs.prompt)
+            if "img2img" in modalities:
+                prompt_ids = (
+                    tokenizer.encode(
+                        PLACEHOLDER_IMAGE_TOKEN_IN_TEXT,
+                        add_special_tokens=False,
+                    )
+                    + prompt_ids
+                )
+
+            image_patch_ids = tokenizer.encode(
+                IMAGE_PATCH_TOKEN,
+                add_special_tokens=False,
+            )
+            if not image_patch_ids or not any(
+                prompt_ids[index : index + len(image_patch_ids)] == image_patch_ids
+                for index in range(len(prompt_ids) - len(image_patch_ids) + 1)
+            ):
+                image_gen_config = getattr(
+                    self.info.ctx.model_config.hf_config,
+                    "image_gen_config",
+                    None,
+                )
+                num_query_tokens = int(
+                    getattr(
+                        image_gen_config,
+                        "num_query_tokens",
+                        DEFAULT_NUM_QUERY_TOKENS,
+                    )
+                )
+                suffix = (
+                    tokenizer.encode("<image>", add_special_tokens=False)
+                    + image_patch_ids * num_query_tokens
+                    + tokenizer.encode("</image>", add_special_tokens=False)
+                )
+                prompt_ids.extend(suffix)
+
+            inputs = replace(inputs, prompt=prompt_ids)
+
+        return super().apply(inputs, timing_ctx)
 
     def _apply_hf_processor_main(
         self,
-        prompt: str | list[int],
         mm_items: MultiModalDataItems,
         hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-        *,
-        enable_hf_prompt_update: bool,
-    ):
-        """Apply Ming-specific prompt rewrites before delegating to vllm's
-        default text / multimodal processing pipeline.
-
-        This is the single point in the base MM-processor call graph where we
-        still have BOTH the full prompt and the request-level
-        ``hf_processor_mm_kwargs`` together. Downstream, vllm's MM caching
-        design splits the work into
-        :meth:`_apply_hf_processor_text_only` (prompt, empty kwargs) and
-        :meth:`_apply_hf_processor_mm_only` (dummy text, real kwargs), so a
-        per-request flag like ``modalities`` cannot reach the real
-        tokenization from inside :meth:`_call_hf_processor`.
-
-        For Ming-flash-omni-2.0, image-generation requests carry
-        ``mm_processor_kwargs["modalities"]`` set to ``["image"]`` (t2i) or
-        ``["img2img"]`` (image edit). We expand the prompt here with the
-        ``<image><imagePatch>*N</image>`` query-token block so the expanded
-        form flows through both the tokenization and (eventual) MM paths.
-        For img2img, we also prepend ``<IMAGE>`` so the thinker-side prompt
-        replacement can locate the reference-image placeholder.
-        """
-        modalities = hf_processor_mm_kwargs.get("modalities") or []
-        is_image_gen = "image" in modalities or "img2img" in modalities
-        if isinstance(prompt, str) and is_image_gen:
-            from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
-                DEFAULT_NUM_QUERY_TOKENS,
-                maybe_expand_image_gen_prompt,
-            )
-
-            # img2img: prepend ``<IMAGE>`` so the ref-image placeholder survives
-            # MM caching (when the ref image is cache-warm, the thinker sees 0
-            # missing items and cannot inject the placeholder itself).
-            if "img2img" in modalities:
-                prompt = "<IMAGE>" + prompt
-
-            ig = getattr(self.info.ctx.model_config.hf_config, "image_gen_config", None)
-            num_query_tokens = getattr(ig, "num_query_tokens", DEFAULT_NUM_QUERY_TOKENS)
-            prompt = maybe_expand_image_gen_prompt(prompt, num_query_tokens=int(num_query_tokens))
-
-        return super()._apply_hf_processor_main(
-            prompt=prompt,
-            mm_items=mm_items,
-            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-            tokenization_kwargs=tokenization_kwargs,
-            enable_hf_prompt_update=enable_hf_prompt_update,
-        )
-
-    def _call_hf_processor(
-        self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        """Call sub-processors for multimodal inputs and tokenize.
+        """Call Ming's modality processors without expanding prompt text.
 
         We call the image/audio sub-processors directly (instead of going
         through `MingFlashOmniProcessor.__call__`) so that the high-level
-        placeholder tokens remain **unexpanded** in the tokenized output.
-
-        NOTE: Image-gen prompt expansion is handled in
-        :meth:`_apply_hf_processor_main` — by the time we get here, vllm's
-        MM caching has already split text/mm processing and ``mm_kwargs``
-        is empty on the text-only branch.
+        placeholder tokens remain under vLLM's prompt-update control.
         """
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
         hf_processor = self.info.get_hf_processor()
-        tokenizer = self.info.get_tokenizer()
 
         data: dict[str, object] = {}
 
@@ -577,11 +569,9 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
             )
             data.update(audio_outputs)
 
-        # Tokenize text with placeholders still intact
-        text_outputs = tokenizer(prompt, return_tensors="pt", **tok_kwargs)
-        data.update(text_outputs)
-
-        return BatchFeature(data=data)
+        processed_data = BatchFeature(data=data)
+        processed_data.update(passthrough_data)
+        return processed_data
 
 
 @MULTIMODAL_REGISTRY.register_processor(

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -23,7 +23,7 @@ import os
 import warnings
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from functools import partial
 from typing import Annotated, Any, Literal, TypeAlias
 
@@ -3403,84 +3403,46 @@ class MiniCPMOMultiModalDataParser(MultiModalDataParser):
 class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45OmniLLMProcessingInfo]):
     """Multimodal processor for MiniCPM-o thinker stage."""
 
-    def _apply_hf_processor_main(
-        self,
-        prompt: str | list[int],
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-        *,
-        enable_hf_prompt_update: bool,
-    ) -> tuple[list[int], BatchFeature, bool]:
-        """
-        MiniCPM-O reimplements this to avoid calling HF processor with text-only
-        when enable_hf_prompt_update=False. The MiniCPM processor asserts
-        len(image_tags) == len(image_sizes) and fails if given placeholder text
-        without corresponding image data (e.g. during profiling/cache-miss path).
-        """
-        use_tts = hf_processor_mm_kwargs.get("use_tts", False)
-        hf_processor_mm_kwargs = {k: v for k, v in hf_processor_mm_kwargs.items() if k != "use_tts"}
-        if isinstance(prompt, str):
-            if use_tts:
-                prompt = prompt + self._TTS_SUFFIX
-            if enable_hf_prompt_update:
-                return self._apply_hf_processor_text_mm(
-                    prompt_text=prompt,
-                    mm_items=mm_items,
-                    hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                    tokenization_kwargs=tokenization_kwargs,
-                )
-            tokenizer = self.info.get_tokenizer()
-            prompt_ids = _encode_tokens(tokenizer, prompt)
-        else:
-            prompt_ids = self._apply_hf_processor_tokens_only(prompt)
-            if use_tts:
-                tokenizer = self.info.get_tokenizer()
-                tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"])
-                prompt_ids = list(prompt_ids) + tts_ids
-
-        mm_processed_data = self._apply_hf_processor_mm_only(
-            mm_items=mm_items,
-            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-            tokenization_kwargs=tokenization_kwargs,
-        )
-
-        return prompt_ids, mm_processed_data, False
-
     _TTS_SUFFIX = "<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>"
 
-    def _call_hf_processor(
+    def apply(self, inputs, timing_ctx):
+        hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs
+        if hf_processor_mm_kwargs.get("use_tts", False):
+            tokenizer = self.info.get_tokenizer()
+            tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"])
+            inputs = replace(
+                inputs,
+                prompt=[*inputs.prompt, *tts_ids],
+                hf_processor_mm_kwargs={
+                    key: value for key, value in hf_processor_mm_kwargs.items() if key != "use_tts"
+                },
+            )
+        return super().apply(inputs, timing_ctx)
+
+    def _apply_hf_processor_main(
         self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
     ) -> BatchFeature:
+        """
+        Process each modality independently because the MiniCPM processor
+        asserts that image tags and image sizes have matching lengths.
+        """
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+
         tokenizer = self.info.get_tokenizer()
-
-        use_tts = mm_kwargs.get("use_tts", False)
-        mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"}
-        if use_tts:
-            prompt = prompt + self._TTS_SUFFIX
-
-        input_ids = torch.tensor([tokenizer.encode(prompt, **tok_kwargs)])
-        mm_inputs = self.process_mm_inputs(mm_data, mm_kwargs, tok_kwargs)
-
-        return BatchFeature(
+        prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
+        input_ids = torch.tensor([tokenizer.encode(prompt_text)])
+        mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs)
+        processed_data = BatchFeature(
             {
                 "input_ids": input_ids,
                 **mm_inputs,
             }
         )
-
-    def _hf_processor_applies_updates(
-        self,
-        prompt_text: str,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> bool:
-        return False
+        processed_data.update(passthrough_data)
+        return processed_data
 
     def get_image_prompt_texts(
         self,
@@ -3528,7 +3490,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         self,
         mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
     ) -> Mapping[str, NestedTensors]:
         if (audios := mm_data.get("audios")) is None:
             return {}
@@ -3540,11 +3501,10 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         if isinstance(parsed_audios, MiniCPMOAudioEmbeddingItems):
             audio_inputs = {}
         else:
-            audio_inputs = self._base_call_hf_processor(
+            audio_inputs = self._call_hf_processor_on_prompts(
                 prompts=[self.info.audio_pattern] * len(parsed_audios),
                 mm_data={"audios": [[audio] for audio in parsed_audios]},
                 mm_kwargs={**mm_kwargs, "chunk_input": True},
-                tok_kwargs=tok_kwargs,
                 out_keys={"audio_features", "audio_feature_lens"},
             )
 
@@ -3569,7 +3529,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         self,
         mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
     ) -> Mapping[str, NestedTensors]:
         if (images := mm_data.get("images")) is None:
             return {}
@@ -3581,11 +3540,10 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         if isinstance(parsed_images, MiniCPMVImageEmbeddingItems):
             image_inputs = {}
         else:
-            image_inputs = self._base_call_hf_processor(
+            image_inputs = self._call_hf_processor_on_prompts(
                 prompts=[self.info.image_pattern] * len(parsed_images),
                 mm_data={"images": [[image] for image in parsed_images]},
                 mm_kwargs=mm_kwargs,
-                tok_kwargs=tok_kwargs,
                 out_keys={"pixel_values", "image_sizes", "tgt_sizes"},
             )
         return image_inputs
@@ -3594,7 +3552,6 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         self,
         mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
     ) -> Mapping[str, NestedTensors]:
         if (videos := mm_data.get("videos")) is None:
             return {}
@@ -3606,14 +3563,13 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         if isinstance(parsed_videos, MiniCPMVVideoEmbeddingItems):
             video_inputs = {}
         else:
-            video_inputs = self._base_call_hf_processor(
+            video_inputs = self._call_hf_processor_on_prompts(
                 prompts=[self.info.image_pattern * len(video) for video in parsed_videos],
                 mm_data={"images": list(parsed_videos)},
                 mm_kwargs={
                     **mm_kwargs,
                     "max_slice_nums": self.info.get_video_max_slice_num(),
                 },
-                tok_kwargs=tok_kwargs,
                 out_keys={"pixel_values", "image_sizes", "tgt_sizes"},
             )
 
@@ -3625,41 +3581,37 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
         self,
         mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
     ) -> Mapping[str, NestedTensors]:
         return {
-            **self.process_images(mm_data, mm_kwargs, tok_kwargs),
-            **self.process_videos(mm_data, mm_kwargs, tok_kwargs),
-            **self.process_audios(mm_data, mm_kwargs, tok_kwargs),
+            **self.process_images(mm_data, mm_kwargs),
+            **self.process_videos(mm_data, mm_kwargs),
+            **self.process_audios(mm_data, mm_kwargs),
         }
 
-    def _base_call_hf_processor(
+    def _call_hf_processor_on_prompts(
         self,
         prompts: list[str],
         mm_data: Mapping[str, Sequence[object]],
         mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
         *,
         out_keys: set[str],
     ) -> dict[str, NestedTensors]:
         mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"}
         # This processor supports zipping prompt and mm_data together
         if self.info.get_model_version() in {(2, 6), (4, 0), (4, 5)}:
-            inputs = super()._call_hf_processor(
-                prompt=prompts,  # type: ignore
-                mm_data=mm_data,
-                mm_kwargs=mm_kwargs,
-                tok_kwargs=tok_kwargs,
+            inputs = self.info.ctx.call_hf_processor(
+                self.info.get_hf_processor(**mm_kwargs),
+                dict(text=prompts, **mm_data),
+                mm_kwargs,
             )
         else:
             inputs = defaultdict[str, list[torch.Tensor]](list)
 
             for i, prompt in enumerate(prompts):
-                inputs_one = super()._call_hf_processor(
-                    prompt=prompt,
-                    mm_data={k: v[i] for k, v in mm_data.items()},
-                    mm_kwargs=mm_kwargs,
-                    tok_kwargs=tok_kwargs,
+                inputs_one = self.info.ctx.call_hf_processor(
+                    self.info.get_hf_processor(**mm_kwargs),
+                    dict(text=prompt, **{key: value[i] for key, value in mm_data.items()}),
+                    mm_kwargs,
                 )
 
                 for k, v in inputs_one.items():
@@ -3722,7 +3674,11 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
             "video": get_video_replacement,
         }
         base_updates = [
-            PromptReplacement(modality=modality, target=pattern, replacement=get_replacement[modality])
+            PromptReplacement(
+                modality=modality,
+                target=tokenizer.encode(pattern, add_special_tokens=False),
+                replacement=get_replacement[modality],
+            )
             for modality, pattern in placeholders
         ]
 
@@ -3747,7 +3703,10 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
             *base_updates,
             PromptReplacement(
                 modality="audio",
-                target=audio_placeholder,
+                target=tokenizer.encode(
+                    audio_placeholder,
+                    add_special_tokens=False,
+                ),
                 replacement=get_audio_replacement,
             ),
         ]

vllm_omni/model_executor/models/minimax_h3/text_encoder.py

@@ -12,6 +12,7 @@ import regex as re
 import torch
 import torch.nn as nn
 from vllm.config import VllmConfig
+from vllm.inputs import mm_input
 from vllm.model_executor.models.qwen3_vl import (
     Qwen3VLDummyInputsBuilder,
     Qwen3VLForConditionalGeneration,
@@ -26,10 +27,6 @@ from vllm_omni.model_executor.models.minimax_h3.conditioning import (
     MINIMAX_H3_PRESENTATION_TASK_KEY,
 )
 from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
-    IMAGE_PAD,
-    VIDEO_PAD,
-    VISION_END,
-    VISION_START,
     build_minimax_h3_presentation,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
@@ -96,58 +93,57 @@ class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor):
 
     def _apply_hf_processor_main(
         self,
-        prompt: str | list[int],
         mm_items: Any,
         hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-        *,
-        enable_hf_prompt_update: bool,
     ):
-        task_value = hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY)
-        if task_value is None:
-            return super()._apply_hf_processor_main(
-                prompt,
-                mm_items,
-                hf_processor_mm_kwargs,
-                tokenization_kwargs,
-                enable_hf_prompt_update=enable_hf_prompt_update,
-            )
-        if not isinstance(prompt, str):
-            raise ValueError("MiniMax H3 presentation requires the original prompt text")
-
-        task = str(task_value)
-        condition_labels = self._condition_labels(hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY))
         base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs)
-        image_count = sum(kind == "image" for kind, _ in condition_labels)
-        video_count = sum(kind == "video" for kind, _ in condition_labels)
-        image_slot = f"{VISION_START}{IMAGE_PAD}{VISION_END}"
-        hf_processor = self.info.get_hf_processor(**base_kwargs)
-        video_slot = (
-            VIDEO_PAD if self._expands_only_video_token(hf_processor) else f"{VISION_START}{VIDEO_PAD}{VISION_END}"
-        )
-        processing_prompt = image_slot * image_count + video_slot * video_count
-        if not processing_prompt:
-            processing_prompt = prompt
-
-        _, processed_data, _ = super()._apply_hf_processor_main(
-            processing_prompt,
+        return super()._apply_hf_processor_main(
             mm_items,
             base_kwargs,
-            tokenization_kwargs,
-            enable_hf_prompt_update=enable_hf_prompt_update,
         )
+
+    def apply(self, inputs: Any, timing_ctx: Any):
+        task_value = inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY)
+        if task_value is None:
+            return super().apply(inputs, timing_ctx)
+
+        prompt_ids = self._postprocess_prompt(inputs.prompt)
+        mm_info = self._cached_apply_hf_processor(inputs, timing_ctx)
+        base_kwargs = self._base_processor_kwargs(inputs.hf_processor_mm_kwargs)
+        tokenizer = self.info.get_tokenizer()
+        prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False)
+        condition_labels = self._condition_labels(inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY))
+        out_mm_data = mm_info.kwargs.get_data()
         image_processor = self.info.get_image_processor(**base_kwargs)
         ids, _ = _build_minimax_h3_presentation(
-            self.info.get_tokenizer(),
-            prompt=prompt,
-            task=task,
+            tokenizer,
+            prompt=prompt_text,
+            task=str(task_value),
             condition_labels=condition_labels,
-            image_grid_thw=processed_data.get("image_grid_thw"),
-            video_grid_thw=processed_data.get("video_grid_thw"),
-            video_timestamps=processed_data.get("timestamps"),
+            image_grid_thw=out_mm_data.get("image_grid_thw"),
+            video_grid_thw=out_mm_data.get("video_grid_thw"),
+            video_timestamps=out_mm_data.get("timestamps"),
             merge_size=int(image_processor.merge_size),
         )
-        return ids.tolist(), processed_data, True
+        prompt_ids = ids.tolist()
+
+        mm_item_counts = inputs.mm_data_items.get_all_counts()
+        self._validate_mm_kwargs(mm_info.kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_info.prompt_updates, mm_item_counts)
+        mm_placeholders = self._find_mm_placeholders(
+            prompt_ids,
+            mm_info.prompt_updates,
+        )
+        self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
+        mm_placeholder_ranges = {
+            modality: [item.to_range() for item in placeholders] for modality, placeholders in mm_placeholders.items()
+        }
+        return mm_input(
+            prompt_token_ids=prompt_ids,
+            mm_kwargs=mm_info.kwargs,
+            mm_hashes=mm_info.hashes,
+            mm_placeholders=mm_placeholder_ranges,
+        )
 
     def _get_prompt_updates(
         self,

vllm_omni/model_executor/models/minimax_music3/talker.py

@@ -37,10 +37,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.qwen3 import Qwen3Model
-from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.v1.outputs import SamplerOutput
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
 from .chunking import chunk_windows
 from .constants import (

vllm_omni/model_executor/models/omnivoice/omnivoice.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 
 import numpy as np
 import torch
@@ -70,6 +71,32 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr
     For auto voice: text only
     """
 
+    _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text"
+
+    def apply(self, inputs: ProcessorInputs, timing_ctx):
+        tokenizer = self.info.get_tokenizer()
+        prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False)
+        prompt_inputs = self._call_hf_processor(
+            prompt_text,
+            {},
+            inputs.hf_processor_mm_kwargs,
+            {},
+        )
+        prompt_ids = prompt_inputs["input_ids"]
+        if isinstance(prompt_ids, torch.Tensor):
+            prompt_ids = prompt_ids.reshape(-1).tolist()
+        else:
+            prompt_ids = list(prompt_ids)
+        inputs = replace(
+            inputs,
+            prompt=prompt_ids,
+            hf_processor_mm_kwargs={
+                **inputs.hf_processor_mm_kwargs,
+                self._OMNI_PROMPT_TEXT_KEY: prompt_text,
+            },
+        )
+        return super().apply(inputs, timing_ctx)
+
     def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None:
         cached_model_dir = getattr(self, "_cached_model_dir", None)
         if cached_model_dir == model_dir:
@@ -178,6 +205,24 @@ class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalPr
         )
         return ft
 
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        mm_kwargs = dict(hf_processor_mm_kwargs)
+        prompt_text = str(mm_kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, ""))
+        valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
+        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processed_data = self._call_hf_processor(
+            prompt_text,
+            mm_data,
+            mm_kwargs,
+            {},
+        )
+        processed_data.update(passthrough_data)
+        return processed_data
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py

@@ -8,7 +8,6 @@ from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     WeightsMapper,
     init_vllm_registered_model,
     maybe_prefix,
@@ -18,6 +17,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
+
 
 class Qwen2_5OmniTalkerForConditionalGeneration(
     nn.Module,

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides."""
 
 import hashlib
@@ -9,6 +12,7 @@ import numpy as np
 import PIL.Image as PILImage
 import torch
 from torch import nn
+from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniThinkerConfig,
 )
@@ -59,7 +63,6 @@ from vllm.model_executor.models.qwen2_audio import (
     _get_feat_extract_output_lengths,
 )
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     WeightsMapper,
     init_vllm_registered_model,
     maybe_prefix,
@@ -84,6 +87,7 @@ from vllm.multimodal.processing.processor import (
 from vllm.sequence import IntermediateTensors
 from vllm.utils.collection_utils import is_list_of
 
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
 )
@@ -360,25 +364,40 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 ):
     """Override to fix use_audio_in_video detection when mm cache returns None."""
 
-    def _call_hf_processor(
+    def _apply_hf_processor_main(
         self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
-    ):
-        mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
-        mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(mm_data, mm_kwargs)
-        hf_inputs = super()._call_hf_processor(
-            prompt=prompt,
-            mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
-            tok_kwargs=tok_kwargs,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        """Apply omni-specific HF kwargs before the upstream Qwen2.5-Omni
+        processor runs.
+
+        Upstream's ``_apply_hf_processor_main`` is fully custom and no longer
+        routes through the generic ``_preprocess_hf_mm_data`` /
+        ``_postprocess_hf_mm_data`` hooks, so the omni video pre-sampling and
+        per-video ``use_audio_in_video`` mask handling have to run here.
+        """
+        valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0})
+        processor_data, _ = self._get_hf_mm_data(valid_mm_items)
+
+        hf_processor_mm_kwargs = _presampled_videos_hf_kwargs(
+            processor_data,
+            hf_processor_mm_kwargs,
         )
-        per_video_mask = mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY)
+        hf_processor_mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(
+            processor_data,
+            hf_processor_mm_kwargs,
+        )
+
+        processed_data = super()._apply_hf_processor_main(
+            mm_items,
+            hf_processor_mm_kwargs,
+        )
+
+        per_video_mask = hf_processor_mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY)
         if per_video_mask is not None:
-            hf_inputs["use_audio_in_video"] = torch.tensor(per_video_mask)
-        return hf_inputs
+            processed_data["use_audio_in_video"] = torch.tensor(per_video_mask)
+        return processed_data
 
     def _get_mm_fields_config(
         self,
@@ -530,17 +549,17 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_qwen2_audio,
             ),
             PromptReplacement(
                 modality="image",
-                target=image_token,
+                target=[image_token_id],
                 replacement=partial(get_replacement_qwen2_vision, modality="image"),
             ),
             PromptReplacement(
                 modality="video",
-                target=video_token,
+                target=[video_token_id],
                 replacement=get_replacement_qwen2_video,
             ),
         ]
@@ -620,7 +639,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
+    ) -> MultiModalProcessingInfo:
         """Cache-aware: processes video/audio pairs as units,
         using pair-specific cache keys. Falls back to standard cache if no pairs."""
 
@@ -673,16 +692,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         )
 
         with timing_ctx.record("apply_hf_processor"):
-            (
-                prompt_ids,
-                mm_missing_processed_data,
-                is_update_applied,
-            ) = self._apply_hf_processor_main(
-                prompt=inputs.prompt,
+            mm_missing_processed_data = self._apply_hf_processor_main(
                 mm_items=mm_missing_data_items,
                 hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-                tokenization_kwargs=inputs.tokenization_kwargs,
-                enable_hf_prompt_update=False,
             )
 
         mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs(
@@ -714,7 +726,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             prompt_updates=mm_prompt_updates,
         )
 
-        return prompt_ids, mm_info, is_update_applied
+        return mm_info
 
     def _derive_audio_from_video_placeholders(
         self,
@@ -785,7 +797,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         prompt_ids: list[int],
         mm_kwargs: MultiModalKwargsItems,
         mm_prompt_updates: MultiModalPromptUpdates,
-        is_update_applied: bool,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
@@ -794,37 +805,27 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates)
         use_audio_in_video = any(video_use_audio_in_video)
 
-        if is_update_applied:
-            mm_placeholders = self._find_mm_placeholders(
+        if use_audio_in_video and "audio" in mm_prompt_updates:
+            filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
+            prompt_ids, mm_placeholders = self._apply_prompt_updates(
+                prompt_ids,
+                filtered_updates,
+            )
+            mm_placeholders = self._derive_audio_from_video_placeholders(
+                mm_placeholders,
+                mm_prompt_updates,
+                video_use_audio_in_video,
+            )
+        else:
+            prompt_ids, mm_placeholders = self._apply_prompt_updates(
                 prompt_ids,
                 mm_prompt_updates,
             )
-            self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-            )
-        else:
-            if use_audio_in_video and "audio" in mm_prompt_updates:
-                filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    filtered_updates,
-                )
-                mm_placeholders = self._derive_audio_from_video_placeholders(
-                    mm_placeholders,
-                    mm_prompt_updates,
-                    video_use_audio_in_video,
-                )
-            else:
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    mm_prompt_updates,
-                )
 
-            self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-            )
+        self._validate_mm_placeholders(
+            mm_placeholders,
+            mm_item_counts,
+        )
 
         return prompt_ids, mm_placeholders
 

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py

@@ -18,7 +18,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead,
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     PPMissingLayer,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -31,6 +30,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
+
 logger = init_logger(__name__)
 
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -23,7 +23,6 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
 from vllm.config import VllmConfig  # type: ignore
 from vllm.logger import init_logger  # type: ignore
 from vllm.model_executor.models.utils import (  # type: ignore
-    AutoWeightsLoader,
     WeightsMapper,
 )
 
@@ -31,6 +30,7 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     Qwen3OmniNestedSupportsQuant,
 )
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 
 logger = init_logger(__name__)
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -12,7 +12,6 @@ from vllm.model_executor.models.interfaces import (
     SupportsPP,
 )
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     WeightsMapper,
     maybe_prefix,
 )
@@ -27,6 +26,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mt
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3MoeLLMForCausalLM,
 )
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 logger = init_logger(__name__)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The Qwen team.
 # Copyright 2023 The vLLM team.
 # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
@@ -31,9 +31,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from packaging.version import Version
 from transformers import PretrainedConfig
-from transformers import __version__ as TRANSFORMERS_VERSION
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeConfig,
@@ -99,7 +97,6 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     _get_feat_extract_output_lengths,
 )
 from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
     WeightsMapper,
     _merge_multimodal_embeddings,
     maybe_prefix,
@@ -124,11 +121,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerMultiModalProcessor,
     _get_request_video_use_audio_in_video,
     _get_video_second_per_grid_t,
-    _presampled_videos_hf_kwargs,
 )
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     Qwen3OmniNestedSupportsQuant,
 )
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
     ComponentQuantizationConfig,
@@ -697,98 +694,82 @@ Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
 class Qwen3OmniMoeThinkerMultiModalProcessor(
     Qwen2_5OmniThinkerMultiModalProcessor,
 ):
-    def _call_hf_processor(
+    def _get_hf_mm_data(
         self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        mm_data = dict(mm_data)
-        mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs))
-        audios = mm_data.pop("audios", [])
+        mm_items: MultiModalDataItems,
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
 
-        def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
-            length = x.shape[-1]
-            if length % hop_length != 0:
-                pad_length = hop_length - (length % hop_length)
-                x = np.pad(x, (0, pad_length), mode="constant", constant_values=0)
-            return x
-
-        # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
-        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
-        hop_length = feature_extractor.hop_length
+        audios = processor_data.get("audios")
         if audios:
-            # NOTE: Qwen3-Omni processor accept "audio"
-            # To make sure the cache works with padding=True, we pre-padded
-            # the audio to multiple of hop_length.
-            mm_data["audio"] = [
-                pad_to_hop_length(audio, hop_length)
-                if isinstance(audio, np.ndarray)
-                else (pad_to_hop_length(audio[0], hop_length), audio[1])
+            feature_extractor = self.info.get_feature_extractor()
+            hop_length = feature_extractor.hop_length
+
+            def pad_to_hop_length(x: np.ndarray) -> np.ndarray:
+                length = x.shape[-1]
+                if length % hop_length != 0:
+                    pad_length = hop_length - (length % hop_length)
+                    x = np.pad(
+                        x,
+                        (0, pad_length),
+                        mode="constant",
+                        constant_values=0,
+                    )
+                return x
+
+            # To make sure the cache works with padding=True, pre-pad audio
+            # to a multiple of the feature extractor's hop length.
+            processor_data = dict(processor_data)
+            processor_data["audios"] = [
+                pad_to_hop_length(audio) if isinstance(audio, np.ndarray) else (pad_to_hop_length(audio[0]), audio[1])
                 for audio in audios
             ]
 
+        return processor_data, passthrough_data
+
+    def _apply_hf_processor_main(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        if mm_items.get_all_counts().get("audio", 0):
             # TODO(Isotr0py): Remove this patch after upstream fix PR
             # released and Transformers version update:
             # https://github.com/huggingface/transformers/pull/41473
-            mm_kwargs = dict(mm_kwargs)
-            tok_kwargs = dict(tok_kwargs)
-            mm_kwargs["audio_kwargs"] = dict(mm_kwargs.get("audio_kwargs") or {})
-            mm_kwargs["text_kwargs"] = dict(mm_kwargs.get("text_kwargs") or {})
-            if Version(TRANSFORMERS_VERSION) < Version("4.58.0"):
-                # Extract audio_sample_rate before restructuring
-                audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None)
+            hf_processor_mm_kwargs = dict(hf_processor_mm_kwargs)
+            hf_processor_mm_kwargs["audio_kwargs"] = dict(hf_processor_mm_kwargs.get("audio_kwargs") or {})
+            hf_processor_mm_kwargs["text_kwargs"] = dict(hf_processor_mm_kwargs.get("text_kwargs") or {})
 
-                # move truncation to audio_kwargs level to avoid conflict
-                # with tok_kwargs
-                mm_kwargs["audio_kwargs"].setdefault("truncation", mm_kwargs.pop("truncation", False))
-                mm_kwargs["text_kwargs"].setdefault("truncation", tok_kwargs.pop("truncation", False))
-
-                # Validate and conditionally pass audio_sample_rate
-                # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's
-                # audio loader already resamples audio to the target rate.
-                # Only pass the value if it matches to avoid unexpected behavior.
-                if audio_sample_rate is not None:
-                    expected_sr = feature_extractor.sampling_rate
-                    if audio_sample_rate != expected_sr:
-                        logger.warning(
-                            "[%s] audio_sample_rate mismatch: user provided %dHz "
-                            "but model expects %dHz. Ignoring user value. "
-                            "vLLM's audio loader already resampled to %dHz.",
-                            self.__class__.__name__,
-                            audio_sample_rate,
-                            expected_sr,
-                            expected_sr,
-                        )
-                    else:
-                        # Sample rate matches, safe to pass
-                        mm_kwargs["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate
-
-        hf_inputs = super()._call_hf_processor(
-            prompt=prompt,
-            mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
-            tok_kwargs=tok_kwargs,
+        processed_data = super()._apply_hf_processor_main(
+            mm_items,
+            hf_processor_mm_kwargs,
         )
 
-        if (
-            "audio_feature_lengths" in hf_inputs
-            and "feature_attention_mask" in hf_inputs
-            and (audios := mm_data.get("audio", []))
-        ):
-            audio_num_frames = []
-            for _, audio in enumerate(audios):
-                audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio)
-                num_frame = (
-                    (audio_length // hop_length) if audio_length % hop_length == 0 else (audio_length // hop_length - 1)
-                )
-                if mm_kwargs.get("truncation", False):
-                    num_frame = min(num_frame, feature_extractor.n_samples // hop_length)
-                audio_num_frames.append(num_frame)
-            hf_inputs["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames]
-            hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames)
-        return hf_inputs
+        if "audio_feature_lengths" in processed_data and "feature_attention_mask" in processed_data:
+            valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0})
+            processor_data, _ = self._get_hf_mm_data(valid_mm_items)
+            audios = processor_data.get("audios", [])
+            if audios:
+                feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
+                hop_length = feature_extractor.hop_length
+
+                audio_num_frames = []
+                for audio in audios:
+                    audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio)
+                    num_frame = (
+                        audio_length // hop_length if audio_length % hop_length == 0 else audio_length // hop_length - 1
+                    )
+                    if hf_processor_mm_kwargs.get("truncation", False):
+                        num_frame = min(
+                            num_frame,
+                            feature_extractor.n_samples // hop_length,
+                        )
+                    audio_num_frames.append(num_frame)
+
+                processed_data["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames]
+                processed_data["audio_feature_lengths"] = torch.tensor(audio_num_frames)
+
+        return processed_data
 
     def _maybe_apply_prompt_updates(
         self,
@@ -796,49 +777,38 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         prompt_ids: list[int],
         mm_kwargs: MultiModalKwargsItems,
         mm_prompt_updates: MultiModalPromptUpdates,
-        is_update_applied: bool,
-    ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
+    ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         """
         Qwen3-Omni reimplements this function to handle `use_audio_in_video`.
         """
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
 
         video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates)
         use_audio_in_video = any(video_use_audio_in_video)
 
-        # normal case with `use_audio_in_video=False`
-        if is_update_applied:
-            mm_placeholders = self._find_mm_placeholders(
+        if use_audio_in_video and "audio" in mm_prompt_updates:
+            filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
+            prompt_ids, mm_placeholders = self._apply_prompt_updates(
+                prompt_ids,
+                filtered_updates,
+            )
+            mm_placeholders = self._derive_audio_from_video_placeholders(
+                mm_placeholders,
+                mm_prompt_updates,
+                video_use_audio_in_video,
+            )
+        else:
+            prompt_ids, mm_placeholders = self._apply_prompt_updates(
                 prompt_ids,
                 mm_prompt_updates,
             )
-            self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-            )
-        else:
-            if use_audio_in_video and "audio" in mm_prompt_updates:
-                filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    filtered_updates,
-                )
-                mm_placeholders = self._derive_audio_from_video_placeholders(
-                    mm_placeholders,
-                    mm_prompt_updates,
-                    video_use_audio_in_video,
-                )
-            else:
-                prompt_ids, mm_placeholders = self._apply_prompt_updates(
-                    prompt_ids,
-                    mm_prompt_updates,
-                )
 
-            self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-            )
+        self._validate_mm_placeholders(
+            mm_placeholders,
+            mm_item_counts,
+        )
 
         return prompt_ids, mm_placeholders
 
@@ -978,17 +948,17 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=[audio_token_id],
                 replacement=get_replacement_qwen2_audio,
             ),
             PromptReplacement(
                 modality="image",
-                target=image_token,
+                target=[image_token_id],
                 replacement=partial(get_replacement_qwen2_vision, modality="image"),
             ),
             PromptReplacement(
                 modality="video",
-                target=video_token,
+                target=[video_token_id],
                 replacement=get_replacement_qwen2_video,
             ),
         ]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -11,9 +11,9 @@ from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader
-from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
 
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -17,12 +17,13 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
 from vllm.model_executor.models.qwen3 import Qwen3Model
-from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper, maybe_prefix
+from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper, maybe_prefix
 from vllm.multimodal.audio import AudioResampler
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.weight_loader import AutoWeightsLoader
 from vllm_omni.utils.speaker_cache import (
     get_speaker_cache,
     iter_custom_voice_profiles,

vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding
 
@@ -19,6 +22,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.utils.rnn import pad_sequence
 from transformers import AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
 from transformers.processing_utils import ProcessorMixin
 from vllm.config import VllmConfig
 from vllm.inputs import MultiModalDataDict
@@ -503,6 +507,10 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
 
         audio_token = getattr(processor, "audio_token", "<audio_patch>")
+        audio_token_ids = self.info.get_tokenizer().encode(
+            audio_token,
+            add_special_tokens=False,
+        )
 
         audio_items = mm_items.get("audio", [])
         num_audio_items = len(audio_items) if audio_items else 0
@@ -548,7 +556,7 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing
         return [
             PromptReplacement(
                 modality="audio",
-                target=audio_token,
+                target=audio_token_ids,
                 replacement=get_replacement_audio,
             )
         ]
@@ -591,6 +599,43 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing
 
         return hf_inputs
 
+    def _preprocess_hf_mm_data(
+        self,
+        mm_data: Mapping[str, object],
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data(
+            mm_data,
+            hf_processor_mm_kwargs,
+        )
+        mm_data = dict(mm_data)
+        audios = mm_data.pop("audios", [])
+        if audios:
+            if isinstance(audios, str):
+                mm_data["audio"] = [audios]
+            elif isinstance(audios, (list, tuple)):
+                mm_data["audio"] = audios
+            else:
+                mm_data["audio"] = [audios]
+        return mm_data, hf_processor_mm_kwargs
+
+    def _postprocess_hf_mm_data(
+        self,
+        mm_data: Mapping[str, object],
+        hf_processor_mm_kwargs: Mapping[str, object],
+        processed_data: BatchFeature,
+    ) -> BatchFeature:
+        processed_data = super()._postprocess_hf_mm_data(
+            mm_data,
+            hf_processor_mm_kwargs,
+            processed_data,
+        )
+        if "audio_mels" not in processed_data:
+            processed_data["audio_mels"] = torch.empty((0, 128, 0))
+        if "audio_lens" not in processed_data:
+            processed_data["audio_lens"] = torch.tensor([], dtype=torch.int32)
+        return processed_data
+
 
 # ============================================================================
 # Main Thinker Model

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import math
 from collections.abc import Iterable, Mapping, Sequence
@@ -871,7 +871,7 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing
         return [
             PromptReplacement(
                 modality="audio",
-                target="",  # Never match the prompt (see below note)
+                target=[audio_id],
                 replacement=get_replacement,
             ),
         ]
@@ -903,11 +903,27 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
-        prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(inputs, timing_ctx)
+    ) -> MultiModalProcessingInfo:
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
 
-        # NOTE: The tokens are already inserted by the chat template
-        return prompt_ids, mm_info, True
+    def _maybe_apply_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        prompt_ids: list[int],
+        mm_kwargs: MultiModalKwargsItems,
+        mm_prompt_updates,
+    ):
+        # Voxtral's Mistral chat template has already inserted the complete
+        # audio-token run, so locate it without applying the replacement again.
+        mm_item_counts = mm_items.get_all_counts()
+        self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
+        self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
+        mm_placeholders = self._find_mm_placeholders(
+            prompt_ids,
+            mm_prompt_updates,
+        )
+        self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
+        return prompt_ids, mm_placeholders
 
 
 @MULTIMODAL_REGISTRY.register_processor(

vllm_omni/model_executor/models/weight_loader.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Compatibility helpers for vLLM model weight loading."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader as _VllmAutoWeightsLoader,
+)
+from vllm.model_executor.models.utils import (
+    WeightsMapper,
+)
+
+
+class AutoWeightsLoader(_VllmAutoWeightsLoader):
+    """Preserve vLLM's pre-0.29 weight-filtering arguments.
+
+    vLLM removed ``skip_prefixes`` and ``skip_substrs`` from
+    :class:`AutoWeightsLoader` in favor of mapper-based filtering. Omni models
+    still use those arguments to exclude checkpoint components owned by a
+    different stage. Convert the filters to a ``WeightsMapper`` so the old
+    behavior is retained without bypassing the current loader implementation.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module,
+        *,
+        skip_prefixes: list[str] | None = None,
+        skip_substrs: list[str] | None = None,
+        ignore_unexpected_prefixes: list[str] | None = None,
+        ignore_unexpected_suffixes: list[str] | None = None,
+    ) -> None:
+        super().__init__(
+            module,
+            ignore_unexpected_prefixes=ignore_unexpected_prefixes,
+            ignore_unexpected_suffixes=ignore_unexpected_suffixes,
+        )
+        self._skip_mapper = WeightsMapper(
+            orig_to_new_prefix={prefix: None for prefix in skip_prefixes or ()},
+            orig_to_new_substr={substring: None for substring in skip_substrs or ()},
+        )
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+        *,
+        mapper: WeightsMapper | None = None,
+    ) -> set[str]:
+        mapper = (mapper or WeightsMapper()) | self._skip_mapper
+        return super().load_weights(weights, mapper=mapper)

vllm_omni/reasoning/step_audio_reasoning_parser.py

@@ -5,7 +5,7 @@ from collections.abc import Iterable, Sequence
 from itertools import islice
 from typing import TYPE_CHECKING
 
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
 from vllm.logger import init_logger
 from vllm.reasoning import ReasoningParser
 from vllm.tokenizers import TokenizerLike

vllm_omni/worker/gpu_ar_model_runner.py

@@ -2089,7 +2089,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 elif self.valid_sampled_token_count_event is not None:
                     assert spec_decode_common_attn_metadata is not None
                     next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded(
-                        self.optimistic_seq_lens_cpu,
                         sampled_token_ids,
                         self.requests,
                         self.input_batch,

vllm_omni/worker/sampling_utils.py

@@ -48,7 +48,7 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in
         if not min_toks:
             continue
         needs_rebuild = False
-        for _, _, stop_tok_ids in min_toks.values():
+        for _, _, stop_tok_ids, _ in min_toks.values():
             oob = [tok for tok in stop_tok_ids if tok >= logits_vocab]
             if not oob:
                 continue
@@ -63,10 +63,19 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in
         if needs_rebuild:
             reqs: list[int] = []
             tok_ids: list[int] = []
-            for index, (_, _, stop_tok_ids) in min_toks.items():
+            restore_reqs: list[int] = []
+            restore_tok_ids: list[int] = []
+            for index, (_, _, stop_tok_ids, uses_structured_output) in min_toks.items():
                 reqs.extend([index] * len(stop_tok_ids))
                 tok_ids.extend(stop_tok_ids)
+                if uses_structured_output:
+                    restore_reqs.extend([index] * len(stop_tok_ids))
+                    restore_tok_ids.extend(stop_tok_ids)
             proc.logits_slice = (
                 proc._device_tensor(reqs, torch.int32),
                 proc._device_tensor(tok_ids, torch.int32),
             )
+            proc.restore_logits_slice = (
+                proc._device_tensor(restore_reqs, torch.int32),
+                proc._device_tensor(restore_tok_ids, torch.int32),
+            )