diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci
index f7c4eaba5..e81055036 100644
--- a/docker/Dockerfile.ci
+++ b/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.
diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py
index 284049347..64bbba94f 100644
--- a/tests/diffusion/layers/test_rope_broadcast.py
+++ b/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] = {}
diff --git a/tests/entrypoints/openai/test_batch_chat_completions.py b/tests/entrypoints/openai/test_batch_chat_completions.py
index 9fc72b2c1..18518b097 100644
--- a/tests/entrypoints/openai/test_batch_chat_completions.py
+++ b/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
diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py
index a9bf76319..9c7eb9d89 100644
--- a/tests/entrypoints/openai_api/test_audio_format.py
+++ b/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"})
diff --git a/tests/entrypoints/openai_api/test_modality_validation.py b/tests/entrypoints/openai_api/test_modality_validation.py
index fde0c0ca1..298f3aaca 100644
--- a/tests/entrypoints/openai_api/test_modality_validation.py
+++ b/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]
diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py
index 2fdf9f646..e493b43bc 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py
+++ b/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,
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index 6a5ebcf7d..3922ab4a7 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/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
diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py
index e382b5e32..428b47567 100644
--- a/tests/entrypoints/openai_api/test_video_stream_handler.py
+++ b/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(
diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py
index 97ed2370b..a5fcba69f 100644
--- a/tests/entrypoints/test_stream_finish_reason.py
+++ b/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"])
diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py
index f1f2710b5..414016720 100644
--- a/tests/helpers/serving_chat.py
+++ b/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
diff --git a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py
index 3e8ee2a63..a76f752c5 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py
+++ b/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="()")
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
diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
index d88d918bc..2683a8d9d 100644
--- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py
+++ b/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"],
diff --git a/tests/model_executor/models/test_weight_loader.py b/tests/model_executor/models/test_weight_loader.py
new file mode 100644
index 000000000..549be7993
--- /dev/null
+++ b/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))
diff --git a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py
index 0e9f23de6..985c64114 100644
--- a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py
+++ b/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
diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py
index 10f0853e0..3313e691f 100644
--- a/vllm_omni/benchmarks/patch/patch.py
+++ b/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],
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index 59d2c9a62..b33ab1fde 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/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,
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index 6598525da..a2f9b0699 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/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
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py
index 6b0433bbe..7228fa655 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py
+++ b/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,
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 948f7e53a..0012039e0 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/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)")
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 78b212f98..5c0204a33 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/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
diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py
index 6e4d0a0e0..f41d7205b 100644
--- a/vllm_omni/entrypoints/duplex/chat_fallback.py
+++ b/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
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index d976eb181..e020be7bf 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/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,
diff --git a/vllm_omni/entrypoints/openai/batch_serving.py b/vllm_omni/entrypoints/openai/batch_serving.py
index bc5ef36ca..f26e68d75 100644
--- a/vllm_omni/entrypoints/openai/batch_serving.py
+++ b/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,
)
diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py
index aad7ea658..94196c4bf 100644
--- a/vllm_omni/entrypoints/openai/realtime_connection.py
+++ b/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
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index bfa241a40..6b50b3887 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/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
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index 88d9f090d..674f2d5c6 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/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
diff --git a/vllm_omni/entrypoints/openai/video_stream_session.py b/vllm_omni/entrypoints/openai/video_stream_session.py
index ead3b1ec5..40bf43977 100644
--- a/vllm_omni/entrypoints/openai/video_stream_session.py
+++ b/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
diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py
index 3da1ef897..6d6fe1d8a 100644
--- a/vllm_omni/inputs/preprocess.py
+++ b/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
diff --git a/vllm_omni/model_executor/models/audex/audex_omni_processing.py b/vllm_omni/model_executor/models/audex/audex_omni_processing.py
index 5e4c53908..52d2bec56 100644
--- a/vllm_omni/model_executor/models/audex/audex_omni_processing.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py
index 495e0037f..aa36533df 100644
--- a/vllm_omni/model_executor/models/bagel/bagel.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
index 84652cada..8ec88b070 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
+++ b/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",
diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py
index efbe7b0d6..4560b2a36 100644
--- a/vllm_omni/model_executor/models/covo_audio/covo_audio.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/gepard/gepard_talker.py b/vllm_omni/model_executor/models/gepard/gepard_talker.py
index 02f5da533..913323b8b 100644
--- a/vllm_omni/model_executor/models/gepard/gepard_talker.py
+++ b/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__)
diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py
index 77aef398b..05ef9830a 100644
--- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py
+++ b/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
diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py
index fb3517036..e57489b82 100644
--- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py
index 8fbac95f7..d4d903733 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py
+++ b/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)
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py
new file mode 100644
index 000000000..7065612d7
--- /dev/null
+++ b/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
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
index 58b67d8aa..82907c3ad 100755
--- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py
index 4309d6177..6e4592dec 100644
--- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
index c90adb467..1be09e2f7 100644
--- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
+++ b/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,
)
]
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py
index d7e31beb3..8f247bfa5 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py
+++ b/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
diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
index 03cf87abb..1338f2606 100644
--- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
+++ b/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("", add_special_tokens=False)
+ + image_patch_ids * num_query_tokens
+ + tokenizer.encode("", 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
- ``*N`` query-token block so the expanded
- form flows through both the tokenization and (eventual) MM paths.
- For img2img, we also prepend ```` 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 ```` 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 = "" + 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(
diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py
index 9dfff5589..b4821c603 100644
--- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py
+++ b/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,
),
]
diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py
index 85583f243..b1933af95 100644
--- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py
index c4fb876fa..48d12eb3c 100644
--- a/vllm_omni/model_executor/models/minimax_music3/talker.py
+++ b/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 (
diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py
index 26a6f6b59..2f95fed86 100644
--- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py
index 35d988895..04e875076 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
index ba073aad2..f606f968e 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+++ b/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
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py
index 91eafd2a8..647cf67d2 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py
+++ b/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__)
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
index 5c3129be9..5f54e2dee 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
+++ b/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__)
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
index 8337f8dad..2d3168458 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+++ b/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__)
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
index bbfea498b..96b9d6e30 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+++ b/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,
),
]
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
index 87df87889..9ba43d0d8 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py
+++ b/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 (
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
index e65334ff5..730221478 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
index b00e24870..71fad0b91 100644
--- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py
+++ b/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_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
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
index 4f63bbec1..0069caca5 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
+++ b/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(
diff --git a/vllm_omni/model_executor/models/weight_loader.py b/vllm_omni/model_executor/models/weight_loader.py
new file mode 100644
index 000000000..a8e898a86
--- /dev/null
+++ b/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)
diff --git a/vllm_omni/reasoning/step_audio_reasoning_parser.py b/vllm_omni/reasoning/step_audio_reasoning_parser.py
index 81ef63496..ddf5a434f 100644
--- a/vllm_omni/reasoning/step_audio_reasoning_parser.py
+++ b/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
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index d293f6312..9a30209ee 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/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,
diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py
index 845881138..cdfe20f7b 100644
--- a/vllm_omni/worker/sampling_utils.py
+++ b/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),
+ )