[Quantization][Qwen2.5-Omni] Support thinker-only ModelOpt NVFP4 W4A4 checkpoints (#5073)

Signed-off-by: Caspian443 <scrisis843@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/config/test_config_factory.py

@@ -1805,6 +1805,7 @@ class TestQwen2_5OmniPipeline:
         assert s.owns_tokenizer is True
         assert s.engine_output_type == "latent"
         assert s.requires_multimodal_data is True
+        assert s.hf_config_name == "thinker_config"
 
     def test_talker(self):
         p = resolve_pipeline_config("qwen2_5_omni")
@@ -1813,6 +1814,7 @@ class TestQwen2_5OmniPipeline:
         s = p.get_stage(1)
         assert isinstance(s, StagePipelineConfig)
         assert s.input_sources == (0,)
+        assert s.hf_config_name == "talker_config"
         assert s.sampling_constraints["stop_token_ids"] == [8294]
         # thinker2talker was removed: qwen2_5_omni has no async_chunk support,
         # so sync_process_input_func always wins and custom_process_input_func
@@ -1829,6 +1831,7 @@ class TestQwen2_5OmniPipeline:
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
         assert s.engine_output_type == "audio"
+        assert s.hf_config_name == "thinker_config"
 
 
 class TestQwen3TTSPipeline:

tests/engine/test_arg_utils.py

@@ -12,7 +12,7 @@ from unittest.mock import Mock
 import pytest
 from omegaconf import OmegaConf
 from pydantic import ValidationError
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 from vllm.engine.arg_utils import EngineArgs
 
 from vllm_omni.config.model import OmniModelConfig
@@ -244,6 +244,28 @@ def test_stage_specific_text_config_override():
     assert omni_config.get_head_size() == talker_head_dim
 
 
+def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
+    """A stage without an embedding-size override keeps the default input width."""
+    vllm_config = EngineArgs().create_model_config()
+    vllm_config.hf_config = Qwen3OmniMoeConfig(enable_audio_output=True)
+
+    omni_config = OmniModelConfig.from_vllm_model_config(
+        vllm_config,
+        model_stage="talker",
+        hf_config_name="talker_config",
+    )
+
+    assert (
+        getattr(
+            omni_config.hf_config.talker_config,
+            "embedding_size",
+            None,
+        )
+        is None
+    )
+    assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
+
+
 def test_stage_configs_path_field():
     """OmniEngineArgs with stage_configs_path should construct without error."""
     args = OmniEngineArgs(stage_configs_path="/some/path.yaml")

vllm_omni/config/model.py

@@ -178,6 +178,14 @@ class OmniModelConfig(ModelConfig):
                 return override
         return super().embedding_size
 
+    def get_inputs_embeds_size(self) -> int:
+        if self.hf_config_name is not None:
+            stage_config = getattr(self.hf_config, self.hf_config_name, None)
+            override = getattr(stage_config, "embedding_size", None)
+            if override is not None:
+                return override
+        return super().get_inputs_embeds_size()
+
     def get_model_arch_config(self):
         # For multi-stage omni models, use a stage-aware convertor so that
         # only the correct stage's quantization config is surfaced.

vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py

@@ -29,6 +29,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             final_output_type="text",
             owns_tokenizer=True,
             requires_multimodal_data=True,
+            hf_config_name="thinker_config",
             engine_output_type="latent",
             sampling_constraints={"detokenize": True},
             custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
@@ -38,6 +39,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             model_stage="talker",
             execution_type=StageExecutionType.LLM_AR,
             input_sources=(0,),
+            hf_config_name="talker_config",
             engine_output_type="latent",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
@@ -53,6 +55,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig(
             input_sources=(1,),
             final_output=True,
             final_output_type="audio",
+            hf_config_name="thinker_config",
             engine_output_type="audio",
             sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
             sampling_constraints={"detokenize": True},
@@ -75,6 +78,7 @@ QWEN2_5_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig(
             final_output_type="text",
             owns_tokenizer=True,
             requires_multimodal_data=True,
+            hf_config_name="thinker_config",
             engine_output_type="latent",
             sampling_constraints={"detokenize": True},
             custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",

vllm_omni/worker/gpu_ar_model_runner.py

@@ -303,7 +303,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
         # each model stage has their own hidden size
         self.hidden_size = self.model_config.hf_text_config.hidden_size
-        self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
+        # The width of embeddings received by this stage. It can differ from
+        # the stage's internal hidden size (e.g., Qwen2.5-Omni Talker receives
+        # 3584-wide Thinker states and projects them to a 896-wide hidden state).
+        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
+        self.inputs_embeds = self._make_buffer(
+            self.max_num_tokens,
+            self.inputs_embeds_size,
+            dtype=self.dtype,
+            numpy=False,
+        )
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)