rebase 2026/03/07

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

vllm_omni/entrypoints/omni_llm.py

@@ -185,6 +185,20 @@ class OmniLLM(LLM):
         self.model_config = self.llm_engine.model_config
         self.input_processor = self.llm_engine.input_processor
 
+        # Parity with upstream LLM for pooling/classify entrypoints
+        chat_template = kwargs.get("chat_template", None)
+        from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
+        from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
+
+        self.chat_template = load_chat_template(chat_template)
+        self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
+        self.init_pooling_io_processors = init_pooling_io_processors(
+            supported_tasks=supported_tasks,
+            model_config=self.model_config,
+            renderer=self.renderer,
+            chat_template_config=self.chat_template_config,
+        )
+
     # ------------------------------------------------------------------
     # Override upstream _preprocess_cmpl so that omni-specific fields
     # (additional_information, prompt_embeds, …) survive the renderer's

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -317,6 +317,33 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int] | None = None,
+    ) -> Mapping[str, int] | None:
+        mm_counts = mm_counts or {}
+        requested_modalities = {m for m, c in mm_counts.items() if c > 0}
+        mm_max_tokens: dict[str, int] = {}
+
+        if requested_modalities & {"image", "video"}:
+            vl_tokens = Qwen2_5_VLProcessingInfo.get_mm_max_tokens_per_item(
+                self,
+                seq_len=seq_len,
+                mm_counts=mm_counts,
+            )
+            mm_max_tokens.update({m: vl_tokens[m] for m in ["image", "video"] if m in requested_modalities})
+
+        if "audio" in requested_modalities:
+            audio_tokens = Qwen2AudioProcessingInfo.get_mm_max_tokens_per_item(
+                self,
+                seq_len=seq_len,
+                mm_counts=mm_counts,
+            )
+            mm_max_tokens["audio"] = audio_tokens["audio"]
+
+        return mm_max_tokens
+
 
 Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder